scala - 矩阵中的火花循环以运行线性回归

Question

我有一个dt如下的火花数据框。BAB 是 ID，我想对 BAB 的每个值使用 AAB 和 AAD 列进行线性回归。

这就是我运行它的方式。通过为每个 BAB 值过滤整个数据帧，它变得非常慢。有没有办法像 3 维矩阵一样循环数据并为每个 BAB 运行回归？所以我只需要通过 BAB 一次。它不必是 spark mllib。任何其他带有 scala 编码的机器学习工具都可以。

val arrColu = Array("AAB", "AAD");
val assFeat = new VectorAssembler().setInputCols(arrColu).setOutputCol("features");

val arrBAB=dt.select("BAB").collect.map(_ (0)).map(x => x.toString);

for (a<-0 to arrBAB.length-1){

val vecDF: DataFrame = assFeat.transform(dt.filter("BAB='"+arrBAB(a)+"'").select("AAB","AAD"));
val lr1=new LinearRegression();
val lr2=lr1.setFeaturesCol("features").setLabelCol("AAD").setFitIntercept(true).
  setMaxIter(10).setRegParam(.3).setElasticNetParam(.8);
val fitD1=lr2.fit(vecDF);
...
}

标签： scalaapache-sparkmatrixlinear-regression

Answer 1

一种方法是将数据框转换为以元组为元素List((BAB1,AAB1,AAD1),(BAB2,AAB2,AAD2),...)的列表，然后对每个人的列表进行切片BAB并在每个切片上运行回归。