如果我的数据包含5个类别(A,B,C,D,E)和客户数据集,其中每个客户可以属于一个,多个或没有类别。如何获取这样的数据集:
id, categories
1 , [A,C]
2 , [B]
3 , []
4 , [D,E]
并将categories列转换为一个热编码向量,如此
id, categories, encoded
1 , [A,C] , [1,0,1,0,0]
2 , [B] , [0,1,0,0,0]
3 , [] , [0,0,0,0,0]
4 , [D,E] , [0,0,0,1,1]
有没有人找到一种简单的方法在火花中做到这一点?
答案 0 :(得分:4)
要获得所需的输出,您可以使用Spark的UDF(用户定义的函数)扩展Stephen Carman答案:
// Prepare training documents from a list of (id, text, label) tuples.
val data = spark.createDataFrame(Seq(
(0L, Seq("A", "B")),
(1L, Seq("B")),
(2L, Seq.empty),
(3L, Seq("D", "E"))
)).toDF("id", "categories")
// Get distinct tags array
val tags = data
.flatMap(r ⇒ r.getAs[Seq[String]]("categories"))
.distinct()
.collect()
.sortWith(_ < _)
val cvmData = new CountVectorizerModel(tags)
.setInputCol("categories")
.setOutputCol("sparseFeatures")
.transform(data)
val asDense = udf((v: Vector) ⇒ v.toDense)
cvmData
.withColumn("features", asDense($"sparseFeatures"))
.select("id", "categories", "features")
.show()
这将为您提供所需的输出
+---+----------+-----------------+
| id|categories| features|
+---+----------+-----------------+
| 0| [A, B]|[1.0,1.0,0.0,0.0]|
| 1| [B]|[0.0,1.0,0.0,0.0]|
| 2| []|[0.0,0.0,0.0,0.0]|
| 3| [D, E]|[0.0,0.0,1.0,1.0]|
+---+----------+-----------------+
答案 1 :(得分:2)
使用CountVectorizerModel
时非常容易做到这一点val df = spark.createDataFrame(Seq(
(1, Seq("A","C")),
(2, Seq("B")),
(3, Seq()),
(4, Seq("D","E")))
).toDF("id", "category")
val cvm = new CountVectorizerModel(Array("A","B","C","D","E"))
.setInputCol("category")
.setOutputCol("features")
cvm.transform(df).show()
/*
+---+--------+-------------------+
| id|category| features|
+---+--------+-------------------+
| 1| [A, C]|(5,[0,2],[1.0,1.0])|
| 2| [B]| (5,[1],[1.0])|
| 3| []| (5,[],[])|
| 4| [D, E]|(5,[3,4],[1.0,1.0])|
+---+--------+-------------------+
*/
这与您想要的完全不同,但功能向量会告诉您数据中存在哪些类别。例如,在第1行中,[0,2]对应于字典的第一个和第三个元素,或者&#34; A&#34;和&#34; C&#34;在那里写的。