1use std::io::Write;
9use std::sync::Arc;
10
11use rustc_hash::FxHashMap;
12
13use crate::directories::DirectoryWriter;
14use crate::dsl::{
15 BinaryDenseVectorConfig, BinaryIndexType, DenseVectorConfig, Field, FieldType,
16 VectorIndexAlter, VectorIndexType,
17};
18use crate::error::{Error, Result};
19use crate::segment::{SegmentFiles, SegmentId, SegmentMeta};
20
21use super::IndexWriter;
22
23const MAX_IVF_CLUSTERS: usize = 1_048_576;
26const MIN_TRAINING_POINTS_PER_CENTROID: usize = 39;
29const COARSE_TRAINING_POINTS_PER_CENTROID: usize = 256;
32const MAX_SAMPLE_READ_BYTES: usize = 64 * 1024 * 1024;
35const MAX_SAMPLE_READ_AMPLIFICATION: usize = 4;
39const MAX_BINARY_REPLENISHMENT_CANDIDATES: usize = 1_000_000;
44const VALIDATION_SAMPLE_DENOMINATOR: usize = 10;
47const MAX_VALIDATION_SAMPLES: usize = 65_536;
48const MAX_VALIDATION_COORDINATE_WORK: usize = 512_000_000;
51const MODEL_SELECTION_SEEDS: [u64; 2] = [42, 0x9e37_79b9_7f4a_7c15];
55const MAX_MULTI_SEED_COORDINATE_WORK: usize = 4_000_000_000;
56const VECTOR_ARTIFACT_PREFIX: &str = "vector_artifact_";
59
60struct TrainedFieldUpdate {
61 field_id: u32,
62 index_type: super::metadata::VectorFieldIndexType,
63 vector_count: usize,
64 num_clusters: usize,
65 centroids_file: String,
66 codebook_file: Option<String>,
67 scann_generation: Option<u64>,
68 scann_artifact_id: Option<u64>,
69}
70
71enum TrainedFieldArtifacts {
72 FloatCentroids(crate::structures::CoarseCentroids),
75 Binary(crate::structures::BinaryCoarseQuantizer),
76 Scann(crate::structures::vector::scann::ScannTrainedArtifact),
77}
78
79struct TrainedFieldModel {
80 update: TrainedFieldUpdate,
81 artifacts: TrainedFieldArtifacts,
82}
83
84#[derive(Clone)]
85enum IvfFieldConfig {
86 Float(DenseVectorConfig),
87 Binary(BinaryDenseVectorConfig),
88}
89
90impl IvfFieldConfig {
91 fn dim(&self) -> usize {
92 match self {
93 Self::Float(config) => config.dim,
94 Self::Binary(config) => config.dim,
95 }
96 }
97
98 fn index_type(&self) -> super::metadata::VectorFieldIndexType {
99 match self {
100 Self::Float(config) => config.index_type.into(),
101 Self::Binary(config) => config.index_type.into(),
102 }
103 }
104
105 fn num_clusters(&self) -> Option<usize> {
106 match self {
107 Self::Float(config) => config.num_clusters,
108 Self::Binary(config) => config.num_clusters,
109 }
110 }
111
112 fn target_vectors(&self) -> Option<u64> {
113 match self {
114 Self::Float(config) => config.target_vectors,
115 Self::Binary(config) => config.target_vectors,
116 }
117 }
118
119 fn uses_target_sized_ivf(&self) -> bool {
120 self.num_clusters().is_none()
121 && self.target_vectors().is_some()
122 && (matches!(self, Self::Float(config) if config.index_type == VectorIndexType::IvfTq)
123 || matches!(self, Self::Binary(config) if config.index_type == BinaryIndexType::Ivf))
124 }
125
126 fn supports_deferred_flat(&self) -> bool {
127 self.uses_target_sized_ivf()
128 || matches!(self, Self::Float(config) if config.index_type == VectorIndexType::Scann)
129 || matches!(self, Self::Binary(config) if config.index_type == BinaryIndexType::Scann)
130 }
131
132 fn optimal_num_clusters(&self, vector_count: usize) -> usize {
133 match self {
134 Self::Float(config) => config.optimal_num_clusters(vector_count),
135 Self::Binary(config) => config.optimal_num_clusters(vector_count),
136 }
137 }
138}
139
140enum TrainingSample {
141 Float(Vec<f32>),
143 Binary(Vec<u8>),
144}
145
146#[derive(Clone, Copy, Debug)]
147struct OccupancyQuality {
148 p95: usize,
149 p99: usize,
150 max: usize,
151 empty: usize,
152 penalty: f64,
153}
154
155#[derive(Clone, Copy, Debug)]
156struct FloatBuildQuality {
157 objective: f64,
158 mean_exact_distortion: f64,
159 mean_routed_distortion: f64,
160 router_recall_at_1: f64,
161 mean_construction_assignments: f64,
162 residual_p50: f32,
163 residual_p95: f32,
164 residual_p99: f32,
165 occupancy: OccupancyQuality,
166}
167
168#[derive(Clone, Copy, Debug, Eq, PartialEq)]
169enum VectorGenerationMode {
170 BuildMissing,
171 RetrainAll,
172}
173
174#[derive(Clone, Copy, Debug, Eq, PartialEq)]
175pub enum AlterVectorIndexState {
176 Built,
177 DeferredFlat,
178 ParametersOnly,
179}
180
181#[derive(Clone, Copy, Debug, Eq, PartialEq)]
182pub struct AlterVectorIndexOutcome {
183 pub publication_generation: u64,
184 pub state: AlterVectorIndexState,
185}
186
187fn same_soar_layout(
188 left: Option<&crate::structures::SoarConfig>,
189 right: Option<&crate::structures::SoarConfig>,
190) -> bool {
191 match (left, right) {
192 (None, None) => true,
193 (Some(left), Some(right)) => {
194 left.num_secondary == right.num_secondary
195 && left.selective == right.selective
196 && left.spill_threshold.to_bits() == right.spill_threshold.to_bits()
197 }
198 _ => false,
199 }
200}
201
202fn alter_requires_rebuild(current: &IvfFieldConfig, target: &IvfFieldConfig) -> bool {
203 match (current, target) {
204 (IvfFieldConfig::Float(current), IvfFieldConfig::Float(target)) => {
205 current.index_type != target.index_type
206 || current.num_clusters != target.num_clusters
207 || (current.num_clusters.is_none()
208 && target.num_clusters.is_none()
209 && current.target_vectors != target.target_vectors)
210 || current.tree_levels != target.tree_levels
211 || current.ivf_routing != target.ivf_routing
212 || current.unit_norm != target.unit_norm
213 || !same_soar_layout(current.soar.as_ref(), target.soar.as_ref())
214 }
215 (IvfFieldConfig::Binary(current), IvfFieldConfig::Binary(target)) => {
216 current.index_type != target.index_type
217 || current.num_clusters != target.num_clusters
218 || (current.num_clusters.is_none()
219 && target.num_clusters.is_none()
220 && current.target_vectors != target.target_vectors)
221 || current.tree_levels != target.tree_levels
222 || current.ivf_routing != target.ivf_routing
223 || !same_soar_layout(current.soar.as_ref(), target.soar.as_ref())
224 }
225 _ => true,
226 }
227}
228
229impl TrainingSample {
230 fn len(&self, dim: usize) -> usize {
231 match self {
232 Self::Float(values) => values.len() / dim,
233 Self::Binary(codes) => codes.len() / dim.div_ceil(8),
234 }
235 }
236}
237
238struct SizeLimitedWriter<'a, W: Write + ?Sized> {
243 inner: &'a mut W,
244 written: usize,
245 limit: usize,
246}
247
248impl<'a, W: Write + ?Sized> SizeLimitedWriter<'a, W> {
249 fn new(inner: &'a mut W, limit: usize) -> Self {
250 Self {
251 inner,
252 written: 0,
253 limit,
254 }
255 }
256}
257
258impl<W: Write + ?Sized> Write for SizeLimitedWriter<'_, W> {
259 fn write(&mut self, buffer: &[u8]) -> std::io::Result<usize> {
260 let next_size = self
261 .written
262 .checked_add(buffer.len())
263 .ok_or_else(|| std::io::Error::other("trained artifact size overflow"))?;
264 if next_size > self.limit {
265 return Err(std::io::Error::new(
266 std::io::ErrorKind::InvalidData,
267 format!(
268 "trained artifact exceeds the {}-byte safety limit",
269 self.limit
270 ),
271 ));
272 }
273 let written = self.inner.write(buffer)?;
274 self.written += written;
275 Ok(written)
276 }
277
278 fn flush(&mut self) -> std::io::Result<()> {
279 self.inner.flush()
280 }
281}
282
283fn validate_explicit_cluster_count(num_clusters: Option<usize>) -> Result<()> {
284 match num_clusters {
285 Some(0) => Err(Error::Schema(
286 "dense vector num_clusters must be at least 1".to_string(),
287 )),
288 Some(value) if value > MAX_IVF_CLUSTERS => Err(Error::Schema(format!(
289 "dense vector num_clusters must not exceed {MAX_IVF_CLUSTERS}, got {value}"
290 ))),
291 _ => Ok(()),
292 }
293}
294
295fn effective_field_num_clusters(
296 config: &IvfFieldConfig,
297 corpus_count: usize,
298 sample_count: usize,
299) -> Result<usize> {
300 if sample_count == 0 {
301 return Err(Error::Schema(
302 "cannot train an IVF vector index without sample vectors".to_string(),
303 ));
304 }
305 validate_explicit_cluster_count(config.num_clusters())?;
306 let centroid_bytes = match config {
307 IvfFieldConfig::Float(config) => config.dim.saturating_mul(size_of::<f32>()),
308 IvfFieldConfig::Binary(config) => config.dim.div_ceil(8),
309 };
310 let artifact_limit = super::metadata::MAX_TRAINED_ARTIFACT_BYTES
311 .saturating_sub(1024)
312 .checked_div(centroid_bytes.max(1))
313 .unwrap_or(0)
314 .max(1);
315 let quality_limit = if config.num_clusters().is_some() {
316 sample_count
317 } else {
318 (sample_count / MIN_TRAINING_POINTS_PER_CENTROID)
319 .max(16)
320 .min(sample_count)
321 };
322 let requested = config.optimal_num_clusters(corpus_count);
323 let stable_automatic_topology =
324 config.num_clusters().is_none() && config.target_vectors().is_some();
325 if (config.num_clusters().is_some() || stable_automatic_topology) && requested > artifact_limit
326 {
327 return Err(Error::Schema(format!(
328 "configured IVF codebook needs {} bytes for {} centroids, exceeding the {}-byte artifact limit",
329 requested.saturating_mul(centroid_bytes),
330 requested,
331 super::metadata::MAX_TRAINED_ARTIFACT_BYTES,
332 )));
333 }
334 if stable_automatic_topology && requested > quality_limit {
335 let required = requested.saturating_mul(MIN_TRAINING_POINTS_PER_CENTROID);
336 return Err(Error::Schema(format!(
337 "target-sized IVF geometry selected {requested} leaves and needs at least {required} training samples (hardcoded {MIN_TRAINING_POINTS_PER_CENTROID} samples/leaf), but the builder can supply {sample_count}"
338 )));
339 }
340 Ok(requested.min(quality_limit).min(artifact_limit))
341}
342
343fn training_sample_limit(
344 max_samples: usize,
345 max_bytes: usize,
346 bytes_per_sample: usize,
347) -> Result<usize> {
348 if max_samples == 0 || max_bytes == 0 || bytes_per_sample == 0 {
349 return Err(Error::Schema(
350 "vector training sample count, memory budget, and vector size must be greater than zero"
351 .into(),
352 ));
353 }
354 let memory_limited = max_bytes / bytes_per_sample;
355 if memory_limited == 0 {
356 return Err(Error::Schema(format!(
357 "vector training memory budget ({max_bytes} bytes) cannot hold one {bytes_per_sample}-byte sample"
358 )));
359 }
360 Ok(max_samples.min(memory_limited))
361}
362
363fn training_sample_bytes(config: &IvfFieldConfig) -> Result<usize> {
364 match config {
365 IvfFieldConfig::Float(config) => config
366 .dim
367 .checked_mul(size_of::<f32>())
368 .ok_or_else(|| Error::Schema("float training vector size overflows".into())),
369 IvfFieldConfig::Binary(config) => Ok(config.dim.div_ceil(8)),
370 }
371}
372
373fn required_scann_training_sample(num_leaves: u32) -> Result<u64> {
374 u64::from(num_leaves)
375 .checked_mul(crate::structures::vector::scann::MIN_PARTITION_TRAINING_POINTS_PER_LEAF)
376 .map(|required| required.max(crate::structures::vector::scann::MIN_POINTS_FOR_PARTITIONING))
377 .ok_or_else(|| Error::Schema("ScaNN minimum training sample overflows u64".into()))
378}
379
380fn final_training_sample_count(
384 config: &IvfFieldConfig,
385 corpus_count: usize,
386 sample_limit: usize,
387) -> Result<usize> {
388 let available = corpus_count.min(sample_limit);
389 if available == 0 {
390 return Ok(0);
391 }
392 let clusters = effective_field_num_clusters(config, corpus_count, available)?;
393 Ok(available.min(clusters.saturating_mul(COARSE_TRAINING_POINTS_PER_CENTROID)))
394}
395
396fn scann_geometry(
397 config: &DenseVectorConfig,
398 corpus_count: usize,
399) -> Result<crate::structures::vector::scann::ScannGeometry> {
400 let dimension = u32::try_from(config.dim)
401 .map_err(|_| Error::Schema("ScaNN dimension exceeds u32".into()))?;
402 let Some(leaves) = config.num_clusters else {
403 let sizing_count = config.target_vectors.unwrap_or(0).max(corpus_count as u64);
404 return crate::structures::vector::scann::derive_geometry_with_levels(
405 sizing_count,
406 dimension,
407 config.tree_levels,
408 )
409 .map_err(|error| Error::Schema(error.to_string()));
410 };
411 let leaves = u32::try_from(leaves)
412 .map_err(|_| Error::Schema("ScaNN num_clusters exceeds u32".into()))?;
413 crate::structures::vector::scann::geometry_for_leaves_with_auto_depth(
414 leaves,
415 dimension,
416 config.tree_levels,
417 )
418 .map_err(|error| Error::Schema(error.to_string()))
419}
420
421fn binary_scann_geometry(
422 config: &BinaryDenseVectorConfig,
423 corpus_count: usize,
424) -> Result<crate::structures::vector::scann::ScannGeometry> {
425 let leaves = config.num_clusters;
426 let Some(leaves) = leaves else {
427 let sizing_count = config.target_vectors.unwrap_or(0).max(corpus_count as u64);
428 return crate::structures::vector::scann::derive_geometry_with_levels(
429 sizing_count,
430 config.dim as u32,
431 config.tree_levels,
432 )
433 .map_err(|error| Error::Schema(error.to_string()));
434 };
435 crate::structures::vector::scann::geometry_for_leaves_with_auto_depth(
436 u32::try_from(leaves)
437 .map_err(|_| Error::Schema("binary ScaNN num_clusters exceeds u32".into()))?,
438 config.dim as u32,
439 config.tree_levels,
440 )
441 .map_err(|error| Error::Schema(error.to_string()))
442}
443
444fn scann_training_sample_count(
445 field: Field,
446 total: usize,
447 limit: usize,
448 bytes_per_sample: usize,
449 geometry: &crate::structures::vector::scann::ScannGeometry,
450 binary: bool,
451) -> Result<usize> {
452 let desired = crate::structures::vector::scann::desired_training_sample(
453 total as u64,
454 geometry.num_leaves,
455 ) as usize;
456 let take = desired.min(limit);
457 let minimum = usize::try_from(required_scann_training_sample(geometry.num_leaves)?)
458 .map_err(|_| Error::Schema("ScaNN minimum sample exceeds usize".into()))?;
459 if take < minimum {
460 let minimum_bytes = minimum
461 .checked_mul(bytes_per_sample)
462 .ok_or_else(|| Error::Schema("ScaNN minimum training memory overflows".into()))?;
463 let kind = if binary { "binary ScaNN" } else { "ScaNN" };
464 return Err(Error::Schema(format!(
465 "{kind} geometry for field {} needs at least {} sampled vectors ({} leaves at the hardcoded {} samples/leaf; {} bytes), but the builder limits allow {}; raise vector_training_memory_bytes/vector_training_max_samples",
466 field.0,
467 minimum,
468 geometry.num_leaves,
469 crate::structures::vector::scann::MIN_PARTITION_TRAINING_POINTS_PER_LEAF,
470 minimum_bytes,
471 take
472 )));
473 }
474 Ok(take)
475}
476
477fn deterministic_sample_ordinals(total: usize, take: usize, seed: u64) -> Vec<usize> {
482 debug_assert!(take <= total);
483 if take == 0 {
484 return Vec::new();
485 }
486 if take == total {
487 return (0..total).collect();
488 }
489 let mut rng = <rand::rngs::StdRng as rand::SeedableRng>::seed_from_u64(seed);
490 let mut ordinals = rand::seq::index::sample(&mut rng, total, take).into_vec();
491 ordinals.sort_unstable();
492 ordinals
493}
494
495fn binary_replenishment_candidate_count(total: usize, take: usize, minimum_usable: usize) -> usize {
496 let reserve = minimum_usable.clamp(1_024, MAX_BINARY_REPLENISHMENT_CANDIDATES);
497 take.saturating_add(reserve).min(total)
498}
499
500fn validation_sample_count(
501 sample_count: usize,
502 num_clusters: usize,
503 values_per_vector: usize,
504) -> usize {
505 if sample_count <= num_clusters {
506 return 0;
507 }
508 let quality_floor = num_clusters.saturating_mul(MIN_TRAINING_POINTS_PER_CENTROID);
509 let minimum_training = if sample_count >= quality_floor {
510 quality_floor
511 } else {
512 num_clusters
513 };
514 let exact_scan_limit = MAX_VALIDATION_COORDINATE_WORK
515 .checked_div(num_clusters.saturating_mul(values_per_vector).max(1))
516 .unwrap_or(0)
517 .max(1);
518 sample_count
519 .div_ceil(VALIDATION_SAMPLE_DENOMINATOR)
520 .clamp(1, MAX_VALIDATION_SAMPLES)
521 .min(exact_scan_limit)
522 .min(sample_count - minimum_training)
523}
524
525fn model_selection_seeds(
526 training_count: usize,
527 num_clusters: usize,
528 dim: usize,
529 has_validation: bool,
530) -> &'static [u64] {
531 if !has_validation {
532 return &MODEL_SELECTION_SEEDS[..1];
533 }
534 let distance_passes = crate::structures::vector::estimated_euclidean_kmeans_distance_multiplier(
535 training_count,
536 num_clusters,
537 25,
538 );
539 let work = training_count
540 .saturating_mul(num_clusters)
541 .saturating_mul(dim)
542 .saturating_mul(distance_passes)
543 .saturating_mul(MODEL_SELECTION_SEEDS.len());
544 if work <= MAX_MULTI_SEED_COORDINATE_WORK {
545 &MODEL_SELECTION_SEEDS
546 } else {
547 &MODEL_SELECTION_SEEDS[..1]
548 }
549}
550
551fn percentile_index(len: usize, percentile: usize) -> usize {
552 debug_assert!(len > 0 && percentile <= 100);
553 (len - 1).saturating_mul(percentile).div_ceil(100)
554}
555
556fn occupancy_quality(mut counts: Vec<usize>, observations: usize) -> OccupancyQuality {
557 if counts.is_empty() {
558 return OccupancyQuality {
559 p95: 0,
560 p99: 0,
561 max: 0,
562 empty: 0,
563 penalty: 0.0,
564 };
565 }
566 counts.sort_unstable();
567 let p95 = counts[percentile_index(counts.len(), 95)];
568 let p99 = counts[percentile_index(counts.len(), 99)];
569 let max = counts.last().copied().unwrap_or(0);
570 let empty = counts.partition_point(|&count| count == 0);
571 let expected = observations as f64 / counts.len() as f64;
572 let denominator = expected.max(1.0);
573 let p99_excess = (p99 as f64 / denominator - 1.0).max(0.0);
574 let max_excess = (max as f64 / denominator - 1.0).max(0.0);
575 let empty_fraction = empty as f64 / counts.len() as f64;
576 let penalty = 0.02 * p99_excess + 0.005 * max_excess + 0.05 * empty_fraction;
579 OccupancyQuality {
580 p95,
581 p99,
582 max,
583 empty,
584 penalty,
585 }
586}
587
588fn float_model_selection_objective(
589 mean_exact_distortion: f64,
590 mean_routed_distortion: f64,
591 occupancy_penalty: f64,
592) -> f64 {
593 let scale = mean_exact_distortion.max(f64::from(f32::EPSILON));
594 let routed_distortion_excess = (mean_routed_distortion - mean_exact_distortion).max(0.0);
595 mean_exact_distortion + scale * occupancy_penalty + routed_distortion_excess
596}
597
598fn partition_contiguous_holdout_suffix<T>(
603 values: &mut [T],
604 values_per_vector: usize,
605 validation_count: usize,
606 seed: u64,
607) -> usize {
608 assert!(values_per_vector > 0);
609 assert_eq!(values.len() % values_per_vector, 0);
610 let sample_count = values.len() / values_per_vector;
611 assert!(validation_count <= sample_count);
612 if validation_count == 0 {
613 return values.len();
614 }
615 let validation_indices =
616 deterministic_sample_ordinals(sample_count, validation_count, seed ^ 0x5641_4c49_4441_5445);
617 let split_row = sample_count - validation_count;
618 let prefix_validation_count = validation_indices.partition_point(|&index| index < split_row);
619 let mut right = sample_count;
620 for &left in &validation_indices[..prefix_validation_count] {
621 loop {
622 right -= 1;
623 if validation_indices.binary_search(&right).is_err() {
624 break;
625 }
626 }
627 debug_assert!(right >= split_row);
628 for component in 0..values_per_vector {
629 values.swap(
630 left * values_per_vector + component,
631 right * values_per_vector + component,
632 );
633 }
634 }
635 split_row * values_per_vector
636}
637
638fn evaluate_float_build_quality(
639 centroids: &crate::structures::CoarseCentroids,
640 validation: &[f32],
641 routing: crate::dsl::IvfRoutingMode,
642) -> Option<FloatBuildQuality> {
643 let dim = centroids.dim;
644 let validation_count = validation.len() / dim;
645 if validation_count == 0 {
646 return None;
647 }
648 let mut occupancy = vec![0usize; centroids.num_clusters as usize];
649 let mut residual_scales = Vec::with_capacity(validation_count);
650 let mut exact_distortion_sum = 0.0f64;
651 let mut routed_distortion_sum = 0.0f64;
652 let mut router_hits = 0usize;
653 let mut construction_assignments = 0usize;
654 let effective_routing = crate::structures::vector::ivf::routing::effective_routing_mode(
655 routing,
656 centroids.num_clusters as usize,
657 );
658 let exact_routing = effective_routing == crate::dsl::IvfRoutingMode::Flat;
659 for vector in validation.chunks_exact(dim) {
660 let (exact_cluster_id, routed_cluster_id) = if exact_routing {
661 if centroids.soar_config.is_some() {
662 let construction_assignment = centroids.assign_with_routing(vector, routing);
666 let exact_cluster_id = construction_assignment.primary_cluster;
667 for cluster_id in construction_assignment.all_clusters() {
668 occupancy[cluster_id as usize] += 1;
669 construction_assignments += 1;
670 }
671 (exact_cluster_id, exact_cluster_id)
672 } else {
673 let exact_cluster_id = centroids.find_nearest(vector);
677 occupancy[exact_cluster_id as usize] += 1;
678 construction_assignments += 1;
679 (exact_cluster_id, exact_cluster_id)
680 }
681 } else {
682 let exact_cluster_id = centroids.find_nearest(vector);
683 let routed_cluster_id = centroids.probe(vector, 1, routing).cluster_ids[0];
684 let construction_assignment = centroids.assign_with_routing(vector, routing);
685 for cluster_id in construction_assignment.all_clusters() {
686 occupancy[cluster_id as usize] += 1;
687 construction_assignments += 1;
688 }
689 (exact_cluster_id, routed_cluster_id)
690 };
691 router_hits += usize::from(routed_cluster_id == exact_cluster_id);
692 let exact_distance = crate::structures::simd::squared_l2_f32(
693 vector,
694 centroids.get_centroid(exact_cluster_id),
695 );
696 let routed_distance = if routed_cluster_id == exact_cluster_id {
697 exact_distance
698 } else {
699 crate::structures::simd::squared_l2_f32(
700 vector,
701 centroids.get_centroid(routed_cluster_id),
702 )
703 };
704 exact_distortion_sum += f64::from(exact_distance);
705 routed_distortion_sum += f64::from(routed_distance);
706 residual_scales.push(exact_distance.max(0.0).sqrt());
707 }
708 residual_scales.sort_unstable_by(f32::total_cmp);
709 let occupancy = occupancy_quality(occupancy, construction_assignments);
710 let mean_exact_distortion = exact_distortion_sum / validation_count as f64;
711 let mean_routed_distortion = routed_distortion_sum / validation_count as f64;
712 let router_recall_at_1 = router_hits as f64 / validation_count as f64;
713 let mean_construction_assignments = construction_assignments as f64 / validation_count as f64;
714 let objective = float_model_selection_objective(
718 mean_exact_distortion,
719 mean_routed_distortion,
720 occupancy.penalty,
721 );
722 Some(FloatBuildQuality {
723 objective,
724 mean_exact_distortion,
725 mean_routed_distortion,
726 router_recall_at_1,
727 mean_construction_assignments,
728 residual_p50: residual_scales[percentile_index(validation_count, 50)],
729 residual_p95: residual_scales[percentile_index(validation_count, 95)],
730 residual_p99: residual_scales[percentile_index(validation_count, 99)],
731 occupancy,
732 })
733}
734
735#[cfg(test)]
742fn effective_ivf_num_clusters(
743 config: &DenseVectorConfig,
744 corpus_count: usize,
745 sample_count: usize,
746) -> Result<usize> {
747 if sample_count == 0 {
748 return Err(Error::Schema(
749 "cannot train an IVF vector index without sample vectors".to_string(),
750 ));
751 }
752
753 effective_field_num_clusters(
754 &IvfFieldConfig::Float(config.clone()),
755 corpus_count,
756 sample_count,
757 )
758}
759
760impl<D: DirectoryWriter + 'static> IndexWriter<D> {
761 pub async fn alter_vector_index(
768 &mut self,
769 field: Field,
770 alter: VectorIndexAlter,
771 ) -> Result<AlterVectorIndexOutcome> {
772 self.commit().await?;
776 let current_generation = self.segment_manager.published_generation();
777 let current_entry = current_generation
778 .schema
779 .get_field_entry(field)
780 .ok_or_else(|| Error::Schema(format!("unknown vector field {}", field.0)))?;
781 let current_config = match current_entry.field_type {
782 FieldType::DenseVector => current_entry
783 .dense_vector_config
784 .clone()
785 .map(IvfFieldConfig::Float),
786 FieldType::BinaryDenseVector => current_entry
787 .binary_dense_vector_config
788 .clone()
789 .map(IvfFieldConfig::Binary),
790 _ => None,
791 }
792 .ok_or_else(|| Error::Schema(format!("field {} is not a vector field", field.0)))?;
793 let candidate_schema = Arc::new(
794 current_generation
795 .schema
796 .with_vector_index_alter(field, alter)
797 .map_err(Error::Schema)?,
798 );
799 let target_entry = candidate_schema
800 .get_field_entry(field)
801 .expect("validated ALTER preserves the field");
802 let target_config = match target_entry.field_type {
803 FieldType::DenseVector => IvfFieldConfig::Float(
804 target_entry
805 .dense_vector_config
806 .clone()
807 .expect("validated dense ALTER has a config"),
808 ),
809 FieldType::BinaryDenseVector => IvfFieldConfig::Binary(
810 target_entry
811 .binary_dense_vector_config
812 .clone()
813 .expect("validated binary ALTER has a config"),
814 ),
815 _ => unreachable!("validated vector ALTER changed field type"),
816 };
817
818 let artifact_update = self.segment_manager.begin_vector_artifact_update().await?;
819 if !alter_requires_rebuild(¤t_config, &target_config) {
820 self.segment_manager
821 .publish_vector_schema_only(&artifact_update, candidate_schema)
822 .await?;
823 drop(artifact_update);
824 return Ok(AlterVectorIndexOutcome {
825 publication_generation: self.segment_manager.publication_id(),
826 state: AlterVectorIndexState::ParametersOnly,
827 });
828 }
829
830 self.cleanup_unreferenced_vector_artifacts().await;
831 let snapshot = self.segment_manager.acquire_snapshot().await;
832 let fields = vec![(field, target_config.clone())];
833 let total_vectors = self
834 .count_vectors_for_training(
835 snapshot.segment_ids(),
836 &fields,
837 false,
838 ¤t_generation.schema,
839 )
840 .await?;
841 let corpus_count = total_vectors.get(&field.0).copied().unwrap_or(0);
842 let mut candidate_metadata = self.segment_manager.read_metadata(Clone::clone).await;
843 candidate_metadata.init_field(field.0, target_config.index_type());
844 let field_metadata = candidate_metadata
845 .vector_fields
846 .get_mut(&field.0)
847 .expect("ALTER initialized vector field metadata");
848 field_metadata.index_type = target_config.index_type();
849 field_metadata.state = super::VectorIndexState::Flat;
850 field_metadata.centroids_file = None;
851 field_metadata.codebook_file = None;
852 field_metadata.artifact_generation = None;
853 field_metadata.artifact_id = None;
854 candidate_metadata.refresh_total_vectors();
855
856 let artifact_generation = SegmentId::new().to_hex();
857 let updates = self
858 .train_fields(
859 snapshot.segment_ids(),
860 &fields,
861 &total_vectors,
862 &artifact_generation,
863 ¤t_generation.schema,
864 )
865 .await?;
866 for update in &updates {
867 if let (Some(generation), Some(artifact_id)) =
868 (update.scann_generation, update.scann_artifact_id)
869 {
870 candidate_metadata.mark_scann_field_built(
871 update.field_id,
872 update.vector_count,
873 update.num_clusters,
874 update.centroids_file.clone(),
875 generation,
876 artifact_id,
877 )?;
878 } else {
879 candidate_metadata.mark_field_built(
880 update.field_id,
881 update.vector_count,
882 update.num_clusters,
883 update.centroids_file.clone(),
884 update.codebook_file.clone(),
885 );
886 }
887 }
888 let built = candidate_metadata.is_field_built(field.0);
889 if !built && !target_config.supports_deferred_flat() && corpus_count > 0 {
890 return Err(Error::Schema(format!(
891 "cannot train target vector index for field {} from {corpus_count} vectors",
892 field.0
893 )));
894 }
895
896 let candidate_trained = super::IndexMetadata::try_load_trained_from_fields(
897 &candidate_metadata.vector_fields,
898 candidate_schema.as_ref(),
899 self.directory.as_ref(),
900 )
901 .await?
902 .map(Arc::new);
903 let finalize_ann_segments = candidate_trained.is_some();
904 let rewrite_trained = candidate_trained
905 .clone()
906 .unwrap_or_else(|| Arc::new(crate::segment::TrainedVectorStructures::default()));
907 let staged = self
908 .segment_manager
909 .stage_vector_generation_with_schema(
910 &artifact_update,
911 snapshot.segment_ids(),
912 &[field.0],
913 rewrite_trained,
914 true,
915 Arc::clone(&candidate_schema),
916 )
917 .await?;
918 self.segment_manager
919 .publish_vector_generation_with_schema(
920 &artifact_update,
921 candidate_schema,
922 candidate_metadata.vector_fields,
923 candidate_trained,
924 staged,
925 )
926 .await?;
927 drop(snapshot);
928 drop(artifact_update);
929 if finalize_ann_segments {
930 self.segment_manager
931 .rewrite_vector_segments(&[field.0])
932 .await?;
933 }
934 self.cleanup_unreferenced_vector_artifacts().await;
935
936 Ok(AlterVectorIndexOutcome {
937 publication_generation: self.segment_manager.publication_id(),
938 state: if built {
939 AlterVectorIndexState::Built
940 } else {
941 AlterVectorIndexState::DeferredFlat
942 },
943 })
944 }
945
946 pub async fn build_vector_index(&self) -> Result<()> {
953 self.build_vector_generation(VectorGenerationMode::BuildMissing)
954 .await
955 }
956
957 pub async fn retrain_vector_index(&self) -> Result<()> {
962 self.build_vector_generation(VectorGenerationMode::RetrainAll)
963 .await
964 }
965
966 async fn build_vector_generation(&self, mode: VectorGenerationMode) -> Result<()> {
967 let artifact_update = self.segment_manager.begin_vector_artifact_update().await?;
968 let generation = self.segment_manager.published_generation();
969 let schema = generation.schema.clone();
970 let dense_fields = Self::get_ivf_vector_fields(&schema);
971 if dense_fields.is_empty() {
972 log::info!(
973 "[vector_training] no dense vector fields configured for ANN indexing: index={}",
974 schema.index_label()
975 );
976 return Ok(());
977 }
978
979 self.cleanup_unreferenced_vector_artifacts().await;
980
981 let fields_to_train = match mode {
982 VectorGenerationMode::BuildMissing => self.get_fields_to_build(&dense_fields).await,
983 VectorGenerationMode::RetrainAll => dense_fields.clone(),
984 };
985 for (_, config) in &fields_to_train {
986 if !matches!(
987 config,
988 IvfFieldConfig::Float(config) if config.index_type == VectorIndexType::Scann
989 ) && !matches!(
990 config,
991 IvfFieldConfig::Binary(config) if config.index_type == BinaryIndexType::Scann
992 ) {
993 validate_explicit_cluster_count(config.num_clusters())?;
994 }
995 }
996
997 let snapshot = self.segment_manager.acquire_snapshot().await;
998 if snapshot.is_empty() {
999 if mode == VectorGenerationMode::RetrainAll {
1000 return Err(Error::Schema(
1001 "cannot retrain vector centroids without committed segments".into(),
1002 ));
1003 }
1004 return Ok(());
1005 }
1006
1007 let mut candidate_metadata = self.segment_manager.read_metadata(Clone::clone).await;
1008 if !fields_to_train.is_empty() {
1009 let total_vectors = self
1010 .count_vectors_for_training(
1011 snapshot.segment_ids(),
1012 &fields_to_train,
1013 mode == VectorGenerationMode::BuildMissing,
1014 &schema,
1015 )
1016 .await?;
1017 let artifact_generation = SegmentId::new().to_hex();
1018 let updates = self
1019 .train_fields(
1020 snapshot.segment_ids(),
1021 &fields_to_train,
1022 &total_vectors,
1023 &artifact_generation,
1024 &schema,
1025 )
1026 .await?;
1027 for update in &updates {
1028 candidate_metadata.init_field(update.field_id, update.index_type);
1029 if let (Some(generation), Some(artifact_id)) =
1030 (update.scann_generation, update.scann_artifact_id)
1031 {
1032 candidate_metadata.mark_scann_field_built(
1033 update.field_id,
1034 update.vector_count,
1035 update.num_clusters,
1036 update.centroids_file.clone(),
1037 generation,
1038 artifact_id,
1039 )?;
1040 } else {
1041 candidate_metadata.mark_field_built(
1042 update.field_id,
1043 update.vector_count,
1044 update.num_clusters,
1045 update.centroids_file.clone(),
1046 update.codebook_file.clone(),
1047 );
1048 }
1049 }
1050 }
1051
1052 let target_field_ids = dense_fields
1053 .iter()
1054 .filter_map(|(field, _)| {
1055 candidate_metadata
1056 .is_field_built(field.0)
1057 .then_some(field.0)
1058 })
1059 .collect::<Vec<_>>();
1060 if target_field_ids.is_empty() {
1061 return Ok(());
1062 }
1063
1064 let candidate_trained = super::IndexMetadata::try_load_trained_from_fields(
1065 &candidate_metadata.vector_fields,
1066 schema.as_ref(),
1067 self.directory.as_ref(),
1068 )
1069 .await?
1070 .map(Arc::new)
1071 .ok_or_else(|| Error::Internal("candidate vector generation has no artifacts".into()))?;
1072
1073 let staged = self
1074 .segment_manager
1075 .stage_vector_generation(
1076 &artifact_update,
1077 snapshot.segment_ids(),
1078 &target_field_ids,
1079 Arc::clone(&candidate_trained),
1080 mode == VectorGenerationMode::RetrainAll,
1081 )
1082 .await?;
1083 self.segment_manager
1084 .publish_vector_generation(
1085 &artifact_update,
1086 candidate_metadata.vector_fields,
1087 candidate_trained,
1088 staged,
1089 )
1090 .await?;
1091
1092 drop(snapshot);
1096 drop(artifact_update);
1097
1098 self.segment_manager
1102 .rewrite_vector_segments(&target_field_ids)
1103 .await?;
1104 self.cleanup_unreferenced_vector_artifacts().await;
1105 log::info!(
1106 "[vector_training] ANN generation {:?} complete: index={} {} field(s)",
1107 mode,
1108 self.schema.index_label(),
1109 target_field_ids.len(),
1110 );
1111 Ok(())
1112 }
1113
1114 async fn train_fields(
1115 &self,
1116 segment_ids: &[String],
1117 fields: &[(Field, IvfFieldConfig)],
1118 total_vectors: &FxHashMap<u32, usize>,
1119 artifact_generation: &str,
1120 schema: &Arc<crate::dsl::Schema>,
1121 ) -> Result<Vec<TrainedFieldUpdate>> {
1122 let training_pool = self.segment_manager.background_cpu_pool();
1123 let index_label = self.schema.index_label();
1124 let mut missing = Vec::new();
1125 let mut updates = Vec::with_capacity(fields.len());
1126 for (field, config) in fields {
1127 let corpus_count = total_vectors.get(&field.0).copied().unwrap_or(0);
1131 if config.uses_target_sized_ivf() {
1132 let leaves = config.optimal_num_clusters(corpus_count);
1133 let required = leaves.saturating_mul(MIN_TRAINING_POINTS_PER_CENTROID);
1134 if corpus_count < required {
1135 log::info!(
1136 "[vector_training] deferring target-sized IVF field {}: index={} has {} vectors; selected {}-leaf topology requires at least {}",
1137 field.0,
1138 index_label,
1139 corpus_count,
1140 leaves,
1141 required,
1142 );
1143 continue;
1144 }
1145 }
1146 if let IvfFieldConfig::Float(scann) = config
1147 && scann.index_type == VectorIndexType::Scann
1148 {
1149 let geometry = if corpus_count == 0 {
1150 None
1151 } else {
1152 Some(scann_geometry(scann, corpus_count)?)
1153 };
1154 let required = geometry.as_ref().map_or(
1155 crate::structures::vector::scann::MIN_POINTS_FOR_PARTITIONING,
1156 |geometry| {
1157 crate::structures::vector::scann::MIN_POINTS_FOR_PARTITIONING.max(
1158 required_scann_training_sample(geometry.num_leaves)
1159 .expect("validated ScaNN leaves fit u64"),
1160 )
1161 },
1162 );
1163 if corpus_count < required as usize {
1164 log::info!(
1165 "[vector_training] deferring ScaNN field {}: index={} has {} vectors; selected geometry requires {} (max of hardcoded partition floor {} and {} samples/leaf)",
1166 field.0,
1167 index_label,
1168 corpus_count,
1169 required,
1170 crate::structures::vector::scann::MIN_POINTS_FOR_PARTITIONING,
1171 crate::structures::vector::scann::MIN_PARTITION_TRAINING_POINTS_PER_LEAF,
1172 );
1173 continue;
1174 }
1175 let geometry = geometry.expect("non-empty ready corpus has ScaNN geometry");
1176 debug_assert!(geometry.centroid_levels > 0);
1177 if scann.soar.is_some() {
1178 return Err(Error::Schema(format!(
1179 "ScaNN field {} enables SOAR, but ScaNN SOAR secondary assignments are not implemented; set soar: null before building",
1180 field.0
1181 )));
1182 }
1183 }
1184 if let IvfFieldConfig::Binary(scann) = config
1185 && scann.index_type == BinaryIndexType::Scann
1186 {
1187 let required = if corpus_count == 0 {
1188 crate::structures::vector::scann::MIN_POINTS_FOR_PARTITIONING
1189 } else {
1190 let geometry = binary_scann_geometry(scann, corpus_count)?;
1191 crate::structures::vector::scann::MIN_POINTS_FOR_PARTITIONING
1192 .max(required_scann_training_sample(geometry.num_leaves)?)
1193 };
1194 if corpus_count < required as usize {
1195 log::info!(
1196 "[vector_training] deferring binary ScaNN field {}: index={} has {} vectors; selected geometry requires {}",
1197 field.0,
1198 index_label,
1199 corpus_count,
1200 required,
1201 );
1202 continue;
1203 }
1204 }
1205 let sampled = self
1206 .collect_training_sample(segment_ids, *field, config, corpus_count, schema)
1207 .await?;
1208 let Some(mut sample) = sampled else {
1209 if matches!(config, IvfFieldConfig::Float(config) if config.index_type == VectorIndexType::Scann)
1210 || matches!(config, IvfFieldConfig::Binary(config) if config.index_type == BinaryIndexType::Scann)
1211 {
1212 log::warn!(
1213 "[vector_training] deferring ScaNN field {}: index={} has no usable sampled vectors after quality filtering",
1214 field.0,
1215 index_label,
1216 );
1217 continue;
1218 }
1219 missing.push(field.0);
1220 continue;
1221 };
1222 let minimum_usable = match config {
1223 IvfFieldConfig::Float(config) if config.index_type == VectorIndexType::Scann => {
1224 usize::try_from(required_scann_training_sample(
1225 scann_geometry(config, corpus_count)?.num_leaves,
1226 )?)
1227 .map_err(|_| Error::Schema("ScaNN minimum sample exceeds usize".into()))?
1228 }
1229 IvfFieldConfig::Binary(config) if config.index_type == BinaryIndexType::Scann => {
1230 usize::try_from(required_scann_training_sample(
1231 binary_scann_geometry(config, corpus_count)?.num_leaves,
1232 )?)
1233 .map_err(|_| {
1234 Error::Schema("binary ScaNN minimum sample exceeds usize".into())
1235 })?
1236 }
1237 _ => 0,
1238 };
1239 let usable = sample.len(config.dim());
1240 if usable < minimum_usable {
1241 log::warn!(
1242 "[vector_training] deferring ScaNN field {}: index={} has {} usable sampled vectors after filtering, selected geometry requires at least {}",
1243 field.0,
1244 index_label,
1245 usable,
1246 minimum_usable,
1247 );
1248 continue;
1249 }
1250 let model = crate::segment::block_in_place_if_multithread(|| {
1251 training_pool.install(|| {
1252 Self::train_field_model(
1253 *field,
1254 config,
1255 &mut sample,
1256 corpus_count,
1257 artifact_generation,
1258 index_label,
1259 )
1260 })
1261 })?;
1262 drop(sample);
1265 updates.push(self.save_trained_field(model).await?);
1266 }
1267 if updates.is_empty() && !fields.is_empty() && !missing.is_empty() {
1268 return Err(Error::Schema(format!(
1269 "cannot train vector centroids: no committed vectors for field(s) {missing:?}"
1270 )));
1271 }
1272 if !missing.is_empty() {
1273 log::info!(
1274 "[vector_training] skipping dense vector field(s) {missing:?}: index={index_label} has no vectors in the current corpus"
1275 );
1276 }
1277 Ok(updates)
1278 }
1279
1280 async fn cleanup_unreferenced_vector_artifacts(&self) {
1285 let referenced = self
1286 .segment_manager
1287 .read_metadata(|metadata| {
1288 metadata
1289 .vector_fields
1290 .values()
1291 .flat_map(|field| {
1292 field
1293 .centroids_file
1294 .iter()
1295 .chain(field.codebook_file.iter())
1296 })
1297 .cloned()
1298 .collect::<std::collections::HashSet<_>>()
1299 })
1300 .await;
1301 let files = match self.directory.list_files(std::path::Path::new("")).await {
1302 Ok(files) => files,
1303 Err(error) => {
1304 log::warn!(
1305 "[trained] index={} failed listing abandoned dense vector artifacts: {error}",
1306 self.schema.index_label()
1307 );
1308 return;
1309 }
1310 };
1311 for path in files {
1312 let Some(name) = path.file_name().and_then(|name| name.to_str()) else {
1313 continue;
1314 };
1315 if !name.starts_with(VECTOR_ARTIFACT_PREFIX)
1316 || referenced.contains(path.to_string_lossy().as_ref())
1317 {
1318 continue;
1319 }
1320 if let Err(error) = self.directory.delete(&path).await
1321 && error.kind() != std::io::ErrorKind::NotFound
1322 {
1323 log::warn!(
1324 "[trained] index={} failed deleting abandoned artifact {path:?}: {error}",
1325 self.schema.index_label()
1326 );
1327 }
1328 }
1329 }
1330
1331 fn reject_ann_fields(ann_fields: &[u32], id_str: &str, field_ids: &[u32]) -> Result<()> {
1336 for &field_id in field_ids {
1337 if ann_fields.binary_search(&field_id).is_ok() {
1338 return Err(Error::Schema(format!(
1339 "metadata-flat field {field_id} already has ANN data in segment {id_str}; \
1340 recreate the index instead of mixing vector generations"
1341 )));
1342 }
1343 }
1344 Ok(())
1345 }
1346
1347 async fn load_training_vectors(
1351 &self,
1352 segment_id: SegmentId,
1353 field_ids: &[u32],
1354 schema: &crate::dsl::Schema,
1355 ) -> Result<crate::segment::reader::loader::VectorsFileData> {
1356 let files = SegmentFiles::new(segment_id.0);
1357 let meta_bytes = self
1358 .directory
1359 .open_read(&files.meta)
1360 .await?
1361 .read_bytes()
1362 .await?;
1363 let meta = SegmentMeta::deserialize(meta_bytes.as_slice())?;
1364 if meta.id != segment_id.0 {
1365 return Err(Error::Corruption(format!(
1366 "segment metadata ID {:032x} does not match file ID {}",
1367 meta.id,
1368 segment_id.to_hex(),
1369 )));
1370 }
1371 crate::segment::reader::loader::load_flat_vectors_file(
1372 self.directory.as_ref(),
1373 &files,
1374 schema,
1375 meta.num_docs,
1376 field_ids,
1377 )
1378 .await
1379 }
1380
1381 fn get_ivf_vector_fields(schema: &crate::dsl::Schema) -> Vec<(Field, IvfFieldConfig)> {
1383 schema
1384 .fields()
1385 .filter_map(|(field, entry)| {
1386 if entry.field_type == FieldType::DenseVector && entry.indexed {
1387 entry
1388 .dense_vector_config
1389 .as_ref()
1390 .filter(|c| c.uses_ivf() || c.index_type == VectorIndexType::Scann)
1393 .map(|c| (field, IvfFieldConfig::Float(c.clone())))
1394 } else if entry.field_type == FieldType::BinaryDenseVector && entry.indexed {
1395 entry
1396 .binary_dense_vector_config
1397 .as_ref()
1398 .filter(|config| {
1399 matches!(
1400 config.index_type,
1401 BinaryIndexType::Ivf | BinaryIndexType::Scann
1402 )
1403 })
1404 .map(|config| (field, IvfFieldConfig::Binary(config.clone())))
1405 } else {
1406 None
1407 }
1408 })
1409 .collect()
1410 }
1411
1412 async fn get_fields_to_build(
1414 &self,
1415 dense_fields: &[(Field, IvfFieldConfig)],
1416 ) -> Vec<(Field, IvfFieldConfig)> {
1417 let field_ids: Vec<u32> = dense_fields.iter().map(|(f, _)| f.0).collect();
1418 let built: Vec<u32> = self
1419 .segment_manager
1420 .read_metadata(|meta| {
1421 field_ids
1422 .iter()
1423 .filter(|fid| meta.is_field_built(**fid))
1424 .copied()
1425 .collect()
1426 })
1427 .await;
1428 dense_fields
1429 .iter()
1430 .filter(|(field, _)| !built.contains(&field.0))
1431 .cloned()
1432 .collect()
1433 }
1434
1435 async fn count_vectors_for_training(
1437 &self,
1438 segment_ids: &[String],
1439 fields_to_build: &[(Field, IvfFieldConfig)],
1440 require_flat_generation: bool,
1441 schema: &crate::dsl::Schema,
1442 ) -> Result<FxHashMap<u32, usize>> {
1443 let mut total_vectors: FxHashMap<u32, usize> = FxHashMap::default();
1444 let field_ids: Vec<u32> = fields_to_build.iter().map(|(field, _)| field.0).collect();
1445
1446 for id_str in segment_ids {
1450 let segment_id = SegmentId::from_hex(id_str)
1451 .ok_or_else(|| Error::Corruption(format!("Invalid segment ID: {}", id_str)))?;
1452 let vectors = self
1453 .load_training_vectors(segment_id, &field_ids, schema)
1454 .await?;
1455
1456 if require_flat_generation {
1457 Self::reject_ann_fields(&vectors.ann_fields, id_str, &field_ids)?;
1458 }
1459
1460 for (field, _) in fields_to_build {
1461 if let Some(flat) = vectors.flat_vectors.get(&field.0) {
1462 let total = total_vectors.entry(field.0).or_default();
1463 *total = total.checked_add(flat.num_vectors).ok_or_else(|| {
1464 Error::Corruption(format!(
1465 "vector count overflows usize for field {}",
1466 field.0,
1467 ))
1468 })?;
1469 }
1470 }
1471 }
1472 Ok(total_vectors)
1473 }
1474
1475 async fn collect_training_sample(
1480 &self,
1481 segment_ids: &[String],
1482 field: Field,
1483 config: &IvfFieldConfig,
1484 total: usize,
1485 schema: &crate::dsl::Schema,
1486 ) -> Result<Option<TrainingSample>> {
1487 if total == 0 {
1488 return Ok(None);
1489 }
1490 let bytes_per_sample = training_sample_bytes(config)?;
1491 let limit = training_sample_limit(
1492 self.config.vector_training_max_samples,
1493 self.config.vector_training_memory_bytes,
1494 bytes_per_sample,
1495 )?;
1496 let mut binary_scann_minimum = None;
1497 let take = match config {
1498 IvfFieldConfig::Float(scann) if scann.index_type == VectorIndexType::Scann => {
1499 let geometry = scann_geometry(scann, total)?;
1500 scann_training_sample_count(
1501 field,
1502 total,
1503 limit,
1504 bytes_per_sample,
1505 &geometry,
1506 false,
1507 )?
1508 }
1509 IvfFieldConfig::Binary(scann) if scann.index_type == BinaryIndexType::Scann => {
1510 let geometry = binary_scann_geometry(scann, total)?;
1511 let take = scann_training_sample_count(
1512 field,
1513 total,
1514 limit,
1515 bytes_per_sample,
1516 &geometry,
1517 true,
1518 )?;
1519 binary_scann_minimum = Some(
1520 usize::try_from(required_scann_training_sample(geometry.num_leaves)?).map_err(
1521 |_| Error::Schema("binary ScaNN minimum sample exceeds usize".into()),
1522 )?,
1523 );
1524 take
1525 }
1526 _ => final_training_sample_count(config, total, limit)?,
1527 };
1528 let sample_seed = 0x4845_524d_4553_4956 ^ field.0 as u64 ^ total as u64;
1529 let candidate_count = binary_scann_minimum.map_or(take, |minimum| {
1530 binary_replenishment_candidate_count(total, take, minimum)
1531 });
1532 let ordinals = deterministic_sample_ordinals(total, candidate_count, sample_seed);
1533
1534 let mut sample = match config {
1535 IvfFieldConfig::Float(config) => TrainingSample::Float(Vec::with_capacity(
1536 take.checked_mul(config.dim)
1537 .ok_or_else(|| Error::Schema("float training sample size overflows".into()))?,
1538 )),
1539 IvfFieldConfig::Binary(_) => TrainingSample::Binary(Vec::with_capacity(
1540 take.checked_mul(bytes_per_sample)
1541 .ok_or_else(|| Error::Schema("binary training sample size overflows".into()))?,
1542 )),
1543 };
1544 let max_read_vectors = (MAX_SAMPLE_READ_BYTES / bytes_per_sample.max(1)).max(1);
1545 let mut zero_codes = 0usize;
1546 let mut ones_codes = 0usize;
1547 let mut surplus_codes = 0usize;
1548 let mut global_offset = 0usize;
1549 let mut cursor = 0usize;
1550 let field_ids = [field.0];
1551
1552 for id_str in segment_ids {
1553 let segment_id = SegmentId::from_hex(id_str)
1554 .ok_or_else(|| Error::Corruption(format!("Invalid segment ID: {id_str}")))?;
1555 let vectors = self
1556 .load_training_vectors(segment_id, &field_ids, schema)
1557 .await?;
1558
1559 let Some(lazy_flat) = vectors.flat_vectors.get(&field.0) else {
1560 continue;
1561 };
1562 let base = global_offset;
1563 let end = base.checked_add(lazy_flat.num_vectors).ok_or_else(|| {
1564 Error::Corruption(format!("vector offset overflows for field {}", field.0))
1565 })?;
1566 global_offset = end;
1567 let first = cursor;
1568 while cursor < ordinals.len() && ordinals[cursor] < end {
1569 cursor += 1;
1570 }
1571 let selected = &ordinals[first..cursor];
1572 let mut run_start = 0;
1573 while run_start < selected.len() {
1574 let mut run_end = run_start + 1;
1575 while run_end < selected.len() {
1576 let selected_count = run_end - run_start + 1;
1577 let span = selected[run_end] - selected[run_start] + 1;
1578 if span > max_read_vectors
1579 || span > selected_count.saturating_mul(MAX_SAMPLE_READ_AMPLIFICATION)
1580 {
1581 break;
1582 }
1583 run_end += 1;
1584 }
1585 let local_start = selected[run_start] - base;
1586 let read_len = selected[run_end - 1] - selected[run_start] + 1;
1587 let bytes = lazy_flat
1588 .read_vectors_batch(local_start, read_len)
1589 .await
1590 .map_err(crate::Error::Io)?;
1591 match &mut sample {
1592 TrainingSample::Binary(codes) => {
1593 let expected = read_len.checked_mul(bytes_per_sample).ok_or_else(|| {
1594 Error::Corruption("binary sample read size overflows".into())
1595 })?;
1596 if bytes.len() != expected {
1597 return Err(Error::Corruption(format!(
1598 "binary sample read returned {} bytes, expected {expected}",
1599 bytes.len(),
1600 )));
1601 }
1602 for &ordinal in &selected[run_start..run_end] {
1603 let relative = ordinal - selected[run_start];
1604 let offset = relative * bytes_per_sample;
1605 let code = &bytes.as_slice()[offset..offset + bytes_per_sample];
1606 if code.iter().all(|&byte| byte == 0) {
1616 zero_codes += 1;
1617 continue;
1618 }
1619 if code.iter().all(|&byte| byte == 0xff) {
1620 ones_codes += 1;
1621 continue;
1622 }
1623 if codes.len() / bytes_per_sample < take {
1624 codes.extend_from_slice(code);
1625 } else {
1626 surplus_codes += 1;
1627 }
1628 }
1629 }
1630 TrainingSample::Float(values) => {
1631 let dim = lazy_flat.dim;
1632 let float_count = read_len.checked_mul(dim).ok_or_else(|| {
1633 Error::Corruption("float sample read size overflows".into())
1634 })?;
1635 let mut decoded = vec![0.0; float_count];
1636 crate::segment::dequantize_raw(
1637 bytes.as_slice(),
1638 lazy_flat.quantization,
1639 decoded.len(),
1640 &mut decoded,
1641 )
1642 .map_err(crate::Error::Io)?;
1643 for &ordinal in &selected[run_start..run_end] {
1644 let relative = ordinal - selected[run_start];
1645 let offset = relative * dim;
1646 values.extend_from_slice(&decoded[offset..offset + dim]);
1647 }
1648 }
1649 }
1650 run_start = run_end;
1651 }
1652 }
1653
1654 let collected = sample.len(config.dim());
1655 if global_offset != total
1658 || cursor != candidate_count
1659 || collected + zero_codes + ones_codes + surplus_codes != candidate_count
1660 {
1661 return Err(Error::Corruption(format!(
1662 "training sample coverage mismatch for field {}: counted={total}, traversed={global_offset}, selected={cursor}, collected={collected}, zero={zero_codes}, ones={ones_codes}, surplus={surplus_codes}",
1663 field.0,
1664 )));
1665 }
1666 if let Some(minimum) = binary_scann_minimum
1667 && collected < minimum
1668 {
1669 return Err(Error::Schema(format!(
1670 "binary ScaNN field {} has only {collected} usable training vectors after excluding {zero_codes} all-zero and {ones_codes} all-ones codes from {candidate_count} deterministic candidates; selected geometry requires at least {minimum}. Fix the binary embedding producer/data or choose a smaller geometry",
1671 field.0,
1672 )));
1673 }
1674 if zero_codes > 0 {
1675 log::warn!(
1676 "[vector_training] index={} field={}: {zero_codes} of {candidate_count} sampled candidates \
1677 ({:.1}%) are all-zero and were excluded from training — they cannot be assigned \
1678 to any leaf, so training on them only wastes centroids",
1679 self.schema.index_label(),
1680 field.0,
1681 100.0 * zero_codes as f64 / candidate_count.max(1) as f64,
1682 );
1683 }
1684 if ones_codes > 0 {
1685 log::warn!(
1686 "[vector_training] index={} field={}: {ones_codes} of {candidate_count} sampled candidates \
1687 ({:.1}%) are all-ones and were excluded from training — training on the \
1688 saturated constant only dedicates centroids to a producer bug",
1689 self.schema.index_label(),
1690 field.0,
1691 100.0 * ones_codes as f64 / candidate_count.max(1) as f64,
1692 );
1693 }
1694 if collected == 0 {
1695 log::warn!(
1696 "[vector_training] index={} field={}: every sampled vector is degenerate \
1697 (all-zero or all-ones); skipping ANN training for this field",
1698 self.schema.index_label(),
1699 field.0,
1700 );
1701 return Ok(None);
1702 }
1703 if collected < total {
1704 log::info!(
1705 "[vector_training] sampled {} / {} dense vectors: index={} field={} (max {} vectors / {} resident)",
1706 collected,
1707 total,
1708 self.schema.index_label(),
1709 field.0,
1710 self.config.vector_training_max_samples,
1711 crate::format_bytes(self.config.vector_training_memory_bytes as u64),
1712 );
1713 }
1714 Ok(Some(sample))
1715 }
1716
1717 fn train_field_model(
1720 field: Field,
1721 config: &IvfFieldConfig,
1722 sample: &mut TrainingSample,
1723 corpus_count: usize,
1724 artifact_generation: &str,
1725 index_label: &str,
1726 ) -> Result<TrainedFieldModel> {
1727 let field_id = field.0;
1728 let dim = config.dim();
1729 let sample_count = sample.len(dim);
1730 if sample_count == 0 || corpus_count == 0 {
1731 return Err(Error::Internal(format!(
1732 "empty training sample for non-empty field {field_id}"
1733 )));
1734 }
1735 let num_clusters = match config {
1736 IvfFieldConfig::Float(config) if config.index_type == VectorIndexType::Scann => {
1737 scann_geometry(config, corpus_count)?.num_leaves as usize
1738 }
1739 IvfFieldConfig::Binary(config) if config.index_type == BinaryIndexType::Scann => {
1740 binary_scann_geometry(config, corpus_count)?.num_leaves as usize
1741 }
1742 _ => effective_field_num_clusters(config, corpus_count, sample_count)?,
1743 };
1744
1745 log::info!(
1746 "[vector_training] training model: index={} field={} with {} sampled / {} total vectors, {} clusters (dim={})",
1747 index_label,
1748 field_id,
1749 sample_count,
1750 corpus_count,
1751 num_clusters,
1752 dim,
1753 );
1754
1755 let centroids_filename =
1756 format!("{VECTOR_ARTIFACT_PREFIX}{artifact_generation}_field_{field_id}_centroids.bin");
1757
1758 let artifacts = match (config, sample) {
1759 (IvfFieldConfig::Float(config), TrainingSample::Float(values))
1760 if config.index_type == VectorIndexType::IvfTq =>
1761 {
1762 values
1763 .chunks_exact_mut(dim)
1764 .for_each(crate::structures::vector::ivf::routing::normalize_cosine_in_place);
1765 let candidate_validation_count =
1766 validation_sample_count(sample_count, num_clusters, dim);
1767 let candidate_training_count = sample_count - candidate_validation_count;
1768 let seeds = model_selection_seeds(
1769 candidate_training_count,
1770 num_clusters,
1771 dim,
1772 candidate_validation_count > 0,
1773 );
1774 let split_seed =
1775 MODEL_SELECTION_SEEDS[0] ^ u64::from(field_id) ^ corpus_count as u64;
1776 let split = if seeds.len() > 1 {
1777 partition_contiguous_holdout_suffix(
1778 values.as_mut_slice(),
1779 dim,
1780 candidate_validation_count,
1781 split_seed,
1782 )
1783 } else {
1784 values.len()
1785 };
1786 let (training_values, validation) = values.as_slice().split_at(split);
1787 let training_count = training_values.len() / dim;
1788 let validation_count = validation.len() / dim;
1789 if seeds.len() > 1 {
1790 log::info!(
1791 "[vector_training] model selection: index={index_label} field={field_id}, {} training + {} held-out vectors, {} deterministic centroid seed(s)",
1792 training_count,
1793 validation_count,
1794 seeds.len(),
1795 );
1796 } else {
1797 log::info!(
1798 "[vector_training] model selection: index={index_label} field={field_id}, all {} sampled vectors with one deterministic \
1799 centroid seed; model-selection holdout disabled",
1800 training_count,
1801 );
1802 }
1803
1804 let mut base_config = crate::structures::CoarseConfig::new(dim, num_clusters)
1805 .with_routing(config.ivf_routing);
1806 if let Some(soar) = config.soar.clone() {
1807 base_config = base_config.with_soar(soar);
1808 }
1809
1810 let mut selected: Option<(
1811 crate::structures::CoarseCentroids,
1812 Option<FloatBuildQuality>,
1813 u64,
1814 )> = None;
1815 for &seed in seeds {
1816 let candidate = crate::structures::CoarseCentroids::train_contiguous(
1817 &base_config.clone().with_seed(seed),
1818 training_values,
1819 training_count,
1820 index_label,
1821 );
1822 let quality =
1823 evaluate_float_build_quality(&candidate, validation, config.ivf_routing);
1824 if let Some(quality) = quality {
1825 log::info!(
1826 "[vector_training] IVF candidate: index={index_label} field={field_id} seed={seed}, objective={:.6}, \
1827 exact/routed_mean_distortion={:.6}/{:.6}, router_recall@1={:.4}, \
1828 construction_postings/vector={:.3}, \
1829 residual_scale[p50/p95/p99]={:.4}/{:.4}/{:.4}, \
1830 construction_occupancy[p95/p99/max/empty]={}/{}/{}/{}",
1831 quality.objective,
1832 quality.mean_exact_distortion,
1833 quality.mean_routed_distortion,
1834 quality.router_recall_at_1,
1835 quality.mean_construction_assignments,
1836 quality.residual_p50,
1837 quality.residual_p95,
1838 quality.residual_p99,
1839 quality.occupancy.p95,
1840 quality.occupancy.p99,
1841 quality.occupancy.max,
1842 quality.occupancy.empty,
1843 );
1844 }
1845 let replace = selected.as_ref().is_none_or(|(_, best, _)| {
1846 quality
1847 .map(|quality| quality.objective)
1848 .unwrap_or(f64::INFINITY)
1849 .total_cmp(
1850 &best
1851 .map(|quality| quality.objective)
1852 .unwrap_or(f64::INFINITY),
1853 )
1854 .is_lt()
1855 });
1856 if replace {
1857 selected = Some((candidate, quality, seed));
1858 }
1859 }
1860 let (mut centroids, quality, seed) =
1861 selected.expect("the fixed centroid seed bank is non-empty");
1862 if let Some(quality) = quality {
1863 log::info!(
1864 "[vector_training] selected IVF seed: index={index_label} field={field_id} seed={seed}, held-out objective {:.6} \
1865 (occupancy penalty {:.4})",
1866 quality.objective,
1867 quality.occupancy.penalty,
1868 );
1869 } else {
1870 log::info!(
1871 "[vector_training] selected IVF seed: index={index_label} field={field_id} seed={seed}, without a model-selection \
1872 holdout",
1873 );
1874 }
1875 centroids.version =
1876 crate::structures::mark_ivf_tq_cosine_generation(centroids.version);
1877 TrainedFieldArtifacts::FloatCentroids(centroids)
1878 }
1879 (IvfFieldConfig::Float(config), TrainingSample::Float(values))
1880 if config.index_type == VectorIndexType::Scann =>
1881 {
1882 if config.soar.is_some() {
1883 return Err(Error::Schema(format!(
1884 "ScaNN field {field_id} enables SOAR, but ScaNN SOAR secondary assignments are not implemented; set soar: null before building"
1885 )));
1886 }
1887 values
1888 .chunks_exact_mut(dim)
1889 .for_each(crate::structures::vector::ivf::routing::normalize_cosine_in_place);
1890 let geometry = scann_geometry(config, corpus_count)?;
1891 let dimensions_per_block = 2usize.min(dim);
1892 let seed = MODEL_SELECTION_SEEDS[0] ^ u64::from(field_id) ^ corpus_count as u64;
1893 let generation = u64::from_str_radix(
1894 &artifact_generation[artifact_generation.len().saturating_sub(16)..],
1895 16,
1896 )
1897 .unwrap_or(seed)
1898 .max(1);
1899 let model = crate::structures::vector::scann::FloatScannModel::train_model(
1900 values,
1901 sample_count,
1902 dim,
1903 &geometry.level_counts,
1904 dimensions_per_block,
1905 25,
1906 seed,
1907 crate::structures::vector::scann::DEFAULT_ANISOTROPIC_THRESHOLD,
1908 )
1909 .map_err(|error| Error::Internal(format!("ScaNN training failed: {error}")))?;
1910 let artifact = crate::structures::vector::scann::ScannTrainedArtifact::new(
1911 generation,
1912 sample_count as u64,
1913 crate::structures::vector::scann::ScannConfig {
1914 dimension: dim as u32,
1915 tree_levels: geometry.centroid_levels,
1916 num_leaves: geometry.num_leaves,
1917 encoding: crate::structures::vector::scann::ScannEncoding::AsymmetricHash {
1918 dimensions_per_block: dimensions_per_block as u16,
1919 bits_per_code: 4,
1920 },
1921 },
1922 model.routing.to_quantized_levels(),
1923 Some(model.codebook.to_artifact()),
1924 )
1925 .map_err(|error| Error::Internal(format!("ScaNN artifact failed: {error}")))?;
1926 TrainedFieldArtifacts::Scann(artifact)
1927 }
1928 (IvfFieldConfig::Binary(config), TrainingSample::Binary(codes))
1929 if config.index_type == BinaryIndexType::Ivf =>
1930 {
1931 let byte_len = dim.div_ceil(8);
1932 let training_count = codes.len() / byte_len;
1933 let mut binary_config = crate::structures::BinaryIvfConfig::new(dim, num_clusters);
1934 binary_config.max_train_samples = training_count;
1935 binary_config.routing = config.ivf_routing;
1936 let quantizer = crate::structures::BinaryCoarseQuantizer::train(
1937 binary_config,
1938 codes,
1939 training_count,
1940 index_label,
1941 )
1942 .map_err(Error::Io)?;
1943 TrainedFieldArtifacts::Binary(quantizer)
1944 }
1945 (IvfFieldConfig::Binary(config), TrainingSample::Binary(codes))
1946 if config.index_type == BinaryIndexType::Scann =>
1947 {
1948 let geometry = binary_scann_geometry(config, corpus_count)?;
1949 let seed = MODEL_SELECTION_SEEDS[0] ^ u64::from(field_id) ^ corpus_count as u64;
1950 let generation = u64::from_str_radix(
1951 &artifact_generation[artifact_generation.len().saturating_sub(16)..],
1952 16,
1953 )
1954 .unwrap_or(seed)
1955 .max(1);
1956 let training = crate::structures::vector::scann::BinaryScannTraining {
1957 dim_bits: config.dim as u32,
1958 geometry,
1959 train_iters: 25,
1960 seed,
1961 };
1962 let model = crate::structures::vector::scann::BinaryScannModel::train(
1963 &training,
1964 codes,
1965 sample_count,
1966 index_label,
1967 )
1968 .map_err(|error| {
1969 Error::Internal(format!("binary ScaNN training failed: {error}"))
1970 })?;
1971 TrainedFieldArtifacts::Scann(
1972 model
1973 .to_artifact(generation, sample_count as u64)
1974 .map_err(|error| {
1975 Error::Internal(format!("binary ScaNN artifact failed: {error}"))
1976 })?,
1977 )
1978 }
1979 _ => {
1980 return Err(Error::Internal(format!(
1981 "training sample kind does not match field {field_id}"
1982 )));
1983 }
1984 };
1985
1986 let actual_num_clusters = match &artifacts {
1987 TrainedFieldArtifacts::FloatCentroids(centroids) => centroids.num_clusters as usize,
1988 TrainedFieldArtifacts::Binary(quantizer) => quantizer.num_clusters as usize,
1989 TrainedFieldArtifacts::Scann(artifact) => artifact.config.num_leaves as usize,
1990 };
1991 let (scann_generation, scann_artifact_id) = match &artifacts {
1992 TrainedFieldArtifacts::Scann(artifact) => {
1993 (Some(artifact.generation), Some(artifact.artifact_id))
1994 }
1995 _ => (None, None),
1996 };
1997 Ok(TrainedFieldModel {
1998 update: TrainedFieldUpdate {
1999 field_id,
2000 index_type: config.index_type(),
2001 vector_count: corpus_count,
2002 num_clusters: actual_num_clusters,
2003 centroids_file: centroids_filename,
2004 codebook_file: None,
2005 scann_generation,
2006 scann_artifact_id,
2007 },
2008 artifacts,
2009 })
2010 }
2011
2012 async fn save_trained_field(&self, model: TrainedFieldModel) -> Result<TrainedFieldUpdate> {
2013 let TrainedFieldModel { update, artifacts } = model;
2014 match artifacts {
2015 TrainedFieldArtifacts::FloatCentroids(centroids) => {
2016 self.save_trained_artifact(¢roids, &update.centroids_file)
2017 .await?;
2018 log::info!(
2019 "[vector_training] saved IVF-TQ coarse artifact: index={} field={} ({} clusters; leaf codec is derived)",
2020 self.schema.index_label(),
2021 update.field_id,
2022 centroids.num_clusters,
2023 );
2024 }
2025 TrainedFieldArtifacts::Binary(quantizer) => {
2026 self.save_trained_artifact(&quantizer, &update.centroids_file)
2027 .await?;
2028 log::info!(
2029 "[vector_training] saved binary IVF artifact: index={} field={} ({} clusters)",
2030 self.schema.index_label(),
2031 update.field_id,
2032 quantizer.num_clusters,
2033 );
2034 }
2035 TrainedFieldArtifacts::Scann(artifact) => {
2036 self.save_scann_artifact(&artifact, &update.centroids_file)
2037 .await?;
2038 log::info!(
2039 "[vector_training] saved ScaNN artifact: index={} field={} ({} leaves, {} levels, generation={})",
2040 self.schema.index_label(),
2041 update.field_id,
2042 artifact.config.num_leaves,
2043 artifact.config.tree_levels,
2044 artifact.generation,
2045 );
2046 }
2047 }
2048 Ok(update)
2049 }
2050
2051 async fn save_scann_artifact(
2052 &self,
2053 artifact: &crate::structures::vector::scann::ScannTrainedArtifact,
2054 filename: &str,
2055 ) -> Result<()> {
2056 let temp_filename = format!("{filename}.tmp");
2057 let temp_path = std::path::Path::new(&temp_filename);
2058 let final_path = std::path::Path::new(filename);
2059 let mut writer = self.directory.streaming_writer(temp_path).await?;
2060 artifact
2061 .write_to(&mut writer)
2062 .map_err(|error| Error::Serialization(error.to_string()))?;
2063 writer.finish()?;
2064 if let Err(error) = self.directory.rename(temp_path, final_path).await {
2065 let _ = self.directory.delete(temp_path).await;
2066 return Err(Error::Io(error));
2067 }
2068 self.directory.sync().await?;
2069 Ok(())
2070 }
2071
2072 async fn save_trained_artifact(
2074 &self,
2075 artifact: &impl serde::Serialize,
2076 filename: &str,
2077 ) -> Result<()> {
2078 let temp_filename = format!("{filename}.tmp");
2079 let temp_path = std::path::Path::new(&temp_filename);
2080 let final_path = std::path::Path::new(filename);
2081 let mut writer = self.directory.streaming_writer(temp_path).await?;
2082 let encode_result = {
2083 let mut limited = SizeLimitedWriter::new(
2084 writer.as_mut(),
2085 super::metadata::MAX_TRAINED_ARTIFACT_BYTES,
2086 );
2087 bincode::serde::encode_into_std_write(
2088 artifact,
2089 &mut limited,
2090 bincode::config::standard(),
2091 )
2092 };
2093 if let Err(error) = encode_result {
2094 drop(writer);
2095 let _ = self.directory.delete(temp_path).await;
2096 return Err(Error::Serialization(format!(
2097 "failed to serialize trained artifact '{filename}': {error}"
2098 )));
2099 }
2100 if let Err(error) = writer.finish() {
2101 let _ = self.directory.delete(temp_path).await;
2102 return Err(Error::Io(error));
2103 }
2104 if let Err(error) = self.directory.rename(temp_path, final_path).await {
2105 let _ = self.directory.delete(temp_path).await;
2106 return Err(Error::Io(error));
2107 }
2108 self.directory.sync().await?;
2109 Ok(())
2110 }
2111}
2112
2113#[cfg(test)]
2114mod tests {
2115 use super::*;
2116
2117 fn ivf_config(num_clusters: Option<usize>) -> DenseVectorConfig {
2118 DenseVectorConfig::ivf_tq(8, num_clusters, 4)
2119 }
2120
2121 fn scann_config(dim: usize, leaves: Option<usize>) -> DenseVectorConfig {
2122 DenseVectorConfig {
2123 dim,
2124 index_type: VectorIndexType::Scann,
2125 quantization: crate::dsl::DenseVectorQuantization::F32,
2126 num_clusters: leaves,
2127 target_vectors: None,
2128 tree_levels: Some(1),
2129 ivf_routing: crate::dsl::IvfRoutingMode::Auto,
2130 nprobe: 1,
2131 unit_norm: true,
2132 soar: None,
2133 }
2134 }
2135
2136 #[test]
2137 fn billion_scale_scann_autopilot_preserves_geometry_across_builder_budgets() {
2138 let defaults = crate::IndexConfig::default();
2139
2140 let mut float = scann_config(1_024, None);
2141 float.tree_levels = None;
2142 let float = IvfFieldConfig::Float(float);
2143 let float_limit = training_sample_limit(
2144 defaults.vector_training_max_samples,
2145 defaults.vector_training_memory_bytes,
2146 training_sample_bytes(&float).unwrap(),
2147 )
2148 .unwrap();
2149 let float_geometry = match &float {
2150 IvfFieldConfig::Float(config) => scann_geometry(config, 1_000_000_000).unwrap(),
2151 _ => unreachable!(),
2152 };
2153 assert_eq!(float_limit, 1_048_576);
2154 assert_eq!(float_geometry.level_counts, [1_000, 1_000_000]);
2155 let float_error = scann_training_sample_count(
2156 Field(7),
2157 1_000_000_000,
2158 float_limit,
2159 training_sample_bytes(&float).unwrap(),
2160 &float_geometry,
2161 false,
2162 )
2163 .unwrap_err()
2164 .to_string();
2165 assert!(float_error.contains("1000000 leaves"));
2166 assert!(float_error.contains("8000000 sampled vectors"));
2167
2168 let mut binary = BinaryDenseVectorConfig::new(2_560);
2169 binary.index_type = BinaryIndexType::Scann;
2170 let binary = IvfFieldConfig::Binary(binary);
2171 let binary_limit = training_sample_limit(
2172 defaults.vector_training_max_samples,
2173 defaults.vector_training_memory_bytes,
2174 training_sample_bytes(&binary).unwrap(),
2175 )
2176 .unwrap();
2177 let binary_geometry = match &binary {
2178 IvfFieldConfig::Binary(config) => binary_scann_geometry(config, 1_000_000_000).unwrap(),
2179 _ => unreachable!(),
2180 };
2181 assert_eq!(binary_limit, 10_000_000);
2182 assert_eq!(binary_geometry.level_counts, [1_000, 1_000_000]);
2183 assert!(
2184 required_scann_training_sample(binary_geometry.num_leaves).unwrap()
2185 <= binary_limit as u64
2186 );
2187 }
2188
2189 #[test]
2190 fn explicit_scann_leaves_have_stable_automatic_depth() {
2191 for leaves in [31_622, 31_623] {
2192 let mut float = scann_config(1_024, Some(leaves));
2193 float.tree_levels = None;
2194 assert_eq!(
2195 scann_geometry(&float, 15_000_000).unwrap().level_counts,
2196 [178, leaves as u32]
2197 );
2198
2199 let mut binary = BinaryDenseVectorConfig::new(2_560);
2200 binary.index_type = BinaryIndexType::Scann;
2201 binary.num_clusters = Some(leaves);
2202 assert_eq!(
2203 binary_scann_geometry(&binary, 15_000_000)
2204 .unwrap()
2205 .level_counts,
2206 [178, leaves as u32]
2207 );
2208 }
2209
2210 let mut float = scann_config(1_024, Some(1_000_000));
2211 float.tree_levels = None;
2212 assert_eq!(
2213 scann_geometry(&float, 15_000_000).unwrap().level_counts,
2214 [1_000, 1_000_000]
2215 );
2216 }
2217
2218 #[test]
2219 fn target_vectors_selects_future_scann_topology_without_faking_readiness() {
2220 let mut binary = BinaryDenseVectorConfig::new(2_560);
2221 binary.index_type = BinaryIndexType::Scann;
2222 binary.target_vectors = Some(1_000_000_000);
2223
2224 let geometry = binary_scann_geometry(&binary, 1_000_000).unwrap();
2225 assert_eq!(geometry.level_counts, [1_000, 1_000_000]);
2226 assert_eq!(
2227 required_scann_training_sample(geometry.num_leaves).unwrap(),
2228 8_000_000
2229 );
2230 let below_floor =
2231 scann_training_sample_count(Field(7), 7_999_999, 8_000_000, 320, &geometry, true)
2232 .unwrap_err()
2233 .to_string();
2234 assert!(below_floor.contains("needs at least 8000000 sampled vectors"));
2235 assert!(below_floor.contains("builder limits allow 7999999"));
2236 assert_eq!(
2237 scann_training_sample_count(Field(7), 8_000_000, 8_000_000, 320, &geometry, true,)
2238 .unwrap(),
2239 8_000_000
2240 );
2241
2242 let mut explicit = binary.clone();
2243 explicit.num_clusters = Some(4_096);
2244 assert_eq!(
2245 binary_scann_geometry(&explicit, 1_000_000)
2246 .unwrap()
2247 .num_leaves,
2248 4_096,
2249 "explicit leaves override target-sized automatic geometry"
2250 );
2251
2252 let mut unhinted = binary.clone();
2253 unhinted.target_vectors = None;
2254 assert_eq!(
2255 binary_scann_geometry(&binary, 2_000_000_000).unwrap(),
2256 binary_scann_geometry(&unhinted, 2_000_000_000).unwrap(),
2257 "target_vectors is a lower bound and cannot shrink live-corpus geometry"
2258 );
2259 }
2260
2261 #[test]
2262 fn target_sized_ivf_geometry_waits_for_its_hardcoded_sample_floor() {
2263 let binary = IvfFieldConfig::Binary(
2264 BinaryDenseVectorConfig::new(256).with_target_vectors(1_000_000_000),
2265 );
2266 let leaves = 31_623;
2267 let required = leaves * MIN_TRAINING_POINTS_PER_CENTROID;
2268
2269 let error = effective_field_num_clusters(&binary, 1_000_000, required - 1)
2270 .unwrap_err()
2271 .to_string();
2272 assert!(error.contains("target-sized IVF geometry"), "{error}");
2273 assert!(error.contains(&required.to_string()), "{error}");
2274 assert_eq!(
2275 effective_field_num_clusters(&binary, 1_000_000, required).unwrap(),
2276 leaves
2277 );
2278 }
2279
2280 #[test]
2281 fn target_vector_alters_rebuild_only_when_the_hint_controls_geometry() {
2282 let binary_auto = IvfFieldConfig::Binary(
2283 BinaryDenseVectorConfig::new(256).with_target_vectors(15_000_000),
2284 );
2285 let binary_auto_changed = IvfFieldConfig::Binary(
2286 BinaryDenseVectorConfig::new(256).with_target_vectors(1_000_000_000),
2287 );
2288 assert!(alter_requires_rebuild(&binary_auto, &binary_auto_changed));
2289
2290 let binary_explicit = IvfFieldConfig::Binary(
2291 BinaryDenseVectorConfig::new(256)
2292 .with_target_vectors(15_000_000)
2293 .with_ivf(Some(4_096), 64),
2294 );
2295 let binary_explicit_changed = IvfFieldConfig::Binary(
2296 BinaryDenseVectorConfig::new(256)
2297 .with_target_vectors(1_000_000_000)
2298 .with_ivf(Some(4_096), 64),
2299 );
2300 assert!(!alter_requires_rebuild(
2301 &binary_explicit,
2302 &binary_explicit_changed
2303 ));
2304 assert!(alter_requires_rebuild(
2305 &binary_explicit_changed,
2306 &binary_auto_changed
2307 ));
2308
2309 let float_explicit = IvfFieldConfig::Float(
2310 DenseVectorConfig::ivf_tq(128, Some(4_096), 64).with_target_vectors(15_000_000),
2311 );
2312 let float_explicit_changed = IvfFieldConfig::Float(
2313 DenseVectorConfig::ivf_tq(128, Some(4_096), 64).with_target_vectors(1_000_000_000),
2314 );
2315 assert!(!alter_requires_rebuild(
2316 &float_explicit,
2317 &float_explicit_changed
2318 ));
2319 }
2320
2321 #[test]
2322 fn explicit_scann_geometry_rejects_an_inadequate_builder_sample() {
2323 let geometry = crate::structures::vector::scann::geometry_for_leaves(20_000, 1).unwrap();
2324 let error =
2325 scann_training_sample_count(Field(7), 1_000_000, 159_999, 4_096, &geometry, false)
2326 .unwrap_err()
2327 .to_string();
2328 assert!(error.contains("needs at least 160000 sampled vectors"));
2329 assert!(error.contains("8 samples/leaf"));
2330 assert!(error.contains("builder limits allow 159999"));
2331 }
2332
2333 #[test]
2334 fn scann_build_rejects_soar_instead_of_silently_ignoring_it() {
2335 let mut config = scann_config(2, Some(2));
2336 config.soar = Some(crate::structures::SoarConfig::default());
2337 let mut sample = TrainingSample::Float(vec![0.5; 100_000 * 2]);
2338 let result = IndexWriter::<crate::directories::RamDirectory>::train_field_model(
2339 Field(7),
2340 &IvfFieldConfig::Float(config),
2341 &mut sample,
2342 100_000,
2343 "00000000000000000000000000000007",
2344 "test",
2345 );
2346 let error = match result {
2347 Err(error) => error,
2348 Ok(_) => panic!("ScaNN SOAR must fail before training"),
2349 };
2350 assert!(error.to_string().contains("SOAR"));
2351 assert!(error.to_string().contains("soar: null"));
2352 }
2353
2354 #[test]
2355 fn effective_clusters_follow_corpus_heuristic_but_fit_sample() {
2356 let config = ivf_config(None);
2357
2358 assert_eq!(
2359 effective_ivf_num_clusters(&config, 1_000_000, 73).unwrap(),
2360 16
2361 );
2362 assert_eq!(
2363 effective_ivf_num_clusters(&config, 10_000, 1_000).unwrap(),
2364 25
2365 );
2366 }
2367
2368 #[test]
2369 fn effective_clusters_clamp_explicit_value_to_sample() {
2370 let config = ivf_config(Some(256));
2371 assert_eq!(
2372 effective_ivf_num_clusters(&config, 1_000_000, 17).unwrap(),
2373 17
2374 );
2375 }
2376
2377 #[test]
2378 fn effective_clusters_reject_invalid_explicit_bounds() {
2379 let zero = effective_ivf_num_clusters(&ivf_config(Some(0)), 10_000, 100)
2380 .unwrap_err()
2381 .to_string();
2382 assert!(zero.contains("at least 1"));
2383
2384 let too_many =
2385 effective_ivf_num_clusters(&ivf_config(Some(MAX_IVF_CLUSTERS + 1)), 10_000, 100)
2386 .unwrap_err()
2387 .to_string();
2388 assert!(too_many.contains("must not exceed 1048576"));
2389 }
2390
2391 #[test]
2392 fn effective_clusters_reject_empty_training_sample() {
2393 let error = effective_ivf_num_clusters(&ivf_config(None), 10_000, 0)
2394 .unwrap_err()
2395 .to_string();
2396 assert!(error.contains("without sample vectors"));
2397 }
2398
2399 #[test]
2400 fn training_sample_limit_honors_both_cli_bounds() {
2401 assert_eq!(training_sample_limit(10_000_000, 4_096, 4).unwrap(), 1_024);
2402 assert_eq!(training_sample_limit(100, 4_096, 4).unwrap(), 100);
2403 let error = training_sample_limit(100, 3, 4).unwrap_err().to_string();
2404 assert!(error.contains("cannot hold one"), "{error}");
2405 }
2406
2407 #[test]
2408 fn final_sample_is_selected_at_the_points_per_centroid_ceiling() {
2409 let config = IvfFieldConfig::Float(DenseVectorConfig::ivf_tq(8, Some(4), 1));
2410 assert_eq!(
2411 final_training_sample_count(&config, 10_000, 10_000).unwrap(),
2412 4 * COARSE_TRAINING_POINTS_PER_CENTROID,
2413 );
2414 assert_eq!(
2415 final_training_sample_count(&config, 10_000, 512).unwrap(),
2416 512,
2417 );
2418 }
2419
2420 #[test]
2421 fn holdout_preserves_the_training_points_per_centroid_floor() {
2422 assert_eq!(validation_sample_count(390, 10, 8), 0);
2423 assert_eq!(validation_sample_count(391, 10, 8), 1);
2424
2425 let config = IvfFieldConfig::Float(ivf_config(None));
2426 let sample_count = 1_000;
2427 let clusters = effective_field_num_clusters(&config, 10_000, sample_count).unwrap();
2428 let held_out = validation_sample_count(sample_count, clusters, config.dim());
2429 assert!(
2430 sample_count - held_out >= clusters.saturating_mul(MIN_TRAINING_POINTS_PER_CENTROID)
2431 );
2432 }
2433
2434 #[test]
2435 fn deterministic_point_sample_is_sorted_unique_and_repeatable() {
2436 let first = deterministic_sample_ordinals(10_000, 1_000, 7);
2437 let repeated = deterministic_sample_ordinals(10_000, 1_000, 7);
2438 let other_seed = deterministic_sample_ordinals(10_000, 1_000, 8);
2439 assert_eq!(first, repeated);
2440 assert_ne!(first, other_seed);
2441 assert_eq!(first.len(), 1_000);
2442 assert!(first.windows(2).all(|pair| pair[0] < pair[1]));
2443 assert!(first.iter().all(|&ordinal| ordinal < 10_000));
2444 }
2445
2446 #[test]
2447 fn binary_scann_quality_filter_has_deterministic_replenishment_capacity() {
2448 assert_eq!(binary_replenishment_candidate_count(101, 100, 100), 101);
2453 assert_eq!(
2454 binary_replenishment_candidate_count(10_000, 100, 100),
2455 1_124
2456 );
2457 assert_eq!(
2458 binary_replenishment_candidate_count(20_000_000, 10_000_000, 4_000_000),
2459 11_000_000
2460 );
2461 }
2462
2463 #[test]
2464 fn model_selection_accounts_for_initialization_and_all_lloyd_passes() {
2465 assert_eq!(model_selection_seeds(1_000, 16, 8, true).len(), 2);
2466 assert_eq!(model_selection_seeds(1_000, 16, 8, false).len(), 1);
2467 assert_eq!(model_selection_seeds(1_800, 1_000, 100, true).len(), 1);
2470 assert_eq!(
2471 model_selection_seeds(MAX_MULTI_SEED_COORDINATE_WORK, 2, 1, true).len(),
2472 1,
2473 );
2474 }
2475
2476 #[test]
2477 fn flat_float_sample_partitions_a_deterministic_holdout_suffix() {
2478 let original: Vec<f32> = (0..20).map(|value| value as f32).collect();
2479 let mut first = original.clone();
2480 let mut repeated = original.clone();
2481 let validation_count = validation_sample_count(10, 2, 2);
2482 let first_split = partition_contiguous_holdout_suffix(&mut first, 2, validation_count, 11);
2483 let repeated_split =
2484 partition_contiguous_holdout_suffix(&mut repeated, 2, validation_count, 11);
2485
2486 assert_eq!(first, repeated);
2487 assert_eq!(first_split, repeated_split);
2488 assert_eq!(first_split, 18);
2489 assert_eq!(first.len(), original.len());
2490 assert_eq!(first[first_split..].len(), 2);
2491
2492 let mut first_components: Vec<u32> =
2493 first.chunks_exact(2).map(|row| row[0] as u32).collect();
2494 first_components.sort_unstable();
2495 assert_eq!(first_components, vec![0, 2, 4, 6, 8, 10, 12, 14, 16, 18]);
2496 assert!(first.chunks_exact(2).all(|row| row[1] == row[0] + 1.0));
2497 }
2498
2499 #[test]
2500 fn occupancy_report_exposes_tail_and_empty_cells() {
2501 let report = occupancy_quality(vec![0, 1, 2, 7], 10);
2502 assert_eq!(report.p95, 7);
2503 assert_eq!(report.p99, 7);
2504 assert_eq!(report.max, 7);
2505 assert_eq!(report.empty, 1);
2506 assert!(report.penalty > 0.0);
2507 }
2508
2509 #[test]
2510 fn model_selection_objective_prices_routed_distortion_excess() {
2511 let objective = float_model_selection_objective(2.0, 3.0, 0.1);
2512 assert!((objective - 3.2).abs() < f64::EPSILON);
2513 assert_eq!(float_model_selection_objective(2.0, 1.5, 0.0), 2.0);
2514 }
2515
2516 #[test]
2517 fn float_quality_reports_exact_query_router_recall() {
2518 let training = [0.0, 0.0, 0.1, 0.0, 10.0, 10.0, 10.1, 10.0];
2519 let config = crate::structures::CoarseConfig::new(2, 2);
2520 let centroids = crate::structures::CoarseCentroids::train_contiguous(
2521 &config,
2522 &training,
2523 4,
2524 "test-index",
2525 );
2526 for routing in [
2527 crate::dsl::IvfRoutingMode::Flat,
2528 crate::dsl::IvfRoutingMode::Auto,
2529 ] {
2530 let quality =
2531 evaluate_float_build_quality(¢roids, &[0.05, 0.0, 10.05, 10.0], routing)
2532 .unwrap();
2533
2534 assert_eq!(quality.router_recall_at_1, 1.0);
2535 assert!(
2536 (quality.mean_exact_distortion - quality.mean_routed_distortion).abs()
2537 < f64::EPSILON
2538 );
2539 assert_eq!(quality.mean_construction_assignments, 1.0);
2540 assert_eq!(quality.occupancy.empty, 0);
2541 }
2542 }
2543
2544 #[test]
2545 fn float_quality_counts_soar_secondary_postings_in_occupancy() {
2546 let training = [0.0, 0.0, 0.1, 0.0, 10.0, 10.0, 10.1, 10.0];
2547 let config = crate::structures::CoarseConfig::new(2, 2)
2548 .with_routing(crate::dsl::IvfRoutingMode::Flat)
2549 .with_soar(crate::structures::SoarConfig::full());
2550 let centroids = crate::structures::CoarseCentroids::train_contiguous(
2551 &config,
2552 &training,
2553 4,
2554 "test-index",
2555 );
2556 let quality = evaluate_float_build_quality(
2557 ¢roids,
2558 &[0.05, 0.0, 10.05, 10.0],
2559 crate::dsl::IvfRoutingMode::Flat,
2560 )
2561 .unwrap();
2562
2563 assert_eq!(quality.mean_construction_assignments, 2.0);
2564 assert_eq!(quality.occupancy.empty, 0);
2565 }
2566
2567 #[test]
2568 fn artifact_writer_enforces_limit_without_writing_past_it() {
2569 let mut output = Vec::new();
2570 let mut writer = SizeLimitedWriter::new(&mut output, 3);
2571 writer.write_all(&[1, 2]).unwrap();
2572 let error = writer.write_all(&[3, 4]).unwrap_err().to_string();
2573 assert!(error.contains("3-byte safety limit"), "{error}");
2574 assert_eq!(output, vec![1, 2]);
2575 }
2576
2577 #[test]
2578 fn ivf_tq_training_marks_generation_normalizes_and_calibrates_default_soar() {
2579 let config = IvfFieldConfig::Float(DenseVectorConfig::ivf_tq(2, Some(1), 1));
2580 let mut sample = TrainingSample::Float(vec![3.0, 4.0, 30.0, 40.0, 300.0, 400.0]);
2581 let model = IndexWriter::<crate::directories::RamDirectory>::train_field_model(
2582 Field(3),
2583 &config,
2584 &mut sample,
2585 3,
2586 "test",
2587 "test-index",
2588 )
2589 .unwrap();
2590 let TrainedFieldArtifacts::FloatCentroids(centroids) = model.artifacts else {
2591 panic!("expected float IVF-TQ centroids");
2592 };
2593
2594 assert!(crate::structures::is_ivf_tq_cosine_generation(
2595 centroids.version
2596 ));
2597 assert!((centroids.centroids[0] - 0.6).abs() < 1e-6);
2598 assert!((centroids.centroids[1] - 0.8).abs() < 1e-6);
2599 let soar = centroids
2600 .soar_config
2601 .as_ref()
2602 .expect("default SOAR should propagate into the trained router");
2603 assert_eq!(soar.num_secondary, 1);
2604 assert!(soar.selective);
2605 assert!(
2606 soar.spill_threshold > 0.0,
2607 "the negative 30% target tag should be replaced by a calibrated threshold"
2608 );
2609 assert_eq!(soar.calibration_target(), None);
2610 }
2611
2612 #[test]
2613 fn explicitly_disabled_soar_stays_off_during_ivf_tq_training() {
2614 let config = IvfFieldConfig::Float(DenseVectorConfig::ivf_tq(2, Some(1), 1).without_soar());
2615 let mut sample = TrainingSample::Float(vec![3.0, 4.0, 30.0, 40.0, 300.0, 400.0]);
2616 let model = IndexWriter::<crate::directories::RamDirectory>::train_field_model(
2617 Field(3),
2618 &config,
2619 &mut sample,
2620 3,
2621 "test-no-soar",
2622 "test-index",
2623 )
2624 .unwrap();
2625 let TrainedFieldArtifacts::FloatCentroids(centroids) = model.artifacts else {
2626 panic!("expected float IVF-TQ centroids");
2627 };
2628
2629 assert!(centroids.soar_config.is_none());
2630 }
2631
2632 use std::path::Path;
2635 use std::sync::atomic::{AtomicBool, Ordering};
2636
2637 use crate::directories::{
2638 Directory, DirectoryWriter as DirectoryWriterTrait, FileHandle, RamDirectory, RangeReadFn,
2639 };
2640 use crate::dsl::{Document, SchemaBuilder};
2641 use crate::index::{IndexConfig, IndexWriter};
2642
2643 const READ_FAIL_DOCS: usize = 5;
2644 const READ_FAIL_DIM: usize = 4;
2645 const VEC_REGION_START: u64 = 16;
2651 const VEC_REGION_END: u64 = VEC_REGION_START + (READ_FAIL_DOCS * READ_FAIL_DIM * 4) as u64;
2652
2653 #[derive(Clone, Default)]
2658 struct VectorReadFailDirectory {
2659 inner: RamDirectory,
2660 fail_vector_reads: Arc<AtomicBool>,
2661 fail_all_vector_reads: Arc<AtomicBool>,
2662 }
2663
2664 #[async_trait::async_trait]
2665 impl Directory for VectorReadFailDirectory {
2666 async fn exists(&self, path: &Path) -> std::io::Result<bool> {
2667 self.inner.exists(path).await
2668 }
2669
2670 async fn file_size(&self, path: &Path) -> std::io::Result<u64> {
2671 self.inner.file_size(path).await
2672 }
2673
2674 async fn open_read(&self, path: &Path) -> std::io::Result<FileHandle> {
2675 self.inner.open_read(path).await
2676 }
2677
2678 async fn read_range(
2679 &self,
2680 path: &Path,
2681 range: std::ops::Range<u64>,
2682 ) -> std::io::Result<crate::directories::OwnedBytes> {
2683 self.inner.read_range(path, range).await
2684 }
2685
2686 async fn list_files(&self, prefix: &Path) -> std::io::Result<Vec<std::path::PathBuf>> {
2687 self.inner.list_files(prefix).await
2688 }
2689
2690 async fn open_lazy(&self, path: &Path) -> std::io::Result<FileHandle> {
2691 let handle = self.inner.open_lazy(path).await?;
2692 if path.extension().is_some_and(|ext| ext == "vectors") {
2693 let armed = Arc::clone(&self.fail_vector_reads);
2694 let fail_all = Arc::clone(&self.fail_all_vector_reads);
2695 let len = handle.len();
2696 let read_fn: RangeReadFn = Arc::new(move |range: std::ops::Range<u64>| {
2697 let handle = handle.clone();
2698 let armed = Arc::clone(&armed);
2699 let fail_all = Arc::clone(&fail_all);
2700 Box::pin(async move {
2701 if fail_all.load(Ordering::SeqCst)
2702 || (armed.load(Ordering::SeqCst)
2703 && range.start >= VEC_REGION_START
2704 && range.end <= VEC_REGION_END)
2705 {
2706 return Err(std::io::Error::other("injected vector data read failure"));
2707 }
2708 handle.read_bytes_range(range).await
2709 })
2710 });
2711 return Ok(FileHandle::lazy(len, read_fn));
2712 }
2713 Ok(handle)
2714 }
2715 }
2716
2717 #[async_trait::async_trait]
2718 impl DirectoryWriterTrait for VectorReadFailDirectory {
2719 async fn write(&self, path: &Path, data: &[u8]) -> std::io::Result<()> {
2720 self.inner.write(path, data).await
2721 }
2722
2723 async fn delete(&self, path: &Path) -> std::io::Result<()> {
2724 self.inner.delete(path).await
2725 }
2726
2727 async fn rename(&self, from: &Path, to: &Path) -> std::io::Result<()> {
2728 self.inner.rename(from, to).await
2729 }
2730
2731 async fn sync(&self) -> std::io::Result<()> {
2732 self.inner.sync().await
2733 }
2734
2735 async fn streaming_writer(
2736 &self,
2737 path: &Path,
2738 ) -> std::io::Result<Box<dyn crate::directories::StreamingWriter>> {
2739 self.inner.streaming_writer(path).await
2740 }
2741 }
2742
2743 #[tokio::test]
2746 async fn build_propagates_vector_read_errors_without_publishing_artifacts() {
2747 let mut sb = SchemaBuilder::default();
2748 let embedding = sb.add_dense_vector_field_with_config(
2749 "embedding",
2750 true,
2751 true,
2752 DenseVectorConfig::ivf_tq(READ_FAIL_DIM, Some(1), 1),
2753 );
2754 let schema = sb.build();
2755
2756 let dir = VectorReadFailDirectory::default();
2757 let config = IndexConfig {
2758 merge_policy: Box::new(crate::merge::NoMergePolicy),
2759 num_indexing_threads: 1,
2760 ..Default::default()
2761 };
2762 let mut writer = IndexWriter::create(dir.clone(), schema, config)
2763 .await
2764 .unwrap();
2765 for i in 0..READ_FAIL_DOCS {
2766 let mut doc = Document::new();
2767 doc.add_dense_vector(embedding, vec![i as f32 + 1.0; READ_FAIL_DIM]);
2768 writer.add_document(doc).unwrap();
2769 }
2770 writer.commit().await.unwrap();
2771 dir.fail_vector_reads.store(true, Ordering::SeqCst);
2772 let error = writer
2773 .build_vector_index()
2774 .await
2775 .expect_err("failed sample collection must fail the build")
2776 .to_string();
2777 assert!(
2778 error.contains("injected vector data read failure"),
2779 "{error}"
2780 );
2781
2782 assert!(
2783 !writer
2784 .segment_manager
2785 .read_metadata(|meta| meta.is_field_built(embedding.0))
2786 .await,
2787 "a failed build must not publish Built metadata"
2788 );
2789 assert!(
2790 writer.segment_manager.trained().is_none(),
2791 "a failed build must not publish trained artifacts"
2792 );
2793 }
2794
2795 #[tokio::test]
2796 async fn retrain_read_failure_keeps_the_complete_published_generation() {
2797 let mut sb = SchemaBuilder::default();
2798 let embedding = sb.add_dense_vector_field_with_config(
2799 "embedding",
2800 true,
2801 true,
2802 DenseVectorConfig::ivf_tq(READ_FAIL_DIM, Some(1), 1),
2803 );
2804 let schema = sb.build();
2805 let dir = VectorReadFailDirectory::default();
2806 let config = IndexConfig {
2807 merge_policy: Box::new(crate::merge::NoMergePolicy),
2808 num_indexing_threads: 1,
2809 ..Default::default()
2810 };
2811 let mut writer = IndexWriter::create(dir.clone(), schema, config)
2812 .await
2813 .unwrap();
2814 for i in 0..READ_FAIL_DOCS {
2815 let mut doc = Document::new();
2816 doc.add_dense_vector(embedding, vec![i as f32 + 1.0; READ_FAIL_DIM]);
2817 writer.add_document(doc).unwrap();
2818 }
2819 writer.commit().await.unwrap();
2820 writer.build_vector_index().await.unwrap();
2821
2822 let old_ids = writer.segment_manager.get_segment_ids().await;
2823 let old_meta = writer
2824 .segment_manager
2825 .read_metadata(|metadata| metadata.get_field_meta(embedding.0).cloned())
2826 .await
2827 .unwrap();
2828 let old_version = writer.segment_manager.trained().unwrap().centroids[&embedding.0].version;
2829
2830 dir.fail_all_vector_reads.store(true, Ordering::SeqCst);
2831 let error = writer
2832 .retrain_vector_index()
2833 .await
2834 .expect_err("failed sample collection must abort the retrain")
2835 .to_string();
2836 assert!(
2837 error.contains("injected vector data read failure"),
2838 "{error}"
2839 );
2840 assert_eq!(writer.segment_manager.get_segment_ids().await, old_ids);
2841 assert_eq!(
2842 writer
2843 .segment_manager
2844 .read_metadata(|metadata| metadata
2845 .get_field_meta(embedding.0)
2846 .map(|field| (field.centroids_file.clone(), field.codebook_file.clone())))
2847 .await,
2848 Some((old_meta.centroids_file, old_meta.codebook_file)),
2849 );
2850 assert_eq!(
2851 writer.segment_manager.trained().unwrap().centroids[&embedding.0].version,
2852 old_version,
2853 );
2854 }
2855
2856 #[tokio::test]
2857 async fn alter_target_sized_ivf_below_hard_floor_publishes_deferred_flat() {
2858 let mut sb = SchemaBuilder::default();
2859 let hash = sb.add_binary_dense_vector_field_with_config(
2860 "hash",
2861 true,
2862 true,
2863 BinaryDenseVectorConfig::new(256).with_ivf(Some(1), 1),
2864 );
2865 let dir = RamDirectory::new();
2866 let config = IndexConfig {
2867 merge_policy: Box::new(crate::merge::NoMergePolicy),
2868 num_indexing_threads: 1,
2869 ..Default::default()
2870 };
2871 let mut writer = IndexWriter::create(dir.clone(), sb.build(), config)
2872 .await
2873 .unwrap();
2874 for row in 0u8..64 {
2875 let mut doc = Document::new();
2876 doc.add_binary_dense_vector(hash, vec![row; 32]);
2877 writer.add_document(doc).unwrap();
2878 }
2879 writer.commit().await.unwrap();
2880 writer.build_vector_index().await.unwrap();
2881
2882 let target = BinaryDenseVectorConfig::new(256).with_target_vectors(1_000_000_000);
2883 let required = target.optimal_num_clusters(64) * MIN_TRAINING_POINTS_PER_CENTROID;
2884 assert!(
2885 64 < required,
2886 "test corpus must stay below the hardcoded floor"
2887 );
2888 let outcome = writer
2889 .alter_vector_index(hash, VectorIndexAlter::Binary(target))
2890 .await
2891 .unwrap();
2892
2893 assert_eq!(outcome.state, AlterVectorIndexState::DeferredFlat);
2894 let generation = writer.segment_manager.published_generation();
2895 let config = generation
2896 .schema
2897 .get_field_entry(hash)
2898 .unwrap()
2899 .binary_dense_vector_config
2900 .as_ref()
2901 .unwrap();
2902 assert_eq!(config.index_type, BinaryIndexType::Ivf);
2903 assert_eq!(config.num_clusters, None);
2904 assert_eq!(config.target_vectors, Some(1_000_000_000));
2905 assert!(
2906 !writer
2907 .segment_manager
2908 .read_metadata(|metadata| metadata.is_field_built(hash.0))
2909 .await
2910 );
2911 assert!(generation.trained_vectors.is_none());
2912 for id in writer.segment_manager.get_segment_ids().await {
2913 let segment = crate::segment::SegmentReader::open(
2914 &dir,
2915 SegmentId::from_hex(&id).unwrap(),
2916 generation.schema.clone(),
2917 16,
2918 )
2919 .await
2920 .unwrap();
2921 assert!(segment.get_vector_index(hash).is_none());
2922 }
2923 }
2924
2925 #[tokio::test]
2926 async fn alter_target_sized_float_ivf_below_hard_floor_publishes_deferred_flat() {
2927 let mut sb = SchemaBuilder::default();
2928 let embedding = sb.add_dense_vector_field_with_config(
2929 "embedding",
2930 true,
2931 true,
2932 DenseVectorConfig::ivf_tq(4, Some(1), 1),
2933 );
2934 let dir = RamDirectory::new();
2935 let config = IndexConfig {
2936 merge_policy: Box::new(crate::merge::NoMergePolicy),
2937 num_indexing_threads: 1,
2938 ..Default::default()
2939 };
2940 let mut writer = IndexWriter::create(dir.clone(), sb.build(), config)
2941 .await
2942 .unwrap();
2943 for row in 0..64 {
2944 let mut doc = Document::new();
2945 doc.add_dense_vector(embedding, vec![row as f32 + 1.0, 1.0, 0.5, 0.25]);
2946 writer.add_document(doc).unwrap();
2947 }
2948 writer.commit().await.unwrap();
2949 writer.build_vector_index().await.unwrap();
2950
2951 let target = DenseVectorConfig::ivf_tq(4, None, 1).with_target_vectors(1_000_000_000);
2952 let outcome = writer
2953 .alter_vector_index(embedding, VectorIndexAlter::Dense(target))
2954 .await
2955 .unwrap();
2956
2957 assert_eq!(outcome.state, AlterVectorIndexState::DeferredFlat);
2958 let generation = writer.segment_manager.published_generation();
2959 let config = generation
2960 .schema
2961 .get_field_entry(embedding)
2962 .unwrap()
2963 .dense_vector_config
2964 .as_ref()
2965 .unwrap();
2966 assert_eq!(config.index_type, VectorIndexType::IvfTq);
2967 assert_eq!(config.num_clusters, None);
2968 assert_eq!(config.target_vectors, Some(1_000_000_000));
2969 assert!(
2970 !writer
2971 .segment_manager
2972 .read_metadata(|metadata| metadata.is_field_built(embedding.0))
2973 .await
2974 );
2975 assert!(generation.trained_vectors.is_none());
2976 for id in writer.segment_manager.get_segment_ids().await {
2977 let segment = crate::segment::SegmentReader::open(
2978 &dir,
2979 SegmentId::from_hex(&id).unwrap(),
2980 generation.schema.clone(),
2981 16,
2982 )
2983 .await
2984 .unwrap();
2985 assert!(segment.get_vector_index(embedding).is_none());
2986 }
2987 }
2988
2989 #[tokio::test]
2990 async fn alter_ivf_scann_deferred_and_back_is_atomic() {
2991 let mut sb = SchemaBuilder::default();
2992 let embedding = sb.add_dense_vector_field_with_config(
2993 "embedding",
2994 true,
2995 true,
2996 DenseVectorConfig::ivf_tq(4, Some(1), 1),
2997 );
2998 let dir = crate::directories::RamDirectory::new();
2999 let config = crate::IndexConfig {
3000 merge_policy: Box::new(crate::merge::NoMergePolicy),
3001 num_indexing_threads: 1,
3002 ..Default::default()
3003 };
3004 let mut writer = IndexWriter::create(dir.clone(), sb.build(), config)
3005 .await
3006 .unwrap();
3007 for row in 0..64 {
3008 let mut doc = Document::new();
3009 doc.add_dense_vector(embedding, vec![row as f32 + 1.0, 1.0, 0.5, 0.25]);
3010 writer.add_document(doc).unwrap();
3011 }
3012 writer.commit().await.unwrap();
3013 writer.build_vector_index().await.unwrap();
3014
3015 for row in 64..72 {
3019 let mut doc = Document::new();
3020 doc.add_dense_vector(embedding, vec![row as f32 + 1.0, 1.0, 0.5, 0.25]);
3021 writer.add_document(doc).unwrap();
3022 }
3023
3024 let deferred = writer
3025 .alter_vector_index(embedding, VectorIndexAlter::Dense(scann_config(4, Some(2))))
3026 .await
3027 .unwrap();
3028 assert_eq!(deferred.state, AlterVectorIndexState::DeferredFlat);
3029 let deferred_generation = writer.segment_manager.published_generation();
3030 assert_eq!(
3031 deferred_generation
3032 .schema
3033 .get_field_entry(embedding)
3034 .unwrap()
3035 .dense_vector_config
3036 .as_ref()
3037 .unwrap()
3038 .index_type,
3039 VectorIndexType::Scann
3040 );
3041 assert!(
3042 !writer
3043 .segment_manager
3044 .read_metadata(|metadata| metadata.is_field_built(embedding.0))
3045 .await
3046 );
3047 assert!(deferred_generation.trained_vectors.is_none());
3048
3049 for row in 72..80 {
3053 let mut doc = Document::new();
3054 doc.add_dense_vector(embedding, vec![row as f32 + 1.0, 1.0, 0.5, 0.25]);
3055 writer.add_document(doc).unwrap();
3056 }
3057 writer.commit().await.unwrap();
3058 assert_eq!(
3059 writer
3060 .schema()
3061 .get_field_entry(embedding)
3062 .unwrap()
3063 .dense_vector_config
3064 .as_ref()
3065 .unwrap()
3066 .index_type,
3067 VectorIndexType::Scann
3068 );
3069 for id in writer.segment_manager.get_segment_ids().await {
3070 let segment = crate::segment::SegmentReader::open(
3071 &dir,
3072 SegmentId::from_hex(&id).unwrap(),
3073 deferred_generation.schema.clone(),
3074 16,
3075 )
3076 .await
3077 .unwrap();
3078 assert!(segment.get_vector_index(embedding).is_none());
3079 }
3080
3081 let rebuilt = writer
3082 .alter_vector_index(
3083 embedding,
3084 VectorIndexAlter::Dense(DenseVectorConfig::ivf_tq(4, Some(1), 1)),
3085 )
3086 .await
3087 .unwrap();
3088 assert_eq!(rebuilt.state, AlterVectorIndexState::Built);
3089 assert!(rebuilt.publication_generation > deferred.publication_generation);
3090 let rebuilt_generation = writer.segment_manager.published_generation();
3091 assert_eq!(
3092 rebuilt_generation
3093 .schema
3094 .get_field_entry(embedding)
3095 .unwrap()
3096 .dense_vector_config
3097 .as_ref()
3098 .unwrap()
3099 .index_type,
3100 VectorIndexType::IvfTq
3101 );
3102 assert!(
3103 rebuilt_generation
3104 .trained_vectors
3105 .as_ref()
3106 .is_some_and(|trained| trained.centroids.contains_key(&embedding.0))
3107 );
3108 for id in writer.segment_manager.get_segment_ids().await {
3109 let segment = crate::segment::SegmentReader::open(
3110 &dir,
3111 SegmentId::from_hex(&id).unwrap(),
3112 rebuilt_generation.schema.clone(),
3113 16,
3114 )
3115 .await
3116 .unwrap();
3117 assert!(matches!(
3118 segment.get_vector_index(embedding),
3119 Some(crate::segment::VectorIndex::IvfTq { .. })
3120 ));
3121 }
3122
3123 let ids_before_parameter_change = writer.segment_manager.get_segment_ids().await;
3124 let parameters_only = writer
3125 .alter_vector_index(
3126 embedding,
3127 VectorIndexAlter::Dense(DenseVectorConfig::ivf_tq(4, Some(1), 7)),
3128 )
3129 .await
3130 .unwrap();
3131 assert_eq!(parameters_only.state, AlterVectorIndexState::ParametersOnly);
3132 assert!(parameters_only.publication_generation > rebuilt.publication_generation);
3133 assert_eq!(
3134 writer.segment_manager.get_segment_ids().await,
3135 ids_before_parameter_change
3136 );
3137 assert_eq!(
3138 writer
3139 .schema()
3140 .get_field_entry(embedding)
3141 .unwrap()
3142 .dense_vector_config
3143 .as_ref()
3144 .unwrap()
3145 .nprobe,
3146 7
3147 );
3148 }
3149}