最近在用Spark MLlib进行特征处理时，对于StringIndexer和IndexToString遇到了点问题，查阅官方文档也没有解决疑惑。无奈之下翻看源码才明白其中一二...这就给大家娓娓道来。
更多内容参考我的大数据学习之路
文档说明
StringIndexer 字符串转索引
StringIndexer可以把字符串的列按照出现频率进行排序，出现次数最高的对应的Index为0。比如下面的列表进行StringIndexer
| id | category |
|---|---|
| 0 | a |
| 1 | b |
| 2 | c |
| 3 | a |
| 4 | a |
| 5 | c |
就可以得到如下：
| id | category | categoryIndex |
|---|---|---|
| 0 | a | 0.0 |
| 1 | b | 2.0 |
| 2 | c | 1.0 |
| 3 | a | 0.0 |
| 4 | a | 0.0 |
| 5 | c | 1.0 |
可以看到出现次数最多的"a"，索引为0；次数最少的"b"索引为2。
针对训练集中没有出现的字符串值，spark提供了几种处理的方法：
- error，直接抛出异常
- skip，跳过该样本数据
- keep，使用一个新的最大索引，来表示所有未出现的值
下面是基于Spark MLlib 2.2.0的代码样例：
package xingoo.ml.features.tranformer
import org.apache.spark.sql.SparkSession
import org.apache.spark.ml.feature.StringIndexer
object StringIndexerTest {
def main(args: Array[String]): Unit = {
val spark = SparkSession.builder().master("local[*]").appName("string-indexer").getOrCreate()
spark.sparkContext.setLogLevel("WARN")
val df = spark.createDataFrame(
Seq((0, "a"), (1, "b"), (2, "c"), (3, "a"), (4, "a"), (5, "c"))
).toDF("id", "category")
val df1 = spark.createDataFrame(
Seq((0, "a"), (1, "b"), (2, "c"), (3, "a"), (4, "e"), (5, "f"))
).toDF("id", "category")
val indexer = new StringIndexer()
.setInputCol("category")
.setOutputCol("categoryIndex")
.setHandleInvalid("keep") //skip keep error
val model = indexer.fit(df)
val indexed = model.transform(df1)
indexed.show(false)
}
}
得到的结果为：
+---+--------+-------------+
|id |category|categoryIndex|
+---+--------+-------------+
|0 |a |0.0 |
|1 |b |2.0 |
|2 |c |1.0 |
|3 |a |0.0 |
|4 |e |3.0 |
|5 |f |3.0 |
+---+--------+-------------+
IndexToString 索引转字符串
这个索引转回字符串要搭配前面的StringIndexer一起使用才行：
package xingoo.ml.features.tranformer
import org.apache.spark.ml.attribute.Attribute
import org.apache.spark.ml.feature.{IndexToString, StringIndexer}
import org.apache.spark.sql.SparkSession
object IndexToString2 {
def main(args: Array[String]): Unit = {
val spark = SparkSession.builder().master("local[*]").appName("dct").getOrCreate()
spark.sparkContext.setLogLevel("WARN")
val df = spark.createDataFrame(Seq(
(0, "a"),
(1, "b"),
(2, "c"),
(3, "a"),
(4, "a"),
(5, "c")
)).toDF("id", "category")
val indexer = new StringIndexer()
.setInputCol("category")
.setOutputCol("categoryIndex")
.fit(df)
val indexed = indexer.transform(df)
println(s"Transformed string column '${indexer.getInputCol}' " +
s"to indexed column '${indexer.getOutputCol}'")
indexed.show()
val inputColSchema = indexed.schema(indexer.getOutputCol)
println(s"StringIndexer will store labels in output column metadata: " +
s"${Attribute.fromStructField(inputColSchema).toString}\n")
val converter = new IndexToString()
.setInputCol("categoryIndex")
.setOutputCol("originalCategory")
val converted = converter.transform(indexed)
println(s"Transformed indexed column '${converter.getInputCol}' back to original string " +
s"column '${converter.getOutputCol}' using labels in metadata")
converted.select("id", "categoryIndex", "originalCategory").show()
}
}
得到的结果如下：
Transformed string column 'category' to indexed column 'categoryIndex'
+---+--------+-------------+
| id|category|categoryIndex|
+---+--------+-------------+
| 0| a| 0.0|
| 1| b| 2.0|
| 2| c| 1.0|
| 3| a| 0.0|
| 4| a| 0.0|
| 5| c| 1.0|
+---+--------+-------------+
StringIndexer will store labels in output column metadata: {"vals":["a","c","b"],"type":"nominal","name":"categoryIndex"}
Transformed indexed column 'categoryIndex' back to original string column 'originalCategory' using labels in metadata
+---+-------------+----------------+
| id|categoryIndex|originalCategory|
+---+-------------+----------------+
| 0| 0.0| a|
| 1| 2.0| b|
| 2| 1.0| c|
| 3| 0.0| a|
| 4| 0.0| a|
| 5| 1.0| c|
+---+-------------+----------------+
使用问题
假如处理的过程很复杂，重新生成了一个DataFrame，此时想要把这个DataFrame基于IndexToString转回原来的字符串怎么办呢? 先来试试看：
package xingoo.ml.features.tranformer
import org.apache.spark.ml.feature.{IndexToString, StringIndexer}
import org.apache.spark.sql.SparkSession
object IndexToString3 {
def main(args: Array[String]): Unit = {
val spark = SparkSession.builder().master("local[*]").appName("dct").getOrCreate()
spark.sparkContext.setLogLevel("WARN")
val df = spark.createDataFrame(Seq(
(0, "a"),
(1, "b"),
(2, "c"),
(3, "a"),
(4, "a"),
(5, "c")
)).toDF("id", "category")
val df2 = spark.createDataFrame(Seq(
(0, 2.0),
(1, 1.0),
(2, 1.0),
(3, 0.0)
)).toDF("id", "index")
val indexer = new StringIndexer()
.setInputCol("category")
.setOutputCol("categoryIndex")
.fit(df)
val indexed = indexer.transform(df)
val converter = new IndexToString()
.setInputCol("categoryIndex")
.setOutputCol("originalCategory")
val converted = converter.transform(df2)
converted.show()
}
}
运行后发现异常：
18/07/05 20:20:32 INFO StateStoreCoordinatorRef: Registered StateStoreCoordinator endpoint
Exception in thread "main" java.lang.IllegalArgumentException: Field "categoryIndex" does not exist.
at org.apache.spark.sql.types.StructType$$anonfun$apply$1.apply(StructType.scala:266)
at org.apache.spark.sql.types.StructType$$anonfun$apply$1.apply(StructType.scala:266)
at scala.collection.MapLike$class.getOrElse(MapLike.scala:128)
at scala.collection.AbstractMap.getOrElse(Map.scala:59)
at org.apache.spark.sql.types.StructType.apply(StructType.scala:265)
at org.apache.spark.ml.feature.IndexToString.transformSchema(StringIndexer.scala:338)
at org.apache.spark.ml.PipelineStage.transformSchema(Pipeline.scala:74)
at org.apache.spark.ml.feature.IndexToString.transform(StringIndexer.scala:352)
at xingoo.ml.features.tranformer.IndexToString3$.main(IndexToString3.scala:37)
at xingoo.ml.features.tranformer.IndexToString3.main(IndexToString3.scala)
这是为什么呢？跟随源码来看吧！
源码剖析
首先我们创建一个DataFrame，获得原始数据：
val df = spark.createDataFrame(Seq(
(0, "a"),
(1, "b"),
(2, "c"),
(3, "a"),
(4, "a"),
(5, "c")
)).toDF("id", "category")
然后创建对应的StringIndexer：
val indexer = new StringIndexer()
.setInputCol("category")
.setOutputCol("categoryIndex")
.setHandleInvalid("skip")
.fit(df)
这里面的fit就是在训练转换器了，进入fit()：
override def fit(dataset: Dataset[_]): StringIndexerModel = {
transformSchema(dataset.schema, logging = true)
// 这里针对需要转换的列先强制转换成字符串，然后遍历统计每个字符串出现的次数
val counts = dataset.na.drop(Array($(inputCol))).select(col($(inputCol)).cast(StringType))
.rdd
.map(_.getString(0))
.countByValue()
// counts是一个map，里面的内容为{a->3, b->1, c->2}
val labels = counts.toSeq.sortBy(-_._2).map(_._1).toArray
// 按照个数大小排序，返回数组，[a, c, b]
// 把这个label保存起来，并返回对应的model（mllib里边的模型都是这个套路，跟sklearn学的）
copyValues(new StringIndexerModel(uid, labels).setParent(this))
}
这样就得到了一个列表，列表里面的内容是[a, c, b]，然后执行transform来进行转换：
val indexed = indexer.transform(df)
这个transform可想而知就是用这个数组对每一行的该列进行转换，但是它其实还做了其他的事情：
override def transform(dataset: Dataset[_]): DataFrame = {
...
// --------
// 通过label生成一个Metadata，这个很关键！！！
// metadata其实是一个map，内容为：
// {"ml_attr":{"vals":["a","c","b"],"type":"nominal","name":"categoryIndex"}}
// --------
val metadata = NominalAttribute.defaultAttr
.withName($(outputCol)).withValues(filteredLabels).toMetadata()
// 如果是skip则过滤一些数据
...
// 下面是针对不同的情况处理转换的列，逻辑很简单
val indexer = udf { label: String =>
...
if (labelToIndex.contains(label)) {
labelToIndex(label) //如果正常，就进行转换
} else if (keepInvalid) {
labels.length // 如果是keep，就返回索引的最大值（即数组的长度）
} else {
... // 如果是error，就抛出异常
}
}
// 保留之前所有的列，新增一个字段，并设置字段的StructField中的Metadata！！！！
// 并设置字段的StructField中的Metadata！！！！
// 并设置字段的StructField中的Metadata！！！！
// 并设置字段的StructField中的Metadata！！！！
filteredDataset.select(col("*"),
indexer(dataset($(inputCol)).cast(StringType)).as($(outputCol), metadata))
}
看到了吗！关键的地方在这里，给新增加的字段的类型StructField设置了一个Metadata。这个Metadata正常都是空的{}，但是这里设置了metadata之后，里面包含了label数组的信息。
接下来看看IndexToString是怎么用的，由于IndexToString是一个Transformer，因此只有一个trasform方法：
override def transform(dataset: Dataset[_]): DataFrame = {
transformSchema(dataset.schema, logging = true)
val inputColSchema = dataset.schema($(inputCol))
// If the labels array is empty use column metadata
// 关键是这里：
// 如果IndexToString设置了labels数组，就直接返回；
// 否则，就读取了传入的DataFrame的StructField中的Metadata
val values = if (!isDefined(labels) || $(labels).isEmpty) {
Attribute.fromStructField(inputColSchema)
.asInstanceOf[NominalAttribute].values.get
} else {
$(labels)
}
// 基于这个values把index转成对应的值
val indexer = udf { index: Double =>
val idx = index.toInt
if (0 <= idx && idx < values.length) {
values(idx)
} else {
throw new SparkException(s"Unseen index: $index ??")
}
}
val outputColName = $(outputCol)
dataset.select(col("*"),
indexer(dataset($(inputCol)).cast(DoubleType)).as(outputColName))
}
了解StringIndexer和IndexToString的原理机制后，就可以作出如下的应对策略了。
1 增加StructField的MetaData信息
val df2 = spark.createDataFrame(Seq(
(0, 2.0),
(1, 1.0),
(2, 1.0),
(3, 0.0)
)).toDF("id", "index").select(col("*"),col("index").as("formated_index", indexed.schema("categoryIndex").metadata))
val converter = new IndexToString()
.setInputCol("formated_index")
.setOutputCol("origin_col")
val converted = converter.transform(df2)
converted.show(false)
+---+-----+--------------+----------+
|id |index|formated_index|origin_col|
+---+-----+--------------+----------+
|0 |2.0 |2.0 |b |
|1 |1.0 |1.0 |c |
|2 |1.0 |1.0 |c |
|3 |0.0 |0.0 |a |
+---+-----+--------------+----------+
2 获取之前StringIndexer后的DataFrame中的Label信息
val df3 = spark.createDataFrame(Seq(
(0, 2.0),
(1, 1.0),
(2, 1.0),
(3, 0.0)
)).toDF("id", "index")
val converter2 = new IndexToString()
.setInputCol("index")
.setOutputCol("origin_col")
.setLabels(indexed.schema("categoryIndex").metadata.getMetadata("ml_attr").getStringArray("vals"))
val converted2 = converter2.transform(df3)
converted2.show(false)
+---+-----+----------+
|id |index|origin_col|
+---+-----+----------+
|0 |2.0 |b |
|1 |1.0 |c |
|2 |1.0 |c |
|3 |0.0 |a |
+---+-----+----------+
两种方法都能得到正确的输出。
完整的代码可以参考github链接：
https://github.com/xinghalo/spark-in-action/blob/master/src/xingoo/ml/features/tranformer/IndexToStringTest.scala
最终还是推荐详细阅读官方文档，不过官方文档真心有些粗糙，想要了解其中的原理，还是得静下心来看看源码。