diff --git a/core/benches/ann_bench.rs b/core/benches/ann_bench.rs index 0d79609..16c15c6 100644 --- a/core/benches/ann_bench.rs +++ b/core/benches/ann_bench.rs @@ -645,6 +645,7 @@ fn index_specs(config: &Config) -> Vec { nlist: config.nlist, metric: MetricType::L2, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }, searches: vec![ivf_search], }, @@ -655,6 +656,7 @@ fn index_specs(config: &Config) -> Vec { nlist: config.nlist, metric: MetricType::L2, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }, searches: vec![ivf_search], }, @@ -668,6 +670,8 @@ fn index_specs(config: &Config) -> Vec { use_opq: false, use_approximate_coarse_assignment: true, canonical_pq_encoding: false, + ivf_train_max_points_per_centroid: 256, + pq_train_max_points_per_centroid: 256, }, searches: vec![ivf_search], }, @@ -679,6 +683,7 @@ fn index_specs(config: &Config) -> Vec { bits: config.rq_bits, metric: MetricType::L2, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }, searches: vec![ivf_search], }, @@ -695,6 +700,7 @@ fn index_specs(config: &Config) -> Vec { raw_vector_encoding: config.diskann_raw_vector_encoding, ..DiskAnnBuildParams::default() }, + pq_train_max_points_per_centroid: 256, }, searches: config .diskann_l_searches diff --git a/core/src/diskann.rs b/core/src/diskann.rs index 30ae782..b668b59 100644 --- a/core/src/diskann.rs +++ b/core/src/diskann.rs @@ -282,6 +282,15 @@ impl DiskAnnIndex { } pub fn train(&mut self, data: &[f32], n: usize) -> io::Result<()> { + self.train_with_config(data, n, &KMeansConfig::default()) + } + + pub fn train_with_config( + &mut self, + data: &[f32], + n: usize, + config: &KMeansConfig, + ) -> io::Result<()> { if n == 0 { return Err(invalid_input( "DiskANN training vector count must be greater than zero", @@ -307,7 +316,7 @@ impl DiskAnnIndex { self.pq.train_hot_start_with_parallelism( &processed, plan.sample_count, - &KMeansConfig::default(), + config, false, plan.parallelism, ); diff --git a/core/src/index.rs b/core/src/index.rs index 9acdee5..6b3e491 100644 --- a/core/src/index.rs +++ b/core/src/index.rs @@ -50,6 +50,7 @@ use crate::ivfsq_io::{ search_batch_ivfsq_reader_filter_range, search_batch_ivfsq_reader_roaring_filter_range, write_ivfsq_index, IVFSQIndexReader, IVF_SQ_MAGIC, }; +use crate::kmeans::KMeansConfig; pub use crate::read_options::{DeploymentProfile, VectorIndexReadPlan, VectorIndexReaderOptions}; use crate::rq::{is_supported_rq_bits, padded_dimension, DEFAULT_RQ_BITS}; use rand::rngs::StdRng; @@ -147,6 +148,7 @@ pub enum VectorIndexConfig { nlist: usize, metric: MetricType, use_approximate_coarse_assignment: bool, + ivf_train_max_points_per_centroid: usize, }, IvfPq { dimension: usize, @@ -158,6 +160,8 @@ pub enum VectorIndexConfig { /// Use canonical expanded-form PQ encoding instead of the default /// transposed direct-L2 encoder. canonical_pq_encoding: bool, + ivf_train_max_points_per_centroid: usize, + pq_train_max_points_per_centroid: usize, }, IvfRq { dimension: usize, @@ -165,12 +169,14 @@ pub enum VectorIndexConfig { bits: usize, metric: MetricType, use_approximate_coarse_assignment: bool, + ivf_train_max_points_per_centroid: usize, }, IvfSq { dimension: usize, nlist: usize, metric: MetricType, use_approximate_coarse_assignment: bool, + ivf_train_max_points_per_centroid: usize, }, DiskAnn { dimension: usize, @@ -178,6 +184,7 @@ pub enum VectorIndexConfig { pq_m: usize, pq_bits: usize, build: DiskAnnBuildParams, + pq_train_max_points_per_centroid: usize, }, } @@ -211,6 +218,8 @@ impl VectorIndexConfig { use_opq, use_approximate_coarse_assignment: true, canonical_pq_encoding: false, + ivf_train_max_points_per_centroid: 256, + pq_train_max_points_per_centroid: 256, }; validate_config(&config)?; Ok(config) @@ -228,6 +237,7 @@ impl VectorIndexConfig { pq_m: infer_pq_m(dimension, pq_bits, DEFAULT_PQ_CODE_RATIO)?, pq_bits, build, + pq_train_max_points_per_centroid: 256, }; validate_config(&config)?; Ok(config) @@ -240,6 +250,7 @@ impl VectorIndexConfig { bits: DEFAULT_RQ_BITS, metric, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }; validate_config(&config)?; Ok(config) @@ -290,6 +301,8 @@ pub struct ResolvedVectorIndexConfig { /// Only meaningful for IVF-PQ. pub canonical_pq_encoding: bool, pub diskann_build: Option, + pub ivf_train_max_points_per_centroid: Option, + pub pq_train_max_points_per_centroid: Option, } impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig { @@ -300,12 +313,14 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig { nlist, metric, use_approximate_coarse_assignment, + ivf_train_max_points_per_centroid, } | VectorIndexConfig::IvfSq { dimension, nlist, metric, use_approximate_coarse_assignment, + ivf_train_max_points_per_centroid, } => Self { index_type: config.index_type(), dimension: *dimension, @@ -318,6 +333,8 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig { use_approximate_coarse_assignment: *use_approximate_coarse_assignment, canonical_pq_encoding: false, diskann_build: None, + ivf_train_max_points_per_centroid: Some(*ivf_train_max_points_per_centroid), + pq_train_max_points_per_centroid: None, }, VectorIndexConfig::IvfPq { dimension, @@ -327,6 +344,8 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig { use_opq, use_approximate_coarse_assignment, canonical_pq_encoding, + ivf_train_max_points_per_centroid, + pq_train_max_points_per_centroid, } => Self { index_type: IndexType::IvfPq, dimension: *dimension, @@ -339,6 +358,8 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig { use_approximate_coarse_assignment: *use_approximate_coarse_assignment, canonical_pq_encoding: *canonical_pq_encoding, diskann_build: None, + ivf_train_max_points_per_centroid: Some(*ivf_train_max_points_per_centroid), + pq_train_max_points_per_centroid: Some(*pq_train_max_points_per_centroid), }, VectorIndexConfig::IvfRq { dimension, @@ -346,6 +367,7 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig { bits, metric, use_approximate_coarse_assignment, + ivf_train_max_points_per_centroid, } => Self { index_type: IndexType::IvfRq, dimension: *dimension, @@ -358,6 +380,8 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig { use_approximate_coarse_assignment: *use_approximate_coarse_assignment, canonical_pq_encoding: false, diskann_build: None, + ivf_train_max_points_per_centroid: Some(*ivf_train_max_points_per_centroid), + pq_train_max_points_per_centroid: None, }, VectorIndexConfig::DiskAnn { dimension, @@ -365,6 +389,7 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig { pq_m, pq_bits, build, + pq_train_max_points_per_centroid, } => Self { index_type: IndexType::DiskAnn, dimension: *dimension, @@ -377,6 +402,8 @@ impl From<&VectorIndexConfig> for ResolvedVectorIndexConfig { use_approximate_coarse_assignment: false, canonical_pq_encoding: false, diskann_build: Some(*build), + ivf_train_max_points_per_centroid: None, + pq_train_max_points_per_centroid: Some(*pq_train_max_points_per_centroid), }, } } @@ -450,6 +477,10 @@ impl VectorIndexBuildPlan { nlist: parse_nlist_options(&mut options, expected_vector_count)?, metric, use_approximate_coarse_assignment, + ivf_train_max_points_per_centroid: parse_training_max_points_per_centroid( + &mut options, + "ivf.train.max-points-per-centroid", + )?, }, IndexType::IvfPq => VectorIndexConfig::IvfPq { dimension, @@ -473,6 +504,14 @@ impl VectorIndexBuildPlan { }, use_approximate_coarse_assignment, canonical_pq_encoding, + ivf_train_max_points_per_centroid: parse_training_max_points_per_centroid( + &mut options, + "ivf.train.max-points-per-centroid", + )?, + pq_train_max_points_per_centroid: parse_training_max_points_per_centroid( + &mut options, + "pq.train.max-points-per-centroid", + )?, }, IndexType::IvfRq => { let explicit_bits = options @@ -494,6 +533,10 @@ impl VectorIndexBuildPlan { bits, metric, use_approximate_coarse_assignment, + ivf_train_max_points_per_centroid: parse_training_max_points_per_centroid( + &mut options, + "ivf.train.max-points-per-centroid", + )?, } } IndexType::IvfSq => VectorIndexConfig::IvfSq { @@ -501,6 +544,10 @@ impl VectorIndexBuildPlan { nlist: parse_nlist_options(&mut options, expected_vector_count)?, metric, use_approximate_coarse_assignment, + ivf_train_max_points_per_centroid: parse_training_max_points_per_centroid( + &mut options, + "ivf.train.max-points-per-centroid", + )?, }, IndexType::DiskAnn => { let pq_bits = match options.optional("pq.bits") { @@ -547,6 +594,10 @@ impl VectorIndexBuildPlan { )?, pq_bits, build, + pq_train_max_points_per_centroid: parse_training_max_points_per_centroid( + &mut options, + "pq.train.max-points-per-centroid", + )?, } } }; @@ -620,6 +671,17 @@ impl ConfigOptions { } } +fn parse_training_max_points_per_centroid( + options: &mut ConfigOptions, + key: &str, +) -> io::Result { + options + .optional(key) + .map(|value| parse_usize_option(key, &value)) + .transpose() + .map(|value| value.unwrap_or(KMeansConfig::default().max_points_per_centroid)) +} + fn parse_nlist_options( options: &mut ConfigOptions, expected_vector_count: Option, @@ -1224,6 +1286,8 @@ pub struct DiskAnnMetadata { pub struct VectorIndexTrainer { writer: VectorIndexWriter, + ivf_training: KMeansConfig, + pq_training: KMeansConfig, training_data: Vec, training_vector_count: usize, training_vectors_seen: usize, @@ -1233,6 +1297,15 @@ pub struct VectorIndexTrainer { impl VectorIndexTrainer { pub fn new(config: VectorIndexConfig) -> io::Result { + let resolved = config.resolved(); + let mut ivf_training = KMeansConfig::default(); + let mut pq_training = KMeansConfig::default(); + if let Some(max_points) = resolved.ivf_train_max_points_per_centroid { + ivf_training.max_points_per_centroid = max_points; + } + if let Some(max_points) = resolved.pq_train_max_points_per_centroid { + pq_training.max_points_per_centroid = max_points; + } let training_sample_limit = match &config { VectorIndexConfig::DiskAnn { dimension, @@ -1240,6 +1313,7 @@ impl VectorIndexTrainer { pq_m, pq_bits, build, + .. } => diskann_training_sample_limit( *dimension, *metric, @@ -1256,6 +1330,8 @@ impl VectorIndexTrainer { let writer = VectorIndexWriter::from_config(config)?; Ok(Self { writer, + ivf_training, + pq_training, training_data: Vec::new(), training_vector_count: 0, training_vectors_seen: 0, @@ -1311,8 +1387,12 @@ impl VectorIndexTrainer { if self.training_vector_count == 0 || self.training_data.is_empty() { return Err(invalid_input("no training vectors added")); } - self.writer - .train_internal(&self.training_data, self.training_vector_count)?; + self.writer.train_internal( + &self.training_data, + self.training_vector_count, + &self.ivf_training, + &self.pq_training, + )?; Ok(VectorIndexTraining { inner: self.writer }) } } @@ -1352,6 +1432,7 @@ impl VectorIndexWriter { nlist, metric, use_approximate_coarse_assignment, + .. } => { let mut index = IVFFlatIndex::new(dimension, nlist, metric); index.set_approximate_coarse_assignment(use_approximate_coarse_assignment); @@ -1362,6 +1443,7 @@ impl VectorIndexWriter { nlist, metric, use_approximate_coarse_assignment, + .. } => { let mut index = IVFSQIndex::new(dimension, nlist, metric); index.set_approximate_coarse_assignment(use_approximate_coarse_assignment); @@ -1375,6 +1457,7 @@ impl VectorIndexWriter { use_opq, use_approximate_coarse_assignment, canonical_pq_encoding, + .. } => { let mut index = IVFPQIndex::new(dimension, nlist, m, metric, use_opq); index.set_approximate_coarse_assignment(use_approximate_coarse_assignment); @@ -1387,6 +1470,7 @@ impl VectorIndexWriter { 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); @@ -1398,6 +1482,7 @@ impl VectorIndexWriter { pq_m, pq_bits, build, + .. } => Self::DiskAnn(DiskAnnIndex::with_pq_bits( dimension, metric, pq_m, pq_bits, build, )), @@ -1424,14 +1509,20 @@ impl VectorIndexWriter { } } - fn train_internal(&mut self, data: &[f32], n: usize) -> io::Result<()> { + fn train_internal( + &mut self, + data: &[f32], + n: usize, + ivf_training: &KMeansConfig, + pq_training: &KMeansConfig, + ) -> io::Result<()> { debug_assert_eq!(Some(data.len()), n.checked_mul(self.dimension())); match self { - Self::IvfFlat(index) => index.train(data, n), - Self::IvfSq(index) => index.train(data, n), - Self::IvfPq(index) => index.train(data, n), - Self::IvfRq(index) => index.train(data, n), - Self::DiskAnn(index) => return index.train(data, n), + Self::IvfFlat(index) => index.train_with_config(data, n, ivf_training), + Self::IvfSq(index) => index.train_with_config(data, n, ivf_training), + Self::IvfPq(index) => index.train_with_config(data, n, ivf_training, pq_training), + Self::IvfRq(index) => index.train_with_config(data, n, ivf_training), + Self::DiskAnn(index) => return index.train_with_config(data, n, pq_training), } Ok(()) } @@ -2375,11 +2466,32 @@ fn validate_config(config: &VectorIndexConfig) -> io::Result<()> { pq_m, pq_bits, build, + .. } => { validate_diskann_config(*dimension, *metric, *pq_m, *pq_bits, *build)?; } _ => {} } + let resolved = config.resolved(); + for (key, max_points, centroids) in [ + ( + "ivf.train.max-points-per-centroid", + resolved.ivf_train_max_points_per_centroid, + config.nlist(), + ), + ( + "pq.train.max-points-per-centroid", + resolved.pq_train_max_points_per_centroid, + 1usize << resolved.pq_bits.unwrap_or(8), + ), + ] { + if let Some(max_points) = max_points { + validate_positive(max_points, key)?; + centroids.checked_mul(max_points).ok_or_else(|| { + invalid_input(format!("{key} times centroid count overflows usize")) + })?; + } + } Ok(()) } @@ -2691,6 +2803,7 @@ mod tests { nlist: 1, metric: MetricType::L2, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }, &[0.0, 1.0], 2, @@ -2720,6 +2833,7 @@ mod tests { build_search_list_size: 16, ..DiskAnnBuildParams::default() }, + pq_train_max_points_per_centroid: 256, }); reader @@ -2744,6 +2858,7 @@ mod tests { metric: MetricType::L2, bits: 4, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }); reader @@ -2768,6 +2883,7 @@ mod tests { metric: MetricType::L2, bits: 4, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }); let queries = [0, nlist - 1] .into_iter() @@ -2813,6 +2929,7 @@ mod tests { build_search_list_size: 16, ..DiskAnnBuildParams::default() }, + pq_train_max_points_per_centroid: 256, }, &data, count, @@ -2922,6 +3039,7 @@ mod tests { nlist: 4, metric: MetricType::L2, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }); roundtrip(VectorIndexConfig::ivf_pq(16, 4, MetricType::L2, false).unwrap()); roundtrip(VectorIndexConfig::IvfRq { @@ -2930,12 +3048,14 @@ mod tests { bits: DEFAULT_RQ_BITS, metric: MetricType::L2, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }); roundtrip(VectorIndexConfig::IvfSq { dimension: 8, nlist: 4, metric: MetricType::L2, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }); roundtrip( VectorIndexConfig::disk_ann( @@ -2975,6 +3095,7 @@ mod tests { nlist: 4, metric: MetricType::L2, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }, VectorIndexConfig::IvfPq { dimension: 16, @@ -2984,6 +3105,8 @@ mod tests { use_opq: false, use_approximate_coarse_assignment: true, canonical_pq_encoding: false, + ivf_train_max_points_per_centroid: 256, + pq_train_max_points_per_centroid: 256, }, VectorIndexConfig::IvfRq { dimension: 8, @@ -2991,12 +3114,14 @@ mod tests { bits: DEFAULT_RQ_BITS, metric: MetricType::L2, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }, VectorIndexConfig::IvfSq { dimension: 8, nlist: 4, metric: MetricType::L2, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }, ] { let d = config.dimension(); @@ -3085,6 +3210,8 @@ mod tests { use_opq: false, use_approximate_coarse_assignment: true, canonical_pq_encoding: false, + ivf_train_max_points_per_centroid: 256, + pq_train_max_points_per_centroid: 256, }) { Ok(_) => panic!("invalid PQ config should be rejected"), Err(err) => err, @@ -3100,6 +3227,7 @@ mod tests { bits: DEFAULT_RQ_BITS, metric: MetricType::L2, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }) .unwrap(); let err = match VectorIndexTrainer::new(VectorIndexConfig::IvfRq { @@ -3108,6 +3236,7 @@ mod tests { bits: 9, metric: MetricType::L2, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }) { Ok(_) => panic!("invalid RQ config should be rejected"), Err(err) => err, @@ -3824,6 +3953,7 @@ mod tests { build_distance: DiskAnnBuildDistance::ProductQuantized, ..DiskAnnBuildParams::default() }, + pq_train_max_points_per_centroid: 256, }; let mut writer = build_writer(config, data, count); writer.add_vectors(ids, data, count).unwrap(); @@ -3929,6 +4059,7 @@ mod tests { raw_vector_encoding: DiskAnnRawVectorEncoding::F32, ..DiskAnnBuildParams::default() }, + pq_train_max_points_per_centroid: 256, }, &data, count, @@ -4053,6 +4184,7 @@ mod tests { pq_m: 2, pq_bits: 8, build: DiskAnnBuildParams::default(), + pq_train_max_points_per_centroid: 256, }) .expect("DiskANN trainer should open"); @@ -4075,6 +4207,7 @@ mod tests { seed: 73, ..DiskAnnBuildParams::default() }, + pq_train_max_points_per_centroid: 256, }; let mut whole = VectorIndexTrainer::new(config()).unwrap(); @@ -4113,6 +4246,7 @@ mod tests { memory_budget_bytes, ..DiskAnnBuildParams::default() }, + pq_train_max_points_per_centroid: 256, }; let trainer = VectorIndexTrainer::new(config).unwrap(); assert!(trainer.training_sample_limit < DISKANN_MAX_PQ_TRAINING_VECTORS); @@ -4135,6 +4269,7 @@ mod tests { pq_m: 2, pq_bits: 8, build: DiskAnnBuildParams::default(), + pq_train_max_points_per_centroid: 256, }, &data, count, @@ -4157,6 +4292,7 @@ mod tests { pq_m: 2, pq_bits: 8, build: DiskAnnBuildParams::default(), + pq_train_max_points_per_centroid: 256, }, &training_data, training_count, @@ -4276,6 +4412,7 @@ mod tests { nlist: 1, metric: MetricType::L2, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }, &[value, 1.0], 2, @@ -4396,6 +4533,105 @@ mod tests { assert!(error.to_string().contains("only valid for IVF-PQ")); } + #[test] + fn training_max_points_per_centroid_defaults_and_validation() { + let base = options(&[ + ("index.type", "ivf_pq"), + ("dimension", "4"), + ("nlist", "2"), + ("metric", "l2"), + ]); + let resolved = VectorIndexConfig::from_options(&base).unwrap().resolved(); + assert_eq!(resolved.ivf_train_max_points_per_centroid, Some(256)); + assert_eq!(resolved.pq_train_max_points_per_centroid, Some(256)); + + for key in [ + "ivf.train.max-points-per-centroid", + "pq.train.max-points-per-centroid", + ] { + for value in ["0", "-1", "1.5", "abc", "", &usize::MAX.to_string()] { + let mut opts = base.clone(); + opts.insert(key.into(), value.into()); + let error = VectorIndexConfig::from_options(&opts).unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::InvalidInput); + assert!(error.to_string().contains(key), "{error}"); + } + } + for (index_type, key) in [ + ("ivf_flat", "pq.train.max-points-per-centroid"), + ("ivf_sq", "pq.train.max-points-per-centroid"), + ("ivf_rq", "pq.train.max-points-per-centroid"), + ("diskann", "ivf.train.max-points-per-centroid"), + ("ivf_pq", "fields.vector.ivf.train.max-points-per-centroid"), + ("ivf_pq", "fields.vector.pq.train.max-points-per-centroid"), + ] { + let mut opts = base.clone(); + opts.insert("index.type".into(), index_type.into()); + if index_type == "diskann" { + opts.remove("nlist"); + } + opts.insert(key.into(), "32".into()); + let error = VectorIndexConfig::from_options(&opts).unwrap_err(); + assert!(error.to_string().contains("unknown vector index option")); + assert!(error.to_string().contains(key)); + } + } + + #[test] + fn training_max_points_per_centroid_controls_ivf_and_pq() { + use crate::kmeans::{kmeans_train, KMeansConfig}; + use crate::pq::ProductQuantizer; + + let n = 600; + let d = 4; + let data = (0..n * d) + .map(|i| ((i * 37 % 997) as f32).sin()) + .collect::>(); + let ivf_config = KMeansConfig { + max_points_per_centroid: 1, + ..KMeansConfig::default() + }; + let pq_config = KMeansConfig { + max_points_per_centroid: 2, + ..KMeansConfig::default() + }; + let expected_centroids = kmeans_train(&ivf_config, &data, n, d, 2); + let mut expected_pq = ProductQuantizer::new(d, 1); + expected_pq.train_with_config(&data, n, &pq_config); + + for index_type in ["ivf_flat", "ivf_sq", "ivf_rq", "ivf_pq", "diskann"] { + let mut opts = options(&[ + ("index.type", index_type), + ("dimension", "4"), + ("metric", "inner_product"), + ]); + if index_type != "diskann" { + opts.insert("nlist".into(), "2".into()); + opts.insert("ivf.train.max-points-per-centroid".into(), "1".into()); + } + if matches!(index_type, "ivf_pq" | "diskann") { + opts.insert("pq.m".into(), "1".into()); + opts.insert("pq.train.max-points-per-centroid".into(), "2".into()); + } + let config = VectorIndexConfig::from_options(&opts).unwrap(); + let training = VectorIndexTrainer::train(config, &data, n).unwrap(); + let centroids = match VectorIndexWriter::new(training) { + VectorIndexWriter::IvfFlat(index) => index.quantizer_centroids().to_vec(), + VectorIndexWriter::IvfSq(index) => index.quantizer_centroids().to_vec(), + VectorIndexWriter::IvfRq(index) => index.quantizer_centroids().to_vec(), + VectorIndexWriter::IvfPq(index) => { + assert_eq!(index.pq.centroids(), expected_pq.centroids()); + index.quantizer_centroids().to_vec() + } + VectorIndexWriter::DiskAnn(index) => { + assert_eq!(index.pq.centroids(), expected_pq.centroids()); + continue; + } + }; + assert_eq!(centroids, expected_centroids, "{index_type}"); + } + } + #[test] fn config_from_options_rejects_unknown_options() { let err = VectorIndexConfig::from_options(&options(&[ @@ -4470,6 +4706,7 @@ mod tests { nlist: 1, metric: MetricType::L2, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }, &[0.0, 1.0], 2, diff --git a/core/src/ivfflat.rs b/core/src/ivfflat.rs index ab1d622..67251d2 100644 --- a/core/src/ivfflat.rs +++ b/core/src/ivfflat.rs @@ -73,9 +73,12 @@ impl IVFFlatIndex { } pub fn train(&mut self, data: &[f32], n: usize) { + self.train_with_config(data, n, &KMeansConfig::default()) + } + + pub fn train_with_config(&mut self, data: &[f32], n: usize, config: &KMeansConfig) { let train_data = self.preprocess_vectors(data, n); - self.quantizer_centroids = - kmeans::kmeans_train(&KMeansConfig::default(), &train_data, n, self.d, self.nlist); + self.quantizer_centroids = kmeans::kmeans_train(config, &train_data, n, self.d, self.nlist); self.coarse_assignment.reset(); } diff --git a/core/src/ivfpq.rs b/core/src/ivfpq.rs index a7b1082..60644e8 100644 --- a/core/src/ivfpq.rs +++ b/core/src/ivfpq.rs @@ -209,6 +209,16 @@ impl IVFPQIndex { } pub fn train(&mut self, data: &[f32], n: usize) { + self.train_with_config(data, n, &KMeansConfig::default(), &KMeansConfig::default()); + } + + pub fn train_with_config( + &mut self, + data: &[f32], + n: usize, + ivf_config: &KMeansConfig, + pq_config: &KMeansConfig, + ) { let d = self.d; let train_data = if self.metric == MetricType::Cosine { @@ -225,7 +235,7 @@ impl IVFPQIndex { // IVF centroids must be trained on projected (rotated) data since // add() and search() assign rotated vectors via preprocess_queries(). let effective_data = if let Some(ref mut opq) = self.opq { - opq.train(&train_data, n, &mut self.pq); + opq.train_with_config(&train_data, n, &mut self.pq, pq_config); let mut projected = vec![0.0f32; n * d]; opq.apply_batch(&train_data, &mut projected, n); projected @@ -233,9 +243,8 @@ impl IVFPQIndex { train_data }; - let km_config = KMeansConfig::default(); self.quantizer_centroids = - kmeans::kmeans_train(&km_config, &effective_data, n, d, self.nlist); + kmeans::kmeans_train(ivf_config, &effective_data, n, d, self.nlist); self.coarse_assignment.reset(); // Retrain PQ on the same assignment distribution that add/search will encode. @@ -259,7 +268,7 @@ impl IVFPQIndex { } else { effective_data }; - self.pq.train(&pq_train_data, n); + self.pq.train_with_config(&pq_train_data, n, pq_config); } /// Add vectors in batches (Faiss-style: batch assign → batch residual → batch encode). diff --git a/core/src/ivfrq.rs b/core/src/ivfrq.rs index 86596ae..036caea 100644 --- a/core/src/ivfrq.rs +++ b/core/src/ivfrq.rs @@ -152,6 +152,10 @@ impl IVFRQIndex { } pub fn train(&mut self, data: &[f32], n: usize) { + self.train_with_config(data, n, &KMeansConfig::default()) + } + + pub fn train_with_config(&mut self, data: &[f32], n: usize, config: &KMeansConfig) { let timing = build_timing_enabled(); let total_started = Instant::now(); let phase_started = Instant::now(); @@ -159,8 +163,7 @@ impl IVFRQIndex { 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_centroids = kmeans::kmeans_train(config, &processed, n, self.d, self.nlist); log_build_timing(timing, "train.kmeans", phase_started); let phase_started = Instant::now(); diff --git a/core/src/ivfsq.rs b/core/src/ivfsq.rs index af69ed9..58e68be 100644 --- a/core/src/ivfsq.rs +++ b/core/src/ivfsq.rs @@ -82,9 +82,12 @@ impl IVFSQIndex { } pub fn train(&mut self, data: &[f32], n: usize) { + self.train_with_config(data, n, &KMeansConfig::default()) + } + + pub fn train_with_config(&mut self, data: &[f32], n: usize, config: &KMeansConfig) { let processed = self.preprocess_vectors(data, n); - self.quantizer_centroids = - kmeans::kmeans_train(&KMeansConfig::default(), &processed, n, self.d, self.nlist); + self.quantizer_centroids = kmeans::kmeans_train(config, &processed, n, self.d, self.nlist); self.coarse_assignment.reset(); let list_ids = self.coarse_assignment.assign( &processed, diff --git a/core/src/kmeans.rs b/core/src/kmeans.rs index 55fd814..0eac266 100644 --- a/core/src/kmeans.rs +++ b/core/src/kmeans.rs @@ -150,6 +150,7 @@ fn kmeans_train_hierarchical( let initial_config = KMeansConfig { niter: config.niter, seed: config.seed, + max_points_per_centroid: config.max_points_per_centroid, ..KMeansConfig::default() }; let initial_centroids = @@ -200,6 +201,7 @@ fn kmeans_train_hierarchical( let sub_config = KMeansConfig { niter: 10, seed: config.seed + finalized.len() as u64, + max_points_per_centroid: config.max_points_per_centroid, ..KMeansConfig::default() }; let sub_centroids = kmeans_train_with_init(&sub_config, &sub_data, sub_n, d, split_k, None); diff --git a/core/src/opq.rs b/core/src/opq.rs index 0ebf46b..29acd57 100644 --- a/core/src/opq.rs +++ b/core/src/opq.rs @@ -57,6 +57,16 @@ impl OPQMatrix { /// Train the OPQ rotation matrix. /// data: flat [n * d]. pub fn train(&mut self, data: &[f32], n: usize, pq: &mut ProductQuantizer) { + self.train_with_config(data, n, pq, &KMeansConfig::default()); + } + + pub fn train_with_config( + &mut self, + data: &[f32], + n: usize, + pq: &mut ProductQuantizer, + config: &KMeansConfig, + ) { let d = self.d; let mut rng = StdRng::seed_from_u64(12345); @@ -122,7 +132,7 @@ impl OPQMatrix { }; let km_config = KMeansConfig { niter: pq_niter, - ..KMeansConfig::default() + ..*config }; let hot_start = iter > 0; pq.train_hot_start(&projected, train_n, &km_config, hot_start); @@ -156,7 +166,7 @@ impl OPQMatrix { // Final PQ training with the learned rotation self.apply_batch(&train_data, &mut projected, train_n); - pq.train_with_config(&projected, train_n, &KMeansConfig::default()); + pq.train_with_config(&projected, train_n, config); self.is_trained = true; } @@ -196,6 +206,39 @@ impl OPQMatrix { mod tests { use super::*; + #[test] + fn training_max_points_per_centroid_applies_to_opq_iterations_and_final_pq() { + let d = 4; + let n = 64; + let mut rng = StdRng::seed_from_u64(42); + let mut data = Vec::new(); + for _ in 0..n / 2 { + let vector = (0..d).map(|_| rng.gen::()).collect::>(); + data.extend_from_slice(&vector); + data.extend(vector.iter().map(|value| -value)); + } + let config = KMeansConfig { + max_points_per_centroid: 1, + ..KMeansConfig::default() + }; + let mut opq = OPQMatrix::new(d, 2); + opq.niter = 2; + let mut pq = ProductQuantizer::with_nbits(d, 2, 4); + opq.train_with_config(&data, n, &mut pq, &config); + + // Paired vectors have zero mean, so OPQ centering leaves this data unchanged. + let mut projected = vec![0.0; n * d]; + opq.apply_batch(&data, &mut projected, n); + let mut expected = ProductQuantizer::with_nbits(d, 2, 4); + expected.train_with_config(&projected, n, &config); + assert_eq!(pq.centroids(), expected.centroids()); + + let mut default_opq = OPQMatrix::new(d, 2); + default_opq.niter = 2; + default_opq.train(&data, n, &mut expected); + assert_ne!(opq.rotation, default_opq.rotation); + } + #[test] fn test_rotation_orthogonality() { let d = 8; diff --git a/docs/api.html b/docs/api.html index 79bedf8..3d3bf35 100644 --- a/docs/api.html +++ b/docs/api.html @@ -49,6 +49,12 @@

Public integration layers

Shared lifecycle

01Create a Trainer
Parse and validate options
02Submit one or more
training batches
03Finish training and
create a one-shot Writer
04Add row IDs / vectors
and write the file
05Detect file magic
and execute searches
  • Vectors are contiguous f32 values; length must equal vector_count × dimension.
  • Training data may arrive in batches. Every IVF trainer keeps a deterministic reservoir of at most max(65,536, 64 × resolved nlist) 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 diskann.memory-budget-bytes. Sampling is independent of batch boundaries.
  • The Python and Java one-shot train helpers infer dimension from the matrix and use its row count for automatic nlist. When the matrix is only a sample, pass the final corpus size as expected-vector-count. Streaming Trainer APIs require a concrete dimension before their first batch.
  • A Writer may receive production vectors in multiple batches. Row-ID count must equal vector count.
  • Readers expose metadata, single-query search, batch search, and Roaring64-filtered variants.
  • Files carry their type and resolved model sections. Callers do not pass index options again when opening a Reader.
+

Training options are parsed by the shared Rust core in every option-map API. Both limits default to 256 and must be positive integers.

+
+ + +
OptionApplies toTraining limit
ivf.train.max-points-per-centroidIVF-FLAT, IVF-SQ, IVF-RQ, IVF-PQAt most nlist × value vectors for coarse K-means; also applied to hierarchical clustering stages.
pq.train.max-points-per-centroidIVF-PQ, DiskANNAt most 2^pq_bits × value vectors per PQ subquantizer, including PQ training inside OPQ. IVF-PQ uses 8-bit codebooks.
+

These limits apply within the Trainer reservoir described above; increasing them does not increase that reservoir. OPQ also retains its 65,536-row input cap, and DiskANN retains its memory budget. The limits affect training only and are not stored in the index file. Pass the bare keys to this library, without a fields.<field-name>. prefix. Paimon integrations must also allow these keys through their option filter.

IVF coarse assignment is approximate by default for large centroid matricesWhen dimension × nlist ≥ 1,000,000, ivf.coarse-assignment=auto 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 nprobe and does not guarantee that a vector is found by a self-query with nprobe=1. Set ivf.coarse-assignment=exact to disable the graph, preserve exact nearest-centroid assignment, and avoid graph startup cost for small non-empty batches. Empty batches never build the graph.
IVF-PQ encoding defaults to autoFor 8-bit PQ (ksub=256) where every subvector has at least four dimensions, ivf.pq-encoding=auto uses the transposed direct-L2 encoder on x86 with AVX2+FMA and on AArch64. Other CPUs use blocked SGEMM for finite codebooks. Unsupported shapes, and non-finite codebooks on the SGEMM fallback, use the canonical encoder. Codes can differ across these backends; NaN and high-dynamic-range inputs also follow the selected backend's arithmetic. Set ivf.pq-encoding=canonical to reproduce ProductQuantizer::encode_batch on the same CPU and runtime backend; it is substantially slower. Neither mode promises byte-identical codes across CPU feature sets. The choice affects builds only, is not stored, and does not change the index format or search path.
@@ -109,6 +115,7 @@

Rust

nlist: 1024, metric: MetricType::L2, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }; let training = VectorIndexTrainer::train( @@ -132,6 +139,7 @@

Rust

Rust · other configurations
VectorIndexConfig::IvfFlat {
     dimension: 128, nlist: 1024, metric: MetricType::L2,
     use_approximate_coarse_assignment: true,
+    ivf_train_max_points_per_centroid: 256,
 };
 VectorIndexConfig::ivf_pq(
     128, 1024, MetricType::L2, false,
@@ -139,10 +147,12 @@ 
Rust

 VectorIndexConfig::IvfRq {
     dimension: 128, nlist: 1024, bits: 4, metric: MetricType::L2,
     use_approximate_coarse_assignment: true,
+    ivf_train_max_points_per_centroid: 256,
 };
 VectorIndexConfig::IvfSq {
     dimension: 128, nlist: 1024, metric: MetricType::L2,
     use_approximate_coarse_assignment: true,
+    ivf_train_max_points_per_centroid: 256,
 };

The IVF-PQ constructor uses the default relative PQ-code budget and resolves a concrete m. In every option-map API, pq.m is optional: pq.code-ratio=0.0625 is the default, and an explicit pq.m takes precedence. Metadata and the on-disk header expose the resolved value. Rust callers select the policy through VectorIndexConfig before training; direct IVF indexes do not expose a post-training policy switch.

diff --git a/docs/diskann.html b/docs/diskann.html index 168973c..1a9ed5f 100644 --- a/docs/diskann.html +++ b/docs/diskann.html @@ -162,6 +162,7 @@

Build parameters

pq.code-ratio0.0625; finite and in (0, 0.25] for 8-bit or (0, 0.125] for 4-bitTarget ratio between resident PQ-code bytes and raw f32-vector bytes. The builder selects the nearest m and distributes dimensions across balanced chunks.Usually reduces PQ error when increased, but grows resident memory and per-candidate lookup work. pq.mOptional expert override; 1..=dimensionConcrete PQ chunk count. Explicit values take precedence over pq.code-ratio; exact chunk offsets are persisted in the self-describing codebook.Use only for a measured override of automatic sizing. Non-divisible dimensions and odd 4-bit values are valid. pq.bits8; must be 4 or 8Centroids and stored bits per PQ chunk. Four-bit codes pack two chunks per byte, use 16-entry query tables, and require a zero high padding nibble when m is odd.Eight bits generally improve graph-navigation recall; four bits reduce codebook, resident codes, training work, and lookup-table size. Rebuild and benchmark both. + pq.train.max-points-per-centroidPositive integer; default 256Maximum training vectors per PQ centroidCaps input at 2^pq.bits × value vectors per subquantizer; the existing 50,000-vector cap and memory budget still apply. diskann.max-degree64; 1..=1023, preserving page-contained raw fallbackMaximum graph out-degree R.May improve connectivity and recall; increases graph bytes, build work, and page density cost. diskann.build-search-list-sizeOmitted: max(100, R); explicit values must be ≥ RCandidate width Lbuild during Vamana construction.Usually improves graph quality while increasing build CPU and per-worker scratch. diskann.alpha1.2; finite and ≥ 1Second-pass robust-prune threshold.Higher values prune candidates less aggressively; validate degree, recall, and graph behavior empirically. @@ -171,6 +172,7 @@

Build parameters

diskann.raw-vector-encodingauto or omitted; preset/budget resolves F32 or F16Controls the persisted rerank-vector element width for both layouts. Compact F32/F16 payloads are exactly 4 × d × N / 2 × d × N bytes with no per-page padding.Explicit F32 preserves original rerank distances. Explicit F16 halves raw-vector I/O but must be recall-tested. diskann.build-distanceauto or omitted; preset resolves PQ or full precisionSelects build-traversal distance. Both modes use full precision for robust pruning and connectivity repair.high_recall uses full precision; balanced/fast presets use PQ guidance. +

Larger training samples increase training work and may improve codebook quality. This option does not raise the existing training sample or memory limits. See the shared training options for sampling and Paimon integration details.

Starting values

  • Start with diskann.build-preset=balanced, the automatic pq.code-ratio=0.0625, and the intended deployment-profile. The balanced preset resolves to 8-bit PQ, R=64, Lbuild=100, alpha 1.2, F16, and PQ-guided construction unless an explicit option changes the representation.
    • diff --git a/docs/ivf-flat.html b/docs/ivf-flat.html index 5dc76b1..d225574 100644 --- a/docs/ivf-flat.html +++ b/docs/ivf-flat.html @@ -64,13 +64,15 @@

    Usage

    nlist: 1024, metric: MetricType::L2, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }; let params = VectorSearchParams::new(10, 16);

    Parameters

    -
    ParameterRequirementPurposeEffect when increased
    dimensionInferred by Java/Python one-shot training; otherwise required and > 0Input dimensionLinearly increases compute and vector payload
    nlistAuto from expected-vector-count, or explicit > 0IVF partition countShorter average lists and a larger centroid table; automatic nprobe follows the resolved value
    metricRequired: l2, inner_product, or cosineTraining, assignment, and search distanceSemantic, not inferred; it must match ground truth
    ivf.coarse-assignmentauto by default; optional exactControls build-time list assignmentauto uses Vamana when dimension × nlist ≥ 1,000,000, trading build speed for possible low-nprobe recall loss and graph startup cost; exact disables it
    top_kQuery-time, > 0Requested resultsIncreases heap and output work
    nprobeAutomatic by default; explicit 1 to nlistLists to probeAuto accounts for K, average list size, and filter selectivity; explicit values remain available for measured overrides
    +
    ParameterRequirementPurposeEffect when increased
    dimensionInferred by Java/Python one-shot training; otherwise required and > 0Input dimensionLinearly increases compute and vector payload
    nlistAuto from expected-vector-count, or explicit > 0IVF partition countShorter average lists and a larger centroid table; automatic nprobe follows the resolved value
    metricRequired: l2, inner_product, or cosineTraining, assignment, and search distanceSemantic, not inferred; it must match ground truth
    ivf.train.max-points-per-centroidPositive integer; default 256Maximum training vectors per coarse centroidCaps coarse K-means input at nlist × value vectors, including hierarchical clustering stages.
    ivf.coarse-assignmentauto by default; optional exactControls build-time list assignmentauto uses Vamana when dimension × nlist ≥ 1,000,000, trading build speed for possible low-nprobe recall loss and graph startup cost; exact disables it
    top_kQuery-time, > 0Requested resultsIncreases heap and output work
    nprobeAutomatic by default; explicit 1 to nlistLists to probeAuto accounts for K, average list size, and filter selectivity; explicit values remain available for measured overrides
    +

    Larger training samples increase training work and may improve centroid quality. These limits apply within the Trainer reservoir of max(65536, 64 × nlist) vectors; increasing them does not enlarge that reservoir. See the shared training options for sampling and Paimon integration details.

    diff --git a/docs/ivf-pq.html b/docs/ivf-pq.html index d39de04..6aa3c22 100644 --- a/docs/ivf-pq.html +++ b/docs/ivf-pq.html @@ -73,7 +73,8 @@

    Usage

    Parameters

    -
    ParameterRequirement / defaultPurposeTuning meaning
    dimensionInferred by Java/Python one-shot training; otherwise > 0Input dimension dThe inferred or explicit pq.m must divide it
    nlistAuto from expected-vector-count, or explicit > 0IVF partition countControls list length, centroid cost, and automatic nprobe
    pq.code-ratioDefault 0.0625; finite and in (0, 0.25]Target ratio between PQ-code bytes and raw f32-vector bytesThe closest valid divisor is selected; max-bytes-per-vector can supply the target instead
    pq.mOptional expert override; > 0; d % m == 0Concrete subspace count and 8-bit code bytesTakes precedence over automatic sizing; use after representative recall measurements
    use-opqAuto enables when target-recall ≥ 0.9; explicit true/false winsLearn and apply an orthogonal rotationTrades training and query work for possible recall improvement
    ivf.coarse-assignmentauto by default; optional exactControls build-time list assignmentauto uses Vamana when dimension × nlist ≥ 1,000,000, trading build speed for possible low-nprobe recall loss and graph startup cost; exact disables it
    ivf.pq-encodingauto by default; optional canonicalControls build-time PQ encodingFor 8-bit PQ with every dsub >= 4, auto uses transposed direct-L2 on x86 with AVX2+FMA and on AArch64, and blocked SGEMM on other CPUs when the codebook is finite; unsupported shapes and non-finite codebooks on the SGEMM fallback use canonical. canonical reproduces ProductQuantizer::encode_batch on the same CPU and runtime backend at substantially higher build cost
    nprobeAutomatic by default; explicit expert overrideLists readAuto accounts for K, average list size, and filter selectivity
    +
    ParameterRequirement / defaultPurposeTuning meaning
    dimensionInferred by Java/Python one-shot training; otherwise > 0Input dimension dThe inferred or explicit pq.m must divide it
    nlistAuto from expected-vector-count, or explicit > 0IVF partition countControls list length, centroid cost, and automatic nprobe
    pq.code-ratioDefault 0.0625; finite and in (0, 0.25]Target ratio between PQ-code bytes and raw f32-vector bytesThe closest valid divisor is selected; max-bytes-per-vector can supply the target instead
    pq.mOptional expert override; > 0; d % m == 0Concrete subspace count and 8-bit code bytesTakes precedence over automatic sizing; use after representative recall measurements
    use-opqAuto enables when target-recall ≥ 0.9; explicit true/false winsLearn and apply an orthogonal rotationTrades training and query work for possible recall improvement
    ivf.train.max-points-per-centroidPositive integer; default 256Maximum training vectors per coarse centroidCaps coarse K-means input at nlist × value vectors, including hierarchical clustering stages.
    pq.train.max-points-per-centroidPositive integer; default 256Maximum training vectors per PQ centroidCaps input at 256 × value vectors per subquantizer, including PQ training inside OPQ.
    ivf.coarse-assignmentauto by default; optional exactControls build-time list assignmentauto uses Vamana when dimension × nlist ≥ 1,000,000, trading build speed for possible low-nprobe recall loss and graph startup cost; exact disables it
    ivf.pq-encodingauto by default; optional canonicalControls build-time PQ encodingFor 8-bit PQ with every dsub >= 4, auto uses transposed direct-L2 on x86 with AVX2+FMA and on AArch64, and blocked SGEMM on other CPUs when the codebook is finite; unsupported shapes and non-finite codebooks on the SGEMM fallback use canonical. canonical reproduces ProductQuantizer::encode_batch on the same CPU and runtime backend at substantially higher build cost
    nprobeAutomatic by default; explicit expert overrideLists readAuto accounts for K, average list size, and filter selectivity
    +

    Larger training samples increase training work and may improve centroid quality. These limits apply within the Trainer reservoir of max(65536, 64 × nlist) vectors; increasing them does not enlarge that reservoir. OPQ also retains its 65,536-vector input cap. See the shared training options for sampling and Paimon integration details.

    All PQ encoding backends produce the same index format. The accelerated automatic paths require 8-bit PQ (ksub=256) and at least four dimensions in every subvector. Eligible shapes use transposed direct squared L2 on x86 with AVX2+FMA and on AArch64; other CPUs use blocked SGEMM when the codebook is finite. Unsupported shapes and non-finite codebooks on the SGEMM fallback use canonical. SGEMM and canonical use the expanded form, so codes can differ near ties; NaN centroids and high-dynamic-range inputs can also behave differently because the expanded form can cancel. The automatic backend therefore is not byte-stable across CPU feature sets. Canonical mode reproduces ProductQuantizer::encode_batch on the same CPU and runtime backend, but its SIMD norm reductions also do not promise byte-identical codes across CPU feature sets.

    diff --git a/docs/ivf-rq.html b/docs/ivf-rq.html index add8457..e33b440 100644 --- a/docs/ivf-rq.html +++ b/docs/ivf-rq.html @@ -56,13 +56,15 @@

    Usage

    bits: 4, metric: MetricType::L2, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }; let params = VectorSearchParams::new(10, 64);

    Parameters

    -
    ParameterRequirement / defaultPurposeGuidance
    dimensionInferred by Java/Python one-shot training; otherwise > 0Logical vector dimensionStorage pads internally to a multiple of 64.
    nlistAuto from expected-vector-count, or explicit > 0IVF partition countCompare the resolved value with the same IVF-FLAT baseline.
    rq.bits1–8; auto from max-bytes-per-vector, otherwise 4Persisted residual level widthHigher values increase recall, file bytes, I/O, and scan work linearly.
    metricRequiredL2 / inner product / cosineSemantic, not inferred; fixed in the file.
    ivf.coarse-assignmentauto by default; optional exactControls build-time list assignmentauto uses Vamana when dimension × nlist ≥ 1,000,000, trading build speed for possible low-nprobe recall loss and graph startup cost; exact disables it.
    nprobeAutomatic by default; explicit expert overrideLists probedAuto accounts for K, average list size, and filter selectivity.
    +
    ParameterRequirement / defaultPurposeGuidance
    dimensionInferred by Java/Python one-shot training; otherwise > 0Logical vector dimensionStorage pads internally to a multiple of 64.
    nlistAuto from expected-vector-count, or explicit > 0IVF partition countCompare the resolved value with the same IVF-FLAT baseline.
    rq.bits1–8; auto from max-bytes-per-vector, otherwise 4Persisted residual level widthHigher values increase recall, file bytes, I/O, and scan work linearly.
    metricRequiredL2 / inner product / cosineSemantic, not inferred; fixed in the file.
    ivf.train.max-points-per-centroidPositive integer; default 256Maximum training vectors per coarse centroidCaps coarse K-means input at nlist × value vectors, including hierarchical clustering stages.
    ivf.coarse-assignmentauto by default; optional exactControls build-time list assignmentauto uses Vamana when dimension × nlist ≥ 1,000,000, trading build speed for possible low-nprobe recall loss and graph startup cost; exact disables it.
    nprobeAutomatic by default; explicit expert overrideLists probedAuto accounts for K, average list size, and filter selectivity.
    +

    Larger training samples increase training work and may improve centroid quality. These limits apply within the Trainer reservoir of max(65536, 64 × nlist) vectors; increasing them does not enlarge that reservoir. See the shared training options for sampling and Paimon integration details.

    No query-side bit widthThe Reader always evaluates the representation stored in the file. Changing rq.bits requires rebuilding the index; this keeps one file's accuracy and cost contract stable.
    diff --git a/docs/ivf-sq.html b/docs/ivf-sq.html index 34a52c9..718cefc 100644 --- a/docs/ivf-sq.html +++ b/docs/ivf-sq.html @@ -37,12 +37,14 @@

    Configuration

    nlist: 1024, metric: MetricType::L2, use_approximate_coarse_assignment: true, + ivf_train_max_points_per_centroid: 256, }; let params = VectorSearchParams::new(10, 16);

    Parameters

    -
    ParameterRequirementEffect
    dimensionInferred by Java/Python one-shot training; otherwise > 0Each vector uses exactly d SQ-code bytes.
    nlistAuto from expected-vector-count, or explicit > 0 and no larger than training countMore lists shorten scans but enlarge centroid and per-list-bound metadata.
    metricRequired: L2, inner product, or cosineSelects preprocessing and the distance kernel.
    ivf.coarse-assignmentauto by default; optional exactauto uses Vamana when dimension × nlist ≥ 1,000,000, trading build speed for possible low-nprobe recall loss and graph startup cost; exact disables it.
    nprobeAutomatic by default; explicit 1 to nlistAuto accounts for K, average list size, and filter selectivity; explicit values provide a measured override.
    +
    ParameterRequirementEffect
    dimensionInferred by Java/Python one-shot training; otherwise > 0Each vector uses exactly d SQ-code bytes.
    nlistAuto from expected-vector-count, or explicit > 0 and no larger than training countMore lists shorten scans but enlarge centroid and per-list-bound metadata.
    metricRequired: L2, inner product, or cosineSelects preprocessing and the distance kernel.
    ivf.train.max-points-per-centroidPositive integer; default 256Caps coarse K-means input at nlist × value vectors, including hierarchical clustering stages.
    ivf.coarse-assignmentauto by default; optional exactauto uses Vamana when dimension × nlist ≥ 1,000,000, trading build speed for possible low-nprobe recall loss and graph startup cost; exact disables it.
    nprobeAutomatic by default; explicit 1 to nlistAuto accounts for K, average list size, and filter selectivity; explicit values provide a measured override.
    +

    Larger training samples increase training work and may improve centroid quality. These limits apply within the Trainer reservoir of max(65536, 64 × nlist) vectors; increasing them does not enlarge that reservoir. See the shared training options for sampling and Paimon integration details.

    The scalar code width is fixed at 8 bits in v1. There is deliberately no sq.bits, graph-width, or search-width option.

    diff --git a/docs/releases.html b/docs/releases.html index 65a2fb2..94b0bb0 100644 --- a/docs/releases.html +++ b/docs/releases.html @@ -52,6 +52,7 @@

    Upcoming: 0.5.0

    IVF-SQ now pools residual training bounds, fuses residual encoding, transposes output in bounded parallel batches, and reuses query heaps with conservative L2 block pruning. Unified readers and language bindings also cache decoded partitions within the existing reader memory budget. The SIFT1M/GIST1M/GloVe benchmarks record build, native query, and recall results with reproducible commands.

    IVSQ v1 files remain compatible. Existing indexes receive reader optimizations without a rebuild; retrain and rebuild to use the new quantization bounds. Set the reader memory budget to zero to disable the SQ cache; direct Rust IVFSQIndexReader::open remains uncached. See reader options and upgrade details.

    Rust IVF API migrationVersion 0.5.0 makes quantizer_centroids private on IVFFlatIndex, IVFPQIndex, IVFSQIndex, and IVFRQIndex. Replace direct reads with quantizer_centroids() and direct assignments with set_quantizer_centroids(...). The setter validates the centroid shape, rejects replacement after vectors are added, and refreshes cached derived state. IVF variants of VectorIndexConfig also require use_approximate_coarse_assignment; set it to true for the automatic 0.5.0 behavior or false for exact nearest-centroid assignment. VectorIndexConfig::IvfPq additionally requires canonical_pq_encoding; use false for the default automatic encoder. Option-map callers can select these policies with ivf.coarse-assignment=auto|exact and ivf.pq-encoding=auto|canonical. The policies are fixed when the writer is created; direct IVF indexes do not expose post-training policy switches. ProductQuantizer::centroids and its derived norm cache are also private in 0.5.0; replace direct codebook reads with centroids() and replacements with set_centroids(...). The setter validates the codebook shape and refreshes the norm and transposed-codebook caches. The PQ layout fields d, m, nbits, dsub, ksub, and chunk_offsets are private; use the same-named read-only methods. To change the layout, construct a new ProductQuantizer so derived caches cannot outlive their layout. These are source-level changes; the stored index format is unchanged.
    +
    Rust training configuration migrationDirect VectorIndexConfig enum literals must now include ivf_train_max_points_per_centroid: 256 for IvfFlat, IvfSq, IvfRq, and IvfPq, and pq_train_max_points_per_centroid: 256 for IvfPq and DiskAnn. These values preserve the existing training defaults; constructors and train methods retain those defaults. Option-map callers can set ivf.train.max-points-per-centroid and pq.train.max-points-per-centroid without a fields.<field-name>. prefix. Paimon integrations must allow these keys through their option filter. The index file format is unchanged. See training options and sampling limits.
    IVF-PQ build encoding changes by defaultFor 8-bit PQ (ksub=256) where every subvector has at least four dimensions, the default automatic encoder uses transposed direct-L2 on x86 with AVX2+FMA and on AArch64. Other CPUs use blocked SGEMM expanded-form encoding for finite codebooks. Unsupported shapes, and non-finite codebooks on the SGEMM fallback, use canonical encoding. Codes can differ across these backends at near ties; NaN centroids and high-dynamic-range finite inputs can also behave differently because the expanded form can cancel. Set ivf.pq-encoding=canonical to reproduce ProductQuantizer::encode_batch on the same CPU and runtime backend. Canonical mode does not promise byte-identical codes across CPU feature sets and is substantially slower; the index format and query path are unchanged.