Apache Spark如何将(多个)索引器和编码器连接成一个管道

时间:2017-02-21 11:24:57

标签: java apache-spark apache-spark-mllib

在我可以在当前数据帧上使用机器学习模型之前,我必须应用几个预处理步骤。这包括使用StringIndexer将多个分类列转换为索引列,并使用OneHotEncoder列进行编码。

我希望将所有转换方法放在管道中以保留它并将其轻松应用于传入数据。但由于索引列尚不存在,因此在编码器上调用fit()方法失败。

如何创建包含所有这些步骤的管道?

1 个答案:

答案 0 :(得分:0)

您可以同时威胁索引器和编码器 - PipelineStage将它们添加到您的管道中并一步完成整个管道。例如:

String INDEX_APPENDIX = "_IDX";
String VECTOR_APPENDIX = "_VEC";
ArrayList<PipelineStage> stages = new ArrayList<>();

for (String column : Arrays.asList("col1", "col2" )) {
        stages.add(new StringIndexer().setInputCol(column).setOutputCol(column + INDEX_APPENDIX));
        stages.add(new OneHotEncoder().setInputCol(column + INDEX_APPENDIX).setOutputCol(column +
                VECTOR_APPENDIX));
    }

Pipeline pipeline = new Pipeline()
    .setStages(stages.toArray(new PipelineStage[stages.size()]));

Dataset<Row> processedDf = pipeline.fit(df).transform(df);