PySpark窗口函数：基于时间与行数双条件收集历史行咨询

需求概述 #

为PySpark数据框新增一列，收集当前行的历史数据，需同时满足两个条件：

• 仅保留与当前行日期相差不超过10天的历史行；
• 最多保留最近2条符合时间条件的历史行。

输入数据示例 #

期望输出示例 #

当前代码局限 #

现有窗口函数仅能按固定行数选取历史行，无法加入时间差过滤逻辑：

window = Window.orderBy('date').partitionBy('id').rowsBetween(-3, -1)

解决方案 #

由于需要同时满足时间范围和行数上限两个动态条件，无法直接使用固定行范围的窗口函数，需结合collect_list、filter和slice函数实现：

完整实现代码 #

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

# 初始化Spark会话
spark = SparkSession.builder.appName("HistoryRowCollection").getOrCreate()

# 模拟输入数据（实际使用时替换为你的数据源）
data = [
    (1, "2023-01-01", 100),
    (1, "2023-05-01", 200),
    (1, "2023-05-02", 300),
    (1, "2023-05-03", 400),
    (1, "2023-05-04", 500)
]
df = spark.createDataFrame(data, ["id", "日期", "值"])

# 确保日期列为Date类型
df = df.withColumn("日期", F.to_date("日期"))

# 定义窗口：按id分区，日期升序，收集当前行之前的所有历史行
window_spec = Window.partitionBy("id").orderBy("日期")
df = df.withColumn(
    "所有历史行",
    F.collect_list(F.struct("id", "日期", "值")).over(window_spec.rowsBetween(Window.unboundedPreceding, -1))
)

# 1. 过滤出日期差≤10天的历史行；2. 取最近2行；3. 格式化为目标字符串
df = df.withColumn(
    "符合条件的历史行",
    F.expr("filter(所有历史行, x -> datediff(日期, x.日期) <= 10)")
).withColumn(
    "历史记录",
    F.expr("array_join(slice(reverse(符合条件的历史行), 1, 2), ', ')")
).withColumn(
    "历史记录",
    F.when(F.size("符合条件的历史行") == 0, "()").otherwise(F.concat(F.lit("("), F.col("历史记录"), F.lit(")")))
)

# 清理中间辅助列
df = df.drop("所有历史行", "符合条件的历史行")

# 查看结果
df.show(truncate=False)

关键步骤说明 #

1. 收集所有历史行：通过窗口函数collect_list收集当前行之前的所有历史数据，存储为结构化数组。
2. 时间条件过滤：用filter函数筛选出与当前行日期差不超过10天的记录。
3. 限制行数：通过reverse反转数组（让最近的行排在首位），再用slice截取前2条，确保最多保留2行。
4. 格式转换：用array_join将数组转为字符串，并处理空列表的特殊显示格式。

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