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}
74
75impl AnnKind {
76 fn from_u8(value: u8) -> io::Result<Self> {
77 match value {
78 1 => Err(invalid_data(
79 "ANN kind 1 (IVF-PQ) is no longer supported; recreate the index \
80 with `ivf_tq` and reindex",
81 )),
82 2 => Ok(Self::BinaryIvf),
83 3 => Ok(Self::TqFlat),
84 4 => Ok(Self::IvfTq),
85 _ => Err(invalid_data(format!("unknown ANN kind {value}"))),
86 }
87 }
88}
89
90fn expected_codes_column_len(kind: AnnKind, count: usize, code_size: usize) -> io::Result<usize> {
94 match kind {
95 AnnKind::BinaryIvf => count
96 .checked_mul(code_size)
97 .ok_or_else(|| invalid_data("ANN code column size overflows usize")),
98 AnnKind::TqFlat => {
100 crate::structures::vector::quantization::tq_codes_column_len_checked(count, code_size)
101 .ok_or_else(|| invalid_data("TQ code column size overflows usize"))
102 }
103 AnnKind::IvfTq => crate::structures::vector::quantization::tq_ivf_codes_column_len_checked(
104 count, code_size,
105 )
106 .ok_or_else(|| invalid_data("IVF-TQ code column size overflows usize")),
107 }
108}
109
110#[derive(Debug, Clone, PartialEq, Eq)]
111pub(crate) struct AnnDiskHeader {
112 pub kind: AnnKind,
113 pub routing: IvfRoutingMode,
114 pub dim: usize,
115 pub code_size: usize,
116 pub num_clusters: u32,
117 pub quantizer_version: u64,
118 pub codebook_version: u64,
119 pub vector_count: usize,
120}
121
122#[derive(Debug)]
123struct AnnRun {
124 cluster_id: u32,
125 doc_base: u32,
126 max_doc_id: u32,
127 count: usize,
128 doc_ids: Range<usize>,
129 ordinals: Range<usize>,
130 codes: Range<usize>,
131}
132
133#[derive(Clone, Copy)]
134struct IvfTqScanTask<'a> {
135 run: &'a AnnRun,
136 cluster_dot: f32,
137 first_block: usize,
138 last_block: usize,
139}
140
141#[derive(Clone, Copy)]
142struct BinaryScanTask<'a> {
143 run: &'a AnnRun,
144 first_index: usize,
145 count: usize,
146}
147
148pub(crate) struct AnnDiskIndex {
151 #[cfg(feature = "native")]
153 heap_pins: crate::segment::pin::HeapPinSet,
154 raw: OwnedBytes,
155 header: AnnDiskHeader,
156 runs: Vec<AnnRun>,
157}
158
159#[derive(Debug, Clone, Copy, PartialEq)]
167pub struct AnnHealth {
168 pub vectors: u64,
170 pub clusters_nonempty: u32,
172 pub clusters_total: u32,
174 pub runs: u32,
176 pub largest_cluster: u32,
178 pub largest_cluster_vectors: u64,
179 pub imbalance: f64,
184 pub payload_bytes: u64,
186}
187
188impl AnnHealth {
189 pub fn fragmentation(&self) -> f64 {
193 if self.clusters_nonempty == 0 {
194 return 0.0;
195 }
196 f64::from(self.runs) / f64::from(self.clusters_nonempty)
197 }
198
199 pub fn largest_cluster_share(&self) -> f64 {
201 if self.vectors == 0 {
202 return 0.0;
203 }
204 self.largest_cluster_vectors as f64 / self.vectors as f64
205 }
206}
207
208const ANN_SKEW_WARN_SHARE: f64 = 0.05;
212const ANN_SKEW_WARN_MIN_VECTORS: u64 = 100_000;
213const ANN_FRAGMENTATION_WARN: f64 = 8.0;
216
217#[derive(Debug, Clone, Copy, PartialEq)]
223pub(crate) struct AnnDocumentCandidate {
224 pub(crate) doc_id: u32,
225 pub(crate) score: f32,
226}
227
228impl AnnDiskIndex {
229 pub(crate) fn open(
230 raw: OwnedBytes,
231 expected_kind: AnnKind,
232 total_docs: u32,
233 ) -> io::Result<Self> {
234 if raw.len() < ANN_HEADER_SIZE + ANN_FOOTER_SIZE {
235 return Err(invalid_data("ANN payload is shorter than header + footer"));
236 }
237 let bytes = raw.as_slice();
238 let mut header_cursor = std::io::Cursor::new(&bytes[..ANN_HEADER_SIZE]);
239 if header_cursor.read_u32::<LittleEndian>()? != ANN_HEADER_MAGIC {
240 return Err(invalid_data("ANN payload has unsupported header magic"));
241 }
242 let kind = AnnKind::from_u8(header_cursor.read_u8()?)?;
243 if kind != expected_kind {
244 return Err(invalid_data(format!(
245 "ANN payload kind {kind:?} does not match expected {expected_kind:?}"
246 )));
247 }
248 let routing = routing_from_u8(header_cursor.read_u8()?)?;
249 if header_cursor.read_u16::<LittleEndian>()? != ANN_DISK_VERSION {
250 return Err(invalid_data("ANN payload has unsupported format version"));
251 }
252 let dim = header_cursor.read_u32::<LittleEndian>()? as usize;
253 let code_size = header_cursor.read_u32::<LittleEndian>()? as usize;
254 let num_clusters = header_cursor.read_u32::<LittleEndian>()?;
255 if header_cursor.read_u32::<LittleEndian>()? != 0 {
256 return Err(invalid_data("ANN header reserved field is non-zero"));
257 }
258 let quantizer_version = header_cursor.read_u64::<LittleEndian>()?;
259 let codebook_version = header_cursor.read_u64::<LittleEndian>()?;
260 let vector_count = usize::try_from(header_cursor.read_u64::<LittleEndian>()?)
261 .map_err(|_| invalid_data("ANN vector count exceeds usize"))?;
262 if header_cursor.read_u64::<LittleEndian>()? != 0 {
263 return Err(invalid_data("ANN header tail is non-zero"));
264 }
265 let header = AnnDiskHeader {
266 kind,
267 routing,
268 dim,
269 code_size,
270 num_clusters,
271 quantizer_version,
272 codebook_version,
273 vector_count,
274 };
275 validate_header(&header)?;
276
277 let footer_start = bytes.len() - ANN_FOOTER_SIZE;
278 let mut footer_cursor = std::io::Cursor::new(&bytes[footer_start..]);
279 let directory_offset = usize::try_from(footer_cursor.read_u64::<LittleEndian>()?)
280 .map_err(|_| invalid_data("ANN directory offset exceeds usize"))?;
281 let num_runs = usize::try_from(footer_cursor.read_u64::<LittleEndian>()?)
282 .map_err(|_| invalid_data("ANN run count exceeds usize"))?;
283 if footer_cursor.read_u32::<LittleEndian>()? != ANN_FOOTER_MAGIC
284 || footer_cursor.read_u32::<LittleEndian>()? != u32::from(ANN_DISK_VERSION)
285 {
286 return Err(invalid_data("ANN payload has unsupported footer"));
287 }
288 if num_runs == 0 {
289 return Err(invalid_data("ANN payload has no cluster runs"));
290 }
291 let directory_len = num_runs
292 .checked_mul(ANN_RUN_SIZE)
293 .ok_or_else(|| invalid_data("ANN directory size overflows usize"))?;
294 if directory_offset < ANN_HEADER_SIZE
295 || directory_offset.checked_add(directory_len) != Some(footer_start)
296 {
297 return Err(invalid_data("ANN directory does not end at the footer"));
298 }
299
300 let mut runs = Vec::with_capacity(num_runs);
301 let mut directory_cursor = std::io::Cursor::new(&bytes[directory_offset..footer_start]);
302 let mut previous_cluster = None;
303 let mut counted_vectors = 0usize;
304 for _ in 0..num_runs {
305 let cluster_id = directory_cursor.read_u32::<LittleEndian>()?;
306 let doc_base = directory_cursor.read_u32::<LittleEndian>()?;
307 let count = directory_cursor.read_u32::<LittleEndian>()? as usize;
308 let max_doc_id = directory_cursor.read_u32::<LittleEndian>()?;
309 let doc_ids_offset = usize::try_from(directory_cursor.read_u64::<LittleEndian>()?)
310 .map_err(|_| invalid_data("ANN doc-ID offset exceeds usize"))?;
311 let ordinals_offset = usize::try_from(directory_cursor.read_u64::<LittleEndian>()?)
312 .map_err(|_| invalid_data("ANN ordinal offset exceeds usize"))?;
313 let codes_offset = usize::try_from(directory_cursor.read_u64::<LittleEndian>()?)
314 .map_err(|_| invalid_data("ANN code offset exceeds usize"))?;
315 let codes_len = usize::try_from(directory_cursor.read_u64::<LittleEndian>()?)
316 .map_err(|_| invalid_data("ANN code length exceeds usize"))?;
317 if count == 0
318 || cluster_id >= num_clusters
319 || previous_cluster.is_some_and(|previous| previous > cluster_id)
320 || doc_base
321 .checked_add(max_doc_id)
322 .is_none_or(|doc_id| doc_id >= total_docs)
323 {
324 return Err(invalid_data("ANN run metadata is invalid"));
325 }
326 previous_cluster = Some(cluster_id);
327 let doc_ids_len = count
328 .checked_mul(std::mem::size_of::<u32>())
329 .ok_or_else(|| invalid_data("ANN doc-ID column size overflows usize"))?;
330 let ordinals_len = count
331 .checked_mul(std::mem::size_of::<u16>())
332 .ok_or_else(|| invalid_data("ANN ordinal column size overflows usize"))?;
333 let expected_codes_len = expected_codes_column_len(kind, count, code_size)?;
334 let doc_ids_end = doc_ids_offset
335 .checked_add(doc_ids_len)
336 .ok_or_else(|| invalid_data("ANN doc-ID range overflows usize"))?;
337 let ordinals_end = ordinals_offset
338 .checked_add(ordinals_len)
339 .ok_or_else(|| invalid_data("ANN ordinal range overflows usize"))?;
340 let codes_end = codes_offset
341 .checked_add(codes_len)
342 .ok_or_else(|| invalid_data("ANN code range overflows usize"))?;
343 if doc_ids_offset < ANN_HEADER_SIZE
344 || ordinals_offset != doc_ids_end
345 || codes_offset != ordinals_end
346 || codes_len != expected_codes_len
347 || codes_end > directory_offset
348 {
349 return Err(invalid_data("ANN run columns are not contiguous/in bounds"));
350 }
351 runs.push(AnnRun {
352 cluster_id,
353 doc_base,
354 max_doc_id,
355 count,
356 doc_ids: doc_ids_offset..doc_ids_end,
357 ordinals: ordinals_offset..ordinals_end,
358 codes: codes_offset..codes_end,
359 });
360 counted_vectors = counted_vectors
361 .checked_add(count)
362 .ok_or_else(|| invalid_data("ANN run vector count overflows usize"))?;
363 }
364 let mut payload_order: Vec<usize> = (0..runs.len()).collect();
365 payload_order.sort_unstable_by_key(|&index| runs[index].doc_ids.start);
366 let mut expected_payload_offset = ANN_HEADER_SIZE;
367 for index in payload_order {
368 let run = &runs[index];
369 if run.doc_ids.start != expected_payload_offset {
370 return Err(invalid_data("ANN payload runs overlap or contain gaps"));
371 }
372 expected_payload_offset = run.codes.end;
373 }
374 if expected_payload_offset != directory_offset || counted_vectors != vector_count {
375 return Err(invalid_data(
376 "ANN payload coverage/vector count is inconsistent",
377 ));
378 }
379
380 #[cfg(feature = "native")]
387 raw.madvise_range(
388 ANN_HEADER_SIZE..directory_offset,
389 if kind == AnnKind::TqFlat {
390 libc::MADV_SEQUENTIAL
391 } else {
392 libc::MADV_RANDOM
393 },
394 );
395
396 Ok(Self {
397 #[cfg(feature = "native")]
398 heap_pins: Default::default(),
399 raw,
400 header,
401 runs,
402 })
403 }
404
405 pub(crate) fn health(&self) -> AnnHealth {
408 let mut vectors = 0u64;
409 let mut clusters_nonempty = 0u32;
410 let mut payload_bytes = 0u64;
411 let mut largest = (0u32, 0u64);
412 let mut sum_squares = 0f64;
413 let mut index = 0usize;
415 while index < self.runs.len() {
416 let cluster_id = self.runs[index].cluster_id;
417 let mut cluster_vectors = 0u64;
418 while index < self.runs.len() && self.runs[index].cluster_id == cluster_id {
419 let run = &self.runs[index];
420 cluster_vectors += run.count as u64;
421 payload_bytes += (run.codes.end - run.codes.start) as u64;
422 index += 1;
423 }
424 vectors += cluster_vectors;
425 clusters_nonempty += 1;
426 sum_squares += (cluster_vectors as f64) * (cluster_vectors as f64);
427 if cluster_vectors > largest.1 {
428 largest = (cluster_id, cluster_vectors);
429 }
430 }
431 let imbalance = if vectors == 0 || clusters_nonempty == 0 {
432 0.0
433 } else {
434 f64::from(clusters_nonempty) * sum_squares / ((vectors as f64) * (vectors as f64))
435 };
436 AnnHealth {
437 vectors,
438 clusters_nonempty,
439 clusters_total: self.header.num_clusters,
440 runs: self.runs.len() as u32,
441 largest_cluster: largest.0,
442 largest_cluster_vectors: largest.1,
443 imbalance,
444 payload_bytes,
445 }
446 }
447
448 pub(crate) fn report_health(&self, index_label: &str, field_id: u32, segment_id: u128) {
453 let health = self.health();
454 let share = health.largest_cluster_share();
455 let fragmentation = health.fragmentation();
456 log::info!(
457 "[ann_health] index={index_label} field={field_id} segment={segment_id:016x}: \
458 vectors={} clusters={}/{} runs={} fragmentation={fragmentation:.2} \
459 imbalance={:.2} largest_leaf={:.2}% payload={}",
460 health.vectors,
461 health.clusters_nonempty,
462 health.clusters_total,
463 health.runs,
464 health.imbalance,
465 100.0 * share,
466 crate::format_bytes(health.payload_bytes),
467 );
468 crate::observe::ann_health(
469 index_label,
470 field_id,
471 health.imbalance,
472 fragmentation,
473 share,
474 );
475 if share >= ANN_SKEW_WARN_SHARE && health.vectors >= ANN_SKEW_WARN_MIN_VECTORS {
476 log::warn!(
477 "[ann_health] index={index_label} field={field_id} segment={segment_id:016x}: \
478 leaf {} holds {:.1}% of {} vectors — every query probing it scans that leaf \
479 in full; degenerate embeddings collapse into one leaf exactly like this",
480 health.largest_cluster,
481 100.0 * share,
482 health.vectors,
483 );
484 }
485 if fragmentation >= ANN_FRAGMENTATION_WARN {
486 log::warn!(
487 "[ann_health] index={index_label} field={field_id} segment={segment_id:016x}: \
488 {fragmentation:.1} extents per probed cluster ({} runs / {} clusters) — \
489 cold probes pay that many seeks; the next merge or vector-generation rewrite \
490 compacts to 1.0",
491 health.runs,
492 health.clusters_nonempty,
493 );
494 }
495 }
496
497 pub(crate) fn header(&self) -> &AnnDiskHeader {
498 &self.header
499 }
500
501 pub(crate) fn estimated_heap_bytes(&self) -> usize {
502 std::mem::size_of::<Self>() + self.runs.capacity() * std::mem::size_of::<AnnRun>()
503 }
504
505 #[cfg(feature = "native")]
506 pub(crate) fn pin_lookup_directory(
507 &mut self,
508 mode: crate::segment::pin::PinMode,
509 remaining: &mut u64,
510 report: &mut crate::segment::pin::PinReport,
511 ) {
512 let before = self.heap_pins.report();
513 self.heap_pins
514 .pin_slice(&self.runs, "ANN cluster-run directory", mode, remaining);
515 let after = self.heap_pins.report();
516 report.intended_bytes += after.intended_bytes - before.intended_bytes;
517 report.pinned_bytes += after.pinned_bytes - before.pinned_bytes;
518 report.skipped_budget_bytes += after.skipped_budget_bytes - before.skipped_budget_bytes;
519 report.failed_bytes += after.failed_bytes - before.failed_bytes;
520 report.heap_copy_bytes += after.heap_copy_bytes - before.heap_copy_bytes;
521 }
522
523 fn cluster_runs(&self, cluster_id: u32) -> &[AnnRun] {
524 let start = self.runs.partition_point(|run| run.cluster_id < cluster_id);
525 let end = self
526 .runs
527 .partition_point(|run| run.cluster_id <= cluster_id);
528 &self.runs[start..end]
529 }
530
531 #[cfg(feature = "native")]
532 fn ivf_tq_scan_tasks<'a>(
533 &'a self,
534 plan: &'a crate::structures::TqIvfQueryPlan,
535 block_bytes: usize,
536 chunk_blocks: usize,
537 ) -> Vec<IvfTqScanTask<'a>> {
538 debug_assert!(chunk_blocks > 0);
539 let mut tasks = Vec::new();
540 for (cluster_id, cluster_dot) in plan.cluster_dots() {
541 for run in self.cluster_runs(cluster_id) {
542 let block_count = run.codes.len() / block_bytes;
543 for first_block in (0..block_count).step_by(chunk_blocks) {
544 tasks.push(IvfTqScanTask {
545 run,
546 cluster_dot,
547 first_block,
548 last_block: (first_block + chunk_blocks).min(block_count),
549 });
550 }
551 }
552 }
553 tasks
554 }
555
556 #[cfg(feature = "native")]
557 fn binary_scan_tasks<'a>(&'a self, cluster_ids: &[u32]) -> Vec<BinaryScanTask<'a>> {
558 let mut tasks = Vec::new();
559 for &cluster_id in cluster_ids {
560 for run in self.cluster_runs(cluster_id) {
561 for first_index in (0..run.count).step_by(BINARY_SCORE_BATCH) {
562 tasks.push(BinaryScanTask {
563 run,
564 first_index,
565 count: BINARY_SCORE_BATCH.min(run.count - first_index),
566 });
567 }
568 }
569 }
570 tasks
571 }
572
573 #[cfg(feature = "native")]
580 fn prefetch_cluster_runs(&self, cluster_ids: &[u32]) {
581 if cluster_ids.is_empty() || !self.raw.is_mmap() {
582 return;
583 }
584 let mut ranges = Vec::with_capacity(cluster_ids.len());
585 for &cluster_id in cluster_ids {
586 ranges.extend(
587 self.cluster_runs(cluster_id)
588 .iter()
589 .map(|run| run.doc_ids.start..run.codes.end),
590 );
591 }
592 coalesce_prefetch_ranges(&mut ranges);
593 for range in ranges {
594 self.raw.madvise_range(range, libc::MADV_WILLNEED);
595 }
596 }
597
598 pub(crate) fn search_tq_combined_documents(
603 &self,
604 k: usize,
605 plan: &crate::structures::TqQueryPlan,
606 combiner: crate::query::MultiValueCombiner,
607 ) -> io::Result<Vec<AnnDocumentCandidate>> {
608 use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_block_bytes};
609
610 combiner
611 .validate()
612 .map_err(|message| io::Error::new(io::ErrorKind::InvalidInput, message))?;
613 if plan.padded_dim() != self.header.code_size * 2 {
614 return Err(io::Error::new(
615 io::ErrorKind::InvalidInput,
616 "TQ query plan does not match the payload dimension",
617 ));
618 }
619 if k == 0 {
620 return Ok(Vec::new());
621 }
622
623 let block_bytes = tq_block_bytes(self.header.code_size);
624 let bytes = self.raw.as_slice();
625 let mut top_documents = BoundedAnnCollector::<true, true>::new(k);
626 let mut ordinal_scores = Vec::new();
627 let mut scores = [0.0f32; TQ_BLOCK_LANES];
628
629 for run in &self.runs {
630 let mut current_doc = None;
631 let codes = &bytes[run.codes.clone()];
632 for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
633 crate::structures::vector::quantization::tq_score_block(plan, block, &mut scores);
634 let lane_base = block_index * TQ_BLOCK_LANES;
635 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
636 for (lane, &score) in scores.iter().enumerate().take(lanes) {
637 let index = lane_base + lane;
638 let doc_id = run_doc_id(bytes, run, index)?;
639 if current_doc.is_some_and(|previous| doc_id < previous) {
640 return Err(invalid_data("flat TQ run is not grouped by document ID"));
641 }
642 if current_doc.is_some_and(|previous| doc_id != previous) {
643 retain_combined_document(
644 &mut top_documents,
645 current_doc.expect("current document is present"),
646 &ordinal_scores,
647 combiner,
648 );
649 ordinal_scores.clear();
650 }
651 current_doc = Some(doc_id);
652 if score.is_finite() {
653 ordinal_scores.push((
654 u32::from(read_u16(bytes, run.ordinals.start + index * 2)),
655 score,
656 ));
657 }
658 }
659 }
660 if let Some(doc_id) = current_doc {
661 retain_combined_document(&mut top_documents, doc_id, &ordinal_scores, combiner);
662 ordinal_scores.clear();
663 }
664 }
665
666 Ok(top_documents
667 .into_sorted_results()
668 .into_iter()
669 .map(|(doc_id, _, score)| AnnDocumentCandidate { doc_id, score })
670 .collect())
671 }
672
673 pub(crate) fn search_ivf_tq_combined_documents(
681 &self,
682 k: usize,
683 plan: &crate::structures::TqIvfQueryPlan,
684 combiner: crate::query::MultiValueCombiner,
685 ) -> io::Result<Vec<AnnDocumentCandidate>> {
686 self.search_ivf_tq_combined_documents_with_tuning(
687 k,
688 plan,
689 combiner,
690 IVF_PARALLEL_SCAN_MIN_POSTINGS,
691 IVF_TQ_PARALLEL_SCAN_CHUNK_BLOCKS,
692 )
693 }
694
695 fn search_ivf_tq_combined_documents_with_tuning(
696 &self,
697 k: usize,
698 plan: &crate::structures::TqIvfQueryPlan,
699 combiner: crate::query::MultiValueCombiner,
700 parallel_min_postings: usize,
701 parallel_chunk_blocks: usize,
702 ) -> io::Result<Vec<AnnDocumentCandidate>> {
703 use crate::structures::vector::quantization::{
704 TQ_BLOCK_LANES, tq_ivf_block_bytes, tq_score_ivf_block,
705 };
706
707 validate_combined_search(combiner)?;
708 let tq_plan = plan.tq_plan();
709 self.validate_ivf_tq_query_plan(plan)?;
710 #[cfg(not(feature = "native"))]
711 let _ = (parallel_min_postings, parallel_chunk_blocks);
712 if k == 0 {
713 return Ok(Vec::new());
714 }
715 #[cfg(feature = "native")]
716 self.prefetch_cluster_runs(&plan.cluster_ids);
717
718 let block_bytes = tq_ivf_block_bytes(self.header.code_size);
719 let bytes = self.raw.as_slice();
720 let posting_count = probed_posting_count(self, &plan.cluster_ids)?;
721
722 #[cfg(feature = "native")]
723 if rayon::current_num_threads() > 1 && posting_count >= parallel_min_postings {
724 use rayon::prelude::*;
725 let tasks = self.ivf_tq_scan_tasks(plan, block_bytes, parallel_chunk_blocks);
726 let ordinal_scores = tasks
727 .par_iter()
728 .try_fold(
729 || Vec::with_capacity(parallel_chunk_blocks * TQ_BLOCK_LANES),
730 |mut ordinal_scores, task| {
731 score_ivf_tq_combined_blocks(
732 bytes,
733 tq_plan,
734 block_bytes,
735 *task,
736 &mut ordinal_scores,
737 )?;
738 Ok::<_, io::Error>(ordinal_scores)
739 },
740 )
741 .try_reduce(Vec::new, |mut left, mut right| {
742 left.append(&mut right);
743 Ok(left)
744 })?;
745 return Ok(combine_scored_ordinals(ordinal_scores, k, combiner));
746 }
747
748 let mut ordinal_scores = Vec::new();
749 ordinal_scores
750 .try_reserve_exact(posting_count)
751 .map_err(|_| invalid_data("IVF-TQ combined score buffer allocation failed"))?;
752 let mut scores = [0.0f32; TQ_BLOCK_LANES];
753 for (cluster_id, cluster_dot) in plan.cluster_dots() {
754 for run in self.cluster_runs(cluster_id) {
755 let codes = &bytes[run.codes.clone()];
756 for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
757 tq_score_ivf_block(tq_plan, block, cluster_dot, &mut scores);
758 let lane_base = block_index * TQ_BLOCK_LANES;
759 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
760 for (lane, &score) in scores.iter().enumerate().take(lanes) {
761 if !score.is_finite() {
762 continue;
763 }
764 let index = lane_base + lane;
765 ordinal_scores.push((
766 run_doc_id(bytes, run, index)?,
767 read_u16(bytes, run.ordinals.start + index * 2),
768 score,
769 ));
770 }
771 }
772 }
773 }
774 Ok(combine_scored_ordinals(ordinal_scores, k, combiner))
775 }
776
777 pub(crate) fn search_binary_combined_documents(
786 &self,
787 k: usize,
788 query: &[u8],
789 cluster_ids: &[u32],
790 combiner: crate::query::MultiValueCombiner,
791 ) -> io::Result<CombinedBinaryCandidates> {
792 self.search_binary_combined_documents_with_tuning(
793 k,
794 query,
795 cluster_ids,
796 combiner,
797 IVF_PARALLEL_SCAN_MIN_POSTINGS,
798 )
799 }
800
801 fn search_binary_combined_documents_with_tuning(
802 &self,
803 k: usize,
804 query: &[u8],
805 cluster_ids: &[u32],
806 combiner: crate::query::MultiValueCombiner,
807 parallel_min_postings: usize,
808 ) -> io::Result<CombinedBinaryCandidates> {
809 validate_combined_search(combiner)?;
810 #[cfg(not(feature = "native"))]
811 let _ = parallel_min_postings;
812 if query.len() != self.header.code_size {
813 return Err(io::Error::new(
814 io::ErrorKind::InvalidInput,
815 "binary ANN query has the wrong byte length",
816 ));
817 }
818 if k == 0 {
819 return Ok((Vec::new(), Vec::new()));
820 }
821 #[cfg(feature = "native")]
822 self.prefetch_cluster_runs(cluster_ids);
823
824 let bytes = self.raw.as_slice();
825 let posting_count = probed_posting_count(self, cluster_ids)?;
826 #[cfg(feature = "native")]
827 if rayon::current_num_threads() > 1 && posting_count >= parallel_min_postings {
828 use rayon::prelude::*;
829 let tasks = self.binary_scan_tasks(cluster_ids);
830 let (ordinal_scores, _) = tasks
831 .par_iter()
832 .try_fold(
833 || {
834 (
835 Vec::with_capacity(BINARY_SCORE_BATCH),
836 vec![0.0f32; BINARY_SCORE_BATCH],
837 )
838 },
839 |(mut ordinal_scores, mut scores), task| {
840 score_binary_task(
841 bytes,
842 query,
843 self.header.dim,
844 self.header.code_size,
845 *task,
846 &mut scores,
847 &mut ordinal_scores,
848 )?;
849 Ok::<_, io::Error>((ordinal_scores, scores))
850 },
851 )
852 .try_reduce(
853 || (Vec::new(), Vec::new()),
854 |(mut left, scores), (mut right, _)| {
855 left.append(&mut right);
856 Ok((left, scores))
857 },
858 )?;
859 return Ok(combine_scored_ordinals_retaining(
860 ordinal_scores,
861 k,
862 combiner,
863 ));
864 }
865
866 let mut score_batch = vec![0.0f32; BINARY_SCORE_BATCH.min(self.header.vector_count)];
867 let mut ordinal_scores = Vec::new();
868 ordinal_scores
869 .try_reserve_exact(posting_count)
870 .map_err(|_| invalid_data("binary combined score buffer allocation failed"))?;
871 score_binary_cluster_runs(
872 self,
873 bytes,
874 query,
875 cluster_ids,
876 &mut score_batch,
877 &mut ordinal_scores,
878 )?;
879 Ok(combine_scored_ordinals_retaining(
880 ordinal_scores,
881 k,
882 combiner,
883 ))
884 }
885
886 pub(crate) fn search_tq_distinct(
889 &self,
890 k: usize,
891 plan: &crate::structures::TqQueryPlan,
892 ) -> io::Result<Vec<(u32, u16, f32)>> {
893 use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_block_bytes};
894
895 if plan.padded_dim() != self.header.code_size * 2 {
896 return Err(io::Error::new(
897 io::ErrorKind::InvalidInput,
898 "TQ query plan does not match the payload dimension",
899 ));
900 }
901 let block_bytes = tq_block_bytes(self.header.code_size);
902 let bytes = self.raw.as_slice();
903
904 #[cfg(feature = "native")]
909 if self.header.vector_count >= TQ_PARALLEL_SCAN_MIN_VECTORS {
910 use rayon::prelude::*;
911 let collector = self
912 .runs
913 .par_iter()
914 .flat_map(|run| {
915 let codes = &bytes[run.codes.clone()];
916 let blocks = codes.len() / block_bytes;
917 (0..blocks.div_ceil(TQ_PARALLEL_SCAN_CHUNK_BLOCKS))
918 .into_par_iter()
919 .map(move |chunk| (run, chunk * TQ_PARALLEL_SCAN_CHUNK_BLOCKS, blocks))
920 })
921 .try_fold(
922 || BoundedAnnCollector::<true, true>::new(k),
923 |mut collector, (run, first_block, total_blocks)| {
924 let codes = &bytes[run.codes.clone()];
925 let last_block =
926 (first_block + TQ_PARALLEL_SCAN_CHUNK_BLOCKS).min(total_blocks);
927 let mut scores = [0.0f32; TQ_BLOCK_LANES];
928 for block_index in first_block..last_block {
929 let block = &codes[block_index * block_bytes..][..block_bytes];
930 crate::structures::vector::quantization::tq_score_block(
931 plan,
932 block,
933 &mut scores,
934 );
935 let lane_base = block_index * TQ_BLOCK_LANES;
936 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
937 for (lane, &score) in scores.iter().enumerate().take(lanes) {
938 let index = lane_base + lane;
939 collector.insert(
940 run_doc_id(bytes, run, index)?,
941 read_u16(bytes, run.ordinals.start + index * 2),
942 score,
943 );
944 }
945 }
946 Ok::<_, io::Error>(collector)
947 },
948 )
949 .try_reduce(
950 || BoundedAnnCollector::<true, true>::new(k),
951 |mut collector, partial| {
952 collector.merge_from(partial);
953 Ok(collector)
954 },
955 )?;
956 return Ok(collector.into_sorted_results());
957 }
958
959 let mut collector = BoundedAnnCollector::<true, true>::new(k);
960 let mut scores = [0.0f32; TQ_BLOCK_LANES];
961 for run in &self.runs {
962 let codes = &bytes[run.codes.clone()];
963 for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
964 crate::structures::vector::quantization::tq_score_block(plan, block, &mut scores);
965 let lane_base = block_index * TQ_BLOCK_LANES;
966 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
967 for (lane, &score) in scores.iter().enumerate().take(lanes) {
968 let index = lane_base + lane;
969 collector.insert(
970 run_doc_id(bytes, run, index)?,
971 read_u16(bytes, run.ordinals.start + index * 2),
972 score,
973 );
974 }
975 }
976 }
977 Ok(collector.into_sorted_results())
978 }
979
980 pub(crate) fn search_ivf_tq_distinct(
988 &self,
989 k: usize,
990 plan: &crate::structures::TqIvfQueryPlan,
991 ) -> io::Result<Vec<(u32, u16, f32)>> {
992 self.search_ivf_tq_distinct_with_tuning(
993 k,
994 plan,
995 IVF_PARALLEL_SCAN_MIN_POSTINGS,
996 IVF_TQ_PARALLEL_SCAN_CHUNK_BLOCKS,
997 )
998 }
999
1000 fn search_ivf_tq_distinct_with_tuning(
1001 &self,
1002 k: usize,
1003 plan: &crate::structures::TqIvfQueryPlan,
1004 parallel_min_postings: usize,
1005 parallel_chunk_blocks: usize,
1006 ) -> io::Result<Vec<(u32, u16, f32)>> {
1007 use crate::structures::vector::quantization::tq_ivf_block_bytes;
1008
1009 let tq_plan = plan.tq_plan();
1010 self.validate_ivf_tq_query_plan(plan)?;
1011 #[cfg(not(feature = "native"))]
1012 let _ = (parallel_min_postings, parallel_chunk_blocks);
1013 if k == 0 {
1014 return Ok(Vec::new());
1015 }
1016 #[cfg(feature = "native")]
1017 self.prefetch_cluster_runs(&plan.cluster_ids);
1018 let block_bytes = tq_ivf_block_bytes(self.header.code_size);
1019 let bytes = self.raw.as_slice();
1020
1021 #[cfg(feature = "native")]
1026 if rayon::current_num_threads() > 1
1027 && probed_posting_count(self, &plan.cluster_ids)? >= parallel_min_postings
1028 {
1029 use rayon::prelude::*;
1030 let tasks = self.ivf_tq_scan_tasks(plan, block_bytes, parallel_chunk_blocks);
1031
1032 if let Some((pilot, remaining)) = tasks.split_first() {
1033 let mut pilot_collector = BoundedAnnCollector::<true, true>::new(k);
1034 let (pilot_pruned, pilot_scored) =
1035 score_ivf_tq_blocks(bytes, tq_plan, block_bytes, *pilot, &mut pilot_collector)?;
1036 let seed = pilot_collector.into_sorted_results();
1037 let seeded_collector = || {
1038 let mut collector = BoundedAnnCollector::<true, true>::new(k);
1039 for &(doc_id, ordinal, score) in &seed {
1040 collector.insert(doc_id, ordinal, score);
1041 }
1042 collector
1043 };
1044 let (collector, pruned_blocks, scored_blocks) = remaining
1045 .par_iter()
1046 .try_fold(
1047 || (seeded_collector(), 0usize, 0usize),
1048 |(mut collector, pruned, scored), task| {
1049 let (task_pruned, task_scored) = score_ivf_tq_blocks(
1050 bytes,
1051 tq_plan,
1052 block_bytes,
1053 *task,
1054 &mut collector,
1055 )?;
1056 Ok::<_, io::Error>((
1057 collector,
1058 pruned + task_pruned,
1059 scored + task_scored,
1060 ))
1061 },
1062 )
1063 .try_reduce(
1064 || (seeded_collector(), 0usize, 0usize),
1065 |(mut collector, left_pruned, left_scored),
1066 (partial, right_pruned, right_scored)| {
1067 collector.merge_from(partial);
1068 Ok((
1069 collector,
1070 left_pruned + right_pruned,
1071 left_scored + right_scored,
1072 ))
1073 },
1074 )?;
1075 log_ivf_tq_pruning(
1076 pilot_pruned + pruned_blocks,
1077 pilot_scored + scored_blocks,
1078 true,
1079 );
1080 return Ok(collector.into_sorted_results());
1081 }
1082 }
1083
1084 let mut collector = BoundedAnnCollector::<true, true>::new(k);
1085 let mut pruned_blocks = 0usize;
1086 let mut scored_blocks = 0usize;
1087 for (cluster_id, cluster_dot) in plan.cluster_dots() {
1088 for run in self.cluster_runs(cluster_id) {
1089 let block_count = run.codes.len() / block_bytes;
1090 let (run_pruned, run_scored) = score_ivf_tq_blocks(
1091 bytes,
1092 tq_plan,
1093 block_bytes,
1094 IvfTqScanTask {
1095 run,
1096 cluster_dot,
1097 first_block: 0,
1098 last_block: block_count,
1099 },
1100 &mut collector,
1101 )?;
1102 pruned_blocks += run_pruned;
1103 scored_blocks += run_scored;
1104 }
1105 }
1106 log_ivf_tq_pruning(pruned_blocks, scored_blocks, false);
1107 Ok(collector.into_sorted_results())
1108 }
1109
1110 fn validate_ivf_tq_query_plan(
1111 &self,
1112 plan: &crate::structures::TqIvfQueryPlan,
1113 ) -> io::Result<()> {
1114 if self.header.kind != AnnKind::IvfTq
1115 || !crate::structures::is_ivf_tq_cosine_generation(self.header.quantizer_version)
1116 {
1117 return Err(invalid_data(
1118 "legacy raw IVF-TQ payloads cannot be searched; rebuild the index",
1119 ));
1120 }
1121 if plan.tq_plan().padded_dim() != self.header.code_size * 2
1122 || plan.quantizer_version != self.header.quantizer_version
1123 || plan.fingerprint != self.header.codebook_version
1124 {
1125 return Err(io::Error::new(
1126 io::ErrorKind::InvalidInput,
1127 "IVF-TQ query plan does not match the payload generation",
1128 ));
1129 }
1130 Ok(())
1131 }
1132
1133 pub(crate) fn search_binary_clusters<const BY_DOCUMENT: bool>(
1134 &self,
1135 query: &[u8],
1136 k: usize,
1137 cluster_ids: &[u32],
1138 ) -> io::Result<Vec<(u32, u16, f32)>> {
1139 self.search_binary_clusters_with_tuning::<BY_DOCUMENT>(
1140 query,
1141 k,
1142 cluster_ids,
1143 IVF_PARALLEL_SCAN_MIN_POSTINGS,
1144 )
1145 }
1146
1147 fn search_binary_clusters_with_tuning<const BY_DOCUMENT: bool>(
1148 &self,
1149 query: &[u8],
1150 k: usize,
1151 cluster_ids: &[u32],
1152 parallel_min_postings: usize,
1153 ) -> io::Result<Vec<(u32, u16, f32)>> {
1154 #[cfg(not(feature = "native"))]
1155 let _ = parallel_min_postings;
1156 if query.len() != self.header.code_size {
1157 return Err(io::Error::new(
1158 io::ErrorKind::InvalidInput,
1159 "binary ANN query has the wrong byte length",
1160 ));
1161 }
1162 if k == 0 {
1163 return Ok(Vec::new());
1164 }
1165 #[cfg(feature = "native")]
1166 self.prefetch_cluster_runs(cluster_ids);
1167 let bytes = self.raw.as_slice();
1168
1169 debug_assert!(
1173 BY_DOCUMENT || {
1174 let mut seen = rustc_hash::FxHashSet::default();
1175 cluster_ids
1176 .iter()
1177 .all(|cluster_id| seen.insert(*cluster_id))
1178 },
1179 "an IVF probe plan must not repeat a cluster",
1180 );
1181
1182 #[cfg(feature = "native")]
1183 if rayon::current_num_threads() > 1
1184 && probed_posting_count(self, cluster_ids)? >= parallel_min_postings
1185 {
1186 use rayon::prelude::*;
1187 let tasks = self.binary_scan_tasks(cluster_ids);
1188 let (collector, _) = tasks
1189 .par_iter()
1190 .try_fold(
1191 || {
1192 (
1193 BoundedAnnCollector::<BY_DOCUMENT, true>::new(k),
1194 vec![0.0f32; BINARY_SCORE_BATCH],
1195 )
1196 },
1197 |(mut collector, mut scores), task| {
1198 score_binary_task(
1199 bytes,
1200 query,
1201 self.header.dim,
1202 self.header.code_size,
1203 *task,
1204 &mut scores,
1205 &mut collector,
1206 )?;
1207 Ok::<_, io::Error>((collector, scores))
1208 },
1209 )
1210 .try_reduce(
1211 || (BoundedAnnCollector::<BY_DOCUMENT, true>::new(k), Vec::new()),
1212 |(mut collector, scores), (partial, _)| {
1213 collector.merge_from(partial);
1214 Ok((collector, scores))
1215 },
1216 )?;
1217 return Ok(collector.into_sorted_results());
1218 }
1219
1220 let mut scores = vec![0.0f32; BINARY_SCORE_BATCH.min(self.header.vector_count)];
1221 if BY_DOCUMENT {
1222 let mut collector = BoundedAnnCollector::<true, true>::new(k);
1223 score_binary_cluster_runs(
1224 self,
1225 bytes,
1226 query,
1227 cluster_ids,
1228 &mut scores,
1229 &mut collector,
1230 )?;
1231 return Ok(collector.into_sorted_results());
1232 }
1233
1234 let mut collector = BoundedUniqueAnnCollector::<true>::new(k);
1236 score_binary_cluster_runs(self, bytes, query, cluster_ids, &mut scores, &mut collector)?;
1237 Ok(collector.into_sorted_results())
1238 }
1239}
1240
1241fn score_ivf_tq_blocks(
1245 bytes: &[u8],
1246 plan: &crate::structures::TqQueryPlan,
1247 block_bytes: usize,
1248 task: IvfTqScanTask<'_>,
1249 collector: &mut BoundedAnnCollector<true, true>,
1250) -> io::Result<(usize, usize)> {
1251 use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_score_ivf_block};
1252
1253 let run = task.run;
1254 let codes = &bytes[run.codes.clone()];
1255 let mut scores = [0.0f32; TQ_BLOCK_LANES];
1256 let mut scored_blocks = 0usize;
1257 for block_index in task.first_block..task.last_block {
1258 let block = &codes[block_index * block_bytes..][..block_bytes];
1259 if let Some(threshold) = collector.pruning_threshold()
1260 && task.cluster_dot + tq_ivf_block_max_scale(block) * TQ_PRUNE_ESTIMATE_BOUND
1261 <= threshold
1262 {
1263 return Ok((task.last_block - block_index, scored_blocks));
1264 }
1265 scored_blocks += 1;
1266 tq_score_ivf_block(plan, block, task.cluster_dot, &mut scores);
1267 let lane_base = block_index * TQ_BLOCK_LANES;
1268 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
1269 for (lane, &score) in scores.iter().enumerate().take(lanes) {
1270 let index = lane_base + lane;
1271 collector.insert(
1272 run_doc_id(bytes, run, index)?,
1273 read_u16(bytes, run.ordinals.start + index * 2),
1274 score,
1275 );
1276 }
1277 }
1278 Ok((0, scored_blocks))
1279}
1280
1281#[cfg(feature = "native")]
1282fn score_ivf_tq_combined_blocks(
1283 bytes: &[u8],
1284 plan: &crate::structures::TqQueryPlan,
1285 block_bytes: usize,
1286 task: IvfTqScanTask<'_>,
1287 ordinal_scores: &mut Vec<(u32, u16, f32)>,
1288) -> io::Result<()> {
1289 use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_score_ivf_block};
1290
1291 let run = task.run;
1292 let codes = &bytes[run.codes.clone()];
1293 let mut scores = [0.0f32; TQ_BLOCK_LANES];
1294 for block_index in task.first_block..task.last_block {
1295 let block = &codes[block_index * block_bytes..][..block_bytes];
1296 tq_score_ivf_block(plan, block, task.cluster_dot, &mut scores);
1297 let lane_base = block_index * TQ_BLOCK_LANES;
1298 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
1299 for (lane, &score) in scores.iter().enumerate().take(lanes) {
1300 if score.is_finite() {
1301 let index = lane_base + lane;
1302 ordinal_scores.push((
1303 run_doc_id(bytes, run, index)?,
1304 read_u16(bytes, run.ordinals.start + index * 2),
1305 score,
1306 ));
1307 }
1308 }
1309 }
1310 Ok(())
1311}
1312
1313fn log_ivf_tq_pruning(pruned_blocks: usize, scored_blocks: usize, parallel: bool) {
1314 if pruned_blocks > 0 {
1315 log::debug!(
1316 "[search_ivf_tq] pruned {pruned_blocks} of {} blocks via scale bounds ({})",
1317 pruned_blocks + scored_blocks,
1318 if parallel { "parallel" } else { "serial" },
1319 );
1320 }
1321}
1322
1323#[cfg(feature = "native")]
1324fn coalesce_prefetch_ranges(ranges: &mut Vec<Range<usize>>) {
1325 if ranges.len() < 2 {
1326 return;
1327 }
1328 ranges.sort_unstable_by_key(|range| range.start);
1329 let mut output_len = 1usize;
1330 for input_index in 1..ranges.len() {
1331 let next_start = ranges[input_index].start;
1332 let next_end = ranges[input_index].end;
1333 let previous = &mut ranges[output_len - 1];
1334 if next_start <= previous.end.saturating_add(PREFETCH_COALESCE_GAP) {
1335 previous.end = previous.end.max(next_end);
1336 } else {
1337 ranges[output_len] = next_start..next_end;
1338 output_len += 1;
1339 }
1340 }
1341 ranges.truncate(output_len);
1342}
1343
1344trait AnnScoreSink {
1345 fn insert_score(&mut self, doc_id: u32, ordinal: u16, score: f32);
1346}
1347
1348impl<const BY_DOCUMENT: bool> AnnScoreSink for BoundedAnnCollector<BY_DOCUMENT, true> {
1349 #[inline]
1350 fn insert_score(&mut self, doc_id: u32, ordinal: u16, score: f32) {
1351 self.insert(doc_id, ordinal, score);
1352 }
1353}
1354
1355impl AnnScoreSink for BoundedUniqueAnnCollector<true> {
1356 #[inline]
1357 fn insert_score(&mut self, doc_id: u32, ordinal: u16, score: f32) {
1358 self.insert(doc_id, ordinal, score);
1359 }
1360}
1361
1362impl AnnScoreSink for Vec<(u32, u16, f32)> {
1363 #[inline]
1364 fn insert_score(&mut self, doc_id: u32, ordinal: u16, score: f32) {
1365 if score.is_finite() {
1366 self.push((doc_id, ordinal, score));
1367 }
1368 }
1369}
1370
1371fn probed_posting_count(index: &AnnDiskIndex, cluster_ids: &[u32]) -> io::Result<usize> {
1372 cluster_ids.iter().try_fold(0usize, |count, &cluster_id| {
1373 index
1374 .cluster_runs(cluster_id)
1375 .iter()
1376 .try_fold(count, |count, run| {
1377 count
1378 .checked_add(run.count)
1379 .ok_or_else(|| invalid_data("ANN probed posting count overflows usize"))
1380 })
1381 })
1382}
1383
1384fn score_binary_cluster_runs(
1385 index: &AnnDiskIndex,
1386 bytes: &[u8],
1387 query: &[u8],
1388 cluster_ids: &[u32],
1389 scores: &mut [f32],
1390 collector: &mut impl AnnScoreSink,
1391) -> io::Result<()> {
1392 for &cluster_id in cluster_ids {
1393 for run in index.cluster_runs(cluster_id) {
1394 score_binary_run(
1395 bytes,
1396 run,
1397 query,
1398 index.header.dim,
1399 index.header.code_size,
1400 scores,
1401 collector,
1402 )?;
1403 }
1404 }
1405 Ok(())
1406}
1407
1408fn score_binary_run(
1409 bytes: &[u8],
1410 run: &AnnRun,
1411 query: &[u8],
1412 dim_bits: usize,
1413 code_size: usize,
1414 scores: &mut [f32],
1415 collector: &mut impl AnnScoreSink,
1416) -> io::Result<()> {
1417 for batch_start in (0..run.count).step_by(BINARY_SCORE_BATCH) {
1418 score_binary_task(
1419 bytes,
1420 query,
1421 dim_bits,
1422 code_size,
1423 BinaryScanTask {
1424 run,
1425 first_index: batch_start,
1426 count: BINARY_SCORE_BATCH.min(run.count - batch_start),
1427 },
1428 scores,
1429 collector,
1430 )?;
1431 }
1432 Ok(())
1433}
1434
1435fn score_binary_task(
1436 bytes: &[u8],
1437 query: &[u8],
1438 dim_bits: usize,
1439 code_size: usize,
1440 task: BinaryScanTask<'_>,
1441 scores: &mut [f32],
1442 collector: &mut impl AnnScoreSink,
1443) -> io::Result<()> {
1444 let code_start = task.run.codes.start + task.first_index * code_size;
1445 let code_end = code_start + task.count * code_size;
1446 crate::structures::simd::batch_hamming_scores(
1447 query,
1448 &bytes[code_start..code_end],
1449 code_size,
1450 dim_bits,
1451 &mut scores[..task.count],
1452 );
1453 for (batch_index, &score) in scores.iter().enumerate().take(task.count) {
1454 let index = task.first_index + batch_index;
1455 collector.insert_score(
1456 run_doc_id(bytes, task.run, index)?,
1457 read_u16(bytes, task.run.ordinals.start + index * 2),
1458 score,
1459 );
1460 }
1461 Ok(())
1462}
1463
1464#[inline]
1465fn retain_combined_document(
1466 collector: &mut BoundedAnnCollector<true, true>,
1467 doc_id: u32,
1468 ordinal_scores: &[(u32, f32)],
1469 combiner: crate::query::MultiValueCombiner,
1470) {
1471 if !ordinal_scores.is_empty() {
1472 collector.insert(doc_id, 0, combiner.combine(ordinal_scores));
1473 }
1474}
1475
1476fn validate_combined_search(combiner: crate::query::MultiValueCombiner) -> io::Result<()> {
1477 combiner
1478 .validate()
1479 .map_err(|message| io::Error::new(io::ErrorKind::InvalidInput, message))
1480}
1481
1482fn combine_scored_ordinals(
1488 scores: Vec<(u32, u16, f32)>,
1489 k: usize,
1490 combiner: crate::query::MultiValueCombiner,
1491) -> Vec<AnnDocumentCandidate> {
1492 combine_scored_ordinals_retaining(scores, k, combiner).0
1493}
1494
1495fn combine_scored_ordinals_retaining(
1501 mut scores: Vec<(u32, u16, f32)>,
1502 k: usize,
1503 combiner: crate::query::MultiValueCombiner,
1504) -> CombinedBinaryCandidates {
1505 if k == 0 || scores.is_empty() {
1506 return (Vec::new(), Vec::new());
1507 }
1508 scores.retain(|entry| entry.2.is_finite());
1509 let by_document_ordinal = |left: &(u32, u16, f32), right: &(u32, u16, f32)| {
1510 left.0.cmp(&right.0).then_with(|| left.1.cmp(&right.1))
1511 };
1512 #[cfg(feature = "native")]
1513 if rayon::current_num_threads() > 1 && scores.len() >= IVF_PARALLEL_SCAN_MIN_POSTINGS {
1514 use rayon::prelude::*;
1515 scores.par_sort_unstable_by(by_document_ordinal);
1516 } else {
1517 scores.sort_unstable_by(by_document_ordinal);
1518 }
1519 #[cfg(not(feature = "native"))]
1520 scores.sort_unstable_by(by_document_ordinal);
1521
1522 let mut unique_len = 0usize;
1527 for read_index in 0..scores.len() {
1528 let candidate = scores[read_index];
1529 if unique_len > 0
1530 && scores[unique_len - 1].0 == candidate.0
1531 && scores[unique_len - 1].1 == candidate.1
1532 {
1533 if candidate.2.total_cmp(&scores[unique_len - 1].2).is_gt() {
1534 scores[unique_len - 1].2 = candidate.2;
1535 }
1536 } else {
1537 scores[unique_len] = candidate;
1538 unique_len += 1;
1539 }
1540 }
1541 scores.truncate(unique_len);
1542
1543 let mut top_documents = BoundedAnnCollector::<true, true>::new(k);
1544 let mut current_doc = None;
1545 let mut ordinal_scores = Vec::new();
1546 for &(doc_id, ordinal, score) in &scores {
1547 if current_doc.is_some_and(|current| current != doc_id) {
1548 retain_combined_document(
1549 &mut top_documents,
1550 current_doc.expect("current document is present"),
1551 &ordinal_scores,
1552 combiner,
1553 );
1554 ordinal_scores.clear();
1555 }
1556 current_doc = Some(doc_id);
1557 ordinal_scores.push((u32::from(ordinal), score));
1558 }
1559 if let Some(doc_id) = current_doc {
1560 retain_combined_document(&mut top_documents, doc_id, &ordinal_scores, combiner);
1561 }
1562
1563 let candidates: Vec<AnnDocumentCandidate> = top_documents
1564 .into_sorted_results()
1565 .into_iter()
1566 .map(|(doc_id, _, score)| AnnDocumentCandidate { doc_id, score })
1567 .collect();
1568
1569 let mut retained_ids: Vec<u32> = candidates
1573 .iter()
1574 .map(|candidate| candidate.doc_id)
1575 .collect();
1576 retained_ids.sort_unstable();
1577 let mut retained = Vec::with_capacity(scores.len().min(retained_ids.len().saturating_mul(4)));
1578 let mut cursor = 0usize;
1579 for entry in scores {
1580 while cursor < retained_ids.len() && retained_ids[cursor] < entry.0 {
1581 cursor += 1;
1582 }
1583 if retained_ids.get(cursor) == Some(&entry.0) {
1584 retained.push(entry);
1585 }
1586 }
1587 (candidates, retained)
1588}
1589
1590#[cfg(feature = "native")]
1591pub(crate) fn write_built_binary_ivf(
1592 index: &BinaryIvfIndex,
1593 routing: IvfRoutingMode,
1594 writer: &mut (impl Write + ?Sized),
1595) -> io::Result<u64> {
1596 let runs: Vec<_> = index
1597 .clusters
1598 .iter()
1599 .map(|(cluster_id, cluster)| BuildRun {
1600 cluster_id: *cluster_id,
1601 doc_ids: &cluster.doc_ids,
1602 ordinals: &cluster.ordinals,
1603 codes: &cluster.codes,
1604 })
1605 .collect();
1606 write_built_runs(
1607 AnnDiskHeader {
1608 kind: AnnKind::BinaryIvf,
1609 routing,
1610 dim: index.dim_bits,
1611 code_size: index.dim_bits.div_ceil(8),
1612 num_clusters: index.num_clusters,
1613 quantizer_version: index.quantizer_version,
1614 codebook_version: 0,
1615 vector_count: index.len(),
1616 },
1617 &runs,
1618 writer,
1619 )
1620}
1621
1622#[cfg(feature = "native")]
1625pub(crate) fn write_built_ivf_tq(
1626 index: &crate::structures::IvfTqIndex,
1627 num_clusters: u32,
1628 writer: &mut (impl Write + ?Sized),
1629) -> io::Result<u64> {
1630 use crate::structures::vector::quantization::{TQ_BLOCK_LANES, tq_pack_ivf_block};
1631
1632 if !crate::structures::is_ivf_tq_cosine_generation(index.centroids_version) {
1633 return Err(invalid_data(
1634 "legacy raw IVF-TQ generations cannot be serialized; rebuild the index",
1635 ));
1636 }
1637 let codec = index.codec();
1638 let padded_dim = codec.padded_dim();
1639 let mut clusters: Vec<_> = index.clusters.iter().collect();
1640 clusters.sort_unstable_by_key(|(cluster_id, _)| **cluster_id);
1641 struct PackedCluster {
1645 cluster_id: u32,
1646 doc_ids: Vec<u32>,
1647 ordinals: Vec<u16>,
1648 codes: Vec<u8>,
1649 }
1650 let packed: Vec<PackedCluster> = clusters
1651 .iter()
1652 .map(|&(&cluster_id, cluster)| {
1653 let count = cluster.doc_ids.len();
1654 let mut order: Vec<usize> = (0..count).collect();
1655 order.sort_by(|&a, &b| {
1656 cluster.scales[b]
1657 .total_cmp(&cluster.scales[a])
1658 .then_with(|| a.cmp(&b))
1659 });
1660 let doc_ids: Vec<u32> = order.iter().map(|&i| cluster.doc_ids[i]).collect();
1661 let ordinals: Vec<u16> = order.iter().map(|&i| cluster.ordinals[i]).collect();
1662 let scales: Vec<f32> = order.iter().map(|&i| cluster.scales[i]).collect();
1663 let gammas: Vec<f32> = order.iter().map(|&i| cluster.gammas[i]).collect();
1664 let mut codes = Vec::with_capacity(
1665 crate::structures::vector::quantization::tq_ivf_codes_column_len_checked(
1666 count,
1667 codec.code_size(),
1668 )
1669 .unwrap_or_default(),
1670 );
1671 for block_start in (0..count).step_by(TQ_BLOCK_LANES) {
1672 let lanes = TQ_BLOCK_LANES.min(count - block_start);
1673 let rows: Vec<&[u8]> = order[block_start..block_start + lanes]
1674 .iter()
1675 .map(|&row| &cluster.rows[row * padded_dim..(row + 1) * padded_dim])
1676 .collect();
1677 tq_pack_ivf_block(
1678 &rows,
1679 &scales[block_start..block_start + lanes],
1680 &gammas[block_start..block_start + lanes],
1681 padded_dim,
1682 &mut codes,
1683 );
1684 }
1685 PackedCluster {
1686 cluster_id,
1687 doc_ids,
1688 ordinals,
1689 codes,
1690 }
1691 })
1692 .collect();
1693 let runs: Vec<_> = packed
1694 .iter()
1695 .map(|cluster| BuildRun {
1696 cluster_id: cluster.cluster_id,
1697 doc_ids: &cluster.doc_ids,
1698 ordinals: &cluster.ordinals,
1699 codes: &cluster.codes,
1700 })
1701 .collect();
1702 write_built_runs(
1703 AnnDiskHeader {
1704 kind: AnnKind::IvfTq,
1705 routing: index.routing,
1706 dim: index.dim,
1707 code_size: codec.code_size(),
1708 num_clusters,
1709 quantizer_version: index.centroids_version,
1710 codebook_version: codec.fingerprint(),
1711 vector_count: index.len(),
1712 },
1713 &runs,
1714 writer,
1715 )
1716}
1717
1718#[cfg(feature = "native")]
1720pub(crate) fn tq_builder_extra_run(builder: &crate::structures::TqFlatBuilder) -> BuildRun<'_> {
1721 BuildRun {
1722 cluster_id: 0,
1723 doc_ids: &builder.doc_ids,
1724 ordinals: &builder.ordinals,
1725 codes: &builder.codes,
1726 }
1727}
1728
1729#[cfg(feature = "native")]
1731pub(crate) fn write_built_tq_flat(
1732 builder: &crate::structures::TqFlatBuilder,
1733 writer: &mut (impl Write + ?Sized),
1734) -> io::Result<u64> {
1735 let codec = builder.codec();
1736 let runs = [BuildRun {
1737 cluster_id: 0,
1738 doc_ids: &builder.doc_ids,
1739 ordinals: &builder.ordinals,
1740 codes: &builder.codes,
1741 }];
1742 write_built_runs(
1743 AnnDiskHeader {
1744 kind: AnnKind::TqFlat,
1745 routing: IvfRoutingMode::Flat,
1746 dim: codec.dim(),
1747 code_size: codec.code_size(),
1748 num_clusters: 1,
1749 quantizer_version: codec.fingerprint(),
1750 codebook_version: 0,
1751 vector_count: builder.len(),
1752 },
1753 &runs,
1754 writer,
1755 )
1756}
1757
1758#[cfg(feature = "native")]
1759pub(crate) struct BuildRun<'a> {
1760 cluster_id: u32,
1761 doc_ids: &'a [u32],
1762 ordinals: &'a [u16],
1763 codes: &'a [u8],
1764}
1765
1766#[cfg(feature = "native")]
1767struct RunRecord {
1768 cluster_id: u32,
1769 doc_base: u32,
1770 count: u32,
1771 max_doc_id: u32,
1772 doc_ids_offset: u64,
1773 ordinals_offset: u64,
1774 codes_offset: u64,
1775 codes_len: u64,
1776}
1777
1778#[cfg(feature = "native")]
1779fn write_built_runs(
1780 header: AnnDiskHeader,
1781 runs: &[BuildRun<'_>],
1782 writer: &mut (impl Write + ?Sized),
1783) -> io::Result<u64> {
1784 if runs.is_empty() || header.vector_count == 0 {
1785 return Err(invalid_data("cannot write an empty ANN payload"));
1786 }
1787 validate_header(&header)?;
1788 write_header(writer, &header)?;
1789 let mut offset = ANN_HEADER_SIZE as u64;
1790 let mut records = Vec::with_capacity(runs.len());
1791 let mut counted = 0usize;
1792 let mut scratch = Vec::new();
1793 let mut previous_cluster = None;
1794 for run in runs {
1795 let count = run.doc_ids.len();
1796 if count == 0
1797 || run.cluster_id >= header.num_clusters
1798 || previous_cluster.is_some_and(|cluster| cluster >= run.cluster_id)
1799 || run.ordinals.len() != count
1800 || run.codes.len() != expected_codes_column_len(header.kind, count, header.code_size)?
1801 {
1802 return Err(invalid_data("ANN build run columns are inconsistent"));
1803 }
1804 previous_cluster = Some(run.cluster_id);
1805 let count_u32 = u32::try_from(count)
1806 .map_err(|_| invalid_data("ANN cluster run exceeds u32 vectors"))?;
1807 let max_doc_id = run.doc_ids.iter().copied().max().unwrap_or(0);
1808 let doc_ids_offset = offset;
1809 write_u32_column(writer, run.doc_ids, &mut scratch)?;
1810 offset = offset
1811 .checked_add(
1812 u64::try_from(count)
1813 .ok()
1814 .and_then(|count| count.checked_mul(4))
1815 .ok_or_else(|| invalid_data("ANN doc-ID output size overflows u64"))?,
1816 )
1817 .ok_or_else(|| invalid_data("ANN output offset overflow"))?;
1818 let ordinals_offset = offset;
1819 write_u16_column(writer, run.ordinals, &mut scratch)?;
1820 offset = offset
1821 .checked_add(
1822 u64::try_from(count)
1823 .ok()
1824 .and_then(|count| count.checked_mul(2))
1825 .ok_or_else(|| invalid_data("ANN ordinal output size overflows u64"))?,
1826 )
1827 .ok_or_else(|| invalid_data("ANN output offset overflow"))?;
1828 let codes_offset = offset;
1829 writer.write_all(run.codes)?;
1830 offset = offset
1831 .checked_add(
1832 u64::try_from(run.codes.len())
1833 .map_err(|_| invalid_data("ANN code output size exceeds u64"))?,
1834 )
1835 .ok_or_else(|| invalid_data("ANN output offset overflow"))?;
1836 records.push(RunRecord {
1837 cluster_id: run.cluster_id,
1838 doc_base: 0,
1839 count: count_u32,
1840 max_doc_id,
1841 doc_ids_offset,
1842 ordinals_offset,
1843 codes_offset,
1844 codes_len: u64::try_from(run.codes.len())
1845 .map_err(|_| invalid_data("ANN code output size exceeds u64"))?,
1846 });
1847 counted = counted
1848 .checked_add(count)
1849 .ok_or_else(|| invalid_data("ANN vector count overflow"))?;
1850 }
1851 if counted != header.vector_count {
1852 return Err(invalid_data("ANN header/build vector counts disagree"));
1853 }
1854 finish_layout(writer, offset, &records)
1855}
1856
1857#[cfg(all(feature = "native", test))]
1860pub(crate) fn write_merged_ann(
1861 sources: &[(&AnnDiskIndex, u32)],
1862 writer: &mut (impl Write + ?Sized),
1863) -> io::Result<u64> {
1864 write_merged_ann_impl(sources, &[], writer, None)
1865}
1866
1867#[cfg(feature = "native")]
1868pub(crate) fn write_merged_ann_cancellable(
1869 sources: &[(&AnnDiskIndex, u32)],
1870 writer: &mut (impl Write + ?Sized),
1871 cancellation: Option<&std::sync::atomic::AtomicBool>,
1872) -> io::Result<u64> {
1873 write_merged_ann_impl(sources, &[], writer, cancellation)
1874}
1875
1876#[cfg(feature = "native")]
1885pub(crate) fn predicted_merge_fragmentation(sources: &[(&AnnDiskIndex, u32)]) -> f64 {
1886 let mut total_runs = 0usize;
1887 let mut cursors: Vec<std::iter::Peekable<std::slice::Iter<'_, AnnRun>>> = sources
1890 .iter()
1891 .map(|(source, _)| {
1892 total_runs += source.runs.len();
1893 source.runs.iter().peekable()
1894 })
1895 .collect();
1896 let mut distinct = 0usize;
1897 while let Some(cluster) = cursors
1898 .iter_mut()
1899 .filter_map(|cursor| cursor.peek().map(|run| run.cluster_id))
1900 .min()
1901 {
1902 distinct += 1;
1903 for cursor in &mut cursors {
1904 while cursor.peek().is_some_and(|run| run.cluster_id == cluster) {
1905 cursor.next();
1906 }
1907 }
1908 }
1909 if distinct == 0 {
1910 0.0
1911 } else {
1912 total_runs as f64 / distinct as f64
1913 }
1914}
1915
1916#[cfg(feature = "native")]
1918const DOC_ID_REWRITE_CHUNK: usize = 64 * 1024;
1919
1920#[cfg(feature = "native")]
1947pub(crate) fn write_compacted_ann_cancellable(
1948 sources: &[(&AnnDiskIndex, u32)],
1949 writer: &mut (impl Write + ?Sized),
1950 cancellation: Option<&std::sync::atomic::AtomicBool>,
1951) -> io::Result<u64> {
1952 let Some((first, _)) = sources.first() else {
1953 return Err(invalid_data("cannot compact an empty ANN source list"));
1954 };
1955 if first.header.kind != AnnKind::BinaryIvf {
1956 return Err(invalid_data(
1957 "ANN run compaction is only defined for binary IVF payloads",
1958 ));
1959 }
1960 let mut header = first.header.clone();
1961 header.vector_count = 0;
1962 for &(source, _) in sources {
1963 if !headers_compatible(&first.header, &source.header) {
1964 return Err(invalid_data(
1965 "ANN compaction sources use incompatible generations",
1966 ));
1967 }
1968 header.vector_count = header
1969 .vector_count
1970 .checked_add(source.header.vector_count)
1971 .ok_or_else(|| invalid_data("compacted ANN vector count overflows usize"))?;
1972 }
1973 validate_header(&header)?;
1974 write_header(writer, &header)?;
1975
1976 let code_size = header.code_size;
1977 let mut offset = ANN_HEADER_SIZE as u64;
1978 let mut records: Vec<RunRecord> = Vec::new();
1979 let mut scratch = Vec::new();
1980 let mut cursors: Vec<usize> = vec![0; sources.len()];
1981
1982 loop {
1983 if cancellation.is_some_and(|flag| flag.load(std::sync::atomic::Ordering::Relaxed)) {
1984 return Err(io::Error::new(
1985 io::ErrorKind::Interrupted,
1986 "ANN compaction cancelled",
1987 ));
1988 }
1989 let Some(cluster_id) = sources
1991 .iter()
1992 .zip(&cursors)
1993 .filter_map(|((source, _), &cursor)| source.runs.get(cursor).map(|run| run.cluster_id))
1994 .min()
1995 else {
1996 break;
1997 };
1998
1999 let mut count = 0usize;
2004 let mut max_doc_id = 0u32;
2005 let doc_ids_offset = offset;
2006 for (source_index, &(source, segment_base)) in sources.iter().enumerate() {
2007 let mut cursor = cursors[source_index];
2008 while let Some(run) = source
2009 .runs
2010 .get(cursor)
2011 .filter(|run| run.cluster_id == cluster_id)
2012 {
2013 let base = run
2014 .doc_base
2015 .checked_add(segment_base)
2016 .ok_or_else(|| invalid_data("compacted ANN document base overflows u32"))?;
2017 let bytes = source.raw.as_slice();
2018 for chunk_start in (0..run.count).step_by(DOC_ID_REWRITE_CHUNK) {
2023 let chunk_end = (chunk_start + DOC_ID_REWRITE_CHUNK).min(run.count);
2024 scratch.clear();
2025 scratch.reserve((chunk_end - chunk_start) * 4);
2026 for index in chunk_start..chunk_end {
2027 let doc_id = run_doc_id_with_base(bytes, run, index, base)?;
2028 max_doc_id = max_doc_id.max(doc_id);
2029 scratch.extend_from_slice(&doc_id.to_le_bytes());
2030 }
2031 writer.write_all(&scratch)?;
2032 if cancellation
2033 .is_some_and(|flag| flag.load(std::sync::atomic::Ordering::Relaxed))
2034 {
2035 return Err(io::Error::new(
2036 io::ErrorKind::Interrupted,
2037 "ANN compaction cancelled",
2038 ));
2039 }
2040 }
2041 offset = checked_advance(offset, run.count * 4)?;
2042 count = count
2043 .checked_add(run.count)
2044 .ok_or_else(|| invalid_data("compacted ANN run count overflows usize"))?;
2045 cursor += 1;
2046 }
2047 }
2048
2049 let ordinals_offset = offset;
2051 for (source_index, &(source, _)) in sources.iter().enumerate() {
2052 let mut cursor = cursors[source_index];
2053 while let Some(run) = source
2054 .runs
2055 .get(cursor)
2056 .filter(|run| run.cluster_id == cluster_id)
2057 {
2058 copy_range(writer, &source.raw, run.ordinals.clone(), cancellation)?;
2059 offset = checked_advance(offset, run.ordinals.len())?;
2060 cursor += 1;
2061 }
2062 }
2063
2064 let codes_offset = offset;
2066 for (source_index, &(source, _)) in sources.iter().enumerate() {
2067 let mut cursor = cursors[source_index];
2068 while let Some(run) = source
2069 .runs
2070 .get(cursor)
2071 .filter(|run| run.cluster_id == cluster_id)
2072 {
2073 copy_range(writer, &source.raw, run.codes.clone(), cancellation)?;
2074 offset = checked_advance(offset, run.codes.len())?;
2075 cursor += 1;
2076 }
2077 }
2078
2079 for (source_index, &(source, _)) in sources.iter().enumerate() {
2081 while source
2082 .runs
2083 .get(cursors[source_index])
2084 .is_some_and(|run| run.cluster_id == cluster_id)
2085 {
2086 cursors[source_index] += 1;
2087 }
2088 }
2089
2090 records.push(RunRecord {
2091 cluster_id,
2092 doc_base: 0,
2093 count: u32::try_from(count)
2094 .map_err(|_| invalid_data("compacted ANN run exceeds u32 vectors"))?,
2095 max_doc_id,
2096 doc_ids_offset,
2097 ordinals_offset,
2098 codes_offset,
2099 codes_len: u64::try_from(expected_codes_column_len(
2100 AnnKind::BinaryIvf,
2101 count,
2102 code_size,
2103 )?)
2104 .map_err(|_| invalid_data("compacted ANN code length exceeds u64"))?,
2105 });
2106 }
2107
2108 if records.is_empty() {
2109 return Err(invalid_data("cannot compact an ANN payload with no runs"));
2110 }
2111 finish_layout(writer, offset, &records)
2112}
2113
2114#[cfg(feature = "native")]
2116fn run_doc_id_with_base(bytes: &[u8], run: &AnnRun, index: usize, base: u32) -> io::Result<u32> {
2117 let local_doc_id = read_u32(bytes, run.doc_ids.start + index * 4);
2118 if local_doc_id > run.max_doc_id {
2119 return Err(invalid_data(
2120 "ANN run contains a document above its declared maximum",
2121 ));
2122 }
2123 base.checked_add(local_doc_id)
2124 .ok_or_else(|| invalid_data("compacted ANN document ID overflows u32"))
2125}
2126
2127#[cfg(feature = "native")]
2131pub(crate) fn write_merged_ann_with_extra(
2132 sources: &[(&AnnDiskIndex, u32)],
2133 extra: &[BuildRun<'_>],
2134 writer: &mut (impl Write + ?Sized),
2135 cancellation: Option<&std::sync::atomic::AtomicBool>,
2136) -> io::Result<u64> {
2137 write_merged_ann_impl(sources, extra, writer, cancellation)
2138}
2139
2140#[cfg(feature = "native")]
2141fn write_merged_ann_impl(
2142 sources: &[(&AnnDiskIndex, u32)],
2143 extra: &[BuildRun<'_>],
2144 writer: &mut (impl Write + ?Sized),
2145 cancellation: Option<&std::sync::atomic::AtomicBool>,
2146) -> io::Result<u64> {
2147 let Some((first, _)) = sources.first() else {
2148 return Err(invalid_data("cannot merge an empty ANN source list"));
2149 };
2150 if first.header.kind == AnnKind::IvfTq
2151 && !crate::structures::is_ivf_tq_cosine_generation(first.header.quantizer_version)
2152 {
2153 return Err(invalid_data(
2154 "legacy raw IVF-TQ generations cannot be merged; rebuild the index",
2155 ));
2156 }
2157 let mut header = first.header.clone();
2158 header.vector_count = 0;
2159 for &(source, _) in sources {
2160 if !headers_compatible(&first.header, &source.header) {
2161 return Err(invalid_data(
2162 "ANN merge sources use incompatible generations",
2163 ));
2164 }
2165 header.vector_count = header
2166 .vector_count
2167 .checked_add(source.header.vector_count)
2168 .ok_or_else(|| invalid_data("merged ANN vector count overflows usize"))?;
2169 }
2170 for run in extra {
2171 if run.doc_ids.is_empty()
2172 || run.cluster_id >= header.num_clusters
2173 || run.ordinals.len() != run.doc_ids.len()
2174 || run.codes.len()
2175 != expected_codes_column_len(header.kind, run.doc_ids.len(), header.code_size)?
2176 {
2177 return Err(invalid_data("extra ANN merge run columns are inconsistent"));
2178 }
2179 header.vector_count = header
2180 .vector_count
2181 .checked_add(run.doc_ids.len())
2182 .ok_or_else(|| invalid_data("merged ANN vector count overflows usize"))?;
2183 }
2184 validate_header(&header)?;
2185 write_header(writer, &header)?;
2186 let mut offset = ANN_HEADER_SIZE as u64;
2187 let run_capacity = sources.iter().try_fold(extra.len(), |count, (source, _)| {
2188 count
2189 .checked_add(source.runs.len())
2190 .ok_or_else(|| invalid_data("merged ANN run count overflows usize"))
2191 })?;
2192 let mut output_payload_starts = Vec::with_capacity(sources.len());
2193 for &(source, _) in sources {
2194 let payload_end = source
2195 .runs
2196 .iter()
2197 .map(|run| run.codes.end)
2198 .max()
2199 .ok_or_else(|| invalid_data("ANN source has no payload runs"))?;
2200 let output_payload_start = offset;
2201 output_payload_starts.push(output_payload_start);
2202 copy_range(
2203 writer,
2204 &source.raw,
2205 ANN_HEADER_SIZE..payload_end,
2206 cancellation,
2207 )?;
2208 offset = checked_advance(offset, payload_end - ANN_HEADER_SIZE)?;
2209 }
2210
2211 let mut extra_records = Vec::with_capacity(extra.len());
2214 let mut scratch = Vec::new();
2215 for run in extra {
2216 let count = run.doc_ids.len();
2217 let doc_ids_offset = offset;
2218 write_u32_column(writer, run.doc_ids, &mut scratch)?;
2219 offset = checked_advance(offset, count * 4)?;
2220 let ordinals_offset = offset;
2221 write_u16_column(writer, run.ordinals, &mut scratch)?;
2222 offset = checked_advance(offset, count * 2)?;
2223 let codes_offset = offset;
2224 writer.write_all(run.codes)?;
2225 offset = checked_advance(offset, run.codes.len())?;
2226 extra_records.push(RunRecord {
2227 cluster_id: run.cluster_id,
2228 doc_base: 0,
2229 count: u32::try_from(count)
2230 .map_err(|_| invalid_data("extra ANN run exceeds u32 vectors"))?,
2231 max_doc_id: run.doc_ids.iter().copied().max().unwrap_or(0),
2232 doc_ids_offset,
2233 ordinals_offset,
2234 codes_offset,
2235 codes_len: u64::try_from(run.codes.len())
2236 .map_err(|_| invalid_data("extra ANN code length exceeds u64"))?,
2237 });
2238 }
2239
2240 let directory_offset = offset;
2246 let mut pending = BinaryHeap::with_capacity(sources.len() + 1);
2247 for (source_index, (source, _)) in sources.iter().enumerate() {
2248 pending.push(Reverse((source.runs[0].cluster_id, source_index, 0usize)));
2249 }
2250 if let Some(first_extra) = extra_records.first() {
2251 pending.push(Reverse((first_extra.cluster_id, sources.len(), 0usize)));
2252 }
2253 let mut written_runs = 0usize;
2254 while let Some(Reverse((_, source_index, run_index))) = pending.pop() {
2255 if source_index == sources.len() {
2256 write_run_record(writer, &extra_records[run_index])?;
2257 written_runs += 1;
2258 if let Some(next) = extra_records.get(run_index + 1) {
2259 pending.push(Reverse((next.cluster_id, source_index, run_index + 1)));
2260 }
2261 continue;
2262 }
2263 let (source, segment_base) = sources[source_index];
2264 let run = &source.runs[run_index];
2265 write_run_record(
2266 writer,
2267 &RunRecord {
2268 cluster_id: run.cluster_id,
2269 doc_base: run
2270 .doc_base
2271 .checked_add(segment_base)
2272 .ok_or_else(|| invalid_data("merged ANN document base overflows u32"))?,
2273 count: u32::try_from(run.count)
2274 .map_err(|_| invalid_data("ANN source run exceeds u32 vectors"))?,
2275 max_doc_id: run.max_doc_id,
2276 doc_ids_offset: relocate_payload_offset(
2277 output_payload_starts[source_index],
2278 run.doc_ids.start,
2279 )?,
2280 ordinals_offset: relocate_payload_offset(
2281 output_payload_starts[source_index],
2282 run.ordinals.start,
2283 )?,
2284 codes_offset: relocate_payload_offset(
2285 output_payload_starts[source_index],
2286 run.codes.start,
2287 )?,
2288 codes_len: u64::try_from(run.codes.len())
2289 .map_err(|_| invalid_data("ANN source code length exceeds u64"))?,
2290 },
2291 )?;
2292 written_runs = written_runs
2293 .checked_add(1)
2294 .ok_or_else(|| invalid_data("merged ANN run count overflows usize"))?;
2295 let next_run_index = run_index + 1;
2296 if let Some(next_run) = source.runs.get(next_run_index) {
2297 pending.push(Reverse((next_run.cluster_id, source_index, next_run_index)));
2298 }
2299 }
2300 if written_runs != run_capacity {
2301 return Err(invalid_data("merged ANN directory lost source runs"));
2302 }
2303 finish_footer(writer, directory_offset, written_runs)
2304}
2305
2306#[cfg(feature = "native")]
2307fn relocate_payload_offset(output_payload_start: u64, source_offset: usize) -> io::Result<u64> {
2308 let relative = source_offset
2309 .checked_sub(ANN_HEADER_SIZE)
2310 .ok_or_else(|| invalid_data("ANN source offset precedes its payload"))?;
2311 output_payload_start
2312 .checked_add(
2313 u64::try_from(relative)
2314 .map_err(|_| invalid_data("ANN source relative offset exceeds u64"))?,
2315 )
2316 .ok_or_else(|| invalid_data("merged ANN payload offset overflows u64"))
2317}
2318
2319#[cfg(feature = "native")]
2320fn headers_compatible(left: &AnnDiskHeader, right: &AnnDiskHeader) -> bool {
2321 left.kind == right.kind
2322 && left.routing == right.routing
2323 && left.dim == right.dim
2324 && left.code_size == right.code_size
2325 && left.num_clusters == right.num_clusters
2326 && left.quantizer_version == right.quantizer_version
2327 && left.codebook_version == right.codebook_version
2328}
2329
2330#[cfg(feature = "native")]
2331fn finish_layout(
2332 writer: &mut (impl Write + ?Sized),
2333 directory_offset: u64,
2334 records: &[RunRecord],
2335) -> io::Result<u64> {
2336 for record in records {
2337 write_run_record(writer, record)?;
2338 }
2339 finish_footer(writer, directory_offset, records.len())
2340}
2341
2342#[cfg(feature = "native")]
2343fn write_run_record(writer: &mut (impl Write + ?Sized), record: &RunRecord) -> io::Result<()> {
2344 writer.write_u32::<LittleEndian>(record.cluster_id)?;
2345 writer.write_u32::<LittleEndian>(record.doc_base)?;
2346 writer.write_u32::<LittleEndian>(record.count)?;
2347 writer.write_u32::<LittleEndian>(record.max_doc_id)?;
2348 writer.write_u64::<LittleEndian>(record.doc_ids_offset)?;
2349 writer.write_u64::<LittleEndian>(record.ordinals_offset)?;
2350 writer.write_u64::<LittleEndian>(record.codes_offset)?;
2351 writer.write_u64::<LittleEndian>(record.codes_len)?;
2352 Ok(())
2353}
2354
2355#[cfg(feature = "native")]
2356fn finish_footer(
2357 writer: &mut (impl Write + ?Sized),
2358 directory_offset: u64,
2359 num_records: usize,
2360) -> io::Result<u64> {
2361 writer.write_u64::<LittleEndian>(directory_offset)?;
2362 writer.write_u64::<LittleEndian>(
2363 u64::try_from(num_records).map_err(|_| invalid_data("ANN run count exceeds u64"))?,
2364 )?;
2365 writer.write_u32::<LittleEndian>(ANN_FOOTER_MAGIC)?;
2366 writer.write_u32::<LittleEndian>(u32::from(ANN_DISK_VERSION))?;
2367 let tail_size = num_records
2368 .checked_mul(ANN_RUN_SIZE)
2369 .and_then(|size| size.checked_add(ANN_FOOTER_SIZE))
2370 .and_then(|size| u64::try_from(size).ok())
2371 .ok_or_else(|| invalid_data("ANN final tail size overflows u64"))?;
2372 directory_offset
2373 .checked_add(tail_size)
2374 .ok_or_else(|| invalid_data("ANN final size overflows u64"))
2375}
2376
2377#[cfg(feature = "native")]
2378fn write_header(writer: &mut (impl Write + ?Sized), header: &AnnDiskHeader) -> io::Result<()> {
2379 writer.write_u32::<LittleEndian>(ANN_HEADER_MAGIC)?;
2380 writer.write_u8(header.kind as u8)?;
2381 writer.write_u8(routing_to_u8(header.routing))?;
2382 writer.write_u16::<LittleEndian>(ANN_DISK_VERSION)?;
2383 writer.write_u32::<LittleEndian>(
2384 u32::try_from(header.dim).map_err(|_| invalid_data("ANN dimension exceeds u32"))?,
2385 )?;
2386 writer.write_u32::<LittleEndian>(
2387 u32::try_from(header.code_size).map_err(|_| invalid_data("ANN code size exceeds u32"))?,
2388 )?;
2389 writer.write_u32::<LittleEndian>(header.num_clusters)?;
2390 writer.write_u32::<LittleEndian>(0)?;
2391 writer.write_u64::<LittleEndian>(header.quantizer_version)?;
2392 writer.write_u64::<LittleEndian>(header.codebook_version)?;
2393 writer.write_u64::<LittleEndian>(
2394 u64::try_from(header.vector_count)
2395 .map_err(|_| invalid_data("ANN vector count exceeds u64"))?,
2396 )?;
2397 writer.write_u64::<LittleEndian>(0)?;
2398 Ok(())
2399}
2400
2401#[cfg(feature = "native")]
2402fn write_u32_column(
2403 writer: &mut (impl Write + ?Sized),
2404 values: &[u32],
2405 scratch: &mut Vec<u8>,
2406) -> io::Result<()> {
2407 for chunk in values.chunks(64 * 1024) {
2408 scratch.clear();
2409 scratch.reserve(chunk.len() * 4);
2410 for value in chunk {
2411 scratch.extend_from_slice(&value.to_le_bytes());
2412 }
2413 writer.write_all(scratch)?;
2414 }
2415 Ok(())
2416}
2417
2418#[cfg(feature = "native")]
2419fn write_u16_column(
2420 writer: &mut (impl Write + ?Sized),
2421 values: &[u16],
2422 scratch: &mut Vec<u8>,
2423) -> io::Result<()> {
2424 for chunk in values.chunks(64 * 1024) {
2425 scratch.clear();
2426 scratch.reserve(chunk.len() * 2);
2427 for value in chunk {
2428 scratch.extend_from_slice(&value.to_le_bytes());
2429 }
2430 writer.write_all(scratch)?;
2431 }
2432 Ok(())
2433}
2434
2435#[cfg(feature = "native")]
2436fn copy_range(
2437 writer: &mut (impl Write + ?Sized),
2438 bytes: &OwnedBytes,
2439 range: Range<usize>,
2440 cancellation: Option<&std::sync::atomic::AtomicBool>,
2441) -> io::Result<()> {
2442 if range.is_empty() {
2443 return Ok(());
2444 }
2445 let range_end = range.end;
2446 let mut chunk_start = range.start;
2447 let first_end = chunk_start.saturating_add(COPY_CHUNK).min(range_end);
2448 bytes.madvise_range(chunk_start..first_end, libc::MADV_WILLNEED);
2449 while chunk_start < range_end {
2450 if cancellation
2451 .is_some_and(|cancelled| cancelled.load(std::sync::atomic::Ordering::Relaxed))
2452 {
2453 return Err(io::Error::new(
2454 io::ErrorKind::Interrupted,
2455 "ANN merge copy cancelled",
2456 ));
2457 }
2458 let chunk_end = chunk_start.saturating_add(COPY_CHUNK).min(range_end);
2459 let next_end = chunk_end.saturating_add(COPY_CHUNK).min(range_end);
2460 if chunk_end < next_end {
2461 bytes.madvise_range(chunk_end..next_end, libc::MADV_WILLNEED);
2464 }
2465 writer.write_all(&bytes.as_slice()[chunk_start..chunk_end])?;
2466 chunk_start = chunk_end;
2467 }
2468 Ok(())
2469}
2470
2471#[cfg(feature = "native")]
2472fn checked_advance(offset: u64, length: usize) -> io::Result<u64> {
2473 offset
2474 .checked_add(
2475 u64::try_from(length).map_err(|_| invalid_data("ANN copy length exceeds u64"))?,
2476 )
2477 .ok_or_else(|| invalid_data("ANN output offset overflows u64"))
2478}
2479
2480fn validate_header(header: &AnnDiskHeader) -> io::Result<()> {
2481 if header.kind == AnnKind::IvfTq
2482 && !crate::structures::is_ivf_tq_cosine_generation(header.quantizer_version)
2483 {
2484 return Err(invalid_data(
2485 "IVF-TQ payload uses an unsupported legacy generation; rebuild the index",
2486 ));
2487 }
2488 if header.dim == 0
2489 || header.code_size == 0
2490 || header.num_clusters == 0
2491 || header.quantizer_version == 0
2492 || header.vector_count == 0
2493 || (header.kind == AnnKind::BinaryIvf
2494 && (header.codebook_version != 0
2495 || !header.dim.is_multiple_of(8)
2496 || header.code_size != header.dim.div_ceil(8)))
2497 || (header.kind == AnnKind::TqFlat
2498 && (header.codebook_version != 0
2499 || header.num_clusters != 1
2500 || header.routing != IvfRoutingMode::Flat
2501 || header.code_size * 2
2502 != crate::structures::vector::quantization::tq_padded_dim(header.dim)))
2503 || (header.kind == AnnKind::IvfTq
2506 && (header.codebook_version == 0
2507 || header.code_size * 2
2508 != crate::structures::vector::quantization::tq_padded_dim(header.dim)))
2509 {
2510 return Err(invalid_data("ANN header contains invalid metadata"));
2511 }
2512 Ok(())
2513}
2514
2515fn read_u32(bytes: &[u8], offset: usize) -> u32 {
2516 u32::from_le_bytes(bytes[offset..offset + 4].try_into().unwrap())
2517}
2518
2519fn read_u16(bytes: &[u8], offset: usize) -> u16 {
2520 u16::from_le_bytes(bytes[offset..offset + 2].try_into().unwrap())
2521}
2522
2523#[inline]
2524fn tq_ivf_block_max_scale(block: &[u8]) -> f32 {
2525 f32::from_le_bytes(
2529 block[..size_of::<f32>()]
2530 .try_into()
2531 .expect("scale is one f32"),
2532 )
2533}
2534
2535fn run_doc_id(bytes: &[u8], run: &AnnRun, index: usize) -> io::Result<u32> {
2536 let local_doc_id = read_u32(bytes, run.doc_ids.start + index * 4);
2537 if local_doc_id > run.max_doc_id {
2538 return Err(invalid_data(
2539 "ANN run contains a document above its declared maximum",
2540 ));
2541 }
2542 run.doc_base
2543 .checked_add(local_doc_id)
2544 .ok_or_else(|| invalid_data("ANN run document ID overflows u32"))
2545}
2546
2547#[cfg(feature = "native")]
2548fn routing_to_u8(routing: IvfRoutingMode) -> u8 {
2549 match routing {
2550 IvfRoutingMode::Auto => 0,
2551 IvfRoutingMode::Flat => 1,
2552 IvfRoutingMode::TwoLevel => 2,
2553 IvfRoutingMode::Hnsw => 3,
2554 }
2555}
2556
2557fn routing_from_u8(value: u8) -> io::Result<IvfRoutingMode> {
2558 match value {
2559 0 => Ok(IvfRoutingMode::Auto),
2560 1 => Ok(IvfRoutingMode::Flat),
2561 2 => Ok(IvfRoutingMode::TwoLevel),
2562 3 => Ok(IvfRoutingMode::Hnsw),
2563 _ => Err(invalid_data(format!("unknown ANN routing mode {value}"))),
2564 }
2565}
2566
2567fn invalid_data(message: impl Into<String>) -> io::Error {
2568 io::Error::new(io::ErrorKind::InvalidData, message.into())
2569}
2570
2571#[cfg(all(test, feature = "native"))]
2572mod tests {
2573 use super::*;
2574
2575 #[test]
2580 fn compacted_merge_matches_byte_copy_and_resets_fragmentation() {
2581 let a0_docs = [0u32, 1, 2];
2585 let a0_ords = [0u16, 0, 1];
2586 let a0_codes = [0x11u8, 0x22, 0x33];
2587 let a5_docs = [3u32, 4];
2588 let a5_ords = [0u16; 2];
2589 let a5_codes = [0x44u8, 0x55];
2590 let a_runs = [
2591 BuildRun {
2592 cluster_id: 0,
2593 doc_ids: &a0_docs,
2594 ordinals: &a0_ords,
2595 codes: &a0_codes,
2596 },
2597 BuildRun {
2598 cluster_id: 5,
2599 doc_ids: &a5_docs,
2600 ordinals: &a5_ords,
2601 codes: &a5_codes,
2602 },
2603 ];
2604 let mut header = binary_header(5);
2605 header.num_clusters = 8;
2606 let mut a_bytes = Vec::new();
2607 write_built_runs(header.clone(), &a_runs, &mut a_bytes).unwrap();
2608 let a = AnnDiskIndex::open(OwnedBytes::new(a_bytes), AnnKind::BinaryIvf, 5).unwrap();
2609
2610 let b0_docs = [0u32, 2];
2611 let b0_ords = [1u16, 0];
2612 let b0_codes = [0x66u8, 0x77];
2613 let b2_docs = [1u32];
2614 let b2_ords = [0u16];
2615 let b2_codes = [0x88u8];
2616 let b_runs = [
2617 BuildRun {
2618 cluster_id: 0,
2619 doc_ids: &b0_docs,
2620 ordinals: &b0_ords,
2621 codes: &b0_codes,
2622 },
2623 BuildRun {
2624 cluster_id: 2,
2625 doc_ids: &b2_docs,
2626 ordinals: &b2_ords,
2627 codes: &b2_codes,
2628 },
2629 ];
2630 let mut header_b = binary_header(3);
2631 header_b.num_clusters = 8;
2632 let mut b_bytes = Vec::new();
2633 write_built_runs(header_b, &b_runs, &mut b_bytes).unwrap();
2634 let b = AnnDiskIndex::open(OwnedBytes::new(b_bytes), AnnKind::BinaryIvf, 3).unwrap();
2635
2636 let mut gen1_bytes = Vec::new();
2638 write_merged_ann(&[(&a, 0), (&b, 5)], &mut gen1_bytes).unwrap();
2639 let gen1 = AnnDiskIndex::open(OwnedBytes::new(gen1_bytes), AnnKind::BinaryIvf, 8).unwrap();
2640
2641 let sources: [(&AnnDiskIndex, u32); 2] = [(&gen1, 0), (&a, 8)];
2643 let predicted = predicted_merge_fragmentation(&sources);
2644 assert!((predicted - 2.0).abs() < 1e-9, "{predicted}");
2646
2647 let mut copied_bytes = Vec::new();
2648 write_merged_ann(&sources, &mut copied_bytes).unwrap();
2649 let copied =
2650 AnnDiskIndex::open(OwnedBytes::new(copied_bytes), AnnKind::BinaryIvf, 13).unwrap();
2651 let mut compacted_bytes = Vec::new();
2652 write_compacted_ann_cancellable(&sources, &mut compacted_bytes, None).unwrap();
2653 let compacted =
2654 AnnDiskIndex::open(OwnedBytes::new(compacted_bytes), AnnKind::BinaryIvf, 13).unwrap();
2655
2656 let copied_health = copied.health();
2658 let compacted_health = compacted.health();
2659 assert!((copied_health.fragmentation() - 2.0).abs() < 1e-9);
2660 assert!((compacted_health.fragmentation() - 1.0).abs() < 1e-9);
2661 assert_eq!(compacted_health.runs, 3, "one run per non-empty cluster");
2662 assert_eq!(copied_health.vectors, compacted_health.vectors);
2663 assert_eq!(copied_health.payload_bytes, compacted_health.payload_bytes);
2664 assert_eq!(
2665 copied_health.largest_cluster_vectors,
2666 compacted_health.largest_cluster_vectors
2667 );
2668
2669 for cluster in 0..8u32 {
2671 let query = [0x5Au8];
2672 let from_copy = copied
2673 .search_binary_clusters::<false>(&query, 16, &[cluster])
2674 .unwrap();
2675 let from_compact = compacted
2676 .search_binary_clusters::<false>(&query, 16, &[cluster])
2677 .unwrap();
2678 assert_eq!(from_copy, from_compact, "cluster {cluster} diverged");
2679 }
2680
2681 assert!(compacted.runs.iter().all(|run| run.doc_base == 0));
2683
2684 let mut generation3 = Vec::new();
2687 write_merged_ann(&[(&compacted, 0), (&b, 13)], &mut generation3).unwrap();
2688 let generation3 =
2689 AnnDiskIndex::open(OwnedBytes::new(generation3), AnnKind::BinaryIvf, 16).unwrap();
2690 assert_eq!(generation3.health().vectors, 16);
2691 let all: Vec<(u32, u16, f32)> = compacted
2693 .search_binary_clusters::<false>(&[0x5A], 32, &[0, 2, 5])
2694 .unwrap();
2695 let mut docs: Vec<u32> = all.iter().map(|&(doc, _, _)| doc).collect();
2696 docs.sort_unstable();
2697 assert_eq!(docs, (0..=12).collect::<Vec<u32>>());
2698 }
2699
2700 #[test]
2703 #[ignore]
2704 fn ann_merge_throughput_byte_copy_vs_compaction() {
2705 let code_size = 320usize;
2706 let clusters = 4_096u32;
2707 let vectors_per_source = 262_144usize;
2708 let sources_count = 4usize;
2709
2710 let mut sources_bytes = Vec::new();
2711 for source_index in 0..sources_count {
2712 let mut per_cluster: Vec<(Vec<u32>, Vec<u16>, Vec<u8>)> = Vec::new();
2713 let vectors_per_cluster = vectors_per_source / clusters as usize;
2714 let mut doc = 0u32;
2715 for cluster in 0..clusters {
2716 let mut docs = Vec::with_capacity(vectors_per_cluster);
2717 let mut ords = Vec::with_capacity(vectors_per_cluster);
2718 let mut codes = Vec::with_capacity(vectors_per_cluster * code_size);
2719 for _ in 0..vectors_per_cluster {
2720 docs.push(doc);
2721 ords.push(0u16);
2722 codes.extend(std::iter::repeat_n(
2723 (doc ^ cluster ^ source_index as u32) as u8,
2724 code_size,
2725 ));
2726 doc += 1;
2727 }
2728 per_cluster.push((docs, ords, codes));
2729 }
2730 let runs: Vec<BuildRun<'_>> = per_cluster
2731 .iter()
2732 .enumerate()
2733 .map(|(cluster, (docs, ords, codes))| BuildRun {
2734 cluster_id: cluster as u32,
2735 doc_ids: docs,
2736 ordinals: ords,
2737 codes,
2738 })
2739 .collect();
2740 let header = AnnDiskHeader {
2741 kind: AnnKind::BinaryIvf,
2742 routing: IvfRoutingMode::Hnsw,
2743 dim: code_size * 8,
2744 code_size,
2745 num_clusters: clusters,
2746 quantizer_version: 42,
2747 codebook_version: 0,
2748 vector_count: vectors_per_source,
2749 };
2750 let mut bytes = Vec::new();
2751 write_built_runs(header, &runs, &mut bytes).unwrap();
2752 sources_bytes.push(bytes);
2753 }
2754 let sources_open: Vec<AnnDiskIndex> = sources_bytes
2755 .iter()
2756 .map(|bytes| {
2757 AnnDiskIndex::open(
2758 OwnedBytes::new(bytes.clone()),
2759 AnnKind::BinaryIvf,
2760 (vectors_per_source * sources_count) as u32,
2761 )
2762 .unwrap()
2763 })
2764 .collect();
2765 let sources: Vec<(&AnnDiskIndex, u32)> = sources_open
2766 .iter()
2767 .enumerate()
2768 .map(|(index, source)| (source, (index * vectors_per_source) as u32))
2769 .collect();
2770 let payload_bytes = sources_bytes.iter().map(Vec::len).sum::<usize>();
2771
2772 let mut out = vec![0u8; payload_bytes + (1 << 20)];
2778 out.clear();
2779 let mut copy_secs = f64::INFINITY;
2782 let mut compact_secs = f64::INFINITY;
2783 let mut compacted_bytes = Vec::new();
2784 for _ in 0..3 {
2785 out.clear();
2786 let start = std::time::Instant::now();
2787 write_merged_ann(&sources, &mut out).unwrap();
2788 copy_secs = copy_secs.min(start.elapsed().as_secs_f64());
2789
2790 out.clear();
2791 let start = std::time::Instant::now();
2792 write_compacted_ann_cancellable(&sources, &mut out, None).unwrap();
2793 compact_secs = compact_secs.min(start.elapsed().as_secs_f64());
2794 compacted_bytes = out.clone();
2795 }
2796 let compacted = AnnDiskIndex::open(
2797 OwnedBytes::new(compacted_bytes),
2798 AnnKind::BinaryIvf,
2799 (vectors_per_source * sources_count) as u32,
2800 )
2801 .unwrap();
2802 assert!((compacted.health().fragmentation() - 1.0).abs() < 1e-9);
2803
2804 let gib = payload_bytes as f64 / (1u64 << 30) as f64;
2805 println!(
2806 "ann merge {:.2} GiB: byte-copy {:.3}s ({:.2} GiB/s), compaction {:.3}s \
2807 ({:.2} GiB/s), overhead {:.1}%",
2808 gib,
2809 copy_secs,
2810 gib / copy_secs,
2811 compact_secs,
2812 gib / compact_secs,
2813 100.0 * (compact_secs - copy_secs) / copy_secs,
2814 );
2815 }
2816
2817 #[test]
2819 fn compaction_refuses_non_binary_payloads() {
2820 let codec = std::sync::Arc::new(crate::structures::TqCodec::new(8));
2826 let mut builder = crate::structures::TqFlatBuilder::new(codec);
2827 builder
2828 .add_batch(
2829 &[(0, 0), (1, 0)],
2830 &[
2831 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,
2833 ],
2834 )
2835 .unwrap();
2836 builder.finish();
2837 let mut bytes = Vec::new();
2838 write_built_tq_flat(&builder, &mut bytes).unwrap();
2839 let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::TqFlat, 2).unwrap();
2840 let error = write_compacted_ann_cancellable(&[(&disk, 0)], &mut Vec::new(), None)
2841 .expect_err("TQ payloads must not compact");
2842 assert!(error.to_string().contains("binary"), "{error}");
2843 }
2844
2845 #[test]
2849 fn ann_health_measures_skew_and_fragmentation() {
2850 let a0_docs = [0u32, 1, 2, 3, 4, 5];
2855 let a0_ords = [0u16; 6];
2856 let a0_codes = [0xAAu8; 6];
2857 let a3_docs = [6u32, 7];
2858 let a3_ords = [0u16; 2];
2859 let a3_codes = [0x0Fu8; 2];
2860 let a_runs = [
2861 BuildRun {
2862 cluster_id: 0,
2863 doc_ids: &a0_docs,
2864 ordinals: &a0_ords,
2865 codes: &a0_codes,
2866 },
2867 BuildRun {
2868 cluster_id: 3,
2869 doc_ids: &a3_docs,
2870 ordinals: &a3_ords,
2871 codes: &a3_codes,
2872 },
2873 ];
2874 let mut header = binary_header(8);
2875 header.num_clusters = 8;
2876 let mut a_bytes = Vec::new();
2877 write_built_runs(header, &a_runs, &mut a_bytes).unwrap();
2878 let a = AnnDiskIndex::open(OwnedBytes::new(a_bytes), AnnKind::BinaryIvf, 8).unwrap();
2879
2880 let b0_docs = [0u32, 1];
2881 let b0_ords = [0u16; 2];
2882 let b0_codes = [0xBBu8; 2];
2883 let b_runs = [BuildRun {
2884 cluster_id: 0,
2885 doc_ids: &b0_docs,
2886 ordinals: &b0_ords,
2887 codes: &b0_codes,
2888 }];
2889 let mut header_b = binary_header(2);
2890 header_b.num_clusters = 8;
2891 let mut b_bytes = Vec::new();
2892 write_built_runs(header_b, &b_runs, &mut b_bytes).unwrap();
2893 let b = AnnDiskIndex::open(OwnedBytes::new(b_bytes), AnnKind::BinaryIvf, 2).unwrap();
2894
2895 let mut merged_bytes = Vec::new();
2896 write_merged_ann(&[(&a, 0), (&b, 8)], &mut merged_bytes).unwrap();
2897 let disk =
2898 AnnDiskIndex::open(OwnedBytes::new(merged_bytes), AnnKind::BinaryIvf, 10).unwrap();
2899
2900 let health = disk.health();
2901 assert_eq!(health.vectors, 10);
2902 assert_eq!(health.clusters_nonempty, 2);
2903 assert_eq!(health.clusters_total, 8);
2904 assert_eq!(health.runs, 3);
2905 assert_eq!(health.largest_cluster, 0);
2906 assert_eq!(health.largest_cluster_vectors, 8);
2907 assert!((health.largest_cluster_share() - 0.8).abs() < 1e-9);
2908 assert!((health.fragmentation() - 1.5).abs() < 1e-9);
2910 assert!(
2912 (health.imbalance - 1.36).abs() < 1e-9,
2913 "{}",
2914 health.imbalance
2915 );
2916 assert_eq!(health.payload_bytes, 10);
2918 }
2919
2920 #[test]
2921 fn ann_health_is_balanced_at_one() {
2922 let docs: Vec<Vec<u32>> = (0..4).map(|c| vec![c * 2, c * 2 + 1]).collect();
2923 let ords = [0u16; 2];
2924 let codes = [0x55u8; 2];
2925 let runs: Vec<BuildRun<'_>> = docs
2926 .iter()
2927 .enumerate()
2928 .map(|(cluster, doc_ids)| BuildRun {
2929 cluster_id: cluster as u32,
2930 doc_ids,
2931 ordinals: &ords,
2932 codes: &codes,
2933 })
2934 .collect();
2935 let mut header = binary_header(8);
2936 header.num_clusters = 4;
2937 let mut bytes = Vec::new();
2938 write_built_runs(header, &runs, &mut bytes).unwrap();
2939 let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::BinaryIvf, 8).unwrap();
2940 let health = disk.health();
2941 assert!((health.imbalance - 1.0).abs() < 1e-9);
2942 assert!((health.fragmentation() - 1.0).abs() < 1e-9);
2943 assert!((health.largest_cluster_share() - 0.25).abs() < 1e-9);
2944 }
2945
2946 fn binary_header(vector_count: usize) -> AnnDiskHeader {
2947 AnnDiskHeader {
2948 kind: AnnKind::BinaryIvf,
2949 routing: IvfRoutingMode::Hnsw,
2950 dim: 8,
2951 code_size: 1,
2952 num_clusters: 2,
2953 quantizer_version: 42,
2954 codebook_version: 0,
2955 vector_count,
2956 }
2957 }
2958
2959 fn payload_end(index: &AnnDiskIndex) -> usize {
2960 index.runs.iter().map(|run| run.codes.end).max().unwrap()
2961 }
2962
2963 #[test]
2964 fn ann_prefetch_ranges_are_sorted_and_only_merge_page_near_extents() {
2965 let mut ranges = vec![
2966 15_000..16_000,
2967 0..1_000,
2968 9_000..10_000,
2969 1_000..2_000,
2970 7_000..8_000,
2971 ];
2972 coalesce_prefetch_ranges(&mut ranges);
2973 assert_eq!(ranges, [0..2_000, 7_000..10_000, 15_000..16_000]);
2974 }
2975
2976 #[test]
2977 fn normal_merge_copies_ann_payload_columns_byte_for_byte() {
2978 let first_doc_0 = [0u32];
2979 let first_doc_1 = [1u32];
2980 let first_ord_0 = [0u16];
2981 let first_ord_1 = [2u16];
2982 let first_code_0 = [0x00u8];
2983 let first_code_1 = [0xffu8];
2984 let first_runs = [
2985 BuildRun {
2986 cluster_id: 0,
2987 doc_ids: &first_doc_0,
2988 ordinals: &first_ord_0,
2989 codes: &first_code_0,
2990 },
2991 BuildRun {
2992 cluster_id: 1,
2993 doc_ids: &first_doc_1,
2994 ordinals: &first_ord_1,
2995 codes: &first_code_1,
2996 },
2997 ];
2998 let mut first_bytes = Vec::new();
2999 write_built_runs(binary_header(2), &first_runs, &mut first_bytes).unwrap();
3000 let first = AnnDiskIndex::open(OwnedBytes::new(first_bytes.clone()), AnnKind::BinaryIvf, 2)
3001 .unwrap();
3002
3003 let second_docs = [0u32, 1u32];
3004 let second_ords = [1u16, 0u16];
3005 let second_codes = [0x0fu8, 0xf0u8];
3006 let second_runs = [BuildRun {
3007 cluster_id: 0,
3008 doc_ids: &second_docs,
3009 ordinals: &second_ords,
3010 codes: &second_codes,
3011 }];
3012 let mut second_bytes = Vec::new();
3013 write_built_runs(binary_header(2), &second_runs, &mut second_bytes).unwrap();
3014 let second =
3015 AnnDiskIndex::open(OwnedBytes::new(second_bytes.clone()), AnnKind::BinaryIvf, 2)
3016 .unwrap();
3017
3018 let mut merged_bytes = Vec::new();
3019 write_merged_ann(&[(&first, 0), (&second, 2)], &mut merged_bytes).unwrap();
3020 let merged =
3021 AnnDiskIndex::open(OwnedBytes::new(merged_bytes.clone()), AnnKind::BinaryIvf, 4)
3022 .unwrap();
3023
3024 let mut expected_payload = first_bytes[ANN_HEADER_SIZE..payload_end(&first)].to_vec();
3025 expected_payload.extend_from_slice(&second_bytes[ANN_HEADER_SIZE..payload_end(&second)]);
3026 assert_eq!(
3027 &merged_bytes[ANN_HEADER_SIZE..payload_end(&merged)],
3028 expected_payload.as_slice(),
3029 "normal merge must not decode or rewrite any corpus-sized ANN column",
3030 );
3031
3032 let mut docs: Vec<u32> = merged
3033 .search_binary_clusters::<false>(&[0], 4, &[0, 1])
3034 .unwrap()
3035 .into_iter()
3036 .map(|result| result.0)
3037 .collect();
3038 docs.sort_unstable();
3039 assert_eq!(docs, [0, 1, 2, 3]);
3040 let serial = merged
3041 .search_binary_clusters_with_tuning::<false>(&[0], 4, &[0, 1], usize::MAX)
3042 .unwrap();
3043 let parallel = rayon::ThreadPoolBuilder::new()
3044 .num_threads(4)
3045 .build()
3046 .unwrap()
3047 .install(|| merged.search_binary_clusters_with_tuning::<false>(&[0], 4, &[0, 1], 1))
3048 .unwrap();
3049 assert_eq!(parallel, serial, "parallel binary top-k changed results");
3050
3051 let mut second_merge_bytes = Vec::new();
3055 write_merged_ann(&[(&merged, 0), (&first, 4)], &mut second_merge_bytes).unwrap();
3056 let second_merge = AnnDiskIndex::open(
3057 OwnedBytes::new(second_merge_bytes.clone()),
3058 AnnKind::BinaryIvf,
3059 6,
3060 )
3061 .unwrap();
3062 let mut expected_second_payload =
3063 merged_bytes[ANN_HEADER_SIZE..payload_end(&merged)].to_vec();
3064 expected_second_payload
3065 .extend_from_slice(&first_bytes[ANN_HEADER_SIZE..payload_end(&first)]);
3066 assert_eq!(
3067 &second_merge_bytes[ANN_HEADER_SIZE..payload_end(&second_merge)],
3068 expected_second_payload.as_slice(),
3069 );
3070 let mut docs: Vec<u32> = second_merge
3071 .search_binary_clusters::<false>(&[0], 6, &[0, 1])
3072 .unwrap()
3073 .into_iter()
3074 .map(|result| result.0)
3075 .collect();
3076 docs.sort_unstable();
3077 assert_eq!(docs, [0, 1, 2, 3, 4, 5]);
3078 }
3079
3080 #[test]
3081 fn legacy_ivf_tq_payload_is_rejected_while_opening() {
3082 let dim = 8;
3083 let marked_version = crate::structures::mark_ivf_tq_cosine_generation(7);
3084 let centroids = crate::structures::CoarseCentroids {
3085 num_clusters: 1,
3086 dim,
3087 centroids: vec![0.0; dim],
3088 version: marked_version,
3089 soar_config: None,
3090 routing_index: None,
3091 };
3092 let mut bytes = crate::segment::ann_build::build_ivf_tq(
3093 dim,
3094 IvfRoutingMode::Flat,
3095 ¢roids,
3096 &[(0, 0)],
3097 &[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
3098 )
3099 .unwrap();
3100
3101 bytes[24..32].copy_from_slice(&7u64.to_le_bytes());
3103 let error = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::IvfTq, 1)
3104 .err()
3105 .expect("legacy IVF-TQ payload must fail while opening")
3106 .to_string();
3107 assert!(error.contains("unsupported legacy generation"), "{error}");
3108 }
3109
3110 #[test]
3111 fn binary_combined_search_deduplicates_soar_and_bounds_document_results() {
3112 let cluster_0_docs = [0u32, 0, 1];
3116 let cluster_0_ordinals = [0u16, 1, 0];
3117 let cluster_0_codes = [0x00u8, 0xff, 0x03];
3118 let cluster_1_docs = [0u32, 1, 2];
3119 let cluster_1_ordinals = [0u16, 1, 0];
3120 let cluster_1_codes = [0x00u8, 0x0c, 0xf0];
3121 let runs = [
3122 BuildRun {
3123 cluster_id: 0,
3124 doc_ids: &cluster_0_docs,
3125 ordinals: &cluster_0_ordinals,
3126 codes: &cluster_0_codes,
3127 },
3128 BuildRun {
3129 cluster_id: 1,
3130 doc_ids: &cluster_1_docs,
3131 ordinals: &cluster_1_ordinals,
3132 codes: &cluster_1_codes,
3133 },
3134 ];
3135 let mut bytes = Vec::new();
3136 write_built_runs(binary_header(6), &runs, &mut bytes).unwrap();
3137 let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::BinaryIvf, 3).unwrap();
3138 let parallel_pool = rayon::ThreadPoolBuilder::new()
3139 .num_threads(4)
3140 .build()
3141 .unwrap();
3142 let serial_documents = disk
3143 .search_binary_clusters_with_tuning::<true>(&[0], 3, &[0, 1], usize::MAX)
3144 .unwrap();
3145 let parallel_documents = parallel_pool
3146 .install(|| disk.search_binary_clusters_with_tuning::<true>(&[0], 3, &[0, 1], 1))
3147 .unwrap();
3148 assert_eq!(parallel_documents, serial_documents);
3149
3150 let sum = crate::query::MultiValueCombiner::Sum;
3154 let (result, probed) = disk
3155 .search_binary_combined_documents(1, &[0], &[0, 1], sum)
3156 .unwrap();
3157 let parallel_sum = parallel_pool
3158 .install(|| disk.search_binary_combined_documents_with_tuning(1, &[0], &[0, 1], sum, 1))
3159 .unwrap();
3160 assert_eq!(parallel_sum, (result.clone(), probed.clone()));
3161 assert_eq!(result.len(), 1, "combined search must honor k");
3162 assert_eq!(
3163 result[0].doc_id, 1,
3164 "SOAR duplicate changed {sum:?} ranking: {result:?}",
3165 );
3166 assert_eq!(
3169 probed
3170 .iter()
3171 .map(|&(doc_id, ordinal, _)| (doc_id, ordinal))
3172 .collect::<Vec<_>>(),
3173 vec![(1, 0), (1, 1)],
3174 "{sum:?}",
3175 );
3176
3177 let smooth_max = crate::query::MultiValueCombiner::default();
3183 let (result, probed) = disk
3184 .search_binary_combined_documents(1, &[0], &[0, 1], smooth_max)
3185 .unwrap();
3186 let parallel_smooth_max = parallel_pool
3187 .install(|| {
3188 disk.search_binary_combined_documents_with_tuning(1, &[0], &[0, 1], smooth_max, 1)
3189 })
3190 .unwrap();
3191 assert_eq!(parallel_smooth_max, (result.clone(), probed.clone()));
3192 assert_eq!(result.len(), 1, "combined search must honor k");
3193 assert_eq!(result[0].doc_id, 0, "{result:?}");
3194 let expected = smooth_max.combine(&[(0, 1.0), (1, 0.0)]);
3195 assert!(
3196 (result[0].score - expected).abs() < 1e-6,
3197 "SOAR duplicate leaked into {smooth_max:?}: got {}, expected {expected}",
3198 result[0].score,
3199 );
3200 assert_eq!(
3201 probed
3202 .iter()
3203 .map(|&(doc_id, ordinal, _)| (doc_id, ordinal))
3204 .collect::<Vec<_>>(),
3205 vec![(0, 0), (0, 1)],
3206 "{smooth_max:?}",
3207 );
3208
3209 let (top_two, probed) = disk
3210 .search_binary_combined_documents(
3211 2,
3212 &[0],
3213 &[0, 1],
3214 crate::query::MultiValueCombiner::Sum,
3215 )
3216 .unwrap();
3217 assert_eq!(
3218 top_two
3219 .iter()
3220 .map(|candidate| candidate.doc_id)
3221 .collect::<Vec<_>>(),
3222 vec![1, 0],
3223 );
3224 assert_eq!(top_two.len(), 2, "full probing must still return at most k");
3225 assert_eq!(
3228 probed
3229 .iter()
3230 .map(|&(doc_id, ordinal, _)| (doc_id, ordinal))
3231 .collect::<Vec<_>>(),
3232 vec![(0, 0), (0, 1), (1, 0), (1, 1)],
3233 );
3234 }
3235
3236 #[test]
3237 fn combined_ordinal_reduction_handles_out_of_order_runs_for_every_combiner() {
3238 let out_of_order_with_duplicate = vec![
3239 (7, 1, 0.4),
3240 (3, 0, 0.8),
3241 (7, 0, 0.6),
3242 (3, 1, 0.2),
3243 (7, 1, 0.5), ];
3245 for combiner in [
3246 crate::query::MultiValueCombiner::Max,
3247 crate::query::MultiValueCombiner::Sum,
3248 crate::query::MultiValueCombiner::Avg,
3249 crate::query::MultiValueCombiner::default(),
3250 crate::query::MultiValueCombiner::WeightedTopK { k: 2, decay: 0.7 },
3251 ] {
3252 let actual = combine_scored_ordinals(out_of_order_with_duplicate.clone(), 2, combiner);
3253 let mut expected = vec![
3254 AnnDocumentCandidate {
3255 doc_id: 3,
3256 score: combiner.combine(&[(0, 0.8), (1, 0.2)]),
3257 },
3258 AnnDocumentCandidate {
3259 doc_id: 7,
3260 score: combiner.combine(&[(0, 0.6), (1, 0.5)]),
3261 },
3262 ];
3263 expected.sort_unstable_by(|left, right| {
3264 right
3265 .score
3266 .total_cmp(&left.score)
3267 .then_with(|| left.doc_id.cmp(&right.doc_id))
3268 });
3269 assert_eq!(actual, expected, "combiner {combiner:?}");
3270 }
3271 }
3272
3273 fn build_tq_payload(dim: usize, count: usize, seed: u64) -> (Vec<u8>, Vec<Vec<f32>>) {
3274 let codec = std::sync::Arc::new(crate::structures::TqCodec::new(dim));
3275 let mut builder = crate::structures::TqFlatBuilder::new(std::sync::Arc::clone(&codec));
3276 let mut state = seed;
3277 let mut vectors = Vec::new();
3278 let mut flat = Vec::new();
3279 for _ in 0..count {
3280 let vector: Vec<f32> = (0..dim)
3281 .map(|_| {
3282 state = state
3283 .wrapping_mul(6364136223846793005)
3284 .wrapping_add(1442695040888963407);
3285 ((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5
3286 })
3287 .collect();
3288 flat.extend_from_slice(&vector);
3289 vectors.push(vector);
3290 }
3291 let labels: Vec<(u32, u16)> = (0..count).map(|index| (index as u32, 0)).collect();
3292 builder.add_batch(&labels, &flat).unwrap();
3293 builder.finish();
3294 let mut bytes = Vec::new();
3295 write_built_tq_flat(&builder, &mut bytes).unwrap();
3296 (bytes, vectors)
3297 }
3298
3299 #[test]
3300 fn tq_combined_scan_ranks_complete_documents_instead_of_individual_values() {
3301 let dim = 8;
3302 let codec = std::sync::Arc::new(crate::structures::TqCodec::new(dim));
3303 let mut builder = crate::structures::TqFlatBuilder::new(std::sync::Arc::clone(&codec));
3304 let labels = [(0u32, 0u16), (1, 0), (1, 1)];
3305 let vectors = [
3306 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,
3309 ];
3310 builder.add_batch(&labels, &vectors).unwrap();
3311 builder.finish();
3312 let mut bytes = Vec::new();
3313 write_built_tq_flat(&builder, &mut bytes).unwrap();
3314 let disk = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::TqFlat, 2).unwrap();
3315 let query = [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0];
3316 let plan = crate::structures::TqQueryPlan::build(&codec, &query);
3317
3318 let max = disk.search_tq_distinct(1, &plan).unwrap();
3319 let sum = disk
3320 .search_tq_combined_documents(1, &plan, crate::query::MultiValueCombiner::Sum)
3321 .unwrap();
3322 let default_combiner = disk
3323 .search_tq_combined_documents(1, &plan, crate::query::MultiValueCombiner::default())
3324 .unwrap();
3325 assert_eq!(max[0].0, 0, "fixture must favor doc 0 by Max: {max:?}");
3326 assert_eq!(
3327 sum[0].doc_id, 1,
3328 "two complete values must make doc 1 win by Sum: {sum:?}"
3329 );
3330 assert_eq!(
3334 default_combiner[0].doc_id, 0,
3335 "default smooth-max must follow the best value: {default_combiner:?}"
3336 );
3337 assert!(sum[0].score > max[0].2);
3338
3339 let (single_bytes, _) = build_tq_payload(dim, 1, 91);
3343 let single = AnnDiskIndex::open(OwnedBytes::new(single_bytes), AnnKind::TqFlat, 1).unwrap();
3344 let mut merged_bytes = Vec::new();
3345 write_merged_ann(&[(&disk, 1), (&single, 0)], &mut merged_bytes).unwrap();
3346 let merged = AnnDiskIndex::open(OwnedBytes::new(merged_bytes), AnnKind::TqFlat, 3).unwrap();
3347 let merged_sum = merged
3348 .search_tq_combined_documents(1, &plan, crate::query::MultiValueCombiner::Sum)
3349 .unwrap();
3350 assert_eq!(merged_sum[0].doc_id, 2);
3351 }
3352
3353 #[test]
3354 fn tq_payload_roundtrip_search_and_pure_copy_merge() {
3355 let dim = 20; let count = 21;
3357 let (bytes, vectors) = build_tq_payload(dim, count, 42);
3358 let index = AnnDiskIndex::open(
3359 OwnedBytes::new(bytes.clone()),
3360 AnnKind::TqFlat,
3361 count as u32,
3362 )
3363 .unwrap();
3364 assert_eq!(index.header().vector_count, count);
3365
3366 let codec = crate::structures::TqCodec::new(dim);
3368 for target in [0usize, 7, 20] {
3369 let plan = crate::structures::TqQueryPlan::build(&codec, &vectors[target]);
3370 let results = index.search_tq_distinct(3, &plan).unwrap();
3371 assert_eq!(
3372 results[0].0, target as u32,
3373 "query duplicating vector {target} must rank it first: {results:?}"
3374 );
3375 }
3376
3377 let (second_bytes, _) = build_tq_payload(dim, 5, 77);
3379 let second =
3380 AnnDiskIndex::open(OwnedBytes::new(second_bytes.clone()), AnnKind::TqFlat, 5).unwrap();
3381 let mut merged_bytes = Vec::new();
3382 write_merged_ann(&[(&index, 0), (&second, count as u32)], &mut merged_bytes).unwrap();
3383 let merged = AnnDiskIndex::open(
3384 OwnedBytes::new(merged_bytes.clone()),
3385 AnnKind::TqFlat,
3386 count as u32 + 5,
3387 )
3388 .unwrap();
3389 let mut expected_payload = bytes[ANN_HEADER_SIZE..payload_end(&index)].to_vec();
3390 expected_payload.extend_from_slice(&second_bytes[ANN_HEADER_SIZE..payload_end(&second)]);
3391 assert_eq!(
3392 &merged_bytes[ANN_HEADER_SIZE..payload_end(&merged)],
3393 expected_payload.as_slice(),
3394 "TQ merge must be a pure byte copy of the source columns",
3395 );
3396 let plan = crate::structures::TqQueryPlan::build(&codec, &vectors[7]);
3397 let results = merged.search_tq_distinct(1, &plan).unwrap();
3398 assert_eq!(results[0].0, 7, "merged payload must keep doc bases");
3399 }
3400
3401 #[test]
3402 fn tq_parallel_fold_matches_a_sequential_scan() {
3403 use crate::structures::vector::quantization::{
3404 TQ_BLOCK_LANES, tq_block_bytes, tq_score_block,
3405 };
3406
3407 let dim = 8;
3410 let count = TQ_PARALLEL_SCAN_MIN_VECTORS;
3411 let (bytes, vectors) = build_tq_payload(dim, count, 87);
3412 let disk =
3413 AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::TqFlat, count as u32).unwrap();
3414 let codec = crate::structures::TqCodec::new(dim);
3415 let plan = crate::structures::TqQueryPlan::build(&codec, &vectors[count / 3]);
3416 let k = 31;
3417
3418 let block_bytes = tq_block_bytes(disk.header().code_size);
3419 assert_eq!(
3420 disk.runs.len(),
3421 1,
3422 "the regression must parallelize chunks inside one run"
3423 );
3424 assert!(
3425 disk.runs[0].codes.len() / block_bytes > TQ_PARALLEL_SCAN_CHUNK_BLOCKS,
3426 "the single run must span multiple parallel chunks"
3427 );
3428 let raw = disk.raw.as_slice();
3429 let mut reference = BoundedAnnCollector::<true, true>::new(k);
3430 let mut scores = [0.0f32; TQ_BLOCK_LANES];
3431 for run in &disk.runs {
3432 let codes = &raw[run.codes.clone()];
3433 for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
3434 tq_score_block(&plan, block, &mut scores);
3435 let lane_base = block_index * TQ_BLOCK_LANES;
3436 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
3437 for (lane, &score) in scores.iter().enumerate().take(lanes) {
3438 let index = lane_base + lane;
3439 reference.insert(
3440 run_doc_id(raw, run, index).unwrap(),
3441 read_u16(raw, run.ordinals.start + index * 2),
3442 score,
3443 );
3444 }
3445 }
3446 }
3447
3448 assert_eq!(
3449 disk.search_tq_distinct(k, &plan).unwrap(),
3450 reference.into_sorted_results(),
3451 );
3452 }
3453
3454 #[test]
3455 fn ivf_tq_scale_pruning_matches_the_unpruned_scan() {
3456 use crate::structures::vector::ivf::{CoarseCentroids, CoarseConfig};
3457 use crate::structures::vector::quantization::{
3458 TQ_BLOCK_LANES, tq_ivf_block_bytes, tq_score_ivf_block,
3459 };
3460 use crate::structures::{IvfTqIndex, TqCodec, TqIvfEncodeScratch, TqIvfQueryPlan};
3461
3462 let dim = 32;
3463 let count = 400usize;
3464 let codec = std::sync::Arc::new(TqCodec::new(dim));
3465 let mut state = 5u64;
3466 let mut next = move || {
3467 state = state
3468 .wrapping_mul(6364136223846793005)
3469 .wrapping_add(1442695040888963407);
3470 ((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5
3471 };
3472 let vectors: Vec<Vec<f32>> = (0..count)
3473 .map(|_| {
3474 let mut v: Vec<f32> = (0..dim).map(|_| next()).collect();
3475 let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
3476 v.iter_mut().for_each(|x| *x /= norm);
3477 v
3478 })
3479 .collect();
3480 let mut centroids = CoarseCentroids::train(&CoarseConfig::new(dim, 8), &vectors, "test");
3481 centroids.version = crate::structures::mark_ivf_tq_cosine_generation(centroids.version);
3482 let mut index = IvfTqIndex::new(
3483 dim,
3484 crate::dsl::IvfRoutingMode::Flat,
3485 centroids.version,
3486 std::sync::Arc::clone(&codec),
3487 );
3488 let mut scratch = TqIvfEncodeScratch::default();
3489 for (i, vector) in vectors.iter().enumerate() {
3490 index.add_vector(
3491 ¢roids,
3492 (i / 2) as u32,
3493 (i % 2) as u16,
3494 vector,
3495 &mut scratch,
3496 );
3497 }
3498 let mut bytes = Vec::new();
3499 write_built_ivf_tq(&index, centroids.num_clusters, &mut bytes).unwrap();
3500 let disk =
3501 AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::IvfTq, count as u32).unwrap();
3502
3503 let block_bytes = tq_ivf_block_bytes(disk.header().code_size);
3506 let raw = disk.raw.as_slice();
3507 for run in &disk.runs {
3508 let codes = &raw[run.codes.clone()];
3509 let mut previous_scale = f32::INFINITY;
3510 for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
3511 let lane_base = block_index * TQ_BLOCK_LANES;
3512 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
3513 let mut block_scales = block[..TQ_BLOCK_LANES * size_of::<f32>()]
3514 .chunks_exact(size_of::<f32>())
3515 .take(lanes)
3516 .map(|lane| f32::from_le_bytes(lane.try_into().unwrap()));
3517 let first_scale = block_scales.next().unwrap();
3518 assert_eq!(tq_ivf_block_max_scale(block), first_scale);
3519 assert!(first_scale <= previous_scale);
3520 previous_scale = first_scale;
3521 for scale in block_scales {
3522 assert!(scale <= previous_scale);
3523 previous_scale = scale;
3524 }
3525 }
3526 }
3527 let k = 10;
3528 let parallel_pool = rayon::ThreadPoolBuilder::new()
3529 .num_threads(4)
3530 .build()
3531 .unwrap();
3532 for query_seed in [1u64, 9, 42] {
3533 let mut qstate = query_seed;
3534 let mut qnext = move || {
3535 qstate = qstate
3536 .wrapping_mul(6364136223846793005)
3537 .wrapping_add(1442695040888963407);
3538 ((qstate >> 33) as f32 / (1u64 << 31) as f32) - 0.5
3539 };
3540 let query: Vec<f32> = (0..dim).map(|_| qnext()).collect();
3541 let plan = TqIvfQueryPlan::build(
3542 ¢roids,
3543 &codec,
3544 &query,
3545 8,
3546 crate::dsl::IvfRoutingMode::Flat,
3547 );
3548 let mut reference = BoundedAnnCollector::<true, true>::new(k);
3551 let mut unpruned_scores = Vec::new();
3552 let mut scores = [0.0f32; TQ_BLOCK_LANES];
3553 for (cluster_id, cluster_dot) in plan.cluster_dots() {
3554 for run in disk.cluster_runs(cluster_id) {
3555 let codes = &raw[run.codes.clone()];
3556 for (block_index, block) in codes.chunks_exact(block_bytes).enumerate() {
3557 tq_score_ivf_block(plan.tq_plan(), block, cluster_dot, &mut scores);
3558 let lane_base = block_index * TQ_BLOCK_LANES;
3559 let lanes = TQ_BLOCK_LANES.min(run.count.saturating_sub(lane_base));
3560 for (lane, &score) in scores.iter().enumerate().take(lanes) {
3561 let idx = lane_base + lane;
3562 reference.insert(
3563 run_doc_id(raw, run, idx).unwrap(),
3564 read_u16(raw, run.ordinals.start + idx * 2),
3565 score,
3566 );
3567 unpruned_scores.push((
3568 run_doc_id(raw, run, idx).unwrap(),
3569 read_u16(raw, run.ordinals.start + idx * 2),
3570 score,
3571 ));
3572 }
3573 }
3574 }
3575 }
3576
3577 let pruned = disk.search_ivf_tq_distinct(k, &plan).unwrap();
3578 let forced_parallel = parallel_pool.install(|| {
3579 disk.search_ivf_tq_distinct_with_tuning(k, &plan, 1, 4)
3582 .unwrap()
3583 });
3584 let reference = reference.into_sorted_results();
3585 assert_eq!(
3586 pruned, reference,
3587 "scale-bound pruning must not change the estimated top-k (seed {query_seed})"
3588 );
3589 assert_eq!(
3590 forced_parallel, reference,
3591 "parallel IVF-TQ pruning must match the unpruned top-k (seed {query_seed})"
3592 );
3593 for combiner in [
3594 crate::query::MultiValueCombiner::Max,
3595 crate::query::MultiValueCombiner::Sum,
3596 crate::query::MultiValueCombiner::Avg,
3597 crate::query::MultiValueCombiner::default(),
3598 crate::query::MultiValueCombiner::WeightedTopK { k: 3, decay: 0.7 },
3599 ] {
3600 let expected = combine_scored_ordinals(unpruned_scores.clone(), k, combiner);
3601 let combined = disk
3602 .search_ivf_tq_combined_documents(k, &plan, combiner)
3603 .unwrap();
3604 let parallel_combined = parallel_pool
3605 .install(|| {
3606 disk.search_ivf_tq_combined_documents_with_tuning(k, &plan, combiner, 1, 4)
3607 })
3608 .unwrap();
3609 assert_eq!(
3610 combined, expected,
3611 "combined IVF-TQ scan diverged from the unpruned reference \
3612 for {combiner:?} (seed {query_seed})",
3613 );
3614 assert_eq!(
3615 parallel_combined, expected,
3616 "parallel combined IVF-TQ scan diverged from the unpruned reference \
3617 for {combiner:?} (seed {query_seed})",
3618 );
3619 assert!(combined.len() <= k);
3620 }
3621 }
3622 }
3623
3624 #[test]
3632 #[ignore]
3633 fn ivf_tq_parallel_scan_benchmark() {
3634 use crate::structures::vector::ivf::CoarseCentroids;
3635 use crate::structures::{IvfTqIndex, TqCodec, TqIvfEncodeScratch, TqIvfQueryPlan};
3636
3637 fn env_usize(name: &str, default: usize) -> usize {
3638 std::env::var(name)
3639 .ok()
3640 .map(|value| {
3641 value
3642 .parse::<usize>()
3643 .unwrap_or_else(|_| panic!("{name} must be a positive integer"))
3644 })
3645 .unwrap_or(default)
3646 .max(1)
3647 }
3648
3649 fn median_ms(samples: &mut [f64]) -> f64 {
3650 samples.sort_unstable_by(f64::total_cmp);
3651 samples[samples.len() / 2]
3652 }
3653
3654 let dim = env_usize("IVF_SCAN_BENCH_DIM", 128);
3655 let count = env_usize("IVF_SCAN_BENCH_DOCS", 262_144);
3656 let iterations = env_usize("IVF_SCAN_BENCH_ITERS", 30);
3657 let threads = env_usize("IVF_SCAN_BENCH_THREADS", num_cpus::get().min(8));
3658 let codec = std::sync::Arc::new(TqCodec::new(dim));
3659 let version = crate::structures::mark_ivf_tq_cosine_generation(0x51ca_0001);
3660 let mut centroid = vec![0.0f32; dim];
3661 centroid[0] = 1.0;
3662 let centroids = CoarseCentroids {
3663 num_clusters: 1,
3664 dim,
3665 centroids: centroid,
3666 version,
3667 soar_config: None,
3668 routing_index: None,
3669 };
3670 let mut index = IvfTqIndex::new(
3671 dim,
3672 crate::dsl::IvfRoutingMode::Flat,
3673 version,
3674 std::sync::Arc::clone(&codec),
3675 );
3676 let mut scratch = TqIvfEncodeScratch::default();
3677 let mut state = 0x5eed_u64;
3678 let mut vector = vec![0.0f32; dim];
3679 for posting_index in 0..count {
3680 for value in &mut vector {
3681 state = state
3682 .wrapping_mul(6364136223846793005)
3683 .wrapping_add(1442695040888963407);
3684 *value = ((state >> 33) as f32 / (1u64 << 31) as f32) - 0.5;
3685 }
3686 let norm = vector.iter().map(|value| value * value).sum::<f32>().sqrt();
3687 vector.iter_mut().for_each(|value| *value /= norm);
3688 index.add_vector(
3689 ¢roids,
3690 u32::try_from(posting_index / 2).unwrap(),
3691 u16::try_from(posting_index % 2).unwrap(),
3692 &vector,
3693 &mut scratch,
3694 );
3695 }
3696 let mut bytes = Vec::new();
3697 write_built_ivf_tq(&index, 1, &mut bytes).unwrap();
3698 let disk = AnnDiskIndex::open(
3699 OwnedBytes::new(bytes),
3700 AnnKind::IvfTq,
3701 u32::try_from(count.div_ceil(2)).unwrap(),
3702 )
3703 .unwrap();
3704 let query = vec![1.0f32; dim];
3705 let plan = TqIvfQueryPlan::build(
3706 ¢roids,
3707 &codec,
3708 &query,
3709 1,
3710 crate::dsl::IvfRoutingMode::Flat,
3711 );
3712 let pool = rayon::ThreadPoolBuilder::new()
3713 .num_threads(threads)
3714 .build()
3715 .unwrap();
3716 let serial_pool = rayon::ThreadPoolBuilder::new()
3717 .num_threads(1)
3718 .build()
3719 .unwrap();
3720
3721 let serial = pool
3722 .install(|| disk.search_ivf_tq_distinct_with_tuning(20, &plan, usize::MAX, 512))
3723 .unwrap();
3724 let parallel = pool
3725 .install(|| disk.search_ivf_tq_distinct_with_tuning(20, &plan, 1, 512))
3726 .unwrap();
3727 assert_eq!(parallel, serial);
3728
3729 let mut serial_ms = Vec::with_capacity(iterations);
3730 let mut parallel_ms = Vec::with_capacity(iterations);
3731 for _ in 0..iterations {
3732 let start = std::time::Instant::now();
3733 std::hint::black_box(
3734 pool.install(|| {
3735 disk.search_ivf_tq_distinct_with_tuning(20, &plan, usize::MAX, 512)
3736 })
3737 .unwrap(),
3738 );
3739 serial_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
3740
3741 let start = std::time::Instant::now();
3742 std::hint::black_box(
3743 pool.install(|| disk.search_ivf_tq_distinct_with_tuning(20, &plan, 1, 512))
3744 .unwrap(),
3745 );
3746 parallel_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
3747 }
3748 let serial_p50 = median_ms(&mut serial_ms);
3749 let parallel_p50 = median_ms(&mut parallel_ms);
3750 println!(
3751 "IVF-TQ distinct scan: postings={count} dim={dim} threads={threads} \
3752 serial_p50={serial_p50:.3}ms parallel_p50={parallel_p50:.3}ms speedup={:.2}x",
3753 serial_p50 / parallel_p50,
3754 );
3755
3756 let combiner = crate::query::MultiValueCombiner::Sum;
3757 let serial_combined = serial_pool
3758 .install(|| {
3759 disk.search_ivf_tq_combined_documents_with_tuning(20, &plan, combiner, 1, 512)
3760 })
3761 .unwrap();
3762 let parallel_combined = pool
3763 .install(|| {
3764 disk.search_ivf_tq_combined_documents_with_tuning(20, &plan, combiner, 1, 512)
3765 })
3766 .unwrap();
3767 assert_eq!(parallel_combined, serial_combined);
3768 let mut serial_combined_ms = Vec::with_capacity(iterations);
3769 let mut parallel_combined_ms = Vec::with_capacity(iterations);
3770 for _ in 0..iterations {
3771 let start = std::time::Instant::now();
3772 std::hint::black_box(
3773 serial_pool
3774 .install(|| {
3775 disk.search_ivf_tq_combined_documents_with_tuning(
3776 20, &plan, combiner, 1, 512,
3777 )
3778 })
3779 .unwrap(),
3780 );
3781 serial_combined_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
3782
3783 let start = std::time::Instant::now();
3784 std::hint::black_box(
3785 pool.install(|| {
3786 disk.search_ivf_tq_combined_documents_with_tuning(20, &plan, combiner, 1, 512)
3787 })
3788 .unwrap(),
3789 );
3790 parallel_combined_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
3791 }
3792 let serial_combined_p50 = median_ms(&mut serial_combined_ms);
3793 let parallel_combined_p50 = median_ms(&mut parallel_combined_ms);
3794 println!(
3795 "IVF-TQ combined scan: postings={count} dim={dim} threads={threads} \
3796 serial_p50={serial_combined_p50:.3}ms parallel_p50={parallel_combined_p50:.3}ms \
3797 speedup={:.2}x",
3798 serial_combined_p50 / parallel_combined_p50,
3799 );
3800 }
3801
3802 #[test]
3806 #[ignore]
3807 fn binary_ivf_parallel_scan_benchmark() {
3808 fn env_usize(name: &str, default: usize) -> usize {
3809 std::env::var(name)
3810 .ok()
3811 .map(|value| value.parse::<usize>().unwrap())
3812 .unwrap_or(default)
3813 .max(1)
3814 }
3815
3816 fn median_ms(samples: &mut [f64]) -> f64 {
3817 samples.sort_unstable_by(f64::total_cmp);
3818 samples[samples.len() / 2]
3819 }
3820
3821 let count = env_usize("BINARY_SCAN_BENCH_POSTINGS", 131_072);
3822 let dim_bits = env_usize("BINARY_SCAN_BENCH_DIM_BITS", 2_048);
3823 assert!(dim_bits.is_multiple_of(8));
3824 let code_size = dim_bits / 8;
3825 let iterations = env_usize("BINARY_SCAN_BENCH_ITERS", 30);
3826 let threads = env_usize("BINARY_SCAN_BENCH_THREADS", num_cpus::get().min(8));
3827 let mut state = 0xb1a4_5eed_u64;
3828 let mut codes = vec![0u8; count * code_size];
3829 for code in &mut codes {
3830 state = state
3831 .wrapping_mul(6364136223846793005)
3832 .wrapping_add(1442695040888963407);
3833 *code = (state >> 56) as u8;
3834 }
3835 let query = vec![0xa5u8; code_size];
3836 let unique_docs: Vec<u32> = (0..count)
3837 .map(|index| u32::try_from(index).unwrap())
3838 .collect();
3839 let unique_ordinals = vec![0u16; count];
3840 let multi_docs: Vec<u32> = (0..count)
3841 .map(|index| u32::try_from(index / 2).unwrap())
3842 .collect();
3843 let multi_ordinals: Vec<u16> = (0..count)
3844 .map(|index| u16::try_from(index % 2).unwrap())
3845 .collect();
3846 let header = AnnDiskHeader {
3847 kind: AnnKind::BinaryIvf,
3848 routing: IvfRoutingMode::Flat,
3849 dim: dim_bits,
3850 code_size,
3851 num_clusters: 1,
3852 quantizer_version: 0xb1a4,
3853 codebook_version: 0,
3854 vector_count: count,
3855 };
3856 let mut unique_bytes = Vec::new();
3857 write_built_runs(
3858 header.clone(),
3859 &[BuildRun {
3860 cluster_id: 0,
3861 doc_ids: &unique_docs,
3862 ordinals: &unique_ordinals,
3863 codes: &codes,
3864 }],
3865 &mut unique_bytes,
3866 )
3867 .unwrap();
3868 let unique_disk = AnnDiskIndex::open(
3869 OwnedBytes::new(unique_bytes),
3870 AnnKind::BinaryIvf,
3871 u32::try_from(count).unwrap(),
3872 )
3873 .unwrap();
3874 let mut multi_bytes = Vec::new();
3875 write_built_runs(
3876 header,
3877 &[BuildRun {
3878 cluster_id: 0,
3879 doc_ids: &multi_docs,
3880 ordinals: &multi_ordinals,
3881 codes: &codes,
3882 }],
3883 &mut multi_bytes,
3884 )
3885 .unwrap();
3886 let multi_disk = AnnDiskIndex::open(
3887 OwnedBytes::new(multi_bytes),
3888 AnnKind::BinaryIvf,
3889 u32::try_from(count.div_ceil(2)).unwrap(),
3890 )
3891 .unwrap();
3892 let pool = rayon::ThreadPoolBuilder::new()
3893 .num_threads(threads)
3894 .build()
3895 .unwrap();
3896 let serial_pool = rayon::ThreadPoolBuilder::new()
3897 .num_threads(1)
3898 .build()
3899 .unwrap();
3900
3901 let serial = pool
3902 .install(|| {
3903 unique_disk.search_binary_clusters_with_tuning::<false>(
3904 &query,
3905 20,
3906 &[0],
3907 usize::MAX,
3908 )
3909 })
3910 .unwrap();
3911 let parallel = pool
3912 .install(|| {
3913 unique_disk.search_binary_clusters_with_tuning::<false>(&query, 20, &[0], 1)
3914 })
3915 .unwrap();
3916 assert_eq!(parallel, serial);
3917 let mut serial_ms = Vec::with_capacity(iterations);
3918 let mut parallel_ms = Vec::with_capacity(iterations);
3919 for _ in 0..iterations {
3920 let start = std::time::Instant::now();
3921 std::hint::black_box(
3922 pool.install(|| {
3923 unique_disk.search_binary_clusters_with_tuning::<false>(
3924 &query,
3925 20,
3926 &[0],
3927 usize::MAX,
3928 )
3929 })
3930 .unwrap(),
3931 );
3932 serial_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
3933
3934 let start = std::time::Instant::now();
3935 std::hint::black_box(
3936 pool.install(|| {
3937 unique_disk.search_binary_clusters_with_tuning::<false>(&query, 20, &[0], 1)
3938 })
3939 .unwrap(),
3940 );
3941 parallel_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
3942 }
3943 let serial_p50 = median_ms(&mut serial_ms);
3944 let parallel_p50 = median_ms(&mut parallel_ms);
3945 println!(
3946 "binary-IVF distinct scan: postings={count} dim_bits={dim_bits} threads={threads} \
3947 serial_p50={serial_p50:.3}ms parallel_p50={parallel_p50:.3}ms speedup={:.2}x",
3948 serial_p50 / parallel_p50,
3949 );
3950
3951 let combiner = crate::query::MultiValueCombiner::Sum;
3952 let serial_combined = serial_pool
3953 .install(|| {
3954 multi_disk.search_binary_combined_documents_with_tuning(
3955 20,
3956 &query,
3957 &[0],
3958 combiner,
3959 1,
3960 )
3961 })
3962 .unwrap();
3963 let parallel_combined = pool
3964 .install(|| {
3965 multi_disk.search_binary_combined_documents_with_tuning(
3966 20,
3967 &query,
3968 &[0],
3969 combiner,
3970 1,
3971 )
3972 })
3973 .unwrap();
3974 assert_eq!(parallel_combined, serial_combined);
3975 let mut serial_combined_ms = Vec::with_capacity(iterations);
3976 let mut parallel_combined_ms = Vec::with_capacity(iterations);
3977 for _ in 0..iterations {
3978 let start = std::time::Instant::now();
3979 std::hint::black_box(
3980 serial_pool
3981 .install(|| {
3982 multi_disk.search_binary_combined_documents_with_tuning(
3983 20,
3984 &query,
3985 &[0],
3986 combiner,
3987 1,
3988 )
3989 })
3990 .unwrap(),
3991 );
3992 serial_combined_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
3993
3994 let start = std::time::Instant::now();
3995 std::hint::black_box(
3996 pool.install(|| {
3997 multi_disk.search_binary_combined_documents_with_tuning(
3998 20,
3999 &query,
4000 &[0],
4001 combiner,
4002 1,
4003 )
4004 })
4005 .unwrap(),
4006 );
4007 parallel_combined_ms.push(start.elapsed().as_secs_f64() * 1_000.0);
4008 }
4009 let serial_combined_p50 = median_ms(&mut serial_combined_ms);
4010 let parallel_combined_p50 = median_ms(&mut parallel_combined_ms);
4011 println!(
4012 "binary-IVF combined scan: postings={count} dim_bits={dim_bits} threads={threads} \
4013 serial_p50={serial_combined_p50:.3}ms parallel_p50={parallel_combined_p50:.3}ms \
4014 speedup={:.2}x",
4015 serial_combined_p50 / parallel_combined_p50,
4016 );
4017 }
4018
4019 #[test]
4020 fn open_rejects_tq_payload_with_inconsistent_geometry() {
4021 let (bytes, _) = build_tq_payload(20, 4, 9);
4022 let mut corrupted = bytes.clone();
4024 corrupted[12..16].copy_from_slice(&15u32.to_le_bytes());
4025 assert!(
4026 AnnDiskIndex::open(OwnedBytes::new(corrupted), AnnKind::TqFlat, 4).is_err(),
4027 "TQ header with code_size != padded_dim/2 must be refused"
4028 );
4029
4030 let (short_bytes, _) = build_tq_payload(20, 4, 9);
4032 let mut wrong_kind = short_bytes.clone();
4033 wrong_kind[4] = AnnKind::BinaryIvf as u8;
4034 assert!(
4035 AnnDiskIndex::open(OwnedBytes::new(wrong_kind), AnnKind::BinaryIvf, 4).is_err(),
4036 "TQ block-padded columns must not validate under another kind"
4037 );
4038
4039 let (legacy_bytes, _) = build_tq_payload(20, 4, 9);
4041 let mut legacy_kind = legacy_bytes.clone();
4042 legacy_kind[4] = 1;
4043 let Err(error) = AnnDiskIndex::open(OwnedBytes::new(legacy_kind), AnnKind::TqFlat, 4)
4044 else {
4045 panic!("retired IVF-PQ payloads must not open");
4046 };
4047 assert!(
4048 error.to_string().contains("IVF-PQ"),
4049 "error must name the retired format: {error}"
4050 );
4051 assert!(AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::TqFlat, 4).is_ok());
4052 }
4053
4054 #[test]
4055 fn open_rejects_old_or_out_of_range_ann_payloads() {
4056 let mut legacy = vec![0u8; ANN_HEADER_SIZE + ANN_FOOTER_SIZE];
4057 legacy[..4].copy_from_slice(b"old!");
4058 assert!(AnnDiskIndex::open(OwnedBytes::new(legacy), AnnKind::BinaryIvf, 1).is_err());
4059
4060 let docs = [0u32];
4061 let ordinals = [0u16];
4062 let codes = [0u8];
4063 let runs = [BuildRun {
4064 cluster_id: 0,
4065 doc_ids: &docs,
4066 ordinals: &ordinals,
4067 codes: &codes,
4068 }];
4069 let mut bytes = Vec::new();
4070 write_built_runs(binary_header(1), &runs, &mut bytes).unwrap();
4071 let footer = bytes.len() - ANN_FOOTER_SIZE;
4072 let directory = usize::try_from(u64::from_le_bytes(
4073 bytes[footer..footer + 8].try_into().unwrap(),
4074 ))
4075 .unwrap();
4076 bytes[directory + 12..directory + 16].copy_from_slice(&10u32.to_le_bytes());
4077 assert!(AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::BinaryIvf, 1).is_err());
4078 }
4079}