\set ECHO none
\pset format unaligned
SET search_path TO provsql_test, provsql;

-- ORDER BY ... LIMIT k over tracked relations keeps, in each possible world,
-- the present rows that fewer than k present rows precede: the filter of
-- row_number() <= k, or of rank() <= k with FETCH ... WITH TIES.  Every row
-- that may be kept is output, annotated with that condition.  LIMIT
-- plain(k) truncates the actual result instead.  Rows 2 and 3 tie on x; y
-- has NULLs.

CREATE TABLE lr_plain(id int, g text, x int, y int, p float8);
INSERT INTO lr_plain VALUES
  (1, 'a', 10, 1, 0.2), (2, 'a', 20, NULL, 0.5), (3, 'a', 20, 3, 0.7),
  (4, 'a', 30, 4, 0.4), (5, 'b', 5, NULL, 0.6), (6, 'b', 7, 6, 0.9);
CREATE TABLE lr AS SELECT id, g, x, y FROM lr_plain;
SELECT add_provenance('lr');
SELECT set_prob(provenance(), p) FROM lr JOIN lr_plain USING (id) \g /dev/null
CREATE SCHEMA lr_world;

-- For each output key k of the query, its probability by ProvSQL over lr, and
-- by running the query over each of the 64 worlds of lr_plain (a view lr of
-- lr_world, first on the search_path, holds the rows of the world).
CREATE FUNCTION lr_check(label text, sql text)
  RETURNS TABLE(q text, exact boolean, probabilities text) LANGUAGE plpgsql AS $f$
DECLARE
  m int;
