-- The widths of the shapes the aggregate scan holds (UNION ALL arms, virtual leaves, lookup-joined tables, key and read columns of a join, value
-- combinations, float8 operations, leaves of a generic aggregate argument, aggregates sharing a distinct set): a shape one past a width keeps the normal
-- plan with the slug shape_width (rows_shape for the aggregates of a rows set) and the limit in the reason; one at the width runs in the scan and gives the rows of the normal plan. Each width is
-- tried at its default (where the shape is cheap to make) and at a value sift.test_limits lowers it to.
CREATE EXTENSION pg_sift;
-- the tuples through the arms are in the price of the scan: on these few rows the normal plan is the cheaper one, and the test asks about the scan
SET sift.agg_force = on;
LOAD 'pg_sift';
SET max_parallel_workers_per_gather = 0;
SET seq_page_cost = 10000;
SET random_page_cost = 10000;
CREATE TABLE events (id int PRIMARY KEY, parent_id int, kind text, ts date, a int, b int, c int, d int, doc jsonb) WITH (autovacuum_enabled = false);
INSERT INTO events
SELECT i, i % 50, 'k' || i % 5, date '2025-01-01' + i % 6, i % 7, i % 5, i % 3, i % 2,
       jsonb_build_object('metrics', (SELECT jsonb_object_agg('m' || j, (i * j) % 10) FROM generate_series(1, 20) j),
                          'parts', (SELECT jsonb_agg(jsonb_build_object('cat', 'cat_' || (i + j) % 6, 'score', (i * j) % 10)) FROM generate_series(1, 1 + i % 4) j))
FROM generate_series(1, 3000) i;
CREATE INDEX events_col ON events USING sift_col (id, parent_id, kind, ts, a, b, c, d, doc) WITH (paths = 'doc.metrics.* float8, doc.parts[]{cat text, score float8}');
NOTICE:  sift_col index "events_col": 1 automatic choice (track_changes)
DETAIL:  track_changes = id
HINT:  Reasons, costs and how to turn each off: SELECT * FROM sift.index_choices('events_col');
-- a dimension with four key columns and seventeen value columns
CREATE TABLE dd (id int NOT NULL, a int NOT NULL, b int NOT NULL, c int NOT NULL, d int NOT NULL, v1 int NOT NULL, v2 int NOT NULL, v3 int NOT NULL, v4 int NOT NULL, v5 int NOT NULL, v6 int NOT NULL, v7 int NOT NULL, v8 int NOT NULL, v9 int NOT NULL, v10 int NOT NULL, v11 int NOT NULL, v12 int NOT NULL, v13 int NOT NULL, v14 int NOT NULL, v15 int NOT NULL, v16 int NOT NULL, v17 int NOT NULL) WITH (autovacuum_enabled = false);
INSERT INTO dd SELECT i, i % 7, i % 5, i % 3, i % 2, i, i + 1, i + 2, i + 3, i + 4, i + 5, i + 6, i + 7, i + 8, i + 9, i + 10, i + 11, i + 12, i + 13, i + 14, i + 15, i + 16 FROM generate_series(0, 49) i;
-- the unique index on the key and the first ORDER BY column leaves no tie to a first-row lookup
CREATE UNIQUE INDEX dd_first ON dd (a, b, v2);
-- seventeen integer columns: 136 pairs of them, each a value combination of its own
CREATE TABLE wide (kind text, c1 int, c2 int, c3 int, c4 int, c5 int, c6 int, c7 int, c8 int, c9 int, c10 int, c11 int, c12 int, c13 int, c14 int, c15 int, c16 int, c17 int) WITH (autovacuum_enabled = false);
INSERT INTO wide SELECT 'k' || i % 5, i % 2, i % 3, i % 5, i % 7, i % 11, i % 13, i % 17, i % 19, i % 23, i % 29, i % 31, i % 37, i % 41, i % 43, i % 47, i % 53, i % 59 FROM generate_series(1, 3000) i;
CREATE INDEX wide_col ON wide USING sift_col (kind, c1, c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14, c15, c16, c17);
NOTICE:  sift_col index "wide_col": 1 automatic choice, 1 advice, 1 limit
DETAIL:  advice: track_changes = the table has no one-column primary key; limit: cube = "c16", "c17";
rollup_columns = kind, c3
HINT:  Reasons, costs and how to turn each off: SELECT * FROM sift.index_choices('wide_col');
VACUUM (FREEZE, ANALYZE) events, dd, wide;
SET join_collapse_limit = 64;
SET from_collapse_limit = 64;
-- the rows of the scan and of the normal plan, the arm's answer and the reason
CREATE FUNCTION same(q text) RETURNS bool LANGUAGE plpgsql AS $$
DECLARE a bigint; b 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_off;
	DROP TABLE r_on, r_off;
	RETURN a = 0 AND n > 0;
