如何对单个DatasetDataFrame的每个分区子集拟合多个Spark ml模型？

Question

我想一次训练多个Spark ML模型，使用单个Dataset的每个子集。一个分区对一个模型）。

如何高效实现？

Scala的循环似乎不好用，因为DAG会复杂很多，但Spark的mapGroups不好用，因为Spark计划会被嵌套。

抽象的例子在下面。

case class Person(name: String, age: Int, sex: String)

val input = spark.read.load(..).as[Person]
val someModel = SomeModel()

// loop pattern
val models1 = input.map(_.age).collect.map{
  key => 
    val subset = input.filter(row => row.age == key)
    someModel.fit(subset)
}

// mapGroups pattern
val models2 = input
  .groupByKey(x => x.age)
  .mapGroups{
    (k, v) => 
      val subset = v.toList.toDS
      someModel.fit(subset)
  }

都不好用。你有什么办法吗？

Answer 1

我尝试了很久，想让这个工作，但是放弃了。我最终使用python库即statsmodels来获得必要的系数输出。如果你像你一样转换到RDD或mapGroups方法，那么它将工作。我试着在1个数据集中对70万个项目运行spark ml glm poisson线性回归，花了10多个小时。