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
25#[cfg(feature = "native")]
26use crate::structures::BinaryIvfIndex;
27use crate::structures::vector::index::{BoundedAnnCollector, BoundedUniqueAnnCollector};
28
29type CombinedBinaryCandidates = (Vec<AnnDocumentCandidate>, Vec<(u32, u16, f32)>);
32
33const ANN_HEADER_MAGIC: u32 = 0x3152_4e41; const ANN_FOOTER_MAGIC: u32 = 0x3146_4e41; const ANN_DISK_VERSION: u16 = 1;
36const ANN_HEADER_SIZE: usize = 56;
37const ANN_RUN_SIZE: usize = 48;
38const ANN_FOOTER_SIZE: usize = 24;
39#[cfg(feature = "native")]
40const COPY_CHUNK: usize = 8 * 1024 * 1024;
41#[cfg(feature = "native")]
42const PREFETCH_COALESCE_GAP: usize = 4 * 1024;
43const BINARY_SCORE_BATCH: usize = 8_192;
44const TQ_PRUNE_ESTIMATE_BOUND: f32 = 1.3;
49#[cfg(feature = "native")]
55const TQ_PARALLEL_SCAN_MIN_VECTORS: usize = 65_536;
56#[cfg(feature = "native")]
57const TQ_PARALLEL_SCAN_CHUNK_BLOCKS: usize = 512;
58const IVF_PARALLEL_SCAN_MIN_POSTINGS: usize = 65_536;
62const IVF_TQ_PARALLEL_SCAN_CHUNK_BLOCKS: usize = 512;
63
64#[derive(Debug, Clone, Copy, PartialEq, Eq)]
65pub(crate) enum AnnKind {
66 BinaryIvf = 2,
69 TqFlat = 3,
71 IvfTq = 4,
73 ScannAh = 5,
75 ScannBinary = 6,
77}
78
79impl AnnKind {
80 fn from_u8(value: u8) -> io::Result<Self> {
81 match value {
82 1 => Err(invalid_data(
83 "ANN kind 1 (IVF-PQ) is no longer supported; recreate the index \
84 with `ivf_tq` and reindex",
85 )),
86 2 => Ok(Self::BinaryIvf),
87 3 => Ok(Self::TqFlat),
88 4 => Ok(Self::IvfTq),
89 5 => Ok(Self::ScannAh),
90 6 => Ok(Self::ScannBinary),
91 _ => Err(invalid_data(format!("unknown ANN kind {value}"))),
92 }
93 }
94}
95
96fn expected_codes_column_len(
100 kind: AnnKind,
101 count: usize,
102 dim: usize,
103 code_size: usize,
104) -> io::Result<usize> {
105 match kind {
106 AnnKind::BinaryIvf | AnnKind::ScannBinary => count
107 .checked_mul(code_size)
108 .ok_or_else(|| invalid_data("ANN code column size overflows usize")),
109 AnnKind::TqFlat => {
111 crate::structures::vector::quantization::tq_codes_column_len_checked(count, code_size)
112 .ok_or_else(|| invalid_data("TQ code column size overflows usize"))
113 }
114 AnnKind::IvfTq => crate::structures::vector::quantization::tq_ivf_codes_column_len_checked(
115 count, code_size,
116 )
117 .ok_or_else(|| invalid_data("IVF-TQ code column size overflows usize")),
118 AnnKind::ScannAh => crate::structures::vector::scann::ScannEncoding::AsymmetricHash {
123 dimensions_per_block: u16::try_from(code_size)
124 .map_err(|_| invalid_data("ScaNN AH block dimension exceeds u16"))?,
125 bits_per_code: 4,
126 }
127 .leaf_code_bytes(
128 u32::try_from(dim).map_err(|_| invalid_data("ScaNN dimension exceeds u32"))?,
129 count,
130 )
131 .map_err(|error| invalid_data(error.to_string())),
132 }
133}
134
135#[derive(Debug, Clone, PartialEq, Eq)]
136pub(crate) struct AnnDiskHeader {
137 pub kind: AnnKind,
138 pub routing: IvfRoutingMode,
139 pub dim: usize,
140 pub code_size: usize,
141 pub num_clusters: u32,
142 pub quantizer_version: u64,
143 pub codebook_version: u64,
144 pub vector_count: usize,
145}
146
147#[derive(Debug)]
148struct AnnRun {
149 cluster_id: u32,
150 doc_base: u32,
151 max_doc_id: u32,
152 count: usize,
153 doc_ids: Range<usize>,
154 ordinals: Range<usize>,
155 codes: Range<usize>,
156}
157
158#[derive(Clone, Copy)]
159struct IvfTqScanTask<'a> {
160 run: &'a AnnRun,
161 cluster_dot: f32,
162 first_block: usize,
163 last_block: usize,
164}
165
166#[derive(Clone, Copy)]
167struct BinaryScanTask<'a> {
168 run: &'a AnnRun,
169 first_index: usize,
170 count: usize,
171}
172
173pub(crate) struct AnnDiskIndex {
176 #[cfg(feature = "native")]
178 heap_pins: crate::segment::pin::HeapPinSet,
179 raw: OwnedBytes,
180 header: AnnDiskHeader,
181 runs: Vec<AnnRun>,
182}
183
184#[derive(Debug, Clone, Copy, PartialEq)]
192pub struct AnnHealth {
193 pub vectors: u64,
195 pub clusters_nonempty: u32,
197 pub clusters_total: u32,
199 pub runs: u32,
201 pub largest_cluster: u32,
203 pub largest_cluster_vectors: u64,
204 pub imbalance: f64,
209 pub payload_bytes: u64,
211}
212
213impl AnnHealth {
214 pub fn fragmentation(&self) -> f64 {
218 if self.clusters_nonempty == 0 {
219 return 0.0;
220 }
221 f64::from(self.runs) / f64::from(self.clusters_nonempty)
222 }
223
224 pub fn largest_cluster_share(&self) -> f64 {
226 if self.vectors == 0 {
227 return 0.0;
228 }
229 self.largest_cluster_vectors as f64 / self.vectors as f64
230 }
231}
232
233const ANN_SKEW_WARN_SHARE: f64 = 0.05;
237const ANN_SKEW_WARN_MIN_VECTORS: u64 = 100_000;
238const ANN_FRAGMENTATION_WARN: f64 = 8.0;
241
242#[derive(Debug, Clone, Copy, PartialEq)]
248pub(crate) struct AnnDocumentCandidate {
249 pub(crate) doc_id: u32,
250 pub(crate) score: f32,
251}
252
253impl AnnDiskIndex {
254 pub(crate) fn open(
255 raw: OwnedBytes,
256 expected_kind: AnnKind,
257 total_docs: u32,
258 ) -> io::Result<Self> {
259 if raw.len() < ANN_HEADER_SIZE + ANN_FOOTER_SIZE {
260 return Err(invalid_data("ANN payload is shorter than header + footer"));
261 }
262 let bytes = raw.as_slice();
263 let mut header_cursor = std::io::Cursor::new(&bytes[..ANN_HEADER_SIZE]);
264 if header_cursor.read_u32::<LittleEndian>()? != ANN_HEADER_MAGIC {
265 return Err(invalid_data("ANN payload has unsupported header magic"));
266 }
267 let kind = AnnKind::from_u8(header_cursor.read_u8()?)?;
268 if kind != expected_kind {
269 return Err(invalid_data(format!(
270 "ANN payload kind {kind:?} does not match expected {expected_kind:?}"
271 )));
272 }
273 let routing = routing_from_u8(header_cursor.read_u8()?)?;
274 if header_cursor.read_u16::<LittleEndian>()? != ANN_DISK_VERSION {
275 return Err(invalid_data("ANN payload has unsupported format version"));
276 }
277 let dim = header_cursor.read_u32::<LittleEndian>()? as usize;
278 let code_size = header_cursor.read_u32::<LittleEndian>()? as usize;
279 let num_clusters = header_cursor.read_u32::<LittleEndian>()?;
280 if header_cursor.read_u32::<LittleEndian>()? != 0 {
281 return Err(invalid_data("ANN header reserved field is non-zero"));
282 }
283 let quantizer_version = header_cursor.read_u64::<LittleEndian>()?;
284 let codebook_version = header_cursor.read_u64::<LittleEndian>()?;
285 let vector_count = usize::try_from(header_cursor.read_u64::<LittleEndian>()?)
286 .map_err(|_| invalid_data("ANN vector count exceeds usize"))?;
287 if header_cursor.read_u64::<LittleEndian>()? != 0 {
288 return Err(invalid_data("ANN header tail is non-zero"));
289 }
290 let header = AnnDiskHeader {
291 kind,
292 routing,
293 dim,
294 code_size,
295 num_clusters,
296 quantizer_version,
297 codebook_version,
298 vector_count,
299 };
300 validate_header(&header)?;
301
302 let footer_start = bytes.len() - ANN_FOOTER_SIZE;
303 let mut footer_cursor = std::io::Cursor::new(&bytes[footer_start..]);
304 let directory_offset = usize::try_from(footer_cursor.read_u64::<LittleEndian>()?)
305 .map_err(|_| invalid_data("ANN directory offset exceeds usize"))?;
306 let num_runs = usize::try_from(footer_cursor.read_u64::<LittleEndian>()?)
307 .map_err(|_| invalid_data("ANN run count exceeds usize"))?;
308 if footer_cursor.read_u32::<LittleEndian>()? != ANN_FOOTER_MAGIC
309 || footer_cursor.read_u32::<LittleEndian>()? != u32::from(ANN_DISK_VERSION)
310 {
311 return Err(invalid_data("ANN payload has unsupported footer"));
312 }
313 if num_runs == 0 {
314 return Err(invalid_data("ANN payload has no cluster runs"));
315 }
316 let directory_len = num_runs
317 .checked_mul(ANN_RUN_SIZE)
318 .ok_or_else(|| invalid_data("ANN directory size overflows usize"))?;
319 if directory_offset < ANN_HEADER_SIZE
320 || directory_offset.checked_add(directory_len) != Some(footer_start)
321 {
322 return Err(invalid_data("ANN directory does not end at the footer"));
323 }
324
325 let mut runs = Vec::with_capacity(num_runs);
326 let mut directory_cursor = std::io::Cursor::new(&bytes[directory_offset..footer_start]);
327 let mut previous_cluster = None;
328 let mut counted_vectors = 0usize;
329 for _ in 0..num_runs {
330 let cluster_id = directory_cursor.read_u32::<LittleEndian>()?;
331 let doc_base = directory_cursor.read_u32::<LittleEndian>()?;
332 let count = directory_cursor.read_u32::<LittleEndian>()? as usize;
333 let max_doc_id = directory_cursor.read_u32::<LittleEndian>()?;
334 let doc_ids_offset = usize::try_from(directory_cursor.read_u64::<LittleEndian>()?)
335 .map_err(|_| invalid_data("ANN doc-ID offset exceeds usize"))?;
336 let ordinals_offset = usize::try_from(directory_cursor.read_u64::<LittleEndian>()?)
337 .map_err(|_| invalid_data("ANN ordinal offset exceeds usize"))?;
338 let codes_offset = usize::try_from(directory_cursor.read_u64::<LittleEndian>()?)
339 .map_err(|_| invalid_data("ANN code offset exceeds usize"))?;
340 let codes_len = usize::try_from(directory_cursor.read_u64::<LittleEndian>()?)
341 .map_err(|_| invalid_data("ANN code length exceeds usize"))?;
342 if count == 0
343 || cluster_id >= num_clusters
344 || previous_cluster.is_some_and(|previous| previous > cluster_id)
345 || doc_base
346 .checked_add(max_doc_id)
347 .is_none_or(|doc_id| doc_id >= total_docs)
348 {
349 return Err(invalid_data("ANN run metadata is invalid"));
350 }
351 previous_cluster = Some(cluster_id);
352 let doc_ids_len = count
353 .checked_mul(std::mem::size_of::<u32>())
354 .ok_or_else(|| invalid_data("ANN doc-ID column size overflows usize"))?;
355 let ordinals_len = count
356 .checked_mul(std::mem::size_of::<u16>())
357 .ok_or_else(|| invalid_data("ANN ordinal column size overflows usize"))?;
358 let expected_codes_len = expected_codes_column_len(kind, count, dim, code_size)?;
359 let doc_ids_end = doc_ids_offset
360 .checked_add(doc_ids_len)
361 .ok_or_else(|| invalid_data("ANN doc-ID range overflows usize"))?;
362 let ordinals_end = ordinals_offset
363 .checked_add(ordinals_len)
364 .ok_or_else(|| invalid_data("ANN ordinal range overflows usize"))?;
365 let codes_end = codes_offset
366 .checked_add(codes_len)
367 .ok_or_else(|| invalid_data("ANN code range overflows usize"))?;
368 if doc_ids_offset < ANN_HEADER_SIZE
369 || ordinals_offset != doc_ids_end
370 || codes_offset != ordinals_end
371 || codes_len != expected_codes_len
372 || codes_end > directory_offset
373 {
374 return Err(invalid_data("ANN run columns are not contiguous/in bounds"));
375 }
376 runs.push(AnnRun {
377 cluster_id,
378 doc_base,
379 max_doc_id,
380 count,
381 doc_ids: doc_ids_offset..doc_ids_end,
382 ordinals: ordinals_offset..ordinals_end,
383 codes: codes_offset..codes_end,
384 });
385 counted_vectors = counted_vectors
386 .checked_add(count)
387 .ok_or_else(|| invalid_data("ANN run vector count overflows usize"))?;
388 }
389 let mut payload_order: Vec<usize> = (0..runs.len()).collect();
390 payload_order.sort_unstable_by_key(|&index| runs[index].doc_ids.start);
391 let mut expected_payload_offset = ANN_HEADER_SIZE;
392 for index in payload_order {
393 let run = &runs[index];
394 if run.doc_ids.start != expected_payload_offset {
395 return Err(invalid_data("ANN payload runs overlap or contain gaps"));
396 }
397 expected_payload_offset = run.codes.end;
398 }
399 if expected_payload_offset != directory_offset || counted_vectors != vector_count {
400 return Err(invalid_data(
401 "ANN payload coverage/vector count is inconsistent",
402 ));
403 }
404
405 #[cfg(feature = "native")]
412 raw.madvise_range(
413 ANN_HEADER_SIZE..directory_offset,
414 if kind == AnnKind::TqFlat {
415 libc::MADV_SEQUENTIAL
416 } else {
417 libc::MADV_RANDOM
418 },
419 );
420
421 Ok(Self {
422 #[cfg(feature = "native")]
423 heap_pins: Default::default(),
424 raw,
425 header,
426 runs,
427 })
428 }
429
430 pub(crate) fn health(&self) -> AnnHealth {
433 let mut vectors = 0u64;
434 let mut clusters_nonempty = 0u32;
435 let mut payload_bytes = 0u64;
436 let mut largest = (0u32, 0u64);
437 let mut sum_squares = 0f64;
438 let mut index = 0usize;
440 while index < self.runs.len() {
441 let cluster_id = self.runs[index].cluster_id;
442 let mut cluster_vectors = 0u64;
443 while index < self.runs.len() && self.runs[index].cluster_id == cluster_id {
444 let run = &self.runs[index];
445 cluster_vectors += run.count as u64;
446 payload_bytes += (run.codes.end - run.codes.start) as u64;
447 index += 1;
448 }
449 vectors += cluster_vectors;
450 clusters_nonempty += 1;
451 sum_squares += (cluster_vectors as f64) * (cluster_vectors as f64);
452 if cluster_vectors > largest.1 {
453 largest = (cluster_id, cluster_vectors);
454 }
455 }
456 let imbalance = if vectors == 0 || clusters_nonempty == 0 {
457 0.0
458 } else {
459 f64::from(clusters_nonempty) * sum_squares / ((vectors as f64) * (vectors as f64))
460 };
461 AnnHealth {
462 vectors,
463 clusters_nonempty,
464 clusters_total: self.header.num_clusters,
465 runs: self.runs.len() as u32,
466 largest_cluster: largest.0,
467 largest_cluster_vectors: largest.1,
468 imbalance,
469 payload_bytes,
470 }
471 }
472
473 pub(crate) fn report_health(&self, index_label: &str, field_id: u32, segment_id: u128) {
478 let health = self.health();
479 let share = health.largest_cluster_share();
480 let fragmentation = health.fragmentation();
481 log::info!(
482 "[ann_health] index={index_label} field={field_id} segment={segment_id:016x}: \
483 vectors={} clusters={}/{} runs={} fragmentation={fragmentation:.2} \
484 imbalance={:.2} largest_leaf={:.2}% payload={}",
485 health.vectors,
486 health.clusters_nonempty,
487 health.clusters_total,
488 health.runs,
489 health.imbalance,
490 100.0 * share,
491 crate::format_bytes(health.payload_bytes),
492 );
493 crate::observe::ann_health(
494 index_label,
495 field_id,
496 health.imbalance,
497 fragmentation,
498 share,
499 );
500 if share >= ANN_SKEW_WARN_SHARE && health.vectors >= ANN_SKEW_WARN_MIN_VECTORS {
501 log::warn!(
502 "[ann_health] index={index_label} field={field_id} segment={segment_id:016x}: \
503 leaf {} holds {:.1}% of {} vectors — every query probing it scans that leaf \
504 in full; degenerate embeddings collapse into one leaf exactly like this",
505 health.largest_cluster,
506 100.0 * share,
507 health.vectors,
508 );
509 }
510 if fragmentation >= ANN_FRAGMENTATION_WARN {
511 log::warn!(
512 "[ann_health] index={index_label} field={field_id} segment={segment_id:016x}: \
513 {fragmentation:.1} extents per probed cluster ({} runs / {} clusters) — \
514 cold probes pay that many seeks; the next merge or vector-generation rewrite \
515 compacts to 1.0",
516 health.runs,
517 health.clusters_nonempty,
518 );
519 }
520 }
521
522 pub(crate) fn header(&self) -> &AnnDiskHeader {
523 &self.header
524 }
525
526 pub(crate) fn validate_scann_generation(
529 &self,
530 config: &crate::structures::vector::scann::ScannConfig,
531 generation: u64,
532 artifact_id: u64,
533 ) -> io::Result<()> {
534 use crate::structures::vector::scann::ScannEncoding;
535
536 let expected_kind = match config.encoding {
537 ScannEncoding::AsymmetricHash { .. } => AnnKind::ScannAh,
538 ScannEncoding::BinaryHamming => AnnKind::ScannBinary,
539 };
540 if self.header.kind != expected_kind
541 || self.header.routing != IvfRoutingMode::Flat
542 || self.header.dim != config.dimension as usize
543 || self.header.num_clusters != config.num_leaves
544 || self.header.quantizer_version != generation
545 || self.header.codebook_version != artifact_id
546 {
547 return Err(invalid_data(
548 "ScaNN ANN payload does not match the global trained generation",
549 ));
550 }
551 match config.encoding {
552 ScannEncoding::AsymmetricHash {
553 dimensions_per_block,
554 bits_per_code: 4,
555 } if self.header.code_size == usize::from(dimensions_per_block) => Ok(()),
556 ScannEncoding::BinaryHamming
557 if self.header.code_size == config.dimension as usize / 8 =>
558 {
559 Ok(())
560 }
561 _ => Err(invalid_data(
562 "ScaNN ANN payload encoding does not match the global trained artifact",
563 )),
564 }
565 }
566
567 #[cfg(feature = "native")]
572 pub(crate) fn validate_scann_posting_count(
573 &self,
574 logical_vectors: usize,
575 soar: Option<&crate::structures::SoarConfig>,
576 ) -> io::Result<()> {
577 let spill_budget = match self.header.kind {
578 AnnKind::ScannAh => 0,
579 AnnKind::ScannBinary
580 if self.header.num_clusters > 1
581 && soar.is_some_and(|config| config.num_secondary > 0) =>
582 {
583 match soar.and_then(crate::structures::SoarConfig::calibration_target) {
584 Some(target_fraction) => {
585 (logical_vectors as f64 * f64::from(target_fraction)).floor() as usize
586 }
587 None => logical_vectors,
588 }
589 }
590 AnnKind::ScannBinary => 0,
591 _ => {
592 return Err(invalid_data(
593 "posting-count validation requires a ScaNN ANN payload",
594 ));
595 }
596 };
597 let maximum = logical_vectors
598 .checked_add(spill_budget)
599 .ok_or_else(|| invalid_data("ScaNN posting-count bound overflows usize"))?;
600 if self.header.vector_count < logical_vectors || self.header.vector_count > maximum {
601 return Err(invalid_data(format!(
602 "ScaNN ANN payload has {} physical postings for {logical_vectors} logical vectors; expected {logical_vectors}..={maximum}",
603 self.header.vector_count,
604 )));
605 }
606 Ok(())
607 }
608
609 pub(crate) fn estimated_heap_bytes(&self) -> usize {
610 std::mem::size_of::<Self>() + self.runs.capacity() * std::mem::size_of::<AnnRun>()
611 }
612
613 #[cfg(feature = "native")]
614 pub(crate) fn pin_lookup_directory(
615 &mut self,
616 mode: crate::segment::pin::PinMode,
617 remaining: &mut u64,
618 report: &mut crate::segment::pin::PinReport,
619 ) {
620 let before = self.heap_pins.report();
621 self.heap_pins
622 .pin_slice(&self.runs, "ANN cluster-run directory", mode, remaining);
623 let after = self.heap_pins.report();
624 report.intended_bytes += after.intended_bytes - before.intended_bytes;
625 report.pinned_bytes += after.pinned_bytes - before.pinned_bytes;
626 report.skipped_budget_bytes += after.skipped_budget_bytes - before.skipped_budget_bytes;
627 report.failed_bytes += after.failed_bytes - before.failed_bytes;
628 report.heap_copy_bytes += after.heap_copy_bytes - before.heap_copy_bytes;
629 }
630
631 fn cluster_runs(&self, cluster_id: u32) -> &[AnnRun] {
632 let start = self.runs.partition_point(|run| run.cluster_id < cluster_id);
633 let end = self
634 .runs
635 .partition_point(|run| run.cluster_id <= cluster_id);
636 &self.runs[start..end]
637 }
638
639 #[cfg(feature = "native")]
640 fn ivf_tq_scan_tasks<'a>(
641 &'a self,
642 plan: &'a crate::structures::TqIvfQueryPlan,
643 block_bytes: usize,
644 chunk_blocks: usize,
645 ) -> Vec<IvfTqScanTask<'a>> {
646 debug_assert!(chunk_blocks > 0);
647 let mut tasks = Vec::new();
648 for (cluster_id, cluster_dot) in plan.cluster_dots() {
649 for run in self.cluster_runs(cluster_id) {
650 let block_count = run.codes.len() / block_bytes;
651 for first_block in (0..block_count).step_by(chunk_blocks) {
652 tasks.push(IvfTqScanTask {
653 run,
654 cluster_dot,
655 first_block,
656 last_block: (first_block + chunk_blocks).min(block_count),
657 });
658 }
659 }
660 }
661 tasks
662 }
663
664 #[cfg(feature = "native")]
665 fn binary_scan_tasks<'a>(&'a self, cluster_ids: &[u32]) -> Vec<BinaryScanTask<'a>> {
666 let mut tasks = Vec::new();
667 for &cluster_id in cluster_ids {
668 for run in self.cluster_runs(cluster_id) {
669 for first_index in (0..run.count).step_by(BINARY_SCORE_BATCH) {
670 tasks.push(BinaryScanTask {
671 run,
672 first_index,
673 count: BINARY_SCORE_BATCH.min(run.count - first_index),
674 });
675 }
676 }
677 }
678 tasks
679 }
680
681 #[cfg(feature = "native")]
688 fn prefetch_cluster_runs(&self, cluster_ids: &[u32]) {
689 if cluster_ids.is_empty() || !self.raw.is_mmap() {
690 return;
691 }
692 let mut ranges = Vec::with_capacity(cluster_ids.len());
693 for &cluster_id in cluster_ids {
694 ranges.extend(
695 self.cluster_runs(cluster_id)
696 .iter()
697 .map(|run| run.doc_ids.start..run.codes.end),
698 );
699 }
700 coalesce_prefetch_ranges(&mut ranges);
701 for range in ranges {
702 self.raw.madvise_range(range, libc::MADV_WILLNEED);
703 }
704 }
705
706 pub(crate) fn search_tq_combined_documents(
711 &self,
712 k: usize,
713 plan: &crate::structures::TqQueryPlan,
714 combiner: crate::query::MultiValueCombiner,
715 ) -> io::Result<Vec<AnnDocumentCandidate>> {
716 use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_block_bytes};
717
718 combiner
719 .validate()
720 .map_err(|message| io::Error::new(io::ErrorKind::InvalidInput, message))?;
721 if plan.padded_dim() != self.header.code_size * 2 {
722 return Err(io::Error::new(
723 io::ErrorKind::InvalidInput,
724 "TQ query plan does not match the payload dimension",
725 ));
726 }
727 if k == 0 {
728 return Ok(Vec::new());
729 }
730
731 let block_bytes = tq_block_bytes(self.header.code_size);
732 let bytes = self.raw.as_slice();
733 let mut top_documents = BoundedAnnCollector::<true, true>::new(k);
734 let mut ordinal_scores = Vec::new();
735 let mut scores = [0.0f32; TQ_BLOCK_LANES];
736
737 for run in &self.runs {
738 let mut current_doc = None;
739 let codes = &bytes[run.codes.clone()];
740 for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
741 crate::structures::vector::quantization::tq_score_block(plan, block, &mut scores);
742 let lane_base = block_index * TQ_BLOCK_LANES;
743 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
744 for (lane, &score) in scores.iter().enumerate().take(lanes) {
745 let index = lane_base + lane;
746 let doc_id = run_doc_id(bytes, run, index)?;
747 if current_doc.is_some_and(|previous| doc_id < previous) {
748 return Err(invalid_data("flat TQ run is not grouped by document ID"));
749 }
750 if current_doc.is_some_and(|previous| doc_id != previous) {
751 retain_combined_document(
752 &mut top_documents,
753 current_doc.expect("current document is present"),
754 &ordinal_scores,
755 combiner,
756 );
757 ordinal_scores.clear();
758 }
759 current_doc = Some(doc_id);
760 if score.is_finite() {
761 ordinal_scores.push((
762 u32::from(read_u16(bytes, run.ordinals.start + index * 2)),
763 score,
764 ));
765 }
766 }
767 }
768 if let Some(doc_id) = current_doc {
769 retain_combined_document(&mut top_documents, doc_id, &ordinal_scores, combiner);
770 ordinal_scores.clear();
771 }
772 }
773
774 Ok(top_documents
775 .into_sorted_results()
776 .into_iter()
777 .map(|(doc_id, _, score)| AnnDocumentCandidate { doc_id, score })
778 .collect())
779 }
780
781 pub(crate) fn search_ivf_tq_combined_documents(
789 &self,
790 k: usize,
791 plan: &crate::structures::TqIvfQueryPlan,
792 combiner: crate::query::MultiValueCombiner,
793 ) -> io::Result<Vec<AnnDocumentCandidate>> {
794 self.search_ivf_tq_combined_documents_with_tuning(
795 k,
796 plan,
797 combiner,
798 IVF_PARALLEL_SCAN_MIN_POSTINGS,
799 IVF_TQ_PARALLEL_SCAN_CHUNK_BLOCKS,
800 )
801 }
802
803 fn search_ivf_tq_combined_documents_with_tuning(
804 &self,
805 k: usize,
806 plan: &crate::structures::TqIvfQueryPlan,
807 combiner: crate::query::MultiValueCombiner,
808 parallel_min_postings: usize,
809 parallel_chunk_blocks: usize,
810 ) -> io::Result<Vec<AnnDocumentCandidate>> {
811 use crate::structures::vector::quantization::{
812 TQ_BLOCK_LANES, tq_ivf_block_bytes, tq_score_ivf_block,
813 };
814
815 validate_combined_search(combiner)?;
816 let tq_plan = plan.tq_plan();
817 self.validate_ivf_tq_query_plan(plan)?;
818 #[cfg(not(feature = "native"))]
819 let _ = (parallel_min_postings, parallel_chunk_blocks);
820 if k == 0 {
821 return Ok(Vec::new());
822 }
823 #[cfg(feature = "native")]
824 self.prefetch_cluster_runs(&plan.cluster_ids);
825
826 let block_bytes = tq_ivf_block_bytes(self.header.code_size);
827 let bytes = self.raw.as_slice();
828 let posting_count = probed_posting_count(self, &plan.cluster_ids)?;
829
830 #[cfg(feature = "native")]
831 if rayon::current_num_threads() > 1 && posting_count >= parallel_min_postings {
832 use rayon::prelude::*;
833 let tasks = self.ivf_tq_scan_tasks(plan, block_bytes, parallel_chunk_blocks);
834 let ordinal_scores = tasks
835 .par_iter()
836 .try_fold(
837 || Vec::with_capacity(parallel_chunk_blocks * TQ_BLOCK_LANES),
838 |mut ordinal_scores, task| {
839 score_ivf_tq_combined_blocks(
840 bytes,
841 tq_plan,
842 block_bytes,
843 *task,
844 &mut ordinal_scores,
845 )?;
846 Ok::<_, io::Error>(ordinal_scores)
847 },
848 )
849 .try_reduce(Vec::new, |mut left, mut right| {
850 left.append(&mut right);
851 Ok(left)
852 })?;
853 return Ok(combine_scored_ordinals(ordinal_scores, k, combiner));
854 }
855
856 let mut ordinal_scores = Vec::new();
857 ordinal_scores
858 .try_reserve_exact(posting_count)
859 .map_err(|_| invalid_data("IVF-TQ combined score buffer allocation failed"))?;
860 let mut scores = [0.0f32; TQ_BLOCK_LANES];
861 for (cluster_id, cluster_dot) in plan.cluster_dots() {
862 for run in self.cluster_runs(cluster_id) {
863 let codes = &bytes[run.codes.clone()];
864 for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
865 tq_score_ivf_block(tq_plan, block, cluster_dot, &mut scores);
866 let lane_base = block_index * TQ_BLOCK_LANES;
867 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
868 for (lane, &score) in scores.iter().enumerate().take(lanes) {
869 if !score.is_finite() {
870 continue;
871 }
872 let index = lane_base + lane;
873 ordinal_scores.push((
874 run_doc_id(bytes, run, index)?,
875 read_u16(bytes, run.ordinals.start + index * 2),
876 score,
877 ));
878 }
879 }
880 }
881 }
882 Ok(combine_scored_ordinals(ordinal_scores, k, combiner))
883 }
884
885 pub(crate) fn search_binary_combined_documents(
894 &self,
895 k: usize,
896 query: &[u8],
897 cluster_ids: &[u32],
898 combiner: crate::query::MultiValueCombiner,
899 ) -> io::Result<CombinedBinaryCandidates> {
900 self.search_binary_combined_documents_with_tuning(
901 k,
902 query,
903 cluster_ids,
904 combiner,
905 IVF_PARALLEL_SCAN_MIN_POSTINGS,
906 )
907 }
908
909 pub(crate) fn search_scann_ah_combined_documents(
912 &self,
913 k: usize,
914 query: &crate::structures::vector::scann::FloatScannQuery,
915 combiner: crate::query::MultiValueCombiner,
916 ) -> io::Result<Vec<AnnDocumentCandidate>> {
917 use crate::structures::vector::scann::{FAST_SCAN_LANES, FastScanQuery};
918
919 validate_combined_search(combiner)?;
920 if self.header.kind != AnnKind::ScannAh {
921 return Err(invalid_data(
922 "ScaNN AH query used with a different ANN payload",
923 ));
924 }
925 if k == 0 {
926 return Ok(Vec::new());
927 }
928 #[cfg(feature = "native")]
929 self.prefetch_cluster_runs(query.routed_leaves());
930 let blocks = self.header.dim.div_ceil(self.header.code_size);
931 let packed_row_bytes = blocks.div_ceil(2);
932 let packed_block_bytes = blocks
933 .checked_mul(FAST_SCAN_LANES / 2)
934 .ok_or_else(|| invalid_data("ScaNN FastScan block size overflows"))?;
935 let fast_query = FastScanQuery::new(query.ah_query());
936 let bytes = self.raw.as_slice();
937 let mut ordinal_scores = Vec::new();
938 for &leaf in query.routed_leaves() {
939 let centroid_dot = query
940 .centroid_dot(leaf)
941 .ok_or_else(|| invalid_data("ScaNN routed leaf has no centroid score"))?;
942 for run in self.cluster_runs(leaf) {
943 let codes = &bytes[run.codes.clone()];
944 let full_blocks = run.count / FAST_SCAN_LANES;
945 for block_index in 0..full_blocks {
946 let start = block_index * packed_block_bytes;
947 let scores = fast_query
948 .score_block(&codes[start..start + packed_block_bytes], centroid_dot)
949 .map_err(|error| invalid_data(error.to_string()))?;
950 for (lane, &score) in scores.iter().enumerate() {
951 let row = block_index * FAST_SCAN_LANES + lane;
952 ordinal_scores.push((
953 run_doc_id(bytes, run, row)?,
954 read_u16(bytes, run.ordinals.start + row * 2),
955 score,
956 ));
957 }
958 }
959 let tail_start = full_blocks * packed_block_bytes;
960 for tail in 0..run.count % FAST_SCAN_LANES {
961 let start = tail_start + tail * packed_row_bytes;
962 let score = query
963 .ah_query()
964 .score_packed(&codes[start..start + packed_row_bytes], centroid_dot)
965 .map_err(|error| invalid_data(error.to_string()))?;
966 let row = full_blocks * FAST_SCAN_LANES + tail;
967 ordinal_scores.push((
968 run_doc_id(bytes, run, row)?,
969 read_u16(bytes, run.ordinals.start + row * 2),
970 score,
971 ));
972 }
973 }
974 }
975 Ok(combine_scored_ordinals(ordinal_scores, k, combiner))
976 }
977
978 fn search_binary_combined_documents_with_tuning(
979 &self,
980 k: usize,
981 query: &[u8],
982 cluster_ids: &[u32],
983 combiner: crate::query::MultiValueCombiner,
984 parallel_min_postings: usize,
985 ) -> io::Result<CombinedBinaryCandidates> {
986 validate_combined_search(combiner)?;
987 #[cfg(not(feature = "native"))]
988 let _ = parallel_min_postings;
989 if query.len() != self.header.code_size {
990 return Err(io::Error::new(
991 io::ErrorKind::InvalidInput,
992 "binary ANN query has the wrong byte length",
993 ));
994 }
995 if k == 0 {
996 return Ok((Vec::new(), Vec::new()));
997 }
998 #[cfg(feature = "native")]
999 self.prefetch_cluster_runs(cluster_ids);
1000
1001 let bytes = self.raw.as_slice();
1002 let posting_count = probed_posting_count(self, cluster_ids)?;
1003 #[cfg(feature = "native")]
1004 if rayon::current_num_threads() > 1 && posting_count >= parallel_min_postings {
1005 use rayon::prelude::*;
1006 let tasks = self.binary_scan_tasks(cluster_ids);
1007 let (ordinal_scores, _) = tasks
1008 .par_iter()
1009 .try_fold(
1010 || {
1011 (
1012 Vec::with_capacity(BINARY_SCORE_BATCH),
1013 vec![0.0f32; BINARY_SCORE_BATCH],
1014 )
1015 },
1016 |(mut ordinal_scores, mut scores), task| {
1017 score_binary_task(
1018 bytes,
1019 query,
1020 self.header.dim,
1021 self.header.code_size,
1022 *task,
1023 &mut scores,
1024 &mut ordinal_scores,
1025 )?;
1026 Ok::<_, io::Error>((ordinal_scores, scores))
1027 },
1028 )
1029 .try_reduce(
1030 || (Vec::new(), Vec::new()),
1031 |(mut left, scores), (mut right, _)| {
1032 left.append(&mut right);
1033 Ok((left, scores))
1034 },
1035 )?;
1036 return Ok(combine_scored_ordinals_retaining(
1037 ordinal_scores,
1038 k,
1039 combiner,
1040 ));
1041 }
1042
1043 let mut score_batch = vec![0.0f32; BINARY_SCORE_BATCH.min(self.header.vector_count)];
1044 let mut ordinal_scores = Vec::new();
1045 ordinal_scores
1046 .try_reserve_exact(posting_count)
1047 .map_err(|_| invalid_data("binary combined score buffer allocation failed"))?;
1048 score_binary_cluster_runs(
1049 self,
1050 bytes,
1051 query,
1052 cluster_ids,
1053 &mut score_batch,
1054 &mut ordinal_scores,
1055 )?;
1056 Ok(combine_scored_ordinals_retaining(
1057 ordinal_scores,
1058 k,
1059 combiner,
1060 ))
1061 }
1062
1063 pub(crate) fn search_tq_distinct(
1066 &self,
1067 k: usize,
1068 plan: &crate::structures::TqQueryPlan,
1069 ) -> io::Result<Vec<(u32, u16, f32)>> {
1070 use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_block_bytes};
1071
1072 if plan.padded_dim() != self.header.code_size * 2 {
1073 return Err(io::Error::new(
1074 io::ErrorKind::InvalidInput,
1075 "TQ query plan does not match the payload dimension",
1076 ));
1077 }
1078 let block_bytes = tq_block_bytes(self.header.code_size);
1079 let bytes = self.raw.as_slice();
1080
1081 #[cfg(feature = "native")]
1086 if self.header.vector_count >= TQ_PARALLEL_SCAN_MIN_VECTORS {
1087 use rayon::prelude::*;
1088 let collector = self
1089 .runs
1090 .par_iter()
1091 .flat_map(|run| {
1092 let codes = &bytes[run.codes.clone()];
1093 let blocks = codes.len() / block_bytes;
1094 (0..blocks.div_ceil(TQ_PARALLEL_SCAN_CHUNK_BLOCKS))
1095 .into_par_iter()
1096 .map(move |chunk| (run, chunk * TQ_PARALLEL_SCAN_CHUNK_BLOCKS, blocks))
1097 })
1098 .try_fold(
1099 || BoundedAnnCollector::<true, true>::new(k),
1100 |mut collector, (run, first_block, total_blocks)| {
1101 let codes = &bytes[run.codes.clone()];
1102 let last_block =
1103 (first_block + TQ_PARALLEL_SCAN_CHUNK_BLOCKS).min(total_blocks);
1104 let mut scores = [0.0f32; TQ_BLOCK_LANES];
1105 for block_index in first_block..last_block {
1106 let block = &codes[block_index * block_bytes..][..block_bytes];
1107 crate::structures::vector::quantization::tq_score_block(
1108 plan,
1109 block,
1110 &mut scores,
1111 );
1112 let lane_base = block_index * TQ_BLOCK_LANES;
1113 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
1114 for (lane, &score) in scores.iter().enumerate().take(lanes) {
1115 let index = lane_base + lane;
1116 collector.insert(
1117 run_doc_id(bytes, run, index)?,
1118 read_u16(bytes, run.ordinals.start + index * 2),
1119 score,
1120 );
1121 }
1122 }
1123 Ok::<_, io::Error>(collector)
1124 },
1125 )
1126 .try_reduce(
1127 || BoundedAnnCollector::<true, true>::new(k),
1128 |mut collector, partial| {
1129 collector.merge_from(partial);
1130 Ok(collector)
1131 },
1132 )?;
1133 return Ok(collector.into_sorted_results());
1134 }
1135
1136 let mut collector = BoundedAnnCollector::<true, true>::new(k);
1137 let mut scores = [0.0f32; TQ_BLOCK_LANES];
1138 for run in &self.runs {
1139 let codes = &bytes[run.codes.clone()];
1140 for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
1141 crate::structures::vector::quantization::tq_score_block(plan, block, &mut scores);
1142 let lane_base = block_index * TQ_BLOCK_LANES;
1143 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
1144 for (lane, &score) in scores.iter().enumerate().take(lanes) {
1145 let index = lane_base + lane;
1146 collector.insert(
1147 run_doc_id(bytes, run, index)?,
1148 read_u16(bytes, run.ordinals.start + index * 2),
1149 score,
1150 );
1151 }
1152 }
1153 }
1154 Ok(collector.into_sorted_results())
1155 }
1156
1157 pub(crate) fn search_ivf_tq_distinct(
1165 &self,
1166 k: usize,
1167 plan: &crate::structures::TqIvfQueryPlan,
1168 ) -> io::Result<Vec<(u32, u16, f32)>> {
1169 self.search_ivf_tq_distinct_with_tuning(
1170 k,
1171 plan,
1172 IVF_PARALLEL_SCAN_MIN_POSTINGS,
1173 IVF_TQ_PARALLEL_SCAN_CHUNK_BLOCKS,
1174 )
1175 }
1176
1177 fn search_ivf_tq_distinct_with_tuning(
1178 &self,
1179 k: usize,
1180 plan: &crate::structures::TqIvfQueryPlan,
1181 parallel_min_postings: usize,
1182 parallel_chunk_blocks: usize,
1183 ) -> io::Result<Vec<(u32, u16, f32)>> {
1184 use crate::structures::vector::quantization::tq_ivf_block_bytes;
1185
1186 let tq_plan = plan.tq_plan();
1187 self.validate_ivf_tq_query_plan(plan)?;
1188 #[cfg(not(feature = "native"))]
1189 let _ = (parallel_min_postings, parallel_chunk_blocks);
1190 if k == 0 {
1191 return Ok(Vec::new());
1192 }
1193 #[cfg(feature = "native")]
1194 self.prefetch_cluster_runs(&plan.cluster_ids);
1195 let block_bytes = tq_ivf_block_bytes(self.header.code_size);
1196 let bytes = self.raw.as_slice();
1197
1198 #[cfg(feature = "native")]
1203 if rayon::current_num_threads() > 1
1204 && probed_posting_count(self, &plan.cluster_ids)? >= parallel_min_postings
1205 {
1206 use rayon::prelude::*;
1207 let tasks = self.ivf_tq_scan_tasks(plan, block_bytes, parallel_chunk_blocks);
1208
1209 if let Some((pilot, remaining)) = tasks.split_first() {
1210 let mut pilot_collector = BoundedAnnCollector::<true, true>::new(k);
1211 let (pilot_pruned, pilot_scored) =
1212 score_ivf_tq_blocks(bytes, tq_plan, block_bytes, *pilot, &mut pilot_collector)?;
1213 let seed = pilot_collector.into_sorted_results();
1214 let seeded_collector = || {
1215 let mut collector = BoundedAnnCollector::<true, true>::new(k);
1216 for &(doc_id, ordinal, score) in &seed {
1217 collector.insert(doc_id, ordinal, score);
1218 }
1219 collector
1220 };
1221 let (collector, pruned_blocks, scored_blocks) = remaining
1222 .par_iter()
1223 .try_fold(
1224 || (seeded_collector(), 0usize, 0usize),
1225 |(mut collector, pruned, scored), task| {
1226 let (task_pruned, task_scored) = score_ivf_tq_blocks(
1227 bytes,
1228 tq_plan,
1229 block_bytes,
1230 *task,
1231 &mut collector,
1232 )?;
1233 Ok::<_, io::Error>((
1234 collector,
1235 pruned + task_pruned,
1236 scored + task_scored,
1237 ))
1238 },
1239 )
1240 .try_reduce(
1241 || (seeded_collector(), 0usize, 0usize),
1242 |(mut collector, left_pruned, left_scored),
1243 (partial, right_pruned, right_scored)| {
1244 collector.merge_from(partial);
1245 Ok((
1246 collector,
1247 left_pruned + right_pruned,
1248 left_scored + right_scored,
1249 ))
1250 },
1251 )?;
1252 log_ivf_tq_pruning(
1253 pilot_pruned + pruned_blocks,
1254 pilot_scored + scored_blocks,
1255 true,
1256 );
1257 return Ok(collector.into_sorted_results());
1258 }
1259 }
1260
1261 let mut collector = BoundedAnnCollector::<true, true>::new(k);
1262 let mut pruned_blocks = 0usize;
1263 let mut scored_blocks = 0usize;
1264 for (cluster_id, cluster_dot) in plan.cluster_dots() {
1265 for run in self.cluster_runs(cluster_id) {
1266 let block_count = run.codes.len() / block_bytes;
1267 let (run_pruned, run_scored) = score_ivf_tq_blocks(
1268 bytes,
1269 tq_plan,
1270 block_bytes,
1271 IvfTqScanTask {
1272 run,
1273 cluster_dot,
1274 first_block: 0,
1275 last_block: block_count,
1276 },
1277 &mut collector,
1278 )?;
1279 pruned_blocks += run_pruned;
1280 scored_blocks += run_scored;
1281 }
1282 }
1283 log_ivf_tq_pruning(pruned_blocks, scored_blocks, false);
1284 Ok(collector.into_sorted_results())
1285 }
1286
1287 fn validate_ivf_tq_query_plan(
1288 &self,
1289 plan: &crate::structures::TqIvfQueryPlan,
1290 ) -> io::Result<()> {
1291 if self.header.kind != AnnKind::IvfTq
1292 || !crate::structures::is_ivf_tq_cosine_generation(self.header.quantizer_version)
1293 {
1294 return Err(invalid_data(
1295 "legacy raw IVF-TQ payloads cannot be searched; rebuild the index",
1296 ));
1297 }
1298 if plan.tq_plan().padded_dim() != self.header.code_size * 2
1299 || plan.quantizer_version != self.header.quantizer_version
1300 || plan.fingerprint != self.header.codebook_version
1301 {
1302 return Err(io::Error::new(
1303 io::ErrorKind::InvalidInput,
1304 "IVF-TQ query plan does not match the payload generation",
1305 ));
1306 }
1307 Ok(())
1308 }
1309
1310 pub(crate) fn search_binary_clusters<const BY_DOCUMENT: bool>(
1311 &self,
1312 query: &[u8],
1313 k: usize,
1314 cluster_ids: &[u32],
1315 ) -> io::Result<Vec<(u32, u16, f32)>> {
1316 self.search_binary_clusters_with_tuning::<BY_DOCUMENT>(
1317 query,
1318 k,
1319 cluster_ids,
1320 IVF_PARALLEL_SCAN_MIN_POSTINGS,
1321 )
1322 }
1323
1324 fn search_binary_clusters_with_tuning<const BY_DOCUMENT: bool>(
1325 &self,
1326 query: &[u8],
1327 k: usize,
1328 cluster_ids: &[u32],
1329 parallel_min_postings: usize,
1330 ) -> io::Result<Vec<(u32, u16, f32)>> {
1331 #[cfg(not(feature = "native"))]
1332 let _ = parallel_min_postings;
1333 if query.len() != self.header.code_size {
1334 return Err(io::Error::new(
1335 io::ErrorKind::InvalidInput,
1336 "binary ANN query has the wrong byte length",
1337 ));
1338 }
1339 if k == 0 {
1340 return Ok(Vec::new());
1341 }
1342 #[cfg(feature = "native")]
1343 self.prefetch_cluster_runs(cluster_ids);
1344 let bytes = self.raw.as_slice();
1345
1346 debug_assert!(
1349 BY_DOCUMENT || self.header.kind == AnnKind::ScannBinary || {
1350 let mut seen = rustc_hash::FxHashSet::default();
1351 cluster_ids
1352 .iter()
1353 .all(|cluster_id| seen.insert(*cluster_id))
1354 },
1355 "an IVF probe plan must not repeat a cluster",
1356 );
1357
1358 #[cfg(feature = "native")]
1359 if rayon::current_num_threads() > 1
1360 && probed_posting_count(self, cluster_ids)? >= parallel_min_postings
1361 {
1362 use rayon::prelude::*;
1363 let tasks = self.binary_scan_tasks(cluster_ids);
1364 let (collector, _) = tasks
1365 .par_iter()
1366 .try_fold(
1367 || {
1368 (
1369 BoundedAnnCollector::<BY_DOCUMENT, true>::new(k),
1370 vec![0.0f32; BINARY_SCORE_BATCH],
1371 )
1372 },
1373 |(mut collector, mut scores), task| {
1374 score_binary_task(
1375 bytes,
1376 query,
1377 self.header.dim,
1378 self.header.code_size,
1379 *task,
1380 &mut scores,
1381 &mut collector,
1382 )?;
1383 Ok::<_, io::Error>((collector, scores))
1384 },
1385 )
1386 .try_reduce(
1387 || (BoundedAnnCollector::<BY_DOCUMENT, true>::new(k), Vec::new()),
1388 |(mut collector, scores), (partial, _)| {
1389 collector.merge_from(partial);
1390 Ok((collector, scores))
1391 },
1392 )?;
1393 return Ok(collector.into_sorted_results());
1394 }
1395
1396 let mut scores = vec![0.0f32; BINARY_SCORE_BATCH.min(self.header.vector_count)];
1397 if BY_DOCUMENT || self.header.kind == AnnKind::ScannBinary {
1398 let mut collector = BoundedAnnCollector::<BY_DOCUMENT, true>::new(k);
1399 score_binary_cluster_runs(
1400 self,
1401 bytes,
1402 query,
1403 cluster_ids,
1404 &mut scores,
1405 &mut collector,
1406 )?;
1407 return Ok(collector.into_sorted_results());
1408 }
1409
1410 let mut collector = BoundedUniqueAnnCollector::<true>::new(k);
1412 score_binary_cluster_runs(self, bytes, query, cluster_ids, &mut scores, &mut collector)?;
1413 Ok(collector.into_sorted_results())
1414 }
1415}
1416
1417fn score_ivf_tq_blocks(
1421 bytes: &[u8],
1422 plan: &crate::structures::TqQueryPlan,
1423 block_bytes: usize,
1424 task: IvfTqScanTask<'_>,
1425 collector: &mut BoundedAnnCollector<true, true>,
1426) -> io::Result<(usize, usize)> {
1427 use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_score_ivf_block};
1428
1429 let run = task.run;
1430 let codes = &bytes[run.codes.clone()];
1431 let mut scores = [0.0f32; TQ_BLOCK_LANES];
1432 let mut scored_blocks = 0usize;
1433 for block_index in task.first_block..task.last_block {
1434 let block = &codes[block_index * block_bytes..][..block_bytes];
1435 if let Some(threshold) = collector.pruning_threshold()
1436 && task.cluster_dot + tq_ivf_block_max_scale(block) * TQ_PRUNE_ESTIMATE_BOUND
1437 <= threshold
1438 {
1439 return Ok((task.last_block - block_index, scored_blocks));
1440 }
1441 scored_blocks += 1;
1442 tq_score_ivf_block(plan, block, task.cluster_dot, &mut scores);
1443 let lane_base = block_index * TQ_BLOCK_LANES;
1444 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
1445 for (lane, &score) in scores.iter().enumerate().take(lanes) {
1446 let index = lane_base + lane;
1447 collector.insert(
1448 run_doc_id(bytes, run, index)?,
1449 read_u16(bytes, run.ordinals.start + index * 2),
1450 score,
1451 );
1452 }
1453 }
1454 Ok((0, scored_blocks))
1455}
1456
1457#[cfg(feature = "native")]
1458fn score_ivf_tq_combined_blocks(
1459 bytes: &[u8],
1460 plan: &crate::structures::TqQueryPlan,
1461 block_bytes: usize,
1462 task: IvfTqScanTask<'_>,
1463 ordinal_scores: &mut Vec<(u32, u16, f32)>,
1464) -> io::Result<()> {
1465 use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_score_ivf_block};
1466
1467 let run = task.run;
1468 let codes = &bytes[run.codes.clone()];
1469 let mut scores = [0.0f32; TQ_BLOCK_LANES];
1470 for block_index in task.first_block..task.last_block {
1471 let block = &codes[block_index * block_bytes..][..block_bytes];
1472 tq_score_ivf_block(plan, block, task.cluster_dot, &mut scores);
1473 let lane_base = block_index * TQ_BLOCK_LANES;
1474 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
1475 for (lane, &score) in scores.iter().enumerate().take(lanes) {
1476 if score.is_finite() {
1477 let index = lane_base + lane;
1478 ordinal_scores.push((
1479 run_doc_id(bytes, run, index)?,
1480 read_u16(bytes, run.ordinals.start + index * 2),
1481 score,
1482 ));
1483 }
1484 }
1485 }
1486 Ok(())
1487}
1488
1489fn log_ivf_tq_pruning(pruned_blocks: usize, scored_blocks: usize, parallel: bool) {
1490 if pruned_blocks > 0 {
1491 log::debug!(
1492 "[search_ivf_tq] pruned {pruned_blocks} of {} blocks via scale bounds ({})",
1493 pruned_blocks + scored_blocks,
1494 if parallel { "parallel" } else { "serial" },
1495 );
1496 }
1497}
1498
1499#[cfg(feature = "native")]
1500fn coalesce_prefetch_ranges(ranges: &mut Vec<Range<usize>>) {
1501 if ranges.len() < 2 {
1502 return;
1503 }
1504 ranges.sort_unstable_by_key(|range| range.start);
1505 let mut output_len = 1usize;
1506 for input_index in 1..ranges.len() {
1507 let next_start = ranges[input_index].start;
1508 let next_end = ranges[input_index].end;
1509 let previous = &mut ranges[output_len - 1];
1510 if next_start <= previous.end.saturating_add(PREFETCH_COALESCE_GAP) {
1511 previous.end = previous.end.max(next_end);
1512 } else {
1513 ranges[output_len] = next_start..next_end;
1514 output_len += 1;
1515 }
1516 }
1517 ranges.truncate(output_len);
1518}
1519
1520trait AnnScoreSink {
1521 fn insert_score(&mut self, doc_id: u32, ordinal: u16, score: f32);
1522}
1523
1524impl<const BY_DOCUMENT: bool> AnnScoreSink for BoundedAnnCollector<BY_DOCUMENT, true> {
1525 #[inline]
1526 fn insert_score(&mut self, doc_id: u32, ordinal: u16, score: f32) {
1527 self.insert(doc_id, ordinal, score);
1528 }
1529}
1530
1531impl AnnScoreSink for BoundedUniqueAnnCollector<true> {
1532 #[inline]
1533 fn insert_score(&mut self, doc_id: u32, ordinal: u16, score: f32) {
1534 self.insert(doc_id, ordinal, score);
1535 }
1536}
1537
1538impl AnnScoreSink for Vec<(u32, u16, f32)> {
1539 #[inline]
1540 fn insert_score(&mut self, doc_id: u32, ordinal: u16, score: f32) {
1541 if score.is_finite() {
1542 self.push((doc_id, ordinal, score));
1543 }
1544 }
1545}
1546
1547fn probed_posting_count(index: &AnnDiskIndex, cluster_ids: &[u32]) -> io::Result<usize> {
1548 cluster_ids.iter().try_fold(0usize, |count, &cluster_id| {
1549 index
1550 .cluster_runs(cluster_id)
1551 .iter()
1552 .try_fold(count, |count, run| {
1553 count
1554 .checked_add(run.count)
1555 .ok_or_else(|| invalid_data("ANN probed posting count overflows usize"))
1556 })
1557 })
1558}
1559
1560fn score_binary_cluster_runs(
1561 index: &AnnDiskIndex,
1562 bytes: &[u8],
1563 query: &[u8],
1564 cluster_ids: &[u32],
1565 scores: &mut [f32],
1566 collector: &mut impl AnnScoreSink,
1567) -> io::Result<()> {
1568 for &cluster_id in cluster_ids {
1569 for run in index.cluster_runs(cluster_id) {
1570 score_binary_run(
1571 bytes,
1572 run,
1573 query,
1574 index.header.dim,
1575 index.header.code_size,
1576 scores,
1577 collector,
1578 )?;
1579 }
1580 }
1581 Ok(())
1582}
1583
1584fn score_binary_run(
1585 bytes: &[u8],
1586 run: &AnnRun,
1587 query: &[u8],
1588 dim_bits: usize,
1589 code_size: usize,
1590 scores: &mut [f32],
1591 collector: &mut impl AnnScoreSink,
1592) -> io::Result<()> {
1593 for batch_start in (0..run.count).step_by(BINARY_SCORE_BATCH) {
1594 score_binary_task(
1595 bytes,
1596 query,
1597 dim_bits,
1598 code_size,
1599 BinaryScanTask {
1600 run,
1601 first_index: batch_start,
1602 count: BINARY_SCORE_BATCH.min(run.count - batch_start),
1603 },
1604 scores,
1605 collector,
1606 )?;
1607 }
1608 Ok(())
1609}
1610
1611fn score_binary_task(
1612 bytes: &[u8],
1613 query: &[u8],
1614 dim_bits: usize,
1615 code_size: usize,
1616 task: BinaryScanTask<'_>,
1617 scores: &mut [f32],
1618 collector: &mut impl AnnScoreSink,
1619) -> io::Result<()> {
1620 let code_start = task.run.codes.start + task.first_index * code_size;
1621 let code_end = code_start + task.count * code_size;
1622 crate::structures::simd::batch_hamming_scores(
1623 query,
1624 &bytes[code_start..code_end],
1625 code_size,
1626 dim_bits,
1627 &mut scores[..task.count],
1628 );
1629 for (batch_index, &score) in scores.iter().enumerate().take(task.count) {
1630 let index = task.first_index + batch_index;
1631 collector.insert_score(
1632 run_doc_id(bytes, task.run, index)?,
1633 read_u16(bytes, task.run.ordinals.start + index * 2),
1634 score,
1635 );
1636 }
1637 Ok(())
1638}
1639
1640#[inline]
1641fn retain_combined_document(
1642 collector: &mut BoundedAnnCollector<true, true>,
1643 doc_id: u32,
1644 ordinal_scores: &[(u32, f32)],
1645 combiner: crate::query::MultiValueCombiner,
1646) {
1647 if !ordinal_scores.is_empty() {
1648 collector.insert(doc_id, 0, combiner.combine(ordinal_scores));
1649 }
1650}
1651
1652fn validate_combined_search(combiner: crate::query::MultiValueCombiner) -> io::Result<()> {
1653 combiner
1654 .validate()
1655 .map_err(|message| io::Error::new(io::ErrorKind::InvalidInput, message))
1656}
1657
1658fn combine_scored_ordinals(
1664 scores: Vec<(u32, u16, f32)>,
1665 k: usize,
1666 combiner: crate::query::MultiValueCombiner,
1667) -> Vec<AnnDocumentCandidate> {
1668 combine_scored_ordinals_retaining(scores, k, combiner).0
1669}
1670
1671fn combine_scored_ordinals_retaining(
1677 mut scores: Vec<(u32, u16, f32)>,
1678 k: usize,
1679 combiner: crate::query::MultiValueCombiner,
1680) -> CombinedBinaryCandidates {
1681 if k == 0 || scores.is_empty() {
1682 return (Vec::new(), Vec::new());
1683 }
1684 scores.retain(|entry| entry.2.is_finite());
1685 let by_document_ordinal = |left: &(u32, u16, f32), right: &(u32, u16, f32)| {
1686 left.0.cmp(&right.0).then_with(|| left.1.cmp(&right.1))
1687 };
1688 #[cfg(feature = "native")]
1689 if rayon::current_num_threads() > 1 && scores.len() >= IVF_PARALLEL_SCAN_MIN_POSTINGS {
1690 use rayon::prelude::*;
1691 scores.par_sort_unstable_by(by_document_ordinal);
1692 } else {
1693 scores.sort_unstable_by(by_document_ordinal);
1694 }
1695 #[cfg(not(feature = "native"))]
1696 scores.sort_unstable_by(by_document_ordinal);
1697
1698 let mut unique_len = 0usize;
1703 for read_index in 0..scores.len() {
1704 let candidate = scores[read_index];
1705 if unique_len > 0
1706 && scores[unique_len - 1].0 == candidate.0
1707 && scores[unique_len - 1].1 == candidate.1
1708 {
1709 if candidate.2.total_cmp(&scores[unique_len - 1].2).is_gt() {
1710 scores[unique_len - 1].2 = candidate.2;
1711 }
1712 } else {
1713 scores[unique_len] = candidate;
1714 unique_len += 1;
1715 }
1716 }
1717 scores.truncate(unique_len);
1718
1719 let mut top_documents = BoundedAnnCollector::<true, true>::new(k);
1720 let mut current_doc = None;
1721 let mut ordinal_scores = Vec::new();
1722 for &(doc_id, ordinal, score) in &scores {
1723 if current_doc.is_some_and(|current| current != doc_id) {
1724 retain_combined_document(
1725 &mut top_documents,
1726 current_doc.expect("current document is present"),
1727 &ordinal_scores,
1728 combiner,
1729 );
1730 ordinal_scores.clear();
1731 }
1732 current_doc = Some(doc_id);
1733 ordinal_scores.push((u32::from(ordinal), score));
1734 }
1735 if let Some(doc_id) = current_doc {
1736 retain_combined_document(&mut top_documents, doc_id, &ordinal_scores, combiner);
1737 }
1738
1739 let candidates: Vec<AnnDocumentCandidate> = top_documents
1740 .into_sorted_results()
1741 .into_iter()
1742 .map(|(doc_id, _, score)| AnnDocumentCandidate { doc_id, score })
1743 .collect();
1744
1745 let mut retained_ids: Vec<u32> = candidates
1749 .iter()
1750 .map(|candidate| candidate.doc_id)
1751 .collect();
1752 retained_ids.sort_unstable();
1753 let mut retained = Vec::with_capacity(scores.len().min(retained_ids.len().saturating_mul(4)));
1754 let mut cursor = 0usize;
1755 for entry in scores {
1756 while cursor < retained_ids.len() && retained_ids[cursor] < entry.0 {
1757 cursor += 1;
1758 }
1759 if retained_ids.get(cursor) == Some(&entry.0) {
1760 retained.push(entry);
1761 }
1762 }
1763 (candidates, retained)
1764}
1765
1766#[cfg(feature = "native")]
1767pub(crate) fn write_built_binary_ivf(
1768 index: &BinaryIvfIndex,
1769 routing: IvfRoutingMode,
1770 writer: &mut (impl Write + ?Sized),
1771) -> io::Result<u64> {
1772 let runs: Vec<_> = index
1773 .clusters
1774 .iter()
1775 .map(|(cluster_id, cluster)| BuildRun {
1776 cluster_id: *cluster_id,
1777 doc_ids: &cluster.doc_ids,
1778 ordinals: &cluster.ordinals,
1779 codes: &cluster.codes,
1780 })
1781 .collect();
1782 write_built_runs(
1783 AnnDiskHeader {
1784 kind: AnnKind::BinaryIvf,
1785 routing,
1786 dim: index.dim_bits,
1787 code_size: index.dim_bits.div_ceil(8),
1788 num_clusters: index.num_clusters,
1789 quantizer_version: index.quantizer_version,
1790 codebook_version: 0,
1791 vector_count: index.len(),
1792 },
1793 &runs,
1794 writer,
1795 )
1796}
1797
1798#[cfg(feature = "native")]
1803#[cfg_attr(not(test), allow(dead_code))]
1804pub(crate) fn write_built_scann(
1805 payload: &crate::structures::vector::scann::ScannSegmentPayload,
1806 writer: &mut (impl Write + ?Sized),
1807) -> io::Result<u64> {
1808 use crate::structures::vector::scann::ScannEncoding;
1809
1810 let (kind, code_size) = match payload.encoding {
1811 ScannEncoding::AsymmetricHash {
1812 dimensions_per_block,
1813 bits_per_code: 4,
1814 } => (AnnKind::ScannAh, usize::from(dimensions_per_block)),
1815 ScannEncoding::AsymmetricHash { .. } => {
1816 return Err(invalid_data(
1817 "ScaNN ANN disk format supports only 4-bit AH codes",
1818 ));
1819 }
1820 ScannEncoding::BinaryHamming => (
1821 AnnKind::ScannBinary,
1822 usize::try_from(payload.dimension)
1823 .map_err(|_| invalid_data("ScaNN dimension exceeds usize"))?
1824 .div_ceil(8),
1825 ),
1826 };
1827 let vector_count = payload.runs().iter().try_fold(0usize, |total, run| {
1828 total
1829 .checked_add(run.row_count as usize)
1830 .ok_or_else(|| invalid_data("ScaNN vector count overflows usize"))
1831 })?;
1832 let header = AnnDiskHeader {
1833 kind,
1834 routing: IvfRoutingMode::Flat,
1837 dim: usize::try_from(payload.dimension)
1838 .map_err(|_| invalid_data("ScaNN dimension exceeds usize"))?,
1839 code_size,
1840 num_clusters: payload.num_leaves,
1841 quantizer_version: payload.generation,
1842 codebook_version: payload.artifact_id,
1843 vector_count,
1844 };
1845 validate_header(&header)?;
1846 write_header(writer, &header)?;
1847
1848 let mut offset = ANN_HEADER_SIZE as u64;
1849 let mut records = Vec::with_capacity(payload.runs().len());
1850 let mut previous_leaf = None;
1851 for run in payload.runs() {
1852 let count = run.row_count as usize;
1853 let expected_docs = count
1854 .checked_mul(4)
1855 .ok_or_else(|| invalid_data("ScaNN doc-ID column size overflows usize"))?;
1856 let expected_ordinals = count
1857 .checked_mul(2)
1858 .ok_or_else(|| invalid_data("ScaNN ordinal column size overflows usize"))?;
1859 if count == 0
1860 || run.leaf_id >= header.num_clusters
1861 || previous_leaf.is_some_and(|leaf| leaf > run.leaf_id)
1862 || run.doc_ids_le.len() != expected_docs
1863 || run.ordinals_le.len() != expected_ordinals
1864 || run.codes.len()
1865 != expected_codes_column_len(header.kind, count, header.dim, header.code_size)?
1866 {
1867 return Err(invalid_data("ScaNN ANN leaf run columns are inconsistent"));
1868 }
1869 previous_leaf = Some(run.leaf_id);
1870 let max_doc_id = run
1871 .doc_ids_le
1872 .chunks_exact(4)
1873 .map(|bytes| u32::from_le_bytes(bytes.try_into().unwrap()))
1874 .max()
1875 .unwrap_or(0);
1876 if run
1877 .doc_base
1878 .checked_add(max_doc_id)
1879 .is_none_or(|doc_id| doc_id >= payload.doc_count)
1880 {
1881 return Err(invalid_data(
1882 "ScaNN ANN leaf run document ID is out of range",
1883 ));
1884 }
1885
1886 let doc_ids_offset = offset;
1887 writer.write_all(&run.doc_ids_le)?;
1888 offset = checked_advance(offset, run.doc_ids_le.len())?;
1889 let ordinals_offset = offset;
1890 writer.write_all(&run.ordinals_le)?;
1891 offset = checked_advance(offset, run.ordinals_le.len())?;
1892 let codes_offset = offset;
1893 writer.write_all(&run.codes)?;
1894 offset = checked_advance(offset, run.codes.len())?;
1895 records.push(RunRecord {
1896 cluster_id: run.leaf_id,
1897 doc_base: run.doc_base,
1898 count: run.row_count,
1899 max_doc_id,
1900 doc_ids_offset,
1901 ordinals_offset,
1902 codes_offset,
1903 codes_len: u64::try_from(run.codes.len())
1904 .map_err(|_| invalid_data("ScaNN code output size exceeds u64"))?,
1905 });
1906 }
1907 if records.is_empty() {
1908 return Err(invalid_data("cannot write an empty ScaNN ANN payload"));
1909 }
1910 finish_layout(writer, offset, &records)
1911}
1912
1913#[cfg(feature = "native")]
1916pub(crate) fn write_built_ivf_tq(
1917 index: &crate::structures::IvfTqIndex,
1918 num_clusters: u32,
1919 writer: &mut (impl Write + ?Sized),
1920) -> io::Result<u64> {
1921 use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_pack_ivf_block};
1922
1923 if !crate::structures::is_ivf_tq_cosine_generation(index.centroids_version) {
1924 return Err(invalid_data(
1925 "legacy raw IVF-TQ generations cannot be serialized; rebuild the index",
1926 ));
1927 }
1928 let codec = index.codec();
1929 let padded_dim = codec.padded_dim();
1930 let mut clusters: Vec<_> = index.clusters.iter().collect();
1931 clusters.sort_unstable_by_key(|(cluster_id, _)| **cluster_id);
1932 struct PackedCluster {
1936 cluster_id: u32,
1937 doc_ids: Vec<u32>,
1938 ordinals: Vec<u16>,
1939 codes: Vec<u8>,
1940 }
1941 let packed: Vec<PackedCluster> = clusters
1942 .iter()
1943 .map(|&(&cluster_id, cluster)| {
1944 let count = cluster.doc_ids.len();
1945 let mut order: Vec<usize> = (0..count).collect();
1946 order.sort_by(|&a, &b| {
1947 cluster.scales[b]
1948 .total_cmp(&cluster.scales[a])
1949 .then_with(|| a.cmp(&b))
1950 });
1951 let doc_ids: Vec<u32> = order.iter().map(|&i| cluster.doc_ids[i]).collect();
1952 let ordinals: Vec<u16> = order.iter().map(|&i| cluster.ordinals[i]).collect();
1953 let scales: Vec<f32> = order.iter().map(|&i| cluster.scales[i]).collect();
1954 let gammas: Vec<f32> = order.iter().map(|&i| cluster.gammas[i]).collect();
1955 let mut codes = Vec::with_capacity(
1956 crate::structures::vector::quantization::tq_ivf_codes_column_len_checked(
1957 count,
1958 codec.code_size(),
1959 )
1960 .unwrap_or_default(),
1961 );
1962 for block_start in (0..count).step_by(TQ_BLOCK_LANES) {
1963 let lanes = TQ_BLOCK_LANES.min(count - block_start);
1964 let rows: Vec<&[u8]> = order[block_start..block_start + lanes]
1965 .iter()
1966 .map(|&row| &cluster.rows[row * padded_dim..(row + 1) * padded_dim])
1967 .collect();
1968 tq_pack_ivf_block(
1969 &rows,
1970 &scales[block_start..block_start + lanes],
1971 &gammas[block_start..block_start + lanes],
1972 padded_dim,
1973 &mut codes,
1974 );
1975 }
1976 PackedCluster {
1977 cluster_id,
1978 doc_ids,
1979 ordinals,
1980 codes,
1981 }
1982 })
1983 .collect();
1984 let runs: Vec<_> = packed
1985 .iter()
1986 .map(|cluster| BuildRun {
1987 cluster_id: cluster.cluster_id,
1988 doc_ids: &cluster.doc_ids,
1989 ordinals: &cluster.ordinals,
1990 codes: &cluster.codes,
1991 })
1992 .collect();
1993 write_built_runs(
1994 AnnDiskHeader {
1995 kind: AnnKind::IvfTq,
1996 routing: index.routing,
1997 dim: index.dim,
1998 code_size: codec.code_size(),
1999 num_clusters,
2000 quantizer_version: index.centroids_version,
2001 codebook_version: codec.fingerprint(),
2002 vector_count: index.len(),
2003 },
2004 &runs,
2005 writer,
2006 )
2007}
2008
2009#[cfg(feature = "native")]
2011pub(crate) fn tq_builder_extra_run(builder: &crate::structures::TqFlatBuilder) -> BuildRun<'_> {
2012 BuildRun {
2013 cluster_id: 0,
2014 doc_ids: &builder.doc_ids,
2015 ordinals: &builder.ordinals,
2016 codes: &builder.codes,
2017 }
2018}
2019
2020#[cfg(feature = "native")]
2022pub(crate) fn write_built_tq_flat(
2023 builder: &crate::structures::TqFlatBuilder,
2024 writer: &mut (impl Write + ?Sized),
2025) -> io::Result<u64> {
2026 let codec = builder.codec();
2027 let runs = [BuildRun {
2028 cluster_id: 0,
2029 doc_ids: &builder.doc_ids,
2030 ordinals: &builder.ordinals,
2031 codes: &builder.codes,
2032 }];
2033 write_built_runs(
2034 AnnDiskHeader {
2035 kind: AnnKind::TqFlat,
2036 routing: IvfRoutingMode::Flat,
2037 dim: codec.dim(),
2038 code_size: codec.code_size(),
2039 num_clusters: 1,
2040 quantizer_version: codec.fingerprint(),
2041 codebook_version: 0,
2042 vector_count: builder.len(),
2043 },
2044 &runs,
2045 writer,
2046 )
2047}
2048
2049#[cfg(feature = "native")]
2050pub(crate) struct BuildRun<'a> {
2051 cluster_id: u32,
2052 doc_ids: &'a [u32],
2053 ordinals: &'a [u16],
2054 codes: &'a [u8],
2055}
2056
2057#[cfg(feature = "native")]
2058struct RunRecord {
2059 cluster_id: u32,
2060 doc_base: u32,
2061 count: u32,
2062 max_doc_id: u32,
2063 doc_ids_offset: u64,
2064 ordinals_offset: u64,
2065 codes_offset: u64,
2066 codes_len: u64,
2067}
2068
2069#[cfg(feature = "native")]
2070fn write_built_runs(
2071 header: AnnDiskHeader,
2072 runs: &[BuildRun<'_>],
2073 writer: &mut (impl Write + ?Sized),
2074) -> io::Result<u64> {
2075 if runs.is_empty() || header.vector_count == 0 {
2076 return Err(invalid_data("cannot write an empty ANN payload"));
2077 }
2078 validate_header(&header)?;
2079 write_header(writer, &header)?;
2080 let mut offset = ANN_HEADER_SIZE as u64;
2081 let mut records = Vec::with_capacity(runs.len());
2082 let mut counted = 0usize;
2083 let mut scratch = Vec::new();
2084 let mut previous_cluster = None;
2085 for run in runs {
2086 let count = run.doc_ids.len();
2087 if count == 0
2088 || run.cluster_id >= header.num_clusters
2089 || previous_cluster.is_some_and(|cluster| cluster >= run.cluster_id)
2090 || run.ordinals.len() != count
2091 || run.codes.len()
2092 != expected_codes_column_len(header.kind, count, header.dim, header.code_size)?
2093 {
2094 return Err(invalid_data("ANN build run columns are inconsistent"));
2095 }
2096 previous_cluster = Some(run.cluster_id);
2097 let count_u32 = u32::try_from(count)
2098 .map_err(|_| invalid_data("ANN cluster run exceeds u32 vectors"))?;
2099 let max_doc_id = run.doc_ids.iter().copied().max().unwrap_or(0);
2100 let doc_ids_offset = offset;
2101 write_u32_column(writer, run.doc_ids, &mut scratch)?;
2102 offset = offset
2103 .checked_add(
2104 u64::try_from(count)
2105 .ok()
2106 .and_then(|count| count.checked_mul(4))
2107 .ok_or_else(|| invalid_data("ANN doc-ID output size overflows u64"))?,
2108 )
2109 .ok_or_else(|| invalid_data("ANN output offset overflow"))?;
2110 let ordinals_offset = offset;
2111 write_u16_column(writer, run.ordinals, &mut scratch)?;
2112 offset = offset
2113 .checked_add(
2114 u64::try_from(count)
2115 .ok()
2116 .and_then(|count| count.checked_mul(2))
2117 .ok_or_else(|| invalid_data("ANN ordinal output size overflows u64"))?,
2118 )
2119 .ok_or_else(|| invalid_data("ANN output offset overflow"))?;
2120 let codes_offset = offset;
2121 writer.write_all(run.codes)?;
2122 offset = offset
2123 .checked_add(
2124 u64::try_from(run.codes.len())
2125 .map_err(|_| invalid_data("ANN code output size exceeds u64"))?,
2126 )
2127 .ok_or_else(|| invalid_data("ANN output offset overflow"))?;
2128 records.push(RunRecord {
2129 cluster_id: run.cluster_id,
2130 doc_base: 0,
2131 count: count_u32,
2132 max_doc_id,
2133 doc_ids_offset,
2134 ordinals_offset,
2135 codes_offset,
2136 codes_len: u64::try_from(run.codes.len())
2137 .map_err(|_| invalid_data("ANN code output size exceeds u64"))?,
2138 });
2139 counted = counted
2140 .checked_add(count)
2141 .ok_or_else(|| invalid_data("ANN vector count overflow"))?;
2142 }
2143 if counted != header.vector_count {
2144 return Err(invalid_data("ANN header/build vector counts disagree"));
2145 }
2146 finish_layout(writer, offset, &records)
2147}
2148
2149#[cfg(all(feature = "native", test))]
2152pub(crate) fn write_merged_ann(
2153 sources: &[(&AnnDiskIndex, u32)],
2154 writer: &mut (impl Write + ?Sized),
2155) -> io::Result<u64> {
2156 write_merged_ann_impl(sources, &[], writer, None)
2157}
2158
2159#[cfg(feature = "native")]
2160pub(crate) fn write_merged_ann_cancellable(
2161 sources: &[(&AnnDiskIndex, u32)],
2162 writer: &mut (impl Write + ?Sized),
2163 cancellation: Option<&std::sync::atomic::AtomicBool>,
2164) -> io::Result<u64> {
2165 write_merged_ann_impl(sources, &[], writer, cancellation)
2166}
2167
2168#[cfg(feature = "native")]
2177pub(crate) fn predicted_merge_fragmentation(sources: &[(&AnnDiskIndex, u32)]) -> f64 {
2178 let mut total_runs = 0usize;
2179 let mut cursors: Vec<std::iter::Peekable<std::slice::Iter<'_, AnnRun>>> = sources
2182 .iter()
2183 .map(|(source, _)| {
2184 total_runs += source.runs.len();
2185 source.runs.iter().peekable()
2186 })
2187 .collect();
2188 let mut distinct = 0usize;
2189 while let Some(cluster) = cursors
2190 .iter_mut()
2191 .filter_map(|cursor| cursor.peek().map(|run| run.cluster_id))
2192 .min()
2193 {
2194 distinct += 1;
2195 for cursor in &mut cursors {
2196 while cursor.peek().is_some_and(|run| run.cluster_id == cluster) {
2197 cursor.next();
2198 }
2199 }
2200 }
2201 if distinct == 0 {
2202 0.0
2203 } else {
2204 total_runs as f64 / distinct as f64
2205 }
2206}
2207
2208#[cfg(feature = "native")]
2210const DOC_ID_REWRITE_CHUNK: usize = 64 * 1024;
2211
2212#[cfg(feature = "native")]
2240pub(crate) fn write_compacted_ann_cancellable(
2241 sources: &[(&AnnDiskIndex, u32)],
2242 writer: &mut (impl Write + ?Sized),
2243 cancellation: Option<&std::sync::atomic::AtomicBool>,
2244) -> io::Result<u64> {
2245 let Some((first, _)) = sources.first() else {
2246 return Err(invalid_data("cannot compact an empty ANN source list"));
2247 };
2248 if !matches!(
2249 first.header.kind,
2250 AnnKind::BinaryIvf | AnnKind::ScannBinary | AnnKind::ScannAh
2251 ) {
2252 return Err(invalid_data(
2253 "ANN run compaction is only defined for exact binary and ScaNN AH payloads",
2254 ));
2255 }
2256 let mut header = first.header.clone();
2257 header.vector_count = 0;
2258 for &(source, _) in sources {
2259 if !headers_compatible(&first.header, &source.header) {
2260 return Err(invalid_data(
2261 "ANN compaction sources use incompatible generations",
2262 ));
2263 }
2264 header.vector_count = header
2265 .vector_count
2266 .checked_add(source.header.vector_count)
2267 .ok_or_else(|| invalid_data("compacted ANN vector count overflows usize"))?;
2268 }
2269 validate_header(&header)?;
2270 write_header(writer, &header)?;
2271
2272 let code_size = header.code_size;
2273 let mut offset = ANN_HEADER_SIZE as u64;
2274 let mut records: Vec<RunRecord> = Vec::new();
2275 let mut scratch = Vec::new();
2276 let mut cursors: Vec<usize> = vec![0; sources.len()];
2277
2278 loop {
2279 if cancellation.is_some_and(|flag| flag.load(std::sync::atomic::Ordering::Relaxed)) {
2280 return Err(io::Error::new(
2281 io::ErrorKind::Interrupted,
2282 "ANN compaction cancelled",
2283 ));
2284 }
2285 let Some(cluster_id) = sources
2287 .iter()
2288 .zip(&cursors)
2289 .filter_map(|((source, _), &cursor)| source.runs.get(cursor).map(|run| run.cluster_id))
2290 .min()
2291 else {
2292 break;
2293 };
2294
2295 let mut count = 0usize;
2300 let mut max_doc_id = 0u32;
2301 let doc_ids_offset = offset;
2302 for (source_index, &(source, segment_base)) in sources.iter().enumerate() {
2303 let mut cursor = cursors[source_index];
2304 while let Some(run) = source
2305 .runs
2306 .get(cursor)
2307 .filter(|run| run.cluster_id == cluster_id)
2308 {
2309 let base = run
2310 .doc_base
2311 .checked_add(segment_base)
2312 .ok_or_else(|| invalid_data("compacted ANN document base overflows u32"))?;
2313 let bytes = source.raw.as_slice();
2314 for chunk_start in (0..run.count).step_by(DOC_ID_REWRITE_CHUNK) {
2319 let chunk_end = (chunk_start + DOC_ID_REWRITE_CHUNK).min(run.count);
2320 scratch.clear();
2321 scratch.reserve((chunk_end - chunk_start) * 4);
2322 for index in chunk_start..chunk_end {
2323 let doc_id = run_doc_id_with_base(bytes, run, index, base)?;
2324 max_doc_id = max_doc_id.max(doc_id);
2325 scratch.extend_from_slice(&doc_id.to_le_bytes());
2326 }
2327 writer.write_all(&scratch)?;
2328 if cancellation
2329 .is_some_and(|flag| flag.load(std::sync::atomic::Ordering::Relaxed))
2330 {
2331 return Err(io::Error::new(
2332 io::ErrorKind::Interrupted,
2333 "ANN compaction cancelled",
2334 ));
2335 }
2336 }
2337 offset = checked_advance(offset, run.count * 4)?;
2338 count = count
2339 .checked_add(run.count)
2340 .ok_or_else(|| invalid_data("compacted ANN run count overflows usize"))?;
2341 cursor += 1;
2342 }
2343 }
2344
2345 let ordinals_offset = offset;
2347 for (source_index, &(source, _)) in sources.iter().enumerate() {
2348 let mut cursor = cursors[source_index];
2349 while let Some(run) = source
2350 .runs
2351 .get(cursor)
2352 .filter(|run| run.cluster_id == cluster_id)
2353 {
2354 copy_range(writer, &source.raw, run.ordinals.clone(), cancellation)?;
2355 offset = checked_advance(offset, run.ordinals.len())?;
2356 cursor += 1;
2357 }
2358 }
2359
2360 let codes_offset = offset;
2363 if header.kind == AnnKind::ScannAh {
2364 let blocks = header.dim.div_ceil(code_size);
2365 let lanes = crate::structures::vector::scann::FAST_SCAN_LANES;
2366 let mut unpacked = Vec::with_capacity(lanes * blocks);
2367 let mut packed = Vec::with_capacity(blocks * lanes / 2);
2368 for (source_index, &(source, _)) in sources.iter().enumerate() {
2369 let mut cursor = cursors[source_index];
2370 while let Some(run) = source
2371 .runs
2372 .get(cursor)
2373 .filter(|run| run.cluster_id == cluster_id)
2374 {
2375 let bytes = &source.raw.as_slice()[run.codes.clone()];
2376 for row in 0..run.count {
2377 unpack_scann_ah_row(bytes, run.count, blocks, row, &mut unpacked)?;
2378 if unpacked.len() == lanes * blocks {
2379 packed.clear();
2380 crate::structures::vector::scann::pack_fast_scan_block(
2381 &unpacked,
2382 blocks,
2383 &mut packed,
2384 )
2385 .map_err(|error| invalid_data(error.to_string()))?;
2386 writer.write_all(&packed)?;
2387 offset = checked_advance(offset, packed.len())?;
2388 unpacked.clear();
2389 if cancellation
2390 .is_some_and(|flag| flag.load(std::sync::atomic::Ordering::Relaxed))
2391 {
2392 return Err(io::Error::new(
2393 io::ErrorKind::Interrupted,
2394 "ANN compaction cancelled",
2395 ));
2396 }
2397 }
2398 }
2399 cursor += 1;
2400 }
2401 }
2402 if !unpacked.is_empty() {
2405 packed.clear();
2406 for row in unpacked.chunks_exact(blocks) {
2407 for pair in row.chunks(2) {
2408 packed.push(pair[0] | (pair.get(1).copied().unwrap_or(0) << 4));
2409 }
2410 }
2411 writer.write_all(&packed)?;
2412 offset = checked_advance(offset, packed.len())?;
2413 }
2414 } else {
2415 for (source_index, &(source, _)) in sources.iter().enumerate() {
2416 let mut cursor = cursors[source_index];
2417 while let Some(run) = source
2418 .runs
2419 .get(cursor)
2420 .filter(|run| run.cluster_id == cluster_id)
2421 {
2422 copy_range(writer, &source.raw, run.codes.clone(), cancellation)?;
2423 offset = checked_advance(offset, run.codes.len())?;
2424 cursor += 1;
2425 }
2426 }
2427 }
2428 let expected_codes_len =
2429 expected_codes_column_len(header.kind, count, header.dim, header.code_size)?;
2430 if offset - codes_offset != expected_codes_len as u64 {
2431 return Err(invalid_data("compacted ANN code column length mismatch"));
2432 }
2433
2434 for (source_index, &(source, _)) in sources.iter().enumerate() {
2436 while source
2437 .runs
2438 .get(cursors[source_index])
2439 .is_some_and(|run| run.cluster_id == cluster_id)
2440 {
2441 cursors[source_index] += 1;
2442 }
2443 }
2444
2445 records.push(RunRecord {
2446 cluster_id,
2447 doc_base: 0,
2448 count: u32::try_from(count)
2449 .map_err(|_| invalid_data("compacted ANN run exceeds u32 vectors"))?,
2450 max_doc_id,
2451 doc_ids_offset,
2452 ordinals_offset,
2453 codes_offset,
2454 codes_len: u64::try_from(expected_codes_column_len(
2455 header.kind,
2456 count,
2457 header.dim,
2458 code_size,
2459 )?)
2460 .map_err(|_| invalid_data("compacted ANN code length exceeds u64"))?,
2461 });
2462 }
2463
2464 if records.is_empty() {
2465 return Err(invalid_data("cannot compact an ANN payload with no runs"));
2466 }
2467 finish_layout(writer, offset, &records)
2468}
2469
2470#[cfg(feature = "native")]
2473fn unpack_scann_ah_row(
2474 codes: &[u8],
2475 count: usize,
2476 blocks: usize,
2477 row: usize,
2478 output: &mut Vec<u8>,
2479) -> io::Result<()> {
2480 let lanes = crate::structures::vector::scann::FAST_SCAN_LANES;
2481 let full_rows = count / lanes * lanes;
2482 let full_block_bytes = blocks
2483 .checked_mul(lanes / 2)
2484 .ok_or_else(|| invalid_data("ScaNN AH block size overflows usize"))?;
2485 let tail_row_bytes = blocks.div_ceil(2);
2486 for block in 0..blocks {
2487 let (byte_offset, high) = if row < full_rows {
2488 let lane = row % lanes;
2489 (
2490 (row / lanes) * full_block_bytes + block * (lanes / 2) + lane / 2,
2491 !lane.is_multiple_of(2),
2492 )
2493 } else {
2494 (
2495 (full_rows / lanes) * full_block_bytes
2496 + (row - full_rows) * tail_row_bytes
2497 + block / 2,
2498 !block.is_multiple_of(2),
2499 )
2500 };
2501 let byte = *codes
2502 .get(byte_offset)
2503 .ok_or_else(|| invalid_data("ScaNN AH row exceeds its code column"))?;
2504 output.push(if high { byte >> 4 } else { byte & 0x0f });
2505 }
2506 Ok(())
2507}
2508
2509#[cfg(feature = "native")]
2511fn run_doc_id_with_base(bytes: &[u8], run: &AnnRun, index: usize, base: u32) -> io::Result<u32> {
2512 let local_doc_id = read_u32(bytes, run.doc_ids.start + index * 4);
2513 if local_doc_id > run.max_doc_id {
2514 return Err(invalid_data(
2515 "ANN run contains a document above its declared maximum",
2516 ));
2517 }
2518 base.checked_add(local_doc_id)
2519 .ok_or_else(|| invalid_data("compacted ANN document ID overflows u32"))
2520}
2521
2522#[cfg(feature = "native")]
2526pub(crate) fn write_merged_ann_with_extra(
2527 sources: &[(&AnnDiskIndex, u32)],
2528 extra: &[BuildRun<'_>],
2529 writer: &mut (impl Write + ?Sized),
2530 cancellation: Option<&std::sync::atomic::AtomicBool>,
2531) -> io::Result<u64> {
2532 write_merged_ann_impl(sources, extra, writer, cancellation)
2533}
2534
2535#[cfg(feature = "native")]
2536fn write_merged_ann_impl(
2537 sources: &[(&AnnDiskIndex, u32)],
2538 extra: &[BuildRun<'_>],
2539 writer: &mut (impl Write + ?Sized),
2540 cancellation: Option<&std::sync::atomic::AtomicBool>,
2541) -> io::Result<u64> {
2542 let Some((first, _)) = sources.first() else {
2543 return Err(invalid_data("cannot merge an empty ANN source list"));
2544 };
2545 if first.header.kind == AnnKind::IvfTq
2546 && !crate::structures::is_ivf_tq_cosine_generation(first.header.quantizer_version)
2547 {
2548 return Err(invalid_data(
2549 "legacy raw IVF-TQ generations cannot be merged; rebuild the index",
2550 ));
2551 }
2552 let mut header = first.header.clone();
2553 header.vector_count = 0;
2554 for &(source, _) in sources {
2555 if !headers_compatible(&first.header, &source.header) {
2556 return Err(invalid_data(
2557 "ANN merge sources use incompatible generations",
2558 ));
2559 }
2560 header.vector_count = header
2561 .vector_count
2562 .checked_add(source.header.vector_count)
2563 .ok_or_else(|| invalid_data("merged ANN vector count overflows usize"))?;
2564 }
2565 for run in extra {
2566 if run.doc_ids.is_empty()
2567 || run.cluster_id >= header.num_clusters
2568 || run.ordinals.len() != run.doc_ids.len()
2569 || run.codes.len()
2570 != expected_codes_column_len(
2571 header.kind,
2572 run.doc_ids.len(),
2573 header.dim,
2574 header.code_size,
2575 )?
2576 {
2577 return Err(invalid_data("extra ANN merge run columns are inconsistent"));
2578 }
2579 header.vector_count = header
2580 .vector_count
2581 .checked_add(run.doc_ids.len())
2582 .ok_or_else(|| invalid_data("merged ANN vector count overflows usize"))?;
2583 }
2584 validate_header(&header)?;
2585 write_header(writer, &header)?;
2586 let mut offset = ANN_HEADER_SIZE as u64;
2587 let run_capacity = sources.iter().try_fold(extra.len(), |count, (source, _)| {
2588 count
2589 .checked_add(source.runs.len())
2590 .ok_or_else(|| invalid_data("merged ANN run count overflows usize"))
2591 })?;
2592 let mut output_payload_starts = Vec::with_capacity(sources.len());
2593 for &(source, _) in sources {
2594 let payload_end = source
2595 .runs
2596 .iter()
2597 .map(|run| run.codes.end)
2598 .max()
2599 .ok_or_else(|| invalid_data("ANN source has no payload runs"))?;
2600 let output_payload_start = offset;
2601 output_payload_starts.push(output_payload_start);
2602 copy_range(
2603 writer,
2604 &source.raw,
2605 ANN_HEADER_SIZE..payload_end,
2606 cancellation,
2607 )?;
2608 offset = checked_advance(offset, payload_end - ANN_HEADER_SIZE)?;
2609 }
2610
2611 let mut extra_records = Vec::with_capacity(extra.len());
2614 let mut scratch = Vec::new();
2615 for run in extra {
2616 let count = run.doc_ids.len();
2617 let doc_ids_offset = offset;
2618 write_u32_column(writer, run.doc_ids, &mut scratch)?;
2619 offset = checked_advance(offset, count * 4)?;
2620 let ordinals_offset = offset;
2621 write_u16_column(writer, run.ordinals, &mut scratch)?;
2622 offset = checked_advance(offset, count * 2)?;
2623 let codes_offset = offset;
2624 writer.write_all(run.codes)?;
2625 offset = checked_advance(offset, run.codes.len())?;
2626 extra_records.push(RunRecord {
2627 cluster_id: run.cluster_id,
2628 doc_base: 0,
2629 count: u32::try_from(count)
2630 .map_err(|_| invalid_data("extra ANN run exceeds u32 vectors"))?,
2631 max_doc_id: run.doc_ids.iter().copied().max().unwrap_or(0),
2632 doc_ids_offset,
2633 ordinals_offset,
2634 codes_offset,
2635 codes_len: u64::try_from(run.codes.len())
2636 .map_err(|_| invalid_data("extra ANN code length exceeds u64"))?,
2637 });
2638 }
2639
2640 let directory_offset = offset;
2646 let mut pending = BinaryHeap::with_capacity(sources.len() + 1);
2647 for (source_index, (source, _)) in sources.iter().enumerate() {
2648 pending.push(Reverse((source.runs[0].cluster_id, source_index, 0usize)));
2649 }
2650 if let Some(first_extra) = extra_records.first() {
2651 pending.push(Reverse((first_extra.cluster_id, sources.len(), 0usize)));
2652 }
2653 let mut written_runs = 0usize;
2654 while let Some(Reverse((_, source_index, run_index))) = pending.pop() {
2655 if source_index == sources.len() {
2656 write_run_record(writer, &extra_records[run_index])?;
2657 written_runs += 1;
2658 if let Some(next) = extra_records.get(run_index + 1) {
2659 pending.push(Reverse((next.cluster_id, source_index, run_index + 1)));
2660 }
2661 continue;
2662 }
2663 let (source, segment_base) = sources[source_index];
2664 let run = &source.runs[run_index];
2665 write_run_record(
2666 writer,
2667 &RunRecord {
2668 cluster_id: run.cluster_id,
2669 doc_base: run
2670 .doc_base
2671 .checked_add(segment_base)
2672 .ok_or_else(|| invalid_data("merged ANN document base overflows u32"))?,
2673 count: u32::try_from(run.count)
2674 .map_err(|_| invalid_data("ANN source run exceeds u32 vectors"))?,
2675 max_doc_id: run.max_doc_id,
2676 doc_ids_offset: relocate_payload_offset(
2677 output_payload_starts[source_index],
2678 run.doc_ids.start,
2679 )?,
2680 ordinals_offset: relocate_payload_offset(
2681 output_payload_starts[source_index],
2682 run.ordinals.start,
2683 )?,
2684 codes_offset: relocate_payload_offset(
2685 output_payload_starts[source_index],
2686 run.codes.start,
2687 )?,
2688 codes_len: u64::try_from(run.codes.len())
2689 .map_err(|_| invalid_data("ANN source code length exceeds u64"))?,
2690 },
2691 )?;
2692 written_runs = written_runs
2693 .checked_add(1)
2694 .ok_or_else(|| invalid_data("merged ANN run count overflows usize"))?;
2695 let next_run_index = run_index + 1;
2696 if let Some(next_run) = source.runs.get(next_run_index) {
2697 pending.push(Reverse((next_run.cluster_id, source_index, next_run_index)));
2698 }
2699 }
2700 if written_runs != run_capacity {
2701 return Err(invalid_data("merged ANN directory lost source runs"));
2702 }
2703 finish_footer(writer, directory_offset, written_runs)
2704}
2705
2706#[cfg(feature = "native")]
2707fn relocate_payload_offset(output_payload_start: u64, source_offset: usize) -> io::Result<u64> {
2708 let relative = source_offset
2709 .checked_sub(ANN_HEADER_SIZE)
2710 .ok_or_else(|| invalid_data("ANN source offset precedes its payload"))?;
2711 output_payload_start
2712 .checked_add(
2713 u64::try_from(relative)
2714 .map_err(|_| invalid_data("ANN source relative offset exceeds u64"))?,
2715 )
2716 .ok_or_else(|| invalid_data("merged ANN payload offset overflows u64"))
2717}
2718
2719#[cfg(feature = "native")]
2720fn headers_compatible(left: &AnnDiskHeader, right: &AnnDiskHeader) -> bool {
2721 left.kind == right.kind
2722 && left.routing == right.routing
2723 && left.dim == right.dim
2724 && left.code_size == right.code_size
2725 && left.num_clusters == right.num_clusters
2726 && left.quantizer_version == right.quantizer_version
2727 && left.codebook_version == right.codebook_version
2728}
2729
2730#[cfg(feature = "native")]
2731fn finish_layout(
2732 writer: &mut (impl Write + ?Sized),
2733 directory_offset: u64,
2734 records: &[RunRecord],
2735) -> io::Result<u64> {
2736 for record in records {
2737 write_run_record(writer, record)?;
2738 }
2739 finish_footer(writer, directory_offset, records.len())
2740}
2741
2742#[cfg(feature = "native")]
2743fn write_run_record(writer: &mut (impl Write + ?Sized), record: &RunRecord) -> io::Result<()> {
2744 writer.write_u32::<LittleEndian>(record.cluster_id)?;
2745 writer.write_u32::<LittleEndian>(record.doc_base)?;
2746 writer.write_u32::<LittleEndian>(record.count)?;
2747 writer.write_u32::<LittleEndian>(record.max_doc_id)?;
2748 writer.write_u64::<LittleEndian>(record.doc_ids_offset)?;
2749 writer.write_u64::<LittleEndian>(record.ordinals_offset)?;
2750 writer.write_u64::<LittleEndian>(record.codes_offset)?;
2751 writer.write_u64::<LittleEndian>(record.codes_len)?;
2752 Ok(())
2753}
2754
2755#[cfg(feature = "native")]
2756fn finish_footer(
2757 writer: &mut (impl Write + ?Sized),
2758 directory_offset: u64,
2759 num_records: usize,
2760) -> io::Result<u64> {
2761 writer.write_u64::<LittleEndian>(directory_offset)?;
2762 writer.write_u64::<LittleEndian>(
2763 u64::try_from(num_records).map_err(|_| invalid_data("ANN run count exceeds u64"))?,
2764 )?;
2765 writer.write_u32::<LittleEndian>(ANN_FOOTER_MAGIC)?;
2766 writer.write_u32::<LittleEndian>(u32::from(ANN_DISK_VERSION))?;
2767 let tail_size = num_records
2768 .checked_mul(ANN_RUN_SIZE)
2769 .and_then(|size| size.checked_add(ANN_FOOTER_SIZE))
2770 .and_then(|size| u64::try_from(size).ok())
2771 .ok_or_else(|| invalid_data("ANN final tail size overflows u64"))?;
2772 directory_offset
2773 .checked_add(tail_size)
2774 .ok_or_else(|| invalid_data("ANN final size overflows u64"))
2775}
2776
2777#[cfg(feature = "native")]
2778fn write_header(writer: &mut (impl Write + ?Sized), header: &AnnDiskHeader) -> io::Result<()> {
2779 writer.write_u32::<LittleEndian>(ANN_HEADER_MAGIC)?;
2780 writer.write_u8(header.kind as u8)?;
2781 writer.write_u8(routing_to_u8(header.routing))?;
2782 writer.write_u16::<LittleEndian>(ANN_DISK_VERSION)?;
2783 writer.write_u32::<LittleEndian>(
2784 u32::try_from(header.dim).map_err(|_| invalid_data("ANN dimension exceeds u32"))?,
2785 )?;
2786 writer.write_u32::<LittleEndian>(
2787 u32::try_from(header.code_size).map_err(|_| invalid_data("ANN code size exceeds u32"))?,
2788 )?;
2789 writer.write_u32::<LittleEndian>(header.num_clusters)?;
2790 writer.write_u32::<LittleEndian>(0)?;
2791 writer.write_u64::<LittleEndian>(header.quantizer_version)?;
2792 writer.write_u64::<LittleEndian>(header.codebook_version)?;
2793 writer.write_u64::<LittleEndian>(
2794 u64::try_from(header.vector_count)
2795 .map_err(|_| invalid_data("ANN vector count exceeds u64"))?,
2796 )?;
2797 writer.write_u64::<LittleEndian>(0)?;
2798 Ok(())
2799}
2800
2801#[cfg(feature = "native")]
2802fn write_u32_column(
2803 writer: &mut (impl Write + ?Sized),
2804 values: &[u32],
2805 scratch: &mut Vec<u8>,
2806) -> io::Result<()> {
2807 for chunk in values.chunks(64 * 1024) {
2808 scratch.clear();
2809 scratch.reserve(chunk.len() * 4);
2810 for value in chunk {
2811 scratch.extend_from_slice(&value.to_le_bytes());
2812 }
2813 writer.write_all(scratch)?;
2814 }
2815 Ok(())
2816}
2817
2818#[cfg(feature = "native")]
2819fn write_u16_column(
2820 writer: &mut (impl Write + ?Sized),
2821 values: &[u16],
2822 scratch: &mut Vec<u8>,
2823) -> io::Result<()> {
2824 for chunk in values.chunks(64 * 1024) {
2825 scratch.clear();
2826 scratch.reserve(chunk.len() * 2);
2827 for value in chunk {
2828 scratch.extend_from_slice(&value.to_le_bytes());
2829 }
2830 writer.write_all(scratch)?;
2831 }
2832 Ok(())
2833}
2834
2835#[cfg(feature = "native")]
2836fn copy_range(
2837 writer: &mut (impl Write + ?Sized),
2838 bytes: &OwnedBytes,
2839 range: Range<usize>,
2840 cancellation: Option<&std::sync::atomic::AtomicBool>,
2841) -> io::Result<()> {
2842 if range.is_empty() {
2843 return Ok(());
2844 }
2845 let range_end = range.end;
2846 let mut chunk_start = range.start;
2847 let first_end = chunk_start.saturating_add(COPY_CHUNK).min(range_end);
2848 bytes.madvise_range(chunk_start..first_end, libc::MADV_WILLNEED);
2849 while chunk_start < range_end {
2850 if cancellation
2851 .is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Relaxed))
2852 {
2853 return Err(io::Error::new(
2854 io::ErrorKind::Interrupted,
2855 "ANN merge copy cancelled",
2856 ));
2857 }
2858 let chunk_end = chunk_start.saturating_add(COPY_CHUNK).min(range_end);
2859 let next_end = chunk_end.saturating_add(COPY_CHUNK).min(range_end);
2860 if chunk_end < next_end {
2861 bytes.madvise_range(chunk_end..next_end, libc::MADV_WILLNEED);
2864 }
2865 writer.write_all(&bytes.as_slice()[chunk_start..chunk_end])?;
2866 chunk_start = chunk_end;
2867 }
2868 Ok(())
2869}
2870
2871#[cfg(feature = "native")]
2872fn checked_advance(offset: u64, length: usize) -> io::Result<u64> {
2873 offset
2874 .checked_add(
2875 u64::try_from(length).map_err(|_| invalid_data("ANN copy length exceeds u64"))?,
2876 )
2877 .ok_or_else(|| invalid_data("ANN output offset overflows u64"))
2878}
2879
2880fn validate_header(header: &AnnDiskHeader) -> io::Result<()> {
2881 if header.kind == AnnKind::IvfTq
2882 && !crate::structures::is_ivf_tq_cosine_generation(header.quantizer_version)
2883 {
2884 return Err(invalid_data(
2885 "IVF-TQ payload uses an unsupported legacy generation; rebuild the index",
2886 ));
2887 }
2888 if header.dim == 0
2889 || header.code_size == 0
2890 || header.num_clusters == 0
2891 || header.quantizer_version == 0
2892 || header.vector_count == 0
2893 || (header.kind == AnnKind::BinaryIvf
2894 && (header.codebook_version != 0
2895 || !header.dim.is_multiple_of(8)
2896 || header.code_size != header.dim.div_ceil(8)))
2897 || (header.kind == AnnKind::ScannBinary
2898 && (header.codebook_version == 0
2899 || header.routing != IvfRoutingMode::Flat
2900 || !header.dim.is_multiple_of(8)
2901 || header.code_size != header.dim / 8))
2902 || (header.kind == AnnKind::TqFlat
2903 && (header.codebook_version != 0
2904 || header.num_clusters != 1
2905 || header.routing != IvfRoutingMode::Flat
2906 || header.code_size * 2
2907 != crate::structures::vector::quantization::tq_padded_dim(header.dim)))
2908 || (header.kind == AnnKind::IvfTq
2911 && (header.codebook_version == 0
2912 || header.code_size * 2
2913 != crate::structures::vector::quantization::tq_padded_dim(header.dim)))
2914 || (header.kind == AnnKind::ScannAh
2915 && (header.codebook_version == 0
2916 || header.routing != IvfRoutingMode::Flat
2917 || u16::try_from(header.code_size).is_err()
2918 || u32::try_from(header.dim).is_err()
2919 || crate::structures::vector::scann::ScannEncoding::AsymmetricHash {
2920 dimensions_per_block: header.code_size as u16,
2921 bits_per_code: 4,
2922 }
2923 .row_code_bytes(u32::try_from(header.dim).unwrap_or(0))
2924 .is_err()))
2925 {
2926 return Err(invalid_data("ANN header contains invalid metadata"));
2927 }
2928 Ok(())
2929}
2930
2931fn read_u32(bytes: &[u8], offset: usize) -> u32 {
2932 u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
2933}
2934
2935fn read_u16(bytes: &[u8], offset: usize) -> u16 {
2936 u16::from_le_bytes(bytes[offset..offset + 2].try_into().unwrap())
2937}
2938
2939#[inline]
2940fn tq_ivf_block_max_scale(block: &[u8]) -> f32 {
2941 f32::from_le_bytes(
2945 block[..size_of::<f32>()]
2946 .try_into()
2947 .expect("scale is one f32"),
2948 )
2949}
2950
2951fn run_doc_id(bytes: &[u8], run: &AnnRun, index: usize) -> io::Result<u32> {
2952 let local_doc_id = read_u32(bytes, run.doc_ids.start + index * 4);
2953 if local_doc_id > run.max_doc_id {
2954 return Err(invalid_data(
2955 "ANN run contains a document above its declared maximum",
2956 ));
2957 }
2958 run.doc_base
2959 .checked_add(local_doc_id)
2960 .ok_or_else(|| invalid_data("ANN run document ID overflows u32"))
2961}
2962
2963#[cfg(feature = "native")]
2964fn routing_to_u8(routing: IvfRoutingMode) -> u8 {
2965 match routing {
2966 IvfRoutingMode::Auto => 0,
2967 IvfRoutingMode::Flat => 1,
2968 IvfRoutingMode::TwoLevel => 2,
2969 IvfRoutingMode::Hnsw => 3,
2970 }
2971}
2972
2973fn routing_from_u8(value: u8) -> io::Result<IvfRoutingMode> {
2974 match value {
2975 0 => Ok(IvfRoutingMode::Auto),
2976 1 => Ok(IvfRoutingMode::Flat),
2977 2 => Ok(IvfRoutingMode::TwoLevel),
2978 3 => Ok(IvfRoutingMode::Hnsw),
2979 _ => Err(invalid_data(format!("unknown ANN routing mode {value}"))),
2980 }
2981}
2982
2983fn invalid_data(message: impl Into<String>) -> io::Error {
2984 io::Error::new(io::ErrorKind::InvalidData, message.into())
2985}
2986
2987#[cfg(all(test, feature = "native"))]
2988mod tests {
2989 use super::*;
2990
2991 #[test]
2996 fn compacted_merge_matches_byte_copy_and_resets_fragmentation() {
2997 let a0_docs = [0u32, 1, 2];
3001 let a0_ords = [0u16, 0, 1];
3002 let a0_codes = [0x11u8, 0x22, 0x33];
3003 let a5_docs = [3u32, 4];
3004 let a5_ords = [0u16; 2];
3005 let a5_codes = [0x44u8, 0x55];
3006 let a_runs = [
3007 BuildRun {
3008 cluster_id: 0,
3009 doc_ids: &a0_docs,
3010 ordinals: &a0_ords,
3011 codes: &a0_codes,
3012 },
3013 BuildRun {
3014 cluster_id: 5,
3015 doc_ids: &a5_docs,
3016 ordinals: &a5_ords,
3017 codes: &a5_codes,
3018 },
3019 ];
3020 let mut header = binary_header(5);
3021 header.num_clusters = 8;
3022 let mut a_bytes = Vec::new();
3023 write_built_runs(header.clone(), &a_runs, &mut a_bytes).unwrap();
3024 let a = AnnDiskIndex::open(OwnedBytes::new(a_bytes), AnnKind::BinaryIvf, 5).unwrap();
3025
3026 let b0_docs = [0u32, 2];
3027 let b0_ords = [1u16, 0];
3028 let b0_codes = [0x66u8, 0x77];
3029 let b2_docs = [1u32];
3030 let b2_ords = [0u16];
3031 let b2_codes = [0x88u8];
3032 let b_runs = [
3033 BuildRun {
3034 cluster_id: 0,
3035 doc_ids: &b0_docs,
3036 ordinals: &b0_ords,
3037 codes: &b0_codes,
3038 },
3039 BuildRun {
3040 cluster_id: 2,
3041 doc_ids: &b2_docs,
3042 ordinals: &b2_ords,
3043 codes: &b2_codes,
3044 },
3045 ];
3046 let mut header_b = binary_header(3);
3047 header_b.num_clusters = 8;
3048 let mut b_bytes = Vec::new();
3049 write_built_runs(header_b, &b_runs, &mut b_bytes).unwrap();
3050 let b = AnnDiskIndex::open(OwnedBytes::new(b_bytes), AnnKind::BinaryIvf, 3).unwrap();
3051
3052 let mut gen1_bytes = Vec::new();
3054 write_merged_ann(&[(&a, 0), (&b, 5)], &mut gen1_bytes).unwrap();
3055 let gen1 = AnnDiskIndex::open(OwnedBytes::new(gen1_bytes), AnnKind::BinaryIvf, 8).unwrap();
3056
3057 let sources: [(&AnnDiskIndex, u32); 2] = [(&gen1, 0), (&a, 8)];
3059 let predicted = predicted_merge_fragmentation(&sources);
3060 assert!((predicted - 2.0).abs() < 1e-9, "{predicted}");
3062
3063 let mut copied_bytes = Vec::new();
3064 write_merged_ann(&sources, &mut copied_bytes).unwrap();
3065 let copied =
3066 AnnDiskIndex::open(OwnedBytes::new(copied_bytes), AnnKind::BinaryIvf, 13).unwrap();
3067 let mut compacted_bytes = Vec::new();
3068 write_compacted_ann_cancellable(&sources, &mut compacted_bytes, None).unwrap();
3069 let compacted =
3070 AnnDiskIndex::open(OwnedBytes::new(compacted_bytes), AnnKind::BinaryIvf, 13).unwrap();
3071
3072 let copied_health = copied.health();
3074 let compacted_health = compacted.health();
3075 assert!((copied_health.fragmentation() - 2.0).abs() < 1e-9);
3076 assert!((compacted_health.fragmentation() - 1.0).abs() < 1e-9);
3077 assert_eq!(compacted_health.runs, 3, "one run per non-empty cluster");
3078 assert_eq!(copied_health.vectors, compacted_health.vectors);
3079 assert_eq!(copied_health.payload_bytes, compacted_health.payload_bytes);
3080 assert_eq!(
3081 copied_health.largest_cluster_vectors,
3082 compacted_health.largest_cluster_vectors
3083 );
3084
3085 for cluster in 0..8u32 {
3087 let query = [0x5Au8];
3088 let from_copy = copied
3089 .search_binary_clusters::<false>(&query, 16, &[cluster])
3090 .unwrap();
3091 let from_compact = compacted
3092 .search_binary_clusters::<false>(&query, 16, &[cluster])
3093 .unwrap();
3094 assert_eq!(from_copy, from_compact, "cluster {cluster} diverged");
3095 }
3096
3097 assert!(compacted.runs.iter().all(|run| run.doc_base == 0));
3099
3100 let mut generation3 = Vec::new();
3103 write_merged_ann(&[(&compacted, 0), (&b, 13)], &mut generation3).unwrap();
3104 let generation3 =
3105 AnnDiskIndex::open(OwnedBytes::new(generation3), AnnKind::BinaryIvf, 16).unwrap();
3106 assert_eq!(generation3.health().vectors, 16);
3107 let all: Vec<(u32, u16, f32)> = compacted
3109 .search_binary_clusters::<false>(&[0x5A], 32, &[0, 2, 5])
3110 .unwrap();
3111 let mut docs: Vec<u32> = all.iter().map(|&(doc, _, _)| doc).collect();
3112 docs.sort_unstable();
3113 assert_eq!(docs, (0..=12).collect::<Vec<u32>>());
3114 }
3115
3116 #[test]
3117 fn compacted_scann_ah_repacks_blocks_across_run_boundaries() {
3118 use crate::structures::vector::scann::{FAST_SCAN_LANES, pack_fast_scan_block};
3119
3120 const DIM: usize = 8;
3121 const DIMS_PER_BLOCK: usize = 2;
3122 const BLOCKS: usize = DIM / DIMS_PER_BLOCK;
3123
3124 let make_rows = |start: usize, count: usize| -> Vec<u8> {
3125 (0..count)
3126 .flat_map(|row| {
3127 (0..BLOCKS).map(move |block| ((start + row * 3 + block * 5) & 0x0f) as u8)
3128 })
3129 .collect()
3130 };
3131 let encode = |rows: &[u8]| -> Vec<u8> {
3132 let count = rows.len() / BLOCKS;
3133 let full_rows = count / FAST_SCAN_LANES * FAST_SCAN_LANES;
3134 let mut codes = Vec::new();
3135 for group in rows[..full_rows * BLOCKS].chunks_exact(FAST_SCAN_LANES * BLOCKS) {
3136 pack_fast_scan_block(group, BLOCKS, &mut codes).unwrap();
3137 }
3138 for row in rows[full_rows * BLOCKS..].chunks_exact(BLOCKS) {
3139 for pair in row.chunks(2) {
3140 codes.push(pair[0] | (pair.get(1).copied().unwrap_or(0) << 4));
3141 }
3142 }
3143 codes
3144 };
3145 let make_source = |rows: &[u8]| {
3146 let count = rows.len() / BLOCKS;
3147 let docs: Vec<u32> = (0..count as u32).collect();
3148 let ordinals = vec![0u16; count];
3149 let codes = encode(rows);
3150 let run = BuildRun {
3151 cluster_id: 0,
3152 doc_ids: &docs,
3153 ordinals: &ordinals,
3154 codes: &codes,
3155 };
3156 let header = AnnDiskHeader {
3157 kind: AnnKind::ScannAh,
3158 routing: IvfRoutingMode::Flat,
3159 dim: DIM,
3160 code_size: DIMS_PER_BLOCK,
3161 num_clusters: 1,
3162 quantizer_version: 41,
3163 codebook_version: 73,
3164 vector_count: count,
3165 };
3166 let mut bytes = Vec::new();
3167 write_built_runs(header, &[run], &mut bytes).unwrap();
3168 AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::ScannAh, count as u32).unwrap()
3169 };
3170
3171 let left_rows = make_rows(1, 35);
3174 let right_rows = make_rows(9, 30);
3175 let left = make_source(&left_rows);
3176 let right = make_source(&right_rows);
3177 let mut bytes = Vec::new();
3178 write_compacted_ann_cancellable(&[(&left, 0), (&right, 35)], &mut bytes, None).unwrap();
3179 let compacted = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::ScannAh, 65).unwrap();
3180
3181 assert_eq!(compacted.runs.len(), 1);
3182 assert!((compacted.health().fragmentation() - 1.0).abs() < 1e-9);
3183 let run = &compacted.runs[0];
3184 let codes = &compacted.raw.as_slice()[run.codes.clone()];
3185 let mut decoded = Vec::new();
3186 for row in 0..run.count {
3187 unpack_scann_ah_row(codes, run.count, BLOCKS, row, &mut decoded).unwrap();
3188 }
3189 let mut expected = left_rows;
3190 expected.extend_from_slice(&right_rows);
3191 assert_eq!(decoded, expected);
3192 for row in 0..65 {
3193 assert_eq!(
3194 run_doc_id(compacted.raw.as_slice(), run, row).unwrap(),
3195 row as u32
3196 );
3197 }
3198 }
3199
3200 #[test]
3203 #[ignore]
3204 fn ann_merge_throughput_byte_copy_vs_compaction() {
3205 let code_size = 320usize;
3206 let clusters = 4_096u32;
3207 let vectors_per_source = 262_144usize;
3208 let sources_count = 4usize;
3209
3210 let mut sources_bytes = Vec::new();
3211 for source_index in 0..sources_count {
3212 let mut per_cluster: Vec<(Vec<u32>, Vec<u16>, Vec<u8>)> = Vec::new();
3213 let vectors_per_cluster = vectors_per_source / clusters as usize;
3214 let mut doc = 0u32;
3215 for cluster in 0..clusters {
3216 let mut docs = Vec::with_capacity(vectors_per_cluster);
3217 let mut ords = Vec::with_capacity(vectors_per_cluster);
3218 let mut codes = Vec::with_capacity(vectors_per_cluster * code_size);
3219 for _ in 0..vectors_per_cluster {
3220 docs.push(doc);
3221 ords.push(0u16);
3222 codes.extend(std::iter::repeat_n(
3223 (doc ^ cluster ^ source_index as u32) as u8,
3224 code_size,
3225 ));
3226 doc += 1;
3227 }
3228 per_cluster.push((docs, ords, codes));
3229 }
3230 let runs: Vec<BuildRun<'_>> = per_cluster
3231 .iter()
3232 .enumerate()
3233 .map(|(cluster, (docs, ords, codes))| BuildRun {
3234 cluster_id: cluster as u32,
3235 doc_ids: docs,
3236 ordinals: ords,
3237 codes,
3238 })
3239 .collect();
3240 let header = AnnDiskHeader {
3241 kind: AnnKind::BinaryIvf,
3242 routing: IvfRoutingMode::Hnsw,
3243 dim: code_size * 8,
3244 code_size,
3245 num_clusters: clusters,
3246 quantizer_version: 42,
3247 codebook_version: 0,
3248 vector_count: vectors_per_source,
3249 };
3250 let mut bytes = Vec::new();
3251 write_built_runs(header, &runs, &mut bytes).unwrap();
3252 sources_bytes.push(bytes);
3253 }
3254 let sources_open: Vec<AnnDiskIndex> = sources_bytes
3255 .iter()
3256 .map(|bytes| {
3257 AnnDiskIndex::open(
3258 OwnedBytes::new(bytes.clone()),
3259 AnnKind::BinaryIvf,
3260 (vectors_per_source * sources_count) as u32,
3261 )
3262 .unwrap()
3263 })
3264 .collect();
3265 let sources: Vec<(&AnnDiskIndex, u32)> = sources_open
3266 .iter()
3267 .enumerate()
3268 .map(|(index, source)| (source, (index * vectors_per_source) as u32))
3269 .collect();
3270 let payload_bytes = sources_bytes.iter().map(Vec::len).sum::<usize>();
3271
3272 let mut out = vec![0u8; payload_bytes + (1 << 20)];
3278 out.clear();
3279 let mut copy_secs = f64::INFINITY;
3282 let mut compact_secs = f64::INFINITY;
3283 let mut compacted_bytes = Vec::new();
3284 for _ in 0..3 {
3285 out.clear();
3286 let start = std::time::Instant::now();
3287 write_merged_ann(&sources, &mut out).unwrap();
3288 copy_secs = copy_secs.min(start.elapsed().as_secs_f64());
3289
3290 out.clear();
3291 let start = std::time::Instant::now();
3292 write_compacted_ann_cancellable(&sources, &mut out, None).unwrap();
3293 compact_secs = compact_secs.min(start.elapsed().as_secs_f64());
3294 compacted_bytes = out.clone();
3295 }
3296 let compacted = AnnDiskIndex::open(
3297 OwnedBytes::new(compacted_bytes),
3298 AnnKind::BinaryIvf,
3299 (vectors_per_source * sources_count) as u32,
3300 )
3301 .unwrap();
3302 assert!((compacted.health().fragmentation() - 1.0).abs() < 1e-9);
3303
3304 let gib = payload_bytes as f64 / (1u64 << 30) as f64;
3305 println!(
3306 "ann merge {:.2} GiB: byte-copy {:.3}s ({:.2} GiB/s), compaction {:.3}s \
3307 ({:.2} GiB/s), overhead {:.1}%",
3308 gib,
3309 copy_secs,
3310 gib / copy_secs,
3311 compact_secs,
3312 gib / compact_secs,
3313 100.0 * (compact_secs - copy_secs) / copy_secs,
3314 );
3315 }
3316
3317 #[test]
3319 fn compaction_refuses_non_binary_payloads() {
3320 let codec = std::sync::Arc::new(crate::structures::TqCodec::new(8));
3326 let mut builder = crate::structures::TqFlatBuilder::new(codec);
3327 builder
3328 .add_batch(
3329 &[(0, 0), (1, 0)],
3330 &[
3331 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,
3333 ],
3334 )
3335 .unwrap();
3336 builder.finish();
3337 let mut bytes = Vec::new();
3338 write_built_tq_flat(&builder, &mut bytes).unwrap();
3339 let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::TqFlat, 2).unwrap();
3340 let error = write_compacted_ann_cancellable(&[(&disk, 0)], &mut Vec::new(), None)
3341 .expect_err("TQ payloads must not compact");
3342 assert!(error.to_string().contains("binary"), "{error}");
3343 }
3344
3345 #[test]
3349 fn ann_health_measures_skew_and_fragmentation() {
3350 let a0_docs = [0u32, 1, 2, 3, 4, 5];
3355 let a0_ords = [0u16; 6];
3356 let a0_codes = [0xAAu8; 6];
3357 let a3_docs = [6u32, 7];
3358 let a3_ords = [0u16; 2];
3359 let a3_codes = [0x0Fu8; 2];
3360 let a_runs = [
3361 BuildRun {
3362 cluster_id: 0,
3363 doc_ids: &a0_docs,
3364 ordinals: &a0_ords,
3365 codes: &a0_codes,
3366 },
3367 BuildRun {
3368 cluster_id: 3,
3369 doc_ids: &a3_docs,
3370 ordinals: &a3_ords,
3371 codes: &a3_codes,
3372 },
3373 ];
3374 let mut header = binary_header(8);
3375 header.num_clusters = 8;
3376 let mut a_bytes = Vec::new();
3377 write_built_runs(header, &a_runs, &mut a_bytes).unwrap();
3378 let a = AnnDiskIndex::open(OwnedBytes::new(a_bytes), AnnKind::BinaryIvf, 8).unwrap();
3379
3380 let b0_docs = [0u32, 1];
3381 let b0_ords = [0u16; 2];
3382 let b0_codes = [0xBBu8; 2];
3383 let b_runs = [BuildRun {
3384 cluster_id: 0,
3385 doc_ids: &b0_docs,
3386 ordinals: &b0_ords,
3387 codes: &b0_codes,
3388 }];
3389 let mut header_b = binary_header(2);
3390 header_b.num_clusters = 8;
3391 let mut b_bytes = Vec::new();
3392 write_built_runs(header_b, &b_runs, &mut b_bytes).unwrap();
3393 let b = AnnDiskIndex::open(OwnedBytes::new(b_bytes), AnnKind::BinaryIvf, 2).unwrap();
3394
3395 let mut merged_bytes = Vec::new();
3396 write_merged_ann(&[(&a, 0), (&b, 8)], &mut merged_bytes).unwrap();
3397 let disk =
3398 AnnDiskIndex::open(OwnedBytes::new(merged_bytes), AnnKind::BinaryIvf, 10).unwrap();
3399
3400 let health = disk.health();
3401 assert_eq!(health.vectors, 10);
3402 assert_eq!(health.clusters_nonempty, 2);
3403 assert_eq!(health.clusters_total, 8);
3404 assert_eq!(health.runs, 3);
3405 assert_eq!(health.largest_cluster, 0);
3406 assert_eq!(health.largest_cluster_vectors, 8);
3407 assert!((health.largest_cluster_share() - 0.8).abs() < 1e-9);
3408 assert!((health.fragmentation() - 1.5).abs() < 1e-9);
3410 assert!(
3412 (health.imbalance - 1.36).abs() < 1e-9,
3413 "{}",
3414 health.imbalance
3415 );
3416 assert_eq!(health.payload_bytes, 10);
3418 }
3419
3420 #[test]
3421 fn ann_health_is_balanced_at_one() {
3422 let docs: Vec<Vec<u32>> = (0..4).map(|c| vec![c * 2, c * 2 + 1]).collect();
3423 let ords = [0u16; 2];
3424 let codes = [0x55u8; 2];
3425 let runs: Vec<BuildRun<'_>> = docs
3426 .iter()
3427 .enumerate()
3428 .map(|(cluster, doc_ids)| BuildRun {
3429 cluster_id: cluster as u32,
3430 doc_ids,
3431 ordinals: &ords,
3432 codes: &codes,
3433 })
3434 .collect();
3435 let mut header = binary_header(8);
3436 header.num_clusters = 4;
3437 let mut bytes = Vec::new();
3438 write_built_runs(header, &runs, &mut bytes).unwrap();
3439 let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::BinaryIvf, 8).unwrap();
3440 let health = disk.health();
3441 assert!((health.imbalance - 1.0).abs() < 1e-9);
3442 assert!((health.fragmentation() - 1.0).abs() < 1e-9);
3443 assert!((health.largest_cluster_share() - 0.25).abs() < 1e-9);
3444 }
3445
3446 fn binary_header(vector_count: usize) -> AnnDiskHeader {
3447 AnnDiskHeader {
3448 kind: AnnKind::BinaryIvf,
3449 routing: IvfRoutingMode::Hnsw,
3450 dim: 8,
3451 code_size: 1,
3452 num_clusters: 2,
3453 quantizer_version: 42,
3454 codebook_version: 0,
3455 vector_count,
3456 }
3457 }
3458
3459 fn payload_end(index: &AnnDiskIndex) -> usize {
3460 index.runs.iter().map(|run| run.codes.end).max().unwrap()
3461 }
3462
3463 #[test]
3464 fn ann_prefetch_ranges_are_sorted_and_only_merge_page_near_extents() {
3465 let mut ranges = vec![
3466 15_000..16_000,
3467 0..1_000,
3468 9_000..10_000,
3469 1_000..2_000,
3470 7_000..8_000,
3471 ];
3472 coalesce_prefetch_ranges(&mut ranges);
3473 assert_eq!(ranges, [0..2_000, 7_000..10_000, 15_000..16_000]);
3474 }
3475
3476 #[test]
3477 fn normal_merge_copies_ann_payload_columns_byte_for_byte() {
3478 let first_doc_0 = [0u32];
3479 let first_doc_1 = [1u32];
3480 let first_ord_0 = [0u16];
3481 let first_ord_1 = [2u16];
3482 let first_code_0 = [0x00u8];
3483 let first_code_1 = [0xffu8];
3484 let first_runs = [
3485 BuildRun {
3486 cluster_id: 0,
3487 doc_ids: &first_doc_0,
3488 ordinals: &first_ord_0,
3489 codes: &first_code_0,
3490 },
3491 BuildRun {
3492 cluster_id: 1,
3493 doc_ids: &first_doc_1,
3494 ordinals: &first_ord_1,
3495 codes: &first_code_1,
3496 },
3497 ];
3498 let mut first_bytes = Vec::new();
3499 write_built_runs(binary_header(2), &first_runs, &mut first_bytes).unwrap();
3500 let first = AnnDiskIndex::open(OwnedBytes::new(first_bytes.clone()), AnnKind::BinaryIvf, 2)
3501 .unwrap();
3502
3503 let second_docs = [0u32, 1u32];
3504 let second_ords = [1u16, 0u16];
3505 let second_codes = [0x0fu8, 0xf0u8];
3506 let second_runs = [BuildRun {
3507 cluster_id: 0,
3508 doc_ids: &second_docs,
3509 ordinals: &second_ords,
3510 codes: &second_codes,
3511 }];
3512 let mut second_bytes = Vec::new();
3513 write_built_runs(binary_header(2), &second_runs, &mut second_bytes).unwrap();
3514 let second =
3515 AnnDiskIndex::open(OwnedBytes::new(second_bytes.clone()), AnnKind::BinaryIvf, 2)
3516 .unwrap();
3517
3518 let mut merged_bytes = Vec::new();
3519 write_merged_ann(&[(&first, 0), (&second, 2)], &mut merged_bytes).unwrap();
3520 let merged =
3521 AnnDiskIndex::open(OwnedBytes::new(merged_bytes.clone()), AnnKind::BinaryIvf, 4)
3522 .unwrap();
3523
3524 let mut expected_payload = first_bytes[ANN_HEADER_SIZE..payload_end(&first)].to_vec();
3525 expected_payload.extend_from_slice(&second_bytes[ANN_HEADER_SIZE..payload_end(&second)]);
3526 assert_eq!(
3527 &merged_bytes[ANN_HEADER_SIZE..payload_end(&merged)],
3528 expected_payload.as_slice(),
3529 "normal merge must not decode or rewrite any corpus-sized ANN column",
3530 );
3531
3532 let mut docs: Vec<u32> = merged
3533 .search_binary_clusters::<false>(&[0], 4, &[0, 1])
3534 .unwrap()
3535 .into_iter()
3536 .map(|result| result.0)
3537 .collect();
3538 docs.sort_unstable();
3539 assert_eq!(docs, [0, 1, 2, 3]);
3540 let serial = merged
3541 .search_binary_clusters_with_tuning::<false>(&[0], 4, &[0, 1], usize::MAX)
3542 .unwrap();
3543 let parallel = rayon::ThreadPoolBuilder::new()
3544 .num_threads(4)
3545 .build()
3546 .unwrap()
3547 .install(|| merged.search_binary_clusters_with_tuning::<false>(&[0], 4, &[0, 1], 1))
3548 .unwrap();
3549 assert_eq!(parallel, serial, "parallel binary top-k changed results");
3550
3551 let mut second_merge_bytes = Vec::new();
3555 write_merged_ann(&[(&merged, 0), (&first, 4)], &mut second_merge_bytes).unwrap();
3556 let second_merge = AnnDiskIndex::open(
3557 OwnedBytes::new(second_merge_bytes.clone()),
3558 AnnKind::BinaryIvf,
3559 6,
3560 )
3561 .unwrap();
3562 let mut expected_second_payload =
3563 merged_bytes[ANN_HEADER_SIZE..payload_end(&merged)].to_vec();
3564 expected_second_payload
3565 .extend_from_slice(&first_bytes[ANN_HEADER_SIZE..payload_end(&first)]);
3566 assert_eq!(
3567 &second_merge_bytes[ANN_HEADER_SIZE..payload_end(&second_merge)],
3568 expected_second_payload.as_slice(),
3569 );
3570 let mut docs: Vec<u32> = second_merge
3571 .search_binary_clusters::<false>(&[0], 6, &[0, 1])
3572 .unwrap()
3573 .into_iter()
3574 .map(|result| result.0)
3575 .collect();
3576 docs.sort_unstable();
3577 assert_eq!(docs, [0, 1, 2, 3, 4, 5]);
3578 }
3579
3580 #[test]
3581 fn legacy_ivf_tq_payload_is_rejected_while_opening() {
3582 let dim = 8;
3583 let marked_version = crate::structures::mark_ivf_tq_cosine_generation(7);
3584 let centroids = crate::structures::CoarseCentroids {
3585 num_clusters: 1,
3586 dim,
3587 centroids: vec![0.0; dim],
3588 version: marked_version,
3589 soar_config: None,
3590 routing_index: None,
3591 };
3592 let mut bytes = crate::segment::ann_build::build_ivf_tq(
3593 dim,
3594 IvfRoutingMode::Flat,
3595 ¢roids,
3596 &[(0, 0)],
3597 &[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
3598 )
3599 .unwrap();
3600
3601 bytes[24..32].copy_from_slice(&7u64.to_le_bytes());
3603 let error = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::IvfTq, 1)
3604 .err()
3605 .expect("legacy IVF-TQ payload must fail while opening")
3606 .to_string();
3607 assert!(error.contains("unsupported legacy generation"), "{error}");
3608 }
3609
3610 #[test]
3611 fn binary_combined_search_deduplicates_soar_and_bounds_document_results() {
3612 let cluster_0_docs = [0u32, 0, 1];
3616 let cluster_0_ordinals = [0u16, 1, 0];
3617 let cluster_0_codes = [0x00u8, 0xff, 0x03];
3618 let cluster_1_docs = [0u32, 1, 2];
3619 let cluster_1_ordinals = [0u16, 1, 0];
3620 let cluster_1_codes = [0x00u8, 0x0c, 0xf0];
3621 let runs = [
3622 BuildRun {
3623 cluster_id: 0,
3624 doc_ids: &cluster_0_docs,
3625 ordinals: &cluster_0_ordinals,
3626 codes: &cluster_0_codes,
3627 },
3628 BuildRun {
3629 cluster_id: 1,
3630 doc_ids: &cluster_1_docs,
3631 ordinals: &cluster_1_ordinals,
3632 codes: &cluster_1_codes,
3633 },
3634 ];
3635 let mut bytes = Vec::new();
3636 write_built_runs(binary_header(6), &runs, &mut bytes).unwrap();
3637 let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::BinaryIvf, 3).unwrap();
3638 let parallel_pool = rayon::ThreadPoolBuilder::new()
3639 .num_threads(4)
3640 .build()
3641 .unwrap();
3642 let serial_documents = disk
3643 .search_binary_clusters_with_tuning::<true>(&[0], 3, &[0, 1], usize::MAX)
3644 .unwrap();
3645 let parallel_documents = parallel_pool
3646 .install(|| disk.search_binary_clusters_with_tuning::<true>(&[0], 3, &[0, 1], 1))
3647 .unwrap();
3648 assert_eq!(parallel_documents, serial_documents);
3649
3650 let sum = crate::query::MultiValueCombiner::Sum;
3654 let (result, probed) = disk
3655 .search_binary_combined_documents(1, &[0], &[0, 1], sum)
3656 .unwrap();
3657 let parallel_sum = parallel_pool
3658 .install(|| disk.search_binary_combined_documents_with_tuning(1, &[0], &[0, 1], sum, 1))
3659 .unwrap();
3660 assert_eq!(parallel_sum, (result.clone(), probed.clone()));
3661 assert_eq!(result.len(), 1, "combined search must honor k");
3662 assert_eq!(
3663 result[0].doc_id, 1,
3664 "SOAR duplicate changed {sum:?} ranking: {result:?}",
3665 );
3666 assert_eq!(
3669 probed
3670 .iter()
3671 .map(|&(doc_id, ordinal, _)| (doc_id, ordinal))
3672 .collect::<Vec<_>>(),
3673 vec![(1, 0), (1, 1)],
3674 "{sum:?}",
3675 );
3676
3677 let smooth_max = crate::query::MultiValueCombiner::default();
3683 let (result, probed) = disk
3684 .search_binary_combined_documents(1, &[0], &[0, 1], smooth_max)
3685 .unwrap();
3686 let parallel_smooth_max = parallel_pool
3687 .install(|| {
3688 disk.search_binary_combined_documents_with_tuning(1, &[0], &[0, 1], smooth_max, 1)
3689 })
3690 .unwrap();
3691 assert_eq!(parallel_smooth_max, (result.clone(), probed.clone()));
3692 assert_eq!(result.len(), 1, "combined search must honor k");
3693 assert_eq!(result[0].doc_id, 0, "{result:?}");
3694 let expected = smooth_max.combine(&[(0, 1.0), (1, 0.0)]);
3695 assert!(
3696 (result[0].score - expected).abs() < 1e-6,
3697 "SOAR duplicate leaked into {smooth_max:?}: got {}, expected {expected}",
3698 result[0].score,
3699 );
3700 assert_eq!(
3701 probed
3702 .iter()
3703 .map(|&(doc_id, ordinal, _)| (doc_id, ordinal))
3704 .collect::<Vec<_>>(),
3705 vec![(0, 0), (0, 1)],
3706 "{smooth_max:?}",
3707 );
3708
3709 let (top_two, probed) = disk
3710 .search_binary_combined_documents(
3711 2,
3712 &[0],
3713 &[0, 1],
3714 crate::query::MultiValueCombiner::Sum,
3715 )
3716 .unwrap();
3717 assert_eq!(
3718 top_two
3719 .iter()
3720 .map(|candidate| candidate.doc_id)
3721 .collect::<Vec<_>>(),
3722 vec![1, 0],
3723 );
3724 assert_eq!(top_two.len(), 2, "full probing must still return at most k");
3725 assert_eq!(
3728 probed
3729 .iter()
3730 .map(|&(doc_id, ordinal, _)| (doc_id, ordinal))
3731 .collect::<Vec<_>>(),
3732 vec![(0, 0), (0, 1), (1, 0), (1, 1)],
3733 );
3734 }
3735
3736 #[test]
3737 fn combined_ordinal_reduction_handles_out_of_order_runs_for_every_combiner() {
3738 let out_of_order_with_duplicate = vec![
3739 (7, 1, 0.4),
3740 (3, 0, 0.8),
3741 (7, 0, 0.6),
3742 (3, 1, 0.2),
3743 (7, 1, 0.5), ];
3745 for combiner in [
3746 crate::query::MultiValueCombiner::Max,
3747 crate::query::MultiValueCombiner::Sum,
3748 crate::query::MultiValueCombiner::Avg,
3749 crate::query::MultiValueCombiner::default(),
3750 crate::query::MultiValueCombiner::WeightedTopK { k: 2, decay: 0.7 },
3751 ] {
3752 let actual = combine_scored_ordinals(out_of_order_with_duplicate.clone(), 2, combiner);
3753 let mut expected = vec![
3754 AnnDocumentCandidate {
3755 doc_id: 3,
3756 score: combiner.combine(&[(0, 0.8), (1, 0.2)]),
3757 },
3758 AnnDocumentCandidate {
3759 doc_id: 7,
3760 score: combiner.combine(&[(0, 0.6), (1, 0.5)]),
3761 },
3762 ];
3763 expected.sort_unstable_by(|left, right| {
3764 right
3765 .score
3766 .total_cmp(&left.score)
3767 .then_with(|| left.doc_id.cmp(&right.doc_id))
3768 });
3769 assert_eq!(actual, expected, "combiner {combiner:?}");
3770 }
3771 }
3772
3773 fn build_tq_payload(dim: usize, count: usize, seed: u64) -> (Vec<u8>, Vec<Vec<f32>>) {
3774 let codec = std::sync::Arc::new(crate::structures::TqCodec::new(dim));
3775 let mut builder = crate::structures::TqFlatBuilder::new(std::sync::Arc::clone(&codec));
3776 let mut state = seed;
3777 let mut vectors = Vec::new();
3778 let mut flat = Vec::new();
3779 for _ in 0..count {
3780 let vector: Vec<f32> = (0..dim)
3781 .map(|_| {
3782 state = state
3783 .wrapping_mul(6364136223846793005)
3784 .wrapping_add(1442695040888963407);
3785 ((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5
3786 })
3787 .collect();
3788 flat.extend_from_slice(&vector);
3789 vectors.push(vector);
3790 }
3791 let labels: Vec<(u32, u16)> = (0..count).map(|index| (index as u32, 0)).collect();
3792 builder.add_batch(&labels, &flat).unwrap();
3793 builder.finish();
3794 let mut bytes = Vec::new();
3795 write_built_tq_flat(&builder, &mut bytes).unwrap();
3796 (bytes, vectors)
3797 }
3798
3799 #[test]
3800 fn tq_combined_scan_ranks_complete_documents_instead_of_individual_values() {
3801 let dim = 8;
3802 let codec = std::sync::Arc::new(crate::structures::TqCodec::new(dim));
3803 let mut builder = crate::structures::TqFlatBuilder::new(std::sync::Arc::clone(&codec));
3804 let labels = [(0u32, 0u16), (1, 0), (1, 1)];
3805 let vectors = [
3806 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,
3809 ];
3810 builder.add_batch(&labels, &vectors).unwrap();
3811 builder.finish();
3812 let mut bytes = Vec::new();
3813 write_built_tq_flat(&builder, &mut bytes).unwrap();
3814 let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::TqFlat, 2).unwrap();
3815 let query = [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0];
3816 let plan = crate::structures::TqQueryPlan::build(&codec, &query);
3817
3818 let max = disk.search_tq_distinct(1, &plan).unwrap();
3819 let sum = disk
3820 .search_tq_combined_documents(1, &plan, crate::query::MultiValueCombiner::Sum)
3821 .unwrap();
3822 let default_combiner = disk
3823 .search_tq_combined_documents(1, &plan, crate::query::MultiValueCombiner::default())
3824 .unwrap();
3825 assert_eq!(max[0].0, 0, "fixture must favor doc 0 by Max: {max:?}");
3826 assert_eq!(
3827 sum[0].doc_id, 1,
3828 "two complete values must make doc 1 win by Sum: {sum:?}"
3829 );
3830 assert_eq!(
3834 default_combiner[0].doc_id, 0,
3835 "default smooth-max must follow the best value: {default_combiner:?}"
3836 );
3837 assert!(sum[0].score > max[0].2);
3838
3839 let (single_bytes, _) = build_tq_payload(dim, 1, 91);
3843 let single = AnnDiskIndex::open(OwnedBytes::new(single_bytes), AnnKind::TqFlat, 1).unwrap();
3844 let mut merged_bytes = Vec::new();
3845 write_merged_ann(&[(&disk, 1), (&single, 0)], &mut merged_bytes).unwrap();
3846 let merged = AnnDiskIndex::open(OwnedBytes::new(merged_bytes), AnnKind::TqFlat, 3).unwrap();
3847 let merged_sum = merged
3848 .search_tq_combined_documents(1, &plan, crate::query::MultiValueCombiner::Sum)
3849 .unwrap();
3850 assert_eq!(merged_sum[0].doc_id, 2);
3851 }
3852
3853 #[test]
3854 fn tq_payload_roundtrip_search_and_pure_copy_merge() {
3855 let dim = 20; let count = 21;
3857 let (bytes, vectors) = build_tq_payload(dim, count, 42);
3858 let index = AnnDiskIndex::open(
3859 OwnedBytes::new(bytes.clone()),
3860 AnnKind::TqFlat,
3861 count as u32,
3862 )
3863 .unwrap();
3864 assert_eq!(index.header().vector_count, count);
3865
3866 let codec = crate::structures::TqCodec::new(dim);
3868 for target in [0usize, 7, 20] {
3869 let plan = crate::structures::TqQueryPlan::build(&codec, &vectors[target]);
3870 let results = index.search_tq_distinct(3, &plan).unwrap();
3871 assert_eq!(
3872 results[0].0, target as u32,
3873 "query duplicating vector {target} must rank it first: {results:?}"
3874 );
3875 }
3876
3877 let (second_bytes, _) = build_tq_payload(dim, 5, 77);
3879 let second =
3880 AnnDiskIndex::open(OwnedBytes::new(second_bytes.clone()), AnnKind::TqFlat, 5).unwrap();
3881 let mut merged_bytes = Vec::new();
3882 write_merged_ann(&[(&index, 0), (&second, count as u32)], &mut merged_bytes).unwrap();
3883 let merged = AnnDiskIndex::open(
3884 OwnedBytes::new(merged_bytes.clone()),
3885 AnnKind::TqFlat,
3886 count as u32 + 5,
3887 )
3888 .unwrap();
3889 let mut expected_payload = bytes[ANN_HEADER_SIZE..payload_end(&index)].to_vec();
3890 expected_payload.extend_from_slice(&second_bytes[ANN_HEADER_SIZE..payload_end(&second)]);
3891 assert_eq!(
3892 &merged_bytes[ANN_HEADER_SIZE..payload_end(&merged)],
3893 expected_payload.as_slice(),
3894 "TQ merge must be a pure byte copy of the source columns",
3895 );
3896 let plan = crate::structures::TqQueryPlan::build(&codec, &vectors[7]);
3897 let results = merged.search_tq_distinct(1, &plan).unwrap();
3898 assert_eq!(results[0].0, 7, "merged payload must keep doc bases");
3899 }
3900
3901 #[test]
3902 fn tq_parallel_fold_matches_a_sequential_scan() {
3903 use crate::structures::vector::quantization::{
3904 TQ_BLOCK_LANES, tq_block_bytes, tq_score_block,
3905 };
3906
3907 let dim = 8;
3910 let count = TQ_PARALLEL_SCAN_MIN_VECTORS;
3911 let (bytes, vectors) = build_tq_payload(dim, count, 87);
3912 let disk =
3913 AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::TqFlat, count as u32).unwrap();
3914 let codec = crate::structures::TqCodec::new(dim);
3915 let plan = crate::structures::TqQueryPlan::build(&codec, &vectors[count / 3]);
3916 let k = 31;
3917
3918 let block_bytes = tq_block_bytes(disk.header().code_size);
3919 assert_eq!(
3920 disk.runs.len(),
3921 1,
3922 "the regression must parallelize chunks inside one run"
3923 );
3924 assert!(
3925 disk.runs[0].codes.len() / block_bytes > TQ_PARALLEL_SCAN_CHUNK_BLOCKS,
3926 "the single run must span multiple parallel chunks"
3927 );
3928 let raw = disk.raw.as_slice();
3929 let mut reference = BoundedAnnCollector::<true, true>::new(k);
3930 let mut scores = [0.0f32; TQ_BLOCK_LANES];
3931 for run in &disk.runs {
3932 let codes = &raw[run.codes.clone()];
3933 for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
3934 tq_score_block(&plan, block, &mut scores);
3935 let lane_base = block_index * TQ_BLOCK_LANES;
3936 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
3937 for (lane, &score) in scores.iter().enumerate().take(lanes) {
3938 let index = lane_base + lane;
3939 reference.insert(
3940 run_doc_id(raw, run, index).unwrap(),
3941 read_u16(raw, run.ordinals.start + index * 2),
3942 score,
3943 );
3944 }
3945 }
3946 }
3947
3948 assert_eq!(
3949 disk.search_tq_distinct(k, &plan).unwrap(),
3950 reference.into_sorted_results(),
3951 );
3952 }
3953
3954 #[test]
3955 fn ivf_tq_scale_pruning_matches_the_unpruned_scan() {
3956 use crate::structures::vector::ivf::{CoarseCentroids, CoarseConfig};
3957 use crate::structures::vector::quantization::{
3958 TQ_BLOCK_LANES, tq_ivf_block_bytes, tq_score_ivf_block,
3959 };
3960 use crate::structures::{IvfTqIndex, TqCodec, TqIvfEncodeScratch, TqIvfQueryPlan};
3961
3962 let dim = 32;
3963 let count = 400usize;
3964 let codec = std::sync::Arc::new(TqCodec::new(dim));
3965 let mut state = 5u64;
3966 let mut next = move || {
3967 state = state
3968 .wrapping_mul(6364136223846793005)
3969 .wrapping_add(1442695040888963407);
3970 ((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5
3971 };
3972 let vectors: Vec<Vec<f32>> = (0..count)
3973 .map(|_| {
3974 let mut v: Vec<f32> = (0..dim).map(|_| next()).collect();
3975 let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
3976 v.iter_mut().for_each(|x| *x /= norm);
3977 v
3978 })
3979 .collect();
3980 let mut centroids = CoarseCentroids::train(&CoarseConfig::new(dim, 8), &vectors, "test");
3981 centroids.version = crate::structures::mark_ivf_tq_cosine_generation(centroids.version);
3982 let mut index = IvfTqIndex::new(
3983 dim,
3984 crate::dsl::IvfRoutingMode::Flat,
3985 centroids.version,
3986 std::sync::Arc::clone(&codec),
3987 );
3988 let mut scratch = TqIvfEncodeScratch::default();
3989 for (i, vector) in vectors.iter().enumerate() {
3990 index.add_vector(
3991 ¢roids,
3992 (i / 2) as u32,
3993 (i % 2) as u16,
3994 vector,
3995 &mut scratch,
3996 );
3997 }
3998 let mut bytes = Vec::new();
3999 write_built_ivf_tq(&index, centroids.num_clusters, &mut bytes).unwrap();
4000 let disk =
4001 AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::IvfTq, count as u32).unwrap();
4002
4003 let block_bytes = tq_ivf_block_bytes(disk.header().code_size);
4006 let raw = disk.raw.as_slice();
4007 for run in &disk.runs {
4008 let codes = &raw[run.codes.clone()];
4009 let mut previous_scale = f32::INFINITY;
4010 for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
4011 let lane_base = block_index * TQ_BLOCK_LANES;
4012 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
4013 let mut block_scales = block[..TQ_BLOCK_LANES * size_of::<f32>()]
4014 .chunks_exact(size_of::<f32>())
4015 .take(lanes)
4016 .map(|lane| f32::from_le_bytes(lane.try_into().unwrap()));
4017 let first_scale = block_scales.next().unwrap();
4018 assert_eq!(tq_ivf_block_max_scale(block), first_scale);
4019 assert!(first_scale <= previous_scale);
4020 previous_scale = first_scale;
4021 for scale in block_scales {
4022 assert!(scale <= previous_scale);
4023 previous_scale = scale;
4024 }
4025 }
4026 }
4027 let k = 10;
4028 let parallel_pool = rayon::ThreadPoolBuilder::new()
4029 .num_threads(4)
4030 .build()
4031 .unwrap();
4032 for query_seed in [1u64, 9, 42] {
4033 let mut qstate = query_seed;
4034 let mut qnext = move || {
4035 qstate = qstate
4036 .wrapping_mul(6364136223846793005)
4037 .wrapping_add(1442695040888963407);
4038 ((qstate >> 33) as f32 / (1u64 << 31) as f32) - 0.5
4039 };
4040 let query: Vec<f32> = (0..dim).map(|_| qnext()).collect();
4041 let plan = TqIvfQueryPlan::build(
4042 ¢roids,
4043 &codec,
4044 &query,
4045 8,
4046 crate::dsl::IvfRoutingMode::Flat,
4047 );
4048 let mut reference = BoundedAnnCollector::<true, true>::new(k);
4051 let mut unpruned_scores = Vec::new();
4052 let mut scores = [0.0f32; TQ_BLOCK_LANES];
4053 for (cluster_id, cluster_dot) in plan.cluster_dots() {
4054 for run in disk.cluster_runs(cluster_id) {
4055 let codes = &raw[run.codes.clone()];
4056 for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
4057 tq_score_ivf_block(plan.tq_plan(), block, cluster_dot, &mut scores);
4058 let lane_base = block_index * TQ_BLOCK_LANES;
4059 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
4060 for (lane, &score) in scores.iter().enumerate().take(lanes) {
4061 let idx = lane_base + lane;
4062 reference.insert(
4063 run_doc_id(raw, run, idx).unwrap(),
4064 read_u16(raw, run.ordinals.start + idx * 2),
4065 score,
4066 );
4067 unpruned_scores.push((
4068 run_doc_id(raw, run, idx).unwrap(),
4069 read_u16(raw, run.ordinals.start + idx * 2),
4070 score,
4071 ));
4072 }
4073 }
4074 }
4075 }
4076
4077 let pruned = disk.search_ivf_tq_distinct(k, &plan).unwrap();
4078 let forced_parallel = parallel_pool.install(|| {
4079 disk.search_ivf_tq_distinct_with_tuning(k, &plan, 1, 4)
4082 .unwrap()
4083 });
4084 let reference = reference.into_sorted_results();
4085 assert_eq!(
4086 pruned, reference,
4087 "scale-bound pruning must not change the estimated top-k (seed {query_seed})"
4088 );
4089 assert_eq!(
4090 forced_parallel, reference,
4091 "parallel IVF-TQ pruning must match the unpruned top-k (seed {query_seed})"
4092 );
4093 for combiner in [
4094 crate::query::MultiValueCombiner::Max,
4095 crate::query::MultiValueCombiner::Sum,
4096 crate::query::MultiValueCombiner::Avg,
4097 crate::query::MultiValueCombiner::default(),
4098 crate::query::MultiValueCombiner::WeightedTopK { k: 3, decay: 0.7 },
4099 ] {
4100 let expected = combine_scored_ordinals(unpruned_scores.clone(), k, combiner);
4101 let combined = disk
4102 .search_ivf_tq_combined_documents(k, &plan, combiner)
4103 .unwrap();
4104 let parallel_combined = parallel_pool
4105 .install(|| {
4106 disk.search_ivf_tq_combined_documents_with_tuning(k, &plan, combiner, 1, 4)
4107 })
4108 .unwrap();
4109 assert_eq!(
4110 combined, expected,
4111 "combined IVF-TQ scan diverged from the unpruned reference \
4112 for {combiner:?} (seed {query_seed})",
4113 );
4114 assert_eq!(
4115 parallel_combined, expected,
4116 "parallel combined IVF-TQ scan diverged from the unpruned reference \
4117 for {combiner:?} (seed {query_seed})",
4118 );
4119 assert!(combined.len() <= k);
4120 }
4121 }
4122 }
4123
4124 #[test]
4132 #[ignore]
4133 fn ivf_tq_parallel_scan_benchmark() {
4134 use crate::structures::vector::ivf::CoarseCentroids;
4135 use crate::structures::{IvfTqIndex, TqCodec, TqIvfEncodeScratch, TqIvfQueryPlan};
4136
4137 fn env_usize(name: &str, default: usize) -> usize {
4138 std::env::var(name)
4139 .ok()
4140 .map(|value| {
4141 value
4142 .parse::<usize>()
4143 .unwrap_or_else(|_| panic!("{name} must be a positive integer"))
4144 })
4145 .unwrap_or(default)
4146 .max(1)
4147 }
4148
4149 fn median_ms(samples: &mut [f64]) -> f64 {
4150 samples.sort_unstable_by(f64::total_cmp);
4151 samples[samples.len() / 2]
4152 }
4153
4154 let dim = env_usize("IVF_SCAN_BENCH_DIM", 128);
4155 let count = env_usize("IVF_SCAN_BENCH_DOCS", 262_144);
4156 let iterations = env_usize("IVF_SCAN_BENCH_ITERS", 30);
4157 let threads = env_usize("IVF_SCAN_BENCH_THREADS", num_cpus::get().min(8));
4158 let codec = std::sync::Arc::new(TqCodec::new(dim));
4159 let version = crate::structures::mark_ivf_tq_cosine_generation(0x51ca_0001);
4160 let mut centroid = vec![0.0f32; dim];
4161 centroid[0] = 1.0;
4162 let centroids = CoarseCentroids {
4163 num_clusters: 1,
4164 dim,
4165 centroids: centroid,
4166 version,
4167 soar_config: None,
4168 routing_index: None,
4169 };
4170 let mut index = IvfTqIndex::new(
4171 dim,
4172 crate::dsl::IvfRoutingMode::Flat,
4173 version,
4174 std::sync::Arc::clone(&codec),
4175 );
4176 let mut scratch = TqIvfEncodeScratch::default();
4177 let mut state = 0x5eed_u64;
4178 let mut vector = vec![0.0f32; dim];
4179 for posting_index in 0..count {
4180 for value in &mut vector {
4181 state = state
4182 .wrapping_mul(6364136223846793005)
4183 .wrapping_add(1442695040888963407);
4184 *value = ((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5;
4185 }
4186 let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
4187 vector.iter_mut().for_each(|value| *value /= norm);
4188 index.add_vector(
4189 ¢roids,
4190 u32::try_from(posting_index / 2).unwrap(),
4191 u16::try_from(posting_index % 2).unwrap(),
4192 &vector,
4193 &mut scratch,
4194 );
4195 }
4196 let mut bytes = Vec::new();
4197 write_built_ivf_tq(&index, 1, &mut bytes).unwrap();
4198 let disk = AnnDiskIndex::open(
4199 OwnedBytes::new(bytes),
4200 AnnKind::IvfTq,
4201 u32::try_from(count.div_ceil(2)).unwrap(),
4202 )
4203 .unwrap();
4204 let query = vec![1.0f32; dim];
4205 let plan = TqIvfQueryPlan::build(
4206 ¢roids,
4207 &codec,
4208 &query,
4209 1,
4210 crate::dsl::IvfRoutingMode::Flat,
4211 );
4212 let pool = rayon::ThreadPoolBuilder::new()
4213 .num_threads(threads)
4214 .build()
4215 .unwrap();
4216 let serial_pool = rayon::ThreadPoolBuilder::new()
4217 .num_threads(1)
4218 .build()
4219 .unwrap();
4220
4221 let serial = pool
4222 .install(|| disk.search_ivf_tq_distinct_with_tuning(20, &plan, usize::MAX, 512))
4223 .unwrap();
4224 let parallel = pool
4225 .install(|| disk.search_ivf_tq_distinct_with_tuning(20, &plan, 1, 512))
4226 .unwrap();
4227 assert_eq!(parallel, serial);
4228
4229 let mut serial_ms = Vec::with_capacity(iterations);
4230 let mut parallel_ms = Vec::with_capacity(iterations);
4231 for _ in 0..iterations {
4232 let start = std::time::Instant::now();
4233 std::hint::black_box(
4234 pool.install(|| {
4235 disk.search_ivf_tq_distinct_with_tuning(20, &plan, usize::MAX, 512)
4236 })
4237 .unwrap(),
4238 );
4239 serial_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
4240
4241 let start = std::time::Instant::now();
4242 std::hint::black_box(
4243 pool.install(|| disk.search_ivf_tq_distinct_with_tuning(20, &plan, 1, 512))
4244 .unwrap(),
4245 );
4246 parallel_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
4247 }
4248 let serial_p50 = median_ms(&mut serial_ms);
4249 let parallel_p50 = median_ms(&mut parallel_ms);
4250 println!(
4251 "IVF-TQ distinct scan: postings={count} dim={dim} threads={threads} \
4252 serial_p50={serial_p50:.3}ms parallel_p50={parallel_p50:.3}ms speedup={:.2}x",
4253 serial_p50 / parallel_p50,
4254 );
4255
4256 let combiner = crate::query::MultiValueCombiner::Sum;
4257 let serial_combined = serial_pool
4258 .install(|| {
4259 disk.search_ivf_tq_combined_documents_with_tuning(20, &plan, combiner, 1, 512)
4260 })
4261 .unwrap();
4262 let parallel_combined = pool
4263 .install(|| {
4264 disk.search_ivf_tq_combined_documents_with_tuning(20, &plan, combiner, 1, 512)
4265 })
4266 .unwrap();
4267 assert_eq!(parallel_combined, serial_combined);
4268 let mut serial_combined_ms = Vec::with_capacity(iterations);
4269 let mut parallel_combined_ms = Vec::with_capacity(iterations);
4270 for _ in 0..iterations {
4271 let start = std::time::Instant::now();
4272 std::hint::black_box(
4273 serial_pool
4274 .install(|| {
4275 disk.search_ivf_tq_combined_documents_with_tuning(
4276 20, &plan, combiner, 1, 512,
4277 )
4278 })
4279 .unwrap(),
4280 );
4281 serial_combined_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
4282
4283 let start = std::time::Instant::now();
4284 std::hint::black_box(
4285 pool.install(|| {
4286 disk.search_ivf_tq_combined_documents_with_tuning(20, &plan, combiner, 1, 512)
4287 })
4288 .unwrap(),
4289 );
4290 parallel_combined_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
4291 }
4292 let serial_combined_p50 = median_ms(&mut serial_combined_ms);
4293 let parallel_combined_p50 = median_ms(&mut parallel_combined_ms);
4294 println!(
4295 "IVF-TQ combined scan: postings={count} dim={dim} threads={threads} \
4296 serial_p50={serial_combined_p50:.3}ms parallel_p50={parallel_combined_p50:.3}ms \
4297 speedup={:.2}x",
4298 serial_combined_p50 / parallel_combined_p50,
4299 );
4300 }
4301
4302 #[test]
4306 #[ignore]
4307 fn binary_ivf_parallel_scan_benchmark() {
4308 fn env_usize(name: &str, default: usize) -> usize {
4309 std::env::var(name)
4310 .ok()
4311 .map(|value| value.parse::<usize>().unwrap())
4312 .unwrap_or(default)
4313 .max(1)
4314 }
4315
4316 fn median_ms(samples: &mut [f64]) -> f64 {
4317 samples.sort_unstable_by(f64::total_cmp);
4318 samples[samples.len() / 2]
4319 }
4320
4321 let count = env_usize("BINARY_SCAN_BENCH_POSTINGS", 131_072);
4322 let dim_bits = env_usize("BINARY_SCAN_BENCH_DIM_BITS", 2_048);
4323 assert!(dim_bits.is_multiple_of(8));
4324 let code_size = dim_bits / 8;
4325 let iterations = env_usize("BINARY_SCAN_BENCH_ITERS", 30);
4326 let threads = env_usize("BINARY_SCAN_BENCH_THREADS", num_cpus::get().min(8));
4327 let mut state = 0xb1a4_5eed_u64;
4328 let mut codes = vec![0u8; count * code_size];
4329 for code in &mut codes {
4330 state = state
4331 .wrapping_mul(6364136223846793005)
4332 .wrapping_add(1442695040888963407);
4333 *code = (state >> 56) as u8;
4334 }
4335 let query = vec![0xa5u8; code_size];
4336 let unique_docs: Vec<u32> = (0..count)
4337 .map(|index| u32::try_from(index).unwrap())
4338 .collect();
4339 let unique_ordinals = vec![0u16; count];
4340 let multi_docs: Vec<u32> = (0..count)
4341 .map(|index| u32::try_from(index / 2).unwrap())
4342 .collect();
4343 let multi_ordinals: Vec<u16> = (0..count)
4344 .map(|index| u16::try_from(index % 2).unwrap())
4345 .collect();
4346 let header = AnnDiskHeader {
4347 kind: AnnKind::BinaryIvf,
4348 routing: IvfRoutingMode::Flat,
4349 dim: dim_bits,
4350 code_size,
4351 num_clusters: 1,
4352 quantizer_version: 0xb1a4,
4353 codebook_version: 0,
4354 vector_count: count,
4355 };
4356 let mut unique_bytes = Vec::new();
4357 write_built_runs(
4358 header.clone(),
4359 &[BuildRun {
4360 cluster_id: 0,
4361 doc_ids: &unique_docs,
4362 ordinals: &unique_ordinals,
4363 codes: &codes,
4364 }],
4365 &mut unique_bytes,
4366 )
4367 .unwrap();
4368 let unique_disk = AnnDiskIndex::open(
4369 OwnedBytes::new(unique_bytes),
4370 AnnKind::BinaryIvf,
4371 u32::try_from(count).unwrap(),
4372 )
4373 .unwrap();
4374 let mut multi_bytes = Vec::new();
4375 write_built_runs(
4376 header,
4377 &[BuildRun {
4378 cluster_id: 0,
4379 doc_ids: &multi_docs,
4380 ordinals: &multi_ordinals,
4381 codes: &codes,
4382 }],
4383 &mut multi_bytes,
4384 )
4385 .unwrap();
4386 let multi_disk = AnnDiskIndex::open(
4387 OwnedBytes::new(multi_bytes),
4388 AnnKind::BinaryIvf,
4389 u32::try_from(count.div_ceil(2)).unwrap(),
4390 )
4391 .unwrap();
4392 let pool = rayon::ThreadPoolBuilder::new()
4393 .num_threads(threads)
4394 .build()
4395 .unwrap();
4396 let serial_pool = rayon::ThreadPoolBuilder::new()
4397 .num_threads(1)
4398 .build()
4399 .unwrap();
4400
4401 let serial = pool
4402 .install(|| {
4403 unique_disk.search_binary_clusters_with_tuning::<false>(
4404 &query,
4405 20,
4406 &[0],
4407 usize::MAX,
4408 )
4409 })
4410 .unwrap();
4411 let parallel = pool
4412 .install(|| {
4413 unique_disk.search_binary_clusters_with_tuning::<false>(&query, 20, &[0], 1)
4414 })
4415 .unwrap();
4416 assert_eq!(parallel, serial);
4417 let mut serial_ms = Vec::with_capacity(iterations);
4418 let mut parallel_ms = Vec::with_capacity(iterations);
4419 for _ in 0..iterations {
4420 let start = std::time::Instant::now();
4421 std::hint::black_box(
4422 pool.install(|| {
4423 unique_disk.search_binary_clusters_with_tuning::<false>(
4424 &query,
4425 20,
4426 &[0],
4427 usize::MAX,
4428 )
4429 })
4430 .unwrap(),
4431 );
4432 serial_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
4433
4434 let start = std::time::Instant::now();
4435 std::hint::black_box(
4436 pool.install(|| {
4437 unique_disk.search_binary_clusters_with_tuning::<false>(&query, 20, &[0], 1)
4438 })
4439 .unwrap(),
4440 );
4441 parallel_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
4442 }
4443 let serial_p50 = median_ms(&mut serial_ms);
4444 let parallel_p50 = median_ms(&mut parallel_ms);
4445 println!(
4446 "binary-IVF distinct scan: postings={count} dim_bits={dim_bits} threads={threads} \
4447 serial_p50={serial_p50:.3}ms parallel_p50={parallel_p50:.3}ms speedup={:.2}x",
4448 serial_p50 / parallel_p50,
4449 );
4450
4451 let combiner = crate::query::MultiValueCombiner::Sum;
4452 let serial_combined = serial_pool
4453 .install(|| {
4454 multi_disk.search_binary_combined_documents_with_tuning(
4455 20,
4456 &query,
4457 &[0],
4458 combiner,
4459 1,
4460 )
4461 })
4462 .unwrap();
4463 let parallel_combined = pool
4464 .install(|| {
4465 multi_disk.search_binary_combined_documents_with_tuning(
4466 20,
4467 &query,
4468 &[0],
4469 combiner,
4470 1,
4471 )
4472 })
4473 .unwrap();
4474 assert_eq!(parallel_combined, serial_combined);
4475 let mut serial_combined_ms = Vec::with_capacity(iterations);
4476 let mut parallel_combined_ms = Vec::with_capacity(iterations);
4477 for _ in 0..iterations {
4478 let start = std::time::Instant::now();
4479 std::hint::black_box(
4480 serial_pool
4481 .install(|| {
4482 multi_disk.search_binary_combined_documents_with_tuning(
4483 20,
4484 &query,
4485 &[0],
4486 combiner,
4487 1,
4488 )
4489 })
4490 .unwrap(),
4491 );
4492 serial_combined_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
4493
4494 let start = std::time::Instant::now();
4495 std::hint::black_box(
4496 pool.install(|| {
4497 multi_disk.search_binary_combined_documents_with_tuning(
4498 20,
4499 &query,
4500 &[0],
4501 combiner,
4502 1,
4503 )
4504 })
4505 .unwrap(),
4506 );
4507 parallel_combined_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
4508 }
4509 let serial_combined_p50 = median_ms(&mut serial_combined_ms);
4510 let parallel_combined_p50 = median_ms(&mut parallel_combined_ms);
4511 println!(
4512 "binary-IVF combined scan: postings={count} dim_bits={dim_bits} threads={threads} \
4513 serial_p50={serial_combined_p50:.3}ms parallel_p50={parallel_combined_p50:.3}ms \
4514 speedup={:.2}x",
4515 serial_combined_p50 / parallel_combined_p50,
4516 );
4517 }
4518
4519 #[test]
4520 fn open_rejects_tq_payload_with_inconsistent_geometry() {
4521 let (bytes, _) = build_tq_payload(20, 4, 9);
4522 let mut corrupted = bytes.clone();
4524 corrupted[12..16].copy_from_slice(&15u32.to_le_bytes());
4525 assert!(
4526 AnnDiskIndex::open(OwnedBytes::new(corrupted), AnnKind::TqFlat, 4).is_err(),
4527 "TQ header with code_size != padded_dim/2 must be refused"
4528 );
4529
4530 let (short_bytes, _) = build_tq_payload(20, 4, 9);
4532 let mut wrong_kind = short_bytes.clone();
4533 wrong_kind[4] = AnnKind::BinaryIvf as u8;
4534 assert!(
4535 AnnDiskIndex::open(OwnedBytes::new(wrong_kind), AnnKind::BinaryIvf, 4).is_err(),
4536 "TQ block-padded columns must not validate under another kind"
4537 );
4538
4539 let (legacy_bytes, _) = build_tq_payload(20, 4, 9);
4541 let mut legacy_kind = legacy_bytes.clone();
4542 legacy_kind[4] = 1;
4543 let Err(error) = AnnDiskIndex::open(OwnedBytes::new(legacy_kind), AnnKind::TqFlat, 4)
4544 else {
4545 panic!("retired IVF-PQ payloads must not open");
4546 };
4547 assert!(
4548 error.to_string().contains("IVF-PQ"),
4549 "error must name the retired format: {error}"
4550 );
4551 assert!(AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::TqFlat, 4).is_ok());
4552 }
4553
4554 #[test]
4555 fn open_rejects_old_or_out_of_range_ann_payloads() {
4556 let mut legacy = vec![0u8; ANN_HEADER_SIZE + ANN_FOOTER_SIZE];
4557 legacy[..4].copy_from_slice(b"old!");
4558 assert!(AnnDiskIndex::open(OwnedBytes::new(legacy), AnnKind::BinaryIvf, 1).is_err());
4559
4560 let docs = [0u32];
4561 let ordinals = [0u16];
4562 let codes = [0u8];
4563 let runs = [BuildRun {
4564 cluster_id: 0,
4565 doc_ids: &docs,
4566 ordinals: &ordinals,
4567 codes: &codes,
4568 }];
4569 let mut bytes = Vec::new();
4570 write_built_runs(binary_header(1), &runs, &mut bytes).unwrap();
4571 let footer = bytes.len() - ANN_FOOTER_SIZE;
4572 let directory = usize::try_from(u64::from_le_bytes(
4573 bytes[footer..footer + 8].try_into().unwrap(),
4574 ))
4575 .unwrap();
4576 bytes[directory + 12..directory + 16].copy_from_slice(&10u32.to_le_bytes());
4577 assert!(AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::BinaryIvf, 1).is_err());
4578 }
4579
4580 fn scann_binary_artifact(
4581 generation: u64,
4582 ) -> crate::structures::vector::scann::ScannTrainedArtifact {
4583 use crate::structures::vector::scann::{
4584 ScannConfig, ScannEncoding, ScannRoutingLevel, ScannTrainedArtifact,
4585 };
4586 ScannTrainedArtifact::new(
4587 generation,
4588 100_000,
4589 ScannConfig {
4590 dimension: 16,
4591 tree_levels: 1,
4592 num_leaves: 2,
4593 encoding: ScannEncoding::BinaryHamming,
4594 },
4595 vec![ScannRoutingLevel {
4596 centroid_count: 2,
4597 centroid_codes: vec![0, 0, 0xff, 0xff],
4598 minimums: Vec::new(),
4599 steps: Vec::new(),
4600 child_offsets: Vec::new(),
4601 }],
4602 None,
4603 )
4604 .unwrap()
4605 }
4606
4607 #[test]
4608 fn scann_binary_ann_round_trip_validates_exact_global_generation() {
4609 use crate::structures::vector::scann::{
4610 ScannEncoding, ScannLeafRun, ScannSegmentPayload, ScannTrainedArtifactView,
4611 };
4612
4613 let artifact = scann_binary_artifact(41);
4614 let primary = ScannLeafRun::from_rows(
4615 0,
4616 0,
4617 &[0, 1],
4618 &[0, 2],
4619 vec![0x12, 0x34, 0xab, 0xcd],
4620 ScannEncoding::BinaryHamming,
4621 16,
4622 )
4623 .unwrap();
4624 let secondary = ScannLeafRun::from_rows(
4625 1,
4626 0,
4627 &[0],
4628 &[0],
4629 vec![0x12, 0x34],
4630 ScannEncoding::BinaryHamming,
4631 16,
4632 )
4633 .unwrap();
4634 let payload = ScannSegmentPayload::new(&artifact, 2, vec![primary, secondary]).unwrap();
4635 let mut bytes = Vec::new();
4636 write_built_scann(&payload, &mut bytes).unwrap();
4637
4638 let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::ScannBinary, 2).unwrap();
4639 assert_eq!(disk.header.quantizer_version, artifact.generation);
4640 assert_eq!(disk.header.codebook_version, artifact.artifact_id);
4641 assert_eq!(disk.header.vector_count, 3);
4642 let artifact_bytes = artifact.to_bytes().unwrap();
4643 let view = ScannTrainedArtifactView::parse(&artifact_bytes).unwrap();
4644 disk.validate_scann_generation(&view.config, view.generation, view.artifact_id)
4645 .unwrap();
4646 let selective = crate::structures::SoarConfig::new().target_spill_fraction(0.5);
4647 disk.validate_scann_posting_count(2, Some(&selective))
4648 .unwrap();
4649 assert!(
4650 disk.validate_scann_posting_count(2, None)
4651 .unwrap_err()
4652 .to_string()
4653 .contains("physical postings")
4654 );
4655 let too_small_budget = crate::structures::SoarConfig::new().target_spill_fraction(0.49);
4656 assert!(
4657 disk.validate_scann_posting_count(2, Some(&too_small_budget))
4658 .is_err()
4659 );
4660
4661 let serial = disk
4662 .search_binary_clusters_with_tuning::<false>(&[0x12, 0x34], 2, &[0, 1], usize::MAX)
4663 .unwrap();
4664 let parallel = rayon::ThreadPoolBuilder::new()
4665 .num_threads(2)
4666 .build()
4667 .unwrap()
4668 .install(|| {
4669 disk.search_binary_clusters_with_tuning::<false>(&[0x12, 0x34], 2, &[0, 1], 1)
4670 })
4671 .unwrap();
4672 assert_eq!(parallel, serial);
4673 assert_eq!(serial.len(), 2, "secondary posting duplicated a result");
4674 assert_eq!(serial[0], (0, 0, 1.0));
4675
4676 let other = scann_binary_artifact(42);
4677 let other_bytes = other.to_bytes().unwrap();
4678 let other_view = ScannTrainedArtifactView::parse(&other_bytes).unwrap();
4679 let error = disk
4680 .validate_scann_generation(
4681 &other_view.config,
4682 other_view.generation,
4683 other_view.artifact_id,
4684 )
4685 .unwrap_err();
4686 assert!(error.to_string().contains("global trained generation"));
4687 }
4688
4689 #[test]
4690 fn scann_ah_ann_round_trip_preserves_fastscan_tail_geometry() {
4691 use crate::structures::vector::scann::{
4692 ScannAhCodebook, ScannConfig, ScannEncoding, ScannLeafRun, ScannRoutingLevel,
4693 ScannSegmentPayload, ScannTrainedArtifact, ScannTrainedArtifactView,
4694 };
4695
4696 let encoding = ScannEncoding::AsymmetricHash {
4697 dimensions_per_block: 2,
4698 bits_per_code: 4,
4699 };
4700 let artifact = ScannTrainedArtifact::new(
4701 51,
4702 100_000,
4703 ScannConfig {
4704 dimension: 8,
4705 tree_levels: 1,
4706 num_leaves: 2,
4707 encoding,
4708 },
4709 vec![ScannRoutingLevel {
4710 centroid_count: 2,
4711 centroid_codes: vec![0; 16],
4712 minimums: vec![0.0; 8],
4713 steps: vec![1.0; 8],
4714 child_offsets: Vec::new(),
4715 }],
4716 Some(ScannAhCodebook {
4717 dimensions_per_block: 2,
4718 centers_per_block: 16,
4719 centers: vec![0.0; 4 * 16 * 2],
4720 }),
4721 )
4722 .unwrap();
4723 let docs: Vec<u32> = (0..33).collect();
4724 let ordinals = vec![0u16; docs.len()];
4725 let codes = vec![0x5a; encoding.leaf_code_bytes(8, docs.len()).unwrap()];
4726 let run = ScannLeafRun::from_rows(0, 0, &docs, &ordinals, codes, encoding, 8).unwrap();
4727 let payload = ScannSegmentPayload::new(&artifact, 33, vec![run]).unwrap();
4728 let mut bytes = Vec::new();
4729 write_built_scann(&payload, &mut bytes).unwrap();
4730
4731 let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::ScannAh, 33).unwrap();
4732 assert_eq!(
4733 disk.runs[0].codes.len(),
4734 encoding.leaf_code_bytes(8, 33).unwrap()
4735 );
4736 let artifact_bytes = artifact.to_bytes().unwrap();
4737 let view = ScannTrainedArtifactView::parse(&artifact_bytes).unwrap();
4738 disk.validate_scann_generation(&view.config, view.generation, view.artifact_id)
4739 .unwrap();
4740 }
4741
4742 #[test]
4743 fn scann_binary_merge_copies_codes_and_rebases_without_retraining() {
4744 use crate::structures::vector::scann::{ScannEncoding, ScannLeafRun, ScannSegmentPayload};
4745
4746 let artifact = scann_binary_artifact(61);
4747 let make = |codes: Vec<u8>| {
4748 let secondary_code = codes[..2].to_vec();
4749 let primary = ScannLeafRun::from_rows(
4750 0,
4751 0,
4752 &[0, 1],
4753 &[0, 0],
4754 codes,
4755 ScannEncoding::BinaryHamming,
4756 16,
4757 )
4758 .unwrap();
4759 let secondary = ScannLeafRun::from_rows(
4760 1,
4761 0,
4762 &[0],
4763 &[0],
4764 secondary_code,
4765 ScannEncoding::BinaryHamming,
4766 16,
4767 )
4768 .unwrap();
4769 let payload = ScannSegmentPayload::new(&artifact, 2, vec![primary, secondary]).unwrap();
4770 let mut bytes = Vec::new();
4771 write_built_scann(&payload, &mut bytes).unwrap();
4772 AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::ScannBinary, 2).unwrap()
4773 };
4774 let left = make(vec![1, 2, 3, 4]);
4775 let right = make(vec![5, 6, 7, 8]);
4776 let left_codes = left.raw.as_slice()[left.runs[0].codes.clone()].to_vec();
4777 let right_codes = right.raw.as_slice()[right.runs[0].codes.clone()].to_vec();
4778
4779 let mut merged_bytes = Vec::new();
4780 write_merged_ann(&[(&left, 0), (&right, 2)], &mut merged_bytes).unwrap();
4781 let merged =
4782 AnnDiskIndex::open(OwnedBytes::new(merged_bytes), AnnKind::ScannBinary, 4).unwrap();
4783 assert_eq!(merged.header.vector_count, 6);
4784 merged
4785 .validate_scann_posting_count(
4786 4,
4787 Some(&crate::structures::SoarConfig::new().target_spill_fraction(0.5)),
4788 )
4789 .unwrap();
4790 assert_eq!(merged.runs[0].doc_base, 0);
4791 assert_eq!(merged.runs[1].doc_base, 2);
4792 assert_eq!(
4793 &merged.raw.as_slice()[merged.runs[0].codes.clone()],
4794 left_codes
4795 );
4796 assert_eq!(
4797 &merged.raw.as_slice()[merged.runs[1].codes.clone()],
4798 right_codes
4799 );
4800 assert_eq!(
4801 merged
4802 .search_binary_clusters::<false>(&[0, 0], 4, &[0, 1])
4803 .unwrap()
4804 .len(),
4805 4,
4806 "merge exposed primary and secondary postings as separate results",
4807 );
4808
4809 let other_artifact = scann_binary_artifact(62);
4810 let other_run = ScannLeafRun::from_rows(
4811 0,
4812 0,
4813 &[0],
4814 &[0],
4815 vec![9, 10],
4816 ScannEncoding::BinaryHamming,
4817 16,
4818 )
4819 .unwrap();
4820 let other_payload = ScannSegmentPayload::new(&other_artifact, 1, vec![other_run]).unwrap();
4821 let mut other_bytes = Vec::new();
4822 write_built_scann(&other_payload, &mut other_bytes).unwrap();
4823 let other =
4824 AnnDiskIndex::open(OwnedBytes::new(other_bytes), AnnKind::ScannBinary, 1).unwrap();
4825 let error = write_merged_ann(&[(&left, 0), (&other, 2)], &mut Vec::new()).unwrap_err();
4826 assert!(error.to_string().contains("incompatible generations"));
4827 }
4828}