"""Filtered scans (trigpoint.iterative_scan). A WHERE above the index makes the executor pull past the index's first k rows. The scan then goes on in batches: k and the cells opened double, the routed order is kept, the opened cells are cascaded again under the larger tau and no heap row is fetched twice. On a 100k synthetic corpus with a `category` column at 10 %, 1 % and 0.1 % selectivity, `WHERE category = $1 ORDER BY v <-> q LIMIT 10` must come back full and close to the filtered exact answer; the caps stop the scan short with no error; strict_order keeps the stream in distance order; `off` is the 0.4 widening (k doubles, the cells opened do not grow); an always-true filter at LIMIT 10/100/1000 is at least as good as the 0.4 scan at the same N. """ import struct import time import numpy as np import pytest psycopg = pytest.importorskip("psycopg") from test_pg_extension import _lit # noqa: E402 N, D, NQ = 100_000, 64, 500 SELECTIVITIES = {"c10": 0.10, "c1": 0.01, "c01": 0.001} # category column -> share of the rows def _connect(): try: c = psycopg.connect(autocommit=True, connect_timeout=3) c.execute("SET statement_timeout = '600s'") c.execute("SET client_min_messages = warning") return c except Exception as e: # noqa: BLE001 pytest.skip(f"postgres not reachable: {e}") @pytest.fixture(scope="module") def conn(): c = _connect() c.execute("CREATE EXTENSION IF NOT EXISTS vector") try: c.execute("CREATE EXTENSION IF NOT EXISTS trigpoint") except Exception as e: # noqa: BLE001 pytest.skip(f"trigpoint extension not installed: {e}") c.execute("LOAD 'trigpoint'") # the custom settings exist once the library is loaded in this backend if c.execute("SELECT count(*) FROM pg_settings WHERE name = 'trigpoint.iterative_scan'").fetchone()[0] == 0: pytest.skip("trigpoint < 0.5: no iterative scan") c.sanitized = "sanitizer build" in c.execute("SELECT version()").fetchone()[0] yield c c.close() def _corpus(): rng = np.random.default_rng(7) cents = rng.normal(size=(256, D)).astype(np.float32) * 2.5 x = (cents[rng.integers(0, 256, N)] + rng.normal(size=(N, D))).astype(np.float32) q = (cents[rng.integers(0, 256, NQ)] + rng.normal(size=(NQ, D))).astype(np.float32) # the category columns: 10 / 100 / 1000 distinct values, uniform, so every value has the stated share cat = {name: rng.integers(0, int(round(1 / s)), N).astype(np.int32) for name, s in SELECTIVITIES.items()} return x, q, cat def _exact(x, q, mask, k): xs = x[mask] ids = np.flatnonzero(mask) d2 = (xs ** 2).sum(1)[None, :] - 2 * q @ xs.T + (q ** 2).sum(1)[:, None] order = np.argsort(d2, axis=1)[:, :k] return ids[order], np.take_along_axis(d2, order, 1) @pytest.fixture(scope="module") def table(conn): x, q, cat = _corpus() cur = conn.cursor() cur.execute("DROP TABLE IF EXISTS tp_filter") cur.execute(f"CREATE TABLE tp_filter (id int PRIMARY KEY, c10 int, c1 int, c01 int, v vector({D})) WITH (autovacuum_enabled = off)") dt = np.dtype([("ncols", ">i2"), ("l1", ">i4"), ("id", ">i4"), ("l2", ">i4"), ("c10", ">i4"), ("l3", ">i4"), ("c1", ">i4"), ("l4", ">i4"), ("c01", ">i4"), ("l5", ">i4"), ("dim", ">u2"), ("unused", ">u2"), ("v", ">f4", (D,))]) rec = np.empty(N, dt) rec["ncols"] = 5 rec["l1"] = rec["l2"] = rec["l3"] = rec["l4"] = 4 rec["id"] = np.arange(N); rec["c10"] = cat["c10"]; rec["c1"] = cat["c1"]; rec["c01"] = cat["c01"] rec["l5"] = 4 + 4 * D; rec["dim"] = D; rec["unused"] = 0; rec["v"] = x with cur.copy("COPY tp_filter (id, c10, c1, c01, v) FROM STDIN (FORMAT BINARY)") as cp: cp.write(b"PGCOPY\n\xff\r\n\0" + struct.pack(">ii", 0, 0)); cp.write(rec.tobytes()); cp.write(struct.pack(">h", -1)) cur.execute("CREATE INDEX tp_filter_ix ON tp_filter USING trigpoint (v vector_l2_ops) WITH (calibration_queries = 500, seed = 1)") cur.execute("ANALYZE tp_filter") yield {"x": x, "q": q, "cat": cat, "cur": cur, "sanitized": conn.sanitized} cur.execute("DROP TABLE IF EXISTS tp_filter") def _settings(cur, mode="relaxed_order", **kw): cur.execute("SET enable_seqscan = off") cur.execute("SET enable_bitmapscan = off") cur.execute("SET enable_indexscan = on") cur.execute(f"SET trigpoint.iterative_scan = {mode}") cur.execute("SET trigpoint.max_scan_tuples = %d" % kw.get("max_scan_tuples", 20000)) cur.execute("SET trigpoint.max_probes = %d" % kw.get("max_probes", 0)) cur.execute("SET trigpoint.exact = %s" % ("on" if kw.get("exact") else "off")) def _filtered(cur, col, val, qv, k=10): cur.execute(f"SELECT id, v <-> %s::vector FROM tp_filter WHERE {col} = %s ORDER BY v <-> %s::vector LIMIT {k}", (_lit(qv), int(val), _lit(qv)), prepare=False) return cur.fetchall() def _unfiltered(cur, qv, k=10): cur.execute(f"SELECT id, v <-> %s::vector FROM tp_filter ORDER BY v <-> %s::vector LIMIT {k}", (_lit(qv), _lit(qv)), prepare=False) return cur.fetchall() def test_filtered_query_uses_the_index(table): cur = table["cur"] _settings(cur) cur.execute("EXPLAIN SELECT id FROM tp_filter WHERE c1 = 3 ORDER BY v <-> %s::vector LIMIT 10", (_lit(table["q"][0]),), prepare=False) plan = "\n".join(r[0] for r in cur.fetchall()) assert "Index Scan using tp_filter_ix" in plan and "Filter: (c1 = 3)" in plan, plan @pytest.mark.parametrize("col,full_floor,recall_floor,lat_x", [("c10", 0.99, 0.95, 10), ("c1", 0.99, 0.95, 30), ("c01", 0.99, 0.90, 120)]) def test_filtered_scan_comes_back_full_and_close_to_exact(table, col, full_floor, recall_floor, lat_x): """The acceptance of the plan: 10 rows for >= 99 % of the queries, recall@10 against the filtered exact answer >= 0.95 at 10 % and 1 %, >= 0.90 at 0.1 %, p50 latency within 10x the unfiltered scan at 10 %. At 1 % and 0.1 % the bound is 30x and 120x: a filter the index cannot see makes the executor look at ~10 / s rows for LIMIT 10 (1,000 and 10,000 against the 10 of the unfiltered scan), so the work itself is 100x and 1,000x; measured on the local loop (64 dims, 100k rows) the scan is at 4x, 18x and 61x the unfiltered 0.4 ms (2026-11-03; 20x at 1 % under the load of a full suite run, hence the margin).""" cur, x, q, cat = table["cur"], table["x"], table["q"], table["cat"][col] nval = int(round(1 / SELECTIVITIES[col])) vals = np.random.default_rng(3).integers(0, nval, NQ) _settings(cur) t_unf = [] for qv in q[:100]: t = time.perf_counter(); _unfiltered(cur, qv); t_unf.append(time.perf_counter() - t) full, recall, lat = 0, [], [] for qv, val in zip(q, vals): gt, _ = _exact(x, qv[None, :], cat == val, 10) t = time.perf_counter(); rows = _filtered(cur, col, val, qv); lat.append(time.perf_counter() - t) got = [r[0] for r in rows] full += len(got) == 10 recall.append(len(set(got) & set(gt[0].tolist())) / 10) p50, p50_unf = np.median(lat) * 1e3, np.median(t_unf) * 1e3 print(f"\n{col}: full {full}/{NQ}, recall@10 {np.mean(recall):.4f}, p50 {p50:.2f} ms vs unfiltered {p50_unf:.2f} ms ({p50 / p50_unf:.1f}x)") assert np.mean(recall) >= recall_floor, np.mean(recall) if not table["sanitized"]: # the sanitizers slow the core's cascade more than the executor: the ratio means nothing there assert p50 <= lat_x * p50_unf, (p50, p50_unf) if col != "c01": assert full >= full_floor * NQ, (full, NQ) return # 0.1 %: the default cap of 20,000 fetched rows holds ~20 matches, so a few queries come back short (6 of 500 on the # local loop); the cap is the operator's knob: at 50,000 the scan comes back full assert full >= 0.98 * NQ, (full, NQ) _settings(cur, max_scan_tuples=50000) full = sum(len(_filtered(cur, col, val, qv)) == 10 for qv, val in zip(q[:200], vals[:200])) print(f"{col} at max_scan_tuples = 50000: full {full}/200") assert full >= full_floor * 200, full def test_always_true_filter_matches_or_beats_the_plain_widening(table): """Filter = always true, iterative on, LIMIT 10 / 100 / 1000 against the exact top-N: recall at least that of the 0.4 widening (iterative_scan = off, k doubled over the same cells) at the same N.""" cur, x, q = table["cur"], table["x"], table["q"] for n in (10, 100, 1000): rec = {} for mode in ("off", "relaxed_order"): _settings(cur, mode) hits = [] for qv in q[:50]: gt, _ = _exact(x, qv[None, :], np.ones(N, bool), n) cur.execute(f"SELECT id FROM tp_filter WHERE id >= 0 ORDER BY v <-> %s::vector LIMIT {n}", (_lit(qv),), prepare=False) got = [r[0] for r in cur.fetchall()] assert len(got) == n, (mode, n, len(got)) hits.append(len(set(got) & set(gt[0].tolist())) / n) rec[mode] = float(np.mean(hits)) print(f"\nLIMIT {n}: recall off {rec['off']:.4f} relaxed_order {rec['relaxed_order']:.4f}") assert rec["relaxed_order"] >= rec["off"] - 1e-9, (n, rec) if n == 10: assert rec["relaxed_order"] == rec["off"], rec # the first batch is the 0.4 computation def test_first_batch_is_the_plain_scan(table): """Unfiltered LIMIT 10 returns the same ids whatever the mode: the first batch is the 0.4 computation.""" cur, q = table["cur"], table["q"] out = {} for mode in ("off", "relaxed_order", "strict_order"): _settings(cur, mode) out[mode] = [[r[0] for r in _unfiltered(cur, qv)] for qv in q[:100]] assert out["off"] == out["relaxed_order"] == out["strict_order"] def test_caps_stop_the_scan_short_without_error(table): cur, q, cat = table["cur"], table["q"], table["cat"]["c01"] _settings(cur, max_scan_tuples=200) short = 0 for qv in q[:50]: rows = _filtered(cur, "c01", 5, qv) short += len(rows) < 10 assert short >= 40, short # 200 fetched rows at 0.1 % hold ~0.2 matches: nearly every query comes back short _settings(cur, max_probes=32) # the cells opened never grow past the first batch's 32 (~3,000 rows: ~3 matches at 0.1 %) short = 0 for qv in q[:50]: rows = _filtered(cur, "c01", 5, qv) short += len(rows) < 10 assert short >= 40, short _settings(cur) assert sum(len(_filtered(cur, "c01", 5, qv)) == 10 for qv in q[:50]) >= 45 # back to the defaults: full but for the cap's few def test_strict_order_is_non_decreasing_and_relaxed_may_not_be(table): cur, q, cat = table["cur"], table["q"], table["cat"]["c1"] _settings(cur, "strict_order") for qv in q[:200]: rows = _filtered(cur, "c1", 7, qv, k=50) d = [r[1] for r in rows] assert d == sorted(d), d # strict_order + exact: every returned prefix is the exact filtered top-N _settings(cur, "strict_order", exact=True) x = table["x"] for qv in q[:20]: gt, _ = _exact(x, qv[None, :], cat == 7, 50) rows = _filtered(cur, "c1", 7, qv, k=50) assert [r[0] for r in rows] == gt[0].tolist() def test_off_is_the_old_widening(table): """iterative_scan = off: the cells opened never grow, so the scan stops after at most ceil(log2(n / k)) + 1 batches with whatever the first cells hold; a 1 % filter at LIMIT 100 comes back short more often than relaxed_order (which opens more cells).""" cur, q = table["cur"], table["q"] short = {} for mode in ("off", "relaxed_order"): _settings(cur, mode) short[mode] = sum(len(_filtered(cur, "c1", 11, qv, k=100)) < 100 for qv in q[:50]) print(f"\nshort at LIMIT 100, 1 %: off {short['off']}/50, relaxed_order {short['relaxed_order']}/50") assert short["relaxed_order"] <= short["off"] assert short["relaxed_order"] <= 1 def test_debug_line_names_the_mode_and_the_batches(conn, table): q = table["q"] with conn.cursor() as cur: notices = [] conn.add_notice_handler(lambda n: notices.append(n.message_primary)) _settings(cur, "relaxed_order") cur.execute("SET client_min_messages = notice") cur.execute("SET trigpoint.debug = on") _filtered(cur, "c01", 1, q[0]) cur.execute("SET trigpoint.debug = off") cur.execute("SET client_min_messages = warning") assert any("batch=1 " in n and "iterative=relaxed_order" in n for n in notices), notices[:3] assert any("batch=2" in n for n in notices), notices[:3] assert any("(last)" in n for n in notices), notices[-1:] def test_cursor_across_batches_sees_every_row_once(conn, table): """A cursor fetched a few rows at a time crosses batch boundaries: no row twice, none skipped versus one shot.""" q = table["q"] with conn.cursor() as cur: _settings(cur, "relaxed_order") one = [r[0] for r in _filtered(cur, "c1", 2, q[1], k=100)] assert len(one) == 100 cur.execute("BEGIN") cur.execute("DECLARE cf CURSOR FOR SELECT id FROM tp_filter WHERE c1 = 2 ORDER BY v <-> %s::vector LIMIT 100", (_lit(q[1]),)) got = [] while True: cur.execute("FETCH 7 FROM cf") rows = cur.fetchall() if not rows: break got += [r[0] for r in rows] cur.execute("COMMIT") assert len(got) == 100 and len(set(got)) == 100 assert got == one