Spark MLlib中的StringIndexer

Question

我在PipelinedRDD中有一列标称值，为了进行分类，我想将其转换为索引编码。

我曾经在pyspark.ml使用StringIndexer ，它非常易于使用。 但是，这次我正在学习如何处理rdd而不是数据帧，并且pyspark.mllib没有这样的东西。

任何帮助表示赞赏。

Answer 1

Spark MLlib中没有StringIndexer ，因此您需要自己完成工作。 首先收集该列的所有可能值，并为每个数字分配一个数字，然后将其另存为字典。 然后，将其应用于原始的rdd值。

下面的代码假定PipelinedRDD每行包含两个值，该值在第一个位置（0）进行转换：

dic = PipelinedRDD.map(lambda x: x[0]).distinct().zipWithIndex().collectAsMap()
PipelinedRDD = PipelinedRDD.map(lambda x: (dic[x[0]], x[1]))

注意 ：这与StringIndexer的Spark实现略有不同，因为它没有考虑值的频率（Spark会将0分配给最出现的值，然后是1，依此类推）。 但是，在大多数情况下，分配给哪个索引不同的字符串并不重要。

---

扩展如果您想精确地模仿StringIndexer的功能（如上面的注释中所述），可以对代码进行略微修改以考虑到这一点

dic = PiplelinedRDD.groupBy('cat_column').count().sort(col("count").desc()).map(lambda x: x[0]).zipWithIndex().collectAsMap()
PipelinedRDD = PipelinedRDD.map(lambda x: (dic[x[0]], x[1]))