Llamar a la función Java/Scala desde una tarea

Antecedentes

Mi pregunta original aquí fue ¿Por qué usar DecisionTreeModel.predict dentro de la función de mapa plantea una excepción? y está relacionado con Cómo generar tuplas de (etiqueta original, etiqueta predicha) en Spark con MLlib?

Cuando usamos Scala API una forma recomendada de obtener predicciones para RDD[LabeledPoint] usando DecisionTreeModel es simplemente mapear sobre RDD:

val labelAndPreds = testData.map { point =>
  val prediction = model.predict(point.features)
  (point.label, prediction)
}

Desafortunadamente un enfoque similar en PySpark no funciona tan bien:

labelsAndPredictions = testData.map(
    lambda lp: (lp.label, model.predict(lp.features))
labelsAndPredictions.first()

Excepción: Parece que está intentando hacer referencia a SparkContext desde una variable de difusión, acción o transforamción. SparkContext solo se puede usar en el controlador, no en el código que ejecuta en workers. Para obtener más información, consulte SPARK-5063.

En lugar de eso la documentación oficial recomienda algo como esto:

predictions = model.predict(testData.map(lambda x: x.features))
labelsAndPredictions = testData.map(lambda lp: lp.label).zip(predictions)

Entonces, ¿qué está pasando aquí? No hay ninguna variable de difusión aquí y Scala API define predict como sigue:

/**
 * Predict values for a single data point using the model trained.
 *
 * @param features array representing a single data point
 * @return Double prediction from the trained model
 */
def predict(features: Vector): Double = {
  topNode.predict(features)
}

/**
 * Predict values for the given data set using the model trained.
 *
 * @param features RDD representing data points to be predicted
 * @return RDD of predictions for each of the given data points
 */
def predict(features: RDD[Vector]): RDD[Double] = {
  features.map(x => predict(x))
}

Así que al menos a primera vista llamar desde acción o transformación no es un problema ya que la predicción parece ser una operación local.

Explicación

Después de algunas investigaciones me di cuenta de que la fuente del problema es un JavaModelWrapper.call método invocado desde DecisionTreeModel.predecir. It acceso SparkContext que se requiere para llamar a la función Java:

callJavaFunc(self._sc, getattr(self._java_model, name), *a)

Pregunta

En el caso de DecisionTreeModel.predict hay un solución recomendada y todo el código requerido ya es parte de la API de Scala, pero ¿hay alguna forma elegante de manejar un problema como este en general?

Las únicas soluciones que se me ocurren en este momento son bastante pesadas:

• empujando todo a JVM ya sea extendiendo las clases Spark a través de Conversiones implícitas o agregando algún tipo de envoltorios
• usando la pasarela Py4j directamente