This is an automated email from the ASF dual-hosted git repository.
jerry-024 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 8dcabf2 perf: accelerate IVF builds with Vamana coarse assignment
(#85)
8dcabf2 is described below
commit 8dcabf208c99707eab3938af0bcc40530fdb4cad
Author: jerry <[email protected]>
AuthorDate: Wed Sep 2 17:30:42 2026 +0800
perf: accelerate IVF builds with Vamana coarse assignment (#85)
---
core/benches/ann_bench.rs | 4 +
core/benches/ivfpq_add_bench.rs | 2 +-
core/benches/ivfpq_batch_reuse_bench.rs | 8 +-
core/benches/ivfpq_filter_scan_bench.rs | 8 +-
core/benches/pq4_bench.rs | 2 +-
core/src/coarse.rs | 205 +++++++++++++++++++++++
core/src/distance.rs | 27 ++-
core/src/index.rs | 132 ++++++++++++++-
core/src/io.rs | 2 +-
core/src/ivfflat.rs | 36 +++-
core/src/ivfflat_io.rs | 26 +--
core/src/ivfpq.rs | 66 +++++++-
core/src/ivfrq.rs | 212 ++++++++++++++++++++----
core/src/ivfrq_io.rs | 24 ++-
core/src/ivfsq.rs | 47 +++++-
core/src/ivfsq_io.rs | 4 +-
core/src/kmeans.rs | 284 ++++++++++++++++++++++++++++----
core/src/lib.rs | 1 +
core/src/vamana.rs | 233 ++++++++++++++++++++++----
core/tests/storage_format_fixtures.rs | 18 +-
docs/api.html | 7 +-
docs/ivf-flat.html | 6 +-
docs/ivf-pq.html | 3 +-
docs/ivf-rq.html | 4 +-
docs/ivf-sq.html | 6 +-
docs/releases.html | 1 +
26 files changed, 1206 insertions(+), 162 deletions(-)
diff --git a/core/benches/ann_bench.rs b/core/benches/ann_bench.rs
index 1cb167f..a62566b 100644
--- a/core/benches/ann_bench.rs
+++ b/core/benches/ann_bench.rs
@@ -644,6 +644,7 @@ fn index_specs(config: &Config) -> Vec<IndexSpec> {
dimension: config.d,
nlist: config.nlist,
metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
},
searches: vec![ivf_search],
},
@@ -653,6 +654,7 @@ fn index_specs(config: &Config) -> Vec<IndexSpec> {
dimension: config.d,
nlist: config.nlist,
metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
},
searches: vec![ivf_search],
},
@@ -664,6 +666,7 @@ fn index_specs(config: &Config) -> Vec<IndexSpec> {
m: config.pq_m,
metric: MetricType::L2,
use_opq: false,
+ use_approximate_coarse_assignment: true,
},
searches: vec![ivf_search],
},
@@ -674,6 +677,7 @@ fn index_specs(config: &Config) -> Vec<IndexSpec> {
nlist: config.nlist,
bits: config.rq_bits,
metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
},
searches: vec![ivf_search],
},
diff --git a/core/benches/ivfpq_add_bench.rs b/core/benches/ivfpq_add_bench.rs
index 6243e6d..3893cb8 100644
--- a/core/benches/ivfpq_add_bench.rs
+++ b/core/benches/ivfpq_add_bench.rs
@@ -113,7 +113,7 @@ fn new_index(
norms: &[f32],
) -> IVFPQIndex {
let mut index = IVFPQIndex::new(case.d, case.nlist, case.m,
MetricType::L2, false);
- index.quantizer_centroids = quantizer_centroids.to_vec();
+ index.set_quantizer_centroids(quantizer_centroids.to_vec());
index.pq.centroids = centroids.to_vec();
index.pq.centroid_norms_cache = norms.to_vec();
index
diff --git a/core/benches/ivfpq_batch_reuse_bench.rs
b/core/benches/ivfpq_batch_reuse_bench.rs
index 1445015..7a67db8 100644
--- a/core/benches/ivfpq_batch_reuse_bench.rs
+++ b/core/benches/ivfpq_batch_reuse_bench.rs
@@ -98,9 +98,11 @@ fn percentile(samples: &[Sample], percentile: usize, value:
impl Fn(&Sample) ->
fn main() {
let mut rng = StdRng::seed_from_u64(42);
let mut index = IVFPQIndex::new(D, NLIST, M, MetricType::InnerProduct,
false);
- index.quantizer_centroids = (0..NLIST * D)
- .map(|_| rng.gen_range(-1.0f32..1.0))
- .collect();
+ index.set_quantizer_centroids(
+ (0..NLIST * D)
+ .map(|_| rng.gen_range(-1.0f32..1.0))
+ .collect(),
+ );
index.pq.centroids = (0..M * index.pq.ksub * index.pq.dsub)
.map(|_| rng.gen_range(-1.0f32..1.0))
.collect();
diff --git a/core/benches/ivfpq_filter_scan_bench.rs
b/core/benches/ivfpq_filter_scan_bench.rs
index 3218e44..a886c51 100644
--- a/core/benches/ivfpq_filter_scan_bench.rs
+++ b/core/benches/ivfpq_filter_scan_bench.rs
@@ -99,9 +99,11 @@ fn main() {
let mut rng = StdRng::seed_from_u64(42);
let mut index = IVFPQIndex::new(D, NLIST, M, MetricType::InnerProduct,
false);
- index.quantizer_centroids = (0..NLIST * D)
- .map(|_| rng.gen_range(-1.0f32..1.0))
- .collect();
+ index.set_quantizer_centroids(
+ (0..NLIST * D)
+ .map(|_| rng.gen_range(-1.0f32..1.0))
+ .collect(),
+ );
index.pq.centroids = (0..M * index.pq.ksub * index.pq.dsub)
.map(|_| rng.gen_range(-1.0f32..1.0))
.collect();
diff --git a/core/benches/pq4_bench.rs b/core/benches/pq4_bench.rs
index 8bd3d27..239089b 100644
--- a/core/benches/pq4_bench.rs
+++ b/core/benches/pq4_bench.rs
@@ -121,7 +121,7 @@ fn main() {
let mut sim_table = vec![0.0f32; m4 * 16];
let query0 = &data[0..d];
// compute residual
- let centroid = &idx4.quantizer_centroids[biggest_list * d..(biggest_list +
1) * d];
+ let centroid = &idx4.quantizer_centroids()[biggest_list * d..(biggest_list
+ 1) * d];
let residual: Vec<f32> = (0..d).map(|j| query0[j] - centroid[j]).collect();
idx4.pq
.compute_distance_table(&residual, MetricType::L2, &mut sim_table);
diff --git a/core/src/coarse.rs b/core/src/coarse.rs
new file mode 100644
index 0000000..3e0dc0a
--- /dev/null
+++ b/core/src/coarse.rs
@@ -0,0 +1,205 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use crate::diskann::{
+ DiskAnnBuildDistance, DiskAnnBuildParams, DiskAnnRawVectorEncoding,
DiskAnnStorageLayout,
+};
+use crate::kmeans;
+use crate::logging::{emit_log, LogLevel};
+use crate::vamana::VamanaGraph;
+use rayon::prelude::*;
+
+const APPROX_ASSIGN_SEARCH_LIST: usize = 15;
+const APPROX_ASSIGN_MIN_CENTROID_VALUES: usize = 1_000_000;
+
+fn use_approximate_assignment(d: usize, nlist: usize) -> bool {
+ d.saturating_mul(nlist) >= APPROX_ASSIGN_MIN_CENTROID_VALUES
+}
+
+pub(crate) struct CoarseAssignment {
+ graph: Option<VamanaGraph>,
+ build_attempted: bool,
+ approximate_enabled: bool,
+}
+
+impl Default for CoarseAssignment {
+ fn default() -> Self {
+ Self {
+ graph: None,
+ build_attempted: false,
+ approximate_enabled: true,
+ }
+ }
+}
+
+impl CoarseAssignment {
+ pub(crate) fn reset(&mut self) {
+ self.graph = None;
+ self.build_attempted = false;
+ }
+
+ pub(crate) fn set_approximate_enabled(&mut self, enabled: bool) {
+ self.reset();
+ self.approximate_enabled = enabled;
+ }
+
+ pub(crate) fn approximate_enabled(&self) -> bool {
+ self.approximate_enabled
+ }
+
+ #[cfg(test)]
+ pub(crate) fn build_attempted(&self) -> bool {
+ self.build_attempted
+ }
+
+ pub(crate) fn prepare(&mut self, centroids: &[f32], nlist: usize, d:
usize) {
+ if self.build_attempted {
+ return;
+ }
+ self.build_attempted = true;
+ if !self.approximate_enabled || !use_approximate_assignment(d, nlist) {
+ return;
+ }
+
+ let params = DiskAnnBuildParams {
+ max_degree: 12,
+ build_search_list_size: APPROX_ASSIGN_SEARCH_LIST,
+ alpha: 1.2,
+ seed: 42,
+ memory_budget_bytes: 1024 * 1024 * 1024,
+ storage_layout: DiskAnnStorageLayout::Compact,
+ raw_vector_encoding: DiskAnnRawVectorEncoding::F32,
+ build_distance: DiskAnnBuildDistance::FullPrecision,
+ };
+ match VamanaGraph::build(centroids, nlist, d, params) {
+ Ok(graph) => self.graph = Some(graph),
+ Err(error) => emit_log(
+ LogLevel::Warn,
+ &format!("automatic approximate coarse assignment disabled:
{error}"),
+ ),
+ }
+ }
+
+ pub(crate) fn assign(
+ &mut self,
+ data: &[f32],
+ n: usize,
+ centroids: &[f32],
+ nlist: usize,
+ d: usize,
+ ) -> Vec<usize> {
+ if n == 0 {
+ return Vec::new();
+ }
+ self.prepare(centroids, nlist, d);
+ let Some(graph) = &self.graph else {
+ return kmeans::find_nearest_batch(data, n, centroids, nlist, d);
+ };
+
+ let mut assignments = vec![0usize; n];
+ let chunk = (n / (rayon::current_num_threads() * 4).max(1)).clamp(16,
1024);
+ assignments.par_chunks_mut(chunk).enumerate().for_each_init(
+ || graph.search_scratch(APPROX_ASSIGN_SEARCH_LIST),
+ |scratch, (chunk_idx, chunk_assignments)| {
+ let row0 = chunk_idx * chunk;
+ for (i, assignment) in
chunk_assignments.iter_mut().enumerate() {
+ let row = row0 + i;
+ *assignment = graph
+ .greedy_search_best_with_scratch(
+ centroids,
+ d,
+ &data[row * d..(row + 1) * d],
+ APPROX_ASSIGN_SEARCH_LIST,
+ scratch,
+ )
+ .map(|node| node.id as usize)
+ .unwrap_or(0);
+ }
+ },
+ );
+ assignments
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ #[test]
+ fn approximate_assignment_depends_on_centroid_values() {
+ assert!(!use_approximate_assignment(768, 1024));
+ assert!(use_approximate_assignment(768, 4096));
+ }
+
+ #[test]
+ fn vamana_coarse_assignment_matches_exact_on_connected_graph() {
+ let d = 4;
+ let nlist = 4;
+ let n = 64;
+ let centroids = (0..nlist)
+ .flat_map(|list| (0..d).map(move |dimension| list as f32 * 10.0 +
dimension as f32))
+ .collect::<Vec<_>>();
+ let data = (0..n)
+ .flat_map(|row| {
+ (0..d).map(move |dimension| {
+ (row % nlist) as f32 * 10.0 + row as f32 * 0.01 +
dimension as f32
+ })
+ })
+ .collect::<Vec<_>>();
+ let expected = kmeans::find_nearest_batch(&data, n, ¢roids, nlist,
d);
+ let adjacency = (0..nlist)
+ .map(|node| {
+ (0..nlist)
+ .filter(|&neighbor| neighbor != node)
+ .map(|neighbor| neighbor as u32)
+ .collect()
+ })
+ .collect();
+ let mut assignment = CoarseAssignment {
+ graph: Some(VamanaGraph::from_adjacency(0, adjacency)),
+ build_attempted: true,
+ approximate_enabled: true,
+ };
+
+ assert_eq!(assignment.assign(&data, n, ¢roids, nlist, d),
expected);
+ assignment.reset();
+ assert!(assignment.graph.is_none());
+ assert!(!assignment.build_attempted);
+ assert!(assignment.approximate_enabled);
+ }
+
+ #[test]
+ fn empty_assignment_does_not_prepare_graph() {
+ let mut assignment = CoarseAssignment::default();
+
+ assert!(assignment.assign(&[], 0, &[], 4096, 256).is_empty());
+ assert!(!assignment.build_attempted);
+ }
+
+ #[test]
+ fn exact_assignment_disables_graph_build_without_changing_reset_policy() {
+ let mut assignment = CoarseAssignment::default();
+ assignment.set_approximate_enabled(false);
+ assignment.prepare(&[], 4096, 256);
+
+ assert!(assignment.build_attempted);
+ assert!(assignment.graph.is_none());
+ assignment.reset();
+ assert!(!assignment.build_attempted);
+ assert!(!assignment.approximate_enabled);
+ }
+}
diff --git a/core/src/distance.rs b/core/src/distance.rs
index 99d6cae..53a13c2 100644
--- a/core/src/distance.rs
+++ b/core/src/distance.rs
@@ -16,6 +16,7 @@
// under the License.
use crate::blas::sgemm_a_bt;
+use rayon::prelude::*;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u32)]
@@ -776,9 +777,15 @@ fn fvec_cosine_distance_with_norms(a: &[f32], b: &[f32],
a_norm: f32, b_norm: f3
pub fn preprocess_vectors(data: &[f32], n: usize, d: usize, metric:
MetricType) -> Vec<f32> {
let mut processed = data[..n * d].to_vec();
- if metric == MetricType::Cosine {
- for i in 0..n {
- fvec_normalize(&mut processed[i * d..(i + 1) * d]);
+ if metric == MetricType::Cosine && d > 0 {
+ if n > 1_000 {
+ processed.par_chunks_mut(d).for_each(|vector| {
+ fvec_normalize(vector);
+ });
+ } else {
+ processed.chunks_mut(d).for_each(|vector| {
+ fvec_normalize(vector);
+ });
}
}
processed
@@ -801,6 +808,20 @@ mod preprocess_tests {
vec![0.6, 0.8]
);
}
+
+ #[test]
+ fn test_preprocess_vectors_normalizes_cosine_in_parallel() {
+ let data = [3.0, 4.0].repeat(1_001);
+ let processed = preprocess_vectors(&data, 1_001, 2,
MetricType::Cosine);
+
+ assert!(processed.chunks_exact(2).all(|vector| vector == [0.6, 0.8]));
+ }
+
+ #[test]
+ fn test_preprocess_vectors_accepts_zero_dimension() {
+ assert!(preprocess_vectors(&[], 1, 0, MetricType::Cosine).is_empty());
+ assert!(preprocess_vectors(&[], 1_001, 0,
MetricType::Cosine).is_empty());
+ }
}
/// Compute result[i] = a[i] + bf * b[i]. Used for precomputed table merging.
diff --git a/core/src/index.rs b/core/src/index.rs
index 60c7762..198eaed 100644
--- a/core/src/index.rs
+++ b/core/src/index.rs
@@ -146,6 +146,7 @@ pub enum VectorIndexConfig {
dimension: usize,
nlist: usize,
metric: MetricType,
+ use_approximate_coarse_assignment: bool,
},
IvfPq {
dimension: usize,
@@ -153,17 +154,20 @@ pub enum VectorIndexConfig {
m: usize,
metric: MetricType,
use_opq: bool,
+ use_approximate_coarse_assignment: bool,
},
IvfRq {
dimension: usize,
nlist: usize,
bits: usize,
metric: MetricType,
+ use_approximate_coarse_assignment: bool,
},
IvfSq {
dimension: usize,
nlist: usize,
metric: MetricType,
+ use_approximate_coarse_assignment: bool,
},
DiskAnn {
dimension: usize,
@@ -202,6 +206,7 @@ impl VectorIndexConfig {
m: infer_uniform_pq_m(dimension, 8, DEFAULT_PQ_CODE_RATIO)?,
metric,
use_opq,
+ use_approximate_coarse_assignment: true,
};
validate_config(&config)?;
Ok(config)
@@ -230,6 +235,7 @@ impl VectorIndexConfig {
nlist,
bits: DEFAULT_RQ_BITS,
metric,
+ use_approximate_coarse_assignment: true,
};
validate_config(&config)?;
Ok(config)
@@ -276,6 +282,7 @@ pub struct ResolvedVectorIndexConfig {
pub pq_bits: Option<usize>,
pub rq_bits: Option<usize>,
pub use_opq: bool,
+ pub use_approximate_coarse_assignment: bool,
pub diskann_build: Option<DiskAnnBuildParams>,
}
@@ -286,11 +293,13 @@ impl From<&VectorIndexConfig> for
ResolvedVectorIndexConfig {
dimension,
nlist,
metric,
+ use_approximate_coarse_assignment,
}
| VectorIndexConfig::IvfSq {
dimension,
nlist,
metric,
+ use_approximate_coarse_assignment,
} => Self {
index_type: config.index_type(),
dimension: *dimension,
@@ -300,6 +309,7 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig
{
pq_bits: None,
rq_bits: None,
use_opq: false,
+ use_approximate_coarse_assignment:
*use_approximate_coarse_assignment,
diskann_build: None,
},
VectorIndexConfig::IvfPq {
@@ -308,6 +318,7 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig
{
m,
metric,
use_opq,
+ use_approximate_coarse_assignment,
} => Self {
index_type: IndexType::IvfPq,
dimension: *dimension,
@@ -317,6 +328,7 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig
{
pq_bits: Some(8),
rq_bits: None,
use_opq: *use_opq,
+ use_approximate_coarse_assignment:
*use_approximate_coarse_assignment,
diskann_build: None,
},
VectorIndexConfig::IvfRq {
@@ -324,6 +336,7 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig
{
nlist,
bits,
metric,
+ use_approximate_coarse_assignment,
} => Self {
index_type: IndexType::IvfRq,
dimension: *dimension,
@@ -333,6 +346,7 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig
{
pq_bits: None,
rq_bits: Some(*bits),
use_opq: false,
+ use_approximate_coarse_assignment:
*use_approximate_coarse_assignment,
diskann_build: None,
},
VectorIndexConfig::DiskAnn {
@@ -350,6 +364,7 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig
{
pq_bits: Some(*pq_bits),
rq_bits: None,
use_opq: false,
+ use_approximate_coarse_assignment: false,
diskann_build: Some(*build),
},
}
@@ -410,12 +425,19 @@ impl VectorIndexBuildPlan {
max_build_seconds,
deployment_profile,
};
+ let use_approximate_coarse_assignment = match index_type {
+ IndexType::IvfFlat | IndexType::IvfPq | IndexType::IvfRq |
IndexType::IvfSq => {
+ parse_ivf_coarse_assignment_option(&mut options)?
+ }
+ IndexType::DiskAnn => true,
+ };
let config = match index_type {
IndexType::IvfFlat => VectorIndexConfig::IvfFlat {
dimension,
nlist: parse_nlist_options(&mut options,
expected_vector_count)?,
metric,
+ use_approximate_coarse_assignment,
},
IndexType::IvfPq => VectorIndexConfig::IvfPq {
dimension,
@@ -437,6 +459,7 @@ impl VectorIndexBuildPlan {
Some(use_opq) => parse_bool_option("use-opq", &use_opq)?,
None => target_recall.is_some_and(|recall| recall >= 0.9),
},
+ use_approximate_coarse_assignment,
},
IndexType::IvfRq => {
let explicit_bits = options
@@ -457,12 +480,14 @@ impl VectorIndexBuildPlan {
nlist: parse_nlist_options(&mut options,
expected_vector_count)?,
bits,
metric,
+ use_approximate_coarse_assignment,
}
}
IndexType::IvfSq => VectorIndexConfig::IvfSq {
dimension,
nlist: parse_nlist_options(&mut options,
expected_vector_count)?,
metric,
+ use_approximate_coarse_assignment,
},
IndexType::DiskAnn => {
let pq_bits = match options.optional("pq.bits") {
@@ -600,6 +625,20 @@ fn parse_nlist_options(
}
}
+fn parse_ivf_coarse_assignment_option(options: &mut ConfigOptions) ->
io::Result<bool> {
+ match options
+ .optional("ivf.coarse-assignment")
+ .as_deref()
+ .map(str::trim)
+ {
+ None | Some("auto") => Ok(true),
+ Some("exact") => Ok(false),
+ Some(_) => Err(invalid_input(
+ "option 'ivf.coarse-assignment' must be auto or exact",
+ )),
+ }
+}
+
fn parse_deployment_profile_option(name: &str, value: &str) ->
io::Result<DeploymentProfile> {
match value.trim() {
"auto" => Ok(DeploymentProfile::Auto),
@@ -1277,25 +1316,45 @@ impl VectorIndexWriter {
dimension,
nlist,
metric,
- } => Self::IvfFlat(IVFFlatIndex::new(dimension, nlist, metric)),
+ use_approximate_coarse_assignment,
+ } => {
+ let mut index = IVFFlatIndex::new(dimension, nlist, metric);
+
index.set_approximate_coarse_assignment(use_approximate_coarse_assignment);
+ Self::IvfFlat(index)
+ }
VectorIndexConfig::IvfSq {
dimension,
nlist,
metric,
- } => Self::IvfSq(IVFSQIndex::new(dimension, nlist, metric)),
+ use_approximate_coarse_assignment,
+ } => {
+ let mut index = IVFSQIndex::new(dimension, nlist, metric);
+
index.set_approximate_coarse_assignment(use_approximate_coarse_assignment);
+ Self::IvfSq(index)
+ }
VectorIndexConfig::IvfPq {
dimension,
nlist,
m,
metric,
use_opq,
- } => Self::IvfPq(IVFPQIndex::new(dimension, nlist, m, metric,
use_opq)),
+ use_approximate_coarse_assignment,
+ } => {
+ let mut index = IVFPQIndex::new(dimension, nlist, m, metric,
use_opq);
+
index.set_approximate_coarse_assignment(use_approximate_coarse_assignment);
+ Self::IvfPq(index)
+ }
VectorIndexConfig::IvfRq {
dimension,
nlist,
bits,
metric,
- } => Self::IvfRq(IVFRQIndex::with_bits(dimension, nlist, bits,
metric)),
+ use_approximate_coarse_assignment,
+ } => {
+ let mut index = IVFRQIndex::with_bits(dimension, nlist, bits,
metric);
+
index.set_approximate_coarse_assignment(use_approximate_coarse_assignment);
+ Self::IvfRq(index)
+ }
VectorIndexConfig::DiskAnn {
dimension,
metric,
@@ -2594,6 +2653,7 @@ mod tests {
dimension: 1,
nlist: 1,
metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
},
&[0.0, 1.0],
2,
@@ -2646,6 +2706,7 @@ mod tests {
nlist: 8,
metric: MetricType::L2,
bits: 4,
+ use_approximate_coarse_assignment: true,
});
reader
@@ -2669,6 +2730,7 @@ mod tests {
nlist,
metric: MetricType::L2,
bits: 4,
+ use_approximate_coarse_assignment: true,
});
let queries = [0, nlist - 1]
.into_iter()
@@ -2822,6 +2884,7 @@ mod tests {
dimension: 8,
nlist: 4,
metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
});
roundtrip(VectorIndexConfig::ivf_pq(16, 4, MetricType::L2,
false).unwrap());
roundtrip(VectorIndexConfig::IvfRq {
@@ -2829,11 +2892,13 @@ mod tests {
nlist: 4,
bits: DEFAULT_RQ_BITS,
metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
});
roundtrip(VectorIndexConfig::IvfSq {
dimension: 8,
nlist: 4,
metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
});
roundtrip(
VectorIndexConfig::disk_ann(
@@ -2872,6 +2937,7 @@ mod tests {
dimension: 8,
nlist: 4,
metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
},
VectorIndexConfig::IvfPq {
dimension: 16,
@@ -2879,17 +2945,20 @@ mod tests {
m: 4,
metric: MetricType::L2,
use_opq: false,
+ use_approximate_coarse_assignment: true,
},
VectorIndexConfig::IvfRq {
dimension: 8,
nlist: 4,
bits: DEFAULT_RQ_BITS,
metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
},
VectorIndexConfig::IvfSq {
dimension: 8,
nlist: 4,
metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
},
] {
let d = config.dimension();
@@ -2976,6 +3045,7 @@ mod tests {
m: 3,
metric: MetricType::L2,
use_opq: false,
+ use_approximate_coarse_assignment: true,
}) {
Ok(_) => panic!("invalid PQ config should be rejected"),
Err(err) => err,
@@ -2990,6 +3060,7 @@ mod tests {
nlist: 4,
bits: DEFAULT_RQ_BITS,
metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
})
.unwrap();
let err = match VectorIndexTrainer::new(VectorIndexConfig::IvfRq {
@@ -2997,6 +3068,7 @@ mod tests {
nlist: 4,
bits: 9,
metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
}) {
Ok(_) => panic!("invalid RQ config should be rejected"),
Err(err) => err,
@@ -4110,6 +4182,7 @@ mod tests {
nlist,
bits,
metric,
+ ..
} => {
assert_eq!(dimension, 8);
assert_eq!(nlist, 4);
@@ -4131,6 +4204,7 @@ mod tests {
dimension,
nlist,
metric,
+ ..
} => {
assert_eq!(dimension, 8);
assert_eq!(nlist, 4);
@@ -4162,6 +4236,7 @@ mod tests {
dimension: 1,
nlist: 1,
metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
},
&[value, 1.0],
2,
@@ -4172,6 +4247,54 @@ mod tests {
}
}
+ #[test]
+ fn ivf_coarse_assignment_option_controls_threshold_enabled_self_recall() {
+ let default = VectorIndexConfig::from_options(&options(&[
+ ("index.type", "ivf_flat"),
+ ("dimension", "8"),
+ ("nlist", "4"),
+ ("metric", "l2"),
+ ]))
+ .unwrap();
+ assert!(default.resolved().use_approximate_coarse_assignment);
+
+ let exact = VectorIndexConfig::from_options(&options(&[
+ ("index.type", "ivf_flat"),
+ ("dimension", "256"),
+ ("nlist", "4096"),
+ ("metric", "l2"),
+ ("ivf.coarse-assignment", "exact"),
+ ]))
+ .unwrap();
+ assert!(!exact.resolved().use_approximate_coarse_assignment);
+
+ let VectorIndexWriter::IvfFlat(mut index) =
VectorIndexWriter::from_config(exact).unwrap()
+ else {
+ unreachable!()
+ };
+ let centroids = (0..index.d * index.nlist)
+ .map(|i| ((i * 17 + i / 11) % 1009) as f32 / 1009.0)
+ .collect::<Vec<_>>();
+ let query = centroids[..index.d].to_vec();
+ index.set_quantizer_centroids(centroids);
+ index.add(&query, &[42], 1);
+ let mut distances = [f32::MAX];
+ let mut labels = [-1];
+ index.search(&query, 1, 1, 1, &mut distances, &mut labels);
+ assert_eq!(labels, [42]);
+ assert_eq!(distances, [0.0]);
+
+ let error = VectorIndexConfig::from_options(&options(&[
+ ("index.type", "ivf_flat"),
+ ("dimension", "8"),
+ ("nlist", "4"),
+ ("metric", "l2"),
+ ("ivf.coarse-assignment", "vamana"),
+ ]))
+ .unwrap_err();
+ assert!(error.to_string().contains("must be auto or exact"));
+ }
+
#[test]
fn config_from_options_rejects_unknown_options() {
let err = VectorIndexConfig::from_options(&options(&[
@@ -4245,6 +4368,7 @@ mod tests {
dimension: 1,
nlist: 1,
metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
},
&[0.0, 1.0],
2,
diff --git a/core/src/io.rs b/core/src/io.rs
index 29fa69a..3a55823 100644
--- a/core/src/io.rs
+++ b/core/src/io.rs
@@ -363,7 +363,7 @@ pub fn write_index(index: &IVFPQIndex, out: &mut dyn
SeekWrite) -> io::Result<()
write_f32_slice(out, &opq.rotation)?;
}
- write_f32_slice(out, &index.quantizer_centroids)?;
+ write_f32_slice(out, index.quantizer_centroids())?;
write_f32_slice(out, &index.pq.centroids)?;
// Compute offsets for inverted lists
diff --git a/core/src/ivfflat.rs b/core/src/ivfflat.rs
index 796dbc2..ab1d622 100644
--- a/core/src/ivfflat.rs
+++ b/core/src/ivfflat.rs
@@ -15,6 +15,7 @@
// specific language governing permissions and limitations
// under the License.
+use crate::coarse::CoarseAssignment;
use crate::distance::{preprocess_vectors, MetricType, QueryDistance};
use crate::ivfpq::RowIdFilter;
use crate::kmeans::{self, KMeansConfig};
@@ -24,9 +25,10 @@ pub struct IVFFlatIndex {
pub d: usize,
pub nlist: usize,
pub metric: MetricType,
- pub quantizer_centroids: Vec<f32>,
+ quantizer_centroids: Vec<f32>,
pub ids: Vec<Vec<i64>>,
pub vectors: Vec<Vec<f32>>,
+ coarse_assignment: CoarseAssignment,
}
impl IVFFlatIndex {
@@ -38,18 +40,48 @@ impl IVFFlatIndex {
quantizer_centroids: Vec::new(),
ids: vec![Vec::new(); nlist],
vectors: vec![Vec::new(); nlist],
+ coarse_assignment: CoarseAssignment::default(),
}
}
+ pub fn quantizer_centroids(&self) -> &[f32] {
+ &self.quantizer_centroids
+ }
+
+ /// Enables automatic Vamana coarse assignment for large centroid matrices.
+ /// Disable it to keep vector assignment exact.
+ pub(crate) fn set_approximate_coarse_assignment(&mut self, enabled: bool) {
+ assert!(
+ self.ids.iter().all(Vec::is_empty),
+ "cannot change coarse assignment after vectors have been added"
+ );
+ self.coarse_assignment.set_approximate_enabled(enabled);
+ }
+
+ pub fn set_quantizer_centroids(&mut self, centroids: Vec<f32>) {
+ assert_eq!(
+ centroids.len(),
+ self.nlist * self.d,
+ "quantizer centroids must hold nlist * d values"
+ );
+ assert!(
+ self.ids.iter().all(Vec::is_empty),
+ "cannot replace quantizer centroids after vectors have been added"
+ );
+ self.quantizer_centroids = centroids;
+ self.coarse_assignment.reset();
+ }
+
pub fn train(&mut self, data: &[f32], n: usize) {
let train_data = self.preprocess_vectors(data, n);
self.quantizer_centroids =
kmeans::kmeans_train(&KMeansConfig::default(), &train_data, n,
self.d, self.nlist);
+ self.coarse_assignment.reset();
}
pub fn add(&mut self, data: &[f32], ids: &[i64], n: usize) {
let processed = self.preprocess_vectors(data, n);
- let list_ids = kmeans::find_nearest_batch(
+ let list_ids = self.coarse_assignment.assign(
&processed,
n,
&self.quantizer_centroids,
diff --git a/core/src/ivfflat_io.rs b/core/src/ivfflat_io.rs
index 99a5a91..7d0fccf 100644
--- a/core/src/ivfflat_io.rs
+++ b/core/src/ivfflat_io.rs
@@ -103,7 +103,7 @@ fn write_ivfflat_index_with_buffer_limit(
write_f32_slice(
out,
- &index.quantizer_centroids,
+ index.quantizer_centroids(),
&mut write_buffer,
buffer_limit,
)?;
@@ -1223,12 +1223,12 @@ fn validate_index_shape(index: &IVFFlatIndex) ->
io::Result<()> {
));
}
let centroid_len = checked_section_size(index.nlist, index.d)?;
- if index.quantizer_centroids.len() != centroid_len {
+ if index.quantizer_centroids().len() != centroid_len {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!(
"centroid length {} does not match nlist*d {}",
- index.quantizer_centroids.len(),
+ index.quantizer_centroids().len(),
centroid_len
),
));
@@ -1441,11 +1441,13 @@ mod tests {
fn balanced_flat_index(d: usize, nlist: usize, rows_per_list: usize) ->
IVFFlatIndex {
let mut index = IVFFlatIndex::new(d, nlist, MetricType::L2);
- index.quantizer_centroids = (0..nlist)
- .flat_map(|list_id| {
- (0..d).map(move |dimension| list_id as f32 * 10.0 + dimension
as f32 * 0.01)
- })
- .collect();
+ index.set_quantizer_centroids(
+ (0..nlist)
+ .flat_map(|list_id| {
+ (0..d).map(move |dimension| list_id as f32 * 10.0 +
dimension as f32 * 0.01)
+ })
+ .collect(),
+ );
for list_id in 0..nlist {
index.ids[list_id] = (0..rows_per_list)
.map(|row| (list_id * rows_per_list + row) as i64)
@@ -1489,7 +1491,7 @@ mod tests {
const TEST_BUDGET: usize = 64;
let mut index = IVFFlatIndex::new(3, 1, MetricType::L2);
- index.quantizer_centroids = vec![1.0, 2.0, 3.0];
+ index.set_quantizer_centroids(vec![1.0, 2.0, 3.0]);
index.ids[0] = vec![50, 10, 40, 20, 30, 60, 0];
index.vectors[0] = index.ids[0]
.iter()
@@ -1518,7 +1520,7 @@ mod tests {
let wide_dimension = TEST_BUDGET / size_of::<f32>() + 1;
let mut wide_index = IVFFlatIndex::new(wide_dimension, 1,
MetricType::L2);
- wide_index.quantizer_centroids = vec![0.0; wide_dimension];
+ wide_index.set_quantizer_centroids(vec![0.0; wide_dimension]);
wide_index.ids[0] = vec![1];
wide_index.vectors[0] = vec![1.0; wide_dimension];
let mut wide = MaxWriteWriter {
@@ -1615,7 +1617,7 @@ mod tests {
#[test]
fn
test_ivfflat_reader_handles_unaligned_vector_suffix_without_copy_contract_change()
{
let mut index = IVFFlatIndex::new(3, 1, MetricType::L2);
- index.quantizer_centroids = vec![0.0; 3];
+ index.set_quantizer_centroids(vec![0.0; 3]);
// These deltas occupy 1, 2, and 3 bytes, so the raw-vector suffix
// starts at byte 18 inside the list payload instead of a f32 boundary.
index.ids[0] = vec![1, 130, 16_515];
@@ -1999,7 +2001,7 @@ mod tests {
#[test]
fn test_ivfflat_writer_validates_shape_before_writing() {
let mut index = IVFFlatIndex::new(2, 1, MetricType::L2);
- index.quantizer_centroids = vec![0.0, 0.0];
+ index.set_quantizer_centroids(vec![0.0, 0.0]);
index.ids[0] = vec![1, 2];
index.vectors[0] = vec![0.0, 0.0];
diff --git a/core/src/ivfpq.rs b/core/src/ivfpq.rs
index 55b5bd5..15e530c 100644
--- a/core/src/ivfpq.rs
+++ b/core/src/ivfpq.rs
@@ -15,6 +15,7 @@
// specific language governing permissions and limitations
// under the License.
+use crate::coarse::CoarseAssignment;
use crate::distance::{
fvec_inner_product, fvec_madd, fvec_normalize, pq_distance_four_codes,
pq_distance_from_table,
MetricType,
@@ -66,7 +67,7 @@ pub struct IVFPQIndex {
pub metric: MetricType,
pub by_residual: bool,
- pub quantizer_centroids: Vec<f32>,
+ quantizer_centroids: Vec<f32>,
pub pq: ProductQuantizer,
pub opq: Option<OPQMatrix>,
@@ -78,6 +79,7 @@ pub struct IVFPQIndex {
precomputed_table: Vec<f32>,
/// Block-layout packed codes for 4-bit FastScan. One per list.
fastscan_codes: Vec<Vec<u8>>,
+ coarse_assignment: CoarseAssignment,
}
impl IVFPQIndex {
@@ -114,9 +116,40 @@ impl IVFPQIndex {
codes: vec![Vec::new(); nlist],
precomputed_table: Vec::new(),
fastscan_codes: Vec::new(),
+ coarse_assignment: CoarseAssignment::default(),
}
}
+ pub fn quantizer_centroids(&self) -> &[f32] {
+ &self.quantizer_centroids
+ }
+
+ /// Enables automatic Vamana coarse assignment for large centroid matrices.
+ /// Disable it to keep vector assignment exact.
+ pub(crate) fn set_approximate_coarse_assignment(&mut self, enabled: bool) {
+ assert!(
+ self.ids.iter().all(Vec::is_empty),
+ "cannot change coarse assignment after vectors have been added"
+ );
+ self.coarse_assignment.set_approximate_enabled(enabled);
+ }
+
+ pub fn set_quantizer_centroids(&mut self, centroids: Vec<f32>) {
+ assert_eq!(
+ centroids.len(),
+ self.nlist * self.d,
+ "quantizer centroids must hold nlist * d values"
+ );
+ assert!(
+ self.ids.iter().all(Vec::is_empty),
+ "cannot replace quantizer centroids after vectors have been added"
+ );
+ self.quantizer_centroids = centroids;
+ self.precomputed_table.clear();
+ self.fastscan_codes.clear();
+ self.coarse_assignment.reset();
+ }
+
/// Create an index with automatic nlist based on target partition size.
/// nlist = max(1, n / target_partition_size), clamped to reasonable
bounds.
pub fn with_target_partition_size(
@@ -135,7 +168,7 @@ impl IVFPQIndex {
/// The new index has empty inverted lists — call `add()` to populate.
/// Used for distributed build: train once globally, then each worker
creates from_trained.
pub fn from_trained(trained: &IVFPQIndex) -> Self {
- IVFPQIndex {
+ let mut index = IVFPQIndex {
d: trained.d,
nlist: trained.nlist,
metric: trained.metric,
@@ -165,7 +198,10 @@ impl IVFPQIndex {
codes: vec![Vec::new(); trained.nlist],
precomputed_table: Vec::new(),
fastscan_codes: Vec::new(),
- }
+ coarse_assignment: CoarseAssignment::default(),
+ };
+
index.set_approximate_coarse_assignment(trained.coarse_assignment.approximate_enabled());
+ index
}
pub fn train(&mut self, data: &[f32], n: usize) {
@@ -196,12 +232,26 @@ impl IVFPQIndex {
let km_config = KMeansConfig::default();
self.quantizer_centroids =
kmeans::kmeans_train(&km_config, &effective_data, n, d,
self.nlist);
+ self.coarse_assignment.reset();
- // Retrain PQ on the exact distribution that add/search will encode.
+ // Retrain PQ on the same assignment distribution that add/search will
encode.
// For OPQ: opq.train() trained PQ on centered data, but add/search
// encode uncentered vectors, so we must retrain here for all metrics.
let pq_train_data = if self.by_residual {
- compute_residuals(&effective_data, n, d,
&self.quantizer_centroids, self.nlist)
+ let assignments = self.coarse_assignment.assign(
+ &effective_data,
+ n,
+ &self.quantizer_centroids,
+ self.nlist,
+ d,
+ );
+ compute_residuals(
+ &effective_data,
+ n,
+ d,
+ &self.quantizer_centroids,
+ &assignments,
+ )
} else {
effective_data
};
@@ -229,7 +279,8 @@ impl IVFPQIndex {
// L2/IP without OPQ borrows the caller's batch instead of copying it.
let processed = self.preprocess_queries(data, n);
let assignments =
- kmeans::find_nearest_batch(&processed, n,
&self.quantizer_centroids, self.nlist, d);
+ self.coarse_assignment
+ .assign(&processed, n, &self.quantizer_centroids, self.nlist,
d);
let to_encode = if self.by_residual {
let mut residuals = vec![0.0f32; n * d];
@@ -2839,10 +2890,9 @@ fn compute_residuals(
n: usize,
d: usize,
centroids: &[f32],
- nlist: usize,
+ assignments: &[usize],
) -> Vec<f32> {
let mut residuals = vec![0.0f32; n * d];
- let assignments = kmeans::find_nearest_batch(data, n, centroids, nlist, d);
residuals
.par_chunks_mut(d)
.enumerate()
diff --git a/core/src/ivfrq.rs b/core/src/ivfrq.rs
index b446780..86596ae 100644
--- a/core/src/ivfrq.rs
+++ b/core/src/ivfrq.rs
@@ -15,9 +15,11 @@
// specific language governing permissions and limitations
// under the License.
+use crate::coarse::CoarseAssignment;
use crate::distance::{fvec_madd, fvec_norm_l2sqr, preprocess_vectors,
MetricType};
use crate::ivfpq::RowIdFilter;
use crate::kmeans::{self, KMeansConfig};
+use crate::logging::{emit_log, LogLevel};
use crate::rq::{
RQEncodeScratch, RQQueryContext, RQRotation, RQVectorFactors,
RaBitQuantizer, DEFAULT_RQ_BITS,
DEFAULT_RQ_ROTATION_ROUNDS, DEFAULT_RQ_ROTATION_SEED,
@@ -25,6 +27,27 @@ use crate::rq::{
use crate::topk::TopKHeap;
use rayon::prelude::*;
use std::borrow::Cow;
+use std::time::{Duration, Instant};
+
+pub(crate) fn build_timing_enabled() -> bool {
+ std::env::var_os("PAIMON_LOG_VECTOR_INDEX_BUILD_TIMING").is_some()
+}
+
+pub(crate) fn log_build_timing(enabled: bool, stage: &str, started: Instant) {
+ log_build_elapsed(enabled, stage, started.elapsed());
+}
+
+pub(crate) fn log_build_elapsed(enabled: bool, stage: &str, elapsed: Duration)
{
+ if enabled {
+ emit_log(
+ LogLevel::Info,
+ &format!(
+ "ivf_rq_build stage={stage} elapsed_ms={:.3}",
+ elapsed.as_secs_f64() * 1_000.0
+ ),
+ );
+ }
+}
pub struct IVFRQIndex {
pub d: usize,
@@ -32,7 +55,7 @@ pub struct IVFRQIndex {
pub nlist: usize,
pub bits: usize,
pub metric: MetricType,
- pub quantizer_centroids: Vec<f32>,
+ quantizer_centroids: Vec<f32>,
pub quantizer_centroid_norms: Vec<f32>,
pub rotated_centroids: Vec<f32>,
pub rotation_seed: u64,
@@ -42,6 +65,7 @@ pub struct IVFRQIndex {
pub factors: Vec<Vec<RQVectorFactors>>,
quantizer: RaBitQuantizer,
rotation: RQRotation,
+ coarse_assignment: CoarseAssignment,
}
impl IVFRQIndex {
@@ -93,44 +117,99 @@ impl IVFRQIndex {
factors: vec![Vec::new(); nlist],
quantizer,
rotation: RQRotation::new(d, rotation_seed, rotation_rounds),
+ coarse_assignment: CoarseAssignment::default(),
}
}
+ pub fn quantizer_centroids(&self) -> &[f32] {
+ &self.quantizer_centroids
+ }
+
+ /// Enables automatic Vamana coarse assignment for large centroid matrices.
+ /// Disable it to keep vector assignment exact.
+ pub(crate) fn set_approximate_coarse_assignment(&mut self, enabled: bool) {
+ assert!(
+ self.ids.iter().all(Vec::is_empty),
+ "cannot change coarse assignment after vectors have been added"
+ );
+ self.coarse_assignment.set_approximate_enabled(enabled);
+ }
+
+ pub fn set_quantizer_centroids(&mut self, centroids: Vec<f32>) {
+ assert_eq!(
+ centroids.len(),
+ self.nlist * self.d,
+ "quantizer centroids must hold nlist * d values"
+ );
+ assert!(
+ self.ids.iter().all(Vec::is_empty),
+ "cannot replace quantizer centroids after vectors have been added"
+ );
+ self.quantizer_centroids = centroids;
+ self.rebuild_centroid_norms();
+ self.rebuild_rotated_centroids();
+ self.coarse_assignment.reset();
+ }
+
pub fn train(&mut self, data: &[f32], n: usize) {
+ let timing = build_timing_enabled();
+ let total_started = Instant::now();
+ let phase_started = Instant::now();
let processed = self.preprocess_vectors(data, n);
+ log_build_timing(timing, "train.preprocess", phase_started);
+
+ let phase_started = Instant::now();
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 {
- let centroid = &self.quantizer_centroids[list_id *
self.d..(list_id + 1) * self.d];
- self.rotation.rotate(
- centroid,
- &mut self.rotated_centroids[list_id * self.padded_d..(list_id
+ 1) * self.padded_d],
- &mut scratch,
- );
- }
+ log_build_timing(timing, "train.kmeans", phase_started);
+
+ let phase_started = Instant::now();
+ self.rebuild_centroid_norms();
+ log_build_timing(timing, "train.centroid_norms", phase_started);
+
+ let phase_started = Instant::now();
+ self.rebuild_rotated_centroids();
+ log_build_timing(timing, "train.rotate_centroids", phase_started);
+
+ let phase_started = Instant::now();
+ self.coarse_assignment.reset();
+ log_build_timing(timing, "train.assign_graph", phase_started);
+ log_build_timing(timing, "train.total", total_started);
}
pub fn add(&mut self, data: &[f32], ids: &[i64], n: usize) {
+ if n == 0 {
+ return;
+ }
+ let timing = build_timing_enabled();
+ let total_started = Instant::now();
+ let phase_started = Instant::now();
+ self.coarse_assignment
+ .prepare(&self.quantizer_centroids, self.nlist, self.d);
+ log_build_timing(timing, "add.assign_graph", phase_started);
+
+ let phase_started = Instant::now();
let processed = self.preprocess_vectors(data, n);
- let list_ids = kmeans::find_nearest_batch(
+ log_build_timing(timing, "add.preprocess", phase_started);
+
+ let phase_started = Instant::now();
+ let list_ids = self.coarse_assignment.assign(
&processed,
n,
&self.quantizer_centroids,
self.nlist,
self.d,
);
+ log_build_timing(timing, "add.assign", phase_started);
+
+ let phase_started = Instant::now();
let mut list_rows = vec![Vec::new(); self.nlist];
for (row, list_id) in list_ids.into_iter().enumerate() {
list_rows[list_id].push(row);
}
+ log_build_timing(timing, "add.group", phase_started);
+ let phase_started = Instant::now();
let d = self.d;
let padded_d = self.padded_d;
let metric = self.metric;
@@ -148,14 +227,14 @@ impl IVFRQIndex {
.zip(output_factors.par_iter_mut())
.zip(list_rows.into_par_iter())
.enumerate()
- .for_each(
- |(list_id, (((list_ids, list_codes), list_factors),
rows))| {
+ .for_each_init(
+ || IVFRQEncodeScratch::new(d, padded_d,
quantizer.code_size()),
+ |scratch, (list_id, (((list_ids, list_codes),
list_factors), rows))| {
append_encoded_rows(
&processed,
ids,
&rows,
d,
- padded_d,
¢roids[list_id * d..(list_id + 1) * d],
&rotated_centroids[list_id * padded_d..(list_id +
1) * padded_d],
metric,
@@ -164,10 +243,12 @@ impl IVFRQIndex {
list_ids,
list_codes,
list_factors,
+ scratch,
);
},
);
} else {
+ let mut scratch = IVFRQEncodeScratch::new(d, padded_d,
quantizer.code_size());
for (list_id, (((list_ids, list_codes), list_factors), rows)) in
output_ids
.iter_mut()
.zip(output_codes.iter_mut())
@@ -180,7 +261,6 @@ impl IVFRQIndex {
ids,
&rows,
d,
- padded_d,
¢roids[list_id * d..(list_id + 1) * d],
&rotated_centroids[list_id * padded_d..(list_id + 1) *
padded_d],
metric,
@@ -189,9 +269,33 @@ impl IVFRQIndex {
list_ids,
list_codes,
list_factors,
+ &mut scratch,
);
}
}
+ log_build_timing(timing, "add.encode", phase_started);
+ log_build_timing(timing, "add.total", total_started);
+ }
+
+ fn rebuild_centroid_norms(&mut self) {
+ self.quantizer_centroid_norms = self
+ .quantizer_centroids
+ .chunks_exact(self.d)
+ .map(fvec_norm_l2sqr)
+ .collect();
+ }
+
+ fn rebuild_rotated_centroids(&mut self) {
+ 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 {
+ let centroid = &self.quantizer_centroids[list_id *
self.d..(list_id + 1) * self.d];
+ self.rotation.rotate(
+ centroid,
+ &mut self.rotated_centroids[list_id * self.padded_d..(list_id
+ 1) * self.padded_d],
+ &mut scratch,
+ );
+ }
}
pub fn total_vectors(&self) -> usize {
@@ -339,13 +443,32 @@ impl IVFRQIndex {
}
}
+struct IVFRQEncodeScratch {
+ residual: Vec<f32>,
+ rotated_residual: Vec<f32>,
+ rotation: Vec<f32>,
+ code: Vec<u8>,
+ encode: RQEncodeScratch,
+}
+
+impl IVFRQEncodeScratch {
+ fn new(d: usize, padded_d: usize, code_size: usize) -> Self {
+ Self {
+ residual: vec![0.0; d],
+ rotated_residual: vec![0.0; padded_d],
+ rotation: vec![0.0; padded_d],
+ code: vec![0; code_size],
+ encode: RQEncodeScratch::new(padded_d),
+ }
+ }
+}
+
#[allow(clippy::too_many_arguments)]
fn append_encoded_rows(
data: &[f32],
input_ids: &[i64],
rows: &[usize],
d: usize,
- padded_d: usize,
centroid: &[f32],
rotated_centroid: &[f32],
metric: MetricType,
@@ -354,29 +477,29 @@ fn append_encoded_rows(
output_ids: &mut Vec<i64>,
output_codes: &mut Vec<u8>,
output_factors: &mut Vec<RQVectorFactors>,
+ scratch: &mut IVFRQEncodeScratch,
) {
let code_size = quantizer.code_size();
output_ids.reserve(rows.len());
output_codes.reserve(rows.len().saturating_mul(code_size));
output_factors.reserve(rows.len());
- let mut residual = vec![0.0f32; d];
- let mut rotated_residual = vec![0.0f32; padded_d];
- let mut rotation_scratch = vec![0.0f32; padded_d];
- let mut code = vec![0u8; code_size];
- let mut encode_scratch = RQEncodeScratch::new(padded_d);
for &row in rows {
let vector = &data[row * d..(row + 1) * d];
- fvec_madd(vector, centroid, -1.0, &mut residual);
- rotation.rotate(&residual, &mut rotated_residual, &mut
rotation_scratch);
+ fvec_madd(vector, centroid, -1.0, &mut scratch.residual);
+ rotation.rotate(
+ &scratch.residual,
+ &mut scratch.rotated_residual,
+ &mut scratch.rotation,
+ );
let factors = quantizer.encode_with_scratch(
- &rotated_residual,
+ &scratch.rotated_residual,
rotated_centroid,
metric,
- &mut code,
- &mut encode_scratch,
+ &mut scratch.code,
+ &mut scratch.encode,
);
output_ids.push(input_ids[row]);
- output_codes.extend_from_slice(&code);
+ output_codes.extend_from_slice(&scratch.code);
output_factors.push(factors);
}
}
@@ -385,6 +508,29 @@ fn append_encoded_rows(
mod tests {
use super::*;
+ #[test]
+ fn ivfrq_replacing_centroids_rebuilds_derived_state() {
+ let data = vec![0.0, 0.0, 10.0, 10.0];
+ let mut index = IVFRQIndex::with_bits(1, 2, 4, MetricType::L2);
+ index.train(&data, 4);
+ index.set_quantizer_centroids(vec![100.0, 7.0]);
+
+ assert_eq!(index.quantizer_centroids(), &[100.0, 7.0]);
+ assert_eq!(index.quantizer_centroid_norms, vec![10_000.0, 49.0]);
+
+ index.add(&[7.0], &[9], 1);
+ assert_eq!(index.ids[1], vec![9]);
+ }
+
+ #[test]
+ fn ivfrq_empty_add_does_not_prepare_coarse_graph() {
+ let mut index = IVFRQIndex::new(256, 4096, MetricType::L2);
+
+ index.add(&[], &[], 0);
+
+ assert!(!index.coarse_assignment.build_attempted());
+ }
+
#[test]
fn ivfrq_only_allocates_preprocessed_vectors_for_cosine() {
let data = vec![3.0, 4.0, 1.0, 2.0];
diff --git a/core/src/ivfrq_io.rs b/core/src/ivfrq_io.rs
index fc8d6d5..1aee33a 100644
--- a/core/src/ivfrq_io.rs
+++ b/core/src/ivfrq_io.rs
@@ -22,7 +22,7 @@ use crate::index_io_util::{
};
use crate::io::{PreadCursor, ReadRequest, SeekRead, SeekWrite};
use crate::ivfpq::RowIdFilter;
-use crate::ivfrq::IVFRQIndex;
+use crate::ivfrq::{build_timing_enabled, log_build_elapsed, log_build_timing,
IVFRQIndex};
use crate::kmeans;
use crate::rq::{
is_supported_rq_bits, padded_dimension, RQCodeFactors, RQQueryContext,
RQQueryTerms,
@@ -104,13 +104,18 @@ struct RQListWritePlan {
}
pub fn write_ivfrq_index(index: &IVFRQIndex, out: &mut dyn SeekWrite) ->
io::Result<()> {
+ let timing = build_timing_enabled();
+ let total_started = std::time::Instant::now();
validate_index_shape(index)?;
let total_vectors = index.ids.iter().try_fold(0i64, |sum, ids| {
sum.checked_add(usize_to_i64(ids.len(), "total vector count")?)
.ok_or_else(|| invalid_input("total vector count exceeds i64"))
})?;
+ let phase_started = std::time::Instant::now();
let write_plans = plan_sorted_lists(index);
+ log_build_timing(timing, "write.plan_lists", phase_started);
+ let phase_started = std::time::Instant::now();
write_u32_le(out, IVF_RQ_MAGIC)?;
write_u32_le(out, IVF_RQ_VERSION)?;
write_i32_le(out, usize_to_i32(index.d, "dimension")?)?;
@@ -126,7 +131,7 @@ pub fn write_ivfrq_index(index: &IVFRQIndex, out: &mut dyn
SeekWrite) -> io::Res
write_u32_le(out, IVF_RQ_ROTATION_TYPE_BLOCK_FHT)?;
write_u32_le(out, IVF_RQ_FACTOR_LAYOUT_COMPACT_V1)?;
- write_f32_slice(out, &index.quantizer_centroids)?;
+ write_f32_slice(out, index.quantizer_centroids())?;
let offset_table_bytes = index
.nlist
@@ -165,12 +170,20 @@ pub fn write_ivfrq_index(index: &IVFRQIndex, out: &mut
dyn SeekWrite) -> io::Res
write_i32_le(out, count)?;
write_i32_le(out, id_bytes)?;
}
+ log_build_timing(timing, "write.metadata", phase_started);
+ let phase_started = std::time::Instant::now();
+ let mut block_lists_elapsed = std::time::Duration::ZERO;
+ let mut output_elapsed = std::time::Duration::ZERO;
for (list_id, plan) in write_plans.into_iter().enumerate() {
if plan.order.is_empty() {
continue;
}
+ let list_started = std::time::Instant::now();
let (blocked_codes, blocked_factors) = block_list(index, list_id,
&plan.order);
+ block_lists_elapsed += list_started.elapsed();
+
+ let output_started = std::time::Instant::now();
write_i64_le(out, plan.base_id)?;
write_i32_le(out, usize_to_i32(plan.id_bytes.len(), "delta ID
bytes")?)?;
write_i32_le(
@@ -180,7 +193,12 @@ pub fn write_ivfrq_index(index: &IVFRQIndex, out: &mut dyn
SeekWrite) -> io::Res
out.write_all(&plan.id_bytes)?;
out.write_all(&blocked_codes)?;
write_f32_slice(out, &blocked_factors)?;
+ output_elapsed += output_started.elapsed();
}
+ log_build_elapsed(timing, "write.block_lists", block_lists_elapsed);
+ log_build_elapsed(timing, "write.output", output_elapsed);
+ log_build_timing(timing, "write.lists", phase_started);
+ log_build_timing(timing, "write.total", total_started);
Ok(())
}
@@ -1171,7 +1189,7 @@ fn validate_index_shape(index: &IVFRQIndex) ->
io::Result<()> {
"IVF-RQ rotation rounds must be {DEFAULT_RQ_ROTATION_ROUNDS}"
)));
}
- if index.quantizer_centroids.len() != checked_section_size(index.nlist,
index.d)?
+ if index.quantizer_centroids().len() != checked_section_size(index.nlist,
index.d)?
|| index.rotated_centroids.len() != checked_section_size(index.nlist,
index.padded_d)?
{
return Err(invalid_input("IVF-RQ centroid storage shape mismatch"));
diff --git a/core/src/ivfsq.rs b/core/src/ivfsq.rs
index 6d16191..c39a62e 100644
--- a/core/src/ivfsq.rs
+++ b/core/src/ivfsq.rs
@@ -17,6 +17,7 @@
//! IVF with per-list, per-dimension 8-bit residual scalar quantization.
+use crate::coarse::CoarseAssignment;
use crate::distance::{fvec_madd, preprocess_vectors, MetricType};
use crate::ivfpq::RowIdFilter;
use crate::kmeans::{self, KMeansConfig};
@@ -29,11 +30,12 @@ pub struct IVFSQIndex {
pub d: usize,
pub nlist: usize,
pub metric: MetricType,
- pub quantizer_centroids: Vec<f32>,
+ quantizer_centroids: Vec<f32>,
pub sq: ScalarQuantizer,
pub list_sqs: Vec<ScalarQuantizer>,
pub ids: Vec<Vec<i64>>,
pub codes: Vec<Vec<u8>>,
+ coarse_assignment: CoarseAssignment,
}
impl IVFSQIndex {
@@ -47,13 +49,43 @@ impl IVFSQIndex {
list_sqs: vec![ScalarQuantizer::new(d); nlist],
ids: vec![Vec::new(); nlist],
codes: vec![Vec::new(); nlist],
+ coarse_assignment: CoarseAssignment::default(),
}
}
+ pub fn quantizer_centroids(&self) -> &[f32] {
+ &self.quantizer_centroids
+ }
+
+ /// Enables automatic Vamana coarse assignment for large centroid matrices.
+ /// Disable it to keep vector assignment exact.
+ pub(crate) fn set_approximate_coarse_assignment(&mut self, enabled: bool) {
+ assert!(
+ self.ids.iter().all(Vec::is_empty),
+ "cannot change coarse assignment after vectors have been added"
+ );
+ self.coarse_assignment.set_approximate_enabled(enabled);
+ }
+
+ pub fn set_quantizer_centroids(&mut self, centroids: Vec<f32>) {
+ assert_eq!(
+ centroids.len(),
+ self.nlist * self.d,
+ "quantizer centroids must hold nlist * d values"
+ );
+ assert!(
+ self.ids.iter().all(Vec::is_empty),
+ "cannot replace quantizer centroids after vectors have been added"
+ );
+ self.quantizer_centroids = centroids;
+ self.coarse_assignment.reset();
+ }
+
pub fn train(&mut self, data: &[f32], n: usize) {
let processed = self.preprocess_vectors(data, n);
self.quantizer_centroids =
kmeans::kmeans_train(&KMeansConfig::default(), &processed, n,
self.d, self.nlist);
+ self.coarse_assignment.reset();
let (list_ids, residuals) = self.assign_residuals(&processed, n);
self.sq.train(&residuals, n);
self.train_list_sqs(&list_ids, &residuals);
@@ -61,7 +93,7 @@ impl IVFSQIndex {
pub fn add(&mut self, data: &[f32], ids: &[i64], n: usize) {
let processed = self.preprocess_vectors(data, n);
- let list_ids = kmeans::find_nearest_batch(
+ let list_ids = self.coarse_assignment.assign(
&processed,
n,
&self.quantizer_centroids,
@@ -215,9 +247,14 @@ impl IVFSQIndex {
}
}
- fn assign_residuals(&self, processed: &[f32], n: usize) -> (Vec<usize>,
Vec<f32>) {
- let list_ids =
- kmeans::find_nearest_batch(processed, n,
&self.quantizer_centroids, self.nlist, self.d);
+ fn assign_residuals(&mut self, processed: &[f32], n: usize) ->
(Vec<usize>, Vec<f32>) {
+ let list_ids = self.coarse_assignment.assign(
+ processed,
+ n,
+ &self.quantizer_centroids,
+ self.nlist,
+ self.d,
+ );
let mut residuals = vec![0.0f32; n * self.d];
for i in 0..n {
let vector = &processed[i * self.d..(i + 1) * self.d];
diff --git a/core/src/ivfsq_io.rs b/core/src/ivfsq_io.rs
index 0bc5819..6fb7c5b 100644
--- a/core/src/ivfsq_io.rs
+++ b/core/src/ivfsq_io.rs
@@ -80,7 +80,7 @@ pub fn write_ivfsq_index(index: &IVFSQIndex, out: &mut dyn
SeekWrite) -> io::Res
write_f32_slice(out, &sq.mins)?;
write_f32_slice(out, &sq.maxs)?;
}
- write_f32_slice(out, &index.quantizer_centroids)?;
+ write_f32_slice(out, index.quantizer_centroids())?;
let offset_table_size = index.nlist.checked_mul(16).ok_or_else(|| {
io::Error::new(
@@ -971,7 +971,7 @@ fn validate_index_shape(index: &IVFSQIndex) ->
io::Result<()> {
"IVF-SQ inverted-list state does not match nlist",
));
}
- if index.quantizer_centroids.len() != checked_section_size(index.nlist,
index.d)? {
+ if index.quantizer_centroids().len() != checked_section_size(index.nlist,
index.d)? {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
"IVF-SQ centroid storage does not match nlist * dimension",
diff --git a/core/src/kmeans.rs b/core/src/kmeans.rs
index 485a431..55fd814 100644
--- a/core/src/kmeans.rs
+++ b/core/src/kmeans.rs
@@ -16,7 +16,7 @@
// under the License.
use crate::blas::sgemm_a_bt;
-use crate::distance::{fvec_l2sqr, fvec_norm_l2sqr};
+use crate::distance::{fvec_l2sqr, fvec_l2sqr_four, fvec_norm_l2sqr};
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
use rayon::prelude::*;
@@ -688,7 +688,23 @@ fn update_centroids(
pub fn find_nearest(point: &[f32], centroids: &[f32], k: usize, d: usize) ->
usize {
let mut best = 0;
let mut best_dist = f32::MAX;
- for c in 0..k {
+ let four_end = k / 4 * 4;
+ for c in (0..four_end).step_by(4) {
+ let distances = fvec_l2sqr_four(
+ point,
+ ¢roids[c * d..(c + 1) * d],
+ ¢roids[(c + 1) * d..(c + 2) * d],
+ ¢roids[(c + 2) * d..(c + 3) * d],
+ ¢roids[(c + 3) * d..(c + 4) * d],
+ );
+ for (offset, dist) in distances.into_iter().enumerate() {
+ if dist < best_dist {
+ best_dist = dist;
+ best = c + offset;
+ }
+ }
+ }
+ for c in four_end..k {
let dist = fvec_l2sqr(point, ¢roids[c * d..(c + 1) * d]);
if dist < best_dist {
best_dist = dist;
@@ -728,16 +744,32 @@ pub fn find_topk(
if nprobe == 0 {
return (Vec::new(), Vec::new());
}
- let mut dists: Vec<(f32, usize)> = (0..k)
- .map(|c| (fvec_l2sqr(point, ¢roids[c * d..(c + 1) * d]), c))
- .collect();
+ let mut dists = Vec::with_capacity(k);
+ let four_end = k / 4 * 4;
+ for c in (0..four_end).step_by(4) {
+ let distances = fvec_l2sqr_four(
+ point,
+ ¢roids[c * d..(c + 1) * d],
+ ¢roids[(c + 1) * d..(c + 2) * d],
+ ¢roids[(c + 2) * d..(c + 3) * d],
+ ¢roids[(c + 3) * d..(c + 4) * d],
+ );
+ dists.extend(
+ distances
+ .into_iter()
+ .enumerate()
+ .map(|(offset, distance)| (distance, c + offset)),
+ );
+ }
+ dists.extend((four_end..k).map(|c| (fvec_l2sqr(point, ¢roids[c * d..(c
+ 1) * d]), c)));
select_topk_prefix(&mut dists, nprobe);
let indices: Vec<usize> = dists[..nprobe].iter().map(|&(_, i)|
i).collect();
let distances: Vec<f32> = dists[..nprobe].iter().map(|&(d, _)|
d).collect();
(indices, distances)
}
-/// Batch find top-nprobe nearest centroids for multiple queries using sgemm.
+/// Batch find top-nprobe nearest centroids using SGEMM, with direct L2
fallback when
+/// cancellation error can dominate the nearest computed distance.
/// Returns (all_indices, all_distances) each of length nq * nprobe.
pub fn find_topk_batch(
queries: &[f32],
@@ -748,7 +780,7 @@ pub fn find_topk_batch(
nprobe: usize,
) -> (Vec<Vec<usize>>, Vec<Vec<f32>>) {
let centroid_norms = (0..k)
- .map(|c| fvec_norm_l2sqr(¢roids[c * d..(c + 1) * d]))
+ .map(|centroid| fvec_norm_l2sqr(¢roids[centroid * d..(centroid +
1) * d]))
.collect::<Vec<_>>();
find_topk_batch_with_centroid_norms(queries, nq, centroids,
¢roid_norms, k, d, nprobe)
}
@@ -767,41 +799,154 @@ pub(crate) fn find_topk_batch_with_centroid_norms(
if nprobe == 0 {
return (vec![Vec::new(); nq], vec![Vec::new(); nq]);
}
-
- if nq == 1 {
- let (indices, distances) = find_topk(&queries[..d], centroids, k, d,
nprobe);
- return (vec![indices], vec![distances]);
+ if nq == 1
+ || nprobe == k
+ || k > MAX_MATRIX_ELEMS
+ || centroid_norms
+ .iter()
+ .any(|&norm| !norm.is_finite() || norm < 0.0)
+ {
+ return (0..nq)
+ .into_par_iter()
+ .map(|query| {
+ find_topk(
+ &queries[query * d..(query + 1) * d],
+ centroids,
+ k,
+ d,
+ nprobe,
+ )
+ })
+ .unzip();
}
- // Precompute norms
- let q_norms: Vec<f32> = (0..nq)
- .map(|i| fvec_norm_l2sqr(&queries[i * d..(i + 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);
-
- // Extract top-nprobe per query
+ let max_centroid_norm = centroid_norms.iter().copied().fold(0.0, f32::max);
+ // Bound the SGEMM result matrix to the same 16 MiB scratch budget as
assignment.
+ let tile_rows = (MAX_MATRIX_ELEMS / k).max(1);
let mut all_indices = Vec::with_capacity(nq);
let mut all_distances = Vec::with_capacity(nq);
+ for query_start in (0..nq).step_by(tile_rows) {
+ let tile_n = tile_rows.min(nq - query_start);
+ let tile_queries = &queries[query_start * d..(query_start + tile_n) *
d];
+ let query_norms = (0..tile_n)
+ .map(|query| fvec_norm_l2sqr(&tile_queries[query * d..(query + 1)
* d]))
+ .collect::<Vec<_>>();
+ let mut inner_products = vec![0.0f32; tile_n * k];
+ sgemm_a_bt(
+ tile_n,
+ k,
+ d,
+ 1.0,
+ tile_queries,
+ centroids,
+ 0.0,
+ &mut inner_products,
+ );
- for qi in 0..nq {
- let row = qi * k;
- let mut dists: Vec<(f32, usize)> = (0..k)
- .map(|c| {
- let dist = q_norms[qi] + centroid_norms[c] - 2.0 *
ip_matrix[row + c];
- (dist.max(0.0), c)
- })
- .collect();
- select_topk_prefix(&mut dists, nprobe);
+ for (query, (&query_norm, inner_products)) in query_norms
+ .iter()
+ .zip(inner_products.chunks_exact(k))
+ .enumerate()
+ {
+ let query = &tile_queries[query * d..(query + 1) * d];
+ if !query_norm.is_finite() {
+ let (indices, distances) = find_topk(query, centroids, k, d,
nprobe);
+ all_indices.push(indices);
+ all_distances.push(distances);
+ continue;
+ }
- all_indices.push(dists[..nprobe].iter().map(|&(_, i)| i).collect());
- all_distances.push(dists[..nprobe].iter().map(|&(d, _)| d).collect());
- }
+ let max_scale =
+ query_norm + max_centroid_norm + 2.0 * (query_norm *
max_centroid_norm).sqrt();
+ let mut all_finite = true;
+ let approximate = centroid_norms
+ .iter()
+ .zip(inner_products)
+ .enumerate()
+ .map(|(centroid, (¢roid_norm, &inner_product))| {
+ let distance = query_norm + centroid_norm - 2.0 *
inner_product;
+ all_finite &= distance.is_finite();
+ (distance.max(0.0), centroid)
+ })
+ .collect::<Vec<_>>();
+ // Clamped non-finite values are never consumed: the complete
query falls back.
+ if !all_finite {
+ let (indices, distances) = find_topk(query, centroids, k, d,
nprobe);
+ all_indices.push(indices);
+ all_distances.push(distances);
+ continue;
+ }
+ let rounding = d as f32 * f32::EPSILON;
+ let error_factor = if rounding < 1.0 {
+ 4.0 * rounding / (1.0 - rounding)
+ } else {
+ f32::INFINITY
+ };
+ let error_bound = if max_scale == 0.0 {
+ 0.0
+ } else {
+ max_scale * error_factor
+ };
+ let (indices, distances) =
+ refine_topk_boundary(query, centroids, k, d, nprobe,
approximate, error_bound);
+ all_indices.push(indices);
+ all_distances.push(distances);
+ }
+ }
(all_indices, all_distances)
}
+fn refine_topk_boundary(
+ query: &[f32],
+ centroids: &[f32],
+ k: usize,
+ d: usize,
+ nprobe: usize,
+ mut approximate: Vec<(f32, usize)>,
+ error_bound: f32,
+) -> (Vec<usize>, Vec<f32>) {
+ select_topk_prefix(&mut approximate, nprobe);
+ let boundary = approximate[nprobe - 1].0;
+ // Values farther than 2E from the approximate boundary cannot cross it
when
+ // every SGEMM-expanded distance has absolute error at most E.
+ let lower = boundary - 2.0 * error_bound;
+ let upper = boundary + 2.0 * error_bound;
+ let mut selected = Vec::with_capacity(nprobe);
+ let mut ambiguous = Vec::new();
+ for &(distance, centroid) in &approximate {
+ if distance < lower {
+ selected.push((
+ fvec_l2sqr(query, ¢roids[centroid * d..(centroid + 1) *
d]),
+ centroid,
+ ));
+ } else if distance <= upper {
+ ambiguous.push(centroid);
+ }
+ }
+
+ if ambiguous.len() == k {
+ return find_topk(query, centroids, k, d, nprobe);
+ }
+ let needed = nprobe - selected.len();
+ let mut ambiguous = ambiguous
+ .into_iter()
+ .map(|centroid| {
+ (
+ fvec_l2sqr(query, ¢roids[centroid * d..(centroid + 1) *
d]),
+ centroid,
+ )
+ })
+ .collect::<Vec<_>>();
+ ambiguous.sort_by(compare_distance_then_index);
+ selected.extend(ambiguous.into_iter().take(needed));
+ selected.sort_by(compare_distance_then_index);
+
+ let indices = selected.iter().map(|&(_, index)| index).collect();
+ let distances = selected.iter().map(|&(distance, _)| distance).collect();
+ (indices, distances)
+}
+
fn select_topk_prefix(dists: &mut [(f32, usize)], nprobe: usize) {
debug_assert!(nprobe > 0 && nprobe <= dists.len());
if nprobe < dists.len() {
@@ -1378,6 +1523,81 @@ mod tests {
}
}
+ #[test]
+ fn test_find_topk_batch_matches_direct_distance_after_translation() {
+ let centroids = vec![1.0e8, 1.0e8, 1.0e8 + 8.0, 1.0e8];
+ let queries = vec![1.0e8 + 8.0, 1.0e8, 1.0e8 + 8.0, 1.0e8];
+ let expected = find_topk(&queries[..2], ¢roids, 2, 2, 1);
+
+ let (indices, distances) = find_topk_batch(&queries, 2, ¢roids, 2,
2, 1);
+
+ assert_eq!((indices[0].clone(), distances[0].clone()), expected);
+ }
+
+ #[test]
+ fn test_find_topk_batch_matches_scalar_for_non_finite_values() {
+ for (queries, centroids) in [
+ (
+ vec![f32::NAN, 0.0, 0.1, 0.0],
+ vec![0.0, 0.0, 1.0, 0.0, 2.0, 0.0],
+ ),
+ (
+ vec![0.0, 0.0, 0.1, 0.0],
+ vec![f32::NAN, 0.0, 1.0, 0.0, 2.0, 0.0],
+ ),
+ ] {
+ let expected = (0..2)
+ .map(|query| find_topk(&queries[query * 2..(query + 1) * 2],
¢roids, 3, 2, 2))
+ .collect::<Vec<_>>();
+ let actual = find_topk_batch(&queries, 2, ¢roids, 3, 2, 2);
+
+ assert_eq!(
+ actual.0,
+ expected
+ .iter()
+ .map(|result| result.0.clone())
+ .collect::<Vec<_>>()
+ );
+ assert!(actual
+ .1
+ .iter()
+ .zip(&expected)
+ .all(|(actual, expected)| actual
+ .iter()
+ .zip(&expected.1)
+ .all(|(actual, expected)| actual.to_bits() ==
expected.to_bits())));
+ }
+ }
+
+ #[test]
+ fn test_find_topk_batch_refines_an_ambiguous_nprobe_boundary() {
+ let query = [0.0];
+ let centroids = [10.0, 10.05, 10.04];
+ let approximate = vec![(100.0, 0), (100.75, 1), (101.05, 2)];
+ let expected = find_topk(&query, ¢roids, 3, 1, 2);
+
+ let actual = refine_topk_boundary(&query, ¢roids, 3, 1, 2,
approximate, 0.3);
+
+ assert_eq!(actual, expected);
+ }
+
+ #[test]
+ fn test_find_topk_batch_tiles_sgemm_scratch() {
+ let k = 4_096;
+ let nq = MAX_MATRIX_ELEMS / k + 1;
+ let centroids = (0..k).map(|centroid| centroid as
f32).collect::<Vec<_>>();
+ let queries = (0..nq).map(|query| query as f32).collect::<Vec<_>>();
+
+ let (indices, distances) = find_topk_batch(&queries, nq, ¢roids,
k, 1, 1);
+
+ assert_eq!(indices.len(), nq);
+ assert!(indices
+ .iter()
+ .enumerate()
+ .all(|(query, result)| result == &[query]));
+ assert!(distances.iter().all(|result| result == &[0.0]));
+ }
+
#[test]
fn test_find_nearest_batch_matches_scalar() {
let d = 5;
diff --git a/core/src/lib.rs b/core/src/lib.rs
index 6d9c64c..8a645a5 100644
--- a/core/src/lib.rs
+++ b/core/src/lib.rs
@@ -20,6 +20,7 @@
pub mod autotune;
pub mod blas;
+pub(crate) mod coarse;
pub mod diskann;
pub mod diskann_io;
pub(crate) mod diskann_search;
diff --git a/core/src/vamana.rs b/core/src/vamana.rs
index d711e89..fe037c7 100644
--- a/core/src/vamana.rs
+++ b/core/src/vamana.rs
@@ -158,6 +158,63 @@ pub(crate) fn estimate_vamana_memory_bytes(
})
}
+fn estimate_sequential_vamana_memory_bytes(
+ node_count: usize,
+ max_degree: usize,
+ search_list_size: usize,
+) -> Option<usize> {
+ let degree = max_degree.min(node_count.saturating_sub(1));
+ let search_list_size = search_list_size.min(node_count);
+ let edge_bytes = degree.checked_mul(size_of::<u32>())?;
+ let nested_graph =
node_count.checked_mul(edge_bytes.checked_add(size_of::<Vec<u32>>())?)?;
+ let compact_graph = node_count
+ .checked_mul(edge_bytes.checked_add(size_of::<u16>())?)?
+ .checked_add(
+ node_count
+ .div_ceil(PARALLEL_ADJACENCY_NODES_PER_SHARD)
+ .checked_mul(size_of::<AdjacencyShard>())?,
+ )?;
+ let conversion_peak = nested_graph.checked_add(compact_graph)?;
+
+ let expected_visited = search_list_size
+ .checked_mul(degree)?
+ .checked_add(1)?
+ .min(node_count);
+ let dense_states = node_count
+ .checked_mul(size_of::<u8>())?
+ .checked_add(expected_visited.checked_mul(size_of::<u32>())?)?;
+ let sparse_states = sparse_table_memory_bytes(expected_visited,
size_of::<u8>())?;
+ let visit_states = if sparse_states
+ .checked_mul(SPARSE_BUILD_VISITED_MIN_MEMORY_SAVINGS)
+ .is_some_and(|threshold| threshold < dense_states)
+ {
+ sparse_states
+ } else {
+ dense_states
+ };
+ let search_scratch = visit_states
+ .checked_add(
+ search_list_size
+ .checked_mul(3)?
+ .checked_mul(size_of::<ScoredNode>())?,
+ )?
+ .checked_add(
+ search_list_size
+ .checked_add(degree)?
+
.checked_mul(size_of::<ScoredNode>().checked_add(size_of::<u32>())?)?,
+ )?
+ .checked_add(search_list_size.checked_mul(size_of::<u32>())?)?
+ .checked_add(degree.checked_mul(2 * size_of::<u32>())?)?;
+ let pass_peak = compact_graph
+ .checked_add(node_count.checked_mul(size_of::<usize>())?)?
+ .checked_add(search_scratch)?;
+ let connectivity_peak = compact_graph
+ .checked_add(node_count.checked_mul(
+ size_of::<bool>() + size_of::<Option<usize>>() + 3 *
size_of::<usize>(),
+ )?)?;
+ Some(conversion_peak.max(pass_peak).max(connectivity_peak))
+}
+
pub(crate) fn estimate_sharded_vamana_memory_bytes(
node_count: usize,
dimension: usize,
@@ -232,6 +289,14 @@ pub(crate) fn estimate_sharded_vamana_memory_bytes(
}
impl VamanaGraph {
+ pub(crate) fn search_scratch(&self, search_list_size: usize) ->
GreedySearchScratch {
+ GreedySearchScratch::new(
+ self.adjacency.len(),
+ self.adjacency.max_degree,
+ search_list_size.min(self.adjacency.len()),
+ )
+ }
+
pub fn from_adjacency(entry_node: u32, adjacency: Vec<Vec<u32>>) -> Self {
let max_degree = adjacency.iter().map(Vec::len).max().unwrap_or(0);
Self {
@@ -531,6 +596,16 @@ impl VamanaGraph {
search_distance: BuildSearchDistance<'_>,
) -> io::Result<(Self, VamanaBuildStats)> {
validate_build_inputs(vectors, count, dimension, params)?;
+ validate_build_memory_budget(
+ estimate_vamana_memory_bytes(
+ count,
+ params.max_degree,
+ params.build_search_list_size,
+ rayon::current_num_threads(),
+ )
+ .map(|estimate|
estimate.build_peak_bytes.max(estimate.remap_peak_bytes)),
+ params.memory_budget_bytes,
+ )?;
let entry_node = centroid_entry(vectors, count, dimension, metric) as
u32;
let degree = params.max_degree.min(count.saturating_sub(1));
let initialization_started = Instant::now();
@@ -590,6 +665,14 @@ impl VamanaGraph {
params: DiskAnnBuildParams,
) -> io::Result<Self> {
validate_build_inputs(vectors, count, dimension, params)?;
+ validate_build_memory_budget(
+ estimate_sequential_vamana_memory_bytes(
+ count,
+ params.max_degree,
+ params.build_search_list_size,
+ ),
+ params.memory_budget_bytes,
+ )?;
let entry_node = centroid_entry(vectors, count, dimension, metric) as
u32;
let mut rng = StdRng::seed_from_u64(params.seed);
let degree = params.max_degree.min(count.saturating_sub(1));
@@ -711,47 +794,76 @@ impl VamanaGraph {
query: &[f32],
search_list_size: usize,
) -> Vec<ScoredNode> {
+ let mut scratch = self.search_scratch(search_list_size);
+ self.greedy_search_with_scratch(
+ vectors,
+ dimension,
+ metric,
+ query,
+ search_list_size,
+ &mut scratch,
+ );
+ scratch.results.into_sorted_vec()
+ }
+
+ pub(crate) fn greedy_search_best_with_scratch(
+ &self,
+ vectors: &[f32],
+ dimension: usize,
+ query: &[f32],
+ search_list_size: usize,
+ scratch: &mut GreedySearchScratch,
+ ) -> Option<ScoredNode> {
+ self.greedy_search_with_scratch(
+ vectors,
+ dimension,
+ MetricType::L2,
+ query,
+ search_list_size,
+ scratch,
+ );
+ scratch.results.iter().min().copied()
+ }
+
+ fn greedy_search_with_scratch(
+ &self,
+ vectors: &[f32],
+ dimension: usize,
+ metric: MetricType,
+ query: &[f32],
+ search_list_size: usize,
+ scratch: &mut GreedySearchScratch,
+ ) {
+ scratch.begin_search();
if search_list_size == 0 || self.adjacency.is_empty() {
- return Vec::new();
+ return;
}
- let mut visited = vec![false; self.adjacency.len()];
- let mut expanded = vec![false; self.adjacency.len()];
let entry = self.entry_node as usize;
- visited[entry] = true;
- let mut results = vec![ScoredNode {
- id: self.entry_node,
- distance: node_distance(vectors, dimension, entry, query, metric),
- }];
-
- loop {
- results.sort_by(scored_node_order);
- results.truncate(search_list_size);
- let Some(current) = results
- .iter()
- .find(|node| !expanded[node.id as usize])
- .copied()
- else {
- break;
- };
- expanded[current.id as usize] = true;
+ scratch.insert_candidate(
+ ScoredNode {
+ id: self.entry_node,
+ distance: node_distance(vectors, dimension, entry, query,
metric),
+ },
+ search_list_size,
+ );
+ while let Some(current) = scratch.pop_nearest_unexpanded() {
+ scratch.mark_expanded(current.id as usize);
for &neighbor in &self.adjacency[current.id as usize] {
let neighbor = neighbor as usize;
- if neighbor >= self.adjacency.len() || visited[neighbor] {
+ if neighbor >= self.adjacency.len() ||
scratch.is_visited(neighbor) {
continue;
}
- visited[neighbor] = true;
- results.push(ScoredNode {
- id: neighbor as u32,
- distance: node_distance(vectors, dimension, neighbor,
query, metric),
- });
+ scratch.insert_candidate(
+ ScoredNode {
+ id: neighbor as u32,
+ distance: node_distance(vectors, dimension, neighbor,
query, metric),
+ },
+ search_list_size,
+ );
}
}
-
- results.sort_by(scored_node_order);
- results.truncate(search_list_size);
- results
}
#[allow(clippy::too_many_arguments)]
@@ -1591,7 +1703,7 @@ enum BuildVisitStates {
Sparse(SparseTable<u8>),
}
-struct GreedySearchScratch {
+pub(crate) struct GreedySearchScratch {
visit_states: BuildVisitStates,
results: BinaryHeap<ScoredNode>,
frontier: BinaryHeap<Reverse<ScoredNode>>,
@@ -2181,6 +2293,22 @@ fn validate_build_inputs(
Ok(())
}
+fn validate_build_memory_budget(estimated_peak: Option<usize>, budget: usize)
-> io::Result<()> {
+ if estimated_peak.is_some_and(|peak| peak <= budget) {
+ return Ok(());
+ }
+ Err(io::Error::new(
+ io::ErrorKind::OutOfMemory,
+ format!(
+ "Vamana build requires {} bytes, exceeding the {budget} byte
memory budget",
+ estimated_peak.map_or_else(
+ || "an unrepresentable amount of memory".to_string(),
+ |peak| peak.to_string()
+ )
+ ),
+ ))
+}
+
fn centroid_entry(vectors: &[f32], count: usize, dimension: usize, metric:
MetricType) -> usize {
let mut centroid = vec![0.0f32; dimension];
for vector in vectors.chunks_exact(dimension) {
@@ -2465,6 +2593,15 @@ mod tests {
);
}
+ #[test]
+ fn vamana_greedy_search_caps_search_list_size_to_node_count() {
+ let graph = VamanaGraph::from_adjacency(0, vec![vec![]]);
+ assert_eq!(graph.greedy_search(&[1.0], 1, &[1.0], usize::MAX).len(),
1);
+
+ let empty = VamanaGraph::from_adjacency(0, vec![]);
+ assert!(empty.greedy_search(&[], 1, &[1.0], usize::MAX).is_empty());
+ }
+
#[test]
fn vamana_prune_removes_occluded_and_duplicate_neighbors() {
let graph = VamanaGraph::from_adjacency(0, vec![vec![], vec![],
vec![], vec![]]);
@@ -2518,6 +2655,40 @@ mod tests {
assert!(graph.is_fully_reachable());
}
+ #[test]
+ fn vamana_build_rejects_insufficient_memory_budget() {
+ let error = VamanaGraph::build(
+ &[0.0, 1.0, 2.0, 3.0],
+ 4,
+ 1,
+ DiskAnnBuildParams {
+ max_degree: 2,
+ build_search_list_size: 2,
+ memory_budget_bytes: 1,
+ ..DiskAnnBuildParams::default()
+ },
+ )
+ .expect_err("build must enforce its memory budget");
+
+ assert_eq!(error.kind(), io::ErrorKind::OutOfMemory);
+ }
+
+ #[test]
+ fn vamana_sequential_budget_uses_clamped_degree() {
+ let params = DiskAnnBuildParams {
+ max_degree: 1_024,
+ build_search_list_size: 15,
+ memory_budget_bytes: 1024 * 1024,
+ ..DiskAnnBuildParams::default()
+ };
+
+ let graph = VamanaGraph::build_sequential(&[0.0, 1.0, 2.0, 3.0], 4, 1,
params)
+ .expect("small sequential graph must fit in 1 MiB");
+
+ assert!(graph.adjacency.iter().all(|neighbors| neighbors.len() <= 3));
+ assert!(graph.is_fully_reachable());
+ }
+
#[test]
fn vamana_parallel_build_is_deterministic_across_worker_counts() {
let dimension = 4;
diff --git a/core/tests/storage_format_fixtures.rs
b/core/tests/storage_format_fixtures.rs
index 3ce9e2f..df3d0e4 100644
--- a/core/tests/storage_format_fixtures.rs
+++ b/core/tests/storage_format_fixtures.rs
@@ -476,14 +476,10 @@ fn write_diskann_fixture(index: DiskAnnIndex) -> Vec<u8> {
}
fn build_ivf_flat_fixture() -> Vec<u8> {
- let index = IVFFlatIndex {
- d: 2,
- nlist: 2,
- metric: MetricType::L2,
- quantizer_centroids: vec![0.0, 0.0, 10.0, 10.0],
- ids: vec![vec![42, 7], vec![99]],
- vectors: vec![vec![1.0, 0.0, 0.0, 0.0], vec![10.0, 10.0]],
- };
+ let mut index = IVFFlatIndex::new(2, 2, MetricType::L2);
+ index.set_quantizer_centroids(vec![0.0, 0.0, 10.0, 10.0]);
+ index.ids = vec![vec![42, 7], vec![99]];
+ index.vectors = vec![vec![1.0, 0.0, 0.0, 0.0], vec![10.0, 10.0]];
let mut buf = Vec::new();
write_ivfflat_index(&index, &mut PosWriter::new(&mut buf)).unwrap();
buf
@@ -491,7 +487,7 @@ fn build_ivf_flat_fixture() -> Vec<u8> {
fn build_ivf_pq_fixture() -> Vec<u8> {
let mut index = IVFPQIndex::new(1, 2, 1, MetricType::L2, false);
- index.quantizer_centroids = vec![0.0, 10.0];
+ index.set_quantizer_centroids(vec![0.0, 10.0]);
index.pq.centroids = (0..index.pq.ksub).map(|code| code as f32 *
0.25).collect();
index.pq.rebuild_norms_cache();
index.ids = vec![vec![20, 10], vec![30]];
@@ -504,7 +500,7 @@ fn build_ivf_pq_fixture() -> Vec<u8> {
fn build_ivf_pq_4bit_fixture() -> Vec<u8> {
let mut index = IVFPQIndex::with_nbits(2, 2, 2, 4, MetricType::L2, false);
- index.quantizer_centroids = vec![0.0, 0.0, 10.0, 10.0];
+ index.set_quantizer_centroids(vec![0.0, 0.0, 10.0, 10.0]);
index.pq.centroids = (0..index.pq.m)
.flat_map(|_| (0..index.pq.ksub).map(|code| code as f32 * 0.5))
.collect();
@@ -534,7 +530,7 @@ fn build_ivf_rq_fixture() -> Vec<u8> {
fn build_ivf_sq_fixture() -> Vec<u8> {
let sq = ScalarQuantizer::with_dimension_bounds(2, vec![0.0, 0.0],
vec![1.0, 1.0]);
let mut index = IVFSQIndex::new(2, 2, MetricType::L2);
- index.quantizer_centroids = vec![0.0, 0.0, 10.0, 10.0];
+ index.set_quantizer_centroids(vec![0.0, 0.0, 10.0, 10.0]);
index.sq = sq.clone();
index.list_sqs = vec![sq; 2];
index.ids = vec![vec![7], vec![99]];
diff --git a/docs/api.html b/docs/api.html
index 080c0ab..0be8fbb 100644
--- a/docs/api.html
+++ b/docs/api.html
@@ -49,6 +49,7 @@
<h2>Shared lifecycle</h2>
<div class="flow" aria-label="Unified API lifecycle"><div
class="flow-step"><small>01</small><strong>Create a Trainer<br>Parse and
validate options</strong></div><div
class="flow-step"><small>02</small><strong>Submit one or more<br>training
batches</strong></div><div class="flow-step"><small>03</small><strong>Finish
training and<br>create a one-shot Writer</strong></div><div
class="flow-step"><small>04</small><strong>Add row IDs / vectors<br>and write
the file</strong></div><di [...]
<ul><li>Vectors are contiguous <code>f32</code> values; length must
equal <code>vector_count × dimension</code>.</li><li>Training data may arrive
in batches. Every IVF trainer keeps a deterministic reservoir of at most
<code>max(65,536, 64 × resolved nlist)</code> vectors. DiskANN starts from a
50,000-row cap and lowers it when necessary so the retained sample, optional
cosine-normalized copy, codebook, and parallel PQ-training scratch fit
<code>diskann.memory-budget-bytes</cod [...]
+ <div class="callout warning"><strong>IVF coarse assignment is
approximate by default for large centroid matrices</strong>When <code>dimension
× nlist ≥ 1,000,000</code>, <code>ivf.coarse-assignment=auto</code> uses a
Vamana graph while training and adding vectors. Search still selects lists by
exact centroid distance, so graph assignment can lower recall at small
<code>nprobe</code> and does not guarantee that a vector is found by a
self-query with <code>nprobe=1</code>. Set <c [...]
</section>
<section class="article-section" id="params">
@@ -105,6 +106,7 @@ let config = VectorIndexConfig::IvfSq {
dimension: 128,
nlist: 1024,
metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
};
let training = VectorIndexTrainer::train(
@@ -127,17 +129,20 @@ let params = VectorSearchParams::automatic(10)
.with_max_initial_filter_expansion_factor(4);</code></pre></div>
<div class="code-block"><span class="code-label">Rust · other
configurations</span><pre><code>VectorIndexConfig::IvfFlat {
dimension: 128, nlist: 1024, metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
};
VectorIndexConfig::ivf_pq(
128, 1024, MetricType::L2, false,
)?;
VectorIndexConfig::IvfRq {
dimension: 128, nlist: 1024, bits: 4, metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
};
VectorIndexConfig::IvfSq {
dimension: 128, nlist: 1024, metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
};</code></pre></div>
- <p>The IVF-PQ constructor uses the default relative PQ-code budget
and resolves a concrete <code>m</code>. In every option-map API,
<code>pq.m</code> is optional: <code>pq.code-ratio=0.0625</code> is the
default, and an explicit <code>pq.m</code> takes precedence. Metadata and the
on-disk header expose the resolved value.</p>
+ <p>The IVF-PQ constructor uses the default relative PQ-code budget
and resolves a concrete <code>m</code>. In every option-map API,
<code>pq.m</code> is optional: <code>pq.code-ratio=0.0625</code> is the
default, and an explicit <code>pq.m</code> takes precedence. Metadata and the
on-disk header expose the resolved value. Rust callers select the policy
through <code>VectorIndexConfig</code> before training; direct IVF indexes do
not expose a post-training policy switch.</p>
</section>
<section class="article-section" id="c">
diff --git a/docs/ivf-flat.html b/docs/ivf-flat.html
index 6631731..aad9289 100644
--- a/docs/ivf-flat.html
+++ b/docs/ivf-flat.html
@@ -25,7 +25,7 @@
<section class="article-section" id="design">
<h2>Design and data flow</h2>
<h3>Build</h3>
- <div class="pipeline"><div class="pipeline-item"><span
class="pipeline-index">1</span><div><h3>Preprocess training
vectors</h3><p>Cosine mode applies L2 normalization. L2 and inner product keep
the input representation.</p></div></div><div class="pipeline-item"><span
class="pipeline-index">2</span><div><h3>Train coarse centroids</h3><p>K-Means
produces <code>nlist × d</code> <code>f32</code> centroid
values.</p></div></div><div class="pipeline-item"><span class="pipeline-index"
[...]
+ <div class="pipeline"><div class="pipeline-item"><span
class="pipeline-index">1</span><div><h3>Preprocess training
vectors</h3><p>Cosine mode applies L2 normalization. L2 and inner product keep
the input representation.</p></div></div><div class="pipeline-item"><span
class="pipeline-index">2</span><div><h3>Train coarse centroids</h3><p>K-Means
produces <code>nlist × d</code> <code>f32</code> centroid
values.</p></div></div><div class="pipeline-item"><span class="pipeline-index"
[...]
<h3>Search</h3>
<ol><li>Apply the same preprocessing used during
construction.</li><li>Measure the query against every IVF centroid and select
the closest <code>nprobe</code> lists.</li><li>Read those list payloads through
the offset table in bounded concurrent multi-range calls.</li><li>Decode row
IDs and compute the true metric against every raw vector. A Roaring filter
skips disallowed IDs.</li><li>For at least 1,048,576 distance components, scan
independent lists on Rayon workers. Batch se [...]
</section>
@@ -46,6 +46,7 @@ options.put("index.type", "ivf_flat");
options.put("dimension", "128");
options.put("nlist", "1024");
options.put("metric", "l2");
+// options.put("ivf.coarse-assignment", "exact"); // optional; default auto
try (VectorIndexTraining training =
VectorIndexTrainer.train(options, trainingVectors, trainingCount);
@@ -62,13 +63,14 @@ try (VectorIndexTraining training =
dimension: 128,
nlist: 1024,
metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
};
let params = VectorSearchParams::new(10, 16);</code></pre></div>
</section>
<section class="article-section" id="parameters">
<h2>Parameters</h2>
- <div
class="table-wrap"><table><thead><tr><th>Parameter</th><th>Requirement</th><th>Purpose</th><th>Effect
when
increased</th></tr></thead><tbody><tr><td><code>dimension</code></td><td>Inferred
by Java/Python one-shot training; otherwise required and > 0</td><td>Input
dimension</td><td>Linearly increases compute and vector
payload</td></tr><tr><td><code>nlist</code></td><td>Auto from
<code>expected-vector-count</code>, or explicit > 0</td><td>IVF partition
count</td><td>S [...]
+ <div
class="table-wrap"><table><thead><tr><th>Parameter</th><th>Requirement</th><th>Purpose</th><th>Effect
when
increased</th></tr></thead><tbody><tr><td><code>dimension</code></td><td>Inferred
by Java/Python one-shot training; otherwise required and > 0</td><td>Input
dimension</td><td>Linearly increases compute and vector
payload</td></tr><tr><td><code>nlist</code></td><td>Auto from
<code>expected-vector-count</code>, or explicit > 0</td><td>IVF partition
count</td><td>S [...]
</section>
<section class="article-section" id="storage">
diff --git a/docs/ivf-pq.html b/docs/ivf-pq.html
index 00050c2..8ceeb65 100644
--- a/docs/ivf-pq.html
+++ b/docs/ivf-pq.html
@@ -47,6 +47,7 @@ options.put("nlist", "1024");
options.put("metric", "l2");
options.put("pq.code-ratio", "0.0625"); // optional; default 0.0625
options.put("use-opq", "false"); // optional; default false
+// options.put("ivf.coarse-assignment", "exact"); // optional; default auto
try (VectorIndexTraining training =
VectorIndexTrainer.train(options, trainingVectors, trainingCount);
@@ -71,7 +72,7 @@ try (VectorIndexTraining training =
<section class="article-section" id="parameters">
<h2>Parameters</h2>
- <div
class="table-wrap"><table><thead><tr><th>Parameter</th><th>Requirement /
default</th><th>Purpose</th><th>Tuning
meaning</th></tr></thead><tbody><tr><td><code>dimension</code></td><td>Inferred
by Java/Python one-shot training; otherwise > 0</td><td>Input dimension
<code>d</code></td><td>The inferred or explicit <code>pq.m</code> must divide
it</td></tr><tr><td><code>nlist</code></td><td>Auto from
<code>expected-vector-count</code>, or explicit > 0</td><td>IVF partitio
[...]
+ <div
class="table-wrap"><table><thead><tr><th>Parameter</th><th>Requirement /
default</th><th>Purpose</th><th>Tuning
meaning</th></tr></thead><tbody><tr><td><code>dimension</code></td><td>Inferred
by Java/Python one-shot training; otherwise > 0</td><td>Input dimension
<code>d</code></td><td>The inferred or explicit <code>pq.m</code> must divide
it</td></tr><tr><td><code>nlist</code></td><td>Auto from
<code>expected-vector-count</code>, or explicit > 0</td><td>IVF partitio
[...]
</section>
<section class="article-section" id="storage">
diff --git a/docs/ivf-rq.html b/docs/ivf-rq.html
index 4d1cdf6..76424f8 100644
--- a/docs/ivf-rq.html
+++ b/docs/ivf-rq.html
@@ -36,6 +36,7 @@ options.put("dimension", "128");
options.put("nlist", "1024");
options.put("rq.bits", "4"); // optional; 4 is the default
options.put("metric", "l2");
+// options.put("ivf.coarse-assignment", "exact"); // optional; default auto
try (VectorIndexTraining training =
VectorIndexTrainer.train(options, trainingVectors, trainingCount);
@@ -54,13 +55,14 @@ try (VectorIndexReader reader = new
VectorIndexReader(vectorIndexInput)) {
nlist: 1024,
bits: 4,
metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
};
let params = VectorSearchParams::new(10, 64);</code></pre></div>
</section>
<section class="article-section" id="parameters">
<h2>Parameters</h2>
- <div
class="table-wrap"><table><thead><tr><th>Parameter</th><th>Requirement /
default</th><th>Purpose</th><th>Guidance</th></tr></thead><tbody><tr><td><code>dimension</code></td><td>Inferred
by Java/Python one-shot training; otherwise > 0</td><td>Logical vector
dimension</td><td>Storage pads internally to a multiple of
64.</td></tr><tr><td><code>nlist</code></td><td>Auto from
<code>expected-vector-count</code>, or explicit > 0</td><td>IVF partition
count</td><td>Compare t [...]
+ <div
class="table-wrap"><table><thead><tr><th>Parameter</th><th>Requirement /
default</th><th>Purpose</th><th>Guidance</th></tr></thead><tbody><tr><td><code>dimension</code></td><td>Inferred
by Java/Python one-shot training; otherwise > 0</td><td>Logical vector
dimension</td><td>Storage pads internally to a multiple of
64.</td></tr><tr><td><code>nlist</code></td><td>Auto from
<code>expected-vector-count</code>, or explicit > 0</td><td>IVF partition
count</td><td>Compare t [...]
<div class="callout warning"><strong>No query-side bit
width</strong>The Reader always evaluates the representation stored in the
file. Changing <code>rq.bits</code> requires rebuilding the index; this keeps
one file's accuracy and cost contract stable.</div>
</section>
diff --git a/docs/ivf-sq.html b/docs/ivf-sq.html
index e16def2..3cf58e0 100644
--- a/docs/ivf-sq.html
+++ b/docs/ivf-sq.html
@@ -29,17 +29,19 @@
<div class="code-block"><span class="code-label">Options
API</span><pre><code>index.type = ivf_sq
dimension = 128
nlist = 1024
-metric = l2</code></pre></div>
+metric = l2
+ivf.coarse-assignment = auto</code></pre></div>
<div class="code-block"><span
class="code-label">Rust</span><pre><code>let config = VectorIndexConfig::IvfSq {
dimension: 128,
nlist: 1024,
metric: MetricType::L2,
+ use_approximate_coarse_assignment: true,
};
let params = VectorSearchParams::new(10, 16);</code></pre></div>
</section>
<section class="article-section" id="parameters">
<h2>Parameters</h2>
- <div
class="table-wrap"><table><thead><tr><th>Parameter</th><th>Requirement</th><th>Effect</th></tr></thead><tbody><tr><td><code>dimension</code></td><td>Inferred
by Java/Python one-shot training; otherwise > 0</td><td>Each vector uses
exactly <code>d</code> SQ-code
bytes.</td></tr><tr><td><code>nlist</code></td><td>Auto from
<code>expected-vector-count</code>, or explicit > 0 and no larger than
training count</td><td>More lists shorten scans but enlarge centroid and per-
[...]
+ <div
class="table-wrap"><table><thead><tr><th>Parameter</th><th>Requirement</th><th>Effect</th></tr></thead><tbody><tr><td><code>dimension</code></td><td>Inferred
by Java/Python one-shot training; otherwise > 0</td><td>Each vector uses
exactly <code>d</code> SQ-code
bytes.</td></tr><tr><td><code>nlist</code></td><td>Auto from
<code>expected-vector-count</code>, or explicit > 0 and no larger than
training count</td><td>More lists shorten scans but enlarge centroid and per-
[...]
<p>The scalar code width is fixed at 8 bits in v1. There is
deliberately no <code>sq.bits</code>, graph-width, or search-width option.</p>
</section>
<section class="article-section" id="storage">
diff --git a/docs/releases.html b/docs/releases.html
index cf04376..a45d83e 100644
--- a/docs/releases.html
+++ b/docs/releases.html
@@ -49,6 +49,7 @@
<section class="article-section" id="upcoming">
<h2>Upcoming: 0.5.0</h2>
<p>The repository is currently developing the 0.5.0 line. Until an
ASF vote passes and the signed source archive appears under Apache downloads,
code and packages from this line are development artifacts rather than an
Apache release.</p>
+ <div class="callout warning"><strong>Rust IVF API
migration</strong>Version 0.5.0 makes <code>quantizer_centroids</code> private
on <code>IVFFlatIndex</code>, <code>IVFPQIndex</code>, <code>IVFSQIndex</code>,
and <code>IVFRQIndex</code>. Replace direct reads with
<code>quantizer_centroids()</code> and direct assignments with
<code>set_quantizer_centroids(...)</code>. The setter validates the centroid
shape, rejects replacement after vectors are added, and refreshes cached deriv
[...]
</section>
<section class="article-section" id="release-040">
