PySpark：分组筛选Timestamp2之后的首个Timestamp1

问题背景 #

现有数据集如下：

ID Timestamp1             Timestamp2
 1  2022-01-01T12:00:00    2022-01-01T12:30:00
 1  2022-01-01T13:00:00    2022-01-01T12:30:00
 1  2022-01-01T14:00:00    2022-01-01T12:30:00
 2  2022-02-02T09:00:00    2022-02-02T10:15:00
 2  2022-02-02T10:00:00    2022-02-02T10:15:00
 2  2022-02-03T11:00:00    2022-02-02T10:15:00

每个唯一ID对应的Timestamp2值一致，需要按ID分组，筛选出每组中晚于Timestamp2的首个Timestamp1，预期结果：

ID Timestamp1             Timestamp2
 1  2022-01-01T13:00:00    2022-01-01T12:30:00
 2  2022-02-03T11:00:00    2022-02-02T10:15:00

两种可行解法对比 #

1. GroupBy 优化方案

你的初始思路方向是对的，但可以简化：直接用Timestamp1 > Timestamp2过滤符合条件的记录，再按ID分组取最小的Timestamp1——因为最小的符合条件的Timestamp1就是每组里首个晚于Timestamp2的那条。

SQL代码示例：

SELECT 
    ID,
    MIN(Timestamp1) AS Timestamp1,
    MAX(Timestamp2) AS Timestamp2  -- 同ID的Timestamp2一致，MIN/MAX效果相同
FROM 
    your_table
WHERE 
    Timestamp1 > Timestamp2
GROUP BY 
    ID;

优势：逻辑直白，执行效率高，大数据量下数据库的聚合优化器能快速处理，没有额外排序开销。

2. 窗口函数（ROW_NUMBER）方案

如果后续有扩展需求（比如要取前N条符合条件的记录），窗口函数会更灵活。思路是：先过滤出符合条件的记录，按ID分组后对每组内的Timestamp1升序排序，取排序后的第一条（ROW_NUMBER=1）。

SQL代码示例：

WITH filtered_ranked AS (
    SELECT 
        ID,
        Timestamp1,
        Timestamp2,
        ROW_NUMBER() OVER (
            PARTITION BY ID 
            ORDER BY Timestamp1 ASC
        ) AS rn
    FROM 
        your_table
    WHERE 
        Timestamp1 > Timestamp2
)
SELECT 
    ID,
    Timestamp1,
    Timestamp2
FROM 
    filtered_ranked
WHERE 
    rn = 1;

优势：灵活性强，调整范围（比如取前2条）只需修改rn <= 2即可，但仅取第一条时，性能不如GroupBy方案。

选择建议 #

• 仅需取每组首个符合条件的记录：优先用GroupBy方案，代码简洁且高效。
• 有复杂排序/多记录筛选需求：用窗口函数方案更适配。

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