BUG #18330: The query planner chooses the wrong plan when using the parallel aggregation function

The following bug has been logged on the website:

Bug reference:      18330
Logged by:          Michael Efremov
Email address:      efremov20081@gmail.com
PostgreSQL version: 15.1
Operating system:   Alpine 12.2.1_git20220924-r4
Description:

Below is an example of queries. It shows that the query planner chooses
sequential aggregation instead of parallel, although limit should not have
affected this.

-- FIRSTLY - create parallel agg function

CREATE TYPE top1_units_weights_state_v1_parallel AS (
  test jsonb
);

CREATE OR REPLACE FUNCTION agg_top1_units_weights_transition_v1_parallel(
  state top1_units_weights_state_v1_parallel,
  test jsonb
) RETURNS top1_units_weights_state_v1_parallel AS $$
begin
  RETURN state;
END;
$$ LANGUAGE plpgsql;


CREATE OR REPLACE FUNCTION agg_top1_units_weights_final_v1_parallel(
  state top1_units_weights_state_v1_parallel
) RETURNS top1_units_weights_state_v1_parallel AS $$
begin
  RAISE NOTICE 'agg_top1_units_weights_final_v1_parallel:% ', state;
  RETURN state;
END;
$$ LANGUAGE plpgsql;

CREATE OR REPLACE FUNCTION agg_top1_units_weights_combinefunc_v1_parallel(
  first top1_units_weights_state_v1_parallel,
  second top1_units_weights_state_v1_parallel
) RETURNS top1_units_weights_state_v1_parallel AS $$
begin
RAISE NOTICE 'agg_top1_units_weights_combinefunc_v1_parallel:% ', first;
  RETURN first;
END;
$$ LANGUAGE plpgsql PARALLEL SAFE;

CREATE OR REPLACE AGGREGATE agg_top1_units_weights_v1_parallel(jsonb) (
  sfunc       = agg_top1_units_weights_transition_v1_parallel,
  stype       = top1_units_weights_state_v1_parallel,
  finalfunc   = agg_top1_units_weights_final_v1_parallel,
  initcond    = '({})',
  COMBINEFUNC = agg_top1_units_weights_combinefunc_v1_parallel, 
  parallel    = SAFE  
);

-- SECONDLY - create test data


CREATE TABLE public.test_jsonb_agg(
    test_jsonb jsonb not null
);

CREATE OR REPLACE FUNCTION random_between(low INT ,high INT)
   RETURNS INT AS
$$
BEGIN
   RETURN floor(random()* (high-low + 1) + low);
END;

CREATE OR REPLACE FUNCTION random_int_array(dim integer, min integer, max
integer) RETURNS integer[] AS $BODY$
begin
    return (select array_agg(random_between(min,max)) from generate_series (0,
dim));
end
$BODY$ LANGUAGE plpgsql;

CREATE OR REPLACE FUNCTION random_jsonb_array_of_array(dim integer, min
integer, max integer) RETURNS jsonb AS $BODY$
declare
 res jsonb = '[]'::jsonb;
 cnt integer = random_between(1,4);
begin
loop
     res = jsonb_build_array(random_int_array(dim,min,max)) || res;
    IF cnt > 3 THEN
        return res;
    END IF;
    cnt = cnt + 1;
END LOOP;
end
$BODY$ LANGUAGE plpgsql;

INSERT into test_jsonb_agg
select
    jsonb_array as jsonb_array
from (
select
    random_jsonb_array_of_array(random_between(1,5), 1, 500) as
jsonb_array
from generate_series(0,100000)
) as data_t

-- THIRDLY - check two times
analyze test_jsonb;

explain (analyze,verbose,timing,costs,buffers)
with test_speed as not MATERIALIZED(
select test_jsonb as col
from test_jsonb_agg  -- table size is 100_000
limit 900000 -- try to comment this line then parallel aggregation plan will
be used. 
) select agg_top1_units_weights_v1_parallel(col)
from test_speed

-> output
Aggregate  (cost=29174.52..29174.53 rows=1 width=32) (actual
time=50.872..50.872 rows=1 loops=1)
  Output: agg_top1_units_weights_v1_parallel(test_jsonb_agg.test_jsonb)
  Buffers: shared hit=2174
  ->  Limit  (cost=0.00..3174.01 rows=100001 width=145) (actual
time=0.007..10.412 rows=100001 loops=1)
        Output: test_jsonb_agg.test_jsonb
        Buffers: shared hit=2174
        ->  Seq Scan on public.test_jsonb_agg  (cost=0.00..3174.01
rows=100001 width=145) (actual time=0.007..4.928 rows=100001 loops=1)
              Output: test_jsonb_agg.test_jsonb
              Buffers: shared hit=2174
Query Identifier: 8741670630168910811
Planning Time: 0.051 ms
Execution Time: 50.920 ms

explain (analyze,verbose,timing,costs,buffers)
with test_speed as not MATERIALIZED(
select test_jsonb as col
from test_jsonb_agg  -- table size is 100_000
--limit 900000 -- try to comment this line then parallel aggregation plan
will be used. 
) select agg_top1_units_weights_v1_parallel(col)
from test_speed

-> output
Finalize Aggregate  (cost=14008.38..14008.39 rows=1 width=32) (actual
time=28.777..32.664 rows=1 loops=1)
  Output: agg_top1_units_weights_v1_parallel(test_jsonb_agg.test_jsonb)
  Buffers: shared hit=2364
  ->  Gather  (cost=14007.42..14007.63 rows=2 width=32) (actual
time=28.619..32.508 rows=3 loops=1)
        Output: (PARTIAL
agg_top1_units_weights_v1_parallel(test_jsonb_agg.test_jsonb))
        Workers Planned: 2
        Workers Launched: 2
        Buffers: shared hit=2364
        ->  Partial Aggregate  (cost=13007.42..13007.43 rows=1 width=32)
(actual time=15.567..15.568 rows=1 loops=3)
              Output: PARTIAL
agg_top1_units_weights_v1_parallel(test_jsonb_agg.test_jsonb)
              Buffers: shared hit=2364
              Worker 0:  actual time=9.197..9.198 rows=1 loops=1
                Buffers: shared hit=501
              Worker 1:  actual time=9.201..9.202 rows=1 loops=1
                Buffers: shared hit=520
              ->  Parallel Seq Scan on public.test_jsonb_agg
(cost=0.00..2590.67 rows=41667 width=145) (actual time=0.008..1.872
rows=33334 loops=3)
                    Output: test_jsonb_agg.test_jsonb
                    Buffers: shared hit=2174
                    Worker 0:  actual time=0.009..1.195 rows=18612 loops=1
                      Buffers: shared hit=406
                    Worker 1:  actual time=0.009..1.205 rows=19462 loops=1
                      Buffers: shared hit=425
Query Identifier: 4368818925053284440
Planning Time: 0.054 ms
Execution Time: 32.718 ms