如何使用StringIndexer生成数字变量？

Question

我希望使用StringIndexer作为对我的数据集中的1000多个类别进行排名的方法，从而生成一个表示相对频率的索引。 然后，我可以将该索引用作模型的数字功能。 不幸的是，默认情况下， StringIndex存储一些元数据，将索引标记为分类，从而迫使我的模型将索引用作类别变量 。

有禁用此功能的某种方法，以便将index变量用作数字变量吗？

编辑：我正在使用字符串索引器作为ML管道中的一个阶段，因此解决方案将需要避免直接操纵数据帧。 另外，我将保存并加载此管道，因此自定义数据转换器可能不切实际。 我怀疑这不可能，因为目前正在编写Spark。

Answer 1

您可以索引数据，然后替换元数据。 假设您的数据如下所示：

import spark.implicits._
import org.apache.spark.ml.feature.StringIndexer

val indexer = new StringIndexer().setInputCol("raw").setOutputCol("indexed")

val df = Seq("a", "b", "b", "c", "c", "c").toDF("raw")
val indexed = indexer.fit(df).transform(df)

我们需要一个NumericAttribute ：

import org.apache.spark.ml.attribute.NumericAttribute

和元数据：

val meta = NumericAttribute.defaultAttr.withName("indexed").toMetadata

最后，我们可以使用as方法替换元数据：

indexed.withColumn("indexed", $"indexed".as("indexed", meta))