Spark DataFrame按组扩展唯一value变体并填充全组行

原数据与需求 #

原DataFrame代码 #

df = spark.createDataFrame(
    [(0, 'foo', '2020-01-01', '2020-02-01'),
     (0, 'bar', '2020-02-01', '2020-03-01'),
     (0, 'foo', '2020-03-01', '2020-04-01'),
     (0, None, '2020-04-01', '2020-05-01'),
     (1, 'bar', '2020-01-01', '2020-02-01'),
     (1, 'foo', '2020-02-01', '2020-03-01'),
     (2, None, '2020-02-01', '2020-03-01'),
     (2, None, '2020-04-01', '2020-07-01')
     ],
    ['group', 'value', 'start', 'end'])

df.show()

需求说明 #

• 按group字段分组，提取每个分组内value列的唯一非None值
• 对每个唯一非None值，生成与原分组行数相同的行，并用该值填充所有行的value列
• 若分组内value全为None，直接保留原分组数据

目标结果 #

group value start      end  
0     foo   2020-01-01 2020-02-01  
0     foo   2020-02-01 2020-03-01      
0     foo   2020-03-01 2020-04-01 
0     foo   2020-04-01 2020-05-01
0     bar   2020-01-01 2020-02-01   
0     bar   2020-02-01 2020-03-01  
0     bar   2020-03-01 2020-04-01
0     bar   2020-04-01 2020-05-01
1     bar   2020-01-01 2020-02-01  
1     bar   2020-02-01 2020-03-01
1     foo   2020-01-01 2020-02-01  
1     foo   2020-02-01 2020-03-01
2     None  2020-02-01 2020-03-01
2     None  2020-04-01 2020-07-01

解决方案 #

实现代码 #

from pyspark.sql import functions as F

# 1. 分组统计每个组的唯一非空value集合，以及非空值数量
group_stats = df.groupBy("group").agg(
    # 收集非空的唯一value，自动去重
    F.collect_set(F.when(F.col("value").isNotNull(), F.col("value"))).alias("unique_values"),
    # 统计非空值的数量，用于判断是否全为None
    F.count(F.when(F.col("value").isNotNull(), 1)).alias("non_null_count")
)

# 2. 关联原表，处理扩展行与值填充
result = df.join(group_stats, on="group", how="left") \
    .withColumn(
        "value",
        F.when(
            # 全为None时保留原value
            F.col("non_null_count") == 0,
            F.col("value")
        ).otherwise(
            # 非空时展开唯一值列表，实现行扩展
            F.explode(F.col("unique_values"))
        )
    ) \
    .drop("unique_values", "non_null_count")  # 删除中间辅助列

# 排序后查看结果
result.orderBy("group", "value", "start").show()

代码说明 #

1. 分组统计阶段：
   • collect_set(F.when(...))：过滤掉value为None的记录，收集每个分组内的唯一非空值，自动去重得到集合。
   • count(F.when(...))：统计分组内非空值的数量，用来判断该分组是否全为None。
2. 关联与扩展阶段：
   • 将原表与分组统计结果按group关联，确保每个原行都能获取到分组的唯一值集合。
   • 使用when-otherwise分支处理两种情况：
     • 当non_null_count为0（全为None），直接保留原value。
     • 否则用explode展开唯一值集合，每个唯一值都会和原分组的每一行生成一条新记录，实现行扩展并填充对应值。
3. 最后删除中间辅助列，排序后即可得到目标结果。

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