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