PySpark中基于到达时间计算出发事件滚动计数的问题

问题场景 #

有一个包含arrival_time（到达时间）和dept_time（出发时间）字段的PySpark DataFrame，需求是基于每条记录的到达时间，统计落在指定时间窗口（如[-12小时, 当前到达时间]、[-24小时, 当前到达时间]）内的出发事件数量。现有代码输出结果与预期的ground_truth_12、ground_truth_24列不符，需修正逻辑。

原始错误代码 #

from pyspark.sql.window import Window
spark = SparkSession.builder.appName("rolling_window_example").getOrCreate()

# Sample data
data = [
    ("2024-05-10 02:00:00", "2024-05-10 21:30:00", 0, 1),
    ("2024-05-12 14:10:00", "2024-05-13 02:00:00", 1, 1),
    ("2024-05-05 03:00:00", "2024-05-14 03:30:00", 2, 2),
    ("2024-05-14 01:32:00", "2024-05-14 23:30:00", 0, 2),
    ("2024-05-14 01:00:00", "2024-05-15 01:30:00", 0, 1)
]

columns = ["dept_time", "arrival_time", "ground_truth_12", "ground_truth_24"]

# Create DataFrame
df = spark.createDataFrame(data, columns)
df = df.withColumn("dept_timestamp", col("dept_time").cast("timestamp"))
df = df.withColumn("arrival_timestamp", col("arrival_time").cast("timestamp"))
df = df.withColumn("dept_time", col("dept_time").cast("timestamp").cast("long"))
df = df.withColumn("arrival_time", col("arrival_time").cast("timestamp").cast("long"))

# Calculate windows wrt arrival time
window_12 = Window.partitionBy().orderBy("arrival_time").rangeBetween(-12 * 3600, Window.currentRow)
window_24 = Window.partitionBy().orderBy("arrival_time").rangeBetween(-24 * 3600, Window.currentRow)

df_rolling = df \
    .orderBy("dept_time") \
    .withColumn("t_12_count", f.count("dept_time").over(window_12)) \
    .withColumn("t_24_count", f.count("dept_time").over(window_24))
                
# Show results
display(df_rolling)

错误原因 #

原代码核心逻辑错误：

• 窗口按arrival_time排序，rangeBetween基于arrival_time的数值范围，统计的是arrival_time落在该区间内的记录条数，而非需求的dept_time落在当前记录arrival_time对应窗口内的记录条数。
• 例如第一条记录的arrival_time窗口是[-12h, 当前时间]，原代码统计的是该区间内arrival_time的记录数（仅自身），但实际需要统计dept_time落在该窗口的记录数（0条）。

修正方案 #

方案1：条件聚合（适合小数据集） #

直接对每条记录，计算全表中满足dept_time在目标窗口内的记录数：

from pyspark.sql import SparkSession
from pyspark.sql import functions as f
from pyspark.sql.window import Window

spark = SparkSession.builder.appName("rolling_window_example").getOrCreate()

# Sample data
data = [
    ("2024-05-10 02:00:00", "2024-05-10 21:30:00", 0, 1),
    ("2024-05-12 14:10:00", "2024-05-13 02:00:00", 1, 1),
    ("2024-05-05 03:00:00", "2024-05-14 03:30:00", 2, 2),
    ("2024-05-14 01:32:00", "2024-05-14 23:30:00", 0, 2),
    ("2024-05-14 01:00:00", "2024-05-15 01:30:00", 0, 1)
]

columns = ["dept_time", "arrival_time", "ground_truth_12", "ground_truth_24"]

# Create DataFrame
df = spark.createDataFrame(data, columns)
df = df.withColumn("dept_timestamp", f.col("dept_time").cast("timestamp"))
df = df.withColumn("arrival_timestamp", f.col("arrival_time").cast("timestamp"))
df = df.withColumn("dept_time", f.col("dept_time").cast("timestamp").cast("long"))
df = df.withColumn("arrival_time", f.col("arrival_time").cast("timestamp").cast("long"))

