1pub(crate) mod bmp;
4pub(crate) mod candidate_lookup;
5pub(crate) mod loader;
6mod term_expansion;
7pub(crate) use term_expansion::ExpandedPosting;
8mod types;
9
10pub use bmp::{BmpDimStats, BmpIndex};
11#[cfg(feature = "native")]
12pub(crate) use types::DimRawData;
13pub use types::{SparseIndex, VectorIndex, VectorOrdinals, VectorSearchResult};
14
15const MAX_PREFIX_TERMS: usize = 1_024;
19const MAX_PREFIX_POSTINGS: u64 = 5_000_000;
20const MAX_DENSE_CANDIDATES_PER_SEGMENT: usize = 20_000;
24const DENSE_SCORE_BATCH: usize = 4_096;
26const BINARY_SCORE_BATCH: usize = 8_192;
27const MAX_VECTOR_SCORE_BATCH_BYTES: usize = 8 * 1024 * 1024;
28
29#[derive(Debug, Clone, Default)]
35pub struct SegmentMemoryStats {
36 pub segment_id: u128,
38 pub num_docs: u32,
40 pub term_dict_cache_bytes: usize,
42 pub posting_integrity_heap_bytes: usize,
44 pub deletion_bytes: usize,
46 pub row_stats_heap_bytes: usize,
48 pub row_stats_file_backed_bytes: u64,
49 pub fast_field_metadata_heap_bytes: usize,
51 pub store_cache_bytes: usize,
53 pub sparse_heap_bytes: usize,
55 pub dense_heap_bytes: usize,
57 pub term_bloom_file_bytes: u64,
59 pub sparse_file_backed_bytes: u64,
61 pub dense_file_backed_bytes: u64,
63 pub pinned_metadata_bytes: u64,
65 pub pin_intended_bytes: u64,
68 pub sparse_pinned_metadata_bytes: u64,
70 pub sparse_pin_intended_bytes: u64,
72 pub dense_pinned_metadata_bytes: u64,
74 pub dense_pin_intended_bytes: u64,
76}
77
78impl SegmentMemoryStats {
79 pub fn estimated_heap_bytes(&self) -> usize {
81 self.deletion_bytes
82 + self.row_stats_heap_bytes
83 + self.fast_field_metadata_heap_bytes
84 + self.term_dict_cache_bytes
85 + self.posting_integrity_heap_bytes
86 + self.store_cache_bytes
87 + self.sparse_heap_bytes
88 + self.dense_heap_bytes
89 }
90
91 pub fn file_backed_bytes(&self) -> u64 {
95 self.term_bloom_file_bytes
96 .saturating_add(self.row_stats_file_backed_bytes)
97 .saturating_add(self.sparse_file_backed_bytes)
98 .saturating_add(self.dense_file_backed_bytes)
99 }
100}
101
102pub(crate) use types::SparseProbeBudget;
103
104use std::cmp::Ordering;
105use std::collections::BinaryHeap;
106use std::sync::Arc;
107
108use rustc_hash::{FxHashMap, FxHashSet};
109
110use super::vector_data::LazyFlatVectorData;
111use crate::directories::{Directory, FileHandle, OwnedBytes};
112use crate::dsl::{DenseVectorQuantization, Document, Field, Schema};
113use crate::observe::DenseAnnScanStats;
114use crate::query::{MAX_DENSE_NPROBE, MAX_DENSE_RERANK_FACTOR};
115use crate::structures::{
116 AsyncSSTableReader, BlockPostingList, CoarseCentroids, SSTableStats, TermInfo,
117};
118use crate::{DocId, Error, Result};
119
120use super::store::{AsyncStoreReader, RawStoreBlock};
121use super::types::{SegmentFiles, SegmentId, SegmentMeta};
122
123pub(crate) fn combine_ordinal_results(
134 raw: impl IntoIterator<Item = (u32, u16, f32)>,
135 combiner: crate::query::MultiValueCombiner,
136 limit: usize,
137) -> Vec<VectorSearchResult> {
138 let mut collected: Vec<(u32, u16, f32)> = raw.into_iter().collect();
139
140 let num_raw = collected.len();
141 if log::log_enabled!(log::Level::Debug) {
142 let mut ids: Vec<u32> = collected.iter().map(|(d, _, _)| *d).collect();
143 ids.sort_unstable();
144 ids.dedup();
145 log::debug!(
146 "combine_ordinal_results: {} raw entries, {} unique docs, combiner={:?}, limit={}",
147 num_raw,
148 ids.len(),
149 combiner,
150 limit
151 );
152 }
153
154 let all_single = collected.iter().all(|&(_, ord, _)| ord == 0);
156 if all_single {
157 let mut results: Vec<VectorSearchResult> = collected
158 .into_iter()
159 .map(|(doc_id, _, score)| VectorSearchResult::single(doc_id, score))
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 return results;
168 }
169
170 collected.sort_by_key(|&(doc_id, _, _)| doc_id);
173 let mut results: Vec<VectorSearchResult> = Vec::new();
174 let mut index = 0;
175 while index < collected.len() {
176 let doc_id = collected[index].0;
177 let mut ordinals = super::VectorOrdinals::new();
178 while index < collected.len() && collected[index].0 == doc_id {
179 let (_, ordinal, score) = collected[index];
180 ordinals.push((ordinal as u32, score));
181 index += 1;
182 }
183 let combined_score = combiner.combine(&ordinals);
184 results.push(VectorSearchResult::with_ordinals(
185 doc_id as DocId,
186 combined_score,
187 ordinals,
188 ));
189 }
190 results.sort_unstable_by(|a, b| {
191 b.score
192 .total_cmp(&a.score)
193 .then_with(|| a.doc_id.cmp(&b.doc_id))
194 });
195 results.truncate(limit);
196 results
197}
198
199struct HeapVectorResult(VectorSearchResult);
204
205impl PartialEq for HeapVectorResult {
206 fn eq(&self, other: &Self) -> bool {
207 self.0.score.to_bits() == other.0.score.to_bits() && self.0.doc_id == other.0.doc_id
208 }
209}
210
211impl Eq for HeapVectorResult {}
212
213impl Ord for HeapVectorResult {
214 fn cmp(&self, other: &Self) -> Ordering {
215 other
218 .0
219 .score
220 .total_cmp(&self.0.score)
221 .then_with(|| self.0.doc_id.cmp(&other.0.doc_id))
222 }
223}
224
225impl PartialOrd for HeapVectorResult {
226 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
227 Some(self.cmp(other))
228 }
229}
230
231struct FlatDocumentCollector {
235 heap: BinaryHeap<HeapVectorResult>,
236 limit: usize,
237 combiner: crate::query::MultiValueCombiner,
238 current_doc: Option<DocId>,
239 current_ordinals: super::VectorOrdinals,
240}
241
242impl FlatDocumentCollector {
243 fn new(limit: usize, combiner: crate::query::MultiValueCombiner) -> Self {
244 Self {
245 heap: BinaryHeap::with_capacity(limit.min(8 * 1024)),
246 limit,
247 combiner,
248 current_doc: None,
249 current_ordinals: super::VectorOrdinals::new(),
250 }
251 }
252
253 fn push(&mut self, doc_id: DocId, ordinal: u16, score: f32) {
254 if self.current_doc.is_some_and(|current| current != doc_id) {
255 self.finish_current();
256 }
257 self.current_doc = Some(doc_id);
258 self.current_ordinals.push((ordinal as u32, score));
259 }
260
261 fn finish_current(&mut self) {
262 let Some(doc_id) = self.current_doc.take() else {
263 return;
264 };
265 let score = self.combiner.combine(&self.current_ordinals);
266 let should_retain = self.heap.len() < self.limit
267 || self.heap.peek().is_some_and(|worst| {
268 HeapVectorResult(VectorSearchResult::with_ordinals(
269 doc_id,
270 score,
271 super::VectorOrdinals::new(),
272 ))
273 .cmp(worst)
274 .is_lt()
275 });
276
277 if !should_retain {
278 self.current_ordinals.clear();
282 return;
283 }
284
285 let ordinals = std::mem::take(&mut self.current_ordinals);
286 let entry = HeapVectorResult(VectorSearchResult::with_ordinals(doc_id, score, ordinals));
287 if self.heap.len() < self.limit {
288 self.heap.push(entry);
289 } else if let Some(mut worst) = self.heap.peek_mut() {
290 let mut evicted = std::mem::replace(&mut worst.0, entry.0);
293 evicted.ordinals.clear();
294 self.current_ordinals = evicted.ordinals;
295 }
296 }
297
298 fn into_results(mut self) -> Vec<VectorSearchResult> {
299 self.finish_current();
300 let mut results: Vec<_> = self.heap.into_iter().map(|entry| entry.0).collect();
301 results.sort_unstable_by(|a, b| {
302 b.score
303 .total_cmp(&a.score)
304 .then_with(|| a.doc_id.cmp(&b.doc_id))
305 });
306 results
307 }
308}
309
310fn combine_grouped_ordinal_results(
313 raw: impl IntoIterator<Item = RawVectorCandidate>,
314 combiner: crate::query::MultiValueCombiner,
315 limit: usize,
316) -> Vec<VectorSearchResult> {
317 let mut collector = FlatDocumentCollector::new(limit, combiner);
318 for (doc_id, ordinal, score) in raw {
319 collector.push(doc_id, ordinal, score);
320 }
321 collector.into_results()
322}
323
324#[derive(Clone, Copy)]
325struct DenseSearchParams {
326 dim: usize,
327 nprobe: usize,
328 unit_norm: bool,
329}
330
331struct PreparedDenseScoreQuery<'a> {
338 query: &'a [f32],
339 query_f16: Vec<u16>,
340 inv_norm_q: f32,
341 quantization: DenseVectorQuantization,
342 dim: usize,
343 unit_norm: bool,
344}
345
346impl<'a> PreparedDenseScoreQuery<'a> {
347 fn new(
348 query: &'a [f32],
349 quantization: DenseVectorQuantization,
350 dim: usize,
351 unit_norm: bool,
352 ) -> Result<Self> {
353 use crate::structures::simd;
354
355 if query.len() != dim {
356 return Err(Error::Query(format!(
357 "dense SIMD query dimension {} does not match vector dimension {dim}",
358 query.len()
359 )));
360 }
361 if quantization == DenseVectorQuantization::Binary {
362 return Err(Error::InvalidFieldType {
363 expected: "non-binary dense vector".to_string(),
364 got: "binary dense vector".to_string(),
365 });
366 }
367
368 let norm_q_sq = simd::dot_product_f32(query, query, dim);
369 let inv_norm_q = if norm_q_sq < f32::EPSILON {
370 0.0
371 } else {
372 simd::fast_inv_sqrt(norm_q_sq)
373 };
374 let query_f16 = if quantization == DenseVectorQuantization::F16 {
375 query.iter().map(|&value| simd::f32_to_f16(value)).collect()
376 } else {
377 Vec::new()
378 };
379
380 Ok(Self {
381 query,
382 query_f16,
383 inv_norm_q,
384 quantization,
385 dim,
386 unit_norm,
387 })
388 }
389
390 fn score_batch(&self, raw: &[u8], scores: &mut [f32]) -> Result<()> {
391 use crate::structures::simd;
392
393 let element_size = match self.quantization {
394 DenseVectorQuantization::F32 => std::mem::size_of::<f32>(),
395 DenseVectorQuantization::F16 => std::mem::size_of::<u16>(),
396 DenseVectorQuantization::UInt8 => 1,
397 DenseVectorQuantization::Binary => {
398 return Err(Error::InvalidFieldType {
399 expected: "non-binary dense vector".to_string(),
400 got: "binary dense vector".to_string(),
401 });
402 }
403 };
404 let required_bytes = scores
405 .len()
406 .checked_mul(self.dim)
407 .and_then(|elements| elements.checked_mul(element_size))
408 .ok_or_else(|| Error::Corruption("dense vector batch byte length overflow".into()))?;
409 if raw.len() < required_bytes {
410 return Err(Error::Corruption(format!(
411 "dense vector batch is truncated: need {required_bytes} bytes, got {}",
412 raw.len()
413 )));
414 }
415 if self.quantization == DenseVectorQuantization::F16
416 && required_bytes > 0
417 && !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>())
418 {
419 return Err(Error::Corruption(
420 "f16 vector data is not 2-byte aligned".to_string(),
421 ));
422 }
423 if self.quantization == DenseVectorQuantization::F32
424 && !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>())
425 {
426 return Err(Error::Corruption(
427 "f32 vector data is not 4-byte aligned".to_string(),
428 ));
429 }
430
431 if self.dim == 0 || scores.is_empty() {
435 return Ok(());
436 }
437
438 match (self.quantization, self.unit_norm) {
439 (DenseVectorQuantization::F32, false) => {
440 let num_floats = scores.len() * self.dim;
441 let vectors: &[f32] =
442 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
443 simd::batch_cosine_scores_precomp(
444 self.query,
445 vectors,
446 self.dim,
447 scores,
448 self.inv_norm_q,
449 );
450 }
451 (DenseVectorQuantization::F32, true) => {
452 let num_floats = scores.len() * self.dim;
453 let vectors: &[f32] =
454 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
455 simd::batch_dot_scores_precomp(
456 self.query,
457 vectors,
458 self.dim,
459 scores,
460 self.inv_norm_q,
461 );
462 }
463 (DenseVectorQuantization::F16, false) => {
464 simd::batch_cosine_scores_f16_precomp(
465 &self.query_f16,
466 raw,
467 self.dim,
468 scores,
469 self.inv_norm_q,
470 );
471 }
472 (DenseVectorQuantization::F16, true) => {
473 simd::batch_dot_scores_f16_precomp(
474 &self.query_f16,
475 raw,
476 self.dim,
477 scores,
478 self.inv_norm_q,
479 );
480 }
481 (DenseVectorQuantization::UInt8, false) => {
482 simd::batch_cosine_scores_u8_precomp(
483 self.query,
484 raw,
485 self.dim,
486 scores,
487 self.inv_norm_q,
488 );
489 }
490 (DenseVectorQuantization::UInt8, true) => {
491 simd::batch_dot_scores_u8_precomp(
492 self.query,
493 raw,
494 self.dim,
495 scores,
496 self.inv_norm_q,
497 );
498 }
499 (DenseVectorQuantization::Binary, _) => unreachable!("validated during preparation"),
500 }
501 Ok(())
502 }
503}
504
505fn checked_dense_fetch_k(k: usize, rerank_factor: f32) -> Result<usize> {
507 if !rerank_factor.is_finite() || !(1.0..=MAX_DENSE_RERANK_FACTOR).contains(&rerank_factor) {
508 return Err(Error::Query(format!(
509 "dense rerank_factor must be finite and in [1, {MAX_DENSE_RERANK_FACTOR}], got {rerank_factor}"
510 )));
511 }
512
513 let fetch = (k as f64) * (rerank_factor as f64);
514 if !fetch.is_finite()
515 || fetch > usize::MAX as f64
516 || fetch > MAX_DENSE_CANDIDATES_PER_SEGMENT as f64
517 {
518 return Err(Error::Query(format!(
519 "dense candidate count exceeds the per-segment maximum of \
520 {MAX_DENSE_CANDIDATES_PER_SEGMENT}: k={k}, rerank_factor={rerank_factor}"
521 )));
522 }
523 Ok(fetch.ceil() as usize)
524}
525
526#[inline]
532fn checked_binary_combined_fetch_k(k: usize) -> Result<usize> {
533 if k > MAX_DENSE_CANDIDATES_PER_SEGMENT {
534 return Err(Error::Query(format!(
535 "binary dense result count exceeds the per-segment maximum of \
536 {MAX_DENSE_CANDIDATES_PER_SEGMENT}: k={k}"
537 )));
538 }
539 Ok(crate::query::max_candidate_limit(k).min(MAX_DENSE_CANDIDATES_PER_SEGMENT))
540}
541
542#[inline]
543fn bounded_vector_score_batch(vector_byte_size: usize, preferred: usize) -> usize {
544 preferred.min((MAX_VECTOR_SCORE_BATCH_BYTES / vector_byte_size.max(1)).max(1))
545}
546
547#[inline]
548fn bounded_rerank_batch(vector_byte_size: usize, preferred: usize, vector_count: usize) -> usize {
549 bounded_vector_score_batch(vector_byte_size, preferred).min(vector_count.max(1))
550}
551
552fn checked_file_range(
553 offset: u64,
554 length: u64,
555 file_length: u64,
556 description: &str,
557) -> Result<std::ops::Range<u64>> {
558 let end = offset
559 .checked_add(length)
560 .ok_or_else(|| Error::Corruption(format!("{description} byte range overflows u64")))?;
561 if end > file_length {
562 return Err(Error::Corruption(format!(
563 "{description} byte range {offset}..{end} exceeds file length {file_length}"
564 )));
565 }
566 Ok(offset..end)
567}
568
569type RawVectorCandidate = (u32, u16, f32);
570type CandidateVectorRef = (DocId, u16, usize); #[derive(Clone, Copy)]
573struct CandidateDocumentRange {
574 doc_id: DocId,
575 start: usize,
576 end: usize,
577}
578
579struct AnnCandidateDocuments {
580 ranges: Vec<CandidateDocumentRange>,
581 vector_count: usize,
582}
583
584fn ann_candidate_document_ranges(
593 ann_results: &[RawVectorCandidate],
594 flat: &LazyFlatVectorData,
595) -> Result<AnnCandidateDocuments> {
596 ann_candidate_document_ranges_from_ids(ann_results.iter().map(|candidate| candidate.0), flat)
597}
598
599fn ann_candidate_document_ranges_from_ids(
600 doc_ids: impl IntoIterator<Item = DocId>,
601 flat: &LazyFlatVectorData,
602) -> Result<AnnCandidateDocuments> {
603 let mut candidate_docs: Vec<DocId> = doc_ids.into_iter().collect();
604 candidate_docs.sort_unstable();
605 candidate_docs.dedup();
606
607 let mut ranges = Vec::with_capacity(candidate_docs.len());
608 let mut vector_count = 0usize;
609 for doc_id in candidate_docs {
610 let (start, count) = flat.flat_indexes_for_doc_range(doc_id);
611 if count == 0 {
612 return Err(Error::Corruption(format!(
613 "ANN candidate document {doc_id} is missing from flat vector storage"
614 )));
615 }
616 vector_count = vector_count
617 .checked_add(count)
618 .ok_or_else(|| Error::Query("ANN candidate vector expansion overflow".to_string()))?;
619 let end = start
620 .checked_add(count)
621 .ok_or_else(|| Error::Corruption("flat vector range overflow".to_string()))?;
622 if end > flat.num_vectors {
623 return Err(Error::Corruption(format!(
624 "flat vector range {start}..{end} for document {doc_id} exceeds {} vectors",
625 flat.num_vectors
626 )));
627 }
628 ranges.push(CandidateDocumentRange { doc_id, start, end });
629 }
630 Ok(AnnCandidateDocuments {
631 ranges,
632 vector_count,
633 })
634}
635
636fn validate_binary_single_value_ann_results(
644 ann_results: Vec<RawVectorCandidate>,
645 flat: &LazyFlatVectorData,
646) -> Result<Vec<RawVectorCandidate>> {
647 let mut seen_docs = FxHashSet::default();
648 let mut validated = Vec::with_capacity(ann_results.len());
649 for (doc_id, ordinal, score) in ann_results {
650 let (start, count) = flat.flat_indexes_for_doc_range(doc_id);
651 if count == 0 {
652 return Err(Error::Corruption(format!(
653 "ANN candidate document {doc_id} is missing from flat vector storage"
654 )));
655 }
656 if count != 1 {
657 return Err(Error::Corruption(format!(
658 "binary ANN single-valued candidate document {doc_id} has {count} flat vectors"
659 )));
660 }
661 let (stored_doc_id, stored_ordinal) = flat.get_doc_id(start);
662 if stored_doc_id != doc_id {
663 return Err(Error::Corruption(format!(
664 "flat vector doc map is not contiguous for document {doc_id}"
665 )));
666 }
667 if stored_ordinal != ordinal {
668 return Err(Error::Corruption(format!(
669 "binary ANN candidate document {doc_id} ordinal {ordinal} is missing from flat vector storage"
670 )));
671 }
672 if seen_docs.insert(doc_id) {
673 validated.push((doc_id, ordinal, score));
674 }
675 }
676 Ok(validated)
677}
678
679struct CandidateVectorCursor<'a> {
680 ranges: &'a [CandidateDocumentRange],
681 range_index: usize,
682 flat_index: usize,
683}
684
685impl<'a> CandidateVectorCursor<'a> {
686 fn new(ranges: &'a [CandidateDocumentRange]) -> Self {
687 Self {
688 ranges,
689 range_index: 0,
690 flat_index: ranges.first().map_or(0, |range| range.start),
691 }
692 }
693
694 fn fill_batch(
698 &mut self,
699 flat: &LazyFlatVectorData,
700 batch: &mut Vec<CandidateVectorRef>,
701 limit: usize,
702 ) -> Result<bool> {
703 batch.clear();
704 while batch.len() < limit && self.range_index < self.ranges.len() {
705 let range = self.ranges[self.range_index];
706 if self.flat_index == range.end {
707 self.range_index += 1;
708 if let Some(next) = self.ranges.get(self.range_index) {
709 self.flat_index = next.start;
710 }
711 continue;
712 }
713 let (stored_doc_id, ordinal) = flat.get_doc_id(self.flat_index);
714 if stored_doc_id != range.doc_id {
715 return Err(Error::Corruption(format!(
716 "flat vector doc map is not contiguous for document {}",
717 range.doc_id
718 )));
719 }
720 batch.push((range.doc_id, ordinal, self.flat_index));
721 self.flat_index += 1;
722 }
723 Ok(!batch.is_empty())
724 }
725}
726
727#[derive(Clone, Copy)]
728struct VectorReadRun {
729 buffer_start: usize,
730 flat_start: usize,
731 count: usize,
732}
733
734fn plan_vector_read_runs(indexes: &[usize], runs: &mut Vec<VectorReadRun>) -> Result<()> {
739 runs.clear();
740 for (buffer_index, &flat_index) in indexes.iter().enumerate() {
741 if let Some(run) = runs.last_mut()
742 && run
743 .flat_start
744 .checked_add(run.count)
745 .is_some_and(|next| next == flat_index)
746 {
747 run.count += 1;
748 continue;
749 }
750 if buffer_index > 0 && flat_index <= indexes[buffer_index - 1] {
751 return Err(Error::Corruption(
752 "candidate flat-vector indexes are not strictly ordered".into(),
753 ));
754 }
755 runs.push(VectorReadRun {
756 buffer_start: buffer_index,
757 flat_start: flat_index,
758 count: 1,
759 });
760 }
761 Ok(())
762}
763
764fn prepare_vector_read_runs(
768 flat: &LazyFlatVectorData,
769 indexes: &[usize],
770 runs: &mut Vec<VectorReadRun>,
771) -> Result<()> {
772 plan_vector_read_runs(indexes, runs)?;
773 #[cfg(feature = "native")]
774 flat.prefetch_vectors(indexes.iter().copied());
775 #[cfg(not(feature = "native"))]
776 let _ = flat;
777 Ok(())
778}
779
780async fn read_vector_runs(
781 flat: &LazyFlatVectorData,
782 indexes: &[usize],
783 runs: &mut Vec<VectorReadRun>,
784 output: &mut [u8],
785) -> Result<()> {
786 prepare_vector_read_runs(flat, indexes, runs)?;
787 let vector_byte_size = flat.vector_byte_size();
788 for run in runs {
789 let bytes = flat
790 .read_vectors_batch(run.flat_start, run.count)
791 .await
792 .map_err(Error::Io)?;
793 let start = run
794 .buffer_start
795 .checked_mul(vector_byte_size)
796 .ok_or_else(|| Error::Query("dense rerank buffer offset overflow".into()))?;
797 let end = start
798 .checked_add(bytes.len())
799 .ok_or_else(|| Error::Query("dense rerank buffer range overflow".into()))?;
800 let destination = output
801 .get_mut(start..end)
802 .ok_or_else(|| Error::Corruption("dense rerank buffer is too short".into()))?;
803 destination.copy_from_slice(bytes.as_slice());
804 }
805 Ok(())
806}
807
808#[cfg(feature = "sync")]
809fn read_vector_runs_sync(
810 flat: &LazyFlatVectorData,
811 indexes: &[usize],
812 runs: &mut Vec<VectorReadRun>,
813 output: &mut [u8],
814) -> Result<()> {
815 prepare_vector_read_runs(flat, indexes, runs)?;
816 let vector_byte_size = flat.vector_byte_size();
817 for run in runs {
818 let bytes = flat
819 .read_vectors_batch_sync(run.flat_start, run.count)
820 .map_err(Error::Io)?;
821 let start = run
822 .buffer_start
823 .checked_mul(vector_byte_size)
824 .ok_or_else(|| Error::Query("dense rerank buffer offset overflow".into()))?;
825 let end = start
826 .checked_add(bytes.len())
827 .ok_or_else(|| Error::Query("dense rerank buffer range overflow".into()))?;
828 let destination = output
829 .get_mut(start..end)
830 .ok_or_else(|| Error::Corruption("dense rerank buffer is too short".into()))?;
831 destination.copy_from_slice(bytes.as_slice());
832 }
833 Ok(())
834}
835
836#[derive(Default)]
837struct DenseRerankStats {
838 vector_count: usize,
839 resolve_elapsed: std::time::Duration,
840 read_elapsed: std::time::Duration,
841 score_elapsed: std::time::Duration,
842}
843
844#[derive(Default)]
850struct DenseScratch {
851 raw: Vec<u8>,
852 scores: Vec<f32>,
853 batch_scores: Vec<f32>,
854 batch: Vec<CandidateVectorRef>,
855 flat_indexes: Vec<usize>,
856 read_runs: Vec<VectorReadRun>,
857 unresolved: Vec<(usize, usize)>,
858}
859
860thread_local! {
861 static DENSE_SCRATCH: std::cell::RefCell<Option<Box<DenseScratch>>> =
862 const { std::cell::RefCell::new(None) };
863}
864
865impl DenseScratch {
866 fn take() -> DenseScratchGuard {
870 let scratch = DENSE_SCRATCH
871 .with(|slot| slot.borrow_mut().take())
872 .unwrap_or_default();
873 DenseScratchGuard(Some(scratch))
874 }
875
876 fn prepare(&mut self, raw_bytes: usize, batch_len: usize) {
879 if self.raw.len() < raw_bytes {
880 self.raw.resize(raw_bytes, 0);
881 }
882 if self.scores.len() < batch_len {
883 self.scores.resize(batch_len, 0.0);
884 }
885 if self.batch_scores.len() < batch_len {
886 self.batch_scores.resize(batch_len, 0.0);
887 }
888 self.batch.clear();
889 self.flat_indexes.clear();
890 self.read_runs.clear();
891 self.unresolved.clear();
892 }
893}
894
895struct DenseScratchGuard(Option<Box<DenseScratch>>);
896
897impl std::ops::Deref for DenseScratchGuard {
898 type Target = DenseScratch;
899 fn deref(&self) -> &DenseScratch {
900 self.0
901 .as_deref()
902 .expect("dense scratch is present until drop")
903 }
904}
905
906impl std::ops::DerefMut for DenseScratchGuard {
907 fn deref_mut(&mut self) -> &mut DenseScratch {
908 self.0
909 .as_deref_mut()
910 .expect("dense scratch is present until drop")
911 }
912}
913
914impl Drop for DenseScratchGuard {
915 fn drop(&mut self) {
916 if let Some(scratch) = self.0.take() {
917 let _ = DENSE_SCRATCH.try_with(|slot| *slot.borrow_mut() = Some(scratch));
919 }
920 }
921}
922
923async fn exact_score_dense_candidate_documents(
924 ann_results: &[RawVectorCandidate],
925 flat: &LazyFlatVectorData,
926 query: &[f32],
927 unit_norm: bool,
928 combiner: crate::query::MultiValueCombiner,
929 limit: usize,
930) -> Result<(Vec<VectorSearchResult>, DenseRerankStats)> {
931 let resolve_started = std::time::Instant::now();
932 let documents = ann_candidate_document_ranges(ann_results, flat)?;
933 let mut stats = DenseRerankStats {
934 vector_count: documents.vector_count,
935 resolve_elapsed: resolve_started.elapsed(),
936 ..Default::default()
937 };
938 let vector_byte_size = flat.vector_byte_size();
939 let batch_len =
940 bounded_rerank_batch(vector_byte_size, DENSE_SCORE_BATCH, documents.vector_count);
941 let raw_capacity = batch_len
942 .checked_mul(vector_byte_size)
943 .ok_or_else(|| Error::Query("dense rerank buffer size overflow".to_string()))?;
944 let prepared_query =
945 PreparedDenseScoreQuery::new(query, flat.quantization, flat.dim, unit_norm)?;
946 let mut scratch = DenseScratch::take();
947 scratch.prepare(raw_capacity, batch_len);
948 let DenseScratch {
949 raw,
950 scores,
951 batch,
952 flat_indexes,
953 read_runs,
954 ..
955 } = &mut *scratch;
956 let mut cursor = CandidateVectorCursor::new(&documents.ranges);
957 let mut collector = FlatDocumentCollector::new(limit, combiner);
958 let mut scored = 0usize;
959
960 while cursor.fill_batch(flat, batch, batch_len)? {
961 flat_indexes.clear();
962 flat_indexes.extend(batch.iter().map(|&(_, _, flat_index)| flat_index));
963 let raw_len = batch
964 .len()
965 .checked_mul(vector_byte_size)
966 .ok_or_else(|| Error::Query("dense rerank buffer size overflow".to_string()))?;
967 let raw = &mut raw[..raw_len];
968
969 let read_started = std::time::Instant::now();
970 read_vector_runs(flat, flat_indexes, read_runs, raw).await?;
971 stats.read_elapsed += read_started.elapsed();
972
973 let score_started = std::time::Instant::now();
974 prepared_query.score_batch(raw, &mut scores[..batch.len()])?;
975 stats.score_elapsed += score_started.elapsed();
976 for (buffer_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
977 collector.push(doc_id, ordinal, scores[buffer_index]);
978 }
979 scored += batch.len();
980 }
981 debug_assert_eq!(scored, documents.vector_count);
982 Ok((collector.into_results(), stats))
983}
984
985#[cfg(feature = "sync")]
986fn exact_score_dense_candidate_documents_sync(
987 ann_results: &[RawVectorCandidate],
988 flat: &LazyFlatVectorData,
989 query: &[f32],
990 unit_norm: bool,
991 combiner: crate::query::MultiValueCombiner,
992 limit: usize,
993) -> Result<Vec<VectorSearchResult>> {
994 let documents = ann_candidate_document_ranges(ann_results, flat)?;
995 let vector_byte_size = flat.vector_byte_size();
996 let batch_len =
997 bounded_rerank_batch(vector_byte_size, DENSE_SCORE_BATCH, documents.vector_count);
998 let raw_capacity = batch_len
999 .checked_mul(vector_byte_size)
1000 .ok_or_else(|| Error::Query("dense rerank buffer size overflow".to_string()))?;
1001 let prepared_query =
1002 PreparedDenseScoreQuery::new(query, flat.quantization, flat.dim, unit_norm)?;
1003 let mut scratch = DenseScratch::take();
1004 scratch.prepare(raw_capacity, batch_len);
1005 let DenseScratch {
1006 raw,
1007 scores,
1008 batch,
1009 flat_indexes,
1010 read_runs,
1011 ..
1012 } = &mut *scratch;
1013 let mut cursor = CandidateVectorCursor::new(&documents.ranges);
1014 let mut collector = FlatDocumentCollector::new(limit, combiner);
1015 let mut scored = 0usize;
1016
1017 while cursor.fill_batch(flat, batch, batch_len)? {
1018 flat_indexes.clear();
1019 flat_indexes.extend(batch.iter().map(|&(_, _, flat_index)| flat_index));
1020 let raw_len = batch
1021 .len()
1022 .checked_mul(vector_byte_size)
1023 .ok_or_else(|| Error::Query("dense rerank buffer size overflow".to_string()))?;
1024 let raw = &mut raw[..raw_len];
1025 read_vector_runs_sync(flat, flat_indexes, read_runs, raw)?;
1026 prepared_query.score_batch(raw, &mut scores[..batch.len()])?;
1027 for (buffer_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
1028 collector.push(doc_id, ordinal, scores[buffer_index]);
1029 }
1030 scored += batch.len();
1031 }
1032 debug_assert_eq!(scored, documents.vector_count);
1033 Ok(collector.into_results())
1034}
1035
1036async fn exact_score_binary_candidate_documents(
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 = sorted_probe_scores(ann_results);
1046 exact_score_binary_resolved_documents(
1047 documents,
1048 &probe_scores,
1049 flat,
1050 query,
1051 dim_bits,
1052 combiner,
1053 limit,
1054 )
1055 .await
1056}
1057
1058async fn exact_score_binary_candidate_document_ids(
1059 candidate_doc_ids: Vec<DocId>,
1060 probed_ordinal_scores: &[(u32, u16, f32)],
1061 flat: &LazyFlatVectorData,
1062 query: &[u8],
1063 dim_bits: usize,
1064 combiner: crate::query::MultiValueCombiner,
1065 limit: usize,
1066) -> Result<Vec<VectorSearchResult>> {
1067 let documents = ann_candidate_document_ranges_from_ids(candidate_doc_ids, flat)?;
1068 let probe_scores = sorted_probe_scores(probed_ordinal_scores);
1071 exact_score_binary_resolved_documents(
1072 documents,
1073 &probe_scores,
1074 flat,
1075 query,
1076 dim_bits,
1077 combiner,
1078 limit,
1079 )
1080 .await
1081}
1082
1083fn sorted_probe_scores(
1087 scores: &[RawVectorCandidate],
1088) -> std::borrow::Cow<'_, [RawVectorCandidate]> {
1089 if scores.is_sorted_by_key(|&(doc_id, ordinal, _)| (doc_id, ordinal)) {
1090 std::borrow::Cow::Borrowed(scores)
1091 } else {
1092 let mut sorted = scores.to_vec();
1093 sorted.sort_unstable_by_key(|&(doc_id, ordinal, _)| (doc_id, ordinal));
1094 std::borrow::Cow::Owned(sorted)
1095 }
1096}
1097
1098struct ProbeScoreCursor<'a> {
1102 scores: &'a [RawVectorCandidate],
1103 position: usize,
1104}
1105
1106impl<'a> ProbeScoreCursor<'a> {
1107 fn new(scores: &'a [RawVectorCandidate]) -> Self {
1108 Self {
1109 scores,
1110 position: 0,
1111 }
1112 }
1113
1114 #[inline]
1117 fn advance_to(&mut self, doc_id: DocId, ordinal: u16) -> Option<f32> {
1118 while let Some(&(probed_doc, probed_ordinal, _)) = self.scores.get(self.position)
1119 && (probed_doc, probed_ordinal) < (doc_id, ordinal)
1120 {
1121 self.position += 1;
1122 }
1123 match self.scores.get(self.position) {
1124 Some(&(probed_doc, probed_ordinal, score))
1125 if probed_doc == doc_id && probed_ordinal == ordinal =>
1126 {
1127 Some(score)
1128 }
1129 _ => None,
1130 }
1131 }
1132}
1133
1134async fn exact_score_binary_resolved_documents(
1135 documents: AnnCandidateDocuments,
1136 probe_scores: &[RawVectorCandidate],
1137 flat: &LazyFlatVectorData,
1138 query: &[u8],
1139 dim_bits: usize,
1140 combiner: crate::query::MultiValueCombiner,
1141 limit: usize,
1142) -> Result<Vec<VectorSearchResult>> {
1143 let vector_byte_size = flat.vector_byte_size();
1144 let batch_len =
1145 bounded_rerank_batch(vector_byte_size, BINARY_SCORE_BATCH, documents.vector_count);
1146 let raw_capacity = batch_len
1147 .checked_mul(vector_byte_size)
1148 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
1149 let mut scratch = DenseScratch::take();
1150 scratch.prepare(raw_capacity, batch_len);
1151 let DenseScratch {
1152 raw,
1153 scores,
1154 batch_scores,
1155 batch,
1156 flat_indexes: unresolved_flat_indexes,
1157 read_runs,
1158 unresolved,
1159 } = &mut *scratch;
1160 let mut cursor = CandidateVectorCursor::new(&documents.ranges);
1161 let mut probe_cursor = ProbeScoreCursor::new(probe_scores);
1162 let mut collector = FlatDocumentCollector::new(limit, combiner);
1163 let mut scored = 0usize;
1164
1165 while cursor.fill_batch(flat, batch, batch_len)? {
1166 unresolved.clear();
1167 for (batch_index, &(doc_id, ordinal, flat_index)) in batch.iter().enumerate() {
1168 match probe_cursor.advance_to(doc_id, ordinal) {
1169 Some(score) => batch_scores[batch_index] = score,
1170 None => unresolved.push((batch_index, flat_index)),
1171 }
1172 }
1173 unresolved_flat_indexes.clear();
1174 unresolved_flat_indexes.extend(unresolved.iter().map(|&(_, flat_index)| flat_index));
1175 let raw_len = unresolved
1176 .len()
1177 .checked_mul(vector_byte_size)
1178 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
1179 let raw = &mut raw[..raw_len];
1180 read_vector_runs(flat, unresolved_flat_indexes, read_runs, raw).await?;
1181 crate::structures::simd::batch_hamming_scores(
1182 query,
1183 raw,
1184 vector_byte_size,
1185 dim_bits,
1186 &mut scores[..unresolved.len()],
1187 );
1188 for (buffer_index, &(batch_index, _)) in unresolved.iter().enumerate() {
1189 batch_scores[batch_index] = scores[buffer_index];
1190 }
1191 for (batch_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
1192 collector.push(doc_id, ordinal, batch_scores[batch_index]);
1193 }
1194 scored += batch.len();
1195 }
1196 debug_assert_eq!(scored, documents.vector_count);
1197 Ok(collector.into_results())
1198}
1199
1200#[cfg(feature = "sync")]
1201fn exact_score_binary_candidate_documents_sync(
1202 ann_results: &[RawVectorCandidate],
1203 flat: &LazyFlatVectorData,
1204 query: &[u8],
1205 dim_bits: usize,
1206 combiner: crate::query::MultiValueCombiner,
1207 limit: usize,
1208) -> Result<Vec<VectorSearchResult>> {
1209 let documents = ann_candidate_document_ranges(ann_results, flat)?;
1210 let probe_scores = sorted_probe_scores(ann_results);
1211 exact_score_binary_resolved_documents_sync(
1212 documents,
1213 &probe_scores,
1214 flat,
1215 query,
1216 dim_bits,
1217 combiner,
1218 limit,
1219 )
1220}
1221
1222#[cfg(feature = "sync")]
1223fn exact_score_binary_candidate_document_ids_sync(
1224 candidate_doc_ids: Vec<DocId>,
1225 probed_ordinal_scores: &[(u32, u16, f32)],
1226 flat: &LazyFlatVectorData,
1227 query: &[u8],
1228 dim_bits: usize,
1229 combiner: crate::query::MultiValueCombiner,
1230 limit: usize,
1231) -> Result<Vec<VectorSearchResult>> {
1232 let documents = ann_candidate_document_ranges_from_ids(candidate_doc_ids, flat)?;
1233 let probe_scores = sorted_probe_scores(probed_ordinal_scores);
1234 exact_score_binary_resolved_documents_sync(
1235 documents,
1236 &probe_scores,
1237 flat,
1238 query,
1239 dim_bits,
1240 combiner,
1241 limit,
1242 )
1243}
1244
1245#[cfg(feature = "sync")]
1246fn exact_score_binary_resolved_documents_sync(
1247 documents: AnnCandidateDocuments,
1248 probe_scores: &[RawVectorCandidate],
1249 flat: &LazyFlatVectorData,
1250 query: &[u8],
1251 dim_bits: usize,
1252 combiner: crate::query::MultiValueCombiner,
1253 limit: usize,
1254) -> Result<Vec<VectorSearchResult>> {
1255 let vector_byte_size = flat.vector_byte_size();
1256 let batch_len =
1257 bounded_rerank_batch(vector_byte_size, BINARY_SCORE_BATCH, documents.vector_count);
1258 let raw_capacity = batch_len
1259 .checked_mul(vector_byte_size)
1260 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
1261 let mut scratch = DenseScratch::take();
1262 scratch.prepare(raw_capacity, batch_len);
1263 let DenseScratch {
1264 raw,
1265 scores,
1266 batch_scores,
1267 batch,
1268 flat_indexes: unresolved_flat_indexes,
1269 read_runs,
1270 unresolved,
1271 } = &mut *scratch;
1272 let mut cursor = CandidateVectorCursor::new(&documents.ranges);
1273 let mut probe_cursor = ProbeScoreCursor::new(probe_scores);
1274 let mut collector = FlatDocumentCollector::new(limit, combiner);
1275 let mut scored = 0usize;
1276
1277 while cursor.fill_batch(flat, batch, batch_len)? {
1278 unresolved.clear();
1279 for (batch_index, &(doc_id, ordinal, flat_index)) in batch.iter().enumerate() {
1280 match probe_cursor.advance_to(doc_id, ordinal) {
1281 Some(score) => batch_scores[batch_index] = score,
1282 None => unresolved.push((batch_index, flat_index)),
1283 }
1284 }
1285 unresolved_flat_indexes.clear();
1286 unresolved_flat_indexes.extend(unresolved.iter().map(|&(_, flat_index)| flat_index));
1287 let raw_len = unresolved
1288 .len()
1289 .checked_mul(vector_byte_size)
1290 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
1291 let raw = &mut raw[..raw_len];
1292 read_vector_runs_sync(flat, unresolved_flat_indexes, read_runs, raw)?;
1293 crate::structures::simd::batch_hamming_scores(
1294 query,
1295 raw,
1296 vector_byte_size,
1297 dim_bits,
1298 &mut scores[..unresolved.len()],
1299 );
1300 for (buffer_index, &(batch_index, _)) in unresolved.iter().enumerate() {
1301 batch_scores[batch_index] = scores[buffer_index];
1302 }
1303 for (batch_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
1304 collector.push(doc_id, ordinal, batch_scores[batch_index]);
1305 }
1306 scored += batch.len();
1307 }
1308 debug_assert_eq!(scored, documents.vector_count);
1309 Ok(collector.into_results())
1310}
1311
1312#[cfg(feature = "sync")]
1316const FLAT_PARALLEL_SCAN_MIN_VECTORS: usize = 65_536;
1317
1318#[cfg(feature = "sync")]
1321fn document_aligned_batches(flat: &LazyFlatVectorData, batch_len: usize) -> Vec<(usize, usize)> {
1322 let n = flat.num_vectors;
1323 let batch_len = batch_len.max(1);
1324 let mut batches = Vec::with_capacity(n.div_ceil(batch_len));
1325 let mut start = 0usize;
1326 while start < n {
1327 let mut end = (start + batch_len).min(n);
1328 while end < n && flat.get_doc_id(end).0 == flat.get_doc_id(end - 1).0 {
1329 end += 1;
1330 }
1331 batches.push((start, end));
1332 start = end;
1333 }
1334 batches
1335}
1336
1337#[cfg(feature = "sync")]
1343fn brute_force_flat_scan_sync(
1344 flat: &LazyFlatVectorData,
1345 visibility: Option<&crate::query::DocBitset>,
1346 query: &[f32],
1347 unit_norm: bool,
1348 limit: usize,
1349 combiner: crate::query::MultiValueCombiner,
1350) -> Result<(Vec<VectorSearchResult>, DenseAnnScanStats)> {
1351 let n = flat.num_vectors;
1352 let batch_len = bounded_vector_score_batch(flat.vector_byte_size(), DENSE_SCORE_BATCH);
1353 let prepared_query =
1354 PreparedDenseScoreQuery::new(query, flat.quantization, flat.dim, unit_norm)?;
1355 let mut stats = DenseAnnScanStats {
1356 posting_count: n,
1357 ..DenseAnnScanStats::default()
1358 };
1359
1360 #[cfg(feature = "native")]
1361 if n >= FLAT_PARALLEL_SCAN_MIN_VECTORS && rayon::current_num_threads() > 1 {
1362 use rayon::prelude::*;
1363 let batches = document_aligned_batches(flat, batch_len);
1364 stats.scored_blocks = batches.len();
1365 stats.parallel = true;
1366 let merge = |mut left: Vec<VectorSearchResult>, mut right: Vec<VectorSearchResult>| {
1367 left.append(&mut right);
1368 left.sort_unstable_by(|a, b| {
1369 b.score
1370 .total_cmp(&a.score)
1371 .then_with(|| a.doc_id.cmp(&b.doc_id))
1372 });
1373 left.truncate(limit);
1374 left
1375 };
1376 let results = batches
1377 .par_iter()
1378 .try_fold(
1379 || {
1380 (
1381 FlatDocumentCollector::new(limit, combiner),
1382 Vec::<f32>::new(),
1383 )
1384 },
1385 |(mut collector, mut scores), &(start, end)| {
1386 let count = end - start;
1387 if scores.len() < count {
1388 scores.resize(count, 0.0);
1389 }
1390 let batch_bytes = flat
1391 .read_vectors_batch_sync(start, count)
1392 .map_err(Error::Io)?;
1393 prepared_query.score_batch(batch_bytes.as_slice(), &mut scores[..count])?;
1394 for (i, &score) in scores.iter().enumerate().take(count) {
1395 let (doc_id, ordinal) = flat.get_doc_id(start + i);
1396 if visibility.is_none_or(|bits| bits.contains(doc_id)) {
1397 collector.push(doc_id, ordinal, score);
1398 }
1399 }
1400 Ok::<_, Error>((collector, scores))
1401 },
1402 )
1403 .map(|folded| folded.map(|(collector, _)| collector.into_results()))
1404 .try_reduce(Vec::new, |left, right| Ok(merge(left, right)))?;
1405 return Ok((results, stats));
1406 }
1407
1408 let mut collector = FlatDocumentCollector::new(limit, combiner);
1409 let mut scratch = DenseScratch::take();
1410 scratch.prepare(0, batch_len);
1411 let scores = &mut scratch.scores;
1412 for batch_start in (0..n).step_by(batch_len) {
1413 let batch_count = batch_len.min(n - batch_start);
1414 let batch_bytes = flat
1415 .read_vectors_batch_sync(batch_start, batch_count)
1416 .map_err(Error::Io)?;
1417 prepared_query.score_batch(batch_bytes.as_slice(), &mut scores[..batch_count])?;
1418 stats.scored_blocks += 1;
1419 for (i, &score) in scores.iter().enumerate().take(batch_count) {
1420 let (doc_id, ordinal) = flat.get_doc_id(batch_start + i);
1421 if visibility.is_none_or(|bits| bits.contains(doc_id)) {
1422 collector.push(doc_id, ordinal, score);
1423 }
1424 }
1425 }
1426 Ok((collector.into_results(), stats))
1427}
1428
1429fn ann_keys_are_unique(
1433 index: &crate::segment::ann_disk::AnnDiskIndex,
1434 flat: &LazyFlatVectorData,
1435) -> bool {
1436 flat.num_vectors == flat.num_docs_with_vectors()
1437 && index.header().vector_count == flat.num_vectors
1438}
1439
1440fn dense_ann_kind_label(ann_index: Option<&VectorIndex>) -> &'static str {
1441 match ann_index {
1442 Some(VectorIndex::BinaryIvf(_)) => "binary_ivf",
1443 Some(VectorIndex::Tq { .. }) => "tq_flat",
1444 Some(VectorIndex::IvfTq { .. }) => "ivf_tq",
1445 Some(VectorIndex::ScannAh(_)) => "scann_ah",
1446 Some(VectorIndex::ScannBinary(_)) => "scann_binary",
1447 None => "flat",
1448 }
1449}
1450
1451fn validate_coarse_centroids(centroids: &CoarseCentroids, dim: usize) -> Result<()> {
1452 let expected = (centroids.num_clusters as usize)
1453 .checked_mul(dim)
1454 .ok_or_else(|| Error::Corruption("coarse centroid size overflow".into()))?;
1455 if centroids.num_clusters == 0
1456 || centroids.dim != dim
1457 || centroids.centroids.len() != expected
1458 || centroids.centroids.iter().any(|value| !value.is_finite())
1459 {
1460 return Err(Error::Corruption(format!(
1461 "invalid coarse centroids: clusters={}, dim={}, values={} (expected dim={dim}, values={expected})",
1462 centroids.num_clusters,
1463 centroids.dim,
1464 centroids.centroids.len()
1465 )));
1466 }
1467 Ok(())
1468}
1469
1470#[derive(Debug, Default)]
1475pub struct DensePlanCache {
1476 pub(crate) tq: std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqQueryPlan>>>,
1477 pub(crate) ivf_tq: std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqIvfQueryPlan>>>,
1478 scann: std::sync::Mutex<Option<ScannPlanCacheEntry>>,
1479}
1480
1481#[derive(Debug)]
1482struct ScannPlanCacheEntry {
1483 artifact_id: u64,
1484 probes: usize,
1485 query_bits: Vec<u32>,
1486 plan: std::sync::Arc<crate::structures::vector::scann::FloatScannQuery>,
1487}
1488
1489#[allow(clippy::too_many_arguments)]
1494fn search_tq_segment(
1495 index: &crate::segment::ann_disk::AnnDiskIndex,
1496 codec: &crate::structures::TqCodec,
1497 query: &[f32],
1498 fetch_k: usize,
1499 document_combiner: Option<crate::query::MultiValueCombiner>,
1500 field: Field,
1501 dim: usize,
1502 plan_cache: Option<&std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqQueryPlan>>>>,
1503 unique_keys: bool,
1504) -> Result<(Vec<RawVectorCandidate>, DenseAnnScanStats)> {
1505 validate_tq_ann(index, codec, dim, field)?;
1506 let plan = cached_tq_query_plan(codec, query, plan_cache)?;
1507 match document_combiner {
1508 Some(combiner) => index
1509 .search_tq_combined_documents(fetch_k, &plan, combiner)
1510 .map(|candidates| {
1511 (
1512 candidates
1513 .into_iter()
1514 .map(|candidate| (candidate.doc_id, 0, 0.0))
1518 .collect(),
1519 DenseAnnScanStats {
1520 posting_count: index.header().vector_count,
1521 ..DenseAnnScanStats::default()
1522 },
1523 )
1524 }),
1525 None => index.search_tq_distinct_with_stats(fetch_k, &plan, unique_keys),
1526 }
1527 .map_err(|error| {
1528 Error::Corruption(format!("invalid TQ payload for field {}: {error}", field.0))
1529 })
1530}
1531
1532fn cached_tq_query_plan(
1535 codec: &crate::structures::TqCodec,
1536 query: &[f32],
1537 plan_cache: Option<&std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqQueryPlan>>>>,
1538) -> Result<std::sync::Arc<crate::structures::TqQueryPlan>> {
1539 Ok(match plan_cache {
1540 Some(cache) => {
1541 let mut cached = cache
1542 .lock()
1543 .map_err(|_| Error::Internal("TQ plan cache is poisoned".into()))?;
1544 match cached.as_ref() {
1545 Some(plan)
1546 if plan.fingerprint() == codec.fingerprint() && plan.matches_query(query) =>
1547 {
1548 std::sync::Arc::clone(plan)
1549 }
1550 _ => {
1551 let plan =
1552 std::sync::Arc::new(crate::structures::TqQueryPlan::build(codec, query));
1553 *cached = Some(std::sync::Arc::clone(&plan));
1554 plan
1555 }
1556 }
1557 }
1558 None => std::sync::Arc::new(crate::structures::TqQueryPlan::build(codec, query)),
1559 })
1560}
1561
1562fn validate_tq_ann(
1563 index: &crate::segment::ann_disk::AnnDiskIndex,
1564 codec: &crate::structures::TqCodec,
1565 dim: usize,
1566 field: Field,
1567) -> Result<()> {
1568 let header = index.header();
1569 if header.dim != dim
1570 || codec.dim() != dim
1571 || header.code_size != codec.code_size()
1572 || header.quantizer_version != codec.fingerprint()
1573 || header.codebook_version != 0
1574 || header.num_clusters != 1
1575 {
1576 return Err(Error::Corruption(format!(
1577 "TQ payload for field {} does not match the codec derived from schema dimension {dim}",
1578 field.0,
1579 )));
1580 }
1581 Ok(())
1582}
1583
1584#[allow(clippy::too_many_arguments)]
1588fn search_ivf_tq_segment(
1589 index: &crate::segment::ann_disk::AnnDiskIndex,
1590 centroids: &CoarseCentroids,
1591 codec: &crate::structures::TqCodec,
1592 query: &[f32],
1593 fetch_k: usize,
1594 document_combiner: Option<crate::query::MultiValueCombiner>,
1595 field: Field,
1596 nprobe: usize,
1597 routing: crate::dsl::IvfRoutingMode,
1598 plan_cache: Option<
1599 &std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqIvfQueryPlan>>>,
1600 >,
1601 unique_keys: bool,
1602) -> Result<(Vec<RawVectorCandidate>, DenseAnnScanStats)> {
1603 let effective_nprobe = nprobe.clamp(1, centroids.num_clusters as usize);
1604 if effective_nprobe != nprobe {
1605 log::debug!(
1606 "[search_ivf_tq] field {}: nprobe {nprobe} clamped to {effective_nprobe} (codebook has {} leaves)",
1607 field.0,
1608 centroids.num_clusters
1609 );
1610 }
1611 let request_fingerprint = crate::structures::TqIvfQueryPlan::request_fingerprint_for(
1612 centroids,
1613 query,
1614 effective_nprobe,
1615 routing,
1616 );
1617 let build = || {
1618 std::sync::Arc::new(crate::structures::TqIvfQueryPlan::build(
1619 centroids,
1620 codec,
1621 query,
1622 effective_nprobe,
1623 routing,
1624 ))
1625 };
1626 let plan = match plan_cache {
1627 Some(cache) => {
1628 let mut cached = cache
1629 .lock()
1630 .map_err(|_| Error::Internal("IVF-TQ plan cache is poisoned".into()))?;
1631 match cached.as_ref() {
1632 Some(plan)
1633 if plan.quantizer_version == centroids.version
1634 && plan.fingerprint == codec.fingerprint()
1635 && plan.request_fingerprint == request_fingerprint
1636 && plan.cluster_ids.len() == effective_nprobe =>
1637 {
1638 std::sync::Arc::clone(plan)
1639 }
1640 _ => {
1641 let plan = build();
1642 *cached = Some(std::sync::Arc::clone(&plan));
1643 plan
1644 }
1645 }
1646 }
1647 None => build(),
1648 };
1649 let candidates = match document_combiner {
1650 Some(combiner) => index
1651 .search_ivf_tq_combined_documents_with_stats(fetch_k, &plan, combiner)
1652 .map(|(documents, stats)| {
1653 (
1654 documents
1655 .into_iter()
1656 .map(|candidate| (candidate.doc_id, 0, 0.0))
1660 .collect(),
1661 stats,
1662 )
1663 }),
1664 None => index.search_ivf_tq_distinct_with_stats(fetch_k, &plan, unique_keys),
1665 };
1666 candidates.map_err(|error| {
1667 Error::Corruption(format!(
1668 "invalid IVF-TQ payload for field {}: {error}",
1669 field.0
1670 ))
1671 })
1672}
1673
1674#[allow(clippy::too_many_arguments)]
1675fn search_scann_ah_segment(
1676 index: &crate::segment::ann_disk::AnnDiskIndex,
1677 artifact: &crate::segment::ScannTrainedArtifactBytes,
1678 query: &[f32],
1679 fetch_k: usize,
1680 combiner: crate::query::MultiValueCombiner,
1681 field: Field,
1682 nprobe: usize,
1683 plan_cache: Option<&std::sync::Mutex<Option<ScannPlanCacheEntry>>>,
1684) -> Result<Vec<RawVectorCandidate>> {
1685 index
1686 .validate_scann_generation(
1687 artifact.config(),
1688 artifact.generation(),
1689 artifact.artifact_id(),
1690 )
1691 .map_err(|error| {
1692 Error::Corruption(format!(
1693 "ScaNN generation mismatch for field {}: {error}",
1694 field.0
1695 ))
1696 })?;
1697 let probes = nprobe.clamp(1, artifact.config().num_leaves as usize);
1698 if probes != nprobe {
1699 log::debug!(
1700 "[search_scann_ah] field {}: nprobe {nprobe} clamped to {probes} (model has {} leaves)",
1701 field.0,
1702 artifact.config().num_leaves
1703 );
1704 }
1705 let build = || {
1706 let mut normalized_query = query.to_vec();
1707 crate::structures::vector::ivf::routing::normalize_cosine_in_place(&mut normalized_query);
1708 artifact
1709 .float_model()
1710 .map_err(Error::Io)?
1711 .prepare_query(&normalized_query, probes)
1712 .map(std::sync::Arc::new)
1713 .map_err(|error| Error::Query(format!("invalid ScaNN query: {error}")))
1714 };
1715 let plan = match plan_cache {
1716 Some(cache) => {
1717 let mut cached = cache
1718 .lock()
1719 .map_err(|_| Error::Internal("ScaNN plan cache is poisoned".into()))?;
1720 match cached.as_ref() {
1721 Some(entry)
1724 if entry.artifact_id == artifact.artifact_id()
1725 && entry.probes == probes
1726 && entry.query_bits.len() == query.len()
1727 && entry
1728 .query_bits
1729 .iter()
1730 .zip(query)
1731 .all(|(&bits, value)| bits == value.to_bits()) =>
1732 {
1733 std::sync::Arc::clone(&entry.plan)
1734 }
1735 _ => {
1736 let plan = build()?;
1737 *cached = Some(ScannPlanCacheEntry {
1738 artifact_id: artifact.artifact_id(),
1739 probes,
1740 query_bits: query.iter().map(|value| value.to_bits()).collect(),
1741 plan: std::sync::Arc::clone(&plan),
1742 });
1743 plan
1744 }
1745 }
1746 }
1747 None => build()?,
1748 };
1749 index
1750 .search_scann_ah_combined_documents(fetch_k, &plan, combiner)
1751 .map(|documents| {
1752 documents
1753 .into_iter()
1754 .map(|candidate| (candidate.doc_id, 0, 0.0))
1755 .collect()
1756 })
1757 .map_err(|error| {
1758 Error::Corruption(format!(
1759 "invalid ScaNN AH payload for field {}: {error}",
1760 field.0
1761 ))
1762 })
1763}
1764
1765fn validate_ivf_tq_ann(
1766 index: &crate::segment::ann_disk::AnnDiskIndex,
1767 centroids: &CoarseCentroids,
1768 codec: &crate::structures::TqCodec,
1769 dim: usize,
1770 routing: crate::dsl::IvfRoutingMode,
1771 field: Field,
1772) -> Result<()> {
1773 let header = index.header();
1774 if !crate::structures::is_ivf_tq_cosine_generation(centroids.version)
1775 || !crate::structures::is_ivf_tq_cosine_generation(header.quantizer_version)
1776 {
1777 return Err(Error::Corruption(format!(
1778 "IVF-TQ field {} uses a legacy unmarked raw-vector generation that cannot \
1779 preserve cosine candidate semantics; rebuild the index with a current \
1780 Summa version",
1781 field.0,
1782 )));
1783 }
1784 if header.dim != dim
1785 || codec.dim() != dim
1786 || header.code_size != codec.code_size()
1787 || header.num_clusters != centroids.num_clusters
1788 || header.quantizer_version != centroids.version
1789 || header.codebook_version != codec.fingerprint()
1790 || header.routing != routing
1791 {
1792 return Err(Error::Corruption(format!(
1793 "IVF-TQ payload for field {} does not match its quantizer/codec generation",
1794 field.0,
1795 )));
1796 }
1797 Ok(())
1798}
1799
1800fn validate_binary_ann(
1801 index: &crate::segment::ann_disk::AnnDiskIndex,
1802 quantizer: &crate::structures::BinaryCoarseQuantizer,
1803 config: &crate::dsl::BinaryDenseVectorConfig,
1804 dim: usize,
1805 field: Field,
1806) -> Result<()> {
1807 let header = index.header();
1808 if header.dim != dim
1809 || header.code_size != config.byte_len()
1810 || header.num_clusters != quantizer.num_clusters
1811 || header.quantizer_version != quantizer.version
1812 || header.codebook_version != 0
1813 || header.routing != config.ivf_routing
1814 || quantizer.dim_bits != dim
1815 {
1816 return Err(Error::Corruption(format!(
1817 "binary IVF field {} does not match its quantizer/schema generation",
1818 field.0,
1819 )));
1820 }
1821 Ok(())
1822}
1823
1824fn binary_probe_clusters(
1825 quantizer: &crate::structures::BinaryCoarseQuantizer,
1826 query: &[u8],
1827 nprobe: usize,
1828 routing: crate::dsl::IvfRoutingMode,
1829 cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
1830) -> Result<std::sync::Arc<[u32]>> {
1831 let effective_nprobe = nprobe.clamp(1, quantizer.num_clusters as usize);
1832 let request_fingerprint = quantizer.request_fingerprint(query, effective_nprobe, routing);
1833 if let Some(cache) = cache {
1834 let mut cached = cache
1835 .lock()
1836 .map_err(|_| Error::Internal("binary IVF probe cache is poisoned".into()))?;
1837 if let Some(plan) = cached.as_ref()
1838 && plan.quantizer_version == quantizer.version
1839 && plan.request_fingerprint == request_fingerprint
1840 && plan.cluster_ids.len() == effective_nprobe
1841 {
1842 return Ok(std::sync::Arc::clone(&plan.cluster_ids));
1843 }
1844 let plan = quantizer
1845 .probe(query, effective_nprobe, routing)
1846 .map_err(|error| Error::Query(format!("binary IVF routing failed: {error}")))?;
1847 let clusters = std::sync::Arc::clone(&plan.cluster_ids);
1848 *cached = Some(plan);
1849 return Ok(clusters);
1850 }
1851 Ok(quantizer
1852 .probe(query, effective_nprobe, routing)
1853 .map_err(|error| Error::Query(format!("binary IVF routing failed: {error}")))?
1854 .cluster_ids)
1855}
1856
1857thread_local! {
1858 static BINARY_SCANN_ROUTING_SCRATCH: std::cell::RefCell<
1862 crate::structures::vector::scann::BinaryScannSearchScratch,
1863 > = std::cell::RefCell::new(Default::default());
1864}
1865
1866fn probe_binary_scann_with_scratch(
1867 model: &crate::structures::vector::scann::QuantizedBinaryScannModelView<'_>,
1868 query: &[u8],
1869 probes: usize,
1870 routing_beam: usize,
1871) -> Result<crate::structures::vector::scann::BinaryScannProbePlan> {
1872 BINARY_SCANN_ROUTING_SCRATCH.with(|cell| {
1873 let mut scratch = cell
1874 .try_borrow_mut()
1875 .map_err(|_| Error::Internal("binary ScaNN routing scratch is busy".into()))?;
1876 model
1877 .probe(query, probes, routing_beam, &mut scratch)
1878 .map_err(|error| Error::Query(format!("binary ScaNN routing failed: {error}")))
1879 })
1880}
1881
1882fn binary_scann_probe_clusters(
1883 model: &crate::structures::vector::scann::QuantizedBinaryScannModelView<'_>,
1884 query: &[u8],
1885 nprobe: usize,
1886 cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
1887) -> Result<std::sync::Arc<[u32]>> {
1888 let probes = nprobe.clamp(1, model.num_leaves() as usize);
1889 let routing_beam = probes.min(64);
1892 let mut fingerprint = 0xcbf2_9ce4_8422_2325u64;
1893 for byte in query.iter().copied().chain((probes as u64).to_le_bytes()) {
1894 fingerprint ^= u64::from(byte);
1895 fingerprint = fingerprint.wrapping_mul(0x0000_0100_0000_01b3);
1896 }
1897 if let Some(cache) = cache {
1898 let mut cached = cache
1899 .lock()
1900 .map_err(|_| Error::Internal("binary ScaNN probe cache is poisoned".into()))?;
1901 if let Some(plan) = cached.as_ref()
1902 && plan.quantizer_version == model.fingerprint()
1903 && plan.request_fingerprint == fingerprint
1904 && plan.cluster_ids.len() == probes
1905 {
1906 return Ok(std::sync::Arc::clone(&plan.cluster_ids));
1907 }
1908 let routed = probe_binary_scann_with_scratch(model, query, probes, routing_beam)?;
1909 let plan =
1910 crate::structures::IvfProbePlan::new(model.fingerprint(), fingerprint, routed.leaf_ids);
1911 let leaves = std::sync::Arc::clone(&plan.cluster_ids);
1912 *cached = Some(plan);
1913 return Ok(leaves);
1914 }
1915 probe_binary_scann_with_scratch(model, query, probes, routing_beam)
1916 .map(|plan| plan.leaf_ids.into())
1917}
1918
1919#[derive(Clone)]
1925pub struct SegmentReader {
1926 row_stats: Arc<FxHashMap<u32, crate::structures::fast_field::FastFieldReader>>,
1927 deletion_meta: Option<super::DeletionMeta>,
1928 alive_docs: Option<Arc<crate::query::DocBitset>>,
1929 meta: SegmentMeta,
1930 term_dict: Arc<AsyncSSTableReader<TermInfo>>,
1932 postings: crate::structures::postings::PostingListReader,
1934 store: Arc<AsyncStoreReader>,
1936 schema: Arc<Schema>,
1937 vector_indexes: FxHashMap<u32, VectorIndex>,
1939 flat_vectors: FxHashMap<u32, LazyFlatVectorData>,
1941 dense_file_backed_bytes: u64,
1943 trained_vectors: Arc<crate::segment::TrainedVectorStructures>,
1945 sparse_indexes: FxHashMap<u32, SparseIndex>,
1947 seismic_indexes: FxHashMap<u32, crate::segment::seismic::SeismicIndex>,
1948 bmp_indexes: FxHashMap<u32, BmpIndex>,
1950 sparse_file_backed_bytes: u64,
1952 fast_fields: Arc<FxHashMap<u32, crate::structures::fast_field::FastFieldReader>>,
1954 chunk_maps: FxHashMap<u32, super::chunk_map::ChunkMap>,
1956 doc_lengths: FxHashMap<u32, super::chunk_map::DocLengths>,
1958 #[cfg(feature = "native")]
1960 dense_pin_report: crate::segment::pin::PinReport,
1961 #[cfg(feature = "native")]
1963 sparse_pin_report: crate::segment::pin::PinReport,
1964}
1965
1966impl SegmentReader {
1967 pub fn check_posting_integrity(&self) -> Result<()> {
1971 self.postings
1972 .check_integrity()
1973 .map_err(|error| Error::Corruption(format!("segment {:032x}: {error}", self.meta.id)))
1974 }
1975
1976 pub async fn open<D: Directory>(
1978 dir: &D,
1979 segment_id: SegmentId,
1980 schema: Arc<Schema>,
1981 term_cache_blocks: usize,
1982 ) -> Result<Self> {
1983 Self::open_with_term_cache_budget(dir, segment_id, schema, term_cache_blocks, None).await
1984 }
1985
1986 pub(crate) async fn open_with_term_cache_budget<D: Directory>(
1987 dir: &D,
1988 segment_id: SegmentId,
1989 schema: Arc<Schema>,
1990 term_cache_blocks: usize,
1991 term_cache_budget_bytes: Option<usize>,
1992 ) -> Result<Self> {
1993 Self::open_with_store_cache(
1994 dir,
1995 segment_id,
1996 schema,
1997 term_cache_blocks,
1998 term_cache_budget_bytes,
1999 dir as *const D as usize,
2000 Arc::new(super::SharedStoreCache::new(0)),
2001 )
2002 .await
2003 }
2004
2005 #[allow(clippy::too_many_arguments)]
2007 pub(crate) async fn open_with_store_cache<D: Directory>(
2008 dir: &D,
2009 segment_id: SegmentId,
2010 schema: Arc<Schema>,
2011 term_cache_blocks: usize,
2012 term_cache_budget_bytes: Option<usize>,
2013 store_cache_directory_namespace: usize,
2014 store_cache: Arc<super::SharedStoreCache>,
2015 ) -> Result<Self> {
2016 let files = SegmentFiles::new(segment_id.0);
2017
2018 let meta_slice = dir.open_read(&files.meta).await?;
2020 let meta_bytes = meta_slice.read_bytes().await?;
2021 let meta = SegmentMeta::deserialize(meta_bytes.as_slice())?;
2022 debug_assert_eq!(meta.id, segment_id.0);
2023
2024 let term_dict_handle = dir.open_lazy(&files.term_dict).await?;
2026 let term_dict = AsyncSSTableReader::open_with_cache_budget(
2027 term_dict_handle,
2028 term_cache_blocks,
2029 term_cache_budget_bytes,
2030 )
2031 .await?;
2032
2033 let positions_handle = loader::open_positions_file(dir, &files, &schema).await?;
2035 let postings = crate::structures::postings::PostingListReader::new(
2036 dir.open_lazy(&files.postings).await?,
2037 positions_handle,
2038 );
2039
2040 let store_handle = dir.open_lazy(&files.store).await?;
2042 let store = AsyncStoreReader::open(
2043 store_handle,
2044 store_cache_directory_namespace,
2045 segment_id.0,
2046 store_cache,
2047 )
2048 .await?;
2049
2050 let vectors_data = loader::load_vectors_file(dir, &files, &schema, meta.num_docs).await?;
2052 let dense_file_backed_bytes = vectors_data.file_backed_bytes;
2053 let vector_indexes = vectors_data.indexes;
2054 let flat_vectors = vectors_data.flat_vectors;
2055
2056 #[cfg(feature = "native")]
2061 for (field_id, lazy_flat) in &flat_vectors {
2062 if vector_indexes.contains_key(field_id) {
2063 lazy_flat.advise_random_access();
2064 }
2065 }
2066
2067 let sparse_data = loader::load_sparse_file(dir, &files, meta.num_docs, &schema).await?;
2069 let sparse_file_backed_bytes = sparse_data.file_backed_bytes;
2070 let sparse_indexes = sparse_data.maxscore_indexes;
2071 let bmp_indexes = sparse_data.bmp_indexes;
2072 let seismic_indexes = sparse_data.seismic_indexes;
2073
2074 let fast_fields = loader::load_fast_fields_file(dir, &files, &schema).await?;
2076
2077 let chunk_file = loader::load_chunk_maps_file(dir, &files, &schema).await?;
2080 let chunk_maps = chunk_file.chunk_maps;
2081 let doc_lengths = chunk_file.doc_lengths;
2082 for (&field, map) in &chunk_maps {
2083 if map.is_document_map()
2084 && (map.num_chunks() != meta.num_docs
2085 || schema.get_field_entry(Field(field)).is_none_or(|entry| {
2086 entry.chunked || entry.field_type != crate::dsl::FieldType::Text
2087 }))
2088 {
2089 return Err(Error::Corruption(
2090 "document map disagrees with segment scoring units".into(),
2091 ));
2092 }
2093 }
2094
2095 {
2097 let mut parts = vec![format!(
2098 "[segment] loaded {:016x}: docs={}",
2099 segment_id.0, meta.num_docs
2100 )];
2101 if !vector_indexes.is_empty() || !flat_vectors.is_empty() {
2102 parts.push(format!(
2103 "dense vectors: {} ANN + {} flat fields",
2104 vector_indexes.len(),
2105 flat_vectors.len()
2106 ));
2107 }
2108 for (field_id, idx) in &sparse_indexes {
2109 parts.push(format!(
2110 "sparse vector field {}: {} dims, ~{}",
2111 field_id,
2112 idx.num_dimensions(),
2113 crate::format_bytes(idx.num_dimensions() as u64 * 24)
2114 ));
2115 }
2116 for (field_id, idx) in &bmp_indexes {
2117 parts.push(format!(
2118 "bmp field {}: {} dims, {} blocks",
2119 field_id,
2120 idx.dims(),
2121 idx.num_blocks
2122 ));
2123 }
2124 if !fast_fields.is_empty() {
2125 parts.push(format!("fast: {} fields", fast_fields.len()));
2126 }
2127 for (field_id, map) in &chunk_maps {
2128 parts.push(format!(
2129 "chunked text field {}: {} chunks",
2130 field_id,
2131 map.num_chunks()
2132 ));
2133 }
2134 log::debug!("{}", parts.join(", "));
2135 }
2136
2137 let row_stats = match dir.open_read(&files.row_stats).await {
2138 Ok(handle) => loader::load_columns(handle).await?,
2139 Err(error) if error.kind() == std::io::ErrorKind::NotFound => FxHashMap::default(),
2140 Err(error) => return Err(error.into()),
2141 };
2142 for column in row_stats.values() {
2143 if column.num_docs != meta.num_docs
2144 || column.multi
2145 || column.column_type != crate::structures::fast_field::FastFieldColumnType::U64
2146 {
2147 return Err(Error::Corruption("invalid row statistics column".into()));
2148 }
2149 }
2150 #[allow(unused_mut)]
2151 let mut reader = Self {
2152 row_stats: Arc::new(row_stats),
2153 deletion_meta: None,
2154 alive_docs: None,
2155 meta,
2156 term_dict: Arc::new(term_dict),
2157 postings,
2158 store: Arc::new(store),
2159 schema,
2160 vector_indexes,
2161 flat_vectors,
2162 dense_file_backed_bytes,
2163 trained_vectors: Arc::new(crate::segment::TrainedVectorStructures::default()),
2164 sparse_indexes,
2165 bmp_indexes,
2166 seismic_indexes,
2167 sparse_file_backed_bytes,
2168 fast_fields: Arc::new(fast_fields),
2169 chunk_maps,
2170 doc_lengths,
2171 #[cfg(feature = "native")]
2172 dense_pin_report: Default::default(),
2173 #[cfg(feature = "native")]
2174 sparse_pin_report: Default::default(),
2175 };
2176
2177 #[cfg(feature = "native")]
2179 reader.apply_pin_policy(&crate::segment::pin::pin_policy().to_owned());
2180
2181 for (&field_id, vector_index) in &reader.vector_indexes {
2186 match vector_index {
2187 VectorIndex::BinaryIvf(index)
2188 | VectorIndex::IvfTq { index, .. }
2189 | VectorIndex::ScannAh(index)
2190 | VectorIndex::ScannBinary(index) => {
2191 index.get().report_health(
2192 reader.schema.index_label(),
2193 field_id,
2194 reader.meta.id,
2195 );
2196 }
2197 VectorIndex::Tq { .. } => {}
2200 }
2201 }
2202
2203 Ok(reader)
2204 }
2205
2206 pub fn ann_health(&self, field: Field) -> Option<crate::segment::ann_disk::AnnHealth> {
2210 match self.vector_indexes.get(&field.0)? {
2211 VectorIndex::BinaryIvf(index)
2212 | VectorIndex::IvfTq { index, .. }
2213 | VectorIndex::ScannAh(index)
2214 | VectorIndex::ScannBinary(index) => Some(index.get().health()),
2215 VectorIndex::Tq { .. } => None,
2216 }
2217 }
2218
2219 #[cfg(feature = "native")]
2229 pub(crate) fn apply_pin_policy(&mut self, policy: &crate::segment::pin::PinPolicy) {
2230 use crate::segment::pin::PinReport;
2231
2232 let disabled = !policy.is_enabled();
2236 let mut remaining = policy.budget_bytes;
2237 let mut dense_report = PinReport::default();
2238 let mut sparse_report = PinReport::default();
2239
2240 if !disabled {
2243 for index in self.vector_indexes.values_mut() {
2244 index.pin_lookup_directory(policy.mode, &mut remaining, &mut dense_report);
2245 }
2246 }
2247 for bmp in self.bmp_indexes.values_mut() {
2249 bmp.pin_block_starts(policy.mode, &mut remaining, &mut sparse_report);
2250 }
2251 for seismic in self.seismic_indexes.values_mut() {
2252 seismic.pin_term_directories(policy.mode, &mut remaining, &mut sparse_report);
2253 }
2254 for sparse in self.sparse_indexes.values_mut() {
2256 sparse.pin_skip_section(policy.mode, &mut remaining, &mut sparse_report);
2257 }
2258 for flat in self.flat_vectors.values_mut() {
2260 flat.pin_doc_ids(policy.mode, &mut remaining, &mut dense_report);
2261 }
2262 for bmp in self.bmp_indexes.values_mut() {
2263 bmp.pin_doc_maps(policy.mode, &mut remaining, &mut sparse_report);
2264 }
2265 for seismic in self.seismic_indexes.values_mut() {
2266 seismic.pin_row_directories(policy.mode, &mut remaining, &mut sparse_report);
2267 }
2268 for bmp in self.bmp_indexes.values_mut() {
2270 bmp.pin_query_hierarchy(policy.mode, &mut remaining, &mut sparse_report);
2271 }
2272
2273 let report = PinReport {
2274 intended_bytes: dense_report
2275 .intended_bytes
2276 .saturating_add(sparse_report.intended_bytes),
2277 pinned_bytes: dense_report
2278 .pinned_bytes
2279 .saturating_add(sparse_report.pinned_bytes),
2280 skipped_budget_bytes: dense_report
2281 .skipped_budget_bytes
2282 .saturating_add(sparse_report.skipped_budget_bytes),
2283 failed_bytes: dense_report
2284 .failed_bytes
2285 .saturating_add(sparse_report.failed_bytes),
2286 heap_copy_bytes: dense_report
2287 .heap_copy_bytes
2288 .saturating_add(sparse_report.heap_copy_bytes),
2289 };
2290 self.dense_pin_report = dense_report;
2291 self.sparse_pin_report = sparse_report;
2292 if disabled {
2293 crate::segment::pin::warn_if_pinning_disabled(
2294 self.schema.index_label(),
2295 self.meta.id,
2296 report.intended_bytes,
2297 );
2298 return;
2299 }
2300 if report.skipped_budget_bytes > 0 || report.failed_bytes > 0 {
2301 log::warn!(
2302 "[pin] index={} segment {:016x}: pinned {}/{} (budget skipped {}, mlock failed {}) — \
2303 raise SUMMA_PIN_METADATA_BUDGET_MB or RLIMIT_MEMLOCK for full coverage",
2304 self.schema.index_label(),
2305 self.meta.id,
2306 crate::format_bytes(report.pinned_bytes),
2307 crate::format_bytes(report.intended_bytes),
2308 crate::format_bytes(report.skipped_budget_bytes),
2309 crate::format_bytes(report.failed_bytes),
2310 );
2311 } else if report.pinned_bytes > 0 {
2312 log::info!(
2313 "[pin] index={} segment {:016x}: pinned {} of hot metadata ({:?})",
2314 self.schema.index_label(),
2315 self.meta.id,
2316 crate::format_bytes(report.pinned_bytes),
2317 policy.mode,
2318 );
2319 }
2320 }
2321
2322 pub fn meta(&self) -> &SegmentMeta {
2328 &self.meta
2329 }
2330
2331 #[cfg(feature = "native")]
2332 pub(crate) fn row_stats(
2333 &self,
2334 ) -> &FxHashMap<u32, crate::structures::fast_field::FastFieldReader> {
2335 &self.row_stats
2336 }
2337
2338 pub async fn open_with_deletions<D: Directory>(
2342 dir: &D,
2343 id: SegmentId,
2344 schema: Arc<Schema>,
2345 term_cache_blocks: usize,
2346 deletions: Option<super::DeletionMeta>,
2347 ) -> Result<Self> {
2348 let mut reader = Self::open(dir, id, schema, term_cache_blocks).await?;
2349 if let Some(meta) = deletions {
2350 reader.load_deletions(dir, meta).await?;
2351 }
2352 Ok(reader)
2353 }
2354
2355 pub(crate) async fn load_deletions<D: Directory>(
2356 &mut self,
2357 dir: &D,
2358 meta: super::DeletionMeta,
2359 ) -> Result<()> {
2360 let bits = meta.load(dir, self.num_docs()).await?;
2361 for index in self.vector_indexes.values_mut() {
2362 index.set_alive_docs(Some(Arc::clone(&bits)));
2363 }
2364 self.alive_docs = Some(bits);
2365 self.deletion_meta = Some(meta);
2366 Ok(())
2367 }
2368
2369 pub(crate) async fn with_deletions<D: Directory>(
2372 &self,
2373 dir: &D,
2374 meta: Option<super::DeletionMeta>,
2375 ) -> Result<Self> {
2376 let alive_docs = match &meta {
2377 Some(meta) => Some(meta.load(dir, self.num_docs()).await?),
2378 None => None,
2379 };
2380 let mut reader = self.clone();
2381 for index in reader.vector_indexes.values_mut() {
2382 index.set_alive_docs(alive_docs.clone());
2383 }
2384 reader.alive_docs = alive_docs;
2385 reader.deletion_meta = meta;
2386 Ok(reader)
2387 }
2388
2389 pub(crate) fn deletion_meta(&self) -> Option<&super::DeletionMeta> {
2390 self.deletion_meta.as_ref()
2391 }
2392
2393 pub(crate) fn alive_docs(&self) -> Option<Arc<crate::query::DocBitset>> {
2394 self.alive_docs.clone()
2395 }
2396
2397 pub fn is_alive(&self, doc_id: u32) -> bool {
2399 doc_id < self.num_docs()
2400 && self
2401 .alive_docs
2402 .as_ref()
2403 .is_none_or(|bits| bits.contains(doc_id))
2404 }
2405
2406 pub fn num_live_docs(&self) -> u32 {
2408 self.num_docs()
2409 - self
2410 .deletion_meta
2411 .as_ref()
2412 .map_or(0, |meta| meta.num_deleted)
2413 }
2414
2415 pub fn num_docs(&self) -> u32 {
2416 self.meta.num_docs
2417 }
2418
2419 pub fn avg_field_len(&self, field: Field) -> f32 {
2421 self.meta.avg_field_len(field)
2422 }
2423
2424 pub fn schema(&self) -> &Schema {
2425 &self.schema
2426 }
2427
2428 pub(crate) fn seismic_index(
2430 &self,
2431 field: Field,
2432 ) -> Option<&crate::segment::seismic::SeismicIndex> {
2433 self.seismic_indexes.get(&field.0)
2434 }
2435
2436 #[cfg(any(feature = "native", test))]
2438 pub(crate) fn seismic_indexes(&self) -> &FxHashMap<u32, crate::segment::seismic::SeismicIndex> {
2439 &self.seismic_indexes
2440 }
2441
2442 pub fn seismic_stats(&self) -> Vec<(u32, crate::segment::SeismicStats)> {
2444 let mut fields: Vec<_> = self
2445 .seismic_indexes
2446 .iter()
2447 .map(|(&field, index)| {
2448 (
2449 field,
2450 crate::segment::SeismicStats {
2451 total_vectors: index.total_vectors(),
2452 dimensions: index.dims(),
2453 nominations: index.nomination_count(),
2454 forward_entries: index.forward_entries(),
2455 clusters: index.cluster_count(),
2456 encoded_bytes: index.encoded_bytes() as u64,
2457 pending_terms: index.pending_terms(),
2458 runs: index.run_count() as u32,
2459 },
2460 )
2461 })
2462 .collect();
2463 fields.sort_unstable_by_key(|(field, _)| *field);
2464 fields
2465 }
2466
2467 pub fn sparse_indexes(&self) -> &FxHashMap<u32, SparseIndex> {
2469 &self.sparse_indexes
2470 }
2471
2472 pub fn sparse_index(&self, field: Field) -> Option<&SparseIndex> {
2474 self.sparse_indexes.get(&field.0)
2475 }
2476
2477 pub fn bmp_index(&self, field: Field) -> Option<&BmpIndex> {
2479 self.bmp_indexes.get(&field.0)
2480 }
2481
2482 pub fn bmp_indexes(&self) -> &FxHashMap<u32, BmpIndex> {
2484 &self.bmp_indexes
2485 }
2486
2487 pub fn vector_indexes(&self) -> &FxHashMap<u32, VectorIndex> {
2489 &self.vector_indexes
2490 }
2491
2492 pub fn flat_vectors(&self) -> &FxHashMap<u32, LazyFlatVectorData> {
2494 &self.flat_vectors
2495 }
2496
2497 pub fn fast_field(
2499 &self,
2500 field_id: u32,
2501 ) -> Option<&crate::structures::fast_field::FastFieldReader> {
2502 self.fast_fields.get(&field_id)
2503 }
2504
2505 pub fn fast_fields(&self) -> &FxHashMap<u32, crate::structures::fast_field::FastFieldReader> {
2507 &self.fast_fields
2508 }
2509
2510 pub fn chunk_map(&self, field: Field) -> Option<&super::chunk_map::ChunkMap> {
2513 self.chunk_maps.get(&field.0)
2514 }
2515
2516 pub fn chunk_maps(&self) -> &FxHashMap<u32, super::chunk_map::ChunkMap> {
2518 &self.chunk_maps
2519 }
2520
2521 pub fn doc_lengths(&self, field: Field) -> Option<&super::chunk_map::DocLengths> {
2524 self.doc_lengths.get(&field.0)
2525 }
2526
2527 pub fn is_chunked_field(&self, field: Field) -> bool {
2530 self.schema
2531 .get_field_entry(field)
2532 .is_some_and(|entry| entry.chunked)
2533 }
2534
2535 pub(crate) fn has_text_mapping(&self, field: Field) -> bool {
2537 self.is_chunked_field(field) || self.chunk_maps.contains_key(&field.0)
2538 }
2539
2540 pub fn num_chunks(&self, field: Field) -> u32 {
2542 self.chunk_maps
2543 .get(&field.0)
2544 .map_or(0, |map| map.num_chunks())
2545 }
2546
2547 pub fn text_corpus_size(&self, field: Field) -> f32 {
2550 if self.is_chunked_field(field) {
2551 self.num_chunks(field) as f32
2552 } else {
2553 self.meta.num_docs as f32
2554 }
2555 }
2556
2557 pub fn has_chunks_file(&self) -> bool {
2559 !self.chunk_maps.is_empty() || !self.doc_lengths.is_empty()
2560 }
2561
2562 pub fn term_dict_stats(&self) -> SSTableStats {
2564 self.term_dict.stats()
2565 }
2566
2567 pub fn memory_stats(&self) -> SegmentMemoryStats {
2569 let term_dict_stats = self.term_dict.stats();
2570
2571 let term_dict_cache_bytes = self.term_dict.cached_bytes();
2575 let store_cache_bytes = self.store.cached_bytes();
2576
2577 let sparse_heap_bytes: usize = self
2580 .seismic_indexes
2581 .values()
2582 .map(|i| i.estimated_heap_bytes())
2583 .sum::<usize>()
2584 + self
2585 .sparse_indexes
2586 .values()
2587 .map(|s| s.estimated_heap_bytes())
2588 .sum::<usize>()
2589 + self
2590 .bmp_indexes
2591 .values()
2592 .map(|b| b.estimated_heap_bytes())
2593 .sum::<usize>();
2594
2595 let dense_heap_bytes: usize = self
2598 .vector_indexes
2599 .values()
2600 .map(|v| v.estimated_heap_bytes())
2601 .sum::<usize>()
2602 + self
2603 .flat_vectors
2604 .values()
2605 .map(LazyFlatVectorData::estimated_heap_bytes)
2606 .sum::<usize>();
2607
2608 #[cfg(feature = "native")]
2609 let (sparse_heap_bytes, dense_heap_bytes) = (
2610 sparse_heap_bytes.saturating_add(
2611 usize::try_from(self.sparse_pin_report.heap_copy_bytes).unwrap_or(usize::MAX),
2612 ),
2613 dense_heap_bytes.saturating_add(
2614 usize::try_from(self.dense_pin_report.heap_copy_bytes).unwrap_or(usize::MAX),
2615 ),
2616 );
2617
2618 #[cfg(feature = "native")]
2619 let (
2620 sparse_pinned_metadata_bytes,
2621 sparse_pin_intended_bytes,
2622 dense_pinned_metadata_bytes,
2623 dense_pin_intended_bytes,
2624 ) = (
2625 self.sparse_pin_report.pinned_bytes,
2626 self.sparse_pin_report.intended_bytes,
2627 self.dense_pin_report.pinned_bytes,
2628 self.dense_pin_report.intended_bytes,
2629 );
2630 #[cfg(not(feature = "native"))]
2631 let (
2632 sparse_pinned_metadata_bytes,
2633 sparse_pin_intended_bytes,
2634 dense_pinned_metadata_bytes,
2635 dense_pin_intended_bytes,
2636 ) = (0u64, 0u64, 0u64, 0u64);
2637
2638 let pinned_metadata_bytes =
2639 sparse_pinned_metadata_bytes.saturating_add(dense_pinned_metadata_bytes);
2640 let pin_intended_bytes = sparse_pin_intended_bytes.saturating_add(dense_pin_intended_bytes);
2641
2642 SegmentMemoryStats {
2643 fast_field_metadata_heap_bytes: self
2644 .fast_fields
2645 .values()
2646 .map(|column| column.block_metadata_bytes())
2647 .sum(),
2648 row_stats_heap_bytes: self
2649 .row_stats
2650 .values()
2651 .map(|column| column.block_metadata_bytes())
2652 .sum(),
2653 row_stats_file_backed_bytes: self
2654 .row_stats
2655 .values()
2656 .map(|column| column.disk_bytes())
2657 .sum(),
2658 segment_id: self.meta.id,
2659 deletion_bytes: self
2660 .alive_docs
2661 .as_ref()
2662 .map_or(0, |bits| bits.bits.len() * 8),
2663 num_docs: self.meta.num_docs,
2664 term_dict_cache_bytes,
2665 posting_integrity_heap_bytes: self.postings.integrity_heap_bytes(),
2666 store_cache_bytes,
2667 sparse_heap_bytes,
2668 dense_heap_bytes,
2669 term_bloom_file_bytes: term_dict_stats.bloom_filter_size as u64,
2670 sparse_file_backed_bytes: self.sparse_file_backed_bytes,
2671 dense_file_backed_bytes: self.dense_file_backed_bytes,
2672 pinned_metadata_bytes,
2673 pin_intended_bytes,
2674 sparse_pinned_metadata_bytes,
2675 sparse_pin_intended_bytes,
2676 dense_pinned_metadata_bytes,
2677 dense_pin_intended_bytes,
2678 }
2679 }
2680
2681 pub(crate) async fn text_doc_freq(&self, field: Field, term: &[u8]) -> Result<u32> {
2684 let mut key = Vec::with_capacity(4 + term.len());
2685 key.extend_from_slice(&field.0.to_le_bytes());
2686 key.extend_from_slice(term);
2687 Ok(self
2688 .term_dict
2689 .get(&key)
2690 .await?
2691 .map_or(0, |info| info.doc_freq()))
2692 }
2693
2694 pub async fn get_postings(
2699 &self,
2700 field: Field,
2701 term: &[u8],
2702 ) -> Result<Option<BlockPostingList>> {
2703 log::debug!(
2704 "SegmentReader::get_postings field={} term_len={}",
2705 field.0,
2706 term.len()
2707 );
2708
2709 let mut key = Vec::with_capacity(4 + term.len());
2711 key.extend_from_slice(&field.0.to_le_bytes());
2712 key.extend_from_slice(term);
2713
2714 let term_info = match self.term_dict.get(&key).await? {
2716 Some(info) => {
2717 log::debug!("SegmentReader::get_postings found term_info");
2718 info
2719 }
2720 None => {
2721 log::debug!("SegmentReader::get_postings term not found");
2722 return Ok(None);
2723 }
2724 };
2725
2726 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
2728 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
2730 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
2731 posting_list.push(doc_id, tf);
2732 }
2733 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
2734 return Ok(Some(block_list));
2735 }
2736
2737 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
2739 Error::Corruption("TermInfo has neither inline nor external data".to_string())
2740 })?;
2741
2742 let range = checked_file_range(
2743 posting_offset,
2744 posting_len,
2745 self.postings.file().len(),
2746 "posting",
2747 )?;
2748 let block_list = self.postings.read(range).await?;
2749
2750 Ok(Some(block_list))
2751 }
2752
2753 pub async fn doc(&self, local_doc_id: DocId) -> Result<Option<Document>> {
2758 self.doc_with_fields(local_doc_id, None).await
2759 }
2760
2761 pub async fn doc_with_fields(
2767 &self,
2768 local_doc_id: DocId,
2769 fields: Option<&rustc_hash::FxHashSet<u32>>,
2770 ) -> Result<Option<Document>> {
2771 if !self.is_alive(local_doc_id) {
2772 return Ok(None);
2773 }
2774 let mut doc = match fields {
2775 Some(set) => {
2776 let field_ids: Vec<u32> = set.iter().copied().collect();
2777 match self
2778 .store
2779 .get_fields(local_doc_id, &self.schema, &field_ids)
2780 .await
2781 {
2782 Ok(Some(d)) => d,
2783 Ok(None) => return Ok(None),
2784 Err(e) => return Err(Error::from(e)),
2785 }
2786 }
2787 None => match self.store.get(local_doc_id, &self.schema).await {
2788 Ok(Some(d)) => d,
2789 Ok(None) => return Ok(None),
2790 Err(e) => return Err(Error::from(e)),
2791 },
2792 };
2793
2794 for (&field_id, lazy_flat) in &self.flat_vectors {
2796 if let Some(set) = fields
2798 && !set.contains(&field_id)
2799 {
2800 continue;
2801 }
2802
2803 let is_binary = lazy_flat.quantization == DenseVectorQuantization::Binary;
2804 let (start, entries) = lazy_flat.flat_indexes_for_doc(local_doc_id);
2805 for (j, &(_doc_id, _ordinal)) in entries.iter().enumerate() {
2806 let flat_idx = start + j;
2807 if is_binary {
2808 let vbs = lazy_flat.vector_byte_size();
2809 let mut raw = vec![0u8; vbs];
2810 match lazy_flat.read_vector_raw_into(flat_idx, &mut raw).await {
2811 Ok(()) => {
2812 doc.add_binary_dense_vector(Field(field_id), raw);
2813 }
2814 Err(e) => {
2815 log::warn!(
2816 "Failed to hydrate binary dense vector field {}: {}",
2817 field_id,
2818 e
2819 );
2820 }
2821 }
2822 } else {
2823 match lazy_flat.get_vector(flat_idx).await {
2824 Ok(vec) => {
2825 doc.add_dense_vector(Field(field_id), vec);
2826 }
2827 Err(e) => {
2828 log::warn!("Failed to hydrate dense vector field {}: {}", field_id, e);
2829 }
2830 }
2831 }
2832 }
2833 }
2834
2835 Ok(Some(doc))
2836 }
2837
2838 pub async fn prefetch_terms(
2840 &self,
2841 field: Field,
2842 start_term: &[u8],
2843 end_term: &[u8],
2844 ) -> Result<()> {
2845 let mut start_key = Vec::with_capacity(4 + start_term.len());
2846 start_key.extend_from_slice(&field.0.to_le_bytes());
2847 start_key.extend_from_slice(start_term);
2848
2849 let mut end_key = Vec::with_capacity(4 + end_term.len());
2850 end_key.extend_from_slice(&field.0.to_le_bytes());
2851 end_key.extend_from_slice(end_term);
2852
2853 self.term_dict.prefetch_range(&start_key, &end_key).await?;
2854 Ok(())
2855 }
2856
2857 pub fn store_has_dict(&self) -> bool {
2859 self.store.has_dict()
2860 }
2861
2862 pub fn store(&self) -> &super::store::AsyncStoreReader {
2864 &self.store
2865 }
2866
2867 pub fn store_raw_blocks(&self) -> Vec<RawStoreBlock> {
2869 self.store.raw_blocks()
2870 }
2871
2872 pub fn store_data_slice(&self) -> &FileHandle {
2874 self.store.data_slice()
2875 }
2876
2877 pub async fn all_terms(&self) -> Result<Vec<(Vec<u8>, TermInfo)>> {
2879 self.term_dict.all_entries().await.map_err(Error::from)
2880 }
2881
2882 pub async fn all_terms_with_stats(&self) -> Result<Vec<(Field, String, u32)>> {
2887 let entries = self.term_dict.all_entries().await?;
2888 let mut result = Vec::with_capacity(entries.len());
2889
2890 for (key, term_info) in entries {
2891 if key.len() > 4 {
2893 let field_id = u32::from_le_bytes([key[0], key[1], key[2], key[3]]);
2894 let term_bytes = &key[4..];
2895 if let Ok(term_str) = std::str::from_utf8(term_bytes) {
2896 result.push((Field(field_id), term_str.to_string(), term_info.doc_freq()));
2897 }
2898 }
2899 }
2900
2901 Ok(result)
2902 }
2903
2904 pub fn term_dict_iter(&self) -> crate::structures::AsyncSSTableIterator<'_, TermInfo> {
2906 self.term_dict.iter()
2907 }
2908
2909 pub async fn prefetch_term_dict(&self) -> crate::Result<()> {
2912 self.term_dict
2913 .prefetch_leading_blocks()
2914 .await
2915 .map_err(crate::Error::from)
2916 }
2917
2918 #[cfg(feature = "native")]
2919 pub(crate) fn posting_file_range(&self, offset: u64, len: u64) -> Result<FileHandle> {
2920 let range = checked_file_range(offset, len, self.postings.file().len(), "posting")?;
2921 Ok(self.postings.file().slice(range))
2922 }
2923
2924 #[cfg(feature = "native")]
2925 pub(crate) fn position_file_range(&self, offset: u64, len: u64) -> Result<FileHandle> {
2926 let handle = self
2927 .postings
2928 .positions_file()
2929 .ok_or_else(|| Error::Corruption("missing position data".into()))?;
2930 Ok(handle.slice(checked_file_range(offset, len, handle.len(), "position")?))
2931 }
2932
2933 pub async fn read_postings(&self, offset: u64, len: u64) -> Result<OwnedBytes> {
2935 let range = checked_file_range(offset, len, self.postings.file().len(), "posting")?;
2936 Ok(self.postings.file().read_bytes_range(range).await?)
2937 }
2938
2939 pub async fn read_position_bytes(&self, offset: u64, len: u64) -> Result<Option<OwnedBytes>> {
2941 let handle = match self.postings.positions_file() {
2942 Some(h) => h,
2943 None => return Ok(None),
2944 };
2945 let range = checked_file_range(offset, len, handle.len(), "position")?;
2946 Ok(Some(handle.read_bytes_range(range).await?))
2947 }
2948
2949 pub fn has_positions_file(&self) -> bool {
2951 self.postings.positions_file().is_some()
2952 }
2953
2954 fn validate_dense_search_request(
2958 &self,
2959 field: Field,
2960 query: &[f32],
2961 nprobe: usize,
2962 rerank_factor: f32,
2963 combiner: crate::query::MultiValueCombiner,
2964 ) -> Result<DenseSearchParams> {
2965 let entry = self
2966 .schema
2967 .get_field_entry(field)
2968 .ok_or_else(|| Error::FieldNotFound(field.0.to_string()))?;
2969 if entry.field_type != crate::dsl::FieldType::DenseVector {
2970 return Err(Error::InvalidFieldType {
2971 expected: "dense_vector".to_string(),
2972 got: format!("{:?}", entry.field_type),
2973 });
2974 }
2975 let config = entry.dense_vector_config.as_ref().ok_or_else(|| {
2976 Error::Schema(format!(
2977 "dense vector field '{}' has no dense vector configuration",
2978 entry.name
2979 ))
2980 })?;
2981
2982 if query.is_empty() {
2983 return Err(Error::Query(format!(
2984 "dense query vector for field '{}' must not be empty",
2985 entry.name
2986 )));
2987 }
2988 if query.len() != config.dim {
2989 return Err(Error::Query(format!(
2990 "dense query vector dimension {} does not match field '{}' dimension {}",
2991 query.len(),
2992 entry.name,
2993 config.dim
2994 )));
2995 }
2996 if let Some((index, value)) = query
2997 .iter()
2998 .enumerate()
2999 .find(|(_, value)| !value.is_finite())
3000 {
3001 return Err(Error::Query(format!(
3002 "dense query vector for field '{}' contains non-finite value {value} at index {index}",
3003 entry.name
3004 )));
3005 }
3006
3007 let nprobe = match (nprobe, config.nprobe) {
3010 (0, 0) => 32,
3011 (0, schema_nprobe) => schema_nprobe,
3012 (query_nprobe, _) => query_nprobe,
3013 };
3014 if nprobe > MAX_DENSE_NPROBE {
3015 return Err(Error::Query(format!(
3016 "dense nprobe must be at most {MAX_DENSE_NPROBE}, got {nprobe}"
3017 )));
3018 }
3019
3020 checked_dense_fetch_k(0, rerank_factor)?;
3023 combiner.validate().map_err(Error::Query)?;
3024
3025 Ok(DenseSearchParams {
3026 dim: config.dim,
3027 nprobe,
3028 unit_norm: config.unit_norm,
3029 })
3030 }
3031
3032 fn validate_binary_search_request(&self, field: Field, query: &[u8]) -> Result<usize> {
3033 let entry = self
3034 .schema
3035 .get_field_entry(field)
3036 .ok_or_else(|| Error::FieldNotFound(field.0.to_string()))?;
3037 if entry.field_type != crate::dsl::FieldType::BinaryDenseVector {
3038 return Err(Error::InvalidFieldType {
3039 expected: "binary_dense_vector".to_string(),
3040 got: format!("{:?}", entry.field_type),
3041 });
3042 }
3043 let config = entry.binary_dense_vector_config.as_ref().ok_or_else(|| {
3044 Error::Schema(format!(
3045 "binary dense vector field '{}' has no configuration",
3046 entry.name
3047 ))
3048 })?;
3049 if config.dim == 0 || !config.dim.is_multiple_of(8) {
3050 return Err(Error::Schema(format!(
3051 "binary dense vector field '{}' has invalid dimension {}",
3052 entry.name, config.dim
3053 )));
3054 }
3055 if query.len() != config.byte_len() {
3056 return Err(Error::Query(format!(
3057 "binary query byte length {} does not match field '{}' byte length {}",
3058 query.len(),
3059 entry.name,
3060 config.byte_len()
3061 )));
3062 }
3063 Ok(config.dim)
3064 }
3065
3066 #[cfg(test)]
3068 fn score_quantized_batch_legacy(
3069 query: &[f32],
3070 raw: &[u8],
3071 quant: crate::dsl::DenseVectorQuantization,
3072 dim: usize,
3073 scores: &mut [f32],
3074 unit_norm: bool,
3075 ) -> Result<()> {
3076 use crate::dsl::DenseVectorQuantization;
3077 use crate::structures::simd;
3078
3079 if query.len() != dim {
3080 return Err(Error::Query(format!(
3081 "dense SIMD query dimension {} does not match vector dimension {dim}",
3082 query.len()
3083 )));
3084 }
3085 let element_size = match quant {
3086 DenseVectorQuantization::F32 => std::mem::size_of::<f32>(),
3087 DenseVectorQuantization::F16 => std::mem::size_of::<u16>(),
3088 DenseVectorQuantization::UInt8 => 1,
3089 DenseVectorQuantization::Binary => {
3090 return Err(Error::InvalidFieldType {
3091 expected: "non-binary dense vector".to_string(),
3092 got: "binary dense vector".to_string(),
3093 });
3094 }
3095 };
3096 let required_bytes = scores
3097 .len()
3098 .checked_mul(dim)
3099 .and_then(|elements| elements.checked_mul(element_size))
3100 .ok_or_else(|| Error::Corruption("dense vector batch byte length overflow".into()))?;
3101 if raw.len() < required_bytes {
3102 return Err(Error::Corruption(format!(
3103 "dense vector batch is truncated: need {required_bytes} bytes, got {}",
3104 raw.len()
3105 )));
3106 }
3107 if quant == DenseVectorQuantization::F16
3108 && required_bytes > 0
3109 && !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>())
3110 {
3111 return Err(Error::Corruption(
3112 "f16 vector data is not 2-byte aligned".to_string(),
3113 ));
3114 }
3115
3116 match (quant, unit_norm) {
3117 (DenseVectorQuantization::F32, false) => {
3118 let num_floats = scores.len() * dim;
3119 if !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()) {
3120 return Err(Error::Corruption(
3121 "f32 vector data is not 4-byte aligned".to_string(),
3122 ));
3123 }
3124 let vectors: &[f32] =
3125 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
3126 simd::batch_cosine_scores(query, vectors, dim, scores);
3127 }
3128 (DenseVectorQuantization::F32, true) => {
3129 let num_floats = scores.len() * dim;
3130 if !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()) {
3131 return Err(Error::Corruption(
3132 "f32 vector data is not 4-byte aligned".to_string(),
3133 ));
3134 }
3135 let vectors: &[f32] =
3136 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
3137 simd::batch_dot_scores(query, vectors, dim, scores);
3138 }
3139 (DenseVectorQuantization::F16, false) => {
3140 simd::batch_cosine_scores_f16(query, raw, dim, scores);
3141 }
3142 (DenseVectorQuantization::F16, true) => {
3143 simd::batch_dot_scores_f16(query, raw, dim, scores);
3144 }
3145 (DenseVectorQuantization::UInt8, false) => {
3146 simd::batch_cosine_scores_u8(query, raw, dim, scores);
3147 }
3148 (DenseVectorQuantization::UInt8, true) => {
3149 simd::batch_dot_scores_u8(query, raw, dim, scores);
3150 }
3151 (DenseVectorQuantization::Binary, _) => unreachable!("validated above"),
3152 }
3153 Ok(())
3154 }
3155
3156 pub async fn search_dense_vector(
3162 &self,
3163 field: Field,
3164 query: &[f32],
3165 k: usize,
3166 nprobe: usize,
3167 rerank_factor: f32,
3168 combiner: crate::query::MultiValueCombiner,
3169 ) -> Result<Vec<VectorSearchResult>> {
3170 self.search_dense_vector_impl(field, query, k, nprobe, rerank_factor, combiner, None)
3171 .await
3172 }
3173
3174 #[allow(clippy::too_many_arguments)]
3175 pub(crate) async fn search_dense_vector_with_probe_cache(
3176 &self,
3177 field: Field,
3178 query: &[f32],
3179 k: usize,
3180 nprobe: usize,
3181 rerank_factor: f32,
3182 combiner: crate::query::MultiValueCombiner,
3183 plan_cache: &DensePlanCache,
3184 ) -> Result<Vec<VectorSearchResult>> {
3185 self.search_dense_vector_impl(
3186 field,
3187 query,
3188 k,
3189 nprobe,
3190 rerank_factor,
3191 combiner,
3192 Some(plan_cache),
3193 )
3194 .await
3195 }
3196
3197 #[allow(clippy::too_many_arguments)]
3198 async fn search_dense_vector_impl(
3199 &self,
3200 field: Field,
3201 query: &[f32],
3202 k: usize,
3203 nprobe: usize,
3204 rerank_factor: f32,
3205 combiner: crate::query::MultiValueCombiner,
3206 plan_cache: Option<&DensePlanCache>,
3207 ) -> Result<Vec<VectorSearchResult>> {
3208 let params =
3209 self.validate_dense_search_request(field, query, nprobe, rerank_factor, combiner)?;
3210 let fetch_k = checked_dense_fetch_k(k, rerank_factor)?;
3211 if k == 0 {
3212 return Ok(Vec::new());
3213 }
3214
3215 let configured_ann_index = self.vector_indexes.get(&field.0);
3216 let lazy_flat = self.flat_vectors.get(&field.0);
3217 if configured_ann_index.is_none() && lazy_flat.is_none() {
3219 return Ok(Vec::new());
3220 }
3221
3222 if configured_ann_index.is_some() && lazy_flat.is_none() {
3223 return Err(Error::Corruption(format!(
3224 "dense ANN field {} is missing flat vector storage",
3225 field.0
3226 )));
3227 }
3228
3229 if let Some(flat) = lazy_flat
3230 && flat.dim != params.dim
3231 {
3232 return Err(Error::Corruption(format!(
3233 "dense vector field {} has schema dimension {} but flat storage dimension {}",
3234 field.0, params.dim, flat.dim
3235 )));
3236 }
3237
3238 let needs_document_aggregation = lazy_flat.is_some_and(|flat| {
3239 flat.num_vectors != flat.num_docs_with_vectors()
3240 && !matches!(combiner, crate::query::MultiValueCombiner::Max)
3241 });
3242 let ann_index = configured_ann_index;
3246
3247 let t0 = std::time::Instant::now();
3249 let mut flat_results = None;
3250 let (results, scan_stats): (Vec<(u32, u16, f32)>, DenseAnnScanStats) = if let Some(index) =
3251 ann_index
3252 {
3253 match index {
3255 VectorIndex::Tq { index: lazy, codec } => {
3256 let flat = lazy_flat.expect("ANN/flat pairing validated above");
3257 search_tq_segment(
3259 lazy.get(),
3260 codec,
3261 query,
3262 fetch_k.min(flat.num_docs_with_vectors()),
3263 needs_document_aggregation.then_some(combiner),
3264 field,
3265 params.dim,
3266 plan_cache.map(|cache| &cache.tq),
3267 ann_keys_are_unique(lazy.get(), flat),
3268 )?
3269 }
3270 VectorIndex::IvfTq { index: lazy, codec } => {
3271 let index = lazy.get();
3272 let centroids =
3273 self.trained_vectors
3274 .centroids
3275 .get(&field.0)
3276 .ok_or_else(|| {
3277 Error::Schema(format!(
3278 "IVF-TQ index requires coarse centroids for field {}",
3279 field.0
3280 ))
3281 })?;
3282 validate_coarse_centroids(centroids, params.dim)?;
3283 let routing = self
3284 .schema
3285 .get_field_entry(field)
3286 .and_then(|entry| entry.dense_vector_config.as_ref())
3287 .map_or(crate::dsl::IvfRoutingMode::Auto, |config| {
3288 config.ivf_routing
3289 });
3290 validate_ivf_tq_ann(index, centroids, codec, params.dim, routing, field)?;
3291 let flat = lazy_flat.expect("ANN/flat pairing validated above");
3292 search_ivf_tq_segment(
3293 index,
3294 centroids,
3295 codec,
3296 query,
3297 fetch_k.min(flat.num_docs_with_vectors()),
3298 needs_document_aggregation.then_some(combiner),
3299 field,
3300 params.nprobe,
3301 routing,
3302 plan_cache.map(|cache| &cache.ivf_tq),
3303 ann_keys_are_unique(index, flat),
3304 )?
3305 }
3306 VectorIndex::BinaryIvf(_) => {
3307 return Err(Error::Query(format!(
3310 "dense vector field '{}' is served by a binary IVF index; use BinaryDenseVectorQuery",
3311 self.schema.get_field_name(field).unwrap_or("?")
3312 )));
3313 }
3314 VectorIndex::ScannAh(lazy) => {
3315 let artifact = self
3316 .trained_vectors
3317 .scann_artifacts
3318 .get(&field.0)
3319 .ok_or_else(|| {
3320 Error::Schema(format!(
3321 "ScaNN field {} has no loaded global artifact",
3322 field.0
3323 ))
3324 })?;
3325 let flat = lazy_flat.expect("ANN/flat pairing validated above");
3326 search_scann_ah_segment(
3327 lazy.get(),
3328 artifact,
3329 query,
3330 fetch_k.min(flat.num_docs_with_vectors()),
3331 combiner,
3332 field,
3333 params.nprobe,
3334 plan_cache.map(|cache| &cache.scann),
3335 )
3336 .map(|candidates| (candidates, DenseAnnScanStats::default()))?
3337 }
3338 VectorIndex::ScannBinary(_) => {
3339 return Err(Error::Corruption(format!(
3340 "binary ScaNN payload was attached to float field {}",
3341 field.0
3342 )));
3343 }
3344 }
3345 } else if let Some(lazy_flat) = lazy_flat {
3346 log::debug!(
3350 "[dense_vector_search] index={} field {}: brute-force on {} vectors (dim={}, quant={:?})",
3351 self.schema.index_label(),
3352 field.0,
3353 lazy_flat.num_vectors,
3354 lazy_flat.dim,
3355 lazy_flat.quantization
3356 );
3357 let dim = lazy_flat.dim;
3358 let n = lazy_flat.num_vectors;
3359 let quant = lazy_flat.quantization;
3360 let batch_len =
3361 bounded_vector_score_batch(lazy_flat.vector_byte_size(), DENSE_SCORE_BATCH);
3362 let mut collector = FlatDocumentCollector::new(fetch_k.min(n), combiner);
3363 let prepared_query = PreparedDenseScoreQuery::new(query, quant, dim, params.unit_norm)?;
3364 let mut flat_stats = DenseAnnScanStats {
3365 posting_count: n,
3366 ..DenseAnnScanStats::default()
3367 };
3368 let mut scratch = DenseScratch::take();
3369 scratch.prepare(0, batch_len);
3370 let scores = &mut scratch.scores;
3371
3372 for batch_start in (0..n).step_by(batch_len) {
3373 let batch_count = batch_len.min(n - batch_start);
3374 let batch_bytes = lazy_flat
3375 .read_vectors_batch(batch_start, batch_count)
3376 .await
3377 .map_err(crate::Error::Io)?;
3378 let raw = batch_bytes.as_slice();
3379
3380 prepared_query.score_batch(raw, &mut scores[..batch_count])?;
3381 flat_stats.scored_blocks += 1;
3382
3383 for (i, &score) in scores.iter().enumerate().take(batch_count) {
3384 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
3385 if self.is_alive(doc_id) {
3386 collector.push(doc_id, ordinal, score);
3387 }
3388 }
3389 }
3390
3391 flat_results = Some(collector.into_results());
3392 (Vec::new(), flat_stats)
3393 } else {
3394 return Ok(Vec::new());
3395 };
3396 let l1_elapsed = t0.elapsed();
3397 {
3398 let kind = dense_ann_kind_label(ann_index);
3399 let field_name = self.schema.get_field_name(field).unwrap_or("?");
3400 crate::observe::dense_l1(
3401 self.schema.index_label(),
3402 field_name,
3403 kind,
3404 l1_elapsed.as_secs_f64(),
3405 flat_results.as_ref().map_or(results.len(), Vec::len),
3406 );
3407 crate::observe::dense_ann_scan(self.schema.index_label(), field_name, kind, scan_stats);
3408 crate::observe::warn_non_finite_dense_scores(
3409 self.schema.index_label(),
3410 field_name,
3411 kind,
3412 scan_stats.non_finite_dropped,
3413 );
3414 }
3415 log::debug!(
3416 "[dense_vector_search] index={} field {}: L1 returned {} candidates in {:.1}ms",
3417 self.schema.index_label(),
3418 field.0,
3419 flat_results.as_ref().map_or(results.len(), Vec::len),
3420 l1_elapsed.as_secs_f64() * 1000.0
3421 );
3422
3423 if let Some(results) = flat_results {
3424 return Ok(results);
3425 }
3426
3427 if ann_index.is_some()
3430 && !results.is_empty()
3431 && let Some(lazy_flat) = lazy_flat
3432 {
3433 let t_rerank = std::time::Instant::now();
3434 let vbs = lazy_flat.vector_byte_size();
3435 let (reranked, stats) = exact_score_dense_candidate_documents(
3436 &results,
3437 lazy_flat,
3438 query,
3439 params.unit_norm,
3440 combiner,
3441 k,
3442 )
3443 .await?;
3444
3445 crate::observe::dense_rerank(
3446 self.schema.index_label(),
3447 self.schema.get_field_name(field).unwrap_or("?"),
3448 t_rerank.elapsed().as_secs_f64(),
3449 stats.resolve_elapsed.as_secs_f64(),
3450 stats.read_elapsed.as_secs_f64(),
3451 stats.vector_count,
3452 );
3453 log::debug!(
3454 "[dense_vector_search] index={} field {}: rerank {} vectors (dim={}, quant={:?}, bytes_per_vector={}): resolve={:.1}ms read={:.1}ms score={:.1}ms",
3455 self.schema.index_label(),
3456 field.0,
3457 stats.vector_count,
3458 lazy_flat.dim,
3459 lazy_flat.quantization,
3460 vbs,
3461 stats.resolve_elapsed.as_secs_f64() * 1000.0,
3462 stats.read_elapsed.as_secs_f64() * 1000.0,
3463 stats.score_elapsed.as_secs_f64() * 1000.0,
3464 );
3465
3466 log::debug!(
3467 "[dense_vector_search] index={} field {}: rerank total={:.1}ms",
3468 self.schema.index_label(),
3469 field.0,
3470 t_rerank.elapsed().as_secs_f64() * 1000.0
3471 );
3472 return Ok(reranked);
3473 }
3474
3475 Ok(combine_grouped_ordinal_results(results, combiner, k))
3476 }
3477
3478 async fn search_binary_dense_vector_impl(
3483 &self,
3484 field: Field,
3485 query: &[u8],
3486 k: usize,
3487 combiner: crate::query::MultiValueCombiner,
3488 probe_cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
3489 ) -> Result<Vec<VectorSearchResult>> {
3490 let schema_dim = self.validate_binary_search_request(field, query)?;
3491 combiner.validate().map_err(Error::Query)?;
3492 if k == 0 {
3493 return Ok(Vec::new());
3494 }
3495 let t0 = crate::observe::Timer::start();
3496 if let Some(VectorIndex::ScannBinary(lazy)) = self.vector_indexes.get(&field.0) {
3497 let artifact = self
3498 .trained_vectors
3499 .scann_artifacts
3500 .get(&field.0)
3501 .ok_or_else(|| {
3502 Error::Schema(format!(
3503 "binary ScaNN field {} has no loaded global artifact",
3504 field.0
3505 ))
3506 })?;
3507 lazy.get()
3508 .validate_scann_generation(
3509 artifact.config(),
3510 artifact.generation(),
3511 artifact.artifact_id(),
3512 )
3513 .map_err(|error| {
3514 Error::Corruption(format!(
3515 "binary ScaNN generation mismatch for field {}: {error}",
3516 field.0
3517 ))
3518 })?;
3519 let config = self
3520 .schema
3521 .get_field_entry(field)
3522 .and_then(|entry| entry.binary_dense_vector_config.as_ref())
3523 .ok_or_else(|| {
3524 Error::Schema(format!(
3525 "binary ScaNN field {} has no schema configuration",
3526 field.0
3527 ))
3528 })?;
3529 let model = artifact.binary_model().map_err(Error::Io)?;
3530 let clusters = binary_scann_probe_clusters(&model, query, config.nprobe, probe_cache)?;
3531 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
3532 Error::Corruption(format!(
3533 "binary ScaNN field {} is missing flat vectors",
3534 field.0
3535 ))
3536 })?;
3537 let candidate_limit =
3538 checked_binary_combined_fetch_k(k)?.min(flat.num_docs_with_vectors());
3539 let (documents, ordinal_scores) = lazy
3540 .get()
3541 .search_binary_combined_documents(candidate_limit, query, &clusters, combiner)
3542 .map_err(|error| {
3543 Error::Corruption(format!(
3544 "invalid binary ScaNN payload for field {}: {error}",
3545 field.0
3546 ))
3547 })?;
3548 let results = exact_score_binary_candidate_document_ids(
3549 documents
3550 .into_iter()
3551 .map(|candidate| candidate.doc_id)
3552 .collect(),
3553 &ordinal_scores,
3554 flat,
3555 query,
3556 schema_dim,
3557 combiner,
3558 k,
3559 )
3560 .await?;
3561 crate::observe::dense_l1(
3562 self.schema.index_label(),
3563 self.schema.get_field_name(field).unwrap_or("?"),
3564 "binary_scann",
3565 t0.secs(),
3566 results.len(),
3567 );
3568 return Ok(results);
3569 }
3570 if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0) {
3571 let ivf = lazy.get();
3572 let config = self
3573 .schema
3574 .get_field_entry(field)
3575 .and_then(|entry| entry.binary_dense_vector_config.as_ref())
3576 .ok_or_else(|| {
3577 Error::Schema(format!(
3578 "binary IVF field {} has no schema configuration",
3579 field.0
3580 ))
3581 })?;
3582 let quantizer = self
3583 .trained_vectors
3584 .binary_quantizers
3585 .get(&field.0)
3586 .ok_or_else(|| {
3587 Error::Schema(format!(
3588 "global binary IVF field {} has no loaded quantizer",
3589 field.0
3590 ))
3591 })?;
3592 validate_binary_ann(ivf, quantizer, config, schema_dim, field)?;
3593 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
3594 Error::Corruption(format!(
3595 "global binary IVF field {} is missing flat vector storage",
3596 field.0
3597 ))
3598 })?;
3599 let single_valued = flat.num_vectors == flat.num_docs_with_vectors();
3600 let clusters = binary_probe_clusters(
3601 quantizer,
3602 query,
3603 config.nprobe,
3604 config.ivf_routing,
3605 probe_cache,
3606 )?;
3607 let results = if !single_valued
3608 && !matches!(combiner, crate::query::MultiValueCombiner::Max)
3609 {
3610 let candidate_limit =
3611 checked_binary_combined_fetch_k(k)?.min(flat.num_docs_with_vectors());
3612 let (candidate_documents, probed_ordinal_scores) = ivf
3613 .search_binary_combined_documents(candidate_limit, query, &clusters, combiner)
3614 .map_err(|error| {
3615 Error::Corruption(format!(
3616 "invalid binary IVF payload for field {}: {error}",
3617 field.0,
3618 ))
3619 })?;
3620 exact_score_binary_candidate_document_ids(
3621 candidate_documents
3622 .into_iter()
3623 .map(|candidate| candidate.doc_id)
3624 .collect(),
3625 &probed_ordinal_scores,
3626 flat,
3627 query,
3628 schema_dim,
3629 combiner,
3630 k,
3631 )
3632 .await?
3633 } else {
3634 let candidate_docs = if single_valued {
3635 k
3636 } else {
3637 checked_binary_combined_fetch_k(k)?
3642 }
3643 .min(flat.num_docs_with_vectors());
3644 let ann_results = if single_valued {
3645 ivf.search_binary_clusters::<false>(query, candidate_docs, &clusters)
3646 } else {
3647 ivf.search_binary_clusters::<true>(query, candidate_docs, &clusters)
3648 }
3649 .map_err(|error| {
3650 Error::Corruption(format!(
3651 "invalid binary IVF payload for field {}: {error}",
3652 field.0,
3653 ))
3654 })?;
3655 if single_valued {
3658 let ann_results = validate_binary_single_value_ann_results(ann_results, flat)?;
3659 combine_ordinal_results(ann_results, combiner, k)
3660 } else {
3661 exact_score_binary_candidate_documents(
3662 &ann_results,
3663 flat,
3664 query,
3665 schema_dim,
3666 combiner,
3667 k,
3668 )
3669 .await?
3670 }
3671 };
3672 crate::observe::dense_l1(
3673 self.schema.index_label(),
3674 self.schema.get_field_name(field).unwrap_or("?"),
3675 "global_binary_ivf",
3676 t0.secs(),
3677 results.len(),
3678 );
3679 return Ok(results);
3680 }
3681 let lazy_flat = match self.flat_vectors.get(&field.0) {
3682 Some(f) => f,
3683 None => return Ok(Vec::new()),
3684 };
3685
3686 let dim_bits = lazy_flat.dim;
3687 let byte_len = lazy_flat.vector_byte_size();
3688 let n = lazy_flat.num_vectors;
3689
3690 if dim_bits != schema_dim {
3691 return Err(Error::Corruption(format!(
3692 "binary vector field {} has schema dimension {} but flat storage dimension {}",
3693 field.0, schema_dim, dim_bits
3694 )));
3695 }
3696
3697 if byte_len != query.len() {
3698 return Err(Error::Schema(format!(
3699 "Binary query vector byte length {} != field byte length {}",
3700 query.len(),
3701 byte_len
3702 )));
3703 }
3704
3705 let batch_len = bounded_vector_score_batch(byte_len, BINARY_SCORE_BATCH);
3706 let mut collector = FlatDocumentCollector::new(k, combiner);
3707 let mut scratch = DenseScratch::take();
3708 scratch.prepare(0, batch_len);
3709 let scores = &mut scratch.scores;
3710
3711 for batch_start in (0..n).step_by(batch_len) {
3712 let batch_count = batch_len.min(n - batch_start);
3713 let batch_bytes = lazy_flat
3714 .read_vectors_batch(batch_start, batch_count)
3715 .await
3716 .map_err(crate::Error::Io)?;
3717 let raw = batch_bytes.as_slice();
3718
3719 crate::structures::simd::batch_hamming_scores(
3720 query,
3721 raw,
3722 byte_len,
3723 dim_bits,
3724 &mut scores[..batch_count],
3725 );
3726
3727 for (i, &score) in scores.iter().enumerate().take(batch_count) {
3728 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
3729 if self.is_alive(doc_id) {
3730 collector.push(doc_id, ordinal, score);
3731 }
3732 }
3733 }
3734
3735 let results = collector.into_results();
3736
3737 crate::observe::dense_l1(
3738 self.schema.index_label(),
3739 self.schema.get_field_name(field).unwrap_or("?"),
3740 "binary_flat",
3741 t0.secs(),
3742 results.len(),
3743 );
3744 Ok(results)
3745 }
3746
3747 pub async fn search_binary_dense_vector(
3748 &self,
3749 field: Field,
3750 query: &[u8],
3751 k: usize,
3752 combiner: crate::query::MultiValueCombiner,
3753 ) -> Result<Vec<VectorSearchResult>> {
3754 self.search_binary_dense_vector_impl(field, query, k, combiner, None)
3755 .await
3756 }
3757
3758 pub(crate) async fn search_binary_dense_vector_with_probe_cache(
3759 &self,
3760 field: Field,
3761 query: &[u8],
3762 k: usize,
3763 combiner: crate::query::MultiValueCombiner,
3764 probe_cache: &std::sync::Mutex<Option<crate::structures::IvfProbePlan>>,
3765 ) -> Result<Vec<VectorSearchResult>> {
3766 self.search_binary_dense_vector_impl(field, query, k, combiner, Some(probe_cache))
3767 .await
3768 }
3769
3770 pub fn coarse_centroids(&self, field_id: u32) -> Option<&Arc<CoarseCentroids>> {
3772 self.trained_vectors.centroids.get(&field_id)
3773 }
3774
3775 pub fn set_trained_vectors(
3776 &mut self,
3777 trained_vectors: Arc<crate::segment::TrainedVectorStructures>,
3778 ) {
3779 self.trained_vectors = trained_vectors;
3780 }
3781
3782 pub fn get_vector_index(&self, field: Field) -> Option<&VectorIndex> {
3784 self.vector_indexes.get(&field.0)
3785 }
3786
3787 pub async fn get_positions(
3792 &self,
3793 field: Field,
3794 term: &[u8],
3795 ) -> Result<Option<crate::structures::TermPositions>> {
3796 let handle = match self.postings.positions_file() {
3798 Some(h) => h,
3799 None => return Ok(None),
3800 };
3801
3802 let mut key = Vec::with_capacity(4 + term.len());
3804 key.extend_from_slice(&field.0.to_le_bytes());
3805 key.extend_from_slice(term);
3806
3807 let term_info = match self.term_dict.get(&key).await? {
3809 Some(info) => info,
3810 None => return Ok(None),
3811 };
3812
3813 let (offset, length) = match term_info.position_info() {
3815 Some((o, l)) => (o, l),
3816 None => return Ok(None),
3817 };
3818
3819 let range = checked_file_range(offset, length, handle.len(), "position list")?;
3823 Ok(Some(self.postings.read_positions(range).await?))
3826 }
3827
3828 pub fn has_positions(&self, field: Field) -> bool {
3830 if let Some(entry) = self.schema.get_field_entry(field) {
3832 entry.positions.is_some()
3833 } else {
3834 false
3835 }
3836 }
3837}
3838
3839#[cfg(feature = "sync")]
3841impl SegmentReader {
3842 pub fn text_doc_freq_sync(&self, field: Field, term: &[u8]) -> Result<u32> {
3845 let mut key = Vec::with_capacity(4 + term.len());
3846 key.extend_from_slice(&field.0.to_le_bytes());
3847 key.extend_from_slice(term);
3848 Ok(self
3849 .term_dict
3850 .get_sync(&key)?
3851 .map_or(0, |info| info.doc_freq()))
3852 }
3853
3854 pub fn get_postings_sync(&self, field: Field, term: &[u8]) -> Result<Option<BlockPostingList>> {
3856 let mut key = Vec::with_capacity(4 + term.len());
3858 key.extend_from_slice(&field.0.to_le_bytes());
3859 key.extend_from_slice(term);
3860
3861 let term_info = match self.term_dict.get_sync(&key)? {
3863 Some(info) => info,
3864 None => return Ok(None),
3865 };
3866
3867 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
3869 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
3870 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
3871 posting_list.push(doc_id, tf);
3872 }
3873 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
3874 return Ok(Some(block_list));
3875 }
3876
3877 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
3879 Error::Corruption("TermInfo has neither inline nor external data".to_string())
3880 })?;
3881
3882 let range = checked_file_range(
3883 posting_offset,
3884 posting_len,
3885 self.postings.file().len(),
3886 "posting",
3887 )?;
3888 let block_list = self.postings.read_sync(range)?;
3889
3890 Ok(Some(block_list))
3891 }
3892
3893 pub fn get_positions_sync(
3895 &self,
3896 field: Field,
3897 term: &[u8],
3898 ) -> Result<Option<crate::structures::TermPositions>> {
3899 let handle = match self.postings.positions_file() {
3900 Some(h) => h,
3901 None => return Ok(None),
3902 };
3903
3904 let mut key = Vec::with_capacity(4 + term.len());
3906 key.extend_from_slice(&field.0.to_le_bytes());
3907 key.extend_from_slice(term);
3908
3909 let term_info = match self.term_dict.get_sync(&key)? {
3911 Some(info) => info,
3912 None => return Ok(None),
3913 };
3914
3915 let (offset, length) = match term_info.position_info() {
3916 Some((o, l)) => (o, l),
3917 None => return Ok(None),
3918 };
3919
3920 let range = checked_file_range(offset, length, handle.len(), "position list")?;
3921 Ok(Some(self.postings.read_positions_sync(range)?))
3922 }
3923
3924 pub fn search_dense_vector_sync(
3927 &self,
3928 field: Field,
3929 query: &[f32],
3930 k: usize,
3931 nprobe: usize,
3932 rerank_factor: f32,
3933 combiner: crate::query::MultiValueCombiner,
3934 ) -> Result<Vec<VectorSearchResult>> {
3935 self.search_dense_vector_sync_impl(field, query, k, nprobe, rerank_factor, combiner, None)
3936 }
3937
3938 #[cfg(feature = "sync")]
3939 #[allow(clippy::too_many_arguments)]
3940 pub(crate) fn search_dense_vector_sync_with_probe_cache(
3941 &self,
3942 field: Field,
3943 query: &[f32],
3944 k: usize,
3945 nprobe: usize,
3946 rerank_factor: f32,
3947 combiner: crate::query::MultiValueCombiner,
3948 plan_cache: &DensePlanCache,
3949 ) -> Result<Vec<VectorSearchResult>> {
3950 self.search_dense_vector_sync_impl(
3951 field,
3952 query,
3953 k,
3954 nprobe,
3955 rerank_factor,
3956 combiner,
3957 Some(plan_cache),
3958 )
3959 }
3960
3961 #[cfg(feature = "sync")]
3962 #[allow(clippy::too_many_arguments)]
3963 fn search_dense_vector_sync_impl(
3964 &self,
3965 field: Field,
3966 query: &[f32],
3967 k: usize,
3968 nprobe: usize,
3969 rerank_factor: f32,
3970 combiner: crate::query::MultiValueCombiner,
3971 plan_cache: Option<&DensePlanCache>,
3972 ) -> Result<Vec<VectorSearchResult>> {
3973 let params =
3974 self.validate_dense_search_request(field, query, nprobe, rerank_factor, combiner)?;
3975 let fetch_k = checked_dense_fetch_k(k, rerank_factor)?;
3976 if k == 0 {
3977 return Ok(Vec::new());
3978 }
3979
3980 let configured_ann_index = self.vector_indexes.get(&field.0);
3981 let lazy_flat = self.flat_vectors.get(&field.0);
3982 if configured_ann_index.is_none() && lazy_flat.is_none() {
3983 return Ok(Vec::new());
3984 }
3985
3986 if configured_ann_index.is_some() && lazy_flat.is_none() {
3987 return Err(Error::Corruption(format!(
3988 "dense ANN field {} is missing flat vector storage",
3989 field.0
3990 )));
3991 }
3992
3993 if let Some(flat) = lazy_flat
3994 && flat.dim != params.dim
3995 {
3996 return Err(Error::Corruption(format!(
3997 "dense vector field {} has schema dimension {} but flat storage dimension {}",
3998 field.0, params.dim, flat.dim
3999 )));
4000 }
4001
4002 let needs_document_aggregation = lazy_flat.is_some_and(|flat| {
4003 flat.num_vectors != flat.num_docs_with_vectors()
4004 && !matches!(combiner, crate::query::MultiValueCombiner::Max)
4005 });
4006 let ann_index = configured_ann_index;
4009
4010 let (results, scan_stats): (Vec<(u32, u16, f32)>, DenseAnnScanStats) = if let Some(index) =
4011 ann_index
4012 {
4013 match index {
4015 VectorIndex::Tq { index: lazy, codec } => {
4016 let flat = lazy_flat.expect("ANN/flat pairing validated above");
4017 search_tq_segment(
4018 lazy.get(),
4019 codec,
4020 query,
4021 fetch_k.min(flat.num_docs_with_vectors()),
4022 needs_document_aggregation.then_some(combiner),
4023 field,
4024 params.dim,
4025 plan_cache.map(|cache| &cache.tq),
4026 ann_keys_are_unique(lazy.get(), flat),
4027 )?
4028 }
4029 VectorIndex::IvfTq { index: lazy, codec } => {
4030 let index = lazy.get();
4031 let centroids =
4032 self.trained_vectors
4033 .centroids
4034 .get(&field.0)
4035 .ok_or_else(|| {
4036 Error::Schema(format!(
4037 "IVF-TQ index requires coarse centroids for field {}",
4038 field.0
4039 ))
4040 })?;
4041 validate_coarse_centroids(centroids, params.dim)?;
4042 let routing = self
4043 .schema
4044 .get_field_entry(field)
4045 .and_then(|entry| entry.dense_vector_config.as_ref())
4046 .map_or(crate::dsl::IvfRoutingMode::Auto, |config| {
4047 config.ivf_routing
4048 });
4049 validate_ivf_tq_ann(index, centroids, codec, params.dim, routing, field)?;
4050 let flat = lazy_flat.expect("ANN/flat pairing validated above");
4051 search_ivf_tq_segment(
4052 index,
4053 centroids,
4054 codec,
4055 query,
4056 fetch_k.min(flat.num_docs_with_vectors()),
4057 needs_document_aggregation.then_some(combiner),
4058 field,
4059 params.nprobe,
4060 routing,
4061 plan_cache.map(|cache| &cache.ivf_tq),
4062 ann_keys_are_unique(index, flat),
4063 )?
4064 }
4065 VectorIndex::BinaryIvf(_) => {
4066 return Err(Error::Query(format!(
4069 "dense vector field '{}' is served by a binary IVF index; use BinaryDenseVectorQuery",
4070 self.schema.get_field_name(field).unwrap_or("?")
4071 )));
4072 }
4073 VectorIndex::ScannAh(lazy) => {
4074 let artifact = self
4075 .trained_vectors
4076 .scann_artifacts
4077 .get(&field.0)
4078 .ok_or_else(|| {
4079 Error::Schema(format!(
4080 "ScaNN field {} has no loaded global artifact",
4081 field.0
4082 ))
4083 })?;
4084 let flat = lazy_flat.expect("ANN/flat pairing validated above");
4085 search_scann_ah_segment(
4086 lazy.get(),
4087 artifact,
4088 query,
4089 fetch_k.min(flat.num_docs_with_vectors()),
4090 combiner,
4091 field,
4092 params.nprobe,
4093 plan_cache.map(|cache| &cache.scann),
4094 )
4095 .map(|candidates| (candidates, DenseAnnScanStats::default()))?
4096 }
4097 VectorIndex::ScannBinary(_) => {
4098 return Err(Error::Corruption(format!(
4099 "binary ScaNN payload was attached to float field {}",
4100 field.0
4101 )));
4102 }
4103 }
4104 } else if let Some(lazy_flat) = lazy_flat {
4105 let (results, flat_stats) = brute_force_flat_scan_sync(
4107 lazy_flat,
4108 self.alive_docs.as_deref(),
4109 query,
4110 params.unit_norm,
4111 fetch_k.min(lazy_flat.num_vectors),
4112 combiner,
4113 )?;
4114 crate::observe::dense_ann_scan(
4115 self.schema.index_label(),
4116 self.schema.get_field_name(field).unwrap_or("?"),
4117 "flat",
4118 flat_stats,
4119 );
4120 return Ok(results);
4121 } else {
4122 return Ok(Vec::new());
4123 };
4124 {
4125 let kind = dense_ann_kind_label(ann_index);
4126 let field_name = self.schema.get_field_name(field).unwrap_or("?");
4127 crate::observe::dense_ann_scan(self.schema.index_label(), field_name, kind, scan_stats);
4128 crate::observe::warn_non_finite_dense_scores(
4129 self.schema.index_label(),
4130 field_name,
4131 kind,
4132 scan_stats.non_finite_dropped,
4133 );
4134 }
4135
4136 if ann_index.is_some()
4138 && !results.is_empty()
4139 && let Some(lazy_flat) = lazy_flat
4140 {
4141 return exact_score_dense_candidate_documents_sync(
4142 &results,
4143 lazy_flat,
4144 query,
4145 params.unit_norm,
4146 combiner,
4147 k,
4148 );
4149 }
4150
4151 Ok(combine_grouped_ordinal_results(results, combiner, k))
4152 }
4153
4154 #[cfg(feature = "sync")]
4159 fn search_binary_dense_vector_sync_impl(
4160 &self,
4161 field: Field,
4162 query: &[u8],
4163 k: usize,
4164 combiner: crate::query::MultiValueCombiner,
4165 probe_cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
4166 ) -> Result<Vec<VectorSearchResult>> {
4167 let schema_dim = self.validate_binary_search_request(field, query)?;
4168 combiner.validate().map_err(Error::Query)?;
4169 if k == 0 {
4170 return Ok(Vec::new());
4171 }
4172 let t0 = crate::observe::Timer::start();
4173 if let Some(VectorIndex::ScannBinary(lazy)) = self.vector_indexes.get(&field.0) {
4174 let artifact = self
4175 .trained_vectors
4176 .scann_artifacts
4177 .get(&field.0)
4178 .ok_or_else(|| {
4179 Error::Schema(format!(
4180 "binary ScaNN field {} has no loaded global artifact",
4181 field.0
4182 ))
4183 })?;
4184 lazy.get()
4185 .validate_scann_generation(
4186 artifact.config(),
4187 artifact.generation(),
4188 artifact.artifact_id(),
4189 )
4190 .map_err(|error| {
4191 Error::Corruption(format!(
4192 "binary ScaNN generation mismatch for field {}: {error}",
4193 field.0
4194 ))
4195 })?;
4196 let config = self
4197 .schema
4198 .get_field_entry(field)
4199 .and_then(|entry| entry.binary_dense_vector_config.as_ref())
4200 .ok_or_else(|| {
4201 Error::Schema(format!(
4202 "binary ScaNN field {} has no schema configuration",
4203 field.0
4204 ))
4205 })?;
4206 let model = artifact.binary_model().map_err(Error::Io)?;
4207 let clusters = binary_scann_probe_clusters(&model, query, config.nprobe, probe_cache)?;
4208 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
4209 Error::Corruption(format!(
4210 "binary ScaNN field {} is missing flat vectors",
4211 field.0
4212 ))
4213 })?;
4214 let candidate_limit =
4215 checked_binary_combined_fetch_k(k)?.min(flat.num_docs_with_vectors());
4216 let (documents, ordinal_scores) = lazy
4217 .get()
4218 .search_binary_combined_documents(candidate_limit, query, &clusters, combiner)
4219 .map_err(|error| {
4220 Error::Corruption(format!(
4221 "invalid binary ScaNN payload for field {}: {error}",
4222 field.0
4223 ))
4224 })?;
4225 let results = exact_score_binary_candidate_document_ids_sync(
4226 documents
4227 .into_iter()
4228 .map(|candidate| candidate.doc_id)
4229 .collect(),
4230 &ordinal_scores,
4231 flat,
4232 query,
4233 schema_dim,
4234 combiner,
4235 k,
4236 )?;
4237 crate::observe::dense_l1(
4238 self.schema.index_label(),
4239 self.schema.get_field_name(field).unwrap_or("?"),
4240 "binary_scann",
4241 t0.secs(),
4242 results.len(),
4243 );
4244 return Ok(results);
4245 }
4246 if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0) {
4247 let ivf = lazy.get();
4248 let config = self
4249 .schema
4250 .get_field_entry(field)
4251 .and_then(|entry| entry.binary_dense_vector_config.as_ref())
4252 .ok_or_else(|| {
4253 Error::Schema(format!(
4254 "binary IVF field {} has no schema configuration",
4255 field.0
4256 ))
4257 })?;
4258 let quantizer = self
4259 .trained_vectors
4260 .binary_quantizers
4261 .get(&field.0)
4262 .ok_or_else(|| {
4263 Error::Schema(format!(
4264 "global binary IVF field {} has no loaded quantizer",
4265 field.0
4266 ))
4267 })?;
4268 validate_binary_ann(ivf, quantizer, config, schema_dim, field)?;
4269 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
4270 Error::Corruption(format!(
4271 "global binary IVF field {} is missing flat vector storage",
4272 field.0
4273 ))
4274 })?;
4275 let single_valued = flat.num_vectors == flat.num_docs_with_vectors();
4276 let clusters = binary_probe_clusters(
4277 quantizer,
4278 query,
4279 config.nprobe,
4280 config.ivf_routing,
4281 probe_cache,
4282 )?;
4283 let results = if !single_valued
4284 && !matches!(combiner, crate::query::MultiValueCombiner::Max)
4285 {
4286 let candidate_limit =
4287 checked_binary_combined_fetch_k(k)?.min(flat.num_docs_with_vectors());
4288 let (candidate_documents, probed_ordinal_scores) = ivf
4289 .search_binary_combined_documents(candidate_limit, query, &clusters, combiner)
4290 .map_err(|error| {
4291 Error::Corruption(format!(
4292 "invalid binary IVF payload for field {}: {error}",
4293 field.0,
4294 ))
4295 })?;
4296 exact_score_binary_candidate_document_ids_sync(
4297 candidate_documents
4298 .into_iter()
4299 .map(|candidate| candidate.doc_id)
4300 .collect(),
4301 &probed_ordinal_scores,
4302 flat,
4303 query,
4304 schema_dim,
4305 combiner,
4306 k,
4307 )?
4308 } else {
4309 let candidate_docs = if single_valued {
4310 k
4311 } else {
4312 checked_binary_combined_fetch_k(k)?
4313 }
4314 .min(flat.num_docs_with_vectors());
4315 let ann_results = if single_valued {
4316 ivf.search_binary_clusters::<false>(query, candidate_docs, &clusters)
4317 } else {
4318 ivf.search_binary_clusters::<true>(query, candidate_docs, &clusters)
4319 }
4320 .map_err(|error| {
4321 Error::Corruption(format!(
4322 "invalid binary IVF payload for field {}: {error}",
4323 field.0,
4324 ))
4325 })?;
4326 if single_valued {
4327 let ann_results = validate_binary_single_value_ann_results(ann_results, flat)?;
4328 combine_ordinal_results(ann_results, combiner, k)
4329 } else {
4330 exact_score_binary_candidate_documents_sync(
4331 &ann_results,
4332 flat,
4333 query,
4334 schema_dim,
4335 combiner,
4336 k,
4337 )?
4338 }
4339 };
4340 crate::observe::dense_l1(
4341 self.schema.index_label(),
4342 self.schema.get_field_name(field).unwrap_or("?"),
4343 "global_binary_ivf",
4344 t0.secs(),
4345 results.len(),
4346 );
4347 return Ok(results);
4348 }
4349 let lazy_flat = match self.flat_vectors.get(&field.0) {
4350 Some(f) => f,
4351 None => return Ok(Vec::new()),
4352 };
4353
4354 let dim_bits = lazy_flat.dim;
4355 let byte_len = lazy_flat.vector_byte_size();
4356 let n = lazy_flat.num_vectors;
4357
4358 if dim_bits != schema_dim {
4359 return Err(Error::Corruption(format!(
4360 "binary vector field {} has schema dimension {} but flat storage dimension {}",
4361 field.0, schema_dim, dim_bits
4362 )));
4363 }
4364
4365 if byte_len != query.len() {
4366 return Err(Error::Schema(format!(
4367 "Binary query vector byte length {} != field byte length {}",
4368 query.len(),
4369 byte_len
4370 )));
4371 }
4372
4373 let batch_len = bounded_vector_score_batch(byte_len, BINARY_SCORE_BATCH);
4374 let mut collector = FlatDocumentCollector::new(k, combiner);
4375 let mut scratch = DenseScratch::take();
4376 scratch.prepare(0, batch_len);
4377 let scores = &mut scratch.scores;
4378
4379 for batch_start in (0..n).step_by(batch_len) {
4380 let batch_count = batch_len.min(n - batch_start);
4381 let batch_bytes = lazy_flat
4382 .read_vectors_batch_sync(batch_start, batch_count)
4383 .map_err(crate::Error::Io)?;
4384 let raw = batch_bytes.as_slice();
4385
4386 crate::structures::simd::batch_hamming_scores(
4387 query,
4388 raw,
4389 byte_len,
4390 dim_bits,
4391 &mut scores[..batch_count],
4392 );
4393
4394 for (i, &score) in scores.iter().enumerate().take(batch_count) {
4395 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
4396 if self.is_alive(doc_id) {
4397 collector.push(doc_id, ordinal, score);
4398 }
4399 }
4400 }
4401
4402 let results = collector.into_results();
4403
4404 crate::observe::dense_l1(
4405 self.schema.index_label(),
4406 self.schema.get_field_name(field).unwrap_or("?"),
4407 "binary_flat",
4408 t0.secs(),
4409 results.len(),
4410 );
4411 Ok(results)
4412 }
4413
4414 #[cfg(feature = "sync")]
4415 pub fn search_binary_dense_vector_sync(
4416 &self,
4417 field: Field,
4418 query: &[u8],
4419 k: usize,
4420 combiner: crate::query::MultiValueCombiner,
4421 ) -> Result<Vec<VectorSearchResult>> {
4422 self.search_binary_dense_vector_sync_impl(field, query, k, combiner, None)
4423 }
4424
4425 #[cfg(feature = "sync")]
4426 pub(crate) fn search_binary_dense_vector_sync_with_probe_cache(
4427 &self,
4428 field: Field,
4429 query: &[u8],
4430 k: usize,
4431 combiner: crate::query::MultiValueCombiner,
4432 probe_cache: &std::sync::Mutex<Option<crate::structures::IvfProbePlan>>,
4433 ) -> Result<Vec<VectorSearchResult>> {
4434 self.search_binary_dense_vector_sync_impl(field, query, k, combiner, Some(probe_cache))
4435 }
4436}
4437
4438#[cfg(test)]
4439mod dense_search_safety_tests {
4440 use super::*;
4441
4442 #[test]
4443 fn dense_fetch_count_rejects_non_finite_and_unbounded_factors() {
4444 for factor in [
4445 f32::NAN,
4446 f32::INFINITY,
4447 f32::NEG_INFINITY,
4448 0.0,
4449 0.5,
4450 2.01,
4451 MAX_DENSE_RERANK_FACTOR + 1.0,
4452 ] {
4453 assert!(
4454 checked_dense_fetch_k(10, factor).is_err(),
4455 "factor={factor}"
4456 );
4457 }
4458 }
4459
4460 fn values_as_bytes<T>(values: &[T]) -> &[u8] {
4461 unsafe {
4462 std::slice::from_raw_parts(values.as_ptr() as *const u8, std::mem::size_of_val(values))
4463 }
4464 }
4465
4466 fn assert_prepared_dense_scores_match_legacy(
4467 quantization: DenseVectorQuantization,
4468 raw: &[u8],
4469 unit_norm: bool,
4470 ) {
4471 const DIM: usize = 4;
4472 const VECTOR_COUNT: usize = 4;
4473 let query = [0.25, -0.5, 0.75, 1.0];
4474 let mut expected = [0.0; VECTOR_COUNT];
4475 SegmentReader::score_quantized_batch_legacy(
4476 &query,
4477 raw,
4478 quantization,
4479 DIM,
4480 &mut expected,
4481 unit_norm,
4482 )
4483 .unwrap();
4484
4485 let prepared = PreparedDenseScoreQuery::new(&query, quantization, DIM, unit_norm).unwrap();
4486 let vector_bytes = DIM
4487 * match quantization {
4488 DenseVectorQuantization::F32 => std::mem::size_of::<f32>(),
4489 DenseVectorQuantization::F16 => std::mem::size_of::<u16>(),
4490 DenseVectorQuantization::UInt8 => 1,
4491 DenseVectorQuantization::Binary => unreachable!(),
4492 };
4493 let split = 2 * vector_bytes;
4494 let mut actual = [0.0; VECTOR_COUNT];
4495 prepared
4496 .score_batch(&raw[..split], &mut actual[..2])
4497 .unwrap();
4498 prepared
4499 .score_batch(&raw[split..], &mut actual[2..])
4500 .unwrap();
4501
4502 assert_eq!(
4503 actual.map(f32::to_bits),
4504 expected.map(f32::to_bits),
4505 "quantization={quantization:?}, unit_norm={unit_norm}"
4506 );
4507 }
4508
4509 #[test]
4510 fn prepared_dense_query_matches_legacy_scoring_across_batches() {
4511 let vectors_f32 = [
4512 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,
4513 -0.25,
4514 ];
4515 let vectors_f16: Vec<u16> = vectors_f32
4516 .iter()
4517 .map(|&value| crate::structures::simd::f32_to_f16(value))
4518 .collect();
4519 let vectors_u8 = [
4520 255, 96, 224, 160, 0, 192, 144, 128, 128, 128, 128, 128, 224, 192, 64, 96,
4521 ];
4522
4523 for unit_norm in [false, true] {
4524 assert_prepared_dense_scores_match_legacy(
4525 DenseVectorQuantization::F32,
4526 values_as_bytes(&vectors_f32),
4527 unit_norm,
4528 );
4529 assert_prepared_dense_scores_match_legacy(
4530 DenseVectorQuantization::F16,
4531 values_as_bytes(&vectors_f16),
4532 unit_norm,
4533 );
4534 assert_prepared_dense_scores_match_legacy(
4535 DenseVectorQuantization::UInt8,
4536 &vectors_u8,
4537 unit_norm,
4538 );
4539 }
4540 }
4541
4542 #[test]
4543 fn prepared_dense_query_preserves_scoring_validation_errors() {
4544 assert!(matches!(
4545 PreparedDenseScoreQuery::new(&[1.0], DenseVectorQuantization::F32, 2, false).err(),
4546 Some(Error::Query(_))
4547 ));
4548 assert!(matches!(
4549 PreparedDenseScoreQuery::new(&[1.0], DenseVectorQuantization::Binary, 1, false).err(),
4550 Some(Error::InvalidFieldType { .. })
4551 ));
4552
4553 let query = [1.0, 2.0];
4554 let prepared =
4555 PreparedDenseScoreQuery::new(&query, DenseVectorQuantization::F32, 2, false).unwrap();
4556 let mut scores = [0.0];
4557 assert!(matches!(
4558 prepared.score_batch(&[0; 7], &mut scores),
4559 Err(Error::Corruption(_))
4560 ));
4561 }
4562
4563 #[test]
4564 fn flat_document_collector_does_not_let_one_multivalue_doc_crowd_out_others() {
4565 let mut collector = FlatDocumentCollector::new(2, crate::query::MultiValueCombiner::Max);
4566 collector.push(1, 0, 1.0);
4567 collector.push(1, 1, 0.9);
4568 collector.push(2, 0, 0.8);
4569
4570 let results = collector.into_results();
4571 assert_eq!(
4572 results
4573 .iter()
4574 .map(|result| result.doc_id)
4575 .collect::<Vec<_>>(),
4576 vec![1, 2]
4577 );
4578 assert_eq!(results[0].ordinals.len(), 2);
4579 }
4580
4581 #[test]
4582 fn flat_document_collector_evicts_by_score_then_doc_id() {
4583 let mut collector = FlatDocumentCollector::new(2, crate::query::MultiValueCombiner::Max);
4584 collector.push(1, 0, 0.5);
4585 collector.push(3, 0, 0.8);
4586 collector.push(2, 0, 0.9);
4587 let results = collector.into_results();
4588 assert_eq!(
4589 results
4590 .iter()
4591 .map(|result| result.doc_id)
4592 .collect::<Vec<_>>(),
4593 vec![2, 3]
4594 );
4595
4596 let mut tied = FlatDocumentCollector::new(1, crate::query::MultiValueCombiner::Max);
4597 tied.push(2, 0, 1.0);
4598 tied.push(1, 0, 1.0);
4599 let results = tied.into_results();
4600 assert_eq!(results[0].doc_id, 1);
4601 }
4602
4603 #[test]
4604 fn dense_fetch_count_rounds_up_and_detects_overflow() {
4605 assert_eq!(checked_dense_fetch_k(3, 1.5).unwrap(), 5);
4606 assert_eq!(checked_dense_fetch_k(10_000, 2.0).unwrap(), 20_000);
4607 assert!(checked_dense_fetch_k(10_001, 2.0).is_err());
4608 assert!(checked_dense_fetch_k(usize::MAX, 2.0).is_err());
4609 }
4610
4611 #[test]
4612 fn binary_combined_fetch_count_uses_shared_bounded_oversampling() {
4613 assert_eq!(checked_binary_combined_fetch_k(3).unwrap(), 6);
4614 assert_eq!(checked_binary_combined_fetch_k(10_000).unwrap(), 20_000);
4615 assert_eq!(checked_binary_combined_fetch_k(10_001).unwrap(), 20_000);
4616 assert_eq!(checked_binary_combined_fetch_k(20_000).unwrap(), 20_000);
4617 assert!(checked_binary_combined_fetch_k(20_001).is_err());
4618 assert!(checked_binary_combined_fetch_k(usize::MAX).is_err());
4619 }
4620
4621 #[cfg(feature = "native")]
4622 #[test]
4623 fn legacy_ivf_tq_generation_is_rejected_while_opening() {
4624 use crate::directories::OwnedBytes;
4625 use crate::dsl::IvfRoutingMode;
4626 use crate::segment::ann_disk::{AnnDiskIndex, AnnKind};
4627
4628 let centroids = CoarseCentroids {
4629 num_clusters: 1,
4630 dim: 2,
4631 centroids: vec![1.0, 0.0],
4632 version: 7,
4633 soar_config: None,
4634 routing_index: None,
4635 };
4636 let mut build_centroids = centroids.clone();
4637 build_centroids.version =
4638 crate::structures::mark_ivf_tq_cosine_generation(build_centroids.version);
4639 let mut bytes = crate::segment::ann_build::build_ivf_tq(
4640 2,
4641 IvfRoutingMode::Flat,
4642 &build_centroids,
4643 &[(0, 0)],
4644 &[1.0, 0.0],
4645 )
4646 .unwrap();
4647 bytes[24..32].copy_from_slice(¢roids.version.to_le_bytes());
4650 let error = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::IvfTq, 1)
4651 .err()
4652 .expect("legacy IVF-TQ payload must fail while opening")
4653 .to_string();
4654 assert!(error.contains("unsupported legacy generation"), "{error}");
4655 }
4656
4657 #[test]
4658 fn rerank_batch_is_capped_by_actual_candidate_vectors() {
4659 assert_eq!(bounded_rerank_batch(3_072, DENSE_SCORE_BATCH, 20), 20);
4660 assert_eq!(
4661 bounded_rerank_batch(3_072, DENSE_SCORE_BATCH, 10_000),
4662 MAX_VECTOR_SCORE_BATCH_BYTES / 3_072
4663 );
4664 assert_eq!(bounded_rerank_batch(3_072, DENSE_SCORE_BATCH, 0), 1);
4665 }
4666
4667 #[test]
4668 fn file_ranges_reject_overflow_and_truncation() {
4669 assert_eq!(checked_file_range(4, 3, 7, "test").unwrap(), 4..7);
4670 assert!(checked_file_range(u64::MAX, 1, u64::MAX, "test").is_err());
4671 assert!(checked_file_range(5, 3, 7, "test").is_err());
4672 }
4673
4674 #[test]
4675 fn shared_tq_plan_cache_rebuilds_for_divergent_query_clones() {
4676 let codec = crate::structures::TqCodec::new(4);
4677 let cache = std::sync::Mutex::new(None);
4678 let original_query = vec![1.0, 2.0, 3.0, 4.0];
4679
4680 let original =
4681 cached_tq_query_plan(&codec, &original_query, Some(&cache)).expect("build plan");
4682 let reused =
4683 cached_tq_query_plan(&codec, &original_query, Some(&cache)).expect("reuse plan");
4684 assert!(
4685 std::sync::Arc::ptr_eq(&original, &reused),
4686 "unchanged queries must share their plan across segments"
4687 );
4688
4689 let mut divergent_clone = original_query.clone();
4690 divergent_clone[0] = -1.0;
4691 let rebuilt =
4692 cached_tq_query_plan(&codec, &divergent_clone, Some(&cache)).expect("rebuild plan");
4693 assert!(
4694 !std::sync::Arc::ptr_eq(&original, &rebuilt),
4695 "a clone with a mutated vector must not reuse stale LUTs"
4696 );
4697 assert!(rebuilt.matches_query(&divergent_clone));
4698 assert!(!rebuilt.matches_query(&original_query));
4699 }
4700
4701 #[test]
4702 fn candidate_vector_reads_coalesce_contiguous_values() {
4703 let mut runs = Vec::new();
4704 plan_vector_read_runs(&[3, 4, 5, 9, 12, 13], &mut runs).unwrap();
4705 assert_eq!(runs.len(), 3);
4706 assert!(matches!(
4707 runs.as_slice(),
4708 [
4709 VectorReadRun {
4710 buffer_start: 0,
4711 flat_start: 3,
4712 count: 3,
4713 },
4714 VectorReadRun {
4715 buffer_start: 3,
4716 flat_start: 9,
4717 count: 1,
4718 },
4719 VectorReadRun {
4720 buffer_start: 4,
4721 flat_start: 12,
4722 count: 2,
4723 },
4724 ]
4725 ));
4726 assert!(plan_vector_read_runs(&[3, 3], &mut runs).is_err());
4727 }
4728
4729 #[tokio::test]
4730 async fn binary_single_value_ann_fast_path_validates_and_deduplicates() {
4731 use crate::directories::{FileHandle, OwnedBytes};
4732 use crate::segment::FlatVectorData;
4733
4734 let mut encoded = Vec::new();
4735 FlatVectorData::serialize_binary_from_bits_streaming(
4736 8,
4737 &[0x0f, 0xf0],
4738 &[(1, 0), (3, 2)],
4739 &mut encoded,
4740 )
4741 .unwrap();
4742 let flat = LazyFlatVectorData::open_with_doc_limit(
4743 FileHandle::from_bytes(OwnedBytes::new(encoded)),
4744 Some(4),
4745 )
4746 .await
4747 .unwrap();
4748 assert_eq!(flat.num_vectors, flat.num_docs_with_vectors());
4749
4750 let validated = validate_binary_single_value_ann_results(
4751 vec![(3, 2, 0.9), (1, 0, 0.8), (3, 2, 0.7)],
4752 &flat,
4753 )
4754 .unwrap();
4755 assert_eq!(validated, vec![(3, 2, 0.9), (1, 0, 0.8)]);
4756
4757 assert!(matches!(
4758 validate_binary_single_value_ann_results(vec![(2, 0, 1.0)], &flat),
4759 Err(Error::Corruption(_))
4760 ));
4761 assert!(matches!(
4762 validate_binary_single_value_ann_results(vec![(3, 0, 1.0)], &flat),
4763 Err(Error::Corruption(_))
4764 ));
4765 }
4766
4767 #[tokio::test]
4768 async fn multivalue_ann_rerank_streams_past_document_candidate_cap() {
4769 use crate::directories::{FileHandle, OwnedBytes};
4770 use crate::segment::FlatVectorData;
4771
4772 const VALUES: usize = MAX_DENSE_CANDIDATES_PER_SEGMENT + 1;
4773 let mut encoded = Vec::new();
4774 let vectors = vec![1.0f32; VALUES];
4775 let doc_ids: Vec<_> = (0..VALUES).map(|ordinal| (0, ordinal as u16)).collect();
4776 FlatVectorData::serialize_binary_from_flat_streaming(
4777 1,
4778 &vectors,
4779 &doc_ids,
4780 DenseVectorQuantization::F32,
4781 &mut encoded,
4782 )
4783 .unwrap();
4784 let flat = LazyFlatVectorData::open_with_doc_limit(
4785 FileHandle::from_bytes(OwnedBytes::new(encoded)),
4786 Some(1),
4787 )
4788 .await
4789 .unwrap();
4790
4791 let (results, stats) = exact_score_dense_candidate_documents(
4792 &[(0, 0, 0.0)],
4793 &flat,
4794 &[1.0],
4795 false,
4796 crate::query::MultiValueCombiner::Max,
4797 1,
4798 )
4799 .await
4800 .unwrap();
4801 assert_eq!(stats.vector_count, VALUES);
4802 assert_eq!(results.len(), 1);
4803 assert_eq!(results[0].ordinals.len(), VALUES);
4804 assert!((results[0].score - 1.0).abs() < 1e-5);
4805 }
4806}