feat: add cosine and inner-product range search for IVF indexes

core/README.md

@@ -22,14 +22,17 @@
 `paimon-vindex-core` contains the Rust implementations and seek-based readers
 for IVF-FLAT, IVF-SQ, IVF-PQ, IVF-RQ, and DiskANN.
 
-The Rust reader supports distance range search for IVF-FLAT, IVF-RQ, and IVF-SQ with
-squared L2, using `DistanceBand`, `VectorRangeSearchParams`, and CSR
-`RangeSearchResult` buffers. All three families support single and batch queries,
+The Rust reader supports distance range search for IVF-FLAT, IVF-SQ, IVF-PQ and
+IVF-RQ with L2, cosine and inner product, using `DistanceBand`,
+`VectorRangeSearchParams`, and CSR `RangeSearchResult` buffers. All four families
+support single and batch queries,
 with or without a serialized Roaring allow-list, and a fixed positive `nprobe`.
 IVF-FLAT tests exact distances; IVF-RQ tests its one-bit or full multi-bit
-estimated distances; IVF-SQ tests scalar-quantized estimates. Results are uncapped
-and unordered. Probing every list removes the IVF coverage gap, but not the
-quantization error of IVF-RQ or IVF-SQ. The range path does not change top-K
+estimated distances; IVF-SQ tests scalar-quantized estimates; IVF-PQ tests
+floating-point ADC estimates for 4-bit and 8-bit codes, with optional residual
+encoding and OPQ. Results are uncapped and unordered. Probing every list removes
+the IVF coverage gap, but not the compressed families' quantization error.
+The range path does not change top-K
 search or the v1 storage format.
 
 See the [range search guide](../docs/range-search.html) for membership,
@@ -42,11 +45,36 @@ abstractions. It does not incorporate source code from Microsoft's
 The implementation supports L2, inner-product, and cosine search with the same
 lower-is-better distance semantics as the IVF indexes.
 
-## IVF-SQ range search
+## Distance semantics
+
+`DistanceBand::new` takes internal, lower-is-better f32 scores: squared L2,
+cosine distance (`1 - cos`, not clamped), or negative inner product.
+`MetricType::public_distance` converts a score to its public f64 predicate value:
+L2 takes the f32 square root before widening, cosine widens unchanged, and inner
+product negates. Returned `distances` always remain in internal units.
+
+Use `DistanceBand::from_endpoints` for public f64 predicates. `Ge`/`Gt` belong on
+the lower side and `Le`/`Lt` on the upper side. L2 resolves square-root rounding;
+cosine searches all finite f32 values (including negative roundoff and signed
+zero); inner product reverses both the side and comparison when negating.
+Endpoints are never rounded to f32 first. Missing sides are structurally
+unbounded, so an inclusive endpoint at `f32::MAX` does not lose that value.
+Out-of-domain cosine/IP predicates become empty or unbounded bands; unrepresentable
+L2 cuts retain the existing `Unsupported` response.
+
+`IndexType::supports_range_search(metric)` and `reader.supports_range_search()`
+report capability without reading list payloads. DiskANN and C/JNI range APIs
+remain unsupported. Bad queries, mismatched metrics, non-finite endpoints,
+malformed filters and zero `nprobe` are errors even for an empty band.
+Non-finite consumed distances or cosine norms return `InvalidData`, not partial
+results. Cosine queries use the existing normalization, including leaving zero
+vectors at zero; normalization overflow is an error on this path.
+
+## IVF-SQ and IVF-PQ range search
 
 Use `DistanceBand` and `VectorRangeSearchParams` with `range_search`,
 `range_search_batch`, or their `*_with_roaring_filter` variants. Bands are
-half-open `[lower, upper)` in squared-L2 units; results use `RangeSearchResult`
+half-open `[lower, upper)` in internal metric units; results use `RangeSearchResult`
 CSR buffers without sorting, padding, or a top-K cap.
 
 IVF-SQ uses the same blocked SIMD estimator as top-K, including the stored
@@ -57,11 +85,32 @@ membership is required. There is no original-vector reranking or top-K fallback.
 
 Batch queries share list reads, reuse the existing partition cache, and evaluate
 the Roaring allow-list once per list row using query-local one-bit-per-row masks.
