Как вписать несколько моделей Spark ml в каждое секционированное подмножество одного набора данных / фрейма данных?

Question

Я хочу обучить несколько моделей Spark ML одновременно, используя каждое подмножество одного набора данных.(один раздел для одной модели.)

Как я могу реализовать это эффективно?

Цикл Scala кажется не очень хорошим, потому что DAG будет очень сложным, но mapGroups Spark также не работает, потому что план 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)
  }

Ни один из них не работает хорошо.У вас есть идеи?