1#[cfg(feature = "native")]
10use std::cmp::Reverse;
11#[cfg(feature = "native")]
12use std::collections::BinaryHeap;
13use std::io;
14#[cfg(feature = "native")]
15use std::io::Write;
16use std::ops::Range;
17
18#[cfg(feature = "native")]
19use byteorder::WriteBytesExt;
20use byteorder::{LittleEndian, ReadBytesExt};
21
22use crate::directories::OwnedBytes;
23use crate::dsl::IvfRoutingMode;
24
25use crate::observe::DenseAnnScanStats;
26#[cfg(feature = "native")]
27use crate::structures::BinaryIvfIndex;
28use crate::structures::vector::index::{BoundedAnnCollector, BoundedUniqueAnnCollector};
29
30type CombinedBinaryCandidates = (Vec<AnnDocumentCandidate>, Vec<(u32, u16, f32)>);
33type ScanResults = (Vec<(u32, u16, f32)>, DenseAnnScanStats);
35
36trait IvfTqTopK: Sized + Send {
41 fn with_k(k: usize) -> Self;
42 fn insert(&mut self, doc_id: u32, ordinal: u16, score: f32);
43 fn pruning_threshold(&self) -> Option<f32>;
45 fn non_finite_dropped(&self) -> usize;
46 #[cfg(feature = "native")]
47 fn merge_from(&mut self, other: Self);
48 fn into_sorted_results(self) -> Vec<(u32, u16, f32)>;
49}
50
51impl IvfTqTopK for BoundedAnnCollector<true, true> {
52 #[inline]
53 fn with_k(k: usize) -> Self {
54 Self::new(k)
55 }
56 #[inline]
57 fn insert(&mut self, doc_id: u32, ordinal: u16, score: f32) {
58 BoundedAnnCollector::insert(self, doc_id, ordinal, score)
59 }
60 #[inline]
61 fn pruning_threshold(&self) -> Option<f32> {
62 BoundedAnnCollector::pruning_threshold(self)
63 }
64 #[inline]
65 fn non_finite_dropped(&self) -> usize {
66 BoundedAnnCollector::non_finite_dropped(self)
67 }
68 #[cfg(feature = "native")]
69 #[inline]
70 fn merge_from(&mut self, other: Self) {
71 BoundedAnnCollector::merge_from(self, other)
72 }
73 #[inline]
74 fn into_sorted_results(self) -> Vec<(u32, u16, f32)> {
75 BoundedAnnCollector::into_sorted_results(self)
76 }
77}
78
79impl IvfTqTopK for BoundedUniqueAnnCollector<true> {
80 #[inline]
81 fn with_k(k: usize) -> Self {
82 Self::new(k)
83 }
84 #[inline]
85 fn insert(&mut self, doc_id: u32, ordinal: u16, score: f32) {
86 BoundedUniqueAnnCollector::insert(self, doc_id, ordinal, score)
87 }
88 #[inline]
89 fn pruning_threshold(&self) -> Option<f32> {
90 BoundedUniqueAnnCollector::pruning_threshold(self)
91 }
92 #[inline]
93 fn non_finite_dropped(&self) -> usize {
94 BoundedUniqueAnnCollector::non_finite_dropped(self)
95 }
96 #[cfg(feature = "native")]
97 #[inline]
98 fn merge_from(&mut self, other: Self) {
99 BoundedUniqueAnnCollector::merge_from(self, other)
100 }
101 #[inline]
102 fn into_sorted_results(self) -> Vec<(u32, u16, f32)> {
103 BoundedUniqueAnnCollector::into_sorted_results(self)
104 }
105}
106
107const ANN_HEADER_MAGIC: u32 = 0x3152_4e41; const ANN_FOOTER_MAGIC: u32 = 0x3146_4e41; const ANN_DISK_VERSION: u16 = 1;
110const ANN_HEADER_SIZE: usize = 56;
111const ANN_RUN_SIZE: usize = 48;
112const ANN_FOOTER_SIZE: usize = 24;
113#[cfg(feature = "native")]
114const COPY_CHUNK: usize = 8 * 1024 * 1024;
115#[cfg(feature = "native")]
116const PREFETCH_COALESCE_GAP: usize = 4 * 1024;
117const BINARY_SCORE_BATCH: usize = 8_192;
118const TQ_PRUNE_ESTIMATE_BOUND: f32 = 1.3;
123#[cfg(feature = "native")]
129const TQ_PARALLEL_SCAN_MIN_VECTORS: usize = 65_536;
130#[cfg(feature = "native")]
131const TQ_PARALLEL_SCAN_CHUNK_BLOCKS: usize = 512;
132const IVF_PARALLEL_SCAN_MIN_POSTINGS: usize = 65_536;
136const IVF_TQ_PARALLEL_SCAN_CHUNK_BLOCKS: usize = 512;
137
138#[derive(Debug, Clone, Copy, PartialEq, Eq)]
139pub(crate) enum AnnKind {
140 BinaryIvf = 2,
143 TqFlat = 3,
145 IvfTq = 4,
147 ScannAh = 5,
149 ScannBinary = 6,
151}
152
153impl AnnKind {
154 fn from_u8(value: u8) -> io::Result<Self> {
155 match value {
156 1 => Err(invalid_data(
157 "ANN kind 1 (IVF-PQ) is no longer supported; recreate the index \
158 with `ivf_tq` and reindex",
159 )),
160 2 => Ok(Self::BinaryIvf),
161 3 => Ok(Self::TqFlat),
162 4 => Ok(Self::IvfTq),
163 5 => Ok(Self::ScannAh),
164 6 => Ok(Self::ScannBinary),
165 _ => Err(invalid_data(format!("unknown ANN kind {value}"))),
166 }
167 }
168}
169
170fn expected_codes_column_len(
174 kind: AnnKind,
175 count: usize,
176 dim: usize,
177 code_size: usize,
178) -> io::Result<usize> {
179 match kind {
180 AnnKind::BinaryIvf | AnnKind::ScannBinary => count
181 .checked_mul(code_size)
182 .ok_or_else(|| invalid_data("ANN code column size overflows usize")),
183 AnnKind::TqFlat => {
185 crate::structures::vector::quantization::tq_codes_column_len_checked(count, code_size)
186 .ok_or_else(|| invalid_data("TQ code column size overflows usize"))
187 }
188 AnnKind::IvfTq => crate::structures::vector::quantization::tq_ivf_codes_column_len_checked(
189 count, code_size,
190 )
191 .ok_or_else(|| invalid_data("IVF-TQ code column size overflows usize")),
192 AnnKind::ScannAh => crate::structures::vector::scann::ScannEncoding::AsymmetricHash {
197 dimensions_per_block: u16::try_from(code_size)
198 .map_err(|_| invalid_data("ScaNN AH block dimension exceeds u16"))?,
199 bits_per_code: 4,
200 }
201 .leaf_code_bytes(
202 u32::try_from(dim).map_err(|_| invalid_data("ScaNN dimension exceeds u32"))?,
203 count,
204 )
205 .map_err(|error| invalid_data(error.to_string())),
206 }
207}
208
209#[derive(Debug, Clone, PartialEq, Eq)]
210pub(crate) struct AnnDiskHeader {
211 pub kind: AnnKind,
212 pub routing: IvfRoutingMode,
213 pub dim: usize,
214 pub code_size: usize,
215 pub num_clusters: u32,
216 pub quantizer_version: u64,
217 pub codebook_version: u64,
218 pub vector_count: usize,
219}
220
221#[derive(Debug)]
222struct AnnRun {
223 cluster_id: u32,
224 doc_base: u32,
225 max_doc_id: u32,
226 count: usize,
227 doc_ids: Range<usize>,
228 ordinals: Range<usize>,
229 codes: Range<usize>,
230}
231
232#[derive(Clone, Copy)]
233struct IvfTqScanTask<'a> {
234 run: &'a AnnRun,
235 cluster_dot: f32,
236 first_block: usize,
237 last_block: usize,
238}
239
240#[derive(Clone, Copy)]
241struct BinaryScanTask<'a> {
242 run: &'a AnnRun,
243 first_index: usize,
244 count: usize,
245}
246
247pub(crate) struct AnnDiskIndex {
250 #[cfg(feature = "native")]
252 heap_pins: crate::segment::pin::HeapPinSet,
253 raw: OwnedBytes,
254 header: AnnDiskHeader,
255 runs: Vec<AnnRun>,
256}
257
258#[derive(Debug, Clone, Copy, PartialEq)]
266pub struct AnnHealth {
267 pub vectors: u64,
269 pub clusters_nonempty: u32,
271 pub clusters_total: u32,
273 pub runs: u32,
275 pub largest_cluster: u32,
277 pub largest_cluster_vectors: u64,
278 pub imbalance: f64,
283 pub payload_bytes: u64,
285}
286
287impl AnnHealth {
288 pub fn fragmentation(&self) -> f64 {
292 if self.clusters_nonempty == 0 {
293 return 0.0;
294 }
295 f64::from(self.runs) / f64::from(self.clusters_nonempty)
296 }
297
298 pub fn largest_cluster_share(&self) -> f64 {
300 if self.vectors == 0 {
301 return 0.0;
302 }
303 self.largest_cluster_vectors as f64 / self.vectors as f64
304 }
305}
306
307const ANN_SKEW_WARN_SHARE: f64 = 0.05;
311const ANN_SKEW_WARN_MIN_VECTORS: u64 = 100_000;
312const ANN_FRAGMENTATION_WARN: f64 = 8.0;
315
316#[derive(Debug, Clone, Copy, PartialEq)]
322pub(crate) struct AnnDocumentCandidate {
323 pub(crate) doc_id: u32,
324 pub(crate) score: f32,
325}
326
327impl AnnDiskIndex {
328 pub(crate) fn open(
329 raw: OwnedBytes,
330 expected_kind: AnnKind,
331 total_docs: u32,
332 ) -> io::Result<Self> {
333 if raw.len() < ANN_HEADER_SIZE + ANN_FOOTER_SIZE {
334 return Err(invalid_data("ANN payload is shorter than header + footer"));
335 }
336 let bytes = raw.as_slice();
337 let mut header_cursor = std::io::Cursor::new(&bytes[..ANN_HEADER_SIZE]);
338 if header_cursor.read_u32::<LittleEndian>()? != ANN_HEADER_MAGIC {
339 return Err(invalid_data("ANN payload has unsupported header magic"));
340 }
341 let kind = AnnKind::from_u8(header_cursor.read_u8()?)?;
342 if kind != expected_kind {
343 return Err(invalid_data(format!(
344 "ANN payload kind {kind:?} does not match expected {expected_kind:?}"
345 )));
346 }
347 let routing = routing_from_u8(header_cursor.read_u8()?)?;
348 if header_cursor.read_u16::<LittleEndian>()? != ANN_DISK_VERSION {
349 return Err(invalid_data("ANN payload has unsupported format version"));
350 }
351 let dim = header_cursor.read_u32::<LittleEndian>()? as usize;
352 let code_size = header_cursor.read_u32::<LittleEndian>()? as usize;
353 let num_clusters = header_cursor.read_u32::<LittleEndian>()?;
354 check_layout_revision(kind, header_cursor.read_u32::<LittleEndian>()?)?;
355 let quantizer_version = header_cursor.read_u64::<LittleEndian>()?;
356 let codebook_version = header_cursor.read_u64::<LittleEndian>()?;
357 let vector_count = usize::try_from(header_cursor.read_u64::<LittleEndian>()?)
358 .map_err(|_| invalid_data("ANN vector count exceeds usize"))?;
359 if header_cursor.read_u64::<LittleEndian>()? != 0 {
360 return Err(invalid_data("ANN header tail is non-zero"));
361 }
362 let header = AnnDiskHeader {
363 kind,
364 routing,
365 dim,
366 code_size,
367 num_clusters,
368 quantizer_version,
369 codebook_version,
370 vector_count,
371 };
372 validate_header(&header)?;
373
374 let footer_start = bytes.len() - ANN_FOOTER_SIZE;
375 let mut footer_cursor = std::io::Cursor::new(&bytes[footer_start..]);
376 let directory_offset = usize::try_from(footer_cursor.read_u64::<LittleEndian>()?)
377 .map_err(|_| invalid_data("ANN directory offset exceeds usize"))?;
378 let num_runs = usize::try_from(footer_cursor.read_u64::<LittleEndian>()?)
379 .map_err(|_| invalid_data("ANN run count exceeds usize"))?;
380 if footer_cursor.read_u32::<LittleEndian>()? != ANN_FOOTER_MAGIC
381 || footer_cursor.read_u32::<LittleEndian>()? != u32::from(ANN_DISK_VERSION)
382 {
383 return Err(invalid_data("ANN payload has unsupported footer"));
384 }
385 if num_runs == 0 {
386 return Err(invalid_data("ANN payload has no cluster runs"));
387 }
388 let directory_len = num_runs
389 .checked_mul(ANN_RUN_SIZE)
390 .ok_or_else(|| invalid_data("ANN directory size overflows usize"))?;
391 if directory_offset < ANN_HEADER_SIZE
392 || directory_offset.checked_add(directory_len) != Some(footer_start)
393 {
394 return Err(invalid_data("ANN directory does not end at the footer"));
395 }
396
397 let mut runs = Vec::with_capacity(num_runs);
398 let mut directory_cursor = std::io::Cursor::new(&bytes[directory_offset..footer_start]);
399 let mut previous_cluster = None;
400 let mut counted_vectors = 0usize;
401 for _ in 0..num_runs {
402 let cluster_id = directory_cursor.read_u32::<LittleEndian>()?;
403 let doc_base = directory_cursor.read_u32::<LittleEndian>()?;
404 let count = directory_cursor.read_u32::<LittleEndian>()? as usize;
405 let max_doc_id = directory_cursor.read_u32::<LittleEndian>()?;
406 let doc_ids_offset = usize::try_from(directory_cursor.read_u64::<LittleEndian>()?)
407 .map_err(|_| invalid_data("ANN doc-ID offset exceeds usize"))?;
408 let ordinals_offset = usize::try_from(directory_cursor.read_u64::<LittleEndian>()?)
409 .map_err(|_| invalid_data("ANN ordinal offset exceeds usize"))?;
410 let codes_offset = usize::try_from(directory_cursor.read_u64::<LittleEndian>()?)
411 .map_err(|_| invalid_data("ANN code offset exceeds usize"))?;
412 let codes_len = usize::try_from(directory_cursor.read_u64::<LittleEndian>()?)
413 .map_err(|_| invalid_data("ANN code length exceeds usize"))?;
414 if count == 0
415 || cluster_id >= num_clusters
416 || previous_cluster.is_some_and(|previous| previous > cluster_id)
417 || doc_base
418 .checked_add(max_doc_id)
419 .is_none_or(|doc_id| doc_id >= total_docs)
420 {
421 return Err(invalid_data("ANN run metadata is invalid"));
422 }
423 previous_cluster = Some(cluster_id);
424 let doc_ids_len = count
425 .checked_mul(std::mem::size_of::<u32>())
426 .ok_or_else(|| invalid_data("ANN doc-ID column size overflows usize"))?;
427 let ordinals_len = count
428 .checked_mul(std::mem::size_of::<u16>())
429 .ok_or_else(|| invalid_data("ANN ordinal column size overflows usize"))?;
430 let expected_codes_len = expected_codes_column_len(kind, count, dim, code_size)?;
431 let doc_ids_end = doc_ids_offset
432 .checked_add(doc_ids_len)
433 .ok_or_else(|| invalid_data("ANN doc-ID range overflows usize"))?;
434 let ordinals_end = ordinals_offset
435 .checked_add(ordinals_len)
436 .ok_or_else(|| invalid_data("ANN ordinal range overflows usize"))?;
437 let codes_end = codes_offset
438 .checked_add(codes_len)
439 .ok_or_else(|| invalid_data("ANN code range overflows usize"))?;
440 if doc_ids_offset < ANN_HEADER_SIZE
441 || ordinals_offset != doc_ids_end
442 || codes_offset != ordinals_end
443 || codes_len != expected_codes_len
444 || codes_end > directory_offset
445 {
446 return Err(invalid_data("ANN run columns are not contiguous/in bounds"));
447 }
448 runs.push(AnnRun {
449 cluster_id,
450 doc_base,
451 max_doc_id,
452 count,
453 doc_ids: doc_ids_offset..doc_ids_end,
454 ordinals: ordinals_offset..ordinals_end,
455 codes: codes_offset..codes_end,
456 });
457 counted_vectors = counted_vectors
458 .checked_add(count)
459 .ok_or_else(|| invalid_data("ANN run vector count overflows usize"))?;
460 }
461 for run in &runs {
465 let column = &bytes[run.doc_ids.clone()];
466 if let Some(local_doc_id) = column
467 .chunks_exact(std::mem::size_of::<u32>())
468 .map(|entry| u32::from_le_bytes([entry[0], entry[1], entry[2], entry[3]]))
469 .find(|&local_doc_id| local_doc_id > run.max_doc_id)
470 {
471 return Err(invalid_data(format!(
472 "ANN run for cluster {} contains local document {local_doc_id} above its declared maximum {}",
473 run.cluster_id, run.max_doc_id
474 )));
475 }
476 }
477 let mut payload_order: Vec<usize> = (0..runs.len()).collect();
478 payload_order.sort_unstable_by_key(|&index| runs[index].doc_ids.start);
479 let mut expected_payload_offset = ANN_HEADER_SIZE;
480 for index in payload_order {
481 let run = &runs[index];
482 if run.doc_ids.start != expected_payload_offset {
483 return Err(invalid_data("ANN payload runs overlap or contain gaps"));
484 }
485 expected_payload_offset = run.codes.end;
486 }
487 if expected_payload_offset != directory_offset || counted_vectors != vector_count {
488 return Err(invalid_data(
489 "ANN payload coverage/vector count is inconsistent",
490 ));
491 }
492
493 #[cfg(feature = "native")]
500 raw.madvise_range(
501 ANN_HEADER_SIZE..directory_offset,
502 if kind == AnnKind::TqFlat {
503 libc::MADV_SEQUENTIAL
504 } else {
505 libc::MADV_RANDOM
506 },
507 );
508
509 Ok(Self {
510 #[cfg(feature = "native")]
511 heap_pins: Default::default(),
512 raw,
513 header,
514 runs,
515 })
516 }
517
518 pub(crate) fn health(&self) -> AnnHealth {
521 let mut vectors = 0u64;
522 let mut clusters_nonempty = 0u32;
523 let mut payload_bytes = 0u64;
524 let mut largest = (0u32, 0u64);
525 let mut sum_squares = 0f64;
526 let mut index = 0usize;
528 while index < self.runs.len() {
529 let cluster_id = self.runs[index].cluster_id;
530 let mut cluster_vectors = 0u64;
531 while index < self.runs.len() && self.runs[index].cluster_id == cluster_id {
532 let run = &self.runs[index];
533 cluster_vectors += run.count as u64;
534 payload_bytes += (run.codes.end - run.codes.start) as u64;
535 index += 1;
536 }
537 vectors += cluster_vectors;
538 clusters_nonempty += 1;
539 sum_squares += (cluster_vectors as f64) * (cluster_vectors as f64);
540 if cluster_vectors > largest.1 {
541 largest = (cluster_id, cluster_vectors);
542 }
543 }
544 let imbalance = if vectors == 0 || clusters_nonempty == 0 {
545 0.0
546 } else {
547 f64::from(clusters_nonempty) * sum_squares / ((vectors as f64) * (vectors as f64))
548 };
549 AnnHealth {
550 vectors,
551 clusters_nonempty,
552 clusters_total: self.header.num_clusters,
553 runs: self.runs.len() as u32,
554 largest_cluster: largest.0,
555 largest_cluster_vectors: largest.1,
556 imbalance,
557 payload_bytes,
558 }
559 }
560
561 pub(crate) fn report_health(&self, index_label: &str, field_id: u32, segment_id: u128) {
566 let health = self.health();
567 let share = health.largest_cluster_share();
568 let fragmentation = health.fragmentation();
569 log::info!(
570 "[ann_health] index={index_label} field={field_id} segment={segment_id:016x}: \
571 vectors={} clusters={}/{} runs={} fragmentation={fragmentation:.2} \
572 imbalance={:.2} largest_leaf={:.2}% payload={}",
573 health.vectors,
574 health.clusters_nonempty,
575 health.clusters_total,
576 health.runs,
577 health.imbalance,
578 100.0 * share,
579 crate::format_bytes(health.payload_bytes),
580 );
581 crate::observe::ann_health(
582 index_label,
583 field_id,
584 health.imbalance,
585 fragmentation,
586 share,
587 );
588 if share >= ANN_SKEW_WARN_SHARE && health.vectors >= ANN_SKEW_WARN_MIN_VECTORS {
589 log::warn!(
590 "[ann_health] index={index_label} field={field_id} segment={segment_id:016x}: \
591 leaf {} holds {:.1}% of {} vectors — every query probing it scans that leaf \
592 in full; degenerate embeddings collapse into one leaf exactly like this",
593 health.largest_cluster,
594 100.0 * share,
595 health.vectors,
596 );
597 }
598 if fragmentation >= ANN_FRAGMENTATION_WARN {
599 log::warn!(
600 "[ann_health] index={index_label} field={field_id} segment={segment_id:016x}: \
601 {fragmentation:.1} extents per probed cluster ({} runs / {} clusters) — \
602 cold probes pay that many seeks; the next merge or vector-generation rewrite \
603 compacts to 1.0",
604 health.runs,
605 health.clusters_nonempty,
606 );
607 }
608 }
609
610 pub(crate) fn header(&self) -> &AnnDiskHeader {
611 &self.header
612 }
613
614 pub(crate) fn validate_scann_generation(
617 &self,
618 config: &crate::structures::vector::scann::ScannConfig,
619 generation: u64,
620 artifact_id: u64,
621 ) -> io::Result<()> {
622 use crate::structures::vector::scann::ScannEncoding;
623
624 let expected_kind = match config.encoding {
625 ScannEncoding::AsymmetricHash { .. } => AnnKind::ScannAh,
626 ScannEncoding::BinaryHamming => AnnKind::ScannBinary,
627 };
628 if self.header.kind != expected_kind
629 || self.header.routing != IvfRoutingMode::Flat
630 || self.header.dim != config.dimension as usize
631 || self.header.num_clusters != config.num_leaves
632 || self.header.quantizer_version != generation
633 || self.header.codebook_version != artifact_id
634 {
635 return Err(invalid_data(
636 "ScaNN ANN payload does not match the global trained generation",
637 ));
638 }
639 match config.encoding {
640 ScannEncoding::AsymmetricHash {
641 dimensions_per_block,
642 bits_per_code: 4,
643 } if self.header.code_size == usize::from(dimensions_per_block) => Ok(()),
644 ScannEncoding::BinaryHamming
645 if self.header.code_size == config.dimension as usize / 8 =>
646 {
647 Ok(())
648 }
649 _ => Err(invalid_data(
650 "ScaNN ANN payload encoding does not match the global trained artifact",
651 )),
652 }
653 }
654
655 #[cfg(feature = "native")]
660 pub(crate) fn validate_scann_posting_count(
661 &self,
662 logical_vectors: usize,
663 soar: Option<&crate::structures::SoarConfig>,
664 ) -> io::Result<()> {
665 let spill_budget = match self.header.kind {
666 AnnKind::ScannAh => 0,
667 AnnKind::ScannBinary
668 if self.header.num_clusters > 1
669 && soar.is_some_and(|config| config.num_secondary > 0) =>
670 {
671 match soar.and_then(crate::structures::SoarConfig::calibration_target) {
672 Some(target_fraction) => {
673 (logical_vectors as f64 * f64::from(target_fraction)).floor() as usize
674 }
675 None => logical_vectors,
676 }
677 }
678 AnnKind::ScannBinary => 0,
679 _ => {
680 return Err(invalid_data(
681 "posting-count validation requires a ScaNN ANN payload",
682 ));
683 }
684 };
685 let maximum = logical_vectors
686 .checked_add(spill_budget)
687 .ok_or_else(|| invalid_data("ScaNN posting-count bound overflows usize"))?;
688 if self.header.vector_count < logical_vectors || self.header.vector_count > maximum {
689 return Err(invalid_data(format!(
690 "ScaNN ANN payload has {} physical postings for {logical_vectors} logical vectors; expected {logical_vectors}..={maximum}",
691 self.header.vector_count,
692 )));
693 }
694 Ok(())
695 }
696
697 pub(crate) fn estimated_heap_bytes(&self) -> usize {
698 std::mem::size_of::<Self>() + self.runs.capacity() * std::mem::size_of::<AnnRun>()
699 }
700
701 #[cfg(feature = "native")]
702 pub(crate) fn pin_lookup_directory(
703 &mut self,
704 mode: crate::segment::pin::PinMode,
705 remaining: &mut u64,
706 report: &mut crate::segment::pin::PinReport,
707 ) {
708 let before = self.heap_pins.report();
709 self.heap_pins
710 .pin_slice(&self.runs, "ANN cluster-run directory", mode, remaining);
711 let after = self.heap_pins.report();
712 report.intended_bytes += after.intended_bytes - before.intended_bytes;
713 report.pinned_bytes += after.pinned_bytes - before.pinned_bytes;
714 report.skipped_budget_bytes += after.skipped_budget_bytes - before.skipped_budget_bytes;
715 report.failed_bytes += after.failed_bytes - before.failed_bytes;
716 report.heap_copy_bytes += after.heap_copy_bytes - before.heap_copy_bytes;
717 }
718
719 fn cluster_runs(&self, cluster_id: u32) -> &[AnnRun] {
720 let start = self.runs.partition_point(|run| run.cluster_id < cluster_id);
721 let end = self
722 .runs
723 .partition_point(|run| run.cluster_id <= cluster_id);
724 &self.runs[start..end]
725 }
726
727 #[cfg(feature = "native")]
728 fn ivf_tq_scan_tasks<'a>(
729 &'a self,
730 plan: &'a crate::structures::TqIvfQueryPlan,
731 block_bytes: usize,
732 chunk_blocks: usize,
733 ) -> Vec<IvfTqScanTask<'a>> {
734 debug_assert!(chunk_blocks > 0);
735 let mut tasks = Vec::new();
736 for (cluster_id, cluster_dot) in plan.cluster_dots() {
737 for run in self.cluster_runs(cluster_id) {
738 let block_count = run.codes.len() / block_bytes;
739 for first_block in (0..block_count).step_by(chunk_blocks) {
740 tasks.push(IvfTqScanTask {
741 run,
742 cluster_dot,
743 first_block,
744 last_block: (first_block + chunk_blocks).min(block_count),
745 });
746 }
747 }
748 }
749 tasks
750 }
751
752 #[cfg(feature = "native")]
753 fn binary_scan_tasks<'a>(&'a self, cluster_ids: &[u32]) -> Vec<BinaryScanTask<'a>> {
754 let mut tasks = Vec::new();
755 for &cluster_id in cluster_ids {
756 for run in self.cluster_runs(cluster_id) {
757 for first_index in (0..run.count).step_by(BINARY_SCORE_BATCH) {
758 tasks.push(BinaryScanTask {
759 run,
760 first_index,
761 count: BINARY_SCORE_BATCH.min(run.count - first_index),
762 });
763 }
764 }
765 }
766 tasks
767 }
768
769 #[cfg(feature = "native")]
776 fn prefetch_cluster_runs(&self, cluster_ids: &[u32]) {
777 if cluster_ids.is_empty() || !self.raw.is_mmap() {
778 return;
779 }
780 let mut ranges = Vec::with_capacity(cluster_ids.len());
781 for &cluster_id in cluster_ids {
782 ranges.extend(
783 self.cluster_runs(cluster_id)
784 .iter()
785 .map(|run| run.doc_ids.start..run.codes.end),
786 );
787 }
788 coalesce_prefetch_ranges(&mut ranges);
789 for range in ranges {
790 self.raw.madvise_range(range, libc::MADV_WILLNEED);
791 }
792 }
793
794 pub(crate) fn search_tq_combined_documents(
799 &self,
800 k: usize,
801 plan: &crate::structures::TqQueryPlan,
802 combiner: crate::query::MultiValueCombiner,
803 ) -> io::Result<Vec<AnnDocumentCandidate>> {
804 use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_block_bytes};
805
806 combiner
807 .validate()
808 .map_err(|message| io::Error::new(io::ErrorKind::InvalidInput, message))?;
809 if plan.padded_dim() != self.header.code_size * 2 {
810 return Err(io::Error::new(
811 io::ErrorKind::InvalidInput,
812 "TQ query plan does not match the payload dimension",
813 ));
814 }
815 if k == 0 {
816 return Ok(Vec::new());
817 }
818
819 let block_bytes = tq_block_bytes(self.header.code_size);
820 let bytes = self.raw.as_slice();
821 let mut top_documents = BoundedAnnCollector::<true, true>::new(k);
822 let mut ordinal_scores = Vec::new();
823 let mut scores = [0.0f32; TQ_BLOCK_LANES];
824
825 for run in &self.runs {
826 let mut current_doc = None;
827 let codes = &bytes[run.codes.clone()];
828 for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
829 crate::structures::vector::quantization::tq_score_block(plan, block, &mut scores);
830 let lane_base = block_index * TQ_BLOCK_LANES;
831 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
832 for (lane, &score) in scores.iter().enumerate().take(lanes) {
833 let index = lane_base + lane;
834 let doc_id = run_doc_id_validated(bytes, run, index);
835 if current_doc.is_some_and(|previous| doc_id < previous) {
836 return Err(invalid_data("flat TQ run is not grouped by document ID"));
837 }
838 if current_doc.is_some_and(|previous| doc_id != previous) {
839 retain_combined_document(
840 &mut top_documents,
841 current_doc.expect("current document is present"),
842 &ordinal_scores,
843 combiner,
844 );
845 ordinal_scores.clear();
846 }
847 current_doc = Some(doc_id);
848 if score.is_finite() {
849 ordinal_scores.push((
850 u32::from(read_u16(bytes, run.ordinals.start + index * 2)),
851 score,
852 ));
853 }
854 }
855 }
856 if let Some(doc_id) = current_doc {
857 retain_combined_document(&mut top_documents, doc_id, &ordinal_scores, combiner);
858 ordinal_scores.clear();
859 }
860 }
861
862 Ok(top_documents
863 .into_sorted_results()
864 .into_iter()
865 .map(|(doc_id, _, score)| AnnDocumentCandidate { doc_id, score })
866 .collect())
867 }
868
869 #[cfg(test)]
877 pub(crate) fn search_ivf_tq_combined_documents(
878 &self,
879 k: usize,
880 plan: &crate::structures::TqIvfQueryPlan,
881 combiner: crate::query::MultiValueCombiner,
882 ) -> io::Result<Vec<AnnDocumentCandidate>> {
883 self.search_ivf_tq_combined_documents_with_stats(k, plan, combiner)
884 .map(|(candidates, _)| candidates)
885 }
886
887 pub(crate) fn search_ivf_tq_combined_documents_with_stats(
891 &self,
892 k: usize,
893 plan: &crate::structures::TqIvfQueryPlan,
894 combiner: crate::query::MultiValueCombiner,
895 ) -> io::Result<(Vec<AnnDocumentCandidate>, DenseAnnScanStats)> {
896 self.search_ivf_tq_combined_documents_impl(
897 k,
898 plan,
899 combiner,
900 IVF_PARALLEL_SCAN_MIN_POSTINGS,
901 IVF_TQ_PARALLEL_SCAN_CHUNK_BLOCKS,
902 )
903 }
904
905 #[cfg(test)]
906 fn search_ivf_tq_combined_documents_with_tuning(
907 &self,
908 k: usize,
909 plan: &crate::structures::TqIvfQueryPlan,
910 combiner: crate::query::MultiValueCombiner,
911 parallel_min_postings: usize,
912 parallel_chunk_blocks: usize,
913 ) -> io::Result<Vec<AnnDocumentCandidate>> {
914 self.search_ivf_tq_combined_documents_impl(
915 k,
916 plan,
917 combiner,
918 parallel_min_postings,
919 parallel_chunk_blocks,
920 )
921 .map(|(candidates, _)| candidates)
922 }
923
924 fn search_ivf_tq_combined_documents_impl(
925 &self,
926 k: usize,
927 plan: &crate::structures::TqIvfQueryPlan,
928 combiner: crate::query::MultiValueCombiner,
929 parallel_min_postings: usize,
930 parallel_chunk_blocks: usize,
931 ) -> io::Result<(Vec<AnnDocumentCandidate>, DenseAnnScanStats)> {
932 use crate::structures::vector::quantization::{
933 TQ_BLOCK_LANES, tq_ivf_block_bytes, tq_score_ivf_block,
934 };
935
936 validate_combined_search(combiner)?;
937 let tq_plan = plan.tq_plan();
938 self.validate_ivf_tq_query_plan(plan)?;
939 #[cfg(not(feature = "native"))]
940 let _ = (parallel_min_postings, parallel_chunk_blocks);
941 if k == 0 {
942 return Ok((Vec::new(), DenseAnnScanStats::default()));
943 }
944 #[cfg(feature = "native")]
945 self.prefetch_cluster_runs(&plan.cluster_ids);
946
947 let block_bytes = tq_ivf_block_bytes(self.header.code_size);
948 let bytes = self.raw.as_slice();
949 let posting_count = probed_posting_count(self, &plan.cluster_ids)?;
950 let mut stats = DenseAnnScanStats {
951 posting_count,
952 ..DenseAnnScanStats::default()
953 };
954
955 #[cfg(feature = "native")]
956 if rayon::current_num_threads() > 1 && posting_count >= parallel_min_postings {
957 use rayon::prelude::*;
958 let tasks = self.ivf_tq_scan_tasks(plan, block_bytes, parallel_chunk_blocks);
959 let (ordinal_scores, scored_blocks, non_finite) = tasks
960 .par_iter()
961 .fold(
962 || {
963 (
964 Vec::with_capacity(parallel_chunk_blocks * TQ_BLOCK_LANES),
965 0usize,
966 0usize,
967 )
968 },
969 |(mut ordinal_scores, scored, non_finite), task| {
970 let (task_scored, task_non_finite) = score_ivf_tq_combined_blocks(
971 bytes,
972 tq_plan,
973 block_bytes,
974 *task,
975 &mut ordinal_scores,
976 );
977 (
978 ordinal_scores,
979 scored + task_scored,
980 non_finite + task_non_finite,
981 )
982 },
983 )
984 .reduce(
985 || (Vec::new(), 0usize, 0usize),
986 |(mut left, left_scored, left_non_finite),
987 (mut right, right_scored, right_non_finite)| {
988 left.append(&mut right);
989 (
990 left,
991 left_scored + right_scored,
992 left_non_finite + right_non_finite,
993 )
994 },
995 );
996 stats.scored_blocks = scored_blocks;
997 stats.non_finite_dropped = non_finite;
998 stats.parallel = true;
999 return Ok((combine_scored_ordinals(ordinal_scores, k, combiner), stats));
1000 }
1001
1002 let mut ordinal_scores = Vec::new();
1003 ordinal_scores
1004 .try_reserve_exact(posting_count)
1005 .map_err(|_| invalid_data("IVF-TQ combined score buffer allocation failed"))?;
1006 let mut scores = [0.0f32; TQ_BLOCK_LANES];
1007 for (cluster_id, cluster_dot) in plan.cluster_dots() {
1008 for run in self.cluster_runs(cluster_id) {
1009 let codes = &bytes[run.codes.clone()];
1010 for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
1011 tq_score_ivf_block(tq_plan, block, cluster_dot, &mut scores);
1012 stats.scored_blocks += 1;
1013 let lane_base = block_index * TQ_BLOCK_LANES;
1014 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
1015 for (lane, &score) in scores.iter().enumerate().take(lanes) {
1016 if !score.is_finite() {
1017 stats.non_finite_dropped += 1;
1018 continue;
1019 }
1020 let index = lane_base + lane;
1021 ordinal_scores.push((
1022 run_doc_id_validated(bytes, run, index),
1023 read_u16(bytes, run.ordinals.start + index * 2),
1024 score,
1025 ));
1026 }
1027 }
1028 }
1029 }
1030 Ok((combine_scored_ordinals(ordinal_scores, k, combiner), stats))
1031 }
1032
1033 pub(crate) fn search_binary_combined_documents(
1042 &self,
1043 k: usize,
1044 query: &[u8],
1045 cluster_ids: &[u32],
1046 combiner: crate::query::MultiValueCombiner,
1047 ) -> io::Result<CombinedBinaryCandidates> {
1048 self.search_binary_combined_documents_with_tuning(
1049 k,
1050 query,
1051 cluster_ids,
1052 combiner,
1053 IVF_PARALLEL_SCAN_MIN_POSTINGS,
1054 )
1055 }
1056
1057 pub(crate) fn search_scann_ah_combined_documents(
1060 &self,
1061 k: usize,
1062 query: &crate::structures::vector::scann::FloatScannQuery,
1063 combiner: crate::query::MultiValueCombiner,
1064 ) -> io::Result<Vec<AnnDocumentCandidate>> {
1065 self.search_scann_ah_combined_documents_with_tuning(
1066 k,
1067 query,
1068 combiner,
1069 IVF_PARALLEL_SCAN_MIN_POSTINGS,
1070 )
1071 }
1072
1073 fn search_scann_ah_combined_documents_with_tuning(
1085 &self,
1086 k: usize,
1087 query: &crate::structures::vector::scann::FloatScannQuery,
1088 combiner: crate::query::MultiValueCombiner,
1089 parallel_min_postings: usize,
1090 ) -> io::Result<Vec<AnnDocumentCandidate>> {
1091 validate_combined_search(combiner)?;
1092 if self.header.kind != AnnKind::ScannAh {
1093 return Err(invalid_data(
1094 "ScaNN AH query used with a different ANN payload",
1095 ));
1096 }
1097 #[cfg(not(feature = "native"))]
1098 let _ = parallel_min_postings;
1099 if k == 0 {
1100 return Ok(Vec::new());
1101 }
1102 #[cfg(feature = "native")]
1103 self.prefetch_cluster_runs(query.routed_leaves());
1104 let blocks = self.header.dim.div_ceil(self.header.code_size);
1105 if query.fast_scan().blocks() != blocks {
1106 return Err(invalid_data(format!(
1107 "ScaNN query has {} AH blocks but the payload stores {blocks}",
1108 query.fast_scan().blocks()
1109 )));
1110 }
1111 let geometry = ScannLeafGeometry {
1112 packed_block_bytes: query.fast_scan().packed_block_bytes(),
1113 packed_row_bytes: blocks.div_ceil(2),
1114 };
1115 let bytes = self.raw.as_slice();
1116 let posting_count = probed_posting_count(self, query.routed_leaves())?;
1117
1118 if matches!(combiner, crate::query::MultiValueCombiner::Max) {
1119 #[cfg(feature = "native")]
1120 if rayon::current_num_threads() > 1 && posting_count >= parallel_min_postings {
1121 use rayon::prelude::*;
1122 let tasks = self.scann_scan_tasks(query, SCANN_PARALLEL_SCAN_CHUNK_ROWS);
1123 let collector = tasks
1124 .par_iter()
1125 .try_fold(
1126 || BoundedAnnCollector::<true, true>::new(k),
1127 |mut collector, task| {
1128 scan_scann_ah_task(bytes, query, geometry, *task, &mut collector)?;
1129 Ok::<_, io::Error>(collector)
1130 },
1131 )
1132 .try_reduce(
1133 || BoundedAnnCollector::<true, true>::new(k),
1134 |mut left, right| {
1135 left.merge_from(right);
1136 Ok(left)
1137 },
1138 )?;
1139 return Ok(scann_document_candidates(collector));
1140 }
1141 let mut collector = BoundedAnnCollector::<true, true>::new(k);
1142 for (leaf, centroid_dot) in query.routed() {
1143 for run in self.cluster_runs(leaf) {
1144 scan_scann_ah_task(
1145 bytes,
1146 query,
1147 geometry,
1148 ScannScanTask {
1149 run,
1150 centroid_dot,
1151 first_row: 0,
1152 count: run.count,
1153 },
1154 &mut collector,
1155 )?;
1156 }
1157 }
1158 return Ok(scann_document_candidates(collector));
1159 }
1160
1161 #[cfg(feature = "native")]
1162 if rayon::current_num_threads() > 1 && posting_count >= parallel_min_postings {
1163 use rayon::prelude::*;
1164 let tasks = self.scann_scan_tasks(query, SCANN_PARALLEL_SCAN_CHUNK_ROWS);
1165 let ordinal_scores = tasks
1166 .par_iter()
1167 .try_fold(
1168 || Vec::with_capacity(SCANN_PARALLEL_SCAN_CHUNK_ROWS),
1169 |mut ordinal_scores, task| {
1170 scan_scann_ah_task(bytes, query, geometry, *task, &mut ordinal_scores)?;
1171 Ok::<_, io::Error>(ordinal_scores)
1172 },
1173 )
1174 .try_reduce(Vec::new, |mut left, mut right| {
1175 left.append(&mut right);
1176 Ok(left)
1177 })?;
1178 return Ok(combine_scored_ordinals(ordinal_scores, k, combiner));
1179 }
1180 let mut ordinal_scores = Vec::new();
1181 ordinal_scores
1182 .try_reserve_exact(posting_count)
1183 .map_err(|_| invalid_data("ScaNN combined score buffer allocation failed"))?;
1184 for (leaf, centroid_dot) in query.routed() {
1185 for run in self.cluster_runs(leaf) {
1186 scan_scann_ah_task(
1187 bytes,
1188 query,
1189 geometry,
1190 ScannScanTask {
1191 run,
1192 centroid_dot,
1193 first_row: 0,
1194 count: run.count,
1195 },
1196 &mut ordinal_scores,
1197 )?;
1198 }
1199 }
1200 Ok(combine_scored_ordinals(ordinal_scores, k, combiner))
1201 }
1202
1203 #[cfg(feature = "native")]
1206 fn scann_scan_tasks<'a>(
1207 &'a self,
1208 query: &crate::structures::vector::scann::FloatScannQuery,
1209 chunk_rows: usize,
1210 ) -> Vec<ScannScanTask<'a>> {
1211 debug_assert!(chunk_rows.is_multiple_of(crate::structures::vector::scann::FAST_SCAN_LANES));
1212 let mut tasks = Vec::new();
1213 for (leaf, centroid_dot) in query.routed() {
1214 for run in self.cluster_runs(leaf) {
1215 for first_row in (0..run.count).step_by(chunk_rows) {
1216 tasks.push(ScannScanTask {
1217 run,
1218 centroid_dot,
1219 first_row,
1220 count: chunk_rows.min(run.count - first_row),
1221 });
1222 }
1223 }
1224 }
1225 tasks
1226 }
1227
1228 fn search_binary_combined_documents_with_tuning(
1229 &self,
1230 k: usize,
1231 query: &[u8],
1232 cluster_ids: &[u32],
1233 combiner: crate::query::MultiValueCombiner,
1234 parallel_min_postings: usize,
1235 ) -> io::Result<CombinedBinaryCandidates> {
1236 validate_combined_search(combiner)?;
1237 #[cfg(not(feature = "native"))]
1238 let _ = parallel_min_postings;
1239 if query.len() != self.header.code_size {
1240 return Err(io::Error::new(
1241 io::ErrorKind::InvalidInput,
1242 "binary ANN query has the wrong byte length",
1243 ));
1244 }
1245 if k == 0 {
1246 return Ok((Vec::new(), Vec::new()));
1247 }
1248 #[cfg(feature = "native")]
1249 self.prefetch_cluster_runs(cluster_ids);
1250
1251 let bytes = self.raw.as_slice();
1252 let posting_count = probed_posting_count(self, cluster_ids)?;
1253 #[cfg(feature = "native")]
1254 if rayon::current_num_threads() > 1 && posting_count >= parallel_min_postings {
1255 use rayon::prelude::*;
1256 let tasks = self.binary_scan_tasks(cluster_ids);
1257 let (ordinal_scores, _) = tasks
1258 .par_iter()
1259 .try_fold(
1260 || {
1261 (
1262 Vec::with_capacity(BINARY_SCORE_BATCH),
1263 vec![0.0f32; BINARY_SCORE_BATCH],
1264 )
1265 },
1266 |(mut ordinal_scores, mut scores), task| {
1267 score_binary_task(
1268 bytes,
1269 query,
1270 self.header.dim,
1271 self.header.code_size,
1272 *task,
1273 &mut scores,
1274 &mut ordinal_scores,
1275 )?;
1276 Ok::<_, io::Error>((ordinal_scores, scores))
1277 },
1278 )
1279 .try_reduce(
1280 || (Vec::new(), Vec::new()),
1281 |(mut left, scores), (mut right, _)| {
1282 left.append(&mut right);
1283 Ok((left, scores))
1284 },
1285 )?;
1286 return Ok(combine_scored_ordinals_retaining(
1287 ordinal_scores,
1288 k,
1289 combiner,
1290 ));
1291 }
1292
1293 let mut score_batch = vec![0.0f32; BINARY_SCORE_BATCH.min(self.header.vector_count)];
1294 let mut ordinal_scores = Vec::new();
1295 ordinal_scores
1296 .try_reserve_exact(posting_count)
1297 .map_err(|_| invalid_data("binary combined score buffer allocation failed"))?;
1298 score_binary_cluster_runs(
1299 self,
1300 bytes,
1301 query,
1302 cluster_ids,
1303 &mut score_batch,
1304 &mut ordinal_scores,
1305 )?;
1306 Ok(combine_scored_ordinals_retaining(
1307 ordinal_scores,
1308 k,
1309 combiner,
1310 ))
1311 }
1312
1313 #[cfg(test)]
1316 pub(crate) fn search_tq_distinct(
1317 &self,
1318 k: usize,
1319 plan: &crate::structures::TqQueryPlan,
1320 ) -> io::Result<Vec<(u32, u16, f32)>> {
1321 self.search_tq_distinct_generic::<BoundedAnnCollector<true, true>>(k, plan)
1322 .map(|(results, _)| results)
1323 }
1324
1325 pub(crate) fn search_tq_distinct_with_stats(
1330 &self,
1331 k: usize,
1332 plan: &crate::structures::TqQueryPlan,
1333 unique_keys: bool,
1334 ) -> io::Result<ScanResults> {
1335 if unique_keys {
1336 self.search_tq_distinct_generic::<BoundedUniqueAnnCollector<true>>(k, plan)
1337 } else {
1338 self.search_tq_distinct_generic::<BoundedAnnCollector<true, true>>(k, plan)
1339 }
1340 }
1341
1342 fn search_tq_distinct_generic<C: IvfTqTopK>(
1343 &self,
1344 k: usize,
1345 plan: &crate::structures::TqQueryPlan,
1346 ) -> io::Result<ScanResults> {
1347 use crate::structures::vector::quantization::{
1348 TQ_BLOCK_LANES, tq_block_bytes, tq_score_block,
1349 };
1350
1351 if plan.padded_dim() != self.header.code_size * 2 {
1352 return Err(io::Error::new(
1353 io::ErrorKind::InvalidInput,
1354 "TQ query plan does not match the payload dimension",
1355 ));
1356 }
1357 let block_bytes = tq_block_bytes(self.header.code_size);
1358 let bytes = self.raw.as_slice();
1359 let mut stats = DenseAnnScanStats {
1360 posting_count: self.header.vector_count,
1361 ..DenseAnnScanStats::default()
1362 };
1363
1364 let score_lanes = |collector: &mut C,
1368 run: &AnnRun,
1369 block_index: usize,
1370 scores: &[f32; TQ_BLOCK_LANES]| {
1371 let lane_base = block_index * TQ_BLOCK_LANES;
1372 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
1373 let mut threshold = collector.pruning_threshold();
1374 for (lane, &score) in scores.iter().enumerate().take(lanes) {
1375 if threshold.is_some_and(|threshold| score < threshold) {
1376 continue;
1377 }
1378 let index = lane_base + lane;
1379 collector.insert(
1380 run_doc_id_validated(bytes, run, index),
1381 read_u16(bytes, run.ordinals.start + index * 2),
1382 score,
1383 );
1384 threshold = collector.pruning_threshold();
1385 }
1386 };
1387
1388 #[cfg(feature = "native")]
1393 if self.header.vector_count >= TQ_PARALLEL_SCAN_MIN_VECTORS {
1394 use rayon::prelude::*;
1395 let (collector, scored_blocks) = self
1396 .runs
1397 .par_iter()
1398 .flat_map(|run| {
1399 let codes = &bytes[run.codes.clone()];
1400 let blocks = codes.len() / block_bytes;
1401 (0..blocks.div_ceil(TQ_PARALLEL_SCAN_CHUNK_BLOCKS))
1402 .into_par_iter()
1403 .map(move |chunk| (run, chunk * TQ_PARALLEL_SCAN_CHUNK_BLOCKS, blocks))
1404 })
1405 .fold(
1406 || (C::with_k(k), 0usize),
1407 |(mut collector, scored), (run, first_block, total_blocks)| {
1408 let codes = &bytes[run.codes.clone()];
1409 let last_block =
1410 (first_block + TQ_PARALLEL_SCAN_CHUNK_BLOCKS).min(total_blocks);
1411 let mut scores = [0.0f32; TQ_BLOCK_LANES];
1412 for block_index in first_block..last_block {
1413 let block = &codes[block_index * block_bytes..][..block_bytes];
1414 tq_score_block(plan, block, &mut scores);
1415 score_lanes(&mut collector, run, block_index, &scores);
1416 }
1417 (collector, scored + (last_block - first_block))
1418 },
1419 )
1420 .reduce(
1421 || (C::with_k(k), 0usize),
1422 |(mut collector, left_scored), (partial, right_scored)| {
1423 collector.merge_from(partial);
1424 (collector, left_scored + right_scored)
1425 },
1426 );
1427 stats.scored_blocks = scored_blocks;
1428 stats.non_finite_dropped = collector.non_finite_dropped();
1429 stats.parallel = true;
1430 return Ok((collector.into_sorted_results(), stats));
1431 }
1432
1433 let mut collector = C::with_k(k);
1434 let mut scores = [0.0f32; TQ_BLOCK_LANES];
1435 for run in &self.runs {
1436 let codes = &bytes[run.codes.clone()];
1437 for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
1438 tq_score_block(plan, block, &mut scores);
1439 stats.scored_blocks += 1;
1440 score_lanes(&mut collector, run, block_index, &scores);
1441 }
1442 }
1443 stats.non_finite_dropped = collector.non_finite_dropped();
1444 Ok((collector.into_sorted_results(), stats))
1445 }
1446
1447 #[cfg(test)]
1455 pub(crate) fn search_ivf_tq_distinct(
1456 &self,
1457 k: usize,
1458 plan: &crate::structures::TqIvfQueryPlan,
1459 ) -> io::Result<Vec<(u32, u16, f32)>> {
1460 self.search_ivf_tq_distinct_generic::<BoundedAnnCollector<true, true>>(
1461 k,
1462 plan,
1463 IVF_PARALLEL_SCAN_MIN_POSTINGS,
1464 IVF_TQ_PARALLEL_SCAN_CHUNK_BLOCKS,
1465 )
1466 .map(|(results, _)| results)
1467 }
1468
1469 pub(crate) fn search_ivf_tq_distinct_with_stats(
1474 &self,
1475 k: usize,
1476 plan: &crate::structures::TqIvfQueryPlan,
1477 unique_keys: bool,
1478 ) -> io::Result<ScanResults> {
1479 if unique_keys {
1480 self.search_ivf_tq_distinct_generic::<BoundedUniqueAnnCollector<true>>(
1481 k,
1482 plan,
1483 IVF_PARALLEL_SCAN_MIN_POSTINGS,
1484 IVF_TQ_PARALLEL_SCAN_CHUNK_BLOCKS,
1485 )
1486 } else {
1487 self.search_ivf_tq_distinct_generic::<BoundedAnnCollector<true, true>>(
1488 k,
1489 plan,
1490 IVF_PARALLEL_SCAN_MIN_POSTINGS,
1491 IVF_TQ_PARALLEL_SCAN_CHUNK_BLOCKS,
1492 )
1493 }
1494 }
1495
1496 #[cfg(test)]
1497 fn search_ivf_tq_distinct_with_tuning(
1498 &self,
1499 k: usize,
1500 plan: &crate::structures::TqIvfQueryPlan,
1501 parallel_min_postings: usize,
1502 parallel_chunk_blocks: usize,
1503 ) -> io::Result<Vec<(u32, u16, f32)>> {
1504 self.search_ivf_tq_distinct_generic::<BoundedAnnCollector<true, true>>(
1505 k,
1506 plan,
1507 parallel_min_postings,
1508 parallel_chunk_blocks,
1509 )
1510 .map(|(results, _)| results)
1511 }
1512
1513 fn search_ivf_tq_distinct_generic<C: IvfTqTopK>(
1514 &self,
1515 k: usize,
1516 plan: &crate::structures::TqIvfQueryPlan,
1517 parallel_min_postings: usize,
1518 parallel_chunk_blocks: usize,
1519 ) -> io::Result<ScanResults> {
1520 use crate::structures::vector::quantization::tq_ivf_block_bytes;
1521
1522 let tq_plan = plan.tq_plan();
1523 self.validate_ivf_tq_query_plan(plan)?;
1524 #[cfg(not(feature = "native"))]
1525 let _ = (parallel_min_postings, parallel_chunk_blocks);
1526 if k == 0 {
1527 return Ok((Vec::new(), DenseAnnScanStats::default()));
1528 }
1529 #[cfg(feature = "native")]
1530 self.prefetch_cluster_runs(&plan.cluster_ids);
1531 let block_bytes = tq_ivf_block_bytes(self.header.code_size);
1532 let bytes = self.raw.as_slice();
1533 let posting_count = probed_posting_count(self, &plan.cluster_ids)?;
1534 let mut stats = DenseAnnScanStats {
1535 posting_count,
1536 ..DenseAnnScanStats::default()
1537 };
1538
1539 #[cfg(feature = "native")]
1544 if rayon::current_num_threads() > 1 && posting_count >= parallel_min_postings {
1545 use rayon::prelude::*;
1546 let tasks = self.ivf_tq_scan_tasks(plan, block_bytes, parallel_chunk_blocks);
1547
1548 if let Some((pilot, remaining)) = tasks.split_first() {
1549 let mut pilot_collector = C::with_k(k);
1550 let (pilot_pruned, pilot_scored) =
1551 score_ivf_tq_blocks(bytes, tq_plan, block_bytes, *pilot, &mut pilot_collector);
1552 let pilot_non_finite = pilot_collector.non_finite_dropped();
1553 let seed = pilot_collector.into_sorted_results();
1554 let seeded_collector = || {
1555 let mut collector = C::with_k(k);
1556 for &(doc_id, ordinal, score) in &seed {
1557 collector.insert(doc_id, ordinal, score);
1558 }
1559 collector
1560 };
1561 let (collector, pruned_blocks, scored_blocks) = remaining
1562 .par_iter()
1563 .fold(
1564 || (seeded_collector(), 0usize, 0usize),
1565 |(mut collector, pruned, scored), task| {
1566 let (task_pruned, task_scored) = score_ivf_tq_blocks(
1567 bytes,
1568 tq_plan,
1569 block_bytes,
1570 *task,
1571 &mut collector,
1572 );
1573 (collector, pruned + task_pruned, scored + task_scored)
1574 },
1575 )
1576 .reduce(
1577 || (seeded_collector(), 0usize, 0usize),
1578 |(mut collector, left_pruned, left_scored),
1579 (partial, right_pruned, right_scored)| {
1580 collector.merge_from(partial);
1581 (
1582 collector,
1583 left_pruned + right_pruned,
1584 left_scored + right_scored,
1585 )
1586 },
1587 );
1588 stats.pruned_blocks = pilot_pruned + pruned_blocks;
1589 stats.scored_blocks = pilot_scored + scored_blocks;
1590 stats.non_finite_dropped = pilot_non_finite + collector.non_finite_dropped();
1591 stats.parallel = true;
1592 log_ivf_tq_pruning(stats.pruned_blocks, stats.scored_blocks, true);
1593 return Ok((collector.into_sorted_results(), stats));
1594 }
1595 }
1596
1597 let mut collector = C::with_k(k);
1598 for (cluster_id, cluster_dot) in plan.cluster_dots() {
1599 for run in self.cluster_runs(cluster_id) {
1600 let block_count = run.codes.len() / block_bytes;
1601 let (run_pruned, run_scored) = score_ivf_tq_blocks(
1602 bytes,
1603 tq_plan,
1604 block_bytes,
1605 IvfTqScanTask {
1606 run,
1607 cluster_dot,
1608 first_block: 0,
1609 last_block: block_count,
1610 },
1611 &mut collector,
1612 );
1613 stats.pruned_blocks += run_pruned;
1614 stats.scored_blocks += run_scored;
1615 }
1616 }
1617 stats.non_finite_dropped = collector.non_finite_dropped();
1618 log_ivf_tq_pruning(stats.pruned_blocks, stats.scored_blocks, false);
1619 Ok((collector.into_sorted_results(), stats))
1620 }
1621
1622 fn validate_ivf_tq_query_plan(
1623 &self,
1624 plan: &crate::structures::TqIvfQueryPlan,
1625 ) -> io::Result<()> {
1626 if self.header.kind != AnnKind::IvfTq
1627 || !crate::structures::is_ivf_tq_cosine_generation(self.header.quantizer_version)
1628 {
1629 return Err(invalid_data(
1630 "legacy raw IVF-TQ payloads cannot be searched; rebuild the index",
1631 ));
1632 }
1633 if plan.tq_plan().padded_dim() != self.header.code_size * 2
1634 || plan.quantizer_version != self.header.quantizer_version
1635 || plan.fingerprint != self.header.codebook_version
1636 {
1637 return Err(io::Error::new(
1638 io::ErrorKind::InvalidInput,
1639 "IVF-TQ query plan does not match the payload generation",
1640 ));
1641 }
1642 Ok(())
1643 }
1644
1645 pub(crate) fn search_binary_clusters<const BY_DOCUMENT: bool>(
1646 &self,
1647 query: &[u8],
1648 k: usize,
1649 cluster_ids: &[u32],
1650 ) -> io::Result<Vec<(u32, u16, f32)>> {
1651 self.search_binary_clusters_with_tuning::<BY_DOCUMENT>(
1652 query,
1653 k,
1654 cluster_ids,
1655 IVF_PARALLEL_SCAN_MIN_POSTINGS,
1656 )
1657 }
1658
1659 fn search_binary_clusters_with_tuning<const BY_DOCUMENT: bool>(
1660 &self,
1661 query: &[u8],
1662 k: usize,
1663 cluster_ids: &[u32],
1664 parallel_min_postings: usize,
1665 ) -> io::Result<Vec<(u32, u16, f32)>> {
1666 #[cfg(not(feature = "native"))]
1667 let _ = parallel_min_postings;
1668 if query.len() != self.header.code_size {
1669 return Err(io::Error::new(
1670 io::ErrorKind::InvalidInput,
1671 "binary ANN query has the wrong byte length",
1672 ));
1673 }
1674 if k == 0 {
1675 return Ok(Vec::new());
1676 }
1677 #[cfg(feature = "native")]
1678 self.prefetch_cluster_runs(cluster_ids);
1679 let bytes = self.raw.as_slice();
1680
1681 debug_assert!(
1684 BY_DOCUMENT || self.header.kind == AnnKind::ScannBinary || {
1685 let mut seen = rustc_hash::FxHashSet::default();
1686 cluster_ids
1687 .iter()
1688 .all(|cluster_id| seen.insert(*cluster_id))
1689 },
1690 "an IVF probe plan must not repeat a cluster",
1691 );
1692
1693 #[cfg(feature = "native")]
1694 if rayon::current_num_threads() > 1
1695 && probed_posting_count(self, cluster_ids)? >= parallel_min_postings
1696 {
1697 use rayon::prelude::*;
1698 let tasks = self.binary_scan_tasks(cluster_ids);
1699 let (collector, _) = tasks
1700 .par_iter()
1701 .try_fold(
1702 || {
1703 (
1704 BoundedAnnCollector::<BY_DOCUMENT, true>::new(k),
1705 vec![0.0f32; BINARY_SCORE_BATCH],
1706 )
1707 },
1708 |(mut collector, mut scores), task| {
1709 score_binary_task(
1710 bytes,
1711 query,
1712 self.header.dim,
1713 self.header.code_size,
1714 *task,
1715 &mut scores,
1716 &mut collector,
1717 )?;
1718 Ok::<_, io::Error>((collector, scores))
1719 },
1720 )
1721 .try_reduce(
1722 || (BoundedAnnCollector::<BY_DOCUMENT, true>::new(k), Vec::new()),
1723 |(mut collector, scores), (partial, _)| {
1724 collector.merge_from(partial);
1725 Ok((collector, scores))
1726 },
1727 )?;
1728 return Ok(collector.into_sorted_results());
1729 }
1730
1731 let mut scores = vec![0.0f32; BINARY_SCORE_BATCH.min(self.header.vector_count)];
1732 if BY_DOCUMENT || self.header.kind == AnnKind::ScannBinary {
1733 let mut collector = BoundedAnnCollector::<BY_DOCUMENT, true>::new(k);
1734 score_binary_cluster_runs(
1735 self,
1736 bytes,
1737 query,
1738 cluster_ids,
1739 &mut scores,
1740 &mut collector,
1741 )?;
1742 return Ok(collector.into_sorted_results());
1743 }
1744
1745 let mut collector = BoundedUniqueAnnCollector::<true>::new(k);
1747 score_binary_cluster_runs(self, bytes, query, cluster_ids, &mut scores, &mut collector)?;
1748 Ok(collector.into_sorted_results())
1749 }
1750}
1751
1752fn score_ivf_tq_blocks<C: IvfTqTopK>(
1758 bytes: &[u8],
1759 plan: &crate::structures::TqQueryPlan,
1760 block_bytes: usize,
1761 task: IvfTqScanTask<'_>,
1762 collector: &mut C,
1763) -> (usize, usize) {
1764 use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_score_ivf_block};
1765
1766 let run = task.run;
1767 let codes = &bytes[run.codes.clone()];
1768 let mut scores = [0.0f32; TQ_BLOCK_LANES];
1769 let mut scored_blocks = 0usize;
1770 let mut threshold = collector.pruning_threshold();
1771 for block_index in task.first_block..task.last_block {
1772 let block = &codes[block_index * block_bytes..][..block_bytes];
1773 if let Some(threshold) = threshold
1774 && task.cluster_dot + tq_ivf_block_max_scale(block) * TQ_PRUNE_ESTIMATE_BOUND
1775 <= threshold
1776 {
1777 return (task.last_block - block_index, scored_blocks);
1778 }
1779 scored_blocks += 1;
1780 tq_score_ivf_block(plan, block, task.cluster_dot, &mut scores);
1781 let lane_base = block_index * TQ_BLOCK_LANES;
1782 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
1783 for (lane, &score) in scores.iter().enumerate().take(lanes) {
1784 if threshold.is_some_and(|threshold| score < threshold) {
1789 continue;
1790 }
1791 let index = lane_base + lane;
1792 collector.insert(
1793 run_doc_id_validated(bytes, run, index),
1794 read_u16(bytes, run.ordinals.start + index * 2),
1795 score,
1796 );
1797 threshold = collector.pruning_threshold();
1798 }
1799 }
1800 (0, scored_blocks)
1801}
1802
1803#[cfg(feature = "native")]
1805fn score_ivf_tq_combined_blocks(
1806 bytes: &[u8],
1807 plan: &crate::structures::TqQueryPlan,
1808 block_bytes: usize,
1809 task: IvfTqScanTask<'_>,
1810 ordinal_scores: &mut Vec<(u32, u16, f32)>,
1811) -> (usize, usize) {
1812 use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_score_ivf_block};
1813
1814 let run = task.run;
1815 let codes = &bytes[run.codes.clone()];
1816 let mut scores = [0.0f32; TQ_BLOCK_LANES];
1817 let mut non_finite = 0usize;
1818 for block_index in task.first_block..task.last_block {
1819 let block = &codes[block_index * block_bytes..][..block_bytes];
1820 tq_score_ivf_block(plan, block, task.cluster_dot, &mut scores);
1821 let lane_base = block_index * TQ_BLOCK_LANES;
1822 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
1823 for (lane, &score) in scores.iter().enumerate().take(lanes) {
1824 if !score.is_finite() {
1825 non_finite += 1;
1826 continue;
1827 }
1828 let index = lane_base + lane;
1829 ordinal_scores.push((
1830 run_doc_id_validated(bytes, run, index),
1831 read_u16(bytes, run.ordinals.start + index * 2),
1832 score,
1833 ));
1834 }
1835 }
1836 (task.last_block - task.first_block, non_finite)
1837}
1838
1839fn log_ivf_tq_pruning(pruned_blocks: usize, scored_blocks: usize, parallel: bool) {
1840 if pruned_blocks > 0 {
1841 log::debug!(
1842 "[search_ivf_tq] pruned {pruned_blocks} of {} blocks via scale bounds ({})",
1843 pruned_blocks + scored_blocks,
1844 if parallel { "parallel" } else { "serial" },
1845 );
1846 }
1847}
1848
1849#[cfg(feature = "native")]
1850fn coalesce_prefetch_ranges(ranges: &mut Vec<Range<usize>>) {
1851 if ranges.len() < 2 {
1852 return;
1853 }
1854 ranges.sort_unstable_by_key(|range| range.start);
1855 let mut output_len = 1usize;
1856 for input_index in 1..ranges.len() {
1857 let next_start = ranges[input_index].start;
1858 let next_end = ranges[input_index].end;
1859 let previous = &mut ranges[output_len - 1];
1860 if next_start <= previous.end.saturating_add(PREFETCH_COALESCE_GAP) {
1861 previous.end = previous.end.max(next_end);
1862 } else {
1863 ranges[output_len] = next_start..next_end;
1864 output_len += 1;
1865 }
1866 }
1867 ranges.truncate(output_len);
1868}
1869
1870trait AnnScoreSink {
1871 fn insert_score(&mut self, doc_id: u32, ordinal: u16, score: f32);
1872}
1873
1874impl<const BY_DOCUMENT: bool> AnnScoreSink for BoundedAnnCollector<BY_DOCUMENT, true> {
1875 #[inline]
1876 fn insert_score(&mut self, doc_id: u32, ordinal: u16, score: f32) {
1877 self.insert(doc_id, ordinal, score);
1878 }
1879}
1880
1881impl AnnScoreSink for BoundedUniqueAnnCollector<true> {
1882 #[inline]
1883 fn insert_score(&mut self, doc_id: u32, ordinal: u16, score: f32) {
1884 self.insert(doc_id, ordinal, score);
1885 }
1886}
1887
1888impl AnnScoreSink for Vec<(u32, u16, f32)> {
1889 #[inline]
1890 fn insert_score(&mut self, doc_id: u32, ordinal: u16, score: f32) {
1891 if score.is_finite() {
1892 self.push((doc_id, ordinal, score));
1893 }
1894 }
1895}
1896
1897fn probed_posting_count(index: &AnnDiskIndex, cluster_ids: &[u32]) -> io::Result<usize> {
1898 cluster_ids.iter().try_fold(0usize, |count, &cluster_id| {
1899 index
1900 .cluster_runs(cluster_id)
1901 .iter()
1902 .try_fold(count, |count, run| {
1903 count
1904 .checked_add(run.count)
1905 .ok_or_else(|| invalid_data("ANN probed posting count overflows usize"))
1906 })
1907 })
1908}
1909
1910fn score_binary_cluster_runs(
1911 index: &AnnDiskIndex,
1912 bytes: &[u8],
1913 query: &[u8],
1914 cluster_ids: &[u32],
1915 scores: &mut [f32],
1916 collector: &mut impl AnnScoreSink,
1917) -> io::Result<()> {
1918 for &cluster_id in cluster_ids {
1919 for run in index.cluster_runs(cluster_id) {
1920 score_binary_run(
1921 bytes,
1922 run,
1923 query,
1924 index.header.dim,
1925 index.header.code_size,
1926 scores,
1927 collector,
1928 )?;
1929 }
1930 }
1931 Ok(())
1932}
1933
1934fn score_binary_run(
1935 bytes: &[u8],
1936 run: &AnnRun,
1937 query: &[u8],
1938 dim_bits: usize,
1939 code_size: usize,
1940 scores: &mut [f32],
1941 collector: &mut impl AnnScoreSink,
1942) -> io::Result<()> {
1943 for batch_start in (0..run.count).step_by(BINARY_SCORE_BATCH) {
1944 score_binary_task(
1945 bytes,
1946 query,
1947 dim_bits,
1948 code_size,
1949 BinaryScanTask {
1950 run,
1951 first_index: batch_start,
1952 count: BINARY_SCORE_BATCH.min(run.count - batch_start),
1953 },
1954 scores,
1955 collector,
1956 )?;
1957 }
1958 Ok(())
1959}
1960
1961fn score_binary_task(
1962 bytes: &[u8],
1963 query: &[u8],
1964 dim_bits: usize,
1965 code_size: usize,
1966 task: BinaryScanTask<'_>,
1967 scores: &mut [f32],
1968 collector: &mut impl AnnScoreSink,
1969) -> io::Result<()> {
1970 let code_start = task.run.codes.start + task.first_index * code_size;
1971 let code_end = code_start + task.count * code_size;
1972 crate::structures::simd::batch_hamming_scores(
1973 query,
1974 &bytes[code_start..code_end],
1975 code_size,
1976 dim_bits,
1977 &mut scores[..task.count],
1978 );
1979 for (batch_index, &score) in scores.iter().enumerate().take(task.count) {
1980 let index = task.first_index + batch_index;
1981 collector.insert_score(
1982 run_doc_id_validated(bytes, task.run, index),
1983 read_u16(bytes, task.run.ordinals.start + index * 2),
1984 score,
1985 );
1986 }
1987 Ok(())
1988}
1989
1990#[inline]
1991fn retain_combined_document(
1992 collector: &mut BoundedAnnCollector<true, true>,
1993 doc_id: u32,
1994 ordinal_scores: &[(u32, f32)],
1995 combiner: crate::query::MultiValueCombiner,
1996) {
1997 if !ordinal_scores.is_empty() {
1998 collector.insert(doc_id, 0, combiner.combine(ordinal_scores));
1999 }
2000}
2001
2002fn validate_combined_search(combiner: crate::query::MultiValueCombiner) -> io::Result<()> {
2003 combiner
2004 .validate()
2005 .map_err(|message| io::Error::new(io::ErrorKind::InvalidInput, message))
2006}
2007
2008pub(crate) static NON_FINITE_ANN_SCORES_DROPPED: std::sync::atomic::AtomicU64 =
2011 std::sync::atomic::AtomicU64::new(0);
2012static NON_FINITE_ANN_SCORE_DROP_EVENTS: std::sync::atomic::AtomicU64 =
2013 std::sync::atomic::AtomicU64::new(0);
2014
2015#[cfg(feature = "native")]
2018const SCANN_PARALLEL_SCAN_CHUNK_ROWS: usize =
2019 512 * crate::structures::vector::scann::FAST_SCAN_LANES;
2020
2021#[derive(Clone, Copy)]
2023struct ScannLeafGeometry {
2024 packed_block_bytes: usize,
2026 packed_row_bytes: usize,
2028}
2029
2030#[derive(Clone, Copy)]
2032struct ScannScanTask<'a> {
2033 run: &'a AnnRun,
2034 centroid_dot: f32,
2035 first_row: usize,
2036 count: usize,
2037}
2038
2039trait ScannPruneThreshold {
2042 fn prune_threshold(&self) -> Option<f32>;
2043}
2044
2045impl ScannPruneThreshold for BoundedAnnCollector<true, true> {
2046 #[inline]
2047 fn prune_threshold(&self) -> Option<f32> {
2048 self.pruning_threshold()
2049 }
2050}
2051
2052impl ScannPruneThreshold for Vec<(u32, u16, f32)> {
2053 #[inline]
2054 fn prune_threshold(&self) -> Option<f32> {
2055 None
2058 }
2059}
2060
2061fn scan_scann_ah_task(
2065 bytes: &[u8],
2066 query: &crate::structures::vector::scann::FloatScannQuery,
2067 geometry: ScannLeafGeometry,
2068 task: ScannScanTask<'_>,
2069 sink: &mut (impl AnnScoreSink + ScannPruneThreshold),
2070) -> io::Result<()> {
2071 use crate::structures::vector::scann::FAST_SCAN_LANES;
2072
2073 let run = task.run;
2074 let codes = &bytes[run.codes.clone()];
2075 let fast = query.fast_scan();
2076 let full_rows = run.count / FAST_SCAN_LANES * FAST_SCAN_LANES;
2077 let end = task.first_row + task.count;
2078 debug_assert!(task.first_row.is_multiple_of(FAST_SCAN_LANES) || task.first_row >= full_rows);
2079 let mut lanes = [0i32; FAST_SCAN_LANES];
2080 let mut row = task.first_row;
2081 while row < full_rows.min(end) {
2082 let start = (row / FAST_SCAN_LANES) * geometry.packed_block_bytes;
2083 fast.accumulate_block(
2084 &codes[start..start + geometry.packed_block_bytes],
2085 &mut lanes,
2086 )
2087 .map_err(|error| invalid_data(error.to_string()))?;
2088 let cutoff = sink
2089 .prune_threshold()
2090 .map(|threshold| fast.lane_cutoff(threshold, task.centroid_dot));
2091 for (lane, &sum) in lanes.iter().enumerate() {
2092 if cutoff.is_some_and(|cutoff| sum <= cutoff) {
2093 continue;
2094 }
2095 let index = row + lane;
2096 sink.insert_score(
2097 run_doc_id(bytes, run, index)?,
2098 read_u16(bytes, run.ordinals.start + index * 2),
2099 fast.lane_score(sum, task.centroid_dot),
2100 );
2101 }
2102 row += FAST_SCAN_LANES;
2103 }
2104 let tail_bytes = (full_rows / FAST_SCAN_LANES) * geometry.packed_block_bytes;
2105 while row < end {
2106 let start = tail_bytes + (row - full_rows) * geometry.packed_row_bytes;
2107 let score = query
2108 .ah_query()
2109 .score_packed(
2110 &codes[start..start + geometry.packed_row_bytes],
2111 task.centroid_dot,
2112 )
2113 .map_err(|error| invalid_data(error.to_string()))?;
2114 sink.insert_score(
2115 run_doc_id(bytes, run, row)?,
2116 read_u16(bytes, run.ordinals.start + row * 2),
2117 score,
2118 );
2119 row += 1;
2120 }
2121 Ok(())
2122}
2123
2124fn scann_document_candidates(
2125 collector: BoundedAnnCollector<true, true>,
2126) -> Vec<AnnDocumentCandidate> {
2127 collector
2128 .into_sorted_results()
2129 .into_iter()
2130 .map(|(doc_id, _, score)| AnnDocumentCandidate { doc_id, score })
2131 .collect()
2132}
2133
2134fn combine_scored_ordinals(
2140 scores: Vec<(u32, u16, f32)>,
2141 k: usize,
2142 combiner: crate::query::MultiValueCombiner,
2143) -> Vec<AnnDocumentCandidate> {
2144 combine_scored_ordinals_retaining(scores, k, combiner).0
2145}
2146
2147fn combine_scored_ordinals_retaining(
2153 mut scores: Vec<(u32, u16, f32)>,
2154 k: usize,
2155 combiner: crate::query::MultiValueCombiner,
2156) -> CombinedBinaryCandidates {
2157 if k == 0 || scores.is_empty() {
2158 return (Vec::new(), Vec::new());
2159 }
2160 let scored = scores.len();
2161 scores.retain(|entry| entry.2.is_finite());
2162 let dropped = scored - scores.len();
2163 if dropped > 0 {
2164 NON_FINITE_ANN_SCORES_DROPPED
2168 .fetch_add(dropped as u64, std::sync::atomic::Ordering::Relaxed);
2169 crate::observe::ann_non_finite_scores_dropped(dropped as u64);
2170 crate::structures::vector::scann::warn_rate_limited(
2171 &NON_FINITE_ANN_SCORE_DROP_EVENTS,
2172 |events| {
2173 format!(
2174 "[ann] dropped {dropped} of {scored} non-finite leaf scores before combining \
2175 ({events} scans affected so far; total dropped {})",
2176 NON_FINITE_ANN_SCORES_DROPPED.load(std::sync::atomic::Ordering::Relaxed)
2177 )
2178 },
2179 );
2180 }
2181 let by_document_ordinal = |left: &(u32, u16, f32), right: &(u32, u16, f32)| {
2182 left.0.cmp(&right.0).then_with(|| left.1.cmp(&right.1))
2183 };
2184 #[cfg(feature = "native")]
2185 if rayon::current_num_threads() > 1 && scores.len() >= IVF_PARALLEL_SCAN_MIN_POSTINGS {
2186 use rayon::prelude::*;
2187 scores.par_sort_unstable_by(by_document_ordinal);
2188 } else {
2189 scores.sort_unstable_by(by_document_ordinal);
2190 }
2191 #[cfg(not(feature = "native"))]
2192 scores.sort_unstable_by(by_document_ordinal);
2193
2194 let mut unique_len = 0usize;
2199 for read_index in 0..scores.len() {
2200 let candidate = scores[read_index];
2201 if unique_len > 0
2202 && scores[unique_len - 1].0 == candidate.0
2203 && scores[unique_len - 1].1 == candidate.1
2204 {
2205 if candidate.2.total_cmp(&scores[unique_len - 1].2).is_gt() {
2206 scores[unique_len - 1].2 = candidate.2;
2207 }
2208 } else {
2209 scores[unique_len] = candidate;
2210 unique_len += 1;
2211 }
2212 }
2213 scores.truncate(unique_len);
2214
2215 let mut top_documents = BoundedAnnCollector::<true, true>::new(k);
2216 let mut current_doc = None;
2217 let mut ordinal_scores = Vec::new();
2218 for &(doc_id, ordinal, score) in &scores {
2219 if current_doc.is_some_and(|current| current != doc_id) {
2220 retain_combined_document(
2221 &mut top_documents,
2222 current_doc.expect("current document is present"),
2223 &ordinal_scores,
2224 combiner,
2225 );
2226 ordinal_scores.clear();
2227 }
2228 current_doc = Some(doc_id);
2229 ordinal_scores.push((u32::from(ordinal), score));
2230 }
2231 if let Some(doc_id) = current_doc {
2232 retain_combined_document(&mut top_documents, doc_id, &ordinal_scores, combiner);
2233 }
2234
2235 let candidates: Vec<AnnDocumentCandidate> = top_documents
2236 .into_sorted_results()
2237 .into_iter()
2238 .map(|(doc_id, _, score)| AnnDocumentCandidate { doc_id, score })
2239 .collect();
2240
2241 let mut retained_ids: Vec<u32> = candidates
2245 .iter()
2246 .map(|candidate| candidate.doc_id)
2247 .collect();
2248 retained_ids.sort_unstable();
2249 let mut retained = Vec::with_capacity(scores.len().min(retained_ids.len().saturating_mul(4)));
2250 let mut cursor = 0usize;
2251 for entry in scores {
2252 while cursor < retained_ids.len() && retained_ids[cursor] < entry.0 {
2253 cursor += 1;
2254 }
2255 if retained_ids.get(cursor) == Some(&entry.0) {
2256 retained.push(entry);
2257 }
2258 }
2259 (candidates, retained)
2260}
2261
2262#[cfg(feature = "native")]
2263pub(crate) fn write_built_binary_ivf(
2264 index: &BinaryIvfIndex,
2265 routing: IvfRoutingMode,
2266 writer: &mut (impl Write + ?Sized),
2267) -> io::Result<u64> {
2268 let runs: Vec<_> = index
2269 .clusters
2270 .iter()
2271 .map(|(cluster_id, cluster)| BuildRun {
2272 cluster_id: *cluster_id,
2273 doc_ids: &cluster.doc_ids,
2274 ordinals: &cluster.ordinals,
2275 codes: &cluster.codes,
2276 })
2277 .collect();
2278 write_built_runs(
2279 AnnDiskHeader {
2280 kind: AnnKind::BinaryIvf,
2281 routing,
2282 dim: index.dim_bits,
2283 code_size: index.dim_bits.div_ceil(8),
2284 num_clusters: index.num_clusters,
2285 quantizer_version: index.quantizer_version,
2286 codebook_version: 0,
2287 vector_count: index.len(),
2288 },
2289 &runs,
2290 writer,
2291 )
2292}
2293
2294#[cfg(feature = "native")]
2299#[cfg_attr(not(test), allow(dead_code))]
2300pub(crate) fn write_built_scann(
2301 payload: &crate::structures::vector::scann::ScannSegmentPayload,
2302 writer: &mut (impl Write + ?Sized),
2303) -> io::Result<u64> {
2304 use crate::structures::vector::scann::ScannEncoding;
2305
2306 let (kind, code_size) = match payload.encoding {
2307 ScannEncoding::AsymmetricHash {
2308 dimensions_per_block,
2309 bits_per_code: 4,
2310 } => (AnnKind::ScannAh, usize::from(dimensions_per_block)),
2311 ScannEncoding::AsymmetricHash { .. } => {
2312 return Err(invalid_data(
2313 "ScaNN ANN disk format supports only 4-bit AH codes",
2314 ));
2315 }
2316 ScannEncoding::BinaryHamming => (
2317 AnnKind::ScannBinary,
2318 usize::try_from(payload.dimension)
2319 .map_err(|_| invalid_data("ScaNN dimension exceeds usize"))?
2320 .div_ceil(8),
2321 ),
2322 };
2323 let vector_count = payload.runs().iter().try_fold(0usize, |total, run| {
2324 total
2325 .checked_add(run.row_count as usize)
2326 .ok_or_else(|| invalid_data("ScaNN vector count overflows usize"))
2327 })?;
2328 let header = AnnDiskHeader {
2329 kind,
2330 routing: IvfRoutingMode::Flat,
2333 dim: usize::try_from(payload.dimension)
2334 .map_err(|_| invalid_data("ScaNN dimension exceeds usize"))?,
2335 code_size,
2336 num_clusters: payload.num_leaves,
2337 quantizer_version: payload.generation,
2338 codebook_version: payload.artifact_id,
2339 vector_count,
2340 };
2341 validate_header(&header)?;
2342 write_header(writer, &header)?;
2343
2344 let mut offset = ANN_HEADER_SIZE as u64;
2345 let mut records = Vec::with_capacity(payload.runs().len());
2346 let mut previous_leaf = None;
2347 for run in payload.runs() {
2348 let count = run.row_count as usize;
2349 let expected_docs = count
2350 .checked_mul(4)
2351 .ok_or_else(|| invalid_data("ScaNN doc-ID column size overflows usize"))?;
2352 let expected_ordinals = count
2353 .checked_mul(2)
2354 .ok_or_else(|| invalid_data("ScaNN ordinal column size overflows usize"))?;
2355 if count == 0
2356 || run.leaf_id >= header.num_clusters
2357 || previous_leaf.is_some_and(|leaf| leaf > run.leaf_id)
2358 || run.doc_ids_le.len() != expected_docs
2359 || run.ordinals_le.len() != expected_ordinals
2360 || run.codes.len()
2361 != expected_codes_column_len(header.kind, count, header.dim, header.code_size)?
2362 {
2363 return Err(invalid_data("ScaNN ANN leaf run columns are inconsistent"));
2364 }
2365 previous_leaf = Some(run.leaf_id);
2366 let max_doc_id = run
2367 .doc_ids_le
2368 .chunks_exact(4)
2369 .map(|bytes| u32::from_le_bytes(bytes.try_into().unwrap()))
2370 .max()
2371 .unwrap_or(0);
2372 if run
2373 .doc_base
2374 .checked_add(max_doc_id)
2375 .is_none_or(|doc_id| doc_id >= payload.doc_count)
2376 {
2377 return Err(invalid_data(
2378 "ScaNN ANN leaf run document ID is out of range",
2379 ));
2380 }
2381
2382 let doc_ids_offset = offset;
2383 writer.write_all(&run.doc_ids_le)?;
2384 offset = checked_advance(offset, run.doc_ids_le.len())?;
2385 let ordinals_offset = offset;
2386 writer.write_all(&run.ordinals_le)?;
2387 offset = checked_advance(offset, run.ordinals_le.len())?;
2388 let codes_offset = offset;
2389 writer.write_all(&run.codes)?;
2390 offset = checked_advance(offset, run.codes.len())?;
2391 records.push(RunRecord {
2392 cluster_id: run.leaf_id,
2393 doc_base: run.doc_base,
2394 count: run.row_count,
2395 max_doc_id,
2396 doc_ids_offset,
2397 ordinals_offset,
2398 codes_offset,
2399 codes_len: u64::try_from(run.codes.len())
2400 .map_err(|_| invalid_data("ScaNN code output size exceeds u64"))?,
2401 });
2402 }
2403 if records.is_empty() {
2404 return Err(invalid_data("cannot write an empty ScaNN ANN payload"));
2405 }
2406 finish_layout(writer, offset, &records)
2407}
2408
2409#[cfg(feature = "native")]
2412pub(crate) fn write_built_ivf_tq(
2413 index: &crate::structures::IvfTqIndex,
2414 num_clusters: u32,
2415 writer: &mut (impl Write + ?Sized),
2416) -> io::Result<u64> {
2417 use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_pack_ivf_block};
2418
2419 if !crate::structures::is_ivf_tq_cosine_generation(index.centroids_version) {
2420 return Err(invalid_data(
2421 "legacy raw IVF-TQ generations cannot be serialized; rebuild the index",
2422 ));
2423 }
2424 let codec = index.codec();
2425 let padded_dim = codec.padded_dim();
2426 let mut clusters: Vec<_> = index.clusters.iter().collect();
2427 clusters.sort_unstable_by_key(|(cluster_id, _)| **cluster_id);
2428 struct PackedCluster {
2432 cluster_id: u32,
2433 doc_ids: Vec<u32>,
2434 ordinals: Vec<u16>,
2435 codes: Vec<u8>,
2436 }
2437 let packed: Vec<PackedCluster> = clusters
2438 .iter()
2439 .map(|&(&cluster_id, cluster)| {
2440 let count = cluster.doc_ids.len();
2441 let mut order: Vec<usize> = (0..count).collect();
2442 order.sort_by(|&a, &b| {
2443 cluster.scales[b]
2444 .total_cmp(&cluster.scales[a])
2445 .then_with(|| a.cmp(&b))
2446 });
2447 let doc_ids: Vec<u32> = order.iter().map(|&i| cluster.doc_ids[i]).collect();
2448 let ordinals: Vec<u16> = order.iter().map(|&i| cluster.ordinals[i]).collect();
2449 let scales: Vec<f32> = order.iter().map(|&i| cluster.scales[i]).collect();
2450 let gammas: Vec<f32> = order.iter().map(|&i| cluster.gammas[i]).collect();
2451 let mut codes = Vec::with_capacity(
2452 crate::structures::vector::quantization::tq_ivf_codes_column_len_checked(
2453 count,
2454 codec.code_size(),
2455 )
2456 .unwrap_or_default(),
2457 );
2458 for block_start in (0..count).step_by(TQ_BLOCK_LANES) {
2459 let lanes = TQ_BLOCK_LANES.min(count - block_start);
2460 let rows: Vec<&[u8]> = order[block_start..block_start + lanes]
2461 .iter()
2462 .map(|&row| &cluster.rows[row * padded_dim..(row + 1) * padded_dim])
2463 .collect();
2464 tq_pack_ivf_block(
2465 &rows,
2466 &scales[block_start..block_start + lanes],
2467 &gammas[block_start..block_start + lanes],
2468 padded_dim,
2469 &mut codes,
2470 );
2471 }
2472 PackedCluster {
2473 cluster_id,
2474 doc_ids,
2475 ordinals,
2476 codes,
2477 }
2478 })
2479 .collect();
2480 let runs: Vec<_> = packed
2481 .iter()
2482 .map(|cluster| BuildRun {
2483 cluster_id: cluster.cluster_id,
2484 doc_ids: &cluster.doc_ids,
2485 ordinals: &cluster.ordinals,
2486 codes: &cluster.codes,
2487 })
2488 .collect();
2489 write_built_runs(
2490 AnnDiskHeader {
2491 kind: AnnKind::IvfTq,
2492 routing: index.routing,
2493 dim: index.dim,
2494 code_size: codec.code_size(),
2495 num_clusters,
2496 quantizer_version: index.centroids_version,
2497 codebook_version: codec.fingerprint(),
2498 vector_count: index.len(),
2499 },
2500 &runs,
2501 writer,
2502 )
2503}
2504
2505#[cfg(feature = "native")]
2507pub(crate) fn tq_builder_extra_run(builder: &crate::structures::TqFlatBuilder) -> BuildRun<'_> {
2508 BuildRun {
2509 cluster_id: 0,
2510 doc_ids: &builder.doc_ids,
2511 ordinals: &builder.ordinals,
2512 codes: &builder.codes,
2513 }
2514}
2515
2516#[cfg(feature = "native")]
2518pub(crate) fn write_built_tq_flat(
2519 builder: &crate::structures::TqFlatBuilder,
2520 writer: &mut (impl Write + ?Sized),
2521) -> io::Result<u64> {
2522 let codec = builder.codec();
2523 let runs = [BuildRun {
2524 cluster_id: 0,
2525 doc_ids: &builder.doc_ids,
2526 ordinals: &builder.ordinals,
2527 codes: &builder.codes,
2528 }];
2529 write_built_runs(
2530 AnnDiskHeader {
2531 kind: AnnKind::TqFlat,
2532 routing: IvfRoutingMode::Flat,
2533 dim: codec.dim(),
2534 code_size: codec.code_size(),
2535 num_clusters: 1,
2536 quantizer_version: codec.fingerprint(),
2537 codebook_version: 0,
2538 vector_count: builder.len(),
2539 },
2540 &runs,
2541 writer,
2542 )
2543}
2544
2545#[cfg(feature = "native")]
2546pub(crate) struct BuildRun<'a> {
2547 cluster_id: u32,
2548 doc_ids: &'a [u32],
2549 ordinals: &'a [u16],
2550 codes: &'a [u8],
2551}
2552
2553#[cfg(feature = "native")]
2554struct RunRecord {
2555 cluster_id: u32,
2556 doc_base: u32,
2557 count: u32,
2558 max_doc_id: u32,
2559 doc_ids_offset: u64,
2560 ordinals_offset: u64,
2561 codes_offset: u64,
2562 codes_len: u64,
2563}
2564
2565#[cfg(feature = "native")]
2566fn write_built_runs(
2567 header: AnnDiskHeader,
2568 runs: &[BuildRun<'_>],
2569 writer: &mut (impl Write + ?Sized),
2570) -> io::Result<u64> {
2571 if runs.is_empty() || header.vector_count == 0 {
2572 return Err(invalid_data("cannot write an empty ANN payload"));
2573 }
2574 validate_header(&header)?;
2575 write_header(writer, &header)?;
2576 let mut offset = ANN_HEADER_SIZE as u64;
2577 let mut records = Vec::with_capacity(runs.len());
2578 let mut counted = 0usize;
2579 let mut scratch = Vec::new();
2580 let mut previous_cluster = None;
2581 for run in runs {
2582 let count = run.doc_ids.len();
2583 if count == 0
2584 || run.cluster_id >= header.num_clusters
2585 || previous_cluster.is_some_and(|cluster| cluster >= run.cluster_id)
2586 || run.ordinals.len() != count
2587 || run.codes.len()
2588 != expected_codes_column_len(header.kind, count, header.dim, header.code_size)?
2589 {
2590 return Err(invalid_data("ANN build run columns are inconsistent"));
2591 }
2592 previous_cluster = Some(run.cluster_id);
2593 let count_u32 = u32::try_from(count)
2594 .map_err(|_| invalid_data("ANN cluster run exceeds u32 vectors"))?;
2595 let max_doc_id = run.doc_ids.iter().copied().max().unwrap_or(0);
2596 let doc_ids_offset = offset;
2597 write_u32_column(writer, run.doc_ids, &mut scratch)?;
2598 offset = offset
2599 .checked_add(
2600 u64::try_from(count)
2601 .ok()
2602 .and_then(|count| count.checked_mul(4))
2603 .ok_or_else(|| invalid_data("ANN doc-ID output size overflows u64"))?,
2604 )
2605 .ok_or_else(|| invalid_data("ANN output offset overflow"))?;
2606 let ordinals_offset = offset;
2607 write_u16_column(writer, run.ordinals, &mut scratch)?;
2608 offset = offset
2609 .checked_add(
2610 u64::try_from(count)
2611 .ok()
2612 .and_then(|count| count.checked_mul(2))
2613 .ok_or_else(|| invalid_data("ANN ordinal output size overflows u64"))?,
2614 )
2615 .ok_or_else(|| invalid_data("ANN output offset overflow"))?;
2616 let codes_offset = offset;
2617 writer.write_all(run.codes)?;
2618 offset = offset
2619 .checked_add(
2620 u64::try_from(run.codes.len())
2621 .map_err(|_| invalid_data("ANN code output size exceeds u64"))?,
2622 )
2623 .ok_or_else(|| invalid_data("ANN output offset overflow"))?;
2624 records.push(RunRecord {
2625 cluster_id: run.cluster_id,
2626 doc_base: 0,
2627 count: count_u32,
2628 max_doc_id,
2629 doc_ids_offset,
2630 ordinals_offset,
2631 codes_offset,
2632 codes_len: u64::try_from(run.codes.len())
2633 .map_err(|_| invalid_data("ANN code output size exceeds u64"))?,
2634 });
2635 counted = counted
2636 .checked_add(count)
2637 .ok_or_else(|| invalid_data("ANN vector count overflow"))?;
2638 }
2639 if counted != header.vector_count {
2640 return Err(invalid_data("ANN header/build vector counts disagree"));
2641 }
2642 finish_layout(writer, offset, &records)
2643}
2644
2645#[cfg(all(feature = "native", test))]
2648pub(crate) fn write_merged_ann(
2649 sources: &[(&AnnDiskIndex, u32)],
2650 writer: &mut (impl Write + ?Sized),
2651) -> io::Result<u64> {
2652 write_merged_ann_impl(sources, &[], writer, None)
2653}
2654
2655#[cfg(feature = "native")]
2656pub(crate) fn write_merged_ann_cancellable(
2657 sources: &[(&AnnDiskIndex, u32)],
2658 writer: &mut (impl Write + ?Sized),
2659 cancellation: Option<&std::sync::atomic::AtomicBool>,
2660) -> io::Result<u64> {
2661 write_merged_ann_impl(sources, &[], writer, cancellation)
2662}
2663
2664#[cfg(feature = "native")]
2673pub(crate) fn predicted_merge_fragmentation(sources: &[(&AnnDiskIndex, u32)]) -> f64 {
2674 let mut total_runs = 0usize;
2675 let mut cursors: Vec<std::iter::Peekable<std::slice::Iter<'_, AnnRun>>> = sources
2678 .iter()
2679 .map(|(source, _)| {
2680 total_runs += source.runs.len();
2681 source.runs.iter().peekable()
2682 })
2683 .collect();
2684 let mut distinct = 0usize;
2685 while let Some(cluster) = cursors
2686 .iter_mut()
2687 .filter_map(|cursor| cursor.peek().map(|run| run.cluster_id))
2688 .min()
2689 {
2690 distinct += 1;
2691 for cursor in &mut cursors {
2692 while cursor.peek().is_some_and(|run| run.cluster_id == cluster) {
2693 cursor.next();
2694 }
2695 }
2696 }
2697 if distinct == 0 {
2698 0.0
2699 } else {
2700 total_runs as f64 / distinct as f64
2701 }
2702}
2703
2704#[cfg(feature = "native")]
2706const DOC_ID_REWRITE_CHUNK: usize = 64 * 1024;
2707
2708#[cfg(feature = "native")]
2736pub(crate) fn write_compacted_ann_cancellable(
2737 sources: &[(&AnnDiskIndex, u32)],
2738 writer: &mut (impl Write + ?Sized),
2739 cancellation: Option<&std::sync::atomic::AtomicBool>,
2740) -> io::Result<u64> {
2741 let Some((first, _)) = sources.first() else {
2742 return Err(invalid_data("cannot compact an empty ANN source list"));
2743 };
2744 if !matches!(
2745 first.header.kind,
2746 AnnKind::BinaryIvf | AnnKind::ScannBinary | AnnKind::ScannAh
2747 ) {
2748 return Err(invalid_data(
2749 "ANN run compaction is only defined for exact binary and ScaNN AH payloads",
2750 ));
2751 }
2752 let mut header = first.header.clone();
2753 header.vector_count = 0;
2754 for &(source, _) in sources {
2755 if !headers_compatible(&first.header, &source.header) {
2756 return Err(invalid_data(
2757 "ANN compaction sources use incompatible generations",
2758 ));
2759 }
2760 header.vector_count = header
2761 .vector_count
2762 .checked_add(source.header.vector_count)
2763 .ok_or_else(|| invalid_data("compacted ANN vector count overflows usize"))?;
2764 }
2765 validate_header(&header)?;
2766 write_header(writer, &header)?;
2767
2768 let code_size = header.code_size;
2769 let mut offset = ANN_HEADER_SIZE as u64;
2770 let mut records: Vec<RunRecord> = Vec::new();
2771 let mut scratch = Vec::new();
2772 let mut cursors: Vec<usize> = vec![0; sources.len()];
2773
2774 loop {
2775 if cancellation.is_some_and(|flag| flag.load(std::sync::atomic::Ordering::Relaxed)) {
2776 return Err(io::Error::new(
2777 io::ErrorKind::Interrupted,
2778 "ANN compaction cancelled",
2779 ));
2780 }
2781 let Some(cluster_id) = sources
2783 .iter()
2784 .zip(&cursors)
2785 .filter_map(|((source, _), &cursor)| source.runs.get(cursor).map(|run| run.cluster_id))
2786 .min()
2787 else {
2788 break;
2789 };
2790
2791 let mut count = 0usize;
2796 let mut max_doc_id = 0u32;
2797 let doc_ids_offset = offset;
2798 for (source_index, &(source, segment_base)) in sources.iter().enumerate() {
2799 let mut cursor = cursors[source_index];
2800 while let Some(run) = source
2801 .runs
2802 .get(cursor)
2803 .filter(|run| run.cluster_id == cluster_id)
2804 {
2805 let base = run
2806 .doc_base
2807 .checked_add(segment_base)
2808 .ok_or_else(|| invalid_data("compacted ANN document base overflows u32"))?;
2809 let bytes = source.raw.as_slice();
2810 for chunk_start in (0..run.count).step_by(DOC_ID_REWRITE_CHUNK) {
2815 let chunk_end = (chunk_start + DOC_ID_REWRITE_CHUNK).min(run.count);
2816 scratch.clear();
2817 scratch.reserve((chunk_end - chunk_start) * 4);
2818 for index in chunk_start..chunk_end {
2819 let doc_id = run_doc_id_with_base(bytes, run, index, base)?;
2820 max_doc_id = max_doc_id.max(doc_id);
2821 scratch.extend_from_slice(&doc_id.to_le_bytes());
2822 }
2823 writer.write_all(&scratch)?;
2824 if cancellation
2825 .is_some_and(|flag| flag.load(std::sync::atomic::Ordering::Relaxed))
2826 {
2827 return Err(io::Error::new(
2828 io::ErrorKind::Interrupted,
2829 "ANN compaction cancelled",
2830 ));
2831 }
2832 }
2833 offset = checked_advance(offset, run.count * 4)?;
2834 count = count
2835 .checked_add(run.count)
2836 .ok_or_else(|| invalid_data("compacted ANN run count overflows usize"))?;
2837 cursor += 1;
2838 }
2839 }
2840
2841 let ordinals_offset = offset;
2843 for (source_index, &(source, _)) in sources.iter().enumerate() {
2844 let mut cursor = cursors[source_index];
2845 while let Some(run) = source
2846 .runs
2847 .get(cursor)
2848 .filter(|run| run.cluster_id == cluster_id)
2849 {
2850 copy_range(writer, &source.raw, run.ordinals.clone(), cancellation)?;
2851 offset = checked_advance(offset, run.ordinals.len())?;
2852 cursor += 1;
2853 }
2854 }
2855
2856 let codes_offset = offset;
2859 if header.kind == AnnKind::ScannAh {
2860 let blocks = header.dim.div_ceil(code_size);
2861 let lanes = crate::structures::vector::scann::FAST_SCAN_LANES;
2862 let mut unpacked = Vec::with_capacity(lanes * blocks);
2863 let mut packed = Vec::with_capacity(blocks * lanes / 2);
2864 for (source_index, &(source, _)) in sources.iter().enumerate() {
2865 let mut cursor = cursors[source_index];
2866 while let Some(run) = source
2867 .runs
2868 .get(cursor)
2869 .filter(|run| run.cluster_id == cluster_id)
2870 {
2871 let bytes = &source.raw.as_slice()[run.codes.clone()];
2872 for row in 0..run.count {
2873 unpack_scann_ah_row(bytes, run.count, blocks, row, &mut unpacked)?;
2874 if unpacked.len() == lanes * blocks {
2875 packed.clear();
2876 crate::structures::vector::scann::pack_fast_scan_block(
2877 &unpacked,
2878 blocks,
2879 &mut packed,
2880 )
2881 .map_err(|error| invalid_data(error.to_string()))?;
2882 writer.write_all(&packed)?;
2883 offset = checked_advance(offset, packed.len())?;
2884 unpacked.clear();
2885 if cancellation
2886 .is_some_and(|flag| flag.load(std::sync::atomic::Ordering::Relaxed))
2887 {
2888 return Err(io::Error::new(
2889 io::ErrorKind::Interrupted,
2890 "ANN compaction cancelled",
2891 ));
2892 }
2893 }
2894 }
2895 cursor += 1;
2896 }
2897 }
2898 if !unpacked.is_empty() {
2901 packed.clear();
2902 for row in unpacked.chunks_exact(blocks) {
2903 for pair in row.chunks(2) {
2904 packed.push(pair[0] | (pair.get(1).copied().unwrap_or(0) << 4));
2905 }
2906 }
2907 writer.write_all(&packed)?;
2908 offset = checked_advance(offset, packed.len())?;
2909 }
2910 } else {
2911 for (source_index, &(source, _)) in sources.iter().enumerate() {
2912 let mut cursor = cursors[source_index];
2913 while let Some(run) = source
2914 .runs
2915 .get(cursor)
2916 .filter(|run| run.cluster_id == cluster_id)
2917 {
2918 copy_range(writer, &source.raw, run.codes.clone(), cancellation)?;
2919 offset = checked_advance(offset, run.codes.len())?;
2920 cursor += 1;
2921 }
2922 }
2923 }
2924 let expected_codes_len =
2925 expected_codes_column_len(header.kind, count, header.dim, header.code_size)?;
2926 if offset - codes_offset != expected_codes_len as u64 {
2927 return Err(invalid_data("compacted ANN code column length mismatch"));
2928 }
2929
2930 for (source_index, &(source, _)) in sources.iter().enumerate() {
2932 while source
2933 .runs
2934 .get(cursors[source_index])
2935 .is_some_and(|run| run.cluster_id == cluster_id)
2936 {
2937 cursors[source_index] += 1;
2938 }
2939 }
2940
2941 records.push(RunRecord {
2942 cluster_id,
2943 doc_base: 0,
2944 count: u32::try_from(count)
2945 .map_err(|_| invalid_data("compacted ANN run exceeds u32 vectors"))?,
2946 max_doc_id,
2947 doc_ids_offset,
2948 ordinals_offset,
2949 codes_offset,
2950 codes_len: u64::try_from(expected_codes_column_len(
2951 header.kind,
2952 count,
2953 header.dim,
2954 code_size,
2955 )?)
2956 .map_err(|_| invalid_data("compacted ANN code length exceeds u64"))?,
2957 });
2958 }
2959
2960 if records.is_empty() {
2961 return Err(invalid_data("cannot compact an ANN payload with no runs"));
2962 }
2963 finish_layout(writer, offset, &records)
2964}
2965
2966#[cfg(feature = "native")]
2970fn unpack_scann_ah_row(
2971 codes: &[u8],
2972 count: usize,
2973 blocks: usize,
2974 row: usize,
2975 output: &mut Vec<u8>,
2976) -> io::Result<()> {
2977 use crate::structures::vector::scann::{
2978 FAST_SCAN_LANES, packed_block_bytes, packed_code_position,
2979 };
2980
2981 let lanes = FAST_SCAN_LANES;
2982 let full_rows = count / lanes * lanes;
2983 let full_block_bytes = packed_block_bytes(blocks)
2984 .ok_or_else(|| invalid_data("ScaNN AH block size overflows usize"))?;
2985 let tail_row_bytes = blocks.div_ceil(2);
2986 for block in 0..blocks {
2987 let (byte_offset, high) = if row < full_rows {
2988 let (offset, high) = packed_code_position(row % lanes, block);
2989 ((row / lanes) * full_block_bytes + offset, high)
2990 } else {
2991 (
2992 (full_rows / lanes) * full_block_bytes
2993 + (row - full_rows) * tail_row_bytes
2994 + block / 2,
2995 !block.is_multiple_of(2),
2996 )
2997 };
2998 let byte = *codes
2999 .get(byte_offset)
3000 .ok_or_else(|| invalid_data("ScaNN AH row exceeds its code column"))?;
3001 output.push(if high { byte >> 4 } else { byte & 0x0f });
3002 }
3003 Ok(())
3004}
3005
3006#[cfg(feature = "native")]
3008fn run_doc_id_with_base(bytes: &[u8], run: &AnnRun, index: usize, base: u32) -> io::Result<u32> {
3009 let local_doc_id = read_u32(bytes, run.doc_ids.start + index * 4);
3010 if local_doc_id > run.max_doc_id {
3011 return Err(invalid_data(
3012 "ANN run contains a document above its declared maximum",
3013 ));
3014 }
3015 base.checked_add(local_doc_id)
3016 .ok_or_else(|| invalid_data("compacted ANN document ID overflows u32"))
3017}
3018
3019#[cfg(feature = "native")]
3023pub(crate) fn write_merged_ann_with_extra(
3024 sources: &[(&AnnDiskIndex, u32)],
3025 extra: &[BuildRun<'_>],
3026 writer: &mut (impl Write + ?Sized),
3027 cancellation: Option<&std::sync::atomic::AtomicBool>,
3028) -> io::Result<u64> {
3029 write_merged_ann_impl(sources, extra, writer, cancellation)
3030}
3031
3032#[cfg(feature = "native")]
3033fn write_merged_ann_impl(
3034 sources: &[(&AnnDiskIndex, u32)],
3035 extra: &[BuildRun<'_>],
3036 writer: &mut (impl Write + ?Sized),
3037 cancellation: Option<&std::sync::atomic::AtomicBool>,
3038) -> io::Result<u64> {
3039 let Some((first, _)) = sources.first() else {
3040 return Err(invalid_data("cannot merge an empty ANN source list"));
3041 };
3042 if first.header.kind == AnnKind::IvfTq
3043 && !crate::structures::is_ivf_tq_cosine_generation(first.header.quantizer_version)
3044 {
3045 return Err(invalid_data(
3046 "legacy raw IVF-TQ generations cannot be merged; rebuild the index",
3047 ));
3048 }
3049 let mut header = first.header.clone();
3050 header.vector_count = 0;
3051 for &(source, _) in sources {
3052 if !headers_compatible(&first.header, &source.header) {
3053 return Err(invalid_data(
3054 "ANN merge sources use incompatible generations",
3055 ));
3056 }
3057 header.vector_count = header
3058 .vector_count
3059 .checked_add(source.header.vector_count)
3060 .ok_or_else(|| invalid_data("merged ANN vector count overflows usize"))?;
3061 }
3062 for run in extra {
3063 if run.doc_ids.is_empty()
3064 || run.cluster_id >= header.num_clusters
3065 || run.ordinals.len() != run.doc_ids.len()
3066 || run.codes.len()
3067 != expected_codes_column_len(
3068 header.kind,
3069 run.doc_ids.len(),
3070 header.dim,
3071 header.code_size,
3072 )?
3073 {
3074 return Err(invalid_data("extra ANN merge run columns are inconsistent"));
3075 }
3076 header.vector_count = header
3077 .vector_count
3078 .checked_add(run.doc_ids.len())
3079 .ok_or_else(|| invalid_data("merged ANN vector count overflows usize"))?;
3080 }
3081 validate_header(&header)?;
3082 write_header(writer, &header)?;
3083 let mut offset = ANN_HEADER_SIZE as u64;
3084 let run_capacity = sources.iter().try_fold(extra.len(), |count, (source, _)| {
3085 count
3086 .checked_add(source.runs.len())
3087 .ok_or_else(|| invalid_data("merged ANN run count overflows usize"))
3088 })?;
3089 let mut output_payload_starts = Vec::with_capacity(sources.len());
3090 for &(source, _) in sources {
3091 let payload_end = source
3092 .runs
3093 .iter()
3094 .map(|run| run.codes.end)
3095 .max()
3096 .ok_or_else(|| invalid_data("ANN source has no payload runs"))?;
3097 let output_payload_start = offset;
3098 output_payload_starts.push(output_payload_start);
3099 copy_range(
3100 writer,
3101 &source.raw,
3102 ANN_HEADER_SIZE..payload_end,
3103 cancellation,
3104 )?;
3105 offset = checked_advance(offset, payload_end - ANN_HEADER_SIZE)?;
3106 }
3107
3108 let mut extra_records = Vec::with_capacity(extra.len());
3111 let mut scratch = Vec::new();
3112 for run in extra {
3113 let count = run.doc_ids.len();
3114 let doc_ids_offset = offset;
3115 write_u32_column(writer, run.doc_ids, &mut scratch)?;
3116 offset = checked_advance(offset, count * 4)?;
3117 let ordinals_offset = offset;
3118 write_u16_column(writer, run.ordinals, &mut scratch)?;
3119 offset = checked_advance(offset, count * 2)?;
3120 let codes_offset = offset;
3121 writer.write_all(run.codes)?;
3122 offset = checked_advance(offset, run.codes.len())?;
3123 extra_records.push(RunRecord {
3124 cluster_id: run.cluster_id,
3125 doc_base: 0,
3126 count: u32::try_from(count)
3127 .map_err(|_| invalid_data("extra ANN run exceeds u32 vectors"))?,
3128 max_doc_id: run.doc_ids.iter().copied().max().unwrap_or(0),
3129 doc_ids_offset,
3130 ordinals_offset,
3131 codes_offset,
3132 codes_len: u64::try_from(run.codes.len())
3133 .map_err(|_| invalid_data("extra ANN code length exceeds u64"))?,
3134 });
3135 }
3136
3137 let directory_offset = offset;
3143 let mut pending = BinaryHeap::with_capacity(sources.len() + 1);
3144 for (source_index, (source, _)) in sources.iter().enumerate() {
3145 pending.push(Reverse((source.runs[0].cluster_id, source_index, 0usize)));
3146 }
3147 if let Some(first_extra) = extra_records.first() {
3148 pending.push(Reverse((first_extra.cluster_id, sources.len(), 0usize)));
3149 }
3150 let mut written_runs = 0usize;
3151 while let Some(Reverse((_, source_index, run_index))) = pending.pop() {
3152 if source_index == sources.len() {
3153 write_run_record(writer, &extra_records[run_index])?;
3154 written_runs += 1;
3155 if let Some(next) = extra_records.get(run_index + 1) {
3156 pending.push(Reverse((next.cluster_id, source_index, run_index + 1)));
3157 }
3158 continue;
3159 }
3160 let (source, segment_base) = sources[source_index];
3161 let run = &source.runs[run_index];
3162 write_run_record(
3163 writer,
3164 &RunRecord {
3165 cluster_id: run.cluster_id,
3166 doc_base: run
3167 .doc_base
3168 .checked_add(segment_base)
3169 .ok_or_else(|| invalid_data("merged ANN document base overflows u32"))?,
3170 count: u32::try_from(run.count)
3171 .map_err(|_| invalid_data("ANN source run exceeds u32 vectors"))?,
3172 max_doc_id: run.max_doc_id,
3173 doc_ids_offset: relocate_payload_offset(
3174 output_payload_starts[source_index],
3175 run.doc_ids.start,
3176 )?,
3177 ordinals_offset: relocate_payload_offset(
3178 output_payload_starts[source_index],
3179 run.ordinals.start,
3180 )?,
3181 codes_offset: relocate_payload_offset(
3182 output_payload_starts[source_index],
3183 run.codes.start,
3184 )?,
3185 codes_len: u64::try_from(run.codes.len())
3186 .map_err(|_| invalid_data("ANN source code length exceeds u64"))?,
3187 },
3188 )?;
3189 written_runs = written_runs
3190 .checked_add(1)
3191 .ok_or_else(|| invalid_data("merged ANN run count overflows usize"))?;
3192 let next_run_index = run_index + 1;
3193 if let Some(next_run) = source.runs.get(next_run_index) {
3194 pending.push(Reverse((next_run.cluster_id, source_index, next_run_index)));
3195 }
3196 }
3197 if written_runs != run_capacity {
3198 return Err(invalid_data("merged ANN directory lost source runs"));
3199 }
3200 finish_footer(writer, directory_offset, written_runs)
3201}
3202
3203#[cfg(feature = "native")]
3204fn relocate_payload_offset(output_payload_start: u64, source_offset: usize) -> io::Result<u64> {
3205 let relative = source_offset
3206 .checked_sub(ANN_HEADER_SIZE)
3207 .ok_or_else(|| invalid_data("ANN source offset precedes its payload"))?;
3208 output_payload_start
3209 .checked_add(
3210 u64::try_from(relative)
3211 .map_err(|_| invalid_data("ANN source relative offset exceeds u64"))?,
3212 )
3213 .ok_or_else(|| invalid_data("merged ANN payload offset overflows u64"))
3214}
3215
3216#[cfg(feature = "native")]
3217fn headers_compatible(left: &AnnDiskHeader, right: &AnnDiskHeader) -> bool {
3218 left.kind == right.kind
3219 && left.routing == right.routing
3220 && left.dim == right.dim
3221 && left.code_size == right.code_size
3222 && left.num_clusters == right.num_clusters
3223 && left.quantizer_version == right.quantizer_version
3224 && left.codebook_version == right.codebook_version
3225}
3226
3227#[cfg(feature = "native")]
3228fn finish_layout(
3229 writer: &mut (impl Write + ?Sized),
3230 directory_offset: u64,
3231 records: &[RunRecord],
3232) -> io::Result<u64> {
3233 for record in records {
3234 write_run_record(writer, record)?;
3235 }
3236 finish_footer(writer, directory_offset, records.len())
3237}
3238
3239#[cfg(feature = "native")]
3240fn write_run_record(writer: &mut (impl Write + ?Sized), record: &RunRecord) -> io::Result<()> {
3241 writer.write_u32::<LittleEndian>(record.cluster_id)?;
3242 writer.write_u32::<LittleEndian>(record.doc_base)?;
3243 writer.write_u32::<LittleEndian>(record.count)?;
3244 writer.write_u32::<LittleEndian>(record.max_doc_id)?;
3245 writer.write_u64::<LittleEndian>(record.doc_ids_offset)?;
3246 writer.write_u64::<LittleEndian>(record.ordinals_offset)?;
3247 writer.write_u64::<LittleEndian>(record.codes_offset)?;
3248 writer.write_u64::<LittleEndian>(record.codes_len)?;
3249 Ok(())
3250}
3251
3252#[cfg(feature = "native")]
3253fn finish_footer(
3254 writer: &mut (impl Write + ?Sized),
3255 directory_offset: u64,
3256 num_records: usize,
3257) -> io::Result<u64> {
3258 writer.write_u64::<LittleEndian>(directory_offset)?;
3259 writer.write_u64::<LittleEndian>(
3260 u64::try_from(num_records).map_err(|_| invalid_data("ANN run count exceeds u64"))?,
3261 )?;
3262 writer.write_u32::<LittleEndian>(ANN_FOOTER_MAGIC)?;
3263 writer.write_u32::<LittleEndian>(u32::from(ANN_DISK_VERSION))?;
3264 let tail_size = num_records
3265 .checked_mul(ANN_RUN_SIZE)
3266 .and_then(|size| size.checked_add(ANN_FOOTER_SIZE))
3267 .and_then(|size| u64::try_from(size).ok())
3268 .ok_or_else(|| invalid_data("ANN final tail size overflows u64"))?;
3269 directory_offset
3270 .checked_add(tail_size)
3271 .ok_or_else(|| invalid_data("ANN final size overflows u64"))
3272}
3273
3274fn expected_layout_revision(kind: AnnKind) -> u32 {
3282 match kind {
3283 AnnKind::ScannAh => u32::from(crate::structures::vector::scann::FAST_SCAN_LAYOUT_VERSION),
3284 AnnKind::BinaryIvf | AnnKind::TqFlat | AnnKind::IvfTq | AnnKind::ScannBinary => 0,
3285 }
3286}
3287
3288fn check_layout_revision(kind: AnnKind, found: u32) -> io::Result<()> {
3289 let expected = expected_layout_revision(kind);
3290 if found == expected {
3291 return Ok(());
3292 }
3293 Err(invalid_data(match kind {
3294 AnnKind::ScannAh => format!(
3295 "ScaNN AH payload uses FastScan layout revision {found}; this build reads revision \
3296 {expected} (docs/fast-scan-layout-v2.md). Rebuild the float ScaNN generation for \
3297 this field before searching it"
3298 ),
3299 _ => format!("ANN header reserved field is non-zero ({found}) for kind {kind:?}"),
3300 }))
3301}
3302
3303#[cfg(feature = "native")]
3304fn write_header(writer: &mut (impl Write + ?Sized), header: &AnnDiskHeader) -> io::Result<()> {
3305 writer.write_u32::<LittleEndian>(ANN_HEADER_MAGIC)?;
3306 writer.write_u8(header.kind as u8)?;
3307 writer.write_u8(routing_to_u8(header.routing))?;
3308 writer.write_u16::<LittleEndian>(ANN_DISK_VERSION)?;
3309 writer.write_u32::<LittleEndian>(
3310 u32::try_from(header.dim).map_err(|_| invalid_data("ANN dimension exceeds u32"))?,
3311 )?;
3312 writer.write_u32::<LittleEndian>(
3313 u32::try_from(header.code_size).map_err(|_| invalid_data("ANN code size exceeds u32"))?,
3314 )?;
3315 writer.write_u32::<LittleEndian>(header.num_clusters)?;
3316 writer.write_u32::<LittleEndian>(expected_layout_revision(header.kind))?;
3317 writer.write_u64::<LittleEndian>(header.quantizer_version)?;
3318 writer.write_u64::<LittleEndian>(header.codebook_version)?;
3319 writer.write_u64::<LittleEndian>(
3320 u64::try_from(header.vector_count)
3321 .map_err(|_| invalid_data("ANN vector count exceeds u64"))?,
3322 )?;
3323 writer.write_u64::<LittleEndian>(0)?;
3324 Ok(())
3325}
3326
3327#[cfg(feature = "native")]
3328fn write_u32_column(
3329 writer: &mut (impl Write + ?Sized),
3330 values: &[u32],
3331 scratch: &mut Vec<u8>,
3332) -> io::Result<()> {
3333 for chunk in values.chunks(64 * 1024) {
3334 scratch.clear();
3335 scratch.reserve(chunk.len() * 4);
3336 for value in chunk {
3337 scratch.extend_from_slice(&value.to_le_bytes());
3338 }
3339 writer.write_all(scratch)?;
3340 }
3341 Ok(())
3342}
3343
3344#[cfg(feature = "native")]
3345fn write_u16_column(
3346 writer: &mut (impl Write + ?Sized),
3347 values: &[u16],
3348 scratch: &mut Vec<u8>,
3349) -> io::Result<()> {
3350 for chunk in values.chunks(64 * 1024) {
3351 scratch.clear();
3352 scratch.reserve(chunk.len() * 2);
3353 for value in chunk {
3354 scratch.extend_from_slice(&value.to_le_bytes());
3355 }
3356 writer.write_all(scratch)?;
3357 }
3358 Ok(())
3359}
3360
3361#[cfg(feature = "native")]
3362fn copy_range(
3363 writer: &mut (impl Write + ?Sized),
3364 bytes: &OwnedBytes,
3365 range: Range<usize>,
3366 cancellation: Option<&std::sync::atomic::AtomicBool>,
3367) -> io::Result<()> {
3368 if range.is_empty() {
3369 return Ok(());
3370 }
3371 let range_end = range.end;
3372 let mut chunk_start = range.start;
3373 let first_end = chunk_start.saturating_add(COPY_CHUNK).min(range_end);
3374 bytes.madvise_range(chunk_start..first_end, libc::MADV_WILLNEED);
3375 while chunk_start < range_end {
3376 if cancellation
3377 .is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Relaxed))
3378 {
3379 return Err(io::Error::new(
3380 io::ErrorKind::Interrupted,
3381 "ANN merge copy cancelled",
3382 ));
3383 }
3384 let chunk_end = chunk_start.saturating_add(COPY_CHUNK).min(range_end);
3385 let next_end = chunk_end.saturating_add(COPY_CHUNK).min(range_end);
3386 if chunk_end < next_end {
3387 bytes.madvise_range(chunk_end..next_end, libc::MADV_WILLNEED);
3390 }
3391 writer.write_all(&bytes.as_slice()[chunk_start..chunk_end])?;
3392 chunk_start = chunk_end;
3393 }
3394 Ok(())
3395}
3396
3397#[cfg(feature = "native")]
3398fn checked_advance(offset: u64, length: usize) -> io::Result<u64> {
3399 offset
3400 .checked_add(
3401 u64::try_from(length).map_err(|_| invalid_data("ANN copy length exceeds u64"))?,
3402 )
3403 .ok_or_else(|| invalid_data("ANN output offset overflows u64"))
3404}
3405
3406fn validate_header(header: &AnnDiskHeader) -> io::Result<()> {
3407 if header.kind == AnnKind::IvfTq
3408 && !crate::structures::is_ivf_tq_cosine_generation(header.quantizer_version)
3409 {
3410 return Err(invalid_data(
3411 "IVF-TQ payload uses an unsupported legacy generation; rebuild the index",
3412 ));
3413 }
3414 if header.dim == 0
3415 || header.code_size == 0
3416 || header.num_clusters == 0
3417 || header.quantizer_version == 0
3418 || header.vector_count == 0
3419 || (header.kind == AnnKind::BinaryIvf
3420 && (header.codebook_version != 0
3421 || !header.dim.is_multiple_of(8)
3422 || header.code_size != header.dim.div_ceil(8)))
3423 || (header.kind == AnnKind::ScannBinary
3424 && (header.codebook_version == 0
3425 || header.routing != IvfRoutingMode::Flat
3426 || !header.dim.is_multiple_of(8)
3427 || header.code_size != header.dim / 8))
3428 || (header.kind == AnnKind::TqFlat
3429 && (header.codebook_version != 0
3430 || header.num_clusters != 1
3431 || header.routing != IvfRoutingMode::Flat
3432 || header.code_size * 2
3433 != crate::structures::vector::quantization::tq_padded_dim(header.dim)))
3434 || (header.kind == AnnKind::IvfTq
3437 && (header.codebook_version == 0
3438 || header.code_size * 2
3439 != crate::structures::vector::quantization::tq_padded_dim(header.dim)))
3440 || (header.kind == AnnKind::ScannAh
3441 && (header.codebook_version == 0
3442 || header.routing != IvfRoutingMode::Flat
3443 || u16::try_from(header.code_size).is_err()
3444 || u32::try_from(header.dim).is_err()
3445 || crate::structures::vector::scann::ScannEncoding::AsymmetricHash {
3446 dimensions_per_block: header.code_size as u16,
3447 bits_per_code: 4,
3448 }
3449 .row_code_bytes(u32::try_from(header.dim).unwrap_or(0))
3450 .is_err()))
3451 {
3452 return Err(invalid_data("ANN header contains invalid metadata"));
3453 }
3454 Ok(())
3455}
3456
3457fn read_u32(bytes: &[u8], offset: usize) -> u32 {
3458 u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
3459}
3460
3461fn read_u16(bytes: &[u8], offset: usize) -> u16 {
3462 u16::from_le_bytes(bytes[offset..offset + 2].try_into().unwrap())
3463}
3464
3465#[inline]
3466fn tq_ivf_block_max_scale(block: &[u8]) -> f32 {
3467 f32::from_le_bytes(
3471 block[..size_of::<f32>()]
3472 .try_into()
3473 .expect("scale is one f32"),
3474 )
3475}
3476
3477#[inline]
3483fn run_doc_id_validated(bytes: &[u8], run: &AnnRun, index: usize) -> u32 {
3484 debug_assert!(index < run.count);
3485 run.doc_base + read_u32(bytes, run.doc_ids.start + index * 4)
3486}
3487
3488fn run_doc_id(bytes: &[u8], run: &AnnRun, index: usize) -> io::Result<u32> {
3489 let local_doc_id = read_u32(bytes, run.doc_ids.start + index * 4);
3490 if local_doc_id > run.max_doc_id {
3491 return Err(invalid_data(
3492 "ANN run contains a document above its declared maximum",
3493 ));
3494 }
3495 run.doc_base
3496 .checked_add(local_doc_id)
3497 .ok_or_else(|| invalid_data("ANN run document ID overflows u32"))
3498}
3499
3500#[cfg(feature = "native")]
3501fn routing_to_u8(routing: IvfRoutingMode) -> u8 {
3502 match routing {
3503 IvfRoutingMode::Auto => 0,
3504 IvfRoutingMode::Flat => 1,
3505 IvfRoutingMode::TwoLevel => 2,
3506 IvfRoutingMode::Hnsw => 3,
3507 }
3508}
3509
3510fn routing_from_u8(value: u8) -> io::Result<IvfRoutingMode> {
3511 match value {
3512 0 => Ok(IvfRoutingMode::Auto),
3513 1 => Ok(IvfRoutingMode::Flat),
3514 2 => Ok(IvfRoutingMode::TwoLevel),
3515 3 => Ok(IvfRoutingMode::Hnsw),
3516 _ => Err(invalid_data(format!("unknown ANN routing mode {value}"))),
3517 }
3518}
3519
3520fn invalid_data(message: impl Into<String>) -> io::Error {
3521 io::Error::new(io::ErrorKind::InvalidData, message.into())
3522}
3523
3524#[cfg(all(test, feature = "native"))]
3525mod tests {
3526 use super::*;
3527
3528 #[test]
3533 fn compacted_merge_matches_byte_copy_and_resets_fragmentation() {
3534 let a0_docs = [0u32, 1, 2];
3538 let a0_ords = [0u16, 0, 1];
3539 let a0_codes = [0x11u8, 0x22, 0x33];
3540 let a5_docs = [3u32, 4];
3541 let a5_ords = [0u16; 2];
3542 let a5_codes = [0x44u8, 0x55];
3543 let a_runs = [
3544 BuildRun {
3545 cluster_id: 0,
3546 doc_ids: &a0_docs,
3547 ordinals: &a0_ords,
3548 codes: &a0_codes,
3549 },
3550 BuildRun {
3551 cluster_id: 5,
3552 doc_ids: &a5_docs,
3553 ordinals: &a5_ords,
3554 codes: &a5_codes,
3555 },
3556 ];
3557 let mut header = binary_header(5);
3558 header.num_clusters = 8;
3559 let mut a_bytes = Vec::new();
3560 write_built_runs(header.clone(), &a_runs, &mut a_bytes).unwrap();
3561 let a = AnnDiskIndex::open(OwnedBytes::new(a_bytes), AnnKind::BinaryIvf, 5).unwrap();
3562
3563 let b0_docs = [0u32, 2];
3564 let b0_ords = [1u16, 0];
3565 let b0_codes = [0x66u8, 0x77];
3566 let b2_docs = [1u32];
3567 let b2_ords = [0u16];
3568 let b2_codes = [0x88u8];
3569 let b_runs = [
3570 BuildRun {
3571 cluster_id: 0,
3572 doc_ids: &b0_docs,
3573 ordinals: &b0_ords,
3574 codes: &b0_codes,
3575 },
3576 BuildRun {
3577 cluster_id: 2,
3578 doc_ids: &b2_docs,
3579 ordinals: &b2_ords,
3580 codes: &b2_codes,
3581 },
3582 ];
3583 let mut header_b = binary_header(3);
3584 header_b.num_clusters = 8;
3585 let mut b_bytes = Vec::new();
3586 write_built_runs(header_b, &b_runs, &mut b_bytes).unwrap();
3587 let b = AnnDiskIndex::open(OwnedBytes::new(b_bytes), AnnKind::BinaryIvf, 3).unwrap();
3588
3589 let mut gen1_bytes = Vec::new();
3591 write_merged_ann(&[(&a, 0), (&b, 5)], &mut gen1_bytes).unwrap();
3592 let gen1 = AnnDiskIndex::open(OwnedBytes::new(gen1_bytes), AnnKind::BinaryIvf, 8).unwrap();
3593
3594 let sources: [(&AnnDiskIndex, u32); 2] = [(&gen1, 0), (&a, 8)];
3596 let predicted = predicted_merge_fragmentation(&sources);
3597 assert!((predicted - 2.0).abs() < 1e-9, "{predicted}");
3599
3600 let mut copied_bytes = Vec::new();
3601 write_merged_ann(&sources, &mut copied_bytes).unwrap();
3602 let copied =
3603 AnnDiskIndex::open(OwnedBytes::new(copied_bytes), AnnKind::BinaryIvf, 13).unwrap();
3604 let mut compacted_bytes = Vec::new();
3605 write_compacted_ann_cancellable(&sources, &mut compacted_bytes, None).unwrap();
3606 let compacted =
3607 AnnDiskIndex::open(OwnedBytes::new(compacted_bytes), AnnKind::BinaryIvf, 13).unwrap();
3608
3609 let copied_health = copied.health();
3611 let compacted_health = compacted.health();
3612 assert!((copied_health.fragmentation() - 2.0).abs() < 1e-9);
3613 assert!((compacted_health.fragmentation() - 1.0).abs() < 1e-9);
3614 assert_eq!(compacted_health.runs, 3, "one run per non-empty cluster");
3615 assert_eq!(copied_health.vectors, compacted_health.vectors);
3616 assert_eq!(copied_health.payload_bytes, compacted_health.payload_bytes);
3617 assert_eq!(
3618 copied_health.largest_cluster_vectors,
3619 compacted_health.largest_cluster_vectors
3620 );
3621
3622 for cluster in 0..8u32 {
3624 let query = [0x5Au8];
3625 let from_copy = copied
3626 .search_binary_clusters::<false>(&query, 16, &[cluster])
3627 .unwrap();
3628 let from_compact = compacted
3629 .search_binary_clusters::<false>(&query, 16, &[cluster])
3630 .unwrap();
3631 assert_eq!(from_copy, from_compact, "cluster {cluster} diverged");
3632 }
3633
3634 assert!(compacted.runs.iter().all(|run| run.doc_base == 0));
3636
3637 let mut generation3 = Vec::new();
3640 write_merged_ann(&[(&compacted, 0), (&b, 13)], &mut generation3).unwrap();
3641 let generation3 =
3642 AnnDiskIndex::open(OwnedBytes::new(generation3), AnnKind::BinaryIvf, 16).unwrap();
3643 assert_eq!(generation3.health().vectors, 16);
3644 let all: Vec<(u32, u16, f32)> = compacted
3646 .search_binary_clusters::<false>(&[0x5A], 32, &[0, 2, 5])
3647 .unwrap();
3648 let mut docs: Vec<u32> = all.iter().map(|&(doc, _, _)| doc).collect();
3649 docs.sort_unstable();
3650 assert_eq!(docs, (0..=12).collect::<Vec<u32>>());
3651 }
3652
3653 #[test]
3654 fn compacted_scann_ah_repacks_blocks_across_run_boundaries() {
3655 use crate::structures::vector::scann::{FAST_SCAN_LANES, pack_fast_scan_block};
3656
3657 const DIM: usize = 8;
3658 const DIMS_PER_BLOCK: usize = 2;
3659 const BLOCKS: usize = DIM / DIMS_PER_BLOCK;
3660
3661 let make_rows = |start: usize, count: usize| -> Vec<u8> {
3662 (0..count)
3663 .flat_map(|row| {
3664 (0..BLOCKS).map(move |block| ((start + row * 3 + block * 5) & 0x0f) as u8)
3665 })
3666 .collect()
3667 };
3668 let encode = |rows: &[u8]| -> Vec<u8> {
3669 let count = rows.len() / BLOCKS;
3670 let full_rows = count / FAST_SCAN_LANES * FAST_SCAN_LANES;
3671 let mut codes = Vec::new();
3672 for group in rows[..full_rows * BLOCKS].chunks_exact(FAST_SCAN_LANES * BLOCKS) {
3673 pack_fast_scan_block(group, BLOCKS, &mut codes).unwrap();
3674 }
3675 for row in rows[full_rows * BLOCKS..].chunks_exact(BLOCKS) {
3676 for pair in row.chunks(2) {
3677 codes.push(pair[0] | (pair.get(1).copied().unwrap_or(0) << 4));
3678 }
3679 }
3680 codes
3681 };
3682 let make_source = |rows: &[u8]| {
3683 let count = rows.len() / BLOCKS;
3684 let docs: Vec<u32> = (0..count as u32).collect();
3685 let ordinals = vec![0u16; count];
3686 let codes = encode(rows);
3687 let run = BuildRun {
3688 cluster_id: 0,
3689 doc_ids: &docs,
3690 ordinals: &ordinals,
3691 codes: &codes,
3692 };
3693 let header = AnnDiskHeader {
3694 kind: AnnKind::ScannAh,
3695 routing: IvfRoutingMode::Flat,
3696 dim: DIM,
3697 code_size: DIMS_PER_BLOCK,
3698 num_clusters: 1,
3699 quantizer_version: 41,
3700 codebook_version: 73,
3701 vector_count: count,
3702 };
3703 let mut bytes = Vec::new();
3704 write_built_runs(header, &[run], &mut bytes).unwrap();
3705 AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::ScannAh, count as u32).unwrap()
3706 };
3707
3708 let left_rows = make_rows(1, 35);
3711 let right_rows = make_rows(9, 30);
3712 let left = make_source(&left_rows);
3713 let right = make_source(&right_rows);
3714 let mut bytes = Vec::new();
3715 write_compacted_ann_cancellable(&[(&left, 0), (&right, 35)], &mut bytes, None).unwrap();
3716 let compacted = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::ScannAh, 65).unwrap();
3717
3718 assert_eq!(compacted.runs.len(), 1);
3719 assert!((compacted.health().fragmentation() - 1.0).abs() < 1e-9);
3720 let run = &compacted.runs[0];
3721 let codes = &compacted.raw.as_slice()[run.codes.clone()];
3722 let mut decoded = Vec::new();
3723 for row in 0..run.count {
3724 unpack_scann_ah_row(codes, run.count, BLOCKS, row, &mut decoded).unwrap();
3725 }
3726 let mut expected = left_rows;
3727 expected.extend_from_slice(&right_rows);
3728 assert_eq!(decoded, expected);
3729 for row in 0..65 {
3730 assert_eq!(
3731 run_doc_id(compacted.raw.as_slice(), run, row).unwrap(),
3732 row as u32
3733 );
3734 }
3735 }
3736
3737 #[test]
3740 #[ignore]
3741 fn ann_merge_throughput_byte_copy_vs_compaction() {
3742 let code_size = 320usize;
3743 let clusters = 4_096u32;
3744 let vectors_per_source = 262_144usize;
3745 let sources_count = 4usize;
3746
3747 let mut sources_bytes = Vec::new();
3748 for source_index in 0..sources_count {
3749 let mut per_cluster: Vec<(Vec<u32>, Vec<u16>, Vec<u8>)> = Vec::new();
3750 let vectors_per_cluster = vectors_per_source / clusters as usize;
3751 let mut doc = 0u32;
3752 for cluster in 0..clusters {
3753 let mut docs = Vec::with_capacity(vectors_per_cluster);
3754 let mut ords = Vec::with_capacity(vectors_per_cluster);
3755 let mut codes = Vec::with_capacity(vectors_per_cluster * code_size);
3756 for _ in 0..vectors_per_cluster {
3757 docs.push(doc);
3758 ords.push(0u16);
3759 codes.extend(std::iter::repeat_n(
3760 (doc ^ cluster ^ source_index as u32) as u8,
3761 code_size,
3762 ));
3763 doc += 1;
3764 }
3765 per_cluster.push((docs, ords, codes));
3766 }
3767 let runs: Vec<BuildRun<'_>> = per_cluster
3768 .iter()
3769 .enumerate()
3770 .map(|(cluster, (docs, ords, codes))| BuildRun {
3771 cluster_id: cluster as u32,
3772 doc_ids: docs,
3773 ordinals: ords,
3774 codes,
3775 })
3776 .collect();
3777 let header = AnnDiskHeader {
3778 kind: AnnKind::BinaryIvf,
3779 routing: IvfRoutingMode::Hnsw,
3780 dim: code_size * 8,
3781 code_size,
3782 num_clusters: clusters,
3783 quantizer_version: 42,
3784 codebook_version: 0,
3785 vector_count: vectors_per_source,
3786 };
3787 let mut bytes = Vec::new();
3788 write_built_runs(header, &runs, &mut bytes).unwrap();
3789 sources_bytes.push(bytes);
3790 }
3791 let sources_open: Vec<AnnDiskIndex> = sources_bytes
3792 .iter()
3793 .map(|bytes| {
3794 AnnDiskIndex::open(
3795 OwnedBytes::new(bytes.clone()),
3796 AnnKind::BinaryIvf,
3797 (vectors_per_source * sources_count) as u32,
3798 )
3799 .unwrap()
3800 })
3801 .collect();
3802 let sources: Vec<(&AnnDiskIndex, u32)> = sources_open
3803 .iter()
3804 .enumerate()
3805 .map(|(index, source)| (source, (index * vectors_per_source) as u32))
3806 .collect();
3807 let payload_bytes = sources_bytes.iter().map(Vec::len).sum::<usize>();
3808
3809 let mut out = vec![0u8; payload_bytes + (1 << 20)];
3815 out.clear();
3816 let mut copy_secs = f64::INFINITY;
3819 let mut compact_secs = f64::INFINITY;
3820 let mut compacted_bytes = Vec::new();
3821 for _ in 0..3 {
3822 out.clear();
3823 let start = std::time::Instant::now();
3824 write_merged_ann(&sources, &mut out).unwrap();
3825 copy_secs = copy_secs.min(start.elapsed().as_secs_f64());
3826
3827 out.clear();
3828 let start = std::time::Instant::now();
3829 write_compacted_ann_cancellable(&sources, &mut out, None).unwrap();
3830 compact_secs = compact_secs.min(start.elapsed().as_secs_f64());
3831 compacted_bytes = out.clone();
3832 }
3833 let compacted = AnnDiskIndex::open(
3834 OwnedBytes::new(compacted_bytes),
3835 AnnKind::BinaryIvf,
3836 (vectors_per_source * sources_count) as u32,
3837 )
3838 .unwrap();
3839 assert!((compacted.health().fragmentation() - 1.0).abs() < 1e-9);
3840
3841 let gib = payload_bytes as f64 / (1u64 << 30) as f64;
3842 println!(
3843 "ann merge {:.2} GiB: byte-copy {:.3}s ({:.2} GiB/s), compaction {:.3}s \
3844 ({:.2} GiB/s), overhead {:.1}%",
3845 gib,
3846 copy_secs,
3847 gib / copy_secs,
3848 compact_secs,
3849 gib / compact_secs,
3850 100.0 * (compact_secs - copy_secs) / copy_secs,
3851 );
3852 }
3853
3854 #[test]
3856 fn compaction_refuses_non_binary_payloads() {
3857 let codec = std::sync::Arc::new(crate::structures::TqCodec::new(8));
3863 let mut builder = crate::structures::TqFlatBuilder::new(codec);
3864 builder
3865 .add_batch(
3866 &[(0, 0), (1, 0)],
3867 &[
3868 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
3870 ],
3871 )
3872 .unwrap();
3873 builder.finish();
3874 let mut bytes = Vec::new();
3875 write_built_tq_flat(&builder, &mut bytes).unwrap();
3876 let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::TqFlat, 2).unwrap();
3877 let error = write_compacted_ann_cancellable(&[(&disk, 0)], &mut Vec::new(), None)
3878 .expect_err("TQ payloads must not compact");
3879 assert!(error.to_string().contains("binary"), "{error}");
3880 }
3881
3882 #[test]
3886 fn ann_health_measures_skew_and_fragmentation() {
3887 let a0_docs = [0u32, 1, 2, 3, 4, 5];
3892 let a0_ords = [0u16; 6];
3893 let a0_codes = [0xAAu8; 6];
3894 let a3_docs = [6u32, 7];
3895 let a3_ords = [0u16; 2];
3896 let a3_codes = [0x0Fu8; 2];
3897 let a_runs = [
3898 BuildRun {
3899 cluster_id: 0,
3900 doc_ids: &a0_docs,
3901 ordinals: &a0_ords,
3902 codes: &a0_codes,
3903 },
3904 BuildRun {
3905 cluster_id: 3,
3906 doc_ids: &a3_docs,
3907 ordinals: &a3_ords,
3908 codes: &a3_codes,
3909 },
3910 ];
3911 let mut header = binary_header(8);
3912 header.num_clusters = 8;
3913 let mut a_bytes = Vec::new();
3914 write_built_runs(header, &a_runs, &mut a_bytes).unwrap();
3915 let a = AnnDiskIndex::open(OwnedBytes::new(a_bytes), AnnKind::BinaryIvf, 8).unwrap();
3916
3917 let b0_docs = [0u32, 1];
3918 let b0_ords = [0u16; 2];
3919 let b0_codes = [0xBBu8; 2];
3920 let b_runs = [BuildRun {
3921 cluster_id: 0,
3922 doc_ids: &b0_docs,
3923 ordinals: &b0_ords,
3924 codes: &b0_codes,
3925 }];
3926 let mut header_b = binary_header(2);
3927 header_b.num_clusters = 8;
3928 let mut b_bytes = Vec::new();
3929 write_built_runs(header_b, &b_runs, &mut b_bytes).unwrap();
3930 let b = AnnDiskIndex::open(OwnedBytes::new(b_bytes), AnnKind::BinaryIvf, 2).unwrap();
3931
3932 let mut merged_bytes = Vec::new();
3933 write_merged_ann(&[(&a, 0), (&b, 8)], &mut merged_bytes).unwrap();
3934 let disk =
3935 AnnDiskIndex::open(OwnedBytes::new(merged_bytes), AnnKind::BinaryIvf, 10).unwrap();
3936
3937 let health = disk.health();
3938 assert_eq!(health.vectors, 10);
3939 assert_eq!(health.clusters_nonempty, 2);
3940 assert_eq!(health.clusters_total, 8);
3941 assert_eq!(health.runs, 3);
3942 assert_eq!(health.largest_cluster, 0);
3943 assert_eq!(health.largest_cluster_vectors, 8);
3944 assert!((health.largest_cluster_share() - 0.8).abs() < 1e-9);
3945 assert!((health.fragmentation() - 1.5).abs() < 1e-9);
3947 assert!(
3949 (health.imbalance - 1.36).abs() < 1e-9,
3950 "{}",
3951 health.imbalance
3952 );
3953 assert_eq!(health.payload_bytes, 10);
3955 }
3956
3957 #[test]
3958 fn ann_health_is_balanced_at_one() {
3959 let docs: Vec<Vec<u32>> = (0..4).map(|c| vec![c * 2, c * 2 + 1]).collect();
3960 let ords = [0u16; 2];
3961 let codes = [0x55u8; 2];
3962 let runs: Vec<BuildRun<'_>> = docs
3963 .iter()
3964 .enumerate()
3965 .map(|(cluster, doc_ids)| BuildRun {
3966 cluster_id: cluster as u32,
3967 doc_ids,
3968 ordinals: &ords,
3969 codes: &codes,
3970 })
3971 .collect();
3972 let mut header = binary_header(8);
3973 header.num_clusters = 4;
3974 let mut bytes = Vec::new();
3975 write_built_runs(header, &runs, &mut bytes).unwrap();
3976 let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::BinaryIvf, 8).unwrap();
3977 let health = disk.health();
3978 assert!((health.imbalance - 1.0).abs() < 1e-9);
3979 assert!((health.fragmentation() - 1.0).abs() < 1e-9);
3980 assert!((health.largest_cluster_share() - 0.25).abs() < 1e-9);
3981 }
3982
3983 fn binary_header(vector_count: usize) -> AnnDiskHeader {
3984 AnnDiskHeader {
3985 kind: AnnKind::BinaryIvf,
3986 routing: IvfRoutingMode::Hnsw,
3987 dim: 8,
3988 code_size: 1,
3989 num_clusters: 2,
3990 quantizer_version: 42,
3991 codebook_version: 0,
3992 vector_count,
3993 }
3994 }
3995
3996 fn payload_end(index: &AnnDiskIndex) -> usize {
3997 index.runs.iter().map(|run| run.codes.end).max().unwrap()
3998 }
3999
4000 #[test]
4001 fn ann_prefetch_ranges_are_sorted_and_only_merge_page_near_extents() {
4002 let mut ranges = vec![
4003 15_000..16_000,
4004 0..1_000,
4005 9_000..10_000,
4006 1_000..2_000,
4007 7_000..8_000,
4008 ];
4009 coalesce_prefetch_ranges(&mut ranges);
4010 assert_eq!(ranges, [0..2_000, 7_000..10_000, 15_000..16_000]);
4011 }
4012
4013 #[test]
4014 fn normal_merge_copies_ann_payload_columns_byte_for_byte() {
4015 let first_doc_0 = [0u32];
4016 let first_doc_1 = [1u32];
4017 let first_ord_0 = [0u16];
4018 let first_ord_1 = [2u16];
4019 let first_code_0 = [0x00u8];
4020 let first_code_1 = [0xffu8];
4021 let first_runs = [
4022 BuildRun {
4023 cluster_id: 0,
4024 doc_ids: &first_doc_0,
4025 ordinals: &first_ord_0,
4026 codes: &first_code_0,
4027 },
4028 BuildRun {
4029 cluster_id: 1,
4030 doc_ids: &first_doc_1,
4031 ordinals: &first_ord_1,
4032 codes: &first_code_1,
4033 },
4034 ];
4035 let mut first_bytes = Vec::new();
4036 write_built_runs(binary_header(2), &first_runs, &mut first_bytes).unwrap();
4037 let first = AnnDiskIndex::open(OwnedBytes::new(first_bytes.clone()), AnnKind::BinaryIvf, 2)
4038 .unwrap();
4039
4040 let second_docs = [0u32, 1u32];
4041 let second_ords = [1u16, 0u16];
4042 let second_codes = [0x0fu8, 0xf0u8];
4043 let second_runs = [BuildRun {
4044 cluster_id: 0,
4045 doc_ids: &second_docs,
4046 ordinals: &second_ords,
4047 codes: &second_codes,
4048 }];
4049 let mut second_bytes = Vec::new();
4050 write_built_runs(binary_header(2), &second_runs, &mut second_bytes).unwrap();
4051 let second =
4052 AnnDiskIndex::open(OwnedBytes::new(second_bytes.clone()), AnnKind::BinaryIvf, 2)
4053 .unwrap();
4054
4055 let mut merged_bytes = Vec::new();
4056 write_merged_ann(&[(&first, 0), (&second, 2)], &mut merged_bytes).unwrap();
4057 let merged =
4058 AnnDiskIndex::open(OwnedBytes::new(merged_bytes.clone()), AnnKind::BinaryIvf, 4)
4059 .unwrap();
4060
4061 let mut expected_payload = first_bytes[ANN_HEADER_SIZE..payload_end(&first)].to_vec();
4062 expected_payload.extend_from_slice(&second_bytes[ANN_HEADER_SIZE..payload_end(&second)]);
4063 assert_eq!(
4064 &merged_bytes[ANN_HEADER_SIZE..payload_end(&merged)],
4065 expected_payload.as_slice(),
4066 "normal merge must not decode or rewrite any corpus-sized ANN column",
4067 );
4068
4069 let mut docs: Vec<u32> = merged
4070 .search_binary_clusters::<false>(&[0], 4, &[0, 1])
4071 .unwrap()
4072 .into_iter()
4073 .map(|result| result.0)
4074 .collect();
4075 docs.sort_unstable();
4076 assert_eq!(docs, [0, 1, 2, 3]);
4077 let serial = merged
4078 .search_binary_clusters_with_tuning::<false>(&[0], 4, &[0, 1], usize::MAX)
4079 .unwrap();
4080 let parallel = rayon::ThreadPoolBuilder::new()
4081 .num_threads(4)
4082 .build()
4083 .unwrap()
4084 .install(|| merged.search_binary_clusters_with_tuning::<false>(&[0], 4, &[0, 1], 1))
4085 .unwrap();
4086 assert_eq!(parallel, serial, "parallel binary top-k changed results");
4087
4088 let mut second_merge_bytes = Vec::new();
4092 write_merged_ann(&[(&merged, 0), (&first, 4)], &mut second_merge_bytes).unwrap();
4093 let second_merge = AnnDiskIndex::open(
4094 OwnedBytes::new(second_merge_bytes.clone()),
4095 AnnKind::BinaryIvf,
4096 6,
4097 )
4098 .unwrap();
4099 let mut expected_second_payload =
4100 merged_bytes[ANN_HEADER_SIZE..payload_end(&merged)].to_vec();
4101 expected_second_payload
4102 .extend_from_slice(&first_bytes[ANN_HEADER_SIZE..payload_end(&first)]);
4103 assert_eq!(
4104 &second_merge_bytes[ANN_HEADER_SIZE..payload_end(&second_merge)],
4105 expected_second_payload.as_slice(),
4106 );
4107 let mut docs: Vec<u32> = second_merge
4108 .search_binary_clusters::<false>(&[0], 6, &[0, 1])
4109 .unwrap()
4110 .into_iter()
4111 .map(|result| result.0)
4112 .collect();
4113 docs.sort_unstable();
4114 assert_eq!(docs, [0, 1, 2, 3, 4, 5]);
4115 }
4116
4117 #[test]
4118 fn legacy_ivf_tq_payload_is_rejected_while_opening() {
4119 let dim = 8;
4120 let marked_version = crate::structures::mark_ivf_tq_cosine_generation(7);
4121 let centroids = crate::structures::CoarseCentroids {
4122 num_clusters: 1,
4123 dim,
4124 centroids: vec![0.0; dim],
4125 version: marked_version,
4126 soar_config: None,
4127 routing_index: None,
4128 };
4129 let mut bytes = crate::segment::ann_build::build_ivf_tq(
4130 dim,
4131 IvfRoutingMode::Flat,
4132 ¢roids,
4133 &[(0, 0)],
4134 &[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
4135 )
4136 .unwrap();
4137
4138 bytes[24..32].copy_from_slice(&7u64.to_le_bytes());
4140 let error = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::IvfTq, 1)
4141 .err()
4142 .expect("legacy IVF-TQ payload must fail while opening")
4143 .to_string();
4144 assert!(error.contains("unsupported legacy generation"), "{error}");
4145 }
4146
4147 #[test]
4148 fn binary_combined_search_deduplicates_soar_and_bounds_document_results() {
4149 let cluster_0_docs = [0u32, 0, 1];
4153 let cluster_0_ordinals = [0u16, 1, 0];
4154 let cluster_0_codes = [0x00u8, 0xff, 0x03];
4155 let cluster_1_docs = [0u32, 1, 2];
4156 let cluster_1_ordinals = [0u16, 1, 0];
4157 let cluster_1_codes = [0x00u8, 0x0c, 0xf0];
4158 let runs = [
4159 BuildRun {
4160 cluster_id: 0,
4161 doc_ids: &cluster_0_docs,
4162 ordinals: &cluster_0_ordinals,
4163 codes: &cluster_0_codes,
4164 },
4165 BuildRun {
4166 cluster_id: 1,
4167 doc_ids: &cluster_1_docs,
4168 ordinals: &cluster_1_ordinals,
4169 codes: &cluster_1_codes,
4170 },
4171 ];
4172 let mut bytes = Vec::new();
4173 write_built_runs(binary_header(6), &runs, &mut bytes).unwrap();
4174 let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::BinaryIvf, 3).unwrap();
4175 let parallel_pool = rayon::ThreadPoolBuilder::new()
4176 .num_threads(4)
4177 .build()
4178 .unwrap();
4179 let serial_documents = disk
4180 .search_binary_clusters_with_tuning::<true>(&[0], 3, &[0, 1], usize::MAX)
4181 .unwrap();
4182 let parallel_documents = parallel_pool
4183 .install(|| disk.search_binary_clusters_with_tuning::<true>(&[0], 3, &[0, 1], 1))
4184 .unwrap();
4185 assert_eq!(parallel_documents, serial_documents);
4186
4187 let sum = crate::query::MultiValueCombiner::Sum;
4191 let (result, probed) = disk
4192 .search_binary_combined_documents(1, &[0], &[0, 1], sum)
4193 .unwrap();
4194 let parallel_sum = parallel_pool
4195 .install(|| disk.search_binary_combined_documents_with_tuning(1, &[0], &[0, 1], sum, 1))
4196 .unwrap();
4197 assert_eq!(parallel_sum, (result.clone(), probed.clone()));
4198 assert_eq!(result.len(), 1, "combined search must honor k");
4199 assert_eq!(
4200 result[0].doc_id, 1,
4201 "SOAR duplicate changed {sum:?} ranking: {result:?}",
4202 );
4203 assert_eq!(
4206 probed
4207 .iter()
4208 .map(|&(doc_id, ordinal, _)| (doc_id, ordinal))
4209 .collect::<Vec<_>>(),
4210 vec![(1, 0), (1, 1)],
4211 "{sum:?}",
4212 );
4213
4214 let smooth_max = crate::query::MultiValueCombiner::default();
4220 let (result, probed) = disk
4221 .search_binary_combined_documents(1, &[0], &[0, 1], smooth_max)
4222 .unwrap();
4223 let parallel_smooth_max = parallel_pool
4224 .install(|| {
4225 disk.search_binary_combined_documents_with_tuning(1, &[0], &[0, 1], smooth_max, 1)
4226 })
4227 .unwrap();
4228 assert_eq!(parallel_smooth_max, (result.clone(), probed.clone()));
4229 assert_eq!(result.len(), 1, "combined search must honor k");
4230 assert_eq!(result[0].doc_id, 0, "{result:?}");
4231 let expected = smooth_max.combine(&[(0, 1.0), (1, 0.0)]);
4232 assert!(
4233 (result[0].score - expected).abs() < 1e-6,
4234 "SOAR duplicate leaked into {smooth_max:?}: got {}, expected {expected}",
4235 result[0].score,
4236 );
4237 assert_eq!(
4238 probed
4239 .iter()
4240 .map(|&(doc_id, ordinal, _)| (doc_id, ordinal))
4241 .collect::<Vec<_>>(),
4242 vec![(0, 0), (0, 1)],
4243 "{smooth_max:?}",
4244 );
4245
4246 let (top_two, probed) = disk
4247 .search_binary_combined_documents(
4248 2,
4249 &[0],
4250 &[0, 1],
4251 crate::query::MultiValueCombiner::Sum,
4252 )
4253 .unwrap();
4254 assert_eq!(
4255 top_two
4256 .iter()
4257 .map(|candidate| candidate.doc_id)
4258 .collect::<Vec<_>>(),
4259 vec![1, 0],
4260 );
4261 assert_eq!(top_two.len(), 2, "full probing must still return at most k");
4262 assert_eq!(
4265 probed
4266 .iter()
4267 .map(|&(doc_id, ordinal, _)| (doc_id, ordinal))
4268 .collect::<Vec<_>>(),
4269 vec![(0, 0), (0, 1), (1, 0), (1, 1)],
4270 );
4271 }
4272
4273 #[test]
4274 fn combined_ordinal_reduction_handles_out_of_order_runs_for_every_combiner() {
4275 let out_of_order_with_duplicate = vec![
4276 (7, 1, 0.4),
4277 (3, 0, 0.8),
4278 (7, 0, 0.6),
4279 (3, 1, 0.2),
4280 (7, 1, 0.5), ];
4282 for combiner in [
4283 crate::query::MultiValueCombiner::Max,
4284 crate::query::MultiValueCombiner::Sum,
4285 crate::query::MultiValueCombiner::Avg,
4286 crate::query::MultiValueCombiner::default(),
4287 crate::query::MultiValueCombiner::WeightedTopK { k: 2, decay: 0.7 },
4288 ] {
4289 let actual = combine_scored_ordinals(out_of_order_with_duplicate.clone(), 2, combiner);
4290 let mut expected = vec![
4291 AnnDocumentCandidate {
4292 doc_id: 3,
4293 score: combiner.combine(&[(0, 0.8), (1, 0.2)]),
4294 },
4295 AnnDocumentCandidate {
4296 doc_id: 7,
4297 score: combiner.combine(&[(0, 0.6), (1, 0.5)]),
4298 },
4299 ];
4300 expected.sort_unstable_by(|left, right| {
4301 right
4302 .score
4303 .total_cmp(&left.score)
4304 .then_with(|| left.doc_id.cmp(&right.doc_id))
4305 });
4306 assert_eq!(actual, expected, "combiner {combiner:?}");
4307 }
4308 }
4309
4310 fn build_tq_payload(dim: usize, count: usize, seed: u64) -> (Vec<u8>, Vec<Vec<f32>>) {
4311 let codec = std::sync::Arc::new(crate::structures::TqCodec::new(dim));
4312 let mut builder = crate::structures::TqFlatBuilder::new(std::sync::Arc::clone(&codec));
4313 let mut state = seed;
4314 let mut vectors = Vec::new();
4315 let mut flat = Vec::new();
4316 for _ in 0..count {
4317 let vector: Vec<f32> = (0..dim)
4318 .map(|_| {
4319 state = state
4320 .wrapping_mul(6364136223846793005)
4321 .wrapping_add(1442695040888963407);
4322 ((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5
4323 })
4324 .collect();
4325 flat.extend_from_slice(&vector);
4326 vectors.push(vector);
4327 }
4328 let labels: Vec<(u32, u16)> = (0..count).map(|index| (index as u32, 0)).collect();
4329 builder.add_batch(&labels, &flat).unwrap();
4330 builder.finish();
4331 let mut bytes = Vec::new();
4332 write_built_tq_flat(&builder, &mut bytes).unwrap();
4333 (bytes, vectors)
4334 }
4335
4336 #[test]
4337 fn tq_combined_scan_ranks_complete_documents_instead_of_individual_values() {
4338 let dim = 8;
4339 let codec = std::sync::Arc::new(crate::structures::TqCodec::new(dim));
4340 let mut builder = crate::structures::TqFlatBuilder::new(std::sync::Arc::clone(&codec));
4341 let labels = [(0u32, 0u16), (1, 0), (1, 1)];
4342 let vectors = [
4343 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.8, 0.6, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.8, -0.6, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
4346 ];
4347 builder.add_batch(&labels, &vectors).unwrap();
4348 builder.finish();
4349 let mut bytes = Vec::new();
4350 write_built_tq_flat(&builder, &mut bytes).unwrap();
4351 let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::TqFlat, 2).unwrap();
4352 let query = [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0];
4353 let plan = crate::structures::TqQueryPlan::build(&codec, &query);
4354
4355 let max = disk.search_tq_distinct(1, &plan).unwrap();
4356 let sum = disk
4357 .search_tq_combined_documents(1, &plan, crate::query::MultiValueCombiner::Sum)
4358 .unwrap();
4359 let default_combiner = disk
4360 .search_tq_combined_documents(1, &plan, crate::query::MultiValueCombiner::default())
4361 .unwrap();
4362 assert_eq!(max[0].0, 0, "fixture must favor doc 0 by Max: {max:?}");
4363 assert_eq!(
4364 sum[0].doc_id, 1,
4365 "two complete values must make doc 1 win by Sum: {sum:?}"
4366 );
4367 assert_eq!(
4371 default_combiner[0].doc_id, 0,
4372 "default smooth-max must follow the best value: {default_combiner:?}"
4373 );
4374 assert!(sum[0].score > max[0].2);
4375
4376 let (single_bytes, _) = build_tq_payload(dim, 1, 91);
4380 let single = AnnDiskIndex::open(OwnedBytes::new(single_bytes), AnnKind::TqFlat, 1).unwrap();
4381 let mut merged_bytes = Vec::new();
4382 write_merged_ann(&[(&disk, 1), (&single, 0)], &mut merged_bytes).unwrap();
4383 let merged = AnnDiskIndex::open(OwnedBytes::new(merged_bytes), AnnKind::TqFlat, 3).unwrap();
4384 let merged_sum = merged
4385 .search_tq_combined_documents(1, &plan, crate::query::MultiValueCombiner::Sum)
4386 .unwrap();
4387 assert_eq!(merged_sum[0].doc_id, 2);
4388 }
4389
4390 #[test]
4391 fn tq_payload_roundtrip_search_and_pure_copy_merge() {
4392 let dim = 20; let count = 21;
4394 let (bytes, vectors) = build_tq_payload(dim, count, 42);
4395 let index = AnnDiskIndex::open(
4396 OwnedBytes::new(bytes.clone()),
4397 AnnKind::TqFlat,
4398 count as u32,
4399 )
4400 .unwrap();
4401 assert_eq!(index.header().vector_count, count);
4402
4403 let codec = crate::structures::TqCodec::new(dim);
4405 for target in [0usize, 7, 20] {
4406 let plan = crate::structures::TqQueryPlan::build(&codec, &vectors[target]);
4407 let results = index.search_tq_distinct(3, &plan).unwrap();
4408 assert_eq!(
4409 results[0].0, target as u32,
4410 "query duplicating vector {target} must rank it first: {results:?}"
4411 );
4412 }
4413
4414 let (second_bytes, _) = build_tq_payload(dim, 5, 77);
4416 let second =
4417 AnnDiskIndex::open(OwnedBytes::new(second_bytes.clone()), AnnKind::TqFlat, 5).unwrap();
4418 let mut merged_bytes = Vec::new();
4419 write_merged_ann(&[(&index, 0), (&second, count as u32)], &mut merged_bytes).unwrap();
4420 let merged = AnnDiskIndex::open(
4421 OwnedBytes::new(merged_bytes.clone()),
4422 AnnKind::TqFlat,
4423 count as u32 + 5,
4424 )
4425 .unwrap();
4426 let mut expected_payload = bytes[ANN_HEADER_SIZE..payload_end(&index)].to_vec();
4427 expected_payload.extend_from_slice(&second_bytes[ANN_HEADER_SIZE..payload_end(&second)]);
4428 assert_eq!(
4429 &merged_bytes[ANN_HEADER_SIZE..payload_end(&merged)],
4430 expected_payload.as_slice(),
4431 "TQ merge must be a pure byte copy of the source columns",
4432 );
4433 let plan = crate::structures::TqQueryPlan::build(&codec, &vectors[7]);
4434 let results = merged.search_tq_distinct(1, &plan).unwrap();
4435 assert_eq!(results[0].0, 7, "merged payload must keep doc bases");
4436 }
4437
4438 #[test]
4445 fn open_refuses_scann_ah_payload_with_pre_v2_fast_scan_layout() {
4446 use crate::structures::vector::scann::{
4447 FAST_SCAN_LANES, FAST_SCAN_LAYOUT_VERSION, pack_fast_scan_block,
4448 };
4449 const BLOCKS: usize = 4;
4450 const DIMS_PER_BLOCK: usize = 2;
4451 const DIM: usize = BLOCKS * DIMS_PER_BLOCK;
4452 const REVISION_OFFSET: usize = 20;
4455
4456 let count = FAST_SCAN_LANES;
4457 let rows: Vec<u8> = (0..count * BLOCKS).map(|i| (i % 16) as u8).collect();
4458 let mut codes = Vec::new();
4459 pack_fast_scan_block(&rows, BLOCKS, &mut codes).unwrap();
4460 let docs: Vec<u32> = (0..count as u32).collect();
4461 let ordinals = vec![0u16; count];
4462 let run = BuildRun {
4463 cluster_id: 0,
4464 doc_ids: &docs,
4465 ordinals: &ordinals,
4466 codes: &codes,
4467 };
4468 let header = AnnDiskHeader {
4469 kind: AnnKind::ScannAh,
4470 routing: IvfRoutingMode::Flat,
4471 dim: DIM,
4472 code_size: DIMS_PER_BLOCK,
4473 num_clusters: 1,
4474 quantizer_version: 1,
4475 codebook_version: 1,
4476 vector_count: count,
4477 };
4478 let mut bytes = Vec::new();
4479 write_built_runs(header, &[run], &mut bytes).unwrap();
4480
4481 let stamped = u32::from_le_bytes(
4482 bytes[REVISION_OFFSET..REVISION_OFFSET + 4]
4483 .try_into()
4484 .unwrap(),
4485 );
4486 assert_eq!(stamped, u32::from(FAST_SCAN_LAYOUT_VERSION));
4487 AnnDiskIndex::open(
4488 OwnedBytes::new(bytes.clone()),
4489 AnnKind::ScannAh,
4490 count as u32,
4491 )
4492 .expect("current layout revision opens");
4493
4494 let mut legacy = bytes;
4496 legacy[REVISION_OFFSET..REVISION_OFFSET + 4].copy_from_slice(&0u32.to_le_bytes());
4497 let error =
4498 match AnnDiskIndex::open(OwnedBytes::new(legacy), AnnKind::ScannAh, count as u32) {
4499 Ok(_) => panic!("a pre-v2 FastScan payload must be refused at open"),
4500 Err(error) => error,
4501 };
4502 let message = error.to_string();
4503 assert!(
4504 message.contains("FastScan layout revision 0") && message.contains("Rebuild"),
4505 "unexpected error: {message}"
4506 );
4507
4508 let (tq, _) = build_tq_payload(8, 40, 5);
4510 let mut tq_bad = tq;
4511 tq_bad[REVISION_OFFSET..REVISION_OFFSET + 4].copy_from_slice(&7u32.to_le_bytes());
4512 let error = AnnDiskIndex::open(OwnedBytes::new(tq_bad), AnnKind::TqFlat, 40)
4513 .err()
4514 .expect("non-zero reserved field must be refused");
4515 assert!(
4516 error.to_string().contains("reserved field is non-zero"),
4517 "unexpected error: {error}"
4518 );
4519 }
4520
4521 #[test]
4522 fn open_rejects_doc_id_column_above_run_maximum() {
4523 let dim = 8;
4524 let count = 40;
4525 let (bytes, _) = build_tq_payload(dim, count, 5);
4526 let disk = AnnDiskIndex::open(
4527 OwnedBytes::new(bytes.clone()),
4528 AnnKind::TqFlat,
4529 count as u32,
4530 )
4531 .unwrap();
4532 let run = &disk.runs[0];
4533 let mut corrupt = bytes;
4536 let offset = run.doc_ids.start + 3 * 4;
4537 corrupt[offset..offset + 4].copy_from_slice(&(run.max_doc_id + 1).to_le_bytes());
4538 let error =
4539 match AnnDiskIndex::open(OwnedBytes::new(corrupt), AnnKind::TqFlat, count as u32) {
4540 Ok(_) => panic!("corrupt doc-ID column must fail at open"),
4541 Err(error) => error,
4542 };
4543 assert!(
4544 error.to_string().contains("above its declared maximum"),
4545 "unexpected error: {error}"
4546 );
4547 }
4548
4549 #[test]
4550 fn tq_parallel_fold_matches_a_sequential_scan() {
4551 use crate::structures::vector::quantization::{
4552 TQ_BLOCK_LANES, tq_block_bytes, tq_score_block,
4553 };
4554
4555 let dim = 8;
4558 let count = TQ_PARALLEL_SCAN_MIN_VECTORS;
4559 let (bytes, vectors) = build_tq_payload(dim, count, 87);
4560 let disk =
4561 AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::TqFlat, count as u32).unwrap();
4562 let codec = crate::structures::TqCodec::new(dim);
4563 let plan = crate::structures::TqQueryPlan::build(&codec, &vectors[count / 3]);
4564 let k = 31;
4565
4566 let block_bytes = tq_block_bytes(disk.header().code_size);
4567 assert_eq!(
4568 disk.runs.len(),
4569 1,
4570 "the regression must parallelize chunks inside one run"
4571 );
4572 assert!(
4573 disk.runs[0].codes.len() / block_bytes > TQ_PARALLEL_SCAN_CHUNK_BLOCKS,
4574 "the single run must span multiple parallel chunks"
4575 );
4576 let raw = disk.raw.as_slice();
4577 let mut reference = BoundedAnnCollector::<true, true>::new(k);
4578 let mut scores = [0.0f32; TQ_BLOCK_LANES];
4579 for run in &disk.runs {
4580 let codes = &raw[run.codes.clone()];
4581 for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
4582 tq_score_block(&plan, block, &mut scores);
4583 let lane_base = block_index * TQ_BLOCK_LANES;
4584 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
4585 for (lane, &score) in scores.iter().enumerate().take(lanes) {
4586 let index = lane_base + lane;
4587 reference.insert(
4588 run_doc_id(raw, run, index).unwrap(),
4589 read_u16(raw, run.ordinals.start + index * 2),
4590 score,
4591 );
4592 }
4593 }
4594 }
4595
4596 assert_eq!(
4597 disk.search_tq_distinct(k, &plan).unwrap(),
4598 reference.into_sorted_results(),
4599 );
4600 }
4601
4602 #[test]
4603 fn ivf_tq_scale_pruning_matches_the_unpruned_scan() {
4604 use crate::structures::vector::ivf::{CoarseCentroids, CoarseConfig};
4605 use crate::structures::vector::quantization::{
4606 TQ_BLOCK_LANES, tq_ivf_block_bytes, tq_score_ivf_block,
4607 };
4608 use crate::structures::{IvfTqIndex, TqCodec, TqIvfEncodeScratch, TqIvfQueryPlan};
4609
4610 let dim = 32;
4611 let count = 400usize;
4612 let codec = std::sync::Arc::new(TqCodec::new(dim));
4613 let mut state = 5u64;
4614 let mut next = move || {
4615 state = state
4616 .wrapping_mul(6364136223846793005)
4617 .wrapping_add(1442695040888963407);
4618 ((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5
4619 };
4620 let vectors: Vec<Vec<f32>> = (0..count)
4621 .map(|_| {
4622 let mut v: Vec<f32> = (0..dim).map(|_| next()).collect();
4623 let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
4624 v.iter_mut().for_each(|x| *x /= norm);
4625 v
4626 })
4627 .collect();
4628 let mut centroids = CoarseCentroids::train(&CoarseConfig::new(dim, 8), &vectors, "test");
4629 centroids.version = crate::structures::mark_ivf_tq_cosine_generation(centroids.version);
4630 let mut index = IvfTqIndex::new(
4631 dim,
4632 crate::dsl::IvfRoutingMode::Flat,
4633 centroids.version,
4634 std::sync::Arc::clone(&codec),
4635 );
4636 let mut scratch = TqIvfEncodeScratch::default();
4637 for (i, vector) in vectors.iter().enumerate() {
4638 index.add_vector(
4639 ¢roids,
4640 (i / 2) as u32,
4641 (i % 2) as u16,
4642 vector,
4643 &mut scratch,
4644 );
4645 }
4646 let mut bytes = Vec::new();
4647 write_built_ivf_tq(&index, centroids.num_clusters, &mut bytes).unwrap();
4648 let disk =
4649 AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::IvfTq, count as u32).unwrap();
4650
4651 let block_bytes = tq_ivf_block_bytes(disk.header().code_size);
4654 let raw = disk.raw.as_slice();
4655 for run in &disk.runs {
4656 let codes = &raw[run.codes.clone()];
4657 let mut previous_scale = f32::INFINITY;
4658 for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
4659 let lane_base = block_index * TQ_BLOCK_LANES;
4660 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
4661 let mut block_scales = block[..TQ_BLOCK_LANES * size_of::<f32>()]
4662 .chunks_exact(size_of::<f32>())
4663 .take(lanes)
4664 .map(|lane| f32::from_le_bytes(lane.try_into().unwrap()));
4665 let first_scale = block_scales.next().unwrap();
4666 assert_eq!(tq_ivf_block_max_scale(block), first_scale);
4667 assert!(first_scale <= previous_scale);
4668 previous_scale = first_scale;
4669 for scale in block_scales {
4670 assert!(scale <= previous_scale);
4671 previous_scale = scale;
4672 }
4673 }
4674 }
4675 let k = 10;
4676 let parallel_pool = rayon::ThreadPoolBuilder::new()
4677 .num_threads(4)
4678 .build()
4679 .unwrap();
4680 for query_seed in [1u64, 9, 42] {
4681 let mut qstate = query_seed;
4682 let mut qnext = move || {
4683 qstate = qstate
4684 .wrapping_mul(6364136223846793005)
4685 .wrapping_add(1442695040888963407);
4686 ((qstate >> 33) as f32 / (1u64 << 31) as f32) - 0.5
4687 };
4688 let query: Vec<f32> = (0..dim).map(|_| qnext()).collect();
4689 let plan = TqIvfQueryPlan::build(
4690 ¢roids,
4691 &codec,
4692 &query,
4693 8,
4694 crate::dsl::IvfRoutingMode::Flat,
4695 );
4696 let mut reference = BoundedAnnCollector::<true, true>::new(k);
4699 let mut unpruned_scores = Vec::new();
4700 let mut scores = [0.0f32; TQ_BLOCK_LANES];
4701 for (cluster_id, cluster_dot) in plan.cluster_dots() {
4702 for run in disk.cluster_runs(cluster_id) {
4703 let codes = &raw[run.codes.clone()];
4704 for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
4705 tq_score_ivf_block(plan.tq_plan(), block, cluster_dot, &mut scores);
4706 let lane_base = block_index * TQ_BLOCK_LANES;
4707 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
4708 for (lane, &score) in scores.iter().enumerate().take(lanes) {
4709 let idx = lane_base + lane;
4710 reference.insert(
4711 run_doc_id(raw, run, idx).unwrap(),
4712 read_u16(raw, run.ordinals.start + idx * 2),
4713 score,
4714 );
4715 unpruned_scores.push((
4716 run_doc_id(raw, run, idx).unwrap(),
4717 read_u16(raw, run.ordinals.start + idx * 2),
4718 score,
4719 ));
4720 }
4721 }
4722 }
4723 }
4724
4725 let pruned = disk.search_ivf_tq_distinct(k, &plan).unwrap();
4726 let forced_parallel = parallel_pool.install(|| {
4727 disk.search_ivf_tq_distinct_with_tuning(k, &plan, 1, 4)
4730 .unwrap()
4731 });
4732 let reference = reference.into_sorted_results();
4733 assert_eq!(
4734 pruned, reference,
4735 "scale-bound pruning must not change the estimated top-k (seed {query_seed})"
4736 );
4737 assert_eq!(
4738 forced_parallel, reference,
4739 "parallel IVF-TQ pruning must match the unpruned top-k (seed {query_seed})"
4740 );
4741 for combiner in [
4742 crate::query::MultiValueCombiner::Max,
4743 crate::query::MultiValueCombiner::Sum,
4744 crate::query::MultiValueCombiner::Avg,
4745 crate::query::MultiValueCombiner::default(),
4746 crate::query::MultiValueCombiner::WeightedTopK { k: 3, decay: 0.7 },
4747 ] {
4748 let expected = combine_scored_ordinals(unpruned_scores.clone(), k, combiner);
4749 let combined = disk
4750 .search_ivf_tq_combined_documents(k, &plan, combiner)
4751 .unwrap();
4752 let parallel_combined = parallel_pool
4753 .install(|| {
4754 disk.search_ivf_tq_combined_documents_with_tuning(k, &plan, combiner, 1, 4)
4755 })
4756 .unwrap();
4757 assert_eq!(
4758 combined, expected,
4759 "combined IVF-TQ scan diverged from the unpruned reference \
4760 for {combiner:?} (seed {query_seed})",
4761 );
4762 assert_eq!(
4763 parallel_combined, expected,
4764 "parallel combined IVF-TQ scan diverged from the unpruned reference \
4765 for {combiner:?} (seed {query_seed})",
4766 );
4767 assert!(combined.len() <= k);
4768 }
4769 }
4770 }
4771
4772 #[test]
4780 #[ignore]
4781 fn ivf_tq_parallel_scan_benchmark() {
4782 use crate::structures::vector::ivf::CoarseCentroids;
4783 use crate::structures::{IvfTqIndex, TqCodec, TqIvfEncodeScratch, TqIvfQueryPlan};
4784
4785 fn env_usize(name: &str, default: usize) -> usize {
4786 std::env::var(name)
4787 .ok()
4788 .map(|value| {
4789 value
4790 .parse::<usize>()
4791 .unwrap_or_else(|_| panic!("{name} must be a positive integer"))
4792 })
4793 .unwrap_or(default)
4794 .max(1)
4795 }
4796
4797 fn median_ms(samples: &mut [f64]) -> f64 {
4798 samples.sort_unstable_by(f64::total_cmp);
4799 samples[samples.len() / 2]
4800 }
4801
4802 let dim = env_usize("IVF_SCAN_BENCH_DIM", 128);
4803 let count = env_usize("IVF_SCAN_BENCH_DOCS", 262_144);
4804 let iterations = env_usize("IVF_SCAN_BENCH_ITERS", 30);
4805 let threads = env_usize("IVF_SCAN_BENCH_THREADS", num_cpus::get().min(8));
4806 let codec = std::sync::Arc::new(TqCodec::new(dim));
4807 let version = crate::structures::mark_ivf_tq_cosine_generation(0x51ca_0001);
4808 let mut centroid = vec![0.0f32; dim];
4809 centroid[0] = 1.0;
4810 let centroids = CoarseCentroids {
4811 num_clusters: 1,
4812 dim,
4813 centroids: centroid,
4814 version,
4815 soar_config: None,
4816 routing_index: None,
4817 };
4818 let mut index = IvfTqIndex::new(
4819 dim,
4820 crate::dsl::IvfRoutingMode::Flat,
4821 version,
4822 std::sync::Arc::clone(&codec),
4823 );
4824 let mut scratch = TqIvfEncodeScratch::default();
4825 let mut state = 0x5eed_u64;
4826 let mut vector = vec![0.0f32; dim];
4827 for posting_index in 0..count {
4828 for value in &mut vector {
4829 state = state
4830 .wrapping_mul(6364136223846793005)
4831 .wrapping_add(1442695040888963407);
4832 *value = ((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5;
4833 }
4834 let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
4835 vector.iter_mut().for_each(|value| *value /= norm);
4836 index.add_vector(
4837 ¢roids,
4838 u32::try_from(posting_index / 2).unwrap(),
4839 u16::try_from(posting_index % 2).unwrap(),
4840 &vector,
4841 &mut scratch,
4842 );
4843 }
4844 let mut bytes = Vec::new();
4845 write_built_ivf_tq(&index, 1, &mut bytes).unwrap();
4846 let disk = AnnDiskIndex::open(
4847 OwnedBytes::new(bytes),
4848 AnnKind::IvfTq,
4849 u32::try_from(count.div_ceil(2)).unwrap(),
4850 )
4851 .unwrap();
4852 let query = vec![1.0f32; dim];
4853 let plan = TqIvfQueryPlan::build(
4854 ¢roids,
4855 &codec,
4856 &query,
4857 1,
4858 crate::dsl::IvfRoutingMode::Flat,
4859 );
4860 let pool = rayon::ThreadPoolBuilder::new()
4861 .num_threads(threads)
4862 .build()
4863 .unwrap();
4864 let serial_pool = rayon::ThreadPoolBuilder::new()
4865 .num_threads(1)
4866 .build()
4867 .unwrap();
4868
4869 let serial = pool
4870 .install(|| disk.search_ivf_tq_distinct_with_tuning(20, &plan, usize::MAX, 512))
4871 .unwrap();
4872 let parallel = pool
4873 .install(|| disk.search_ivf_tq_distinct_with_tuning(20, &plan, 1, 512))
4874 .unwrap();
4875 assert_eq!(parallel, serial);
4876
4877 let mut serial_ms = Vec::with_capacity(iterations);
4878 let mut parallel_ms = Vec::with_capacity(iterations);
4879 for _ in 0..iterations {
4880 let start = std::time::Instant::now();
4881 std::hint::black_box(
4882 pool.install(|| {
4883 disk.search_ivf_tq_distinct_with_tuning(20, &plan, usize::MAX, 512)
4884 })
4885 .unwrap(),
4886 );
4887 serial_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
4888
4889 let start = std::time::Instant::now();
4890 std::hint::black_box(
4891 pool.install(|| disk.search_ivf_tq_distinct_with_tuning(20, &plan, 1, 512))
4892 .unwrap(),
4893 );
4894 parallel_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
4895 }
4896 let serial_p50 = median_ms(&mut serial_ms);
4897 let parallel_p50 = median_ms(&mut parallel_ms);
4898 println!(
4899 "IVF-TQ distinct scan: postings={count} dim={dim} threads={threads} \
4900 serial_p50={serial_p50:.3}ms parallel_p50={parallel_p50:.3}ms speedup={:.2}x",
4901 serial_p50 / parallel_p50,
4902 );
4903
4904 let combiner = crate::query::MultiValueCombiner::Sum;
4905 let serial_combined = serial_pool
4906 .install(|| {
4907 disk.search_ivf_tq_combined_documents_with_tuning(20, &plan, combiner, 1, 512)
4908 })
4909 .unwrap();
4910 let parallel_combined = pool
4911 .install(|| {
4912 disk.search_ivf_tq_combined_documents_with_tuning(20, &plan, combiner, 1, 512)
4913 })
4914 .unwrap();
4915 assert_eq!(parallel_combined, serial_combined);
4916 let mut serial_combined_ms = Vec::with_capacity(iterations);
4917 let mut parallel_combined_ms = Vec::with_capacity(iterations);
4918 for _ in 0..iterations {
4919 let start = std::time::Instant::now();
4920 std::hint::black_box(
4921 serial_pool
4922 .install(|| {
4923 disk.search_ivf_tq_combined_documents_with_tuning(
4924 20, &plan, combiner, 1, 512,
4925 )
4926 })
4927 .unwrap(),
4928 );
4929 serial_combined_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
4930
4931 let start = std::time::Instant::now();
4932 std::hint::black_box(
4933 pool.install(|| {
4934 disk.search_ivf_tq_combined_documents_with_tuning(20, &plan, combiner, 1, 512)
4935 })
4936 .unwrap(),
4937 );
4938 parallel_combined_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
4939 }
4940 let serial_combined_p50 = median_ms(&mut serial_combined_ms);
4941 let parallel_combined_p50 = median_ms(&mut parallel_combined_ms);
4942 println!(
4943 "IVF-TQ combined scan: postings={count} dim={dim} threads={threads} \
4944 serial_p50={serial_combined_p50:.3}ms parallel_p50={parallel_combined_p50:.3}ms \
4945 speedup={:.2}x",
4946 serial_combined_p50 / parallel_combined_p50,
4947 );
4948 }
4949
4950 #[test]
4954 #[ignore]
4955 fn binary_ivf_parallel_scan_benchmark() {
4956 fn env_usize(name: &str, default: usize) -> usize {
4957 std::env::var(name)
4958 .ok()
4959 .map(|value| value.parse::<usize>().unwrap())
4960 .unwrap_or(default)
4961 .max(1)
4962 }
4963
4964 fn median_ms(samples: &mut [f64]) -> f64 {
4965 samples.sort_unstable_by(f64::total_cmp);
4966 samples[samples.len() / 2]
4967 }
4968
4969 let count = env_usize("BINARY_SCAN_BENCH_POSTINGS", 131_072);
4970 let dim_bits = env_usize("BINARY_SCAN_BENCH_DIM_BITS", 2_048);
4971 assert!(dim_bits.is_multiple_of(8));
4972 let code_size = dim_bits / 8;
4973 let iterations = env_usize("BINARY_SCAN_BENCH_ITERS", 30);
4974 let threads = env_usize("BINARY_SCAN_BENCH_THREADS", num_cpus::get().min(8));
4975 let mut state = 0xb1a4_5eed_u64;
4976 let mut codes = vec![0u8; count * code_size];
4977 for code in &mut codes {
4978 state = state
4979 .wrapping_mul(6364136223846793005)
4980 .wrapping_add(1442695040888963407);
4981 *code = (state >> 56) as u8;
4982 }
4983 let query = vec![0xa5u8; code_size];
4984 let unique_docs: Vec<u32> = (0..count)
4985 .map(|index| u32::try_from(index).unwrap())
4986 .collect();
4987 let unique_ordinals = vec![0u16; count];
4988 let multi_docs: Vec<u32> = (0..count)
4989 .map(|index| u32::try_from(index / 2).unwrap())
4990 .collect();
4991 let multi_ordinals: Vec<u16> = (0..count)
4992 .map(|index| u16::try_from(index % 2).unwrap())
4993 .collect();
4994 let header = AnnDiskHeader {
4995 kind: AnnKind::BinaryIvf,
4996 routing: IvfRoutingMode::Flat,
4997 dim: dim_bits,
4998 code_size,
4999 num_clusters: 1,
5000 quantizer_version: 0xb1a4,
5001 codebook_version: 0,
5002 vector_count: count,
5003 };
5004 let mut unique_bytes = Vec::new();
5005 write_built_runs(
5006 header.clone(),
5007 &[BuildRun {
5008 cluster_id: 0,
5009 doc_ids: &unique_docs,
5010 ordinals: &unique_ordinals,
5011 codes: &codes,
5012 }],
5013 &mut unique_bytes,
5014 )
5015 .unwrap();
5016 let unique_disk = AnnDiskIndex::open(
5017 OwnedBytes::new(unique_bytes),
5018 AnnKind::BinaryIvf,
5019 u32::try_from(count).unwrap(),
5020 )
5021 .unwrap();
5022 let mut multi_bytes = Vec::new();
5023 write_built_runs(
5024 header,
5025 &[BuildRun {
5026 cluster_id: 0,
5027 doc_ids: &multi_docs,
5028 ordinals: &multi_ordinals,
5029 codes: &codes,
5030 }],
5031 &mut multi_bytes,
5032 )
5033 .unwrap();
5034 let multi_disk = AnnDiskIndex::open(
5035 OwnedBytes::new(multi_bytes),
5036 AnnKind::BinaryIvf,
5037 u32::try_from(count.div_ceil(2)).unwrap(),
5038 )
5039 .unwrap();
5040 let pool = rayon::ThreadPoolBuilder::new()
5041 .num_threads(threads)
5042 .build()
5043 .unwrap();
5044 let serial_pool = rayon::ThreadPoolBuilder::new()
5045 .num_threads(1)
5046 .build()
5047 .unwrap();
5048
5049 let serial = pool
5050 .install(|| {
5051 unique_disk.search_binary_clusters_with_tuning::<false>(
5052 &query,
5053 20,
5054 &[0],
5055 usize::MAX,
5056 )
5057 })
5058 .unwrap();
5059 let parallel = pool
5060 .install(|| {
5061 unique_disk.search_binary_clusters_with_tuning::<false>(&query, 20, &[0], 1)
5062 })
5063 .unwrap();
5064 assert_eq!(parallel, serial);
5065 let mut serial_ms = Vec::with_capacity(iterations);
5066 let mut parallel_ms = Vec::with_capacity(iterations);
5067 for _ in 0..iterations {
5068 let start = std::time::Instant::now();
5069 std::hint::black_box(
5070 pool.install(|| {
5071 unique_disk.search_binary_clusters_with_tuning::<false>(
5072 &query,
5073 20,
5074 &[0],
5075 usize::MAX,
5076 )
5077 })
5078 .unwrap(),
5079 );
5080 serial_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
5081
5082 let start = std::time::Instant::now();
5083 std::hint::black_box(
5084 pool.install(|| {
5085 unique_disk.search_binary_clusters_with_tuning::<false>(&query, 20, &[0], 1)
5086 })
5087 .unwrap(),
5088 );
5089 parallel_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
5090 }
5091 let serial_p50 = median_ms(&mut serial_ms);
5092 let parallel_p50 = median_ms(&mut parallel_ms);
5093 println!(
5094 "binary-IVF distinct scan: postings={count} dim_bits={dim_bits} threads={threads} \
5095 serial_p50={serial_p50:.3}ms parallel_p50={parallel_p50:.3}ms speedup={:.2}x",
5096 serial_p50 / parallel_p50,
5097 );
5098
5099 let combiner = crate::query::MultiValueCombiner::Sum;
5100 let serial_combined = serial_pool
5101 .install(|| {
5102 multi_disk.search_binary_combined_documents_with_tuning(
5103 20,
5104 &query,
5105 &[0],
5106 combiner,
5107 1,
5108 )
5109 })
5110 .unwrap();
5111 let parallel_combined = pool
5112 .install(|| {
5113 multi_disk.search_binary_combined_documents_with_tuning(
5114 20,
5115 &query,
5116 &[0],
5117 combiner,
5118 1,
5119 )
5120 })
5121 .unwrap();
5122 assert_eq!(parallel_combined, serial_combined);
5123 let mut serial_combined_ms = Vec::with_capacity(iterations);
5124 let mut parallel_combined_ms = Vec::with_capacity(iterations);
5125 for _ in 0..iterations {
5126 let start = std::time::Instant::now();
5127 std::hint::black_box(
5128 serial_pool
5129 .install(|| {
5130 multi_disk.search_binary_combined_documents_with_tuning(
5131 20,
5132 &query,
5133 &[0],
5134 combiner,
5135 1,
5136 )
5137 })
5138 .unwrap(),
5139 );
5140 serial_combined_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
5141
5142 let start = std::time::Instant::now();
5143 std::hint::black_box(
5144 pool.install(|| {
5145 multi_disk.search_binary_combined_documents_with_tuning(
5146 20,
5147 &query,
5148 &[0],
5149 combiner,
5150 1,
5151 )
5152 })
5153 .unwrap(),
5154 );
5155 parallel_combined_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
5156 }
5157 let serial_combined_p50 = median_ms(&mut serial_combined_ms);
5158 let parallel_combined_p50 = median_ms(&mut parallel_combined_ms);
5159 println!(
5160 "binary-IVF combined scan: postings={count} dim_bits={dim_bits} threads={threads} \
5161 serial_p50={serial_combined_p50:.3}ms parallel_p50={parallel_combined_p50:.3}ms \
5162 speedup={:.2}x",
5163 serial_combined_p50 / parallel_combined_p50,
5164 );
5165 }
5166
5167 #[test]
5168 fn open_rejects_tq_payload_with_inconsistent_geometry() {
5169 let (bytes, _) = build_tq_payload(20, 4, 9);
5170 let mut corrupted = bytes.clone();
5172 corrupted[12..16].copy_from_slice(&15u32.to_le_bytes());
5173 assert!(
5174 AnnDiskIndex::open(OwnedBytes::new(corrupted), AnnKind::TqFlat, 4).is_err(),
5175 "TQ header with code_size != padded_dim/2 must be refused"
5176 );
5177
5178 let (short_bytes, _) = build_tq_payload(20, 4, 9);
5180 let mut wrong_kind = short_bytes.clone();
5181 wrong_kind[4] = AnnKind::BinaryIvf as u8;
5182 assert!(
5183 AnnDiskIndex::open(OwnedBytes::new(wrong_kind), AnnKind::BinaryIvf, 4).is_err(),
5184 "TQ block-padded columns must not validate under another kind"
5185 );
5186
5187 let (legacy_bytes, _) = build_tq_payload(20, 4, 9);
5189 let mut legacy_kind = legacy_bytes.clone();
5190 legacy_kind[4] = 1;
5191 let Err(error) = AnnDiskIndex::open(OwnedBytes::new(legacy_kind), AnnKind::TqFlat, 4)
5192 else {
5193 panic!("retired IVF-PQ payloads must not open");
5194 };
5195 assert!(
5196 error.to_string().contains("IVF-PQ"),
5197 "error must name the retired format: {error}"
5198 );
5199 assert!(AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::TqFlat, 4).is_ok());
5200 }
5201
5202 #[test]
5203 fn open_rejects_old_or_out_of_range_ann_payloads() {
5204 let mut legacy = vec![0u8; ANN_HEADER_SIZE + ANN_FOOTER_SIZE];
5205 legacy[..4].copy_from_slice(b"old!");
5206 assert!(AnnDiskIndex::open(OwnedBytes::new(legacy), AnnKind::BinaryIvf, 1).is_err());
5207
5208 let docs = [0u32];
5209 let ordinals = [0u16];
5210 let codes = [0u8];
5211 let runs = [BuildRun {
5212 cluster_id: 0,
5213 doc_ids: &docs,
5214 ordinals: &ordinals,
5215 codes: &codes,
5216 }];
5217 let mut bytes = Vec::new();
5218 write_built_runs(binary_header(1), &runs, &mut bytes).unwrap();
5219 let footer = bytes.len() - ANN_FOOTER_SIZE;
5220 let directory = usize::try_from(u64::from_le_bytes(
5221 bytes[footer..footer + 8].try_into().unwrap(),
5222 ))
5223 .unwrap();
5224 bytes[directory + 12..directory + 16].copy_from_slice(&10u32.to_le_bytes());
5225 assert!(AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::BinaryIvf, 1).is_err());
5226 }
5227
5228 fn scann_binary_artifact(
5229 generation: u64,
5230 ) -> crate::structures::vector::scann::ScannTrainedArtifact {
5231 use crate::structures::vector::scann::{
5232 ScannConfig, ScannEncoding, ScannRoutingLevel, ScannTrainedArtifact,
5233 };
5234 ScannTrainedArtifact::new(
5235 generation,
5236 100_000,
5237 ScannConfig {
5238 dimension: 16,
5239 tree_levels: 1,
5240 num_leaves: 2,
5241 encoding: ScannEncoding::BinaryHamming,
5242 },
5243 vec![ScannRoutingLevel {
5244 centroid_count: 2,
5245 centroid_codes: vec![0, 0, 0xff, 0xff],
5246 minimums: Vec::new(),
5247 steps: Vec::new(),
5248 child_offsets: Vec::new(),
5249 }],
5250 None,
5251 )
5252 .unwrap()
5253 }
5254
5255 #[test]
5256 fn scann_binary_ann_round_trip_validates_exact_global_generation() {
5257 use crate::structures::vector::scann::{
5258 ScannEncoding, ScannLeafRun, ScannSegmentPayload, ScannTrainedArtifactView,
5259 };
5260
5261 let artifact = scann_binary_artifact(41);
5262 let primary = ScannLeafRun::from_rows(
5263 0,
5264 0,
5265 &[0, 1],
5266 &[0, 2],
5267 vec![0x12, 0x34, 0xab, 0xcd],
5268 ScannEncoding::BinaryHamming,
5269 16,
5270 )
5271 .unwrap();
5272 let secondary = ScannLeafRun::from_rows(
5273 1,
5274 0,
5275 &[0],
5276 &[0],
5277 vec![0x12, 0x34],
5278 ScannEncoding::BinaryHamming,
5279 16,
5280 )
5281 .unwrap();
5282 let payload = ScannSegmentPayload::new(&artifact, 2, vec![primary, secondary]).unwrap();
5283 let mut bytes = Vec::new();
5284 write_built_scann(&payload, &mut bytes).unwrap();
5285
5286 let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::ScannBinary, 2).unwrap();
5287 assert_eq!(disk.header.quantizer_version, artifact.generation);
5288 assert_eq!(disk.header.codebook_version, artifact.artifact_id);
5289 assert_eq!(disk.header.vector_count, 3);
5290 let artifact_bytes = artifact.to_bytes().unwrap();
5291 let view = ScannTrainedArtifactView::parse(&artifact_bytes).unwrap();
5292 disk.validate_scann_generation(&view.config, view.generation, view.artifact_id)
5293 .unwrap();
5294 let selective = crate::structures::SoarConfig::new().target_spill_fraction(0.5);
5295 disk.validate_scann_posting_count(2, Some(&selective))
5296 .unwrap();
5297 assert!(
5298 disk.validate_scann_posting_count(2, None)
5299 .unwrap_err()
5300 .to_string()
5301 .contains("physical postings")
5302 );
5303 let too_small_budget = crate::structures::SoarConfig::new().target_spill_fraction(0.49);
5304 assert!(
5305 disk.validate_scann_posting_count(2, Some(&too_small_budget))
5306 .is_err()
5307 );
5308
5309 let serial = disk
5310 .search_binary_clusters_with_tuning::<false>(&[0x12, 0x34], 2, &[0, 1], usize::MAX)
5311 .unwrap();
5312 let parallel = rayon::ThreadPoolBuilder::new()
5313 .num_threads(2)
5314 .build()
5315 .unwrap()
5316 .install(|| {
5317 disk.search_binary_clusters_with_tuning::<false>(&[0x12, 0x34], 2, &[0, 1], 1)
5318 })
5319 .unwrap();
5320 assert_eq!(parallel, serial);
5321 assert_eq!(serial.len(), 2, "secondary posting duplicated a result");
5322 assert_eq!(serial[0], (0, 0, 1.0));
5323
5324 let other = scann_binary_artifact(42);
5325 let other_bytes = other.to_bytes().unwrap();
5326 let other_view = ScannTrainedArtifactView::parse(&other_bytes).unwrap();
5327 let error = disk
5328 .validate_scann_generation(
5329 &other_view.config,
5330 other_view.generation,
5331 other_view.artifact_id,
5332 )
5333 .unwrap_err();
5334 assert!(error.to_string().contains("global trained generation"));
5335 }
5336
5337 #[test]
5338 fn scann_ah_ann_round_trip_preserves_fastscan_tail_geometry() {
5339 use crate::structures::vector::scann::{
5340 ScannAhCodebook, ScannConfig, ScannEncoding, ScannLeafRun, ScannRoutingLevel,
5341 ScannSegmentPayload, ScannTrainedArtifact, ScannTrainedArtifactView,
5342 };
5343
5344 let encoding = ScannEncoding::AsymmetricHash {
5345 dimensions_per_block: 2,
5346 bits_per_code: 4,
5347 };
5348 let artifact = ScannTrainedArtifact::new(
5349 51,
5350 100_000,
5351 ScannConfig {
5352 dimension: 8,
5353 tree_levels: 1,
5354 num_leaves: 2,
5355 encoding,
5356 },
5357 vec![ScannRoutingLevel {
5358 centroid_count: 2,
5359 centroid_codes: vec![0; 16],
5360 minimums: vec![0.0; 8],
5361 steps: vec![1.0; 8],
5362 child_offsets: Vec::new(),
5363 }],
5364 Some(ScannAhCodebook {
5365 dimensions_per_block: 2,
5366 centers_per_block: 16,
5367 centers: vec![0.0; 4 * 16 * 2],
5368 }),
5369 )
5370 .unwrap();
5371 let docs: Vec<u32> = (0..33).collect();
5372 let ordinals = vec![0u16; docs.len()];
5373 let codes = vec![0x5a; encoding.leaf_code_bytes(8, docs.len()).unwrap()];
5374 let run = ScannLeafRun::from_rows(0, 0, &docs, &ordinals, codes, encoding, 8).unwrap();
5375 let payload = ScannSegmentPayload::new(&artifact, 33, vec![run]).unwrap();
5376 let mut bytes = Vec::new();
5377 write_built_scann(&payload, &mut bytes).unwrap();
5378
5379 let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::ScannAh, 33).unwrap();
5380 assert_eq!(
5381 disk.runs[0].codes.len(),
5382 encoding.leaf_code_bytes(8, 33).unwrap()
5383 );
5384 let artifact_bytes = artifact.to_bytes().unwrap();
5385 let view = ScannTrainedArtifactView::parse(&artifact_bytes).unwrap();
5386 disk.validate_scann_generation(&view.config, view.generation, view.artifact_id)
5387 .unwrap();
5388 }
5389
5390 #[test]
5391 fn scann_binary_merge_copies_codes_and_rebases_without_retraining() {
5392 use crate::structures::vector::scann::{ScannEncoding, ScannLeafRun, ScannSegmentPayload};
5393
5394 let artifact = scann_binary_artifact(61);
5395 let make = |codes: Vec<u8>| {
5396 let secondary_code = codes[..2].to_vec();
5397 let primary = ScannLeafRun::from_rows(
5398 0,
5399 0,
5400 &[0, 1],
5401 &[0, 0],
5402 codes,
5403 ScannEncoding::BinaryHamming,
5404 16,
5405 )
5406 .unwrap();
5407 let secondary = ScannLeafRun::from_rows(
5408 1,
5409 0,
5410 &[0],
5411 &[0],
5412 secondary_code,
5413 ScannEncoding::BinaryHamming,
5414 16,
5415 )
5416 .unwrap();
5417 let payload = ScannSegmentPayload::new(&artifact, 2, vec![primary, secondary]).unwrap();
5418 let mut bytes = Vec::new();
5419 write_built_scann(&payload, &mut bytes).unwrap();
5420 AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::ScannBinary, 2).unwrap()
5421 };
5422 let left = make(vec![1, 2, 3, 4]);
5423 let right = make(vec![5, 6, 7, 8]);
5424 let left_codes = left.raw.as_slice()[left.runs[0].codes.clone()].to_vec();
5425 let right_codes = right.raw.as_slice()[right.runs[0].codes.clone()].to_vec();
5426
5427 let mut merged_bytes = Vec::new();
5428 write_merged_ann(&[(&left, 0), (&right, 2)], &mut merged_bytes).unwrap();
5429 let merged =
5430 AnnDiskIndex::open(OwnedBytes::new(merged_bytes), AnnKind::ScannBinary, 4).unwrap();
5431 assert_eq!(merged.header.vector_count, 6);
5432 merged
5433 .validate_scann_posting_count(
5434 4,
5435 Some(&crate::structures::SoarConfig::new().target_spill_fraction(0.5)),
5436 )
5437 .unwrap();
5438 assert_eq!(merged.runs[0].doc_base, 0);
5439 assert_eq!(merged.runs[1].doc_base, 2);
5440 assert_eq!(
5441 &merged.raw.as_slice()[merged.runs[0].codes.clone()],
5442 left_codes
5443 );
5444 assert_eq!(
5445 &merged.raw.as_slice()[merged.runs[1].codes.clone()],
5446 right_codes
5447 );
5448 assert_eq!(
5449 merged
5450 .search_binary_clusters::<false>(&[0, 0], 4, &[0, 1])
5451 .unwrap()
5452 .len(),
5453 4,
5454 "merge exposed primary and secondary postings as separate results",
5455 );
5456
5457 let other_artifact = scann_binary_artifact(62);
5458 let other_run = ScannLeafRun::from_rows(
5459 0,
5460 0,
5461 &[0],
5462 &[0],
5463 vec![9, 10],
5464 ScannEncoding::BinaryHamming,
5465 16,
5466 )
5467 .unwrap();
5468 let other_payload = ScannSegmentPayload::new(&other_artifact, 1, vec![other_run]).unwrap();
5469 let mut other_bytes = Vec::new();
5470 write_built_scann(&other_payload, &mut other_bytes).unwrap();
5471 let other =
5472 AnnDiskIndex::open(OwnedBytes::new(other_bytes), AnnKind::ScannBinary, 1).unwrap();
5473 let error = write_merged_ann(&[(&left, 0), (&other, 2)], &mut Vec::new()).unwrap_err();
5474 assert!(error.to_string().contains("incompatible generations"));
5475 }
5476}