在DuckDB中无需UDF，能否基于前序值按列更新数据？

示例数据 #

核心逻辑回顾 #

按D列排序后，A列的计算规则：

• 第一行A值固定为0
• 后续行：
  • 若前一行A=0且当前B=1 → 当前A=1
  • 若前一行A=1且当前C=1 → 当前A=0
  • 其他情况继承前一行A的值

DuckDB 实现方法 #

方法1：递归CTE #

DuckDB支持递归CTE，可以逐行迭代计算A值，这是最直观的实现方式：

WITH RECURSIVE ranked_data AS (
    -- 按D排序并添加行号，用于递归关联
    SELECT 
        B, C, D,
        ROW_NUMBER() OVER (ORDER BY D) AS rn
    FROM your_table
),
recursive_a AS (
    -- 初始化第一行A值
    SELECT 
        rn, B, C, D,
        0::INTEGER AS A
    FROM ranked_data
    WHERE rn = 1
    
    UNION ALL
    
    -- 递归计算后续每一行的A值
    SELECT 
        rd.rn, rd.B, rd.C, rd.D,
        CASE
            WHEN ra.A = 0 AND rd.B = 1 THEN 1
            WHEN ra.A = 1 AND rd.C = 1 THEN 0
            ELSE ra.A
        END AS A
    FROM ranked_data rd
    JOIN recursive_a ra ON rd.rn = ra.rn + 1
)
SELECT B, C, A, D
FROM recursive_a
ORDER BY D;

方法2：数组聚合 + list_reduce #

如果递归CTE性能达不到预期，可以用数组聚合将B、C按D排序后转为数组，再通过list_reduce逐元素计算A值：

WITH sorted_arrays AS (
    SELECT 
        array_agg(B ORDER BY D) AS B_arr,
        array_agg(C ORDER BY D) AS C_arr,
        array_agg(D ORDER BY D) AS D_arr
    FROM your_table
),
computed_a AS (
    SELECT
        list_reduce(
            generate_subscripts(B_arr, 1),
            [0::INTEGER],
            (acc, idx) -> 
                acc || CASE
                    WHEN acc[idx-1] = 0 AND B_arr[idx] = 1 THEN 1
                    WHEN acc[idx-1] = 1 AND C_arr[idx] = 1 THEN 0
                    ELSE acc[idx-1]
                END
        ) AS A_arr,
        B_arr, C_arr, D_arr
    FROM sorted_arrays
)
SELECT
    unnest(B_arr) AS B,
    unnest(C_arr) AS C,
    unnest(A_arr) AS A,
    unnest(D_arr) AS D
FROM computed_a;

方法3：自定义UDF（可选） #

如果上述方法都不满足需求，可以编写UDF模拟R中的循环逻辑：

-- 创建计算A列的UDF
CREATE OR REPLACE FUNCTION compute_a(B_arr INTEGER[], C_arr INTEGER[])
RETURNS INTEGER[]
LANGUAGE plpgsql
AS $$
DECLARE
    len INTEGER;
    A_arr INTEGER[];
    i INTEGER;
BEGIN
    len := array_length(B_arr, 1);
    A_arr := array_fill(0, ARRAY[len]);
    
    FOR i IN 2..len LOOP
        A_arr[i] := A_arr[i-1];
        IF A_arr[i-1] = 0 AND B_arr[i] = 1 THEN
            A_arr[i] := 1;
        ELSIF A_arr[i-1] = 1 AND C_arr[i] = 1 THEN
            A_arr[i] := 0;
        END IF;
    END LOOP;
    
    RETURN A_arr;
END;
$$;

-- 使用UDF计算结果
WITH sorted_data AS (
    SELECT
        array_agg(B ORDER BY D) AS B_arr,
        array_agg(C ORDER BY D) AS C_arr,
        array_agg(D ORDER BY D) AS D_arr
    FROM your_table
)
SELECT
    unnest(B_arr) AS B,
    unnest(C_arr) AS C,
    unnest(compute_a(B_arr, C_arr)) AS A,
    unnest(D_arr) AS D
FROM sorted_data;

其他SQL变体实现 #

PostgreSQL #

递归CTE写法与DuckDB完全一致，也支持相同的数组UDF逻辑。

BigQuery #

使用递归CTE实现：

WITH RECURSIVE ranked_data AS (
    SELECT 
        B, C, D,
        ROW_NUMBER() OVER (ORDER BY D) AS rn
    FROM your_table
),
recursive_a AS (
    SELECT rn, B, C, D, 0 AS A
    FROM ranked_data WHERE rn = 1
    UNION ALL
    SELECT 
        rd.rn, rd.B, rd.C, rd.D,
        CASE
            WHEN ra.A = 0 AND rd.B = 1 THEN 1
            WHEN ra.A = 1 AND rd.C = 1 THEN 0
            ELSE ra.A
        END AS A
    FROM ranked_data rd
    JOIN recursive_a ra ON rd.rn = ra.rn + 1
)
SELECT B, C, A, D FROM recursive_a ORDER BY D;

Spark SQL #

递归CTE实现逻辑相同：

WITH RECURSIVE recursive_a AS (
    SELECT 
        B, C, D,
        0 AS A,
        ROW_NUMBER() OVER (ORDER BY D) AS rn
    FROM your_table
    UNION ALL
    SELECT 
        rd.B, rd.C, rd.D,
        CASE
            WHEN ra.A = 0 AND rd.B = 1 THEN 1
            WHEN ra.A = 1 AND rd.C = 1 THEN 0
            ELSE ra.A
        END AS A,
        rd.rn
    FROM your_table rd
    JOIN recursive_a ra ON rd.rn = ra.rn + 1
)
SELECT B, C, A, D FROM recursive_a ORDER BY D;

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