如何用优雅的Spark SQL关联方式生成指定的TABLE_RESULT表？

问题 #

现有两张源表结构及测试数据如下：

TABLE_1 #

CREATE TABLE TABLE_1(
CALL_ID INT,
CALL_DATE DATE);
INSERT INTO TABLE_1(CALL_ID, CALL_DATE)
VALUES (1, '2022-10-22'),
(2, '2022-10-31'),
(3, '2022-11-04');

TABLE_2 #

CREATE TABLE TABLE_2(
PROD_ID INT,
PROD_DATE DATE);
INSERT INTO TABLE_2(PROD_ID, PROD_DATE)
VALUES (1, '2022-10-25'),
(2, '2022-11-17');

期望生成的目标表TABLE_RESULT结构及示例数据：

CREATE TABLE TABLE_RESULT(
CALL_ID INT,
CALL_DATE DATE,
PROD_ID INT,
PROD_DATE DATE);
INSERT INTO TABLE_RESULT(CALL_ID, CALL_DATE, PROD_ID, PROD_DATE)
VALUES (1, '2022-10-22', 1, '2022-10-25'),
(2, '2022-10-31', NULL, NULL),
(3, '2022-11-04', 2, '2022-11-17');

需求：用Spark SQL实现关联逻辑，将TABLE_1的每条记录匹配TABLE_2中最近的晚于其CALL_DATE的PROD_DATE记录，无匹配时对应PROD字段为NULL，生成上述目标表。此场景为小数据量，需合适的实现方案。

解决方案 #

针对小数据量场景，可通过左连接+窗口函数的方式优雅实现需求，具体步骤如下：

实现思路 #

1. 先将TABLE_1与TABLE_2做左连接，筛选出PROD_DATE > CALL_DATE的有效匹配记录；
2. 对每个CALL_ID分组，用窗口函数对匹配到的记录按PROD_DATE与CALL_DATE的时间差升序排序，标记出每个分组中时间差最小（即最近）的那条记录；
3. 最后过滤出标记为第一条的记录，同时保留TABLE_1中无匹配的原始记录。

Spark SQL 代码 #

WITH matched_records AS (
    SELECT 
        t1.CALL_ID,
        t1.CALL_DATE,
        t2.PROD_ID,
        t2.PROD_DATE,
        -- 计算时间差，用于排序
        DATEDIFF(t2.PROD_DATE, t1.CALL_DATE) AS date_diff,
        -- 按CALL_ID分组，按时间差升序排序，标记最近的记录
        ROW_NUMBER() OVER (PARTITION BY t1.CALL_ID ORDER BY DATEDIFF(t2.PROD_DATE, t1.CALL_DATE) ASC) AS rn
    FROM TABLE_1 t1
    LEFT JOIN TABLE_2 t2 ON t2.PROD_DATE > t1.CALL_DATE
)
SELECT 
    CALL_ID,
    CALL_DATE,
    CASE WHEN rn = 1 THEN PROD_ID ELSE NULL END AS PROD_ID,
    CASE WHEN rn = 1 THEN PROD_DATE ELSE NULL END AS PROD_DATE
FROM matched_records
WHERE rn = 1 OR rn IS NULL
ORDER BY CALL_ID;

代码说明 #

• CTE部分：通过左连接获取所有可能的匹配（包括无匹配的记录），同时计算时间差并给每个CALL_ID下的匹配记录排序；
• 主查询部分：筛选出每个CALL_ID下排序为1的最近匹配记录，以及无匹配的原始记录，最终输出符合要求的结果。

由于是小数据量，左连接的开销完全可控，窗口函数的计算效率也能满足需求。

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