PySpark中高效计算事件间关联DataFrame时长总和的方法

问题背景 #

我有两个PySpark DataFrame：

• df1：存储事件列表，包含当前事件时间datetime和前一事件时间prev_datetime
• df2：存储带duration字段的事件，需计算其中时间落在df1每行prev_datetime与datetime区间内的duration总和，新增到df1中

示例数据（实际使用Unix时间戳） #

df1：

df2：

需求示例：Second事件对应的总和为60（B和C的duration之和）

之前尝试的merge+groupby在大数据集下过慢，直接跨DF引用列的代码无法运行，需要高效实现方案。

解决方案 #

直接跨DF引用列会失败，因为Spark需要显式关联两个DataFrame。根据df2的数据量大小，分两种高效方案：

方案一：df2为小表时（广播优化+区间Join+聚合） #

如果df2数据量较小，使用广播（Broadcast）将df2分发到所有Executor，避免Shuffle，大幅提升Join效率：

from pyspark.sql import functions as F
from pyspark.sql.functions import broadcast

# 1. 转换时间为Unix时间戳（兼容实际场景）
df1 = df1.withColumn("datetime_ts", F.col("datetime").cast("long")) \
         .withColumn("prev_datetime_ts", F.col("prev_datetime").cast("long"))

df2 = df2.withColumn("datetime_ts", F.col("datetime").cast("long"))

# 2. 广播df2并执行区间Join
joined_df = df1.join(
    broadcast(df2),
    # 匹配区间：df2时间在prev_datetime和datetime之间；处理prev_datetime为null的边界情况
    (F.col("datetime_ts") > F.col("prev_datetime_ts")) & (F.col("datetime_ts") < F.col("datetime_ts")) |
    (F.col("prev_datetime_ts").isNull() & (F.col("datetime_ts") < F.col("datetime_ts"))),
    how="left"
)

# 3. 按df1的事件维度聚合，计算duration总和
result_df = joined_df.groupBy(
    df1["event"], df1["datetime"], df1["prev_datetime"]
).agg(
    F.coalesce(F.sum("duration"), F.lit(0)).alias("total_duration")
)

方案二：df2为大表时（累积和+差值计算） #

如果df2数据量极大，全量Join会产生大量中间数据，此时可以通过累积和+区间差值的方式避免全量关联：

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

# 1. 转换时间为Unix时间戳
df1 = df1.withColumn("datetime_ts", F.col("datetime").cast("long")) \
         .withColumn("prev_datetime_ts", F.col("prev_datetime").cast("long"))

df2 = df2.withColumn("datetime_ts", F.col("datetime").cast("long"))

# 2. 对df2按时间排序，计算累积duration总和
window = Window.orderBy("datetime_ts")
df2_cumulative = df2.withColumn(
    "cumulative_duration",
    F.sum("duration").over(window)
)

# 添加初始行（处理prev_datetime为null的情况，确保区间从最早时间开始）
min_ts = df2_cumulative.select(F.min("datetime_ts")).first()[0]
dummy_row = spark.createDataFrame(
    [("dummy", min_ts - 1, 0, 0)],
    df2_cumulative.columns
)
df2_cumulative = df2_cumulative.union(dummy_row)

# 3. 匹配df1每行的datetime对应的最大累积和
df1_with_current = df1.join(
    df2_cumulative,
    df1["datetime_ts"] >= df2_cumulative["datetime_ts"],
    how="left"
).groupBy(
    df1["event"], df1["datetime"], df1["prev_datetime"], df1["datetime_ts"], df1["prev_datetime_ts"]
).agg(
    F.max("cumulative_duration").alias("current_total")
)

# 4. 匹配df1每行的prev_datetime对应的最大累积和
df1_with_prev = df1_with_current.join(
    df2_cumulative,
    df1_with_current["prev_datetime_ts"] >= df2_cumulative["datetime_ts"],
    how="left"
).groupBy(
    df1_with_current["event"], df1_with_current["datetime"], df1_with_current["prev_datetime"]
).agg(
    F.max("current_total").alias("current_total"),
    F.coalesce(F.max("cumulative_duration"), F.lit(0)).alias("prev_total")
)

# 5. 计算区间内的duration总和
result_df = df1_with_prev.withColumn(
    "total_duration",
    F.col("current_total") - F.col("prev_total")
).drop("current_total", "prev_total")

为什么之前的代码失败？ #

你尝试的代码直接在withColumn中引用df2的列，但Spark的DataFrame是分布式数据集，两个无关联的DF无法直接跨数据集引用列，必须通过join操作建立关联关系后才能进行计算。

结果验证 #

按示例数据运行后，result_df的输出如下：

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