如何在PySpark中实现基于日期的多类别补全逻辑？

核心逻辑拆解 #

针对每个id：

1. 计算每个类别首次出现的日期，取这些日期的最大值作为临界日期（比如id=1的临界日期是2022-06-13，即所有类别都已首次出现的日期）
2. 临界日期之前的记录：只保留原始数据中实际存在的(id, date, category)组合
3. 临界日期及之后的记录：补全所有类别，即使该日期下原数据没有对应类别的记录

代码实现 #

假设原始DataFrame结构为：id: int, date: date, category: string, value: int（若没有value列，可自行调整代码）

1. 计算每个(id, category)的首次出现日期 #

from pyspark.sql import functions as F

# 分组计算每个id下每个类别的首次出现日期
first_occur_df = df.groupBy("id", "category")\
                   .agg(F.min("date").alias("first_date"))

2. 确定每个id的临界日期 #

# 对每个id，取所有类别首次出现日期的最大值作为临界日期
cutoff_df = first_occur_df.groupBy("id")\
                          .agg(F.max("first_date").alias("cutoff_date"))

3. 生成所有需要的(id, date, category)组合 #

先提取每个id的所有日期，再和所有类别做笛卡尔积，生成完整的组合：

# 提取每个id的所有唯一日期
id_dates_df = df.groupBy("id")\
                .agg(F.collect_set("date").alias("dates"))\
                .withColumn("date", F.explode("dates"))\
                .drop("dates")

# 提取所有唯一类别（如果类别固定为A、B，可直接用下方注释的代码）
categories = df.select("category").distinct().rdd.flatMap(lambda x: x).collect()
categories_df = spark.createDataFrame([(categories,)], ["categories"])\
                     .withColumn("category", F.explode("categories"))\
                     .drop("categories")
# 固定类别时用这行替代：
# categories_df = spark.createDataFrame([("A",), ("B",)], ["category"])

# 生成id、日期、类别的笛卡尔积
full_combinations_df = id_dates_df.crossJoin(categories_df)

4. 关联数据并过滤有效记录 #

# 关联临界日期
full_combinations_with_cutoff = full_combinations_df.join(cutoff_df, on="id", how="left")

# 关联原始数据，获取对应value（无数据则为null）
final_df = full_combinations_with_cutoff.join(
    df, 
    on=["id", "date", "category"], 
    how="left"
).withColumn(
    # 过滤规则：临界日期前只保留原数据存在的记录，之后保留所有组合
    "valid_record",
    F.when(F.col("date") < F.col("cutoff_date"), F.col("value").isNotNull())
     .otherwise(F.lit(True))
).filter(F.col("valid_record"))\
 .drop("valid_record", "cutoff_date")

5. 排序输出（可选） #

final_df.orderBy("id", "date", "category").show()

特殊情况处理 #

如果原始数据没有value列，需调整过滤条件，直接判断原数据是否存在对应组合：

final_df = full_combinations_with_cutoff.join(
    df, 
    on=["id", "date", "category"], 
    how="left"
).withColumn(
    "valid_record",
    F.when(F.col("date") < F.col("cutoff_date"), 
           F.exists(
               df.select("id", "date", "category"),
               F.struct(F.col("id"), F.col("date"), F.col("category")) == F.struct(F.col("id"), F.col("date"), F.col("category"))
           ))
     .otherwise(F.lit(True))
).filter(F.col("valid_record"))\
 .drop("valid_record", "cutoff_date")

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