PySpark 3.2：按Version分组用众数填充Color空值，求无需Join的优化方案

问题概述 #

使用PySpark 3.2版本（无法升级，也无法使用内置mode聚合函数），需要按version字段分组，将分组内color字段的空值用该分组的color众数填充。当前已通过计算分组众数+Join关联的方式实现，现寻求无需Join的更优方案。

当前实现代码（修正原代码语法错误后）：

from pyspark.sql import functions as F
from pyspark.sql.window import Window

window_spec = Window.partitionBy(F.col('version')).orderBy(F.col('count_granularity').desc())
window_granularity = Window.partitionBy(F.col('version'), F.col('color'))

mode = df.withColumn('count_granularity', F.sum(F.lit(1)).over(window_granularity)) \
         .withColumn('rank', F.row_number().over(window_spec)) \
         .filter(F.col('rank') == 1) \
         .withColumnRenamed('color', 'mode') \
         .select('version', 'mode')

df = df.join(mode, on='version', how='left') \
       .withColumn('mode_color', F.when(df.color.isNull(), mode.mode).otherwise(df.color)) \
       .drop('color', 'mode') \
       .withColumnRenamed('mode_color', 'color')

无需Join的优化方案 #

可以通过嵌套窗口函数直接在原DataFrame中计算分组众数并完成空值填充，全程避免Join操作，减少数据shuffle开销。

实现代码 #

from pyspark.sql import functions as F
from pyspark.sql.window import Window

# 1. 计算每个(version, color)组合的出现次数
count_window = Window.partitionBy('version', 'color')
# 2. 按version分组，按出现次数降序排序（若次数相同，按color字典序稳定排序）
rank_window = Window.partitionBy('version').orderBy(F.col('color_count').desc(), F.col('color'))
# 3. 提取每个version分组的众数
mode_window = Window.partitionBy('version')

df_filled = df.withColumn('color_count', F.count('*').over(count_window)) \
              .withColumn('rank', F.row_number().over(rank_window)) \
              .withColumn('group_mode', F.first('color').over(mode_window)) \
              # 4. 用众数填充空值
              .withColumn('color', F.when(F.col('color').isNull(), F.col('group_mode')).otherwise(F.col('color'))) \
              .drop('color_count', 'rank', 'group_mode')

方案细节说明 #

• 用F.count('*').over(count_window)替代原代码的F.sum(F.lit(1))，逻辑更直观高效
• rank_window中添加F.col('color')排序，确保当分组存在多个众数时，结果稳定可预期
• 通过F.first('color').over(mode_window)直接在原DataFrame的每个version分组内提取众数，无需额外生成众数表再Join
• 直接在原color字段上完成空值替换，减少中间列的创建与删除操作

注意点 #

• 若业务允许分组存在多个众数时随机选择，可去掉rank_window中的F.col('color')排序项
• 该方案全程基于窗口函数操作，避免了Join带来的跨表shuffle，在大数据量场景下性能更优

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