-- count(DISTINCT), sum(DISTINCT) and avg(DISTINCT) of numeric values (a numeric or bigint column or an expression of one) with 0, 2 and 4 workers: a
-- worker sends the pairs (group, value) of its set with the numbers themselves, since the ids of its table of them are its own, and the leader
-- gives them its own ids and counts or adds each number once per group. The rows are those of the normal plan, with NULLs, NaN, duplicates,
-- a FILTER, a small work_mem (the sets go through a sort), the leader out of the scan, and a global aggregate. mode: parallel (workers planned
-- in the SiftAgg node), serial (SiftAgg alone), none (the normal plan).
CREATE EXTENSION pg_sift;
CREATE EXTENSION pg_sift_helpers;
SET client_min_messages = warning;
SET sift.agg_preagg = off;

CREATE TABLE facts (id int PRIMARY KEY, grp int, kind text, amount numeric(10,2), price numeric, qty int, big bigint, wide bigint, f float8, ts timestamptz, u uuid) WITH (autovacuum_enabled = false);
INSERT INTO facts
SELECT i, i % 7, (ARRAY['cat_a', 'cat_b', 'cat_c', 'cat_d'])[1 + i % 4],
       CASE WHEN i % 29 = 0 THEN NULL WHEN i % 997 = 0 THEN 'NaN'::numeric ELSE ((hashint4(i) & 2147483647) % 3000)::numeric / 10 END,
       CASE WHEN i % 31 = 0 THEN NULL ELSE ((hashint4(i + 1) & 2147483647) % 500)::numeric / 8 END,
       CASE WHEN i % 13 = 0 THEN NULL ELSE (hashint4(i + 2) & 2147483647) % 400 END,
       (hashint4(i + 3) & 2147483647) % 50000 * 100003::bigint,
       CASE WHEN i % 11 = 0 THEN NULL WHEN i % 503 = 0 THEN 9223372036854775807 WHEN i % 509 = 0 THEN -9223372036854775807 - 1
            ELSE (hashint4(i + 4) & 2147483647) % 30000 - 15000 END,
       CASE WHEN i % 23 = 0 THEN NULL WHEN i % 101 = 0 THEN 'NaN' WHEN i % 103 = 0 THEN 'Infinity' WHEN i % 107 = 0 THEN '-0'::float8 WHEN i % 109 = 0 THEN 0 ELSE ((hashint4(i + 5) & 2147483647) % 2000) / 8.0 END,
       CASE WHEN i % 37 = 0 THEN NULL WHEN i % 211 = 0 THEN 'infinity' WHEN i % 223 = 0 THEN '-infinity' ELSE timestamptz '2024-01-01' + ((hashint4(i + 6) & 2147483647) % 20000) * interval '1 minute' END,
       CASE WHEN i % 41 = 0 THEN NULL ELSE md5(((hashint4(i + 7) & 2147483647) % 3000)::text)::uuid END
FROM generate_series(1, 20000) i;
CREATE INDEX facts_col ON facts USING sift_col (id, grp, kind, amount, price, qty, big, wide, f, ts, u);
VACUUM ANALYZE facts;

SET parallel_setup_cost = 0;
SET parallel_tuple_cost = 0;
SET min_parallel_table_scan_size = 0;
SET min_parallel_index_scan_size = 0;
SET seq_page_cost = 10000;

CREATE TABLE shapes (n serial, name text, q text);
INSERT INTO shapes (name, q) VALUES
	('count amount',      'SELECT grp, count(DISTINCT amount) FROM facts GROUP BY grp'),
	('count global',      'SELECT count(DISTINCT amount), count(DISTINCT price) FROM facts'),
	('count expression',  'SELECT grp, count(DISTINCT amount * 3), count(DISTINCT price + 1) FROM facts GROUP BY grp'),
	('count with int',    'SELECT kind, grp, count(DISTINCT amount), count(DISTINCT qty), count(*) FROM facts GROUP BY kind, grp'),
	('sum amount',        'SELECT grp, sum(DISTINCT amount) FROM facts GROUP BY grp'),
	('avg amount',        'SELECT kind, avg(DISTINCT amount), avg(DISTINCT price) FROM facts GROUP BY kind'),
	('sum global',        'SELECT sum(DISTINCT amount), avg(DISTINCT amount), count(DISTINCT amount) FROM facts'),
	('sum expression',    'SELECT grp, sum(DISTINCT amount + price) FROM facts GROUP BY grp'),
	('filter',            'SELECT grp, count(DISTINCT amount) FILTER (WHERE qty > 100), sum(DISTINCT amount) FILTER (WHERE kind = ''cat_a'') FROM facts GROUP BY grp'),
	('where',             'SELECT grp, count(DISTINCT amount), sum(DISTINCT price) FROM facts WHERE qty < 200 AND kind <> ''cat_b'' GROUP BY grp'),
	('nothing',           'SELECT grp, count(DISTINCT amount), sum(DISTINCT amount) FROM facts WHERE id < 0 GROUP BY grp'),
	('all null group',    'SELECT qty IS NULL AS q, count(DISTINCT amount), sum(DISTINCT amount) FROM facts GROUP BY 1'),
	('count bigint',      'SELECT grp, count(DISTINCT big) FROM facts GROUP BY grp'),
	('bigint global',     'SELECT count(DISTINCT wide), count(DISTINCT big), count(DISTINCT id::bigint) FROM facts'),
	('bigint extremes',   'SELECT kind, count(DISTINCT wide), count(wide) FROM facts GROUP BY kind'),
	('bigint filter',     'SELECT kind, count(DISTINCT wide) FILTER (WHERE qty > 100), count(DISTINCT big) FROM facts GROUP BY kind'),
	('bigint expression', 'SELECT grp, count(DISTINCT big + 1), count(DISTINCT big * 2) FROM facts GROUP BY grp'),
	('count float',       'SELECT grp, count(DISTINCT f) FROM facts GROUP BY grp'),
	('count timestamptz', 'SELECT grp, count(DISTINCT ts) FROM facts GROUP BY grp'),
	('leaf global',       'SELECT count(DISTINCT f), count(DISTINCT ts), count(DISTINCT u) FROM facts'),
	('leaf filter',       'SELECT kind, count(DISTINCT f) FILTER (WHERE qty > 100), count(DISTINCT ts) FILTER (WHERE grp < 3) FROM facts GROUP BY kind'),
	('leaf mixed',        'SELECT kind, count(DISTINCT f), sum(DISTINCT amount), count(DISTINCT qty), count(DISTINCT ts) FROM facts GROUP BY kind');

