-- `WHERE rn <= N` over `row_number() OVER (PARTITION BY k ORDER BY o)` is the first N rows of each group: the scan runs it as the group's arrays of each
-- column, ordered by o, cut at N and unnested with their ordinal (rn is the ordinal), where the window sorts every row. The answers are the normal plan's where
-- the ORDER BY has no ties; with ties the rows kept are N of the tie, in the window and in the arrays alike. Anything else stays as written.
CREATE EXTENSION pg_sift;
SET max_parallel_workers_per_gather = 0;
SET seq_page_cost = 1000;
SET random_page_cost = 1000;
CREATE TABLE events (id int PRIMARY KEY, grp int NOT NULL, sub int, ts int, amount int NOT NULL, tie int NOT NULL, txt text, num numeric(10,2), tags int[]) WITH (autovacuum_enabled = false);
INSERT INTO events SELECT i, i % 37, CASE WHEN i % 5 = 0 THEN NULL ELSE i % 3 END, CASE WHEN i % 7 = 0 THEN NULL ELSE (i * 31) % 1000 END, i % 11, i % 4,
       CASE WHEN i % 6 = 0 THEN NULL ELSE 't' || i % 13 END, CASE WHEN i % 9 = 0 THEN NULL ELSE i / 7.0 END, ARRAY[i, i + 1]
FROM generate_series(1, 5000) i;
CREATE INDEX events_col ON events USING sift_col (id, grp, sub, ts, amount, tie, txt, num);
NOTICE:  sift_col index "events_col": 2 automatic choices (rollup_columns, track_changes)
DETAIL:  rollup_columns = sub, tie; track_changes = id
HINT:  Reasons, costs and how to turn each off: SELECT * FROM sift.index_choices('events_col');
CREATE TABLE plain_events (id int PRIMARY KEY, grp int NOT NULL, ts int) WITH (autovacuum_enabled = false);
INSERT INTO plain_events SELECT i, i % 37, i FROM generate_series(1, 500) i;
VACUUM (FREEZE, ANALYZE) events;
VACUUM (FREEZE, ANALYZE) plain_events;
\ir test/agg_helpers.sql
\set ECHO none
CREATE FUNCTION check_query(name text, q text) RETURNS TABLE (label text, in_scan bool, same bool, rows bigint) LANGUAGE plpgsql AS $$
BEGIN
	RETURN QUERY SELECT name, runs_in_scan(q), outcome(q, false) = outcome(q, true), (SELECT count(*) FROM jsonb_array_elements(outcome(q, true)::jsonb));
END $$;
CREATE FUNCTION window_why(q text) RETURNS text LANGUAGE sql AS $$
	SELECT string_agg(DISTINCT substring(reason from '\(([^()]*)\)$'), '; ') FROM sift.explain_why(q) WHERE slug = 'window_first'
$$;
-- the bound, in its forms; rn is read: it is 1 to N
SELECT * FROM check_query('rn <= 3', $q$ SELECT * FROM (SELECT id, grp, amount, row_number() OVER (PARTITION BY grp ORDER BY ts DESC NULLS LAST, id) rn FROM events) t WHERE rn <= 3 $q$);
  label  | in_scan | same | rows 
---------+---------+------+------
 rn <= 3 | t       | t    |  111
(1 row)

SELECT * FROM check_query('rn < 4', $q$ SELECT * FROM (SELECT id, grp, row_number() OVER (PARTITION BY grp ORDER BY ts DESC NULLS LAST, id) rn FROM events) t WHERE rn < 4 $q$);
 label  | in_scan | same | rows 
--------+---------+------+------
 rn < 4 | t       | t    |  111
(1 row)

SELECT * FROM check_query('3 >= rn', $q$ SELECT * FROM (SELECT id, grp, row_number() OVER (PARTITION BY grp ORDER BY ts DESC NULLS LAST, id) rn FROM events) t WHERE 3 >= rn $q$);
  label  | in_scan | same | rows 