-Large lists stream in bounded chunks; scan scratch is reused, and a finite upper cut
-allows entire SQ blocks to stop after their partial distances reach that cut.
+Large lists stream in bounded chunks; scan scratch is reused, and under L2 a
+finite upper cut allows entire SQ blocks to stop after their partial distances
+reach that cut.
 Result memory still grows with the number of hits. Cache hits are excluded from
-`call_stats().list_reads()`. Range support does not extend to other metrics,
-IVF-PQ, DiskANN, or language bindings in this change.
+`call_stats().list_reads()`. Cosine and IP never prune partial sums and report zero
+`early_abandoned`; L2 pruning is unchanged. Filters exclude negative row IDs in
+the Roaring API, while direct `RowIdFilter` implementations may admit them.
+
+IVF-PQ sums full f32 subvector distances in subquantizer order rather than using
+top-K's quantized FastScan tables.
+L2 uses ADC squared distances; cosine uses half the ADC squared distance after
+query normalization; IP uses negative estimated dot product. The cosine score
+is a unit-vector surrogate, not a re-normalized exact distance to decoded codes.
+Like IVF-RQ's cosine estimator, it can differ from exact cosine even for zero
+queries. IVF-FLAT and IVF-SQ define cosine distance involving a zero vector as 1.
+PQ and RQ always evaluate complete estimates, without early abandonment.
+Shared lists are read once per call, oversized PQ lists stream in bounded chunks,
+and the PQ batch filter is evaluated once per list row, not once per query.
+Query lookup tables are allocated lazily within an 8 MiB cache cap, with reusable
+worker scratch beyond that budget. Residual tables are reused across chunks of
+the same list and rebuilt for a different list. Large batches scan queries in
+parallel; membership does not depend on worker count, list size, batch size or
+`optimize_for_search`. Non-finite transformed queries, coarse distances (even to
+unselected lists), or consumed estimates return `InvalidData` without partial results.
+PQ range scores need not be bit-identical to top-K scores: the range path does
+not use top-K table quantization, L2 table expansion, or its cosine score scale.
 
 The crate ships its [normative v1 storage-format specification](STORAGE_FORMAT.md)
 and byte-exact fixtures. Project documentation, language bindings, and

core/src/collect.rs

