Postgres custom aggregation returns null when parallelized - postgresql

Question

I have created a custom aggregate in postgres 11.3 and it works when parallel if off. When I mark it as parallel = safe, it returns null.
Could someone point me in the direction of where to start looking or how do I debug a parallel aggregation in postgres? In non parallel aggregation I can insert the state at each record into a temporary table, but inserts are not allowed in parallel queries...
Here's the aggregate:
CREATE OR REPLACE FUNCTION array_sort(ANYARRAY)
RETURNS ANYARRAY LANGUAGE SQL
AS $$
SELECT ARRAY(SELECT unnest($1) ORDER BY 1)
$$;
create type _stats_agg_accum_type AS (
cnt bigint,
q double precision[],
n double precision[],
np double precision[],
dn double precision[]
);
create type _stats_agg_result_type AS (
count bigint,
q25 double precision,
q50 double precision,
q75 double precision
);
create or replace function _stats_agg_p2_parabolic(_stats_agg_accum_type, double precision, double precision)
returns double precision AS '
DECLARE
a alias for $1;
i alias for $2;
d alias for $3;
BEGIN
RETURN a.q[i] + d / (a.n[i + 1] - a.n[i - 1]) * ((a.n[i] - a.n[i - 1] + d) * (a.q[i + 1] - a.q[i]) / (a.n[i + 1] - a.n[i]) + (a.n[i + 1] - a.n[i] - d) * (a.q[i] - a.q[i - 1]) / (a.n[i] - a.n[i - 1]));
END;
'
language plpgsql;
create or replace function _stats_agg_p2_linear(_stats_agg_accum_type, double precision, double precision)
returns double precision AS '
DECLARE
a alias for $1;
i alias for $2;
d alias for $3;
BEGIN
return a.q[i] + d * (a.q[i + d] - a.q[i]) / (a.n[i + d] - a.n[i]);
END;
'
language plpgsql;
create or replace function _stats_agg_accumulator(_stats_agg_accum_type, double precision)
returns _stats_agg_accum_type AS '
DECLARE
a ALIAS FOR $1;
x alias for $2;
k int;
d double precision;
qp double precision;
BEGIN
a.cnt = a.cnt + 1;
if a.cnt <= 5 then
a.q = array_append(a.q, x);
if a.cnt = 5 then
a.q = array_sort(a.q);
end if;
return a;
end if;
case
when x < a.q[1] then
a.q[1] = x;
k = 1;
when x >= a.q[1] and x < a.q[2] then
k = 1;
when x >= a.q[2] and x < a.q[3] then
k = 2;
when x >= a.q[3] and x < a.q[4] then
k = 3;
when x >= a.q[4] and x <= a.q[5] then
k = 4;
when x > a.q[5] then
a.q[5] = x;
k = 4;
end case;
for ii in 1..5 loop
if ii > k then
a.n[ii] = a.n[ii] + 1;
end if;
a.np[ii] = a.np[ii] + a.dn[ii];
end loop;
for ii in 2..4 loop
d = a.np[ii] - a.n[ii];
if (d >= 1 and a.n[ii+1] - a.n[ii] > 1) or (d <= -1 and a.n[ii-1] - a.n[ii] < -1) then
d = sign(d);
qp = _stats_agg_p2_parabolic(a, ii, d);
if qp > a.q[ii-1] and qp < a.q[ii+1] then
a.q[ii] = qp;
else
a.q[ii] = _stats_agg_p2_linear(a, ii, d);
end if;
a.n[ii] = a.n[ii] + d;
end if;
end loop;
return a;
END;
'
language plpgsql;
create or replace function _stats_agg_combiner(_stats_agg_accum_type, _stats_agg_accum_type)
returns _stats_agg_accum_type AS '
DECLARE
a alias for $1;
b alias for $2;
c _stats_agg_accum_type;
BEGIN
c.cnt = a.cnt + b.cnt;
c.q[2] = (a.q[2] + b.q[2]) / 2;
c.q[3] = (a.q[3] + b.q[3]) / 2;
c.q[4] = (a.q[4] + b.q[4]) / 2;
RETURN c;
END;
'
strict language plpgsql;
create or replace function _stats_agg_finalizer(_stats_agg_accum_type)
returns _stats_agg_result_type AS '
BEGIN
RETURN row(
$1.cnt,
$1.q[2],
$1.q[3],
$1.q[4]
);
END;
'
language plpgsql;
create aggregate stats_agg(double precision) (
sfunc = _stats_agg_accumulator,
stype = _stats_agg_accum_type,
finalfunc = _stats_agg_finalizer,
combinefunc = _stats_agg_combiner,
--parallel = safe,
initcond = '(0, {}, "{1,2,3,4,5}", "{1,2,3,4,5}", "{0,0.25,0.5,0.75,1}")'
);
Here's the setup and run code:
--CREATE TABLE temp (val double precision);
--insert into temp (val) select i from generate_series(0, 150000) as t(i);
select (stats_agg(val)).* from temp;
The expected result as follows and it works when run in parallel = unsafe
150001, 37500, 75000, 112500
In parallel = safe I get nulls:
150001, null, null, null

Answer 1

The problem is in the _stats_agg_combiner function. The function definition includes the strict keyword so there is no need to check for null input values.
In this specific aggregate, the _stats_agg_accum_type includes multiple arrays and the _stats_agg_combiner function requires that these arrays be filled with a minimum of 5 entries. This assumes that each new _stats_agg_accum_type instance processes at a minimum 5 records before being passed to the _stats_agg_combiner function.
Tests were being done on a table with 150k records and an assumption that each instance would therefore receive at a minimum 5 records. For whatever reason, this is an incorrect assumption. Regardless of the number of workers used (tested with 1-4) there is always at least one instance which processed exactly 0 records.
The solution was to add support for a _stats_agg_accum_type instance that had processed zero records and this had an array length of 0. See code below.
create or replace function _stats_agg_combiner(_stats_agg_accum_type, _stats_agg_accum_type)
returns _stats_agg_accum_type AS '
DECLARE
a alias for $1;
b alias for $2;
c _stats_agg_accum_type;
addA boolean;
addB boolean;
BEGIN
addA = a.cnt <= 5;
addB = b.cnt <= 5;
if addA and not addB then
c = b;
elsif addB and not addA then
c = a;
else
c.cnt = a.cnt + b.cnt;
for ii in 2..4 loop
c.q[ii] = (a.q[ii] + b.q[ii]) / 2;
end loop;
end if;
for ii in 1..5 loop
if addA and ii <= a.cnt then
c = _stats_agg_accumulator(c, a.q[ii]);
end if;
if addB and ii <= b.cnt then
c = _stats_agg_accumulator(c, b.q[ii]);
end if;
end loop;
RETURN c;
END;
'
language plpgsql strict;