新利国际官方导航新利国际官方导航


18luck新利

Spark MLlib特征处理 之 StringIndexer、IndexToString使用说明以及源码剖析

最近在用Spark MLlib进行特征处理时,对于StringIndexer和IndexToString遇到了点问题,查阅官方文档也没有解决疑惑。无奈之下翻看源码才明白其中一二...这就给大家娓娓道来。

更多内容参考我的大数据学习之路

文档说明

StringIndexer 字符串转索引

StringIndexer可以把字符串的列按照出现频率进行排序,出现次数最高的对应的Index为0。比如下面的列表进行StringIndexer

idcategory
0a
1b
2c
3a
4a
5c

就可以得到如下:

idcategorycategoryIndex
0a0.0
1b2.0
2c1.0
3a0.0
4a0.0
5c1.0

可以看到出现次数最多的"a",索引为0;次数最少的"b"索引为2。

针对训练集中没有出现的字符串值,spark提供了几种处理的方法:

error,直接抛出异常skip,跳过该样本数据keep,使用一个新的最大索引,来表示所有未出现的值

下面是基于Spark MLlib 2.2.0的代码样例:

package xingoo.ml.features.tranformerimport org.apache.spark.sql.SparkSessionimport org.apache.spark.ml.feature.StringIndexerobject 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("a"), "b"), "c"), "a"), "a"), "c")) ).toDF("id", "category") val df1 = spark.createDataFrame( Seq("a"), "b"), "c"), "a"), "e"), "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.tranformerimport org.apache.spark.ml.attribute.Attributeimport org.apache.spark.ml.feature.{IndexToString, StringIndexer}import org.apache.spark.sql.SparkSessionobject 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( "a"), "b"), "c"), "a"), "a"), "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}") 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.tranformerimport org.apache.spark.ml.feature.{IndexToString, StringIndexer}import org.apache.spark.sql.SparkSessionobject 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( "a"), "b"), "c"), "a"), "a"), "c") )).toDF("id", "category") val df2 = spark.createDataFrame(Seq( 2.0), 1.0), 1.0), 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 endpointException 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( "a"), "b"), "c"), "a"), "a"), "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( 2.0), 1.0), 1.0), 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( 2.0), 1.0), 1.0), 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

最终还是推荐详细阅读官方文档,不过官方文档真心有些粗糙,想要了解其中的原理,还是得静下心来看看源码。

欢迎阅读本文章: 申清

新利国际官方线路

18luck新利