CREATE FUNCTION rows_of(q text, push bool, workers int, leader bool DEFAULT true) RETURNS text LANGUAGE sql AS $$ SELECT sift_helpers.rows_agg_with_workers_and_leader(q, push, workers, leader) $$;

CREATE FUNCTION mode_of(q text, workers int) RETURNS text LANGUAGE sql AS $$ SELECT sift_helpers.plan_mode_workers(q, workers) $$;

-- the mode with 4 workers, then whether 0, 2 and 4 workers (and the leader out of the scan) give the rows of the normal plan
SELECT s.name, mode_of(s.q, 4) AS mode,
       rows_of(s.q, true, 0) = rows_of(s.q, false, 0) AS scan_w0,
       rows_of(s.q, true, 2) = rows_of(s.q, false, 0) AS scan_w2,
       rows_of(s.q, true, 4) = rows_of(s.q, false, 0) AS scan_w4,
       rows_of(s.q, true, 3, false) = rows_of(s.q, false, 0) AS scan_w3_no_leader
FROM shapes s ORDER BY s.n;

-- the sets of a worker go through a sort when work_mem is small: the same rows
SET work_mem = '64kB';
SET hash_mem_multiplier = 1;
SELECT s.name, rows_of(s.q, true, 2) = rows_of(s.q, false, 0) AS small_work_mem_w2, rows_of(s.q, true, 4) = rows_of(s.q, false, 0) AS small_work_mem_w4
FROM shapes s WHERE s.name IN ('count amount', 'count global', 'sum amount', 'sum expression') ORDER BY s.n;
RESET work_mem;
RESET hash_mem_multiplier;

-- a scan that reads rows in the write buffer and rows of updated tuples
INSERT INTO facts SELECT 20000 + i, i % 7, 'cat_a', (i % 400)::numeric / 4, i % 9, i % 300, i, i - 350, (i % 90) / 4.0, timestamptz '2024-01-01' + i * interval '7 minutes', md5((i % 500)::text)::uuid FROM generate_series(1, 700) i;
UPDATE facts SET amount = amount + 0.5 WHERE id % 17 = 0;
DELETE FROM facts WHERE id % 61 = 0;
SELECT s.name, rows_of(s.q, true, 0) = rows_of(s.q, false, 0) AS dirty_w0, rows_of(s.q, true, 2) = rows_of(s.q, false, 0) AS dirty_w2,
       rows_of(s.q, true, 4) = rows_of(s.q, false, 0) AS dirty_w4
FROM shapes s ORDER BY s.n;

-- what stays on the leader says why in EXPLAIN VERBOSE: the aggregate whose partial state the leader cannot merge
CREATE FUNCTION serial_because(q text) RETURNS text LANGUAGE plpgsql AS $$
DECLARE l text; r text := '';
BEGIN
	PERFORM set_config('sift.enable_agg', 'on', true);
	PERFORM set_config('max_parallel_workers_per_gather', '4', true);
	FOR l IN EXECUTE 'EXPLAIN (COSTS OFF, VERBOSE) ' || q LOOP
		IF l LIKE '%Serial Because%' THEN r := btrim(l); END IF;
	END LOOP;
	RETURN r;
END $$;
SELECT mode_of('SELECT grp, string_agg(DISTINCT kind, '','') FROM facts GROUP BY grp', 4) AS string_agg_mode,
       serial_because('SELECT grp, string_agg(DISTINCT kind, '','') FROM facts GROUP BY grp') AS string_agg_why,
       serial_because('SELECT grp, count(DISTINCT amount) FROM facts GROUP BY grp') AS count_distinct_why;

DROP TABLE facts, shapes;
DROP EXTENSION pg_sift_helpers;
DROP EXTENSION pg_sift;
DROP FUNCTION mode_of, rows_of, serial_because;