检索 Spark Mllib StringIndexer 列映射

Posted 2023-04-15

【中文标题】检索 Spark Mllib StringIndexer 列映射

【英文标题】：Retrieve Spark Mllib StringIndexer column mapping

【发布时间】：2017-04-23 19:09:47

【问题描述】：

如何从经过训练的 Spark MLlib StringIndexerModel 中获取映射？

val stringIndexer = new StringIndexer()
    .setInputCol("myCol")
    .setOutputCol("myColIdx")
val stringIndexerModel = stringIndexer.fit(data)
val res = stringIndexerModel.transform(data)

上面的代码将在myCol 中的我的DataFrame 映射值中添加一个myColIdx，并将其添加到基于值频率的索引中。即最频繁的值 -> 0，第二频繁的值 -> 1，等等......

如何从模型中检索该映射？如果我对模型进行序列化/反序列化，映射是否稳定（即，我是否保证在转换后得到相同的结果）？

【参考方案1】：

StringIndexerModel 使用 labels 属性公开映射：

stringIndexerModel.labels: Array[String]

其中值对应于连续标签，例如：

val data = Seq("foo", "bar", "foo", "bar", "foobar", "bar").toDF("myCol")

你会关注labels:

import org.apache.spark.ml.feature.IndexToString

Array(bar, foo, foobar)

bar 索引为 0.0，foo 索引为 1.0，foobar 索引为 2.0。这是模型的属性，当模型为saved时保留。

在Pipeline 中使用时，您还可以使用IndexToString，它将使用列元数据将索引映射回标签。

indexToString.transform(stringIndexerModel.transform(data)).show

+------+--------+-------------+
| myCol|myColIdx|myColReversed|
+------+--------+-------------+
|   foo|     1.0|          foo|
|   bar|     0.0|          bar|
|   foo|     1.0|          foo|
|   bar|     0.0|          bar|
|foobar|     2.0|       foobar|
|   bar|     0.0|          bar|
+------+--------+-------------+