Forma óptima de crear una canalización de ml en Apache Spark para conjuntos de datos con un gran número de columnas

Estoy trabajando con Spark 2.1.1 en un conjunto de datos con ~2000 características y tratando de crear una tubería de ML básica, que consiste en algunos Transformadores y un Clasificador.

Supongamos, en aras de la simplicidad, que la tubería con la que estoy trabajando consiste en un VectorAssembler, StringIndexer y un Clasificador, que sería un caso de uso bastante común.

// Pipeline elements
val assmbleFeatures: VectorAssembler = new VectorAssembler()
  .setInputCols(featureColumns)
  .setOutputCol("featuresRaw")

val labelIndexer: StringIndexer = new StringIndexer()
  .setInputCol("TARGET")
  .setOutputCol("indexedLabel")

// Train a RandomForest model.
val rf: RandomForestClassifier = new RandomForestClassifier()
  .setLabelCol("indexedLabel")
  .setFeaturesCol("featuresRaw")
  .setMaxBins(30)

// add the params, unique to this classifier
val paramGrid = new ParamGridBuilder()
  .addGrid(rf.numTrees, Array(5))
  .addGrid(rf.maxDepth, Array(5))
  .build()

// Treat the Pipeline as an Estimator, to jointly choose parameters for all Pipeline stages.
val evaluator = new BinaryClassificationEvaluator()
  .setMetricName("areaUnderROC")
  .setLabelCol("indexedLabel")

Si los pasos de la tubería están separados en una tubería de transformador (VectorAssembler + StringIndexer) y un segundo clasificador canalización, y si las columnas innecesarias se dejan caer entre ambas canalizaciones, el entrenamiento tiene éxito. Esto significa que para reutilizar los modelos, dos modelos de Pipelin tienen que ser guardados después de la formación y un paso intermedio de preprocesamiento tiene que ser introducido.

// Split indexers and forest in two Pipelines.
val prePipeline = new Pipeline().setStages(Array(labelIndexer, assmbleFeatures)).fit(dfTrain)
// Transform data and drop all columns, except those needed for training 
val dfTrainT = prePipeline.transform(dfTrain)
val columnsToDrop = dfTrainT.columns.filter(col => !Array("featuresRaw", "indexedLabel").contains(col))
val dfTrainRdy = dfTrainT.drop(columnsToDrop:_*)

val mainPipeline = new Pipeline().setStages(Array(rf))

val cv = new CrossValidator()
  .setEstimator(mainPipeline)
  .setEvaluator(evaluator)
  .setEstimatorParamMaps(paramGrid)
  .setNumFolds(2)

val bestModel = cv.fit(dfTrainRdy).bestModel.asInstanceOf[PipelineModel]

La solución (en mi humilde opinión) mucho más limpia sería fusionar todas las etapas de la tubería en una sola tubería.

val pipeline = new Pipeline()
  .setStages(Array(labelIndexer, assmbleFeatures, rf))

val cv = new CrossValidator()
  .setEstimator(pipeline)
  .setEvaluator(evaluator)
  .setEstimatorParamMaps(paramGrid)
  .setNumFolds(2)

// This will fail! 
val bestModel = cv.fit(dfTrain).bestModel.asInstanceOf[PipelineModel]

Sin embargo, poner todas las PipelineStages en una tubería conduce a la siguiente excepción, probablemente debido al problema este PR finalmente resolverá:

ERROR CodeGenerator: failed to compile: org.codehaus.janino.JaninoRuntimeException: Grupo constante para la organización de clases.apache.chispa.SQL.catalizador.expresiones.GeneratedClass Specif SpecificUnsafeProjection has grown past JVM limit of 0xFFFF

La razón de esto es que el VectorAssembler duplica efectivamente (en este ejemplo) la cantidad de datos en el DataFrame, ya que no hay ningún transformador que pueda eliminar las columnas innecesarias. (Ver spark pipeline vector assembler drop otras columnas )

Para el ejemplo funciona en el conjunto de datos golub y los siguientes pasos de preprocesamiento son necesarios:

import org.apache.spark.sql.types.DoubleType
import org.apache.spark.ml.classification.RandomForestClassifier
import org.apache.spark.ml.{Pipeline, PipelineModel, PipelineStage}
import org.apache.spark.ml.evaluation.BinaryClassificationEvaluator
import org.apache.spark.ml.feature._
import org.apache.spark.sql._
import org.apache.spark.ml.tuning.{CrossValidator, ParamGridBuilder}

val df = spark.read.option("header", true).option("inferSchema", true).csv("/path/to/dataset/golub_merged.csv").drop("_c0").repartition(100)

// Those steps are necessary, otherwise training would fail either way
val colsToDrop = df.columns.take(5000)
val dfValid = df.withColumn("TARGET", df("TARGET_REAL").cast(DoubleType)).drop("TARGET_REAL").drop(colsToDrop:_*)

// Split df in train and test sets
val Array(dfTrain, dfTest) = dfValid.randomSplit(Array(0.7, 0.3))

// Feature columns are columns except "TARGET"
val featureColumns = dfTrain.columns.filter(col => col != "TARGET")

Como soy nuevo en Spark, no estoy seguro de cuál sería la mejor manera de resolver este problema. ¿Podría sugerirlo?..

1. para crear un nuevo transformador, que cae columnas y que se puede incorporar en la tubería?
2. dividir ambas tuberías e introducir el paso intermedio
3. ¿algo más? :)

O me estoy perdiendo algo importante (pasos de la tubería, relaciones públicas, etc.) que resolvería este problema?

Editar:

Implementé un nuevo Transformador DroppingVectorAssembler, que elimina columnas innecesarias, sin embargo, se lanza la misma excepción.

Además de eso, establecer spark.sql.codegen.wholeStage a false no resuelve el problema.