---------+---------+------+------
 3 >= rn | t       | t    |  111
(1 row)

SELECT * FROM check_query('4 > rn', $q$ SELECT * FROM (SELECT id, grp, row_number() OVER (PARTITION BY grp ORDER BY ts DESC NULLS LAST, id) rn FROM events) t WHERE 4 > rn $q$);
 label  | in_scan | same | rows 
--------+---------+------+------
 4 > rn | t       | t    |  111
(1 row)

SELECT * FROM check_query('rn <= 2::bigint AND rn <= 5', $q$ SELECT * FROM (SELECT id, grp, row_number() OVER (PARTITION BY grp ORDER BY ts DESC NULLS LAST, id) rn FROM events) t WHERE rn <= 2::bigint AND rn <= 5 $q$);
            label            | in_scan | same | rows 
-----------------------------+---------+------+------
 rn <= 2::bigint AND rn <= 5 | t       | t    |   74
(1 row)

SELECT * FROM check_query('N past the size of the groups', $q$ SELECT * FROM (SELECT id, grp, row_number() OVER (PARTITION BY grp ORDER BY id) rn FROM events) t WHERE rn <= 100000 $q$);
             label             | in_scan | same | rows 
-------------------------------+---------+------+------
 N past the size of the groups | t       | t    | 5000
(1 row)

SELECT * FROM check_query('rn = 1 AND rn <= 3 is the first row', $q$ SELECT * FROM (SELECT id, grp, row_number() OVER (PARTITION BY grp ORDER BY id) rn FROM events) t WHERE rn = 1 AND rn <= 3 $q$);
                label                | in_scan | same | rows 
-------------------------------------+---------+------+------
 rn = 1 AND rn <= 3 is the first row | t       | t    |   37
(1 row)

SELECT * FROM check_query('rn read in expressions and ordered by', $q$
 SELECT id, grp, rn, rn * 10 AS tenth FROM (SELECT id, grp, row_number() OVER (PARTITION BY grp ORDER BY id DESC) rn FROM events) t WHERE rn <= 4 AND rn <> 3 ORDER BY rn DESC, grp LIMIT 20 $q$);
                 label                 | in_scan | same | rows 
---------------------------------------+---------+------+------
 rn read in expressions and ordered by | t       | t    |   20
(1 row)

-- the columns the arrays hold: text, numeric, NULLs, expressions, a column twice
SELECT * FROM check_query('columns of several types', $q$ SELECT * FROM (SELECT id, grp, txt, num, ts, amount * 2 AS twice, txt || '!' AS bang, row_number() OVER (PARTITION BY grp ORDER BY id) rn FROM events) t WHERE rn <= 3 $q$);
          label           | in_scan | same | rows 
--------------------------+---------+------+------
 columns of several types | t       | t    |  111
(1 row)

SELECT * FROM check_query('a column twice', $q$ SELECT * FROM (SELECT id, id AS again, grp, row_number() OVER (PARTITION BY grp ORDER BY id) rn FROM events) t WHERE rn <= 3 $q$);
     label      | in_scan | same | rows 
----------------+---------+------+------
 a column twice | t       | t    |  111
(1 row)

-- the window's shapes: NULL partition keys, two keys, orders with NULLS FIRST / LAST and DESC, an order key that is a partition key, a named window, a condition of the sub-select
SELECT * FROM check_query('NULL partition key', $q$ SELECT * FROM (SELECT id, sub, row_number() OVER (PARTITION BY sub ORDER BY id DESC) rn FROM events) t WHERE rn <= 4 $q$);
       label        | in_scan | same | rows 
--------------------+---------+------+------
 NULL partition key | t       | t    |   16
(1 row)

SELECT * FROM check_query('two partition keys', $q$ SELECT * FROM (SELECT id, grp, sub, row_number() OVER (PARTITION BY grp, sub ORDER BY id DESC) rn FROM events) t WHERE rn <= 2 $q$);
       label        | in_scan | same | rows 
