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

Question

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

是否有某种方法可以禁用此功能，以便索引变量可以用作数字变量？

编辑：我使用字符串索引器作为 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))