-- The shapes of an event analytics workload: documents with an array of parts held by an index expression, a lookup of the entities that an
-- ARRAY sub-select returns, EXISTS and nested jsonpath tests over the parts, and per-item arms over a text leaf with the cast in the aggregate
-- (docs/AGG.md, "Workload shapes"). Every statement is run with the pushdown and with sift.enable_agg = off: the rows must be the same.
CREATE EXTENSION pg_sift;
CREATE EXTENSION pg_sift_helpers;
SET max_parallel_workers_per_gather = 0;
-- the table is small: page costs that keep the btree and bitmap plans from looking cheap as well
SET seq_page_cost = 10000; SET random_page_cost = 10000;
CREATE TABLE ents (id int PRIMARY KEY, type text NOT NULL);
INSERT INTO ents SELECT i, (ARRAY['site', 'cat_c', 'cat_a', 'cat_b'])[1 + i % 4] FROM generate_series(1, 60) i;
-- the parts of an item: a category, a sign and a score that is a number, a numeric string or text
CREATE FUNCTION op_doc(i int) RETURNS jsonb LANGUAGE sql IMMUTABLE AS $$
	SELECT CASE WHEN i % 301 = 0 THEN '[]'::jsonb ELSE jsonb_build_object('lang', (ARRAY['en', 'de', 'fr'])[1 + i % 3], 'parts', (
		SELECT coalesce(jsonb_agg(jsonb_build_object(
			'category', (ARRAY['catg_j', 'catg_d', 'catg_m', 'catg_h'])[1 + (i * j) % 4],
			'dir_a', (ARRAY['up', 'down', 'flat'])[1 + (i + j) % 3],
			'score', CASE (i + j) % 11 WHEN 0 THEN to_jsonb('n/a'::text) WHEN 1 THEN to_jsonb(((i * j) % 10)::text) ELSE to_jsonb(((i * j) % 100) / 10.0) END) ORDER BY j), '[]')
		FROM generate_series(1, i % 5) j)) END $$;
-- the same with the scores as text (what a client that keeps the values as they are reads); a document that is no object is NULL
CREATE FUNCTION op_txt(doc jsonb) RETURNS jsonb LANGUAGE sql IMMUTABLE AS $$
	SELECT CASE WHEN jsonb_typeof(doc) = 'object' THEN jsonb_set(doc, '{parts}', (
		SELECT coalesce(jsonb_agg(o || jsonb_build_object('score', o ->> 'score') ORDER BY ord), '[]')
		FROM jsonb_array_elements(doc -> 'parts') WITH ORDINALITY t(o, ord))) END $$;
CREATE TABLE ev (id int PRIMARY KEY, dim_a_id int, deleted_at timestamptz, doc jsonb) WITH (autovacuum_enabled = false);
INSERT INTO ev SELECT i, 1 + (i * 7) % 60, CASE WHEN i % 17 = 0 THEN now() END, op_doc(i) FROM generate_series(1, 6000) i;
CREATE INDEX ev_col ON ev USING sift_col (id, dim_a_id, deleted_at, (op_txt(doc)))
	WITH (paths = 'op_txt.lang astext, op_txt.parts[]{category astext, dir_a astext, score astext}');
NOTICE:  sift_col index "ev_col": 1 automatic choice (track_changes)
DETAIL:  track_changes = id
HINT:  Reasons, costs and how to turn each off: SELECT * FROM sift.index_choices('ev_col');
VACUUM (FREEZE, ANALYZE) ev;
ANALYZE ents;
-- the shape of the plan: the number of SiftAgg nodes and the aggregates one pushes
CREATE FUNCTION plan_of(q text) RETURNS text LANGUAGE plpgsql AS $$
DECLARE l text; n int := 0; r text := '';
BEGIN
	FOR l IN EXECUTE 'EXPLAIN (COSTS OFF) ' || q LOOP
		IF l LIKE '%Custom Scan (SiftAgg)%' THEN n := n + 1; END IF;
		IF l LIKE '%Pushed Aggregates:%' THEN r := r || trim(l) || '; '; END IF;
	END LOOP;
	RETURN CASE n WHEN 0 THEN 'normal' WHEN 1 THEN 'SiftAgg' ELSE n || ' SiftAgg' END || CASE WHEN r = '' THEN '' ELSE ' (' || r || ')' END;
END $$;
CREATE FUNCTION same(q text) RETURNS TABLE (plan text, result_rows bigint, vs_normal_plan bigint) LANGUAGE sql AS $$ SELECT * FROM sift_helpers.agg_plan_and_differences(plan_of(q), q) $$;
-- x = ANY (ARRAY (sub-select)): the sub-select runs once and its rows are a lookup table of the leaf; a sub-select that reads the row, is volatile or
-- may raise keeps the normal plan
SELECT * FROM same($$SELECT count(*) AS n, sum(id) AS s FROM ev WHERE deleted_at IS NULL AND dim_a_id = ANY (ARRAY(SELECT id FROM ents WHERE type = 'site'))$$);
                       plan                       | result_rows | vs_normal_plan 
