Aplicando IndexToString ao vetor de recursos no Spark

Contexto: Eu tenho um quadro de dados onde todos os valores categóricos foram indexados usando StringIndexer.

val categoricalColumns = df.schema.collect { case StructField(name, StringType, nullable, meta) => name }    

val categoryIndexers = categoricalColumns.map {
  col => new StringIndexer().setInputCol(col).setOutputCol(s"${col}Indexed") 
}

Em seguida, usei o VectorAssembler para vetorizar todas as colunas de recursos (incluindo as categóricas indexadas).

val assembler = new VectorAssembler()
    .setInputCols(dfIndexed.columns.diff(List("label") ++ categoricalColumns))
    .setOutputCol("features")

Depois de aplicar o classificador e algumas etapas adicionais, termino com um quadro de dados que possui rótulo, recursos e previsão. Gostaria de expandir meu vetor de recursos para separar colunas, a fim de converter os valores indexados em seu formato String original.

val categoryConverters = categoricalColumns.zip(categoryIndexers).map {
colAndIndexer => new IndexToString().setInputCol(s"${colAndIndexer._1}Indexed").setOutputCol(colAndIndexer._1).setLabels(colAndIndexer._2.fit(df).labels)
}

Pergunta, questão: Tem algumasimples maneira de fazer isso, ou é a melhor abordagem para, de alguma forma, anexar a coluna de previsão ao quadro de dados de teste?

O que eu tentei:

val featureSlicers = categoricalColumns.map {
  col => new VectorSlicer().setInputCol("features").setOutputCol(s"${col}Indexed").setNames(Array(s"${col}Indexed"))
}

A aplicação disso fornece as colunas que eu quero, mas elas estão no formato Vetor (como deve ser feito) e não no tipo Duplo.

Editar: A saída desejada é o quadro de dados original (ou seja, façanha categórica, res como String não indexada) com uma coluna adicional indicando o rótulo previsto (que no meu caso é 0 ou 1).

Por exemplo, digamos que a saída do meu classificador tenha algo parecido com isto:

+-----+---------+----------+
|label| features|prediction|
+-----+---------+----------+
|  1.0|[0.0,3.0]|       1.0|
+-----+---------+----------+

Ao aplicar o VectorSlicer em cada recurso, eu obteria:

+-----+---------+----------+-------------+-------------+
|label| features|prediction|statusIndexed|artistIndexed|
+-----+---------+----------+-------------+-------------+
|  1.0|[0.0,3.0]|       1.0|        [0.0]|        [3.0]|
+-----+---------+----------+-------------+-------------+

O que é ótimo, mas eu preciso de:

+-----+---------+----------+-------------+-------------+
|label| features|prediction|statusIndexed|artistIndexed|
+-----+---------+----------+-------------+-------------+
|  1.0|[0.0,3.0]|       1.0|         0.0 |         3.0 |
+-----+---------+----------+-------------+-------------+

Para poder usar IndexToString e convertê-lo em:

+-----+---------+----------+-------------+-------------+
|label| features|prediction|    status   |    artist   |
+-----+---------+----------+-------------+-------------+
|  1.0|[0.0,3.0]|       1.0|        good |  Pink Floyd |
+-----+---------+----------+-------------+-------------+

ou até:

+-----+----------+-------------+-------------+
|label|prediction|    status   |    artist   |
+-----+----------+-------------+-------------+
|  1.0|       1.0|        good |  Pink Floyd |
+-----+----------+-------------+-------------+

questionAnswers(1)

yourAnswerToTheQuestion