-- lionvec, the nearest-neighbour index of pgvector's `vector` (DESIGN.md §44).
--
-- Every ORDER BY below is run through the index (the sequential scan and
-- the sort disabled, so that nothing else can be chosen) and through a
-- sequential scan and sort (the index disabled), and the two answers are
-- compared: the distances in order, and the rows.  pg_lion.vector_epsilon
-- is raised for the comparisons, so that no bound failure can make a row
-- come back late (§44.4), and pg_lion.vector_margin, so that every
-- partition that could hold a nearer row is read, and pg_lion.vector_patience
-- is off, so that the scan never pauses before it has; the default epsilon's
-- own rate is the unit tests', and the default margin's and patience's
-- recall is measured at the end.
--
-- Every index spills by default (DESIGN.md §44.3): each row with a vector
-- is in two partitions, so the comparisons are also of a scan that must
-- return each row once.  vt_spill() reads the CODE items off the pages to
-- show where the copies are.
\set VERBOSITY terse
SET client_min_messages = warning;
CREATE EXTENSION IF NOT EXISTS pg_lion;
CREATE EXTENSION IF NOT EXISTS vector;
CREATE EXTENSION IF NOT EXISTS pg_lion_vector;
CREATE EXTENSION IF NOT EXISTS pageinspect;
RESET client_min_messages;
SET synchronous_commit = on;
SET max_parallel_workers_per_gather = 0;

-- every operator class validates
SELECT opcname, amvalidate(c.oid)
  FROM pg_opclass c JOIN pg_am a ON a.oid = c.opcmethod
 WHERE a.amname = 'lionvec' ORDER BY opcname;

-- the build's phases, as pg_stat_progress_create_index names them (DESIGN.md
-- §44.5); what a build reports in them only another session could see, by a
-- race, so it is left alone, as core's tests leave theirs
SELECT n, pg_indexam_progress_phasename(a.oid, n) AS phase
  FROM pg_am a, generate_series(1, 8) n
 WHERE a.amname = 'lionvec' ORDER BY n;

/*
 * vt_vec() is a pseudo-random vector of n coordinates in (-scale, scale),
 * the same on every platform: from hashint8extended(), not random().
 */
CREATE FUNCTION vt_vec(seed bigint, n int, scale float8 DEFAULT 1.0)
RETURNS vector LANGUAGE sql IMMUTABLE
-- an index expression below calls it, and CREATE INDEX runs it under a safe
-- search_path from PostgreSQL 17 on, where vector would not be found
SET search_path FROM CURRENT AS $$
	SELECT array_agg(((hashint8extended(seed * 1000 + j, 42) % 100000)::float8
					  / 100000.0 * scale)::real ORDER BY j)::vector
	  FROM generate_series(1, n) j
$$;

/*
 * vec_cmp() runs q - which returns columns id and d, ORDER BY d - through
 * the index and through a sequential scan and sort, and compares the
 * distances in order (to 4 places), then the rows with a distance, sorted
 * (rows of equal distance may come in either order), then the count of rows
 * without one (NULL vectors: any of them may come).  It says whether the
 * first plan really was an index scan.
 */
CREATE FUNCTION vec_cmp(q text) RETURNS text
LANGUAGE plpgsql AS $$
DECLARE
	ln text;
	r record;
	used boolean := false;
	ad text[] := '{}';
	bd text[] := '{}';
	ai int[] := '{}';
	bi int[] := '{}';
BEGIN
	PERFORM set_config('enable_seqscan', 'off', true);
	PERFORM set_config('enable_bitmapscan', 'off', true);
	PERFORM set_config('enable_sort', 'off', true);
	FOR ln IN EXECUTE 'EXPLAIN (COSTS OFF) ' || q LOOP
		IF ln ~ 'Index Scan using' THEN
			used := true;
		END IF;
	END LOOP;
	FOR r IN EXECUTE q LOOP
		ad := ad || coalesce(round(r.d::numeric, 4)::text, 'NULL');
		IF r.d IS NOT NULL THEN
			ai := ai || r.id;
		END IF;
	END LOOP;
	PERFORM set_config('enable_seqscan', 'on', true);
	PERFORM set_config('enable_bitmapscan', 'on', true);
	PERFORM set_config('enable_sort', 'on', true);
	PERFORM set_config('enable_indexscan', 'off', true);
	FOR r IN EXECUTE q LOOP
		bd := bd || coalesce(round(r.d::numeric, 4)::text, 'NULL');
		IF r.d IS NOT NULL THEN
			bi := bi || r.id;
		END IF;
	END LOOP;
	PERFORM set_config('enable_indexscan', 'on', true);

	IF ad IS DISTINCT FROM bd THEN
		RETURN format('MISMATCH: %s rows through the index, %s by the sort; distances %s / %s',
					  coalesce(array_length(ad, 1), 0), coalesce(array_length(bd, 1), 0),
					  ad[1:3], bd[1:3]);
	END IF;
	IF (SELECT array_agg(x ORDER BY x) FROM unnest(ai) x) IS DISTINCT FROM
	   (SELECT array_agg(x ORDER BY x) FROM unnest(bi) x) THEN
		RETURN format('MISMATCH: other rows at the same distances');
	END IF;
	RETURN format('%s, %s rows',
				  CASE WHEN used THEN 'index' ELSE 'no index' END,
				  coalesce(array_length(ad, 1), 0));
END $$;

/*
 * vt_items() reads a lionvec index's CODE items off its pages (pageinspect):
 * each one's block, the chain the block is in - its partition's number,
 * from the PART chain's heads, or -1 for the rows in no partition, -2 for
 * the NULL rows - its row's TID, its flags - 1 a NULL row, 2 a cosine zero
 * vector, 4 a row's copy in its second partition (DESIGN.md §44.2) - and the
 * rest of its bytes.  vt_spill() sums them up with the meta page's format
 * version and flags (1: the index spills): the rows (the items that are not
 * copies), the copies, the NULL rows and zeros, and whether every row has
 * one item of its own, in a chain, and at most one copy, the same but for
 * the flag and in another partition's chain.  vt_strays() counts the items
 * of rows the table does not have.  vt_refines() reads the REFINE items:
 * the rows' and the centroids' (whose TIDs are (partition, 0)), and whether
 * the rows' are one for each CODE item of a row with a vector that is not a
 * copy, by TID.  The offsets assume a little-endian server.
 */
