PySpark列表列转独热编码时pivot列超10000值的报错解决

问题场景 #

现有如下PySpark DataFrame，包含ID列和存储商品列表的Items列：

df1=spark.createDataFrame([(7573  , ['Apple', 'zucchini', 'eggplant', 'beer', 'spaghetti']),(2331  , ['beer', 'eggplant', 'Strawberry', 'cherryberry', 'Chicken','Fish','Meat']),(5541  , ['Cola', 'Pizza', 'Strawberry', 'cheese', 'Chicken','Yogurt','Meat'])],('ID'  , 'Items'))
df1.show(truncate=False)

需求是将Items中的每个商品转为单独列，标记对应ID是否购买该商品（购买记1，未购买记0），目标结果如下：

+----+-----+-------+----+----+----+-----+----------+------+----+------+-----------+--------+---------+--------+
|  ID|Apple|Chicken|Cola|Fish|Meat|Pizza|Strawberry|Yogurt|beer|cheese|cherryberry|eggplant|spaghetti|zucchini|
+----+-----+-------+----+----+----+-----+----------+------+----+------+-----------+--------+---------+--------+
|2331|    0|      1|   0|   1|   1|    0|         1|     0|   1|     0|          1|       1|        0|       0|
|5541|    0|      1|   1|   0|   1|    1|         1|     1|   0|     1|          0|       0|        0|       0|
|7573|    1|      0|   0|   0|   0|    0|         0|     0|   1|     0|          0|       1|        1|       1|
+----+-----+-------+----+----+----+-----+----------+------+----+------+-----------+--------+---------+--------+

最初使用以下代码实现需求：

df1 = df1.withColumn('exploded', F.explode('Items')).groupBy("ID").pivot("exploded").agg(F.lit(1)).na.fill(0).show()

但处理大数据集时触发报错：The pivot column exploded has more than 10000 distinct values。

解决方法 #

1. 调整Spark配置放宽Pivot基数限制 #

PySpark默认限制pivot列的不同值数量为10000，可以通过修改spark.sql.pivotMaxValues参数提高阈值。可以在创建SparkSession时设置，也可以运行时动态调整：

# 创建SparkSession时配置
spark = SparkSession.builder \
    .appName("pivot-high-cardinality") \
    .config("spark.sql.pivotMaxValues", 20000)  # 根据实际商品数量调整数值
    .getOrCreate()

# 运行时动态设置
spark.conf.set("spark.sql.pivotMaxValues", 20000)

注意：如果商品数量过大（如几十万级），生成的宽表列数过多会导致内存压力激增，这种场景不建议使用该方法。

2. 分批次处理后合并 #

当商品数量极大时，可将商品列表分批次处理，再通过ID合并结果，避免一次性生成过多列：

from pyspark.sql import functions as F

# 获取所有唯一商品
all_items = [row[0] for row in df1.select(F.explode("Items")).distinct().collect()]
batch_size = 5000  # 自定义每个批次处理的商品数量
batches = [all_items[i:i+batch_size] for i in range(0, len(all_items), batch_size)]

# 初始化结果为仅包含ID的DataFrame
result_df = df1.select("ID")

for batch in batches:
    # 处理当前批次的商品，生成对应列
    batch_df = df1.withColumn("exploded", F.explode("Items")) \
        .filter(F.col("exploded").isin(batch)) \
        .groupBy("ID") \
        .pivot("exploded") \
        .agg(F.lit(1)) \
        .na.fill(0)
    # 按ID合并当前批次结果到总结果
    result_df = result_df.join(batch_df, on="ID", how="left").na.fill(0)

result_df.show(truncate=False)

该方法通过拆分批次降低单步内存开销，但需注意多次join操作的性能损耗，建议根据集群资源调整批次大小。

3. 改用稀疏格式存储（高基数场景推荐） #

若商品数量极多，生成宽表并不实用，推荐改用稀疏格式存储，避免大量列的内存占用：

# 方式1：长表格式（ID + 商品 + 购买标记）
sparse_df = df1.withColumn("exploded", F.explode("Items")) \
    .select("ID", "exploded", F.lit(1).alias("purchased"))
sparse_df.show()

# 方式2：用Map类型存储每个ID的购买状态
map_df = df1.withColumn("item_map", F.map_from_entries(
    F.transform("Items", lambda x: F.struct(F.lit(x).alias("key"), F.lit(1).alias("value")))
)).select("ID", "item_map")
map_df.show(truncate=False)

后续查询特定商品的购买状态时，可通过getItem方法提取：

# 查询Apple的购买状态
map_df.select("ID", F.col("item_map").getItem("Apple").alias("Apple")).na.fill(0).show()

内容的提问来源于stack exchange，提问作者Amit BenDavid