--------------------+---------+------+------
 two partition keys | t       | t    |  296
(1 row)

SELECT o AS "order", (check_query('order ' || o, format($q$ SELECT * FROM (SELECT id, sub, ts, row_number() OVER (PARTITION BY sub ORDER BY ts %s, id) rn FROM events) t WHERE rn <= 3 $q$, o))).*
FROM unnest(ARRAY['', 'NULLS FIRST', 'DESC', 'DESC NULLS LAST']) o;
      order      |         label         | in_scan | same | rows 
-----------------+-----------------------+---------+------+------
                 | order                 | t       | t    |   12
 NULLS FIRST     | order NULLS FIRST     | t       | t    |   12
 DESC            | order DESC            | t       | t    |   12
 DESC NULLS LAST | order DESC NULLS LAST | t       | t    |   12
(4 rows)

SELECT * FROM check_query('two order keys', $q$ SELECT * FROM (SELECT id, grp, row_number() OVER (PARTITION BY grp ORDER BY amount DESC, tie, id) rn FROM events) t WHERE rn <= 3 $q$);
     label      | in_scan | same | rows 
----------------+---------+------+------
 two order keys | t       | t    |  111
(1 row)

SELECT * FROM check_query('order key is a partition key', $q$ SELECT * FROM (SELECT id, grp, row_number() OVER (PARTITION BY grp ORDER BY grp, id) rn FROM events) t WHERE rn <= 3 $q$);
            label             | in_scan | same | rows 
------------------------------+---------+------+------
 order key is a partition key | t       | t    |  111
(1 row)

SELECT * FROM check_query('a named window', $q$ SELECT * FROM (SELECT id, grp, row_number() OVER w rn FROM events WINDOW w AS (PARTITION BY grp ORDER BY id DESC)) t WHERE rn <= 3 $q$);
     label      | in_scan | same | rows 
----------------+---------+------+------
 a named window | t       | t    |  111
(1 row)

SELECT * FROM check_query('a condition of the sub-select and of the outer query', $q$
 SELECT * FROM (SELECT id, grp, amount, row_number() OVER (PARTITION BY grp ORDER BY amount DESC, id) rn FROM events WHERE ts > 100) t WHERE rn <= 3 AND amount > 5 $q$);
                        label                         | in_scan | same | rows 
------------------------------------------------------+---------+------+------
 a condition of the sub-select and of the outer query | t       | t    |  111
(1 row)

SELECT * FROM check_query('CTE', $q$ WITH ranked AS (SELECT id, grp, row_number() OVER (PARTITION BY grp ORDER BY id DESC) rn FROM events) SELECT id, grp, rn FROM ranked WHERE rn <= 2 $q$);
 label | in_scan | same | rows 
-------+---------+------+------
 CTE   | t       | t    |   74
(1 row)

SELECT * FROM check_query('under a join and an aggregate', $q$
 SELECT t.grp, count(*) n, sum(t.amount) s FROM (SELECT id, grp, amount, row_number() OVER (PARTITION BY grp ORDER BY id DESC) rn FROM events) t JOIN events e ON e.id = t.id WHERE t.rn <= 5 GROUP BY t.grp $q$);
             label             | in_scan | same | rows 
-------------------------------+---------+------+------
 under a join and an aggregate | t       | t    |   37
(1 row)

-- the result keeps the type modifier of the column (an array element has none)
SET enable_seqscan = off;
SET enable_indexscan = off;
SET enable_bitmapscan = off;
CREATE TEMP TABLE kept AS SELECT * FROM (SELECT id, num, row_number() OVER (PARTITION BY grp ORDER BY id) rn FROM events) t WHERE rn <= 3;
RESET enable_seqscan;
RESET enable_indexscan;
RESET enable_bitmapscan;
SELECT format_type(atttypid, atttypmod) AS num_type FROM pg_attribute WHERE attrelid = 'kept'::regclass AND attname = 'num';
   num_type    