--------------------------------------------------+-------------+----------------
 SiftAgg (Pushed Aggregates: count(*), sum(id); ) |           1 |              0
(1 row)

SELECT * FROM same($$SELECT dim_a_id % 3 AS g, count(*) AS n FROM ev WHERE dim_a_id = ANY (ARRAY(SELECT id FROM ents WHERE type IN ('cat_a', 'cat_b'))) GROUP BY 1$$);
                  plan                   | result_rows | vs_normal_plan 
-----------------------------------------+-------------+----------------
 SiftAgg (Pushed Aggregates: count(*); ) |           3 |              0
(1 row)

SELECT * FROM same($$SELECT count(*) AS n FROM ev WHERE dim_a_id = ANY (ARRAY(SELECT id FROM ents WHERE type = 'none'))$$);
                  plan                   | result_rows | vs_normal_plan 
-----------------------------------------+-------------+----------------
 SiftAgg (Pushed Aggregates: count(*); ) |           1 |              0
(1 row)

SELECT * FROM same($$SELECT count(*) AS n FROM ev WHERE dim_a_id = ANY (ARRAY(SELECT id FROM ents WHERE id = ev.id % 60))$$);
  plan  | result_rows | vs_normal_plan 
--------+-------------+----------------
 normal |           1 |              0
(1 row)

SELECT * FROM same($$SELECT count(*) AS n FROM ev WHERE dim_a_id = ANY (ARRAY(SELECT id FROM ents WHERE random() < 2))$$);
  plan  | result_rows | vs_normal_plan 
--------+-------------+----------------
 normal |           1 |              0
(1 row)

SELECT plan_of($$SELECT count(*) AS n FROM ev WHERE dim_a_id = ANY (ARRAY(SELECT 60 / (id - 7) FROM ents))$$) AS raising_division;
 raising_division 
------------------
 normal
(1 row)

-- EXISTS over the parts of an index expression, with the members of one element and the other elements not counting
SELECT * FROM same($$SELECT count(*) FILTER (WHERE EXISTS (SELECT 1 FROM jsonb_to_recordset(op_txt(doc) -> 'parts') AS o(category text, dir_a text)
	WHERE o.category = 'catg_j' AND o.dir_a = 'down')) AS catg_j_negative, count(*) AS n FROM ev$$);
                       plan                        | result_rows | vs_normal_plan 
---------------------------------------------------+-------------+----------------
 SiftAgg (Pushed Aggregates: count(*), count(*); ) |           1 |              0
(1 row)

SELECT * FROM same($$SELECT op_txt(doc) ->> 'lang' AS lang, count(*) AS n FROM ev
	WHERE EXISTS (SELECT 1 FROM jsonb_to_recordset(op_txt(doc) -> 'parts') AS o(category text) WHERE o.category IN ('catg_m', 'catg_h')) GROUP BY 1$$);
                  plan                   | result_rows | vs_normal_plan 
-----------------------------------------+-------------+----------------
 SiftAgg (Pushed Aggregates: count(*); ) |           3 |              0
(1 row)

SELECT * FROM same($$SELECT count(*) AS n FROM ev WHERE NOT EXISTS (SELECT 1 FROM jsonb_to_recordset(op_txt(doc) -> 'parts') AS o(category text) WHERE o.category = 'catg_d')$$);
                  plan                   | result_rows | vs_normal_plan 
-----------------------------------------+-------------+----------------
 SiftAgg (Pushed Aggregates: count(*); ) |           1 |              0
(1 row)

-- nested jsonpath tests over the parts
SELECT * FROM same($$SELECT count(*) AS n FROM ev WHERE jsonb_path_exists(op_txt(doc), '$.parts[*] ? (@.category == "catg_j" && @.dir_a == "down")')$$);
                  plan                   | result_rows | vs_normal_plan 
-----------------------------------------+-------------+----------------
 SiftAgg (Pushed Aggregates: count(*); ) |           1 |              0
(1 row)

SELECT * FROM same($$SELECT count(*) AS n FROM ev WHERE op_txt(doc) @? '$.parts[*] ? (@.dir_a != "flat")'$$);
  plan  | result_rows | vs_normal_plan 
--------+-------------+----------------
 normal |           1 |              0
(1 row)

-- a score is text in the leaf: a number literal is not the text of a string score, so the comparison keeps the normal plan
SELECT plan_of($$SELECT count(*) AS n FROM ev WHERE jsonb_path_exists(op_txt(doc), '$.parts[*] ? (@.score == 7)')$$) AS number_against_text_leaf;
 number_against_text_leaf 
--------------------------
 normal
(1 row)

-- a strict path raises for a document that has no such key: the heap raises, the scan must not answer instead
SELECT plan_of($$SELECT count(*) AS n FROM ev WHERE jsonb_path_exists(op_txt(doc), 'strict $.parts[*] ? (@.category == "catg_d")')$$) AS strict_path;
 strict_path 
-------------
 normal
(1 row)

