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