是否可以在 Spark ML 中为随机森林制作通用训练管道？

Question

我刚开始使用 Spark 和 Spark ML，我发现它比 Python 和 sklearn 难得多。

开发时间要长得多，所以我想知道制作一个适用于任何（足够小的）数据集并训练随机森林分类器的通用管道是否可行。理想情况下，我会创建一个像

这样的函数

def trainClassifier(df: DataFrame, labelColumn: String) {
  ...
}

Spark 的大量开发时间花在了将列编码为数字列，然后从特征中形成向量，以便 Spark ML 的随机森林可以实际使用它。所以最后写成

这样的行

val indexer = new StringIndexer()
                   .setInputCol("category")
                   .setOutputCol("categoryIndex")
                   .fit(df)

val indexed = indexer.transform(df)

val encoder = new OneHotEncoder()
                   .setInputCol("categoryIndex")
                   .setOutputCol("categoryVec")

val encoded = encoder.transform(indexed)

所以我的问题更像是一个设计问题（如果合适，请引导我到不同的站点）关于如何编写适用于任何 DataFrame 的通用分类训练函数，但它是还有一个关于 Spark 的问题，我在问这种事情在 Spark 中是否可行（所以这是一个 API 问题，所以它更适合 Whosebug）？

编辑：我的意思是我没有指定列并为每个新数据框手动转换列。我想要一个函数 trainClassifier ，它将接收具有不同列和不同列类型的各种数据帧。遍历除 labelColumn 之外的所有列并将它们一起编译成分类器可以使用的特征向量的东西。

Answer 1

您可以创建自定义管道：

val start = "category"; // can be parameter of method or function
val indexer = new StringIndexer()
               .setInputCol(start )
               .setOutputCol(start + "Index")
               .fit(df)

val encoder = new OneHotEncoder()
               .setInputCol(encoder.outputCol)
               .setOutputCol(start  + "encoded")

这些步骤可以在 returns Array[Stage] - Array(indexer, encoder) 的函数中进行。现在你可以编写一些函数来连接所有数组并创建管道：

val randomForest = ... 

val pipeline = new Pipeline()
    .setStages(allStepsArray(indexer , encoder , randomForest))

然后你可以在 Pipeline 上调用 fit 甚至像 link:

那样构建 CrossValidator

val model = pipeline.fit(testData)

是否可以在 Spark ML 中为随机森林制作通用训练管道？

Is it possible to make a generic training pipeline for Random Forest in Spark ML?

scala

machine-learning

apache-spark

apache-spark-mllib