---------------
 numeric(10,2)
(1 row)

DROP TABLE kept;
-- ties: each group gives the N rows that come first in the order, whichever of the tie; rn is 1 to the number of rows it gave
CREATE FUNCTION first_n_of_tie(q text, n int) RETURNS TABLE (groups bigint, wrong_tie_values bigint, wrong_rn bigint) LANGUAGE plpgsql AS $$
BEGIN
	PERFORM set_config('sift.enable_agg', 'on', true);
	RETURN QUERY EXECUTE format($f$
		WITH got AS (%s),
		want AS (SELECT grp, array_agg(tie ORDER BY tie) FILTER (WHERE r <= %s) AS ties FROM (SELECT grp, tie, row_number() OVER (PARTITION BY grp ORDER BY tie) r FROM events) x GROUP BY grp),
		have AS (SELECT grp, array_agg(tie ORDER BY tie) AS ties, array_agg(rn ORDER BY rn) AS rns FROM got GROUP BY grp)
		SELECT count(*), count(*) FILTER (WHERE have.ties IS DISTINCT FROM want.ties),
		       count(*) FILTER (WHERE have.rns IS DISTINCT FROM (SELECT array_agg(i::bigint) FROM generate_series(1, cardinality(want.ties)) i))
		FROM have JOIN want USING (grp)$f$, q, n);
END $$;
SELECT * FROM first_n_of_tie($q$ SELECT grp, id, tie, rn FROM (SELECT id, grp, tie, row_number() OVER (PARTITION BY grp ORDER BY tie) rn FROM events) t WHERE rn <= 3 $q$, 3);
 groups | wrong_tie_values | wrong_rn 
--------+------------------+----------
     37 |                0 |        0
(1 row)

SELECT * FROM first_n_of_tie($q$ SELECT grp, id, tie, rn FROM (SELECT id, grp, tie, row_number() OVER (PARTITION BY grp ORDER BY tie) rn FROM events) t WHERE rn <= 40 $q$, 40);
 groups | wrong_tie_values | wrong_rn 
--------+------------------+----------
     37 |                0 |        0
(1 row)

SELECT runs_in_scan($q$ SELECT grp, id, tie, rn FROM (SELECT id, grp, tie, row_number() OVER (PARTITION BY grp ORDER BY tie) rn FROM events) t WHERE rn <= 3 $q$);
 runs_in_scan 
--------------
 t
(1 row)

-- as written: another window function, nothing to cut but the partition and rn, a column that is an array, a table with no column index, a grouped sub-select
SELECT * FROM check_query('another window function', $q$ SELECT * FROM (SELECT id, grp, row_number() OVER (PARTITION BY grp ORDER BY id DESC) rn, count(*) OVER (PARTITION BY grp) n FROM events) t WHERE rn <= 3 $q$);
          label          | in_scan | same | rows 
-------------------------+---------+------+------
 another window function | f       | t    |  111
(1 row)

SELECT * FROM check_query('only the partition and rn', $q$ SELECT * FROM (SELECT grp, row_number() OVER (PARTITION BY grp ORDER BY id DESC) rn FROM events) t WHERE rn <= 3 $q$);
           label           | in_scan | same | rows 
---------------------------+---------+------+------
 only the partition and rn | f       | t    |  111
(1 row)

SELECT window_why($q$ SELECT * FROM (SELECT grp, row_number() OVER (PARTITION BY grp ORDER BY id DESC) rn FROM events) t WHERE rn <= 3 $q$);
                            window_why                             
-------------------------------------------------------------------
 the sub-select reads nothing but the partition and the row number
(1 row)

SELECT * FROM check_query('an array column', $q$ SELECT * FROM (SELECT id, grp, tags, row_number() OVER (PARTITION BY grp ORDER BY id DESC) rn FROM events) t WHERE rn <= 3 $q$);
      label      | in_scan | same | rows 
