Skip to content

feat: add IVF-PQ L2 distance range search - #107

Open
JunRuiLee wants to merge 2 commits into
apache:mainfrom
JunRuiLee:feat/ivfpq-l2-range-search
Open

JunRuiLee wants to merge 2 commits into
apache:mainfrom
JunRuiLee:feat/ivfpq-l2-range-search

Conversation

@JunRuiLee

@JunRuiLee JunRuiLee commented Sep 17, 2026

Copy link
Copy Markdown
Contributor

Summary

Part of #97.

Add IVF-PQ L2 distance-range search to all four Rust reader entry points: single and batch queries, with and without a serialized Roaring allow-list.

  • Reuse DistanceBand, VectorRangeSearchParams, RangeSearchResult, and RangeCollector.
  • Support PQ4/PQ8, residual/non-residual encoding, and OPQ.
  • Share list I/O across queries, scan sufficiently large batches in parallel, and stream oversized lists through the bounded reader.
  • Preserve existing top-K behavior and the v1 storage format.

Search semantics

Membership is half-open [lower, upper) in squared-L2 estimated-distance space. Both code widths use floating-point ADC: direct query-to-codebook subvector distances, after optional OPQ and coarse residual subtraction, summed in subquantizer order.

The range path does not use u8 FastScan tables, precomputed norm identities, top-K truncation, or original-vector reranking. Membership is independent of list size, batch size, and optimize_for_search. Range estimates can differ from top-K distances; probing every list removes the IVF coverage gap, not PQ quantization error.

Results are uncapped and unordered. Each filter-eligible row is fully evaluated, so PQ early_abandoned is always zero. Non-finite transformed queries, any coarse distance (including unselected centroids), and consumed PQ estimates fail with InvalidData. Filtered-out and unprobed rows are not evaluated. Query shape, metric, probe width, and serialized filters are validated before empty-band shortcuts. Collector failures propagate without returning partial results or reading subsequent streamed chunks.

Performance and memory

Each unique non-empty selected list is read once per call and shared across queries. Large batches use Rayon with query-owned collectors and reusable worker-local scratch. Small workloads remain serial; fully excluded lists/chunks skip distance-table construction.

Non-residual float lookup tables are reused across lists. Residual tables are reused across chunks of the same oversized list and invalidated when the list changes; ordinary residual lists use worker scratch. Query-table caching is lazy and capped at 8 MiB per call, with scratch fallback beyond the cache allowance. This cap covers cached LUTs only, not scratch, inputs, I/O buffers, or uncapped results. The 8 MiB allowance and 8,192 row-query parallel threshold are internal safeguards rather than globally optimal tuning claims.

Local measurements

Apple M2 Pro, macOS/AArch64, Rust 1.95.0, release library and rustc -O probe. The probe uses an in-memory reader, preloaded metadata, warmed thread pools, no row filter, full probing, and an empty-result band [0, 0.0001) to isolate query processing from result materialization. Fixture construction is outside the measurement. Results are median latency over five calls following one warmup, with allocator instrumentation enabled. These synthetic observations are not production-throughput or disk-I/O claims.

Workload PQ bits Threads Median latency (ms)
Dense batch: d=768, m=96, 32 lists x 2,048 rows, 8 queries, residual 4 1 38.834
Same dense batch 4 4 10.947
Same dense batch 8 1 57.423
Same dense batch 8 4 23.325
Thin lists: d=768, m=96, 64 lists x 4 rows, 16 queries, non-residual 8 1 1.265

Allocator-requested bytes for the median-latency sample:

Workload Cumulative allocation (MiB) Extra peak live allocation (MiB)
Dense PQ8 batch, 4 threads 6.993 6.981
Thin PQ8 lists, 1 thread 1.629 1.593

LUT reuse trades bounded live memory for less repeated work; less allocation churn does not imply lower peak memory. These figures measure requested heap bytes, not RSS. The shared result-container memory behavior is outside this PR's scope.

Scope

The change is limited to Rust IVF-PQ L2 range search. Other metrics, language bindings, storage-format changes, top-K scan changes, PQ serialization changes, and shared result-container restructuring are out of scope. distance.rs and the PQ writer are unchanged. The io.rs change adds fallible propagation to the streamed-reader callback while retaining its top-K wrapper. No storage migration is required.

Tests and validation

Thirteen PQ range integration tests cover both code widths, an independent decoded-code oracle, dense OPQ and residual modes, exact boundaries and adjacent bands, non-finite data, cancellation and overflow, sparse/dense/empty filters, statistics, batch/single and query-permutation equivalence, optimization invariance, and unchanged top-K results. Four focused unit tests cover multi-worker scanning, bounded cached/uncached equivalence and residual invalidation, no table allocation for excluded rows, and parallel collector-error propagation. An index larger than 64 MiB checks bounded streaming, shared physical reads, tail-code layout, cache-budget overflow, and immediate collector-error propagation.

Local validation for this revision on macOS/AArch64 with Rust 1.95.0:

  • cargo test --offline --locked --workspace: 647 passed, 2 existing ignored.
  • cargo test --offline --locked --release --workspace: 647 passed, 2 existing ignored.
  • The range integration suite contains 65 passing tests across the supported families.
  • cargo fmt --all -- --check: passed.
  • cargo clippy --offline --locked --all-targets --workspace -- -D warnings: passed.
  • python3 tools/check_license_headers.py: passed.
  • git diff --check: passed.
  • Repository CI: all nine checks passed for this revision, including Rust, C/C++, JNI, Python build, and Linux/macOS/Windows wheels.

Validation caveat: local runs have encountered worker-observation assertion failures in the unchanged IVF-Flat and IVF-SQ tests. The reported full debug/release runs and repository CI pass; historical baseline flakiness has not been established.

Comment thread core/src/ivfpq.rs
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants