-- The plan the planner picks when a CTE hands the scan a few parent ids. PostgreSQL counts the CTE at the rows its sub-select
-- keeps (a row for one parent), and the normal plan is then a probe of the btree for each key. The aggregate over the elements of a
-- jsonb document was charged the excess of the table's rows over that count (and the keys over it) as a lost estimate, which made every
-- request for one or a few parents take the scan at a price above the normal plan's, and lose. A few keys keep the normal plan; a
-- long list of them, where the probes outnumber the scan's start, keeps the scan. Every statement is also run with sift.enable_agg = off:
-- the rows must be the same. A CTE of ids that PostgreSQL takes through a join estimates at a fraction of the smaller of the lists it joins (ten ids semi-joined
-- to forty: one) is not a request for one parent: the normal plan then reads the rows of ten of them and the scan is taken.
CREATE EXTENSION pg_sift;
SET client_min_messages = warning;
SET max_parallel_workers_per_gather = 0;
CREATE TABLE fk_facts (id int PRIMARY KEY, parent_id int NOT NULL, ts date NOT NULL, deleted_at timestamptz, metrics_a jsonb NOT NULL) WITH (autovacuum_enabled = false);
INSERT INTO fk_facts SELECT i, 1 + (i * 7919) % 600, date '2024-01-01' + (i % 700), CASE WHEN i % 97 = 0 THEN now() END,
	jsonb_build_object('mhead', 1 + i % 10) || CASE WHEN i % 3 = 0 THEN jsonb_build_object('quality', 1 + i % 7) ELSE '{}' END
	FROM generate_series(1, 60000) i;
CREATE INDEX fk_facts_parent ON fk_facts (parent_id, ts, deleted_at);
CREATE INDEX fk_facts_col ON fk_facts USING sift_col (id, parent_id, ts, deleted_at, metrics_a) WITH (paths = 'metrics_a.* float8');
-- scope 1: one parent, 2: three, 3: forty. The window is a month for the one parent (the normal plan 124 against the scan's 331, a gap of 62%: a year was 361 against 422, 14%, a tie
-- that the statistics could turn) and a year for the others (the scan 474, 999 and 623 against 1,020, 8,016 and 1,967: 54%, 88% and 68%)
CREATE TABLE fk_scope (scope_id int NOT NULL, parent_id int NOT NULL);
INSERT INTO fk_scope SELECT 1, 11;
INSERT INTO fk_scope SELECT 2, p FROM generate_series(21, 23) p;
INSERT INTO fk_scope SELECT 3, p FROM generate_series(101, 140) p;
CREATE INDEX ON fk_scope (scope_id);
CREATE TABLE fk_parents (id int PRIMARY KEY, name text);
INSERT INTO fk_parents SELECT i, 'p' || i FROM generate_series(1, 600) i;
CREATE TABLE fk_group (grp int NOT NULL, parent_id int NOT NULL);
INSERT INTO fk_group SELECT 4, p FROM generate_series(201, 240) p;
CREATE INDEX ON fk_group (grp);
VACUUM (FREEZE, ANALYZE) fk_facts, fk_scope, fk_parents, fk_group;
CREATE FUNCTION fk_q(sub int) RETURNS text LANGUAGE sql AS $$ SELECT format($f$
WITH parent_ids AS MATERIALIZED (%s),
     parents AS MATERIALIZED (SELECT p.id FROM fk_parents p JOIN parent_ids i ON i.id = p.id)
SELECT split_part(r.key, '_', 1) AS kind, round(avg((r.value #>> '{}')::float8)::numeric, 6) AS score, count(*) AS n
FROM fk_facts v JOIN parents p ON p.id = v.parent_id, jsonb_each(v.metrics_a) r
WHERE v.deleted_at IS NULL AND v.ts >= date '2024-03-01' AND v.ts < date '2024-03-01' + %s
GROUP BY 1 ORDER BY 1$f$, CASE WHEN sub < 4 THEN format('SELECT DISTINCT parent_id AS id FROM fk_scope WHERE scope_id = %s', sub)
       ELSE $s$SELECT DISTINCT p.id FROM fk_parents p WHERE p.id = ANY ('{201,205,210,215,220,225,230,233,236,239}'::int[])
       AND p.id IN (SELECT parent_id FROM fk_group WHERE grp = 4)$s$ END, CASE WHEN sub = 1 THEN 31 ELSE 365 END) $$;
CREATE FUNCTION fk_plan(sub int, agg text DEFAULT 'on') RETURNS TABLE (plan text, total_cost float8) LANGUAGE plpgsql AS $$
DECLARE j jsonb;
BEGIN
	PERFORM set_config('sift.enable_agg', agg, true);
	EXECUTE 'EXPLAIN (FORMAT JSON) ' || fk_q(sub) INTO j;
	plan := CASE WHEN j::text LIKE '%"SiftAgg"%' THEN 'SiftAgg' ELSE 'normal' END;
	total_cost := (j -> 0 -> 'Plan' ->> 'Total Cost')::float8;
	RETURN NEXT;
END $$;
CREATE FUNCTION fk_same(sub int) RETURNS bigint LANGUAGE plpgsql AS $$
DECLARE n bigint;
BEGIN
	EXECUTE format('CREATE TEMP TABLE r_on ON COMMIT DROP AS %s', fk_q(sub));
	SET LOCAL sift.enable_agg = off;
	EXECUTE format('CREATE TEMP TABLE r_off ON COMMIT DROP AS %s', fk_q(sub));
	RESET sift.enable_agg;
	n := (SELECT count(*) FROM ((TABLE r_on EXCEPT TABLE r_off) UNION ALL (TABLE r_off EXCEPT TABLE r_on)) d);
	DROP TABLE r_on, r_off;
	RETURN n;
END $$;
-- the plan taken, and whether it costs no more than 1.2 times the plan with the scan off
SELECT sub AS scope, (SELECT CASE WHEN sub < 4 THEN (SELECT count(*) FROM fk_scope WHERE scope_id = sub) ELSE 10 END) AS parents, a.plan,
       a.total_cost <= 1.2 * n.total_cost AS within_normal_cost, fk_same(sub) AS rows_unlike_normal
FROM generate_series(1, 4) sub, LATERAL fk_plan(sub, 'on') a, LATERAL fk_plan(sub, 'off') n;
 scope | parents |  plan   | within_normal_cost | rows_unlike_normal 
-------+---------+---------+--------------------+--------------------
     1 |       1 | normal  | t                  |                  0
     2 |       3 | SiftAgg | t                  |                  0
     3 |      40 | SiftAgg | t                  |                  0
     4 |      10 | SiftAgg | f                  |                  0
(4 rows)

DROP FUNCTION fk_q(int), fk_plan(int, text), fk_same(int);
DROP TABLE fk_facts, fk_scope, fk_parents, fk_group;
DROP EXTENSION pg_sift;