END $$;
CREATE FUNCTION probe(q text, want_arm text) RETURNS TABLE (used bool, slug text, detail text) LANGUAGE sql AS $$
	SELECT w.used, w.slug, substring(w.reason from 'the limit [a-z_]+ is [0-9]+') FROM sift.explain_why(q) w WHERE w.arm = want_arm
$$;
-- the tables the scan read as first-row lookups (a sub-select with more ORDER BY columns than the width runs once per key instead)
CREATE FUNCTION first_row_lookups(q text) RETURNS int LANGUAGE plpgsql AS $$
DECLARE l text; n int := 0;
BEGIN
	FOR l IN EXECUTE 'EXPLAIN (ANALYZE, COSTS OFF, TIMING OFF, SUMMARY OFF) ' || q LOOP
		IF l ~ '(\d+) first row tables' THEN n := substring(l from '(\d+) first row tables')::int; END IF;
	END LOOP;
	RETURN n;
END $$;
-- the shapes: n of something
CREATE FUNCTION shape(kind text, n int) RETURNS text LANGUAGE sql AS $$
	SELECT CASE kind
	WHEN 'virtual_leaves' THEN
		(SELECT format('SELECT kind, %s FROM events GROUP BY kind', string_agg(format('sum((doc->''metrics''->>''m%s'')::float8) AS s%s', j, j, j), ', ')) FROM generate_series(1, n) j)
	WHEN 'arms' THEN
		(SELECT format('SELECT category, sp.k, avg(score) AS value FROM generate_series(1, 20) sp(k) JOIN (%s) m ON m.parent_id = sp.k GROUP BY 1, 2',
				string_agg(format($q$SELECT e.id, e.parent_id, o.cat AS category, avg(o.score) FILTER (WHERE o.score >= %s) AS score
				           FROM events e CROSS JOIN LATERAL jsonb_to_recordset(e.doc['parts']) o(cat text, score float8) GROUP BY e.id, e.parent_id, o.cat$q$, j), ' UNION ALL '))
		 FROM generate_series(1, n) j)
	WHEN 'joins' THEN
		(SELECT format('SELECT %s, count(*) FROM events e %s GROUP BY 1', 'd1.v1 % 3', string_agg(format('JOIN dd d%s ON d%s.id = e.parent_id', j, j), ' ') ) FROM generate_series(1, n) j)
	WHEN 'combos' THEN
		(SELECT format('SELECT kind, %s FROM wide GROUP BY kind', string_agg(format('sum((c%s * c%s)::numeric / 3) AS s%s_%s', i, j, i, j), ', ' ORDER BY i, j))
		 FROM (SELECT i, j FROM generate_series(1, 17) i, generate_series(1, 17) j WHERE i < j ORDER BY i, j LIMIT n) p)
	WHEN 'first_order_cols' THEN
		(SELECT format('SELECT d.v1 %% 3, count(*) FROM events e LEFT JOIN LATERAL (SELECT x.v1 FROM dd x WHERE x.a = e.a AND x.b = e.b ORDER BY %s LIMIT 1) d ON true GROUP BY 1',
				string_agg('x.' || k, ', ')) FROM unnest((ARRAY['v2', 'c', 'd', 'v3', 'v4', 'v5', 'v6', 'v7', 'v8', 'v9'])[1:n]) k)
	WHEN 'join_keys' THEN
		(SELECT format('SELECT d.v1 %% 3, count(*) FROM events e JOIN dd d ON %s GROUP BY 1', string_agg(format('d.%s = e.%s', k, k), ' AND ')) FROM unnest((ARRAY['a', 'b', 'c', 'd'])[1:n]) k)
	WHEN 'join_cols' THEN
		(SELECT format('SELECT e.kind, sum(%s) FROM events e JOIN dd d ON d.id = e.parent_id GROUP BY 1', string_agg(format('d.v%s', j), ' + ')) FROM generate_series(1, n) j)
	WHEN 'float_ops' THEN
		(SELECT format('SELECT kind, sum(%s) FROM events GROUP BY kind', 'a::float8' || string_agg(' + 1.5', '')) FROM generate_series(1, n - 1) j)
	WHEN 'rows_shared' THEN
		(SELECT format($q$SELECT sp.k, %s FROM generate_series(1, 20) sp(k)
		                  LEFT JOIN (SELECT e.id, e.parent_id, e.a, o.cat, avg(o.score) AS s
		                             FROM events e CROSS JOIN LATERAL jsonb_to_recordset(e.doc['parts']) o(cat text, score float8) GROUP BY e.id, e.parent_id, e.a, o.cat) m
		                  ON m.parent_id = sp.k GROUP BY sp.k$q$,
				string_agg(format('count(DISTINCT m.id) FILTER (WHERE m.a %% %s = 0) AS c%s', j + 1, j), ', ')) FROM generate_series(1, n) j)
	END
