-- EXISTS over the elements of an array that reads a column of a joined table (or a lookup per element) is a test of the joined tuples:
-- PostgreSQL cannot pull it up, the scan resolves each distinct element once against the lookup and keeps the rows with one that passes.
CREATE EXTENSION pg_sift;
SET client_min_messages = warning;
SET max_parallel_workers_per_gather = 0;
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, derived_val text);
INSERT INTO mappings VALUES ('a', 'catg_l', '', 'catg_j', NULL), ('a', 'catg_l', '', 'duplicate', NULL), ('b', 'catg_l', '', 'catg_k', NULL),
 ('b', 'catg_a', 'x', 'catg_b', 'catg_i'), ('a', 'nul', NULL, 'never', NULL), ('a', 'catg_a', 'x', NULL, NULL), (NULL, 'catg_l', '', 'catg_j', NULL);
CREATE TABLE dates (day date PRIMARY KEY);
INSERT INTO dates SELECT d::date FROM generate_series('2026-01-01'::date, '2026-01-12', '1 day') d;
CREATE TABLE facts (id int, dim_a_id int, day date, document jsonb, metrics jsonb, payload text);
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,
 CASE WHEN i % 17 = 0 THEN '{}'::jsonb ELSE jsonb_build_object('m_total', i % 11, 'catg_j_catg_i', (i * 3) % 11) END,
 repeat('payload', 30)
FROM generate_series(1, 30000) i;
CREATE INDEX facts_col ON facts USING sift_col (id, dim_a_id, day, document, metrics, jsonb_exists(document, 'flag_a'))
 WITH (paths = 'document.parts[]{category text,val text,score float8,dir_a text}, metrics.* float8');
