Spark MLib 每日一算法 —— StringIndexer

2018-02-04

StringIndexer

StringIndexer是一个Estimator，用来将某个文本属性的值转化成数字编码index，以便后续其他适用于数字编码的算法使用。编码规则是对该文本属性每个出现的属性值label给出从0～label数量-1的数字，出现频率越高的label，给出的编码数字就越小。由此看出StringIndexer是需要根据已有训练集来进行fit的。

参数

inputCol: 需要给出编码的文本属性名称
outputCol: 编码后的属性名称 (默认值: strIdx_######_output)
handleInvalid: 表示在transform期间如何处理在fit训练集时没有碰到过的属性值label或者Null值。 (默认值: error)
可选项1. skip-放弃有未知label的row；
可选项2. error-抛出异常org.apache.spark.SparkException: Unseen label: xx；
可选项3. keep处理成以label数量为数字的编码.

例子

val training = spark.createDataFrame(
  Seq((0, "a"), (1, "b"), (2, "c"), (3, "a"), (4, "a"), (5, "c"))
).toDF("id", "category")
training.show()
//    +---+--------+
//    | id|category|
//    +---+--------+
//    |  0|       a|
//    |  1|       b|
//    |  2|       c|
//    |  3|       a|
//    |  4|       a|
//    |  5|       c|
//    +---+--------+

val indexer = new StringIndexer()
  .setInputCol("category")
  .setOutputCol("categoryIndex")
  
val model = indexer.fit(training)
model.transform(training).show()
//    +---+--------+-------------+
//    | 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|
//    +---+--------+-------------+

val test = spark.createDataFrame(
  Seq((0, "c"), (1, "d"), (2, "e"), (3, "c"), (4, "b"))
).toDF("id", "category")

try {
  model.transform(test).show()
} catch {  //error from label d & e.
  case ex: SparkException => println(ex)
}

model.setHandleInvalid("keep")
model.transform(test).show()
//    +---+--------+-------------+
//    | id|category|categoryIndex|
//    +---+--------+-------------+
//    |  0|       c|          1.0|
//    |  1|       d|          3.0|
//    |  2|       e|          3.0|
//    |  3|       c|          1.0|
//    |  4|       b|          2.0|
//    +---+--------+-------------+