BEGIN
  CREATE TEMP TABLE lr_tracked(k text, p float8);
  CREATE TEMP TABLE lr_brute(k text, pw float8);
  SET LOCAL search_path TO provsql_test, provsql;
  EXECUTE format('INSERT INTO lr_tracked SELECT k::text, probability_evaluate(provenance())
                  FROM (%s) u', sql);
  FOR m IN 0 .. 63 LOOP
    EXECUTE format('CREATE OR REPLACE VIEW lr_world.lr AS
                    SELECT id, g, x, y FROM provsql_test.lr_plain
                    WHERE %s & (1 << (id - 1)) <> 0', m);
    SET LOCAL search_path TO lr_world, provsql_test, provsql;
    EXECUTE format(
      'INSERT INTO lr_brute
       SELECT k::text, (SELECT exp(sum(ln(CASE WHEN %1$s & (1 << (id - 1)) <> 0
                                         THEN p ELSE 1 - p END)))
                        FROM provsql_test.lr_plain)
       FROM (%2$s) u', m, sql);
    SET LOCAL search_path TO provsql_test, provsql;
  END LOOP;
  RETURN QUERY
    SELECT label,
           bool_and(round(coalesce(t.p, 0)::numeric, 9) =
                    round(coalesce(b.p, 0)::numeric, 9)),
           string_agg(coalesce(t.k, b.k) || '=' ||
                      round(coalesce(t.p, 0)::numeric, 4),
                      ' ' ORDER BY coalesce(t.k, b.k))
    FROM (SELECT t.k, sum(t.p) AS p FROM lr_tracked t GROUP BY t.k) t
    FULL JOIN (SELECT b.k, sum(b.pw) AS p FROM lr_brute b GROUP BY b.k) b
      USING (k);
  DROP TABLE lr_tracked, lr_brute;
END
$f$;

SELECT * FROM lr_check('top 2',
  'SELECT id AS k FROM lr ORDER BY x DESC, id LIMIT 2');
SELECT * FROM lr_check('first',
  'SELECT id AS k FROM lr ORDER BY x, id LIMIT 1');
SELECT * FROM lr_check('OFFSET and LIMIT',
  'SELECT id AS k FROM lr ORDER BY x, id OFFSET 1 LIMIT 2');
SELECT * FROM lr_check('OFFSET alone',
  'SELECT id AS k FROM lr ORDER BY x, id OFFSET 2');
SELECT * FROM lr_check('after a selection, NULLs first',
  'SELECT id AS k FROM lr WHERE g = ''a'' ORDER BY y NULLS FIRST, id LIMIT 2');
SELECT * FROM lr_check('several keys',
  'SELECT id AS k FROM lr ORDER BY g, x DESC, id LIMIT 3');
-- DISTINCT ON (g) keeps the first row of each group: the filter of
-- row_number() OVER (PARTITION BY g ORDER BY ...) <= 1, checked world by
-- world; with a LIMIT on top, and the ORDER BY key of the rows kept.
SELECT * FROM lr_check('DISTINCT ON',
  'SELECT DISTINCT ON (g) id AS k FROM lr ORDER BY g, x DESC, id');
SELECT * FROM lr_check('DISTINCT ON, then LIMIT',
  'SELECT DISTINCT ON (g) id AS k FROM lr ORDER BY g, x, id LIMIT 1');
SELECT * FROM lr_check('DISTINCT ON a key other than the first',
  'SELECT DISTINCT ON (g) g AS k FROM lr WHERE y IS NOT NULL ORDER BY g, id');
SELECT * FROM lr_check('top 1 per group, LATERAL',
  'SELECT gs.g || '':'' || u.id AS k
   FROM (SELECT DISTINCT g FROM provsql_test.lr_plain) gs,
        LATERAL (SELECT id FROM lr w WHERE w.g = gs.g ORDER BY x DESC, id LIMIT 1) u');
SELECT * FROM lr_check('in a CTE, then filtered',
  'WITH top AS (SELECT id, x FROM lr ORDER BY x DESC, id LIMIT 3)
   SELECT id AS k FROM top WHERE x < 30');
SELECT * FROM lr_check('arms of UNION ALL',
  'SELECT id AS k FROM (SELECT id FROM lr ORDER BY x, id LIMIT 1) a
   UNION ALL SELECT id FROM (SELECT id FROM lr ORDER BY x DESC, id LIMIT 1) b');
SELECT * FROM lr_check('arm of UNION',
  '(SELECT id AS k FROM lr ORDER BY x DESC, id LIMIT 1)
   UNION SELECT id FROM lr WHERE g = ''b''');
SELECT * FROM lr_check('two conditions on one count',
  'SELECT g AS k FROM (SELECT g, count(*) AS c FROM lr GROUP BY g) u
   WHERE c > 1 AND c <= 3');
SELECT * FROM lr_check('after a filter on a rank',
  'SELECT id AS k FROM (SELECT id, g, x,
                        row_number() OVER (PARTITION BY g ORDER BY x DESC, id) AS rn
                        FROM lr) t
   WHERE rn = 1 ORDER BY id DESC LIMIT 1');
SELECT * FROM lr_check('a window over rows filtered on a rank',
  'SELECT id AS k FROM (SELECT id, count(*) OVER (ORDER BY id DESC) AS c
                        FROM (SELECT id FROM (SELECT id, g, x,
                              row_number() OVER (PARTITION BY g ORDER BY x DESC, id) AS rn
                              FROM lr) t WHERE rn = 1) f) z
   WHERE c <= 1');
SELECT * FROM lr_check('joined back',
  'SELECT t.id || ''-'' || w.id AS k
   FROM (SELECT id, g FROM lr ORDER BY x DESC, id LIMIT 2) t
   JOIN lr w ON w.g = t.g AND w.id < t.id');

-- The output: every candidate row, its provenance and probability including
-- the condition; SELECT * keeps the columns of the table.
CREATE TABLE lr_top AS
  SELECT *, probability_evaluate(provenance()) AS p
  FROM lr ORDER BY x DESC, id LIMIT 2;
SELECT remove_provenance('lr_top');
SELECT id, g, x, y, round(p::numeric, 4) AS p FROM lr_top ORDER BY id;

-- LIMIT plain(k): the first k rows of the actual result, with the
-- provenance of the full result (the probability of the row itself).
CREATE TABLE lr_actual AS
  SELECT id, probability_evaluate(provenance()) AS p
  FROM lr ORDER BY x DESC, id LIMIT plain(2);
SELECT remove_provenance('lr_actual');
SELECT id, round(p::numeric, 4) AS p FROM lr_actual ORDER BY id;
CREATE TABLE lr_actual2 AS
  SELECT id FROM lr ORDER BY x, id OFFSET plain(1) FETCH FIRST plain(2) ROWS ONLY;
SELECT remove_provenance('lr_actual2');
SELECT id FROM lr_actual2 ORDER BY id;

-- LIMIT over ties: row_number(), tracked as rank(), with a warning.
CREATE TABLE lr_ties AS SELECT id FROM lr ORDER BY x DESC LIMIT 2;
SELECT remove_provenance('lr_ties');
SELECT count(*) AS candidates FROM lr_ties;

-- A subquery with a provenance column, filtered on a column after it.
CREATE TABLE lr_star AS
  SELECT id, rn FROM (SELECT *, row_number() OVER (ORDER BY x DESC, id) AS rn
                      FROM lr) u
  WHERE rn <= 2;
SELECT remove_provenance('lr_star');
SELECT id, rn FROM lr_star ORDER BY id;

-- A top 3 over 40 rows: a COUNT comparison over independent rows,
-- evaluated by the Poisson-binomial pre-pass.
CREATE TABLE lr_many AS
  SELECT i AS id, (i * 7919) % 1000 AS score FROM generate_series(1, 40) i;
SELECT add_provenance('lr_many');
SELECT set_prob(provenance(), 0.3 + (id % 5) / 10.0) FROM lr_many \g /dev/null
SET provsql.verbose_level = 5;
CREATE TABLE lr_many_top AS
  SELECT id, probability_evaluate(provenance()) AS p
  FROM (SELECT id FROM lr_many ORDER BY score DESC LIMIT 3) u
  WHERE id = 7;
-- OFFSET m LIMIT k compares the rank twice: the two comparisons on the one
-- count make a range, evaluated by the same pre-pass.
CREATE TABLE lr_many_range AS
  SELECT id, probability_evaluate(provenance()) AS p
  FROM (SELECT id FROM lr_many ORDER BY score DESC OFFSET 10 LIMIT 5) u
  WHERE id = 7;
RESET provsql.verbose_level;
SELECT remove_provenance('lr_many_top');
SELECT id, round(p::numeric, 12) AS p FROM lr_many_top;
SELECT remove_provenance('lr_many_range');
SELECT id, round(p::numeric, 12) AS p FROM lr_many_range;

DROP TABLE lr_top, lr_actual, lr_actual2, lr_ties, lr_star, lr_many,
  lr_many_top, lr_many_range;

-- FETCH ... WITH TIES (PostgreSQL 13+): rank() <= k.
SELECT current_setting('server_version_num')::int >= 130000 AS pg_has_with_ties
\gset
\if :pg_has_with_ties
SELECT * FROM lr_check('WITH TIES',
  'SELECT id AS k FROM lr ORDER BY x FETCH FIRST 2 ROWS WITH TIES');
SELECT * FROM lr_check('WITH TIES, first',
  'SELECT id AS k FROM lr ORDER BY x DESC FETCH FIRST 1 ROW WITH TIES');
\else
\echo limit_rank: WITH TIES skipped on PostgreSQL < 13
\endif

-- An aggregate over a rank-filtered LIMIT displays the value of the rows
-- the LIMIT keeps, as plain SQL does, not of every candidate row.
CREATE TABLE lr_t AS
  SELECT count(*) AS c, sum(x) AS s
  FROM (SELECT x FROM lr ORDER BY x, id LIMIT 2) t;
SELECT remove_provenance('lr_t');
SELECT 'DISPLAY limit' AS q, c, s FROM lr_t;
DROP TABLE lr_t;
SELECT 'PLAIN limit' AS q, count(*) AS c, sum(x) AS s
FROM (SELECT x FROM lr_plain ORDER BY x, id LIMIT 2) t;

DROP FUNCTION lr_check(text, text);
-- Without ProvSQL's schema in the search_path: the operators the rewriting
-- builds on its own types (the rank as a count plus one, arithmetic on an
-- aggregate) are looked up in that schema.
SET search_path TO provsql_test;
CREATE TABLE lr_nopath AS SELECT id FROM lr ORDER BY x DESC, id LIMIT 2;
CREATE TABLE lr_nopath_agg AS SELECT g, count(*) + 1 AS c FROM lr GROUP BY g;
SELECT provsql.remove_provenance('lr_nopath');
SELECT provsql.remove_provenance('lr_nopath_agg');
SELECT count(*) AS candidates FROM lr_nopath;
SELECT g, c::text AS c, pg_typeof(c) AS type FROM lr_nopath_agg ORDER BY g;
DROP TABLE lr_nopath, lr_nopath_agg;
SET search_path TO provsql_test, provsql;

DROP VIEW lr_world.lr;
DROP SCHEMA lr_world;
DROP TABLE lr, lr_plain;


-- The top-k of an aggregation whose sort key is NOT a bare aggregate of the
-- block: an aggregate column of a derived table, and a scalar subquery that
-- counts -- the way the query is usually written (103 such queries in the
-- SEDE corpus, prevalence-ac).  Both are the filter of a rank like the bare
-- form, so the answer holds every row that is in the top k in SOME world, each
-- with the provenance of being there.
--
-- What blocked them: the lowering wraps the query it ranks in a pass-through
-- subquery, and the rank-over-an-aggregate rewrite looked for the aggregation
-- directly below the window, so the wrapper hid it.
--
-- Over four rows at one half, two of them of the first group: the first group
-- is in the top two of every world where it has a row (3/4), the others in
-- every world where theirs is there (1/2).
CREATE TABLE lr_v(id int, userid int);
INSERT INTO lr_v VALUES (1,1),(2,1),(3,2),(4,3);
SELECT add_provenance('lr_v');
DO $$ BEGIN PERFORM set_prob(provenance(), 0.5) FROM lr_v; END $$;
CREATE TABLE lr_top AS
  SELECT userid, round(probability_evaluate(provenance())::numeric, 6) AS p
  FROM (SELECT userid, count(*) AS n FROM lr_v GROUP BY userid) t
  ORDER BY n DESC LIMIT 2;
SELECT remove_provenance('lr_top');
SELECT userid, p FROM lr_top ORDER BY userid;
DROP TABLE lr_top;
CREATE TABLE lr_top AS
  SELECT userid,
         round(probability_evaluate(provenance())::numeric, 6) AS p
  FROM (SELECT DISTINCT userid FROM lr_v) u
  ORDER BY (SELECT count(*) FROM lr_v WHERE lr_v.userid = u.userid) DESC
  LIMIT 2;
SELECT remove_provenance('lr_top');
SELECT userid, p FROM lr_top ORDER BY userid;
DROP TABLE lr_top;
-- An ORDER BY of SEVERAL keys, the aggregate one and a tie-breaker beside it,
-- which is how a top-k is written where two groups can have the same count:
-- the rank is the lexicographic one, so a group that ties on the count is
-- before or after by the second key rather than tying.  Two rows in the first
-- group and one in the second, each present with probability one half: the
-- first group is in the top one of every world where it has a row (3/4), and
-- the second only in the world where its row is there and neither of the
-- others is (1/8) -- where the count alone would leave it the rank of a tie in
-- the two worlds where the counts are equal (3/8).  And because the keys hold
-- the grouping column, no two groups can tie on all of them: row_number() is
-- the rank it is tracked as, so the warning that they may differ is not
-- raised.
CREATE TABLE lr_tie(g int, d int);
INSERT INTO lr_tie VALUES (1,10),(1,10),(2,20);
SELECT add_provenance('lr_tie');
DO $$ BEGIN PERFORM set_prob(provenance(), 0.5) FROM lr_tie; END $$;
CREATE TABLE lr_top AS
  SELECT g, round(probability_evaluate(provenance())::numeric, 6) AS p
  FROM (SELECT g, count(*) AS n FROM lr_tie GROUP BY g) t
  ORDER BY n DESC, g ASC LIMIT 1;
SELECT remove_provenance('lr_top');
SELECT g, p FROM lr_top ORDER BY g;
DROP TABLE lr_top;
-- The same without the tie-breaker, where the tie stands.
CREATE TABLE lr_top AS
  SELECT g, round(probability_evaluate(provenance())::numeric, 6) AS p
  FROM (SELECT g, count(*) AS n FROM lr_tie GROUP BY g) t
  ORDER BY n DESC LIMIT 1;
SELECT remove_provenance('lr_top');
SELECT g, p FROM lr_top ORDER BY g;
DROP TABLE lr_top;
-- A sort key that reads an aggregate of a DERIVED TABLE while the block that
-- truncates groups as well: on PostgreSQL 18 such a key is a Var of the
-- virtual RTE_GROUP entry holding the grouping expressions, not of the derived
-- table, so following it is what tells this apart from a cut of the actual
-- result (prevalence-ac measured this shape refused on the corpus, where it is
-- "GROUP BY b.badgecount ORDER BY b.badgecount DESC LIMIT 10" over a derived
-- table that counts).  Two rows of the first key and one of the second, each
-- present with probability one half, top-1 by the derived count: the first key
-- holds it in every world where it has a row, 3/4, while the second holds it
-- in the worlds where its row is there and the first has at most one, 3/8.  On
-- PostgreSQL 18 the group-by on that count explodes the first key into one row
-- per value the count takes, 1/4 for the two rows and 1/2 for one of them;
-- earlier versions read the key through the derived table and keep one row,
-- with a warning of their own.  The two are the same rows in every world, which
-- is what is summed per key here.
CREATE TABLE lr_b(k int, v int);
INSERT INTO lr_b VALUES (1,1),(1,2),(2,3);
SELECT add_provenance('lr_b');
DO $$ BEGIN PERFORM set_prob(provenance(), 0.5) FROM lr_b; END $$;
SET client_min_messages = error;
CREATE TABLE lr_top AS
  SELECT b.k, round(probability_evaluate(provenance())::numeric, 6) AS p
  FROM (SELECT k, count(*) AS c FROM lr_b GROUP BY k) b
  GROUP BY b.k, b.c ORDER BY b.c DESC LIMIT 1;
RESET client_min_messages;
SELECT remove_provenance('lr_top');
SELECT k, sum(p) AS p FROM lr_top GROUP BY k ORDER BY k;
DROP TABLE lr_top;
SELECT remove_provenance('lr_b');
DROP TABLE lr_b;
SELECT remove_provenance('lr_tie');
DROP TABLE lr_tie;
SELECT remove_provenance('lr_v');
DROP TABLE lr_v;

-- ORDER BY ... LIMIT k over a block of at most one row -- a scalar
-- aggregation, or a product of such (TPC-DS 16, 90, 92, 94, 96) -- neither
-- orders nor truncates in any world: it is dropped, so the values are tracked
-- and no window the query does not have is reported.
CREATE TABLE lr_one(k int, x int);
INSERT INTO lr_one VALUES (1, 1), (1, 2), (2, 3);
SELECT add_provenance('lr_one');
CREATE TABLE lr_one_r AS
  SELECT count(DISTINCT x) AS n, sum(x) AS s FROM lr_one
  ORDER BY count(DISTINCT x) LIMIT 100;
SELECT remove_provenance('lr_one_r');
SELECT n::text AS n, s::text AS s FROM lr_one_r;
DROP TABLE lr_one_r;
CREATE TABLE lr_one_r AS
  SELECT a::numeric / b::numeric AS r
  FROM (SELECT count(*) AS a FROM lr_one) x,
       (SELECT count(*) AS b FROM lr_one WHERE k = 1) y
  ORDER BY r LIMIT 100;
SELECT remove_provenance('lr_one_r');
SELECT round(r::text::numeric, 4) AS r FROM lr_one_r;
DROP TABLE lr_one_r;
-- The same with no ORDER BY: which rows a LIMIT keeps is left open by SQL only
-- where there are more rows than it keeps, so no warning here.
CREATE TABLE lr_one_r AS SELECT count(x) AS n FROM lr_one LIMIT 10;
SELECT remove_provenance('lr_one_r');
SELECT n::text AS n FROM lr_one_r;
DROP TABLE lr_one_r;
SELECT remove_provenance('lr_one');
DROP TABLE lr_one;

-- A LIMIT over a block another LIMIT already bounds by no more rows cuts
-- nothing either: the middle LIMIT 3 below is dropped, and the other two are
-- read in every world.  Rows at one half; the probabilities are those of the
-- enumeration of the 32 worlds.
CREATE TABLE lr_b(id int);
INSERT INTO lr_b VALUES (1), (2), (3), (4), (5);
SELECT add_provenance('lr_b');
DO $$ BEGIN PERFORM set_prob(provenance(), 0.5) FROM lr_b; END $$;
CREATE TABLE lr_b_r AS
  SELECT id, round(probability_evaluate(provenance())::numeric, 6) AS p
  FROM (SELECT * FROM (SELECT * FROM lr_b ORDER BY id LIMIT 3) a
        ORDER BY id DESC LIMIT 3) b
  ORDER BY id LIMIT 2;
SELECT remove_provenance('lr_b_r');
SELECT id, p FROM lr_b_r ORDER BY id;
DROP TABLE lr_b_r;
SELECT remove_provenance('lr_b');
DROP TABLE lr_b;

-- A key the WHERE sets to constants bounds its relation at one row in every
-- world, a world being a part of the instance: LIMIT 1 then cuts nothing
-- (what an ORM writes for the first row), and no warning is due.  A key that
-- is only partly set does not, and the LIMIT stays a truncation.
CREATE TABLE lr_k(a int, b int, v int, PRIMARY KEY (a, b));
INSERT INTO lr_k VALUES (1, 1, 10), (1, 2, 20);
SELECT add_provenance('lr_k');
CREATE TABLE lr_k_r AS SELECT v FROM lr_k WHERE a = 1 AND b = 2 LIMIT 1;
SELECT remove_provenance('lr_k_r');
SELECT v FROM lr_k_r;
DROP TABLE lr_k_r;
CREATE TABLE lr_k_r AS SELECT v FROM lr_k WHERE a = 1 LIMIT 1;
SELECT remove_provenance('lr_k_r');
SELECT count(*) AS n FROM lr_k_r;
DROP TABLE lr_k_r;
SELECT remove_provenance('lr_k');
DROP TABLE lr_k;