如何在 PostgreSQL 中使用数组和循环加速自定义窗口函数？

Question

我目前正在学习 UDF，并在下面编写了 PostgreSQL UDF 来计算平均偏差 (MAD)。它是任何窗口的平均值和当前值之间的平均绝对差。 在 python pandas/numpy 中，要找到 MAD，我们可以这样写：

series_mad = abs(series - series.mean()).mean()

其中 series 是一组数字，而 series_mad 是表示该系列的 MAD 的单个数值。

我正在尝试使用 Windows 和 UDF 在 PostgreSQL 中编写此代码。到目前为止，这就是我所拥有的：

CREATE TYPE misc_tuple AS (
    arr_store numeric[],
    ma_period integer
);

CREATE OR REPLACE FUNCTION mad_func(prev misc_tuple, curr numeric, ma_period integer)
    RETURNS misc_tuple AS $$
    BEGIN
        IF curr is null THEN
            RETURN (null::numeric[], -1);
        ELSEIF prev.arr_store is null THEN
            RETURN (ARRAY[curr]::numeric[], ma_period);
        ELSE
            -- accumulate new values in array
            prev.arr_store := array_append(prev.arr_store, curr);
            RETURN prev;
        END IF;
    END;
    $$ LANGUAGE plpgsql;

CREATE OR REPLACE FUNCTION mad_final(prev misc_tuple)
    RETURNS numeric AS $$
    DECLARE
        total_len integer;
        count numeric;
        mad_val numeric;
        mean_val numeric;
    BEGIN
        count := 0;
        mad_val := 0;
        mean_val := 0;
        total_len := array_length(prev.arr_store, 1);
        -- first loop to find the mean of the series
        FOR i IN greatest(1,total_len-prev.ma_period+1)..total_len
        LOOP 
            mean_val := mean_val + prev.arr_store[i];
            count := count + 1;
        END LOOP;
        mean_val := mean_val/NULLIF(count,0);
        -- second loop to subtract mean from each value 
        FOR i IN greatest(1,total_len-prev.ma_period+1)..total_len
        LOOP 
            mad_val := mad_val + abs(prev.arr_store[i]-mean_val);
        END LOOP;
        RETURN mad_val/NULLIF(count, 0);
    END;
    $$ LANGUAGE plpgsql;

CREATE OR REPLACE AGGREGATE mad(numeric, integer) (
    SFUNC = mad_func,
    STYPE = misc_tuple,
    FINALFUNC = mad_final
);

这是我测试性能的方式：

-- find rolling 12-period MAD
SELECT x,
       mad(x, 12) OVER (ROWS 12-1 PRECEDING)
FROM generate_series(0,1000000) as g(x);

目前，在我的桌面（i5 4670、3.4 GHz、16 GB RAM）上大约需要 45-50 秒。我仍在学习 UDF，所以我不确定我还能对我的函数做些什么来让它更快。我还有其他一些类似的 UDF - 但不使用数组的 UDF 在相同的 1m 行上花费

我可以对此 UDF 进行任何更改以使其更快吗？

Answer 1

您的示例代码不起作用，您在类型定义中有一个额外的逗号，并且您在其中一个函数中使用了未定义的变量 cnt。

您为什么将 12 既指定为聚合本身的参数，又指定为 ROWS PRECEDING 中的参数？这似乎是多余的。

您与 numpy 的比较似乎不太恰当，因为那不是滑动窗口函数。

我还有一些其他类似的 UDF - 但是不使用数组的 UDF，它们在相同的 1m 行上花费

它们是否也用作滑动窗口功能？也是用plpgsql写的？你能展示一个和它的用法吗？

pl/pgsql 通常不是一种非常有效的语言，尤其是在处理大型数组时。尽管在您的使用中，数组永远不会变得非常大，所以我认为这不会是一个特别大的问题。

提高效率的一种方法是用 C 而不是 pl/pgsql 编写代码，使用 INTERNAL 数据类型而不是 SQL 复合类型。

另一种改进这种特殊用法的方法（大量窗口，每个窗口都很小）可能是为此聚合实现MINVFUNC function and friends，这样它就不必每次都从头开始重新启动聚合行。

这是一个反函数示例，它根本不改变输出，但确实将运行时间缩短了大约一半：

CREATE OR REPLACE FUNCTION mad_invfunc(prev misc_tuple, curr numeric, ma_period integer)
    RETURNS misc_tuple AS $$
    BEGIN
            -- remove prev value
            prev.arr_store := prev.arr_store[2:];
            RETURN prev;
    END;
    $$ LANGUAGE plpgsql;
CREATE OR REPLACE AGGREGATE mad(numeric, integer) (
    SFUNC = mad_func,
    STYPE = misc_tuple,
    FINALFUNC = mad_final,
    MSFUNC = mad_func,
    MSTYPE = misc_tuple,
    MFINALFUNC = mad_final,
    MINVFUNC = mad_invfunc
);

如果我将类型从numeric 更改为double precision，它们会再次将运行时间缩短一半。因此，虽然数组上的循环可能效率不高，但仅使用 12 个成员的窗口时，它们并不是主要瓶颈。