CREATE FUNCTION vt_u16(p bytea, o int) RETURNS int LANGUAGE sql IMMUTABLE AS
$$ SELECT get_byte(p, o) | (get_byte(p, o + 1) << 8) $$;
CREATE FUNCTION vt_u32(p bytea, o int) RETURNS bigint LANGUAGE sql IMMUTABLE AS
$$ SELECT vt_u16(p, o)::bigint | (vt_u16(p, o + 2)::bigint << 16) $$;
CREATE FUNCTION vt_items(idx regclass)
RETURNS TABLE (blk int, chain int, row_tid tid, flags int, rest bytea)
LANGUAGE plpgsql AS $$
DECLARE
	bs int := current_setting('block_size')::int;
	nb int := pg_relation_size(idx) / current_setting('block_size')::int;
	chains int[] := array_fill(NULL::int, ARRAY[nb]);
	m bytea := get_raw_page(idx::text, 0);
	p bytea;
	b int;
	pb bigint;
	h bigint;
	i int;
	lp bigint;
	o int;
BEGIN
	-- the chains' blocks: each partition's from the head in its PART item
	-- (the meta page's PART head at 52), then the REST and NULL chains' (60, 68)
	pb := vt_u32(m, 52);
	WHILE pb <> 4294967295 LOOP
		p := get_raw_page(idx::text, pb::int);
		FOR i IN 1 .. (vt_u16(p, 12) - 24) / 4 LOOP
			lp := vt_u16(p, 20 + 4 * i)::bigint | (vt_u16(p, 22 + 4 * i)::bigint << 16);
			CONTINUE WHEN (lp >> 15) & 3 <> 1;
			o := (lp & 32767)::int;
			h := vt_u32(p, o + 4);
			WHILE h <> 4294967295 LOOP
				chains[h + 1] := vt_u32(p, o)::int;
				h := vt_u32(get_raw_page(idx::text, h::int), bs - 16);
			END LOOP;
		END LOOP;
		pb := vt_u32(p, bs - 16);
	END LOOP;
	FOR i IN 1 .. 2 LOOP
		h := vt_u32(m, 52 + 8 * i);
		WHILE h <> 4294967295 LOOP
			chains[h + 1] := -i;
			h := vt_u32(get_raw_page(idx::text, h::int), bs - 16);
		END LOOP;
	END LOOP;

	FOR b IN 1 .. nb - 1 LOOP
		p := get_raw_page(idx::text, b);
		-- the special space's last two words: the page's type (2, CODE) and lionvec's id
		CONTINUE WHEN vt_u16(p, bs - 4) <> 2 OR vt_u16(p, bs - 2) <> 65420;
		FOR i IN 1 .. (vt_u16(p, 12) - 24) / 4 LOOP
			lp := vt_u16(p, 20 + 4 * i)::bigint | (vt_u16(p, 22 + 4 * i)::bigint << 16);
			CONTINUE WHEN (lp >> 15) & 3 <> 1;		-- not LP_NORMAL
			o := (lp & 32767)::int;
			blk := b;
			chain := chains[b + 1];
			row_tid := format('(%s,%s)', (vt_u16(p, o)::bigint << 16) | vt_u16(p, o + 2),
							  vt_u16(p, o + 4))::tid;
			flags := vt_u16(p, o + 6);
			rest := substr(p, o + 1, 6) || substr(p, o + 9, (lp >> 17)::int - 8);
			RETURN NEXT;
		END LOOP;
	END LOOP;
END $$;
CREATE FUNCTION vt_spill(idx regclass) RETURNS text
LANGUAGE sql AS $$
	WITH it AS (SELECT * FROM vt_items(idx)),
	r AS (SELECT row_tid, count(*) FILTER (WHERE flags & 4 = 0) AS own,
				 count(*) FILTER (WHERE flags & 4 <> 0) AS copies,
				 count(DISTINCT rest) AS bodies, count(DISTINCT chain) AS chains,
				 bool_and(chain IS NOT NULL) AS chained,
				 bool_and(flags & 4 = 0 OR chain >= 0) AS in_parts,
				 bool_and(flags & 3 = 0) AS vec
			FROM it GROUP BY row_tid)
	SELECT format('format %s, flags %s: %s rows, %s copies, %s NULL or zero; %s',
				  vt_u16(get_raw_page(idx::text, 0), 28),
				  vt_u16(get_raw_page(idx::text, 0), 34),
				  (SELECT count(*) FROM it WHERE flags & 4 = 0),
				  (SELECT count(*) FROM it WHERE flags & 4 <> 0),
				  (SELECT count(*) FROM it WHERE flags & 3 <> 0),
				  CASE WHEN (SELECT bool_and(own = 1 AND chained AND (copies = 0 OR
							(copies = 1 AND vec AND bodies = 1 AND chains = 2 AND in_parts))) FROM r)
					   THEN 'each row once, any copy the same in another partition'
					   ELSE 'MALFORMED' END)
$$;
CREATE FUNCTION vt_refines(idx regclass) RETURNS text
LANGUAGE plpgsql AS $$
DECLARE
	bs int := current_setting('block_size')::int;
	p bytea;
	b int;
	i int;
	lp bigint;
	o int;
	tids tid[] := '{}';
	ncent int := 0;
	same boolean;
BEGIN
	FOR b IN 1 .. pg_relation_size(idx) / bs - 1 LOOP
		p := get_raw_page(idx::text, b);
		CONTINUE WHEN vt_u16(p, bs - 4) <> 3 OR vt_u16(p, bs - 2) <> 65420;
		FOR i IN 1 .. (vt_u16(p, 12) - 24) / 4 LOOP
			lp := vt_u16(p, 20 + 4 * i)::bigint | (vt_u16(p, 22 + 4 * i)::bigint << 16);
			CONTINUE WHEN (lp >> 15) & 3 <> 1;
			o := (lp & 32767)::int;
			IF vt_u16(p, o + 4) = 0 THEN
				ncent := ncent + 1;
			ELSE
				tids := tids || format('(%s,%s)', (vt_u16(p, o)::bigint << 16) | vt_u16(p, o + 2),
									   vt_u16(p, o + 4))::tid;
			END IF;
		END LOOP;
	END LOOP;
	SELECT (SELECT array_agg(t ORDER BY t) FROM unnest(tids) t) IS NOT DISTINCT FROM
		   (SELECT array_agg(row_tid ORDER BY row_tid) FROM vt_items(idx) WHERE flags = 0)
	  INTO same;
	RETURN format('%s rows'' refine items, %s centroids''; %s', cardinality(tids), ncent,
				  CASE WHEN same THEN 'one a row with a vector' ELSE 'MALFORMED' END);
END $$;
CREATE FUNCTION vt_strays(idx regclass) RETURNS bigint
LANGUAGE plpgsql AS $$
DECLARE
	n bigint;
BEGIN
	EXECUTE format('SELECT count(*) FROM vt_items(%L) WHERE row_tid NOT IN (SELECT ctid FROM %s)',
				   idx, (SELECT indrelid::regclass FROM pg_index WHERE indexrelid = idx)) INTO n;
	RETURN n;
END $$;

/*
 * vt_regions() reads the REGION chain of an index of 64 partitions or more
 * (DESIGN.md §44.2) from the meta page's head (at 96) and count (100): each
 * item's number, its run of partitions - the first and how many - and where
 * the run's first PART item is (block and offset); and whether they number
 * from 0 up, the runs follow one another and cover the meta page's
 * partitions (56), and each run's PART item is its first partition's (a
 * REGION page's type is 5, a PART page's 4).  vt_placed() compares where
 * inserts put the rows of an index's table whose ids are past off with
 * where the build put the rows whose vectors they copy (id - off): how many
 * are in the same partition, and how many of their copies too.
 */
CREATE FUNCTION vt_regions(idx regclass) RETURNS text
LANGUAGE plpgsql AS $$
DECLARE
	bs int := current_setting('block_size')::int;
	m bytea := get_raw_page(idx::text, 0);
	nparts bigint := vt_u32(m, 56);
	rb bigint := vt_u32(m, 96);
	p bytea;
	q bytea;
	i int;
	lp bigint;
	o int;
	n int := 0;
	total bigint := 0;
	cnt bigint;
	ok boolean := true;
BEGIN
	WHILE rb <> 4294967295 LOOP
		p := get_raw_page(idx::text, rb::int);
		ok := ok AND vt_u16(p, bs - 4) = 5;
		FOR i IN 1 .. (vt_u16(p, 12) - 24) / 4 LOOP
			lp := vt_u16(p, 20 + 4 * i)::bigint | (vt_u16(p, 22 + 4 * i)::bigint << 16);
			CONTINUE WHEN (lp >> 15) & 3 <> 1;
			o := (lp & 32767)::int;
			cnt := vt_u32(p, o + 8);
			ok := ok AND vt_u32(p, o) = n AND vt_u32(p, o + 4) = total;
			IF cnt = 0 THEN
				ok := ok AND vt_u32(p, o + 12) = 4294967295;
			ELSE
				q := get_raw_page(idx::text, vt_u32(p, o + 12)::int);
				lp := vt_u16(q, 20 + 4 * vt_u16(p, o + 16))::bigint |
					(vt_u16(q, 22 + 4 * vt_u16(p, o + 16))::bigint << 16);
				ok := ok AND vt_u16(q, bs - 4) = 4 AND vt_u32(q, (lp & 32767)::int) = total;
			END IF;
			total := total + cnt;
			n := n + 1;
		END LOOP;
		rb := vt_u32(p, bs - 16);
	END LOOP;
	ok := ok AND n = vt_u32(m, 100) AND (n = 0 OR total = nparts);
	RETURN format('format %s: %s partitions in %s regions; %s', vt_u16(m, 28), nparts, n,
				  CASE WHEN NOT ok THEN 'MALFORMED'
					   WHEN n = 0 THEN 'no REGION chain'
					   ELSE 'runs in order, each at its first PART item' END);
END $$;
CREATE FUNCTION vt_placed(idx regclass, off int)
RETURNS TABLE ("rows" bigint, "same partition" bigint, "same copy's" bigint)
LANGUAGE plpgsql AS $$
BEGIN
	RETURN QUERY EXECUTE format($q$
		WITH it AS (SELECT row_tid, chain, flags & 4 <> 0 AS copy FROM vt_items(%L)),
		t AS (SELECT id, ctid FROM %s)
		SELECT count(*), count(*) FILTER (WHERE a.chain = b.chain),
			   count(*) FILTER (WHERE ac.chain = bc.chain)
		  FROM t n JOIN t o ON o.id = n.id - %s
		  JOIN it a ON a.row_tid = n.ctid AND NOT a.copy
		  JOIN it b ON b.row_tid = o.ctid AND NOT b.copy
		  LEFT JOIN it ac ON ac.row_tid = n.ctid AND ac.copy
		  LEFT JOIN it bc ON bc.row_tid = o.ctid AND bc.copy
		 WHERE n.id > %s$q$,
		idx, (SELECT indrelid::regclass FROM pg_index WHERE indexrelid = idx), off, off);
END $$;

-- vt_ids() runs q, which returns a column id, and returns the ids.
CREATE FUNCTION vt_ids(q text) RETURNS SETOF int
LANGUAGE plpgsql AS $$ BEGIN RETURN QUERY EXECUTE q; END $$;

-- 3000 rows in ten clusters, five NULL vectors and one of zeros; indexed in
-- 20 partitions, which a table this small would not be by itself
CREATE TABLE vt (id int PRIMARY KEY, cat int, v vector(40))
	WITH (autovacuum_enabled = off);
INSERT INTO vt
SELECT i, i % 10, vt_vec(-(i % 10) - 1, 40) + vt_vec(i, 40, 0.3)
  FROM generate_series(1, 3000) i;
INSERT INTO vt SELECT i, NULL, NULL FROM generate_series(3001, 3005) i;
INSERT INTO vt VALUES (3006, NULL, array_fill(0, ARRAY[40])::real[]::vector);

CREATE INDEX vt_l2 ON vt USING lionvec (v) WITH (partitions = 20);
CREATE INDEX vt_ip ON vt USING lionvec (v vector_ip_ops) WITH (partitions = 20);
CREATE INDEX vt_cos ON vt USING lionvec (v vector_cosine_ops) WITH (partitions = 20);
VACUUM (ANALYZE) vt;

-- each row with a vector in two partitions, its nearest and SOAR's second:
-- two CODE items naming one REFINE item; NULL rows, and a cosine index's zero
-- vector, once
SELECT i AS index, vt_spill(i) FROM unnest('{vt_l2,vt_ip,vt_cos}'::regclass[]) i;

SET pg_lion.vector_epsilon = 6;
SET pg_lion.vector_margin = 100;
SET pg_lion.vector_patience = 0;

-- every metric, a query vector in a cluster and between them, any LIMIT
CREATE TABLE vt_cases AS
SELECT op, q, k
  FROM (VALUES ('<->'), ('<#>'), ('<=>')) o(op),
	   (VALUES ('(SELECT v FROM vt WHERE id = 7)'),
			   ('(SELECT v FROM vt WHERE id = 2999)'),
			   ('vt_vec(77, 40, 0.5)')) q(q),
	   (VALUES (1), (10), (100), (3006)) k(k);
SELECT op, k, vec_cmp(format('SELECT id, v %s %s AS d FROM vt ORDER BY v %s %s LIMIT %s',
							 op, q, op, q, k)) AS result
  FROM vt_cases ORDER BY op, q, k;

-- no LIMIT: every row, the NULL vectors last
SELECT vec_cmp('SELECT id, v <-> vt_vec(5, 40) AS d FROM vt ORDER BY v <-> vt_vec(5, 40)');

-- each row once, though it is in two partitions: past the table's rows, in
-- order and by a walk (a NULL query's, and a cosine zero query's)
SET enable_seqscan = off;
SET enable_bitmapscan = off;
SET enable_sort = off;
SELECT op, count(*), count(DISTINCT id)
  FROM (VALUES ('<->'), ('<#>'), ('<=>')) o(op),
	   LATERAL vt_ids(format('SELECT id FROM vt ORDER BY v %s vt_vec(77, 40, 0.5) LIMIT 5000', op)) id
 GROUP BY op ORDER BY op;
SELECT count(*), count(DISTINCT id)
  FROM (SELECT id FROM vt ORDER BY v <-> (SELECT v FROM vt WHERE id = 3001)) s;
SELECT count(*), count(DISTINCT id)
  FROM (SELECT id FROM vt ORDER BY v <=> array_fill(0, ARRAY[40])::real[]::vector) s;
RESET enable_seqscan;
RESET enable_bitmapscan;
RESET enable_sort;

-- ties: 1,200 rows of 30 vectors, 40 rows each at every distance, in 12
-- partitions - each row once, at its distance, at LIMITs that end a tie
-- (within one, either plan may take any of its rows)
CREATE TABLE vtt (id int, v vector(40));
INSERT INTO vtt SELECT i, vt_vec(i % 30, 40) FROM generate_series(1, 1200) i;
CREATE INDEX vtt_l2 ON vtt USING lionvec (v) WITH (partitions = 12);
SELECT vt_spill('vtt_l2');
SELECT k, vec_cmp(format('SELECT id, v <-> vt_vec(3, 40) AS d FROM vtt ORDER BY v <-> vt_vec(3, 40) LIMIT %s', k))
  FROM (VALUES (40), (80), (600), (1200), (5000)) k(k) ORDER BY k;
SET enable_seqscan = off;
SET enable_sort = off;
SELECT count(*), count(DISTINCT id)
  FROM (SELECT id FROM vtt ORDER BY v <-> vt_vec(3, 40) LIMIT 5000) s;
RESET enable_seqscan;
RESET enable_sort;
DROP TABLE vtt;

-- a NULL query (a NULL row's: a constant NULL is folded away): every row,
-- each at a NULL distance
SELECT vec_cmp('SELECT id, v <-> (SELECT v FROM vt WHERE id = 3001) AS d FROM vt
				ORDER BY v <-> (SELECT v FROM vt WHERE id = 3001) LIMIT 10');
SELECT vec_cmp('SELECT id, v <=> (SELECT v FROM vt WHERE id = 3001) AS d FROM vt
				ORDER BY v <=> (SELECT v FROM vt WHERE id = 3001)');

-- the cosine of the zero vector is NaN, after every number and before NULL
SELECT vec_cmp('SELECT id, v <=> vt_vec(9, 40) AS d FROM vt ORDER BY v <=> vt_vec(9, 40) OFFSET 2995');
SELECT vec_cmp(format('SELECT id, v <=> %L::vector AS d FROM vt ORDER BY v <=> %L::vector',
					  array_fill(0, ARRAY[40])::real[]::vector,
					  array_fill(0, ARRAY[40])::real[]::vector));

-- a query of other dimensions is pgvector's error
SET enable_seqscan = off;
SELECT id FROM vt ORDER BY v <-> '[1,2,3]' LIMIT 1;
-- with no ORDER BY the index is never scanned, even with sequential scans
-- off (from 18 it was, as an index-only scan, which failed)
SELECT count(*) FROM vt;
RESET enable_seqscan;

-- the scan as it goes on past a pass: a batch of one, and passes cut to
-- work_mem (with room in the progress hash, which hash_mem bounds, for every
-- row: the next test is the hash's)
SET pg_lion.vector_batch = 1;
SELECT vec_cmp('SELECT id, v <-> vt_vec(11, 40) AS d FROM vt ORDER BY v <-> vt_vec(11, 40) LIMIT 300');
SET work_mem = '64kB';
SET hash_mem_multiplier = 64;
SELECT vec_cmp('SELECT id, v <#> vt_vec(12, 40) AS d FROM vt ORDER BY v <#> vt_vec(12, 40) LIMIT 2000');
SELECT vec_cmp('SELECT id, v <-> vt_vec(12, 40) AS d FROM vt ORDER BY v <-> vt_vec(12, 40) LIMIT 3006');
RESET hash_mem_multiplier;
RESET work_mem;
RESET pg_lion.vector_batch;

-- a scan with no LIMIT whose rows outgrow the progress hash's hash_mem
-- (128 kB: 1,843 entries, for 3,001 rows) drops what it no longer needs,
-- and what it can work out again, and returns every row, in order
SET work_mem = '128kB';
SET hash_mem_multiplier = 1;
SELECT vec_cmp('SELECT id, v <-> vt_vec(13, 40) AS d FROM vt ORDER BY v <-> vt_vec(13, 40)');
-- with a work_mem of 64 kB, whose shorter passes need more of the hash, it
-- cannot go on without growing past hash_mem: by default, an ERROR
SET work_mem = '64kB';
SET hash_mem_multiplier = 2;
SET enable_seqscan = off;
SET enable_sort = off;
SHOW pg_lion.vector_overflow;
SELECT count(*) FROM (SELECT id, v <-> vt_vec(13, 40) AS d FROM vt ORDER BY v <-> vt_vec(13, 40)) s;
-- pg_lion.vector_overflow = stop ends the scan there instead, with a
-- NOTICE: what it returned is the start of the sort's order, each row once
SET pg_lion.vector_overflow = stop;
CREATE TEMP TABLE vt_cut AS
SELECT row_number() OVER () AS n, id, d
  FROM (SELECT id, v <-> vt_vec(13, 40) AS d FROM vt ORDER BY v <-> vt_vec(13, 40)) s;
RESET enable_seqscan;
RESET enable_sort;
-- the reference is the sort's order: exact, whatever plan PostgreSQL would
-- pick for the window (on 17+ it can be the lionvec scan)
SET enable_indexscan = off;
SELECT count(*) < 3006 AS stopped, count(*) > 0 AS "some rows",
	   count(*) = count(*) FILTER (WHERE c.id = s.id) AS "the sort's first rows"
  FROM vt_cut c
  JOIN (SELECT row_number() OVER (ORDER BY v <-> vt_vec(13, 40), id) AS n, id
		  FROM vt WHERE v IS NOT NULL) s USING (n);
RESET enable_indexscan;
DROP TABLE vt_cut;
RESET pg_lion.vector_overflow;
SET pg_lion.vector_overflow = 'grow';
RESET hash_mem_multiplier;
RESET work_mem;

-- inserts after the build: a new cluster, far from the mean, which goes to
-- the partitions nearest it
INSERT INTO vt
SELECT i, 10, vt_vec(-11, 40, 4.0) + vt_vec(i, 40, 0.3)
  FROM generate_series(3007, 3600) i;
-- each in two partitions, as the build would put it
SELECT i AS index, vt_spill(i) FROM unnest('{vt_l2,vt_ip,vt_cos}'::regclass[]) i;
SELECT op, k, vec_cmp(format('SELECT id, v %s %s AS d FROM vt ORDER BY v %s %s LIMIT %s',
							 op, q, op, q, k)) AS result
  FROM vt_cases WHERE k IN (10, 3006) ORDER BY op, q, k;
SELECT vec_cmp('SELECT id, v <-> vt_vec(-11, 40, 4.0) AS d FROM vt ORDER BY v <-> vt_vec(-11, 40, 4.0) LIMIT 50');

-- updates, deletes and VACUUM; then more inserts into what it freed
UPDATE vt SET v = vt_vec(id + 100000, 40) WHERE id % 7 = 0;
DELETE FROM vt WHERE id % 3 = 0;
-- the dead rows' REFINE items are there until VACUUM
SELECT i AS index, vt_refines(i) FROM unnest('{vt_l2,vt_ip,vt_cos}'::regclass[]) i;
VACUUM vt;
-- both CODE items of each dead row gone, and its one REFINE item
SELECT i AS index, vt_spill(i), vt_strays(i), vt_refines(i)
  FROM unnest('{vt_l2,vt_ip,vt_cos}'::regclass[]) i;
SELECT op, k, vec_cmp(format('SELECT id, v %s %s AS d FROM vt ORDER BY v %s %s LIMIT %s',
							 op, q, op, q, k)) AS result
  FROM vt_cases WHERE k IN (10, 3006) ORDER BY op, q, k;
INSERT INTO vt
SELECT i, i % 10, vt_vec(-(i % 10) - 1, 40) + vt_vec(i, 40, 0.3)
  FROM generate_series(4000, 4500) i;
SELECT vec_cmp('SELECT id, v <-> vt_vec(13, 40) AS d FROM vt ORDER BY v <-> vt_vec(13, 40) LIMIT 200');

-- every row deleted: nothing comes back, NULLs included
CREATE TABLE vt2 AS SELECT * FROM vt WHERE id < 100 OR v IS NULL;
CREATE INDEX vt2_l2 ON vt2 USING lionvec (v);
DELETE FROM vt2;
VACUUM vt2;
SELECT vec_cmp('SELECT id, v <-> vt_vec(1, 40) AS d FROM vt2 ORDER BY v <-> vt_vec(1, 40)');

-- REINDEX, and a concurrent build
REINDEX INDEX vt_l2;
SELECT vt_spill('vt_l2');
SELECT vec_cmp('SELECT id, v <-> vt_vec(14, 40) AS d FROM vt ORDER BY v <-> vt_vec(14, 40) LIMIT 100');
-- spill is the build's: off takes a REINDEX, after which each row is in one
-- partition, and the scan returns the same
ALTER INDEX vt_l2 SET (spill = off);
SELECT vt_spill('vt_l2');
REINDEX INDEX vt_l2;
SELECT vt_spill('vt_l2');
SELECT vec_cmp('SELECT id, v <-> vt_vec(14, 40) AS d FROM vt ORDER BY v <-> vt_vec(14, 40) LIMIT 100');
INSERT INTO vt SELECT i, 11, vt_vec(i, 40) FROM generate_series(5000, 5049) i;
SELECT vt_spill('vt_l2');
SELECT vec_cmp('SELECT id, v <-> vt_vec(14, 40) AS d FROM vt ORDER BY v <-> vt_vec(14, 40) LIMIT 300');
DELETE FROM vt WHERE id >= 5000;
VACUUM vt;
ALTER INDEX vt_l2 RESET (spill);
REINDEX INDEX vt_l2;
SELECT vt_spill('vt_l2');
DROP INDEX vt_l2;
CREATE INDEX CONCURRENTLY vt_l2 ON vt USING lionvec (v) WITH (partitions = 20);
SELECT vec_cmp('SELECT id, v <-> vt_vec(15, 40) AS d FROM vt ORDER BY v <-> vt_vec(15, 40) LIMIT 100');

-- an expression index, and a partial one
CREATE INDEX vt_half ON vt USING lionvec ((v * vt_vec(16, 40)));
SELECT vec_cmp('SELECT id, (v * vt_vec(16, 40)) <-> vt_vec(17, 40) AS d FROM vt
				ORDER BY (v * vt_vec(16, 40)) <-> vt_vec(17, 40) LIMIT 50');
DROP INDEX vt_half;
CREATE INDEX vt_cat3 ON vt USING lionvec (v) WHERE cat = 3;
SELECT vec_cmp('SELECT id, v <-> vt_vec(18, 40) AS d FROM vt WHERE cat = 3
				ORDER BY v <-> vt_vec(18, 40) LIMIT 50');
DROP INDEX vt_cat3;

-- a scrollable cursor goes backward over a Material: the scan never does
SET enable_seqscan = off;
SET enable_sort = off;
BEGIN;
DECLARE vc SCROLL CURSOR FOR SELECT id FROM vt ORDER BY v <-> vt_vec(19, 40) LIMIT 5;
FETCH 3 FROM vc;
FETCH BACKWARD 2 FROM vc;
COMMIT;
RESET enable_sort;
RESET enable_seqscan;

-- a table empty at the build, whose column has no dimensions: the first row
-- sets them, and a row of others is pgvector's error
CREATE TABLE ve (id int, v vector);
CREATE INDEX ve_l2 ON ve USING lionvec (v);
CREATE INDEX ve_ip ON ve USING lionvec (v vector_ip_ops);
INSERT INTO ve VALUES (1, '[1,2,3]'), (2, '[3,2,1]'), (3, NULL), (4, '[0,0,1]');
INSERT INTO ve VALUES (5, '[1,2,3,4]');
-- no partitions, so no second one
SELECT vt_spill('ve_l2');
SET enable_seqscan = off;
EXPLAIN (COSTS OFF) SELECT id FROM ve ORDER BY v <-> '[1,1,2]' LIMIT 2;
SELECT id, v <-> '[1,1,2]' AS d FROM ve ORDER BY v <-> '[1,1,2]';
SELECT id, v <#> '[1,1,2]' AS d FROM ve ORDER BY v <#> '[1,1,2]';
SELECT id FROM ve ORDER BY v <-> '[1,1]';
RESET enable_seqscan;

-- unlogged: the init fork (ambuildempty)
CREATE UNLOGGED TABLE vu (id int, v vector(3));
CREATE INDEX vu_l2 ON vu USING lionvec (v);
INSERT INTO vu SELECT i, vt_vec(i, 3) FROM generate_series(1, 200) i;
SELECT vec_cmp('SELECT id, v <-> ''[0.1,0.2,0.3]'' AS d FROM vu ORDER BY v <-> ''[0.1,0.2,0.3]'' LIMIT 20');
-- one partition (a table under 10,000 rows): no second one
REINDEX INDEX vu_l2;
SELECT vt_spill('vu_l2');

-- more partitions than rows: as many as the rows, one each; and none, of
-- a table empty at the build, whose rows all go to the chain of rows in no
-- partition
DROP INDEX vt2_l2;
CREATE INDEX vt2_many ON vt2 USING lionvec (v) WITH (partitions = 500);
INSERT INTO vt2 SELECT * FROM vt WHERE id < 300 OR v IS NULL;
-- the REST chain takes no copies
SELECT vt_spill('vt2_many');
SELECT vec_cmp('SELECT id, v <-> vt_vec(20, 40) AS d FROM vt2 ORDER BY v <-> vt_vec(20, 40) LIMIT 20');
REINDEX INDEX vt2_many;
SELECT vt_spill('vt2_many');
SELECT vec_cmp('SELECT id, v <-> vt_vec(20, 40) AS d FROM vt2 ORDER BY v <-> vt_vec(20, 40) LIMIT 20');
SELECT vec_cmp('SELECT id, v <-> vt_vec(20, 40) AS d FROM vt2 ORDER BY v <-> vt_vec(20, 40)');
DROP INDEX vt2_many;

-- fewer partitions than the nine nearest SOAR ranks: every centroid a
-- candidate, at the build and at an insert; of two, the copy goes to the
-- other
CREATE INDEX vt2_two ON vt2 USING lionvec (v) WITH (partitions = 2);
CREATE INDEX vt2_five ON vt2 USING lionvec (v vector_ip_ops) WITH (partitions = 5);
INSERT INTO vt2
SELECT i, i % 10, vt_vec(-(i % 10) - 1, 40) + vt_vec(i, 40, 0.3)
  FROM generate_series(6000, 6099) i;
SELECT i AS index, vt_spill(i), vt_refines(i) FROM unnest('{vt2_two,vt2_five}'::regclass[]) i;
SELECT vec_cmp('SELECT id, v <-> vt_vec(21, 40) AS d FROM vt2 ORDER BY v <-> vt_vec(21, 40) LIMIT 60');
SELECT vec_cmp('SELECT id, v <#> vt_vec(21, 40) AS d FROM vt2 ORDER BY v <#> vt_vec(21, 40)');
DROP INDEX vt2_two, vt2_five;

-- regions (DESIGN.md §44.3): an index of fewer than 64 partitions, one
-- partition included, has none; of 100, ceil(sqrt(100)) = 10, the
-- partitions numbered region by region, so that each region's are a run of
-- the PART chain
SELECT i AS index, vt_regions(i) FROM unnest('{vt_l2,vt_ip,vt_cos,vu_l2}'::regclass[]) i;
CREATE TABLE vr (id int PRIMARY KEY, v vector(40)) WITH (autovacuum_enabled = off);
INSERT INTO vr
SELECT i, vt_vec(-(i % 25) - 1, 40) + vt_vec(i, 40, 0.3)
  FROM generate_series(1, 6400) i;
CREATE INDEX vr_l2 ON vr USING lionvec (v) WITH (partitions = 100);
CREATE INDEX vr_ip ON vr USING lionvec (v vector_ip_ops) WITH (partitions = 100);
CREATE INDEX vr_cos ON vr USING lionvec (v vector_cosine_ops) WITH (partitions = 100);
SELECT i AS index, vt_regions(i), vt_spill(i), vt_refines(i)
  FROM unnest('{vr_l2,vr_ip,vr_cos}'::regclass[]) i;
-- an insert reads the REGION chain and the runs of the regions nearest the
-- row alone, and puts it where the build put a row of the same vector -
-- its copy too, but for the few whose SOAR choice the centroids decoded
-- from their codes (all an insert has of them) turn
INSERT INTO vr SELECT id + 100000, v FROM vr WHERE id <= 1000;
SELECT i AS index, p."rows", p."same partition" = p."rows" AS "the build's partition",
	   p."same copy's" >= p."rows" * 0.9 AS "the build's copy's, mostly"
  FROM unnest('{vr_l2,vr_ip,vr_cos}'::regclass[]) i, LATERAL vt_placed(i, 100000) p;
SELECT op, q, k, vec_cmp(format('SELECT id, v %s %s AS d FROM vr ORDER BY v %s %s LIMIT %s',
								op, q, op, q, k)) AS result
  FROM (VALUES ('<->'), ('<#>'), ('<=>')) o(op),
	   (VALUES ('vt_vec(-3, 40)'), ('vt_vec(23, 40, 2.0)')) c(q),
	   (VALUES (10), (300)) l(k)
 ORDER BY op, q, k;
DROP TABLE vr;

/*
 * vec_recall() runs q - which returns a column id - through the index and
 * through a sequential scan and sort, and returns the share of the second's
 * rows the first has.
 */
CREATE FUNCTION vec_recall(q text) RETURNS float8
LANGUAGE plpgsql AS $$
DECLARE
	a int[];
	b int[];
BEGIN
	PERFORM set_config('enable_seqscan', 'off', true);
	PERFORM set_config('enable_bitmapscan', 'off', true);
	PERFORM set_config('enable_sort', 'off', true);
	EXECUTE 'SELECT array_agg(id) FROM (' || q || ') s' INTO a;
	PERFORM set_config('enable_seqscan', 'on', true);
	PERFORM set_config('enable_bitmapscan', 'on', true);
	PERFORM set_config('enable_sort', 'on', true);
	PERFORM set_config('enable_indexscan', 'off', true);
	EXECUTE 'SELECT array_agg(id) FROM (' || q || ') s' INTO b;
	PERFORM set_config('enable_indexscan', 'on', true);
	RETURN (SELECT count(*) FROM unnest(a) x WHERE x = ANY (b))::float8 /
		greatest(cardinality(b), 1);
END $$;

-- the default margin and patience read the partitions near the query only,
-- and miss few of its nearest rows; a margin below zero reads fewer, and
-- misses more
RESET pg_lion.vector_margin;
RESET pg_lion.vector_patience;
SELECT op, avg(vec_recall(format('SELECT id FROM vt ORDER BY v %s %s LIMIT 10',
								 op, q))) >= 0.9 AS "most of the nearest"
  FROM (VALUES ('<->'), ('<#>'), ('<=>')) o(op),
	   (SELECT format('(vt_vec(-%s, 40) + vt_vec(%s, 40, 0.3))', s % 10 + 1, s + 9000) AS q
		  FROM generate_series(1, 20) s) q
 GROUP BY op ORDER BY op;
SET pg_lion.vector_margin = -10;
SELECT vec_recall('SELECT id FROM vt ORDER BY v <-> vt_vec(77, 40, 0.5) LIMIT 100') < 1 AS "fewer";
RESET pg_lion.vector_margin;

-- spill against none, over two tables of the first 3,000 rows: run to the
-- margin it finds as many of the nearest, and at a margin below zero,
-- which leaves rows in partitions not probed, more - the rows near a
-- partition's edge are in the next one too.  The pause is off: 30 of its
-- units are 820 of these rows, five of the twenty partitions, and in so
-- few, each two thirds of a batch of 100, the spilling index's nearest
-- come early, enter its hundred least upper bounds early, and the pause
-- then stops it at four partitions on <#> where the other runs to seven
CREATE TABLE vs (id int, v vector(40));
INSERT INTO vs SELECT i, vt_vec(-(i % 10) - 1, 40) + vt_vec(i, 40, 0.3)
  FROM generate_series(1, 3000) i;
CREATE TABLE vn AS SELECT * FROM vs;
CREATE INDEX vs_l2 ON vs USING lionvec (v) WITH (partitions = 20);
CREATE INDEX vs_ip ON vs USING lionvec (v vector_ip_ops) WITH (partitions = 20);
CREATE INDEX vs_cos ON vs USING lionvec (v vector_cosine_ops) WITH (partitions = 20);
CREATE INDEX vn_l2 ON vn USING lionvec (v) WITH (partitions = 20, spill = off);
CREATE INDEX vn_ip ON vn USING lionvec (v vector_ip_ops) WITH (partitions = 20, spill = off);
CREATE INDEX vn_cos ON vn USING lionvec (v vector_cosine_ops) WITH (partitions = 20, spill = off);
VACUUM ANALYZE vs, vn;
SELECT vt_spill('vs_l2') AS spill, vt_spill('vn_l2') AS "no spill";
CREATE FUNCTION vt_recall_of(tbl text, op text) RETURNS float8
LANGUAGE sql AS $$
	SELECT avg(vec_recall(format('SELECT id FROM %s ORDER BY v %s %s LIMIT 10', tbl, op, q)))
	  FROM (SELECT format('(vt_vec(-%s, 40) + vt_vec(%s, 40, 0.3))', s % 10 + 1, s + 9000) AS q
			  FROM generate_series(1, 20) s) q
$$;
SET pg_lion.vector_patience = 0;
SELECT op, vt_recall_of('vs', op) >= vt_recall_of('vn', op) AS "as many"
  FROM (VALUES ('<->'), ('<#>'), ('<=>')) o(op) ORDER BY op;
SET pg_lion.vector_margin = -2;
SELECT op, vt_recall_of('vs', op) > vt_recall_of('vn', op) AS "more"
  FROM (VALUES ('<->'), ('<#>'), ('<=>')) o(op) ORDER BY op;
RESET pg_lion.vector_margin;
RESET pg_lion.vector_patience;
DROP FUNCTION vt_recall_of(text, text);
DROP TABLE vs, vn;

-- a scan that pauses at once, every partition within its margin: it still
-- fills its LIMIT, in order, from the rows it holds and then the partitions
-- after the pause
SET pg_lion.vector_margin = 100;
SET pg_lion.vector_patience = 1;
SET enable_seqscan = off;
SET enable_sort = off;
SELECT cardinality(a) AS n, a = (SELECT array_agg(x ORDER BY x) FROM unnest(a) x) AS "in order"
  FROM (SELECT array_agg(d) AS a
		  FROM (SELECT v <-> vt_vec(77, 40, 0.5) AS d FROM vt
				 ORDER BY v <-> vt_vec(77, 40, 0.5) LIMIT 300) l) s;
RESET pg_lion.vector_margin;

-- ... and at a margin the patience does not grow with, it pauses at each
-- tier - ten rows, twenty, forty, eighty, then L's hundred - and in the
-- REST's batch after them: its rows still come in order, each once
SELECT cardinality(a) AS n, (SELECT count(DISTINCT x) FROM unnest(i) x) AS ids,
	   a = (SELECT array_agg(x ORDER BY x) FROM unnest(a) x) AS "in order"
  FROM (SELECT array_agg(d) AS a, array_agg(id) AS i
		  FROM (SELECT id, v <-> vt_vec(77, 40, 0.5) AS d FROM vt
				 ORDER BY v <-> vt_vec(77, 40, 0.5) LIMIT 300) l) s;
RESET enable_seqscan;
RESET enable_sort;
RESET pg_lion.vector_patience;

-- an operator class of one's own, whose operator returns float4: the scan
-- reads the operator's distances, and hands them back, as float4
CREATE FUNCTION vt_l2_f4(vector, vector) RETURNS float4
LANGUAGE sql IMMUTABLE STRICT SET search_path FROM CURRENT
AS $$ SELECT l2_distance($1, $2)::float4 $$;
CREATE OPERATOR <%%> (LEFTARG = vector, RIGHTARG = vector,
					   FUNCTION = vt_l2_f4, COMMUTATOR = <%%>);
CREATE OPERATOR CLASS vt_l2_f4_ops FOR TYPE vector USING lionvec AS
	OPERATOR 1 <%%> (vector, vector) FOR ORDER BY float_ops;
SELECT amvalidate(oid) FROM pg_opclass WHERE opcname = 'vt_l2_f4_ops';
CREATE INDEX vt_f4 ON vt USING lionvec (v vt_l2_f4_ops) WITH (partitions = 20);
SELECT vec_cmp('SELECT id, v <%%> vt_vec(14, 40) AS d FROM vt ORDER BY v <%%> vt_vec(14, 40) LIMIT 50');
SET enable_seqscan = off;
SET enable_sort = off;
SELECT pg_typeof(d), d = (v <%%> vt_vec(14, 40)) AS "the operator's"
  FROM (SELECT v, v <%%> vt_vec(14, 40) AS d FROM vt
		 ORDER BY v <%%> vt_vec(14, 40) LIMIT 1) s;
RESET enable_seqscan;
RESET enable_sort;
DROP INDEX vt_f4;

-- what lionvec refuses
CREATE TABLE vbad (i int, v vector(9000), w vector(3)[]);
CREATE INDEX ON vbad USING lionvec (i);
-- an operator class over another type, which every read of a row or a
-- query would take for a vector: refused before any is read
CREATE FUNCTION vbad_dist(int, int) RETURNS float8 LANGUAGE sql IMMUTABLE
AS $$ SELECT abs($1 - $2)::float8 $$;
CREATE OPERATOR <^^> (LEFTARG = int, RIGHTARG = int, FUNCTION = vbad_dist);
CREATE OPERATOR CLASS vbad_int_ops FOR TYPE int USING lionvec AS
	OPERATOR 1 <^^> (int, int) FOR ORDER BY float_ops;
SELECT amvalidate(oid) FROM pg_opclass WHERE opcname = 'vbad_int_ops';
INSERT INTO vbad (i) VALUES (1), (2);
CREATE INDEX ON vbad USING lionvec (i vbad_int_ops);
-- ... and one whose operator returns neither float8 nor float4
CREATE FUNCTION vt_l2_int(vector, vector) RETURNS int
LANGUAGE sql IMMUTABLE STRICT SET search_path FROM CURRENT
AS $$ SELECT l2_distance($1, $2)::int $$;
CREATE OPERATOR <||> (LEFTARG = vector, RIGHTARG = vector, FUNCTION = vt_l2_int);
CREATE OPERATOR CLASS vt_l2_int_ops FOR TYPE vector USING lionvec AS
	OPERATOR 1 <||> (vector, vector) FOR ORDER BY integer_ops;
SELECT amvalidate(oid) FROM pg_opclass WHERE opcname = 'vt_l2_int_ops';
CREATE INDEX vt_int ON vt USING lionvec (v vt_l2_int_ops) WITH (partitions = 20);
SET enable_seqscan = off;
SET enable_sort = off;
SELECT id FROM vt ORDER BY v <||> vt_vec(14, 40) LIMIT 1;
RESET enable_seqscan;
RESET enable_sort;
DROP INDEX vt_int;
DROP OPERATOR FAMILY vt_l2_int_ops USING lionvec;
DROP OPERATOR <||> (vector, vector);
DROP FUNCTION vt_l2_int(vector, vector);
DROP OPERATOR FAMILY vbad_int_ops USING lionvec;
DROP OPERATOR <^^> (int, int);
DROP FUNCTION vbad_dist(int, int);
DROP OPERATOR FAMILY vt_l2_f4_ops USING lionvec;
DROP OPERATOR <%%> (vector, vector);
DROP FUNCTION vt_l2_f4(vector, vector);
CREATE INDEX ON vbad USING lionvec (v);
CREATE INDEX ON vbad USING lionvec (w);
CREATE INDEX ON vbad USING lionvec (v, v);
CREATE INDEX ON vt USING lionvec (v) WITH (partitions = -1);
CREATE INDEX ON vt USING lionvec (v) WITH (partitions = 262145);
CREATE INDEX ON vt USING lionvec (v) WITH (spill = maybe);
ALTER INDEX vt_l2 SET (spill = 2);
SET pg_lion.vector_epsilon = 0.5;
SET pg_lion.vector_batch = 0;
SET pg_lion.vector_margin = 1001;
SET pg_lion.vector_patience = -1;

DROP TABLE vt, vt2, vt_cases, ve, vu, vbad;
DROP FUNCTION vec_cmp(text);
DROP FUNCTION vec_recall(text);
DROP FUNCTION vt_ids(text);
DROP FUNCTION vt_placed(regclass, int);
DROP FUNCTION vt_regions(regclass);
DROP FUNCTION vt_strays(regclass);
DROP FUNCTION vt_refines(regclass);
DROP FUNCTION vt_spill(regclass);
DROP FUNCTION vt_items(regclass);
DROP FUNCTION vt_u32(bytea, int);
DROP FUNCTION vt_u16(bytea, int);
DROP FUNCTION vt_vec(bigint, int, float8);