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, OwnedBytes};
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<OwnedBytes> {
2719 let range = checked_file_range(offset, len, self.postings_handle.len(), "posting")?;
2720 Ok(self.postings_handle.read_bytes_range(range).await?)
2721 }
2722
2723 pub async fn read_position_bytes(&self, offset: u64, len: u64) -> Result<Option<OwnedBytes>> {
2725 let handle = match &self.positions_handle {
2726 Some(h) => h,
2727 None => return Ok(None),
2728 };
2729 let range = checked_file_range(offset, len, handle.len(), "position")?;
2730 Ok(Some(handle.read_bytes_range(range).await?))
2731 }
2732
2733 pub fn has_positions_file(&self) -> bool {
2735 self.positions_handle.is_some()
2736 }
2737
2738 fn validate_dense_search_request(
2742 &self,
2743 field: Field,
2744 query: &[f32],
2745 nprobe: usize,
2746 rerank_factor: f32,
2747 combiner: crate::query::MultiValueCombiner,
2748 ) -> Result<DenseSearchParams> {
2749 let entry = self
2750 .schema
2751 .get_field_entry(field)
2752 .ok_or_else(|| Error::FieldNotFound(field.0.to_string()))?;
2753 if entry.field_type != crate::dsl::FieldType::DenseVector {
2754 return Err(Error::InvalidFieldType {
2755 expected: "dense_vector".to_string(),
2756 got: format!("{:?}", entry.field_type),
2757 });
2758 }
2759 let config = entry.dense_vector_config.as_ref().ok_or_else(|| {
2760 Error::Schema(format!(
2761 "dense vector field '{}' has no dense vector configuration",
2762 entry.name
2763 ))
2764 })?;
2765
2766 if query.is_empty() {
2767 return Err(Error::Query(format!(
2768 "dense query vector for field '{}' must not be empty",
2769 entry.name
2770 )));
2771 }
2772 if query.len() != config.dim {
2773 return Err(Error::Query(format!(
2774 "dense query vector dimension {} does not match field '{}' dimension {}",
2775 query.len(),
2776 entry.name,
2777 config.dim
2778 )));
2779 }
2780 if let Some((index, value)) = query
2781 .iter()
2782 .enumerate()
2783 .find(|(_, value)| !value.is_finite())
2784 {
2785 return Err(Error::Query(format!(
2786 "dense query vector for field '{}' contains non-finite value {value} at index {index}",
2787 entry.name
2788 )));
2789 }
2790
2791 let nprobe = match (nprobe, config.nprobe) {
2794 (0, 0) => 32,
2795 (0, schema_nprobe) => schema_nprobe,
2796 (query_nprobe, _) => query_nprobe,
2797 };
2798 if nprobe > MAX_DENSE_NPROBE {
2799 return Err(Error::Query(format!(
2800 "dense nprobe must be at most {MAX_DENSE_NPROBE}, got {nprobe}"
2801 )));
2802 }
2803
2804 checked_dense_fetch_k(0, rerank_factor)?;
2807 combiner.validate().map_err(Error::Query)?;
2808
2809 Ok(DenseSearchParams {
2810 dim: config.dim,
2811 nprobe,
2812 unit_norm: config.unit_norm,
2813 })
2814 }
2815
2816 fn validate_binary_search_request(&self, field: Field, query: &[u8]) -> Result<usize> {
2817 let entry = self
2818 .schema
2819 .get_field_entry(field)
2820 .ok_or_else(|| Error::FieldNotFound(field.0.to_string()))?;
2821 if entry.field_type != crate::dsl::FieldType::BinaryDenseVector {
2822 return Err(Error::InvalidFieldType {
2823 expected: "binary_dense_vector".to_string(),
2824 got: format!("{:?}", entry.field_type),
2825 });
2826 }
2827 let config = entry.binary_dense_vector_config.as_ref().ok_or_else(|| {
2828 Error::Schema(format!(
2829 "binary dense vector field '{}' has no configuration",
2830 entry.name
2831 ))
2832 })?;
2833 if config.dim == 0 || !config.dim.is_multiple_of(8) {
2834 return Err(Error::Schema(format!(
2835 "binary dense vector field '{}' has invalid dimension {}",
2836 entry.name, config.dim
2837 )));
2838 }
2839 if query.len() != config.byte_len() {
2840 return Err(Error::Query(format!(
2841 "binary query byte length {} does not match field '{}' byte length {}",
2842 query.len(),
2843 entry.name,
2844 config.byte_len()
2845 )));
2846 }
2847 Ok(config.dim)
2848 }
2849
2850 #[cfg(test)]
2852 fn score_quantized_batch_legacy(
2853 query: &[f32],
2854 raw: &[u8],
2855 quant: crate::dsl::DenseVectorQuantization,
2856 dim: usize,
2857 scores: &mut [f32],
2858 unit_norm: bool,
2859 ) -> Result<()> {
2860 use crate::dsl::DenseVectorQuantization;
2861 use crate::structures::simd;
2862
2863 if query.len() != dim {
2864 return Err(Error::Query(format!(
2865 "dense SIMD query dimension {} does not match vector dimension {dim}",
2866 query.len()
2867 )));
2868 }
2869 let element_size = match quant {
2870 DenseVectorQuantization::F32 => std::mem::size_of::<f32>(),
2871 DenseVectorQuantization::F16 => std::mem::size_of::<u16>(),
2872 DenseVectorQuantization::UInt8 => 1,
2873 DenseVectorQuantization::Binary => {
2874 return Err(Error::InvalidFieldType {
2875 expected: "non-binary dense vector".to_string(),
2876 got: "binary dense vector".to_string(),
2877 });
2878 }
2879 };
2880 let required_bytes = scores
2881 .len()
2882 .checked_mul(dim)
2883 .and_then(|elements| elements.checked_mul(element_size))
2884 .ok_or_else(|| Error::Corruption("dense vector batch byte length overflow".into()))?;
2885 if raw.len() < required_bytes {
2886 return Err(Error::Corruption(format!(
2887 "dense vector batch is truncated: need {required_bytes} bytes, got {}",
2888 raw.len()
2889 )));
2890 }
2891 if quant == DenseVectorQuantization::F16
2892 && required_bytes > 0
2893 && !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>())
2894 {
2895 return Err(Error::Corruption(
2896 "f16 vector data is not 2-byte aligned".to_string(),
2897 ));
2898 }
2899
2900 match (quant, unit_norm) {
2901 (DenseVectorQuantization::F32, false) => {
2902 let num_floats = scores.len() * dim;
2903 if !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()) {
2904 return Err(Error::Corruption(
2905 "f32 vector data is not 4-byte aligned".to_string(),
2906 ));
2907 }
2908 let vectors: &[f32] =
2909 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
2910 simd::batch_cosine_scores(query, vectors, dim, scores);
2911 }
2912 (DenseVectorQuantization::F32, true) => {
2913 let num_floats = scores.len() * dim;
2914 if !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()) {
2915 return Err(Error::Corruption(
2916 "f32 vector data is not 4-byte aligned".to_string(),
2917 ));
2918 }
2919 let vectors: &[f32] =
2920 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
2921 simd::batch_dot_scores(query, vectors, dim, scores);
2922 }
2923 (DenseVectorQuantization::F16, false) => {
2924 simd::batch_cosine_scores_f16(query, raw, dim, scores);
2925 }
2926 (DenseVectorQuantization::F16, true) => {
2927 simd::batch_dot_scores_f16(query, raw, dim, scores);
2928 }
2929 (DenseVectorQuantization::UInt8, false) => {
2930 simd::batch_cosine_scores_u8(query, raw, dim, scores);
2931 }
2932 (DenseVectorQuantization::UInt8, true) => {
2933 simd::batch_dot_scores_u8(query, raw, dim, scores);
2934 }
2935 (DenseVectorQuantization::Binary, _) => unreachable!("validated above"),
2936 }
2937 Ok(())
2938 }
2939
2940 pub async fn search_dense_vector(
2946 &self,
2947 field: Field,
2948 query: &[f32],
2949 k: usize,
2950 nprobe: usize,
2951 rerank_factor: f32,
2952 combiner: crate::query::MultiValueCombiner,
2953 ) -> Result<Vec<VectorSearchResult>> {
2954 self.search_dense_vector_impl(field, query, k, nprobe, rerank_factor, combiner, None)
2955 .await
2956 }
2957
2958 #[allow(clippy::too_many_arguments)]
2959 pub(crate) async fn search_dense_vector_with_probe_cache(
2960 &self,
2961 field: Field,
2962 query: &[f32],
2963 k: usize,
2964 nprobe: usize,
2965 rerank_factor: f32,
2966 combiner: crate::query::MultiValueCombiner,
2967 plan_cache: &DensePlanCache,
2968 ) -> Result<Vec<VectorSearchResult>> {
2969 self.search_dense_vector_impl(
2970 field,
2971 query,
2972 k,
2973 nprobe,
2974 rerank_factor,
2975 combiner,
2976 Some(plan_cache),
2977 )
2978 .await
2979 }
2980
2981 #[allow(clippy::too_many_arguments)]
2982 async fn search_dense_vector_impl(
2983 &self,
2984 field: Field,
2985 query: &[f32],
2986 k: usize,
2987 nprobe: usize,
2988 rerank_factor: f32,
2989 combiner: crate::query::MultiValueCombiner,
2990 plan_cache: Option<&DensePlanCache>,
2991 ) -> Result<Vec<VectorSearchResult>> {
2992 let params =
2993 self.validate_dense_search_request(field, query, nprobe, rerank_factor, combiner)?;
2994 let fetch_k = checked_dense_fetch_k(k, rerank_factor)?;
2995 if k == 0 {
2996 return Ok(Vec::new());
2997 }
2998
2999 let configured_ann_index = self.vector_indexes.get(&field.0);
3000 let lazy_flat = self.flat_vectors.get(&field.0);
3001 if configured_ann_index.is_none() && lazy_flat.is_none() {
3003 return Ok(Vec::new());
3004 }
3005
3006 if configured_ann_index.is_some() && lazy_flat.is_none() {
3007 return Err(Error::Corruption(format!(
3008 "dense ANN field {} is missing flat vector storage",
3009 field.0
3010 )));
3011 }
3012
3013 if let Some(flat) = lazy_flat
3014 && flat.dim != params.dim
3015 {
3016 return Err(Error::Corruption(format!(
3017 "dense vector field {} has schema dimension {} but flat storage dimension {}",
3018 field.0, params.dim, flat.dim
3019 )));
3020 }
3021
3022 let needs_document_aggregation = lazy_flat.is_some_and(|flat| {
3023 flat.num_vectors != flat.num_docs_with_vectors()
3024 && !matches!(combiner, crate::query::MultiValueCombiner::Max)
3025 });
3026 let ann_index = configured_ann_index;
3030
3031 let t0 = std::time::Instant::now();
3033 let mut flat_results = None;
3034 let (results, scan_stats): (Vec<(u32, u16, f32)>, DenseAnnScanStats) = if let Some(index) =
3035 ann_index
3036 {
3037 match index {
3039 VectorIndex::Tq { index: lazy, codec } => {
3040 let flat = lazy_flat.expect("ANN/flat pairing validated above");
3041 search_tq_segment(
3043 lazy.get(),
3044 codec,
3045 query,
3046 fetch_k.min(flat.num_docs_with_vectors()),
3047 needs_document_aggregation.then_some(combiner),
3048 field,
3049 params.dim,
3050 plan_cache.map(|cache| &cache.tq),
3051 ann_keys_are_unique(lazy.get(), flat),
3052 )?
3053 }
3054 VectorIndex::IvfTq { index: lazy, codec } => {
3055 let index = lazy.get();
3056 let centroids =
3057 self.trained_vectors
3058 .centroids
3059 .get(&field.0)
3060 .ok_or_else(|| {
3061 Error::Schema(format!(
3062 "IVF-TQ index requires coarse centroids for field {}",
3063 field.0
3064 ))
3065 })?;
3066 validate_coarse_centroids(centroids, params.dim)?;
3067 let routing = self
3068 .schema
3069 .get_field_entry(field)
3070 .and_then(|entry| entry.dense_vector_config.as_ref())
3071 .map_or(crate::dsl::IvfRoutingMode::Auto, |config| {
3072 config.ivf_routing
3073 });
3074 validate_ivf_tq_ann(index, centroids, codec, params.dim, routing, field)?;
3075 let flat = lazy_flat.expect("ANN/flat pairing validated above");
3076 search_ivf_tq_segment(
3077 index,
3078 centroids,
3079 codec,
3080 query,
3081 fetch_k.min(flat.num_docs_with_vectors()),
3082 needs_document_aggregation.then_some(combiner),
3083 field,
3084 params.nprobe,
3085 routing,
3086 plan_cache.map(|cache| &cache.ivf_tq),
3087 ann_keys_are_unique(index, flat),
3088 )?
3089 }
3090 VectorIndex::BinaryIvf(_) => {
3091 return Err(Error::Query(format!(
3094 "dense vector field '{}' is served by a binary IVF index; use BinaryDenseVectorQuery",
3095 self.schema.get_field_name(field).unwrap_or("?")
3096 )));
3097 }
3098 VectorIndex::ScannAh(lazy) => {
3099 let artifact = self
3100 .trained_vectors
3101 .scann_artifacts
3102 .get(&field.0)
3103 .ok_or_else(|| {
3104 Error::Schema(format!(
3105 "ScaNN field {} has no loaded global artifact",
3106 field.0
3107 ))
3108 })?;
3109 let flat = lazy_flat.expect("ANN/flat pairing validated above");
3110 search_scann_ah_segment(
3111 lazy.get(),
3112 artifact,
3113 query,
3114 fetch_k.min(flat.num_docs_with_vectors()),
3115 combiner,
3116 field,
3117 params.nprobe,
3118 plan_cache.map(|cache| &cache.scann),
3119 )
3120 .map(|candidates| (candidates, DenseAnnScanStats::default()))?
3121 }
3122 VectorIndex::ScannBinary(_) => {
3123 return Err(Error::Corruption(format!(
3124 "binary ScaNN payload was attached to float field {}",
3125 field.0
3126 )));
3127 }
3128 }
3129 } else if let Some(lazy_flat) = lazy_flat {
3130 log::debug!(
3134 "[dense_vector_search] index={} field {}: brute-force on {} vectors (dim={}, quant={:?})",
3135 self.schema.index_label(),
3136 field.0,
3137 lazy_flat.num_vectors,
3138 lazy_flat.dim,
3139 lazy_flat.quantization
3140 );
3141 let dim = lazy_flat.dim;
3142 let n = lazy_flat.num_vectors;
3143 let quant = lazy_flat.quantization;
3144 let batch_len =
3145 bounded_vector_score_batch(lazy_flat.vector_byte_size(), DENSE_SCORE_BATCH);
3146 let mut collector = FlatDocumentCollector::new(fetch_k.min(n), combiner);
3147 let prepared_query = PreparedDenseScoreQuery::new(query, quant, dim, params.unit_norm)?;
3148 let mut flat_stats = DenseAnnScanStats {
3149 posting_count: n,
3150 ..DenseAnnScanStats::default()
3151 };
3152 let mut scratch = DenseScratch::take();
3153 scratch.prepare(0, batch_len);
3154 let scores = &mut scratch.scores;
3155
3156 for batch_start in (0..n).step_by(batch_len) {
3157 let batch_count = batch_len.min(n - batch_start);
3158 let batch_bytes = lazy_flat
3159 .read_vectors_batch(batch_start, batch_count)
3160 .await
3161 .map_err(crate::Error::Io)?;
3162 let raw = batch_bytes.as_slice();
3163
3164 prepared_query.score_batch(raw, &mut scores[..batch_count])?;
3165 flat_stats.scored_blocks += 1;
3166
3167 for (i, &score) in scores.iter().enumerate().take(batch_count) {
3168 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
3169 collector.push(doc_id, ordinal, score);
3170 }
3171 }
3172
3173 flat_results = Some(collector.into_results());
3174 (Vec::new(), flat_stats)
3175 } else {
3176 return Ok(Vec::new());
3177 };
3178 let l1_elapsed = t0.elapsed();
3179 {
3180 let kind = dense_ann_kind_label(ann_index);
3181 let field_name = self.schema.get_field_name(field).unwrap_or("?");
3182 crate::observe::dense_l1(
3183 self.schema.index_label(),
3184 field_name,
3185 kind,
3186 l1_elapsed.as_secs_f64(),
3187 flat_results.as_ref().map_or(results.len(), Vec::len),
3188 );
3189 crate::observe::dense_ann_scan(self.schema.index_label(), field_name, kind, scan_stats);
3190 crate::observe::warn_non_finite_dense_scores(
3191 self.schema.index_label(),
3192 field_name,
3193 kind,
3194 scan_stats.non_finite_dropped,
3195 );
3196 }
3197 log::debug!(
3198 "[dense_vector_search] index={} field {}: L1 returned {} candidates in {:.1}ms",
3199 self.schema.index_label(),
3200 field.0,
3201 flat_results.as_ref().map_or(results.len(), Vec::len),
3202 l1_elapsed.as_secs_f64() * 1000.0
3203 );
3204
3205 if let Some(results) = flat_results {
3206 return Ok(results);
3207 }
3208
3209 if ann_index.is_some()
3212 && !results.is_empty()
3213 && let Some(lazy_flat) = lazy_flat
3214 {
3215 let t_rerank = std::time::Instant::now();
3216 let vbs = lazy_flat.vector_byte_size();
3217 let (reranked, stats) = exact_score_dense_candidate_documents(
3218 &results,
3219 lazy_flat,
3220 query,
3221 params.unit_norm,
3222 combiner,
3223 k,
3224 )
3225 .await?;
3226
3227 crate::observe::dense_rerank(
3228 self.schema.index_label(),
3229 self.schema.get_field_name(field).unwrap_or("?"),
3230 t_rerank.elapsed().as_secs_f64(),
3231 stats.resolve_elapsed.as_secs_f64(),
3232 stats.read_elapsed.as_secs_f64(),
3233 stats.vector_count,
3234 );
3235 log::debug!(
3236 "[dense_vector_search] index={} field {}: rerank {} vectors (dim={}, quant={:?}, bytes_per_vector={}): resolve={:.1}ms read={:.1}ms score={:.1}ms",
3237 self.schema.index_label(),
3238 field.0,
3239 stats.vector_count,
3240 lazy_flat.dim,
3241 lazy_flat.quantization,
3242 vbs,
3243 stats.resolve_elapsed.as_secs_f64() * 1000.0,
3244 stats.read_elapsed.as_secs_f64() * 1000.0,
3245 stats.score_elapsed.as_secs_f64() * 1000.0,
3246 );
3247
3248 log::debug!(
3249 "[dense_vector_search] index={} field {}: rerank total={:.1}ms",
3250 self.schema.index_label(),
3251 field.0,
3252 t_rerank.elapsed().as_secs_f64() * 1000.0
3253 );
3254 return Ok(reranked);
3255 }
3256
3257 Ok(combine_grouped_ordinal_results(results, combiner, k))
3258 }
3259
3260 async fn search_binary_dense_vector_impl(
3265 &self,
3266 field: Field,
3267 query: &[u8],
3268 k: usize,
3269 combiner: crate::query::MultiValueCombiner,
3270 probe_cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
3271 ) -> Result<Vec<VectorSearchResult>> {
3272 let schema_dim = self.validate_binary_search_request(field, query)?;
3273 combiner.validate().map_err(Error::Query)?;
3274 if k == 0 {
3275 return Ok(Vec::new());
3276 }
3277 let t0 = crate::observe::Timer::start();
3278 if let Some(VectorIndex::ScannBinary(lazy)) = self.vector_indexes.get(&field.0) {
3279 let artifact = self
3280 .trained_vectors
3281 .scann_artifacts
3282 .get(&field.0)
3283 .ok_or_else(|| {
3284 Error::Schema(format!(
3285 "binary ScaNN field {} has no loaded global artifact",
3286 field.0
3287 ))
3288 })?;
3289 lazy.get()
3290 .validate_scann_generation(
3291 artifact.config(),
3292 artifact.generation(),
3293 artifact.artifact_id(),
3294 )
3295 .map_err(|error| {
3296 Error::Corruption(format!(
3297 "binary ScaNN generation mismatch for field {}: {error}",
3298 field.0
3299 ))
3300 })?;
3301 let config = self
3302 .schema
3303 .get_field_entry(field)
3304 .and_then(|entry| entry.binary_dense_vector_config.as_ref())
3305 .ok_or_else(|| {
3306 Error::Schema(format!(
3307 "binary ScaNN field {} has no schema configuration",
3308 field.0
3309 ))
3310 })?;
3311 let model = artifact.binary_model().map_err(Error::Io)?;
3312 let clusters = binary_scann_probe_clusters(&model, query, config.nprobe, probe_cache)?;
3313 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
3314 Error::Corruption(format!(
3315 "binary ScaNN field {} is missing flat vectors",
3316 field.0
3317 ))
3318 })?;
3319 let candidate_limit =
3320 checked_binary_combined_fetch_k(k)?.min(flat.num_docs_with_vectors());
3321 let (documents, ordinal_scores) = lazy
3322 .get()
3323 .search_binary_combined_documents(candidate_limit, query, &clusters, combiner)
3324 .map_err(|error| {
3325 Error::Corruption(format!(
3326 "invalid binary ScaNN payload for field {}: {error}",
3327 field.0
3328 ))
3329 })?;
3330 let results = exact_score_binary_candidate_document_ids(
3331 documents
3332 .into_iter()
3333 .map(|candidate| candidate.doc_id)
3334 .collect(),
3335 &ordinal_scores,
3336 flat,
3337 query,
3338 schema_dim,
3339 combiner,
3340 k,
3341 )
3342 .await?;
3343 crate::observe::dense_l1(
3344 self.schema.index_label(),
3345 self.schema.get_field_name(field).unwrap_or("?"),
3346 "binary_scann",
3347 t0.secs(),
3348 results.len(),
3349 );
3350 return Ok(results);
3351 }
3352 if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0) {
3353 let ivf = lazy.get();
3354 let config = self
3355 .schema
3356 .get_field_entry(field)
3357 .and_then(|entry| entry.binary_dense_vector_config.as_ref())
3358 .ok_or_else(|| {
3359 Error::Schema(format!(
3360 "binary IVF field {} has no schema configuration",
3361 field.0
3362 ))
3363 })?;
3364 let quantizer = self
3365 .trained_vectors
3366 .binary_quantizers
3367 .get(&field.0)
3368 .ok_or_else(|| {
3369 Error::Schema(format!(
3370 "global binary IVF field {} has no loaded quantizer",
3371 field.0
3372 ))
3373 })?;
3374 validate_binary_ann(ivf, quantizer, config, schema_dim, field)?;
3375 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
3376 Error::Corruption(format!(
3377 "global binary IVF field {} is missing flat vector storage",
3378 field.0
3379 ))
3380 })?;
3381 let single_valued = flat.num_vectors == flat.num_docs_with_vectors();
3382 let clusters = binary_probe_clusters(
3383 quantizer,
3384 query,
3385 config.nprobe,
3386 config.ivf_routing,
3387 probe_cache,
3388 )?;
3389 let results = if !single_valued
3390 && !matches!(combiner, crate::query::MultiValueCombiner::Max)
3391 {
3392 let candidate_limit =
3393 checked_binary_combined_fetch_k(k)?.min(flat.num_docs_with_vectors());
3394 let (candidate_documents, probed_ordinal_scores) = ivf
3395 .search_binary_combined_documents(candidate_limit, query, &clusters, combiner)
3396 .map_err(|error| {
3397 Error::Corruption(format!(
3398 "invalid binary IVF payload for field {}: {error}",
3399 field.0,
3400 ))
3401 })?;
3402 exact_score_binary_candidate_document_ids(
3403 candidate_documents
3404 .into_iter()
3405 .map(|candidate| candidate.doc_id)
3406 .collect(),
3407 &probed_ordinal_scores,
3408 flat,
3409 query,
3410 schema_dim,
3411 combiner,
3412 k,
3413 )
3414 .await?
3415 } else {
3416 let candidate_docs = if single_valued {
3417 k
3418 } else {
3419 checked_binary_combined_fetch_k(k)?
3424 }
3425 .min(flat.num_docs_with_vectors());
3426 let ann_results = if single_valued {
3427 ivf.search_binary_clusters::<false>(query, candidate_docs, &clusters)
3428 } else {
3429 ivf.search_binary_clusters::<true>(query, candidate_docs, &clusters)
3430 }
3431 .map_err(|error| {
3432 Error::Corruption(format!(
3433 "invalid binary IVF payload for field {}: {error}",
3434 field.0,
3435 ))
3436 })?;
3437 if single_valued {
3440 let ann_results = validate_binary_single_value_ann_results(ann_results, flat)?;
3441 combine_ordinal_results(ann_results, combiner, k)
3442 } else {
3443 exact_score_binary_candidate_documents(
3444 &ann_results,
3445 flat,
3446 query,
3447 schema_dim,
3448 combiner,
3449 k,
3450 )
3451 .await?
3452 }
3453 };
3454 crate::observe::dense_l1(
3455 self.schema.index_label(),
3456 self.schema.get_field_name(field).unwrap_or("?"),
3457 "global_binary_ivf",
3458 t0.secs(),
3459 results.len(),
3460 );
3461 return Ok(results);
3462 }
3463 let lazy_flat = match self.flat_vectors.get(&field.0) {
3464 Some(f) => f,
3465 None => return Ok(Vec::new()),
3466 };
3467
3468 let dim_bits = lazy_flat.dim;
3469 let byte_len = lazy_flat.vector_byte_size();
3470 let n = lazy_flat.num_vectors;
3471
3472 if dim_bits != schema_dim {
3473 return Err(Error::Corruption(format!(
3474 "binary vector field {} has schema dimension {} but flat storage dimension {}",
3475 field.0, schema_dim, dim_bits
3476 )));
3477 }
3478
3479 if byte_len != query.len() {
3480 return Err(Error::Schema(format!(
3481 "Binary query vector byte length {} != field byte length {}",
3482 query.len(),
3483 byte_len
3484 )));
3485 }
3486
3487 let batch_len = bounded_vector_score_batch(byte_len, BINARY_SCORE_BATCH);
3488 let mut collector = FlatDocumentCollector::new(k, combiner);
3489 let mut scratch = DenseScratch::take();
3490 scratch.prepare(0, batch_len);
3491 let scores = &mut scratch.scores;
3492
3493 for batch_start in (0..n).step_by(batch_len) {
3494 let batch_count = batch_len.min(n - batch_start);
3495 let batch_bytes = lazy_flat
3496 .read_vectors_batch(batch_start, batch_count)
3497 .await
3498 .map_err(crate::Error::Io)?;
3499 let raw = batch_bytes.as_slice();
3500
3501 crate::structures::simd::batch_hamming_scores(
3502 query,
3503 raw,
3504 byte_len,
3505 dim_bits,
3506 &mut scores[..batch_count],
3507 );
3508
3509 for (i, &score) in scores.iter().enumerate().take(batch_count) {
3510 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
3511 collector.push(doc_id, ordinal, score);
3512 }
3513 }
3514
3515 let results = collector.into_results();
3516
3517 crate::observe::dense_l1(
3518 self.schema.index_label(),
3519 self.schema.get_field_name(field).unwrap_or("?"),
3520 "binary_flat",
3521 t0.secs(),
3522 results.len(),
3523 );
3524 Ok(results)
3525 }
3526
3527 pub async fn search_binary_dense_vector(
3528 &self,
3529 field: Field,
3530 query: &[u8],
3531 k: usize,
3532 combiner: crate::query::MultiValueCombiner,
3533 ) -> Result<Vec<VectorSearchResult>> {
3534 self.search_binary_dense_vector_impl(field, query, k, combiner, None)
3535 .await
3536 }
3537
3538 pub(crate) async fn search_binary_dense_vector_with_probe_cache(
3539 &self,
3540 field: Field,
3541 query: &[u8],
3542 k: usize,
3543 combiner: crate::query::MultiValueCombiner,
3544 probe_cache: &std::sync::Mutex<Option<crate::structures::IvfProbePlan>>,
3545 ) -> Result<Vec<VectorSearchResult>> {
3546 self.search_binary_dense_vector_impl(field, query, k, combiner, Some(probe_cache))
3547 .await
3548 }
3549
3550 pub fn coarse_centroids(&self, field_id: u32) -> Option<&Arc<CoarseCentroids>> {
3552 self.trained_vectors.centroids.get(&field_id)
3553 }
3554
3555 pub fn set_trained_vectors(
3556 &mut self,
3557 trained_vectors: Arc<crate::segment::TrainedVectorStructures>,
3558 ) {
3559 self.trained_vectors = trained_vectors;
3560 }
3561
3562 pub fn get_vector_index(&self, field: Field) -> Option<&VectorIndex> {
3564 self.vector_indexes.get(&field.0)
3565 }
3566
3567 pub async fn get_positions(
3572 &self,
3573 field: Field,
3574 term: &[u8],
3575 ) -> Result<Option<crate::structures::TermPositions>> {
3576 let handle = match &self.positions_handle {
3578 Some(h) => h,
3579 None => return Ok(None),
3580 };
3581
3582 let mut key = Vec::with_capacity(4 + term.len());
3584 key.extend_from_slice(&field.0.to_le_bytes());
3585 key.extend_from_slice(term);
3586
3587 let term_info = match self.term_dict.get(&key).await? {
3589 Some(info) => info,
3590 None => return Ok(None),
3591 };
3592
3593 let (offset, length) = match term_info.position_info() {
3595 Some((o, l)) => (o, l),
3596 None => return Ok(None),
3597 };
3598
3599 let range = checked_file_range(offset, length, handle.len(), "position list")?;
3603 let data = handle.read_bytes_range(range).await?;
3606 Ok(Some(crate::structures::TermPositions::open(data)?))
3607 }
3608
3609 pub fn has_positions(&self, field: Field) -> bool {
3611 if let Some(entry) = self.schema.get_field_entry(field) {
3613 entry.positions.is_some()
3614 } else {
3615 false
3616 }
3617 }
3618}
3619
3620#[cfg(feature = "sync")]
3622impl SegmentReader {
3623 pub fn text_doc_freq_sync(&self, field: Field, term: &[u8]) -> Result<u32> {
3626 let mut key = Vec::with_capacity(4 + term.len());
3627 key.extend_from_slice(&field.0.to_le_bytes());
3628 key.extend_from_slice(term);
3629 Ok(self
3630 .term_dict
3631 .get_sync(&key)?
3632 .map_or(0, |info| info.doc_freq()))
3633 }
3634
3635 pub fn get_postings_sync(&self, field: Field, term: &[u8]) -> Result<Option<BlockPostingList>> {
3637 let mut key = Vec::with_capacity(4 + term.len());
3639 key.extend_from_slice(&field.0.to_le_bytes());
3640 key.extend_from_slice(term);
3641
3642 let term_info = match self.term_dict.get_sync(&key)? {
3644 Some(info) => info,
3645 None => return Ok(None),
3646 };
3647
3648 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
3650 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
3651 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
3652 posting_list.push(doc_id, tf);
3653 }
3654 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
3655 return Ok(Some(block_list));
3656 }
3657
3658 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
3660 Error::Corruption("TermInfo has neither inline nor external data".to_string())
3661 })?;
3662
3663 let range = checked_file_range(
3664 posting_offset,
3665 posting_len,
3666 self.postings_handle.len(),
3667 "posting",
3668 )?;
3669 let posting_bytes = self.postings_handle.read_bytes_range_sync(range)?;
3670 let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
3671
3672 Ok(Some(block_list))
3673 }
3674
3675 pub fn get_prefix_postings_sync(
3677 &self,
3678 field: Field,
3679 prefix: &[u8],
3680 ) -> Result<Vec<BlockPostingList>> {
3681 if prefix.is_empty() {
3682 return Err(Error::Query("prefix must not be empty".into()));
3683 }
3684 let mut key_prefix = Vec::with_capacity(4 + prefix.len());
3685 key_prefix.extend_from_slice(&field.0.to_le_bytes());
3686 key_prefix.extend_from_slice(prefix);
3687
3688 let (entries, truncated) = self
3689 .term_dict
3690 .prefix_scan_limited_sync(&key_prefix, MAX_PREFIX_TERMS)?;
3691 if truncated {
3692 return Err(Error::Query(format!(
3693 "prefix expands to more than {MAX_PREFIX_TERMS} terms"
3694 )));
3695 }
3696 let posting_count: u64 = entries
3697 .iter()
3698 .map(|(_, term_info)| term_info.doc_freq() as u64)
3699 .sum();
3700 if posting_count > MAX_PREFIX_POSTINGS {
3701 return Err(Error::Query(format!(
3702 "prefix expands to {posting_count} postings (maximum {MAX_PREFIX_POSTINGS})"
3703 )));
3704 }
3705 let mut results = Vec::with_capacity(entries.len());
3706
3707 for (_key, term_info) in entries {
3708 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
3709 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
3710 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
3711 posting_list.push(doc_id, tf);
3712 }
3713 results.push(BlockPostingList::from_posting_list(&posting_list)?);
3714 } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
3715 let range = checked_file_range(
3716 posting_offset,
3717 posting_len,
3718 self.postings_handle.len(),
3719 "prefix posting",
3720 )?;
3721 let posting_bytes = self.postings_handle.read_bytes_range_sync(range)?;
3722 results.push(BlockPostingList::deserialize_zero_copy(posting_bytes)?);
3723 }
3724 }
3725
3726 Ok(results)
3727 }
3728
3729 pub fn get_positions_sync(
3731 &self,
3732 field: Field,
3733 term: &[u8],
3734 ) -> Result<Option<crate::structures::TermPositions>> {
3735 let handle = match &self.positions_handle {
3736 Some(h) => h,
3737 None => return Ok(None),
3738 };
3739
3740 let mut key = Vec::with_capacity(4 + term.len());
3742 key.extend_from_slice(&field.0.to_le_bytes());
3743 key.extend_from_slice(term);
3744
3745 let term_info = match self.term_dict.get_sync(&key)? {
3747 Some(info) => info,
3748 None => return Ok(None),
3749 };
3750
3751 let (offset, length) = match term_info.position_info() {
3752 Some((o, l)) => (o, l),
3753 None => return Ok(None),
3754 };
3755
3756 let range = checked_file_range(offset, length, handle.len(), "position list")?;
3757 let data = handle.read_bytes_range_sync(range)?;
3758 let pos_list = crate::structures::TermPositions::open(data)?;
3759 Ok(Some(pos_list))
3760 }
3761
3762 pub fn search_dense_vector_sync(
3765 &self,
3766 field: Field,
3767 query: &[f32],
3768 k: usize,
3769 nprobe: usize,
3770 rerank_factor: f32,
3771 combiner: crate::query::MultiValueCombiner,
3772 ) -> Result<Vec<VectorSearchResult>> {
3773 self.search_dense_vector_sync_impl(field, query, k, nprobe, rerank_factor, combiner, None)
3774 }
3775
3776 #[cfg(feature = "sync")]
3777 #[allow(clippy::too_many_arguments)]
3778 pub(crate) fn search_dense_vector_sync_with_probe_cache(
3779 &self,
3780 field: Field,
3781 query: &[f32],
3782 k: usize,
3783 nprobe: usize,
3784 rerank_factor: f32,
3785 combiner: crate::query::MultiValueCombiner,
3786 plan_cache: &DensePlanCache,
3787 ) -> Result<Vec<VectorSearchResult>> {
3788 self.search_dense_vector_sync_impl(
3789 field,
3790 query,
3791 k,
3792 nprobe,
3793 rerank_factor,
3794 combiner,
3795 Some(plan_cache),
3796 )
3797 }
3798
3799 #[cfg(feature = "sync")]
3800 #[allow(clippy::too_many_arguments)]
3801 fn search_dense_vector_sync_impl(
3802 &self,
3803 field: Field,
3804 query: &[f32],
3805 k: usize,
3806 nprobe: usize,
3807 rerank_factor: f32,
3808 combiner: crate::query::MultiValueCombiner,
3809 plan_cache: Option<&DensePlanCache>,
3810 ) -> Result<Vec<VectorSearchResult>> {
3811 let params =
3812 self.validate_dense_search_request(field, query, nprobe, rerank_factor, combiner)?;
3813 let fetch_k = checked_dense_fetch_k(k, rerank_factor)?;
3814 if k == 0 {
3815 return Ok(Vec::new());
3816 }
3817
3818 let configured_ann_index = self.vector_indexes.get(&field.0);
3819 let lazy_flat = self.flat_vectors.get(&field.0);
3820 if configured_ann_index.is_none() && lazy_flat.is_none() {
3821 return Ok(Vec::new());
3822 }
3823
3824 if configured_ann_index.is_some() && lazy_flat.is_none() {
3825 return Err(Error::Corruption(format!(
3826 "dense ANN field {} is missing flat vector storage",
3827 field.0
3828 )));
3829 }
3830
3831 if let Some(flat) = lazy_flat
3832 && flat.dim != params.dim
3833 {
3834 return Err(Error::Corruption(format!(
3835 "dense vector field {} has schema dimension {} but flat storage dimension {}",
3836 field.0, params.dim, flat.dim
3837 )));
3838 }
3839
3840 let needs_document_aggregation = lazy_flat.is_some_and(|flat| {
3841 flat.num_vectors != flat.num_docs_with_vectors()
3842 && !matches!(combiner, crate::query::MultiValueCombiner::Max)
3843 });
3844 let ann_index = configured_ann_index;
3847
3848 let (results, scan_stats): (Vec<(u32, u16, f32)>, DenseAnnScanStats) = if let Some(index) =
3849 ann_index
3850 {
3851 match index {
3853 VectorIndex::Tq { index: lazy, codec } => {
3854 let flat = lazy_flat.expect("ANN/flat pairing validated above");
3855 search_tq_segment(
3856 lazy.get(),
3857 codec,
3858 query,
3859 fetch_k.min(flat.num_docs_with_vectors()),
3860 needs_document_aggregation.then_some(combiner),
3861 field,
3862 params.dim,
3863 plan_cache.map(|cache| &cache.tq),
3864 ann_keys_are_unique(lazy.get(), flat),
3865 )?
3866 }
3867 VectorIndex::IvfTq { index: lazy, codec } => {
3868 let index = lazy.get();
3869 let centroids =
3870 self.trained_vectors
3871 .centroids
3872 .get(&field.0)
3873 .ok_or_else(|| {
3874 Error::Schema(format!(
3875 "IVF-TQ index requires coarse centroids for field {}",
3876 field.0
3877 ))
3878 })?;
3879 validate_coarse_centroids(centroids, params.dim)?;
3880 let routing = self
3881 .schema
3882 .get_field_entry(field)
3883 .and_then(|entry| entry.dense_vector_config.as_ref())
3884 .map_or(crate::dsl::IvfRoutingMode::Auto, |config| {
3885 config.ivf_routing
3886 });
3887 validate_ivf_tq_ann(index, centroids, codec, params.dim, routing, field)?;
3888 let flat = lazy_flat.expect("ANN/flat pairing validated above");
3889 search_ivf_tq_segment(
3890 index,
3891 centroids,
3892 codec,
3893 query,
3894 fetch_k.min(flat.num_docs_with_vectors()),
3895 needs_document_aggregation.then_some(combiner),
3896 field,
3897 params.nprobe,
3898 routing,
3899 plan_cache.map(|cache| &cache.ivf_tq),
3900 ann_keys_are_unique(index, flat),
3901 )?
3902 }
3903 VectorIndex::BinaryIvf(_) => {
3904 return Err(Error::Query(format!(
3907 "dense vector field '{}' is served by a binary IVF index; use BinaryDenseVectorQuery",
3908 self.schema.get_field_name(field).unwrap_or("?")
3909 )));
3910 }
3911 VectorIndex::ScannAh(lazy) => {
3912 let artifact = self
3913 .trained_vectors
3914 .scann_artifacts
3915 .get(&field.0)
3916 .ok_or_else(|| {
3917 Error::Schema(format!(
3918 "ScaNN field {} has no loaded global artifact",
3919 field.0
3920 ))
3921 })?;
3922 let flat = lazy_flat.expect("ANN/flat pairing validated above");
3923 search_scann_ah_segment(
3924 lazy.get(),
3925 artifact,
3926 query,
3927 fetch_k.min(flat.num_docs_with_vectors()),
3928 combiner,
3929 field,
3930 params.nprobe,
3931 plan_cache.map(|cache| &cache.scann),
3932 )
3933 .map(|candidates| (candidates, DenseAnnScanStats::default()))?
3934 }
3935 VectorIndex::ScannBinary(_) => {
3936 return Err(Error::Corruption(format!(
3937 "binary ScaNN payload was attached to float field {}",
3938 field.0
3939 )));
3940 }
3941 }
3942 } else if let Some(lazy_flat) = lazy_flat {
3943 let (results, flat_stats) = brute_force_flat_scan_sync(
3945 lazy_flat,
3946 query,
3947 params.unit_norm,
3948 fetch_k.min(lazy_flat.num_vectors),
3949 combiner,
3950 )?;
3951 crate::observe::dense_ann_scan(
3952 self.schema.index_label(),
3953 self.schema.get_field_name(field).unwrap_or("?"),
3954 "flat",
3955 flat_stats,
3956 );
3957 return Ok(results);
3958 } else {
3959 return Ok(Vec::new());
3960 };
3961 {
3962 let kind = dense_ann_kind_label(ann_index);
3963 let field_name = self.schema.get_field_name(field).unwrap_or("?");
3964 crate::observe::dense_ann_scan(self.schema.index_label(), field_name, kind, scan_stats);
3965 crate::observe::warn_non_finite_dense_scores(
3966 self.schema.index_label(),
3967 field_name,
3968 kind,
3969 scan_stats.non_finite_dropped,
3970 );
3971 }
3972
3973 if ann_index.is_some()
3975 && !results.is_empty()
3976 && let Some(lazy_flat) = lazy_flat
3977 {
3978 return exact_score_dense_candidate_documents_sync(
3979 &results,
3980 lazy_flat,
3981 query,
3982 params.unit_norm,
3983 combiner,
3984 k,
3985 );
3986 }
3987
3988 Ok(combine_grouped_ordinal_results(results, combiner, k))
3989 }
3990
3991 #[cfg(feature = "sync")]
3996 fn search_binary_dense_vector_sync_impl(
3997 &self,
3998 field: Field,
3999 query: &[u8],
4000 k: usize,
4001 combiner: crate::query::MultiValueCombiner,
4002 probe_cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
4003 ) -> Result<Vec<VectorSearchResult>> {
4004 let schema_dim = self.validate_binary_search_request(field, query)?;
4005 combiner.validate().map_err(Error::Query)?;
4006 if k == 0 {
4007 return Ok(Vec::new());
4008 }
4009 let t0 = crate::observe::Timer::start();
4010 if let Some(VectorIndex::ScannBinary(lazy)) = self.vector_indexes.get(&field.0) {
4011 let artifact = self
4012 .trained_vectors
4013 .scann_artifacts
4014 .get(&field.0)
4015 .ok_or_else(|| {
4016 Error::Schema(format!(
4017 "binary ScaNN field {} has no loaded global artifact",
4018 field.0
4019 ))
4020 })?;
4021 lazy.get()
4022 .validate_scann_generation(
4023 artifact.config(),
4024 artifact.generation(),
4025 artifact.artifact_id(),
4026 )
4027 .map_err(|error| {
4028 Error::Corruption(format!(
4029 "binary ScaNN generation mismatch for field {}: {error}",
4030 field.0
4031 ))
4032 })?;
4033 let config = self
4034 .schema
4035 .get_field_entry(field)
4036 .and_then(|entry| entry.binary_dense_vector_config.as_ref())
4037 .ok_or_else(|| {
4038 Error::Schema(format!(
4039 "binary ScaNN field {} has no schema configuration",
4040 field.0
4041 ))
4042 })?;
4043 let model = artifact.binary_model().map_err(Error::Io)?;
4044 let clusters = binary_scann_probe_clusters(&model, query, config.nprobe, probe_cache)?;
4045 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
4046 Error::Corruption(format!(
4047 "binary ScaNN field {} is missing flat vectors",
4048 field.0
4049 ))
4050 })?;
4051 let candidate_limit =
4052 checked_binary_combined_fetch_k(k)?.min(flat.num_docs_with_vectors());
4053 let (documents, ordinal_scores) = lazy
4054 .get()
4055 .search_binary_combined_documents(candidate_limit, query, &clusters, combiner)
4056 .map_err(|error| {
4057 Error::Corruption(format!(
4058 "invalid binary ScaNN payload for field {}: {error}",
4059 field.0
4060 ))
4061 })?;
4062 let results = exact_score_binary_candidate_document_ids_sync(
4063 documents
4064 .into_iter()
4065 .map(|candidate| candidate.doc_id)
4066 .collect(),
4067 &ordinal_scores,
4068 flat,
4069 query,
4070 schema_dim,
4071 combiner,
4072 k,
4073 )?;
4074 crate::observe::dense_l1(
4075 self.schema.index_label(),
4076 self.schema.get_field_name(field).unwrap_or("?"),
4077 "binary_scann",
4078 t0.secs(),
4079 results.len(),
4080 );
4081 return Ok(results);
4082 }
4083 if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0) {
4084 let ivf = lazy.get();
4085 let config = self
4086 .schema
4087 .get_field_entry(field)
4088 .and_then(|entry| entry.binary_dense_vector_config.as_ref())
4089 .ok_or_else(|| {
4090 Error::Schema(format!(
4091 "binary IVF field {} has no schema configuration",
4092 field.0
4093 ))
4094 })?;
4095 let quantizer = self
4096 .trained_vectors
4097 .binary_quantizers
4098 .get(&field.0)
4099 .ok_or_else(|| {
4100 Error::Schema(format!(
4101 "global binary IVF field {} has no loaded quantizer",
4102 field.0
4103 ))
4104 })?;
4105 validate_binary_ann(ivf, quantizer, config, schema_dim, field)?;
4106 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
4107 Error::Corruption(format!(
4108 "global binary IVF field {} is missing flat vector storage",
4109 field.0
4110 ))
4111 })?;
4112 let single_valued = flat.num_vectors == flat.num_docs_with_vectors();
4113 let clusters = binary_probe_clusters(
4114 quantizer,
4115 query,
4116 config.nprobe,
4117 config.ivf_routing,
4118 probe_cache,
4119 )?;
4120 let results = if !single_valued
4121 && !matches!(combiner, crate::query::MultiValueCombiner::Max)
4122 {
4123 let candidate_limit =
4124 checked_binary_combined_fetch_k(k)?.min(flat.num_docs_with_vectors());
4125 let (candidate_documents, probed_ordinal_scores) = ivf
4126 .search_binary_combined_documents(candidate_limit, query, &clusters, combiner)
4127 .map_err(|error| {
4128 Error::Corruption(format!(
4129 "invalid binary IVF payload for field {}: {error}",
4130 field.0,
4131 ))
4132 })?;
4133 exact_score_binary_candidate_document_ids_sync(
4134 candidate_documents
4135 .into_iter()
4136 .map(|candidate| candidate.doc_id)
4137 .collect(),
4138 &probed_ordinal_scores,
4139 flat,
4140 query,
4141 schema_dim,
4142 combiner,
4143 k,
4144 )?
4145 } else {
4146 let candidate_docs = if single_valued {
4147 k
4148 } else {
4149 checked_binary_combined_fetch_k(k)?
4150 }
4151 .min(flat.num_docs_with_vectors());
4152 let ann_results = if single_valued {
4153 ivf.search_binary_clusters::<false>(query, candidate_docs, &clusters)
4154 } else {
4155 ivf.search_binary_clusters::<true>(query, candidate_docs, &clusters)
4156 }
4157 .map_err(|error| {
4158 Error::Corruption(format!(
4159 "invalid binary IVF payload for field {}: {error}",
4160 field.0,
4161 ))
4162 })?;
4163 if single_valued {
4164 let ann_results = validate_binary_single_value_ann_results(ann_results, flat)?;
4165 combine_ordinal_results(ann_results, combiner, k)
4166 } else {
4167 exact_score_binary_candidate_documents_sync(
4168 &ann_results,
4169 flat,
4170 query,
4171 schema_dim,
4172 combiner,
4173 k,
4174 )?
4175 }
4176 };
4177 crate::observe::dense_l1(
4178 self.schema.index_label(),
4179 self.schema.get_field_name(field).unwrap_or("?"),
4180 "global_binary_ivf",
4181 t0.secs(),
4182 results.len(),
4183 );
4184 return Ok(results);
4185 }
4186 let lazy_flat = match self.flat_vectors.get(&field.0) {
4187 Some(f) => f,
4188 None => return Ok(Vec::new()),
4189 };
4190
4191 let dim_bits = lazy_flat.dim;
4192 let byte_len = lazy_flat.vector_byte_size();
4193 let n = lazy_flat.num_vectors;
4194
4195 if dim_bits != schema_dim {
4196 return Err(Error::Corruption(format!(
4197 "binary vector field {} has schema dimension {} but flat storage dimension {}",
4198 field.0, schema_dim, dim_bits
4199 )));
4200 }
4201
4202 if byte_len != query.len() {
4203 return Err(Error::Schema(format!(
4204 "Binary query vector byte length {} != field byte length {}",
4205 query.len(),
4206 byte_len
4207 )));
4208 }
4209
4210 let batch_len = bounded_vector_score_batch(byte_len, BINARY_SCORE_BATCH);
4211 let mut collector = FlatDocumentCollector::new(k, combiner);
4212 let mut scratch = DenseScratch::take();
4213 scratch.prepare(0, batch_len);
4214 let scores = &mut scratch.scores;
4215
4216 for batch_start in (0..n).step_by(batch_len) {
4217 let batch_count = batch_len.min(n - batch_start);
4218 let batch_bytes = lazy_flat
4219 .read_vectors_batch_sync(batch_start, batch_count)
4220 .map_err(crate::Error::Io)?;
4221 let raw = batch_bytes.as_slice();
4222
4223 crate::structures::simd::batch_hamming_scores(
4224 query,
4225 raw,
4226 byte_len,
4227 dim_bits,
4228 &mut scores[..batch_count],
4229 );
4230
4231 for (i, &score) in scores.iter().enumerate().take(batch_count) {
4232 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
4233 collector.push(doc_id, ordinal, score);
4234 }
4235 }
4236
4237 let results = collector.into_results();
4238
4239 crate::observe::dense_l1(
4240 self.schema.index_label(),
4241 self.schema.get_field_name(field).unwrap_or("?"),
4242 "binary_flat",
4243 t0.secs(),
4244 results.len(),
4245 );
4246 Ok(results)
4247 }
4248
4249 #[cfg(feature = "sync")]
4250 pub fn search_binary_dense_vector_sync(
4251 &self,
4252 field: Field,
4253 query: &[u8],
4254 k: usize,
4255 combiner: crate::query::MultiValueCombiner,
4256 ) -> Result<Vec<VectorSearchResult>> {
4257 self.search_binary_dense_vector_sync_impl(field, query, k, combiner, None)
4258 }
4259
4260 #[cfg(feature = "sync")]
4261 pub(crate) fn search_binary_dense_vector_sync_with_probe_cache(
4262 &self,
4263 field: Field,
4264 query: &[u8],
4265 k: usize,
4266 combiner: crate::query::MultiValueCombiner,
4267 probe_cache: &std::sync::Mutex<Option<crate::structures::IvfProbePlan>>,
4268 ) -> Result<Vec<VectorSearchResult>> {
4269 self.search_binary_dense_vector_sync_impl(field, query, k, combiner, Some(probe_cache))
4270 }
4271}
4272
4273#[cfg(test)]
4274mod dense_search_safety_tests {
4275 use super::*;
4276
4277 #[test]
4278 fn dense_fetch_count_rejects_non_finite_and_unbounded_factors() {
4279 for factor in [
4280 f32::NAN,
4281 f32::INFINITY,
4282 f32::NEG_INFINITY,
4283 0.0,
4284 0.5,
4285 2.01,
4286 MAX_DENSE_RERANK_FACTOR + 1.0,
4287 ] {
4288 assert!(
4289 checked_dense_fetch_k(10, factor).is_err(),
4290 "factor={factor}"
4291 );
4292 }
4293 }
4294
4295 fn values_as_bytes<T>(values: &[T]) -> &[u8] {
4296 unsafe {
4297 std::slice::from_raw_parts(values.as_ptr() as *const u8, std::mem::size_of_val(values))
4298 }
4299 }
4300
4301 fn assert_prepared_dense_scores_match_legacy(
4302 quantization: DenseVectorQuantization,
4303 raw: &[u8],
4304 unit_norm: bool,
4305 ) {
4306 const DIM: usize = 4;
4307 const VECTOR_COUNT: usize = 4;
4308 let query = [0.25, -0.5, 0.75, 1.0];
4309 let mut expected = [0.0; VECTOR_COUNT];
4310 SegmentReader::score_quantized_batch_legacy(
4311 &query,
4312 raw,
4313 quantization,
4314 DIM,
4315 &mut expected,
4316 unit_norm,
4317 )
4318 .unwrap();
4319
4320 let prepared = PreparedDenseScoreQuery::new(&query, quantization, DIM, unit_norm).unwrap();
4321 let vector_bytes = DIM
4322 * match quantization {
4323 DenseVectorQuantization::F32 => std::mem::size_of::<f32>(),
4324 DenseVectorQuantization::F16 => std::mem::size_of::<u16>(),
4325 DenseVectorQuantization::UInt8 => 1,
4326 DenseVectorQuantization::Binary => unreachable!(),
4327 };
4328 let split = 2 * vector_bytes;
4329 let mut actual = [0.0; VECTOR_COUNT];
4330 prepared
4331 .score_batch(&raw[..split], &mut actual[..2])
4332 .unwrap();
4333 prepared
4334 .score_batch(&raw[split..], &mut actual[2..])
4335 .unwrap();
4336
4337 assert_eq!(
4338 actual.map(f32::to_bits),
4339 expected.map(f32::to_bits),
4340 "quantization={quantization:?}, unit_norm={unit_norm}"
4341 );
4342 }
4343
4344 #[test]
4345 fn prepared_dense_query_matches_legacy_scoring_across_batches() {
4346 let vectors_f32 = [
4347 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,
4348 -0.25,
4349 ];
4350 let vectors_f16: Vec<u16> = vectors_f32
4351 .iter()
4352 .map(|&value| crate::structures::simd::f32_to_f16(value))
4353 .collect();
4354 let vectors_u8 = [
4355 255, 96, 224, 160, 0, 192, 144, 128, 128, 128, 128, 128, 224, 192, 64, 96,
4356 ];
4357
4358 for unit_norm in [false, true] {
4359 assert_prepared_dense_scores_match_legacy(
4360 DenseVectorQuantization::F32,
4361 values_as_bytes(&vectors_f32),
4362 unit_norm,
4363 );
4364 assert_prepared_dense_scores_match_legacy(
4365 DenseVectorQuantization::F16,
4366 values_as_bytes(&vectors_f16),
4367 unit_norm,
4368 );
4369 assert_prepared_dense_scores_match_legacy(
4370 DenseVectorQuantization::UInt8,
4371 &vectors_u8,
4372 unit_norm,
4373 );
4374 }
4375 }
4376
4377 #[test]
4378 fn prepared_dense_query_preserves_scoring_validation_errors() {
4379 assert!(matches!(
4380 PreparedDenseScoreQuery::new(&[1.0], DenseVectorQuantization::F32, 2, false).err(),
4381 Some(Error::Query(_))
4382 ));
4383 assert!(matches!(
4384 PreparedDenseScoreQuery::new(&[1.0], DenseVectorQuantization::Binary, 1, false).err(),
4385 Some(Error::InvalidFieldType { .. })
4386 ));
4387
4388 let query = [1.0, 2.0];
4389 let prepared =
4390 PreparedDenseScoreQuery::new(&query, DenseVectorQuantization::F32, 2, false).unwrap();
4391 let mut scores = [0.0];
4392 assert!(matches!(
4393 prepared.score_batch(&[0; 7], &mut scores),
4394 Err(Error::Corruption(_))
4395 ));
4396 }
4397
4398 #[test]
4399 fn flat_document_collector_does_not_let_one_multivalue_doc_crowd_out_others() {
4400 let mut collector = FlatDocumentCollector::new(2, crate::query::MultiValueCombiner::Max);
4401 collector.push(1, 0, 1.0);
4402 collector.push(1, 1, 0.9);
4403 collector.push(2, 0, 0.8);
4404
4405 let results = collector.into_results();
4406 assert_eq!(
4407 results
4408 .iter()
4409 .map(|result| result.doc_id)
4410 .collect::<Vec<_>>(),
4411 vec![1, 2]
4412 );
4413 assert_eq!(results[0].ordinals.len(), 2);
4414 }
4415
4416 #[test]
4417 fn flat_document_collector_evicts_by_score_then_doc_id() {
4418 let mut collector = FlatDocumentCollector::new(2, crate::query::MultiValueCombiner::Max);
4419 collector.push(1, 0, 0.5);
4420 collector.push(3, 0, 0.8);
4421 collector.push(2, 0, 0.9);
4422 let results = collector.into_results();
4423 assert_eq!(
4424 results
4425 .iter()
4426 .map(|result| result.doc_id)
4427 .collect::<Vec<_>>(),
4428 vec![2, 3]
4429 );
4430
4431 let mut tied = FlatDocumentCollector::new(1, crate::query::MultiValueCombiner::Max);
4432 tied.push(2, 0, 1.0);
4433 tied.push(1, 0, 1.0);
4434 let results = tied.into_results();
4435 assert_eq!(results[0].doc_id, 1);
4436 }
4437
4438 #[test]
4439 fn dense_fetch_count_rounds_up_and_detects_overflow() {
4440 assert_eq!(checked_dense_fetch_k(3, 1.5).unwrap(), 5);
4441 assert_eq!(checked_dense_fetch_k(10_000, 2.0).unwrap(), 20_000);
4442 assert!(checked_dense_fetch_k(10_001, 2.0).is_err());
4443 assert!(checked_dense_fetch_k(usize::MAX, 2.0).is_err());
4444 }
4445
4446 #[test]
4447 fn binary_combined_fetch_count_uses_shared_bounded_oversampling() {
4448 assert_eq!(checked_binary_combined_fetch_k(3).unwrap(), 6);
4449 assert_eq!(checked_binary_combined_fetch_k(10_000).unwrap(), 20_000);
4450 assert_eq!(checked_binary_combined_fetch_k(10_001).unwrap(), 20_000);
4451 assert_eq!(checked_binary_combined_fetch_k(20_000).unwrap(), 20_000);
4452 assert!(checked_binary_combined_fetch_k(20_001).is_err());
4453 assert!(checked_binary_combined_fetch_k(usize::MAX).is_err());
4454 }
4455
4456 #[cfg(feature = "native")]
4457 #[test]
4458 fn legacy_ivf_tq_generation_is_rejected_while_opening() {
4459 use crate::directories::OwnedBytes;
4460 use crate::dsl::IvfRoutingMode;
4461 use crate::segment::ann_disk::{AnnDiskIndex, AnnKind};
4462
4463 let centroids = CoarseCentroids {
4464 num_clusters: 1,
4465 dim: 2,
4466 centroids: vec![1.0, 0.0],
4467 version: 7,
4468 soar_config: None,
4469 routing_index: None,
4470 };
4471 let mut build_centroids = centroids.clone();
4472 build_centroids.version =
4473 crate::structures::mark_ivf_tq_cosine_generation(build_centroids.version);
4474 let mut bytes = crate::segment::ann_build::build_ivf_tq(
4475 2,
4476 IvfRoutingMode::Flat,
4477 &build_centroids,
4478 &[(0, 0)],
4479 &[1.0, 0.0],
4480 )
4481 .unwrap();
4482 bytes[24..32].copy_from_slice(¢roids.version.to_le_bytes());
4485 let error = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::IvfTq, 1)
4486 .err()
4487 .expect("legacy IVF-TQ payload must fail while opening")
4488 .to_string();
4489 assert!(error.contains("unsupported legacy generation"), "{error}");
4490 }
4491
4492 #[test]
4493 fn rerank_batch_is_capped_by_actual_candidate_vectors() {
4494 assert_eq!(bounded_rerank_batch(3_072, DENSE_SCORE_BATCH, 20), 20);
4495 assert_eq!(
4496 bounded_rerank_batch(3_072, DENSE_SCORE_BATCH, 10_000),
4497 MAX_VECTOR_SCORE_BATCH_BYTES / 3_072
4498 );
4499 assert_eq!(bounded_rerank_batch(3_072, DENSE_SCORE_BATCH, 0), 1);
4500 }
4501
4502 #[test]
4503 fn file_ranges_reject_overflow_and_truncation() {
4504 assert_eq!(checked_file_range(4, 3, 7, "test").unwrap(), 4..7);
4505 assert!(checked_file_range(u64::MAX, 1, u64::MAX, "test").is_err());
4506 assert!(checked_file_range(5, 3, 7, "test").is_err());
4507 }
4508
4509 #[test]
4510 fn shared_tq_plan_cache_rebuilds_for_divergent_query_clones() {
4511 let codec = crate::structures::TqCodec::new(4);
4512 let cache = std::sync::Mutex::new(None);
4513 let original_query = vec![1.0, 2.0, 3.0, 4.0];
4514
4515 let original =
4516 cached_tq_query_plan(&codec, &original_query, Some(&cache)).expect("build plan");
4517 let reused =
4518 cached_tq_query_plan(&codec, &original_query, Some(&cache)).expect("reuse plan");
4519 assert!(
4520 std::sync::Arc::ptr_eq(&original, &reused),
4521 "unchanged queries must share their plan across segments"
4522 );
4523
4524 let mut divergent_clone = original_query.clone();
4525 divergent_clone[0] = -1.0;
4526 let rebuilt =
4527 cached_tq_query_plan(&codec, &divergent_clone, Some(&cache)).expect("rebuild plan");
4528 assert!(
4529 !std::sync::Arc::ptr_eq(&original, &rebuilt),
4530 "a clone with a mutated vector must not reuse stale LUTs"
4531 );
4532 assert!(rebuilt.matches_query(&divergent_clone));
4533 assert!(!rebuilt.matches_query(&original_query));
4534 }
4535
4536 #[test]
4537 fn candidate_vector_reads_coalesce_contiguous_values() {
4538 let mut runs = Vec::new();
4539 plan_vector_read_runs(&[3, 4, 5, 9, 12, 13], &mut runs).unwrap();
4540 assert_eq!(runs.len(), 3);
4541 assert!(matches!(
4542 runs.as_slice(),
4543 [
4544 VectorReadRun {
4545 buffer_start: 0,
4546 flat_start: 3,
4547 count: 3,
4548 },
4549 VectorReadRun {
4550 buffer_start: 3,
4551 flat_start: 9,
4552 count: 1,
4553 },
4554 VectorReadRun {
4555 buffer_start: 4,
4556 flat_start: 12,
4557 count: 2,
4558 },
4559 ]
4560 ));
4561 assert!(plan_vector_read_runs(&[3, 3], &mut runs).is_err());
4562 }
4563
4564 #[tokio::test]
4565 async fn binary_single_value_ann_fast_path_validates_and_deduplicates() {
4566 use crate::directories::{FileHandle, OwnedBytes};
4567 use crate::segment::FlatVectorData;
4568
4569 let mut encoded = Vec::new();
4570 FlatVectorData::serialize_binary_from_bits_streaming(
4571 8,
4572 &[0x0f, 0xf0],
4573 &[(1, 0), (3, 2)],
4574 &mut encoded,
4575 )
4576 .unwrap();
4577 let flat = LazyFlatVectorData::open_with_doc_limit(
4578 FileHandle::from_bytes(OwnedBytes::new(encoded)),
4579 Some(4),
4580 )
4581 .await
4582 .unwrap();
4583 assert_eq!(flat.num_vectors, flat.num_docs_with_vectors());
4584
4585 let validated = validate_binary_single_value_ann_results(
4586 vec![(3, 2, 0.9), (1, 0, 0.8), (3, 2, 0.7)],
4587 &flat,
4588 )
4589 .unwrap();
4590 assert_eq!(validated, vec![(3, 2, 0.9), (1, 0, 0.8)]);
4591
4592 assert!(matches!(
4593 validate_binary_single_value_ann_results(vec![(2, 0, 1.0)], &flat),
4594 Err(Error::Corruption(_))
4595 ));
4596 assert!(matches!(
4597 validate_binary_single_value_ann_results(vec![(3, 0, 1.0)], &flat),
4598 Err(Error::Corruption(_))
4599 ));
4600 }
4601
4602 #[tokio::test]
4603 async fn multivalue_ann_rerank_streams_past_document_candidate_cap() {
4604 use crate::directories::{FileHandle, OwnedBytes};
4605 use crate::segment::FlatVectorData;
4606
4607 const VALUES: usize = MAX_DENSE_CANDIDATES_PER_SEGMENT + 1;
4608 let mut encoded = Vec::new();
4609 let vectors = vec![1.0f32; VALUES];
4610 let doc_ids: Vec<_> = (0..VALUES).map(|ordinal| (0, ordinal as u16)).collect();
4611 FlatVectorData::serialize_binary_from_flat_streaming(
4612 1,
4613 &vectors,
4614 &doc_ids,
4615 DenseVectorQuantization::F32,
4616 &mut encoded,
4617 )
4618 .unwrap();
4619 let flat = LazyFlatVectorData::open_with_doc_limit(
4620 FileHandle::from_bytes(OwnedBytes::new(encoded)),
4621 Some(1),
4622 )
4623 .await
4624 .unwrap();
4625
4626 let (results, stats) = exact_score_dense_candidate_documents(
4627 &[(0, 0, 0.0)],
4628 &flat,
4629 &[1.0],
4630 false,
4631 crate::query::MultiValueCombiner::Max,
4632 1,
4633 )
4634 .await
4635 .unwrap();
4636 assert_eq!(stats.vector_count, VALUES);
4637 assert_eq!(results.len(), 1);
4638 assert_eq!(results[0].ordinals.len(), VALUES);
4639 assert!((results[0].score - 1.0).abs() < 1e-5);
4640 }
4641}