-- Differential test of the plan agg_flat_key_lookup covers: a flattened lateral sub-select under a CTE of keys that reads another table with a column index.
-- Every statement runs with the aggregate scan off and on and the rows are compared as sorted jsonb text (a total order of the whole row), for groups of 100, 80, 1
-- and no members, for entries without a member, without scores or with empty scores, labels that map to NULL, and after updates, deletes and inserts of the entries,
-- the links and the labels between the runs.
-- Before PostgreSQL 18 the cast of a JSON null score raises an error in the normal plan too, so the scan refuses the statements that read one (error_parity) and
-- the groups with such entries answer with the error in both runs (agg_flat_key_lookup_diff_1.out).
CREATE EXTENSION pg_sift;
SET client_min_messages = warning;	-- the NOTICE of the order array gives bytes a row that follow the sample of ANALYZE
SET max_parallel_workers_per_gather = 0;
CREATE TABLE sources (id int PRIMARY KEY, enabled bool);
INSERT INTO sources SELECT i, i <> 7 FROM generate_series(1, 20) i;
CREATE TABLE labels (foreign_id int, source_key text, label text);
INSERT INTO labels SELECT s, k, upper(k) FROM generate_series(1, 20) s, unnest(ARRAY['alpha', 'beta', 'gamma']) k WHERE (s + length(k)) % 3 <> 0;
-- the table the CTE of keys picks from: it has a column index, so the plan that reads it is an ordinary read of an indexed table
CREATE TABLE dim_b (id int PRIMARY KEY, kind text);
INSERT INTO dim_b SELECT i, CASE WHEN i % 7 = 0 THEN 'x' ELSE 'a' END FROM generate_series(1, 3000) i;
CREATE INDEX members_col ON dim_b USING sift_col (id, kind);
CREATE TABLE links (grp int, dim_b_id int);
INSERT INTO links SELECT 1, i FROM generate_series(1, 100) i;
INSERT INTO links SELECT 2, i FROM generate_series(301, 380) i;
CREATE UNIQUE INDEX links_idx ON links (grp, dim_b_id);
CREATE TABLE assigns (sub int, grp int, dim_b_id int);
INSERT INTO assigns SELECT 9, 1, NULL;
INSERT INTO assigns SELECT 9, NULL, i FROM generate_series(101, 400) i;
-- the first 100 members hold 1,200 rows each, the others 20
CREATE TABLE entries (id serial, dim_b_id int, foreign_id int, day date, metrics_a jsonb);
INSERT INTO entries (dim_b_id, foreign_id, day, metrics_a)
SELECT m.id, 1 + (m.id * 7 + g) % 20, date '2024-01-01' + (g * 37 + m.id * 11) % 900,
       jsonb_build_object('alpha', (g % 10)::float, 'beta_x', ((g * 3) % 10)::float) || CASE WHEN g % 3 = 0 THEN jsonb_build_object('gamma', (g % 7)::float) ELSE '{}' END
FROM dim_b m, LATERAL generate_series(1, CASE WHEN m.id <= 100 THEN 1200 ELSE 20 END) g;
CREATE INDEX entries_col ON entries USING sift_col (id, dim_b_id, foreign_id, day, metrics_a) WITH (key_cut = 'dim_b_id', paths = 'metrics_a.* float8');
CREATE INDEX entries_member ON entries (dim_b_id);
VACUUM (FREEZE, ANALYZE) entries;
ANALYZE dim_b; ANALYZE links; ANALYZE labels; ANALYZE sources; ANALYZE assigns;
-- a group of one member, and one whose members the subscription does not reach
INSERT INTO links VALUES (4, 5);
INSERT INTO links SELECT 5, i FROM generate_series(2001, 2010) i;
-- NULLs and empty documents the first test has none of: entries without a member, without a document, with an empty one, a label that maps to NULL
INSERT INTO entries (dim_b_id, foreign_id, day, metrics_a)
SELECT CASE WHEN g % 5 = 0 THEN NULL ELSE 1 + g % 100 END, 1 + g % 20, date '2025-06-01' + g % 40,
       CASE g % 4 WHEN 0 THEN NULL WHEN 1 THEN '{}'::jsonb ELSE jsonb_build_object('alpha', g % 10, 'gamma', NULL) END
