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