-- UNION ALL arms that test an EXISTS over the elements of an array (a lookup per element of the entity's kind) share one pass over the index:
-- the test and its lookup run once for the arms together, and an arm that has a test the others lack applies it to its own set.
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 dim_a (id int PRIMARY KEY, kind text);
INSERT INTO dim_a SELECT i, CASE WHEN i % 7 = 0 THEN NULL WHEN i % 2 = 0 THEN 'a' ELSE 'b' END FROM generate_series(0, 48) i;
CREATE TABLE mappings (kind text, category text, val text, derived_category text);
INSERT INTO mappings VALUES ('a', 'catg_l', '', 'catg_j'), ('a', 'catg_l', '', 'duplicate'), ('b', 'catg_l', '', 'catg_k'), ('b', 'catg_a', 'x', 'catg_b'),
 ('a', 'nul', NULL, 'never'), ('a', 'catg_a', 'x', NULL), (NULL, 'catg_l', '', 'catg_j');
CREATE TABLE facts (id int, dim_a_id int, day date, document jsonb) WITH (autovacuum_enabled = false);
INSERT INTO facts
SELECT i, i % 52, '2026-01-01'::date + (i % 9),
 CASE WHEN i % 29 = 0 THEN NULL
      WHEN i % 31 = 0 THEN '{}'::jsonb
      ELSE jsonb_build_object('parts', CASE WHEN i % 7 = 0 THEN '[]'::jsonb
   WHEN i % 3 = 0 THEN '[{"category":"catg_l","val":"","dir_a":"up","score":8}]'::jsonb
   ELSE '[{"category":"catg_l","val":"","dir_a":"up","score":9},{"category":"catg_a","val":"x","dir_a":"down","score":3},{"category":"zzz","val":"","dir_a":"flat","score":null},{"category":"nul","val":null,"dir_a":"flat","score":1}]'::jsonb END)
   || CASE WHEN i % 5 = 0 THEN '{"flag_a":false}'::jsonb WHEN i % 11 = 0 THEN '{"flag_a":null}'::jsonb WHEN i % 13 = 0 THEN '{"flag_a":true}'::jsonb ELSE '{}'::jsonb END END
FROM generate_series(1, 6000) i;
CREATE INDEX facts_col ON facts USING sift_col (id, dim_a_id, day, document, jsonb_exists(document, 'flag_a'))
 WITH (paths = 'document.parts[]{category text,val text,score float8,dir_a text}');
NOTICE:  sift_col index "facts_col": 1 advice (track_changes)
DETAIL:  advice: track_changes = the table has no one-column primary key
HINT:  Reasons, costs and how to turn each off: SELECT * FROM sift.index_choices('facts_col');
VACUUM (FREEZE, ANALYZE) facts;
ANALYZE dim_a;
ANALYZE mappings;
-- the most arms one scan served in the plan; the rows must be those of the normal plan
CREATE FUNCTION arms_shared(q text) RETURNS int LANGUAGE sql AS $$ SELECT sift_helpers.agg_arms_in_one_pass(q) $$;
CREATE FUNCTION same(q text) RETURNS bool LANGUAGE sql AS $$ SELECT sift_helpers.agg_equals_normal_all(q) $$;
-- the test of an arm as the condition of its WHERE: the categories its elements are mapped to
\set cats_jb '(SELECT 1 FROM LATERAL (SELECT (CASE WHEN jsonb_exists(f.document, ''flag_a'') THEN o.category ELSE m.derived_category END) AS category FROM jsonb_to_recordset(f.document -> ''parts'') AS o(category text, val text, score float8, dir_a text) LEFT JOIN LATERAL (SELECT lm.derived_category FROM mappings lm WHERE lm.kind = e.kind AND lm.category = o.category AND lm.val = o.val AND lm.derived_category IS NOT NULL AND NOT jsonb_exists(f.document, ''flag_a'') LIMIT 1) m ON true WHERE o.score IS NOT NULL) x WHERE x.category IN (''catg_j'', ''catg_b''))'
\set cats_k '(SELECT 1 FROM LATERAL (SELECT (CASE WHEN jsonb_exists(f.document, ''flag_a'') THEN o.category ELSE m.derived_category END) AS category FROM jsonb_to_recordset(f.document -> ''parts'') AS o(category text, val text, score float8, dir_a text) LEFT JOIN LATERAL (SELECT lm.derived_category FROM mappings lm WHERE lm.kind = e.kind AND lm.category = o.category AND lm.val = o.val AND lm.derived_category IS NOT NULL AND NOT jsonb_exists(f.document, ''flag_a'') LIMIT 1) m ON true WHERE o.score IS NOT NULL) x WHERE x.category IN (''catg_k''))'
\set from 'FROM facts f JOIN dim_a e ON e.id = f.dim_a_id WHERE '
-- three arms with the same test and keys of their own
\set same_test 'SELECT e.kind AS a, NULL::date AS b, count(*) AS n, avg(f.id) AS m ' :from 'EXISTS ' :cats_jb ' GROUP BY e.kind UNION ALL SELECT e.kind, f.day, count(*), avg(f.id) ' :from 'EXISTS ' :cats_jb ' GROUP BY e.kind, f.day UNION ALL SELECT NULL, f.day, count(*), avg(f.id) ' :from 'EXISTS ' :cats_jb ' GROUP BY f.day'
SELECT arms_shared(:'same_test') AS same_test_shared, same(:'same_test') AS same_test_same;
 same_test_shared | same_test_same 
------------------+----------------
                3 | t
(1 row)

-- a test the other arms do not have, and other tests per arm (each applies to the rows of its own arm)
\set one_test 'SELECT e.kind AS a, NULL::date AS b, count(*) AS n, avg(f.id) AS m ' :from 'EXISTS ' :cats_jb ' GROUP BY e.kind UNION ALL SELECT e.kind, f.day, count(*), avg(f.id) ' :from 'true GROUP BY e.kind, f.day'
SELECT arms_shared(:'one_test') AS one_test_shared, same(:'one_test') AS one_test_same;
 one_test_shared | one_test_same 
-----------------+---------------
               0 | t
(1 row)

\set other_tests 'SELECT e.kind AS a, NULL::date AS b, count(*) AS n, avg(f.id) AS m ' :from 'EXISTS ' :cats_jb ' GROUP BY e.kind UNION ALL SELECT e.kind, f.day, count(*), avg(f.id) ' :from 'EXISTS ' :cats_k ' GROUP BY e.kind, f.day UNION ALL SELECT NULL, f.day, count(*), avg(f.id) ' :from 'f.id > 100 GROUP BY f.day'
SELECT arms_shared(:'other_tests') AS other_tests_shared, same(:'other_tests') AS other_tests_same;
 other_tests_shared | other_tests_same 
--------------------+------------------
                  2 | t
(1 row)

-- an arm without keys is not a set of the pass: it runs on its own, and its rows are the same
\set no_group 'SELECT count(*) AS n, avg(f.id) AS m, NULL::text AS a ' :from 'EXISTS ' :cats_jb ' UNION ALL SELECT count(*), avg(f.id), e.kind ' :from 'EXISTS ' :cats_jb ' GROUP BY e.kind'
SELECT same(:'no_group') AS no_group_same;
 no_group_same 
---------------
 t
(1 row)

DROP FUNCTION arms_shared(text), same(text);
DROP TABLE facts, dim_a, mappings;
DROP EXTENSION pg_sift_helpers;
DROP EXTENSION pg_sift;