FROM generate_series(1, 3000) g;
INSERT INTO labels VALUES (3, 'alpha', NULL), (4, 'gamma', 'ALPHA');
ANALYZE entries; ANALYZE labels;
\ir test/agg_helpers.sql
-- the report of a group: members the subscription reaches in the group, the entries of two periods, labels of the unnested scores, a spine of the periods
CREATE FUNCTION fk_period(g int) RETURNS text LANGUAGE sql AS $f$ SELECT format($q$
 WITH picked AS MATERIALIZED (SELECT DISTINCT m.id FROM dim_b m
   WHERE m.id IN (SELECT coalesce(a.dim_b_id, l.dim_b_id) FROM assigns a LEFT JOIN links l ON l.grp = a.grp WHERE a.sub = 9 AND coalesce(a.dim_b_id, l.dim_b_id) IS NOT NULL)
     AND m.id IN (SELECT dim_b_id FROM links WHERE grp IN (%s)) AND m.kind = 'a'),
 periods AS (SELECT * FROM (VALUES ('current', daterange('2025-06-01', '2025-06-30', '[]')), ('previous', daterange('2024-06-01', '2024-06-30', '[]'))) p(period, range)),
 parsed AS (
   SELECT e.id, e.dim_b_id, p.period, s.category, CAST(s.value AS float) AS score
   FROM entries e JOIN periods p ON e.day <@ p.range JOIN picked ON picked.id = e.dim_b_id
   JOIN sources src ON src.id = e.foreign_id AND src.enabled IS TRUE
   INNER JOIN LATERAL (SELECT coalesce(l.label, split_part(r.key, '_', 1)) AS category, r.value
                       FROM jsonb_each(e.metrics_a) r LEFT JOIN labels l ON l.foreign_id = e.foreign_id AND l.source_key = r.key) s ON true)
 SELECT d.period, x.category, count(DISTINCT x.id) AS entries, count(DISTINCT x.dim_b_id) AS dim_b, count(x.score) AS metrics_a, round(avg(x.score)::numeric, 6) AS score,
        round(sum(x.score)::numeric, 6) AS total
 FROM periods d LEFT JOIN parsed x ON x.period = d.period
 GROUP BY d.period, x.category
$q$, g) $f$;
-- the same members, one level: no spine, the filter on the day in the WHERE, a total by label
CREATE FUNCTION fk_total(g int) RETURNS text LANGUAGE sql AS $f$ SELECT format($q$
 WITH picked AS MATERIALIZED (SELECT DISTINCT m.id FROM dim_b m
   WHERE m.id IN (SELECT coalesce(a.dim_b_id, l.dim_b_id) FROM assigns a LEFT JOIN links l ON l.grp = a.grp WHERE a.sub = 9 AND coalesce(a.dim_b_id, l.dim_b_id) IS NOT NULL)
     AND m.id IN (SELECT dim_b_id FROM links WHERE grp IN (%s)) AND m.kind = 'a')
 SELECT coalesce(l.label, split_part(r.key, '_', 1)) AS category, count(*) AS n, count(DISTINCT e.id) AS entries, sum(CAST(r.value AS float)) AS total
 FROM entries e JOIN picked ON picked.id = e.dim_b_id
 INNER JOIN LATERAL (SELECT key, value FROM jsonb_each(e.metrics_a)) r ON true
 LEFT JOIN labels l ON l.foreign_id = e.foreign_id AND l.source_key = r.key
 WHERE e.day >= '2025-06-01' AND e.day < '2025-07-01'
 GROUP BY 1
$q$, g) $f$;
CREATE FUNCTION fk_check(step text) RETURNS TABLE (label text, same bool, rows int) LANGUAGE plpgsql AS $f$
DECLARE g int; q text;
BEGIN
  FOREACH g IN ARRAY ARRAY[1, 2, 3, 4, 5] LOOP
    FOREACH q IN ARRAY ARRAY[fk_period(g), fk_total(g)] LOOP
      label := step || ' group ' || g || CASE WHEN q = fk_period(g) THEN ' periods' ELSE ' total' END;
      same := outcome(q, false) = outcome(q, true);
      rows := CASE WHEN outcome(q, true) LIKE '[%' THEN (SELECT count(*) FROM jsonb_array_elements(outcome(q, true)::jsonb)) WHEN outcome(q, true) = 'no rows' THEN 0 ELSE -1 END;
      RETURN NEXT;
    END LOOP;
  END LOOP;
END $f$;
SELECT * FROM fk_check('initial');
-- the first group is the one the scan takes over
SELECT pushed(fk_period(1)) AS group_1_pushed;
-- updates and deletes of the entries and of the links, the labels changed, between the runs
UPDATE entries SET day = day + 30 WHERE id % 7 = 0;
DELETE FROM entries WHERE id % 11 = 0;
UPDATE entries SET metrics_a = NULL WHERE id % 13 = 0;
UPDATE entries SET metrics_a = metrics_a || '{"beta_x": null}' WHERE id % 17 = 0 AND metrics_a IS NOT NULL;
SELECT * FROM fk_check('after dml');
DELETE FROM links WHERE grp = 2 AND dim_b_id % 2 = 0;
INSERT INTO links VALUES (4, 6), (5, 101), (3, 2);
UPDATE labels SET label = lower(label) WHERE foreign_id % 3 = 0;
DELETE FROM labels WHERE foreign_id = 5;
UPDATE assigns SET grp = 4 WHERE grp = 1;
SELECT * FROM fk_check('after links');
UPDATE assigns SET grp = 1 WHERE grp = 4;
INSERT INTO entries (dim_b_id, foreign_id, day, metrics_a)
SELECT 1 + g % 100, 1 + g % 20, date '2025-06-01' + g % 30, jsonb_build_object('alpha', g % 9, 'beta_x', (g * 3) % 10) FROM generate_series(1, 500) g;
SELECT * FROM fk_check('after insert');
VACUUM (ANALYZE) entries;
SELECT * FROM fk_check('after vacuum');
DROP FUNCTION fk_check(text), fk_period(int), fk_total(int), outcome(text, bool), runs_in_scan(text), pushed, runq;
DROP TABLE entries, assigns, links, dim_b, labels, sources;
DROP EXTENSION pg_sift;