This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/paimon-vector-index.git
The following commit(s) were added to refs/heads/main by this push:
new 4529c29 Optimize IVF-RQ scan and diagnostics (#66)
4529c29 is described below
commit 4529c2996a73915b97369c6790cd2e42672b30a6
Author: Jingsong Lee <[email protected]>
AuthorDate: Thu Jul 30 15:46:39 2026 +0800
Optimize IVF-RQ scan and diagnostics (#66)
---
core/benches/ann_bench.rs | 61 +++++++-
core/src/index.rs | 30 ++++
core/src/ivfrq.rs | 33 +++--
core/src/ivfrq_io.rs | 365 ++++++++++++++++++++++++++++++++++++++++------
core/src/kmeans.rs | 22 ++-
core/src/rq.rs | 360 ++++++++++++++++++++++++++++++++++++++++++++-
core/src/topk.rs | 41 +++++-
docs/index.html | 20 ++-
8 files changed, 866 insertions(+), 66 deletions(-)
diff --git a/core/benches/ann_bench.rs b/core/benches/ann_bench.rs
index 5d14976..1cb167f 100644
--- a/core/benches/ann_bench.rs
+++ b/core/benches/ann_bench.rs
@@ -884,6 +884,7 @@ fn run_query_case(
);
reset_stats(&stats);
let mut latencies = Vec::with_capacity(config.nq);
+ let mut sequential_rq_stats =
paimon_vindex_core::index::IVFRQSearchStats::default();
let mut virtual_sequential_elapsed = Duration::ZERO;
let sequential_started = Instant::now();
for (query_index, query) in
dataset.queries.chunks_exact(config.d).enumerate() {
@@ -895,6 +896,9 @@ fn run_query_case(
format!("sequential query {query_index} failed: {error}"),
)
})?;
+ if let Some(stats) = reader.ivfrq_search_stats() {
+ sequential_rq_stats.merge(stats);
+ }
let elapsed = started.elapsed();
let elapsed = if storage.virtualize_sequential_latency {
add_fixed_round_latency(
@@ -940,7 +944,44 @@ fn run_query_case(
let batch_started = Instant::now();
let (result_ids, _) = batch_reader.search_batch(&dataset.queries,
config.nq, search)?;
let batch_elapsed = batch_started.elapsed();
+ let rq_stats = batch_reader.ivfrq_search_stats().unwrap_or_default();
let batch_stats = snapshot_stats(&batch_stats);
+ let rq_refine_ratio = if rq_stats.eligible_vectors == 0 {
+ 0.0
+ } else {
+ rq_stats.refined_vectors as f64 / rq_stats.eligible_vectors as f64
+ };
+ if index.name == "IVF_RQ" {
+ eprintln!(
+ "IVF-RQ sequential scan stats: scanned={} refined={} ({:.2}%)
final={} seeded_lists={} parallel_list_tasks={}",
+ sequential_rq_stats.scanned_vectors,
+ sequential_rq_stats.refined_vectors,
+ if sequential_rq_stats.eligible_vectors == 0 {
+ 0.0
+ } else {
+ sequential_rq_stats.refined_vectors as f64
+ / sequential_rq_stats.eligible_vectors as f64
+ * 100.0
+ },
+ sequential_rq_stats.final_distance_evaluations,
+ sequential_rq_stats.seeded_lists,
+ sequential_rq_stats.parallel_list_tasks,
+ );
+ eprintln!(
+ "IVF-RQ batch scan stats: scanned={} eligible={} refined={}
({:.2}%) final={} refined_coarse_lookups={} extra_plane_lookups={}
fastscan_blocks={} scalar_blocks={} seeded_lists={} parallel_list_tasks={}",
+ rq_stats.scanned_vectors,
+ rq_stats.eligible_vectors,
+ rq_stats.refined_vectors,
+ rq_refine_ratio * 100.0,
+ rq_stats.final_distance_evaluations,
+ rq_stats.refined_coarse_byte_lookups,
+ rq_stats.extra_plane_byte_lookups,
+ rq_stats.fastscan_blocks,
+ rq_stats.scalar_blocks,
+ rq_stats.seeded_lists,
+ rq_stats.parallel_list_tasks,
+ );
+ }
let recall_at_10 = recall_at_k(&result_ids, ground_truth, config.k);
let p50 = percentile(&latencies, 50);
@@ -955,7 +996,7 @@ fn run_query_case(
_ => "none",
};
println!(
-
"{dataset},{index},{storage},{n},{train_n},{raw_dataset_bytes},{nq},{d},{k},{nlist},{nprobe},{pq_m},{rq_bits},{build_distance},{raw_vector_encoding},{l_search},{build_ms},{train_ms},{add_ms},{write_ms},{peak_rss_bytes},{optimize_ms},{optimize_rounds},{optimize_ranges},{optimize_bytes},{file_bytes},{recall:.4},{first_us},{p50_us},{p95_us},{sequential_qps:.2},{seq_rounds},{seq_ranges},{seq_bytes},{batch_ms},{batch_qps:.2},{batch_rounds},{batch_ranges},{batch_bytes}",
+
"{dataset},{index},{storage},{n},{train_n},{raw_dataset_bytes},{nq},{d},{k},{nlist},{nprobe},{pq_m},{rq_bits},{build_distance},{raw_vector_encoding},{l_search},{build_ms},{train_ms},{add_ms},{write_ms},{peak_rss_bytes},{optimize_ms},{optimize_rounds},{optimize_ranges},{optimize_bytes},{file_bytes},{recall:.4},{first_us},{p50_us},{p95_us},{sequential_qps:.2},{seq_rounds},{seq_ranges},{seq_bytes},{batch_ms},{batch_qps:.2},{batch_rounds},{batch_ranges},{batch_bytes},{rq_seq_scanned}
[...]
dataset = config.dataset_name,
index = index.name,
storage = storage.name,
@@ -998,6 +1039,22 @@ fn run_query_case(
batch_rounds = batch_stats.rounds,
batch_ranges = batch_stats.ranges,
batch_bytes = batch_stats.bytes,
+ rq_seq_scanned = sequential_rq_stats.scanned_vectors,
+ rq_seq_refined = sequential_rq_stats.refined_vectors,
+ rq_seq_final = sequential_rq_stats.final_distance_evaluations,
+ rq_seq_seeded_lists = sequential_rq_stats.seeded_lists,
+ rq_seq_parallel_list_tasks = sequential_rq_stats.parallel_list_tasks,
+ rq_scanned = rq_stats.scanned_vectors,
+ rq_eligible = rq_stats.eligible_vectors,
+ rq_refined = rq_stats.refined_vectors,
+ rq_refine_ratio = rq_refine_ratio,
+ rq_final = rq_stats.final_distance_evaluations,
+ rq_refined_coarse_lookups = rq_stats.refined_coarse_byte_lookups,
+ rq_extra_plane_lookups = rq_stats.extra_plane_byte_lookups,
+ rq_fastscan_blocks = rq_stats.fastscan_blocks,
+ rq_scalar_blocks = rq_stats.scalar_blocks,
+ rq_seeded_lists = rq_stats.seeded_lists,
+ rq_parallel_list_tasks = rq_stats.parallel_list_tasks,
);
Ok(())
}
@@ -1006,7 +1063,7 @@ struct CsvRow;
impl CsvRow {
fn header() -> &'static str {
-
"dataset,index,storage,n,train_n,raw_dataset_bytes,nq,d,k,nlist,nprobe,pq_m,rq_bits,diskann_build_distance,diskann_raw_vector_encoding,l_search,build_ms,train_ms,add_ms,write_ms,peak_rss_bytes,optimize_ms,optimize_pread_rounds,optimize_pread_ranges,optimize_pread_bytes,file_bytes,recall_at_10,first_query_us,p50_query_us,p95_query_us,sequential_qps,sequential_pread_rounds,sequential_pread_ranges,sequential_pread_bytes,batch_ms,batch_qps,batch_pread_rounds,batch_pread_ranges,batch_
[...]
+
"dataset,index,storage,n,train_n,raw_dataset_bytes,nq,d,k,nlist,nprobe,pq_m,rq_bits,diskann_build_distance,diskann_raw_vector_encoding,l_search,build_ms,train_ms,add_ms,write_ms,peak_rss_bytes,optimize_ms,optimize_pread_rounds,optimize_pread_ranges,optimize_pread_bytes,file_bytes,recall_at_10,first_query_us,p50_query_us,p95_query_us,sequential_qps,sequential_pread_rounds,sequential_pread_ranges,sequential_pread_bytes,batch_ms,batch_qps,batch_pread_rounds,batch_pread_ranges,batch_
[...]
}
}
diff --git a/core/src/index.rs b/core/src/index.rs
index 39a6efc..41d16bc 100644
--- a/core/src/index.rs
+++ b/core/src/index.rs
@@ -39,6 +39,7 @@ use crate::ivfpq::{
search_with_reader, search_with_reader_roaring_filter, IVFPQIndex,
};
use crate::ivfrq::IVFRQIndex;
+pub use crate::ivfrq_io::IVFRQSearchStats;
use crate::ivfrq_io::{
search_batch_ivfrq_reader, search_batch_ivfrq_reader_roaring_filter,
write_ivfrq_index,
IVFRQIndexReader, IVF_RQ_MAGIC,
@@ -1458,6 +1459,13 @@ impl<R: SeekRead> VectorIndexReader<R> {
}
}
+ pub fn ivfrq_search_stats(&self) -> Option<IVFRQSearchStats> {
+ match self {
+ Self::IvfRq(reader) => Some(reader.last_search_stats()),
+ _ => None,
+ }
+ }
+
pub fn read_plan(&self) -> Option<VectorIndexReadPlan> {
match self {
Self::DiskAnn(reader) => Some(reader.vector_read_plan()),
@@ -2354,6 +2362,28 @@ mod tests {
assert!(stats.rerank_unique_windows >= 1);
}
+ #[test]
+ fn ivfrq_search_stats_are_available_through_unified_reader() {
+ let dimension = 64;
+ let (mut reader, data) = build_reader(VectorIndexConfig::IvfRq {
+ dimension,
+ nlist: 8,
+ metric: MetricType::L2,
+ bits: 4,
+ });
+
+ reader
+ .search(&data[..dimension], VectorSearchParams::new(3, 8))
+ .unwrap();
+
+ let stats = reader.ivfrq_search_stats().expect("IVF-RQ diagnostics");
+ assert_eq!(stats.query_count, 1);
+ assert_eq!(stats.scanned_vectors, 512);
+ assert_eq!(stats.eligible_vectors, 512);
+ assert!(stats.refined_vectors > 0);
+ assert!(reader.diskann_search_stats().is_none());
+ }
+
#[test]
fn
diskann_batch_search_overlaps_graph_reads_and_centralizes_filtered_rerank() {
let dimension = 8;
diff --git a/core/src/ivfrq.rs b/core/src/ivfrq.rs
index a2b8b3f..b446780 100644
--- a/core/src/ivfrq.rs
+++ b/core/src/ivfrq.rs
@@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.
-use crate::distance::{fvec_madd, preprocess_vectors, MetricType};
+use crate::distance::{fvec_madd, fvec_norm_l2sqr, preprocess_vectors,
MetricType};
use crate::ivfpq::RowIdFilter;
use crate::kmeans::{self, KMeansConfig};
use crate::rq::{
@@ -33,6 +33,7 @@ pub struct IVFRQIndex {
pub bits: usize,
pub metric: MetricType,
pub quantizer_centroids: Vec<f32>,
+ pub quantizer_centroid_norms: Vec<f32>,
pub rotated_centroids: Vec<f32>,
pub rotation_seed: u64,
pub rotation_rounds: u32,
@@ -83,6 +84,7 @@ impl IVFRQIndex {
bits,
metric,
quantizer_centroids: Vec::new(),
+ quantizer_centroid_norms: Vec::new(),
rotated_centroids: Vec::new(),
rotation_seed,
rotation_rounds,
@@ -98,6 +100,11 @@ impl IVFRQIndex {
let processed = self.preprocess_vectors(data, n);
self.quantizer_centroids =
kmeans::kmeans_train(&KMeansConfig::default(), &processed, n,
self.d, self.nlist);
+ self.quantizer_centroid_norms = self
+ .quantizer_centroids
+ .chunks_exact(self.d)
+ .map(fvec_norm_l2sqr)
+ .collect();
self.rotated_centroids = vec![0.0; self.nlist * self.padded_d];
let mut scratch = vec![0.0; self.padded_d];
for list_id in 0..self.nlist {
@@ -230,10 +237,11 @@ impl IVFRQIndex {
result_labels: &mut [i64],
) {
let processed_queries = self.preprocess_vectors(queries, nq);
- let (all_probe_indices, _) = kmeans::find_topk_batch(
+ let (all_probe_indices, all_probe_distances) =
kmeans::find_topk_batch_with_centroid_norms(
&processed_queries,
nq,
&self.quantizer_centroids,
+ &self.quantizer_centroid_norms,
self.nlist,
self.d,
nprobe,
@@ -246,9 +254,18 @@ impl IVFRQIndex {
self.rotation
.rotate(query, &mut rotated_query, &mut rotation_scratch);
let query_context =
self.quantizer.prepare_query(rotated_query.clone());
+ let query_norm_sqr = fvec_norm_l2sqr(query);
let mut heap = TopKHeap::new(k);
- for &list_id in &all_probe_indices[qi] {
- self.scan_list(&query_context, list_id, filter, &mut heap);
+ for (&list_id, &coarse_distance) in
+ all_probe_indices[qi].iter().zip(&all_probe_distances[qi])
+ {
+ let query_terms =
self.quantizer.query_terms_from_coarse_distance(
+ coarse_distance,
+ query_norm_sqr,
+ self.quantizer_centroid_norms[list_id],
+ self.metric,
+ );
+ self.scan_list(&query_context, query_terms, list_id, filter,
&mut heap);
}
let sorted = heap.into_sorted();
@@ -273,20 +290,14 @@ impl IVFRQIndex {
}
}
- pub(crate) fn rotated_centroid(&self, list_id: usize) -> &[f32] {
- &self.rotated_centroids[list_id * self.padded_d..(list_id + 1) *
self.padded_d]
- }
-
fn scan_list(
&self,
query_context: &RQQueryContext,
+ query_terms: crate::rq::RQQueryTerms,
list_id: usize,
filter: Option<&dyn RowIdFilter>,
heap: &mut TopKHeap,
) {
- let query_terms =
- self.quantizer
- .query_terms(query_context, self.rotated_centroid(list_id),
self.metric);
let code_size = self.code_size();
for (local_idx, &id) in self.ids[list_id].iter().enumerate() {
if filter.map(|f| !f.contains(id)).unwrap_or(false) {
diff --git a/core/src/ivfrq_io.rs b/core/src/ivfrq_io.rs
index f23dcc5..5990d49 100644
--- a/core/src/ivfrq_io.rs
+++ b/core/src/ivfrq_io.rs
@@ -15,7 +15,7 @@
// specific language governing permissions and limitations
// under the License.
-use crate::distance::{preprocess_vectors, MetricType};
+use crate::distance::{fvec_norm_l2sqr, preprocess_vectors, MetricType};
use crate::index_io_util::{
decode_delta_varint_ids, encode_delta_varint_ids, pread_batched_payloads,
validate_search_inputs,
@@ -25,8 +25,8 @@ use crate::ivfpq::RowIdFilter;
use crate::ivfrq::IVFRQIndex;
use crate::kmeans;
use crate::rq::{
- is_supported_rq_bits, padded_dimension, RQCodeFactors, RQQueryContext,
RQRotation,
- RQVectorFactors, RaBitQuantizer, DEFAULT_RQ_ROTATION_ROUNDS,
RQ_SCAN_BLOCK_SIZE,
+ is_supported_rq_bits, padded_dimension, RQCodeFactors, RQQueryContext,
RQQueryTerms,
+ RQRotation, RQVectorFactors, RaBitQuantizer, DEFAULT_RQ_ROTATION_ROUNDS,
RQ_SCAN_BLOCK_SIZE,
};
use crate::topk::TopKHeap;
use rayon::prelude::*;
@@ -48,6 +48,54 @@ const SUPPORTED_FLAGS: u32 = REQUIRED_FLAGS;
const FACTOR_BYTES: usize = 4;
const MAX_RQ_BATCH_READ_BYTES: usize = 64 * 1024 * 1024;
const PARALLEL_RQ_SCAN_MIN_CANDIDATES: usize = 8 * 1024;
+const PARALLEL_RQ_SEED_LISTS: usize = 1;
+const PARALLEL_RQ_SEED_VECTORS: usize = RQ_SCAN_BLOCK_SIZE;
+const FASTSCAN_MIN_PADDED_DIMENSION: usize = 256;
+
+#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
+pub struct IVFRQSearchStats {
+ pub query_count: usize,
+ pub scanned_vectors: usize,
+ pub eligible_vectors: usize,
+ pub coarse_distance_evaluations: usize,
+ pub refined_vectors: usize,
+ pub refined_coarse_byte_lookups: usize,
+ pub extra_plane_byte_lookups: usize,
+ pub final_distance_evaluations: usize,
+ pub heap_admissions: usize,
+ pub fastscan_blocks: usize,
+ pub scalar_blocks: usize,
+ pub seeded_lists: usize,
+ pub parallel_list_tasks: usize,
+}
+
+impl IVFRQSearchStats {
+ pub fn merge(&mut self, other: Self) {
+ self.query_count = self.query_count.saturating_add(other.query_count);
+ self.scanned_vectors =
self.scanned_vectors.saturating_add(other.scanned_vectors);
+ self.eligible_vectors =
self.eligible_vectors.saturating_add(other.eligible_vectors);
+ self.coarse_distance_evaluations = self
+ .coarse_distance_evaluations
+ .saturating_add(other.coarse_distance_evaluations);
+ self.refined_vectors =
self.refined_vectors.saturating_add(other.refined_vectors);
+ self.refined_coarse_byte_lookups = self
+ .refined_coarse_byte_lookups
+ .saturating_add(other.refined_coarse_byte_lookups);
+ self.extra_plane_byte_lookups = self
+ .extra_plane_byte_lookups
+ .saturating_add(other.extra_plane_byte_lookups);
+ self.final_distance_evaluations = self
+ .final_distance_evaluations
+ .saturating_add(other.final_distance_evaluations);
+ self.heap_admissions =
self.heap_admissions.saturating_add(other.heap_admissions);
+ self.fastscan_blocks =
self.fastscan_blocks.saturating_add(other.fastscan_blocks);
+ self.scalar_blocks =
self.scalar_blocks.saturating_add(other.scalar_blocks);
+ self.seeded_lists =
self.seeded_lists.saturating_add(other.seeded_lists);
+ self.parallel_list_tasks = self
+ .parallel_list_tasks
+ .saturating_add(other.parallel_list_tasks);
+ }
+}
struct RQListWritePlan {
order: Vec<usize>,
@@ -200,12 +248,13 @@ pub struct IVFRQIndexReader<R: SeekRead> {
pub rotation_type: u32,
pub factor_layout: u32,
pub quantizer_centroids: Vec<f32>,
- pub rotated_centroids: Vec<f32>,
+ pub quantizer_centroid_norms: Vec<f32>,
pub list_offsets: Vec<i64>,
pub list_counts: Vec<i32>,
pub list_id_bytes_lens: Vec<i32>,
quantizer: RaBitQuantizer,
rotation: RQRotation,
+ last_search_stats: IVFRQSearchStats,
loaded: bool,
}
@@ -313,16 +362,21 @@ impl<R: SeekRead> IVFRQIndexReader<R> {
rotation_type,
factor_layout,
quantizer_centroids: Vec::new(),
- rotated_centroids: Vec::new(),
+ quantizer_centroid_norms: Vec::new(),
list_offsets: Vec::new(),
list_counts: Vec::new(),
list_id_bytes_lens: Vec::new(),
quantizer: RaBitQuantizer::new(d, num_bits),
rotation: RQRotation::new(d, rotation_seed, rotation_rounds),
+ last_search_stats: IVFRQSearchStats::default(),
loaded: false,
})
}
+ pub fn last_search_stats(&self) -> IVFRQSearchStats {
+ self.last_search_stats
+ }
+
pub fn ensure_loaded(&mut self) -> io::Result<()> {
if self.loaded {
return Ok(());
@@ -360,15 +414,11 @@ impl<R: SeekRead> IVFRQIndexReader<R> {
)));
}
- self.rotated_centroids = vec![0.0; self.nlist * self.padded_d];
- let mut scratch = vec![0.0; self.padded_d];
- for list_id in 0..self.nlist {
- self.rotation.rotate(
- &self.quantizer_centroids[list_id * self.d..(list_id + 1) *
self.d],
- &mut self.rotated_centroids[list_id * self.padded_d..(list_id
+ 1) * self.padded_d],
- &mut scratch,
- );
- }
+ self.quantizer_centroid_norms = self
+ .quantizer_centroids
+ .chunks_exact(self.d)
+ .map(fvec_norm_l2sqr)
+ .collect();
self.loaded = true;
Ok(())
}
@@ -586,14 +636,38 @@ pub fn search_batch_ivfrq_reader_filter<R: SeekRead>(
reader.ensure_loaded()?;
validate_search_inputs(queries, nq, reader.d, k, nprobe)?;
let processed = preprocess_vectors(queries, nq, reader.d, reader.metric);
- let (all_probe_indices, _) = kmeans::find_topk_batch(
+ let (all_probe_indices, all_probe_distances) =
kmeans::find_topk_batch_with_centroid_norms(
&processed,
nq,
&reader.quantizer_centroids,
+ &reader.quantizer_centroid_norms,
reader.nlist,
reader.d,
nprobe,
);
+ let query_norms = processed
+ .chunks_exact(reader.d)
+ .map(fvec_norm_l2sqr)
+ .collect::<Vec<_>>();
+ let all_query_terms = all_probe_indices
+ .iter()
+ .zip(&all_probe_distances)
+ .enumerate()
+ .map(|(query_index, (indices, distances))| {
+ indices
+ .iter()
+ .zip(distances)
+ .map(|(&list_id, &distance)| {
+ reader.quantizer.query_terms_from_coarse_distance(
+ distance,
+ query_norms[query_index],
+ reader.quantizer_centroid_norms[list_id],
+ reader.metric,
+ )
+ })
+ .collect::<Vec<_>>()
+ })
+ .collect::<Vec<_>>();
let mut query_contexts = Vec::with_capacity(nq);
let mut rotated = vec![0.0; reader.padded_d];
@@ -615,6 +689,11 @@ pub fn search_batch_ivfrq_reader_filter<R: SeekRead>(
}
let mut heaps: Vec<TopKHeap> = (0..nq).map(|_| TopKHeap::new(k)).collect();
+ let mut query_stats = vec![IVFRQSearchStats::default(); nq];
+ let mut aggregate_stats = IVFRQSearchStats {
+ query_count: nq,
+ ..IVFRQSearchStats::default()
+ };
let mut batch_start = 0;
while batch_start < unique_lists.len() {
let count = reader.batch_read_end(&unique_lists[batch_start..])?;
@@ -625,34 +704,66 @@ pub fn search_batch_ivfrq_reader_filter<R: SeekRead>(
list_positions[list.list_id] = position;
}
let quantizer = &reader.quantizer;
- let metric = reader.metric;
- let rotated_centroids = &reader.rotated_centroids;
- let padded_d = reader.padded_d;
let candidate_count = loaded_lists
.iter()
.map(|list| list.ids.len())
.sum::<usize>();
if nq == 1 && candidate_count >= PARALLEL_RQ_SCAN_MIN_CANDIDATES {
+ let mut seeded_lists = 0usize;
+ if PARALLEL_RQ_SEED_LISTS > 0 {
+ for list in loaded_lists
+ .iter()
+ .filter(|list| !list.ids.is_empty())
+ .take(PARALLEL_RQ_SEED_LISTS)
+ {
+ let probe_position = all_probe_indices[0]
+ .iter()
+ .position(|&probe| probe == list.list_id)
+ .expect("loaded list must belong to the query probe
set");
+ scan_blocked_list(
+ list,
+ quantizer,
+ &query_contexts[0],
+ all_query_terms[0][probe_position],
+ filter,
+ &mut heaps[0],
+ &mut aggregate_stats,
+ PARALLEL_RQ_SEED_VECTORS,
+ );
+ seeded_lists += 1;
+ }
+ }
+ aggregate_stats.seeded_lists =
+ aggregate_stats.seeded_lists.saturating_add(seeded_lists);
+ let seeded_threshold =
heaps[0].worst_distance().unwrap_or(f32::INFINITY);
let per_list_results = loaded_lists
.par_iter()
.map(|list| {
- let mut heap = TopKHeap::new(k);
+ let mut heap = TopKHeap::with_max_distance(k,
seeded_threshold);
+ let mut stats = IVFRQSearchStats::default();
let list_id = list.list_id;
- let rotated_centroid =
- &rotated_centroids[list_id * padded_d..(list_id + 1) *
padded_d];
+ let probe_position = all_probe_indices[0]
+ .iter()
+ .position(|&probe| probe == list_id)
+ .expect("loaded list must belong to the query probe
set");
scan_blocked_list(
list,
quantizer,
- metric,
&query_contexts[0],
- rotated_centroid,
+ all_query_terms[0][probe_position],
filter,
&mut heap,
+ &mut stats,
+ usize::MAX,
);
- heap.into_sorted()
+ (heap.into_sorted(), stats)
})
.collect::<Vec<_>>();
- for results in per_list_results {
+ aggregate_stats.parallel_list_tasks = aggregate_stats
+ .parallel_list_tasks
+ .saturating_add(per_list_results.len());
+ for (results, stats) in per_list_results {
+ aggregate_stats.merge(stats);
for (distance, row_id) in results {
heaps[0].push(distance, row_id);
}
@@ -660,29 +771,35 @@ pub fn search_batch_ivfrq_reader_filter<R: SeekRead>(
} else {
heaps
.par_iter_mut()
+ .zip(query_stats.par_iter_mut())
.enumerate()
- .for_each(|(query_index, heap)| {
- for &list_id in &all_probe_indices[query_index] {
+ .for_each(|(query_index, (heap, stats))| {
+ for (probe_position, &list_id) in
+ all_probe_indices[query_index].iter().enumerate()
+ {
let position = list_positions[list_id];
if position == usize::MAX {
continue;
}
- let rotated_centroid =
- &rotated_centroids[list_id * padded_d..(list_id +
1) * padded_d];
scan_blocked_list(
&loaded_lists[position],
quantizer,
- metric,
&query_contexts[query_index],
- rotated_centroid,
+ all_query_terms[query_index][probe_position],
filter,
heap,
+ stats,
+ usize::MAX,
);
}
});
}
batch_start = batch_end;
}
+ for stats in query_stats {
+ aggregate_stats.merge(stats);
+ }
+ reader.last_search_stats = aggregate_stats;
let mut result_ids = vec![-1; nq * k];
let mut result_distances = vec![f32::MAX; nq * k];
@@ -711,14 +828,14 @@ pub fn search_batch_ivfrq_reader_roaring_filter<R:
SeekRead>(
fn scan_blocked_list(
list: &RQReadList,
quantizer: &RaBitQuantizer,
- metric: MetricType,
query: &RQQueryContext,
- rotated_centroid: &[f32],
+ query_terms: RQQueryTerms,
filter: Option<&dyn RowIdFilter>,
heap: &mut TopKHeap,
+ stats: &mut IVFRQSearchStats,
+ vector_limit: usize,
) {
let blocked_codes = list.blocked_codes();
- let query_terms = quantizer.query_terms(query, rotated_centroid, metric);
let plane_size = quantizer.plane_size();
let bits = quantizer.bits();
let code_size = quantizer.code_size();
@@ -726,14 +843,29 @@ fn scan_blocked_list(
let query_sum = quantizer.query_sum(query);
let center = ((1usize << bits) - 1) as f32 * 0.5;
+ let scan_end = list.ids.len().min(vector_limit);
let mut block_start = 0usize;
- while block_start < list.ids.len() {
- let lanes = (list.ids.len() - block_start).min(RQ_SCAN_BLOCK_SIZE);
+ while block_start < scan_end {
+ let lanes = (scan_end - block_start).min(RQ_SCAN_BLOCK_SIZE);
+ stats.scanned_vectors = stats.scanned_vectors.saturating_add(lanes);
let code_block_start = block_start * code_size;
let factor_block_start = block_start * factor_fields;
let mut allowed = [true; RQ_SCAN_BLOCK_SIZE];
let mut coarse_unsigned = [0.0f32; RQ_SCAN_BLOCK_SIZE];
- if let Some(filter) = filter {
+ let used_fastscan = bits > 1
+ && quantizer.padded_dimension() >= FASTSCAN_MIN_PADDED_DIMENSION
+ && filter.is_none()
+ && lanes == RQ_SCAN_BLOCK_SIZE;
+ if used_fastscan {
+ quantizer.fastscan_coarse_block(
+ query,
+ &blocked_codes
+ [code_block_start..code_block_start + plane_size *
RQ_SCAN_BLOCK_SIZE],
+ &mut coarse_unsigned,
+ );
+ stats.fastscan_blocks = stats.fastscan_blocks.saturating_add(1);
+ } else if let Some(filter) = filter {
+ stats.scalar_blocks = stats.scalar_blocks.saturating_add(1);
for lane in 0..lanes {
allowed[lane] = filter.contains(list.ids[block_start + lane]);
}
@@ -750,6 +882,7 @@ fn scan_blocked_list(
}
}
} else {
+ stats.scalar_blocks = stats.scalar_blocks.saturating_add(1);
for byte_idx in 0..plane_size {
let byte_start = code_block_start + byte_idx * lanes;
for lane in 0..lanes {
@@ -761,6 +894,10 @@ fn scan_blocked_list(
}
}
}
+ let eligible = allowed[..lanes].iter().filter(|&&value| value).count();
+ stats.eligible_vectors =
stats.eligible_vectors.saturating_add(eligible);
+ stats.coarse_distance_evaluations =
+ stats.coarse_distance_evaluations.saturating_add(eligible);
let mut refine = [false; RQ_SCAN_BLOCK_SIZE];
if bits == 1 {
@@ -775,6 +912,7 @@ fn scan_blocked_list(
query_terms,
);
if heap.should_consider(distance) {
+ stats.heap_admissions =
stats.heap_admissions.saturating_add(1);
heap.push(distance, list.ids[block_start + lane]);
}
}
@@ -792,9 +930,17 @@ fn scan_blocked_list(
factors,
query_terms,
);
- refine[lane] =
- heap.should_consider(quantizer.lower_bound(estimate, factors,
query_terms));
+ let mut lower = quantizer.lower_bound(estimate, factors,
query_terms);
+ if used_fastscan {
+ lower -= factors.f_rescale.abs() *
quantizer.fastscan_ip_error(query);
+ }
+ refine[lane] = heap.should_consider(lower);
}
+ let refined = refine[..lanes].iter().filter(|&&value| value).count();
+ stats.refined_vectors = stats.refined_vectors.saturating_add(refined);
+ stats.extra_plane_byte_lookups = stats
+ .extra_plane_byte_lookups
+ .saturating_add(refined.saturating_mul(bits -
1).saturating_mul(plane_size));
let mut full_unsigned = [0.0f32; RQ_SCAN_BLOCK_SIZE];
let sign_weight = (1usize << (bits - 1)) as f32;
@@ -823,12 +969,49 @@ fn scan_blocked_list(
continue;
}
let factors = read_block_factor(list, factor_block_start, lanes,
3, lane, false);
- let distance = quantizer.estimate(
+ let mut distance = quantizer.estimate(
full_unsigned[lane] - center * query_sum,
factors,
query_terms,
);
+ if used_fastscan {
+ let distance_error =
+ factors.f_rescale.abs() * sign_weight *
quantizer.fastscan_ip_error(query);
+ if !heap.should_consider(distance - distance_error) {
+ continue;
+ }
+ let mut exact_coarse = 0.0f32;
+ stats.refined_coarse_byte_lookups =
+
stats.refined_coarse_byte_lookups.saturating_add(plane_size);
+ for byte_idx in 0..plane_size {
+ exact_coarse += quantizer.byte_subset_sum(
+ query,
+ byte_idx,
+ blocked_codes[code_block_start + byte_idx * lanes +
lane],
+ );
+ }
+ let mut exact_unsigned = sign_weight * exact_coarse;
+ stats.extra_plane_byte_lookups = stats
+ .extra_plane_byte_lookups
+ .saturating_add((bits - 1).saturating_mul(plane_size));
+ for plane in 1..bits {
+ let weight = (1usize << (bits - 1 - plane)) as f32;
+ let plane_start = code_block_start + plane * plane_size *
lanes;
+ for byte_idx in 0..plane_size {
+ exact_unsigned += weight
+ * quantizer.byte_subset_sum(
+ query,
+ byte_idx,
+ blocked_codes[plane_start + byte_idx * lanes +
lane],
+ );
+ }
+ }
+ distance =
+ quantizer.estimate(exact_unsigned - center * query_sum,
factors, query_terms);
+ }
+ stats.final_distance_evaluations =
stats.final_distance_evaluations.saturating_add(1);
if heap.should_consider(distance) {
+ stats.heap_admissions =
stats.heap_admissions.saturating_add(1);
heap.push(distance, list.ids[block_start + lane]);
}
}
@@ -1144,7 +1327,7 @@ mod tests {
#[test]
fn ivfrq_unfiltered_scan_matches_all_rows_filter() {
- let d = 64;
+ let d = FASTSCAN_MIN_PADDED_DIMENSION;
let nlist = 4;
let n = 256;
let nq = 8;
@@ -1178,6 +1361,54 @@ mod tests {
.unwrap();
assert_eq!(unfiltered, filtered);
+ assert!(unfiltered_reader.last_search_stats().fastscan_blocks > 0);
+ assert_eq!(filtered_reader.last_search_stats().fastscan_blocks, 0);
+ }
+
+ #[test]
+ fn ivfrq_search_stats_report_two_stage_filtering_work() {
+ let d = FASTSCAN_MIN_PADDED_DIMENSION;
+ let nlist = 4;
+ let n = 256;
+ let nq = 8;
+ let data = (0..n)
+ .flat_map(|row| {
+ let cluster = (row % nlist) as f32 * 40.0;
+ (0..d).map(move |dimension| cluster + row as f32 * 0.003 +
dimension as f32 * 0.01)
+ })
+ .collect::<Vec<_>>();
+ let ids = (0..n as i64).collect::<Vec<_>>();
+ let mut index = IVFRQIndex::with_bits(d, nlist, 4, MetricType::L2);
+ index.train(&data, n);
+ index.add(&data, &ids, n);
+ let mut bytes = Vec::new();
+ write_ivfrq_index(&index, &mut PosWriter::new(&mut bytes)).unwrap();
+
+ let mut reader = IVFRQIndexReader::open(Cursor::new(bytes)).unwrap();
+ search_batch_ivfrq_reader(&mut reader, &data[..nq * d], nq, 10,
nlist).unwrap();
+
+ let stats = reader.last_search_stats();
+ assert_eq!(stats.query_count, nq);
+ assert_eq!(stats.scanned_vectors, nq * n);
+ assert_eq!(stats.eligible_vectors, nq * n);
+ assert_eq!(stats.coarse_distance_evaluations, nq * n);
+ assert!(stats.refined_vectors > 0);
+ assert!(stats.refined_vectors <= stats.eligible_vectors);
+ assert!(stats.final_distance_evaluations > 0);
+ assert!(stats.final_distance_evaluations <= stats.refined_vectors);
+ assert_eq!(
+ stats.extra_plane_byte_lookups,
+ (stats.refined_vectors + stats.refined_coarse_byte_lookups /
index.plane_size())
+ * (index.bits - 1)
+ * index.plane_size()
+ );
+ assert!(stats.heap_admissions > 0);
+ assert!(stats.fastscan_blocks > 0);
+ assert!(stats.refined_coarse_byte_lookups > 0);
+ assert_eq!(
+ stats.refined_coarse_byte_lookups,
+ stats.final_distance_evaluations * index.plane_size()
+ );
}
#[test]
@@ -1276,6 +1507,58 @@ mod tests {
assert_eq!(stats.lock().unwrap().max_ranges_per_batch, nlist);
}
+ #[test]
+ fn ivfrq_seeded_parallel_single_query_matches_sequential_index() {
+ let d = 16;
+ let nlist = 8;
+ let n = 9_216;
+ let k = 10;
+ let data = (0..n)
+ .flat_map(|row| {
+ let cluster = (row % nlist) as f32 * 50.0;
+ (0..d).map(move |dimension| cluster + row as f32 * 0.0007 +
dimension as f32 * 0.01)
+ })
+ .collect::<Vec<_>>();
+ let ids = (0..n as i64).collect::<Vec<_>>();
+ let mut index = IVFRQIndex::with_bits(d, nlist, 4, MetricType::L2);
+ index.train(&data, n);
+ index.add(&data, &ids, n);
+
+ let query = &data[137 * d..138 * d];
+ let mut expected_ids = vec![-1; k];
+ let mut expected_distances = vec![f32::MAX; k];
+ index.search(
+ query,
+ 1,
+ k,
+ nlist,
+ &mut expected_distances,
+ &mut expected_ids,
+ );
+
+ let mut bytes = Vec::new();
+ write_ivfrq_index(&index, &mut PosWriter::new(&mut bytes)).unwrap();
+ let mut reader = IVFRQIndexReader::open(Cursor::new(bytes)).unwrap();
+ let (actual_ids, actual_distances) = rayon::ThreadPoolBuilder::new()
+ .num_threads(4)
+ .build()
+ .unwrap()
+ .install(|| reader.search(query, k, nlist).unwrap());
+
+ assert_eq!(actual_ids, expected_ids);
+ for (actual, expected) in
actual_distances.iter().zip(expected_distances) {
+ assert!((actual - expected).abs() <= 1e-3);
+ }
+ let stats = reader.last_search_stats();
+ assert_eq!(stats.seeded_lists, 1);
+ assert_eq!(stats.parallel_list_tasks, nlist);
+ assert_eq!(
+ stats.scanned_vectors,
+ n + PARALLEL_RQ_SEED_VECTORS,
+ "the seed prefix is intentionally rescanned by its parallel list
task"
+ );
+ }
+
#[test]
fn ivfrq_batch_respects_reader_range_capability() {
let d = 64;
diff --git a/core/src/kmeans.rs b/core/src/kmeans.rs
index 8d03e8f..ede8cb7 100644
--- a/core/src/kmeans.rs
+++ b/core/src/kmeans.rs
@@ -537,6 +537,22 @@ pub fn find_topk_batch(
d: usize,
nprobe: usize,
) -> (Vec<Vec<usize>>, Vec<Vec<f32>>) {
+ let centroid_norms = (0..k)
+ .map(|c| fvec_norm_l2sqr(¢roids[c * d..(c + 1) * d]))
+ .collect::<Vec<_>>();
+ find_topk_batch_with_centroid_norms(queries, nq, centroids,
¢roid_norms, k, d, nprobe)
+}
+
+pub(crate) fn find_topk_batch_with_centroid_norms(
+ queries: &[f32],
+ nq: usize,
+ centroids: &[f32],
+ centroid_norms: &[f32],
+ k: usize,
+ d: usize,
+ nprobe: usize,
+) -> (Vec<Vec<usize>>, Vec<Vec<f32>>) {
+ debug_assert_eq!(centroid_norms.len(), k);
let nprobe = nprobe.min(k);
if nprobe == 0 {
return (vec![Vec::new(); nq], vec![Vec::new(); nq]);
@@ -551,10 +567,6 @@ pub fn find_topk_batch(
let q_norms: Vec<f32> = (0..nq)
.map(|i| fvec_norm_l2sqr(&queries[i * d..(i + 1) * d]))
.collect();
- let c_norms: Vec<f32> = (0..k)
- .map(|c| fvec_norm_l2sqr(¢roids[c * d..(c + 1) * d]))
- .collect();
-
// Batch inner products: ip[nq × k] = queries[nq × d] · centroids[k × d]^T
let mut ip_matrix = vec![0.0f32; nq * k];
sgemm_a_bt(nq, k, d, 1.0, queries, centroids, 0.0, &mut ip_matrix);
@@ -567,7 +579,7 @@ pub fn find_topk_batch(
let row = qi * k;
let mut dists: Vec<(f32, usize)> = (0..k)
.map(|c| {
- let dist = q_norms[qi] + c_norms[c] - 2.0 * ip_matrix[row + c];
+ let dist = q_norms[qi] + centroid_norms[c] - 2.0 *
ip_matrix[row + c];
(dist.max(0.0), c)
})
.collect();
diff --git a/core/src/rq.rs b/core/src/rq.rs
index 3377586..2bf8074 100644
--- a/core/src/rq.rs
+++ b/core/src/rq.rs
@@ -176,6 +176,10 @@ pub struct RQQueryContext {
rotated_query: Vec<f32>,
sum: f32,
byte_subset_sums: Vec<f32>,
+ fastscan_lut: Vec<u8>,
+ fastscan_offset: f32,
+ fastscan_scale: f32,
+ fastscan_ip_error: f32,
}
#[derive(Debug, Clone, Copy)]
@@ -307,10 +311,61 @@ impl RaBitQuantizer {
lut[pattern] = lut[previous] + rotated_query[dim_base + bit];
}
}
+ let nibble_groups = self.padded_d / 4;
+ let mut nibble_subset_sums = vec![0.0f32; nibble_groups * 16];
+ let mut group_mins = vec![0.0f32; nibble_groups];
+ let mut max_group_range = 0.0f32;
+ for group in 0..nibble_groups {
+ let dim_base = group * 4;
+ let lut = &mut nibble_subset_sums[group * 16..(group + 1) * 16];
+ for pattern in 1..16usize {
+ let bit = pattern.trailing_zeros() as usize;
+ let previous = pattern & (pattern - 1);
+ lut[pattern] = lut[previous] + rotated_query[dim_base + bit];
+ }
+ let min = lut.iter().copied().fold(f32::INFINITY, f32::min);
+ let max = lut.iter().copied().fold(f32::NEG_INFINITY, f32::max);
+ group_mins[group] = min;
+ max_group_range = max_group_range.max(max - min);
+ }
+
+ let fastscan_scale = max_group_range / u8::MAX as f32;
+ let mut fastscan_lut = vec![0u8; nibble_subset_sums.len()];
+ let mut fastscan_ip_error = 0.0f32;
+ for group in 0..nibble_groups {
+ let min = group_mins[group];
+ let exact = &nibble_subset_sums[group * 16..(group + 1) * 16];
+ let quantized = &mut fastscan_lut[group * 16..(group + 1) * 16];
+ let mut group_error = 0.0f32;
+ for (exact, quantized) in exact.iter().zip(quantized) {
+ let value = if fastscan_scale <= f32::EPSILON {
+ 0
+ } else {
+ ((*exact - min) / fastscan_scale)
+ .round()
+ .clamp(0.0, u8::MAX as f32) as u8
+ };
+ *quantized = value;
+ let reconstructed = min + fastscan_scale * value as f32;
+ group_error = group_error.max((reconstructed - *exact).abs());
+ }
+ fastscan_ip_error += group_error;
+ }
+ let fastscan_offset = group_mins.iter().sum::<f32>();
+ let rounding_guard = (fastscan_offset.abs()
+ + fastscan_scale * u8::MAX as f32 * nibble_groups as f32
+ + rotated_query.iter().map(|value| value.abs()).sum::<f32>())
+ * f32::EPSILON
+ * 8.0;
+ fastscan_ip_error += rounding_guard;
RQQueryContext {
rotated_query,
sum,
byte_subset_sums,
+ fastscan_lut,
+ fastscan_offset,
+ fastscan_scale,
+ fastscan_ip_error,
}
}
@@ -344,11 +399,36 @@ impl RaBitQuantizer {
}
}
+ pub fn query_terms_from_coarse_distance(
+ &self,
+ residual_norm_sqr: f32,
+ query_norm_sqr: f32,
+ centroid_norm_sqr: f32,
+ metric: MetricType,
+ ) -> RQQueryTerms {
+ let residual_norm_sqr = residual_norm_sqr.max(0.0);
+ let g_error = residual_norm_sqr.sqrt();
+ match metric {
+ MetricType::L2 => RQQueryTerms {
+ g_add: residual_norm_sqr,
+ g_error,
+ },
+ MetricType::Cosine => RQQueryTerms {
+ g_add: 0.5 * residual_norm_sqr,
+ g_error,
+ },
+ MetricType::InnerProduct => RQQueryTerms {
+ g_add: -0.5 * (query_norm_sqr + centroid_norm_sqr -
residual_norm_sqr),
+ g_error,
+ },
+ }
+ }
+
pub fn unsigned_plane_inner_product(&self, context: &RQQueryContext,
plane_code: &[u8]) -> f32 {
debug_assert!(plane_code.len() >= self.plane_size);
let mut result = 0.0f32;
for byte_idx in 0..self.plane_size {
- result += context.byte_subset_sums[byte_idx * 256 +
plane_code[byte_idx] as usize];
+ result += self.byte_subset_sum(context, byte_idx,
plane_code[byte_idx]);
}
result
}
@@ -362,6 +442,56 @@ impl RaBitQuantizer {
context.byte_subset_sums[byte_idx * 256 + pattern as usize]
}
+ pub(crate) fn fastscan_ip_error(&self, context: &RQQueryContext) -> f32 {
+ context.fastscan_ip_error
+ }
+
+ pub(crate) fn fastscan_coarse_block(
+ &self,
+ context: &RQQueryContext,
+ blocked_codes: &[u8],
+ output: &mut [f32; RQ_SCAN_BLOCK_SIZE],
+ ) {
+ debug_assert!(blocked_codes.len() >= self.plane_size *
RQ_SCAN_BLOCK_SIZE);
+ #[cfg(target_arch = "x86_64")]
+ {
+ if is_x86_feature_detected!("avx2") {
+ unsafe {
+ fastscan_coarse_block_avx2(
+ &context.fastscan_lut,
+ blocked_codes,
+ self.plane_size,
+ context.fastscan_offset,
+ context.fastscan_scale,
+ output,
+ );
+ }
+ return;
+ }
+ }
+ #[cfg(target_arch = "aarch64")]
+ unsafe {
+ fastscan_coarse_block_neon(
+ &context.fastscan_lut,
+ blocked_codes,
+ self.plane_size,
+ context.fastscan_offset,
+ context.fastscan_scale,
+ output,
+ );
+ return;
+ }
+ #[allow(unreachable_code)]
+ fastscan_coarse_block_scalar(
+ &context.fastscan_lut,
+ blocked_codes,
+ self.plane_size,
+ context.fastscan_offset,
+ context.fastscan_scale,
+ output,
+ );
+ }
+
pub(crate) fn query_sum(&self, context: &RQQueryContext) -> f32 {
context.sum
}
@@ -402,6 +532,145 @@ impl RaBitQuantizer {
}
}
+fn fastscan_coarse_block_scalar(
+ lut: &[u8],
+ blocked_codes: &[u8],
+ plane_size: usize,
+ offset: f32,
+ scale: f32,
+ output: &mut [f32; RQ_SCAN_BLOCK_SIZE],
+) {
+ let mut quantized = [0u32; RQ_SCAN_BLOCK_SIZE];
+ for byte_idx in 0..plane_size {
+ let code =
+ &blocked_codes[byte_idx * RQ_SCAN_BLOCK_SIZE..(byte_idx + 1) *
RQ_SCAN_BLOCK_SIZE];
+ let lut = &lut[byte_idx * 32..(byte_idx + 1) * 32];
+ for lane in 0..RQ_SCAN_BLOCK_SIZE {
+ let value = code[lane];
+ quantized[lane] += lut[(value & 0x0F) as usize] as u32;
+ quantized[lane] += lut[16 + (value >> 4) as usize] as u32;
+ }
+ }
+ for lane in 0..RQ_SCAN_BLOCK_SIZE {
+ output[lane] = offset + scale * quantized[lane] as f32;
+ }
+}
+
+#[cfg(target_arch = "aarch64")]
+#[target_feature(enable = "neon")]
+unsafe fn fastscan_coarse_block_neon(
+ lut: &[u8],
+ blocked_codes: &[u8],
+ plane_size: usize,
+ offset: f32,
+ scale: f32,
+ output: &mut [f32; RQ_SCAN_BLOCK_SIZE],
+) {
+ use std::arch::aarch64::*;
+
+ const MAX_BYTES_PER_U16_CHUNK: usize = 120;
+ let mut totals = [0u32; RQ_SCAN_BLOCK_SIZE];
+ for chunk_start in (0..plane_size).step_by(MAX_BYTES_PER_U16_CHUNK) {
+ let chunk_end = (chunk_start +
MAX_BYTES_PER_U16_CHUNK).min(plane_size);
+ let mut accumulators = [vdupq_n_u16(0); 4];
+ for byte_idx in chunk_start..chunk_end {
+ let lut_low = vld1q_u8(lut.as_ptr().add(byte_idx * 32));
+ let lut_high = vld1q_u8(lut.as_ptr().add(byte_idx * 32 + 16));
+ let code_ptr = blocked_codes.as_ptr().add(byte_idx *
RQ_SCAN_BLOCK_SIZE);
+ for half in 0..2 {
+ let code = vld1q_u8(code_ptr.add(half * 16));
+ let low = vandq_u8(code, vdupq_n_u8(0x0F));
+ let high = vshrq_n_u8(code, 4);
+ let low_values = vqtbl1q_u8(lut_low, low);
+ let high_values = vqtbl1q_u8(lut_high, high);
+ accumulators[half * 2] = vaddq_u16(
+ accumulators[half * 2],
+ vaddq_u16(
+ vmovl_u8(vget_low_u8(low_values)),
+ vmovl_u8(vget_low_u8(high_values)),
+ ),
+ );
+ accumulators[half * 2 + 1] = vaddq_u16(
+ accumulators[half * 2 + 1],
+ vaddq_u16(
+ vmovl_u8(vget_high_u8(low_values)),
+ vmovl_u8(vget_high_u8(high_values)),
+ ),
+ );
+ }
+ }
+ let mut chunk = [0u16; RQ_SCAN_BLOCK_SIZE];
+ for (part, accumulator) in accumulators.into_iter().enumerate() {
+ vst1q_u16(chunk.as_mut_ptr().add(part * 8), accumulator);
+ }
+ for lane in 0..RQ_SCAN_BLOCK_SIZE {
+ totals[lane] += chunk[lane] as u32;
+ }
+ }
+ for lane in 0..RQ_SCAN_BLOCK_SIZE {
+ output[lane] = offset + scale * totals[lane] as f32;
+ }
+}
+
+#[cfg(target_arch = "x86_64")]
+#[target_feature(enable = "avx2")]
+unsafe fn fastscan_coarse_block_avx2(
+ lut: &[u8],
+ blocked_codes: &[u8],
+ plane_size: usize,
+ offset: f32,
+ scale: f32,
+ output: &mut [f32; RQ_SCAN_BLOCK_SIZE],
+) {
+ use std::arch::x86_64::*;
+
+ const MAX_BYTES_PER_U16_CHUNK: usize = 120;
+ let mut totals = [0u32; RQ_SCAN_BLOCK_SIZE];
+ for chunk_start in (0..plane_size).step_by(MAX_BYTES_PER_U16_CHUNK) {
+ let chunk_end = (chunk_start +
MAX_BYTES_PER_U16_CHUNK).min(plane_size);
+ let mut low_accumulator = _mm256_setzero_si256();
+ let mut high_accumulator = _mm256_setzero_si256();
+ for byte_idx in chunk_start..chunk_end {
+ let lut_low = _mm256_broadcastsi128_si256(_mm_loadu_si128(
+ lut.as_ptr().add(byte_idx * 32) as *const __m128i,
+ ));
+ let lut_high = _mm256_broadcastsi128_si256(_mm_loadu_si128(
+ lut.as_ptr().add(byte_idx * 32 + 16) as *const __m128i,
+ ));
+ let code = _mm256_loadu_si256(
+ blocked_codes.as_ptr().add(byte_idx * RQ_SCAN_BLOCK_SIZE) as
*const __m256i,
+ );
+ let low = _mm256_and_si256(code, _mm256_set1_epi8(0x0F));
+ let high = _mm256_and_si256(_mm256_srli_epi16(code, 4),
_mm256_set1_epi8(0x0F));
+ let low_values = _mm256_shuffle_epi8(lut_low, low);
+ let high_values = _mm256_shuffle_epi8(lut_high, high);
+ low_accumulator = _mm256_add_epi16(
+ low_accumulator,
+ _mm256_add_epi16(
+ _mm256_cvtepu8_epi16(_mm256_castsi256_si128(low_values)),
+ _mm256_cvtepu8_epi16(_mm256_castsi256_si128(high_values)),
+ ),
+ );
+ high_accumulator = _mm256_add_epi16(
+ high_accumulator,
+ _mm256_add_epi16(
+ _mm256_cvtepu8_epi16(_mm256_extracti128_si256(low_values,
1)),
+ _mm256_cvtepu8_epi16(_mm256_extracti128_si256(high_values,
1)),
+ ),
+ );
+ }
+ let mut chunk = [0u16; RQ_SCAN_BLOCK_SIZE];
+ _mm256_storeu_si256(chunk.as_mut_ptr() as *mut __m256i,
low_accumulator);
+ _mm256_storeu_si256(chunk.as_mut_ptr().add(16) as *mut __m256i,
high_accumulator);
+ for lane in 0..RQ_SCAN_BLOCK_SIZE {
+ totals[lane] += chunk[lane] as u32;
+ }
+ }
+ for lane in 0..RQ_SCAN_BLOCK_SIZE {
+ output[lane] = offset + scale * totals[lane] as f32;
+ }
+}
+
pub fn padded_dimension(d: usize) -> usize {
d.max(1).div_ceil(RQ_ROTATION_BLOCK_SIZE) * RQ_ROTATION_BLOCK_SIZE
}
@@ -591,6 +860,95 @@ mod tests {
);
}
+ #[test]
+ fn
coarse_distance_query_terms_match_rotated_vector_terms_for_every_metric() {
+ let d = 70;
+ let quantizer = RaBitQuantizer::new(d, 4);
+ let rotation = RQRotation::new(d, 91, DEFAULT_RQ_ROTATION_ROUNDS);
+ let query = (0..d)
+ .map(|index| ((index * 13 % 43) as f32 - 21.0) * 0.07)
+ .collect::<Vec<_>>();
+ let centroid = (0..d)
+ .map(|index| ((index * 17 % 37) as f32 - 18.0) * 0.05)
+ .collect::<Vec<_>>();
+ let mut rotated_query = vec![0.0; rotation.padded_dimension()];
+ let mut rotated_centroid = vec![0.0; rotation.padded_dimension()];
+ let mut scratch = vec![0.0; rotation.padded_dimension()];
+ rotation.rotate(&query, &mut rotated_query, &mut scratch);
+ rotation.rotate(¢roid, &mut rotated_centroid, &mut scratch);
+ let context = quantizer.prepare_query(rotated_query);
+ let residual_norm_sqr = query
+ .iter()
+ .zip(¢roid)
+ .map(|(&query, ¢roid)| {
+ let residual = query - centroid;
+ residual * residual
+ })
+ .sum::<f32>();
+ let query_norm_sqr = query.iter().map(|value| value *
value).sum::<f32>();
+ let centroid_norm_sqr = centroid.iter().map(|value| value *
value).sum::<f32>();
+
+ for metric in [MetricType::L2, MetricType::Cosine,
MetricType::InnerProduct] {
+ let rotated_terms = quantizer.query_terms(&context,
&rotated_centroid, metric);
+ let reused_terms = quantizer.query_terms_from_coarse_distance(
+ residual_norm_sqr,
+ query_norm_sqr,
+ centroid_norm_sqr,
+ metric,
+ );
+ assert!(
+ (rotated_terms.g_add - reused_terms.g_add).abs() <= 1e-4,
+ "{metric:?} g_add mismatch: {:?} vs {:?}",
+ rotated_terms,
+ reused_terms
+ );
+ assert!(
+ (rotated_terms.g_error - reused_terms.g_error).abs() <= 1e-4,
+ "{metric:?} g_error mismatch: {:?} vs {:?}",
+ rotated_terms,
+ reused_terms
+ );
+ }
+ }
+
+ #[test]
+ fn fastscan_coarse_sum_stays_inside_its_reported_error() {
+ let d = 960;
+ let quantizer = RaBitQuantizer::new(d, 4);
+ let query = (0..d)
+ .map(|index| ((index * 31 % 101) as f32 - 50.0) * 0.013)
+ .collect::<Vec<_>>();
+ let context = quantizer.prepare_query(query);
+ let mut blocked_codes = vec![0u8; quantizer.plane_size() *
RQ_SCAN_BLOCK_SIZE];
+ for byte_idx in 0..quantizer.plane_size() {
+ for lane in 0..RQ_SCAN_BLOCK_SIZE {
+ blocked_codes[byte_idx * RQ_SCAN_BLOCK_SIZE + lane] =
+ ((byte_idx * 37 + lane * 19 + 11) & 0xFF) as u8;
+ }
+ }
+
+ let mut approximate = [0.0f32; RQ_SCAN_BLOCK_SIZE];
+ quantizer.fastscan_coarse_block(&context, &blocked_codes, &mut
approximate);
+ for lane in 0..RQ_SCAN_BLOCK_SIZE {
+ let exact = (0..quantizer.plane_size())
+ .map(|byte_idx| {
+ quantizer.byte_subset_sum(
+ &context,
+ byte_idx,
+ blocked_codes[byte_idx * RQ_SCAN_BLOCK_SIZE + lane],
+ )
+ })
+ .sum::<f32>();
+ assert!(
+ (approximate[lane] - exact).abs() <=
quantizer.fastscan_ip_error(&context),
+ "lane {lane}: approximate={} exact={} error_bound={}",
+ approximate[lane],
+ exact,
+ quantizer.fastscan_ip_error(&context)
+ );
+ }
+ }
+
#[test]
fn four_bit_estimator_is_exact_for_the_encoded_vector() {
let d = 64;
diff --git a/core/src/topk.rs b/core/src/topk.rs
index f883ef7..ae7bfc6 100644
--- a/core/src/topk.rs
+++ b/core/src/topk.rs
@@ -19,6 +19,7 @@ use std::collections::HashMap;
pub(crate) struct TopKHeap {
k: usize,
+ max_distance: f32,
data: Vec<(f32, i64)>,
positions: HashMap<i64, usize>,
}
@@ -27,13 +28,23 @@ impl TopKHeap {
pub(crate) fn new(k: usize) -> Self {
Self {
k,
+ max_distance: f32::INFINITY,
+ data: Vec::with_capacity(k),
+ positions: HashMap::with_capacity(k),
+ }
+ }
+
+ pub(crate) fn with_max_distance(k: usize, max_distance: f32) -> Self {
+ Self {
+ k,
+ max_distance,
data: Vec::with_capacity(k),
positions: HashMap::with_capacity(k),
}
}
pub(crate) fn push(&mut self, dist: f32, id: i64) {
- if self.k == 0 {
+ if self.k == 0 || dist >= self.max_distance {
return;
}
if let Some(&idx) = self.positions.get(&id) {
@@ -59,7 +70,20 @@ impl TopKHeap {
}
pub(crate) fn should_consider(&self, lower_bound: f32) -> bool {
- self.k > 0 && (self.data.len() < self.k || lower_bound <
self.data[0].0)
+ self.k > 0
+ && if self.data.len() < self.k {
+ lower_bound < self.max_distance
+ } else {
+ lower_bound < self.data[0].0
+ }
+ }
+
+ pub(crate) fn is_full(&self) -> bool {
+ self.k > 0 && self.data.len() == self.k
+ }
+
+ pub(crate) fn worst_distance(&self) -> Option<f32> {
+ self.is_full().then(|| self.data[0].0)
}
pub(crate) fn into_sorted(mut self) -> Vec<(f32, i64)> {
@@ -132,4 +156,17 @@ mod tests {
assert_eq!(heap.into_sorted(), vec![(1.0, 1), (4.0, 3), (5.0, 4)]);
}
+
+ #[test]
+ fn test_topk_heap_applies_seeded_max_distance_before_it_is_full() {
+ let mut heap = TopKHeap::with_max_distance(3, 5.0);
+
+ assert!(!heap.should_consider(5.0));
+ assert!(!heap.should_consider(7.0));
+ heap.push(7.0, 1);
+ heap.push(4.0, 2);
+ heap.push(3.0, 3);
+
+ assert_eq!(heap.into_sorted(), vec![(3.0, 3), (4.0, 2)]);
+ }
}
diff --git a/docs/index.html b/docs/index.html
index e4844c6..685e09b 100644
--- a/docs/index.html
+++ b/docs/index.html
@@ -138,7 +138,7 @@
<tr><td>Reader/I/O model</td><td>Automatic read plan, Reader budget,
simulated latency</td><td>Latency-derived local/remote/object-store plans, 4
GiB automatically partitioned budget, 0/2/20 ms per read round</td><td>Fixed by
the current benchmark implementation; <code>ANN_STORAGE_CASES</code> selects a
focused subset. DiskANN range reads use an I/O pool independent of query
workers; all five indexes were refreshed after their current reader, storage,
and parallel-scan work.</td></tr>
</tbody></table></div>
<div class="callout"><strong>Equal relative PQ budget</strong>The
default <code>pq.code-ratio=0.0625</code> automatically resolves SIFT to
<code>pq.m=32</code>, GIST to <code>pq.m=240</code>, and GloVe to
<code>pq.m=25</code>. Every code occupies 6.25% as many bytes as its raw
<code>f32</code> vector and leaves four dimensions per PQ sub-vector. The
concrete value is persisted in index metadata; use explicit <code>pq.m</code>
only as an expert override.</div>
- <p>The run was recorded on 25 July 2026 using an Apple M4 Pro with 12
logical CPUs and 48 GiB RAM, a release build with Rust 1.95, real APFS files
with warm operating-system pages, and the automatic 4 GiB DiskANN Reader
budget. The reproduction command pins Rayon to 12 workers instead of relying on
automatic host parallelism. The modeled serving profiles add 2 ms or 20 ms per
positional-read round while executing all ranges in that round concurrently.
DiskANN's benchmark adapter [...]
+ <p>The cross-index run was recorded on 25 July 2026 using an Apple M4
Pro with 12 logical CPUs and 48 GiB RAM, a release build with Rust 1.95, real
APFS files with warm operating-system pages, and the automatic 4 GiB DiskANN
Reader budget. The IVF-RQ staged A/B and rebased IVF-PQ warm-local refresh were
recorded on 30 July on the same host and toolchain. The reproduction command
pins Rayon to 12 workers instead of relying on automatic host parallelism. The
modeled serving profile [...]
<h4>Build, file, and peak process memory</h4>
<div class="table-wrap"><table><thead><tr><th>Index</th><th>SIFT
build</th><th>SIFT file / RSS</th><th>GIST build</th><th>GIST file /
RSS</th><th>GloVe build</th><th>GloVe file / RSS</th></tr></thead><tbody>
@@ -149,17 +149,29 @@
<tr><td>IVF-FLAT</td><td>4.05 s</td><td>0.479 / 1.83
GiB</td><td>24.9 s</td><td>3.582 / 12.74 GiB</td><td>4.10 s</td><td>0.443 /
1.60 GiB</td></tr>
</tbody></table></div>
<div class="callout"><strong>IVF-SQ build and scan refresh</strong>The
add path now borrows L2/IP input, assigns rows once, and encodes lists in
parallel with one residual scratch vector per active list task instead of
materializing an additional <code>N × d</code> residual matrix. In the
immediately preceding same-machine run, SIFT/GIST/GloVe peak RSS was 1.81 /
12.92 / 1.60 GiB; it is now 0.79 / 5.09 / 0.71 GiB. A Top-K threshold fast path
skips hash work for candidates that ca [...]
- <div class="callout"><strong>Final IVF-PQ read/write
review</strong>The Reader keeps the v1 code suffix zero-copy and aligned,
orders selected lists by persisted offset, decodes the common one-byte row-ID
delta without entering the generic varint loop, and initializes each distance
accumulator from the first PQ column instead of zero-filling it in a separate
pass. In three-run same-file A/B medians, the row-ID fast path alone changed
SIFT/GIST/GloVe batch throughput from 7,528 / [...]
+ <div class="callout"><strong>30 July IVF-PQ batch-table reuse
refresh</strong>The rebased Reader retains the v1 zero-copy/transposed-code,
ordered-list, one-byte row-ID, and first-column accumulation fast paths. For
large 8-bit residual-L2 batches, the default <code>Auto</code> mode now factors
each distance table into reusable per-list and per-query components when the
reuse heuristic and 64 MiB working-memory guard both pass; small or unsuitable
batches keep the direct path, an [...]
<div class="callout"><strong>Latest IVF-FLAT storage and scan
review</strong>The v1 writer retains only sort permutations and encoded IDs,
materializes one sorted raw-vector list at a time, and reproduced all three
prior public files byte-for-byte. The Reader now receives list bytes directly
into an <code>f32</code>-aligned allocation; an internal prefix of at most
three bytes keeps the raw-vector suffix aligned despite variable-length row
IDs, so search no longer allocates and d [...]
<div class="callout"><strong>Pre-release storage changes produced
measurable savings</strong>The compact IVF-RQ factor layout removes one unused
F32 value per row: SIFT/GIST/GloVe files fell from 90.3 / 509.8 / 106.8 MB to
86.3 / 505.8 / 102.1 MB without changing the estimator; the refreshed Recall@10
values are 0.9148 / 0.9039 / 0.8203. Changing the balanced DiskANN default from
F32 to F16 rerank vectors reduced the same three files from 0.599 / 3.877 /
0.617 GiB to 0.361 / 2.08 [...]
- <div class="callout"><strong>Latest IVF-RQ build and scan
refresh</strong>The add path assigns rows once and encodes independent lists in
parallel with list-local residual, rotation, code, and factor scratch.
SIFT/GIST/GloVe add time fell from 5.96 / 39.11 / 6.48 seconds to 3.61 / 21.51
/ 3.76 seconds; complete build time is now 3.92 / 23.5 / 4.03 seconds. The
unfiltered scanner now removes the per-lane filter test from coarse-code
accumulation, and the final Top-K threshold avoi [...]
+ <div class="callout"><strong>30 July IVF-RQ staged scan
A/B</strong>The four changes below were measured against one retained v1 index
per corpus with the same 1,000 queries, <code>nlist=1,024</code>,
<code>nprobe=64</code>, four stored RQ bits, 12 Rayon workers, and warm APFS
pages. SIFT and GloVe use seven interleaved runs. Because GIST-960 showed
thermal drift during long stage sweeps, its reported wall-clock changes use
five baseline/final pairs with alternating execution ord [...]
+ <div class="table-wrap"><table><thead><tr><th>IVF-RQ
change</th><th>SIFT1M result</th><th>GIST1M result</th><th>GloVe-100
result</th></tr></thead><tbody>
+ <tr><td>Block-aggregated scan statistics</td><td>70.91% of batch
candidates reached full bit-plane refinement</td><td>95.88% reached bit-plane
refinement, but the complete approximate-distance bound reduced exact coarse
reevaluation to 1.45%</td><td>96.03% reached full bit-plane refinement</td></tr>
+ <tr><td>Reuse IVF centroid distance</td><td>Removes 8.2 million
repeated rotated centroid terms per 1,000-query run</td><td>Removes 61.4
million repeated terms; end-to-end movement stayed inside run variance because
code scan dominates</td><td>Removes 8.2 million repeated padded terms per
run</td></tr>
+ <tr><td>16-entry FastScan LUT + NEON/AVX2</td><td>Deliberately
bypassed below padded dimension 256</td><td>Versus the optimized scalar
scanner: P95 −8.3%, sequential QPS +6.6%, batch QPS +17.8% in paired
medians</td><td>Deliberately bypassed below padded dimension 256</td></tr>
+ <tr><td>32-vector single-query seed threshold</td><td>Refinement
work −4.3%; P95 −0.9% versus no seed</td><td>Exact final evaluations −33.1%;
paired sequential QPS +1.3% while P95 was neutral</td><td>Refinement work
−0.32%; P95 −2.3% versus no seed</td></tr>
+ </tbody></table></div>
+ <p>The first statistics prototype incremented counters inside
byte-lookup loops and regressed GIST, so it was not retained. The final
implementation derives lookup counts once per block or admitted candidate,
keeps per-list statistics thread-local, and merges them after parallel work.
Centroid reuse derives the RQ query terms from the distances already produced
by IVF probing and stores only centroid norms in the Reader. FastScan quantizes
two 16-entry nibble tables, evaluates 32 [...]
+ <div class="table-wrap"><table><thead><tr><th>Same-file
endpoint</th><th>SIFT P95 / sequential QPS / batch QPS</th><th>GIST P95 /
sequential QPS / batch QPS</th><th>GloVe P95 / sequential QPS / batch
QPS</th></tr></thead><tbody>
+ <tr><td>Pre-change baseline</td><td>1.101 ms / 1,020 /
2,706</td><td>5.599 ms / 203 / 325</td><td>1.071 ms / 1,041 / 2,855</td></tr>
+ <tr><td>Final scanner</td><td>1.081 ms / 1,042 / 2,743</td><td>5.190
ms / 218 / 399</td><td>1.046 ms / 1,057 / 2,839</td></tr>
+ </tbody></table></div>
+ <p>Against its paired baseline, the final GIST scanner improved median
P95 by 8.7%, sequential QPS by 8.1%, and batch QPS by 20.0%. SIFT improved P95
by 1.8%, sequential QPS by 2.1%, and batch QPS by 1.4%. GloVe improved P95 by
2.3% and sequential QPS by 1.6%; its 0.6% batch-QPS decrease is treated as
noise, not an improvement. File bytes, bytes read, storage format, and
compressed-domain ranking are unchanged.</p>
<div class="callout"><strong>DiskANN read-path review</strong>DiskANN
converts F16 rerank vectors and accumulates L2 directly in one AArch64 NEON
loop; its local profile coalesces 16 KiB windows. The final compact-layout
batch rerank groups candidate windows with a hash table, then sorts only the
unique windows into deterministic I/O order. Against the immediately preceding
ordered-map control, median local batch time changed from 117 / 223 / 162 ms to
111 / 215 / 159 ms on SIFT/ [...]
<div class="callout"><strong>Final open-source cross-check and format
decision</strong><a
href="https://github.com/facebookresearch/faiss/wiki/Fast-accumulation-of-PQ-and-AQ-codes-%28FastScan%29";>Faiss
FastScan</a> still trades down to 4-bit lookup tables and a 32-row layout, so
it is not a transparent replacement for the published 8-bit IVF-PQ v1. Faiss
Panorama's additional level-oriented energy data was not needed to keep the
existing IVF-FLAT v1 progressive cutoff. Lance's pa [...]
<p>DiskANN spends almost all build time constructing one global graph:
about 18× IVF-FLAT on SIFT, 28× on GIST, and 37× on GloVe. The balanced F16
default makes its files about 42% smaller than IVF-FLAT on GIST and 11% smaller
on GloVe, but they remain much larger than the compact IVF encodings because
persisted rerank vectors, resident codes, and graph edges are all material.
Peak RSS remains below the raw IVF writers because the DiskANN writer does not
retain a second full raw- [...]
<h4>Warm local-storage result</h4>
<div class="table-wrap"><table><thead><tr><th>Index /
search</th><th>SIFT Recall / P95 / batch QPS / read</th><th>GIST Recall / P95 /
batch QPS / read</th><th>GloVe Recall / P95 / batch QPS /
read</th></tr></thead><tbody>
- <tr><td>IVF-PQ</td><td>0.7142 / 0.71 ms / 8,353 / 2.18
MiB</td><td>0.7410 / 2.19 ms / 1,094 / 17.84 MiB</td><td>0.5819 / 0.64 ms /
9,330 / 1.84 MiB</td></tr>
+ <tr><td>IVF-PQ</td><td>0.7142 / 0.72 ms / 7,899 / 2.18
MiB</td><td>0.7410 / 2.47 ms / 950 / 17.84 MiB</td><td>0.5819 / 0.64 ms / 8,048
/ 1.84 MiB</td></tr>
<tr><td>IVF-SQ</td><td>0.8627 / 0.79 ms / 11,082 / 8.38
MiB</td><td>0.8577 / 3.56 ms / 1,502 / 70.95 MiB</td><td>0.8036 / 0.71 ms /
12,962 / 6.99 MiB</td></tr>
<tr><td>IVF-RQ</td><td>0.9148 / 1.20 ms / 3,074 / 5.54
MiB</td><td>0.9039 / 4.41 ms / 444 / 37.02 MiB</td><td>0.8203 / 1.23 ms / 2,917
/ 5.89 MiB</td></tr>
<tr><td>DiskANN</td><td>0.9915 / 1.50 ms / 9,009 / 0.66
MiB</td><td>0.9336 / 1.83 ms / 4,651 / 0.83 MiB</td><td>0.8355 / 1.90 ms /
6,289 / 0.96 MiB</td></tr>