-- per-item arms under an outer aggregate over a text leaf with the cast in the aggregate (guarded: a score may be text): one fused scan
SELECT * FROM same($$WITH b AS (SELECT r.id AS rid, i.category AS k, CASE WHEN i.score ~ '^[0-9.]+$' THEN i.score::float8 END AS n
	FROM ev r, jsonb_to_recordset(op_txt(r.doc) -> 'parts') AS i(category text, score text) WHERE r.deleted_at IS NULL)
	SELECT k, avg(a) AS value, count(DISTINCT rid) AS items, round(sum(s)::numeric, 6) AS total
	FROM (SELECT rid, k, avg(n) AS a, sum(n) AS s FROM b GROUP BY rid, k UNION ALL SELECT rid, NULL::text, avg(n), sum(n) FROM b GROUP BY rid) x GROUP BY k$$);
                                    plan                                    | result_rows | vs_normal_plan 
----------------------------------------------------------------------------+-------------+----------------
 SiftAgg (Pushed Aggregates: avg(inner), count(DISTINCT row), sum(inner); ) |           5 |              0
(1 row)

-- the arms share a WHERE that keeps the cast off the text scores: it is the scan's own, and the scan casts only what passes it
SELECT * FROM same($$WITH b AS (SELECT r.id AS rid, i.category AS k, i.score AS n
	FROM ev r, jsonb_to_recordset(op_txt(r.doc) -> 'parts') AS i(category text, score text) WHERE r.deleted_at IS NULL)
	SELECT k, avg(a) AS value, count(DISTINCT rid) AS items
	FROM (SELECT rid, k, avg(n::float8) AS a FROM b WHERE n ~ '^[0-9.]+$' GROUP BY rid, k UNION ALL SELECT rid, NULL::text, avg(n::float8) FROM b WHERE n ~ '^[0-9.]+$' GROUP BY rid) x GROUP BY k$$);
                              plan                              | result_rows | vs_normal_plan 
----------------------------------------------------------------+-------------+----------------
 SiftAgg (Pushed Aggregates: avg(inner), count(DISTINCT row); ) |           5 |              0
(1 row)

-- a WHERE of each arm's own (they differ) would have the fused scan cast every item, the text scores too: not fused (an arm may be its own scan)
SELECT * FROM same($$WITH b AS (SELECT r.id AS rid, i.category AS k, i.score AS n
	FROM ev r, jsonb_to_recordset(op_txt(r.doc) -> 'parts') AS i(category text, score text) WHERE r.deleted_at IS NULL)
	SELECT k, avg(a) AS value, count(DISTINCT rid) AS items
	FROM (SELECT rid, k, avg(n::float8) AS a FROM b WHERE n ~ '^[0-9.]+$' GROUP BY rid, k
		  UNION ALL SELECT rid, NULL::text, avg(n::float8) FROM b WHERE n ~ '^[0-9.]+$' AND k <> 'catg_m' GROUP BY rid) x GROUP BY k$$);
                                                          plan                                                          | result_rows | vs_normal_plan 
------------------------------------------------------------------------------------------------------------------------+-------------+----------------
 2 SiftAgg (Pushed Aggregates: avg((i.score)::double precision); Pushed Aggregates: avg((i.score)::double precision); ) |           5 |              0
(1 row)

-- the plain cast raises for a text score, in both plans
SELECT * FROM same($$WITH b AS (SELECT r.id AS rid, i.category AS k, i.score AS n
	FROM ev r, jsonb_to_recordset(op_txt(r.doc) -> 'parts') AS i(category text, score text) WHERE r.deleted_at IS NULL)
	SELECT k, avg(a) AS value, count(DISTINCT rid) AS items
	FROM (SELECT rid, k, avg(n::float8) AS a FROM b GROUP BY rid, k UNION ALL SELECT rid, NULL::text, avg(n::float8) FROM b GROUP BY rid) x GROUP BY k$$);
ERROR:  invalid input syntax for type double precision: "n/a"
CONTEXT:  SQL statement "CREATE TEMP TABLE r_on ON COMMIT DROP AS WITH b AS (SELECT r.id AS rid, i.category AS k, i.score AS n
	FROM ev r, jsonb_to_recordset(op_txt(r.doc) -> 'parts') AS i(category text, score text) WHERE r.deleted_at IS NULL)
	SELECT k, avg(a) AS value, count(DISTINCT rid) AS items
	FROM (SELECT rid, k, avg(n::float8) AS a FROM b GROUP BY rid, k UNION ALL SELECT rid, NULL::text, avg(n::float8) FROM b GROUP BY rid) x GROUP BY k"
PL/pgSQL function sift_helpers.agg_plan_and_differences(text,text) line 4 at EXECUTE
SQL function "same" statement 1
SET sift.enable_agg = off;
SELECT count(*) AS n FROM ev r, jsonb_to_recordset(op_txt(r.doc) -> 'parts') AS i(category text, score text) WHERE i.score::float8 > 5;
ERROR:  invalid input syntax for type double precision: "n/a"
RESET sift.enable_agg;
DROP TABLE ev;
DROP TABLE ents;
DROP FUNCTION same(text), plan_of(text), op_txt(jsonb), op_doc(int);
DROP EXTENSION pg_sift_helpers;
DROP EXTENSION pg_sift;