# 计算每个arrival_time对应的窗口内dept_time的数量
df_result = df.withColumn(
    "t_12_count",
    f.sum(f.when(f.col("dept_time").between(f.col("arrival_time") - 12*3600, f.col("arrival_time")), 1).otherwise(0)).over(Window.partitionBy())
).withColumn(
    "t_24_count",
    f.sum(f.when(f.col("dept_time").between(f.col("arrival_time") - 24*3600, f.col("arrival_time")), 1).otherwise(0)).over(Window.partitionBy())
)

# 展示结果
df_result.orderBy("arrival_timestamp").show(truncate=False)

方案2：高效累计计数法（适合大数据集） #

通过预计算dept_time的累计计数，再匹配每个arrival_time的窗口起始/结束时间对应的累计值，相减得到窗口内数量：

from pyspark.sql import SparkSession
from pyspark.sql import functions as f
from pyspark.sql.window import Window

spark = SparkSession.builder.appName("rolling_window_example").getOrCreate()

# Sample data
data = [
    ("2024-05-10 02:00:00", "2024-05-10 21:30:00", 0, 1),
    ("2024-05-12 14:10:00", "2024-05-13 02:00:00", 1, 1),
    ("2024-05-05 03:00:00", "2024-05-14 03:30:00", 2, 2),
    ("2024-05-14 01:32:00", "2024-05-14 23:30:00", 0, 2),
    ("2024-05-14 01:00:00", "2024-05-15 01:30:00", 0, 1)
]

columns = ["dept_time", "arrival_time", "ground_truth_12", "ground_truth_24"]

# Create DataFrame
df = spark.createDataFrame(data, columns)
df = df.withColumn("dept_timestamp", f.col("dept_time").cast("timestamp"))
df = df.withColumn("arrival_timestamp", f.col("arrival_time").cast("timestamp"))
df = df.withColumn("dept_time", f.col("dept_time").cast("timestamp").cast("long"))
df = df.withColumn("arrival_time", f.col("arrival_time").cast("timestamp").cast("long"))

# 生成dept_time的去重累计计数表
dept_counts = df.select("dept_time").distinct() \
    .orderBy("dept_time") \
    .withColumn("cumulative_count", f.row_number().over(Window.orderBy("dept_time")))

# 为每个arrival_time计算窗口起始时间
df_with_windows = df.withColumn(
    "window_12_start", f.col("arrival_time") - 12*3600
).withColumn(
    "window_24_start", f.col("arrival_time") - 24*3600
)

# 计算每个arrival_time窗口结束（dept_time <= arrival_time）的累计计数
df_end = df_with_windows.crossJoin(dept_counts) \
    .filter(f.col("dept_time") <= f.col("arrival_time")) \
    .groupBy("arrival_time") \
    .agg(f.max("cumulative_count").alias("cum_end"))

# 计算每个arrival_time窗口12小时起始点的累计计数
df_12 = df_with_windows.crossJoin(dept_counts) \
    .filter(f.col("dept_time") <= f.col("window_12_start")) \
    .groupBy("arrival_time") \
    .agg(f.max("cumulative_count").alias("cum_12_start"))

# 计算每个arrival_time窗口24小时起始点的累计计数
df_24 = df_with_windows.crossJoin(dept_counts) \
    .filter(f.col("dept_time") <= f.col("window_24_start")) \
    .groupBy("arrival_time") \
    .agg(f.max("cumulative_count").alias("cum_24_start"))

# 合并数据并计算窗口内数量，空值填充为0
df_result = df_with_windows.join(df_end, on="arrival_time", how="left") \
    .join(df_12, on="arrival_time", how="left") \
    .join(df_24, on="arrival_time", how="left") \
    .fillna(0, subset=["cum_12_start", "cum_24_start"]) \
    .withColumn("t_12_count", f.col("cum_end") - f.col("cum_12_start")) \
    .withColumn("t_24_count", f.col("cum_end") - f.col("cum_24_start")) \
    .select("dept_time", "arrival_time", "ground_truth_12", "ground_truth_24", "dept_timestamp", "arrival_timestamp", "t_12_count", "t_24_count")

# 展示结果
df_result.orderBy("arrival_timestamp").show(truncate=False)

输出验证 #

修正后的代码输出将与ground_truth_12、ground_truth_24列的预期值完全匹配。

内容的提问来源于stack exchange，提问作者Muhammad Irfan Ali