矩陣中的火花循環以運行線性回歸

Question

我有一個如下的火花數據框dt 。 BAB 是 ID，我想對 BAB 的每個值使用 AAB 和 AAD 列進行線性回歸。

這就是我運行它的方式。 通過為每個 BAB 值過濾整個 dataframe，它變得非常慢。 有沒有辦法像 3 維矩陣一樣循環數據並為每個 BAB 運行回歸？ 所以我只需要通過 BAB 一次 go 。 它不必是 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);
...
}

Answer 1

一種方法是將數據幀轉換為以元組為元素的List((BAB1,AAB1,AAD1),(BAB2,AAB2,AAD2),...) ，然后將列表 w.r.t 切片每個單獨的BAB和在每個切片上運行回歸。