-----------------+---------+------+------
 an array column | f       | t    |  111
(1 row)

SELECT window_why($q$ SELECT * FROM (SELECT id, grp, tags, row_number() OVER (PARTITION BY grp ORDER BY id DESC) rn FROM events) t WHERE rn <= 3 $q$);
                            window_why                            
------------------------------------------------------------------
 a column the sub-select reads is in no column index of its table
(1 row)

SELECT * FROM check_query('no column index', $q$ SELECT * FROM (SELECT id, grp, row_number() OVER (PARTITION BY grp ORDER BY id DESC) rn FROM plain_events) t WHERE rn <= 3 $q$);
      label      | in_scan | same | rows 
-----------------+---------+------+------
 no column index | f       | t    |  111
(1 row)

SELECT * FROM check_query('a grouped sub-select', $q$
 SELECT * FROM (SELECT grp, c, row_number() OVER (PARTITION BY c % 3 ORDER BY c DESC, grp) rn FROM (SELECT grp, count(*) c FROM events GROUP BY grp) g) t WHERE rn <= 3 $q$);
        label         | in_scan | same | rows 
----------------------+---------+------+------
 a grouped sub-select | t       | t    |    6
(1 row)

SELECT * FROM check_query('the second row only', $q$ SELECT * FROM (SELECT id, grp, row_number() OVER (PARTITION BY grp ORDER BY id DESC) rn FROM events) t WHERE rn = 2 $q$);
        label        | in_scan | same | rows 
---------------------+---------+------+------
 the second row only | f       | t    |   37
(1 row)

-- the memory decides it: the scan keeps the first rows of every group in memory, which a sort of the window spills to disk, so the scan that runs for 37 groups
-- of 135 rows is refused for 1,000 groups of 5 when the memory of a hash aggregate is 128 kB
CREATE TABLE ranked (id int PRIMARY KEY, k37 int NOT NULL, k1000 int NOT NULL, ts int NOT NULL, amount int NOT NULL) WITH (autovacuum_enabled = false);
INSERT INTO ranked SELECT i, i % 37, i % 1000, (i * 31) % 1000, i % 11 FROM generate_series(1, 5000) i;
CREATE INDEX ranked_col ON ranked USING sift_col (id, k37, k1000, ts, amount);
NOTICE:  sift_col index "ranked_col": 2 automatic choices (rollup_columns, track_changes)
DETAIL:  rollup_columns = k37, amount; track_changes = id
HINT:  Reasons, costs and how to turn each off: SELECT * FROM sift.index_choices('ranked_col');
VACUUM (FREEZE, ANALYZE) ranked;
CREATE FUNCTION plan_in_scan(q text) RETURNS bool LANGUAGE plpgsql AS $$
DECLARE line text; found bool := false;
BEGIN
	FOR line IN EXECUTE 'EXPLAIN (COSTS OFF) ' || q LOOP found := found OR line LIKE '%Custom Scan (SiftAgg)%'; END LOOP;
	RETURN found;
END $$;
SET seq_page_cost = 64;
SET work_mem = '64kB';
SELECT k AS "partition key", plan_in_scan(format($q$ SELECT * FROM (SELECT id, amount, row_number() OVER (PARTITION BY %s ORDER BY ts DESC, id) rn FROM ranked) t WHERE rn <= 3 $q$, k)) AS in_scan
FROM unnest(ARRAY['k37', 'k1000']) k;
 partition key | in_scan 
---------------+---------
 k37           | t
 k1000         | f
(2 rows)

DROP FUNCTION plan_in_scan(text);
DROP TABLE ranked;
DROP FUNCTION first_n_of_tie(text, int);
DROP TABLE events, plain_events;
DROP FUNCTION window_why(text), check_query(text,text), outcome(text,boolean), runs_in_scan(text), runq(text,boolean), pushed(text);
DROP EXTENSION pg_sift;