From 9b7e0aad3c19c83ad78e0516419dae6e1ac54fc7 Mon Sep 17 00:00:00 2001 From: spenc-r Date: Thu, 21 May 2026 17:25:30 +0000 Subject: [PATCH 1/4] feat(vchordrq): estimate filtered candidate selectivity Use the planner's filtered row estimate as filter selectivity for deciding how many IVFFlat candidates are needed to satisfy LIMIT after heap-side quals. Report indexSelectivity from the expected retrieved candidate count instead of the final survivor fraction, so cost_index() prices heap fetch and qpqual evaluation against candidate rows rather than every table row or only the rows that pass filters. Clamp the candidate-processing term by node_count, and tighten the cost estimator sqllogictest coverage around btree alternatives, candidate-cost floors, and extreme selectivity. --- src/index/vchordrq/am/mod.rs | 58 ++-- tests/vchordrq/cost_estimator.slt | 426 ++++++++++++++++++++++++++++++ 2 files changed, 467 insertions(+), 17 deletions(-) create mode 100644 tests/vchordrq/cost_estimator.slt diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index 5040b587..f995c600 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -292,20 +292,31 @@ pub unsafe extern "C-unwind" fn amcostestimate( *index_pages = 1.0; return; } - let selectivity = { - use pgrx::pg_sys::{ - JoinType, add_predicate_to_index_quals, clauselist_selectivity, - get_quals_from_indexclauses, + // Vchordrq indexes only the vector column, so ordinary filters on + // other columns are heap-side quals rather than index quals. We use + // the planner's filtered row estimate to decide how many distance + // candidates we likely need to produce enough survivors for LIMIT. + // + // Keep this separate from the value returned as `indexSelectivity`: + // `cost_index()` interprets that output as the fraction of parent + // table rows the index scan retrieves, not the fraction surviving all + // filters. With LIMIT, the retrieved fraction is better represented + // by the candidate count computed below. + let (total_rows, filter_selectivity) = { + let baserel = (*index_opt_info).rel; + let total_rows = (*baserel).tuples; + let param_info = (*path).path.param_info; + let filtered_rows = if !param_info.is_null() { + (*param_info).ppi_rows + } else { + (*baserel).rows }; - let index_quals = get_quals_from_indexclauses((*path).indexclauses); - let selectivity_quals = add_predicate_to_index_quals(index_opt_info, index_quals); - clauselist_selectivity( - root, - selectivity_quals, - (*(*index_opt_info).rel).relid as _, - JoinType::JOIN_INNER, - std::ptr::null_mut(), - ) + let filter_selectivity = if total_rows > 0.0 { + (filtered_rows / total_rows).clamp(1e-9, 1.0) + } else { + 1.0 + }; + (total_rows, filter_selectivity) }; // index exists if !(*index_opt_info).hypothetical { @@ -367,18 +378,31 @@ pub unsafe extern "C-unwind" fn amcostestimate( pages += cost.cells[0] as f64; pages }; - let next_count = - f64::max(1.0, (*root).limit_tuples) * f64::min(1000.0, 1.0 / selectivity); + // `next_count` represents candidates we expect to process to + // surface `limit_tuples` survivors after filter rejection. Clamp + // by `node_count` so the estimate cannot exceed the candidates + // the IVF visits at the configured probe count. + let next_count = if (*root).limit_tuples > 0.0 { + ((*root).limit_tuples * f64::min(1000.0, 1.0 / filter_selectivity)) + .min(node_count) + } else { + node_count + }; + let scan_selectivity = if total_rows > 0.0 { + (next_count / total_rows).clamp(1e-9, 1.0) + } else { + 1.0 + }; *index_startup_cost = 0.001 * node_count; *index_total_cost = 0.001 * node_count + next_count; - *index_selectivity = selectivity; + *index_selectivity = scan_selectivity; *index_correlation = 0.0; *index_pages = page_count; return; } *index_startup_cost = 0.0; *index_total_cost = 0.0; - *index_selectivity = selectivity; + *index_selectivity = filter_selectivity; *index_correlation = 0.0; *index_pages = 1.0; } diff --git a/tests/vchordrq/cost_estimator.slt b/tests/vchordrq/cost_estimator.slt new file mode 100644 index 00000000..0f4ad266 --- /dev/null +++ b/tests/vchordrq/cost_estimator.slt @@ -0,0 +1,426 @@ +# Tests for amcostestimate in src/index/vchordrq/am/mod.rs. +# +# The cost estimator must: +# 1. Use the baserel's planner-computed row estimate to derive filter +# selectivity, then report index selectivity from the candidate count +# that the scan is expected to retrieve. +# 2. Cap the candidate-processing term (`next_count`) by the IVF candidate +# budget (`node_count`), so absurdly small selectivity cannot inflate +# the reported cost past the work the index actually does. +# 3. Degrade safely when planner stats are missing or extreme. +# +# --------------------------------------------------------------------------- +# Setup +# --------------------------------------------------------------------------- + +statement ok +SET enable_seqscan = off; + +statement ok +SET max_parallel_workers_per_gather = 0; + +statement ok +SET random_page_cost = 1.1; + +statement ok +CREATE TABLE cost_test ( + id int PRIMARY KEY, + category int NOT NULL, + v vector(3) NOT NULL +); + +# 10 000 rows, 100 evenly-sized categories. category=k matches 1% of rows. +# Vectors are deterministic per row so the test is reproducible. +statement ok +INSERT INTO cost_test +SELECT i, + (i % 100), + ARRAY[ + (i % 97) / 97.0, + (i % 89) / 89.0, + (i % 83) / 83.0 + ]::real[]::vector +FROM generate_series(1, 10000) i; + +statement ok +CREATE INDEX cost_test_v ON cost_test USING vchordrq (v vector_l2_ops); + +statement ok +CREATE INDEX cost_test_cat ON cost_test (category); + +# A predicate with an explicitly high procost that PG can't optimize away. +# Three things matter here: +# * `LANGUAGE plpgsql` — prevents SQL inlining (an inlined `IS NOT NULL` +# check on a PRIMARY KEY column gets folded to `true` and dropped from +# the plan, defeating the test). +# * `COST 10000` — sets `pg_proc.procost` so cost_index() charges +# 10000 * cpu_operator_cost ≈ 25 cost-units per tuple of filter eval. +# This mimics PostGIS ST_DWithin, JSONB containment, expensive regex, +# etc. — the real-world shapes where the bug bites. +# * Body is `mod($1,7) <> 99` so the result is always true but PG can't +# prove that across the function boundary, so the filter survives +# constant-folding. +statement ok +CREATE OR REPLACE FUNCTION slow_true(int) RETURNS boolean + LANGUAGE plpgsql IMMUTABLE PARALLEL SAFE COST 10000 AS +$$ BEGIN RETURN mod($1, 7) <> 99; END $$; + +statement ok +ANALYZE cost_test; + +# Helper: returns the index name used at the top of the plan, or NULL if +# no Index Scan / Bitmap Index Scan appears. Robust across PG versions. +statement ok +CREATE OR REPLACE FUNCTION plan_top_index(query text) RETURNS text + LANGUAGE plpgsql AS $$ +DECLARE + line text; + m text[]; +BEGIN + FOR line IN EXECUTE 'EXPLAIN (FORMAT TEXT, COSTS OFF) ' || query LOOP + -- 'Index Scan using on ...' or 'Index Only Scan using on ...' + m := regexp_match(line, 'Index (?:Only )?Scan using (\S+) on '); + IF m IS NOT NULL THEN RETURN m[1]; END IF; + -- 'Bitmap Index Scan on ' + m := regexp_match(line, 'Bitmap Index Scan on (\S+)'); + IF m IS NOT NULL THEN RETURN m[1]; END IF; + END LOOP; + RETURN NULL; +END $$; + +# Helper: returns the top plan node's total cost. +statement ok +CREATE OR REPLACE FUNCTION top_total_cost(query text) RETURNS double precision + LANGUAGE plpgsql AS $$ +DECLARE + rec text; + m text[]; +BEGIN + FOR rec IN EXECUTE 'EXPLAIN ' || query LOOP + m := regexp_match(rec, '\.\.([0-9]+(?:\.[0-9]+)?) rows='); + IF m IS NOT NULL THEN RETURN m[1]::double precision; END IF; + END LOOP; + RETURN NULL; +END $$; + +# Helper: returns total cost for the first plan node matching `pat`. +statement ok +CREATE OR REPLACE FUNCTION plan_node_total_cost(query text, pat text) RETURNS double precision + LANGUAGE plpgsql AS $$ +DECLARE + rec text; + m text[]; +BEGIN + FOR rec IN EXECUTE 'EXPLAIN ' || query LOOP + IF rec LIKE pat THEN + m := regexp_match(rec, '\.\.([0-9]+(?:\.[0-9]+)?) rows='); + IF m IS NOT NULL THEN RETURN m[1]::double precision; END IF; + END IF; + END LOOP; + RETURN NULL; +END $$; + +# --------------------------------------------------------------------------- +# Case 1: pure ORDER BY, no filter. +# Vchord must win. (Passes pre- and post-fix; sanity check.) +# --------------------------------------------------------------------------- + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +); +---- +cost_test_v + +# --------------------------------------------------------------------------- +# Case 2: tight filter (selectivity 1%) with an alternative index. +# Btree must still beat vchord — the IVF candidate budget is too large +# to be competitive at this selectivity. (Passes pre- and post-fix.) +# This guards against over-correction: the fix must not make vchord win +# on tight filters, which was the original motivation for PR #234. +# --------------------------------------------------------------------------- + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE category = 42 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +); +---- +cost_test_cat + +statement ok +DROP INDEX cost_test_cat; + +# --------------------------------------------------------------------------- +# Case 3: selective heap filter combined with an expensive predicate and no +# supporting non-vector index. Pre-fix: index_selectivity=1.0 charges +# filter eval against every table row, so seq scan + sort can look cheaper. +# Post-fix: vchord estimates the candidates needed for the LIMIT and wins. +# The node-cost floor guards against under-pricing the scan as if it fetched +# only the final surviving rows. +# --------------------------------------------------------------------------- + +statement ok +SET enable_seqscan = on; + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE category = 42 AND slow_true(id) + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +); +---- +cost_test_v + +query B +SELECT plan_node_total_cost( + 'SELECT id FROM cost_test WHERE category = 42 AND slow_true(id) + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10', + '%Index Scan using cost_test_v on cost_test%' +) > 10000; +---- +t + +statement ok +SET enable_seqscan = off; + +# --------------------------------------------------------------------------- +# Case 4: filter on a column with no supporting non-vector index, combined +# with the expensive predicate. The only alternatives are seq scan (disabled) +# or vchord. Vchord must be chosen. Pre-fix this still picks vchord because +# seq is disabled, but the *reported* cost must drop enough that LIMIT-aware +# planners above this node (joins, gather, partitions) don't overestimate. +# We assert plan shape + that the cost is bounded below 1e9 (disable_cost). +# --------------------------------------------------------------------------- + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE slow_true(id) + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +); +---- +cost_test_v + +# disable_cost in PG 14-18 is 1.0e10; a healthy estimate is well below that. +query B +SELECT top_total_cost( + 'SELECT id FROM cost_test WHERE slow_true(id) + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +) < 1e9; +---- +t + +# --------------------------------------------------------------------------- +# Case 5: no LIMIT clause. PlannerInfo.limit_tuples is -1, so the +# estimator should fall back to the IVF candidate budget instead of doing +# LIMIT/selectivity math. +# --------------------------------------------------------------------------- + +query B +SELECT top_total_cost( + 'SELECT id FROM cost_test ORDER BY v <-> ''[0,0,0]''::vector' +) IS NOT NULL; +---- +t + +# --------------------------------------------------------------------------- +# Case 6: LIMIT larger than the number of expected survivors. +# The `next_count.min(node_count)` clamp must hold — the AM cannot claim +# to process more candidates than the IVF visits. +# --------------------------------------------------------------------------- + +query B +SELECT top_total_cost( + 'SELECT id FROM cost_test WHERE category = 42 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 100000' +) IS NOT NULL; +---- +t + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE category + 0 = 42 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 100000' +); +---- +cost_test_v + +query B +SELECT plan_node_total_cost( + 'SELECT id FROM cost_test WHERE category + 0 = 42 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 100000', + '%Index Scan using cost_test_v on cost_test%' +) < 100000; +---- +t + +# --------------------------------------------------------------------------- +# Case 7: near-zero filter selectivity. The fix clamps to 1e-9 instead of +# 0 so 1/selectivity does not produce +Inf. +# --------------------------------------------------------------------------- + +query B +SELECT top_total_cost( + 'SELECT id FROM cost_test WHERE category = -1 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +) IS NOT NULL; +---- +t + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE category = -1 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +); +---- +cost_test_v + +query B +SELECT plan_node_total_cost( + 'SELECT id FROM cost_test WHERE category = -1 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10', + '%Index Scan using cost_test_v on cost_test%' +) < 100000; +---- +t + +# --------------------------------------------------------------------------- +# Case 8: empty table. baserel->tuples is 0, fix falls back to +# selectivity = 1.0. Plan and result must not crash. +# --------------------------------------------------------------------------- + +statement ok +CREATE TABLE cost_test_empty (id int, v vector(3)); + +statement ok +CREATE INDEX ON cost_test_empty USING vchordrq (v vector_l2_ops); + +statement ok +ANALYZE cost_test_empty; + +query I +SELECT count(*) FROM ( + SELECT id FROM cost_test_empty ORDER BY v <-> '[0,0,0]'::vector LIMIT 10 +) t; +---- +0 + +# --------------------------------------------------------------------------- +# Case 9: never-ANALYZE'd table. baserel->tuples may be negative or zero +# in the planner's view. The fallback in the fix returns 1.0, matching the +# pre-fix behavior exactly so we don't regress on cold tables. +# --------------------------------------------------------------------------- + +statement ok +CREATE TABLE cost_test_cold (id int, v vector(3)); + +statement ok +INSERT INTO cost_test_cold +SELECT i, ARRAY[(i%7)/7.0, (i%11)/11.0, (i%13)/13.0]::real[]::vector +FROM generate_series(1, 100) i; + +statement ok +CREATE INDEX ON cost_test_cold USING vchordrq (v vector_l2_ops); + +# Intentionally no ANALYZE. + +query B +SELECT top_total_cost( + 'SELECT id FROM cost_test_cold ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +) IS NOT NULL; +---- +t + +# --------------------------------------------------------------------------- +# Case 10: partial vchord index. The filter-selectivity estimate reflects +# baserestrictinfo, including clauses that overlap the partial-index +# predicate. We verify that this builds, plans, and runs. +# --------------------------------------------------------------------------- + +statement ok +CREATE TABLE cost_test_partial (id int, kind int, v vector(3)); + +statement ok +INSERT INTO cost_test_partial +SELECT i, (i % 3), ARRAY[(i%7)/7.0, (i%11)/11.0, (i%13)/13.0]::real[]::vector +FROM generate_series(1, 2000) i; + +statement ok +CREATE INDEX cost_test_partial_v ON cost_test_partial USING vchordrq (v vector_l2_ops) + WHERE kind = 0; + +statement ok +ANALYZE cost_test_partial; + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test_partial WHERE kind = 0 + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 5' +); +---- +cost_test_partial_v + +# --------------------------------------------------------------------------- +# Case 11: vchordrq.enable_scan = off must still disable the path +# regardless of selectivity. Re-enable seqscan inside this case so there is +# a non-disabled alternative for the planner to fall back to — otherwise on +# PG18 every path looks "disabled" and the planner picks vchord by default. +# Use IS DISTINCT FROM so a seq-scan top (plan_top_index returns NULL) +# still compares as "not vchord." +# --------------------------------------------------------------------------- + +statement ok +SET enable_seqscan = on; + +statement ok +SET vchordrq.enable_scan = off; + +query B +SELECT plan_top_index( + 'SELECT id FROM cost_test WHERE category = 42 AND slow_true(id) + ORDER BY v <-> ''[0,0,0]''::vector LIMIT 10' +) IS DISTINCT FROM 'cost_test_v'; +---- +t + +statement ok +RESET vchordrq.enable_scan; + +statement ok +RESET enable_seqscan; + +# --------------------------------------------------------------------------- +# Cleanup +# --------------------------------------------------------------------------- + +statement ok +DROP TABLE cost_test; + +statement ok +DROP TABLE cost_test_empty; + +statement ok +DROP TABLE cost_test_cold; + +statement ok +DROP TABLE cost_test_partial; + +statement ok +DROP FUNCTION slow_true(int); + +statement ok +DROP FUNCTION plan_top_index(text); + +statement ok +DROP FUNCTION plan_node_total_cost(text, text); + +statement ok +DROP FUNCTION top_total_cost(text); + +statement ok +RESET enable_seqscan; + +statement ok +RESET max_parallel_workers_per_gather; + +statement ok +RESET random_page_cost; From 695a4a2d46d23b0f26ed86fac2f240318a76622c Mon Sep 17 00:00:00 2001 From: Amey Pawar <138877912+ameyypawar@users.noreply.github.com> Date: Tue, 30 Jun 2026 17:57:45 +0530 Subject: [PATCH 2/4] fix(vchordg): avoid panic when validating an empty alpha option `validate_alpha` indexes `alpha[0]` after the sorted/range checks, both of which pass on an empty slice (`[].is_sorted()` and `[].iter().all(..)` are true). So an empty `alpha` (e.g. `WITH (options='[index] alpha = []')`) panics with an index-out-of-bounds instead of returning a ValidationError. Use `alpha.first()` so the empty case yields the existing "`alpha` should contain `1.0`" error. Behavior is unchanged for every non-empty input. Add unit tests covering the empty slice plus valid/invalid cases. --- crates/vchordg/src/types.rs | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/crates/vchordg/src/types.rs b/crates/vchordg/src/types.rs index 62ec3cc4..5a5a6169 100644 --- a/crates/vchordg/src/types.rs +++ b/crates/vchordg/src/types.rs @@ -56,7 +56,7 @@ impl VchordgIndexOptions { if !alpha.iter().all(|x| (1.0..2.0).contains(x)) { return Err(ValidationError::new("alpha is too large or too small")); } - if alpha[0] != 1.0 { + if alpha.first() != Some(&1.0) { return Err(ValidationError::new("`alpha` should contain `1.0`")); } Ok(()) @@ -144,3 +144,25 @@ impl Structure { self.children.is_empty() } } + +#[cfg(test)] +mod tests { + use super::VchordgIndexOptions; + + #[test] + fn validate_alpha_handles_empty_without_panicking() { + let empty: &[f32] = &[]; + assert!(VchordgIndexOptions::validate_alpha(empty).is_err()); + } + + #[test] + fn validate_alpha_accepts_default() { + assert!(VchordgIndexOptions::validate_alpha(&[1.0, 1.2]).is_ok()); + } + + #[test] + fn validate_alpha_rejects_unsorted_or_out_of_range() { + assert!(VchordgIndexOptions::validate_alpha(&[1.2, 1.0]).is_err()); + assert!(VchordgIndexOptions::validate_alpha(&[2.0]).is_err()); + } +} From 24d081cd56257b6aac8124249cbd2505d40f95d8 Mon Sep 17 00:00:00 2001 From: Amey Pawar <138877912+ameyypawar@users.noreply.github.com> Date: Thu, 9 Jul 2026 07:52:57 +0530 Subject: [PATCH 3/4] fix(gucs): report an error instead of panicking on an empty `probes` entry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The `vchordrq.probes` parser panics via `.expect("empty probes")` when an entry is empty (e.g. `SET vchordrq.probes = ',5'` or `'1,,2'`), even though the sibling match arm already reports malformed input cleanly with `pgrx::error!`. Handle the empty entry the same way — a proper error rather than a panic. Behavior is unchanged for all valid input. --- src/index/gucs.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/index/gucs.rs b/src/index/gucs.rs index 8c0e9fd3..f2724593 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -388,7 +388,10 @@ pub unsafe fn vchordrq_probes(index: pgrx::pg_sys::Relation) -> Vec { for &c in value.to_bytes() { match c { b' ' => continue, - b',' => result.push(current.take().expect("empty probes")), + b',' => match current.take() { + Some(value) => result.push(value), + None => pgrx::error!("empty entry in probes"), + }, b'0'..=b'9' => { if let Some(x) = current.as_mut() { *x = *x * 10 + (c - b'0') as u32; From 93ae7f7360f4de891e3b482be4f6aba4feebbafc Mon Sep 17 00:00:00 2001 From: jasonlee Date: Mon, 13 Jul 2026 18:40:47 +0800 Subject: [PATCH 4/4] feat(vchordrq): add exact TileMaxSim reranking --- Cargo.lock | 1 + Cargo.toml | 3 + crates/vchordrq/src/build.rs | 1 + crates/vchordrq/src/bulkdelete.rs | 28 +- crates/vchordrq/src/cost.rs | 8 +- crates/vchordrq/src/lib.rs | 6 + crates/vchordrq/src/maxsim_cost.rs | 206 +++ crates/vchordrq/src/statistics.rs | 30 + crates/vchordrq/src/tuples.rs | 151 +- crates/vchordrq/src/types.rs | 34 + devtools/test_tilemaxsim_reference_sidecar.py | 383 +++++ devtools/tilemaxsim_reference_sidecar.py | 780 +++++++++ services/Dockerfile.tilemaxsim | 12 + services/benchmark_tilemaxsim_cuda.py | 152 ++ services/build_tilemaxsim_tensor_cache.py | 235 +++ services/test_tilemaxsim_cuda_sidecar.py | 313 ++++ services/tilemaxsim_cuda_sidecar.py | 802 +++++++++ src/datatype/mod.rs | 14 + src/datatype/operators_halfvec.rs | 8 + src/datatype/operators_rabitq4.rs | 8 + src/datatype/operators_rabitq8.rs | 8 + src/datatype/operators_vector.rs | 8 + src/index/fetcher.rs | 108 +- src/index/gucs.rs | 139 ++ src/index/vchordrq/am/am_build.rs | 31 +- src/index/vchordrq/am/mod.rs | 89 +- src/index/vchordrq/dispatch.rs | 27 +- src/index/vchordrq/opclass.rs | 8 + src/index/vchordrq/scanners/maxsim.rs | 1345 ++++++++------- .../vchordrq/scanners/maxsim/candidate.rs | 192 +++ .../vchordrq/scanners/maxsim/external.rs | 551 ++++++ src/index/vchordrq/scanners/maxsim/gpu.rs | 1495 +++++++++++++++++ src/index/vchordrq/scanners/maxsim/rerank.rs | 315 ++++ src/index/vchordrq/scanners/maxsim/search.rs | 578 +++++++ src/index/vchordrq/scanners/mod.rs | 8 + src/sql/finalize.sql | 821 +++++++++ tests/vchordrq/cost_estimator.slt | 171 ++ tests/vchordrq/maxsim_correctness.slt | 221 +++ tests/vchordrq/maxsim_source_registry.slt | 341 ++++ 39 files changed, 8933 insertions(+), 698 deletions(-) create mode 100644 crates/vchordrq/src/maxsim_cost.rs create mode 100644 crates/vchordrq/src/statistics.rs create mode 100644 devtools/test_tilemaxsim_reference_sidecar.py create mode 100644 devtools/tilemaxsim_reference_sidecar.py create mode 100644 services/Dockerfile.tilemaxsim create mode 100644 services/benchmark_tilemaxsim_cuda.py create mode 100644 services/build_tilemaxsim_tensor_cache.py create mode 100644 services/test_tilemaxsim_cuda_sidecar.py create mode 100644 services/tilemaxsim_cuda_sidecar.py create mode 100644 src/index/vchordrq/scanners/maxsim/candidate.rs create mode 100644 src/index/vchordrq/scanners/maxsim/external.rs create mode 100644 src/index/vchordrq/scanners/maxsim/gpu.rs create mode 100644 src/index/vchordrq/scanners/maxsim/rerank.rs create mode 100644 src/index/vchordrq/scanners/maxsim/search.rs create mode 100644 tests/vchordrq/maxsim_correctness.slt create mode 100644 tests/vchordrq/maxsim_source_registry.slt diff --git a/Cargo.lock b/Cargo.lock index f324f59e..80722f2f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1937,6 +1937,7 @@ dependencies = [ "index", "index_accessor", "k_means", + "libc", "mimalloc", "paste", "pgrx", diff --git a/Cargo.toml b/Cargo.toml index 5868e4d6..32654e1f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -55,6 +55,9 @@ zerocopy.workspace = true [target.'cfg(all(any(target_arch = "x86_64", target_arch = "aarch64"), any(target_os = "linux", target_os = "macos")))'.dependencies] mimalloc = { version = "0.1.49", features = ["local_dynamic_tls"] } +[target.'cfg(unix)'.dependencies] +libc = "0.2.182" + [lints] workspace = true diff --git a/crates/vchordrq/src/build.rs b/crates/vchordrq/src/build.rs index a7371353..d622fbc6 100644 --- a/crates/vchordrq/src/build.rs +++ b/crates/vchordrq/src/build.rs @@ -119,6 +119,7 @@ pub fn build( height_of_root: structures.len() as u32, is_residual, rerank_in_heap: vchordrq_options.rerank_in_table, + indexed_vectors: Some(0), centroids_first: centroids.first(), vectors_first: vectors, centroid_prefetch: pointer_of_centroids diff --git a/crates/vchordrq/src/bulkdelete.rs b/crates/vchordrq/src/bulkdelete.rs index ae54ec54..0aae88fe 100644 --- a/crates/vchordrq/src/bulkdelete.rs +++ b/crates/vchordrq/src/bulkdelete.rs @@ -25,7 +25,8 @@ pub fn bulkdelete( index: &R, check: impl Fn(), callback: impl Fn(NonZero) -> bool, -) where +) -> u64 +where R::Page: Page, { let meta_guard = index.read(0); @@ -38,6 +39,8 @@ pub fn bulkdelete( drop(meta_guard); + let mut live = 0_u64; + let step = |state: State| { let mut results = Vec::new(); for first in state { @@ -64,19 +67,21 @@ pub fn bulkdelete( while current != u32::MAX { check(); let read = index.read(current); - let flag = 'flag: { + let (flag, page_live) = 'scan: { + let mut page_live = 0_u64; for i in 1..=read.len() { let bytes = read.get(i).expect("data corruption"); let tuple = FrozenTuple::deserialize_ref(bytes); if let FrozenTupleReader::_0(tuple) = tuple { for p in tuple.payload().iter() { if Some(true) == p.map(&callback) { - break 'flag true; + break 'scan (true, 0); } + page_live += u64::from(p.is_some()); } } } - false + (false, page_live) }; if flag { drop(read); @@ -89,9 +94,12 @@ pub fn bulkdelete( if Some(true) == p.map(&callback) { *p = None; } + live += u64::from(p.is_some()); } } } + } else { + live += page_live; } current = directory.next().unwrap_or(u32::MAX); } @@ -101,16 +109,18 @@ pub fn bulkdelete( while current != u32::MAX { check(); let read = index.read(current); - let flag = 'flag: { + let (flag, page_live) = 'scan: { + let mut page_live = 0_u64; for i in 1..=read.len() { let bytes = read.get(i).expect("data corruption"); let tuple = AppendableTuple::deserialize_ref(bytes); let p = tuple.payload(); if Some(true) == p.map(&callback) { - break 'flag true; + break 'scan (true, 0); } + page_live += u64::from(p.is_some()); } - false + (false, page_live) }; if flag { drop(read); @@ -122,14 +132,18 @@ pub fn bulkdelete( if Some(true) == p.map(&callback) { *p = None; } + live += u64::from(p.is_some()); } current = write.get_opaque().next; } else { + live += page_live; current = read.get_opaque().next; } } } } + + live } pub fn bulkdelete_vectors( diff --git a/crates/vchordrq/src/cost.rs b/crates/vchordrq/src/cost.rs index 60579401..c8e40252 100644 --- a/crates/vchordrq/src/cost.rs +++ b/crates/vchordrq/src/cost.rs @@ -18,6 +18,7 @@ use index::relation::{Page, RelationRead}; pub struct Cost { pub dim: u32, pub cells: Vec, + pub indexed_vectors: Option, } #[must_use] @@ -27,8 +28,13 @@ pub fn cost(index: &R) -> Cost { let meta_tuple = MetaTuple::deserialize_ref(meta_bytes); let dim = meta_tuple.dim(); let cells = meta_tuple.cells().to_vec(); + let indexed_vectors = meta_tuple.indexed_vectors(); drop(meta_guard); - Cost { dim, cells } + Cost { + dim, + cells, + indexed_vectors, + } } diff --git a/crates/vchordrq/src/lib.rs b/crates/vchordrq/src/lib.rs index d0efef05..b60c2e3b 100644 --- a/crates/vchordrq/src/lib.rs +++ b/crates/vchordrq/src/lib.rs @@ -24,9 +24,11 @@ mod freepages; mod insert; mod linked_vec; mod maintain; +mod maxsim_cost; mod prewarm; mod rerank; mod search; +mod statistics; mod tape; mod tape_writer; mod tuples; @@ -43,9 +45,13 @@ pub use cost::cost; pub use fast_heap::FastHeap; pub use insert::{InsertChooser, insert, insert_vector}; pub use maintain::{MaintainChooser, maintain}; +pub use maxsim_cost::{ + MaxsimCostBackend, MaxsimCostEstimate, MaxsimCostInput, estimate_maxsim_cost, +}; pub use prewarm::prewarm; pub use rerank::{how, rerank_heap, rerank_index}; pub use search::{default_search, maxsim_search}; +pub use statistics::set_indexed_vectors; use zerocopy::{FromBytes, Immutable, IntoBytes, KnownLayout}; diff --git a/crates/vchordrq/src/maxsim_cost.rs b/crates/vchordrq/src/maxsim_cost.rs new file mode 100644 index 00000000..ca843d10 --- /dev/null +++ b/crates/vchordrq/src/maxsim_cost.rs @@ -0,0 +1,206 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the terms of the ELv2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +#[derive(Clone, Copy, Debug)] +pub enum MaxsimCostBackend { + CoarseOnly, + CpuExact, + Gpu, + Auto, +} + +#[derive(Clone, Copy, Debug)] +pub struct MaxsimCostInput { + pub heap_rows: f64, + pub index_tokens: f64, + pub token_nodes_per_query: f64, + pub base_index_pages: f64, + pub dimension: u32, + pub element_bits: u32, + pub query_tokens: u32, + pub limit_tuples: Option, + pub filter_selectivity: f64, + pub candidate_limit: Option, + pub backend: MaxsimCostBackend, +} + +#[derive(Clone, Copy, Debug)] +pub struct MaxsimCostEstimate { + pub startup_cost: f64, + pub total_cost: f64, + pub selectivity: f64, + pub index_pages: f64, +} + +pub fn estimate_maxsim_cost(input: MaxsimCostInput) -> MaxsimCostEstimate { + let heap_rows = input.heap_rows.max(1.0); + let index_tokens = input.index_tokens.max(heap_rows); + let query_tokens = f64::from(input.query_tokens.max(1)); + let average_document_tokens = (index_tokens / heap_rows).clamp(1.0, 65_536.0); + let token_visits = input.token_nodes_per_query.max(1.0) * query_tokens; + + // We do not yet persist page-level candidate statistics. Until then, use a + // conservative occupancy estimate: token visits are sampled from the token + // index, and a page is a candidate when at least one of its average tokens + // is visited. This is intentionally bounded by the heap row count. + let token_visit_fraction = (token_visits / index_tokens).clamp(0.0, 1.0); + let candidate_probability = 1.0 - (1.0 - token_visit_fraction).powf(average_document_tokens); + let generated_pages = (heap_rows * candidate_probability).clamp(1.0, heap_rows); + + let filtered_limit = input + .limit_tuples + .map(|limit| limit.max(1.0) / input.filter_selectivity.clamp(1e-9, 1.0)); + let exact_candidate_count = input + .candidate_limit + .map_or(generated_pages, |limit| { + f64::from(limit).min(generated_pages) + }) + .clamp(1.0, heap_rows); + let returned_pages = match input.backend { + MaxsimCostBackend::CoarseOnly => filtered_limit + .unwrap_or(generated_pages) + .min(generated_pages), + MaxsimCostBackend::CpuExact | MaxsimCostBackend::Gpu | MaxsimCostBackend::Auto => { + exact_candidate_count + } + }; + + // Candidate generation and aggregation are eager in the current scanner, + // so their full work belongs to startup cost even when SQL has a small + // LIMIT. The constants are deliberately conservative placeholders until + // committed corpus benchmarks replace them with fitted values. + let search_cost = 0.001 * token_visits; + let aggregation_cost = 0.01 * token_visits + 0.05 * generated_pages; + let exact_components = exact_candidate_count + * average_document_tokens + * query_tokens + * f64::from(input.dimension.max(1)); + let tensor_bytes = exact_candidate_count + * average_document_tokens + * f64::from(input.dimension.max(1)) + * f64::from(input.element_bits.max(1)) + / 8.0; + let cpu_exact_cost = exact_candidate_count + exact_components * 1e-6; + let gpu_exact_cost = 5.0 + tensor_bytes * 1e-7 + exact_components * 5e-8; + let backend_cost = match input.backend { + MaxsimCostBackend::CoarseOnly => 0.0, + MaxsimCostBackend::CpuExact => cpu_exact_cost, + MaxsimCostBackend::Gpu => gpu_exact_cost, + // Price a small but nonzero fallback risk. Runtime still performs a + // complete CPU rerank on every GPU failure. + MaxsimCostBackend::Auto => gpu_exact_cost + 0.05 * cpu_exact_cost, + }; + let startup_cost = search_cost + aggregation_cost + backend_cost; + let total_cost = startup_cost + returned_pages; + let selectivity = (returned_pages / heap_rows).clamp(1e-9, 1.0); + let index_pages = + input.base_index_pages.max(1.0) * (1.0 + 0.25 * (query_tokens - 1.0).max(0.0)); + + MaxsimCostEstimate { + startup_cost, + total_cost, + selectivity, + index_pages, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn input(backend: MaxsimCostBackend) -> MaxsimCostInput { + MaxsimCostInput { + heap_rows: 34_054.0, + index_tokens: 25_438_338.0, + token_nodes_per_query: 10_000.0, + base_index_pages: 20_000.0, + dimension: 320, + element_bits: 16, + query_tokens: 32, + limit_tuples: Some(20.0), + filter_selectivity: 1.0, + candidate_limit: Some(256), + backend, + } + } + + #[test] + fn maxsim_is_never_zero_cost() { + for backend in [ + MaxsimCostBackend::CoarseOnly, + MaxsimCostBackend::CpuExact, + MaxsimCostBackend::Gpu, + MaxsimCostBackend::Auto, + ] { + let estimate = estimate_maxsim_cost(input(backend)); + assert!(estimate.startup_cost > 0.0); + assert!(estimate.total_cost >= estimate.startup_cost); + assert!((1e-9..=1.0).contains(&estimate.selectivity)); + assert!(estimate.index_pages >= 1.0); + } + } + + #[test] + fn query_token_count_increases_eager_work() { + let one = estimate_maxsim_cost(MaxsimCostInput { + query_tokens: 1, + ..input(MaxsimCostBackend::CpuExact) + }); + let many = estimate_maxsim_cost(MaxsimCostInput { + query_tokens: 64, + ..input(MaxsimCostBackend::CpuExact) + }); + assert!(many.startup_cost > one.startup_cost); + assert!(many.index_pages > one.index_pages); + } + + #[test] + fn exact_candidate_limit_bounds_rows_and_cost() { + let small = estimate_maxsim_cost(MaxsimCostInput { + candidate_limit: Some(128), + ..input(MaxsimCostBackend::CpuExact) + }); + let large = estimate_maxsim_cost(MaxsimCostInput { + candidate_limit: Some(2048), + ..input(MaxsimCostBackend::CpuExact) + }); + assert!(small.selectivity < large.selectivity); + assert!(small.startup_cost < large.startup_cost); + } + + #[test] + fn auto_prices_more_than_gpu_for_fallback_risk() { + let gpu = estimate_maxsim_cost(input(MaxsimCostBackend::Gpu)); + let auto = estimate_maxsim_cost(input(MaxsimCostBackend::Auto)); + assert!(auto.startup_cost > gpu.startup_cost); + } + + #[test] + fn missing_stats_remain_finite() { + let estimate = estimate_maxsim_cost(MaxsimCostInput { + heap_rows: -1.0, + index_tokens: 0.0, + token_nodes_per_query: 0.0, + base_index_pages: 0.0, + filter_selectivity: 0.0, + limit_tuples: None, + candidate_limit: None, + ..input(MaxsimCostBackend::CoarseOnly) + }); + assert!(estimate.startup_cost.is_finite()); + assert!(estimate.total_cost.is_finite()); + assert!(estimate.selectivity.is_finite()); + assert!(estimate.index_pages.is_finite()); + } +} diff --git a/crates/vchordrq/src/statistics.rs b/crates/vchordrq/src/statistics.rs new file mode 100644 index 00000000..f651bec5 --- /dev/null +++ b/crates/vchordrq/src/statistics.rs @@ -0,0 +1,30 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the terms of the ELv2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::tuples::{MetaTuple, WithWriter}; +use index::relation::{Page, RelationWrite}; + +/// Store the number of live vector nodes observed by the latest complete +/// build or vacuum pass. +/// +/// This is deliberately refreshed in bulk instead of on every insert. A +/// per-insert update would serialize all writers on the metapage, which is a +/// poor tradeoff for a planner statistic. Like PostgreSQL's relation +/// statistics, the value may be stale between maintenance passes. +pub fn set_indexed_vectors(index: &R, indexed_vectors: u64) { + let mut meta_guard = index.write(0, false); + let meta_bytes = meta_guard.get_mut(1).expect("data corruption"); + let mut meta_tuple = MetaTuple::deserialize_mut(meta_bytes); + meta_tuple.set_indexed_vectors(indexed_vectors); +} diff --git a/crates/vchordrq/src/tuples.rs b/crates/vchordrq/src/tuples.rs index 8438c1ae..85dee786 100644 --- a/crates/vchordrq/src/tuples.rs +++ b/crates/vchordrq/src/tuples.rs @@ -21,6 +21,8 @@ pub const ALIGN: usize = 8; pub type Tag = u64; const MAGIC: Tag = Tag::from_ne_bytes(*b"vchordrq"); const VERSION: u64 = 1001; +const STATISTICS_VERSION: u16 = 1; +const MAX_INDEXED_VECTORS: u64 = (1_u64 << 48) - 1; #[inline(always)] fn tag(source: &[u8]) -> Tag { @@ -55,12 +57,13 @@ struct MetaTupleHeader { rerank_in_heap: Bool, cells_s: u16, cells_e: u16, - _padding_0: [Padding; 2], + statistics_version: u16, centroids_first: u32, vectors_first_s: u16, vectors_first_e: u16, freepages_first: u32, - _padding_1: [Padding; 6], + indexed_vectors_low: u32, + indexed_vectors_high: u16, // tree centroid_prefetch_s: u16, centroid_prefetch_e: u16, @@ -69,11 +72,24 @@ struct MetaTupleHeader { first: u32, } +// Statistics deliberately replace the old 2-byte and 6-byte padding regions. +// Keep these assertions in non-test builds: changing any offset would require +// an index format version bump and REINDEX instead of the compatibility path. +const _: () = { + assert!(size_of::() == 56); + assert!(std::mem::offset_of!(MetaTupleHeader, statistics_version) == 22); + assert!(std::mem::offset_of!(MetaTupleHeader, centroids_first) == 24); + assert!(std::mem::offset_of!(MetaTupleHeader, indexed_vectors_low) == 36); + assert!(std::mem::offset_of!(MetaTupleHeader, indexed_vectors_high) == 40); + assert!(std::mem::offset_of!(MetaTupleHeader, centroid_prefetch_s) == 42); +}; + pub struct MetaTuple { pub dim: u32, pub height_of_root: u32, pub is_residual: bool, pub rerank_in_heap: bool, + pub indexed_vectors: Option, pub cells: Vec, pub centroids_first: u32, pub vectors_first: Vec, @@ -94,6 +110,7 @@ impl Tuple for MetaTuple { height_of_root, is_residual, rerank_in_heap, + indexed_vectors, cells, centroids_first, vectors_first, @@ -103,6 +120,12 @@ impl Tuple for MetaTuple { centroid_norm, first, } => { + if let Some(indexed_vectors) = indexed_vectors { + assert!( + *indexed_vectors <= MAX_INDEXED_VECTORS, + "indexed vector count exceeds the on-disk 48-bit limit" + ); + } buffer.extend((MAGIC as Tag).to_ne_bytes()); buffer.extend(std::iter::repeat_n(0, size_of::())); // cells @@ -136,17 +159,20 @@ impl Tuple for MetaTuple { rerank_in_heap: (*rerank_in_heap).into(), cells_s, cells_e, + statistics_version: indexed_vectors + .map(|_| STATISTICS_VERSION) + .unwrap_or(0), centroids_first: *centroids_first, vectors_first_s, vectors_first_e, freepages_first: *freepages_first, + indexed_vectors_low: indexed_vectors.unwrap_or(0) as u32, + indexed_vectors_high: (indexed_vectors.unwrap_or(0) >> 32) as u16, centroid_prefetch_s, centroid_prefetch_e, centroid_head: *centroid_head, centroid_norm: *centroid_norm, first: *first, - _padding_0: Default::default(), - _padding_1: Default::default(), } .as_bytes(), ); @@ -186,6 +212,28 @@ impl WithReader for MetaTuple { } } +impl WithWriter for MetaTuple { + type Writer<'a> = MetaTupleWriter<'a>; + + fn deserialize_mut(source: &mut [u8]) -> MetaTupleWriter<'_> { + let tag = tag(source); + match tag { + MAGIC => { + let mut checker = MutChecker::new(source); + let header: &mut MetaTupleHeader = checker.prefix(size_of::()); + if VERSION != header.version { + panic!( + "deserialization: bad version number; {}", + "after upgrading VectorChord, please use REINDEX to rebuild the index." + ); + } + MetaTupleWriter { header } + } + _ => panic!("deserialization: bad magic number"), + } + } +} + #[derive(Debug, Clone, Copy)] pub struct MetaTupleReader<'a> { header: &'a MetaTupleHeader, @@ -207,6 +255,16 @@ impl<'a> MetaTupleReader<'a> { pub fn rerank_in_heap(self) -> bool { self.header.rerank_in_heap.into() } + pub fn indexed_vectors(self) -> Option { + match self.header.statistics_version { + 0 => None, + STATISTICS_VERSION => Some( + u64::from(self.header.indexed_vectors_low) + | (u64::from(self.header.indexed_vectors_high) << 32), + ), + _ => panic!("deserialization: unsupported statistics version"), + } + } pub fn cells(self) -> &'a [u32] { self.cells } @@ -233,6 +291,23 @@ impl<'a> MetaTupleReader<'a> { } } +#[derive(Debug)] +pub struct MetaTupleWriter<'a> { + header: &'a mut MetaTupleHeader, +} + +impl MetaTupleWriter<'_> { + pub fn set_indexed_vectors(&mut self, indexed_vectors: u64) { + assert!( + indexed_vectors <= MAX_INDEXED_VECTORS, + "indexed vector count exceeds the on-disk 48-bit limit" + ); + self.header.statistics_version = STATISTICS_VERSION; + self.header.indexed_vectors_low = indexed_vectors as u32; + self.header.indexed_vectors_high = (indexed_vectors >> 32) as u16; + } +} + #[repr(C, align(8))] #[derive(Debug, Clone, FromBytes, IntoBytes, Immutable, KnownLayout)] struct FreepagesTupleHeader { @@ -1485,3 +1560,71 @@ impl AppendableTupleWriter<'_> { &mut self.header.payload } } + +#[cfg(test)] +mod tests { + use super::*; + + fn meta(indexed_vectors: Option) -> MetaTuple { + MetaTuple { + dim: 3, + height_of_root: 1, + is_residual: false, + rerank_in_heap: false, + indexed_vectors, + cells: vec![4], + centroids_first: 1, + vectors_first: vec![2], + freepages_first: 3, + centroid_prefetch: vec![4], + centroid_head: 0, + centroid_norm: 1.0, + first: 5, + } + } + + #[test] + fn meta_statistics_reuse_padding_without_moving_fields() { + assert_eq!(size_of::(), 56); + assert_eq!( + std::mem::offset_of!(MetaTupleHeader, statistics_version), + 22 + ); + assert_eq!(std::mem::offset_of!(MetaTupleHeader, centroids_first), 24); + assert_eq!( + std::mem::offset_of!(MetaTupleHeader, indexed_vectors_low), + 36 + ); + assert_eq!( + std::mem::offset_of!(MetaTupleHeader, indexed_vectors_high), + 40 + ); + assert_eq!( + std::mem::offset_of!(MetaTupleHeader, centroid_prefetch_s), + 42 + ); + } + + #[test] + fn meta_statistics_roundtrip_full_48_bit_range() { + for expected in [0, 1, u32::MAX as u64 + 1, MAX_INDEXED_VECTORS] { + let bytes = meta(Some(expected)).serialize(); + assert_eq!( + MetaTuple::deserialize_ref(&bytes).indexed_vectors(), + Some(expected) + ); + } + } + + #[test] + fn old_meta_padding_reads_as_missing_statistics_and_can_be_upgraded() { + let mut bytes = meta(None).serialize(); + assert_eq!(MetaTuple::deserialize_ref(&bytes).indexed_vectors(), None); + + MetaTuple::deserialize_mut(&mut bytes).set_indexed_vectors(42); + assert_eq!( + MetaTuple::deserialize_ref(&bytes).indexed_vectors(), + Some(42) + ); + } +} diff --git a/crates/vchordrq/src/types.rs b/crates/vchordrq/src/types.rs index c41ba501..d9b85cbc 100644 --- a/crates/vchordrq/src/types.rs +++ b/crates/vchordrq/src/types.rs @@ -12,12 +12,14 @@ // // Copyright (c) 2025-2026 TensorChord Inc. +use distance::Distance; use serde::{Deserialize, Serialize}; use simd::f16; use validator::{Validate, ValidationError}; use vector::rabitq4::{Rabitq4Borrowed, Rabitq4Owned}; use vector::rabitq8::{Rabitq8Borrowed, Rabitq8Owned}; use vector::vect::{VectBorrowed, VectOwned}; +use vector::{VectorBorrowed, VectorOwned}; #[derive(Debug, Clone, Serialize, Deserialize, Validate)] #[serde(deny_unknown_fields)] @@ -61,6 +63,38 @@ pub enum OwnedVector { Rabitq4(Rabitq4Owned), } +impl OwnedVector { + pub fn dim(&self) -> u32 { + match self { + Self::Vecf32(vector) => vector.as_borrowed().dim(), + Self::Vecf16(vector) => vector.as_borrowed().dim(), + Self::Rabitq8(vector) => vector.as_borrowed().dim(), + Self::Rabitq4(vector) => vector.as_borrowed().dim(), + } + } + + pub fn operator_dot(&self, rhs: &Self) -> Option { + if self.dim() != rhs.dim() { + return None; + } + match (self, rhs) { + (Self::Vecf32(lhs), Self::Vecf32(rhs)) => { + Some(lhs.as_borrowed().operator_dot(rhs.as_borrowed())) + } + (Self::Vecf16(lhs), Self::Vecf16(rhs)) => { + Some(lhs.as_borrowed().operator_dot(rhs.as_borrowed())) + } + (Self::Rabitq8(lhs), Self::Rabitq8(rhs)) => { + Some(lhs.as_borrowed().operator_dot(rhs.as_borrowed())) + } + (Self::Rabitq4(lhs), Self::Rabitq4(rhs)) => { + Some(lhs.as_borrowed().operator_dot(rhs.as_borrowed())) + } + _ => None, + } + } +} + #[derive(Debug, Clone, Copy)] pub enum BorrowedVector<'a> { Vecf32(VectBorrowed<'a, f32>), diff --git a/devtools/test_tilemaxsim_reference_sidecar.py b/devtools/test_tilemaxsim_reference_sidecar.py new file mode 100644 index 00000000..ee5ce666 --- /dev/null +++ b/devtools/test_tilemaxsim_reference_sidecar.py @@ -0,0 +1,383 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the terms of the ELv2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +from __future__ import annotations + +import hashlib +import os +import socket +import stat +import struct +import tempfile +import threading +import time +import unittest +from pathlib import Path + +try: + from . import tilemaxsim_reference_sidecar as sidecar +except ImportError: + import tilemaxsim_reference_sidecar as sidecar + + +def request_frame( + request_id: int, + dtype: int, + query: list[list[float]], + candidates: list[tuple[int, list[list[float]]]], +) -> bytes: + dimension = len(query[0]) + code = "f" if dtype == sidecar.DTYPE_F32 else "e" + body = bytearray( + sidecar.REQUEST_FIXED.pack( + dimension, + len(query), + len(candidates), + dtype, + sidecar.SCORING_SUM_QUERY_MAX_DOCUMENT_DOT, + 0, + ) + ) + body.extend(struct.pack(f"<{len(query) * dimension}{code}", *sum(query, []))) + for candidate_id, tensor in candidates: + body.extend(sidecar.CANDIDATE_FIXED.pack(candidate_id, len(tensor))) + body.extend(struct.pack(f"<{len(tensor) * dimension}{code}", *sum(tensor, []))) + return ( + sidecar.HEADER.pack( + sidecar.MAGIC, + sidecar.VERSION, + sidecar.REQUEST_KIND, + request_id, + len(body), + ) + + body + ) + + +def external_request_frame( + request_id: int, + dtype: int, + query: list[list[float]], + model_contract_id: str, + candidates: list[tuple[int, str, list[list[float]]]], +) -> tuple[bytes, dict[str, bytes]]: + dimension = len(query[0]) + code = "f" if dtype == sidecar.DTYPE_F32 else "e" + contract = model_contract_id.encode() + body = bytearray( + sidecar.EXTERNAL_REQUEST_FIXED.pack( + dimension, + len(query), + len(candidates), + dtype, + sidecar.SCORING_SUM_QUERY_MAX_DOCUMENT_DOT, + 0, + len(contract), + ) + ) + body.extend(contract) + body.extend(struct.pack(f"<{len(query) * dimension}{code}", *sum(query, []))) + objects = {} + for candidate_id, tensor_ref, tensor in candidates: + payload = struct.pack(f"<{len(tensor) * dimension}{code}", *sum(tensor, [])) + objects[tensor_ref] = payload + reference = tensor_ref.encode() + checksum = f"sha256:{hashlib.sha256(payload).hexdigest()}".encode() + body.extend( + sidecar.EXTERNAL_CANDIDATE_FIXED.pack( + candidate_id, len(tensor), len(reference), len(checksum) + ) + ) + body.extend(reference) + body.extend(checksum) + return ( + sidecar.HEADER.pack( + sidecar.MAGIC, + sidecar.EXTERNAL_VERSION, + sidecar.REQUEST_KIND, + request_id, + len(body), + ) + + body, + objects, + ) + + +def decode_response(frame: bytes) -> tuple[int, int, list[tuple[int, float]] | str]: + magic, version, kind, request_id, body_len = sidecar.HEADER.unpack_from(frame) + assert magic == sidecar.MAGIC + assert version in (sidecar.VERSION, sidecar.EXTERNAL_VERSION) + assert kind == sidecar.RESPONSE_KIND + assert len(frame) == sidecar.HEADER.size + body_len + status, count_or_length = sidecar.RESPONSE_FIXED.unpack_from( + frame, sidecar.HEADER.size + ) + offset = sidecar.HEADER.size + sidecar.RESPONSE_FIXED.size + if status: + return request_id, status, frame[offset : offset + count_or_length].decode() + results = [] + for _ in range(count_or_length): + results.append(sidecar.RESULT.unpack_from(frame, offset)) + offset += sidecar.RESULT.size + assert offset == len(frame) + return request_id, status, results + + +class ReferenceSidecarTest(unittest.TestCase): + def test_f32_exact_scores_and_opaque_ids(self) -> None: + frame = request_frame( + 41, + sidecar.DTYPE_F32, + [[1.0, 0.0], [0.0, 1.0]], + [ + (17, [[1.0, 0.0], [0.0, 1.0]]), + (3, [[0.5, 0.5]]), + ], + ) + request_id, status, results = decode_response(sidecar.process_frame(frame)) + + self.assertEqual(request_id, 41) + self.assertEqual(status, 0) + self.assertEqual(results, [(17, 2.0), (3, 1.0)]) + + def test_f16_exact_scores(self) -> None: + frame = request_frame( + 42, + sidecar.DTYPE_F16, + [[1.0, 0.0], [0.0, 1.0]], + [(0, [[0.75, 0.0], [0.0, 0.5]])], + ) + _, status, results = decode_response(sidecar.process_frame(frame)) + + self.assertEqual(status, 0) + self.assertEqual(results, [(0, 1.25)]) + + def test_shared_raw_decoder_preserves_inline_payloads(self) -> None: + frame = request_frame( + 420, + sidecar.DTYPE_F16, + [[1.0, 0.0]], + [(11, [[0.5, 0.25]])], + ) + request = sidecar.parse_request_frame(frame) + self.assertIsInstance(request, sidecar.InlineTensorRequest) + assert isinstance(request, sidecar.InlineTensorRequest) + self.assertEqual(request.request_id, 420) + self.assertEqual(request.query_rows, 1) + self.assertEqual(request.dimension, 2) + self.assertEqual([item.candidate_id for item in request.candidates], [11]) + self.assertEqual(len(request.query_payload), 4) + self.assertEqual(len(request.candidates[0].payload), 4) + + def test_duplicate_id_and_non_finite_input_fail_closed(self) -> None: + duplicate = request_frame( + 43, + sidecar.DTYPE_F32, + [[1.0]], + [(7, [[1.0]]), (7, [[2.0]])], + ) + _, status, message = decode_response(sidecar.process_frame(duplicate)) + self.assertEqual(status, sidecar.STATUS_INVALID_REQUEST) + self.assertIn("duplicate candidate ID", message) + + non_finite = request_frame( + 44, + sidecar.DTYPE_F32, + [[float("nan")]], + [(0, [[1.0]])], + ) + _, status, message = decode_response(sidecar.process_frame(non_finite)) + self.assertEqual(status, sidecar.STATUS_INVALID_REQUEST) + self.assertIn("non-finite", message) + + def test_truncated_and_oversized_frames_fail_closed(self) -> None: + valid = request_frame(45, sidecar.DTYPE_F32, [[1.0]], [(0, [[1.0]])]) + _, status, message = decode_response(sidecar.process_frame(valid[:-1])) + self.assertEqual(status, sidecar.STATUS_INVALID_REQUEST) + self.assertIn("length mismatch", message) + + _, status, message = decode_response( + sidecar.process_frame(valid, sidecar.Limits(max_request_bytes=32)) + ) + self.assertEqual(status, sidecar.STATUS_RESOURCE_LIMIT) + self.assertIn("byte limit", message) + + def test_header_reserved_trailing_and_token_limit_fail_closed(self) -> None: + valid = request_frame(47, sidecar.DTYPE_F32, [[1.0]], [(9, [[1.0]])]) + invalid_frames = [] + for offset, value in ((0, 0), (4, 3), (6, 2)): + invalid = bytearray(valid) + invalid[offset] = value + invalid_frames.append(bytes(invalid)) + + reserved = bytearray(valid) + reserved[sidecar.HEADER.size + 14] = 1 + invalid_frames.append(bytes(reserved)) + + trailing = bytearray(valid) + trailing.extend(b"x") + struct.pack_into(" None: + frame, objects = external_request_frame( + 48, + sidecar.DTYPE_F16, + [[1.0, 0.0], [0.0, 1.0]], + "colqwen@immutable-revision", + [ + (77, "object://fixture/page-1", [[1.0, 0.0], [0.0, 1.0]]), + (4, "object://fixture/page-2", [[0.5, 0.5]]), + ], + ) + seen = [] + + def resolver(request: sidecar.ExternalTensorRequest) -> bytes: + seen.append(request) + return objects[request.tensor_ref] + + request_id, status, results = decode_response( + sidecar.process_frame(frame, resolver=resolver) + ) + + self.assertEqual(request_id, 48) + self.assertEqual(status, 0) + self.assertEqual(results, [(77, 2.0), (4, 1.0)]) + self.assertEqual( + {request.model_contract_id for request in seen}, + {"colqwen@immutable-revision"}, + ) + self.assertEqual({request.tensor_ref for request in seen}, set(objects)) + + parsed = sidecar.parse_request_frame(frame) + self.assertIsInstance(parsed, sidecar.ParsedExternalTensorRequest) + assert isinstance(parsed, sidecar.ParsedExternalTensorRequest) + self.assertEqual(parsed.model_contract_id, "colqwen@immutable-revision") + self.assertEqual( + [candidate.candidate_id for candidate in parsed.candidates], [77, 4] + ) + + def test_external_v2_fails_closed_without_resolver_or_on_checksum_mismatch( + self, + ) -> None: + frame, objects = external_request_frame( + 49, + sidecar.DTYPE_F32, + [[1.0]], + "contract@1", + [(0, "object://immutable/page", [[2.0]])], + ) + _, status, message = decode_response(sidecar.process_frame(frame)) + self.assertEqual(status, sidecar.STATUS_COMPUTE_ERROR) + self.assertIn("resolver is not configured", message) + + def corrupt_resolver(request: sidecar.ExternalTensorRequest) -> bytes: + return objects[request.tensor_ref] + b"x" + + _, status, message = decode_response( + sidecar.process_frame(frame, resolver=corrupt_resolver) + ) + self.assertEqual(status, sidecar.STATUS_INVALID_REQUEST) + self.assertIn("byte length", message) + + def test_external_v2_validates_complete_control_frame_before_resolution( + self, + ) -> None: + frame, objects = external_request_frame( + 50, + sidecar.DTYPE_F32, + [[1.0, 0.0]], + "contract@1", + [(0, "object://immutable/page", [[1.0, 0.0]])], + ) + invalid = bytearray(frame) + contract_length = len("contract@1") + candidate_offset = ( + sidecar.HEADER.size + + sidecar.EXTERNAL_REQUEST_FIXED.size + + contract_length + + 8 + ) + reference_offset = candidate_offset + sidecar.EXTERNAL_CANDIDATE_FIXED.size + invalid[reference_offset] = 0 + called = False + + def resolver(request: sidecar.ExternalTensorRequest) -> bytes: + nonlocal called + called = True + return objects[request.tensor_ref] + + _, status, message = decode_response( + sidecar.process_frame(bytes(invalid), resolver=resolver) + ) + self.assertEqual(status, sidecar.STATUS_INVALID_REQUEST) + self.assertIn("control characters", message) + self.assertFalse(called) + + def test_unix_socket_end_to_end(self) -> None: + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "tilemaxsim.sock" + thread = threading.Thread( + target=sidecar.serve, + args=(path, sidecar.Limits()), + kwargs={"once": True}, + daemon=True, + ) + thread.start() + for _ in range(100): + if path.exists() and stat.S_IMODE(path.stat().st_mode) == 0o600: + break + time.sleep(0.01) + else: + self.fail("sidecar socket was not created with mode 0600") + self.assertEqual(stat.S_IMODE(path.stat().st_mode), 0o600) + + frame = request_frame( + 46, + sidecar.DTYPE_F32, + [[1.0, 0.0], [0.0, 1.0]], + [(5, [[1.0, 0.0], [0.0, 1.0]])], + ) + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: + connection.connect(os.fspath(path)) + connection.sendall(frame) + header = connection.recv(sidecar.HEADER.size) + while len(header) < sidecar.HEADER.size: + header += connection.recv(sidecar.HEADER.size - len(header)) + body_len = sidecar.HEADER.unpack(header)[4] + body = b"" + while len(body) < body_len: + body += connection.recv(body_len - len(body)) + thread.join(timeout=2) + self.assertFalse(thread.is_alive()) + + request_id, status, results = decode_response(header + body) + self.assertEqual(request_id, 46) + self.assertEqual(status, 0) + self.assertEqual(results, [(5, 2.0)]) + self.assertFalse(path.exists()) + + +if __name__ == "__main__": + unittest.main() diff --git a/devtools/tilemaxsim_reference_sidecar.py b/devtools/tilemaxsim_reference_sidecar.py new file mode 100644 index 00000000..8d9f7ee1 --- /dev/null +++ b/devtools/tilemaxsim_reference_sidecar.py @@ -0,0 +1,780 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the terms of the ELv2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +"""CPU reference implementation of the VectorChord TileMaxSim IPC sidecar. + +This executable is intentionally simple and single-threaded. It is a protocol +oracle and end-to-end development aid, not a production or GPU implementation. +The CLI serves inline v1 requests. External-descriptor v2 requests require an +explicit resolver injected by a test or embedding application and fail closed +when no resolver is configured. +""" + +from __future__ import annotations + +import argparse +import hashlib +import hmac +import math +import os +import signal +import socket +import stat +import struct +import threading +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Iterable + +MAGIC = b"VCTM" +VERSION = 1 +EXTERNAL_VERSION = 2 +REQUEST_KIND = 1 +RESPONSE_KIND = 2 +SCORING_SUM_QUERY_MAX_DOCUMENT_DOT = 1 +DTYPE_F32 = 1 +DTYPE_F16 = 2 + +HEADER = struct.Struct("<4sHHQQ") +REQUEST_FIXED = struct.Struct(" None: + super().__init__(message) + self.status = status + + +@dataclass(frozen=True) +class Limits: + max_request_bytes: int = 64 * 1024 * 1024 + max_batch_tokens: int = 1_000_000 + max_tensor_bytes: int = 1024 * 1024 * 1024 + max_candidates: int = 65_536 + + +@dataclass(frozen=True) +class ExternalTensorRequest: + model_contract_id: str + tensor_ref: str + rows: int + dimension: int + dtype: int + checksum: str + + +@dataclass(frozen=True) +class InlineTensorCandidate: + candidate_id: int + rows: int + payload: bytes + + +@dataclass(frozen=True) +class InlineTensorRequest: + request_id: int + dimension: int + query_rows: int + dtype: int + query_payload: bytes + candidates: tuple[InlineTensorCandidate, ...] + + +@dataclass(frozen=True) +class ExternalTensorCandidate: + candidate_id: int + descriptor: ExternalTensorRequest + + +@dataclass(frozen=True) +class ParsedExternalTensorRequest: + request_id: int + dimension: int + query_rows: int + dtype: int + model_contract_id: str + query_payload: bytes + candidates: tuple[ExternalTensorCandidate, ...] + + +ParsedRequest = InlineTensorRequest | ParsedExternalTensorRequest + + +# A reference resolver returns the exact row-major scalar bytes whose SHA-256 +# digest is stored in the descriptor. Production resolvers may adapt richer +# immutable object formats before returning this canonical tensor payload. +ExternalTensorResolver = Callable[[ExternalTensorRequest], bytes] + + +class Reader: + def __init__(self, payload: bytes) -> None: + self.payload = payload + self.offset = 0 + + def take(self, count: int) -> bytes: + end = self.offset + count + if count < 0 or end > len(self.payload): + raise SidecarError(STATUS_INVALID_REQUEST, "truncated request") + chunk = self.payload[self.offset : end] + self.offset = end + return chunk + + def unpack(self, layout: struct.Struct) -> tuple: + return layout.unpack(self.take(layout.size)) + + def finish(self) -> None: + if self.offset != len(self.payload): + raise SidecarError(STATUS_INVALID_REQUEST, "trailing request bytes") + + +def checked_elements(rows: int, dimension: int) -> int: + if rows <= 0 or dimension <= 0: + raise SidecarError( + STATUS_INVALID_REQUEST, "tensor rows and dimension must be positive" + ) + elements = rows * dimension + if elements > (1 << 63) - 1: + raise SidecarError(STATUS_RESOURCE_LIMIT, "tensor shape is too large") + return elements + + +def dtype_size(dtype: int) -> int: + if dtype == DTYPE_F32: + return 4 + if dtype == DTYPE_F16: + return 2 + raise SidecarError(STATUS_INVALID_REQUEST, "unsupported tensor dtype") + + +def checked_tensor_bytes(rows: int, dimension: int, dtype: int) -> int: + return checked_elements(rows, dimension) * dtype_size(dtype) + + +def validate_finite_tensor_payload( + payload: bytes, rows: int, dimension: int, dtype: int +) -> None: + expected = checked_tensor_bytes(rows, dimension, dtype) + if len(payload) != expected: + raise SidecarError( + STATUS_INVALID_REQUEST, "tensor byte length does not match its shape" + ) + code = "f" if dtype == DTYPE_F32 else "e" + try: + values = struct.iter_unpack(f"<{code}", payload) + if any(not math.isfinite(value[0]) for value in values): + raise SidecarError( + STATUS_INVALID_REQUEST, "tensor contains non-finite value" + ) + except struct.error as error: + raise SidecarError(STATUS_INVALID_REQUEST, str(error)) from error + + +def read_text(reader: Reader, length: int, maximum: int, field: str) -> str: + if length <= 0 or length > maximum: + raise SidecarError(STATUS_RESOURCE_LIMIT, f"invalid {field} length") + try: + value = reader.take(length).decode("utf-8") + except UnicodeDecodeError as error: + raise SidecarError(STATUS_INVALID_REQUEST, f"{field} is not UTF-8") from error + if any(ord(character) < 32 or ord(character) == 127 for character in value): + raise SidecarError( + STATUS_INVALID_REQUEST, f"{field} contains control characters" + ) + return value + + +def read_tensor( + reader: Reader, rows: int, dimension: int, dtype: int +) -> list[tuple[float, ...]]: + elements = checked_elements(rows, dimension) + if dtype == DTYPE_F32: + code = "f" + element_size = 4 + elif dtype == DTYPE_F16: + code = "e" + element_size = 2 + else: + raise SidecarError(STATUS_INVALID_REQUEST, "unsupported tensor dtype") + raw = reader.take(elements * element_size) + try: + values = struct.unpack(f"<{elements}{code}", raw) + except struct.error as error: + raise SidecarError(STATUS_INVALID_REQUEST, str(error)) from error + if not all(math.isfinite(value) for value in values): + raise SidecarError(STATUS_INVALID_REQUEST, "tensor contains non-finite value") + return [ + tuple(values[offset : offset + dimension]) + for offset in range(0, elements, dimension) + ] + + +def decode_resolved_tensor( + payload: bytes, request: ExternalTensorRequest +) -> list[tuple[float, ...]]: + expected_bytes = checked_tensor_bytes( + request.rows, request.dimension, request.dtype + ) + if len(payload) != expected_bytes: + raise SidecarError( + STATUS_INVALID_REQUEST, + "resolved tensor byte length does not match descriptor", + ) + expected_checksum = f"sha256:{hashlib.sha256(payload).hexdigest()}" + if not hmac.compare_digest(expected_checksum, request.checksum): + raise SidecarError(STATUS_INVALID_REQUEST, "resolved tensor checksum mismatch") + reader = Reader(payload) + tensor = read_tensor(reader, request.rows, request.dimension, request.dtype) + reader.finish() + return tensor + + +def tilemaxsim( + query: list[tuple[float, ...]], document: list[tuple[float, ...]] +) -> float: + if not query or not document: + raise SidecarError(STATUS_INVALID_REQUEST, "tensor must not be empty") + score = 0.0 + try: + for query_vector in query: + best = -math.inf + for document_vector in document: + dot = math.fsum( + left * right + for left, right in zip(query_vector, document_vector, strict=True) + ) + best = max(best, dot) + score += best + except (OverflowError, ValueError) as error: + raise SidecarError(STATUS_COMPUTE_ERROR, str(error)) from error + if not math.isfinite(score): + raise SidecarError(STATUS_COMPUTE_ERROR, "TileMaxSim result is non-finite") + return score + + +def success_response( + request_id: int, + results: Iterable[tuple[int, float]], + version: int = VERSION, +) -> bytes: + results = list(results) + body = bytearray(RESPONSE_FIXED.pack(0, len(results))) + for candidate_id, similarity in results: + body.extend(RESULT.pack(candidate_id, similarity)) + return HEADER.pack(MAGIC, version, RESPONSE_KIND, request_id, len(body)) + body + + +def error_response( + request_id: int, + status_code: int, + message: str, + version: int = VERSION, +) -> bytes: + encoded = message.encode("utf-8", errors="replace")[:MAX_ERROR_BYTES] + body = ERROR_FIXED.pack(status_code or STATUS_COMPUTE_ERROR, len(encoded)) + encoded + return HEADER.pack(MAGIC, version, RESPONSE_KIND, request_id, len(body)) + body + + +def validate_request_fixed( + dimension: int, + query_rows: int, + candidate_count: int, + dtype: int, + scoring: int, + reserved: int, + limits: Limits, +) -> None: + if dimension == 0 or dimension > 60_000: + raise SidecarError(STATUS_INVALID_REQUEST, "invalid tensor dimension") + if query_rows == 0: + raise SidecarError(STATUS_INVALID_REQUEST, "query tensor is empty") + if candidate_count > limits.max_candidates: + raise SidecarError(STATUS_RESOURCE_LIMIT, "too many candidates") + dtype_size(dtype) + if scoring != SCORING_SUM_QUERY_MAX_DOCUMENT_DOT: + raise SidecarError(STATUS_INVALID_REQUEST, "unsupported scoring function") + if reserved != 0: + raise SidecarError(STATUS_INVALID_REQUEST, "reserved field must be zero") + if query_rows > limits.max_batch_tokens: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds token limit") + if checked_tensor_bytes(query_rows, dimension, dtype) > limits.max_tensor_bytes: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds tensor byte limit") + + +def process_inline_request(reader: Reader, limits: Limits) -> list[tuple[int, float]]: + ( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + ) = reader.unpack(REQUEST_FIXED) + validate_request_fixed( + dimension, query_rows, candidate_count, dtype, scoring, reserved, limits + ) + + query = read_tensor(reader, query_rows, dimension, dtype) + total_tokens = query_rows + total_tensor_bytes = checked_tensor_bytes(query_rows, dimension, dtype) + candidates: list[tuple[int, list[tuple[float, ...]]]] = [] + candidate_ids: set[int] = set() + for _ in range(candidate_count): + candidate_id, rows = reader.unpack(CANDIDATE_FIXED) + if candidate_id in candidate_ids: + raise SidecarError(STATUS_INVALID_REQUEST, "duplicate candidate ID") + candidate_ids.add(candidate_id) + total_tokens += rows + total_tensor_bytes += checked_tensor_bytes(rows, dimension, dtype) + if total_tokens > limits.max_batch_tokens: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds token limit") + if total_tensor_bytes > limits.max_tensor_bytes: + raise SidecarError( + STATUS_RESOURCE_LIMIT, "request exceeds tensor byte limit" + ) + candidates.append((candidate_id, read_tensor(reader, rows, dimension, dtype))) + reader.finish() + return [ + (candidate_id, tilemaxsim(query, document)) + for candidate_id, document in candidates + ] + + +def process_external_request( + reader: Reader, + limits: Limits, + resolver: ExternalTensorResolver | None, +) -> list[tuple[int, float]]: + ( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + contract_length, + ) = reader.unpack(EXTERNAL_REQUEST_FIXED) + validate_request_fixed( + dimension, query_rows, candidate_count, dtype, scoring, reserved, limits + ) + model_contract_id = read_text(reader, contract_length, 512, "model contract") + query = read_tensor(reader, query_rows, dimension, dtype) + total_tokens = query_rows + total_tensor_bytes = checked_tensor_bytes(query_rows, dimension, dtype) + descriptors: list[tuple[int, ExternalTensorRequest]] = [] + candidate_ids: set[int] = set() + for _ in range(candidate_count): + candidate_id, rows, reference_length, checksum_length = reader.unpack( + EXTERNAL_CANDIDATE_FIXED + ) + if candidate_id in candidate_ids: + raise SidecarError(STATUS_INVALID_REQUEST, "duplicate candidate ID") + candidate_ids.add(candidate_id) + tensor_ref = read_text(reader, reference_length, 4096, "tensor reference") + checksum = read_text(reader, checksum_length, 512, "tensor checksum") + digest = checksum.removeprefix("sha256:") + if ( + not checksum.startswith("sha256:") + or len(digest) != 64 + or any(character not in "0123456789abcdef" for character in digest) + ): + raise SidecarError( + STATUS_INVALID_REQUEST, + "tensor checksum must be a lowercase sha256 digest", + ) + total_tokens += rows + total_tensor_bytes += checked_tensor_bytes(rows, dimension, dtype) + if total_tokens > limits.max_batch_tokens: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds token limit") + if total_tensor_bytes > limits.max_tensor_bytes: + raise SidecarError( + STATUS_RESOURCE_LIMIT, "request exceeds tensor byte limit" + ) + descriptors.append( + ( + candidate_id, + ExternalTensorRequest( + model_contract_id=model_contract_id, + tensor_ref=tensor_ref, + rows=rows, + dimension=dimension, + dtype=dtype, + checksum=checksum, + ), + ) + ) + reader.finish() + + # Validate the complete control frame before performing any external I/O. + if resolver is None: + raise SidecarError( + STATUS_COMPUTE_ERROR, "external tensor resolver is not configured" + ) + results = [] + for candidate_id, descriptor in descriptors: + payload = resolver(descriptor) + if not isinstance(payload, bytes): + raise SidecarError( + STATUS_COMPUTE_ERROR, + "external tensor resolver returned a non-bytes value", + ) + document = decode_resolved_tensor(payload, descriptor) + results.append((candidate_id, tilemaxsim(query, document))) + return results + + +def parse_request_frame( + frame: bytes, + limits: Limits = Limits(), + *, + validate_finite: bool = True, +) -> ParsedRequest: + """Decode and validate a request without resolving or computing tensors. + + The production CUDA sidecar uses this function so the protocol oracle and + deployable executor share one strict wire decoder. All control data and the + inline query are validated before a v2 resolver performs external I/O. + """ + + if len(frame) < HEADER.size: + raise SidecarError(STATUS_INVALID_REQUEST, "truncated frame header") + magic, version, kind, request_id, body_len = HEADER.unpack_from(frame) + if magic != MAGIC: + raise SidecarError(STATUS_INVALID_REQUEST, "invalid frame magic") + if version not in (VERSION, EXTERNAL_VERSION): + raise SidecarError(STATUS_INVALID_REQUEST, "unsupported protocol version") + if kind != REQUEST_KIND: + raise SidecarError(STATUS_INVALID_REQUEST, "unexpected message kind") + if body_len != len(frame) - HEADER.size: + raise SidecarError(STATUS_INVALID_REQUEST, "request length mismatch") + if len(frame) > limits.max_request_bytes: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds byte limit") + + reader = Reader(frame[HEADER.size :]) + if version == VERSION: + ( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + ) = reader.unpack(REQUEST_FIXED) + validate_request_fixed( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + limits, + ) + query_payload = reader.take(checked_tensor_bytes(query_rows, dimension, dtype)) + if validate_finite: + validate_finite_tensor_payload(query_payload, query_rows, dimension, dtype) + total_tokens = query_rows + total_tensor_bytes = len(query_payload) + candidate_ids: set[int] = set() + candidates = [] + for _ in range(candidate_count): + candidate_id, rows = reader.unpack(CANDIDATE_FIXED) + if candidate_id in candidate_ids: + raise SidecarError(STATUS_INVALID_REQUEST, "duplicate candidate ID") + candidate_ids.add(candidate_id) + payload_bytes = checked_tensor_bytes(rows, dimension, dtype) + total_tokens += rows + total_tensor_bytes += payload_bytes + if total_tokens > limits.max_batch_tokens: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds token limit") + if total_tensor_bytes > limits.max_tensor_bytes: + raise SidecarError( + STATUS_RESOURCE_LIMIT, "request exceeds tensor byte limit" + ) + payload = reader.take(payload_bytes) + if validate_finite: + validate_finite_tensor_payload(payload, rows, dimension, dtype) + candidates.append(InlineTensorCandidate(candidate_id, rows, payload)) + reader.finish() + return InlineTensorRequest( + request_id, + dimension, + query_rows, + dtype, + query_payload, + tuple(candidates), + ) + + ( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + contract_length, + ) = reader.unpack(EXTERNAL_REQUEST_FIXED) + validate_request_fixed( + dimension, + query_rows, + candidate_count, + dtype, + scoring, + reserved, + limits, + ) + model_contract_id = read_text(reader, contract_length, 512, "model contract") + query_payload = reader.take(checked_tensor_bytes(query_rows, dimension, dtype)) + if validate_finite: + validate_finite_tensor_payload(query_payload, query_rows, dimension, dtype) + total_tokens = query_rows + total_tensor_bytes = len(query_payload) + candidate_ids = set() + candidates = [] + for _ in range(candidate_count): + candidate_id, rows, reference_length, checksum_length = reader.unpack( + EXTERNAL_CANDIDATE_FIXED + ) + if candidate_id in candidate_ids: + raise SidecarError(STATUS_INVALID_REQUEST, "duplicate candidate ID") + candidate_ids.add(candidate_id) + tensor_ref = read_text(reader, reference_length, 4096, "tensor reference") + checksum = read_text(reader, checksum_length, 512, "tensor checksum") + digest = checksum.removeprefix("sha256:") + if ( + not checksum.startswith("sha256:") + or len(digest) != 64 + or any(character not in "0123456789abcdef" for character in digest) + ): + raise SidecarError( + STATUS_INVALID_REQUEST, + "tensor checksum must be a lowercase sha256 digest", + ) + total_tokens += rows + total_tensor_bytes += checked_tensor_bytes(rows, dimension, dtype) + if total_tokens > limits.max_batch_tokens: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds token limit") + if total_tensor_bytes > limits.max_tensor_bytes: + raise SidecarError( + STATUS_RESOURCE_LIMIT, "request exceeds tensor byte limit" + ) + candidates.append( + ExternalTensorCandidate( + candidate_id, + ExternalTensorRequest( + model_contract_id=model_contract_id, + tensor_ref=tensor_ref, + rows=rows, + dimension=dimension, + dtype=dtype, + checksum=checksum, + ), + ) + ) + reader.finish() + return ParsedExternalTensorRequest( + request_id, + dimension, + query_rows, + dtype, + model_contract_id, + query_payload, + tuple(candidates), + ) + + +def process_frame( + frame: bytes, + limits: Limits = Limits(), + resolver: ExternalTensorResolver | None = None, +) -> bytes: + request_id = 0 + response_version = VERSION + try: + if len(frame) < HEADER.size: + raise SidecarError(STATUS_INVALID_REQUEST, "truncated frame header") + magic, version, kind, request_id, body_len = HEADER.unpack_from(frame) + if version in (VERSION, EXTERNAL_VERSION): + response_version = version + if magic != MAGIC: + raise SidecarError(STATUS_INVALID_REQUEST, "invalid frame magic") + if version not in (VERSION, EXTERNAL_VERSION): + raise SidecarError(STATUS_INVALID_REQUEST, "unsupported protocol version") + if kind != REQUEST_KIND: + raise SidecarError(STATUS_INVALID_REQUEST, "unexpected message kind") + if body_len != len(frame) - HEADER.size: + raise SidecarError(STATUS_INVALID_REQUEST, "request length mismatch") + if len(frame) > limits.max_request_bytes: + raise SidecarError(STATUS_RESOURCE_LIMIT, "request exceeds byte limit") + + reader = Reader(frame[HEADER.size :]) + if version == VERSION: + results = process_inline_request(reader, limits) + else: + results = process_external_request(reader, limits, resolver) + return success_response(request_id, results, response_version) + except SidecarError as error: + return error_response(request_id, error.status, str(error), response_version) + except Exception as error: # Keep protocol failures inside the response boundary. + return error_response( + request_id, STATUS_COMPUTE_ERROR, str(error), response_version + ) + + +def receive_exact(connection: socket.socket, count: int) -> bytes: + chunks = bytearray() + while len(chunks) < count: + chunk = connection.recv(count - len(chunks)) + if not chunk: + raise SidecarError( + STATUS_INVALID_REQUEST, "connection closed during request" + ) + chunks.extend(chunk) + return bytes(chunks) + + +def handle_connection( + connection: socket.socket, + limits: Limits, + resolver: ExternalTensorResolver | None = None, +) -> None: + request_id = 0 + response_version = VERSION + try: + header = receive_exact(connection, HEADER.size) + _, version, _, request_id, body_len = HEADER.unpack(header) + if version in (VERSION, EXTERNAL_VERSION): + response_version = version + if body_len > limits.max_request_bytes - HEADER.size: + response = error_response( + request_id, + STATUS_RESOURCE_LIMIT, + "request exceeds byte limit", + response_version, + ) + else: + body = receive_exact(connection, body_len) + response = process_frame(header + body, limits, resolver) + except SidecarError as error: + response = error_response( + request_id, error.status, str(error), response_version + ) + except Exception as error: + response = error_response( + request_id, STATUS_COMPUTE_ERROR, str(error), response_version + ) + connection.sendall(response) + + +def remove_stale_socket(path: Path) -> None: + try: + mode = path.lstat().st_mode + except FileNotFoundError: + return + if not stat.S_ISSOCK(mode): + raise RuntimeError(f"refusing to replace non-socket path: {path}") + path.unlink() + + +def serve( + socket_path: Path, + limits: Limits, + socket_mode: int = 0o600, + once: bool = False, + stop: threading.Event | None = None, + resolver: ExternalTensorResolver | None = None, +) -> None: + stop = stop or threading.Event() + remove_stale_socket(socket_path) + listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + try: + listener.bind(os.fspath(socket_path)) + os.chmod(socket_path, socket_mode) + listener.listen(16) + listener.settimeout(0.25) + bound_identity = socket_path.lstat().st_dev, socket_path.lstat().st_ino + while not stop.is_set(): + try: + connection, _ = listener.accept() + except TimeoutError: + continue + with connection: + handle_connection(connection, limits, resolver) + if once: + break + finally: + listener.close() + try: + current = socket_path.lstat() + if (current.st_dev, current.st_ino) == bound_identity: + socket_path.unlink() + except (FileNotFoundError, UnboundLocalError): + pass + + +def parse_mode(value: str) -> int: + mode = int(value, 8) + if mode < 0 or mode > 0o777: + raise argparse.ArgumentTypeError("socket mode must be between 000 and 777") + return mode + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--socket", required=True, type=Path) + parser.add_argument("--socket-mode", type=parse_mode, default=0o600) + parser.add_argument("--max-request-bytes", type=int, default=64 * 1024 * 1024) + parser.add_argument("--max-batch-tokens", type=int, default=1_000_000) + parser.add_argument("--max-tensor-bytes", type=int, default=1024 * 1024 * 1024) + parser.add_argument("--max-candidates", type=int, default=65_536) + parser.add_argument("--once", action="store_true") + args = parser.parse_args() + limits = Limits( + max_request_bytes=args.max_request_bytes, + max_batch_tokens=args.max_batch_tokens, + max_tensor_bytes=args.max_tensor_bytes, + max_candidates=args.max_candidates, + ) + if ( + min( + limits.max_request_bytes, + limits.max_batch_tokens, + limits.max_tensor_bytes, + limits.max_candidates, + ) + <= 0 + ): + parser.error("all limits must be positive") + + stop = threading.Event() + + def request_stop(_signum: int, _frame: object) -> None: + stop.set() + + signal.signal(signal.SIGINT, request_stop) + signal.signal(signal.SIGTERM, request_stop) + serve(args.socket, limits, args.socket_mode, args.once, stop) + + +if __name__ == "__main__": + main() diff --git a/services/Dockerfile.tilemaxsim b/services/Dockerfile.tilemaxsim new file mode 100644 index 00000000..9f58b88d --- /dev/null +++ b/services/Dockerfile.tilemaxsim @@ -0,0 +1,12 @@ +ARG BASE_IMAGE +FROM ${BASE_IMAGE} + +WORKDIR /opt/vectorchord + +COPY devtools/tilemaxsim_reference_sidecar.py devtools/tilemaxsim_reference_sidecar.py +COPY services/tilemaxsim_cuda_sidecar.py services/tilemaxsim_cuda_sidecar.py + +ENV PYTHONPATH=/opt/vectorchord \ + PYTHONUNBUFFERED=1 + +ENTRYPOINT ["python3", "-m", "services.tilemaxsim_cuda_sidecar"] diff --git a/services/benchmark_tilemaxsim_cuda.py b/services/benchmark_tilemaxsim_cuda.py new file mode 100644 index 00000000..698e4b84 --- /dev/null +++ b/services/benchmark_tilemaxsim_cuda.py @@ -0,0 +1,152 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +"""Reproducible synthetic load probe for the CUDA TileMaxSim executor.""" + +from __future__ import annotations + +import argparse +import json +import math +import statistics +import time + +import torch + +from devtools import tilemaxsim_reference_sidecar as protocol +from services.tilemaxsim_cuda_sidecar import TorchTileMaxsimEngine, positive_int + + +def percentile(samples: list[float], fraction: float) -> float: + ordered = sorted(samples) + index = max(0, math.ceil(fraction * len(ordered)) - 1) + return ordered[index] + + +def canonical_payload(tensor: torch.Tensor, dtype: int) -> bytes: + scalar_dtype = torch.float32 if dtype == protocol.DTYPE_F32 else torch.float16 + return tensor.to(dtype=scalar_dtype).contiguous().numpy().tobytes() + + +def normalized_tensor( + shape: tuple[int, ...], generator: torch.Generator +) -> torch.Tensor: + tensor = torch.randn(shape, dtype=torch.float32, generator=generator) + return torch.nn.functional.normalize(tensor, dim=-1) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--dtype", choices=("f16", "f32"), default="f16") + parser.add_argument("--dimension", type=positive_int, default=320) + parser.add_argument("--query-rows", type=positive_int, default=32) + parser.add_argument("--document-rows", type=positive_int, default=747) + parser.add_argument("--candidates", type=positive_int, default=128) + parser.add_argument("--warmup", type=positive_int, default=3) + parser.add_argument("--iterations", type=positive_int, default=10) + parser.add_argument("--seed", type=int, default=20260713) + parser.add_argument( + "--max-device-bytes", type=positive_int, default=8 * 1024 * 1024 * 1024 + ) + parser.add_argument("--allow-tf32", action="store_true") + args = parser.parse_args() + + dtype = protocol.DTYPE_F32 if args.dtype == "f32" else protocol.DTYPE_F16 + generator = torch.Generator(device="cpu").manual_seed(args.seed) + query = normalized_tensor((args.query_rows, args.dimension), generator) + document_tensors = normalized_tensor( + (args.candidates, args.document_rows, args.dimension), generator + ) + query_payload = canonical_payload(query, dtype) + documents = [ + ( + candidate_id, + args.document_rows, + canonical_payload(document_tensors[candidate_id], dtype), + ) + for candidate_id in range(args.candidates) + ] + del document_tensors + + engine = TorchTileMaxsimEngine( + args.device, args.max_device_bytes, args.allow_tf32, 1 + ) + if engine.device.type == "cuda": + torch.cuda.reset_peak_memory_stats(engine.device) + + total_samples: list[float] = [] + queue_samples: list[float] = [] + compute_samples: list[float] = [] + score_checksum = 0.0 + for iteration in range(args.warmup + args.iterations): + started = time.perf_counter() + results, queue_ms, compute_ms = engine.score( + query_payload, + args.query_rows, + args.dimension, + dtype, + documents, + time.monotonic() + 300, + lambda: False, + ) + total_ms = (time.perf_counter() - started) * 1000.0 + if iteration >= args.warmup: + total_samples.append(total_ms) + queue_samples.append(queue_ms) + compute_samples.append(compute_ms) + score_checksum = math.fsum(score for _, score in results) + + output = { + "benchmark": "tilemaxsim_cuda_synthetic_v1", + "device": str(engine.device), + "device_name": ( + torch.cuda.get_device_name(engine.device) + if engine.device.type == "cuda" + else "cpu" + ), + "torch_version": torch.__version__, + "dtype": args.dtype, + "dimension": args.dimension, + "query_rows": args.query_rows, + "document_rows": args.document_rows, + "candidates": args.candidates, + "candidate_tokens": args.candidates * args.document_rows, + "seed": args.seed, + "warmup": args.warmup, + "iterations": args.iterations, + "allow_tf32": args.allow_tf32, + "max_device_bytes": args.max_device_bytes, + "latency_ms": { + "mean": round(statistics.fmean(total_samples), 3), + "p50": round(percentile(total_samples, 0.50), 3), + "p95": round(percentile(total_samples, 0.95), 3), + "p99": round(percentile(total_samples, 0.99), 3), + "queue_mean": round(statistics.fmean(queue_samples), 3), + "compute_mean": round(statistics.fmean(compute_samples), 3), + }, + "score_checksum": score_checksum, + } + if engine.device.type == "cuda": + output["cuda_peak_allocated_bytes"] = torch.cuda.max_memory_allocated( + engine.device + ) + output["cuda_peak_reserved_bytes"] = torch.cuda.max_memory_reserved( + engine.device + ) + print(json.dumps(output, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/services/build_tilemaxsim_tensor_cache.py b/services/build_tilemaxsim_tensor_cache.py new file mode 100644 index 00000000..a48a1728 --- /dev/null +++ b/services/build_tilemaxsim_tensor_cache.py @@ -0,0 +1,235 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +"""Publish NPY page tensors into the sidecar's immutable SHA-256 cache.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import sys +import tempfile +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import numpy as np + +from services.tilemaxsim_cuda_sidecar import positive_int + + +def canonical_tensor(path: Path, expected_rows: int, expected_dim: int) -> np.ndarray: + tensor = np.load(path, mmap_mode="r", allow_pickle=False) + if tensor.ndim != 2 or tensor.shape != (expected_rows, expected_dim): + raise ValueError( + f"{path}: expected shape {(expected_rows, expected_dim)}, got {tensor.shape}" + ) + if tensor.dtype == np.dtype("float16"): + little_dtype = np.dtype(" str: + digest = hashlib.sha256() + with path.open("rb") as stream: + while chunk := stream.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def write_payload(path: Path, payload: memoryview, digest: str, fsync: bool) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + if path.exists(): + if path.stat().st_size != len(payload): + raise ValueError(f"existing cache payload has wrong size: {path}") + if file_digest(path) != digest: + raise ValueError(f"existing cache payload has wrong checksum: {path}") + return + descriptor, temporary_name = tempfile.mkstemp( + prefix=f".{path.name}.", suffix=".tmp", dir=path.parent + ) + try: + with os.fdopen(descriptor, "wb", closefd=True) as stream: + stream.write(payload) + stream.flush() + if fsync: + os.fsync(stream.fileno()) + try: + os.link(temporary_name, path) + except FileExistsError: + if path.stat().st_size != len(payload): + raise ValueError(f"concurrent cache payload has wrong size: {path}") + if file_digest(path) != digest: + raise ValueError(f"concurrent cache payload has wrong checksum: {path}") + if fsync: + directory_fd = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY) + try: + os.fsync(directory_fd) + finally: + os.close(directory_fd) + finally: + try: + os.unlink(temporary_name) + except FileNotFoundError: + pass + + +def process_record( + record: dict[str, object], + source_root: Path, + cache_root: Path, + fsync: bool, + dry_run: bool, +) -> dict[str, object]: + page_key = record.get("page_key") + relative = record.get("embedding_file") + rows = record.get("n_tokens") + dimension = record.get("dim") + if not isinstance(page_key, str) or not page_key: + raise ValueError("manifest record has no page_key") + if not isinstance(relative, str) or not relative: + raise ValueError(f"manifest page {page_key} has no embedding_file") + if not isinstance(rows, int) or rows <= 0: + raise ValueError(f"manifest page {page_key} has invalid n_tokens") + if not isinstance(dimension, int) or dimension <= 0: + raise ValueError(f"manifest page {page_key} has invalid dim") + source = (source_root / relative).resolve(strict=True) + try: + source.relative_to(source_root) + except ValueError as error: + raise ValueError( + f"embedding path escapes the source root: {relative}" + ) from error + tensor = canonical_tensor(source, rows, dimension) + payload = memoryview(tensor).cast("B") + digest = hashlib.sha256(payload).hexdigest() + destination = cache_root / digest[:2] / f"{digest}.bin" + if not dry_run: + write_payload(destination, payload, digest, fsync) + dtype_name = "float16" if tensor.dtype == np.dtype("float16") else "float32" + return { + "page_key": page_key, + "tensor_ref": f"sha256://{digest}", + "tensor_rows": rows, + "tensor_dim": dimension, + "tensor_dtype": dtype_name, + "tensor_checksum": f"sha256:{digest}", + "canonical_bytes": len(payload), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--manifest", required=True, type=Path) + parser.add_argument("--cache-root", required=True, type=Path) + parser.add_argument("--descriptor-manifest", required=True, type=Path) + parser.add_argument("--workers", type=positive_int, default=4) + parser.add_argument("--no-fsync", action="store_true") + parser.add_argument("--dry-run", action="store_true") + args = parser.parse_args() + + source_root = args.manifest.resolve(strict=True).parent + cache_root = args.cache_root.resolve() + if not cache_root.is_absolute(): + parser.error("--cache-root must be absolute") + records = [] + with args.manifest.open(encoding="utf-8") as stream: + for line_number, line in enumerate(stream, 1): + try: + record = json.loads(line) + except json.JSONDecodeError as error: + raise ValueError( + f"invalid JSON at manifest line {line_number}" + ) from error + if not isinstance(record, dict): + raise ValueError(f"manifest line {line_number} is not an object") + records.append(record) + if not records: + raise ValueError("manifest is empty") + + cache_root.mkdir(parents=True, exist_ok=True) + descriptors = [] + with ThreadPoolExecutor(max_workers=args.workers) as workers: + results = workers.map( + lambda record: process_record( + record, + source_root, + cache_root, + not args.no_fsync, + args.dry_run, + ), + records, + ) + for completed, item in enumerate(results, 1): + descriptors.append(item) + if completed % 1000 == 0 or completed == len(records): + print( + json.dumps( + {"event": "tensor_cache_progress", "completed": completed}, + separators=(",", ":"), + ), + file=sys.stderr, + flush=True, + ) + + output_parent = args.descriptor_manifest.parent + output_parent.mkdir(parents=True, exist_ok=True) + descriptor, temporary_name = tempfile.mkstemp( + prefix=f".{args.descriptor_manifest.name}.", + suffix=".tmp", + dir=output_parent, + text=True, + ) + try: + with os.fdopen(descriptor, "w", encoding="utf-8", closefd=True) as stream: + for item in descriptors: + stream.write(json.dumps(item, separators=(",", ":"), sort_keys=True)) + stream.write("\n") + stream.flush() + os.fsync(stream.fileno()) + os.replace(temporary_name, args.descriptor_manifest) + finally: + try: + os.unlink(temporary_name) + except FileNotFoundError: + pass + + total_bytes = sum(int(item["canonical_bytes"]) for item in descriptors) + print( + json.dumps( + { + "pages": len(descriptors), + "canonical_bytes": total_bytes, + "cache_root": os.fspath(cache_root), + "descriptor_manifest": os.fspath(args.descriptor_manifest), + "dry_run": args.dry_run, + }, + sort_keys=True, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/services/test_tilemaxsim_cuda_sidecar.py b/services/test_tilemaxsim_cuda_sidecar.py new file mode 100644 index 00000000..093245e3 --- /dev/null +++ b/services/test_tilemaxsim_cuda_sidecar.py @@ -0,0 +1,313 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +from __future__ import annotations + +import hashlib +import os +import socket +import stat +import struct +import tempfile +import threading +import time +import unittest +from pathlib import Path + +import numpy as np +import torch + +from devtools import tilemaxsim_reference_sidecar as protocol +from devtools.test_tilemaxsim_reference_sidecar import ( + decode_response, + external_request_frame, + request_frame, +) +from services import tilemaxsim_cuda_sidecar as cuda_sidecar +from services.build_tilemaxsim_tensor_cache import process_record + + +class CapturingMetrics(cuda_sidecar.JsonMetrics): + def __init__(self) -> None: + super().__init__() + self.events: list[dict[str, object]] = [] + + def emit(self, fields: dict[str, object]) -> None: + with self.lock: + self.events.append(fields.copy()) + + +def write_content_addressed(root: Path, payload: bytes) -> tuple[str, str]: + digest = hashlib.sha256(payload).hexdigest() + directory = root / digest[:2] + directory.mkdir(parents=True, exist_ok=True) + (directory / f"{digest}.bin").write_bytes(payload) + return f"sha256://{digest}", f"sha256:{digest}" + + +class CudaSidecarTest(unittest.TestCase): + def test_cache_builder_publishes_resolver_compatible_payload(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + source_root = root / "source" + cache_root = root / "cache" + source_root.mkdir() + tensor = np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=" None: + cuda_sidecar.validate_finite_payload( + struct.pack("<2e", 1.0, 0.0), 1, 2, protocol.DTYPE_F16 + ) + with self.assertRaisesRegex(protocol.SidecarError, "non-finite"): + cuda_sidecar.validate_finite_payload( + struct.pack("<2f", 1.0, float("nan")), + 1, + 2, + protocol.DTYPE_F32, + ) + + def test_content_addressed_resolver_validates_and_caches(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + payload = struct.pack("<4e", 1.0, 0.0, 0.0, 1.0) + tensor_ref, checksum = write_content_addressed(root, payload) + resolver = cuda_sidecar.ContentAddressedResolver({"model@1": root}, 1024) + try: + request = protocol.ExternalTensorRequest( + "model@1", + tensor_ref, + 2, + 2, + protocol.DTYPE_F16, + checksum, + ) + first = resolver.resolve(request) + second = resolver.resolve(request) + self.assertEqual(first.payload, payload) + self.assertFalse(first.cache_hit) + self.assertTrue(second.cache_hit) + + bad = protocol.ExternalTensorRequest( + "model@1", + tensor_ref, + 2, + 2, + protocol.DTYPE_F16, + "sha256:" + "0" * 64, + ) + with self.assertRaisesRegex(protocol.SidecarError, "disagree"): + resolver.resolve(bad) + finally: + resolver.close() + + def test_content_addressed_resolver_rejects_symlink(self) -> None: + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) / "root" + root.mkdir() + payload = struct.pack(" None: + query = [[1.0, 0.0], [0.0, 1.0]] + candidates = [ + (17, [[1.0, 0.0], [0.0, 1.0]]), + (3, [[0.5, 0.5], [0.25, 0.25]]), + ] + frame = request_frame(41, protocol.DTYPE_F32, query, candidates) + parsed = protocol.parse_request_frame(frame) + self.assertIsInstance(parsed, protocol.InlineTensorRequest) + assert isinstance(parsed, protocol.InlineTensorRequest) + documents = [ + (candidate.candidate_id, candidate.rows, candidate.payload) + for candidate in parsed.candidates + ] + # 64 bytes fits one candidate but not both, exercising internal + # all-or-nothing device chunking. + engine = cuda_sidecar.TorchTileMaxsimEngine("cpu", 64, False, 1) + results, _, _ = engine.score( + parsed.query_payload, + parsed.query_rows, + parsed.dimension, + parsed.dtype, + documents, + time.monotonic() + 2, + lambda: False, + ) + _, status, oracle = decode_response(protocol.process_frame(frame)) + self.assertEqual(status, 0) + self.assertEqual(results, oracle) + + def test_compute_capacity_wait_uses_overall_deadline(self) -> None: + engine = cuda_sidecar.TorchTileMaxsimEngine("cpu", 1024, False, 1) + self.assertTrue(engine.compute_slots.acquire(blocking=False)) + try: + started = time.monotonic() + with self.assertRaisesRegex(protocol.SidecarError, "CUDA capacity"): + engine.score( + struct.pack(" None: + query = [[1.0, 0.0, 0.5], [0.0, 1.0, -0.25]] + candidates = [ + (7, [[1.0, 0.0, 0.5], [0.0, 1.0, -0.25]]), + (2, [[0.5, 0.5, 0.0], [-0.5, 0.25, 1.0]]), + ] + frame = request_frame(52, protocol.DTYPE_F16, query, candidates) + parsed = protocol.parse_request_frame(frame) + assert isinstance(parsed, protocol.InlineTensorRequest) + engine = cuda_sidecar.TorchTileMaxsimEngine("cuda:0", 1024 * 1024, False, 1) + results, _, _ = engine.score( + parsed.query_payload, + parsed.query_rows, + parsed.dimension, + parsed.dtype, + [ + (candidate.candidate_id, candidate.rows, candidate.payload) + for candidate in parsed.candidates + ], + time.monotonic() + 5, + lambda: False, + ) + _, status, oracle = decode_response(protocol.process_frame(frame)) + self.assertEqual(status, 0) + assert isinstance(oracle, list) + self.assertEqual([item[0] for item in results], [item[0] for item in oracle]) + for (_, actual), (_, expected) in zip(results, oracle, strict=True): + self.assertAlmostEqual(actual, expected, places=5) + + def test_v2_unix_socket_end_to_end(self) -> None: + with tempfile.TemporaryDirectory() as directory: + directory_path = Path(directory) + root = directory_path / "tensors" + root.mkdir() + tensor = [[1.0, 0.0], [0.0, 1.0]] + payload = struct.pack("<4e", *sum(tensor, [])) + tensor_ref, _ = write_content_addressed(root, payload) + frame, _ = external_request_frame( + 61, + protocol.DTYPE_F16, + [[1.0, 0.0], [0.0, 1.0]], + "model@1", + [(9, tensor_ref, tensor)], + ) + resolver = cuda_sidecar.ContentAddressedResolver({"model@1": root}, 1024) + metrics = CapturingMetrics() + service = cuda_sidecar.TileMaxsimService( + protocol.Limits(), + resolver, + cuda_sidecar.TorchTileMaxsimEngine("cpu", 1024 * 1024, False, 1), + 2000, + metrics, + ) + socket_path = directory_path / "tilemaxsim.sock" + stop = threading.Event() + thread = threading.Thread( + target=cuda_sidecar.serve, + args=(socket_path, 0o600, 4, 2, service, stop), + kwargs={"once": True}, + daemon=True, + ) + thread.start() + for _ in range(100): + if socket_path.exists(): + break + time.sleep(0.01) + else: + self.fail("CUDA sidecar socket was not created") + self.assertEqual(stat.S_IMODE(socket_path.stat().st_mode), 0o600) + + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as connection: + connection.connect(os.fspath(socket_path)) + connection.sendall(frame) + header = protocol.receive_exact(connection, protocol.HEADER.size) + body_len = protocol.HEADER.unpack(header)[4] + response = header + protocol.receive_exact(connection, body_len) + thread.join(timeout=3) + resolver.close() + self.assertFalse(thread.is_alive()) + _, status, results = decode_response(response) + self.assertEqual(status, 0) + self.assertEqual(results, [(9, 2.0)]) + request_events = [ + event + for event in metrics.events + if event.get("event") == "tilemaxsim_request" + ] + self.assertEqual(len(request_events), 1) + self.assertEqual(request_events[0]["source"], "content_addressed") + self.assertEqual(request_events[0]["status"], "ok") + + +if __name__ == "__main__": + unittest.main() diff --git a/services/tilemaxsim_cuda_sidecar.py b/services/tilemaxsim_cuda_sidecar.py new file mode 100644 index 00000000..ce744119 --- /dev/null +++ b/services/tilemaxsim_cuda_sidecar.py @@ -0,0 +1,802 @@ +# This software is licensed under a dual license model: +# +# GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +# distribute this software under the terms of the AGPLv3. +# +# Elastic License v2 (ELv2): You may also use, modify, and distribute this +# software under the Elastic License v2, which has specific restrictions. +# +# We welcome any commercial collaboration or support. For inquiries +# regarding the licenses, please contact us at: +# vectorchord-inquiry@tensorchord.ai +# +# Copyright (c) 2025-2026 TensorChord Inc. + +"""Bounded CUDA executor for VectorChord TileMaxSim IPC v1 and v2. + +Version 1 consumes inline tensors. Version 2 resolves canonical tensor payloads +from an operations-configured, per-model content-addressed cache. Applications may +populate that cache from any object store; object-store credentials and routing +never enter PostgreSQL or this protocol. +""" + +from __future__ import annotations + +import argparse +import hashlib +import hmac +import json +import math +import os +import select +import signal +import socket +import stat +import struct +import sys +import threading +import time +from collections import OrderedDict +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Iterable + +import numpy as np +import torch + +if __package__ in (None, ""): + sys.path.insert(0, os.fspath(Path(__file__).resolve().parents[1])) + +from devtools import tilemaxsim_reference_sidecar as protocol + + +def validate_finite_payload( + payload: bytes, rows: int, dimension: int, dtype: int +) -> None: + expected = protocol.checked_tensor_bytes(rows, dimension, dtype) + if len(payload) != expected: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "tensor byte length does not match its shape", + ) + scalar_dtype = " None: + self.maximum_bytes = maximum_bytes + self.current_bytes = 0 + self.entries: OrderedDict[tuple[object, ...], bytes] = OrderedDict() + self.lock = threading.Lock() + + def get(self, key: tuple[object, ...]) -> bytes | None: + if self.maximum_bytes == 0: + return None + with self.lock: + payload = self.entries.get(key) + if payload is not None: + self.entries.move_to_end(key) + return payload + + def put(self, key: tuple[object, ...], payload: bytes) -> None: + if self.maximum_bytes == 0 or len(payload) > self.maximum_bytes: + return + with self.lock: + previous = self.entries.pop(key, None) + if previous is not None: + self.current_bytes -= len(previous) + self.entries[key] = payload + self.current_bytes += len(payload) + while self.current_bytes > self.maximum_bytes: + _, evicted = self.entries.popitem(last=False) + self.current_bytes -= len(evicted) + + +class ContentAddressedResolver: + """Resolve ``sha256://`` inside an allowlisted model cache root. + + A payload with digest ``abcdef...`` is stored as + ``/ab/abcdef....bin``. Directory and file symlinks are + rejected with ``openat(O_NOFOLLOW)``. The digest in the reference, the + registered checksum, the exact byte length, and the file content must all + agree before a payload is returned. + """ + + def __init__(self, roots: dict[str, Path], cache_bytes: int) -> None: + self.root_fds: dict[str, int] = {} + try: + for contract, path in roots.items(): + self.root_fds[contract] = os.open( + os.fspath(path), + os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW, + ) + except Exception: + self.close() + raise + self.cache = PayloadCache(cache_bytes) + + def close(self) -> None: + for descriptor in self.root_fds.values(): + os.close(descriptor) + self.root_fds.clear() + + @staticmethod + def _digest(request: protocol.ExternalTensorRequest) -> str: + prefix = "sha256://" + if not request.tensor_ref.startswith(prefix): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "unsupported tensor reference; expected sha256://", + ) + digest = request.tensor_ref[len(prefix) :] + if len(digest) != 64 or any( + character not in "0123456789abcdef" for character in digest + ): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "invalid content-addressed tensor reference", + ) + if not hmac.compare_digest(request.checksum, f"sha256:{digest}"): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "tensor reference and checksum disagree", + ) + return digest + + @staticmethod + def _read_exact_file(root_fd: int, digest: str, expected_bytes: int) -> bytes: + directory_fd = -1 + payload_fd = -1 + try: + directory_fd = os.open( + digest[:2], + os.O_RDONLY | os.O_DIRECTORY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=root_fd, + ) + payload_fd = os.open( + f"{digest}.bin", + os.O_RDONLY | os.O_CLOEXEC | os.O_NOFOLLOW, + dir_fd=directory_fd, + ) + metadata = os.fstat(payload_fd) + if not stat.S_ISREG(metadata.st_mode): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor is not a regular file", + ) + if metadata.st_size != expected_bytes: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor byte length does not match descriptor", + ) + chunks = bytearray() + while len(chunks) < expected_bytes: + chunk = os.read( + payload_fd, min(1024 * 1024, expected_bytes - len(chunks)) + ) + if not chunk: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor file ended early", + ) + chunks.extend(chunk) + if os.read(payload_fd, 1): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "resolved tensor file grew during read", + ) + return bytes(chunks) + except FileNotFoundError as error: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, "content-addressed tensor is missing" + ) from error + except OSError as error: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + f"content-addressed tensor read failed: {error.strerror}", + ) from error + finally: + if payload_fd >= 0: + os.close(payload_fd) + if directory_fd >= 0: + os.close(directory_fd) + + def resolve(self, request: protocol.ExternalTensorRequest) -> ResolvedPayload: + root_fd = self.root_fds.get(request.model_contract_id) + if root_fd is None: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "model contract has no configured tensor cache root", + ) + digest = self._digest(request) + expected_bytes = protocol.checked_tensor_bytes( + request.rows, request.dimension, request.dtype + ) + key = ( + request.model_contract_id, + digest, + request.rows, + request.dimension, + request.dtype, + ) + cached = self.cache.get(key) + if cached is not None: + return ResolvedPayload(cached, True) + payload = self._read_exact_file(root_fd, digest, expected_bytes) + actual = hashlib.sha256(payload).hexdigest() + if not hmac.compare_digest(actual, digest): + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, "resolved tensor checksum mismatch" + ) + validate_finite_payload(payload, request.rows, request.dimension, request.dtype) + self.cache.put(key, payload) + return ResolvedPayload(payload, False) + + +class TorchTileMaxsimEngine: + def __init__( + self, + device_name: str, + max_device_bytes: int, + allow_tf32: bool, + max_cuda_inflight: int, + ) -> None: + self.device = torch.device(device_name) + if self.device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError( + "CUDA was requested but torch.cuda.is_available() is false" + ) + if self.device.type not in ("cuda", "cpu"): + raise RuntimeError("device must be CUDA or CPU") + self.max_device_bytes = max_device_bytes + self.compute_slots = threading.BoundedSemaphore(max_cuda_inflight) + if self.device.type == "cuda": + torch.backends.cuda.matmul.allow_tf32 = allow_tf32 + torch.backends.cudnn.allow_tf32 = allow_tf32 + with torch.inference_mode(): + left = torch.zeros((1, 1), dtype=torch.float32, device=self.device) + _ = left @ left + torch.cuda.synchronize(self.device) + + @staticmethod + def _cpu_tensor( + payload: bytes, rows: int, dimension: int, dtype: int + ) -> torch.Tensor: + scalar_dtype = torch.float32 if dtype == protocol.DTYPE_F32 else torch.float16 + # bytearray gives torch a writable, owned buffer; clone detaches the + # resulting tensor before that temporary buffer leaves scope. + tensor = torch.frombuffer(bytearray(payload), dtype=scalar_dtype).reshape( + rows, dimension + ) + if scalar_dtype == torch.float32: + return tensor.clone() + return tensor.to(dtype=torch.float32) + + def _groups( + self, + query_rows: int, + dimension: int, + documents: list[tuple[int, int, bytes]], + ) -> Iterable[list[tuple[int, int, bytes]]]: + query_bytes = query_rows * dimension * 4 + group: list[tuple[int, int, bytes]] = [] + group_rows = 0 + for document in documents: + rows = document[1] + next_rows = group_rows + rows + # Device residency includes the f32 query, f32 documents, and the + # q-by-total-document-token similarity matrix. + required = ( + query_bytes + next_rows * dimension * 4 + query_rows * next_rows * 4 + ) + if required > self.max_device_bytes and group: + yield group + group = [] + group_rows = 0 + next_rows = rows + required = query_bytes + rows * dimension * 4 + query_rows * rows * 4 + if required > self.max_device_bytes: + raise protocol.SidecarError( + protocol.STATUS_RESOURCE_LIMIT, + "one candidate exceeds the CUDA device-byte limit", + ) + group.append(document) + group_rows = next_rows + if group: + yield group + + def score( + self, + query_payload: bytes, + query_rows: int, + dimension: int, + dtype: int, + documents: list[tuple[int, int, bytes]], + deadline: float, + cancelled: Callable[[], bool], + ) -> tuple[list[tuple[int, float]], float, float]: + if not documents: + return [], 0.0, 0.0 + query_cpu = self._cpu_tensor(query_payload, query_rows, dimension, dtype) + results: list[tuple[int, float]] = [] + queue_started = time.monotonic() + remaining = deadline - time.monotonic() + if remaining <= 0 or not self.compute_slots.acquire(timeout=remaining): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "request deadline expired while waiting for CUDA capacity", + ) + queue_ms = (time.monotonic() - queue_started) * 1000.0 + compute_started = time.monotonic() + try: + with torch.inference_mode(): + if time.monotonic() >= deadline: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, "request deadline expired" + ) + query_device = query_cpu.to(self.device) + for group in self._groups(query_rows, dimension, documents): + if cancelled(): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, "request peer disconnected" + ) + if time.monotonic() >= deadline: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, "request deadline expired" + ) + cpu_documents = [ + self._cpu_tensor(payload, rows, dimension, dtype) + for _, rows, payload in group + ] + document_device = torch.cat(cpu_documents).to(self.device) + similarities = query_device @ document_device.transpose(0, 1) + scores = [] + offset = 0 + for _, rows, _ in group: + scores.append( + similarities[:, offset : offset + rows] + .amax(dim=1) + .sum(dtype=torch.float32) + ) + offset += rows + host_scores = torch.stack(scores).to(device="cpu").tolist() + for (candidate_id, _, _), score in zip( + group, host_scores, strict=True + ): + if not math.isfinite(score): + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "TileMaxSim result is non-finite", + ) + results.append((candidate_id, score)) + if self.device.type == "cuda": + torch.cuda.synchronize(self.device) + finally: + self.compute_slots.release() + return ( + results, + queue_ms, + (time.monotonic() - compute_started) * 1000.0, + ) + + +class JsonMetrics: + def __init__(self) -> None: + self.lock = threading.Lock() + + def emit(self, fields: dict[str, object]) -> None: + with self.lock: + print(json.dumps(fields, separators=(",", ":"), sort_keys=True), flush=True) + + +class TileMaxsimService: + def __init__( + self, + limits: protocol.Limits, + resolver: ContentAddressedResolver, + engine: TorchTileMaxsimEngine, + request_timeout_ms: int, + metrics: JsonMetrics, + ) -> None: + self.limits = limits + self.resolver = resolver + self.engine = engine + self.request_timeout_seconds = request_timeout_ms / 1000.0 + self.metrics = metrics + + @staticmethod + def _peer_disconnected(connection: socket.socket) -> bool: + poller = select.poll() + poller.register(connection, select.POLLHUP | select.POLLERR | select.POLLNVAL) + return bool(poller.poll(0)) + + @staticmethod + def _receive_exact_until( + connection: socket.socket, count: int, deadline: float + ) -> bytes: + chunks = bytearray() + while len(chunks) < count: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("request deadline expired during socket read") + connection.settimeout(remaining) + chunk = connection.recv(count - len(chunks)) + if not chunk: + raise protocol.SidecarError( + protocol.STATUS_INVALID_REQUEST, + "connection closed during request", + ) + chunks.extend(chunk) + return bytes(chunks) + + def process_frame( + self, + frame: bytes, + connection: socket.socket, + deadline: float, + peer_credentials: tuple[int, int, int] | None, + ) -> bytes: + request_id = 0 + version = protocol.VERSION + started = time.monotonic() + metrics: dict[str, object] = {"event": "tilemaxsim_request"} + if peer_credentials is not None: + metrics["peer_pid"], metrics["peer_uid"], metrics["peer_gid"] = ( + peer_credentials + ) + try: + if len(frame) >= protocol.HEADER.size: + _, wire_version, _, request_id, _ = protocol.HEADER.unpack_from(frame) + if wire_version in (protocol.VERSION, protocol.EXTERNAL_VERSION): + version = wire_version + request = protocol.parse_request_frame( + frame, self.limits, validate_finite=False + ) + metrics.update( + request_id=request.request_id, + protocol_version=version, + query_rows=request.query_rows, + dimension=request.dimension, + candidate_count=len(request.candidates), + ) + validate_finite_payload( + request.query_payload, + request.query_rows, + request.dimension, + request.dtype, + ) + resolve_started = time.monotonic() + cache_hits = 0 + documents: list[tuple[int, int, bytes]] = [] + if isinstance(request, protocol.InlineTensorRequest): + for candidate in request.candidates: + validate_finite_payload( + candidate.payload, + candidate.rows, + request.dimension, + request.dtype, + ) + documents = [ + (candidate.candidate_id, candidate.rows, candidate.payload) + for candidate in request.candidates + ] + metrics["source"] = "inline" + else: + metrics["source"] = "content_addressed" + for candidate in request.candidates: + if time.monotonic() >= deadline: + raise protocol.SidecarError( + protocol.STATUS_COMPUTE_ERROR, + "request deadline expired during tensor resolution", + ) + resolved = self.resolver.resolve(candidate.descriptor) + cache_hits += int(resolved.cache_hit) + documents.append( + ( + candidate.candidate_id, + candidate.descriptor.rows, + resolved.payload, + ) + ) + metrics["cache_hits"] = cache_hits + metrics["resolve_ms"] = round( + (time.monotonic() - resolve_started) * 1000.0, 3 + ) + metrics["document_tokens"] = sum(rows for _, rows, _ in documents) + results, queue_ms, compute_ms = self.engine.score( + request.query_payload, + request.query_rows, + request.dimension, + request.dtype, + documents, + deadline, + lambda: self._peer_disconnected(connection), + ) + metrics["queue_ms"] = round(queue_ms, 3) + metrics["compute_ms"] = round(compute_ms, 3) + metrics["status"] = "ok" + return protocol.success_response(request.request_id, results, version) + except protocol.SidecarError as error: + metrics.update(status="error", error_class=type(error).__name__) + return protocol.error_response( + request_id, error.status, str(error), version + ) + except torch.OutOfMemoryError: + if self.engine.device.type == "cuda": + torch.cuda.empty_cache() + metrics.update(status="error", error_class="CudaOutOfMemory") + return protocol.error_response( + request_id, + protocol.STATUS_RESOURCE_LIMIT, + "CUDA out of memory", + version, + ) + except Exception as error: + metrics.update(status="error", error_class=type(error).__name__) + return protocol.error_response( + request_id, + protocol.STATUS_COMPUTE_ERROR, + f"TileMaxSim compute failed: {error}", + version, + ) + finally: + metrics["total_ms"] = round((time.monotonic() - started) * 1000.0, 3) + self.metrics.emit(metrics) + + def handle(self, connection: socket.socket) -> None: + deadline = time.monotonic() + self.request_timeout_seconds + request_id = 0 + version = protocol.VERSION + peer_credentials = None + if hasattr(socket, "SO_PEERCRED"): + try: + raw_credentials = connection.getsockopt( + socket.SOL_SOCKET, socket.SO_PEERCRED, struct.calcsize("3i") + ) + peer_credentials = struct.unpack("3i", raw_credentials) + except OSError: + pass + try: + header = self._receive_exact_until( + connection, protocol.HEADER.size, deadline + ) + _, wire_version, _, request_id, body_len = protocol.HEADER.unpack(header) + if wire_version in (protocol.VERSION, protocol.EXTERNAL_VERSION): + version = wire_version + if body_len > self.limits.max_request_bytes - protocol.HEADER.size: + response = protocol.error_response( + request_id, + protocol.STATUS_RESOURCE_LIMIT, + "request exceeds byte limit", + version, + ) + else: + body = self._receive_exact_until(connection, body_len, deadline) + response = self.process_frame( + header + body, connection, deadline, peer_credentials + ) + remaining = deadline - time.monotonic() + if remaining <= 0: + raise TimeoutError("request deadline expired before socket write") + connection.settimeout(remaining) + connection.sendall(response) + except (BrokenPipeError, ConnectionResetError): + return + except (TimeoutError, socket.timeout): + try: + connection.sendall( + protocol.error_response( + request_id, + protocol.STATUS_COMPUTE_ERROR, + "request deadline expired during socket I/O", + version, + ) + ) + except OSError: + pass + except Exception as error: + try: + connection.sendall( + protocol.error_response( + request_id, + protocol.STATUS_COMPUTE_ERROR, + str(error), + version, + ) + ) + except OSError: + pass + + +def serve( + socket_path: Path, + socket_mode: int, + backlog: int, + max_inflight: int, + service: TileMaxsimService, + stop: threading.Event, + once: bool = False, +) -> None: + protocol.remove_stale_socket(socket_path) + listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + slots = threading.BoundedSemaphore(max_inflight) + workers = ThreadPoolExecutor( + max_workers=max_inflight, thread_name_prefix="tilemaxsim" + ) + active = 0 + active_lock = threading.Lock() + + def handle(connection: socket.socket) -> None: + nonlocal active + try: + with connection: + service.handle(connection) + finally: + with active_lock: + active -= 1 + slots.release() + + try: + listener.bind(os.fspath(socket_path)) + os.chmod(socket_path, socket_mode) + listener.listen(backlog) + listener.settimeout(0.25) + bound_identity = socket_path.lstat().st_dev, socket_path.lstat().st_ino + service.metrics.emit( + { + "event": "tilemaxsim_ready", + "device": str(service.engine.device), + "max_inflight": max_inflight, + "socket": os.fspath(socket_path), + } + ) + accepted = 0 + while not stop.is_set(): + if not slots.acquire(timeout=0.25): + continue + try: + connection, _ = listener.accept() + except TimeoutError: + slots.release() + continue + with active_lock: + active += 1 + current_active = active + service.metrics.emit( + {"event": "tilemaxsim_accept", "inflight": current_active} + ) + workers.submit(handle, connection) + accepted += 1 + if once and accepted == 1: + break + finally: + listener.close() + workers.shutdown(wait=True, cancel_futures=False) + try: + current = socket_path.lstat() + if (current.st_dev, current.st_ino) == bound_identity: + socket_path.unlink() + except (FileNotFoundError, UnboundLocalError): + pass + + +def parse_mode(value: str) -> int: + mode = int(value, 8) + if mode < 0 or mode > 0o777: + raise argparse.ArgumentTypeError("socket mode must be between 000 and 777") + return mode + + +def positive_int(value: str) -> int: + parsed = int(value) + if parsed <= 0: + raise argparse.ArgumentTypeError("value must be positive") + return parsed + + +def nonnegative_int(value: str) -> int: + parsed = int(value) + if parsed < 0: + raise argparse.ArgumentTypeError("value must be nonnegative") + return parsed + + +def contract_roots( + values: list[str], parser: argparse.ArgumentParser +) -> dict[str, Path]: + roots = {} + for value in values: + if "=" not in value: + parser.error("--contract-root must be MODEL_CONTRACT_ID=/absolute/path") + contract, raw_path = value.split("=", 1) + path = Path(raw_path) + if not contract or not path.is_absolute(): + parser.error("--contract-root must contain a nonempty ID and absolute path") + if contract in roots: + parser.error(f"duplicate --contract-root for {contract!r}") + roots[contract] = path + return roots + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--socket", required=True, type=Path) + parser.add_argument("--socket-mode", type=parse_mode, default=0o600) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--contract-root", action="append", default=[]) + parser.add_argument( + "--max-request-bytes", type=positive_int, default=64 * 1024 * 1024 + ) + parser.add_argument("--max-batch-tokens", type=positive_int, default=1_000_000) + parser.add_argument( + "--max-tensor-bytes", type=positive_int, default=1024 * 1024 * 1024 + ) + parser.add_argument("--max-candidates", type=positive_int, default=65_536) + parser.add_argument( + "--max-device-bytes", type=positive_int, default=8 * 1024 * 1024 * 1024 + ) + parser.add_argument( + "--cache-bytes", type=nonnegative_int, default=8 * 1024 * 1024 * 1024 + ) + parser.add_argument("--request-timeout-ms", type=positive_int, default=2000) + parser.add_argument("--max-inflight", type=positive_int, default=8) + parser.add_argument("--max-cuda-inflight", type=positive_int, default=1) + parser.add_argument("--backlog", type=positive_int, default=64) + parser.add_argument("--allow-tf32", action="store_true") + parser.add_argument("--once", action="store_true") + args = parser.parse_args() + + roots = contract_roots(args.contract_root, parser) + limits = protocol.Limits( + max_request_bytes=args.max_request_bytes, + max_batch_tokens=args.max_batch_tokens, + max_tensor_bytes=args.max_tensor_bytes, + max_candidates=args.max_candidates, + ) + resolver = ContentAddressedResolver(roots, args.cache_bytes) + metrics = JsonMetrics() + try: + engine = TorchTileMaxsimEngine( + args.device, + args.max_device_bytes, + args.allow_tf32, + args.max_cuda_inflight, + ) + service = TileMaxsimService( + limits, resolver, engine, args.request_timeout_ms, metrics + ) + stop = threading.Event() + + def request_stop(_signum: int, _frame: object) -> None: + stop.set() + + signal.signal(signal.SIGINT, request_stop) + signal.signal(signal.SIGTERM, request_stop) + serve( + args.socket, + args.socket_mode, + args.backlog, + args.max_inflight, + service, + stop, + args.once, + ) + finally: + resolver.close() + + +if __name__ == "__main__": + main() diff --git a/src/datatype/mod.rs b/src/datatype/mod.rs index 3725f050..ac550103 100644 --- a/src/datatype/mod.rs +++ b/src/datatype/mod.rs @@ -29,3 +29,17 @@ mod text_rabitq8; pub mod typmod; mod typmod_rabitq4; mod typmod_rabitq8; + +pub(crate) const MAX_MAXSIM_VECTORS: usize = u16::MAX as usize + 1; + +pub(crate) fn validate_maxsim_array_len(len: usize) { + if len == 0 { + pgrx::error!("MaxSim arrays must contain at least one vector"); + } + if len > MAX_MAXSIM_VECTORS { + pgrx::error!( + "MaxSim arrays cannot contain more than {} vectors", + MAX_MAXSIM_VECTORS + ); + } +} diff --git a/src/datatype/operators_halfvec.rs b/src/datatype/operators_halfvec.rs index 0ae307fe..f20ebc8d 100644 --- a/src/datatype/operators_halfvec.rs +++ b/src/datatype/operators_halfvec.rs @@ -95,12 +95,20 @@ fn _vchord_halfvec_operator_maxsim( lhs: Array<'_, HalfvecInput<'_>>, rhs: Array<'_, HalfvecInput<'_>>, ) -> f32 { + super::validate_maxsim_array_len(lhs.len()); + super::validate_maxsim_array_len(rhs.len()); + if lhs.iter().any(|x| x.is_none()) || rhs.iter().any(|x| x.is_none()) { + pgrx::error!("MaxSim arrays must not contain NULL vectors"); + } let mut maxsim = 0.0f32; for rhs in rhs.iter().flatten() { let mut d = f32::INFINITY; for lhs in lhs.iter().flatten() { let lhs = lhs.as_borrowed(); let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } d = d.min(VectBorrowed::operator_dot(lhs, rhs).to_f32()); } maxsim += d; diff --git a/src/datatype/operators_rabitq4.rs b/src/datatype/operators_rabitq4.rs index bf0ff920..d1980fe9 100644 --- a/src/datatype/operators_rabitq4.rs +++ b/src/datatype/operators_rabitq4.rs @@ -125,12 +125,20 @@ fn _vchord_rabitq4_operator_maxsim( lhs: Array<'_, Rabitq4Input<'_>>, rhs: Array<'_, Rabitq4Input<'_>>, ) -> f32 { + super::validate_maxsim_array_len(lhs.len()); + super::validate_maxsim_array_len(rhs.len()); + if lhs.iter().any(|x| x.is_none()) || rhs.iter().any(|x| x.is_none()) { + pgrx::error!("MaxSim arrays must not contain NULL vectors"); + } let mut maxsim = 0.0f32; for rhs in rhs.iter().flatten() { let mut d = f32::INFINITY; for lhs in lhs.iter().flatten() { let lhs = lhs.as_borrowed(); let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } d = d.min(Rabitq4Borrowed::operator_dot(lhs, rhs).to_f32()); } maxsim += d; diff --git a/src/datatype/operators_rabitq8.rs b/src/datatype/operators_rabitq8.rs index d8d3cb72..7cac9aac 100644 --- a/src/datatype/operators_rabitq8.rs +++ b/src/datatype/operators_rabitq8.rs @@ -125,12 +125,20 @@ fn _vchord_rabitq8_operator_maxsim( lhs: Array<'_, Rabitq8Input<'_>>, rhs: Array<'_, Rabitq8Input<'_>>, ) -> f32 { + super::validate_maxsim_array_len(lhs.len()); + super::validate_maxsim_array_len(rhs.len()); + if lhs.iter().any(|x| x.is_none()) || rhs.iter().any(|x| x.is_none()) { + pgrx::error!("MaxSim arrays must not contain NULL vectors"); + } let mut maxsim = 0.0f32; for rhs in rhs.iter().flatten() { let mut d = f32::INFINITY; for lhs in lhs.iter().flatten() { let lhs = lhs.as_borrowed(); let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } d = d.min(Rabitq8Borrowed::operator_dot(lhs, rhs).to_f32()); } maxsim += d; diff --git a/src/datatype/operators_vector.rs b/src/datatype/operators_vector.rs index 76b3bf8c..c125a18b 100644 --- a/src/datatype/operators_vector.rs +++ b/src/datatype/operators_vector.rs @@ -95,12 +95,20 @@ fn _vchord_vector_operator_maxsim( lhs: Array<'_, VectorInput<'_>>, rhs: Array<'_, VectorInput<'_>>, ) -> f32 { + super::validate_maxsim_array_len(lhs.len()); + super::validate_maxsim_array_len(rhs.len()); + if lhs.iter().any(|x| x.is_none()) || rhs.iter().any(|x| x.is_none()) { + pgrx::error!("MaxSim arrays must not contain NULL vectors"); + } let mut maxsim = 0.0f32; for rhs in rhs.iter().flatten() { let mut d = f32::INFINITY; for lhs in lhs.iter().flatten() { let lhs = lhs.as_borrowed(); let rhs = rhs.as_borrowed(); + if lhs.dim() != rhs.dim() { + pgrx::error!("dimension is not matched"); + } d = d.min(VectBorrowed::operator_dot(lhs, rhs).to_f32()); } maxsim += d; diff --git a/src/index/fetcher.rs b/src/index/fetcher.rs index 704edfc8..fa836316 100644 --- a/src/index/fetcher.rs +++ b/src/index/fetcher.rs @@ -24,6 +24,21 @@ pub trait FilterableTuple: Tuple { pub trait Tuple { fn build(&mut self) -> (&[Datum; 32], &[bool; 32]); + + /// Read one user attribute from the already fetched heap slot. + /// + /// Callers that may expose the value outside PostgreSQL must call + /// [`FilterableTuple::filter`] first. Attribute numbers are PostgreSQL + /// one-based `attnum` values, not zero-based Rust indexes. + #[allow(dead_code, reason = "reserved for the optional Phase 3C heap source")] + fn attribute(&mut self, attnum: i16) -> Option; +} + +#[derive(Clone, Copy)] +#[allow(dead_code, reason = "reserved for the optional Phase 3C heap source")] +pub struct TupleAttribute { + pub datum: Datum, + pub is_null: bool, } pub trait Fetcher { @@ -52,6 +67,7 @@ pub struct HeapFetcher { heap_relation: pgrx::pg_sys::Relation, snapshot: pgrx::pg_sys::Snapshot, heapfetch: *mut pgrx::pg_sys::IndexFetchTableData, + owns_heapfetch: bool, slot: *mut pgrx::pg_sys::TupleTableSlot, values: [Datum; 32], is_nulls: [bool; 32], @@ -77,6 +93,7 @@ impl HeapFetcher { heap_relation, snapshot, heapfetch, + owns_heapfetch: false, slot: pgrx::pg_sys::table_slot_create(heap_relation, std::ptr::null_mut()), values: [Datum::null(); 32], is_nulls: [true; 32], @@ -84,6 +101,44 @@ impl HeapFetcher { } } } + + /// Create a heap fetch state that is not owned by an `IndexScanDesc`. + /// + /// This is used by the restricted external MaxSim executor to resolve the + /// root TIDs stored in the index through HOT chains before SQL-visible + /// descriptor projection. The table AM owns the rules for doing that; + /// looking up `ctid` directly does not follow a HOT chain. + pub unsafe fn new_standalone( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + snapshot: pgrx::pg_sys::Snapshot, + ) -> Self { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; + + unsafe { + let table_am = (*heap_relation).rd_tableam; + if table_am.is_null() { + panic!("unknown heap access method"); + } + let index_fetch_begin = (*table_am) + .index_fetch_begin + .expect("unsupported heap access method"); + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + let heapfetch = pg_guard_ffi_boundary(|| index_fetch_begin(heap_relation)); + if heapfetch.is_null() { + panic!("heap access method returned a null index fetch state"); + } + let mut fetcher = Self::new( + index_relation, + heap_relation, + snapshot, + heapfetch, + std::ptr::null_mut(), + ); + fetcher.owns_heapfetch = true; + fetcher + } + } } impl Drop for HeapFetcher { @@ -93,6 +148,16 @@ impl Drop for HeapFetcher { // free common resources pgrx::pg_sys::ExecDropSingleTupleTableSlot(self.slot); pgrx::pg_sys::FreeExecutorState(self.estate); + if self.owns_heapfetch { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; + + let table_am = (*self.heap_relation).rd_tableam; + let index_fetch_end = (*table_am) + .index_fetch_end + .expect("unsupported heap access method"); + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + pg_guard_ffi_boundary(|| index_fetch_end(self.heapfetch)); + } } } } @@ -147,7 +212,14 @@ impl Fetcher for HeapFetcher { false }; if found { - Some(HeapTuple { this: self }) + // The heap table AM rewrites the requested root TID to the + // snapshot-visible HOT-chain member. The slot itself keeps + // the rewritten index TID as well, but carrying it explicitly + // avoids depending on slot representation details. + Some(HeapTuple { + this: self, + current_ctid: ctid, + }) } else { None } @@ -157,6 +229,16 @@ impl Fetcher for HeapFetcher { pub struct HeapTuple<'a> { this: &'a mut HeapFetcher, + current_ctid: ItemPointerData, +} + +impl HeapTuple<'_> { + /// Return the physical TID of the tuple version materialized in the slot. + /// This may differ from the root TID supplied to `Fetcher::fetch` after a + /// HOT update. + pub fn ctid(&self) -> ItemPointerData { + self.current_ctid + } } impl Tuple for HeapTuple<'_> { @@ -175,6 +257,30 @@ impl Tuple for HeapTuple<'_> { (&this.values, &this.is_nulls) } } + + fn attribute(&mut self, attnum: i16) -> Option { + unsafe { + use pgrx::pg_sys::ffi::pg_guard_ffi_boundary; + + let slot = self.this.slot; + let tuple_descriptor = (*slot).tts_tupleDescriptor; + if attnum <= 0 + || tuple_descriptor.is_null() + || i32::from(attnum) > (*tuple_descriptor).natts + { + return None; + } + #[allow(ffi_unwind_calls, reason = "protected by pg_guard_ffi_boundary")] + pg_guard_ffi_boundary(|| { + pgrx::pg_sys::slot_getsomeattrs_int(slot, i32::from(attnum)); + }); + let offset = usize::try_from(attnum - 1).ok()?; + Some(TupleAttribute { + datum: *(*slot).tts_values.add(offset), + is_null: *(*slot).tts_isnull.add(offset), + }) + } + } } impl FilterableTuple for HeapTuple<'_> { diff --git a/src/index/gucs.rs b/src/index/gucs.rs index f2724593..26cde468 100644 --- a/src/index/gucs.rs +++ b/src/index/gucs.rs @@ -27,6 +27,18 @@ pub enum PostgresIo { ReadStream, } +#[derive(Debug, Clone, Copy, PostgresGucEnum)] +pub enum PostgresMaxsimBackend { + #[name = c"coarse_only"] + CoarseOnly, + #[name = c"cpu_exact"] + CpuExact, + #[name = c"gpu"] + Gpu, + #[name = c"auto"] + Auto, +} + static VCHORDRQ_QUERY_SAMPLING_ENABLE: GucSetting = GucSetting::::new(false); static VCHORDRQ_QUERY_SAMPLING_MAX_RECORDS: GucSetting = GucSetting::::new(0); @@ -78,6 +90,24 @@ static VCHORDRQ_MAXSIM_THRESHOLD: GucSetting = GucSetting::::new(0); static mut VCHORDRQ_MAXSIM_THRESHOLD_CONFIG: *mut pgrx::pg_sys::config_generic = core::ptr::null_mut(); +static VCHORDRQ_MAXSIM_CANDIDATE_LIMIT: GucSetting = GucSetting::::new(-1); + +static VCHORDRQ_MAXSIM_PLANNER_QUERY_TOKENS: GucSetting = GucSetting::::new(32); + +static VCHORDRQ_MAXSIM_PLANNER_DOCUMENT_TOKENS: GucSetting = GucSetting::::new(256); + +static VCHORDRQ_MAXSIM_BACKEND: GucSetting = + GucSetting::::new(PostgresMaxsimBackend::CoarseOnly); + +static VCHORDRQ_MAXSIM_GPU_ENDPOINT: GucSetting> = + GucSetting::>::new(Some(c"")); + +static VCHORDRQ_MAXSIM_GPU_TIMEOUT_MS: GucSetting = GucSetting::::new(2000); + +static VCHORDRQ_MAXSIM_GPU_MAX_BATCH_TOKENS: GucSetting = GucSetting::::new(1_000_000); + +static VCHORDRQ_MAXSIM_GPU_MAX_BATCH_BYTES: GucSetting = GucSetting::::new(1_073_741_824); + static VCHORDRQ_PREFILTER: GucSetting = GucSetting::::new(false); static VCHORDRQ_IO_SEARCH: GucSetting = GucSetting::::new( @@ -151,6 +181,82 @@ pub fn init() { GucContext::Userset, GucFlags::default(), ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_candidate_limit", + c"Maximum number of page candidates produced by MaxSim aggregation.", + c"Use -1 for no page-candidate limit.", + &VCHORDRQ_MAXSIM_CANDIDATE_LIMIT, + -1, + i32::MAX, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_planner_query_tokens", + c"Expected MaxSim query-token count used by the planner.", + c"Set this to the measured deployment average until expression statistics are available.", + &VCHORDRQ_MAXSIM_PLANNER_QUERY_TOKENS, + 1, + 65_536, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_planner_document_tokens", + c"Fallback MaxSim document-token count used by the planner.", + c"Used for indexes that predate the native indexed-vector statistic; set it to the measured deployment average.", + &VCHORDRQ_MAXSIM_PLANNER_DOCUMENT_TOKENS, + 1, + 65_536, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_enum_guc( + c"vchordrq.maxsim_backend", + c"Page-level MaxSim rerank backend.", + c"GPU and auto modes require the native sidecar integration.", + &VCHORDRQ_MAXSIM_BACKEND, + GucContext::Userset, + GucFlags::default(), + ); + GucRegistry::define_string_guc( + c"vchordrq.maxsim_gpu_endpoint", + c"Unix-socket endpoint for the TileMaxSim sidecar.", + c"An empty endpoint disables GPU transport.", + &VCHORDRQ_MAXSIM_GPU_ENDPOINT, + GucContext::Suset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_gpu_timeout_ms", + c"Overall TileMaxSim sidecar deadline in milliseconds.", + c"The deadline covers connection, request write, and response read.", + &VCHORDRQ_MAXSIM_GPU_TIMEOUT_MS, + 1, + 600_000, + GucContext::Suset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_gpu_max_batch_tokens", + c"Maximum query plus document tokens in one TileMaxSim request.", + c"Requests exceeding the limit fail before connecting to the sidecar.", + &VCHORDRQ_MAXSIM_GPU_MAX_BATCH_TOKENS, + 1, + i32::MAX, + GucContext::Suset, + GucFlags::default(), + ); + GucRegistry::define_int_guc( + c"vchordrq.maxsim_gpu_max_batch_bytes", + c"Maximum encoded TileMaxSim request size in bytes.", + c"Requests exceeding the limit fail before connecting to the sidecar.", + &VCHORDRQ_MAXSIM_GPU_MAX_BATCH_BYTES, + 1024, + i32::MAX, + GucContext::Suset, + GucFlags::default(), + ); GucRegistry::define_bool_guc( c"vchordrq.prefilter", c"`prefilter` argument of vchordrq.", @@ -475,6 +581,39 @@ pub fn vchordrq_maxsim_threshold(index: pgrx::pg_sys::Relation) -> u32 { } } +pub fn vchordrq_maxsim_candidate_limit() -> Option { + let value = VCHORDRQ_MAXSIM_CANDIDATE_LIMIT.get(); + if value < 0 { None } else { Some(value as u32) } +} + +pub fn vchordrq_maxsim_planner_query_tokens() -> u32 { + VCHORDRQ_MAXSIM_PLANNER_QUERY_TOKENS.get() as u32 +} + +pub fn vchordrq_maxsim_planner_document_tokens() -> u32 { + VCHORDRQ_MAXSIM_PLANNER_DOCUMENT_TOKENS.get() as u32 +} + +pub fn vchordrq_maxsim_backend() -> PostgresMaxsimBackend { + VCHORDRQ_MAXSIM_BACKEND.get() +} + +pub fn vchordrq_maxsim_gpu_endpoint() -> Option { + VCHORDRQ_MAXSIM_GPU_ENDPOINT.get() +} + +pub fn vchordrq_maxsim_gpu_timeout_ms() -> u32 { + VCHORDRQ_MAXSIM_GPU_TIMEOUT_MS.get() as u32 +} + +pub fn vchordrq_maxsim_gpu_max_batch_tokens() -> u32 { + VCHORDRQ_MAXSIM_GPU_MAX_BATCH_TOKENS.get() as u32 +} + +pub fn vchordrq_maxsim_gpu_max_batch_bytes() -> u32 { + VCHORDRQ_MAXSIM_GPU_MAX_BATCH_BYTES.get() as u32 +} + pub fn vchordrq_prefilter() -> bool { VCHORDRQ_PREFILTER.get() } diff --git a/src/index/vchordrq/am/am_build.rs b/src/index/vchordrq/am/am_build.rs index 647d2478..95b4bc45 100644 --- a/src/index/vchordrq/am/am_build.rs +++ b/src/index/vchordrq/am/am_build.rs @@ -349,9 +349,10 @@ pub unsafe extern "C-unwind" fn ambuild( reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); order }, - |indtuples| { + |indtuples, indexed_vectors| { reporter.tuples_done(indtuples); reporter.tuples_total(indtuples); + vchordrq::set_indexed_vectors(&index, indexed_vectors); // enter the barrier let shared = leader.vchordrqshared; pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); @@ -429,9 +430,10 @@ pub unsafe extern "C-unwind" fn ambuild( || { reporter.phase(BuildPhase::from_code(BuildPhaseCode::Inserting)); }, - |indtuples| { + |indtuples, indexed_vectors| { reporter.tuples_done(indtuples); reporter.tuples_total(indtuples); + vchordrq::set_indexed_vectors(&index, indexed_vectors); reporter.phase(BuildPhase::from_code(BuildPhaseCode::Compacting)); }, || {}, @@ -460,6 +462,7 @@ struct VchordrqShared { barrier_enter_0: i32, nparticipants: u32, indtuples: u64, + indexed_vectors: u64, barrier_leave_0: bool, barrier_enter_1: i32, barrier_leave_1: bool, @@ -721,6 +724,7 @@ impl VchordrqLeader { barrier_leave_2: false, mutex: std::mem::zeroed(), indtuples: 0, + indexed_vectors: 0, }); pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).condvar_barrier_enter_0); pgrx::pg_sys::ConditionVariableInit(&raw mut (*vchordrqshared).condvar_barrier_leave_0); @@ -875,7 +879,7 @@ pub unsafe extern "C-unwind" fn vchordrq_parallel_build_main( pgrx::pg_sys::ConditionVariableCancelSleep(); order }, - |_| { + |_, _| { // enter the barrier let shared = vchordrqshared; pgrx::pg_sys::SpinLockAcquire(&raw mut (*shared).mutex); @@ -941,7 +945,7 @@ unsafe fn parallel_build( vchordrqcached: *const u8, mut callback: impl FnMut(u64), sync_0: impl FnOnce() -> u32, - sync_1: impl FnOnce(u64), + sync_1: impl FnOnce(u64, u64), sync_2: impl FnOnce(), ) { use vchordrq_cached::VchordrqCachedReader; @@ -989,6 +993,7 @@ unsafe fn parallel_build( let store = value .and_then(|x| unsafe { opfamily.store(x) }) .unwrap_or_default(); + let indexed_vectors = store.len() as u64; for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); @@ -1010,6 +1015,7 @@ unsafe fn parallel_build( { pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); (*vchordrqshared).indtuples += 1; + (*vchordrqshared).indexed_vectors += indexed_vectors; indtuples = (*vchordrqshared).indtuples; pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); } @@ -1029,6 +1035,7 @@ unsafe fn parallel_build( let store = value .and_then(|x| unsafe { opfamily.store(x) }) .unwrap_or_default(); + let indexed_vectors = store.len() as u64; for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); @@ -1050,6 +1057,7 @@ unsafe fn parallel_build( { pgrx::pg_sys::SpinLockAcquire(&raw mut (*vchordrqshared).mutex); (*vchordrqshared).indtuples += 1; + (*vchordrqshared).indexed_vectors += indexed_vectors; indtuples = (*vchordrqshared).indtuples; pgrx::pg_sys::SpinLockRelease(&raw mut (*vchordrqshared).mutex); } @@ -1065,7 +1073,13 @@ unsafe fn parallel_build( drop(index); - sync_1(unsafe { (*vchordrqshared).indtuples }); + let (indtuples, indexed_vectors) = unsafe { + ( + (*vchordrqshared).indtuples, + (*vchordrqshared).indexed_vectors, + ) + }; + sync_1(indtuples, indexed_vectors); let index = unsafe { PostgresRelation::new(index_relation) }; @@ -1085,7 +1099,7 @@ unsafe fn sequential_build( vchordrqcached: &[u8], mut callback: impl FnMut(u64), sync_0: impl FnOnce(), - sync_1: impl FnOnce(u64), + sync_1: impl FnOnce(u64, u64), sync_2: impl FnOnce(), ) { use vchordrq_cached::VchordrqCachedReader; @@ -1125,6 +1139,7 @@ unsafe fn sequential_build( let index = unsafe { BufferedPostgresRelation::new(index_relation) }; let mut indtuples = 0; + let mut indexed_vectors = 0_u64; match cached { VchordrqCachedReader::_0(_) => { traverser.traverse(true, |tuple: &mut dyn crate::index::traverse::Tuple| { @@ -1134,6 +1149,7 @@ unsafe fn sequential_build( let store = value .and_then(|x| unsafe { opfamily.store(x) }) .unwrap_or_default(); + indexed_vectors += store.len() as u64; for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); @@ -1166,6 +1182,7 @@ unsafe fn sequential_build( let store = value .and_then(|x| unsafe { opfamily.store(x) }) .unwrap_or_default(); + indexed_vectors += store.len() as u64; for (vector, extra) in store { let key = ctid_to_key(ctid); let payload = kv_to_pointer((key, extra)); @@ -1194,7 +1211,7 @@ unsafe fn sequential_build( drop(index); - sync_1(indtuples); + sync_1(indtuples, indexed_vectors); let index = unsafe { PostgresRelation::new(index_relation) }; diff --git a/src/index/vchordrq/am/mod.rs b/src/index/vchordrq/am/mod.rs index f995c600..976692b6 100644 --- a/src/index/vchordrq/am/mod.rs +++ b/src/index/vchordrq/am/mod.rs @@ -322,28 +322,13 @@ pub unsafe extern "C-unwind" fn amcostestimate( if !(*index_opt_info).hypothetical { let relation = Index::open((*index_opt_info).indexoid, pgrx::pg_sys::NoLock as _); let opfamily = opfamily(relation.raw()); - if !matches!( + let is_maxsim = matches!( opfamily, - Opfamily::HalfvecCosine - | Opfamily::HalfvecIp - | Opfamily::HalfvecL2 - | Opfamily::VectorCosine - | Opfamily::VectorIp - | Opfamily::VectorL2 - | Opfamily::Rabitq8Cosine - | Opfamily::Rabitq8Ip - | Opfamily::Rabitq8L2 - | Opfamily::Rabitq4Cosine - | Opfamily::Rabitq4Ip - | Opfamily::Rabitq4L2 - ) { - *index_startup_cost = 0.0; - *index_total_cost = 0.0; - *index_selectivity = 1.0; - *index_correlation = 0.0; - *index_pages = 1.0; - return; - } + Opfamily::VectorMaxsim + | Opfamily::HalfvecMaxsim + | Opfamily::Rabitq8Maxsim + | Opfamily::Rabitq4Maxsim + ); let index = PostgresRelation::::new(relation.raw()); let probes = gucs::vchordrq_probes(relation.raw()); let cost = vchordrq::cost(&index); @@ -354,15 +339,25 @@ pub unsafe extern "C-unwind" fn amcostestimate( probes.len() ); } + let estimated_index_vectors = if is_maxsim { + cost.indexed_vectors.map_or_else( + || { + (*index_opt_info).tuples.max(total_rows).max(1.0) + * f64::from(gucs::vchordrq_maxsim_planner_document_tokens()) + }, + |indexed_vectors| indexed_vectors as f64, + ) + } else { + (*index_opt_info).tuples.max(0.0) + }; let node_count = { - let tuples = (*index_opt_info).tuples as u32; let mut count = 0.0; - let r = cost.cells.iter().copied().rev(); - let numerator = std::iter::once(1).chain(probes.clone()); + let r = cost.cells.iter().copied().rev().map(f64::from); + let numerator = std::iter::once(1.0).chain(probes.iter().copied().map(f64::from)); let denumerator = r.clone(); - let scale = r.skip(1).chain(std::iter::once(tuples)); + let scale = r.skip(1).chain(std::iter::once(estimated_index_vectors)); for (scale, (numerator, denumerator)) in scale.zip(numerator.zip(denumerator)) { - count += (scale as f64) * 1.0f64.min((numerator as f64) / (denumerator as f64)); + count += scale * 1.0f64.min(numerator / denumerator); } count }; @@ -378,13 +373,41 @@ pub unsafe extern "C-unwind" fn amcostestimate( pages += cost.cells[0] as f64; pages }; + if is_maxsim { + let backend = match gucs::vchordrq_maxsim_backend() { + gucs::PostgresMaxsimBackend::CoarseOnly => { + vchordrq::MaxsimCostBackend::CoarseOnly + } + gucs::PostgresMaxsimBackend::CpuExact => vchordrq::MaxsimCostBackend::CpuExact, + gucs::PostgresMaxsimBackend::Gpu => vchordrq::MaxsimCostBackend::Gpu, + gucs::PostgresMaxsimBackend::Auto => vchordrq::MaxsimCostBackend::Auto, + }; + let estimate = vchordrq::estimate_maxsim_cost(vchordrq::MaxsimCostInput { + heap_rows: total_rows, + index_tokens: estimated_index_vectors, + token_nodes_per_query: node_count, + base_index_pages: page_count, + dimension: cost.dim, + element_bits: opfamily.vector_kind().number_of_bits_of_an_elements(), + query_tokens: gucs::vchordrq_maxsim_planner_query_tokens(), + limit_tuples: ((*root).limit_tuples > 0.0).then_some((*root).limit_tuples), + filter_selectivity, + candidate_limit: gucs::vchordrq_maxsim_candidate_limit(), + backend, + }); + *index_startup_cost = estimate.startup_cost; + *index_total_cost = estimate.total_cost; + *index_selectivity = estimate.selectivity; + *index_correlation = 0.0; + *index_pages = estimate.index_pages; + return; + } // `next_count` represents candidates we expect to process to // surface `limit_tuples` survivors after filter rejection. Clamp // by `node_count` so the estimate cannot exceed the candidates // the IVF visits at the configured probe count. let next_count = if (*root).limit_tuples > 0.0 { - ((*root).limit_tuples * f64::min(1000.0, 1.0 / filter_selectivity)) - .min(node_count) + ((*root).limit_tuples * f64::min(1000.0, 1.0 / filter_selectivity)).min(node_count) } else { node_count }; @@ -482,7 +505,9 @@ pub unsafe extern "C-unwind" fn ambulkdelete( pg_guard_ffi_boundary(|| callback(&mut ctid, callback_state)) } }; - crate::index::vchordrq::dispatch::bulkdelete(opfamily, &index, check, callback); + let indexed_vectors = + crate::index::vchordrq::dispatch::bulkdelete(opfamily, &index, check, callback); + vchordrq::set_indexed_vectors(&index, indexed_vectors); stats } @@ -538,6 +563,12 @@ pub unsafe extern "C-unwind" fn amrescan( max_scan_tuples: gucs::vchordrq_max_scan_tuples(), maxsim_refine: gucs::vchordrq_maxsim_refine((*scan).indexRelation), maxsim_threshold: gucs::vchordrq_maxsim_threshold((*scan).indexRelation), + maxsim_candidate_limit: gucs::vchordrq_maxsim_candidate_limit(), + maxsim_backend: gucs::vchordrq_maxsim_backend(), + maxsim_gpu_endpoint: gucs::vchordrq_maxsim_gpu_endpoint(), + maxsim_gpu_timeout_ms: gucs::vchordrq_maxsim_gpu_timeout_ms(), + maxsim_gpu_max_batch_tokens: gucs::vchordrq_maxsim_gpu_max_batch_tokens(), + maxsim_gpu_max_batch_bytes: gucs::vchordrq_maxsim_gpu_max_batch_bytes(), io_search: gucs::vchordrq_io_search(), io_rerank: gucs::vchordrq_io_rerank(), prefilter: gucs::vchordrq_prefilter(), diff --git a/src/index/vchordrq/dispatch.rs b/src/index/vchordrq/dispatch.rs index 01d70a87..36baee63 100644 --- a/src/index/vchordrq/dispatch.rs +++ b/src/index/vchordrq/dispatch.rs @@ -71,42 +71,51 @@ pub fn bulkdelete( index: &R, check: impl Fn(), callback: impl Fn(NonZero) -> bool, -) where +) -> u64 +where R: RelationRead + RelationWrite, R::Page: Page, { match (opfamily.vector_kind(), opfamily.distance_kind()) { (VectorKind::Vecf32, DistanceKind::L2S) => { - vchordrq::bulkdelete::<_, Op, L2S>>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op, L2S>>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op, L2S>>(index, &check, &callback); + live } (VectorKind::Vecf32, DistanceKind::Dot) => { - vchordrq::bulkdelete::<_, Op, Dot>>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op, Dot>>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op, Dot>>(index, &check, &callback); + live } (VectorKind::Vecf16, DistanceKind::L2S) => { - vchordrq::bulkdelete::<_, Op, L2S>>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op, L2S>>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op, L2S>>(index, &check, &callback); + live } (VectorKind::Vecf16, DistanceKind::Dot) => { - vchordrq::bulkdelete::<_, Op, Dot>>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op, Dot>>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op, Dot>>(index, &check, &callback); + live } (VectorKind::Rabitq8, DistanceKind::L2S) => { - vchordrq::bulkdelete::<_, Op>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); + live } (VectorKind::Rabitq8, DistanceKind::Dot) => { - vchordrq::bulkdelete::<_, Op>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); + live } (VectorKind::Rabitq4, DistanceKind::L2S) => { - vchordrq::bulkdelete::<_, Op>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); + live } (VectorKind::Rabitq4, DistanceKind::Dot) => { - vchordrq::bulkdelete::<_, Op>(index, &check, &callback); + let live = vchordrq::bulkdelete::<_, Op>(index, &check, &callback); vchordrq::bulkdelete_vectors::<_, Op>(index, &check, &callback); + live } } } diff --git a/src/index/vchordrq/opclass.rs b/src/index/vchordrq/opclass.rs index 7c51f5db..075b1075 100644 --- a/src/index/vchordrq/opclass.rs +++ b/src/index/vchordrq/opclass.rs @@ -92,6 +92,7 @@ impl Opfamily { Self::VectorMaxsim => { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for (i, vector) in vectors.iter_deny_null().enumerate() { result.push(( @@ -105,6 +106,7 @@ impl Opfamily { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for (i, vector) in vectors.iter_deny_null().enumerate() { result.push(( @@ -118,6 +120,7 @@ impl Opfamily { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for (i, vector) in vectors.iter_deny_null().enumerate() { result.push(( @@ -131,6 +134,7 @@ impl Opfamily { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for (i, vector) in vectors.iter_deny_null().enumerate() { result.push(( @@ -203,6 +207,7 @@ impl Opfamily { Self::VectorL2 | Self::VectorIp | Self::VectorCosine | Self::VectorMaxsim => { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for vector in vectors.iter_deny_null() { result.push(self.input(BorrowedVector::Vecf32(vector.as_borrowed()))); @@ -213,6 +218,7 @@ impl Opfamily { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for vector in vectors.iter_deny_null() { result.push(self.input(BorrowedVector::Vecf16(vector.as_borrowed()))); @@ -223,6 +229,7 @@ impl Opfamily { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for vector in vectors.iter_deny_null() { result.push(self.input(BorrowedVector::Rabitq8(vector.as_borrowed()))); @@ -233,6 +240,7 @@ impl Opfamily { let vectors = unsafe { pgrx::datum::Array::::from_datum(datum, false).unwrap() }; + crate::datatype::validate_maxsim_array_len(vectors.len()); let mut result = Vec::with_capacity(vectors.len()); for vector in vectors.iter_deny_null() { result.push(self.input(BorrowedVector::Rabitq4(vector.as_borrowed()))); diff --git a/src/index/vchordrq/scanners/maxsim.rs b/src/index/vchordrq/scanners/maxsim.rs index 1c51c97f..238dbe2e 100644 --- a/src/index/vchordrq/scanners/maxsim.rs +++ b/src/index/vchordrq/scanners/maxsim.rs @@ -12,7 +12,17 @@ // // Copyright (c) 2025-2026 TensorChord Inc. +mod candidate; +mod external; +mod gpu; +mod rerank; +mod search; + +use self::candidate::{LegacyPageCandidateGenerator, PageCandidateGenerator}; +use self::gpu::{GpuTileMaxsimBackend, UnixSocketTransport, report_gpu_fallback}; +use self::rerank::{CpuExactMaxsimBackend, ExactMaxsimBackend, HeapArrayTensorSource}; use crate::index::fetcher::*; +use crate::index::gucs::PostgresMaxsimBackend; use crate::index::scanners::{Io, SearchBuilder}; use crate::index::vchordrq::dispatch::*; use crate::index::vchordrq::filter::filter; @@ -21,7 +31,6 @@ use crate::index::vchordrq::scanners::SearchOptions; use crate::recorder::Recorder; use always_equal::AlwaysEqual; use dary_heap::QuaternaryHeap as Heap; -use distance::Distance; use index::bump::Bump; use index::packed::PackedRefMut8; use index::prefetcher::*; @@ -29,8 +38,8 @@ use index::relation::{Hints, Page, RelationPrefetch, RelationRead, RelationReadS use index_accessor::Dot; use simd::f16; use std::cmp::Reverse; -use std::collections::BinaryHeap; use std::num::NonZero; +use std::time::Duration; use vchordrq::types::{DistanceKind, OwnedVector, VectorKind}; use vchordrq::{RerankMethod, how, maxsim_search, rerank_index}; use vector::VectorOwned; @@ -99,10 +108,33 @@ impl SearchBuilder for MaxsimBuilder { } let maxsim_refine = options.maxsim_refine; let maxsim_threshold = options.maxsim_threshold; + let maxsim_candidate_limit = options + .maxsim_candidate_limit + .map_or(usize::MAX, |value| value as usize); + let has_maxsim_candidate_limit = options.maxsim_candidate_limit.is_some(); + let maxsim_backend = options.maxsim_backend; + let maxsim_gpu_endpoint = options + .maxsim_gpu_endpoint + .as_deref() + .map(|endpoint| endpoint.to_string_lossy().into_owned()) + .unwrap_or_default(); + let maxsim_gpu_timeout = Duration::from_millis(options.maxsim_gpu_timeout_ms as u64); + let maxsim_gpu_max_batch_tokens = options.maxsim_gpu_max_batch_tokens as usize; + let maxsim_gpu_max_batch_bytes = options.maxsim_gpu_max_batch_bytes as usize; + if !matches!(maxsim_backend, PostgresMaxsimBackend::CoarseOnly) + && (!has_maxsim_candidate_limit || maxsim_candidate_limit == 0) + { + pgrx::error!("exact MaxSim requires a positive vchordrq.maxsim_candidate_limit"); + } let opfamily = self.opfamily; let Some(vectors) = vectors else { return Box::new(std::iter::empty()) as Box>; }; + let expected_dim = vchordrq::cost(index).dim; + if vectors.iter().any(|vector| vector.dim() != expected_dim) { + pgrx::error!("dimension is not matched"); + } + let rerank_query = vectors.clone(); let method = how(index); if !matches!(method, RerankMethod::Index) { pgrx::error!("maxsim search with rerank_in_table is not supported"); @@ -124,669 +156,702 @@ impl SearchBuilder for MaxsimBuilder { _, AlwaysEqual, _, _)>>, )| (rough, payload); - let iter: Box> = match opfamily.vector_kind() { - VectorKind::Vecf32 => { - type Op = vchordrq::operator::Op, Dot>; - let unprojected = vectors - .into_iter() - .map(|vector| { - if let OwnedVector::Vecf32(vector) = vector { - vector - } else { - unreachable!() - } - }) - .collect::>(); - let projected = unprojected - .iter() - .map(|vector| RandomProject::project(vector.as_borrowed())) - .collect::>(); - Box::new((0..n).map(move |i| { - let (results, estimation_by_threshold) = match options.io_search { - Io::Plain => maxsim_search::<_, Op>( - index, - projected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_plain_prefetcher.clone(), - ), - Io::Simple => maxsim_search::<_, Op>( - index, - projected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_simple_prefetcher.clone(), - ), - Io::Stream => maxsim_search::<_, Op>( - index, - projected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_stream_prefetcher.clone(), - ), - }; - let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); - if maxsim_refine != 0 && !results.is_empty() { - let sequence = Heap::from(results); - match (options.io_rerank, options.prefilter) { - (Io::Plain, false) => { - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Plain, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, false) => { - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Stream, false) => { - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + let mut token_searches = { + let fetcher = &mut fetcher; + let iter: Box> = match opfamily.vector_kind() { + VectorKind::Vecf32 => { + type Op = vchordrq::operator::Op, Dot>; + let unprojected = vectors + .into_iter() + .map(|vector| { + if let OwnedVector::Vecf32(vector) = vector { + vector + } else { + unreachable!() } - (Io::Stream, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + }) + .collect::>(); + let projected = unprojected + .iter() + .map(|vector| RandomProject::project(vector.as_borrowed())) + .collect::>(); + Box::new((0..n).map(move |i| { + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let sequence = Heap::from(results); + match (options.io_rerank, options.prefilter) { + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } } - } - } else { - let rough_iter = results.into_iter(); - rough_set.extend(rough_iter.map(rough_map)); - } - (accu_set, rough_set, estimation_by_threshold) - })) - } - VectorKind::Vecf16 => { - type Op = vchordrq::operator::Op, Dot>; - let unprojected = vectors - .into_iter() - .map(|vector| { - if let OwnedVector::Vecf16(vector) = vector { - vector } else { - unreachable!() + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); } - }) - .collect::>(); - let projected = unprojected - .iter() - .map(|vector| RandomProject::project(vector.as_borrowed())) - .collect::>(); - Box::new((0..n).map(move |i| { - let (results, estimation_by_threshold) = match options.io_search { - Io::Plain => maxsim_search::<_, Op>( - index, - projected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_plain_prefetcher.clone(), - ), - Io::Simple => maxsim_search::<_, Op>( - index, - projected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_simple_prefetcher.clone(), - ), - Io::Stream => maxsim_search::<_, Op>( - index, - projected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_stream_prefetcher.clone(), - ), - }; - let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); - if maxsim_refine != 0 && !results.is_empty() { - let sequence = Heap::from(results); - match (options.io_rerank, options.prefilter) { - (Io::Plain, false) => { - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Plain, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, false) => { - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Stream, false) => { - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + (accu_set, rough_set, estimation_by_threshold) + })) + } + VectorKind::Vecf16 => { + type Op = vchordrq::operator::Op, Dot>; + let unprojected = vectors + .into_iter() + .map(|vector| { + if let OwnedVector::Vecf16(vector) = vector { + vector + } else { + unreachable!() } - (Io::Stream, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + }) + .collect::>(); + let projected = unprojected + .iter() + .map(|vector| RandomProject::project(vector.as_borrowed())) + .collect::>(); + Box::new((0..n).map(move |i| { + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + projected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let sequence = Heap::from(results); + match (options.io_rerank, options.prefilter) { + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } } - } - } else { - let rough_iter = results.into_iter(); - rough_set.extend(rough_iter.map(rough_map)); - } - (accu_set, rough_set, estimation_by_threshold) - })) - } - VectorKind::Rabitq8 => { - type Op = vchordrq::operator::Op; - let unprojected = vectors - .into_iter() - .map(|vector| { - if let OwnedVector::Rabitq8(vector) = vector { - vector } else { - unreachable!() + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); } - }) - .collect::>(); - Box::new((0..n).map(move |i| { - let (results, estimation_by_threshold) = match options.io_search { - Io::Plain => maxsim_search::<_, Op>( - index, - unprojected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_plain_prefetcher.clone(), - ), - Io::Simple => maxsim_search::<_, Op>( - index, - unprojected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_simple_prefetcher.clone(), - ), - Io::Stream => maxsim_search::<_, Op>( - index, - unprojected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_stream_prefetcher.clone(), - ), - }; - let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); - if maxsim_refine != 0 && !results.is_empty() { - let sequence = Heap::from(results); - match (options.io_rerank, options.prefilter) { - (Io::Plain, false) => { - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Plain, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, false) => { - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Stream, false) => { - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + (accu_set, rough_set, estimation_by_threshold) + })) + } + VectorKind::Rabitq8 => { + type Op = vchordrq::operator::Op; + let unprojected = vectors + .into_iter() + .map(|vector| { + if let OwnedVector::Rabitq8(vector) = vector { + vector + } else { + unreachable!() } - (Io::Stream, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + }) + .collect::>(); + Box::new((0..n).map(move |i| { + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let sequence = Heap::from(results); + match (options.io_rerank, options.prefilter) { + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } } - } - } else { - let rough_iter = results.into_iter(); - rough_set.extend(rough_iter.map(rough_map)); - } - (accu_set, rough_set, estimation_by_threshold) - })) - } - VectorKind::Rabitq4 => { - type Op = vchordrq::operator::Op; - let unprojected = vectors - .into_iter() - .map(|vector| { - if let OwnedVector::Rabitq4(vector) = vector { - vector } else { - unreachable!() + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); } - }) - .collect::>(); - Box::new((0..n).map(move |i| { - let (results, estimation_by_threshold) = match options.io_search { - Io::Plain => maxsim_search::<_, Op>( - index, - unprojected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_plain_prefetcher.clone(), - ), - Io::Simple => maxsim_search::<_, Op>( - index, - unprojected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_simple_prefetcher.clone(), - ), - Io::Stream => maxsim_search::<_, Op>( - index, - unprojected[i].as_borrowed(), - options.probes.clone(), - options.epsilon, - maxsim_threshold, - bump, - make_h1_plain_prefetcher.clone(), - make_h0_stream_prefetcher.clone(), - ), - }; - let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); - if maxsim_refine != 0 && !results.is_empty() { - let sequence = Heap::from(results); - match (options.io_rerank, options.prefilter) { - (Io::Plain, false) => { - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Plain, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = PlainPrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, false) => { - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Simple, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = SimplePrefetcher::new(index, sequence); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); - } - (Io::Stream, false) => { - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + (accu_set, rough_set, estimation_by_threshold) + })) + } + VectorKind::Rabitq4 => { + type Op = vchordrq::operator::Op; + let unprojected = vectors + .into_iter() + .map(|vector| { + if let OwnedVector::Rabitq4(vector) = vector { + vector + } else { + unreachable!() } - (Io::Stream, true) => { - let predicate = - id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { - let (key, _) = pointer_to_kv(*pointer); - let Some(mut tuple) = fetcher.fetch(key) else { - return false; - }; - tuple.filter() - }); - let sequence = filter(sequence, predicate); - let prefetcher = - StreamPrefetcher::new(index, sequence, rerank_hints); - let mut reranker = - rerank_index::(unprojected[i].clone(), prefetcher); - accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); - let (rough_iter, accu_iter) = reranker.finish(); - accu_set.extend(accu_iter.map(accu_map)); - rough_set.extend(rough_iter.into_iter().map(rough_map)); + }) + .collect::>(); + Box::new((0..n).map(move |i| { + let (results, estimation_by_threshold) = match options.io_search { + Io::Plain => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_plain_prefetcher.clone(), + ), + Io::Simple => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_simple_prefetcher.clone(), + ), + Io::Stream => maxsim_search::<_, Op>( + index, + unprojected[i].as_borrowed(), + options.probes.clone(), + options.epsilon, + maxsim_threshold, + bump, + make_h1_plain_prefetcher.clone(), + make_h0_stream_prefetcher.clone(), + ), + }; + let (mut accu_set, mut rough_set) = (Vec::new(), Vec::new()); + if maxsim_refine != 0 && !results.is_empty() { + let sequence = Heap::from(results); + match (options.io_rerank, options.prefilter) { + (Io::Plain, false) => { + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Plain, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = PlainPrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, false) => { + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Simple, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = SimplePrefetcher::new(index, sequence); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, false) => { + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } + (Io::Stream, true) => { + let predicate = + id_0(|(_, AlwaysEqual(PackedRefMut8((pointer, _, _))))| { + let (key, _) = pointer_to_kv(*pointer); + let Some(mut tuple) = fetcher.fetch(key) else { + return false; + }; + tuple.filter() + }); + let sequence = filter(sequence, predicate); + let prefetcher = + StreamPrefetcher::new(index, sequence, rerank_hints); + let mut reranker = rerank_index::( + unprojected[i].clone(), + prefetcher, + ); + accu_set.extend(reranker.by_ref().take(maxsim_refine as _)); + let (rough_iter, accu_iter) = reranker.finish(); + accu_set.extend(accu_iter.map(accu_map)); + rough_set.extend(rough_iter.into_iter().map(rough_map)); + } } + } else { + let rough_iter = results.into_iter(); + rough_set.extend(rough_iter.map(rough_map)); } - } else { - let rough_iter = results.into_iter(); - rough_set.extend(rough_iter.map(rough_map)); - } - (accu_set, rough_set, estimation_by_threshold) + (accu_set, rough_set, estimation_by_threshold) + })) + } + }; + iter + }; + let mut candidates = + LegacyPageCandidateGenerator.generate(n, &mut token_searches, maxsim_candidate_limit); + drop(token_searches); + let iter: Box> = match maxsim_backend { + PostgresMaxsimBackend::CoarseOnly => Box::new(candidates.map(|candidate| { + let distance = candidate.approximate_distance.to_f32(); + let recheck = false; + (distance, candidate.heap_key, recheck) + })), + PostgresMaxsimBackend::CpuExact => { + let mut source = HeapArrayTensorSource::new(&mut fetcher, opfamily); + let exact = + CpuExactMaxsimBackend.rerank(&rerank_query, &mut candidates, &mut source); + let exact = exact.unwrap_or_else(|error| pgrx::error!("{error}")); + Box::new(exact.map(|result| { + let distance = result.distance.to_f32(); + let recheck = false; + (distance, result.heap_key, recheck) })) } - }; - let mut updates = Vec::new(); - let mut estimations = Vec::new(); - for (query_id, (accu_set, rough_set, estimation_by_threshold)) in iter.enumerate() { - updates.reserve(accu_set.len() + rough_set.len()); - let is_empty = accu_set.is_empty() && rough_set.is_empty(); - let mut estimation_by_scope = Distance::NEG_INFINITY; - for (distance, payload) in accu_set { - estimation_by_scope = std::cmp::max(estimation_by_scope, distance); - let (key, _) = pointer_to_kv(payload); - updates.push((key, query_id, distance)); + PostgresMaxsimBackend::Gpu => { + let candidates = candidates.collect::>(); + let mut candidate_iter = candidates.iter().copied(); + let mut source = HeapArrayTensorSource::new(&mut fetcher, opfamily); + let transport = UnixSocketTransport::new(maxsim_gpu_endpoint); + let mut backend = GpuTileMaxsimBackend::new( + transport, + maxsim_gpu_timeout, + maxsim_gpu_max_batch_tokens, + maxsim_gpu_max_batch_bytes, + ); + let exact = backend.rerank(&rerank_query, &mut candidate_iter, &mut source); + let exact = exact.unwrap_or_else(|error| pgrx::error!("{error}")); + Box::new(exact.map(|result| { + let distance = result.distance.to_f32(); + let recheck = false; + (distance, result.heap_key, recheck) + })) } - for (distance, payload) in rough_set { - let (key, _) = pointer_to_kv(payload); - updates.push((key, query_id, distance)); + PostgresMaxsimBackend::Auto => { + let candidates = candidates.collect::>(); + let mut candidate_iter = candidates.iter().copied(); + let mut source = HeapArrayTensorSource::new(&mut fetcher, opfamily); + let transport = UnixSocketTransport::new(maxsim_gpu_endpoint); + let mut backend = GpuTileMaxsimBackend::new( + transport, + maxsim_gpu_timeout, + maxsim_gpu_max_batch_tokens, + maxsim_gpu_max_batch_bytes, + ); + let exact = backend.rerank(&rerank_query, &mut candidate_iter, &mut source); + let exact = match exact { + Ok(exact) => exact, + Err(error) => { + report_gpu_fallback(&error); + let mut candidate_iter = candidates.iter().copied(); + let mut source = HeapArrayTensorSource::new(&mut fetcher, opfamily); + CpuExactMaxsimBackend + .rerank(&rerank_query, &mut candidate_iter, &mut source) + .unwrap_or_else(|error| pgrx::error!("{error}")) + } + }; + Box::new(exact.map(|result| { + let distance = result.distance.to_f32(); + let recheck = false; + (distance, result.heap_key, recheck) + })) } - estimations.push(if !is_empty { - std::cmp::max(estimation_by_scope, estimation_by_threshold) - } else { - Distance::ZERO - }); - } - updates.sort_unstable_by_key(|&(key, ..)| key); - let iter = updates - .chunk_by(|(kl, ..), (kr, ..)| kl == kr) - .map(|chunk| { - let key = chunk[0].0; - let mut value = vec![None; n]; - for &(_, query_id, distance) in chunk { - let this = value[query_id].get_or_insert(Distance::INFINITY); - *this = std::cmp::min(*this, distance); - } - let mut maxsim = 0.0f32; - for (query_id, distance) in value.into_iter().enumerate() { - let d = distance.unwrap_or(estimations[query_id]); - maxsim += Distance::to_f32(d); - } - (Reverse(Distance::from_f32(maxsim)), AlwaysEqual(key)) - }) - .collect::>() - .into_iter_sorted_polyfill() - .map(|(Reverse(distance), AlwaysEqual(key))| { - let distance = distance.to_f32(); - let recheck = false; - (distance, key, recheck) - }); - let iter: Box> = Box::new(iter); - let iter = if let Some(max_scan_tuples) = options.max_scan_tuples { - Box::new(iter.take(max_scan_tuples as _)) - } else { - iter }; #[allow(clippy::let_and_return)] iter } } -// Emulate unstable library feature `binary_heap_into_iter_sorted`. -// See https://github.com/rust-lang/rust/issues/59278. - -trait IntoIterSortedPolyfill { - fn into_iter_sorted_polyfill(self) -> IntoIterSorted; -} - -impl IntoIterSortedPolyfill for BinaryHeap { - fn into_iter_sorted_polyfill(self) -> IntoIterSorted { - IntoIterSorted { inner: self } - } -} - -#[derive(Clone, Debug)] -struct IntoIterSorted { - inner: BinaryHeap, -} - -impl Iterator for IntoIterSorted { - type Item = T; - - #[inline] - fn next(&mut self) -> Option { - self.inner.pop() - } - - #[inline] - fn size_hint(&self) -> (usize, Option) { - let exact = self.inner.len(); - (exact, Some(exact)) - } -} - -impl ExactSizeIterator for IntoIterSorted {} - -impl std::iter::FusedIterator for IntoIterSorted {} - #[inline(always)] pub fn id_0(f: F) -> F where diff --git a/src/index/vchordrq/scanners/maxsim/candidate.rs b/src/index/vchordrq/scanners/maxsim/candidate.rs new file mode 100644 index 00000000..9aa3e22f --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/candidate.rs @@ -0,0 +1,192 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use crate::index::fetcher::pointer_to_kv; +use always_equal::AlwaysEqual; +use distance::Distance; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use std::num::NonZero; + +pub(super) type HeapKey = [u16; 3]; +pub(super) type TokenCandidate = (Distance, NonZero); +pub(super) type TokenSearchResult = (Vec, Vec, Distance); + +#[derive(Clone, Copy, Debug)] +pub(super) struct PageCandidate { + pub approximate_distance: Distance, + pub heap_key: HeapKey, +} + +pub(super) trait PageCandidateGenerator { + type Candidates: Iterator; + + fn generate( + &mut self, + query_count: usize, + token_searches: &mut dyn Iterator, + candidate_limit: usize, + ) -> Self::Candidates; +} + +#[derive(Default)] +pub(super) struct LegacyPageCandidateGenerator; + +impl PageCandidateGenerator for LegacyPageCandidateGenerator { + type Candidates = LegacyPageCandidates; + + fn generate( + &mut self, + query_count: usize, + token_searches: &mut dyn Iterator, + candidate_limit: usize, + ) -> Self::Candidates { + let mut updates = Vec::new(); + let mut estimations = Vec::with_capacity(query_count); + for (query_id, (accu_set, rough_set, estimation_by_threshold)) in token_searches.enumerate() + { + updates.reserve(accu_set.len() + rough_set.len()); + let is_empty = accu_set.is_empty() && rough_set.is_empty(); + let mut estimation_by_scope = Distance::NEG_INFINITY; + for (distance, payload) in accu_set { + estimation_by_scope = std::cmp::max(estimation_by_scope, distance); + let (key, _) = pointer_to_kv(payload); + updates.push((key, query_id, distance)); + } + for (distance, payload) in rough_set { + let (key, _) = pointer_to_kv(payload); + updates.push((key, query_id, distance)); + } + estimations.push(if !is_empty { + std::cmp::max(estimation_by_scope, estimation_by_threshold) + } else { + Distance::ZERO + }); + } + debug_assert_eq!(estimations.len(), query_count); + updates.sort_unstable_by_key(|&(key, ..)| key); + let inner = updates + .chunk_by(|(kl, ..), (kr, ..)| kl == kr) + .map(|chunk| { + let key = chunk[0].0; + let mut value = vec![None; query_count]; + for &(_, query_id, distance) in chunk { + let this = value[query_id].get_or_insert(Distance::INFINITY); + *this = std::cmp::min(*this, distance); + } + let mut maxsim = 0.0f32; + for (query_id, distance) in value.into_iter().enumerate() { + let distance = distance.unwrap_or(estimations[query_id]); + maxsim += distance.to_f32(); + } + (Reverse(Distance::from_f32(maxsim)), AlwaysEqual(key)) + }) + .collect::>(); + LegacyPageCandidates { + inner, + remaining: candidate_limit, + } + } +} + +pub(super) struct LegacyPageCandidates { + inner: BinaryHeap<(Reverse, AlwaysEqual)>, + remaining: usize, +} + +impl Iterator for LegacyPageCandidates { + type Item = PageCandidate; + + fn next(&mut self) -> Option { + if self.remaining == 0 { + return None; + } + let (Reverse(approximate_distance), AlwaysEqual(heap_key)) = self.inner.pop()?; + self.remaining -= 1; + Some(PageCandidate { + approximate_distance, + heap_key, + }) + } + + fn size_hint(&self) -> (usize, Option) { + let exact = self.inner.len().min(self.remaining); + (exact, Some(exact)) + } +} + +impl ExactSizeIterator for LegacyPageCandidates {} +impl std::iter::FusedIterator for LegacyPageCandidates {} + +#[cfg(test)] +mod tests { + use super::*; + use crate::index::fetcher::kv_to_pointer; + + fn distance(value: f32) -> Distance { + Distance::from_f32(value) + } + + #[test] + fn aggregates_tokens_by_page_and_preserves_distance_order() { + let page_1 = [0, 0, 1]; + let page_2 = [0, 0, 2]; + let mut token_searches = vec![ + ( + vec![ + (distance(-1.0), kv_to_pointer((page_1, 0))), + (distance(-0.5), kv_to_pointer((page_2, 0))), + ], + vec![(distance(-0.4), kv_to_pointer((page_2, 1)))], + distance(-0.75), + ), + ( + vec![(distance(-2.0), kv_to_pointer((page_1, 1)))], + vec![], + distance(-1.0), + ), + ] + .into_iter(); + let candidates = LegacyPageCandidateGenerator + .generate(2, &mut token_searches, usize::MAX) + .collect::>(); + + assert_eq!(candidates.len(), 2); + assert_eq!(candidates[0].heap_key, page_1); + assert_eq!(candidates[0].approximate_distance.to_f32(), -3.0); + assert_eq!(candidates[1].heap_key, page_2); + assert_eq!(candidates[1].approximate_distance.to_f32(), -1.5); + } + + #[test] + fn candidate_limit_is_applied_after_global_ordering() { + let page_1 = [0, 0, 1]; + let page_2 = [0, 0, 2]; + let mut token_searches = vec![( + vec![ + (distance(-1.0), kv_to_pointer((page_1, 0))), + (distance(-2.0), kv_to_pointer((page_2, 0))), + ], + vec![], + Distance::ZERO, + )] + .into_iter(); + let candidates = LegacyPageCandidateGenerator + .generate(1, &mut token_searches, 1) + .collect::>(); + + assert_eq!(candidates.len(), 1); + assert_eq!(candidates[0].heap_key, page_2); + } +} diff --git a/src/index/vchordrq/scanners/maxsim/external.rs b/src/index/vchordrq/scanners/maxsim/external.rs new file mode 100644 index 00000000..0c8d1eb0 --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/external.rs @@ -0,0 +1,551 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): This software is also available under the ELv2, +// which has specific restrictions. +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use super::candidate::PageCandidate; +use super::rerank::RerankError; +use crate::index::fetcher::{Fetcher, FilterableTuple, Tuple}; +use pgrx::datum::{FromDatum, IntoDatum}; +use std::ffi::CString; + +const MAX_MODEL_CONTRACT_BYTES: usize = 512; +const MAX_TENSOR_REF_BYTES: usize = 4096; +const MAX_CHECKSUM_BYTES: usize = 512; +const MAX_TENSOR_ROWS: u32 = 65_536; +const MAX_TENSOR_DIMENSION: u32 = 60_000; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum ExternalTensorDtype { + F32, + F16, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum ExternalTensorStorage { + SameHeap, + DescriptorRelation, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) struct ExternalTensorColumns { + pub model_contract: i16, + pub public_id: i16, + pub tensor_ref: i16, + pub tensor_rows: i16, + pub tensor_dimension: i16, + pub tensor_dtype: i16, + pub tensor_checksum: i16, +} + +impl ExternalTensorColumns { + pub(super) fn validate(self) -> Result { + let columns = [ + self.model_contract, + self.public_id, + self.tensor_ref, + self.tensor_rows, + self.tensor_dimension, + self.tensor_dtype, + self.tensor_checksum, + ]; + if columns.iter().any(|column| *column <= 0) { + return Err(RerankError::InvalidDescriptor( + "registered attribute number is invalid", + )); + } + for (index, column) in columns.iter().enumerate() { + if columns[..index].contains(column) { + return Err(RerankError::InvalidDescriptor( + "registered descriptor columns are not distinct", + )); + } + } + Ok(self) + } +} + +#[derive(Clone, Debug)] +pub(super) struct ExternalTensorDescriptor { + pub candidate: PageCandidate, + /// Stable application identifier. It stays inside PostgreSQL and is never + /// encoded into the sidecar request. + #[allow( + dead_code, + reason = "returned by the score API through a separate visible-row map" + )] + pub public_id: i64, + pub tensor_ref: String, + pub rows: u32, + pub dimension: u32, + pub dtype: ExternalTensorDtype, + pub checksum: String, +} + +pub(super) trait CandidateTensorDescriptorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError>; +} + +#[allow(dead_code, reason = "reserved for the optional Phase 3C heap source")] +pub(super) struct HeapTensorRefSource<'a, F> { + fetcher: &'a mut F, + model_contract_id: String, + columns: ExternalTensorColumns, +} + +#[derive(Clone, Debug)] +pub(super) struct ExternalTensorSourceBinding { + pub index_oid: pgrx::pg_sys::Oid, + pub heap_oid: pgrx::pg_sys::Oid, + pub descriptor_oid: Option, + pub storage: ExternalTensorStorage, + pub model_contract_id: String, + #[allow( + dead_code, + reason = "physical attnums are consumed by the optional Phase 3C heap source" + )] + pub columns: Option, + pub column_names: ExternalTensorColumnNames, +} + +#[derive(Clone, Debug)] +pub(super) struct ExternalTensorColumnNames { + pub model_contract: String, + pub public_id: String, + pub descriptor_public_id: Option, + pub tensor_ref: String, + pub tensor_rows: String, + pub tensor_dimension: String, + pub tensor_dtype: String, + pub tensor_checksum: String, +} + +/// Resolve the privilege-aware SQL registry boundary. Same-heap bindings also +/// resolve the physical attribute numbers consumed by [`HeapTensorRefSource`]; +/// independent descriptor relations are projected later through SPI. +/// +/// The SECURITY DEFINER SQL function performs ownership/SELECT checks and +/// revalidates the live index, heap relation, opclass, column types, and NOT +/// NULL constraints. This Rust layer deliberately calls that function rather +/// than reading the private registry table with extension privileges. +pub(super) fn resolve_external_tensor_source( + index_oid: pgrx::pg_sys::Oid, +) -> Result { + pgrx::spi::Spi::connect(|client| { + let schema_rows = client + .select( + "SELECT n.nspname::text AS schema_name + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'", + Some(1), + &[], + ) + .map_err(registry_error)?; + if schema_rows.is_empty() { + return Err(RerankError::Registry( + "vchord extension schema is unavailable".into(), + )); + } + let schema_name = schema_rows + .first() + .get_by_name::("schema_name") + .map_err(registry_error)? + .ok_or_else(|| RerankError::Registry("vchord extension schema is NULL".into()))?; + let resolver = pgrx::spi::quote_qualified_identifier( + schema_name, + "vchordrq_maxsim_source_info".to_string(), + ); + let query = format!( + "SELECT registered_index::oid AS index_oid, + heap_relation::oid AS heap_oid, + descriptor_relation::oid AS descriptor_oid, + model_contract_id, + source_storage, + model_contract_column::text AS model_contract_column, + public_id_column::text AS public_id_column, + descriptor_public_id_column::text AS descriptor_public_id_column, + tensor_ref_column::text AS tensor_ref_column, + tensor_rows_column::text AS tensor_rows_column, + tensor_dim_column::text AS tensor_dim_column, + tensor_dtype_column::text AS tensor_dtype_column, + tensor_checksum_column::text AS tensor_checksum_column + FROM {resolver}($1::regclass)" + ); + let prepared = client + .prepare(query.as_str(), &pgrx::oids_of![pgrx::pg_sys::Oid]) + .map_err(registry_error)?; + let rows = client + .select(&prepared, Some(1), &[index_oid.into()]) + .map_err(registry_error)?; + if rows.is_empty() { + return Err(RerankError::Registry( + "registered MaxSim tensor source resolution returned no row".into(), + )); + } + let row = rows.first(); + let resolved_index_oid = required_column::(&row, "index_oid")?; + let heap_oid = required_column::(&row, "heap_oid")?; + let descriptor_oid = optional_column::(&row, "descriptor_oid")?; + let model_contract_id = required_column::(&row, "model_contract_id")?; + let storage = required_column::(&row, "source_storage")?; + if resolved_index_oid != index_oid { + return Err(RerankError::Registry( + "registered MaxSim tensor source resolved a different index".into(), + )); + } + let storage = match storage.as_str() { + "external_ref" if descriptor_oid.is_none() => ExternalTensorStorage::SameHeap, + "external_relation" if descriptor_oid.is_some() => { + ExternalTensorStorage::DescriptorRelation + } + "heap_array" => { + return Err(RerankError::Registry( + "registered MaxSim tensor source is not external".into(), + )); + } + _ => { + return Err(RerankError::Registry( + "registered external MaxSim tensor source is inconsistent".into(), + )); + } + }; + + let column_names = ExternalTensorColumnNames { + model_contract: required_column::(&row, "model_contract_column")?, + public_id: required_column::(&row, "public_id_column")?, + descriptor_public_id: optional_column::(&row, "descriptor_public_id_column")?, + tensor_ref: required_column::(&row, "tensor_ref_column")?, + tensor_rows: required_column::(&row, "tensor_rows_column")?, + tensor_dimension: required_column::(&row, "tensor_dim_column")?, + tensor_dtype: required_column::(&row, "tensor_dtype_column")?, + tensor_checksum: required_column::(&row, "tensor_checksum_column")?, + }; + let columns = match storage { + ExternalTensorStorage::SameHeap => Some( + ExternalTensorColumns { + model_contract: resolve_attnum(heap_oid, &column_names.model_contract)?, + public_id: resolve_attnum(heap_oid, &column_names.public_id)?, + tensor_ref: resolve_attnum(heap_oid, &column_names.tensor_ref)?, + tensor_rows: resolve_attnum(heap_oid, &column_names.tensor_rows)?, + tensor_dimension: resolve_attnum(heap_oid, &column_names.tensor_dimension)?, + tensor_dtype: resolve_attnum(heap_oid, &column_names.tensor_dtype)?, + tensor_checksum: resolve_attnum(heap_oid, &column_names.tensor_checksum)?, + } + .validate()?, + ), + ExternalTensorStorage::DescriptorRelation => { + if column_names.descriptor_public_id.is_none() { + return Err(RerankError::Registry( + "registered descriptor relation has no public ID column".into(), + )); + } + None + } + }; + validate_model_contract(&model_contract_id)?; + Ok(ExternalTensorSourceBinding { + index_oid, + heap_oid, + descriptor_oid, + storage, + model_contract_id, + columns, + column_names, + }) + }) +} + +fn registry_error(error: impl std::fmt::Display) -> RerankError { + RerankError::Registry(error.to_string()) +} + +fn required_column(row: &pgrx::spi::SpiTupleTable<'_>, name: &str) -> Result +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name) + .map_err(registry_error)? + .ok_or_else(|| RerankError::Registry(format!("registry resolver returned NULL {name}"))) +} + +fn optional_column( + row: &pgrx::spi::SpiTupleTable<'_>, + name: &str, +) -> Result, RerankError> +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name).map_err(registry_error) +} + +fn resolve_attnum(heap_oid: pgrx::pg_sys::Oid, column_name: &str) -> Result { + let column_name = CString::new(column_name) + .map_err(|_| RerankError::Registry("registry column contains a NUL byte".into()))?; + let attnum = unsafe { pgrx::pg_sys::get_attnum(heap_oid, column_name.as_ptr()) }; + if attnum <= 0 { + return Err(RerankError::Registry( + "registered descriptor column disappeared during resolution".into(), + )); + } + Ok(attnum) +} + +#[allow(dead_code, reason = "reserved for the optional Phase 3C heap source")] +impl<'a, F> HeapTensorRefSource<'a, F> { + pub(super) fn new( + fetcher: &'a mut F, + model_contract_id: String, + columns: ExternalTensorColumns, + ) -> Result { + validate_model_contract(&model_contract_id)?; + Ok(Self { + fetcher, + model_contract_id, + columns: columns.validate()?, + }) + } +} + +impl CandidateTensorDescriptorSource for HeapTensorRefSource<'_, F> { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + let Some(mut tuple) = self.fetcher.fetch(candidate.heap_key) else { + return Ok(None); + }; + + // This is the data-exfiltration boundary: no descriptor value may be + // read before PostgreSQL has accepted the active base-relation qual. + if !tuple.filter() { + return Ok(None); + } + + let model_contract = read_attribute::(&mut tuple, self.columns.model_contract)?; + if model_contract != self.model_contract_id { + return Err(RerankError::ModelContractMismatch); + } + let public_id = read_attribute::(&mut tuple, self.columns.public_id)?; + let tensor_ref = read_attribute::(&mut tuple, self.columns.tensor_ref)?; + let rows = read_attribute::(&mut tuple, self.columns.tensor_rows)?; + let dimension = read_attribute::(&mut tuple, self.columns.tensor_dimension)?; + let dtype = read_attribute::(&mut tuple, self.columns.tensor_dtype)?; + let checksum = read_attribute::(&mut tuple, self.columns.tensor_checksum)?; + + validate_descriptor( + candidate, public_id, tensor_ref, rows, dimension, dtype, checksum, + ) + .map(Some) + } +} + +#[allow(dead_code, reason = "reserved for the optional Phase 3C heap source")] +fn read_attribute(tuple: &mut impl Tuple, attnum: i16) -> Result { + let attribute = tuple + .attribute(attnum) + .ok_or(RerankError::InvalidDescriptor( + "registered attribute is unavailable", + ))?; + unsafe { T::from_datum(attribute.datum, attribute.is_null) }.ok_or( + RerankError::InvalidDescriptor("registered descriptor value is NULL or malformed"), + ) +} + +fn validate_model_contract(value: &str) -> Result<(), RerankError> { + if value.is_empty() || value.len() > MAX_MODEL_CONTRACT_BYTES || contains_control(value) { + return Err(RerankError::InvalidDescriptor( + "model contract is empty, oversized, or contains control characters", + )); + } + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +pub(super) fn validate_descriptor( + candidate: PageCandidate, + public_id: i64, + tensor_ref: String, + rows: i32, + dimension: i32, + dtype: String, + checksum: String, +) -> Result { + if tensor_ref.is_empty() + || tensor_ref.len() > MAX_TENSOR_REF_BYTES + || contains_control(&tensor_ref) + { + return Err(RerankError::InvalidDescriptor( + "tensor reference is empty, oversized, or contains control characters", + )); + } + let rows = u32::try_from(rows) + .ok() + .filter(|rows| (1..=MAX_TENSOR_ROWS).contains(rows)) + .ok_or(RerankError::InvalidDescriptor( + "tensor row count is invalid", + ))?; + let dimension = u32::try_from(dimension) + .ok() + .filter(|dimension| (1..=MAX_TENSOR_DIMENSION).contains(dimension)) + .ok_or(RerankError::InvalidDescriptor( + "tensor dimension is invalid", + ))?; + let dtype = match dtype.as_str() { + "float32" => ExternalTensorDtype::F32, + "float16" => ExternalTensorDtype::F16, + _ => { + return Err(RerankError::InvalidDescriptor( + "tensor dtype must be float16 or float32", + )); + } + }; + if checksum.len() > MAX_CHECKSUM_BYTES || !is_sha256_checksum(&checksum) { + return Err(RerankError::InvalidDescriptor( + "tensor checksum must be a lowercase sha256 digest", + )); + } + Ok(ExternalTensorDescriptor { + candidate, + public_id, + tensor_ref, + rows, + dimension, + dtype, + checksum, + }) +} + +fn contains_control(value: &str) -> bool { + value.chars().any(char::is_control) +} + +fn is_sha256_checksum(value: &str) -> bool { + value.strip_prefix("sha256:").is_some_and(|digest| { + digest.len() == 64 + && digest + .bytes() + .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase()) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::index::fetcher::TupleAttribute; + use distance::Distance; + + fn candidate() -> PageCandidate { + PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: [0, 0, 1], + } + } + + fn columns() -> ExternalTensorColumns { + ExternalTensorColumns { + model_contract: 1, + public_id: 2, + tensor_ref: 3, + tensor_rows: 4, + tensor_dimension: 5, + tensor_dtype: 6, + tensor_checksum: 7, + } + } + + struct RejectingFetcher; + struct RejectingTuple; + + impl Tuple for RejectingTuple { + fn build(&mut self) -> (&[pgrx::pg_sys::Datum; 32], &[bool; 32]) { + panic!("external descriptor sources do not build index expressions") + } + + fn attribute(&mut self, _attnum: i16) -> Option { + panic!("a rejected tuple must not expose descriptor attributes") + } + } + + impl FilterableTuple for RejectingTuple { + fn filter(&mut self) -> bool { + false + } + } + + impl Fetcher for RejectingFetcher { + type Tuple<'a> = RejectingTuple; + + fn fetch(&mut self, _key: [u16; 3]) -> Option> { + Some(RejectingTuple) + } + } + + #[test] + fn source_filters_before_reading_descriptor_attributes() { + let mut fetcher = RejectingFetcher; + let mut source = + HeapTensorRefSource::new(&mut fetcher, "contract@1".into(), columns()).unwrap(); + assert!(source.fetch(candidate()).unwrap().is_none()); + } + + #[test] + fn descriptor_validation_accepts_bounded_immutable_metadata() { + let descriptor = validate_descriptor( + candidate(), + 42, + "s3://immutable-bucket/page-42.tensor".into(), + 747, + 320, + "float16".into(), + format!("sha256:{}", "a".repeat(64)), + ) + .unwrap(); + assert_eq!(descriptor.public_id, 42); + assert_eq!(descriptor.rows, 747); + assert_eq!(descriptor.dimension, 320); + assert_eq!(descriptor.dtype, ExternalTensorDtype::F16); + } + + #[test] + fn descriptor_validation_rejects_ambiguous_or_unbounded_metadata() { + assert!(matches!( + columns_with_duplicate().validate(), + Err(RerankError::InvalidDescriptor(_)) + )); + for (rows, dimension, dtype, checksum) in [ + (0, 320, "float16", format!("sha256:{}", "a".repeat(64))), + (1, 0, "float16", format!("sha256:{}", "a".repeat(64))), + (1, 320, "bf16", format!("sha256:{}", "a".repeat(64))), + (1, 320, "float16", "sha256:short".into()), + ] { + assert!(matches!( + validate_descriptor( + candidate(), + 1, + "s3://immutable/tensor".into(), + rows, + dimension, + dtype.into(), + checksum, + ), + Err(RerankError::InvalidDescriptor(_)) + )); + } + } + + fn columns_with_duplicate() -> ExternalTensorColumns { + ExternalTensorColumns { + public_id: 1, + ..columns() + } + } +} diff --git a/src/index/vchordrq/scanners/maxsim/gpu.rs b/src/index/vchordrq/scanners/maxsim/gpu.rs new file mode 100644 index 00000000..af8c5455 --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/gpu.rs @@ -0,0 +1,1495 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use super::candidate::{HeapKey, PageCandidate}; +use super::external::{ + CandidateTensorDescriptorSource, ExternalTensorDescriptor, ExternalTensorDtype, +}; +use super::rerank::{CandidateTensorSource, ExactMaxsimBackend, RerankError, RerankResults}; +use distance::Distance; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use std::mem::size_of_val; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; +use vchordrq::types::OwnedVector; + +const MAGIC: &[u8; 4] = b"VCTM"; +const VERSION: u16 = 1; +const EXTERNAL_VERSION: u16 = 2; +const REQUEST_KIND: u16 = 1; +const RESPONSE_KIND: u16 = 2; +const HEADER_LEN: usize = 24; +const MAX_REMOTE_ERROR_BYTES: usize = 64 * 1024; + +static NEXT_REQUEST_ID: AtomicU64 = AtomicU64::new(1); +static LAST_FALLBACK_WARNING_SECONDS: AtomicU64 = AtomicU64::new(0); + +pub(super) fn report_gpu_fallback(error: &RerankError) { + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map_or(0, |duration| duration.as_secs()); + let should_warn = LAST_FALLBACK_WARNING_SECONDS + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |last| { + (last == 0 || now.saturating_sub(last) >= 60).then_some(now) + }) + .is_ok(); + if should_warn { + pgrx::warning!("GPU MaxSim failed; using cpu_exact: {error}"); + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum TensorDtype { + F32 = 1, + F16 = 2, +} + +pub(super) trait TileMaxsimTransport { + fn round_trip( + &mut self, + request: &[u8], + timeout: Duration, + max_response_bytes: usize, + ) -> Result, RerankError>; +} + +pub(super) struct GpuTileMaxsimBackend { + transport: T, + timeout: Duration, + max_batch_tokens: usize, + max_batch_bytes: usize, +} + +impl GpuTileMaxsimBackend { + pub fn new( + transport: T, + timeout: Duration, + max_batch_tokens: usize, + max_batch_bytes: usize, + ) -> Self { + Self { + transport, + timeout, + max_batch_tokens, + max_batch_bytes, + } + } +} + +impl ExactMaxsimBackend for GpuTileMaxsimBackend { + type Results = RerankResults; + + fn rerank( + &mut self, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + ) -> Result { + let request_id = NEXT_REQUEST_ID.fetch_add(1, Ordering::Relaxed); + let encoded = encode_request( + request_id, + query, + candidates, + source, + self.max_batch_tokens, + self.max_batch_bytes, + )?; + if encoded.heap_keys.is_empty() { + return Ok(RerankResults { + inner: BinaryHeap::new(), + }); + } + let max_response_bytes = HEADER_LEN + .checked_add(8) + .and_then(|size| size.checked_add(encoded.heap_keys.len().checked_mul(8)?)) + .map(|size| size.max(HEADER_LEN + 8 + MAX_REMOTE_ERROR_BYTES)) + .ok_or(RerankError::RequestTooLarge)?; + let response = + self.transport + .round_trip(&encoded.frame, self.timeout, max_response_bytes)?; + decode_response(&response, request_id, &encoded.heap_keys) + } +} + +/// GPU backend for Phase 3B external tensor descriptors. +/// +/// This codec is intentionally separate from [`ExactMaxsimBackend`]: an +/// external full tensor may have different logical values from the indexed +/// sketch, so it cannot be substituted into ordinary `@#` execution. +pub(super) struct GpuExternalTileMaxsimBackend { + transport: T, + model_contract_id: String, + timeout: Duration, + max_batch_tokens: usize, + max_batch_bytes: usize, +} + +impl GpuExternalTileMaxsimBackend { + pub(super) fn new( + transport: T, + model_contract_id: String, + timeout: Duration, + max_batch_tokens: usize, + max_batch_bytes: usize, + ) -> Self { + Self { + transport, + model_contract_id, + timeout, + max_batch_tokens, + max_batch_bytes, + } + } +} + +impl GpuExternalTileMaxsimBackend { + pub(super) fn rerank( + &mut self, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + ) -> Result { + let request_id = NEXT_REQUEST_ID.fetch_add(1, Ordering::Relaxed); + let encoded = encode_external_request( + request_id, + &self.model_contract_id, + query, + candidates, + source, + self.max_batch_tokens, + self.max_batch_bytes, + )?; + if encoded.heap_keys.is_empty() { + return Ok(RerankResults { + inner: BinaryHeap::new(), + }); + } + let max_response_bytes = HEADER_LEN + .checked_add(8) + .and_then(|size| size.checked_add(encoded.heap_keys.len().checked_mul(8)?)) + .map(|size| size.max(HEADER_LEN + 8 + MAX_REMOTE_ERROR_BYTES)) + .ok_or(RerankError::RequestTooLarge)?; + let response = + self.transport + .round_trip(&encoded.frame, self.timeout, max_response_bytes)?; + decode_response_for_version(&response, EXTERNAL_VERSION, request_id, &encoded.heap_keys) + } +} + +struct EncodedRequest { + frame: Vec, + heap_keys: Vec, +} + +fn encode_external_request( + request_id: u64, + model_contract_id: &str, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + max_batch_tokens: usize, + max_batch_bytes: usize, +) -> Result { + const MAX_MODEL_CONTRACT_BYTES: usize = 512; + const MAX_TENSOR_REF_BYTES: usize = 4096; + const MAX_CHECKSUM_BYTES: usize = 512; + + if model_contract_id.is_empty() + || model_contract_id.len() > MAX_MODEL_CONTRACT_BYTES + || model_contract_id.chars().any(char::is_control) + { + return Err(RerankError::InvalidDescriptor( + "model contract is empty, oversized, or contains control characters", + )); + } + let (dtype, dimension) = tensor_metadata(query)?; + let external_dtype = match dtype { + TensorDtype::F32 => ExternalTensorDtype::F32, + TensorDtype::F16 => ExternalTensorDtype::F16, + }; + let query_rows = u32::try_from(query.len()).map_err(|_| RerankError::RequestTooLarge)?; + let mut total_tokens = query.len(); + if total_tokens > max_batch_tokens { + return Err(RerankError::RequestTooLarge); + } + let mut declared_tensor_bytes = tensor_bytes(query_rows, dimension, dtype)?; + if declared_tensor_bytes > max_batch_bytes { + return Err(RerankError::RequestTooLarge); + } + + let mut writer = BoundedWriter::new(max_batch_bytes); + writer.zeros(HEADER_LEN)?; + writer.u32(dimension)?; + writer.u32(query_rows)?; + let candidate_count_offset = writer.len(); + writer.u32(0)?; + writer.u8(dtype as u8)?; + writer.u8(1)?; // sum_query_max_document_dot + writer.u16(0)?; + writer + .u32(u32::try_from(model_contract_id.len()).map_err(|_| RerankError::RequestTooLarge)?)?; + writer.bytes(model_contract_id.as_bytes())?; + encode_tensor_values(&mut writer, query, dtype)?; + + let mut heap_keys = Vec::new(); + for candidate in candidates { + let Some(descriptor) = source.fetch(candidate)? else { + continue; + }; + validate_external_for_request( + &descriptor, + dimension, + external_dtype, + MAX_TENSOR_REF_BYTES, + MAX_CHECKSUM_BYTES, + )?; + total_tokens = total_tokens + .checked_add(descriptor.rows as usize) + .ok_or(RerankError::RequestTooLarge)?; + if total_tokens > max_batch_tokens { + return Err(RerankError::RequestTooLarge); + } + declared_tensor_bytes = declared_tensor_bytes + .checked_add(tensor_bytes(descriptor.rows, dimension, dtype)?) + .ok_or(RerankError::RequestTooLarge)?; + if declared_tensor_bytes > max_batch_bytes { + return Err(RerankError::RequestTooLarge); + } + + let candidate_id = + u32::try_from(heap_keys.len()).map_err(|_| RerankError::RequestTooLarge)?; + writer.u32(candidate_id)?; + writer.u32(descriptor.rows)?; + writer.u32( + u32::try_from(descriptor.tensor_ref.len()).map_err(|_| RerankError::RequestTooLarge)?, + )?; + writer.u32( + u32::try_from(descriptor.checksum.len()).map_err(|_| RerankError::RequestTooLarge)?, + )?; + writer.bytes(descriptor.tensor_ref.as_bytes())?; + writer.bytes(descriptor.checksum.as_bytes())?; + heap_keys.push(descriptor.candidate.heap_key); + } + + let candidate_count = + u32::try_from(heap_keys.len()).map_err(|_| RerankError::RequestTooLarge)?; + writer.patch_u32(candidate_count_offset, candidate_count); + let body_len = writer + .len() + .checked_sub(HEADER_LEN) + .ok_or_else(|| RerankError::Protocol("invalid request length".into()))?; + writer.patch_bytes(0, MAGIC); + writer.patch_u16(4, EXTERNAL_VERSION); + writer.patch_u16(6, REQUEST_KIND); + writer.patch_u64(8, request_id); + writer.patch_u64( + 16, + u64::try_from(body_len).map_err(|_| RerankError::RequestTooLarge)?, + ); + Ok(EncodedRequest { + frame: writer.finish(), + heap_keys, + }) +} + +fn tensor_bytes(rows: u32, dimension: u32, dtype: TensorDtype) -> Result { + let scalar_bytes = match dtype { + TensorDtype::F32 => 4usize, + TensorDtype::F16 => 2usize, + }; + usize::try_from(rows) + .ok() + .and_then(|rows| rows.checked_mul(dimension as usize)) + .and_then(|elements| elements.checked_mul(scalar_bytes)) + .ok_or(RerankError::RequestTooLarge) +} + +fn validate_external_for_request( + descriptor: &ExternalTensorDescriptor, + dimension: u32, + dtype: ExternalTensorDtype, + max_tensor_ref_bytes: usize, + max_checksum_bytes: usize, +) -> Result<(), RerankError> { + if descriptor.rows == 0 + || descriptor.rows > 65_536 + || dimension > 60_000 + || descriptor.dimension != dimension + || descriptor.dtype != dtype + { + return Err(RerankError::TensorMismatch); + } + if descriptor.tensor_ref.is_empty() + || descriptor.tensor_ref.len() > max_tensor_ref_bytes + || descriptor.tensor_ref.chars().any(char::is_control) + { + return Err(RerankError::InvalidDescriptor( + "tensor reference is empty, oversized, or contains control characters", + )); + } + if descriptor.checksum.len() > max_checksum_bytes + || !descriptor + .checksum + .strip_prefix("sha256:") + .is_some_and(|digest| { + digest.len() == 64 + && digest + .bytes() + .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase()) + }) + { + return Err(RerankError::InvalidDescriptor( + "tensor checksum must be a lowercase sha256 digest", + )); + } + Ok(()) +} + +fn encode_request( + request_id: u64, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + max_batch_tokens: usize, + max_batch_bytes: usize, +) -> Result { + let (dtype, dimension) = tensor_metadata(query)?; + let query_rows = u32::try_from(query.len()).map_err(|_| RerankError::RequestTooLarge)?; + let mut total_tokens = query.len(); + if total_tokens > max_batch_tokens { + return Err(RerankError::RequestTooLarge); + } + + let mut writer = BoundedWriter::new(max_batch_bytes); + writer.zeros(HEADER_LEN)?; + writer.u32(dimension)?; + writer.u32(query_rows)?; + let candidate_count_offset = writer.len(); + writer.u32(0)?; + writer.u8(dtype as u8)?; + writer.u8(1)?; // sum_query_max_document_dot + writer.u16(0)?; + encode_tensor_values(&mut writer, query, dtype)?; + + let mut heap_keys = Vec::new(); + for candidate in candidates { + let Some(tensor) = source.fetch(candidate)? else { + continue; + }; + let (candidate_dtype, candidate_dimension) = tensor_metadata(&tensor.vectors)?; + if candidate_dtype != dtype || candidate_dimension != dimension { + return Err(RerankError::TensorMismatch); + } + total_tokens = total_tokens + .checked_add(tensor.vectors.len()) + .ok_or(RerankError::RequestTooLarge)?; + if total_tokens > max_batch_tokens { + return Err(RerankError::RequestTooLarge); + } + let candidate_id = + u32::try_from(heap_keys.len()).map_err(|_| RerankError::RequestTooLarge)?; + let rows = u32::try_from(tensor.vectors.len()).map_err(|_| RerankError::RequestTooLarge)?; + writer.u32(candidate_id)?; + writer.u32(rows)?; + encode_tensor_values(&mut writer, &tensor.vectors, dtype)?; + heap_keys.push(tensor.candidate.heap_key); + } + let candidate_count = + u32::try_from(heap_keys.len()).map_err(|_| RerankError::RequestTooLarge)?; + writer.patch_u32(candidate_count_offset, candidate_count); + let body_len = writer + .len() + .checked_sub(HEADER_LEN) + .ok_or_else(|| RerankError::Protocol("invalid request length".into()))?; + let body_len = u64::try_from(body_len).map_err(|_| RerankError::RequestTooLarge)?; + writer.patch_bytes(0, MAGIC); + writer.patch_u16(4, VERSION); + writer.patch_u16(6, REQUEST_KIND); + writer.patch_u64(8, request_id); + writer.patch_u64(16, body_len); + Ok(EncodedRequest { + frame: writer.finish(), + heap_keys, + }) +} + +fn tensor_metadata(vectors: &[OwnedVector]) -> Result<(TensorDtype, u32), RerankError> { + let Some(first) = vectors.first() else { + return Err(RerankError::TensorMismatch); + }; + let dtype = match first { + OwnedVector::Vecf32(_) => TensorDtype::F32, + OwnedVector::Vecf16(_) => TensorDtype::F16, + OwnedVector::Rabitq8(_) | OwnedVector::Rabitq4(_) => { + return Err(RerankError::UnsupportedTensorKind); + } + }; + let dimension = first.dim(); + for vector in vectors { + let this_dtype = match vector { + OwnedVector::Vecf32(_) => TensorDtype::F32, + OwnedVector::Vecf16(_) => TensorDtype::F16, + OwnedVector::Rabitq8(_) | OwnedVector::Rabitq4(_) => { + return Err(RerankError::UnsupportedTensorKind); + } + }; + if this_dtype != dtype || vector.dim() != dimension { + return Err(RerankError::TensorMismatch); + } + } + Ok((dtype, dimension)) +} + +fn encode_tensor_values( + writer: &mut BoundedWriter, + vectors: &[OwnedVector], + dtype: TensorDtype, +) -> Result<(), RerankError> { + for vector in vectors { + match (dtype, vector) { + (TensorDtype::F32, OwnedVector::Vecf32(vector)) => { + for value in vector.slice() { + writer.bytes(&value.to_le_bytes())?; + } + } + (TensorDtype::F16, OwnedVector::Vecf16(vector)) => { + for value in vector.slice() { + writer.u16(value.to_bits())?; + } + } + _ => return Err(RerankError::TensorMismatch), + } + } + Ok(()) +} + +fn decode_response( + frame: &[u8], + request_id: u64, + heap_keys: &[HeapKey], +) -> Result { + decode_response_for_version(frame, VERSION, request_id, heap_keys) +} + +fn decode_response_for_version( + frame: &[u8], + expected_version: u16, + request_id: u64, + heap_keys: &[HeapKey], +) -> Result { + let mut cursor = Cursor::new(frame); + if cursor.bytes(4)? != MAGIC { + return Err(RerankError::Protocol("invalid magic".into())); + } + if cursor.u16()? != expected_version { + return Err(RerankError::Protocol("unsupported version".into())); + } + if cursor.u16()? != RESPONSE_KIND { + return Err(RerankError::Protocol("unexpected message kind".into())); + } + if cursor.u64()? != request_id { + return Err(RerankError::Protocol("request ID mismatch".into())); + } + let body_len = usize::try_from(cursor.u64()?) + .map_err(|_| RerankError::Protocol("response body is too large".into()))?; + if body_len != frame.len().saturating_sub(HEADER_LEN) { + return Err(RerankError::Protocol("response length mismatch".into())); + } + let status = cursor.u32()?; + if status != 0 { + let length = usize::try_from(cursor.u32()?) + .map_err(|_| RerankError::Protocol("remote error is too large".into()))?; + if length > MAX_REMOTE_ERROR_BYTES { + return Err(RerankError::Protocol("remote error is too large".into())); + } + let message = std::str::from_utf8(cursor.bytes(length)?) + .map_err(|_| RerankError::Protocol("remote error is not UTF-8".into()))?; + cursor.finish()?; + return Err(RerankError::Remote(message.into())); + } + let result_count = usize::try_from(cursor.u32()?) + .map_err(|_| RerankError::Protocol("result count is too large".into()))?; + if result_count != heap_keys.len() { + return Err(RerankError::Protocol("partial result set".into())); + } + let mut seen = vec![false; heap_keys.len()]; + let mut results = BinaryHeap::new(); + for _ in 0..result_count { + let candidate_id = usize::try_from(cursor.u32()?) + .map_err(|_| RerankError::Protocol("candidate ID is too large".into()))?; + let Some(heap_key) = heap_keys.get(candidate_id).copied() else { + return Err(RerankError::Protocol("unknown candidate ID".into())); + }; + if std::mem::replace(&mut seen[candidate_id], true) { + return Err(RerankError::Protocol("duplicate candidate ID".into())); + } + let similarity = f32::from_bits(cursor.u32()?); + if !similarity.is_finite() { + return Err(RerankError::Protocol("non-finite similarity".into())); + } + let distance = Distance::from_f32(-similarity); + results.push((Reverse(distance), Reverse(heap_key))); + } + cursor.finish()?; + if seen.iter().any(|seen| !seen) { + return Err(RerankError::Protocol("partial result set".into())); + } + Ok(RerankResults { inner: results }) +} + +struct BoundedWriter { + bytes: Vec, + limit: usize, +} + +impl BoundedWriter { + fn new(limit: usize) -> Self { + Self { + bytes: Vec::new(), + limit, + } + } + + fn len(&self) -> usize { + self.bytes.len() + } + + fn ensure(&self, additional: usize) -> Result<(), RerankError> { + let size = self + .bytes + .len() + .checked_add(additional) + .ok_or(RerankError::RequestTooLarge)?; + if size > self.limit { + return Err(RerankError::RequestTooLarge); + } + Ok(()) + } + + fn zeros(&mut self, count: usize) -> Result<(), RerankError> { + self.ensure(count)?; + self.bytes.resize(self.bytes.len() + count, 0); + Ok(()) + } + + fn bytes(&mut self, bytes: &[u8]) -> Result<(), RerankError> { + self.ensure(bytes.len())?; + self.bytes.extend_from_slice(bytes); + Ok(()) + } + + fn u8(&mut self, value: u8) -> Result<(), RerankError> { + self.bytes(&[value]) + } + + fn u16(&mut self, value: u16) -> Result<(), RerankError> { + self.bytes(&value.to_le_bytes()) + } + + fn u32(&mut self, value: u32) -> Result<(), RerankError> { + self.bytes(&value.to_le_bytes()) + } + + fn patch_bytes(&mut self, offset: usize, bytes: &[u8]) { + self.bytes[offset..offset + bytes.len()].copy_from_slice(bytes); + } + + fn patch_u16(&mut self, offset: usize, value: u16) { + self.patch_bytes(offset, &value.to_le_bytes()); + } + + fn patch_u32(&mut self, offset: usize, value: u32) { + self.patch_bytes(offset, &value.to_le_bytes()); + } + + fn patch_u64(&mut self, offset: usize, value: u64) { + self.patch_bytes(offset, &value.to_le_bytes()); + } + + fn finish(self) -> Vec { + self.bytes + } +} + +struct Cursor<'a> { + bytes: &'a [u8], + offset: usize, +} + +impl<'a> Cursor<'a> { + fn new(bytes: &'a [u8]) -> Self { + Self { bytes, offset: 0 } + } + + fn bytes(&mut self, count: usize) -> Result<&'a [u8], RerankError> { + let end = self + .offset + .checked_add(count) + .ok_or_else(|| RerankError::Protocol("message offset overflow".into()))?; + let bytes = self + .bytes + .get(self.offset..end) + .ok_or_else(|| RerankError::Protocol("truncated message".into()))?; + self.offset = end; + Ok(bytes) + } + + fn u16(&mut self) -> Result { + Ok(u16::from_le_bytes(self.bytes(2)?.try_into().unwrap())) + } + + fn u32(&mut self) -> Result { + Ok(u32::from_le_bytes(self.bytes(4)?.try_into().unwrap())) + } + + fn u64(&mut self) -> Result { + Ok(u64::from_le_bytes(self.bytes(8)?.try_into().unwrap())) + } + + fn finish(self) -> Result<(), RerankError> { + if self.offset != self.bytes.len() { + return Err(RerankError::Protocol("trailing response bytes".into())); + } + Ok(()) + } +} + +pub(super) struct UnixSocketTransport { + endpoint: String, +} + +impl UnixSocketTransport { + pub fn new(endpoint: String) -> Self { + Self { endpoint } + } +} + +#[cfg(unix)] +impl TileMaxsimTransport for UnixSocketTransport { + fn round_trip( + &mut self, + request: &[u8], + timeout: Duration, + max_response_bytes: usize, + ) -> Result, RerankError> { + use std::time::Instant; + + if self.endpoint.is_empty() { + return Err(RerankError::Transport("endpoint is empty".into())); + } + let deadline = Instant::now() + .checked_add(timeout) + .ok_or_else(|| RerankError::Transport("invalid timeout".into()))?; + let mut stream = connect_interruptible(&self.endpoint, deadline)?; + let poll = remaining_until(deadline)?.min(Duration::from_millis(50)); + stream + .set_read_timeout(Some(poll)) + .map_err(|error| RerankError::Transport(error.to_string()))?; + stream + .set_write_timeout(Some(poll)) + .map_err(|error| RerankError::Transport(error.to_string()))?; + + write_interruptible(&mut stream, request, deadline)?; + let mut header = [0u8; HEADER_LEN]; + read_interruptible(&mut stream, &mut header, deadline)?; + let body_len = usize::try_from(u64::from_le_bytes(header[16..24].try_into().unwrap())) + .map_err(|_| RerankError::Protocol("response body is too large".into()))?; + let response_len = HEADER_LEN + .checked_add(body_len) + .ok_or_else(|| RerankError::Protocol("response length overflow".into()))?; + if response_len > max_response_bytes { + return Err(RerankError::Protocol( + "response exceeds configured limit".into(), + )); + } + let mut response = Vec::with_capacity(response_len); + response.extend_from_slice(&header); + response.resize(response_len, 0); + read_interruptible(&mut stream, &mut response[HEADER_LEN..], deadline)?; + Ok(response) + } +} + +#[cfg(unix)] +fn connect_interruptible( + endpoint: &str, + deadline: std::time::Instant, +) -> Result { + use std::os::fd::{AsRawFd, FromRawFd, OwnedFd}; + + let (address, address_len) = unix_socket_address(endpoint)?; + let raw_fd = unsafe { libc::socket(libc::AF_UNIX, libc::SOCK_STREAM, 0) }; + if raw_fd < 0 { + return Err(last_transport_error()); + } + let fd = unsafe { OwnedFd::from_raw_fd(raw_fd) }; + update_fd_flag( + fd.as_raw_fd(), + libc::F_GETFD, + libc::F_SETFD, + libc::FD_CLOEXEC, + true, + )?; + update_fd_flag( + fd.as_raw_fd(), + libc::F_GETFL, + libc::F_SETFL, + libc::O_NONBLOCK, + true, + )?; + + let connected = unsafe { + libc::connect( + fd.as_raw_fd(), + (&raw const address).cast::(), + address_len, + ) + } == 0; + if !connected { + let error = std::io::Error::last_os_error(); + let raw_error = error.raw_os_error(); + if raw_error != Some(libc::EINPROGRESS) + && raw_error != Some(libc::EAGAIN) + && raw_error != Some(libc::EWOULDBLOCK) + { + return Err(RerankError::Transport(error.to_string())); + } + wait_for_connect(fd.as_raw_fd(), deadline)?; + } + + update_fd_flag( + fd.as_raw_fd(), + libc::F_GETFL, + libc::F_SETFL, + libc::O_NONBLOCK, + false, + )?; + Ok(std::os::unix::net::UnixStream::from(fd)) +} + +#[cfg(unix)] +fn unix_socket_address( + endpoint: &str, +) -> Result<(libc::sockaddr_un, libc::socklen_t), RerankError> { + let path = endpoint.as_bytes(); + let mut address = unsafe { std::mem::zeroed::() }; + if path.contains(&0) { + return Err(RerankError::Transport( + "endpoint contains a NUL byte".into(), + )); + } + if path.len() >= address.sun_path.len() { + return Err(RerankError::Transport("endpoint path is too long".into())); + } + address.sun_family = libc::AF_UNIX as libc::sa_family_t; + unsafe { + std::ptr::copy_nonoverlapping( + path.as_ptr().cast::(), + address.sun_path.as_mut_ptr(), + path.len(), + ); + } + let length = std::mem::offset_of!(libc::sockaddr_un, sun_path) + .checked_add(path.len()) + .and_then(|length| length.checked_add(1)) + .and_then(|length| libc::socklen_t::try_from(length).ok()) + .ok_or_else(|| RerankError::Transport("endpoint path is too long".into()))?; + #[cfg(any( + target_os = "aix", + target_os = "dragonfly", + target_os = "freebsd", + target_os = "haiku", + target_os = "hurd", + target_os = "ios", + target_os = "macos", + target_os = "netbsd", + target_os = "openbsd", + target_os = "tvos", + target_os = "visionos", + target_os = "watchos" + ))] + { + address.sun_len = u8::try_from(length) + .map_err(|_| RerankError::Transport("endpoint path is too long".into()))?; + } + Ok((address, length)) +} + +#[cfg(unix)] +fn update_fd_flag( + fd: std::os::fd::RawFd, + get_command: libc::c_int, + set_command: libc::c_int, + flag: libc::c_int, + enabled: bool, +) -> Result<(), RerankError> { + let current = unsafe { libc::fcntl(fd, get_command) }; + if current < 0 { + return Err(last_transport_error()); + } + let updated = if enabled { + current | flag + } else { + current & !flag + }; + if unsafe { libc::fcntl(fd, set_command, updated) } < 0 { + return Err(last_transport_error()); + } + Ok(()) +} + +#[cfg(unix)] +fn wait_for_connect( + fd: std::os::fd::RawFd, + deadline: std::time::Instant, +) -> Result<(), RerankError> { + loop { + pgrx::check_for_interrupts!(); + let remaining = remaining_until(deadline)?; + let timeout_ms = remaining.min(Duration::from_millis(50)).as_millis().max(1) as libc::c_int; + let mut poll_fd = libc::pollfd { + fd, + events: libc::POLLOUT, + revents: 0, + }; + let result = unsafe { libc::poll(&mut poll_fd, 1, timeout_ms) }; + if result == 0 { + continue; + } + if result < 0 { + let error = std::io::Error::last_os_error(); + if error.kind() == std::io::ErrorKind::Interrupted { + continue; + } + return Err(RerankError::Transport(error.to_string())); + } + let mut socket_error = 0; + let mut socket_error_len = size_of_val(&socket_error) as libc::socklen_t; + if unsafe { + libc::getsockopt( + fd, + libc::SOL_SOCKET, + libc::SO_ERROR, + (&raw mut socket_error).cast(), + &raw mut socket_error_len, + ) + } < 0 + { + return Err(last_transport_error()); + } + if socket_error != 0 { + return Err(RerankError::Transport( + std::io::Error::from_raw_os_error(socket_error).to_string(), + )); + } + return Ok(()); + } +} + +#[cfg(unix)] +fn remaining_until(deadline: std::time::Instant) -> Result { + deadline + .checked_duration_since(std::time::Instant::now()) + .filter(|remaining| !remaining.is_zero()) + .ok_or_else(|| RerankError::Transport("request timed out".into())) +} + +#[cfg(unix)] +fn last_transport_error() -> RerankError { + RerankError::Transport(std::io::Error::last_os_error().to_string()) +} + +#[cfg(unix)] +fn write_interruptible( + stream: &mut std::os::unix::net::UnixStream, + mut bytes: &[u8], + deadline: std::time::Instant, +) -> Result<(), RerankError> { + use std::io::Write; + + while !bytes.is_empty() { + match stream.write(bytes) { + Ok(0) => return Err(RerankError::Transport("connection closed".into())), + Ok(count) => bytes = &bytes[count..], + Err(error) + if matches!( + error.kind(), + std::io::ErrorKind::Interrupted + | std::io::ErrorKind::WouldBlock + | std::io::ErrorKind::TimedOut + ) => {} + Err(error) => return Err(RerankError::Transport(error.to_string())), + } + pgrx::check_for_interrupts!(); + if std::time::Instant::now() >= deadline { + return Err(RerankError::Transport("request timed out".into())); + } + } + Ok(()) +} + +#[cfg(unix)] +fn read_interruptible( + stream: &mut std::os::unix::net::UnixStream, + mut bytes: &mut [u8], + deadline: std::time::Instant, +) -> Result<(), RerankError> { + use std::io::Read; + + while !bytes.is_empty() { + match stream.read(bytes) { + Ok(0) => return Err(RerankError::Transport("connection closed".into())), + Ok(count) => bytes = &mut bytes[count..], + Err(error) + if matches!( + error.kind(), + std::io::ErrorKind::Interrupted + | std::io::ErrorKind::WouldBlock + | std::io::ErrorKind::TimedOut + ) => {} + Err(error) => return Err(RerankError::Transport(error.to_string())), + } + pgrx::check_for_interrupts!(); + if std::time::Instant::now() >= deadline { + return Err(RerankError::Transport("request timed out".into())); + } + } + Ok(()) +} + +#[cfg(not(unix))] +impl TileMaxsimTransport for UnixSocketTransport { + fn round_trip( + &mut self, + _request: &[u8], + _timeout: Duration, + _max_response_bytes: usize, + ) -> Result, RerankError> { + if self.endpoint.is_empty() { + return Err(RerankError::Transport("endpoint is empty".into())); + } + Err(RerankError::Transport( + "Unix sockets are not supported on this platform".into(), + )) + } +} + +#[cfg(test)] +mod tests { + use super::super::external::{ + CandidateTensorDescriptorSource, ExternalTensorDescriptor, ExternalTensorDtype, + }; + use super::super::rerank::CandidateTensor; + use super::*; + use std::collections::BTreeMap; + use vector::vect::VectOwned; + + struct MockTensorSource(BTreeMap>); + + impl CandidateTensorSource for MockTensorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + Ok(Some(CandidateTensor { + candidate, + vectors: self + .0 + .remove(&candidate.heap_key) + .ok_or(RerankError::TensorMismatch)?, + })) + } + } + + struct MockDescriptorSource(BTreeMap); + + impl CandidateTensorDescriptorSource for MockDescriptorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + Ok(Some( + self.0 + .remove(&candidate.heap_key) + .ok_or(RerankError::TensorMismatch)?, + )) + } + } + + struct MockTransport { + similarities: Vec<(u32, f32)>, + } + + impl TileMaxsimTransport for MockTransport { + fn round_trip( + &mut self, + request: &[u8], + _timeout: Duration, + _max_response_bytes: usize, + ) -> Result, RerankError> { + let request_id = u64::from_le_bytes(request[8..16].try_into().unwrap()); + let version = u16::from_le_bytes(request[4..6].try_into().unwrap()); + Ok(success_response_with_version( + version, + request_id, + &self.similarities, + )) + } + } + + fn vector(values: &[f32]) -> OwnedVector { + OwnedVector::Vecf32(VectOwned::new(values.to_vec())) + } + + fn half_vector(values: &[f32]) -> OwnedVector { + OwnedVector::Vecf16(VectOwned::new( + values.iter().copied().map(simd::f16::from_f32).collect(), + )) + } + + fn success_response(request_id: u64, similarities: &[(u32, f32)]) -> Vec { + success_response_with_version(VERSION, request_id, similarities) + } + + fn success_response_with_version( + version: u16, + request_id: u64, + similarities: &[(u32, f32)], + ) -> Vec { + let body_len = 8 + similarities.len() * 8; + let mut response = Vec::with_capacity(HEADER_LEN + body_len); + response.extend_from_slice(MAGIC); + response.extend_from_slice(&version.to_le_bytes()); + response.extend_from_slice(&RESPONSE_KIND.to_le_bytes()); + response.extend_from_slice(&request_id.to_le_bytes()); + response.extend_from_slice(&(body_len as u64).to_le_bytes()); + response.extend_from_slice(&0u32.to_le_bytes()); + response.extend_from_slice(&(similarities.len() as u32).to_le_bytes()); + for (candidate_id, similarity) in similarities { + response.extend_from_slice(&candidate_id.to_le_bytes()); + response.extend_from_slice(&similarity.to_bits().to_le_bytes()); + } + response + } + + fn external_descriptor( + candidate: PageCandidate, + public_id: i64, + tensor_ref: &str, + rows: u32, + dimension: u32, + dtype: ExternalTensorDtype, + ) -> ExternalTensorDescriptor { + ExternalTensorDescriptor { + candidate, + public_id, + tensor_ref: tensor_ref.into(), + rows, + dimension, + dtype, + checksum: format!("sha256:{}", "a".repeat(64)), + } + } + + fn error_response(request_id: u64, message: &str) -> Vec { + let body_len = 8 + message.len(); + let mut response = Vec::with_capacity(HEADER_LEN + body_len); + response.extend_from_slice(MAGIC); + response.extend_from_slice(&VERSION.to_le_bytes()); + response.extend_from_slice(&RESPONSE_KIND.to_le_bytes()); + response.extend_from_slice(&request_id.to_le_bytes()); + response.extend_from_slice(&(body_len as u64).to_le_bytes()); + response.extend_from_slice(&1u32.to_le_bytes()); + response.extend_from_slice(&(message.len() as u32).to_le_bytes()); + response.extend_from_slice(message.as_bytes()); + response + } + + #[test] + fn gpu_backend_maps_positive_similarity_to_ascending_distance() { + let page_1 = [0, 0, 1]; + let page_2 = [0, 0, 2]; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![ + PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page_1, + }, + PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page_2, + }, + ] + .into_iter(); + let mut source = MockTensorSource(BTreeMap::from([ + (page_1, vec![vector(&[1.0, 0.0])]), + (page_2, vec![vector(&[0.5, 0.0])]), + ])); + let transport = MockTransport { + similarities: vec![(0, 1.0), (1, 2.0)], + }; + let results = GpuTileMaxsimBackend::new(transport, Duration::from_secs(1), 100, 4096) + .rerank(&query, &mut candidates, &mut source) + .unwrap() + .collect::>(); + + assert_eq!(results[0].heap_key, page_2); + assert_eq!(results[0].distance.to_f32(), -2.0); + assert_eq!(results[1].heap_key, page_1); + assert_eq!(results[1].distance.to_f32(), -1.0); + } + + #[test] + fn response_rejects_partial_and_duplicate_results() { + let keys = [[0, 0, 1], [0, 0, 2]]; + let partial = success_response(7, &[(0, 1.0)]); + assert!(matches!( + decode_response(&partial, 7, &keys), + Err(RerankError::Protocol(_)) + )); + + let duplicate = success_response(7, &[(0, 1.0), (0, 2.0)]); + assert!(matches!( + decode_response(&duplicate, 7, &keys), + Err(RerankError::Protocol(_)) + )); + } + + #[test] + fn response_ids_may_arrive_out_of_order() { + let keys = [[0, 0, 1], [0, 0, 2]]; + let response = success_response(7, &[(1, 0.5), (0, 1.0)]); + let results = decode_response(&response, 7, &keys) + .unwrap() + .collect::>(); + + assert_eq!(results[0].heap_key, keys[0]); + assert_eq!(results[0].distance.to_f32(), -1.0); + assert_eq!(results[1].heap_key, keys[1]); + assert_eq!(results[1].distance.to_f32(), -0.5); + } + + #[test] + fn response_rejects_unknown_non_finite_and_trailing_results() { + let keys = [[0, 0, 1], [0, 0, 2]]; + let unknown = success_response(7, &[(0, 1.0), (2, 2.0)]); + assert!(matches!( + decode_response(&unknown, 7, &keys), + Err(RerankError::Protocol(_)) + )); + + let non_finite = success_response(7, &[(0, f32::NAN), (1, 2.0)]); + assert!(matches!( + decode_response(&non_finite, 7, &keys), + Err(RerankError::Protocol(_)) + )); + + let mut trailing = success_response(7, &[(0, 1.0), (1, 2.0)]); + trailing.push(0); + assert!(matches!( + decode_response(&trailing, 7, &keys), + Err(RerankError::Protocol(_)) + )); + } + + #[test] + fn response_rejects_invalid_header_fields() { + let keys = [[0, 0, 1]]; + let valid = success_response(7, &[(0, 1.0)]); + + for (offset, replacement) in [(0, 0u8), (4, 2), (6, 1), (8, 8), (16, 0)] { + let mut invalid = valid.clone(); + invalid[offset] = replacement; + assert!(matches!( + decode_response(&invalid, 7, &keys), + Err(RerankError::Protocol(_)) + )); + } + } + + #[test] + fn response_surfaces_remote_error() { + let response = error_response(7, "CUDA queue is unavailable"); + assert!(matches!( + decode_response(&response, 7, &[]), + Err(RerankError::Remote(message)) if message == "CUDA queue is unavailable" + )); + } + + #[test] + fn request_frame_is_versioned_and_length_prefixed() { + let page = [0, 0, 1]; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }] + .into_iter(); + let mut source = MockTensorSource(BTreeMap::from([(page, vec![vector(&[0.5, 0.0])])])); + let encoded = encode_request(9, &query, &mut candidates, &mut source, 100, 4096).unwrap(); + + assert_eq!(&encoded.frame[0..4], MAGIC); + assert_eq!( + u16::from_le_bytes(encoded.frame[4..6].try_into().unwrap()), + VERSION + ); + assert_eq!( + u16::from_le_bytes(encoded.frame[6..8].try_into().unwrap()), + REQUEST_KIND + ); + assert_eq!( + u64::from_le_bytes(encoded.frame[8..16].try_into().unwrap()), + 9 + ); + assert_eq!( + u64::from_le_bytes(encoded.frame[16..24].try_into().unwrap()) as usize, + encoded.frame.len() - HEADER_LEN + ); + assert_eq!( + u32::from_le_bytes(encoded.frame[32..36].try_into().unwrap()), + 1 + ); + assert_eq!(encoded.heap_keys, vec![page]); + } + + #[test] + fn external_request_encodes_contract_and_opaque_descriptor_ids() { + let page = [0, 0, 7]; + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }; + let query = vec![half_vector(&[1.0, -0.5])]; + let tensor_ref = "s3://immutable/page-9001.tensor"; + let mut candidates = vec![candidate].into_iter(); + let mut source = MockDescriptorSource(BTreeMap::from([( + page, + external_descriptor(candidate, 9001, tensor_ref, 2, 2, ExternalTensorDtype::F16), + )])); + let contract = "colqwen@immutable-revision"; + let encoded = encode_external_request( + 19, + contract, + &query, + &mut candidates, + &mut source, + 100, + 4096, + ) + .unwrap(); + + assert_eq!(&encoded.frame[0..4], MAGIC); + assert_eq!( + u16::from_le_bytes(encoded.frame[4..6].try_into().unwrap()), + EXTERNAL_VERSION + ); + assert_eq!( + u32::from_le_bytes(encoded.frame[32..36].try_into().unwrap()), + 1 + ); + let contract_len = u32::from_le_bytes(encoded.frame[40..44].try_into().unwrap()) as usize; + assert_eq!(&encoded.frame[44..44 + contract_len], contract.as_bytes()); + let candidate_offset = 44 + contract_len + 4; // one 2-D f16 query row + assert_eq!( + u32::from_le_bytes( + encoded.frame[candidate_offset..candidate_offset + 4] + .try_into() + .unwrap() + ), + 0 + ); + assert_eq!( + u32::from_le_bytes( + encoded.frame[candidate_offset + 4..candidate_offset + 8] + .try_into() + .unwrap() + ), + 2 + ); + let reference_len = u32::from_le_bytes( + encoded.frame[candidate_offset + 8..candidate_offset + 12] + .try_into() + .unwrap(), + ) as usize; + let reference_offset = candidate_offset + 16; + assert_eq!( + &encoded.frame[reference_offset..reference_offset + reference_len], + tensor_ref.as_bytes() + ); + assert_eq!(encoded.heap_keys, vec![page]); + } + + #[test] + fn external_backend_maps_scores_without_exposing_public_ids() { + let page = [0, 0, 8]; + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![candidate].into_iter(); + let mut source = MockDescriptorSource(BTreeMap::from([( + page, + external_descriptor( + candidate, + i64::MAX, + "object://immutable/tensor", + 4, + 2, + ExternalTensorDtype::F32, + ), + )])); + let transport = MockTransport { + similarities: vec![(0, 3.5)], + }; + let results = GpuExternalTileMaxsimBackend::new( + transport, + "contract@1".into(), + Duration::from_secs(1), + 100, + 4096, + ) + .rerank(&query, &mut candidates, &mut source) + .unwrap() + .collect::>(); + + assert_eq!(results.len(), 1); + assert_eq!(results[0].heap_key, page); + assert_eq!(results[0].distance.to_f32(), -3.5); + } + + #[test] + fn external_request_rejects_shape_and_declared_payload_overflow() { + let page = [0, 0, 9]; + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }; + let query = vec![half_vector(&[1.0, 0.0])]; + let make_source = |dimension, rows| { + MockDescriptorSource(BTreeMap::from([( + page, + external_descriptor( + candidate, + 9, + "object://immutable/tensor", + rows, + dimension, + ExternalTensorDtype::F16, + ), + )])) + }; + + let mut candidates = vec![candidate].into_iter(); + assert!(matches!( + encode_external_request( + 1, + "contract@1", + &query, + &mut candidates, + &mut make_source(3, 1), + 100, + 4096, + ), + Err(RerankError::TensorMismatch) + )); + + let mut candidates = vec![candidate].into_iter(); + assert!(matches!( + encode_external_request( + 1, + "contract@1", + &query, + &mut candidates, + &mut make_source(2, 100), + 1000, + 64, + ), + Err(RerankError::RequestTooLarge) + )); + } + + #[test] + fn request_frame_encodes_f16_tensor_bits() { + let page = [0, 0, 1]; + let query = vec![half_vector(&[1.0, -0.5])]; + let mut candidates = vec![PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }] + .into_iter(); + let mut source = + MockTensorSource(BTreeMap::from([(page, vec![half_vector(&[0.25, 2.0])])])); + let encoded = encode_request(9, &query, &mut candidates, &mut source, 100, 4096).unwrap(); + + assert_eq!(encoded.frame[36], TensorDtype::F16 as u8); + assert_eq!( + u16::from_le_bytes(encoded.frame[40..42].try_into().unwrap()), + simd::f16::from_f32(1.0).to_bits() + ); + assert_eq!( + u16::from_le_bytes(encoded.frame[42..44].try_into().unwrap()), + simd::f16::from_f32(-0.5).to_bits() + ); + } + + #[test] + fn request_limits_are_enforced_before_transport() { + let page = [0, 0, 1]; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: page, + }] + .into_iter(); + let mut source = MockTensorSource(BTreeMap::from([(page, vec![vector(&[1.0, 0.0])])])); + assert!(matches!( + encode_request(1, &query, &mut candidates, &mut source, 1, 4096), + Err(RerankError::RequestTooLarge) + )); + } + + #[cfg(unix)] + #[test] + fn unix_socket_address_is_length_bounded_and_nul_terminated() { + let (address, length) = unix_socket_address("/tmp/vectorchord.sock").unwrap(); + let path_offset = std::mem::offset_of!(libc::sockaddr_un, sun_path); + + assert_eq!(address.sun_family, libc::AF_UNIX as libc::sa_family_t); + assert_eq!( + length as usize, + path_offset + "/tmp/vectorchord.sock".len() + 1 + ); + assert_eq!( + &address.sun_path[.."/tmp/vectorchord.sock".len()], + "/tmp/vectorchord.sock" + .as_bytes() + .iter() + .map(|byte| *byte as libc::c_char) + .collect::>() + ); + assert_eq!(address.sun_path["/tmp/vectorchord.sock".len()], 0); + + let too_long = "x".repeat(address.sun_path.len()); + assert!(matches!( + unix_socket_address(&too_long), + Err(RerankError::Transport(message)) if message == "endpoint path is too long" + )); + assert!(matches!( + unix_socket_address("invalid\0path"), + Err(RerankError::Transport(message)) if message == "endpoint contains a NUL byte" + )); + } +} diff --git a/src/index/vchordrq/scanners/maxsim/rerank.rs b/src/index/vchordrq/scanners/maxsim/rerank.rs new file mode 100644 index 00000000..151ae4ea --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/rerank.rs @@ -0,0 +1,315 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use super::candidate::{HeapKey, PageCandidate}; +use crate::index::fetcher::{Fetcher, FilterableTuple, Tuple}; +use crate::index::vchordrq::opclass::Opfamily; +use distance::Distance; +use std::cmp::Reverse; +use std::collections::BinaryHeap; +use vchordrq::types::OwnedVector; + +pub(super) struct CandidateTensor { + pub candidate: PageCandidate, + pub vectors: Vec, +} + +pub(super) trait CandidateTensorSource { + fn fetch(&mut self, candidate: PageCandidate) -> Result, RerankError>; +} + +pub(super) struct HeapArrayTensorSource<'a, F> { + fetcher: &'a mut F, + opfamily: Opfamily, +} + +impl<'a, F> HeapArrayTensorSource<'a, F> { + pub fn new(fetcher: &'a mut F, opfamily: Opfamily) -> Self { + Self { fetcher, opfamily } + } +} + +impl CandidateTensorSource for HeapArrayTensorSource<'_, F> { + fn fetch(&mut self, candidate: PageCandidate) -> Result, RerankError> { + let Some(mut tuple) = self.fetcher.fetch(candidate.heap_key) else { + return Ok(None); + }; + // Exact sources are the last boundary before a tensor may leave the + // executor process. Re-evaluate the active base-relation scan qual + // here even if token reranking already prefiltered some hits. This + // keeps same-relation quals ahead of CPU/GPU tensor access and + // also covers configurations with token-level refine disabled. + if !tuple.filter() { + return Ok(None); + } + let (values, is_nulls) = tuple.build(); + if is_nulls[0] { + return Err(RerankError::TensorMismatch); + } + let vectors = + unsafe { self.opfamily.input_vectors(values[0]) }.ok_or(RerankError::TensorMismatch)?; + Ok(Some(CandidateTensor { candidate, vectors })) + } +} + +pub(super) trait ExactMaxsimBackend { + type Results: Iterator; + + fn rerank( + &mut self, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + ) -> Result; +} + +#[derive(Debug)] +pub(super) enum RerankError { + TensorMismatch, + ModelContractMismatch, + InvalidDescriptor(&'static str), + Registry(String), + Configuration(&'static str), + UnsupportedTensorKind, + RequestTooLarge, + Transport(String), + Protocol(String), + Remote(String), +} + +impl std::fmt::Display for RerankError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::TensorMismatch => write!(f, "MaxSim tensor kind or dimension is not matched"), + Self::ModelContractMismatch => write!(f, "MaxSim model contract is not matched"), + Self::InvalidDescriptor(message) => { + write!(f, "invalid external MaxSim tensor descriptor: {message}") + } + Self::Registry(message) => { + write!(f, "MaxSim tensor-source registry error: {message}") + } + Self::Configuration(message) => write!(f, "MaxSim configuration error: {message}"), + Self::UnsupportedTensorKind => { + write!(f, "GPU MaxSim supports only vector and halfvec tensors") + } + Self::RequestTooLarge => write!(f, "GPU MaxSim request exceeds configured limits"), + Self::Transport(message) => write!(f, "GPU MaxSim transport error: {message}"), + Self::Protocol(message) => write!(f, "GPU MaxSim protocol error: {message}"), + Self::Remote(message) => write!(f, "GPU MaxSim sidecar error: {message}"), + } + } +} + +impl std::error::Error for RerankError {} + +#[derive(Default)] +pub(super) struct CpuExactMaxsimBackend; + +impl ExactMaxsimBackend for CpuExactMaxsimBackend { + type Results = RerankResults; + + fn rerank( + &mut self, + query: &[OwnedVector], + candidates: &mut dyn Iterator, + source: &mut S, + ) -> Result { + let mut results = BinaryHeap::new(); + for candidate in candidates { + let Some(tensor) = source.fetch(candidate)? else { + continue; + }; + let Some(distance) = exact_maxsim_distance(query, &tensor.vectors) else { + return Err(RerankError::TensorMismatch); + }; + results.push((Reverse(distance), Reverse(tensor.candidate.heap_key))); + } + Ok(RerankResults { inner: results }) + } +} + +fn exact_maxsim_distance(query: &[OwnedVector], document: &[OwnedVector]) -> Option { + if query.is_empty() || document.is_empty() { + return None; + } + let mut maxsim = 0.0f32; + for query_vector in query { + let mut best = Distance::INFINITY; + for document_vector in document { + let distance = document_vector.operator_dot(query_vector)?; + best = std::cmp::min(best, distance); + } + maxsim += best.to_f32(); + } + Some(Distance::from_f32(maxsim)) +} + +#[derive(Clone, Copy, Debug)] +pub(super) struct RerankedPage { + pub distance: Distance, + pub heap_key: HeapKey, +} + +pub(super) struct RerankResults { + pub(super) inner: BinaryHeap<(Reverse, Reverse)>, +} + +impl Iterator for RerankResults { + type Item = RerankedPage; + + fn next(&mut self) -> Option { + let (Reverse(distance), Reverse(heap_key)) = self.inner.pop()?; + Some(RerankedPage { distance, heap_key }) + } + + fn size_hint(&self) -> (usize, Option) { + let exact = self.inner.len(); + (exact, Some(exact)) + } +} + +impl ExactSizeIterator for RerankResults {} +impl std::iter::FusedIterator for RerankResults {} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::BTreeMap; + use vector::vect::VectOwned; + + struct MockTensorSource(BTreeMap>); + + impl CandidateTensorSource for MockTensorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + Ok(Some(CandidateTensor { + candidate, + vectors: self + .0 + .remove(&candidate.heap_key) + .ok_or(RerankError::TensorMismatch)?, + })) + } + } + + fn vector(values: &[f32]) -> OwnedVector { + OwnedVector::Vecf32(VectOwned::new(values.to_vec())) + } + + struct RejectingFetcher; + + struct RejectingTuple; + + impl Tuple for RejectingTuple { + fn build(&mut self) -> (&[pgrx::pg_sys::Datum; 32], &[bool; 32]) { + panic!("a rejected tuple must not materialize its tensor") + } + + fn attribute(&mut self, _attnum: i16) -> Option { + panic!("a rejected tuple must not expose heap attributes") + } + } + + impl FilterableTuple for RejectingTuple { + fn filter(&mut self) -> bool { + false + } + } + + impl Fetcher for RejectingFetcher { + type Tuple<'a> = RejectingTuple; + + fn fetch(&mut self, _key: HeapKey) -> Option> { + Some(RejectingTuple) + } + } + + #[test] + fn heap_source_applies_scan_qual_before_materializing_tensor() { + let mut fetcher = RejectingFetcher; + let mut source = HeapArrayTensorSource::new(&mut fetcher, Opfamily::VectorMaxsim); + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: [0, 0, 1], + }; + + assert!(source.fetch(candidate).unwrap().is_none()); + } + + #[test] + fn cpu_backend_reorders_candidates_by_exact_page_maxsim() { + let page_1 = [0, 0, 1]; + let page_2 = [0, 0, 2]; + let query = vec![vector(&[1.0, 0.0]), vector(&[0.0, 1.0])]; + let mut candidates = vec![ + PageCandidate { + approximate_distance: Distance::from_f32(-2.0), + heap_key: page_2, + }, + PageCandidate { + approximate_distance: Distance::from_f32(-1.0), + heap_key: page_1, + }, + ] + .into_iter(); + let mut source = MockTensorSource(BTreeMap::from([ + (page_1, vec![vector(&[1.0, 0.0]), vector(&[0.0, 1.0])]), + (page_2, vec![vector(&[0.5, 0.5])]), + ])); + let results = CpuExactMaxsimBackend + .rerank(&query, &mut candidates, &mut source) + .unwrap() + .collect::>(); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].heap_key, page_1); + assert_eq!(results[0].distance.to_f32(), -2.0); + assert_eq!(results[1].heap_key, page_2); + assert_eq!(results[1].distance.to_f32(), -1.0); + } + + #[test] + fn exact_ties_are_ordered_by_heap_key() { + let page_1 = [0, 0, 1]; + let page_2 = [0, 0, 2]; + let query = vec![vector(&[1.0, 0.0])]; + let mut candidates = vec![ + PageCandidate { + approximate_distance: Distance::from_f32(0.0), + heap_key: page_2, + }, + PageCandidate { + approximate_distance: Distance::from_f32(0.0), + heap_key: page_1, + }, + ] + .into_iter(); + let mut source = MockTensorSource(BTreeMap::from([ + (page_1, vec![vector(&[1.0, 0.0])]), + (page_2, vec![vector(&[1.0, 0.0])]), + ])); + + let results = CpuExactMaxsimBackend + .rerank(&query, &mut candidates, &mut source) + .unwrap() + .collect::>(); + + assert_eq!( + results.iter().map(|page| page.heap_key).collect::>(), + vec![page_1, page_2] + ); + } +} diff --git a/src/index/vchordrq/scanners/maxsim/search.rs b/src/index/vchordrq/scanners/maxsim/search.rs new file mode 100644 index 00000000..915fb87f --- /dev/null +++ b/src/index/vchordrq/scanners/maxsim/search.rs @@ -0,0 +1,578 @@ +// This software is licensed under a dual license model: +// +// GNU Affero General Public License v3 (AGPLv3): You may use, modify, and +// distribute this software under the terms of the AGPLv3. +// +// Elastic License v2 (ELv2): You may also use, modify, and distribute this +// software under the Elastic License v2, which has specific restrictions. +// +// We welcome any commercial collaboration or support. For inquiries +// regarding the licenses, please contact us at: +// vectorchord-inquiry@tensorchord.ai +// +// Copyright (c) 2025-2026 TensorChord Inc. + +use super::MaxsimBuilder; +use super::candidate::{HeapKey, PageCandidate}; +use super::external::{ + CandidateTensorDescriptorSource, ExternalTensorDescriptor, ExternalTensorSourceBinding, + ExternalTensorStorage, resolve_external_tensor_source, validate_descriptor, +}; +use super::gpu::{GpuExternalTileMaxsimBackend, UnixSocketTransport}; +use super::rerank::RerankError; +use crate::index::fetcher::{ + Fetcher, FilterableTuple, HeapFetcher, Tuple, TupleAttribute, ctid_to_key, +}; +use crate::index::gucs::{self, PostgresMaxsimBackend}; +use crate::index::scanners::SearchBuilder; +use crate::index::storage::PostgresRelation; +use crate::index::vchordrq::opclass::{Opfamily, opfamily}; +use crate::index::vchordrq::scanners::SearchOptions; +use crate::recorder::DefaultRecorder; +use distance::Distance; +use pgrx::datum::{DatumWithOid, FromDatum}; +use pgrx::iter::TableIterator; +use pgrx::{AnyArray, IntoDatum, name}; +use std::collections::{BTreeMap, BTreeSet}; +use std::time::Duration; + +const MAX_EXPLICIT_CANDIDATES: i32 = 65_536; + +/// Restricted Phase 3B search surface. Candidate generation reads only the +/// named index. Descriptor projection happens later through SPI, under the +/// caller's normal SELECT privileges and active MVCC snapshot. +#[pgrx::pg_extern(sql = "")] +fn _vchordrq_maxsim_search_external( + index_oid: pgrx::pg_sys::Oid, + query: AnyArray, + candidate_limit: i32, + top_k: i32, +) -> TableIterator<'static, (name!(public_id, i64), name!(similarity, f32))> { + let rows = execute_external_search(index_oid, query, candidate_limit, top_k) + .unwrap_or_else(|error| pgrx::error!("{error}")); + TableIterator::new(rows) +} + +fn execute_external_search( + index_oid: pgrx::pg_sys::Oid, + query: AnyArray, + candidate_limit: i32, + top_k: i32, +) -> Result, RerankError> { + validate_search_limits(candidate_limit, top_k)?; + if !matches!(gucs::vchordrq_maxsim_backend(), PostgresMaxsimBackend::Gpu) { + return Err(RerankError::Configuration( + "external MaxSim search currently requires vchordrq.maxsim_backend = 'gpu'", + )); + } + + let binding = resolve_external_tensor_source(index_oid)?; + let index_lock = RelationLock::open(index_oid, pgrx::pg_sys::AccessShareLock as _)?; + let heap_lock = RelationLock::open(binding.heap_oid, pgrx::pg_sys::AccessShareLock as _)?; + let descriptor_lock = binding + .descriptor_oid + .map(|oid| RelationLock::open(oid, pgrx::pg_sys::AccessShareLock as _)) + .transpose()?; + if binding.index_oid != index_lock.oid() || binding.heap_oid != heap_lock.oid() { + return Err(RerankError::Registry( + "registered MaxSim tensor source changed during execution".into(), + )); + } + if descriptor_lock.as_ref().map(RelationLock::oid) != binding.descriptor_oid { + return Err(RerankError::Registry( + "registered descriptor relation changed during execution".into(), + )); + } + + let opfamily = unsafe { opfamily(index_lock.raw()) }; + if !matches!(opfamily, Opfamily::VectorMaxsim | Opfamily::HalfvecMaxsim) { + return Err(RerankError::UnsupportedTensorKind); + } + let indexed_type = unsafe { pgrx::pg_sys::get_atttype(index_oid, 1) }; + if indexed_type != query.oid() { + return Err(RerankError::TensorMismatch); + } + let query_vectors = + unsafe { opfamily.input_vectors(query.datum()) }.ok_or(RerankError::TensorMismatch)?; + + // The registry resolution happens before the index read, and this + // privilege-only SELECT is planned/executed before any candidate CTID is + // generated. It fails early when the caller cannot project the registered + // descriptor columns. The same query shape is used for the actual fetch. + preflight_descriptor_access(&binding)?; + + let candidates = generate_candidates( + index_lock.raw(), + opfamily, + query.datum(), + candidate_limit as u32, + )?; + let resolved = + resolve_visible_candidates(index_lock.raw(), heap_lock.raw(), candidates.into_iter())?; + let (mut source, public_ids) = load_visible_descriptors(&binding, &resolved)?; + let visible_candidates = resolved + .into_iter() + .filter_map(|resolved| { + public_ids + .contains_key(&resolved.candidate.heap_key) + .then_some(resolved.candidate) + }) + .collect::>(); + let endpoint = gucs::vchordrq_maxsim_gpu_endpoint() + .map(|endpoint| endpoint.to_string_lossy().into_owned()) + .unwrap_or_default(); + let transport = UnixSocketTransport::new(endpoint); + let mut backend = GpuExternalTileMaxsimBackend::new( + transport, + binding.model_contract_id, + Duration::from_millis(gucs::vchordrq_maxsim_gpu_timeout_ms() as u64), + gucs::vchordrq_maxsim_gpu_max_batch_tokens() as usize, + gucs::vchordrq_maxsim_gpu_max_batch_bytes() as usize, + ); + let mut candidate_iter = visible_candidates.into_iter(); + let exact = backend.rerank(&query_vectors, &mut candidate_iter, &mut source)?; + let mut rows = exact + .map(|result| { + let public_id = public_ids.get(&result.heap_key).copied().ok_or_else(|| { + RerankError::Protocol("sidecar result has no visible public ID".into()) + })?; + Ok((result.distance, public_id)) + }) + .collect::, RerankError>>()?; + rows.sort_unstable_by(|(left_distance, left_id), (right_distance, right_id)| { + left_distance + .cmp(right_distance) + .then_with(|| left_id.cmp(right_id)) + }); + rows.truncate(top_k as usize); + Ok(rows + .into_iter() + .map(|(distance, public_id)| (public_id, -distance.to_f32())) + .collect::>() + .into_iter()) +} + +#[derive(Clone, Copy)] +struct ResolvedCandidate { + candidate: PageCandidate, + current_ctid: pgrx::pg_sys::ItemPointerData, +} + +fn resolve_visible_candidates( + index_relation: pgrx::pg_sys::Relation, + heap_relation: pgrx::pg_sys::Relation, + candidates: impl Iterator, +) -> Result, RerankError> { + let snapshot = unsafe { pgrx::pg_sys::GetActiveSnapshot() }; + if snapshot.is_null() { + return Err(RerankError::Configuration( + "external MaxSim search requires an active MVCC snapshot", + )); + } + if unsafe { (*snapshot).snapshot_type } != pgrx::pg_sys::SnapshotType::SNAPSHOT_MVCC { + return Err(RerankError::Configuration( + "external MaxSim search requires an MVCC snapshot", + )); + } + + let mut fetcher = + unsafe { HeapFetcher::new_standalone(index_relation, heap_relation, snapshot) }; + let mut resolved = Vec::new(); + for candidate in candidates { + pgrx::check_for_interrupts!(); + if let Some(tuple) = fetcher.fetch(candidate.heap_key) { + resolved.push(ResolvedCandidate { + candidate, + current_ctid: tuple.ctid(), + }); + } + } + Ok(resolved) +} + +fn validate_search_limits(candidate_limit: i32, top_k: i32) -> Result<(), RerankError> { + if !(1..=MAX_EXPLICIT_CANDIDATES).contains(&candidate_limit) { + return Err(RerankError::Configuration( + "candidate_limit must be between 1 and 65536", + )); + } + if top_k <= 0 || top_k > candidate_limit { + return Err(RerankError::Configuration( + "top_k must be positive and no greater than candidate_limit", + )); + } + Ok(()) +} + +fn generate_candidates( + index_relation: pgrx::pg_sys::Relation, + opfamily: Opfamily, + query: pgrx::pg_sys::Datum, + candidate_limit: u32, +) -> Result, RerankError> { + let index = unsafe { PostgresRelation::::new(index_relation) }; + let mut builder = MaxsimBuilder::new(opfamily); + unsafe { builder.add(3, Some(query)) }; + let options = SearchOptions { + epsilon: unsafe { gucs::vchordrq_epsilon(index_relation) }, + probes: unsafe { gucs::vchordrq_probes(index_relation) }, + max_scan_tuples: None, + maxsim_refine: gucs::vchordrq_maxsim_refine(index_relation), + maxsim_threshold: gucs::vchordrq_maxsim_threshold(index_relation), + maxsim_candidate_limit: Some(candidate_limit), + maxsim_backend: PostgresMaxsimBackend::CoarseOnly, + maxsim_gpu_endpoint: None, + maxsim_gpu_timeout_ms: 1, + maxsim_gpu_max_batch_tokens: 1, + maxsim_gpu_max_batch_bytes: 1, + io_search: gucs::vchordrq_io_search(), + io_rerank: gucs::vchordrq_io_rerank(), + // General same-relation quals would require the optional Phase 3C + // CustomScan. The restricted function applies PostgreSQL row + // visibility during the following SPI descriptor fetch. + prefilter: false, + }; + let bump = bumpalo::Bump::new(); + let recorder = DefaultRecorder { + enable: false, + rate: None, + max_records: 0, + index: unsafe { (*index_relation).rd_id.to_u32() }, + }; + let candidates = builder + .build(&index, options, NoHeapFetcher, &bump, recorder) + .map(|(distance, heap_key, _)| PageCandidate { + approximate_distance: Distance::from_f32(distance), + heap_key, + }) + .collect(); + Ok(candidates) +} + +fn preflight_descriptor_access(binding: &ExternalTensorSourceBinding) -> Result<(), RerankError> { + let query = descriptor_query(binding, true)?; + pgrx::spi::Spi::connect(|client| { + let privilege = client + .prepare( + "SELECT pg_catalog.has_table_privilege($1, 'SELECT') AS allowed", + &pgrx::oids_of![pgrx::pg_sys::Oid], + ) + .map_err(registry_error)?; + for relation_oid in std::iter::once(binding.heap_oid).chain(binding.descriptor_oid) { + let allowed = client + .select(&privilege, Some(1), &[relation_oid.into()]) + .map_err(registry_error)? + .first() + .get_by_name::("allowed") + .map_err(registry_error)? + .unwrap_or(false); + if !allowed { + return Err(RerankError::Registry( + "table-level SELECT privilege is required for every external MaxSim relation" + .into(), + )); + } + } + client + .select(query.as_str(), Some(1), &[]) + .map(|_| ()) + .map_err(registry_error) + }) +} + +fn load_visible_descriptors( + binding: &ExternalTensorSourceBinding, + candidates: &[ResolvedCandidate], +) -> Result<(MaterializedDescriptorSource, BTreeMap), RerankError> { + if candidates.is_empty() { + return Ok((MaterializedDescriptorSource::default(), BTreeMap::new())); + } + let mut candidates_by_key = BTreeMap::new(); + for resolved in candidates { + if candidates_by_key + .insert(ctid_to_key(resolved.current_ctid), resolved.candidate) + .is_some() + { + return Err(RerankError::Registry( + "multiple index candidates resolved to the same visible CTID".into(), + )); + } + } + let ctids = candidates + .iter() + .map(|resolved| resolved.current_ctid) + .collect::>(); + let query = descriptor_query(binding, false)?; + let (descriptors, public_ids) = pgrx::spi::Spi::connect(|client| { + let tid_array_oid = unsafe { pgrx::pg_sys::get_array_type(pgrx::pg_sys::TIDOID) }; + let prepared = client + .prepare( + query.as_str(), + &[ + pgrx::pg_sys::PgOid::from(tid_array_oid), + pgrx::pg_sys::PgOid::from(pgrx::pg_sys::TEXTOID), + ], + ) + .map_err(registry_error)?; + let args: [DatumWithOid<'_>; 2] = [ctids.into(), binding.model_contract_id.clone().into()]; + let rows = client + .select(&prepared, Some(candidates.len() as _), &args) + .map_err(registry_error)?; + let mut descriptors = BTreeMap::new(); + let mut public_ids = BTreeMap::new(); + let mut unique_public_ids = BTreeSet::new(); + for row in rows { + let ctid = required_heap_column::(&row, "heap_tid")?; + let heap_key = ctid_to_key(ctid); + let candidate = candidates_by_key.get(&heap_key).copied().ok_or_else(|| { + RerankError::Registry("descriptor query returned an unknown CTID".into()) + })?; + let public_id = required_heap_column::(&row, "public_id")?; + if !unique_public_ids.insert(public_id) { + return Err(RerankError::InvalidDescriptor( + "public IDs are not unique in the visible candidate batch", + )); + } + let descriptor = validate_descriptor( + candidate, + public_id, + required_heap_column::(&row, "tensor_ref")?, + required_heap_column::(&row, "tensor_rows")?, + required_heap_column::(&row, "tensor_dimension")?, + required_heap_column::(&row, "tensor_dtype")?, + required_heap_column::(&row, "tensor_checksum")?, + )?; + if descriptors.insert(candidate.heap_key, descriptor).is_some() { + return Err(RerankError::Registry( + "descriptor query returned a duplicate CTID".into(), + )); + } + public_ids.insert(candidate.heap_key, public_id); + } + Ok((descriptors, public_ids)) + })?; + Ok((MaterializedDescriptorSource(descriptors), public_ids)) +} + +fn descriptor_query( + binding: &ExternalTensorSourceBinding, + preflight: bool, +) -> Result { + let heap_relation = relation_name(binding.heap_oid)?; + let names = &binding.column_names; + let model_contract = pgrx::spi::quote_identifier(&names.model_contract); + let public_id = pgrx::spi::quote_identifier(&names.public_id); + let tensor_ref = pgrx::spi::quote_identifier(&names.tensor_ref); + let tensor_rows = pgrx::spi::quote_identifier(&names.tensor_rows); + let tensor_dimension = pgrx::spi::quote_identifier(&names.tensor_dimension); + let tensor_dtype = pgrx::spi::quote_identifier(&names.tensor_dtype); + let tensor_checksum = pgrx::spi::quote_identifier(&names.tensor_checksum); + let predicate = if preflight { + "false".to_string() + } else { + format!("h.ctid = ANY($1) AND h.{model_contract} = $2") + }; + match binding.storage { + ExternalTensorStorage::SameHeap => Ok(format!( + "SELECT h.ctid AS heap_tid, + h.{model_contract} AS model_contract, + h.{public_id} AS public_id, + h.{tensor_ref} AS tensor_ref, + h.{tensor_rows} AS tensor_rows, + h.{tensor_dimension} AS tensor_dimension, + h.{tensor_dtype} AS tensor_dtype, + h.{tensor_checksum} AS tensor_checksum + FROM ONLY {heap_relation} AS h + WHERE {predicate}" + )), + ExternalTensorStorage::DescriptorRelation => { + let descriptor_oid = binding.descriptor_oid.ok_or_else(|| { + RerankError::Registry("registered descriptor relation is missing".into()) + })?; + let descriptor_relation = relation_name(descriptor_oid)?; + let descriptor_public_id = pgrx::spi::quote_identifier( + names.descriptor_public_id.as_deref().ok_or_else(|| { + RerankError::Registry( + "registered descriptor public ID column is missing".into(), + ) + })?, + ); + Ok(format!( + "SELECT h.ctid AS heap_tid, + h.{model_contract} AS model_contract, + h.{public_id} AS public_id, + d.{tensor_ref} AS tensor_ref, + d.{tensor_rows} AS tensor_rows, + d.{tensor_dimension} AS tensor_dimension, + d.{tensor_dtype} AS tensor_dtype, + d.{tensor_checksum} AS tensor_checksum + FROM ONLY {heap_relation} AS h + LEFT JOIN ONLY {descriptor_relation} AS d + ON d.{descriptor_public_id} = h.{public_id} + WHERE {predicate}" + )) + } + } +} + +fn relation_name(relation_oid: pgrx::pg_sys::Oid) -> Result { + pgrx::spi::Spi::connect(|client| { + let prepared = client + .prepare( + "SELECT n.nspname::text AS schema_name, c.relname::text AS relation_name + FROM pg_catalog.pg_class AS c + JOIN pg_catalog.pg_namespace AS n ON n.oid = c.relnamespace + WHERE c.oid = $1", + &pgrx::oids_of![pgrx::pg_sys::Oid], + ) + .map_err(registry_error)?; + let rows = client + .select(&prepared, Some(1), &[relation_oid.into()]) + .map_err(registry_error)?; + if rows.is_empty() { + return Err(RerankError::Registry( + "registered relation disappeared".into(), + )); + } + let row = rows.first(); + Ok(pgrx::spi::quote_qualified_identifier( + required_column::(&row, "schema_name")?, + required_column::(&row, "relation_name")?, + )) + }) +} + +fn required_column(row: &pgrx::spi::SpiTupleTable<'_>, name: &str) -> Result +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name) + .map_err(registry_error)? + .ok_or(RerankError::InvalidDescriptor( + "descriptor query returned NULL", + )) +} + +fn required_heap_column( + row: &pgrx::spi::SpiHeapTupleData<'_>, + name: &str, +) -> Result +where + T: FromDatum + IntoDatum, +{ + row.get_by_name::(name) + .map_err(registry_error)? + .ok_or(RerankError::InvalidDescriptor( + "descriptor query returned NULL", + )) +} + +fn registry_error(error: impl std::fmt::Display) -> RerankError { + RerankError::Registry(error.to_string()) +} + +#[derive(Default)] +struct MaterializedDescriptorSource(BTreeMap); + +impl CandidateTensorDescriptorSource for MaterializedDescriptorSource { + fn fetch( + &mut self, + candidate: PageCandidate, + ) -> Result, RerankError> { + Ok(self.0.remove(&candidate.heap_key)) + } +} + +struct RelationLock { + raw: pgrx::pg_sys::Relation, + lockmode: pgrx::pg_sys::LOCKMODE, +} + +impl RelationLock { + fn open(oid: pgrx::pg_sys::Oid, lockmode: pgrx::pg_sys::LOCKMODE) -> Result { + let raw = unsafe { pgrx::pg_sys::relation_open(oid, lockmode) }; + if raw.is_null() { + return Err(RerankError::Registry("relation open returned NULL".into())); + } + Ok(Self { raw, lockmode }) + } + + fn raw(&self) -> pgrx::pg_sys::Relation { + self.raw + } + + fn oid(&self) -> pgrx::pg_sys::Oid { + unsafe { (*self.raw).rd_id } + } +} + +impl Drop for RelationLock { + fn drop(&mut self) { + unsafe { pgrx::pg_sys::relation_close(self.raw, self.lockmode) }; + } +} + +struct NoHeapFetcher; +struct NoHeapTuple; + +impl Tuple for NoHeapTuple { + fn build(&mut self) -> (&[pgrx::pg_sys::Datum; 32], &[bool; 32]) { + unreachable!("restricted external candidate generation must not read the heap") + } + + fn attribute(&mut self, _attnum: i16) -> Option { + unreachable!("restricted external candidate generation must not read the heap") + } +} + +impl FilterableTuple for NoHeapTuple { + fn filter(&mut self) -> bool { + unreachable!("restricted external candidate generation must not read the heap") + } +} + +impl Fetcher for NoHeapFetcher { + type Tuple<'a> = NoHeapTuple; + + fn fetch(&mut self, _key: HeapKey) -> Option> { + unreachable!("restricted external candidate generation must not read the heap") + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn explicit_search_limits_are_bounded() { + assert!(validate_search_limits(256, 10).is_ok()); + for (candidates, top_k) in [(0, 1), (65_537, 1), (1, 0), (10, 11)] { + assert!(matches!( + validate_search_limits(candidates, top_k), + Err(RerankError::Configuration(_)) + )); + } + } + + #[test] + fn materialized_source_only_returns_visible_descriptors_once() { + let candidate = PageCandidate { + approximate_distance: Distance::ZERO, + heap_key: [0, 0, 1], + }; + let descriptor = validate_descriptor( + candidate, + 42, + "s3://immutable/tensor".into(), + 1, + 2, + "float16".into(), + format!("sha256:{}", "a".repeat(64)), + ) + .unwrap(); + let mut source = + MaterializedDescriptorSource(BTreeMap::from([(candidate.heap_key, descriptor)])); + assert_eq!(source.fetch(candidate).unwrap().unwrap().public_id, 42); + assert!(source.fetch(candidate).unwrap().is_none()); + } +} diff --git a/src/index/vchordrq/scanners/mod.rs b/src/index/vchordrq/scanners/mod.rs index b345da37..c5e392bc 100644 --- a/src/index/vchordrq/scanners/mod.rs +++ b/src/index/vchordrq/scanners/mod.rs @@ -15,7 +15,9 @@ mod default; mod maxsim; +use crate::index::gucs::PostgresMaxsimBackend; use crate::index::scanners::Io; +use std::ffi::CString; pub use default::DefaultBuilder; pub use maxsim::MaxsimBuilder; @@ -27,6 +29,12 @@ pub struct SearchOptions { pub max_scan_tuples: Option, pub maxsim_refine: u32, pub maxsim_threshold: u32, + pub maxsim_candidate_limit: Option, + pub maxsim_backend: PostgresMaxsimBackend, + pub maxsim_gpu_endpoint: Option, + pub maxsim_gpu_timeout_ms: u32, + pub maxsim_gpu_max_batch_tokens: u32, + pub maxsim_gpu_max_batch_bytes: u32, pub io_search: Io, pub io_rerank: Io, pub prefilter: bool, diff --git a/src/sql/finalize.sql b/src/sql/finalize.sql index a4f8f38a..a7e7af4e 100644 --- a/src/sql/finalize.sql +++ b/src/sql/finalize.sql @@ -610,3 +610,824 @@ FROM WHERE am.amname = 'vchordrq' ) AS index_oids CROSS JOIN LATERAL vchordrq_sampled_queries(index_oids.oid::regclass) AS record; + +-- Phase 3B tensor-source bindings + +CREATE TABLE _vchordrq_maxsim_sources ( + index_oid oid PRIMARY KEY, + heap_oid oid NOT NULL, + model_contract_id text NOT NULL + CHECK ( + model_contract_id OPERATOR(pg_catalog.<>) ''::text + AND pg_catalog.length(model_contract_id) OPERATOR(pg_catalog.<=) 512 + ), + storage text NOT NULL CHECK ( + storage OPERATOR(pg_catalog.=) ANY ( + ARRAY['heap_array', 'external_ref', 'external_relation']::text[] + ) + ), + model_contract_attnum smallint NOT NULL CHECK ( + model_contract_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + public_id_attnum smallint NOT NULL CHECK ( + public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + ), + descriptor_oid oid, + descriptor_public_id_attnum smallint, + tensor_ref_attnum smallint, + tensor_rows_attnum smallint, + tensor_dim_attnum smallint, + tensor_dtype_attnum smallint, + tensor_checksum_attnum smallint, + registered_by oid NOT NULL, + registered_at timestamptz NOT NULL DEFAULT pg_catalog.clock_timestamp(), + CHECK ( + ( + storage OPERATOR(pg_catalog.=) 'heap_array'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + AND tensor_ref_attnum IS NULL + AND tensor_rows_attnum IS NULL + AND tensor_dim_attnum IS NULL + AND tensor_dtype_attnum IS NULL + AND tensor_checksum_attnum IS NULL + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_ref'::text + AND descriptor_oid IS NULL + AND descriptor_public_id_attnum IS NULL + AND tensor_ref_attnum IS NOT NULL + AND tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_rows_attnum IS NOT NULL + AND tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dim_attnum IS NOT NULL + AND tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dtype_attnum IS NOT NULL + AND tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_checksum_attnum IS NOT NULL + AND tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + OR + ( + storage OPERATOR(pg_catalog.=) 'external_relation'::text + AND descriptor_oid IS NOT NULL + AND descriptor_public_id_attnum IS NOT NULL + AND descriptor_public_id_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_ref_attnum IS NOT NULL + AND tensor_ref_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_rows_attnum IS NOT NULL + AND tensor_rows_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dim_attnum IS NOT NULL + AND tensor_dim_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_dtype_attnum IS NOT NULL + AND tensor_dtype_attnum OPERATOR(pg_catalog.>) 0::smallint + AND tensor_checksum_attnum IS NOT NULL + AND tensor_checksum_attnum OPERATOR(pg_catalog.>) 0::smallint + ) + ) +); + +REVOKE ALL ON TABLE _vchordrq_maxsim_sources FROM PUBLIC; + +CREATE FUNCTION vchordrq_register_maxsim_source( + index_relation regclass, + model_contract_id text, + storage text, + model_contract_column name, + public_id_column name, + tensor_ref_column name DEFAULT NULL, + tensor_rows_column name DEFAULT NULL, + tensor_dim_column name DEFAULT NULL, + tensor_dtype_column name DEFAULT NULL, + tensor_checksum_column name DEFAULT NULL, + descriptor_relation regclass DEFAULT NULL, + descriptor_public_id_column name DEFAULT NULL +) +RETURNS void +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + heap_oid oid; + index_owner oid; + descriptor_oid oid; + descriptor_owner oid; + tensor_relation_oid oid; + normalized_storage text; + attnum smallint; + atttypid oid; + attnotnull boolean; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + descriptor_id_is_unique boolean; +BEGIN + IF index_relation IS NULL THEN + RAISE EXCEPTION 'index_relation must not be NULL'; + END IF; + IF model_contract_id IS NULL + OR btrim(model_contract_id) = '' + OR length(model_contract_id) > 512 THEN + RAISE EXCEPTION 'model_contract_id must contain between 1 and 512 characters'; + END IF; + model_contract_id := btrim(model_contract_id); + IF model_contract_column IS NULL OR public_id_column IS NULL THEN + RAISE EXCEPTION 'model_contract_column and public_id_column must not be NULL'; + END IF; + + normalized_storage := lower(btrim(storage)); + IF normalized_storage IS NULL + OR normalized_storage NOT IN ('heap_array', 'external_ref', 'external_relation') THEN + RAISE EXCEPTION 'storage must be heap_array, external_ref, or external_relation'; + END IF; + + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + IF caller_oid IS NULL THEN + RAISE EXCEPTION 'could not resolve caller role'; + END IF; + + SELECT x.indrelid, i.relowner + INTO heap_oid, index_owner + FROM pg_catalog.pg_index AS x + JOIN pg_catalog.pg_class AS i ON i.oid = x.indexrelid + JOIN pg_catalog.pg_class AS h ON h.oid = x.indrelid + JOIN pg_catalog.pg_am AS am ON am.oid = i.relam + JOIN pg_catalog.pg_opclass AS opc ON opc.oid = x.indclass[0] + WHERE x.indexrelid = index_relation::oid + AND i.relkind = 'i' + AND h.relkind IN ('r', 'm') + AND am.amname = 'vchordrq' + AND opc.opcname IN ( + 'vector_maxsim_ops', + 'halfvec_maxsim_ops', + 'rabitq8_maxsim_ops', + 'rabitq4_maxsim_ops' + ) + AND x.indisvalid + AND x.indisready + AND x.indnatts = 1 + AND x.indnkeyatts = 1; + IF heap_oid IS NULL THEN + RAISE EXCEPTION 'relation % is not a valid single-key vchordrq MaxSim index', + index_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the index owner may register its MaxSim tensor source'; + END IF; + + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = heap_oid + AND a.attname = model_contract_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'model contract column % must be a NOT NULL text column', + model_contract_column; + END IF; + model_contract_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = heap_oid + AND a.attname = public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'bigint'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'public ID column % must be a NOT NULL bigint column', + public_id_column; + END IF; + public_id_attnum := attnum; + + IF normalized_storage = 'external_relation' THEN + IF descriptor_relation IS NULL OR descriptor_public_id_column IS NULL THEN + RAISE EXCEPTION 'external_relation sources require descriptor_relation and descriptor_public_id_column'; + END IF; + SELECT c.oid, c.relowner + INTO descriptor_oid, descriptor_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = descriptor_relation::oid + AND c.relkind IN ('r', 'm'); + IF descriptor_oid IS NULL THEN + RAISE EXCEPTION 'descriptor relation % must be a table or materialized view', + descriptor_relation; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, descriptor_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the descriptor relation owner may register it as a MaxSim tensor source'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = descriptor_oid + AND a.attname = descriptor_public_id_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'bigint'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'descriptor public ID column % must be a NOT NULL bigint column', + descriptor_public_id_column; + END IF; + descriptor_public_id_attnum := attnum; + + SELECT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL + ) INTO descriptor_id_is_unique; + IF NOT descriptor_id_is_unique THEN + RAISE EXCEPTION 'descriptor public ID column % must have a non-partial single-key unique index', + descriptor_public_id_column; + END IF; + tensor_relation_oid := descriptor_oid; + ELSE + IF descriptor_relation IS NOT NULL OR descriptor_public_id_column IS NOT NULL THEN + RAISE EXCEPTION '% sources must not specify a descriptor relation', normalized_storage; + END IF; + tensor_relation_oid := heap_oid; + END IF; + + IF normalized_storage = 'heap_array' THEN + IF tensor_ref_column IS NOT NULL + OR tensor_rows_column IS NOT NULL + OR tensor_dim_column IS NOT NULL + OR tensor_dtype_column IS NOT NULL + OR tensor_checksum_column IS NOT NULL THEN + RAISE EXCEPTION 'heap_array sources must not specify external tensor columns'; + END IF; + ELSE + IF tensor_ref_column IS NULL + OR tensor_rows_column IS NULL + OR tensor_dim_column IS NULL + OR tensor_dtype_column IS NULL + OR tensor_checksum_column IS NULL THEN + RAISE EXCEPTION 'external tensor sources require ref, rows, dim, dtype, and checksum columns'; + END IF; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_ref_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor ref column % must be a NOT NULL text column', + tensor_ref_column; + END IF; + tensor_ref_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_rows_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor rows column % must be a NOT NULL integer column', + tensor_rows_column; + END IF; + tensor_rows_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dim_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'integer'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dim column % must be a NOT NULL integer column', + tensor_dim_column; + END IF; + tensor_dim_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_dtype_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor dtype column % must be a NOT NULL text column', + tensor_dtype_column; + END IF; + tensor_dtype_attnum := attnum; + + attnum := NULL; + SELECT a.attnum, a.atttypid, a.attnotnull + INTO attnum, atttypid, attnotnull + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid + AND a.attname = tensor_checksum_column + AND a.attnum > 0 + AND NOT a.attisdropped; + IF attnum IS NULL OR atttypid <> 'text'::regtype OR NOT attnotnull THEN + RAISE EXCEPTION 'tensor checksum column % must be a NOT NULL text column', + tensor_checksum_column; + END IF; + tensor_checksum_attnum := attnum; + + IF (CASE WHEN normalized_storage = 'external_ref' THEN 7 ELSE 6 END) <> ( + SELECT count(DISTINCT u.attnum) + FROM pg_catalog.unnest(ARRAY[ + CASE WHEN normalized_storage = 'external_ref' THEN model_contract_attnum ELSE descriptor_public_id_attnum END, + CASE WHEN normalized_storage = 'external_ref' THEN public_id_attnum ELSE NULL END, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + ]) AS u(attnum) + ) THEN + RAISE EXCEPTION 'tensor source columns must be distinct'; + END IF; + END IF; + + EXECUTE pg_catalog.format( + 'INSERT INTO %I._vchordrq_maxsim_sources ( + index_oid, heap_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum, registered_by + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14) + ON CONFLICT (index_oid) DO UPDATE SET + heap_oid = EXCLUDED.heap_oid, + model_contract_id = EXCLUDED.model_contract_id, + storage = EXCLUDED.storage, + model_contract_attnum = EXCLUDED.model_contract_attnum, + public_id_attnum = EXCLUDED.public_id_attnum, + descriptor_oid = EXCLUDED.descriptor_oid, + descriptor_public_id_attnum = EXCLUDED.descriptor_public_id_attnum, + tensor_ref_attnum = EXCLUDED.tensor_ref_attnum, + tensor_rows_attnum = EXCLUDED.tensor_rows_attnum, + tensor_dim_attnum = EXCLUDED.tensor_dim_attnum, + tensor_dtype_attnum = EXCLUDED.tensor_dtype_attnum, + tensor_checksum_attnum = EXCLUDED.tensor_checksum_attnum, + registered_by = EXCLUDED.registered_by, + registered_at = pg_catalog.clock_timestamp()', + ext_schema + ) USING + index_relation::oid, + heap_oid, + model_contract_id, + normalized_storage, + model_contract_attnum, + public_id_attnum, + descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum, + caller_oid; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_register_maxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_register_maxsim_source( + regclass, text, text, name, name, name, name, name, name, name, regclass, name +) TO PUBLIC; + +CREATE FUNCTION vchordrq_unregister_maxsim_source(index_relation regclass) +RETURNS boolean +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + index_owner oid; + removed_count bigint; +BEGIN + IF index_relation IS NULL THEN + RETURN false; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + SELECT c.relowner INTO index_owner + FROM pg_catalog.pg_class AS c + WHERE c.oid = index_relation::oid AND c.relkind = 'i'; + IF index_owner IS NULL + OR NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') THEN + RAISE EXCEPTION 'only the index owner may unregister its MaxSim tensor source'; + END IF; + + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_maxsim_sources WHERE index_oid = $1', + ext_schema + ) USING index_relation::oid; + GET DIAGNOSTICS removed_count = ROW_COUNT; + RETURN removed_count > 0; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_unregister_maxsim_source(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_unregister_maxsim_source(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_maxsim_source_info(index_relation regclass) +RETURNS TABLE( + registered_index regclass, + heap_relation regclass, + model_contract_id text, + source_storage text, + model_contract_column name, + public_id_column name, + descriptor_relation regclass, + descriptor_public_id_column name, + tensor_ref_column name, + tensor_rows_column name, + tensor_dim_column name, + tensor_dtype_column name, + tensor_checksum_column name +) +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + caller_oid oid; + bound_heap_oid oid; + live_heap_oid oid; + index_owner oid; + bound_model_contract_id text; + bound_storage text; + bound_descriptor_oid oid; + tensor_relation_oid oid; + model_contract_attnum smallint; + public_id_attnum smallint; + descriptor_public_id_attnum smallint; + tensor_ref_attnum smallint; + tensor_rows_attnum smallint; + tensor_dim_attnum smallint; + tensor_dtype_attnum smallint; + tensor_checksum_attnum smallint; + valid_columns bigint; + expected_columns bigint; + resolved_model_contract_column name; + resolved_public_id_column name; + resolved_descriptor_public_id_column name; + resolved_tensor_ref_column name; + resolved_tensor_rows_column name; + resolved_tensor_dim_column name; + resolved_tensor_dtype_column name; + resolved_tensor_checksum_column name; +BEGIN + IF index_relation IS NULL THEN + RAISE EXCEPTION 'index_relation must not be NULL'; + END IF; + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RAISE EXCEPTION 'vchord is not installed'; + END IF; + SELECT r.oid INTO caller_oid + FROM pg_catalog.pg_roles AS r + WHERE r.rolname = session_user; + + EXECUTE pg_catalog.format( + 'SELECT heap_oid, model_contract_id, storage, + model_contract_attnum, public_id_attnum, + descriptor_oid, descriptor_public_id_attnum, + tensor_ref_attnum, tensor_rows_attnum, tensor_dim_attnum, + tensor_dtype_attnum, tensor_checksum_attnum + FROM %I._vchordrq_maxsim_sources + WHERE index_oid = $1', + ext_schema + ) INTO + bound_heap_oid, + bound_model_contract_id, + bound_storage, + model_contract_attnum, + public_id_attnum, + bound_descriptor_oid, + descriptor_public_id_attnum, + tensor_ref_attnum, + tensor_rows_attnum, + tensor_dim_attnum, + tensor_dtype_attnum, + tensor_checksum_attnum + USING index_relation::oid; + IF bound_heap_oid IS NULL THEN + RAISE EXCEPTION 'MaxSim tensor source is not registered for index %', + index_relation; + END IF; + + SELECT x.indrelid, i.relowner + INTO live_heap_oid, index_owner + FROM pg_catalog.pg_index AS x + JOIN pg_catalog.pg_class AS i ON i.oid = x.indexrelid + JOIN pg_catalog.pg_class AS h ON h.oid = x.indrelid + JOIN pg_catalog.pg_am AS am ON am.oid = i.relam + JOIN pg_catalog.pg_opclass AS opc ON opc.oid = x.indclass[0] + WHERE x.indexrelid = index_relation::oid + AND i.relkind = 'i' + AND h.relkind IN ('r', 'm') + AND am.amname = 'vchordrq' + AND opc.opcname IN ( + 'vector_maxsim_ops', + 'halfvec_maxsim_ops', + 'rabitq8_maxsim_ops', + 'rabitq4_maxsim_ops' + ) + AND x.indisvalid + AND x.indisready + AND x.indnatts = 1 + AND x.indnkeyatts = 1; + IF live_heap_oid IS NULL OR live_heap_oid <> bound_heap_oid THEN + RAISE EXCEPTION 'registered MaxSim tensor source is stale or invalid'; + END IF; + IF NOT pg_catalog.pg_has_role(caller_oid, index_owner, 'USAGE') + AND NOT pg_catalog.has_table_privilege(caller_oid, live_heap_oid, 'SELECT') THEN + RAISE EXCEPTION 'permission denied for registered MaxSim tensor source'; + END IF; + + IF bound_storage NOT IN ('heap_array', 'external_ref', 'external_relation') THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid storage'; + END IF; + + SELECT count(*) + INTO valid_columns + FROM ( + VALUES + (model_contract_attnum, 'text'::regtype), + (public_id_attnum, 'bigint'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = bound_heap_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped; + IF valid_columns <> 2 THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid heap columns'; + END IF; + + IF bound_storage = 'heap_array' THEN + IF bound_descriptor_oid IS NOT NULL + OR descriptor_public_id_attnum IS NOT NULL + OR tensor_ref_attnum IS NOT NULL + OR tensor_rows_attnum IS NOT NULL + OR tensor_dim_attnum IS NOT NULL + OR tensor_dtype_attnum IS NOT NULL + OR tensor_checksum_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid heap_array binding'; + END IF; + tensor_relation_oid := NULL; + ELSIF bound_storage = 'external_ref' THEN + IF bound_descriptor_oid IS NOT NULL OR descriptor_public_id_attnum IS NOT NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid external_ref binding'; + END IF; + tensor_relation_oid := bound_heap_oid; + ELSE + IF bound_descriptor_oid IS NULL OR descriptor_public_id_attnum IS NULL THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid external_relation binding'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_class AS c + WHERE c.oid = bound_descriptor_oid + AND c.relkind IN ('r', 'm'); + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor relation is stale or invalid'; + END IF; + IF NOT pg_catalog.has_table_privilege(caller_oid, bound_descriptor_oid, 'SELECT') THEN + RAISE EXCEPTION 'SELECT privilege on the registered descriptor relation is required'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_descriptor_oid + AND a.attnum = descriptor_public_id_attnum + AND a.atttypid = 'bigint'::regtype + AND a.attnotnull + AND NOT a.attisdropped; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor public ID column is invalid'; + END IF; + PERFORM 1 + FROM pg_catalog.pg_index AS x + WHERE x.indrelid = bound_descriptor_oid + AND x.indisunique + AND x.indisvalid + AND x.indisready + AND x.indnkeyatts = 1 + AND x.indkey[0] = descriptor_public_id_attnum + AND x.indexprs IS NULL + AND x.indpred IS NULL; + IF NOT FOUND THEN + RAISE EXCEPTION 'registered MaxSim descriptor public ID is no longer unique'; + END IF; + tensor_relation_oid := bound_descriptor_oid; + END IF; + + expected_columns := CASE WHEN bound_storage = 'heap_array' THEN 0 ELSE 5 END; + SELECT count(*) + INTO valid_columns + FROM ( + VALUES + (tensor_ref_attnum, 'text'::regtype), + (tensor_rows_attnum, 'integer'::regtype), + (tensor_dim_attnum, 'integer'::regtype), + (tensor_dtype_attnum, 'text'::regtype), + (tensor_checksum_attnum, 'text'::regtype) + ) AS expected(attnum, atttypid) + JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = tensor_relation_oid + AND a.attnum = expected.attnum + AND a.atttypid = expected.atttypid + AND a.attnotnull + AND NOT a.attisdropped + WHERE expected.attnum IS NOT NULL; + IF valid_columns <> expected_columns THEN + RAISE EXCEPTION 'registered MaxSim tensor source has invalid descriptor columns'; + END IF; + + SELECT a.attname INTO resolved_model_contract_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_heap_oid AND a.attnum = model_contract_attnum; + SELECT a.attname INTO resolved_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_heap_oid AND a.attnum = public_id_attnum; + SELECT a.attname INTO resolved_descriptor_public_id_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = bound_descriptor_oid AND a.attnum = descriptor_public_id_attnum; + SELECT a.attname INTO resolved_tensor_ref_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_ref_attnum; + SELECT a.attname INTO resolved_tensor_rows_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_rows_attnum; + SELECT a.attname INTO resolved_tensor_dim_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dim_attnum; + SELECT a.attname INTO resolved_tensor_dtype_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_dtype_attnum; + SELECT a.attname INTO resolved_tensor_checksum_column + FROM pg_catalog.pg_attribute AS a + WHERE a.attrelid = tensor_relation_oid AND a.attnum = tensor_checksum_attnum; + + RETURN QUERY SELECT + index_relation, + bound_heap_oid::regclass, + bound_model_contract_id, + bound_storage, + resolved_model_contract_column, + resolved_public_id_column, + bound_descriptor_oid::regclass, + resolved_descriptor_public_id_column, + resolved_tensor_ref_column, + resolved_tensor_rows_column, + resolved_tensor_dim_column, + resolved_tensor_dtype_column, + resolved_tensor_checksum_column; +END; +$$; + +REVOKE ALL ON FUNCTION vchordrq_maxsim_source_info(regclass) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_maxsim_source_info(regclass) TO PUBLIC; + +CREATE FUNCTION vchordrq_maxsim_search( + index_relation regclass, + query anyarray, + candidate_limit integer, + top_k integer +) +RETURNS TABLE(public_id bigint, similarity real) +STRICT +VOLATILE +PARALLEL UNSAFE +LANGUAGE c +AS 'MODULE_PATHNAME', '_vchordrq_maxsim_search_external_wrapper'; + +COMMENT ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) +IS 'Restricted Phase 3B external-tensor MaxSim search; returns exact similarity under caller MVCC, SELECT privileges, and PostgreSQL row visibility.'; + +REVOKE ALL ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) FROM PUBLIC; +GRANT EXECUTE ON FUNCTION vchordrq_maxsim_search(regclass, anyarray, integer, integer) TO PUBLIC; + +CREATE FUNCTION _vchordrq_maxsim_source_sql_drop() +RETURNS event_trigger +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +AS $$ +DECLARE + ext_schema name; + registry regclass; +BEGIN + SELECT n.nspname + INTO ext_schema + FROM pg_catalog.pg_extension AS e + JOIN pg_catalog.pg_namespace AS n ON n.oid = e.extnamespace + WHERE e.extname = 'vchord'; + IF ext_schema IS NULL THEN + RETURN; + END IF; + registry := pg_catalog.to_regclass( + pg_catalog.format('%I._vchordrq_maxsim_sources', ext_schema) + ); + IF registry IS NULL THEN + RETURN; + END IF; + + EXECUTE pg_catalog.format( + 'DELETE FROM %I._vchordrq_maxsim_sources AS s + USING pg_catalog.pg_event_trigger_dropped_objects() AS d + WHERE d.objid = s.index_oid + OR ( + d.objid = s.heap_oid + AND ( + d.objsubid = 0 + OR d.objsubid = s.model_contract_attnum + OR d.objsubid = s.public_id_attnum + OR ( + s.storage = ''external_ref'' + AND d.objsubid IN ( + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + ) + ) + OR ( + d.objid = s.descriptor_oid + AND ( + d.objsubid = 0 + OR d.objsubid = s.descriptor_public_id_attnum + OR d.objsubid IN ( + s.tensor_ref_attnum, + s.tensor_rows_attnum, + s.tensor_dim_attnum, + s.tensor_dtype_attnum, + s.tensor_checksum_attnum + ) + ) + )', + ext_schema + ); +END; +$$; + +REVOKE ALL ON FUNCTION _vchordrq_maxsim_source_sql_drop() FROM PUBLIC; + +CREATE EVENT TRIGGER _vchordrq_maxsim_source_sql_drop +ON sql_drop +EXECUTE FUNCTION _vchordrq_maxsim_source_sql_drop(); diff --git a/tests/vchordrq/cost_estimator.slt b/tests/vchordrq/cost_estimator.slt index 0f4ad266..87cf85e7 100644 --- a/tests/vchordrq/cost_estimator.slt +++ b/tests/vchordrq/cost_estimator.slt @@ -388,6 +388,174 @@ RESET vchordrq.enable_scan; statement ok RESET enable_seqscan; +# --------------------------------------------------------------------------- +# Case 12: MaxSim must have a nonzero, query-token-aware, backend-aware cost. +# The old special branch returned total_cost=0 and selectivity=1 for every +# MaxSim path, hiding all token expansion and exact rerank work. +# --------------------------------------------------------------------------- + +statement ok +CREATE TABLE cost_test_maxsim ( + id int PRIMARY KEY, + v vector(3)[] NOT NULL +); + +statement ok +INSERT INTO cost_test_maxsim +SELECT i, ARRAY[ + '[1,0,0]'::vector, + '[0,1,0]'::vector, + '[0,0,1]'::vector, + '[0.5,0.5,0]'::vector +] +FROM generate_series(1, 1000) i; + +statement ok +CREATE INDEX cost_test_maxsim_v +ON cost_test_maxsim +USING vchordrq (v vector_maxsim_ops) +WITH (options = $$ +[build.internal] +lists = [4] +$$); + +statement ok +ANALYZE cost_test_maxsim; + +statement ok +SET enable_seqscan = off; + +statement ok +SET vchordrq.probes = '4'; + +statement ok +SET vchordrq.maxsim_planner_document_tokens = 4; + +statement ok +CREATE TEMP TABLE maxsim_cost_observations ( + name text PRIMARY KEY, + value double precision NOT NULL +); + +statement ok +SET vchordrq.maxsim_backend = 'coarse_only'; + +statement ok +SET vchordrq.maxsim_planner_query_tokens = 1; + +# A newly built index has a native indexed-vector count. The document-token +# GUC is only a compatibility fallback for pre-statistics indexes, so changing +# it must not change this index's cost. +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'native_document_4', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +statement ok +SET vchordrq.maxsim_planner_document_tokens = 4096; + +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'native_document_4096', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +query B +SELECT + (SELECT value FROM maxsim_cost_observations WHERE name = 'native_document_4') + = + (SELECT value FROM maxsim_cost_observations WHERE name = 'native_document_4096'); +---- +t + +statement ok +SET vchordrq.maxsim_planner_document_tokens = 4; + +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'query_1', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +statement ok +SET vchordrq.maxsim_planner_query_tokens = 64; + +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'query_64', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +query B +SELECT + (SELECT value FROM maxsim_cost_observations WHERE name = 'query_64') + > + (SELECT value FROM maxsim_cost_observations WHERE name = 'query_1'); +---- +t + +statement ok +SET vchordrq.maxsim_planner_query_tokens = 32; + +statement ok +SET vchordrq.maxsim_candidate_limit = 16; + +statement ok +SET vchordrq.maxsim_backend = 'cpu_exact'; + +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'cpu_16', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +statement ok +SET vchordrq.maxsim_candidate_limit = 256; + +statement ok +INSERT INTO maxsim_cost_observations +SELECT 'cpu_256', top_total_cost( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); + +query B +SELECT + (SELECT value FROM maxsim_cost_observations WHERE name = 'cpu_256') + > + (SELECT value FROM maxsim_cost_observations WHERE name = 'cpu_16'); +---- +t + +query T +SELECT plan_top_index( + 'SELECT id FROM cost_test_maxsim + ORDER BY v @# ARRAY[''[1,0,0]''::vector] LIMIT 10' +); +---- +cost_test_maxsim_v + +statement ok +RESET vchordrq.maxsim_backend; + +statement ok +RESET vchordrq.maxsim_candidate_limit; + +statement ok +RESET vchordrq.maxsim_planner_query_tokens; + +statement ok +RESET vchordrq.maxsim_planner_document_tokens; + +statement ok +RESET vchordrq.probes; + # --------------------------------------------------------------------------- # Cleanup # --------------------------------------------------------------------------- @@ -404,6 +572,9 @@ DROP TABLE cost_test_cold; statement ok DROP TABLE cost_test_partial; +statement ok +DROP TABLE cost_test_maxsim; + statement ok DROP FUNCTION slow_true(int); diff --git a/tests/vchordrq/maxsim_correctness.slt b/tests/vchordrq/maxsim_correctness.slt new file mode 100644 index 00000000..a5d9652f --- /dev/null +++ b/tests/vchordrq/maxsim_correctness.slt @@ -0,0 +1,221 @@ +# Deterministic MaxSim semantics and input-boundary coverage. + +statement ok +SET enable_seqscan TO off; + +# @# is a distance: negative late-interaction similarity, ordered ascending. +query I +SELECT round(( + ARRAY['[1,0]'::vector, '[0,1]'::vector] + @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +)::numeric, 3); +---- +-2.000 + +# MaxSim is asymmetric: the right-hand array contains query vectors. +query I +SELECT round(( + ARRAY['[1,0]'::vector, '[0,1]'::vector] + @# ARRAY['[1,1]'::vector] +)::numeric, 3); +---- +-1.000 + +query I +SELECT round(( + ARRAY['[1,1]'::vector] + @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +)::numeric, 3); +---- +-2.000 + +# halfvec uses the same sign and orientation contract. +query I +SELECT round(( + ARRAY['[1,0]'::halfvec, '[0,1]'::halfvec] + @# ARRAY['[1,0]'::halfvec, '[0,1]'::halfvec] +)::numeric, 3); +---- +-2.000 + +statement error MaxSim arrays must contain at least one vector +SELECT ARRAY[]::vector[] @# ARRAY['[1,0]'::vector]; + +statement error MaxSim arrays must contain at least one vector +SELECT ARRAY['[1,0]'::vector] @# ARRAY[]::vector[]; + +statement error MaxSim arrays must not contain NULL vectors +SELECT ARRAY['[1,0]'::vector, NULL::vector] + @# ARRAY['[1,0]'::vector]; + +statement error dimension is not matched +SELECT ARRAY['[1,0]'::vector] @# ARRAY['[1,0,0]'::vector]; + +statement error MaxSim arrays cannot contain more than 65536 vectors +SELECT ARRAY['[1]'::vector] + @# ARRAY(SELECT '[1]'::vector FROM generate_series(1, 65537)); + +# Indexing an empty document array must fail instead of silently omitting it. +statement ok +CREATE TABLE maxsim_invalid_document (val vector(2)[]); + +statement ok +INSERT INTO maxsim_invalid_document VALUES (ARRAY[]::vector[]); + +statement error MaxSim arrays must contain at least one vector +CREATE INDEX ON maxsim_invalid_document +USING vchordrq (val vector_maxsim_ops); + +statement ok +TRUNCATE maxsim_invalid_document; + +statement ok +INSERT INTO maxsim_invalid_document +SELECT ARRAY( + SELECT '[1,0]'::vector + FROM generate_series(1, 65537) +); + +statement error MaxSim arrays cannot contain more than 65536 vectors +CREATE INDEX ON maxsim_invalid_document +USING vchordrq (val vector_maxsim_ops); + +statement ok +DROP TABLE maxsim_invalid_document; + +# Lock down current ranking with a non-flat index before the Phase 3 refactor. +statement ok +CREATE TABLE maxsim_deterministic ( + id integer primary key, + val vector(2)[] not null +); + +statement ok +INSERT INTO maxsim_deterministic VALUES + (1, ARRAY['[1,0]'::vector, '[0,1]'::vector]), + (2, ARRAY['[0.8,0]'::vector, '[0,0.8]'::vector]), + (3, ARRAY['[0.5,0.5]'::vector]), + (4, ARRAY['[-1,0]'::vector, '[0,-1]'::vector]); + +statement ok +CREATE INDEX maxsim_deterministic_idx +ON maxsim_deterministic +USING vchordrq (val vector_maxsim_ops) +WITH (options = $$ +[build.internal] +lists = [2] +$$); + +statement ok +SET vchordrq.probes = '2'; + +statement ok +SET vchordrq.maxsim_refine = 100; + +statement ok +SET vchordrq.maxsim_threshold = 0; + +query I +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 4; +---- +1 +2 +3 +4 + +statement ok +SET vchordrq.maxsim_candidate_limit = 2; + +query I +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 4; +---- +1 +2 + +statement ok +SET vchordrq.maxsim_candidate_limit = -1; + +statement ok +SET vchordrq.maxsim_backend = 'cpu_exact'; + +statement error exact MaxSim requires a positive vchordrq.maxsim_candidate_limit +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 1; + +statement ok +SET vchordrq.maxsim_candidate_limit = 4; + +query I +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 4; +---- +1 +2 +3 +4 + +query I +SELECT round((val @# ARRAY['[1,0]'::vector, '[0,1]'::vector])::numeric, 3) +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 4; +---- +-2.000 +-1.600 +-1.000 +0.000 + +statement ok +SET vchordrq.maxsim_backend = 'gpu'; + +statement error GPU MaxSim transport error: endpoint is empty +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 1; + +statement ok +SET vchordrq.maxsim_backend = 'auto'; + +query I +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0]'::vector, '[0,1]'::vector] +LIMIT 4; +---- +1 +2 +3 +4 + +statement ok +SET vchordrq.maxsim_backend = 'coarse_only'; + +statement ok +SET vchordrq.maxsim_candidate_limit = -1; + +# An empty query must also fail on the index scan path. +statement error MaxSim arrays must contain at least one vector +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY[]::vector[] +LIMIT 1; + +statement error dimension is not matched +SELECT id +FROM maxsim_deterministic +ORDER BY val @# ARRAY['[1,0,0]'::vector] +LIMIT 1; + +statement ok +DROP TABLE maxsim_deterministic; diff --git a/tests/vchordrq/maxsim_source_registry.slt b/tests/vchordrq/maxsim_source_registry.slt new file mode 100644 index 00000000..c408a469 --- /dev/null +++ b/tests/vchordrq/maxsim_source_registry.slt @@ -0,0 +1,341 @@ +# Phase 3B tensor-source registration must bind to a real MaxSim index by +# relation/attribute OID, reject incompatible descriptors, and fail closed when +# a bound index/table column is dropped. + +statement ok +CREATE TABLE maxsim_source_test ( + id bigint PRIMARY KEY, + model_contract text NOT NULL, + embedding vector(2)[] NOT NULL, + single_embedding vector(2) NOT NULL, + tensor_ref text NOT NULL, + tensor_rows integer NOT NULL, + tensor_dim integer NOT NULL, + tensor_dtype text NOT NULL, + tensor_checksum text NOT NULL, + application_note text +); + +statement ok +INSERT INTO maxsim_source_test VALUES ( + 1, + 'colqwen@test', + ARRAY['[1,0]'::vector, '[0,1]'::vector], + '[1,0]'::vector, + 'tensor://1', + 2, + 2, + 'float32', + 'sha256:test' +); + +statement ok +CREATE INDEX maxsim_source_test_idx +ON maxsim_source_test +USING vchordrq (embedding vector_maxsim_ops) +WITH (options = $$ +[build.internal] +lists = [1] +$$); + +statement ok +CREATE INDEX maxsim_source_wrong_idx +ON maxsim_source_test +USING vchordrq (single_embedding vector_l2_ops) +WITH (options = $$ +[build.internal] +lists = [1] +$$); + +statement error not a valid single-key vchordrq MaxSim index +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_wrong_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'heap_array', + model_contract_column => 'model_contract', + public_id_column => 'id' +); + +statement error tensor rows column tensor_dtype must be a NOT NULL integer column +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'external_ref', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_dtype', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum' +); + +statement ok +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => ' colqwen@test ', + storage => 'EXTERNAL_REF', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_rows', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum' +); + +query TT +SELECT s.storage, s.model_contract_id +FROM _vchordrq_maxsim_sources AS s +WHERE s.index_oid = 'maxsim_source_test_idx'::regclass; +---- +external_ref colqwen@test + +query TTT +SELECT source_storage, tensor_ref_column::text, public_id_column::text +FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); +---- +external_ref tensor_ref id + +# Production external tensors may live in a compact descriptor relation keyed +# by the stable application public ID, avoiding a rewrite of the indexed heap. +statement ok +CREATE TABLE maxsim_descriptor_no_unique ( + public_id bigint NOT NULL, + tensor_ref text NOT NULL, + tensor_rows integer NOT NULL, + tensor_dim integer NOT NULL, + tensor_dtype text NOT NULL, + tensor_checksum text NOT NULL +); + +statement error must have a non-partial single-key unique index +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'external_relation', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_rows', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum', + descriptor_relation => 'maxsim_descriptor_no_unique'::regclass, + descriptor_public_id_column => 'public_id' +); + +statement ok +CREATE TABLE maxsim_descriptor_test ( + public_id bigint PRIMARY KEY, + tensor_ref text NOT NULL, + tensor_rows integer NOT NULL, + tensor_dim integer NOT NULL, + tensor_dtype text NOT NULL, + tensor_checksum text NOT NULL +); + +statement ok +INSERT INTO maxsim_descriptor_test VALUES ( + 1, 'tensor://1', 2, 2, 'float32', 'sha256:test' +); + +statement ok +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'external_relation', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_rows', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum', + descriptor_relation => 'maxsim_descriptor_test'::regclass, + descriptor_public_id_column => 'public_id' +); + +query TTTT +SELECT source_storage, descriptor_relation::text, + descriptor_public_id_column::text, tensor_ref_column::text +FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); +---- +external_relation maxsim_descriptor_test public_id tensor_ref + +statement ok +ALTER TABLE maxsim_descriptor_test RENAME tensor_ref TO tensor_location; + +query T +SELECT tensor_ref_column::text +FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); +---- +tensor_location + +statement ok +ALTER TABLE maxsim_descriptor_test DROP COLUMN tensor_checksum; + +query I +SELECT count(*) +FROM _vchordrq_maxsim_sources +WHERE index_oid = 'maxsim_source_test_idx'::regclass; +---- +0 + +statement ok +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'external_ref', + model_contract_column => 'model_contract', + public_id_column => 'id', + tensor_ref_column => 'tensor_ref', + tensor_rows_column => 'tensor_rows', + tensor_dim_column => 'tensor_dim', + tensor_dtype_column => 'tensor_dtype', + tensor_checksum_column => 'tensor_checksum' +); + +# Unbound application columns are outside the registry contract. Dropping one +# must not invalidate an otherwise live source binding. +statement ok +ALTER TABLE maxsim_source_test DROP COLUMN application_note; + +query T +SELECT model_contract_id +FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); +---- +colqwen@test + +# The explicit score surface validates bounded work and exact query/index type +# before any sidecar access. +statement error candidate_limit must be between 1 and 65536 +SELECT * FROM vchordrq_maxsim_search( + 'maxsim_source_test_idx'::regclass, + ARRAY['[1,0]'::vector], + 0, + 1 +); + +statement error top_k must be positive and no greater than candidate_limit +SELECT * FROM vchordrq_maxsim_search( + 'maxsim_source_test_idx'::regclass, + ARRAY['[1,0]'::vector], + 1, + 2 +); + +statement ok +SET vchordrq.maxsim_backend = 'gpu'; + +statement error MaxSim tensor kind or dimension is not matched +SELECT * FROM vchordrq_maxsim_search( + 'maxsim_source_test_idx'::regclass, + ARRAY['[1,0]'::halfvec], + 1, + 1 +); + +statement ok +SET vchordrq.maxsim_backend = 'coarse_only'; + +query T +SELECT a.attname +FROM _vchordrq_maxsim_sources AS s +JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = s.heap_oid AND a.attnum = s.tensor_ref_attnum +WHERE s.index_oid = 'maxsim_source_test_idx'::regclass; +---- +tensor_ref + +# Attribute numbers survive renames, so bindings do not depend on SQL text. +statement ok +ALTER TABLE maxsim_source_test RENAME tensor_ref TO tensor_location; + +query T +SELECT a.attname +FROM _vchordrq_maxsim_sources AS s +JOIN pg_catalog.pg_attribute AS a + ON a.attrelid = s.heap_oid AND a.attnum = s.tensor_ref_attnum +WHERE s.index_oid = 'maxsim_source_test_idx'::regclass; +---- +tensor_location + +query T +SELECT tensor_ref_column::text +FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); +---- +tensor_location + +# Type-altering DDL does not drop an attribute, so runtime resolution must +# revalidate the complete descriptor and fail closed. +statement ok +ALTER TABLE maxsim_source_test +ALTER COLUMN tensor_dtype TYPE varchar(16); + +statement error registered MaxSim tensor source has invalid descriptor columns +SELECT * FROM vchordrq_maxsim_source_info('maxsim_source_test_idx'::regclass); + +statement ok +ALTER TABLE maxsim_source_test +ALTER COLUMN tensor_dtype TYPE text; + +# Dropping a bound descriptor column invalidates the complete binding. +statement ok +ALTER TABLE maxsim_source_test DROP COLUMN tensor_checksum; + +query I +SELECT count(*) +FROM _vchordrq_maxsim_sources +WHERE index_oid = 'maxsim_source_test_idx'::regclass; +---- +0 + +statement ok +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'heap_array', + model_contract_column => 'model_contract', + public_id_column => 'id' +); + +query B +SELECT tensor_ref_attnum IS NULL +FROM _vchordrq_maxsim_sources +WHERE index_oid = 'maxsim_source_test_idx'::regclass; +---- +t + +query B +SELECT vchordrq_unregister_maxsim_source('maxsim_source_test_idx'::regclass); +---- +t + +query B +SELECT vchordrq_unregister_maxsim_source('maxsim_source_test_idx'::regclass); +---- +f + +statement ok +SELECT vchordrq_register_maxsim_source( + index_relation => 'maxsim_source_test_idx'::regclass, + model_contract_id => 'colqwen@test', + storage => 'heap_array', + model_contract_column => 'model_contract', + public_id_column => 'id' +); + +statement ok +DROP INDEX maxsim_source_test_idx; + +query I +SELECT count(*) FROM _vchordrq_maxsim_sources; +---- +0 + +statement ok +DROP TABLE maxsim_source_test; + +statement ok +DROP TABLE maxsim_descriptor_test, maxsim_descriptor_no_unique;