@@ -36,6 +36,8 @@ use crate::range::{Bound, DistanceBand};
 
 /// The interface a scan kernel uses to hand candidate rows to a collector.
 pub(crate) trait Collector {
+    const VALIDATE_COSINE_INPUTS: bool = false;
+
     /// Called when the scan kernel rejects a row against [`cutoff`] instead of
     /// delivering it.
     ///
@@ -140,6 +142,8 @@ impl RangeCollector {
 }
 
 impl Collector for RangeCollector {
+    const VALIDATE_COSINE_INPUTS: bool = true;
+
     #[inline]
     fn note_abandoned(&mut self) {
         self.scanned += 1;
@@ -222,7 +226,7 @@ mod tests {
     }
 
     #[test]
-    fn an_uncertified_metric_never_prunes_on_a_partial_sum() {
+    fn a_non_l2_metric_never_prunes_on_a_partial_sum() {
         // A partial cosine or inner-product accumulation does not bound the full
         // value, so the cutoff must stay infinite no matter what the band says.
         for metric in [MetricType::Cosine, MetricType::InnerProduct] {

core/src/distance.rs

@@ -27,6 +27,17 @@ pub enum MetricType {
 }
 
 impl MetricType {
+    /// Converts a finite internal score to the public predicate value.
+    /// L2 takes an f32 square root before widening; cosine is `1 - cos`
+    /// without clamping; inner product reverses the internal score's sign.
+    pub fn public_distance(self, distance: f32) -> f64 {
+        f64::from(match self {
+            Self::L2 => distance.sqrt(),
+            Self::Cosine => distance,
+            Self::InnerProduct => -distance,
+        })
+    }
+
     pub fn from_code(code: u32) -> Option<Self> {
         match code {
             0 => Some(MetricType::L2),

core/src/index.rs

@@ -121,6 +121,18 @@ pub enum IndexType {
 }
 
 impl IndexType {
+    /// Range membership is exact for IVF-FLAT and estimated for SQ, PQ and RQ.
+    /// DiskANN does not support range search. This does not imply full IVF coverage.
+    pub fn supports_range_search(self, metric: MetricType) -> bool {
+        matches!(
+            self,
+            Self::IvfFlat | Self::IvfSq | Self::IvfPq | Self::IvfRq
+        ) && matches!(
+            metric,
+            MetricType::L2 | MetricType::Cosine | MetricType::InnerProduct
+        )
+    }
+
     pub fn from_code(code: u32) -> Option<Self> {
         match code {
             0 => Some(Self::IvfFlat),
@@ -1567,6 +1579,12 @@ pub enum VectorIndexReader<R: SeekRead> {
 }
 
 impl<R: SeekRead> VectorIndexReader<R> {
+    /// Reports family/metric capability without loading list payloads.
+    pub fn supports_range_search(&self) -> bool {
+        let metadata = self.metadata();
+        metadata.index_type.supports_range_search(metadata.metric)
+    }
+
     pub fn open(reader: R) -> io::Result<Self> {
         Self::open_with_options(reader, VectorIndexReaderOptions::default())
     }
@@ -1868,7 +1886,8 @@ impl<R: SeekRead> VectorIndexReader<R> {
     /// Distance range search. For the contract see
     /// [`IVFFlatIndexReader::range_search`] (exact distances),
     /// [`IVFRQIndexReader::range_search`] (RQ estimates), and
-    /// [`IVFSQIndexReader::range_search`] (SQ estimates). Only L2 is supported.
+    /// [`IVFSQIndexReader::range_search`] (SQ estimates), and
+    /// [`IVFPQIndexReader::range_search`] (PQ estimates). All three metrics are supported.
     ///
     /// The empty-band short-circuit lives **inside each family's reader**, so a
     /// family that cannot do range search at all still fails loud for every
@@ -1884,7 +1903,7 @@ impl<R: SeekRead> VectorIndexReader<R> {
             Self::IvfFlat(reader) => reader.range_search(query, params),
             Self::IvfRq(reader) => reader.range_search(query, params),
             Self::IvfSq(reader) => reader.range_search(query, params),
-            Self::IvfPq(_) => Err(range_unsupported("ivf_pq")),
+            Self::IvfPq(reader) => reader.range_search(query, params),
             Self::DiskAnn(_) => Err(range_unsupported("diskann")),
         }
     }
@@ -1909,7 +1928,7 @@ impl<R: SeekRead> VectorIndexReader<R> {
             Self::IvfFlat(reader) => reader.range_search_with_filter(query, params, Some(&filter)),
             Self::IvfRq(reader) => reader.range_search_with_filter(query, params, Some(&filter)),
             Self::IvfSq(reader) => reader.range_search_with_filter(query, params, Some(&filter)),
-            Self::IvfPq(_) => Err(range_unsupported("ivf_pq")),
+            Self::IvfPq(reader) => reader.range_search_with_filter(query, params, Some(&filter)),
             Self::DiskAnn(_) => Err(range_unsupported("diskann")),
         }
     }
@@ -1928,7 +1947,7 @@ impl<R: SeekRead> VectorIndexReader<R> {
             Self::IvfFlat(reader) => reader.range_search_batch(queries, query_count, params),
             Self::IvfRq(reader) => reader.range_search_batch(queries, query_count, params),
             Self::IvfSq(reader) => reader.range_search_batch(queries, query_count, params),
-            Self::IvfPq(_) => Err(range_unsupported("ivf_pq")),
+            Self::IvfPq(reader) => reader.range_search_batch(queries, query_count, params),
             Self::DiskAnn(_) => Err(range_unsupported("diskann")),
         }
     }
@@ -1954,7 +1973,9 @@ impl<R: SeekRead> VectorIndexReader<R> {
             Self::IvfSq(reader) => {
                 reader.range_search_batch_with_filter(queries, query_count, params, Some(&filter))
             }
-            Self::IvfPq(_) => Err(range_unsupported("ivf_pq")),
+            Self::IvfPq(reader) => {
+                reader.range_search_batch_with_filter(queries, query_count, params, Some(&filter))
+            }
             Self::DiskAnn(_) => Err(range_unsupported("diskann")),
         }
     }
@@ -2677,7 +2698,7 @@ fn validate_query(query: &[f32], dimension: usize) -> io::Result<()> {
     validate_finite_values(query, dimension, "query")
 }
 
-/// IVF-Flat and IVF-RQ implement range search. The other families return
+/// The IVF families implement range search. DiskANN returns
 /// `Unsupported`, meaning "we cannot serve this request, please fall back",
 /// rather than "the call has a bug". For DiskANN the reason is a lasting one:
 /// graph traversal is inherently k-oriented and has no natural radius

core/src/io.rs

@@ -934,6 +934,17 @@ impl<R: SeekRead> IVFPQIndexReader<R> {
         &mut self,
         list_id: usize,
         mut consume: impl FnMut(&ProductQuantizer, &[i64], &[u8]),
+    ) -> io::Result<()> {
+        self.try_for_each_streamed_list_chunk(list_id, |pq, ids, codes| {
+            consume(pq, ids, codes);
+            Ok(())
+        })
+    }
+
+    pub(crate) fn try_for_each_streamed_list_chunk(
+        &mut self,
+        list_id: usize,
+        mut consume: impl FnMut(&ProductQuantizer, &[i64], &[u8]) -> io::Result<()>,
     ) -> io::Result<()> {
         self.ensure_loaded()?;
         let count = self.list_counts[list_id] as usize;
@@ -1004,7 +1015,7 @@ impl<R: SeekRead> IVFPQIndexReader<R> {
                     .pread(&mut [ReadRequest::new(chunk_offset, payload.codes_mut())])?;
             }
             let row_end = row_start + chunk_rows;
-            consume(&self.pq, &ids[row_start..row_end], payload.codes());
+            consume(&self.pq, &ids[row_start..row_end], payload.codes())?;
             row_start = row_end;
         }
         Ok(())

core/src/ivfflat_io.rs

@@ -26,7 +26,10 @@ use crate::io::{ReadRequest, SeekRead, SeekWrite};
 use crate::ivfflat::IVFFlatIndex;
 use crate::ivfpq::RowIdFilter;
 use crate::kmeans;
-use crate::range::{RangeResultBuilder, RangeSearchResult, VectorRangeSearchParams};
+use crate::range::{
+    checked_cosine_norm, prepare_range_queries, range_probe_lists, RangeResultBuilder,
+    RangeSearchResult, VectorRangeSearchParams,
+};
 use rayon::prelude::*;
 use roaring::RoaringTreemap;
 use std::io;
@@ -882,12 +885,8 @@ impl<R: SeekRead> IVFFlatIndexReader<R> {
         // centroid array.
         self.ensure_loaded()?;
 
-        // Note for metric certification: unlike the top-K path this does not
-        // apply `fvec_normalize` for cosine. That is currently unreachable,
-        // because `params.validate` rejects every non-L2 metric above, but
-        // whoever certifies cosine must add the normalization here as well as
-        // relaxing `ensure_certified_metric`.
-        //
+        let processed = prepare_range_queries(queries, self.d, self.metric)?;
+        let queries = processed.as_ref();
         // Ranked probe selection, one group per query.
         //
         // This relies on `find_topk_batch` selecting the same centroids for a
@@ -899,14 +898,14 @@ impl<R: SeekRead> IVFFlatIndexReader<R> {
         // its error bound leaves the ranking ambiguous, which is what makes the
         // two agree. `kmeans` pins that property with a test; if it is ever
         // relaxed, this caller needs an exact helper of its own again.
-        let (probe_lists, _coarse_distances) = kmeans::find_topk_batch(
+        let probe_lists = range_probe_lists(
             queries,
-            nq,
             &self.quantizer_centroids,
-            self.nlist,
             self.d,
+            self.nlist,
             nprobe,
-        );
+            self.metric,
+        )?;
         for (qi, lists) in probe_lists.iter().enumerate() {
             builder.record_lists_probed(qi, lists.len());
         }
@@ -1344,6 +1343,11 @@ fn scan_flat_rows<C: Collector>(
             continue;
         }
         let vector = &vectors[local_idx * d..(local_idx + 1) * d];
+        let vector_norm = if C::VALIDATE_COSINE_INPUTS && metric == MetricType::Cosine {
+            Some(checked_cosine_norm(vector)?)
+        } else {
+            None
+        };
         let distance = if metric == MetricType::L2 {
             // One traversal, whatever the outcome: the kernel abandons against
             // the collector's cutoff and otherwise hands back the value
@@ -1359,7 +1363,7 @@ fn scan_flat_rows<C: Collector>(
                 }
             }
         } else {
-            distance_context.distance_to(vector, None)
+            distance_context.distance_to(vector, vector_norm)
         };
         collector.push(id, distance)?;
     }