-- The cells of per-row aggregates that group by more than the outer keys (a parent id, an element's tag) live for one row, so the scan numbers
-- them in the row and not in the group table. Every query runs with and without the pushdown; the answers must be the same: cells that
-- are one per row, several per row (a row with more cells than the row's table starts with, and one that spans chunks), HAVING on the cells,
-- grouping sets, an arm without hidden keys next to one with them, and the distinct ids of a hidden key.
CREATE EXTENSION pg_sift;
CREATE EXTENSION pg_sift_helpers;
SET max_parallel_workers_per_gather = 0;
SET seq_page_cost = 1000;
SET random_page_cost = 1000;

CREATE TABLE items (id int PRIMARY KEY, parent_id int, region text, doc jsonb) WITH (autovacuum_enabled = false);
INSERT INTO items
SELECT i, CASE WHEN i % 31 = 0 THEN NULL ELSE i % 600 END, CASE WHEN i % 17 = 0 THEN NULL ELSE 'r' || i % 9 END,
       CASE WHEN i % 13 = 0 THEN NULL
            ELSE jsonb_build_object('parts', (SELECT jsonb_agg(jsonb_build_object('k', (array['a', 'b', 'c'])[1 + (i + j) % 3], 'tag', CASE WHEN (i + j) % 7 = 0 THEN NULL ELSE 't' || (i * j) % 5 END,
                                                                                  'n', (i * j) % 10))
                                              FROM generate_series(1, 1 + i % 4) j)) END
FROM generate_series(1, 6000) i;
-- rows with forty cells in a group, and one that spans three chunks of tuples
INSERT INTO items
SELECT 6000 + i, i, 'wide', jsonb_build_object('parts', (SELECT jsonb_agg(jsonb_build_object('k', (array['a', 'b'])[1 + j % 2], 'tag', 'w' || j, 'n', j % 10))
                                                         FROM generate_series(1, CASE WHEN i = 1 THEN 9000 ELSE 40 END) j))
FROM generate_series(1, 4) i;
CREATE INDEX items_col ON items USING sift_col (id, parent_id, region, doc) WITH (paths = 'doc.parts[]{k text, tag text, n int4}');
VACUUM (FREEZE, ANALYZE) items;

-- the outer aggregates, the arms' extra group column and its HAVING, and the WHERE of the scan
CREATE FUNCTION report(aggs text, hid text DEFAULT 'parent_id', arm_having text DEFAULT '', tail text DEFAULT '', filt text DEFAULT 'true') RETURNS text LANGUAGE sql AS $$
	SELECT format($q$WITH base AS (SELECT t.id, t.parent_id, t.region, p.k, p.tag, p.n
	                                FROM items t CROSS JOIN LATERAL jsonb_to_recordset(t.doc -> 'parts') AS p(k text, tag text, n int) WHERE %4$s)
	                 SELECT k, %1$s FROM (SELECT id, %2$s, k, sum(n) AS s, count(*) AS c FROM base GROUP BY id, %2$s, k %3$s
	                                      UNION ALL
	                                      SELECT id, %2$s, NULL::text, sum(n), count(*) FROM base GROUP BY id, %2$s) q
	                 GROUP BY k %5$s$q$, aggs, hid, arm_having, filt, tail)
$$;

CREATE FUNCTION pushed(q text) RETURNS bool LANGUAGE sql AS $$ SELECT sift_helpers.plan_is_agg_scan(q) $$;

CREATE FUNCTION same(q text) RETURNS bool LANGUAGE plpgsql AS $$
DECLARE a bigint; n bigint;
BEGIN
	SET sift.enable_agg = on;
	EXECUTE format('CREATE TEMP TABLE r_on AS %s', q);
	SET sift.enable_agg = off;
	EXECUTE format('CREATE TEMP TABLE r_off AS %s', q);
	RESET sift.enable_agg;
	SELECT count(*) INTO a FROM ((TABLE r_on EXCEPT TABLE r_off) UNION ALL (TABLE r_off EXCEPT TABLE r_on)) d;
	SELECT count(*) INTO n FROM r_on;
	DROP TABLE r_on, r_off;
	RETURN a = 0 AND n > 0 AND pushed(q);
END $$;

-- a hidden key that is a column of the row: one cell per row and outer group
SELECT same(report('sum(s) AS total, sum(c) AS cells, count(*) AS n, count(DISTINCT id) AS ids, count(DISTINCT parent_id) AS parents')) AS row_key;
SELECT same(report('sum(s) AS total, count(DISTINCT region) AS regions', 'region')) AS row_text_key;
-- count(id) of the row ids counts the cells' rows as count(*) does, alone, with a filter, and with a key that splits the cells
SELECT same(report('count(id) AS counted, count(*) AS n, sum(s) AS total')) AS count_row_id;
SELECT same(report('count(id) AS counted, count(id) FILTER (WHERE s > 9) AS big, count(DISTINCT id) AS ids', 'tag')) AS count_row_id_filter;
-- a hidden key that is an element's tag: several cells in a row, NULL tags among them
SELECT same(report('sum(s) AS total, sum(c) AS cells, count(*) AS n, count(DISTINCT id) AS ids, count(DISTINCT tag) AS tags', 'tag')) AS element_key;
-- both, and the arms' HAVING and WHERE
SELECT same(report('sum(s) AS total, count(*) AS n, count(DISTINCT parent_id) AS parents, count(DISTINCT tag) AS tags', 'parent_id, tag')) AS two_keys;
SELECT same(report('sum(s) AS total, count(*) AS n, count(DISTINCT tag) AS tags', 'tag', 'HAVING sum(n) > 2')) AS arm_having;
SELECT same(report('sum(s) AS total, count(*) AS n, count(DISTINCT parent_id) AS parents', 'parent_id', '', '', 'n > 3 AND t.parent_id < 400')) AS where_of_arms;
-- the outer HAVING, and an outer filter on the cells
SELECT same(report('count(DISTINCT parent_id) AS parents, count(*) FILTER (WHERE s > 9) AS big', 'parent_id', '', 'HAVING count(DISTINCT parent_id) > 100')) AS having_outer;
SELECT same(report('count(DISTINCT parent_id) AS parents, count(*) FILTER (WHERE s > 9) AS big', 'parent_id')) AS cell_filter;
-- grouping sets over the outer key
SELECT same(replace(report('sum(s) AS total, count(*) AS n, count(DISTINCT parent_id) AS parents, count(DISTINCT tag) AS tags', 'parent_id, tag'), 'GROUP BY k', 'GROUP BY ROLLUP (k)')) AS rollup;
-- an arm that has no hidden key next to one that has: its cells are the outer groups
SELECT same($q$WITH base AS (SELECT t.id, t.parent_id, p.k, p.n FROM items t CROSS JOIN LATERAL jsonb_to_recordset(t.doc -> 'parts') AS p(k text, n int))
                SELECT k, sum(s) AS total, count(*) AS n, count(DISTINCT id) AS ids, count(DISTINCT parent_id) AS parents
                FROM (SELECT id, parent_id, k, sum(n) AS s FROM base GROUP BY id, parent_id, k
                      UNION ALL
                      SELECT id, NULL::int, k, sum(n) FROM base GROUP BY id, k) q GROUP BY k$q$) AS mixed_arms;
-- more distinct ids than hash_mem holds
SET work_mem = '64kB';
SET hash_mem_multiplier = 1;
SELECT same(report('sum(s) AS total, count(DISTINCT parent_id) AS parents, count(DISTINCT id) AS ids, count(DISTINCT tag) AS tags', 'parent_id, tag')) AS small_memory;
RESET work_mem;
RESET hash_mem_multiplier;
-- the plan shows the cells' keys
CREATE FUNCTION cell_lines(q text) RETURNS SETOF text LANGUAGE plpgsql AS $$
DECLARE l text;
BEGIN
	FOR l IN EXECUTE 'EXPLAIN (COSTS OFF) ' || q LOOP
		IF l ~ 'Cell Keys' THEN RETURN NEXT trim(l); END IF;
	END LOOP;
END $$;
SELECT * FROM cell_lines(report('sum(s) AS total, count(DISTINCT parent_id) AS parents'));

DROP FUNCTION cell_lines(text), same(text), pushed(text), report(text, text, text, text, text);
-- the shared database of a plain installcheck goes on with the next test
SET client_min_messages = warning;
DROP SCHEMA public CASCADE;
CREATE SCHEMA public;
GRANT USAGE ON SCHEMA public TO PUBLIC;
RESET client_min_messages;
DROP EXTENSION pg_sift_helpers;
DROP EXTENSION pg_sift;