$$;
-- width, arm of explain_why, its default, and the value sift.test_limits lowers it to (NULL: the limit is not read through the registry yet)
CREATE TABLE widths (kind text, limit_id text, arm text, dflt int, lowered int);
INSERT INTO widths VALUES
	('virtual_leaves', 'virtual_leaves', 'single', 16, 3),
	('arms', 'arms', 'two_level', 16, 3),
	('joins', 'joins', 'single', 32, 3),
	('combos', 'combos', 'single', 128, 3),
	('join_keys', 'join_keys', 'single', 3, NULL),
	('join_cols', 'join_cols', 'single', 16, 3),
	('first_order_cols', 'first_order_cols', 'single', 8, 3),
	('float_ops', 'float_ops', 'single', 16, NULL),
	('rows_shared', 'rows_shared', 'rows', 32, 3);
-- first_order_cols is a width too: past it the sub-select runs once per key and the scan reads no first-row lookup (the detail counts them).
-- float_ops and join_keys are widths too, but a shape past them does not keep the normal plan: PostgreSQL evaluates the longer float8 expression for the distinct values of its
-- leaves, and the equality columns of a join beyond the third key by the leaves they read; the rows are those of the normal plan at both sizes
-- every width at N and at N + 1: N runs in the scan (the rows of the normal plan), N + 1 keeps the normal plan and names the width
CREATE FUNCTION at_width(lower_it bool) RETURNS TABLE (kind text, n int, used bool, slug text, detail text, rows_ok bool) LANGUAGE plpgsql AS $$
DECLARE w record; q text; p record; i int;
BEGIN
	FOR w IN SELECT * FROM widths WHERE NOT lower_it OR widths.lowered IS NOT NULL ORDER BY kind LOOP
		IF lower_it THEN
			EXECUTE format('SET sift.test_limits = %L', w.limit_id || '=' || w.lowered);
		END IF;
		FOREACH i IN ARRAY ARRAY[CASE WHEN lower_it THEN w.lowered ELSE w.dflt END, CASE WHEN lower_it THEN w.lowered ELSE w.dflt END + 1] LOOP
			q := shape(w.kind, i);
			SELECT * INTO p FROM probe(q, w.arm);
			kind := w.kind; n := i; used := p.used; slug := p.slug; detail := p.detail;
			IF w.kind = 'first_order_cols' THEN detail := format('first-row lookups: %s', first_row_lookups(q)); END IF;
			rows_ok := CASE WHEN p.used THEN same(q) END;
			RETURN NEXT;
		END LOOP;
		RESET sift.test_limits;
	END LOOP;
END $$;
SELECT * FROM at_width(false);
       kind       |  n  | used |    slug     |             detail             | rows_ok 
------------------+-----+------+-------------+--------------------------------+---------
 arms             |  16 | t    |             |                                | t
 arms             |  17 | f    | shape_width | the limit arms is 16           | 
 combos           | 128 | t    |             |                                | t
 combos           | 129 | f    | shape_width | the limit combos is 128        | 
 first_order_cols |   8 | t    |             | first-row lookups: 1           | t
 first_order_cols |   9 | t    |             | first-row lookups: 0           | t
 float_ops        |  16 | t    |             |                                | t
 float_ops        |  17 | t    |             |                                | t
 join_cols        |  16 | t    |             |                                | t
 join_cols        |  17 | f    | shape_width | the limit join_cols is 16      | 
 join_keys        |   3 | t    |             |                                | t
 join_keys        |   4 | t    |             |                                | t
 joins            |  32 | t    |             |                                | t
 joins            |  33 | f    | shape_width | the limit joins is 32          | 
 rows_shared      |  32 | t    |             |                                | t
 rows_shared      |  33 | f    | rows_shape  | the limit rows_shared is 32    | 
 virtual_leaves   |  16 | t    |             |                                | t
 virtual_leaves   |  17 | f    | shape_width | the limit virtual_leaves is 16 | 
(18 rows)

SELECT * FROM at_width(true);
       kind       | n | used |    slug     |            detail             | rows_ok 
------------------+---+------+-------------+-------------------------------+---------
 arms             | 3 | t    |             |                               | t
 arms             | 4 | f    | shape_width | the limit arms is 3           | 
 combos           | 3 | t    |             |                               | t
 combos           | 4 | f    | shape_width | the limit combos is 3         | 
 first_order_cols | 3 | t    |             | first-row lookups: 1          | t
 first_order_cols | 4 | t    |             | first-row lookups: 0          | t
 join_cols        | 3 | t    |             |                               | t
 join_cols        | 4 | f    | shape_width | the limit join_cols is 3      | 
 joins            | 3 | t    |             |                               | t
 joins            | 4 | f    | shape_width | the limit joins is 3          | 
 rows_shared      | 3 | t    |             |                               | t
 rows_shared      | 4 | f    | rows_shape  | the limit rows_shared is 3    | 
 virtual_leaves   | 3 | t    |             |                               | t
 virtual_leaves   | 4 | f    | shape_width | the limit virtual_leaves is 3 | 
(14 rows)

DROP FUNCTION at_width(bool), shape(text, int), first_row_lookups(text), probe(text, text), same(text);
DROP TABLE widths, wide, dd, events;
DROP EXTENSION pg_sift;