基于JSONB的PostgreSQL线性插值查询实现需求

问题背景 #

现有如下PostgreSQL数据表及数据：

CREATE TABLE sample_table (
    id1 text,
    id2 text,
    data jsonb
);
INSERT INTO sample_table (id1, id2, data)
VALUES
    ('row1', 'abc', '[[0.1, 1.0, 0.2], [0.2, 2.0, 0.4], [0.3, 3.0, 0.6], [0.4, 4.0, 0.8]]'::jsonb),
    ('row2', 'def', '[[0.2, 1.5, 0.3], [0.3, 2.5, 0.5], [0.4, 3.5, 0.7], [0.5, 4.5, 0.9]]'::jsonb),
    ('row3', 'ghi', '[[0.3, 1.2, 0.4], [0.4, 2.2, 0.6], [0.5, 3.2, 0.8], [0.6, 4.2, 1.0]]'::jsonb);

每行由id1和id2唯一标识。需要针对数值2.5，匹配data中每个子数组的第二个元素，找到其所在区间后，利用对应子数组的第三个元素进行线性插值计算。例如：

• row1：0.4 + 0.5*(0.6-0.4) = 0.5
• row2：0.5 + 0*(0.7-0.5) = 0.5
• row3：0.6 + 0.3*(0.8-0.6) = 0.66

最终需要得到包含id1、id2及插值结果的数据集。

解决方案SQL #

WITH expanded_data AS (
    SELECT
        id1,
        id2,
        -- 解析子数组的三个元素为浮点型
        (elem->>0)::float AS val1,
        (elem->>1)::float AS val2,
        (elem->>2)::float AS val3,
        -- 用窗口函数获取同一行组内前一个子数组的val2和val3
        LAG((elem->>1)::float) OVER (PARTITION BY id1, id2 ORDER BY (elem->>1)::float) AS prev_val2,
        LAG((elem->>2)::float) OVER (PARTITION BY id1, id2 ORDER BY (elem->>1)::float) AS prev_val3
    FROM sample_table,
         jsonb_array_elements(data) AS elem
),
interval_matches AS (
    SELECT
        id1,
        id2,
        prev_val2,
        val2 AS curr_val2,
        prev_val3,
        val3 AS curr_val3
    FROM expanded_data
    -- 筛选目标值2.5所在的区间：前一个val2 ≤ 2.5 ≤ 当前val2
    WHERE prev_val2 <= 2.5 AND val2 >= 2.5
    -- 单独处理目标值正好等于某个val2的情况（如row2）
    UNION ALL
    SELECT
        id1,
        id2,
        val2 AS prev_val2,
        val2 AS curr_val2,
        val3 AS prev_val3,
        val3 AS curr_val3
    FROM expanded_data
    WHERE val2 = 2.5
)
SELECT
    id1,
    id2,
    -- 计算线性插值，处理区间端点的特殊情况
    CASE
        WHEN curr_val2 = prev_val2 THEN prev_val3
        ELSE prev_val3 + (2.5 - prev_val2)/(curr_val2 - prev_val2) * (curr_val3 - prev_val3)
    END AS data
FROM interval_matches;

查询结果 #

执行上述SQL后，将得到符合预期的结果：

id1  | id2 | data  
------+-----+-------
 row1 | abc |   0.5
 row2 | def |   0.5
 row3 | ghi |  0.66

逻辑说明 #

1. expanded_data CTE：把每个JSONB数组拆分为单独的行，解析子数组元素为浮点型，同时用LAG窗口函数获取同一行组内前一个子数组的关键值，为区间判断做准备。
2. interval_matches CTE：筛选出目标值2.5对应的区间，同时兼容目标值正好匹配某个子数组第二个元素的场景。
3. 最终查询：套用线性插值公式计算结果，当目标值正好落在区间端点时，直接返回对应的第三个元素值。

内容的提问来源于stack exchange，提问作者Guyon van Rooij