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