如何在PySpark中高效实现Pandas列转二进制特征字典功能？

核心性能问题说明 #

你之前的实现性能差的核心原因是两个：

• 逐列触发collect操作，每列都要全量扫描一次数据集，100列就要扫描100次全表
• RDD转换本身开销远高于Spark SQL内置算子，加上全量数据拉取到Driver端的IO开销，最终耗时被指数放大

优化后实现代码 #

核心逻辑是仅触发2次全表扫描，所有中间计算都在集群分布式执行：

from pyspark.sql import functions as F

def binary_dict_spark(spark_df, feat_list):
    dict_feature = dict()
    # 第一步：一次扫描全表，拿到所有目标列的唯一值列表，避免逐列查询
    agg_exprs = [F.collect_set(col).alias(col) for col in feat_list]
    col_unique_vals = spark_df.agg(*agg_exprs).collect()[0].asDict()
    
    # 第二步：为每个列的每个唯一值生成对应的0/1标识列
    select_exprs = []
    col_new_names = []
    for col in feat_list:
        # 不需要空值的话可以加过滤：unique_vals = [v for v in col_unique_vals[col] if v is not None]
        unique_vals = col_unique_vals[col]
        for val in unique_vals:
            new_col_name = f"{col}_{val}"
            # 用when生成0/1列，底层直接走Spark SQL优化，分布式执行无额外开销
            select_exprs.append(F.when(F.col(col) == val, 1).otherwise(0).alias(new_col_name))
            col_new_names.append(new_col_name)
    
    # 第三步：仅一次collect拉取所有0/1列的结果到Driver端
    result_df = spark_df.select(*select_exprs)
    result_rows = result_df.collect()
    
    # 转换为要求的字典格式
    for col_name in col_new_names:
        dict_feature[col_name] = [row[col_name] for row in result_rows]
    
    return dict_feature

注：如果需要严格保证输出的0/1列表顺序和原始数据行顺序一致，可以在执行函数前先给原DataFrame加序索引：spark_df = spark_df.withColumn("idx", F.monotonically_increasing_id()).orderBy("idx")

适配Azure Synapse的额外优化建议 #

• 如果你的数据集超过100万行，建议提前调整Driver端内存配置，可在Synapse Spark池配置中设置spark.driver.memory到足够大小，避免最终collect时OOM
• 如果单列唯一值超过1000个，建议先过滤低频次唯一值再计算，避免生成的字典过大占用过多内存
• PySpark 3.1默认开启自适应执行，无需额外配置即可自动优化shuffle并行度

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