VACUUM (FREEZE, ANALYZE) facts;
ANALYZE dates;
ANALYZE dim_a;
ANALYZE mappings;
\ir test/agg_check_query.sql
\set ECHO none
-- The lookup-per-element shape: a lookup per element chosen by the entity's kind, the unmapped category of a document that is not mapped, the element's own otherwise.
SELECT * FROM check_query($q$
 SELECT e.kind, count(*) AS n, avg(f.id) AS mean FROM facts f JOIN dim_a e ON e.id = f.dim_a_id
 WHERE EXISTS (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'))
 GROUP BY e.kind
$q$);
 pushed | eligible | differences 
--------+----------+-------------
 t      | t        |           0
(1 row)

-- Another element condition beside it, and the aggregates over the filtered rows.
SELECT * FROM check_query($q$
 SELECT e.kind, count(*) AS n, count(DISTINCT f.id) AS items, sum(f.id) AS total FROM facts f JOIN dim_a e ON e.id = f.dim_a_id
 WHERE f.id > 100 AND EXISTS (SELECT 1 FROM LATERAL (
   SELECT (CASE WHEN jsonb_exists(f.document, 'flag_a') THEN o.category ELSE m.derived_category END) AS category, o.dir_a
   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') AND x.dir_a IN ('up', 'flat'))
 GROUP BY e.kind
$q$);
 pushed | eligible | differences 
--------+----------+-------------
 t      | t        |           0
(1 row)

-- The element is compared with a column of the joined table, without a lookup.
SELECT * FROM check_query($q$
 SELECT e.kind, count(*) AS n FROM facts f JOIN dim_a e ON e.id = f.dim_a_id
 WHERE EXISTS (SELECT 1 FROM jsonb_to_recordset(f.document -> 'parts') AS o(category text, val text, score float8, dir_a text)
   WHERE o.score IS NOT NULL AND o.category = CASE e.kind WHEN 'a' THEN 'catg_l' ELSE 'catg_a' END)
 GROUP BY e.kind
$q$);
 pushed | eligible | differences 
--------+----------+-------------
 t      | t        |           0
(1 row)

-- Two of them, and a count of the elements of another array of the same rows (the tuples repeat the row).
SELECT * FROM check_query($q$
 SELECT r.key, count(*) AS n, avg(r.value::float8) AS mean FROM facts f JOIN dim_a e ON e.id = f.dim_a_id CROSS JOIN LATERAL jsonb_each(f.metrics) r
 WHERE EXISTS (SELECT 1 FROM jsonb_to_recordset(f.document -> 'parts') AS o(category text, val text, score float8, dir_a text)
   WHERE o.category = CASE e.kind WHEN 'a' THEN 'catg_l' ELSE 'catg_a' END)
   AND EXISTS (SELECT 1 FROM jsonb_to_recordset(f.document -> 'parts') AS o(category text, val text, score float8, dir_a text)
   WHERE o.dir_a = CASE e.kind WHEN 'a' THEN 'up' ELSE 'down' END)
 GROUP BY r.key
$q$);
 pushed | eligible | differences 
--------+----------+-------------
 t      | t        |           0
(1 row)

-- The entities with no row, the ones with a NULL kind and the documents with no array: the same rows as PostgreSQL's.
SELECT * FROM check_query($q$
 SELECT e.id IS NULL AS no_entity, e.kind, count(*) AS n FROM facts f LEFT JOIN dim_a e ON e.id = f.dim_a_id
 WHERE EXISTS (SELECT 1 FROM jsonb_to_recordset(f.document -> 'parts') AS o(category text, val text, score float8, dir_a text)
   WHERE o.category = 'catg_l' AND (e.kind IS NULL OR o.dir_a <> 'down'))
 GROUP BY 1, 2
$q$);
 pushed | eligible | differences 
--------+----------+-------------
 t      | t        |           0
(1 row)

-- The test in the nullable input of a preserved side: the buckets that no row reaches stay.
SELECT * FROM check_query($q$
 SELECT d.day, count(DISTINCT f.id) AS n, avg(f.id) AS mean FROM dates d
 LEFT JOIN (SELECT f.id, f.day FROM facts f JOIN dim_a e ON e.id = f.dim_a_id
   WHERE f.id % 4 <> 3 AND EXISTS (SELECT 1 FROM jsonb_to_recordset(f.document -> 'parts') AS o(category text, val text, score float8, dir_a text)
     WHERE o.category = CASE e.kind WHEN 'a' THEN 'catg_l' ELSE 'catg_a' END)) f ON f.day = d.day
 GROUP BY d.day
$q$);
 pushed | eligible | differences 
--------+----------+-------------
 t      | t        |           0
(1 row)

-- NOT EXISTS, and an EXISTS inside an aggregate's FILTER, stay with PostgreSQL.
SELECT * FROM check_query($q$
 SELECT e.kind, count(*) AS n FROM facts f JOIN dim_a e ON e.id = f.dim_a_id
 WHERE NOT EXISTS (SELECT 1 FROM jsonb_to_recordset(f.document -> 'parts') AS o(category text, val text, score float8, dir_a text)
   WHERE o.category = CASE e.kind WHEN 'a' THEN 'catg_l' ELSE 'catg_a' END)
 GROUP BY e.kind
$q$);
 pushed | eligible | differences 
--------+----------+-------------
 f      | f        |           0
(1 row)

SELECT * FROM check_query($q$
 SELECT e.kind, count(*) FILTER (WHERE EXISTS (SELECT 1 FROM jsonb_to_recordset(f.document -> 'parts') AS o(category text, val text, score float8, dir_a text)
   WHERE o.category = CASE e.kind WHEN 'a' THEN 'catg_l' ELSE 'catg_a' END)) FROM facts f JOIN dim_a e ON e.id = f.dim_a_id
 GROUP BY e.kind
$q$);
 pushed | eligible | differences 
--------+----------+-------------
 f      | f        |           0
(1 row)

-- EXPLAIN: the test, the lookup it resolves per element, and what the scan runs.
EXPLAIN (COSTS OFF)
SELECT e.kind, count(*) FROM facts f JOIN dim_a e ON e.id = f.dim_a_id
WHERE EXISTS (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'))
GROUP BY e.kind;
                                                                                                                                    QUERY PLAN                                                                                                                                     
-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
 Custom Scan (SiftAgg)
   Index: facts_col
   Filters: EXISTS document.parts[] (lookup output, document.parts[].category, document.parts[].val, jsonb_exists(f.document, 'flag_a'::text): (CASE WHEN jsonb_exists(f.document, 'flag_a'::text) THEN o.category ELSE m.derived_category END = ANY ('{catg_j,catg_b}'::text[])))
   Group Keys: dim_a_id: e.kind
   Pushed Aggregates: count(*)
   Lookup Joins: dim_a, m (left, per element)
   Name: m (left, per element)
   ->  Limit
         ->  Result
               One-Time Filter: (NOT (InitPlan input).col4)
               ->  Seq Scan on mappings lm
                     Filter: ((derived_category IS NOT NULL) AND (kind = (InitPlan input).col1) AND (category = (InitPlan input).col2) AND (val = (InitPlan input).col3))
(12 rows)

-- Errors come from PostgreSQL's evaluation, for the rows and elements it reaches: an expression that may raise keeps the normal plan, and the
-- row that raises is the same in both. The element that raises is after the one that passes in the rows below, which EXISTS never evaluates.
CREATE TABLE raising (id int, dim_a_id int, document jsonb);
INSERT INTO raising
SELECT i, i % 52, jsonb_build_object('parts', jsonb_build_array(jsonb_build_object('category', 'catg_l', 'val', '1'),
  jsonb_build_object('category', 'catg_a', 'val', CASE WHEN i % 1000 = 0 THEN 'x' ELSE '2' END)))
FROM generate_series(1, 20000) i;
CREATE INDEX raising_col ON raising USING sift_col (id, dim_a_id, document) WITH (paths = 'document.parts[]{category text,val text}');
VACUUM (FREEZE, ANALYZE) raising;
CREATE FUNCTION check_error(q text) RETURNS TABLE (pushed bool, same bool, message text) LANGUAGE plpgsql AS $$
DECLARE on_msg text; off_msg text; line text;
BEGIN
  SET LOCAL sift.enable_agg = off;
  BEGIN EXECUTE q; off_msg := 'no error'; EXCEPTION WHEN OTHERS THEN off_msg := SQLERRM; END;
  SET LOCAL sift.enable_agg = on;
  BEGIN EXECUTE q; on_msg := 'no error'; EXCEPTION WHEN OTHERS THEN on_msg := SQLERRM; END;
  pushed := false;
  FOR line IN EXECUTE 'EXPLAIN (COSTS OFF) ' || q LOOP pushed := pushed OR line LIKE '%Custom Scan (SiftAgg)%'; END LOOP;
  same := on_msg = off_msg;
  message := on_msg;
  RETURN NEXT;
END $$;
-- the cast reaches the element that raises (no element of the row passes before it)
SELECT * FROM check_error($q$
 SELECT e.kind, count(*) FROM raising r JOIN dim_a e ON e.id = r.dim_a_id
 WHERE EXISTS (SELECT 1 FROM jsonb_to_recordset(r.document -> 'parts') AS o(category text, val text)
   WHERE o.val::int > 100 AND o.category <> e.kind) GROUP BY e.kind
$q$);
 pushed | same |                  message                   
--------+------+--------------------------------------------
 f      | t    | invalid input syntax for type integer: "x"
(1 row)

-- the first element passes: the second, which raises, is not evaluated
SELECT * FROM check_error($q$
 SELECT e.kind, count(*) FROM raising r JOIN dim_a e ON e.id = r.dim_a_id
 WHERE EXISTS (SELECT 1 FROM jsonb_to_recordset(r.document -> 'parts') AS o(category text, val text)
   WHERE (o.category = 'catg_l' OR o.val::int > 0) AND e.kind IS DISTINCT FROM 'zz') GROUP BY e.kind
$q$);
 pushed | same | message  
--------+------+----------
 f      | t    | no error
(1 row)

-- the row that raises is dropped by another condition before the test
SELECT * FROM check_error($q$
 SELECT e.kind, count(*) FROM raising r JOIN dim_a e ON e.id = r.dim_a_id
 WHERE r.id % 1000 <> 0 AND EXISTS (SELECT 1 FROM jsonb_to_recordset(r.document -> 'parts') AS o(category text, val text)
   WHERE o.val::int > 100 AND o.category <> e.kind) GROUP BY e.kind
$q$);
 pushed | same | message  
--------+------+----------
 f      | t    | no error
(1 row)

DROP FUNCTION check_query(text), check_error(text);
DROP TABLE facts, raising, dim_a, mappings, dates;
DROP EXTENSION pg_sift;