1pub(crate) mod bmp;
4pub(crate) mod loader;
5mod types;
6
7pub use bmp::{BmpDimStats, BmpIndex};
8#[cfg(feature = "native")]
9pub(crate) use types::DimRawData;
10pub use types::{SparseIndex, VectorIndex, VectorSearchResult};
11
12const MAX_PREFIX_TERMS: usize = 1_024;
16const MAX_PREFIX_POSTINGS: u64 = 5_000_000;
17const MAX_DENSE_CANDIDATES_PER_SEGMENT: usize = 20_000;
21const DENSE_SCORE_BATCH: usize = 4_096;
23const BINARY_SCORE_BATCH: usize = 8_192;
24const MAX_VECTOR_SCORE_BATCH_BYTES: usize = 8 * 1024 * 1024;
25
26#[derive(Debug, Clone, Default)]
32pub struct SegmentMemoryStats {
33 pub segment_id: u128,
35 pub num_docs: u32,
37 pub term_dict_cache_bytes: usize,
39 pub store_cache_bytes: usize,
41 pub sparse_heap_bytes: usize,
43 pub dense_heap_bytes: usize,
45 pub term_bloom_file_bytes: u64,
47 pub sparse_file_backed_bytes: u64,
49 pub dense_file_backed_bytes: u64,
51 pub pinned_metadata_bytes: u64,
53 pub pin_intended_bytes: u64,
56 pub sparse_pinned_metadata_bytes: u64,
58 pub sparse_pin_intended_bytes: u64,
60 pub dense_pinned_metadata_bytes: u64,
62 pub dense_pin_intended_bytes: u64,
64}
65
66impl SegmentMemoryStats {
67 pub fn estimated_heap_bytes(&self) -> usize {
69 self.term_dict_cache_bytes
70 + self.store_cache_bytes
71 + self.sparse_heap_bytes
72 + self.dense_heap_bytes
73 }
74
75 pub fn file_backed_bytes(&self) -> u64 {
79 self.term_bloom_file_bytes
80 .saturating_add(self.sparse_file_backed_bytes)
81 .saturating_add(self.dense_file_backed_bytes)
82 }
83}
84
85use std::cmp::Ordering;
86use std::collections::BinaryHeap;
87use std::sync::Arc;
88
89use rustc_hash::{FxHashMap, FxHashSet};
90
91use super::vector_data::LazyFlatVectorData;
92use crate::directories::{Directory, FileHandle};
93use crate::dsl::{DenseVectorQuantization, Document, Field, Schema};
94use crate::query::{MAX_DENSE_NPROBE, MAX_DENSE_RERANK_FACTOR};
95use crate::structures::{
96 AsyncSSTableReader, BlockPostingList, CoarseCentroids, SSTableStats, TermInfo,
97};
98use crate::{DocId, Error, Result};
99
100use super::store::{AsyncStoreReader, RawStoreBlock};
101use super::types::{SegmentFiles, SegmentId, SegmentMeta};
102
103pub(crate) fn combine_ordinal_results(
109 raw: impl IntoIterator<Item = (u32, u16, f32)>,
110 combiner: crate::query::MultiValueCombiner,
111 limit: usize,
112) -> Vec<VectorSearchResult> {
113 let collected: Vec<(u32, u16, f32)> = raw.into_iter().collect();
114
115 let num_raw = collected.len();
116 if log::log_enabled!(log::Level::Debug) {
117 let mut ids: Vec<u32> = collected.iter().map(|(d, _, _)| *d).collect();
118 ids.sort_unstable();
119 ids.dedup();
120 log::debug!(
121 "combine_ordinal_results: {} raw entries, {} unique docs, combiner={:?}, limit={}",
122 num_raw,
123 ids.len(),
124 combiner,
125 limit
126 );
127 }
128
129 let all_single = collected.iter().all(|&(_, ord, _)| ord == 0);
131 if all_single {
132 let mut results: Vec<VectorSearchResult> = collected
133 .into_iter()
134 .map(|(doc_id, _, score)| VectorSearchResult::new(doc_id, score, vec![(0, score)]))
135 .collect();
136 results.sort_unstable_by(|a, b| {
137 b.score
138 .total_cmp(&a.score)
139 .then_with(|| a.doc_id.cmp(&b.doc_id))
140 });
141 results.truncate(limit);
142 return results;
143 }
144
145 let mut doc_ordinals: rustc_hash::FxHashMap<DocId, Vec<(u32, f32)>> =
147 rustc_hash::FxHashMap::default();
148 for (doc_id, ordinal, score) in collected {
149 doc_ordinals
150 .entry(doc_id as DocId)
151 .or_default()
152 .push((ordinal as u32, score));
153 }
154 let mut results: Vec<VectorSearchResult> = doc_ordinals
155 .into_iter()
156 .map(|(doc_id, ordinals)| {
157 let combined_score = combiner.combine(&ordinals);
158 VectorSearchResult::new(doc_id, combined_score, ordinals)
159 })
160 .collect();
161 results.sort_unstable_by(|a, b| {
162 b.score
163 .total_cmp(&a.score)
164 .then_with(|| a.doc_id.cmp(&b.doc_id))
165 });
166 results.truncate(limit);
167 results
168}
169
170struct HeapVectorResult(VectorSearchResult);
175
176impl PartialEq for HeapVectorResult {
177 fn eq(&self, other: &Self) -> bool {
178 self.0.score.to_bits() == other.0.score.to_bits() && self.0.doc_id == other.0.doc_id
179 }
180}
181
182impl Eq for HeapVectorResult {}
183
184impl Ord for HeapVectorResult {
185 fn cmp(&self, other: &Self) -> Ordering {
186 other
189 .0
190 .score
191 .total_cmp(&self.0.score)
192 .then_with(|| self.0.doc_id.cmp(&other.0.doc_id))
193 }
194}
195
196impl PartialOrd for HeapVectorResult {
197 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
198 Some(self.cmp(other))
199 }
200}
201
202struct FlatDocumentCollector {
206 heap: BinaryHeap<HeapVectorResult>,
207 limit: usize,
208 combiner: crate::query::MultiValueCombiner,
209 current_doc: Option<DocId>,
210 current_ordinals: Vec<(u32, f32)>,
211}
212
213impl FlatDocumentCollector {
214 fn new(limit: usize, combiner: crate::query::MultiValueCombiner) -> Self {
215 Self {
216 heap: BinaryHeap::with_capacity(limit.min(8 * 1024)),
217 limit,
218 combiner,
219 current_doc: None,
220 current_ordinals: Vec::new(),
221 }
222 }
223
224 fn push(&mut self, doc_id: DocId, ordinal: u16, score: f32) {
225 if self.current_doc.is_some_and(|current| current != doc_id) {
226 self.finish_current();
227 }
228 self.current_doc = Some(doc_id);
229 self.current_ordinals.push((ordinal as u32, score));
230 }
231
232 fn finish_current(&mut self) {
233 let Some(doc_id) = self.current_doc.take() else {
234 return;
235 };
236 let score = self.combiner.combine(&self.current_ordinals);
237 let should_retain = self.heap.len() < self.limit
238 || self.heap.peek().is_some_and(|worst| {
239 HeapVectorResult(VectorSearchResult::new(doc_id, score, Vec::new()))
240 .cmp(worst)
241 .is_lt()
242 });
243
244 if !should_retain {
245 self.current_ordinals.clear();
249 return;
250 }
251
252 let ordinals = std::mem::take(&mut self.current_ordinals);
253 let entry = HeapVectorResult(VectorSearchResult::new(doc_id, score, ordinals));
254 if self.heap.len() < self.limit {
255 self.heap.push(entry);
256 } else if let Some(mut worst) = self.heap.peek_mut() {
257 let mut evicted = std::mem::replace(&mut worst.0, entry.0);
260 evicted.ordinals.clear();
261 self.current_ordinals = evicted.ordinals;
262 }
263 }
264
265 fn into_results(mut self) -> Vec<VectorSearchResult> {
266 self.finish_current();
267 let mut results: Vec<_> = self.heap.into_iter().map(|entry| entry.0).collect();
268 results.sort_unstable_by(|a, b| {
269 b.score
270 .total_cmp(&a.score)
271 .then_with(|| a.doc_id.cmp(&b.doc_id))
272 });
273 results
274 }
275}
276
277fn combine_grouped_ordinal_results(
280 raw: impl IntoIterator<Item = RawVectorCandidate>,
281 combiner: crate::query::MultiValueCombiner,
282 limit: usize,
283) -> Vec<VectorSearchResult> {
284 let mut collector = FlatDocumentCollector::new(limit, combiner);
285 for (doc_id, ordinal, score) in raw {
286 collector.push(doc_id, ordinal, score);
287 }
288 collector.into_results()
289}
290
291#[derive(Clone, Copy)]
292struct DenseSearchParams {
293 dim: usize,
294 nprobe: usize,
295 unit_norm: bool,
296}
297
298struct PreparedDenseScoreQuery<'a> {
305 query: &'a [f32],
306 query_f16: Vec<u16>,
307 inv_norm_q: f32,
308 quantization: DenseVectorQuantization,
309 dim: usize,
310 unit_norm: bool,
311}
312
313impl<'a> PreparedDenseScoreQuery<'a> {
314 fn new(
315 query: &'a [f32],
316 quantization: DenseVectorQuantization,
317 dim: usize,
318 unit_norm: bool,
319 ) -> Result<Self> {
320 use crate::structures::simd;
321
322 if query.len() != dim {
323 return Err(Error::Query(format!(
324 "dense SIMD query dimension {} does not match vector dimension {dim}",
325 query.len()
326 )));
327 }
328 if quantization == DenseVectorQuantization::Binary {
329 return Err(Error::InvalidFieldType {
330 expected: "non-binary dense vector".to_string(),
331 got: "binary dense vector".to_string(),
332 });
333 }
334
335 let norm_q_sq = simd::dot_product_f32(query, query, dim);
336 let inv_norm_q = if norm_q_sq < f32::EPSILON {
337 0.0
338 } else {
339 simd::fast_inv_sqrt(norm_q_sq)
340 };
341 let query_f16 = if quantization == DenseVectorQuantization::F16 {
342 query.iter().map(|&value| simd::f32_to_f16(value)).collect()
343 } else {
344 Vec::new()
345 };
346
347 Ok(Self {
348 query,
349 query_f16,
350 inv_norm_q,
351 quantization,
352 dim,
353 unit_norm,
354 })
355 }
356
357 fn score_batch(&self, raw: &[u8], scores: &mut [f32]) -> Result<()> {
358 use crate::structures::simd;
359
360 let element_size = match self.quantization {
361 DenseVectorQuantization::F32 => std::mem::size_of::<f32>(),
362 DenseVectorQuantization::F16 => std::mem::size_of::<u16>(),
363 DenseVectorQuantization::UInt8 => 1,
364 DenseVectorQuantization::Binary => {
365 return Err(Error::InvalidFieldType {
366 expected: "non-binary dense vector".to_string(),
367 got: "binary dense vector".to_string(),
368 });
369 }
370 };
371 let required_bytes = scores
372 .len()
373 .checked_mul(self.dim)
374 .and_then(|elements| elements.checked_mul(element_size))
375 .ok_or_else(|| Error::Corruption("dense vector batch byte length overflow".into()))?;
376 if raw.len() < required_bytes {
377 return Err(Error::Corruption(format!(
378 "dense vector batch is truncated: need {required_bytes} bytes, got {}",
379 raw.len()
380 )));
381 }
382 if self.quantization == DenseVectorQuantization::F16
383 && required_bytes > 0
384 && !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>())
385 {
386 return Err(Error::Corruption(
387 "f16 vector data is not 2-byte aligned".to_string(),
388 ));
389 }
390 if self.quantization == DenseVectorQuantization::F32
391 && !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>())
392 {
393 return Err(Error::Corruption(
394 "f32 vector data is not 4-byte aligned".to_string(),
395 ));
396 }
397
398 if self.dim == 0 || scores.is_empty() {
402 return Ok(());
403 }
404
405 match (self.quantization, self.unit_norm) {
406 (DenseVectorQuantization::F32, false) => {
407 let num_floats = scores.len() * self.dim;
408 let vectors: &[f32] =
409 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
410 simd::batch_cosine_scores_precomp(
411 self.query,
412 vectors,
413 self.dim,
414 scores,
415 self.inv_norm_q,
416 );
417 }
418 (DenseVectorQuantization::F32, true) => {
419 let num_floats = scores.len() * self.dim;
420 let vectors: &[f32] =
421 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
422 simd::batch_dot_scores_precomp(
423 self.query,
424 vectors,
425 self.dim,
426 scores,
427 self.inv_norm_q,
428 );
429 }
430 (DenseVectorQuantization::F16, false) => {
431 simd::batch_cosine_scores_f16_precomp(
432 &self.query_f16,
433 raw,
434 self.dim,
435 scores,
436 self.inv_norm_q,
437 );
438 }
439 (DenseVectorQuantization::F16, true) => {
440 simd::batch_dot_scores_f16_precomp(
441 &self.query_f16,
442 raw,
443 self.dim,
444 scores,
445 self.inv_norm_q,
446 );
447 }
448 (DenseVectorQuantization::UInt8, false) => {
449 simd::batch_cosine_scores_u8_precomp(
450 self.query,
451 raw,
452 self.dim,
453 scores,
454 self.inv_norm_q,
455 );
456 }
457 (DenseVectorQuantization::UInt8, true) => {
458 simd::batch_dot_scores_u8_precomp(
459 self.query,
460 raw,
461 self.dim,
462 scores,
463 self.inv_norm_q,
464 );
465 }
466 (DenseVectorQuantization::Binary, _) => unreachable!("validated during preparation"),
467 }
468 Ok(())
469 }
470}
471
472fn checked_dense_fetch_k(k: usize, rerank_factor: f32) -> Result<usize> {
474 if !rerank_factor.is_finite() || !(1.0..=MAX_DENSE_RERANK_FACTOR).contains(&rerank_factor) {
475 return Err(Error::Query(format!(
476 "dense rerank_factor must be finite and in [1, {MAX_DENSE_RERANK_FACTOR}], got {rerank_factor}"
477 )));
478 }
479
480 let fetch = (k as f64) * (rerank_factor as f64);
481 if !fetch.is_finite()
482 || fetch > usize::MAX as f64
483 || fetch > MAX_DENSE_CANDIDATES_PER_SEGMENT as f64
484 {
485 return Err(Error::Query(format!(
486 "dense candidate count exceeds the per-segment maximum of \
487 {MAX_DENSE_CANDIDATES_PER_SEGMENT}: k={k}, rerank_factor={rerank_factor}"
488 )));
489 }
490 Ok(fetch.ceil() as usize)
491}
492
493#[inline]
499fn checked_binary_combined_fetch_k(k: usize) -> Result<usize> {
500 if k > MAX_DENSE_CANDIDATES_PER_SEGMENT {
501 return Err(Error::Query(format!(
502 "binary dense result count exceeds the per-segment maximum of \
503 {MAX_DENSE_CANDIDATES_PER_SEGMENT}: k={k}"
504 )));
505 }
506 Ok(crate::query::max_candidate_limit(k).min(MAX_DENSE_CANDIDATES_PER_SEGMENT))
507}
508
509#[inline]
510fn bounded_vector_score_batch(vector_byte_size: usize, preferred: usize) -> usize {
511 preferred.min((MAX_VECTOR_SCORE_BATCH_BYTES / vector_byte_size.max(1)).max(1))
512}
513
514#[inline]
515fn bounded_rerank_batch(vector_byte_size: usize, preferred: usize, vector_count: usize) -> usize {
516 bounded_vector_score_batch(vector_byte_size, preferred).min(vector_count.max(1))
517}
518
519fn checked_file_range(
520 offset: u64,
521 length: u64,
522 file_length: u64,
523 description: &str,
524) -> Result<std::ops::Range<u64>> {
525 let end = offset
526 .checked_add(length)
527 .ok_or_else(|| Error::Corruption(format!("{description} byte range overflows u64")))?;
528 if end > file_length {
529 return Err(Error::Corruption(format!(
530 "{description} byte range {offset}..{end} exceeds file length {file_length}"
531 )));
532 }
533 Ok(offset..end)
534}
535
536type RawVectorCandidate = (u32, u16, f32);
537type CandidateVectorRef = (DocId, u16, usize); #[derive(Clone, Copy)]
540struct CandidateDocumentRange {
541 doc_id: DocId,
542 start: usize,
543 end: usize,
544}
545
546struct AnnCandidateDocuments {
547 ranges: Vec<CandidateDocumentRange>,
548 vector_count: usize,
549}
550
551fn ann_candidate_document_ranges(
560 ann_results: &[RawVectorCandidate],
561 flat: &LazyFlatVectorData,
562) -> Result<AnnCandidateDocuments> {
563 ann_candidate_document_ranges_from_ids(ann_results.iter().map(|candidate| candidate.0), flat)
564}
565
566fn ann_candidate_document_ranges_from_ids(
567 doc_ids: impl IntoIterator<Item = DocId>,
568 flat: &LazyFlatVectorData,
569) -> Result<AnnCandidateDocuments> {
570 let mut candidate_docs: Vec<DocId> = doc_ids.into_iter().collect();
571 candidate_docs.sort_unstable();
572 candidate_docs.dedup();
573
574 let mut ranges = Vec::with_capacity(candidate_docs.len());
575 let mut vector_count = 0usize;
576 for doc_id in candidate_docs {
577 let (start, count) = flat.flat_indexes_for_doc_range(doc_id);
578 if count == 0 {
579 return Err(Error::Corruption(format!(
580 "ANN candidate document {doc_id} is missing from flat vector storage"
581 )));
582 }
583 vector_count = vector_count
584 .checked_add(count)
585 .ok_or_else(|| Error::Query("ANN candidate vector expansion overflow".to_string()))?;
586 let end = start
587 .checked_add(count)
588 .ok_or_else(|| Error::Corruption("flat vector range overflow".to_string()))?;
589 if end > flat.num_vectors {
590 return Err(Error::Corruption(format!(
591 "flat vector range {start}..{end} for document {doc_id} exceeds {} vectors",
592 flat.num_vectors
593 )));
594 }
595 ranges.push(CandidateDocumentRange { doc_id, start, end });
596 }
597 Ok(AnnCandidateDocuments {
598 ranges,
599 vector_count,
600 })
601}
602
603fn validate_binary_single_value_ann_results(
611 ann_results: Vec<RawVectorCandidate>,
612 flat: &LazyFlatVectorData,
613) -> Result<Vec<RawVectorCandidate>> {
614 let mut seen_docs = FxHashSet::default();
615 let mut validated = Vec::with_capacity(ann_results.len());
616 for (doc_id, ordinal, score) in ann_results {
617 let (start, count) = flat.flat_indexes_for_doc_range(doc_id);
618 if count == 0 {
619 return Err(Error::Corruption(format!(
620 "ANN candidate document {doc_id} is missing from flat vector storage"
621 )));
622 }
623 if count != 1 {
624 return Err(Error::Corruption(format!(
625 "binary ANN single-valued candidate document {doc_id} has {count} flat vectors"
626 )));
627 }
628 let (stored_doc_id, stored_ordinal) = flat.get_doc_id(start);
629 if stored_doc_id != doc_id {
630 return Err(Error::Corruption(format!(
631 "flat vector doc map is not contiguous for document {doc_id}"
632 )));
633 }
634 if stored_ordinal != ordinal {
635 return Err(Error::Corruption(format!(
636 "binary ANN candidate document {doc_id} ordinal {ordinal} is missing from flat vector storage"
637 )));
638 }
639 if seen_docs.insert(doc_id) {
640 validated.push((doc_id, ordinal, score));
641 }
642 }
643 Ok(validated)
644}
645
646struct CandidateVectorCursor<'a> {
647 ranges: &'a [CandidateDocumentRange],
648 range_index: usize,
649 flat_index: usize,
650}
651
652impl<'a> CandidateVectorCursor<'a> {
653 fn new(ranges: &'a [CandidateDocumentRange]) -> Self {
654 Self {
655 ranges,
656 range_index: 0,
657 flat_index: ranges.first().map_or(0, |range| range.start),
658 }
659 }
660
661 fn fill_batch(
665 &mut self,
666 flat: &LazyFlatVectorData,
667 batch: &mut Vec<CandidateVectorRef>,
668 limit: usize,
669 ) -> Result<bool> {
670 batch.clear();
671 while batch.len() < limit && self.range_index < self.ranges.len() {
672 let range = self.ranges[self.range_index];
673 if self.flat_index == range.end {
674 self.range_index += 1;
675 if let Some(next) = self.ranges.get(self.range_index) {
676 self.flat_index = next.start;
677 }
678 continue;
679 }
680 let (stored_doc_id, ordinal) = flat.get_doc_id(self.flat_index);
681 if stored_doc_id != range.doc_id {
682 return Err(Error::Corruption(format!(
683 "flat vector doc map is not contiguous for document {}",
684 range.doc_id
685 )));
686 }
687 batch.push((range.doc_id, ordinal, self.flat_index));
688 self.flat_index += 1;
689 }
690 Ok(!batch.is_empty())
691 }
692}
693
694#[derive(Clone, Copy)]
695struct VectorReadRun {
696 buffer_start: usize,
697 flat_start: usize,
698 count: usize,
699}
700
701fn plan_vector_read_runs(indexes: &[usize], runs: &mut Vec<VectorReadRun>) -> Result<()> {
706 runs.clear();
707 for (buffer_index, &flat_index) in indexes.iter().enumerate() {
708 if let Some(run) = runs.last_mut()
709 && run
710 .flat_start
711 .checked_add(run.count)
712 .is_some_and(|next| next == flat_index)
713 {
714 run.count += 1;
715 continue;
716 }
717 if buffer_index > 0 && flat_index <= indexes[buffer_index - 1] {
718 return Err(Error::Corruption(
719 "candidate flat-vector indexes are not strictly ordered".into(),
720 ));
721 }
722 runs.push(VectorReadRun {
723 buffer_start: buffer_index,
724 flat_start: flat_index,
725 count: 1,
726 });
727 }
728 Ok(())
729}
730
731fn prepare_vector_read_runs(
735 flat: &LazyFlatVectorData,
736 indexes: &[usize],
737 runs: &mut Vec<VectorReadRun>,
738) -> Result<()> {
739 plan_vector_read_runs(indexes, runs)?;
740 #[cfg(feature = "native")]
741 flat.prefetch_vectors(indexes.iter().copied());
742 #[cfg(not(feature = "native"))]
743 let _ = flat;
744 Ok(())
745}
746
747async fn read_vector_runs(
748 flat: &LazyFlatVectorData,
749 indexes: &[usize],
750 runs: &mut Vec<VectorReadRun>,
751 output: &mut [u8],
752) -> Result<()> {
753 prepare_vector_read_runs(flat, indexes, runs)?;
754 let vector_byte_size = flat.vector_byte_size();
755 for run in runs {
756 let bytes = flat
757 .read_vectors_batch(run.flat_start, run.count)
758 .await
759 .map_err(Error::Io)?;
760 let start = run
761 .buffer_start
762 .checked_mul(vector_byte_size)
763 .ok_or_else(|| Error::Query("dense rerank buffer offset overflow".into()))?;
764 let end = start
765 .checked_add(bytes.len())
766 .ok_or_else(|| Error::Query("dense rerank buffer range overflow".into()))?;
767 let destination = output
768 .get_mut(start..end)
769 .ok_or_else(|| Error::Corruption("dense rerank buffer is too short".into()))?;
770 destination.copy_from_slice(bytes.as_slice());
771 }
772 Ok(())
773}
774
775#[cfg(feature = "sync")]
776fn read_vector_runs_sync(
777 flat: &LazyFlatVectorData,
778 indexes: &[usize],
779 runs: &mut Vec<VectorReadRun>,
780 output: &mut [u8],
781) -> Result<()> {
782 prepare_vector_read_runs(flat, indexes, runs)?;
783 let vector_byte_size = flat.vector_byte_size();
784 for run in runs {
785 let bytes = flat
786 .read_vectors_batch_sync(run.flat_start, run.count)
787 .map_err(Error::Io)?;
788 let start = run
789 .buffer_start
790 .checked_mul(vector_byte_size)
791 .ok_or_else(|| Error::Query("dense rerank buffer offset overflow".into()))?;
792 let end = start
793 .checked_add(bytes.len())
794 .ok_or_else(|| Error::Query("dense rerank buffer range overflow".into()))?;
795 let destination = output
796 .get_mut(start..end)
797 .ok_or_else(|| Error::Corruption("dense rerank buffer is too short".into()))?;
798 destination.copy_from_slice(bytes.as_slice());
799 }
800 Ok(())
801}
802
803#[derive(Default)]
804struct DenseRerankStats {
805 vector_count: usize,
806 resolve_elapsed: std::time::Duration,
807 read_elapsed: std::time::Duration,
808 score_elapsed: std::time::Duration,
809}
810
811async fn exact_score_dense_candidate_documents(
812 ann_results: &[RawVectorCandidate],
813 flat: &LazyFlatVectorData,
814 query: &[f32],
815 unit_norm: bool,
816 combiner: crate::query::MultiValueCombiner,
817 limit: usize,
818) -> Result<(Vec<VectorSearchResult>, DenseRerankStats)> {
819 let resolve_started = std::time::Instant::now();
820 let documents = ann_candidate_document_ranges(ann_results, flat)?;
821 let mut stats = DenseRerankStats {
822 vector_count: documents.vector_count,
823 resolve_elapsed: resolve_started.elapsed(),
824 ..Default::default()
825 };
826 let vector_byte_size = flat.vector_byte_size();
827 let batch_len =
828 bounded_rerank_batch(vector_byte_size, DENSE_SCORE_BATCH, documents.vector_count);
829 let raw_capacity = batch_len
830 .checked_mul(vector_byte_size)
831 .ok_or_else(|| Error::Query("dense rerank buffer size overflow".to_string()))?;
832 let mut raw = vec![0u8; raw_capacity];
833 let mut scores = vec![0.0f32; batch_len];
834 let mut batch = Vec::with_capacity(batch_len);
835 let mut flat_indexes = Vec::with_capacity(batch_len);
836 let mut read_runs = Vec::new();
837 let mut cursor = CandidateVectorCursor::new(&documents.ranges);
838 let mut collector = FlatDocumentCollector::new(limit, combiner);
839 let mut scored = 0usize;
840 let prepared_query =
841 PreparedDenseScoreQuery::new(query, flat.quantization, flat.dim, unit_norm)?;
842
843 while cursor.fill_batch(flat, &mut batch, batch_len)? {
844 flat_indexes.clear();
845 flat_indexes.extend(batch.iter().map(|&(_, _, flat_index)| flat_index));
846 let raw_len = batch
847 .len()
848 .checked_mul(vector_byte_size)
849 .ok_or_else(|| Error::Query("dense rerank buffer size overflow".to_string()))?;
850 let raw = &mut raw[..raw_len];
851
852 let read_started = std::time::Instant::now();
853 read_vector_runs(flat, &flat_indexes, &mut read_runs, raw).await?;
854 stats.read_elapsed += read_started.elapsed();
855
856 let score_started = std::time::Instant::now();
857 prepared_query.score_batch(raw, &mut scores[..batch.len()])?;
858 stats.score_elapsed += score_started.elapsed();
859 for (buffer_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
860 collector.push(doc_id, ordinal, scores[buffer_index]);
861 }
862 scored += batch.len();
863 }
864 debug_assert_eq!(scored, documents.vector_count);
865 Ok((collector.into_results(), stats))
866}
867
868#[cfg(feature = "sync")]
869fn exact_score_dense_candidate_documents_sync(
870 ann_results: &[RawVectorCandidate],
871 flat: &LazyFlatVectorData,
872 query: &[f32],
873 unit_norm: bool,
874 combiner: crate::query::MultiValueCombiner,
875 limit: usize,
876) -> Result<Vec<VectorSearchResult>> {
877 let documents = ann_candidate_document_ranges(ann_results, flat)?;
878 let vector_byte_size = flat.vector_byte_size();
879 let batch_len =
880 bounded_rerank_batch(vector_byte_size, DENSE_SCORE_BATCH, documents.vector_count);
881 let raw_capacity = batch_len
882 .checked_mul(vector_byte_size)
883 .ok_or_else(|| Error::Query("dense rerank buffer size overflow".to_string()))?;
884 let mut raw = vec![0u8; raw_capacity];
885 let mut scores = vec![0.0f32; batch_len];
886 let mut batch = Vec::with_capacity(batch_len);
887 let mut flat_indexes = Vec::with_capacity(batch_len);
888 let mut read_runs = Vec::new();
889 let mut cursor = CandidateVectorCursor::new(&documents.ranges);
890 let mut collector = FlatDocumentCollector::new(limit, combiner);
891 let mut scored = 0usize;
892 let prepared_query =
893 PreparedDenseScoreQuery::new(query, flat.quantization, flat.dim, unit_norm)?;
894
895 while cursor.fill_batch(flat, &mut batch, batch_len)? {
896 flat_indexes.clear();
897 flat_indexes.extend(batch.iter().map(|&(_, _, flat_index)| flat_index));
898 let raw_len = batch
899 .len()
900 .checked_mul(vector_byte_size)
901 .ok_or_else(|| Error::Query("dense rerank buffer size overflow".to_string()))?;
902 let raw = &mut raw[..raw_len];
903 read_vector_runs_sync(flat, &flat_indexes, &mut read_runs, raw)?;
904 prepared_query.score_batch(raw, &mut scores[..batch.len()])?;
905 for (buffer_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
906 collector.push(doc_id, ordinal, scores[buffer_index]);
907 }
908 scored += batch.len();
909 }
910 debug_assert_eq!(scored, documents.vector_count);
911 Ok(collector.into_results())
912}
913
914async fn exact_score_binary_candidate_documents(
915 ann_results: &[RawVectorCandidate],
916 flat: &LazyFlatVectorData,
917 query: &[u8],
918 dim_bits: usize,
919 combiner: crate::query::MultiValueCombiner,
920 limit: usize,
921) -> Result<Vec<VectorSearchResult>> {
922 let documents = ann_candidate_document_ranges(ann_results, flat)?;
923 let probe_scores: FxHashMap<(DocId, u16), f32> = ann_results
924 .iter()
925 .map(|&(doc_id, ordinal, score)| ((doc_id, ordinal), score))
926 .collect();
927 exact_score_binary_resolved_documents(
928 documents,
929 &probe_scores,
930 flat,
931 query,
932 dim_bits,
933 combiner,
934 limit,
935 )
936 .await
937}
938
939async fn exact_score_binary_candidate_document_ids(
940 candidate_doc_ids: Vec<DocId>,
941 probed_ordinal_scores: &[(u32, u16, f32)],
942 flat: &LazyFlatVectorData,
943 query: &[u8],
944 dim_bits: usize,
945 combiner: crate::query::MultiValueCombiner,
946 limit: usize,
947) -> Result<Vec<VectorSearchResult>> {
948 let documents = ann_candidate_document_ranges_from_ids(candidate_doc_ids, flat)?;
949 let probe_scores = binary_probe_score_map(probed_ordinal_scores);
952 exact_score_binary_resolved_documents(
953 documents,
954 &probe_scores,
955 flat,
956 query,
957 dim_bits,
958 combiner,
959 limit,
960 )
961 .await
962}
963
964fn binary_probe_score_map(
965 probed_ordinal_scores: &[(u32, u16, f32)],
966) -> FxHashMap<(DocId, u16), f32> {
967 probed_ordinal_scores
968 .iter()
969 .map(|&(doc_id, ordinal, score)| ((doc_id, ordinal), score))
970 .collect()
971}
972
973async fn exact_score_binary_resolved_documents(
974 documents: AnnCandidateDocuments,
975 probe_scores: &FxHashMap<(DocId, u16), f32>,
976 flat: &LazyFlatVectorData,
977 query: &[u8],
978 dim_bits: usize,
979 combiner: crate::query::MultiValueCombiner,
980 limit: usize,
981) -> Result<Vec<VectorSearchResult>> {
982 let vector_byte_size = flat.vector_byte_size();
983 let batch_len =
984 bounded_rerank_batch(vector_byte_size, BINARY_SCORE_BATCH, documents.vector_count);
985 let raw_capacity = batch_len
986 .checked_mul(vector_byte_size)
987 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
988 let mut raw = vec![0u8; raw_capacity];
989 let mut scores = vec![0.0f32; batch_len];
990 let mut batch_scores = vec![0.0f32; batch_len];
991 let mut batch = Vec::with_capacity(batch_len);
992 let mut unresolved = Vec::with_capacity(batch_len);
993 let mut unresolved_flat_indexes = Vec::with_capacity(batch_len);
994 let mut read_runs = Vec::new();
995 let mut cursor = CandidateVectorCursor::new(&documents.ranges);
996 let mut collector = FlatDocumentCollector::new(limit, combiner);
997 let mut scored = 0usize;
998
999 while cursor.fill_batch(flat, &mut batch, batch_len)? {
1000 unresolved.clear();
1001 for (batch_index, &(doc_id, ordinal, flat_index)) in batch.iter().enumerate() {
1002 if let Some(&score) = probe_scores.get(&(doc_id, ordinal)) {
1003 batch_scores[batch_index] = score;
1004 } else {
1005 unresolved.push((batch_index, flat_index));
1006 }
1007 }
1008 unresolved_flat_indexes.clear();
1009 unresolved_flat_indexes.extend(unresolved.iter().map(|&(_, flat_index)| flat_index));
1010 let raw_len = unresolved
1011 .len()
1012 .checked_mul(vector_byte_size)
1013 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
1014 let raw = &mut raw[..raw_len];
1015 read_vector_runs(flat, &unresolved_flat_indexes, &mut read_runs, raw).await?;
1016 crate::structures::simd::batch_hamming_scores(
1017 query,
1018 raw,
1019 vector_byte_size,
1020 dim_bits,
1021 &mut scores[..unresolved.len()],
1022 );
1023 for (buffer_index, &(batch_index, _)) in unresolved.iter().enumerate() {
1024 batch_scores[batch_index] = scores[buffer_index];
1025 }
1026 for (batch_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
1027 collector.push(doc_id, ordinal, batch_scores[batch_index]);
1028 }
1029 scored += batch.len();
1030 }
1031 debug_assert_eq!(scored, documents.vector_count);
1032 Ok(collector.into_results())
1033}
1034
1035#[cfg(feature = "sync")]
1036fn exact_score_binary_candidate_documents_sync(
1037 ann_results: &[RawVectorCandidate],
1038 flat: &LazyFlatVectorData,
1039 query: &[u8],
1040 dim_bits: usize,
1041 combiner: crate::query::MultiValueCombiner,
1042 limit: usize,
1043) -> Result<Vec<VectorSearchResult>> {
1044 let documents = ann_candidate_document_ranges(ann_results, flat)?;
1045 let probe_scores: FxHashMap<(DocId, u16), f32> = ann_results
1046 .iter()
1047 .map(|&(doc_id, ordinal, score)| ((doc_id, ordinal), score))
1048 .collect();
1049 exact_score_binary_resolved_documents_sync(
1050 documents,
1051 &probe_scores,
1052 flat,
1053 query,
1054 dim_bits,
1055 combiner,
1056 limit,
1057 )
1058}
1059
1060#[cfg(feature = "sync")]
1061fn exact_score_binary_candidate_document_ids_sync(
1062 candidate_doc_ids: Vec<DocId>,
1063 probed_ordinal_scores: &[(u32, u16, f32)],
1064 flat: &LazyFlatVectorData,
1065 query: &[u8],
1066 dim_bits: usize,
1067 combiner: crate::query::MultiValueCombiner,
1068 limit: usize,
1069) -> Result<Vec<VectorSearchResult>> {
1070 let documents = ann_candidate_document_ranges_from_ids(candidate_doc_ids, flat)?;
1071 let probe_scores = binary_probe_score_map(probed_ordinal_scores);
1072 exact_score_binary_resolved_documents_sync(
1073 documents,
1074 &probe_scores,
1075 flat,
1076 query,
1077 dim_bits,
1078 combiner,
1079 limit,
1080 )
1081}
1082
1083#[cfg(feature = "sync")]
1084fn exact_score_binary_resolved_documents_sync(
1085 documents: AnnCandidateDocuments,
1086 probe_scores: &FxHashMap<(DocId, u16), f32>,
1087 flat: &LazyFlatVectorData,
1088 query: &[u8],
1089 dim_bits: usize,
1090 combiner: crate::query::MultiValueCombiner,
1091 limit: usize,
1092) -> Result<Vec<VectorSearchResult>> {
1093 let vector_byte_size = flat.vector_byte_size();
1094 let batch_len =
1095 bounded_rerank_batch(vector_byte_size, BINARY_SCORE_BATCH, documents.vector_count);
1096 let raw_capacity = batch_len
1097 .checked_mul(vector_byte_size)
1098 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
1099 let mut raw = vec![0u8; raw_capacity];
1100 let mut scores = vec![0.0f32; batch_len];
1101 let mut batch_scores = vec![0.0f32; batch_len];
1102 let mut batch = Vec::with_capacity(batch_len);
1103 let mut unresolved = Vec::with_capacity(batch_len);
1104 let mut unresolved_flat_indexes = Vec::with_capacity(batch_len);
1105 let mut read_runs = Vec::new();
1106 let mut cursor = CandidateVectorCursor::new(&documents.ranges);
1107 let mut collector = FlatDocumentCollector::new(limit, combiner);
1108 let mut scored = 0usize;
1109
1110 while cursor.fill_batch(flat, &mut batch, batch_len)? {
1111 unresolved.clear();
1112 for (batch_index, &(doc_id, ordinal, flat_index)) in batch.iter().enumerate() {
1113 if let Some(&score) = probe_scores.get(&(doc_id, ordinal)) {
1114 batch_scores[batch_index] = score;
1115 } else {
1116 unresolved.push((batch_index, flat_index));
1117 }
1118 }
1119 unresolved_flat_indexes.clear();
1120 unresolved_flat_indexes.extend(unresolved.iter().map(|&(_, flat_index)| flat_index));
1121 let raw_len = unresolved
1122 .len()
1123 .checked_mul(vector_byte_size)
1124 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
1125 let raw = &mut raw[..raw_len];
1126 read_vector_runs_sync(flat, &unresolved_flat_indexes, &mut read_runs, raw)?;
1127 crate::structures::simd::batch_hamming_scores(
1128 query,
1129 raw,
1130 vector_byte_size,
1131 dim_bits,
1132 &mut scores[..unresolved.len()],
1133 );
1134 for (buffer_index, &(batch_index, _)) in unresolved.iter().enumerate() {
1135 batch_scores[batch_index] = scores[buffer_index];
1136 }
1137 for (batch_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
1138 collector.push(doc_id, ordinal, batch_scores[batch_index]);
1139 }
1140 scored += batch.len();
1141 }
1142 debug_assert_eq!(scored, documents.vector_count);
1143 Ok(collector.into_results())
1144}
1145
1146fn validate_coarse_centroids(centroids: &CoarseCentroids, dim: usize) -> Result<()> {
1147 let expected = (centroids.num_clusters as usize)
1148 .checked_mul(dim)
1149 .ok_or_else(|| Error::Corruption("coarse centroid size overflow".into()))?;
1150 if centroids.num_clusters == 0
1151 || centroids.dim != dim
1152 || centroids.centroids.len() != expected
1153 || centroids.centroids.iter().any(|value| !value.is_finite())
1154 {
1155 return Err(Error::Corruption(format!(
1156 "invalid coarse centroids: clusters={}, dim={}, values={} (expected dim={dim}, values={expected})",
1157 centroids.num_clusters,
1158 centroids.dim,
1159 centroids.centroids.len()
1160 )));
1161 }
1162 Ok(())
1163}
1164
1165#[derive(Debug, Default)]
1170pub struct DensePlanCache {
1171 pub(crate) tq: std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqQueryPlan>>>,
1172 pub(crate) ivf_tq: std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqIvfQueryPlan>>>,
1173 scann: std::sync::Mutex<Option<ScannPlanCacheEntry>>,
1174}
1175
1176#[derive(Debug)]
1177struct ScannPlanCacheEntry {
1178 artifact_id: u64,
1179 probes: usize,
1180 query_bits: Vec<u32>,
1181 plan: std::sync::Arc<crate::structures::vector::scann::FloatScannQuery>,
1182}
1183
1184#[allow(clippy::too_many_arguments)]
1189fn search_tq_segment(
1190 index: &crate::segment::ann_disk::AnnDiskIndex,
1191 codec: &crate::structures::TqCodec,
1192 query: &[f32],
1193 fetch_k: usize,
1194 document_combiner: Option<crate::query::MultiValueCombiner>,
1195 field: Field,
1196 dim: usize,
1197 plan_cache: Option<&std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqQueryPlan>>>>,
1198) -> Result<Vec<RawVectorCandidate>> {
1199 validate_tq_ann(index, codec, dim, field)?;
1200 let plan = cached_tq_query_plan(codec, query, plan_cache)?;
1201 match document_combiner {
1202 Some(combiner) => index
1203 .search_tq_combined_documents(fetch_k, &plan, combiner)
1204 .map(|candidates| {
1205 candidates
1206 .into_iter()
1207 .map(|candidate| (candidate.doc_id, 0, 0.0))
1211 .collect()
1212 }),
1213 None => index.search_tq_distinct(fetch_k, &plan),
1214 }
1215 .map_err(|error| {
1216 Error::Corruption(format!("invalid TQ payload for field {}: {error}", field.0))
1217 })
1218}
1219
1220fn cached_tq_query_plan(
1223 codec: &crate::structures::TqCodec,
1224 query: &[f32],
1225 plan_cache: Option<&std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqQueryPlan>>>>,
1226) -> Result<std::sync::Arc<crate::structures::TqQueryPlan>> {
1227 Ok(match plan_cache {
1228 Some(cache) => {
1229 let mut cached = cache
1230 .lock()
1231 .map_err(|_| Error::Internal("TQ plan cache is poisoned".into()))?;
1232 match cached.as_ref() {
1233 Some(plan)
1234 if plan.fingerprint() == codec.fingerprint() && plan.matches_query(query) =>
1235 {
1236 std::sync::Arc::clone(plan)
1237 }
1238 _ => {
1239 let plan =
1240 std::sync::Arc::new(crate::structures::TqQueryPlan::build(codec, query));
1241 *cached = Some(std::sync::Arc::clone(&plan));
1242 plan
1243 }
1244 }
1245 }
1246 None => std::sync::Arc::new(crate::structures::TqQueryPlan::build(codec, query)),
1247 })
1248}
1249
1250fn validate_tq_ann(
1251 index: &crate::segment::ann_disk::AnnDiskIndex,
1252 codec: &crate::structures::TqCodec,
1253 dim: usize,
1254 field: Field,
1255) -> Result<()> {
1256 let header = index.header();
1257 if header.dim != dim
1258 || codec.dim() != dim
1259 || header.code_size != codec.code_size()
1260 || header.quantizer_version != codec.fingerprint()
1261 || header.codebook_version != 0
1262 || header.num_clusters != 1
1263 {
1264 return Err(Error::Corruption(format!(
1265 "TQ payload for field {} does not match the codec derived from schema dimension {dim}",
1266 field.0,
1267 )));
1268 }
1269 Ok(())
1270}
1271
1272#[allow(clippy::too_many_arguments)]
1276fn search_ivf_tq_segment(
1277 index: &crate::segment::ann_disk::AnnDiskIndex,
1278 centroids: &CoarseCentroids,
1279 codec: &crate::structures::TqCodec,
1280 query: &[f32],
1281 fetch_k: usize,
1282 document_combiner: Option<crate::query::MultiValueCombiner>,
1283 field: Field,
1284 nprobe: usize,
1285 routing: crate::dsl::IvfRoutingMode,
1286 plan_cache: Option<
1287 &std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqIvfQueryPlan>>>,
1288 >,
1289) -> Result<Vec<RawVectorCandidate>> {
1290 let effective_nprobe = nprobe.clamp(1, centroids.num_clusters as usize);
1291 let request_fingerprint = crate::structures::TqIvfQueryPlan::request_fingerprint_for(
1292 centroids,
1293 query,
1294 effective_nprobe,
1295 routing,
1296 );
1297 let build = || {
1298 std::sync::Arc::new(crate::structures::TqIvfQueryPlan::build(
1299 centroids,
1300 codec,
1301 query,
1302 effective_nprobe,
1303 routing,
1304 ))
1305 };
1306 let plan = match plan_cache {
1307 Some(cache) => {
1308 let mut cached = cache
1309 .lock()
1310 .map_err(|_| Error::Internal("IVF-TQ plan cache is poisoned".into()))?;
1311 match cached.as_ref() {
1312 Some(plan)
1313 if plan.quantizer_version == centroids.version
1314 && plan.fingerprint == codec.fingerprint()
1315 && plan.request_fingerprint == request_fingerprint
1316 && plan.cluster_ids.len() == effective_nprobe =>
1317 {
1318 std::sync::Arc::clone(plan)
1319 }
1320 _ => {
1321 let plan = build();
1322 *cached = Some(std::sync::Arc::clone(&plan));
1323 plan
1324 }
1325 }
1326 }
1327 None => build(),
1328 };
1329 let candidates = match document_combiner {
1330 Some(combiner) => index
1331 .search_ivf_tq_combined_documents(fetch_k, &plan, combiner)
1332 .map(|documents| {
1333 documents
1334 .into_iter()
1335 .map(|candidate| (candidate.doc_id, 0, 0.0))
1339 .collect()
1340 }),
1341 None => index.search_ivf_tq_distinct(fetch_k, &plan),
1342 };
1343 candidates.map_err(|error| {
1344 Error::Corruption(format!(
1345 "invalid IVF-TQ payload for field {}: {error}",
1346 field.0
1347 ))
1348 })
1349}
1350
1351#[allow(clippy::too_many_arguments)]
1352fn search_scann_ah_segment(
1353 index: &crate::segment::ann_disk::AnnDiskIndex,
1354 artifact: &crate::segment::ScannTrainedArtifactBytes,
1355 query: &[f32],
1356 fetch_k: usize,
1357 combiner: crate::query::MultiValueCombiner,
1358 field: Field,
1359 nprobe: usize,
1360 plan_cache: Option<&std::sync::Mutex<Option<ScannPlanCacheEntry>>>,
1361) -> Result<Vec<RawVectorCandidate>> {
1362 index
1363 .validate_scann_generation(
1364 artifact.config(),
1365 artifact.generation(),
1366 artifact.artifact_id(),
1367 )
1368 .map_err(|error| {
1369 Error::Corruption(format!(
1370 "ScaNN generation mismatch for field {}: {error}",
1371 field.0
1372 ))
1373 })?;
1374 let probes = nprobe.clamp(1, artifact.config().num_leaves as usize);
1375 let query_bits = query
1376 .iter()
1377 .map(|value| value.to_bits())
1378 .collect::<Vec<_>>();
1379 let build = || {
1380 let mut normalized_query = query.to_vec();
1381 crate::structures::vector::ivf::routing::normalize_cosine_in_place(&mut normalized_query);
1382 artifact
1383 .float_model()
1384 .map_err(Error::Io)?
1385 .prepare_query(&normalized_query, probes)
1386 .map(std::sync::Arc::new)
1387 .map_err(|error| Error::Query(format!("invalid ScaNN query: {error}")))
1388 };
1389 let plan = match plan_cache {
1390 Some(cache) => {
1391 let mut cached = cache
1392 .lock()
1393 .map_err(|_| Error::Internal("ScaNN plan cache is poisoned".into()))?;
1394 match cached.as_ref() {
1395 Some(entry)
1396 if entry.artifact_id == artifact.artifact_id()
1397 && entry.probes == probes
1398 && entry.query_bits == query_bits =>
1399 {
1400 std::sync::Arc::clone(&entry.plan)
1401 }
1402 _ => {
1403 let plan = build()?;
1404 *cached = Some(ScannPlanCacheEntry {
1405 artifact_id: artifact.artifact_id(),
1406 probes,
1407 query_bits,
1408 plan: std::sync::Arc::clone(&plan),
1409 });
1410 plan
1411 }
1412 }
1413 }
1414 None => build()?,
1415 };
1416 index
1417 .search_scann_ah_combined_documents(fetch_k, &plan, combiner)
1418 .map(|documents| {
1419 documents
1420 .into_iter()
1421 .map(|candidate| (candidate.doc_id, 0, 0.0))
1422 .collect()
1423 })
1424 .map_err(|error| {
1425 Error::Corruption(format!(
1426 "invalid ScaNN AH payload for field {}: {error}",
1427 field.0
1428 ))
1429 })
1430}
1431
1432fn validate_ivf_tq_ann(
1433 index: &crate::segment::ann_disk::AnnDiskIndex,
1434 centroids: &CoarseCentroids,
1435 codec: &crate::structures::TqCodec,
1436 dim: usize,
1437 routing: crate::dsl::IvfRoutingMode,
1438 field: Field,
1439) -> Result<()> {
1440 let header = index.header();
1441 if !crate::structures::is_ivf_tq_cosine_generation(centroids.version)
1442 || !crate::structures::is_ivf_tq_cosine_generation(header.quantizer_version)
1443 {
1444 return Err(Error::Corruption(format!(
1445 "IVF-TQ field {} uses a legacy unmarked raw-vector generation that cannot \
1446 preserve cosine candidate semantics; rebuild the index with a current \
1447 Hermes version",
1448 field.0,
1449 )));
1450 }
1451 if header.dim != dim
1452 || codec.dim() != dim
1453 || header.code_size != codec.code_size()
1454 || header.num_clusters != centroids.num_clusters
1455 || header.quantizer_version != centroids.version
1456 || header.codebook_version != codec.fingerprint()
1457 || header.routing != routing
1458 {
1459 return Err(Error::Corruption(format!(
1460 "IVF-TQ payload for field {} does not match its quantizer/codec generation",
1461 field.0,
1462 )));
1463 }
1464 Ok(())
1465}
1466
1467fn validate_binary_ann(
1468 index: &crate::segment::ann_disk::AnnDiskIndex,
1469 quantizer: &crate::structures::BinaryCoarseQuantizer,
1470 config: &crate::dsl::BinaryDenseVectorConfig,
1471 dim: usize,
1472 field: Field,
1473) -> Result<()> {
1474 let header = index.header();
1475 if header.dim != dim
1476 || header.code_size != config.byte_len()
1477 || header.num_clusters != quantizer.num_clusters
1478 || header.quantizer_version != quantizer.version
1479 || header.codebook_version != 0
1480 || header.routing != config.ivf_routing
1481 || quantizer.dim_bits != dim
1482 {
1483 return Err(Error::Corruption(format!(
1484 "binary IVF field {} does not match its quantizer/schema generation",
1485 field.0,
1486 )));
1487 }
1488 Ok(())
1489}
1490
1491fn binary_probe_clusters(
1492 quantizer: &crate::structures::BinaryCoarseQuantizer,
1493 query: &[u8],
1494 nprobe: usize,
1495 routing: crate::dsl::IvfRoutingMode,
1496 cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
1497) -> Result<std::sync::Arc<[u32]>> {
1498 let effective_nprobe = nprobe.clamp(1, quantizer.num_clusters as usize);
1499 let request_fingerprint = quantizer.request_fingerprint(query, effective_nprobe, routing);
1500 if let Some(cache) = cache {
1501 let mut cached = cache
1502 .lock()
1503 .map_err(|_| Error::Internal("binary IVF probe cache is poisoned".into()))?;
1504 if let Some(plan) = cached.as_ref()
1505 && plan.quantizer_version == quantizer.version
1506 && plan.request_fingerprint == request_fingerprint
1507 && plan.cluster_ids.len() == effective_nprobe
1508 {
1509 return Ok(std::sync::Arc::clone(&plan.cluster_ids));
1510 }
1511 let plan = quantizer.probe(query, effective_nprobe, routing);
1512 let clusters = std::sync::Arc::clone(&plan.cluster_ids);
1513 *cached = Some(plan);
1514 return Ok(clusters);
1515 }
1516 Ok(quantizer
1517 .probe(query, effective_nprobe, routing)
1518 .cluster_ids)
1519}
1520
1521fn binary_scann_probe_clusters(
1522 model: &crate::structures::vector::scann::QuantizedBinaryScannModelView<'_>,
1523 query: &[u8],
1524 nprobe: usize,
1525 cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
1526) -> Result<std::sync::Arc<[u32]>> {
1527 let probes = nprobe.clamp(1, model.num_leaves() as usize);
1528 let routing_beam = probes.min(64);
1531 let mut fingerprint = 0xcbf2_9ce4_8422_2325u64;
1532 for byte in query.iter().copied().chain((probes as u64).to_le_bytes()) {
1533 fingerprint ^= u64::from(byte);
1534 fingerprint = fingerprint.wrapping_mul(0x0000_0100_0000_01b3);
1535 }
1536 if let Some(cache) = cache {
1537 let mut cached = cache
1538 .lock()
1539 .map_err(|_| Error::Internal("binary ScaNN probe cache is poisoned".into()))?;
1540 if let Some(plan) = cached.as_ref()
1541 && plan.quantizer_version == model.fingerprint()
1542 && plan.request_fingerprint == fingerprint
1543 && plan.cluster_ids.len() == probes
1544 {
1545 return Ok(std::sync::Arc::clone(&plan.cluster_ids));
1546 }
1547 let mut scratch = crate::structures::vector::scann::BinaryScannSearchScratch::default();
1548 let routed = model
1549 .probe(query, probes, routing_beam, &mut scratch)
1550 .map_err(|error| Error::Query(format!("binary ScaNN routing failed: {error}")))?;
1551 let plan =
1552 crate::structures::IvfProbePlan::new(model.fingerprint(), fingerprint, routed.leaf_ids);
1553 let leaves = std::sync::Arc::clone(&plan.cluster_ids);
1554 *cached = Some(plan);
1555 return Ok(leaves);
1556 }
1557 let mut scratch = crate::structures::vector::scann::BinaryScannSearchScratch::default();
1558 model
1559 .probe(query, probes, routing_beam, &mut scratch)
1560 .map(|plan| plan.leaf_ids.into())
1561 .map_err(|error| Error::Query(format!("binary ScaNN routing failed: {error}")))
1562}
1563
1564pub struct SegmentReader {
1570 meta: SegmentMeta,
1571 term_dict: Arc<AsyncSSTableReader<TermInfo>>,
1573 postings_handle: FileHandle,
1575 store: Arc<AsyncStoreReader>,
1577 schema: Arc<Schema>,
1578 vector_indexes: FxHashMap<u32, VectorIndex>,
1580 flat_vectors: FxHashMap<u32, LazyFlatVectorData>,
1582 dense_file_backed_bytes: u64,
1584 trained_vectors: Arc<crate::segment::TrainedVectorStructures>,
1586 sparse_indexes: FxHashMap<u32, SparseIndex>,
1588 bmp_indexes: FxHashMap<u32, BmpIndex>,
1590 sparse_file_backed_bytes: u64,
1592 positions_handle: Option<FileHandle>,
1594 fast_fields: FxHashMap<u32, crate::structures::fast_field::FastFieldReader>,
1596 #[cfg(feature = "native")]
1598 dense_pin_report: crate::segment::pin::PinReport,
1599 #[cfg(feature = "native")]
1601 sparse_pin_report: crate::segment::pin::PinReport,
1602}
1603
1604impl SegmentReader {
1605 pub async fn open<D: Directory>(
1607 dir: &D,
1608 segment_id: SegmentId,
1609 schema: Arc<Schema>,
1610 term_cache_blocks: usize,
1611 ) -> Result<Self> {
1612 Self::open_with_store_cache(
1613 dir,
1614 segment_id,
1615 schema,
1616 term_cache_blocks,
1617 dir as *const D as usize,
1618 Arc::new(super::SharedStoreCache::new(0)),
1619 )
1620 .await
1621 }
1622
1623 pub(crate) async fn open_with_store_cache<D: Directory>(
1625 dir: &D,
1626 segment_id: SegmentId,
1627 schema: Arc<Schema>,
1628 term_cache_blocks: usize,
1629 store_cache_directory_namespace: usize,
1630 store_cache: Arc<super::SharedStoreCache>,
1631 ) -> Result<Self> {
1632 let files = SegmentFiles::new(segment_id.0);
1633
1634 let meta_slice = dir.open_read(&files.meta).await?;
1636 let meta_bytes = meta_slice.read_bytes().await?;
1637 let meta = SegmentMeta::deserialize(meta_bytes.as_slice())?;
1638 debug_assert_eq!(meta.id, segment_id.0);
1639
1640 let term_dict_handle = dir.open_lazy(&files.term_dict).await?;
1642 let term_dict = AsyncSSTableReader::open(term_dict_handle, term_cache_blocks).await?;
1643
1644 let postings_handle = dir.open_lazy(&files.postings).await?;
1646
1647 let store_handle = dir.open_lazy(&files.store).await?;
1649 let store = AsyncStoreReader::open(
1650 store_handle,
1651 store_cache_directory_namespace,
1652 segment_id.0,
1653 store_cache,
1654 )
1655 .await?;
1656
1657 let vectors_data = loader::load_vectors_file(dir, &files, &schema, meta.num_docs).await?;
1659 let dense_file_backed_bytes = vectors_data.file_backed_bytes;
1660 let vector_indexes = vectors_data.indexes;
1661 let flat_vectors = vectors_data.flat_vectors;
1662
1663 #[cfg(feature = "native")]
1668 for (field_id, lazy_flat) in &flat_vectors {
1669 if vector_indexes.contains_key(field_id) {
1670 lazy_flat.advise_random_access();
1671 }
1672 }
1673
1674 let sparse_data = loader::load_sparse_file(dir, &files, meta.num_docs, &schema).await?;
1676 let sparse_file_backed_bytes = sparse_data.file_backed_bytes;
1677 let sparse_indexes = sparse_data.maxscore_indexes;
1678 let bmp_indexes = sparse_data.bmp_indexes;
1679
1680 let positions_handle = loader::open_positions_file(dir, &files, &schema).await?;
1682
1683 let fast_fields = loader::load_fast_fields_file(dir, &files, &schema).await?;
1685
1686 {
1688 let mut parts = vec![format!(
1689 "[segment] loaded {:016x}: docs={}",
1690 segment_id.0, meta.num_docs
1691 )];
1692 if !vector_indexes.is_empty() || !flat_vectors.is_empty() {
1693 parts.push(format!(
1694 "dense vectors: {} ANN + {} flat fields",
1695 vector_indexes.len(),
1696 flat_vectors.len()
1697 ));
1698 }
1699 for (field_id, idx) in &sparse_indexes {
1700 parts.push(format!(
1701 "sparse vector field {}: {} dims, ~{}",
1702 field_id,
1703 idx.num_dimensions(),
1704 crate::format_bytes(idx.num_dimensions() as u64 * 24)
1705 ));
1706 }
1707 for (field_id, idx) in &bmp_indexes {
1708 parts.push(format!(
1709 "bmp field {}: {} dims, {} blocks",
1710 field_id,
1711 idx.dims(),
1712 idx.num_blocks
1713 ));
1714 }
1715 if !fast_fields.is_empty() {
1716 parts.push(format!("fast: {} fields", fast_fields.len()));
1717 }
1718 log::debug!("{}", parts.join(", "));
1719 }
1720
1721 #[allow(unused_mut)]
1722 let mut reader = Self {
1723 meta,
1724 term_dict: Arc::new(term_dict),
1725 postings_handle,
1726 store: Arc::new(store),
1727 schema,
1728 vector_indexes,
1729 flat_vectors,
1730 dense_file_backed_bytes,
1731 trained_vectors: Arc::new(crate::segment::TrainedVectorStructures::default()),
1732 sparse_indexes,
1733 bmp_indexes,
1734 sparse_file_backed_bytes,
1735 positions_handle,
1736 fast_fields,
1737 #[cfg(feature = "native")]
1738 dense_pin_report: Default::default(),
1739 #[cfg(feature = "native")]
1740 sparse_pin_report: Default::default(),
1741 };
1742
1743 #[cfg(feature = "native")]
1745 reader.apply_pin_policy(&crate::segment::pin::pin_policy().to_owned());
1746
1747 for (&field_id, vector_index) in &reader.vector_indexes {
1752 match vector_index {
1753 VectorIndex::BinaryIvf(index)
1754 | VectorIndex::IvfTq { index, .. }
1755 | VectorIndex::ScannAh(index)
1756 | VectorIndex::ScannBinary(index) => {
1757 index.get().report_health(
1758 reader.schema.index_label(),
1759 field_id,
1760 reader.meta.id,
1761 );
1762 }
1763 VectorIndex::Tq { .. } => {}
1766 }
1767 }
1768
1769 Ok(reader)
1770 }
1771
1772 pub fn ann_health(&self, field: Field) -> Option<crate::segment::ann_disk::AnnHealth> {
1776 match self.vector_indexes.get(&field.0)? {
1777 VectorIndex::BinaryIvf(index)
1778 | VectorIndex::IvfTq { index, .. }
1779 | VectorIndex::ScannAh(index)
1780 | VectorIndex::ScannBinary(index) => Some(index.get().health()),
1781 VectorIndex::Tq { .. } => None,
1782 }
1783 }
1784
1785 #[cfg(feature = "native")]
1795 pub(crate) fn apply_pin_policy(&mut self, policy: &crate::segment::pin::PinPolicy) {
1796 use crate::segment::pin::PinReport;
1797
1798 if !policy.is_enabled() {
1799 return;
1800 }
1801 let mut remaining = policy.budget_bytes;
1802 let mut dense_report = PinReport::default();
1803 let mut sparse_report = PinReport::default();
1804
1805 for index in self.vector_indexes.values_mut() {
1807 index.pin_lookup_directory(policy.mode, &mut remaining, &mut dense_report);
1808 }
1809 for bmp in self.bmp_indexes.values_mut() {
1811 bmp.pin_block_starts(policy.mode, &mut remaining, &mut sparse_report);
1812 }
1813 for sparse in self.sparse_indexes.values_mut() {
1815 sparse.pin_skip_section(policy.mode, &mut remaining, &mut sparse_report);
1816 }
1817 for flat in self.flat_vectors.values_mut() {
1819 flat.pin_doc_ids(policy.mode, &mut remaining, &mut dense_report);
1820 }
1821 for bmp in self.bmp_indexes.values_mut() {
1822 bmp.pin_doc_maps(policy.mode, &mut remaining, &mut sparse_report);
1823 }
1824 for bmp in self.bmp_indexes.values_mut() {
1826 bmp.pin_query_hierarchy(policy.mode, &mut remaining, &mut sparse_report);
1827 }
1828
1829 let report = PinReport {
1830 intended_bytes: dense_report
1831 .intended_bytes
1832 .saturating_add(sparse_report.intended_bytes),
1833 pinned_bytes: dense_report
1834 .pinned_bytes
1835 .saturating_add(sparse_report.pinned_bytes),
1836 skipped_budget_bytes: dense_report
1837 .skipped_budget_bytes
1838 .saturating_add(sparse_report.skipped_budget_bytes),
1839 failed_bytes: dense_report
1840 .failed_bytes
1841 .saturating_add(sparse_report.failed_bytes),
1842 heap_copy_bytes: dense_report
1843 .heap_copy_bytes
1844 .saturating_add(sparse_report.heap_copy_bytes),
1845 };
1846 if report.skipped_budget_bytes > 0 || report.failed_bytes > 0 {
1847 log::warn!(
1848 "[pin] index={} segment {:016x}: pinned {}/{} (budget skipped {}, mlock failed {}) — \
1849 raise HERMES_PIN_METADATA_BUDGET_MB or RLIMIT_MEMLOCK for full coverage",
1850 self.schema.index_label(),
1851 self.meta.id,
1852 crate::format_bytes(report.pinned_bytes),
1853 crate::format_bytes(report.intended_bytes),
1854 crate::format_bytes(report.skipped_budget_bytes),
1855 crate::format_bytes(report.failed_bytes),
1856 );
1857 } else if report.pinned_bytes > 0 {
1858 log::info!(
1859 "[pin] index={} segment {:016x}: pinned {} of hot metadata ({:?})",
1860 self.schema.index_label(),
1861 self.meta.id,
1862 crate::format_bytes(report.pinned_bytes),
1863 policy.mode,
1864 );
1865 }
1866 self.dense_pin_report = dense_report;
1867 self.sparse_pin_report = sparse_report;
1868 }
1869
1870 pub fn meta(&self) -> &SegmentMeta {
1876 &self.meta
1877 }
1878
1879 pub fn num_docs(&self) -> u32 {
1880 self.meta.num_docs
1881 }
1882
1883 pub fn avg_field_len(&self, field: Field) -> f32 {
1885 self.meta.avg_field_len(field)
1886 }
1887
1888 pub fn schema(&self) -> &Schema {
1889 &self.schema
1890 }
1891
1892 pub fn sparse_indexes(&self) -> &FxHashMap<u32, SparseIndex> {
1894 &self.sparse_indexes
1895 }
1896
1897 pub fn sparse_index(&self, field: Field) -> Option<&SparseIndex> {
1899 self.sparse_indexes.get(&field.0)
1900 }
1901
1902 pub fn bmp_index(&self, field: Field) -> Option<&BmpIndex> {
1904 self.bmp_indexes.get(&field.0)
1905 }
1906
1907 pub fn bmp_indexes(&self) -> &FxHashMap<u32, BmpIndex> {
1909 &self.bmp_indexes
1910 }
1911
1912 pub fn vector_indexes(&self) -> &FxHashMap<u32, VectorIndex> {
1914 &self.vector_indexes
1915 }
1916
1917 pub fn flat_vectors(&self) -> &FxHashMap<u32, LazyFlatVectorData> {
1919 &self.flat_vectors
1920 }
1921
1922 pub fn fast_field(
1924 &self,
1925 field_id: u32,
1926 ) -> Option<&crate::structures::fast_field::FastFieldReader> {
1927 self.fast_fields.get(&field_id)
1928 }
1929
1930 pub fn fast_fields(&self) -> &FxHashMap<u32, crate::structures::fast_field::FastFieldReader> {
1932 &self.fast_fields
1933 }
1934
1935 pub fn term_dict_stats(&self) -> SSTableStats {
1937 self.term_dict.stats()
1938 }
1939
1940 pub fn memory_stats(&self) -> SegmentMemoryStats {
1942 let term_dict_stats = self.term_dict.stats();
1943
1944 let term_dict_cache_bytes = self.term_dict.cached_bytes();
1948 let store_cache_bytes = self.store.cached_bytes();
1949
1950 let sparse_heap_bytes: usize = self
1953 .sparse_indexes
1954 .values()
1955 .map(|s| s.estimated_heap_bytes())
1956 .sum::<usize>()
1957 + self
1958 .bmp_indexes
1959 .values()
1960 .map(|b| b.estimated_heap_bytes())
1961 .sum::<usize>();
1962
1963 let dense_heap_bytes: usize = self
1966 .vector_indexes
1967 .values()
1968 .map(|v| v.estimated_heap_bytes())
1969 .sum::<usize>()
1970 + self
1971 .flat_vectors
1972 .values()
1973 .map(LazyFlatVectorData::estimated_heap_bytes)
1974 .sum::<usize>();
1975
1976 #[cfg(feature = "native")]
1977 let (sparse_heap_bytes, dense_heap_bytes) = (
1978 sparse_heap_bytes.saturating_add(
1979 usize::try_from(self.sparse_pin_report.heap_copy_bytes).unwrap_or(usize::MAX),
1980 ),
1981 dense_heap_bytes.saturating_add(
1982 usize::try_from(self.dense_pin_report.heap_copy_bytes).unwrap_or(usize::MAX),
1983 ),
1984 );
1985
1986 #[cfg(feature = "native")]
1987 let (
1988 sparse_pinned_metadata_bytes,
1989 sparse_pin_intended_bytes,
1990 dense_pinned_metadata_bytes,
1991 dense_pin_intended_bytes,
1992 ) = (
1993 self.sparse_pin_report.pinned_bytes,
1994 self.sparse_pin_report.intended_bytes,
1995 self.dense_pin_report.pinned_bytes,
1996 self.dense_pin_report.intended_bytes,
1997 );
1998 #[cfg(not(feature = "native"))]
1999 let (
2000 sparse_pinned_metadata_bytes,
2001 sparse_pin_intended_bytes,
2002 dense_pinned_metadata_bytes,
2003 dense_pin_intended_bytes,
2004 ) = (0u64, 0u64, 0u64, 0u64);
2005
2006 let pinned_metadata_bytes =
2007 sparse_pinned_metadata_bytes.saturating_add(dense_pinned_metadata_bytes);
2008 let pin_intended_bytes = sparse_pin_intended_bytes.saturating_add(dense_pin_intended_bytes);
2009
2010 SegmentMemoryStats {
2011 segment_id: self.meta.id,
2012 num_docs: self.meta.num_docs,
2013 term_dict_cache_bytes,
2014 store_cache_bytes,
2015 sparse_heap_bytes,
2016 dense_heap_bytes,
2017 term_bloom_file_bytes: term_dict_stats.bloom_filter_size as u64,
2018 sparse_file_backed_bytes: self.sparse_file_backed_bytes,
2019 dense_file_backed_bytes: self.dense_file_backed_bytes,
2020 pinned_metadata_bytes,
2021 pin_intended_bytes,
2022 sparse_pinned_metadata_bytes,
2023 sparse_pin_intended_bytes,
2024 dense_pinned_metadata_bytes,
2025 dense_pin_intended_bytes,
2026 }
2027 }
2028
2029 pub async fn get_postings(
2034 &self,
2035 field: Field,
2036 term: &[u8],
2037 ) -> Result<Option<BlockPostingList>> {
2038 log::debug!(
2039 "SegmentReader::get_postings field={} term_len={}",
2040 field.0,
2041 term.len()
2042 );
2043
2044 let mut key = Vec::with_capacity(4 + term.len());
2046 key.extend_from_slice(&field.0.to_le_bytes());
2047 key.extend_from_slice(term);
2048
2049 let term_info = match self.term_dict.get(&key).await? {
2051 Some(info) => {
2052 log::debug!("SegmentReader::get_postings found term_info");
2053 info
2054 }
2055 None => {
2056 log::debug!("SegmentReader::get_postings term not found");
2057 return Ok(None);
2058 }
2059 };
2060
2061 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
2063 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
2065 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
2066 posting_list.push(doc_id, tf);
2067 }
2068 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
2069 return Ok(Some(block_list));
2070 }
2071
2072 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
2074 Error::Corruption("TermInfo has neither inline nor external data".to_string())
2075 })?;
2076
2077 let range = checked_file_range(
2078 posting_offset,
2079 posting_len,
2080 self.postings_handle.len(),
2081 "posting",
2082 )?;
2083 let posting_bytes = self.postings_handle.read_bytes_range(range).await?;
2084 let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
2085
2086 Ok(Some(block_list))
2087 }
2088
2089 pub async fn get_prefix_postings(
2091 &self,
2092 field: Field,
2093 prefix: &[u8],
2094 ) -> Result<Vec<BlockPostingList>> {
2095 if prefix.is_empty() {
2096 return Err(Error::Query("prefix must not be empty".into()));
2097 }
2098 let mut key_prefix = Vec::with_capacity(4 + prefix.len());
2100 key_prefix.extend_from_slice(&field.0.to_le_bytes());
2101 key_prefix.extend_from_slice(prefix);
2102
2103 let (entries, truncated) = self
2104 .term_dict
2105 .prefix_scan_limited(&key_prefix, MAX_PREFIX_TERMS)
2106 .await?;
2107 if truncated {
2108 return Err(Error::Query(format!(
2109 "prefix expands to more than {MAX_PREFIX_TERMS} terms"
2110 )));
2111 }
2112 let posting_count: u64 = entries
2113 .iter()
2114 .map(|(_, term_info)| term_info.doc_freq() as u64)
2115 .sum();
2116 if posting_count > MAX_PREFIX_POSTINGS {
2117 return Err(Error::Query(format!(
2118 "prefix expands to {posting_count} postings (maximum {MAX_PREFIX_POSTINGS})"
2119 )));
2120 }
2121 let mut results = Vec::with_capacity(entries.len());
2122
2123 for (_key, term_info) in entries {
2124 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
2125 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
2126 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
2127 posting_list.push(doc_id, tf);
2128 }
2129 results.push(BlockPostingList::from_posting_list(&posting_list)?);
2130 } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
2131 let range = checked_file_range(
2132 posting_offset,
2133 posting_len,
2134 self.postings_handle.len(),
2135 "prefix posting",
2136 )?;
2137 let posting_bytes = self.postings_handle.read_bytes_range(range).await?;
2138 results.push(BlockPostingList::deserialize_zero_copy(posting_bytes)?);
2139 }
2140 }
2141
2142 Ok(results)
2143 }
2144
2145 pub async fn doc(&self, local_doc_id: DocId) -> Result<Option<Document>> {
2150 self.doc_with_fields(local_doc_id, None).await
2151 }
2152
2153 pub async fn doc_with_fields(
2159 &self,
2160 local_doc_id: DocId,
2161 fields: Option<&rustc_hash::FxHashSet<u32>>,
2162 ) -> Result<Option<Document>> {
2163 let mut doc = match fields {
2164 Some(set) => {
2165 let field_ids: Vec<u32> = set.iter().copied().collect();
2166 match self
2167 .store
2168 .get_fields(local_doc_id, &self.schema, &field_ids)
2169 .await
2170 {
2171 Ok(Some(d)) => d,
2172 Ok(None) => return Ok(None),
2173 Err(e) => return Err(Error::from(e)),
2174 }
2175 }
2176 None => match self.store.get(local_doc_id, &self.schema).await {
2177 Ok(Some(d)) => d,
2178 Ok(None) => return Ok(None),
2179 Err(e) => return Err(Error::from(e)),
2180 },
2181 };
2182
2183 for (&field_id, lazy_flat) in &self.flat_vectors {
2185 if let Some(set) = fields
2187 && !set.contains(&field_id)
2188 {
2189 continue;
2190 }
2191
2192 let is_binary = lazy_flat.quantization == DenseVectorQuantization::Binary;
2193 let (start, entries) = lazy_flat.flat_indexes_for_doc(local_doc_id);
2194 for (j, &(_doc_id, _ordinal)) in entries.iter().enumerate() {
2195 let flat_idx = start + j;
2196 if is_binary {
2197 let vbs = lazy_flat.vector_byte_size();
2198 let mut raw = vec![0u8; vbs];
2199 match lazy_flat.read_vector_raw_into(flat_idx, &mut raw).await {
2200 Ok(()) => {
2201 doc.add_binary_dense_vector(Field(field_id), raw);
2202 }
2203 Err(e) => {
2204 log::warn!(
2205 "Failed to hydrate binary dense vector field {}: {}",
2206 field_id,
2207 e
2208 );
2209 }
2210 }
2211 } else {
2212 match lazy_flat.get_vector(flat_idx).await {
2213 Ok(vec) => {
2214 doc.add_dense_vector(Field(field_id), vec);
2215 }
2216 Err(e) => {
2217 log::warn!("Failed to hydrate dense vector field {}: {}", field_id, e);
2218 }
2219 }
2220 }
2221 }
2222 }
2223
2224 Ok(Some(doc))
2225 }
2226
2227 pub async fn prefetch_terms(
2229 &self,
2230 field: Field,
2231 start_term: &[u8],
2232 end_term: &[u8],
2233 ) -> Result<()> {
2234 let mut start_key = Vec::with_capacity(4 + start_term.len());
2235 start_key.extend_from_slice(&field.0.to_le_bytes());
2236 start_key.extend_from_slice(start_term);
2237
2238 let mut end_key = Vec::with_capacity(4 + end_term.len());
2239 end_key.extend_from_slice(&field.0.to_le_bytes());
2240 end_key.extend_from_slice(end_term);
2241
2242 self.term_dict.prefetch_range(&start_key, &end_key).await?;
2243 Ok(())
2244 }
2245
2246 pub fn store_has_dict(&self) -> bool {
2248 self.store.has_dict()
2249 }
2250
2251 pub fn store(&self) -> &super::store::AsyncStoreReader {
2253 &self.store
2254 }
2255
2256 pub fn store_raw_blocks(&self) -> Vec<RawStoreBlock> {
2258 self.store.raw_blocks()
2259 }
2260
2261 pub fn store_data_slice(&self) -> &FileHandle {
2263 self.store.data_slice()
2264 }
2265
2266 pub async fn all_terms(&self) -> Result<Vec<(Vec<u8>, TermInfo)>> {
2268 self.term_dict.all_entries().await.map_err(Error::from)
2269 }
2270
2271 pub async fn all_terms_with_stats(&self) -> Result<Vec<(Field, String, u32)>> {
2276 let entries = self.term_dict.all_entries().await?;
2277 let mut result = Vec::with_capacity(entries.len());
2278
2279 for (key, term_info) in entries {
2280 if key.len() > 4 {
2282 let field_id = u32::from_le_bytes([key[0], key[1], key[2], key[3]]);
2283 let term_bytes = &key[4..];
2284 if let Ok(term_str) = std::str::from_utf8(term_bytes) {
2285 result.push((Field(field_id), term_str.to_string(), term_info.doc_freq()));
2286 }
2287 }
2288 }
2289
2290 Ok(result)
2291 }
2292
2293 pub fn term_dict_iter(&self) -> crate::structures::AsyncSSTableIterator<'_, TermInfo> {
2295 self.term_dict.iter()
2296 }
2297
2298 pub async fn prefetch_term_dict(&self) -> crate::Result<()> {
2302 self.term_dict
2303 .prefetch_all_data_bulk()
2304 .await
2305 .map_err(crate::Error::from)
2306 }
2307
2308 pub async fn read_postings(&self, offset: u64, len: u64) -> Result<Vec<u8>> {
2310 let range = checked_file_range(offset, len, self.postings_handle.len(), "posting")?;
2311 let bytes = self.postings_handle.read_bytes_range(range).await?;
2312 Ok(bytes.to_vec())
2313 }
2314
2315 pub async fn read_position_bytes(&self, offset: u64, len: u64) -> Result<Option<Vec<u8>>> {
2317 let handle = match &self.positions_handle {
2318 Some(h) => h,
2319 None => return Ok(None),
2320 };
2321 let range = checked_file_range(offset, len, handle.len(), "position")?;
2322 let bytes = handle.read_bytes_range(range).await?;
2323 Ok(Some(bytes.to_vec()))
2324 }
2325
2326 pub fn has_positions_file(&self) -> bool {
2328 self.positions_handle.is_some()
2329 }
2330
2331 fn validate_dense_search_request(
2335 &self,
2336 field: Field,
2337 query: &[f32],
2338 nprobe: usize,
2339 rerank_factor: f32,
2340 combiner: crate::query::MultiValueCombiner,
2341 ) -> Result<DenseSearchParams> {
2342 let entry = self
2343 .schema
2344 .get_field_entry(field)
2345 .ok_or_else(|| Error::FieldNotFound(field.0.to_string()))?;
2346 if entry.field_type != crate::dsl::FieldType::DenseVector {
2347 return Err(Error::InvalidFieldType {
2348 expected: "dense_vector".to_string(),
2349 got: format!("{:?}", entry.field_type),
2350 });
2351 }
2352 let config = entry.dense_vector_config.as_ref().ok_or_else(|| {
2353 Error::Schema(format!(
2354 "dense vector field '{}' has no dense vector configuration",
2355 entry.name
2356 ))
2357 })?;
2358
2359 if query.is_empty() {
2360 return Err(Error::Query(format!(
2361 "dense query vector for field '{}' must not be empty",
2362 entry.name
2363 )));
2364 }
2365 if query.len() != config.dim {
2366 return Err(Error::Query(format!(
2367 "dense query vector dimension {} does not match field '{}' dimension {}",
2368 query.len(),
2369 entry.name,
2370 config.dim
2371 )));
2372 }
2373 if let Some((index, value)) = query
2374 .iter()
2375 .enumerate()
2376 .find(|(_, value)| !value.is_finite())
2377 {
2378 return Err(Error::Query(format!(
2379 "dense query vector for field '{}' contains non-finite value {value} at index {index}",
2380 entry.name
2381 )));
2382 }
2383
2384 let nprobe = match (nprobe, config.nprobe) {
2387 (0, 0) => 32,
2388 (0, schema_nprobe) => schema_nprobe,
2389 (query_nprobe, _) => query_nprobe,
2390 };
2391 if nprobe > MAX_DENSE_NPROBE {
2392 return Err(Error::Query(format!(
2393 "dense nprobe must be at most {MAX_DENSE_NPROBE}, got {nprobe}"
2394 )));
2395 }
2396
2397 checked_dense_fetch_k(0, rerank_factor)?;
2400 combiner.validate().map_err(Error::Query)?;
2401
2402 Ok(DenseSearchParams {
2403 dim: config.dim,
2404 nprobe,
2405 unit_norm: config.unit_norm,
2406 })
2407 }
2408
2409 fn validate_binary_search_request(&self, field: Field, query: &[u8]) -> Result<usize> {
2410 let entry = self
2411 .schema
2412 .get_field_entry(field)
2413 .ok_or_else(|| Error::FieldNotFound(field.0.to_string()))?;
2414 if entry.field_type != crate::dsl::FieldType::BinaryDenseVector {
2415 return Err(Error::InvalidFieldType {
2416 expected: "binary_dense_vector".to_string(),
2417 got: format!("{:?}", entry.field_type),
2418 });
2419 }
2420 let config = entry.binary_dense_vector_config.as_ref().ok_or_else(|| {
2421 Error::Schema(format!(
2422 "binary dense vector field '{}' has no configuration",
2423 entry.name
2424 ))
2425 })?;
2426 if config.dim == 0 || !config.dim.is_multiple_of(8) {
2427 return Err(Error::Schema(format!(
2428 "binary dense vector field '{}' has invalid dimension {}",
2429 entry.name, config.dim
2430 )));
2431 }
2432 if query.len() != config.byte_len() {
2433 return Err(Error::Query(format!(
2434 "binary query byte length {} does not match field '{}' byte length {}",
2435 query.len(),
2436 entry.name,
2437 config.byte_len()
2438 )));
2439 }
2440 Ok(config.dim)
2441 }
2442
2443 #[cfg(test)]
2445 fn score_quantized_batch_legacy(
2446 query: &[f32],
2447 raw: &[u8],
2448 quant: crate::dsl::DenseVectorQuantization,
2449 dim: usize,
2450 scores: &mut [f32],
2451 unit_norm: bool,
2452 ) -> Result<()> {
2453 use crate::dsl::DenseVectorQuantization;
2454 use crate::structures::simd;
2455
2456 if query.len() != dim {
2457 return Err(Error::Query(format!(
2458 "dense SIMD query dimension {} does not match vector dimension {dim}",
2459 query.len()
2460 )));
2461 }
2462 let element_size = match quant {
2463 DenseVectorQuantization::F32 => std::mem::size_of::<f32>(),
2464 DenseVectorQuantization::F16 => std::mem::size_of::<u16>(),
2465 DenseVectorQuantization::UInt8 => 1,
2466 DenseVectorQuantization::Binary => {
2467 return Err(Error::InvalidFieldType {
2468 expected: "non-binary dense vector".to_string(),
2469 got: "binary dense vector".to_string(),
2470 });
2471 }
2472 };
2473 let required_bytes = scores
2474 .len()
2475 .checked_mul(dim)
2476 .and_then(|elements| elements.checked_mul(element_size))
2477 .ok_or_else(|| Error::Corruption("dense vector batch byte length overflow".into()))?;
2478 if raw.len() < required_bytes {
2479 return Err(Error::Corruption(format!(
2480 "dense vector batch is truncated: need {required_bytes} bytes, got {}",
2481 raw.len()
2482 )));
2483 }
2484 if quant == DenseVectorQuantization::F16
2485 && required_bytes > 0
2486 && !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>())
2487 {
2488 return Err(Error::Corruption(
2489 "f16 vector data is not 2-byte aligned".to_string(),
2490 ));
2491 }
2492
2493 match (quant, unit_norm) {
2494 (DenseVectorQuantization::F32, false) => {
2495 let num_floats = scores.len() * dim;
2496 if !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()) {
2497 return Err(Error::Corruption(
2498 "f32 vector data is not 4-byte aligned".to_string(),
2499 ));
2500 }
2501 let vectors: &[f32] =
2502 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
2503 simd::batch_cosine_scores(query, vectors, dim, scores);
2504 }
2505 (DenseVectorQuantization::F32, true) => {
2506 let num_floats = scores.len() * dim;
2507 if !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()) {
2508 return Err(Error::Corruption(
2509 "f32 vector data is not 4-byte aligned".to_string(),
2510 ));
2511 }
2512 let vectors: &[f32] =
2513 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
2514 simd::batch_dot_scores(query, vectors, dim, scores);
2515 }
2516 (DenseVectorQuantization::F16, false) => {
2517 simd::batch_cosine_scores_f16(query, raw, dim, scores);
2518 }
2519 (DenseVectorQuantization::F16, true) => {
2520 simd::batch_dot_scores_f16(query, raw, dim, scores);
2521 }
2522 (DenseVectorQuantization::UInt8, false) => {
2523 simd::batch_cosine_scores_u8(query, raw, dim, scores);
2524 }
2525 (DenseVectorQuantization::UInt8, true) => {
2526 simd::batch_dot_scores_u8(query, raw, dim, scores);
2527 }
2528 (DenseVectorQuantization::Binary, _) => unreachable!("validated above"),
2529 }
2530 Ok(())
2531 }
2532
2533 pub async fn search_dense_vector(
2539 &self,
2540 field: Field,
2541 query: &[f32],
2542 k: usize,
2543 nprobe: usize,
2544 rerank_factor: f32,
2545 combiner: crate::query::MultiValueCombiner,
2546 ) -> Result<Vec<VectorSearchResult>> {
2547 self.search_dense_vector_impl(field, query, k, nprobe, rerank_factor, combiner, None)
2548 .await
2549 }
2550
2551 #[allow(clippy::too_many_arguments)]
2552 pub(crate) async fn search_dense_vector_with_probe_cache(
2553 &self,
2554 field: Field,
2555 query: &[f32],
2556 k: usize,
2557 nprobe: usize,
2558 rerank_factor: f32,
2559 combiner: crate::query::MultiValueCombiner,
2560 plan_cache: &DensePlanCache,
2561 ) -> Result<Vec<VectorSearchResult>> {
2562 self.search_dense_vector_impl(
2563 field,
2564 query,
2565 k,
2566 nprobe,
2567 rerank_factor,
2568 combiner,
2569 Some(plan_cache),
2570 )
2571 .await
2572 }
2573
2574 #[allow(clippy::too_many_arguments)]
2575 async fn search_dense_vector_impl(
2576 &self,
2577 field: Field,
2578 query: &[f32],
2579 k: usize,
2580 nprobe: usize,
2581 rerank_factor: f32,
2582 combiner: crate::query::MultiValueCombiner,
2583 plan_cache: Option<&DensePlanCache>,
2584 ) -> Result<Vec<VectorSearchResult>> {
2585 let params =
2586 self.validate_dense_search_request(field, query, nprobe, rerank_factor, combiner)?;
2587 let fetch_k = checked_dense_fetch_k(k, rerank_factor)?;
2588 if k == 0 {
2589 return Ok(Vec::new());
2590 }
2591
2592 let configured_ann_index = self.vector_indexes.get(&field.0);
2593 let lazy_flat = self.flat_vectors.get(&field.0);
2594 if configured_ann_index.is_none() && lazy_flat.is_none() {
2596 return Ok(Vec::new());
2597 }
2598
2599 if configured_ann_index.is_some() && lazy_flat.is_none() {
2600 return Err(Error::Corruption(format!(
2601 "dense ANN field {} is missing flat vector storage",
2602 field.0
2603 )));
2604 }
2605
2606 if let Some(flat) = lazy_flat
2607 && flat.dim != params.dim
2608 {
2609 return Err(Error::Corruption(format!(
2610 "dense vector field {} has schema dimension {} but flat storage dimension {}",
2611 field.0, params.dim, flat.dim
2612 )));
2613 }
2614
2615 let needs_document_aggregation = lazy_flat.is_some_and(|flat| {
2616 flat.num_vectors != flat.num_docs_with_vectors()
2617 && !matches!(combiner, crate::query::MultiValueCombiner::Max)
2618 });
2619 let ann_index = configured_ann_index;
2623
2624 let t0 = std::time::Instant::now();
2626 let mut flat_results = None;
2627 let results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
2628 match index {
2630 VectorIndex::Tq { index: lazy, codec } => {
2631 let flat = lazy_flat.expect("ANN/flat pairing validated above");
2632 search_tq_segment(
2634 lazy.get(),
2635 codec,
2636 query,
2637 fetch_k.min(flat.num_docs_with_vectors()),
2638 needs_document_aggregation.then_some(combiner),
2639 field,
2640 params.dim,
2641 plan_cache.map(|cache| &cache.tq),
2642 )?
2643 }
2644 VectorIndex::IvfTq { index: lazy, codec } => {
2645 let index = lazy.get();
2646 let centroids =
2647 self.trained_vectors
2648 .centroids
2649 .get(&field.0)
2650 .ok_or_else(|| {
2651 Error::Schema(format!(
2652 "IVF-TQ index requires coarse centroids for field {}",
2653 field.0
2654 ))
2655 })?;
2656 validate_coarse_centroids(centroids, params.dim)?;
2657 let routing = self
2658 .schema
2659 .get_field_entry(field)
2660 .and_then(|entry| entry.dense_vector_config.as_ref())
2661 .map_or(crate::dsl::IvfRoutingMode::Auto, |config| {
2662 config.ivf_routing
2663 });
2664 validate_ivf_tq_ann(index, centroids, codec, params.dim, routing, field)?;
2665 let flat = lazy_flat.expect("ANN/flat pairing validated above");
2666 search_ivf_tq_segment(
2667 index,
2668 centroids,
2669 codec,
2670 query,
2671 fetch_k.min(flat.num_docs_with_vectors()),
2672 needs_document_aggregation.then_some(combiner),
2673 field,
2674 params.nprobe,
2675 routing,
2676 plan_cache.map(|cache| &cache.ivf_tq),
2677 )?
2678 }
2679 VectorIndex::BinaryIvf(_) => {
2680 Vec::new()
2682 }
2683 VectorIndex::ScannAh(lazy) => {
2684 let artifact = self
2685 .trained_vectors
2686 .scann_artifacts
2687 .get(&field.0)
2688 .ok_or_else(|| {
2689 Error::Schema(format!(
2690 "ScaNN field {} has no loaded global artifact",
2691 field.0
2692 ))
2693 })?;
2694 let flat = lazy_flat.expect("ANN/flat pairing validated above");
2695 search_scann_ah_segment(
2696 lazy.get(),
2697 artifact,
2698 query,
2699 fetch_k.min(flat.num_docs_with_vectors()),
2700 combiner,
2701 field,
2702 params.nprobe,
2703 plan_cache.map(|cache| &cache.scann),
2704 )?
2705 }
2706 VectorIndex::ScannBinary(_) => {
2707 return Err(Error::Corruption(format!(
2708 "binary ScaNN payload was attached to float field {}",
2709 field.0
2710 )));
2711 }
2712 }
2713 } else if let Some(lazy_flat) = lazy_flat {
2714 log::debug!(
2718 "[dense_vector_search] index={} field {}: brute-force on {} vectors (dim={}, quant={:?})",
2719 self.schema.index_label(),
2720 field.0,
2721 lazy_flat.num_vectors,
2722 lazy_flat.dim,
2723 lazy_flat.quantization
2724 );
2725 let dim = lazy_flat.dim;
2726 let n = lazy_flat.num_vectors;
2727 let quant = lazy_flat.quantization;
2728 let batch_len =
2729 bounded_vector_score_batch(lazy_flat.vector_byte_size(), DENSE_SCORE_BATCH);
2730 let mut collector = FlatDocumentCollector::new(fetch_k.min(n), combiner);
2731 let mut scores = vec![0f32; batch_len];
2732 let prepared_query = PreparedDenseScoreQuery::new(query, quant, dim, params.unit_norm)?;
2733
2734 for batch_start in (0..n).step_by(batch_len) {
2735 let batch_count = batch_len.min(n - batch_start);
2736 let batch_bytes = lazy_flat
2737 .read_vectors_batch(batch_start, batch_count)
2738 .await
2739 .map_err(crate::Error::Io)?;
2740 let raw = batch_bytes.as_slice();
2741
2742 prepared_query.score_batch(raw, &mut scores[..batch_count])?;
2743
2744 for (i, &score) in scores.iter().enumerate().take(batch_count) {
2745 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
2746 collector.push(doc_id, ordinal, score);
2747 }
2748 }
2749
2750 flat_results = Some(collector.into_results());
2751 Vec::new()
2752 } else {
2753 return Ok(Vec::new());
2754 };
2755 let l1_elapsed = t0.elapsed();
2756 {
2757 let kind = match ann_index {
2758 Some(VectorIndex::BinaryIvf(_)) => "binary_ivf",
2759 Some(VectorIndex::Tq { .. }) => "tq_flat",
2760 Some(VectorIndex::IvfTq { .. }) => "ivf_tq",
2761 Some(VectorIndex::ScannAh(_)) => "scann_ah",
2762 Some(VectorIndex::ScannBinary(_)) => "scann_binary",
2763 None => "flat",
2764 };
2765 crate::observe::dense_l1(
2766 self.schema.index_label(),
2767 self.schema.get_field_name(field).unwrap_or("?"),
2768 kind,
2769 l1_elapsed.as_secs_f64(),
2770 flat_results.as_ref().map_or(results.len(), Vec::len),
2771 );
2772 }
2773 log::debug!(
2774 "[dense_vector_search] index={} field {}: L1 returned {} candidates in {:.1}ms",
2775 self.schema.index_label(),
2776 field.0,
2777 flat_results.as_ref().map_or(results.len(), Vec::len),
2778 l1_elapsed.as_secs_f64() * 1000.0
2779 );
2780
2781 if let Some(results) = flat_results {
2782 return Ok(results);
2783 }
2784
2785 if ann_index.is_some()
2788 && !results.is_empty()
2789 && let Some(lazy_flat) = lazy_flat
2790 {
2791 let t_rerank = std::time::Instant::now();
2792 let vbs = lazy_flat.vector_byte_size();
2793 let (reranked, stats) = exact_score_dense_candidate_documents(
2794 &results,
2795 lazy_flat,
2796 query,
2797 params.unit_norm,
2798 combiner,
2799 k,
2800 )
2801 .await?;
2802
2803 crate::observe::dense_rerank(
2804 self.schema.index_label(),
2805 self.schema.get_field_name(field).unwrap_or("?"),
2806 t_rerank.elapsed().as_secs_f64(),
2807 stats.resolve_elapsed.as_secs_f64(),
2808 stats.read_elapsed.as_secs_f64(),
2809 stats.vector_count,
2810 );
2811 log::debug!(
2812 "[dense_vector_search] index={} field {}: rerank {} vectors (dim={}, quant={:?}, bytes_per_vector={}): resolve={:.1}ms read={:.1}ms score={:.1}ms",
2813 self.schema.index_label(),
2814 field.0,
2815 stats.vector_count,
2816 lazy_flat.dim,
2817 lazy_flat.quantization,
2818 vbs,
2819 stats.resolve_elapsed.as_secs_f64() * 1000.0,
2820 stats.read_elapsed.as_secs_f64() * 1000.0,
2821 stats.score_elapsed.as_secs_f64() * 1000.0,
2822 );
2823
2824 log::debug!(
2825 "[dense_vector_search] index={} field {}: rerank total={:.1}ms",
2826 self.schema.index_label(),
2827 field.0,
2828 t_rerank.elapsed().as_secs_f64() * 1000.0
2829 );
2830 return Ok(reranked);
2831 }
2832
2833 Ok(combine_grouped_ordinal_results(results, combiner, k))
2834 }
2835
2836 async fn search_binary_dense_vector_impl(
2841 &self,
2842 field: Field,
2843 query: &[u8],
2844 k: usize,
2845 combiner: crate::query::MultiValueCombiner,
2846 probe_cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
2847 ) -> Result<Vec<VectorSearchResult>> {
2848 let schema_dim = self.validate_binary_search_request(field, query)?;
2849 combiner.validate().map_err(Error::Query)?;
2850 if k == 0 {
2851 return Ok(Vec::new());
2852 }
2853 let t0 = crate::observe::Timer::start();
2854 if let Some(VectorIndex::ScannBinary(lazy)) = self.vector_indexes.get(&field.0) {
2855 let artifact = self
2856 .trained_vectors
2857 .scann_artifacts
2858 .get(&field.0)
2859 .ok_or_else(|| {
2860 Error::Schema(format!(
2861 "binary ScaNN field {} has no loaded global artifact",
2862 field.0
2863 ))
2864 })?;
2865 lazy.get()
2866 .validate_scann_generation(
2867 artifact.config(),
2868 artifact.generation(),
2869 artifact.artifact_id(),
2870 )
2871 .map_err(|error| {
2872 Error::Corruption(format!(
2873 "binary ScaNN generation mismatch for field {}: {error}",
2874 field.0
2875 ))
2876 })?;
2877 let config = self
2878 .schema
2879 .get_field_entry(field)
2880 .and_then(|entry| entry.binary_dense_vector_config.as_ref())
2881 .ok_or_else(|| {
2882 Error::Schema(format!(
2883 "binary ScaNN field {} has no schema configuration",
2884 field.0
2885 ))
2886 })?;
2887 let model = artifact.binary_model().map_err(Error::Io)?;
2888 let clusters = binary_scann_probe_clusters(&model, query, config.nprobe, probe_cache)?;
2889 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
2890 Error::Corruption(format!(
2891 "binary ScaNN field {} is missing flat vectors",
2892 field.0
2893 ))
2894 })?;
2895 let candidate_limit =
2896 checked_binary_combined_fetch_k(k)?.min(flat.num_docs_with_vectors());
2897 let (documents, ordinal_scores) = lazy
2898 .get()
2899 .search_binary_combined_documents(candidate_limit, query, &clusters, combiner)
2900 .map_err(|error| {
2901 Error::Corruption(format!(
2902 "invalid binary ScaNN payload for field {}: {error}",
2903 field.0
2904 ))
2905 })?;
2906 let results = exact_score_binary_candidate_document_ids(
2907 documents
2908 .into_iter()
2909 .map(|candidate| candidate.doc_id)
2910 .collect(),
2911 &ordinal_scores,
2912 flat,
2913 query,
2914 schema_dim,
2915 combiner,
2916 k,
2917 )
2918 .await?;
2919 crate::observe::dense_l1(
2920 self.schema.index_label(),
2921 self.schema.get_field_name(field).unwrap_or("?"),
2922 "binary_scann",
2923 t0.secs(),
2924 results.len(),
2925 );
2926 return Ok(results);
2927 }
2928 if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0) {
2929 let ivf = lazy.get();
2930 let config = self
2931 .schema
2932 .get_field_entry(field)
2933 .and_then(|entry| entry.binary_dense_vector_config.as_ref())
2934 .ok_or_else(|| {
2935 Error::Schema(format!(
2936 "binary IVF field {} has no schema configuration",
2937 field.0
2938 ))
2939 })?;
2940 let quantizer = self
2941 .trained_vectors
2942 .binary_quantizers
2943 .get(&field.0)
2944 .ok_or_else(|| {
2945 Error::Schema(format!(
2946 "global binary IVF field {} has no loaded quantizer",
2947 field.0
2948 ))
2949 })?;
2950 validate_binary_ann(ivf, quantizer, config, schema_dim, field)?;
2951 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
2952 Error::Corruption(format!(
2953 "global binary IVF field {} is missing flat vector storage",
2954 field.0
2955 ))
2956 })?;
2957 let single_valued = flat.num_vectors == flat.num_docs_with_vectors();
2958 let clusters = binary_probe_clusters(
2959 quantizer,
2960 query,
2961 config.nprobe,
2962 config.ivf_routing,
2963 probe_cache,
2964 )?;
2965 let results = if !single_valued
2966 && !matches!(combiner, crate::query::MultiValueCombiner::Max)
2967 {
2968 let candidate_limit =
2969 checked_binary_combined_fetch_k(k)?.min(flat.num_docs_with_vectors());
2970 let (candidate_documents, probed_ordinal_scores) = ivf
2971 .search_binary_combined_documents(candidate_limit, query, &clusters, combiner)
2972 .map_err(|error| {
2973 Error::Corruption(format!(
2974 "invalid binary IVF payload for field {}: {error}",
2975 field.0,
2976 ))
2977 })?;
2978 exact_score_binary_candidate_document_ids(
2979 candidate_documents
2980 .into_iter()
2981 .map(|candidate| candidate.doc_id)
2982 .collect(),
2983 &probed_ordinal_scores,
2984 flat,
2985 query,
2986 schema_dim,
2987 combiner,
2988 k,
2989 )
2990 .await?
2991 } else {
2992 let candidate_docs = if single_valued {
2993 k
2994 } else {
2995 checked_binary_combined_fetch_k(k)?
3000 }
3001 .min(flat.num_docs_with_vectors());
3002 let ann_results = if single_valued {
3003 ivf.search_binary_clusters::<false>(query, candidate_docs, &clusters)
3004 } else {
3005 ivf.search_binary_clusters::<true>(query, candidate_docs, &clusters)
3006 }
3007 .map_err(|error| {
3008 Error::Corruption(format!(
3009 "invalid binary IVF payload for field {}: {error}",
3010 field.0,
3011 ))
3012 })?;
3013 if single_valued {
3016 let ann_results = validate_binary_single_value_ann_results(ann_results, flat)?;
3017 combine_ordinal_results(ann_results, combiner, k)
3018 } else {
3019 exact_score_binary_candidate_documents(
3020 &ann_results,
3021 flat,
3022 query,
3023 schema_dim,
3024 combiner,
3025 k,
3026 )
3027 .await?
3028 }
3029 };
3030 crate::observe::dense_l1(
3031 self.schema.index_label(),
3032 self.schema.get_field_name(field).unwrap_or("?"),
3033 "global_binary_ivf",
3034 t0.secs(),
3035 results.len(),
3036 );
3037 return Ok(results);
3038 }
3039 let lazy_flat = match self.flat_vectors.get(&field.0) {
3040 Some(f) => f,
3041 None => return Ok(Vec::new()),
3042 };
3043
3044 let dim_bits = lazy_flat.dim;
3045 let byte_len = lazy_flat.vector_byte_size();
3046 let n = lazy_flat.num_vectors;
3047
3048 if dim_bits != schema_dim {
3049 return Err(Error::Corruption(format!(
3050 "binary vector field {} has schema dimension {} but flat storage dimension {}",
3051 field.0, schema_dim, dim_bits
3052 )));
3053 }
3054
3055 if byte_len != query.len() {
3056 return Err(Error::Schema(format!(
3057 "Binary query vector byte length {} != field byte length {}",
3058 query.len(),
3059 byte_len
3060 )));
3061 }
3062
3063 let batch_len = bounded_vector_score_batch(byte_len, BINARY_SCORE_BATCH);
3064 let mut collector = FlatDocumentCollector::new(k, combiner);
3065 let mut scores = vec![0f32; batch_len];
3066
3067 for batch_start in (0..n).step_by(batch_len) {
3068 let batch_count = batch_len.min(n - batch_start);
3069 let batch_bytes = lazy_flat
3070 .read_vectors_batch(batch_start, batch_count)
3071 .await
3072 .map_err(crate::Error::Io)?;
3073 let raw = batch_bytes.as_slice();
3074
3075 crate::structures::simd::batch_hamming_scores(
3076 query,
3077 raw,
3078 byte_len,
3079 dim_bits,
3080 &mut scores[..batch_count],
3081 );
3082
3083 for (i, &score) in scores.iter().enumerate().take(batch_count) {
3084 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
3085 collector.push(doc_id, ordinal, score);
3086 }
3087 }
3088
3089 let results = collector.into_results();
3090
3091 crate::observe::dense_l1(
3092 self.schema.index_label(),
3093 self.schema.get_field_name(field).unwrap_or("?"),
3094 "binary_flat",
3095 t0.secs(),
3096 results.len(),
3097 );
3098 Ok(results)
3099 }
3100
3101 pub async fn search_binary_dense_vector(
3102 &self,
3103 field: Field,
3104 query: &[u8],
3105 k: usize,
3106 combiner: crate::query::MultiValueCombiner,
3107 ) -> Result<Vec<VectorSearchResult>> {
3108 self.search_binary_dense_vector_impl(field, query, k, combiner, None)
3109 .await
3110 }
3111
3112 pub(crate) async fn search_binary_dense_vector_with_probe_cache(
3113 &self,
3114 field: Field,
3115 query: &[u8],
3116 k: usize,
3117 combiner: crate::query::MultiValueCombiner,
3118 probe_cache: &std::sync::Mutex<Option<crate::structures::IvfProbePlan>>,
3119 ) -> Result<Vec<VectorSearchResult>> {
3120 self.search_binary_dense_vector_impl(field, query, k, combiner, Some(probe_cache))
3121 .await
3122 }
3123
3124 pub fn coarse_centroids(&self, field_id: u32) -> Option<&Arc<CoarseCentroids>> {
3126 self.trained_vectors.centroids.get(&field_id)
3127 }
3128
3129 pub fn set_trained_vectors(
3130 &mut self,
3131 trained_vectors: Arc<crate::segment::TrainedVectorStructures>,
3132 ) {
3133 self.trained_vectors = trained_vectors;
3134 }
3135
3136 pub fn get_vector_index(&self, field: Field) -> Option<&VectorIndex> {
3138 self.vector_indexes.get(&field.0)
3139 }
3140
3141 pub async fn get_positions(
3146 &self,
3147 field: Field,
3148 term: &[u8],
3149 ) -> Result<Option<crate::structures::PositionPostingList>> {
3150 let handle = match &self.positions_handle {
3152 Some(h) => h,
3153 None => return Ok(None),
3154 };
3155
3156 let mut key = Vec::with_capacity(4 + term.len());
3158 key.extend_from_slice(&field.0.to_le_bytes());
3159 key.extend_from_slice(term);
3160
3161 let term_info = match self.term_dict.get(&key).await? {
3163 Some(info) => info,
3164 None => return Ok(None),
3165 };
3166
3167 let (offset, length) = match term_info.position_info() {
3169 Some((o, l)) => (o, l),
3170 None => return Ok(None),
3171 };
3172
3173 let range = checked_file_range(offset, length, handle.len(), "position list")?;
3177 let slice = handle.slice(range);
3178 let data = slice.read_bytes().await?;
3179
3180 let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
3182
3183 Ok(Some(pos_list))
3184 }
3185
3186 pub fn has_positions(&self, field: Field) -> bool {
3188 if let Some(entry) = self.schema.get_field_entry(field) {
3190 entry.positions.is_some()
3191 } else {
3192 false
3193 }
3194 }
3195}
3196
3197#[cfg(feature = "sync")]
3199impl SegmentReader {
3200 pub fn get_postings_sync(&self, field: Field, term: &[u8]) -> Result<Option<BlockPostingList>> {
3202 let mut key = Vec::with_capacity(4 + term.len());
3204 key.extend_from_slice(&field.0.to_le_bytes());
3205 key.extend_from_slice(term);
3206
3207 let term_info = match self.term_dict.get_sync(&key)? {
3209 Some(info) => info,
3210 None => return Ok(None),
3211 };
3212
3213 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
3215 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
3216 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
3217 posting_list.push(doc_id, tf);
3218 }
3219 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
3220 return Ok(Some(block_list));
3221 }
3222
3223 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
3225 Error::Corruption("TermInfo has neither inline nor external data".to_string())
3226 })?;
3227
3228 let range = checked_file_range(
3229 posting_offset,
3230 posting_len,
3231 self.postings_handle.len(),
3232 "posting",
3233 )?;
3234 let posting_bytes = self.postings_handle.read_bytes_range_sync(range)?;
3235 let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
3236
3237 Ok(Some(block_list))
3238 }
3239
3240 pub fn get_prefix_postings_sync(
3242 &self,
3243 field: Field,
3244 prefix: &[u8],
3245 ) -> Result<Vec<BlockPostingList>> {
3246 if prefix.is_empty() {
3247 return Err(Error::Query("prefix must not be empty".into()));
3248 }
3249 let mut key_prefix = Vec::with_capacity(4 + prefix.len());
3250 key_prefix.extend_from_slice(&field.0.to_le_bytes());
3251 key_prefix.extend_from_slice(prefix);
3252
3253 let (entries, truncated) = self
3254 .term_dict
3255 .prefix_scan_limited_sync(&key_prefix, MAX_PREFIX_TERMS)?;
3256 if truncated {
3257 return Err(Error::Query(format!(
3258 "prefix expands to more than {MAX_PREFIX_TERMS} terms"
3259 )));
3260 }
3261 let posting_count: u64 = entries
3262 .iter()
3263 .map(|(_, term_info)| term_info.doc_freq() as u64)
3264 .sum();
3265 if posting_count > MAX_PREFIX_POSTINGS {
3266 return Err(Error::Query(format!(
3267 "prefix expands to {posting_count} postings (maximum {MAX_PREFIX_POSTINGS})"
3268 )));
3269 }
3270 let mut results = Vec::with_capacity(entries.len());
3271
3272 for (_key, term_info) in entries {
3273 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
3274 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
3275 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
3276 posting_list.push(doc_id, tf);
3277 }
3278 results.push(BlockPostingList::from_posting_list(&posting_list)?);
3279 } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
3280 let range = checked_file_range(
3281 posting_offset,
3282 posting_len,
3283 self.postings_handle.len(),
3284 "prefix posting",
3285 )?;
3286 let posting_bytes = self.postings_handle.read_bytes_range_sync(range)?;
3287 results.push(BlockPostingList::deserialize_zero_copy(posting_bytes)?);
3288 }
3289 }
3290
3291 Ok(results)
3292 }
3293
3294 pub fn get_positions_sync(
3296 &self,
3297 field: Field,
3298 term: &[u8],
3299 ) -> Result<Option<crate::structures::PositionPostingList>> {
3300 let handle = match &self.positions_handle {
3301 Some(h) => h,
3302 None => return Ok(None),
3303 };
3304
3305 let mut key = Vec::with_capacity(4 + term.len());
3307 key.extend_from_slice(&field.0.to_le_bytes());
3308 key.extend_from_slice(term);
3309
3310 let term_info = match self.term_dict.get_sync(&key)? {
3312 Some(info) => info,
3313 None => return Ok(None),
3314 };
3315
3316 let (offset, length) = match term_info.position_info() {
3317 Some((o, l)) => (o, l),
3318 None => return Ok(None),
3319 };
3320
3321 let range = checked_file_range(offset, length, handle.len(), "position list")?;
3322 let slice = handle.slice(range);
3323 let data = slice.read_bytes_sync()?;
3324
3325 let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
3326 Ok(Some(pos_list))
3327 }
3328
3329 pub fn search_dense_vector_sync(
3332 &self,
3333 field: Field,
3334 query: &[f32],
3335 k: usize,
3336 nprobe: usize,
3337 rerank_factor: f32,
3338 combiner: crate::query::MultiValueCombiner,
3339 ) -> Result<Vec<VectorSearchResult>> {
3340 self.search_dense_vector_sync_impl(field, query, k, nprobe, rerank_factor, combiner, None)
3341 }
3342
3343 #[cfg(feature = "sync")]
3344 #[allow(clippy::too_many_arguments)]
3345 pub(crate) fn search_dense_vector_sync_with_probe_cache(
3346 &self,
3347 field: Field,
3348 query: &[f32],
3349 k: usize,
3350 nprobe: usize,
3351 rerank_factor: f32,
3352 combiner: crate::query::MultiValueCombiner,
3353 plan_cache: &DensePlanCache,
3354 ) -> Result<Vec<VectorSearchResult>> {
3355 self.search_dense_vector_sync_impl(
3356 field,
3357 query,
3358 k,
3359 nprobe,
3360 rerank_factor,
3361 combiner,
3362 Some(plan_cache),
3363 )
3364 }
3365
3366 #[cfg(feature = "sync")]
3367 #[allow(clippy::too_many_arguments)]
3368 fn search_dense_vector_sync_impl(
3369 &self,
3370 field: Field,
3371 query: &[f32],
3372 k: usize,
3373 nprobe: usize,
3374 rerank_factor: f32,
3375 combiner: crate::query::MultiValueCombiner,
3376 plan_cache: Option<&DensePlanCache>,
3377 ) -> Result<Vec<VectorSearchResult>> {
3378 let params =
3379 self.validate_dense_search_request(field, query, nprobe, rerank_factor, combiner)?;
3380 let fetch_k = checked_dense_fetch_k(k, rerank_factor)?;
3381 if k == 0 {
3382 return Ok(Vec::new());
3383 }
3384
3385 let configured_ann_index = self.vector_indexes.get(&field.0);
3386 let lazy_flat = self.flat_vectors.get(&field.0);
3387 if configured_ann_index.is_none() && lazy_flat.is_none() {
3388 return Ok(Vec::new());
3389 }
3390
3391 if configured_ann_index.is_some() && lazy_flat.is_none() {
3392 return Err(Error::Corruption(format!(
3393 "dense ANN field {} is missing flat vector storage",
3394 field.0
3395 )));
3396 }
3397
3398 if let Some(flat) = lazy_flat
3399 && flat.dim != params.dim
3400 {
3401 return Err(Error::Corruption(format!(
3402 "dense vector field {} has schema dimension {} but flat storage dimension {}",
3403 field.0, params.dim, flat.dim
3404 )));
3405 }
3406
3407 let needs_document_aggregation = lazy_flat.is_some_and(|flat| {
3408 flat.num_vectors != flat.num_docs_with_vectors()
3409 && !matches!(combiner, crate::query::MultiValueCombiner::Max)
3410 });
3411 let ann_index = configured_ann_index;
3414
3415 let results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
3416 match index {
3418 VectorIndex::Tq { index: lazy, codec } => {
3419 let flat = lazy_flat.expect("ANN/flat pairing validated above");
3420 search_tq_segment(
3421 lazy.get(),
3422 codec,
3423 query,
3424 fetch_k.min(flat.num_docs_with_vectors()),
3425 needs_document_aggregation.then_some(combiner),
3426 field,
3427 params.dim,
3428 plan_cache.map(|cache| &cache.tq),
3429 )?
3430 }
3431 VectorIndex::IvfTq { index: lazy, codec } => {
3432 let index = lazy.get();
3433 let centroids =
3434 self.trained_vectors
3435 .centroids
3436 .get(&field.0)
3437 .ok_or_else(|| {
3438 Error::Schema(format!(
3439 "IVF-TQ index requires coarse centroids for field {}",
3440 field.0
3441 ))
3442 })?;
3443 validate_coarse_centroids(centroids, params.dim)?;
3444 let routing = self
3445 .schema
3446 .get_field_entry(field)
3447 .and_then(|entry| entry.dense_vector_config.as_ref())
3448 .map_or(crate::dsl::IvfRoutingMode::Auto, |config| {
3449 config.ivf_routing
3450 });
3451 validate_ivf_tq_ann(index, centroids, codec, params.dim, routing, field)?;
3452 let flat = lazy_flat.expect("ANN/flat pairing validated above");
3453 search_ivf_tq_segment(
3454 index,
3455 centroids,
3456 codec,
3457 query,
3458 fetch_k.min(flat.num_docs_with_vectors()),
3459 needs_document_aggregation.then_some(combiner),
3460 field,
3461 params.nprobe,
3462 routing,
3463 plan_cache.map(|cache| &cache.ivf_tq),
3464 )?
3465 }
3466 VectorIndex::BinaryIvf(_) => {
3467 Vec::new()
3469 }
3470 VectorIndex::ScannAh(lazy) => {
3471 let artifact = self
3472 .trained_vectors
3473 .scann_artifacts
3474 .get(&field.0)
3475 .ok_or_else(|| {
3476 Error::Schema(format!(
3477 "ScaNN field {} has no loaded global artifact",
3478 field.0
3479 ))
3480 })?;
3481 let flat = lazy_flat.expect("ANN/flat pairing validated above");
3482 search_scann_ah_segment(
3483 lazy.get(),
3484 artifact,
3485 query,
3486 fetch_k.min(flat.num_docs_with_vectors()),
3487 combiner,
3488 field,
3489 params.nprobe,
3490 plan_cache.map(|cache| &cache.scann),
3491 )?
3492 }
3493 VectorIndex::ScannBinary(_) => {
3494 return Err(Error::Corruption(format!(
3495 "binary ScaNN payload was attached to float field {}",
3496 field.0
3497 )));
3498 }
3499 }
3500 } else if let Some(lazy_flat) = lazy_flat {
3501 let dim = lazy_flat.dim;
3503 let n = lazy_flat.num_vectors;
3504 let quant = lazy_flat.quantization;
3505 let batch_len =
3506 bounded_vector_score_batch(lazy_flat.vector_byte_size(), DENSE_SCORE_BATCH);
3507 let mut collector = FlatDocumentCollector::new(fetch_k.min(n), combiner);
3508 let mut scores = vec![0f32; batch_len];
3509 let prepared_query = PreparedDenseScoreQuery::new(query, quant, dim, params.unit_norm)?;
3510
3511 for batch_start in (0..n).step_by(batch_len) {
3512 let batch_count = batch_len.min(n - batch_start);
3513 let batch_bytes = lazy_flat
3514 .read_vectors_batch_sync(batch_start, batch_count)
3515 .map_err(crate::Error::Io)?;
3516 let raw = batch_bytes.as_slice();
3517
3518 prepared_query.score_batch(raw, &mut scores[..batch_count])?;
3519
3520 for (i, &score) in scores.iter().enumerate().take(batch_count) {
3521 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
3522 collector.push(doc_id, ordinal, score);
3523 }
3524 }
3525
3526 return Ok(collector.into_results());
3527 } else {
3528 return Ok(Vec::new());
3529 };
3530
3531 if ann_index.is_some()
3533 && !results.is_empty()
3534 && let Some(lazy_flat) = lazy_flat
3535 {
3536 return exact_score_dense_candidate_documents_sync(
3537 &results,
3538 lazy_flat,
3539 query,
3540 params.unit_norm,
3541 combiner,
3542 k,
3543 );
3544 }
3545
3546 Ok(combine_grouped_ordinal_results(results, combiner, k))
3547 }
3548
3549 #[cfg(feature = "sync")]
3554 fn search_binary_dense_vector_sync_impl(
3555 &self,
3556 field: Field,
3557 query: &[u8],
3558 k: usize,
3559 combiner: crate::query::MultiValueCombiner,
3560 probe_cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
3561 ) -> Result<Vec<VectorSearchResult>> {
3562 let schema_dim = self.validate_binary_search_request(field, query)?;
3563 combiner.validate().map_err(Error::Query)?;
3564 if k == 0 {
3565 return Ok(Vec::new());
3566 }
3567 let t0 = crate::observe::Timer::start();
3568 if let Some(VectorIndex::ScannBinary(lazy)) = self.vector_indexes.get(&field.0) {
3569 let artifact = self
3570 .trained_vectors
3571 .scann_artifacts
3572 .get(&field.0)
3573 .ok_or_else(|| {
3574 Error::Schema(format!(
3575 "binary ScaNN field {} has no loaded global artifact",
3576 field.0
3577 ))
3578 })?;
3579 lazy.get()
3580 .validate_scann_generation(
3581 artifact.config(),
3582 artifact.generation(),
3583 artifact.artifact_id(),
3584 )
3585 .map_err(|error| {
3586 Error::Corruption(format!(
3587 "binary ScaNN generation mismatch for field {}: {error}",
3588 field.0
3589 ))
3590 })?;
3591 let config = self
3592 .schema
3593 .get_field_entry(field)
3594 .and_then(|entry| entry.binary_dense_vector_config.as_ref())
3595 .ok_or_else(|| {
3596 Error::Schema(format!(
3597 "binary ScaNN field {} has no schema configuration",
3598 field.0
3599 ))
3600 })?;
3601 let model = artifact.binary_model().map_err(Error::Io)?;
3602 let clusters = binary_scann_probe_clusters(&model, query, config.nprobe, probe_cache)?;
3603 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
3604 Error::Corruption(format!(
3605 "binary ScaNN field {} is missing flat vectors",
3606 field.0
3607 ))
3608 })?;
3609 let candidate_limit =
3610 checked_binary_combined_fetch_k(k)?.min(flat.num_docs_with_vectors());
3611 let (documents, ordinal_scores) = lazy
3612 .get()
3613 .search_binary_combined_documents(candidate_limit, query, &clusters, combiner)
3614 .map_err(|error| {
3615 Error::Corruption(format!(
3616 "invalid binary ScaNN payload for field {}: {error}",
3617 field.0
3618 ))
3619 })?;
3620 let results = exact_score_binary_candidate_document_ids_sync(
3621 documents
3622 .into_iter()
3623 .map(|candidate| candidate.doc_id)
3624 .collect(),
3625 &ordinal_scores,
3626 flat,
3627 query,
3628 schema_dim,
3629 combiner,
3630 k,
3631 )?;
3632 crate::observe::dense_l1(
3633 self.schema.index_label(),
3634 self.schema.get_field_name(field).unwrap_or("?"),
3635 "binary_scann",
3636 t0.secs(),
3637 results.len(),
3638 );
3639 return Ok(results);
3640 }
3641 if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0) {
3642 let ivf = lazy.get();
3643 let config = self
3644 .schema
3645 .get_field_entry(field)
3646 .and_then(|entry| entry.binary_dense_vector_config.as_ref())
3647 .ok_or_else(|| {
3648 Error::Schema(format!(
3649 "binary IVF field {} has no schema configuration",
3650 field.0
3651 ))
3652 })?;
3653 let quantizer = self
3654 .trained_vectors
3655 .binary_quantizers
3656 .get(&field.0)
3657 .ok_or_else(|| {
3658 Error::Schema(format!(
3659 "global binary IVF field {} has no loaded quantizer",
3660 field.0
3661 ))
3662 })?;
3663 validate_binary_ann(ivf, quantizer, config, schema_dim, field)?;
3664 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
3665 Error::Corruption(format!(
3666 "global binary IVF field {} is missing flat vector storage",
3667 field.0
3668 ))
3669 })?;
3670 let single_valued = flat.num_vectors == flat.num_docs_with_vectors();
3671 let clusters = binary_probe_clusters(
3672 quantizer,
3673 query,
3674 config.nprobe,
3675 config.ivf_routing,
3676 probe_cache,
3677 )?;
3678 let results = if !single_valued
3679 && !matches!(combiner, crate::query::MultiValueCombiner::Max)
3680 {
3681 let candidate_limit =
3682 checked_binary_combined_fetch_k(k)?.min(flat.num_docs_with_vectors());
3683 let (candidate_documents, probed_ordinal_scores) = ivf
3684 .search_binary_combined_documents(candidate_limit, query, &clusters, combiner)
3685 .map_err(|error| {
3686 Error::Corruption(format!(
3687 "invalid binary IVF payload for field {}: {error}",
3688 field.0,
3689 ))
3690 })?;
3691 exact_score_binary_candidate_document_ids_sync(
3692 candidate_documents
3693 .into_iter()
3694 .map(|candidate| candidate.doc_id)
3695 .collect(),
3696 &probed_ordinal_scores,
3697 flat,
3698 query,
3699 schema_dim,
3700 combiner,
3701 k,
3702 )?
3703 } else {
3704 let candidate_docs = if single_valued {
3705 k
3706 } else {
3707 checked_binary_combined_fetch_k(k)?
3708 }
3709 .min(flat.num_docs_with_vectors());
3710 let ann_results = if single_valued {
3711 ivf.search_binary_clusters::<false>(query, candidate_docs, &clusters)
3712 } else {
3713 ivf.search_binary_clusters::<true>(query, candidate_docs, &clusters)
3714 }
3715 .map_err(|error| {
3716 Error::Corruption(format!(
3717 "invalid binary IVF payload for field {}: {error}",
3718 field.0,
3719 ))
3720 })?;
3721 if single_valued {
3722 let ann_results = validate_binary_single_value_ann_results(ann_results, flat)?;
3723 combine_ordinal_results(ann_results, combiner, k)
3724 } else {
3725 exact_score_binary_candidate_documents_sync(
3726 &ann_results,
3727 flat,
3728 query,
3729 schema_dim,
3730 combiner,
3731 k,
3732 )?
3733 }
3734 };
3735 crate::observe::dense_l1(
3736 self.schema.index_label(),
3737 self.schema.get_field_name(field).unwrap_or("?"),
3738 "global_binary_ivf",
3739 t0.secs(),
3740 results.len(),
3741 );
3742 return Ok(results);
3743 }
3744 let lazy_flat = match self.flat_vectors.get(&field.0) {
3745 Some(f) => f,
3746 None => return Ok(Vec::new()),
3747 };
3748
3749 let dim_bits = lazy_flat.dim;
3750 let byte_len = lazy_flat.vector_byte_size();
3751 let n = lazy_flat.num_vectors;
3752
3753 if dim_bits != schema_dim {
3754 return Err(Error::Corruption(format!(
3755 "binary vector field {} has schema dimension {} but flat storage dimension {}",
3756 field.0, schema_dim, dim_bits
3757 )));
3758 }
3759
3760 if byte_len != query.len() {
3761 return Err(Error::Schema(format!(
3762 "Binary query vector byte length {} != field byte length {}",
3763 query.len(),
3764 byte_len
3765 )));
3766 }
3767
3768 let batch_len = bounded_vector_score_batch(byte_len, BINARY_SCORE_BATCH);
3769 let mut collector = FlatDocumentCollector::new(k, combiner);
3770 let mut scores = vec![0f32; batch_len];
3771
3772 for batch_start in (0..n).step_by(batch_len) {
3773 let batch_count = batch_len.min(n - batch_start);
3774 let batch_bytes = lazy_flat
3775 .read_vectors_batch_sync(batch_start, batch_count)
3776 .map_err(crate::Error::Io)?;
3777 let raw = batch_bytes.as_slice();
3778
3779 crate::structures::simd::batch_hamming_scores(
3780 query,
3781 raw,
3782 byte_len,
3783 dim_bits,
3784 &mut scores[..batch_count],
3785 );
3786
3787 for (i, &score) in scores.iter().enumerate().take(batch_count) {
3788 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
3789 collector.push(doc_id, ordinal, score);
3790 }
3791 }
3792
3793 let results = collector.into_results();
3794
3795 crate::observe::dense_l1(
3796 self.schema.index_label(),
3797 self.schema.get_field_name(field).unwrap_or("?"),
3798 "binary_flat",
3799 t0.secs(),
3800 results.len(),
3801 );
3802 Ok(results)
3803 }
3804
3805 #[cfg(feature = "sync")]
3806 pub fn search_binary_dense_vector_sync(
3807 &self,
3808 field: Field,
3809 query: &[u8],
3810 k: usize,
3811 combiner: crate::query::MultiValueCombiner,
3812 ) -> Result<Vec<VectorSearchResult>> {
3813 self.search_binary_dense_vector_sync_impl(field, query, k, combiner, None)
3814 }
3815
3816 #[cfg(feature = "sync")]
3817 pub(crate) fn search_binary_dense_vector_sync_with_probe_cache(
3818 &self,
3819 field: Field,
3820 query: &[u8],
3821 k: usize,
3822 combiner: crate::query::MultiValueCombiner,
3823 probe_cache: &std::sync::Mutex<Option<crate::structures::IvfProbePlan>>,
3824 ) -> Result<Vec<VectorSearchResult>> {
3825 self.search_binary_dense_vector_sync_impl(field, query, k, combiner, Some(probe_cache))
3826 }
3827}
3828
3829#[cfg(test)]
3830mod dense_search_safety_tests {
3831 use super::*;
3832
3833 #[test]
3834 fn dense_fetch_count_rejects_non_finite_and_unbounded_factors() {
3835 for factor in [
3836 f32::NAN,
3837 f32::INFINITY,
3838 f32::NEG_INFINITY,
3839 0.0,
3840 0.5,
3841 2.01,
3842 MAX_DENSE_RERANK_FACTOR + 1.0,
3843 ] {
3844 assert!(
3845 checked_dense_fetch_k(10, factor).is_err(),
3846 "factor={factor}"
3847 );
3848 }
3849 }
3850
3851 fn values_as_bytes<T>(values: &[T]) -> &[u8] {
3852 unsafe {
3853 std::slice::from_raw_parts(values.as_ptr() as *const u8, std::mem::size_of_val(values))
3854 }
3855 }
3856
3857 fn assert_prepared_dense_scores_match_legacy(
3858 quantization: DenseVectorQuantization,
3859 raw: &[u8],
3860 unit_norm: bool,
3861 ) {
3862 const DIM: usize = 4;
3863 const VECTOR_COUNT: usize = 4;
3864 let query = [0.25, -0.5, 0.75, 1.0];
3865 let mut expected = [0.0; VECTOR_COUNT];
3866 SegmentReader::score_quantized_batch_legacy(
3867 &query,
3868 raw,
3869 quantization,
3870 DIM,
3871 &mut expected,
3872 unit_norm,
3873 )
3874 .unwrap();
3875
3876 let prepared = PreparedDenseScoreQuery::new(&query, quantization, DIM, unit_norm).unwrap();
3877 let vector_bytes = DIM
3878 * match quantization {
3879 DenseVectorQuantization::F32 => std::mem::size_of::<f32>(),
3880 DenseVectorQuantization::F16 => std::mem::size_of::<u16>(),
3881 DenseVectorQuantization::UInt8 => 1,
3882 DenseVectorQuantization::Binary => unreachable!(),
3883 };
3884 let split = 2 * vector_bytes;
3885 let mut actual = [0.0; VECTOR_COUNT];
3886 prepared
3887 .score_batch(&raw[..split], &mut actual[..2])
3888 .unwrap();
3889 prepared
3890 .score_batch(&raw[split..], &mut actual[2..])
3891 .unwrap();
3892
3893 assert_eq!(
3894 actual.map(f32::to_bits),
3895 expected.map(f32::to_bits),
3896 "quantization={quantization:?}, unit_norm={unit_norm}"
3897 );
3898 }
3899
3900 #[test]
3901 fn prepared_dense_query_matches_legacy_scoring_across_batches() {
3902 let vectors_f32 = [
3903 0.5, -0.25, 0.75, 1.0, -1.0, 0.5, 0.25, 0.125, 0.0, 0.0, 0.0, 0.0, 0.75, 0.5, -0.5,
3904 -0.25,
3905 ];
3906 let vectors_f16: Vec<u16> = vectors_f32
3907 .iter()
3908 .map(|&value| crate::structures::simd::f32_to_f16(value))
3909 .collect();
3910 let vectors_u8 = [
3911 255, 96, 224, 160, 0, 192, 144, 128, 128, 128, 128, 128, 224, 192, 64, 96,
3912 ];
3913
3914 for unit_norm in [false, true] {
3915 assert_prepared_dense_scores_match_legacy(
3916 DenseVectorQuantization::F32,
3917 values_as_bytes(&vectors_f32),
3918 unit_norm,
3919 );
3920 assert_prepared_dense_scores_match_legacy(
3921 DenseVectorQuantization::F16,
3922 values_as_bytes(&vectors_f16),
3923 unit_norm,
3924 );
3925 assert_prepared_dense_scores_match_legacy(
3926 DenseVectorQuantization::UInt8,
3927 &vectors_u8,
3928 unit_norm,
3929 );
3930 }
3931 }
3932
3933 #[test]
3934 fn prepared_dense_query_preserves_scoring_validation_errors() {
3935 assert!(matches!(
3936 PreparedDenseScoreQuery::new(&[1.0], DenseVectorQuantization::F32, 2, false).err(),
3937 Some(Error::Query(_))
3938 ));
3939 assert!(matches!(
3940 PreparedDenseScoreQuery::new(&[1.0], DenseVectorQuantization::Binary, 1, false).err(),
3941 Some(Error::InvalidFieldType { .. })
3942 ));
3943
3944 let query = [1.0, 2.0];
3945 let prepared =
3946 PreparedDenseScoreQuery::new(&query, DenseVectorQuantization::F32, 2, false).unwrap();
3947 let mut scores = [0.0];
3948 assert!(matches!(
3949 prepared.score_batch(&[0; 7], &mut scores),
3950 Err(Error::Corruption(_))
3951 ));
3952 }
3953
3954 #[test]
3955 fn flat_document_collector_does_not_let_one_multivalue_doc_crowd_out_others() {
3956 let mut collector = FlatDocumentCollector::new(2, crate::query::MultiValueCombiner::Max);
3957 collector.push(1, 0, 1.0);
3958 collector.push(1, 1, 0.9);
3959 collector.push(2, 0, 0.8);
3960
3961 let results = collector.into_results();
3962 assert_eq!(
3963 results
3964 .iter()
3965 .map(|result| result.doc_id)
3966 .collect::<Vec<_>>(),
3967 vec![1, 2]
3968 );
3969 assert_eq!(results[0].ordinals.len(), 2);
3970 }
3971
3972 #[test]
3973 fn flat_document_collector_evicts_by_score_then_doc_id() {
3974 let mut collector = FlatDocumentCollector::new(2, crate::query::MultiValueCombiner::Max);
3975 collector.push(1, 0, 0.5);
3976 collector.push(3, 0, 0.8);
3977 collector.push(2, 0, 0.9);
3978 let results = collector.into_results();
3979 assert_eq!(
3980 results
3981 .iter()
3982 .map(|result| result.doc_id)
3983 .collect::<Vec<_>>(),
3984 vec![2, 3]
3985 );
3986
3987 let mut tied = FlatDocumentCollector::new(1, crate::query::MultiValueCombiner::Max);
3988 tied.push(2, 0, 1.0);
3989 tied.push(1, 0, 1.0);
3990 let results = tied.into_results();
3991 assert_eq!(results[0].doc_id, 1);
3992 }
3993
3994 #[test]
3995 fn dense_fetch_count_rounds_up_and_detects_overflow() {
3996 assert_eq!(checked_dense_fetch_k(3, 1.5).unwrap(), 5);
3997 assert_eq!(checked_dense_fetch_k(10_000, 2.0).unwrap(), 20_000);
3998 assert!(checked_dense_fetch_k(10_001, 2.0).is_err());
3999 assert!(checked_dense_fetch_k(usize::MAX, 2.0).is_err());
4000 }
4001
4002 #[test]
4003 fn binary_combined_fetch_count_uses_shared_bounded_oversampling() {
4004 assert_eq!(checked_binary_combined_fetch_k(3).unwrap(), 6);
4005 assert_eq!(checked_binary_combined_fetch_k(10_000).unwrap(), 20_000);
4006 assert_eq!(checked_binary_combined_fetch_k(10_001).unwrap(), 20_000);
4007 assert_eq!(checked_binary_combined_fetch_k(20_000).unwrap(), 20_000);
4008 assert!(checked_binary_combined_fetch_k(20_001).is_err());
4009 assert!(checked_binary_combined_fetch_k(usize::MAX).is_err());
4010 }
4011
4012 #[cfg(feature = "native")]
4013 #[test]
4014 fn legacy_ivf_tq_generation_is_rejected_while_opening() {
4015 use crate::directories::OwnedBytes;
4016 use crate::dsl::IvfRoutingMode;
4017 use crate::segment::ann_disk::{AnnDiskIndex, AnnKind};
4018
4019 let centroids = CoarseCentroids {
4020 num_clusters: 1,
4021 dim: 2,
4022 centroids: vec![1.0, 0.0],
4023 version: 7,
4024 soar_config: None,
4025 routing_index: None,
4026 };
4027 let mut build_centroids = centroids.clone();
4028 build_centroids.version =
4029 crate::structures::mark_ivf_tq_cosine_generation(build_centroids.version);
4030 let mut bytes = crate::segment::ann_build::build_ivf_tq(
4031 2,
4032 IvfRoutingMode::Flat,
4033 &build_centroids,
4034 &[(0, 0)],
4035 &[1.0, 0.0],
4036 )
4037 .unwrap();
4038 bytes[24..32].copy_from_slice(¢roids.version.to_le_bytes());
4041 let error = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::IvfTq, 1)
4042 .err()
4043 .expect("legacy IVF-TQ payload must fail while opening")
4044 .to_string();
4045 assert!(error.contains("unsupported legacy generation"), "{error}");
4046 }
4047
4048 #[test]
4049 fn rerank_batch_is_capped_by_actual_candidate_vectors() {
4050 assert_eq!(bounded_rerank_batch(3_072, DENSE_SCORE_BATCH, 20), 20);
4051 assert_eq!(
4052 bounded_rerank_batch(3_072, DENSE_SCORE_BATCH, 10_000),
4053 MAX_VECTOR_SCORE_BATCH_BYTES / 3_072
4054 );
4055 assert_eq!(bounded_rerank_batch(3_072, DENSE_SCORE_BATCH, 0), 1);
4056 }
4057
4058 #[test]
4059 fn file_ranges_reject_overflow_and_truncation() {
4060 assert_eq!(checked_file_range(4, 3, 7, "test").unwrap(), 4..7);
4061 assert!(checked_file_range(u64::MAX, 1, u64::MAX, "test").is_err());
4062 assert!(checked_file_range(5, 3, 7, "test").is_err());
4063 }
4064
4065 #[test]
4066 fn shared_tq_plan_cache_rebuilds_for_divergent_query_clones() {
4067 let codec = crate::structures::TqCodec::new(4);
4068 let cache = std::sync::Mutex::new(None);
4069 let original_query = vec![1.0, 2.0, 3.0, 4.0];
4070
4071 let original =
4072 cached_tq_query_plan(&codec, &original_query, Some(&cache)).expect("build plan");
4073 let reused =
4074 cached_tq_query_plan(&codec, &original_query, Some(&cache)).expect("reuse plan");
4075 assert!(
4076 std::sync::Arc::ptr_eq(&original, &reused),
4077 "unchanged queries must share their plan across segments"
4078 );
4079
4080 let mut divergent_clone = original_query.clone();
4081 divergent_clone[0] = -1.0;
4082 let rebuilt =
4083 cached_tq_query_plan(&codec, &divergent_clone, Some(&cache)).expect("rebuild plan");
4084 assert!(
4085 !std::sync::Arc::ptr_eq(&original, &rebuilt),
4086 "a clone with a mutated vector must not reuse stale LUTs"
4087 );
4088 assert!(rebuilt.matches_query(&divergent_clone));
4089 assert!(!rebuilt.matches_query(&original_query));
4090 }
4091
4092 #[test]
4093 fn candidate_vector_reads_coalesce_contiguous_values() {
4094 let mut runs = Vec::new();
4095 plan_vector_read_runs(&[3, 4, 5, 9, 12, 13], &mut runs).unwrap();
4096 assert_eq!(runs.len(), 3);
4097 assert!(matches!(
4098 runs.as_slice(),
4099 [
4100 VectorReadRun {
4101 buffer_start: 0,
4102 flat_start: 3,
4103 count: 3,
4104 },
4105 VectorReadRun {
4106 buffer_start: 3,
4107 flat_start: 9,
4108 count: 1,
4109 },
4110 VectorReadRun {
4111 buffer_start: 4,
4112 flat_start: 12,
4113 count: 2,
4114 },
4115 ]
4116 ));
4117 assert!(plan_vector_read_runs(&[3, 3], &mut runs).is_err());
4118 }
4119
4120 #[tokio::test]
4121 async fn binary_single_value_ann_fast_path_validates_and_deduplicates() {
4122 use crate::directories::{FileHandle, OwnedBytes};
4123 use crate::segment::FlatVectorData;
4124
4125 let mut encoded = Vec::new();
4126 FlatVectorData::serialize_binary_from_bits_streaming(
4127 8,
4128 &[0x0f, 0xf0],
4129 &[(1, 0), (3, 2)],
4130 &mut encoded,
4131 )
4132 .unwrap();
4133 let flat = LazyFlatVectorData::open_with_doc_limit(
4134 FileHandle::from_bytes(OwnedBytes::new(encoded)),
4135 Some(4),
4136 )
4137 .await
4138 .unwrap();
4139 assert_eq!(flat.num_vectors, flat.num_docs_with_vectors());
4140
4141 let validated = validate_binary_single_value_ann_results(
4142 vec![(3, 2, 0.9), (1, 0, 0.8), (3, 2, 0.7)],
4143 &flat,
4144 )
4145 .unwrap();
4146 assert_eq!(validated, vec![(3, 2, 0.9), (1, 0, 0.8)]);
4147
4148 assert!(matches!(
4149 validate_binary_single_value_ann_results(vec![(2, 0, 1.0)], &flat),
4150 Err(Error::Corruption(_))
4151 ));
4152 assert!(matches!(
4153 validate_binary_single_value_ann_results(vec![(3, 0, 1.0)], &flat),
4154 Err(Error::Corruption(_))
4155 ));
4156 }
4157
4158 #[tokio::test]
4159 async fn multivalue_ann_rerank_streams_past_document_candidate_cap() {
4160 use crate::directories::{FileHandle, OwnedBytes};
4161 use crate::segment::FlatVectorData;
4162
4163 const VALUES: usize = MAX_DENSE_CANDIDATES_PER_SEGMENT + 1;
4164 let mut encoded = Vec::new();
4165 let vectors = vec![1.0f32; VALUES];
4166 let doc_ids: Vec<_> = (0..VALUES).map(|ordinal| (0, ordinal as u16)).collect();
4167 FlatVectorData::serialize_binary_from_flat_streaming(
4168 1,
4169 &vectors,
4170 &doc_ids,
4171 DenseVectorQuantization::F32,
4172 &mut encoded,
4173 )
4174 .unwrap();
4175 let flat = LazyFlatVectorData::open_with_doc_limit(
4176 FileHandle::from_bytes(OwnedBytes::new(encoded)),
4177 Some(1),
4178 )
4179 .await
4180 .unwrap();
4181
4182 let (results, stats) = exact_score_dense_candidate_documents(
4183 &[(0, 0, 0.0)],
4184 &flat,
4185 &[1.0],
4186 false,
4187 crate::query::MultiValueCombiner::Max,
4188 1,
4189 )
4190 .await
4191 .unwrap();
4192 assert_eq!(stats.vector_count, VALUES);
4193 assert_eq!(results.len(), 1);
4194 assert_eq!(results[0].ordinals.len(), VALUES);
4195 assert!((results[0].score - 1.0).abs() < 1e-5);
4196 }
4197}