-- A statement that names the entities it reads with IN (SELECT ..) in a CTE, joins them to a large table by its key (`FROM keys JOIN dim ON dim.id = keys.id`), and reads the events of those entities
-- with the keys a lateral jsonb_each unnests, is a join to a large table whose only partner is the keys. Made a lateral probe of the keys' rows (what a table of ten thousand rows or more
-- gets), the keys were left with no join condition of their own once the statement was flattened, and the scan found no order to look the joins up in (join_condition): every statement of the
-- shape ran the normal plan as soon as the table had ten thousand rows. The planner's own nested loop probes the table by the keys, and the scan looks up the rows of the table by the keys of the events.
CREATE EXTENSION pg_sift;
SET client_min_messages = warning;
SET jit = off;
SET max_parallel_workers_per_gather = 0;
-- at this size the normal plan costs less than the scan: the scan is forced, so that the test shows whether the statement can be planned as a scan at all
SET sift.agg_force = on;

-- a dimension of 12,000 rows, of which the 40 entities that have events are used
CREATE TABLE dim (id int PRIMARY KEY, kind text NOT NULL, label text NOT NULL) WITH (autovacuum_enabled = false);
INSERT INTO dim SELECT i, 'k' || i % 4, 'entity ' || i FROM generate_series(1, 12000) i;
CREATE TABLE dim_b (grp_id int NOT NULL, dim_id int NOT NULL, gone timestamptz) WITH (autovacuum_enabled = false);
CREATE INDEX dim_b_key ON dim_b (grp_id, dim_id);
INSERT INTO dim_b SELECT 1, i, NULL FROM generate_series(1, 11) i;
INSERT INTO dim_b SELECT 1, 12, now();
INSERT INTO dim_b SELECT 2, i, NULL FROM generate_series(10, 40) i;
CREATE TABLE kinds (id int PRIMARY KEY, enabled bool NOT NULL);
INSERT INTO kinds VALUES (1, true), (2, true), (3, false);
CREATE TABLE links (id int PRIMARY KEY, url text NOT NULL);
INSERT INTO links SELECT i, 'u' || i FROM generate_series(0, 49) i;
-- the names of the keys of a kind that are shown under another
CREATE TABLE key_labels (kind_id int, name text, shown text, UNIQUE (kind_id, name));
INSERT INTO key_labels VALUES (1, 'a', 'alpha'), (2, 'b', 'beta');
CREATE TABLE ev (id int PRIMARY KEY, dim_id int NOT NULL, kind_id int NOT NULL, link_id int, d date NOT NULL, flags jsonb NOT NULL) WITH (autovacuum_enabled = false);
INSERT INTO ev SELECT i, 1 + i % 40, 1 + i % 3, CASE WHEN i % 4 = 0 THEN NULL ELSE i % 50 END, date '2026-01-01' + i % 200, jsonb_build_object('a', i % 5, 'b', i % 3) FROM generate_series(1, 6000) i;
CREATE INDEX ev_col ON ev USING sift_col (id, dim_id, kind_id, link_id, d, flags) WITH (track_changes = 'off', paths = 'flags.* float8');
VACUUM (FREEZE, ANALYZE) dim, dim_b, kinds, links, key_labels, ev;

-- the entities from the keys a group names (a CTE, or a sub-select with a DISTINCT, which is not pulled up) or from a list, then the keys of their events
CREATE FUNCTION keyed_q(form text, grp int) RETURNS text LANGUAGE sql AS $$
	SELECT format('WITH %s,
		evs AS (SELECT ev.*, ents.kind FROM ev INNER JOIN ents ON ents.id = ev.dim_id LEFT JOIN links ON links.id = ev.link_id INNER JOIN kinds ON kinds.id = ev.kind_id AND kinds.enabled IS TRUE WHERE ev.d >= ''2026-02-01''),
		parsed AS (SELECT evs.id, evs.kind, m.label, (m.value #>> ''{}'')::float AS v FROM evs INNER JOIN LATERAL (SELECT coalesce(key_labels.shown, r.key) AS label, r.value FROM jsonb_each(evs.flags) AS r(key, value) LEFT JOIN key_labels ON key_labels.kind_id = evs.kind_id AND key_labels.name = r.key) m ON true)
		SELECT label, kind, round(avg(v)::numeric, 6), count(DISTINCT id) FROM parsed GROUP BY label, kind ORDER BY label, kind',
		CASE form
			WHEN 'cte' THEN format('ids AS MATERIALIZED (SELECT DISTINCT dim.id FROM dim WHERE dim.id IN (SELECT dim_id FROM dim_b WHERE gone IS NULL AND grp_id IN (%s))),
				ents AS (SELECT dim.*, ''base'' AS tag FROM ids AS pids INNER JOIN dim ON dim.id = pids.id)', grp)
			WHEN 'subselect' THEN format('ents AS (SELECT dim.* FROM (SELECT DISTINCT dim_id AS id FROM dim_b WHERE gone IS NULL AND grp_id IN (%s)) ids INNER JOIN dim ON dim.id = ids.id)', grp)
			WHEN 'list' THEN 'ents AS (SELECT dim.* FROM dim WHERE dim.id IN (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11))'
		END)
$$;

-- the scans of a plan
CREATE FUNCTION scans(sql text) RETURNS int LANGUAGE plpgsql AS $$
DECLARE l text; n int := 0; BEGIN
	FOR l IN EXECUTE 'EXPLAIN (COSTS OFF) ' || sql LOOP
		n := n + (l ~ 'Custom Scan \(SiftAgg\)')::int;
	END LOOP;
	RETURN n;
END $$;

-- the rows of the statement, as a text, through the scan or not
CREATE FUNCTION answer(sql text, scan bool) RETURNS text LANGUAGE plpgsql AS $$
DECLARE r text; BEGIN
	PERFORM set_config('sift.enable_agg', CASE WHEN scan THEN 'on' ELSE 'off' END, true);
	EXECUTE format('SELECT string_agg(x::text, E''\n'' ORDER BY x::text) FROM (%s) x', sql) INTO r;
	RETURN r;
END $$;

CREATE FUNCTION check_form(form text, grp int, OUT test text, OUT in_the_scan int, OUT same_rows_as_without bool, OUT has_rows bool) LANGUAGE plpgsql AS $$
DECLARE q text := keyed_q(form, grp); a text; b text; BEGIN
	test := form || ' of group ' || grp;
	PERFORM set_config('sift.enable_agg', 'on', true);
	in_the_scan := scans(q);
	a := answer(q, true);
	b := answer(q, false);
	same_rows_as_without := a = b;
	has_rows := length(a) > 0;
END $$;

-- the keys of a group, and of another: the scan runs, and the rows are those of the plan without it
SELECT * FROM check_form('list', 1);
SELECT * FROM check_form('cte', 1);
SELECT * FROM check_form('cte', 2);
SELECT * FROM check_form('subselect', 1);
SELECT * FROM check_form('subselect', 2);

-- a dimension below the size a probe is made for (ten thousand rows) is joined as written
DELETE FROM dim WHERE id > 9000;
VACUUM (FREEZE, ANALYZE) dim;
SELECT * FROM check_form('cte', 1);

DROP FUNCTION check_form(text, int);
DROP FUNCTION answer(text, bool);
DROP FUNCTION scans(text);
DROP FUNCTION keyed_q(text, int);
DROP TABLE ev, key_labels, links, kinds, dim_b, dim;
DROP EXTENSION pg_sift;