Как получить коэффициент для отображения индекса из RFormula / RFormulaModel в Apache Spark?

Question

Все,

У меня есть простой фрейм данных, как показано ниже

Я использую API-интерфейс RFormula для создания матрицы моделикак показано ниже

val formula = "dep ~ indep"
val rF = new RFormula().setFormula(formula).setFeaturesCol("features").setLabelCol("label")
val rfModel = rF.fit(df)

, где rfModel имеет тип RFormulaModel.Согласно документам здесь

сопоставление категориальной переменной "indep" должно быть доступно для доступа из этого объекта в качестве pipeModel, но, похоже, это закрытый член.

Мой вопрос: как я могу получить метки и соответствующие индексы от объекта RFormulaModel?Я знаю, что могу использовать метаданные преобразованного кадра данных и выполнять манипуляции со строками, но есть ли простой способ сделать это?

Спасибо за любую помощь!

Answer 1

В pyspark вы можете извлечь метаданные из целевого фрейма данных, например

for attr in df.schema["X"].metadata["ml_attr"]["attrs"]["numeric"]:
    print(f"idx: {attr['idx']}, name: {attr['name']}")

Надеюсь, это применимо и к Scala!

Answer 2

Придумал хак, где мне пришлось записать RFormulaModel на диск, а затем прочитать часть pipeModel обратно как PipelineModel.Оттуда у меня есть доступ к этапам StringIndexerModel, как показано здесь

import org.apache.spark.ml.PipelineModel
import org.apache.spark.ml.feature.StringIndexerModel

rfModel.write.overwrite.save("/rfModel")
val pModel = PipelineModel.read.load("/rfModel/pipelineModel")

val strIndexers = pModel.stages.filter(stage => stage.isInstanceOf[StringIndexerModel])
val labelMaps = strIndexers.map(e  => { val i = e.asInstanceOf[StringIndexerModel]; (i.getInputCol, i.labels)})