基于时间间隔的PySpark DataFrame Multiindexing透视实现问询

现有如下PySpark DataFrame：

+------+--------------+---------------+----------+----------+
|A_ID  |B_ID          |C_ID           | BEGIN    |   END    |
+------+--------------+---------------+----------+----------+
| 55264|     12       |         4     |2001-01-01|2012-08-05|
| 54897|     12       |         4     |2001-11-01|2012-11-30|
|  8222|     12       |         5     |2001-08-01|2012-12-31|
| 25001|     12       |         4     |2001-10-01|2015-08-30|
| 40001|     12       |         5     |2001-11-01|2012-12-31|
|  8217|     12       |         5     |2001-05-01|2020-12-31|
| 40002|     12       |         5     |2001-05-01|2009-05-01|
| 25002|     12       |         4     |2001-10-01|2020-05-30|
|  8073|     13       |         3     |2002-05-05|2003-05-04|
...

需对该DataFrame执行带多层索引的透视操作，要求：

• 垂直索引为B_ID和A_ID（多层结构）
• 水平维度为从最早BEGIN日期开始的周，命名为w1、w2...
• 单元格值规则：
  • 0：该周无对应C_ID
  • 1：该周存在唯一C_ID
  • 2：该周存在多个唯一C_ID
• 同时包含C_ID的累计天数信息

期望输出格式如下：

+-----+-----+-----+-----+-----+
|    Weeks  | w1  | w2  | w3  | ....
+-----+-----+
|B_ID | A_ID|
+-----+-----+------------------
| 12  |55264|  0  |  1  |  1  |
|     |82226|  2  |  1  |  0  |
|     |80732|
|     |55264|
|     |40001|
|     |54897|       etc...
| 13  |80732|
|     |32444|
...

1. 确定全局最早日期并计算周序号 #

先找到所有BEGIN中的最早日期，以此为基准计算每条记录对应的周序号：

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

# 获取全局最早BEGIN日期
min_begin = df.select(F.min("BEGIN")).first()[0]

# 计算每条记录起始日期对应的周序号（w1, w2...）
df_with_weeks = df.withColumn(
    "start_week",
    F.concat(F.lit("w"), F.floor(F.datediff(F.col("BEGIN"), min_begin)/7) + 1)
).withColumn(
    "end_week",
    F.concat(F.lit("w"), F.floor(F.datediff(F.col("END"), min_begin)/7) + 1)
)

2. 展开每条记录覆盖的所有周 #

由于单条记录的时间区间可能跨多个周，需要生成该区间内所有周的记录：

# 生成每条记录覆盖的周序列并展开
df_expanded = df_with_weeks.withColumn(
    "week_list",
    F.expr("sequence(start_week, end_week)")
).select(
    "A_ID", "B_ID", "C_ID", "BEGIN", "END",
    F.explode("week_list").alias("week")
)

3. 计算每个(A_ID, B_ID, week)组合的状态与累计天数 #

对每个组合统计唯一C_ID数量，同时计算该周内的有效天数：

# 统计每个周的C_ID状态和累计天数
agg_df = df_expanded.groupBy("A_ID", "B_ID", "week").agg(
    # 确定状态：1=唯一C_ID，2=多个C_ID
    F.when(F.countDistinct("C_ID") == 1, 1)
     .when(F.countDistinct("C_ID") >= 2, 2)
     .otherwise(0).alias("status"),
    # 计算该周内的有效天数：取周边界与记录时间区间的交集天数
    F.sum(
        F.datediff(
            F.least(F.col("END"), F.date_add(min_begin, (F.regexp_extract("week", r"\d+", 0).cast("int")*7)-1)),
            F.greatest(F.col("BEGIN"), F.date_add(min_begin, (F.regexp_extract("week", r"\d+", 0).cast("int")-1)*7))
        ) + 1
    ).alias("cum_days")
)

4. 透视生成目标格式 #

以B_ID和A_ID为分组键，透视周列，并将缺失的周填充为0：

# 获取所有周的有序列表
sorted_weeks = sorted(agg_df.select("week").distinct().rdd.flatMap(lambda x: x).collect())

# 透视操作，将周转为列
pivot_df = agg_df.groupBy("B_ID", "A_ID").pivot("week", sorted_weeks).agg(
    F.first("status").alias("status"),
    F.first("cum_days").alias("cum_days")
).fillna(0)

# 可选：将状态和天数合并为结构化字段
pivot_df = pivot_df.select(
    "B_ID", "A_ID",
    *[F.struct(F.col(f"{w}_status").alias("status"), F.col(f"{w}_cum_days").alias("cum_days")).alias(w) for w in sorted_weeks]
)

5. 调整多层索引展示（可选） #

如果需要严格匹配多层索引的可视化格式，可转换为Pandas DataFrame后设置索引：

pd_df = pivot_df.toPandas().set_index(["B_ID", "A_ID"])

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