在 pyspark 中聚合 One-Hot Encoded 功能

Question

我在 python 方面经验丰富，但对 pyspark 完全陌生。我的数据框包含大约 50M 行，具有几个分类特征。对于每个功能，我都对它们进行了 One-Hot 编码。这是代码的简化但具有代表性的示例。

从 pyspark.ml.feature 导入 StringIndexer, OneHotEncoder 从 pyspark.ml 导入管道

df = sc.parallelize([
     (1, 'grocery'),
     (1, 'drinks'),
     (1, 'bakery'),
     (2, 'grocery'),
     (3, 'bakery'),
     (3, 'bakery'),
 ]).toDF(["id", "category"])

indexer = StringIndexer(inputCol='category', outputCol='categoryIndex')
encoder = OneHotEncoder(inputCol='categoryIndex', outputCol='categoryVec')

pipe = Pipeline(stages = [indexer, encoder])

newDF = pipe.fit(df).transform(df)

输出

+---+--------+-------------+-------------+
| id|category|categoryIndex|  categoryVec|
+---+--------+-------------+-------------+
|  1| grocery|          1.0|(2,[1],[1.0])|
|  1|  drinks|          2.0|    (2,[],[])|
|  1|  bakery|          0.0|(2,[0],[1.0])|
|  2| grocery|          1.0|(2,[1],[1.0])|
|  3|  bakery|          0.0|(2,[0],[1.0])|
|  3|  bakery|          0.0|(2,[0],[1.0])|
+---+--------+-------------+-------------+

我现在想对“id”进行分组并将“categoryVec”列与总和聚合，这样我就可以得到每个 id 的一行，并带有一个向量，该向量指示客户正在购物的（可能是几个）类别中的哪一个. 在 pandas 中，这只是将 sum/mean 应用于pd.get_dummies() 步骤中生成的每一列的情况，但在这里似乎并不那么简单。

然后我会将输出传递给 ML 算法，因此我需要能够在输出上使用 VectorAssembler 或类似工具。

哦，我真的需要一个 pyspark 解决方案。

非常感谢您的帮助！

Answer 1

您可以为此使用 Counvectorizer。它将类别索引数组转换为编码向量。

from pyspark.ml.feature import CountVectorizer
from pyspark.ml import Pipeline
from pyspark.sql.window import Window
from pyspark.sql import functions as F


df = sc.parallelize([
     (1, 'grocery'),
     (1, 'drinks'),
     (1, 'bakery'),
     (2, 'grocery'),
     (3, 'bakery'),
     (3, 'bakery'),
 ]).toDF(["id", "category"]) \
   .groupBy('id') \
   .agg(F.collect_list('category').alias('categoryIndexes'))

cv = CountVectorizer(inputCol='categoryIndexes', outputCol='categoryVec')

transformed_df = cv.fit(df).transform(df)
transformed_df.show()

结果：

+---+--------------------+--------------------+
| id|     categoryIndexes|         categoryVec|
+---+--------------------+--------------------+
|  1|[grocery, drinks,...|(3,[0,1,2],[1.0,1...|
|  3|    [bakery, bakery]|       (3,[0],[2.0])|
|  2|           [grocery]|       (3,[1],[1.0])|
+---+--------------------+--------------------+