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 flat: &LazyFlatVectorData,
942 query: &[u8],
943 dim_bits: usize,
944 combiner: crate::query::MultiValueCombiner,
945 limit: usize,
946) -> Result<Vec<VectorSearchResult>> {
947 let documents = ann_candidate_document_ranges_from_ids(candidate_doc_ids, flat)?;
948 let probe_scores = FxHashMap::default();
949 exact_score_binary_resolved_documents(
950 documents,
951 &probe_scores,
952 flat,
953 query,
954 dim_bits,
955 combiner,
956 limit,
957 )
958 .await
959}
960
961async fn exact_score_binary_resolved_documents(
962 documents: AnnCandidateDocuments,
963 probe_scores: &FxHashMap<(DocId, u16), f32>,
964 flat: &LazyFlatVectorData,
965 query: &[u8],
966 dim_bits: usize,
967 combiner: crate::query::MultiValueCombiner,
968 limit: usize,
969) -> Result<Vec<VectorSearchResult>> {
970 let vector_byte_size = flat.vector_byte_size();
971 let batch_len =
972 bounded_rerank_batch(vector_byte_size, BINARY_SCORE_BATCH, documents.vector_count);
973 let raw_capacity = batch_len
974 .checked_mul(vector_byte_size)
975 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
976 let mut raw = vec![0u8; raw_capacity];
977 let mut scores = vec![0.0f32; batch_len];
978 let mut batch_scores = vec![0.0f32; batch_len];
979 let mut batch = Vec::with_capacity(batch_len);
980 let mut unresolved = Vec::with_capacity(batch_len);
981 let mut unresolved_flat_indexes = Vec::with_capacity(batch_len);
982 let mut read_runs = Vec::new();
983 let mut cursor = CandidateVectorCursor::new(&documents.ranges);
984 let mut collector = FlatDocumentCollector::new(limit, combiner);
985 let mut scored = 0usize;
986
987 while cursor.fill_batch(flat, &mut batch, batch_len)? {
988 unresolved.clear();
989 for (batch_index, &(doc_id, ordinal, flat_index)) in batch.iter().enumerate() {
990 if let Some(&score) = probe_scores.get(&(doc_id, ordinal)) {
991 batch_scores[batch_index] = score;
992 } else {
993 unresolved.push((batch_index, flat_index));
994 }
995 }
996 unresolved_flat_indexes.clear();
997 unresolved_flat_indexes.extend(unresolved.iter().map(|&(_, flat_index)| flat_index));
998 let raw_len = unresolved
999 .len()
1000 .checked_mul(vector_byte_size)
1001 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
1002 let raw = &mut raw[..raw_len];
1003 read_vector_runs(flat, &unresolved_flat_indexes, &mut read_runs, raw).await?;
1004 crate::structures::simd::batch_hamming_scores(
1005 query,
1006 raw,
1007 vector_byte_size,
1008 dim_bits,
1009 &mut scores[..unresolved.len()],
1010 );
1011 for (buffer_index, &(batch_index, _)) in unresolved.iter().enumerate() {
1012 batch_scores[batch_index] = scores[buffer_index];
1013 }
1014 for (batch_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
1015 collector.push(doc_id, ordinal, batch_scores[batch_index]);
1016 }
1017 scored += batch.len();
1018 }
1019 debug_assert_eq!(scored, documents.vector_count);
1020 Ok(collector.into_results())
1021}
1022
1023#[cfg(feature = "sync")]
1024fn exact_score_binary_candidate_documents_sync(
1025 ann_results: &[RawVectorCandidate],
1026 flat: &LazyFlatVectorData,
1027 query: &[u8],
1028 dim_bits: usize,
1029 combiner: crate::query::MultiValueCombiner,
1030 limit: usize,
1031) -> Result<Vec<VectorSearchResult>> {
1032 let documents = ann_candidate_document_ranges(ann_results, flat)?;
1033 let probe_scores: FxHashMap<(DocId, u16), f32> = ann_results
1034 .iter()
1035 .map(|&(doc_id, ordinal, score)| ((doc_id, ordinal), score))
1036 .collect();
1037 exact_score_binary_resolved_documents_sync(
1038 documents,
1039 &probe_scores,
1040 flat,
1041 query,
1042 dim_bits,
1043 combiner,
1044 limit,
1045 )
1046}
1047
1048#[cfg(feature = "sync")]
1049fn exact_score_binary_candidate_document_ids_sync(
1050 candidate_doc_ids: Vec<DocId>,
1051 flat: &LazyFlatVectorData,
1052 query: &[u8],
1053 dim_bits: usize,
1054 combiner: crate::query::MultiValueCombiner,
1055 limit: usize,
1056) -> Result<Vec<VectorSearchResult>> {
1057 let documents = ann_candidate_document_ranges_from_ids(candidate_doc_ids, flat)?;
1058 let probe_scores = FxHashMap::default();
1059 exact_score_binary_resolved_documents_sync(
1060 documents,
1061 &probe_scores,
1062 flat,
1063 query,
1064 dim_bits,
1065 combiner,
1066 limit,
1067 )
1068}
1069
1070#[cfg(feature = "sync")]
1071fn exact_score_binary_resolved_documents_sync(
1072 documents: AnnCandidateDocuments,
1073 probe_scores: &FxHashMap<(DocId, u16), f32>,
1074 flat: &LazyFlatVectorData,
1075 query: &[u8],
1076 dim_bits: usize,
1077 combiner: crate::query::MultiValueCombiner,
1078 limit: usize,
1079) -> Result<Vec<VectorSearchResult>> {
1080 let vector_byte_size = flat.vector_byte_size();
1081 let batch_len =
1082 bounded_rerank_batch(vector_byte_size, BINARY_SCORE_BATCH, documents.vector_count);
1083 let raw_capacity = batch_len
1084 .checked_mul(vector_byte_size)
1085 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
1086 let mut raw = vec![0u8; raw_capacity];
1087 let mut scores = vec![0.0f32; batch_len];
1088 let mut batch_scores = vec![0.0f32; batch_len];
1089 let mut batch = Vec::with_capacity(batch_len);
1090 let mut unresolved = Vec::with_capacity(batch_len);
1091 let mut unresolved_flat_indexes = Vec::with_capacity(batch_len);
1092 let mut read_runs = Vec::new();
1093 let mut cursor = CandidateVectorCursor::new(&documents.ranges);
1094 let mut collector = FlatDocumentCollector::new(limit, combiner);
1095 let mut scored = 0usize;
1096
1097 while cursor.fill_batch(flat, &mut batch, batch_len)? {
1098 unresolved.clear();
1099 for (batch_index, &(doc_id, ordinal, flat_index)) in batch.iter().enumerate() {
1100 if let Some(&score) = probe_scores.get(&(doc_id, ordinal)) {
1101 batch_scores[batch_index] = score;
1102 } else {
1103 unresolved.push((batch_index, flat_index));
1104 }
1105 }
1106 unresolved_flat_indexes.clear();
1107 unresolved_flat_indexes.extend(unresolved.iter().map(|&(_, flat_index)| flat_index));
1108 let raw_len = unresolved
1109 .len()
1110 .checked_mul(vector_byte_size)
1111 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
1112 let raw = &mut raw[..raw_len];
1113 read_vector_runs_sync(flat, &unresolved_flat_indexes, &mut read_runs, raw)?;
1114 crate::structures::simd::batch_hamming_scores(
1115 query,
1116 raw,
1117 vector_byte_size,
1118 dim_bits,
1119 &mut scores[..unresolved.len()],
1120 );
1121 for (buffer_index, &(batch_index, _)) in unresolved.iter().enumerate() {
1122 batch_scores[batch_index] = scores[buffer_index];
1123 }
1124 for (batch_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
1125 collector.push(doc_id, ordinal, batch_scores[batch_index]);
1126 }
1127 scored += batch.len();
1128 }
1129 debug_assert_eq!(scored, documents.vector_count);
1130 Ok(collector.into_results())
1131}
1132
1133fn validate_coarse_centroids(centroids: &CoarseCentroids, dim: usize) -> Result<()> {
1134 let expected = (centroids.num_clusters as usize)
1135 .checked_mul(dim)
1136 .ok_or_else(|| Error::Corruption("coarse centroid size overflow".into()))?;
1137 if centroids.num_clusters == 0
1138 || centroids.dim != dim
1139 || centroids.centroids.len() != expected
1140 || centroids.centroids.iter().any(|value| !value.is_finite())
1141 {
1142 return Err(Error::Corruption(format!(
1143 "invalid coarse centroids: clusters={}, dim={}, values={} (expected dim={dim}, values={expected})",
1144 centroids.num_clusters,
1145 centroids.dim,
1146 centroids.centroids.len()
1147 )));
1148 }
1149 Ok(())
1150}
1151
1152#[derive(Debug, Default)]
1157pub struct DensePlanCache {
1158 pub(crate) tq: std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqQueryPlan>>>,
1159 pub(crate) ivf_tq: std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqIvfQueryPlan>>>,
1160}
1161
1162#[allow(clippy::too_many_arguments)]
1167fn search_tq_segment(
1168 index: &crate::segment::ann_disk::AnnDiskIndex,
1169 codec: &crate::structures::TqCodec,
1170 query: &[f32],
1171 fetch_k: usize,
1172 document_combiner: Option<crate::query::MultiValueCombiner>,
1173 field: Field,
1174 dim: usize,
1175 plan_cache: Option<&std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqQueryPlan>>>>,
1176) -> Result<Vec<RawVectorCandidate>> {
1177 validate_tq_ann(index, codec, dim, field)?;
1178 let plan = cached_tq_query_plan(codec, query, plan_cache)?;
1179 match document_combiner {
1180 Some(combiner) => index
1181 .search_tq_combined_documents(fetch_k, &plan, combiner)
1182 .map(|candidates| {
1183 candidates
1184 .into_iter()
1185 .map(|candidate| (candidate.doc_id, 0, 0.0))
1189 .collect()
1190 }),
1191 None => index.search_tq_distinct(fetch_k, &plan),
1192 }
1193 .map_err(|error| {
1194 Error::Corruption(format!("invalid TQ payload for field {}: {error}", field.0))
1195 })
1196}
1197
1198fn cached_tq_query_plan(
1201 codec: &crate::structures::TqCodec,
1202 query: &[f32],
1203 plan_cache: Option<&std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqQueryPlan>>>>,
1204) -> Result<std::sync::Arc<crate::structures::TqQueryPlan>> {
1205 Ok(match plan_cache {
1206 Some(cache) => {
1207 let mut cached = cache
1208 .lock()
1209 .map_err(|_| Error::Internal("TQ plan cache is poisoned".into()))?;
1210 match cached.as_ref() {
1211 Some(plan)
1212 if plan.fingerprint() == codec.fingerprint() && plan.matches_query(query) =>
1213 {
1214 std::sync::Arc::clone(plan)
1215 }
1216 _ => {
1217 let plan =
1218 std::sync::Arc::new(crate::structures::TqQueryPlan::build(codec, query));
1219 *cached = Some(std::sync::Arc::clone(&plan));
1220 plan
1221 }
1222 }
1223 }
1224 None => std::sync::Arc::new(crate::structures::TqQueryPlan::build(codec, query)),
1225 })
1226}
1227
1228fn validate_tq_ann(
1229 index: &crate::segment::ann_disk::AnnDiskIndex,
1230 codec: &crate::structures::TqCodec,
1231 dim: usize,
1232 field: Field,
1233) -> Result<()> {
1234 let header = index.header();
1235 if header.dim != dim
1236 || codec.dim() != dim
1237 || header.code_size != codec.code_size()
1238 || header.quantizer_version != codec.fingerprint()
1239 || header.codebook_version != 0
1240 || header.num_clusters != 1
1241 {
1242 return Err(Error::Corruption(format!(
1243 "TQ payload for field {} does not match the codec derived from schema dimension {dim}",
1244 field.0,
1245 )));
1246 }
1247 Ok(())
1248}
1249
1250#[allow(clippy::too_many_arguments)]
1254fn search_ivf_tq_segment(
1255 index: &crate::segment::ann_disk::AnnDiskIndex,
1256 centroids: &CoarseCentroids,
1257 codec: &crate::structures::TqCodec,
1258 query: &[f32],
1259 fetch_k: usize,
1260 document_combiner: Option<crate::query::MultiValueCombiner>,
1261 field: Field,
1262 nprobe: usize,
1263 routing: crate::dsl::IvfRoutingMode,
1264 plan_cache: Option<
1265 &std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqIvfQueryPlan>>>,
1266 >,
1267) -> Result<Vec<RawVectorCandidate>> {
1268 let effective_nprobe = nprobe.clamp(1, centroids.num_clusters as usize);
1269 let request_fingerprint = crate::structures::TqIvfQueryPlan::request_fingerprint_for(
1270 centroids,
1271 query,
1272 effective_nprobe,
1273 routing,
1274 );
1275 let build = || {
1276 std::sync::Arc::new(crate::structures::TqIvfQueryPlan::build(
1277 centroids,
1278 codec,
1279 query,
1280 effective_nprobe,
1281 routing,
1282 ))
1283 };
1284 let plan = match plan_cache {
1285 Some(cache) => {
1286 let mut cached = cache
1287 .lock()
1288 .map_err(|_| Error::Internal("IVF-TQ plan cache is poisoned".into()))?;
1289 match cached.as_ref() {
1290 Some(plan)
1291 if plan.quantizer_version == centroids.version
1292 && plan.fingerprint == codec.fingerprint()
1293 && plan.request_fingerprint == request_fingerprint
1294 && plan.cluster_ids.len() == effective_nprobe =>
1295 {
1296 std::sync::Arc::clone(plan)
1297 }
1298 _ => {
1299 let plan = build();
1300 *cached = Some(std::sync::Arc::clone(&plan));
1301 plan
1302 }
1303 }
1304 }
1305 None => build(),
1306 };
1307 let candidates = match document_combiner {
1308 Some(combiner) => index
1309 .search_ivf_tq_combined_documents(fetch_k, &plan, combiner)
1310 .map(|documents| {
1311 documents
1312 .into_iter()
1313 .map(|candidate| (candidate.doc_id, 0, 0.0))
1317 .collect()
1318 }),
1319 None => index.search_ivf_tq_distinct(fetch_k, &plan),
1320 };
1321 candidates.map_err(|error| {
1322 Error::Corruption(format!(
1323 "invalid IVF-TQ payload for field {}: {error}",
1324 field.0
1325 ))
1326 })
1327}
1328
1329fn validate_ivf_tq_ann(
1330 index: &crate::segment::ann_disk::AnnDiskIndex,
1331 centroids: &CoarseCentroids,
1332 codec: &crate::structures::TqCodec,
1333 dim: usize,
1334 routing: crate::dsl::IvfRoutingMode,
1335 field: Field,
1336) -> Result<()> {
1337 let header = index.header();
1338 if !crate::structures::is_ivf_tq_cosine_generation(centroids.version)
1339 || !crate::structures::is_ivf_tq_cosine_generation(header.quantizer_version)
1340 {
1341 return Err(Error::Corruption(format!(
1342 "IVF-TQ field {} uses a legacy unmarked raw-vector generation that cannot \
1343 preserve cosine candidate semantics; rebuild the index with a current \
1344 Hermes version",
1345 field.0,
1346 )));
1347 }
1348 if header.dim != dim
1349 || codec.dim() != dim
1350 || header.code_size != codec.code_size()
1351 || header.num_clusters != centroids.num_clusters
1352 || header.quantizer_version != centroids.version
1353 || header.codebook_version != codec.fingerprint()
1354 || header.routing != routing
1355 {
1356 return Err(Error::Corruption(format!(
1357 "IVF-TQ payload for field {} does not match its quantizer/codec generation",
1358 field.0,
1359 )));
1360 }
1361 Ok(())
1362}
1363
1364fn validate_binary_ann(
1365 index: &crate::segment::ann_disk::AnnDiskIndex,
1366 quantizer: &crate::structures::BinaryCoarseQuantizer,
1367 config: &crate::dsl::BinaryDenseVectorConfig,
1368 dim: usize,
1369 field: Field,
1370) -> Result<()> {
1371 let header = index.header();
1372 if header.dim != dim
1373 || header.code_size != config.byte_len()
1374 || header.num_clusters != quantizer.num_clusters
1375 || header.quantizer_version != quantizer.version
1376 || header.codebook_version != 0
1377 || header.routing != config.ivf_routing
1378 || quantizer.dim_bits != dim
1379 {
1380 return Err(Error::Corruption(format!(
1381 "binary IVF field {} does not match its quantizer/schema generation",
1382 field.0,
1383 )));
1384 }
1385 Ok(())
1386}
1387
1388fn binary_probe_clusters(
1389 quantizer: &crate::structures::BinaryCoarseQuantizer,
1390 query: &[u8],
1391 nprobe: usize,
1392 routing: crate::dsl::IvfRoutingMode,
1393 cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
1394) -> Result<std::sync::Arc<[u32]>> {
1395 let effective_nprobe = nprobe.clamp(1, quantizer.num_clusters as usize);
1396 let request_fingerprint = crate::structures::vector::ivf::routing::binary_probe_fingerprint(
1397 query,
1398 effective_nprobe,
1399 routing,
1400 );
1401 if let Some(cache) = cache {
1402 let mut cached = cache
1403 .lock()
1404 .map_err(|_| Error::Internal("binary IVF probe cache is poisoned".into()))?;
1405 if let Some(plan) = cached.as_ref()
1406 && plan.quantizer_version == quantizer.version
1407 && plan.request_fingerprint == request_fingerprint
1408 && plan.cluster_ids.len() == effective_nprobe
1409 {
1410 return Ok(std::sync::Arc::clone(&plan.cluster_ids));
1411 }
1412 let plan = quantizer.probe(query, effective_nprobe, routing);
1413 let clusters = std::sync::Arc::clone(&plan.cluster_ids);
1414 *cached = Some(plan);
1415 return Ok(clusters);
1416 }
1417 Ok(quantizer
1418 .probe(query, effective_nprobe, routing)
1419 .cluster_ids)
1420}
1421
1422pub struct SegmentReader {
1428 meta: SegmentMeta,
1429 term_dict: Arc<AsyncSSTableReader<TermInfo>>,
1431 postings_handle: FileHandle,
1433 store: Arc<AsyncStoreReader>,
1435 schema: Arc<Schema>,
1436 vector_indexes: FxHashMap<u32, VectorIndex>,
1438 flat_vectors: FxHashMap<u32, LazyFlatVectorData>,
1440 dense_file_backed_bytes: u64,
1442 trained_vectors: Arc<crate::segment::TrainedVectorStructures>,
1444 sparse_indexes: FxHashMap<u32, SparseIndex>,
1446 bmp_indexes: FxHashMap<u32, BmpIndex>,
1448 sparse_file_backed_bytes: u64,
1450 positions_handle: Option<FileHandle>,
1452 fast_fields: FxHashMap<u32, crate::structures::fast_field::FastFieldReader>,
1454 #[cfg(feature = "native")]
1456 dense_pin_report: crate::segment::pin::PinReport,
1457 #[cfg(feature = "native")]
1459 sparse_pin_report: crate::segment::pin::PinReport,
1460}
1461
1462impl SegmentReader {
1463 pub async fn open<D: Directory>(
1465 dir: &D,
1466 segment_id: SegmentId,
1467 schema: Arc<Schema>,
1468 term_cache_blocks: usize,
1469 ) -> Result<Self> {
1470 Self::open_with_store_cache(
1471 dir,
1472 segment_id,
1473 schema,
1474 term_cache_blocks,
1475 dir as *const D as usize,
1476 Arc::new(super::SharedStoreCache::new(0)),
1477 )
1478 .await
1479 }
1480
1481 pub(crate) async fn open_with_store_cache<D: Directory>(
1483 dir: &D,
1484 segment_id: SegmentId,
1485 schema: Arc<Schema>,
1486 term_cache_blocks: usize,
1487 store_cache_directory_namespace: usize,
1488 store_cache: Arc<super::SharedStoreCache>,
1489 ) -> Result<Self> {
1490 let files = SegmentFiles::new(segment_id.0);
1491
1492 let meta_slice = dir.open_read(&files.meta).await?;
1494 let meta_bytes = meta_slice.read_bytes().await?;
1495 let meta = SegmentMeta::deserialize(meta_bytes.as_slice())?;
1496 debug_assert_eq!(meta.id, segment_id.0);
1497
1498 let term_dict_handle = dir.open_lazy(&files.term_dict).await?;
1500 let term_dict = AsyncSSTableReader::open(term_dict_handle, term_cache_blocks).await?;
1501
1502 let postings_handle = dir.open_lazy(&files.postings).await?;
1504
1505 let store_handle = dir.open_lazy(&files.store).await?;
1507 let store = AsyncStoreReader::open(
1508 store_handle,
1509 store_cache_directory_namespace,
1510 segment_id.0,
1511 store_cache,
1512 )
1513 .await?;
1514
1515 let vectors_data = loader::load_vectors_file(dir, &files, &schema, meta.num_docs).await?;
1517 let dense_file_backed_bytes = vectors_data.file_backed_bytes;
1518 let vector_indexes = vectors_data.indexes;
1519 let flat_vectors = vectors_data.flat_vectors;
1520
1521 #[cfg(feature = "native")]
1526 for (field_id, lazy_flat) in &flat_vectors {
1527 if vector_indexes.contains_key(field_id) {
1528 lazy_flat.advise_random_access();
1529 }
1530 }
1531
1532 let sparse_data = loader::load_sparse_file(dir, &files, meta.num_docs, &schema).await?;
1534 let sparse_file_backed_bytes = sparse_data.file_backed_bytes;
1535 let sparse_indexes = sparse_data.maxscore_indexes;
1536 let bmp_indexes = sparse_data.bmp_indexes;
1537
1538 let positions_handle = loader::open_positions_file(dir, &files, &schema).await?;
1540
1541 let fast_fields = loader::load_fast_fields_file(dir, &files, &schema).await?;
1543
1544 {
1546 let mut parts = vec![format!(
1547 "[segment] loaded {:016x}: docs={}",
1548 segment_id.0, meta.num_docs
1549 )];
1550 if !vector_indexes.is_empty() || !flat_vectors.is_empty() {
1551 parts.push(format!(
1552 "dense vectors: {} ANN + {} flat fields",
1553 vector_indexes.len(),
1554 flat_vectors.len()
1555 ));
1556 }
1557 for (field_id, idx) in &sparse_indexes {
1558 parts.push(format!(
1559 "sparse vector field {}: {} dims, ~{}",
1560 field_id,
1561 idx.num_dimensions(),
1562 crate::format_bytes(idx.num_dimensions() as u64 * 24)
1563 ));
1564 }
1565 for (field_id, idx) in &bmp_indexes {
1566 parts.push(format!(
1567 "bmp field {}: {} dims, {} blocks",
1568 field_id,
1569 idx.dims(),
1570 idx.num_blocks
1571 ));
1572 }
1573 if !fast_fields.is_empty() {
1574 parts.push(format!("fast: {} fields", fast_fields.len()));
1575 }
1576 log::debug!("{}", parts.join(", "));
1577 }
1578
1579 #[allow(unused_mut)]
1580 let mut reader = Self {
1581 meta,
1582 term_dict: Arc::new(term_dict),
1583 postings_handle,
1584 store: Arc::new(store),
1585 schema,
1586 vector_indexes,
1587 flat_vectors,
1588 dense_file_backed_bytes,
1589 trained_vectors: Arc::new(crate::segment::TrainedVectorStructures::default()),
1590 sparse_indexes,
1591 bmp_indexes,
1592 sparse_file_backed_bytes,
1593 positions_handle,
1594 fast_fields,
1595 #[cfg(feature = "native")]
1596 dense_pin_report: Default::default(),
1597 #[cfg(feature = "native")]
1598 sparse_pin_report: Default::default(),
1599 };
1600
1601 #[cfg(feature = "native")]
1603 reader.apply_pin_policy(&crate::segment::pin::pin_policy().to_owned());
1604
1605 Ok(reader)
1606 }
1607
1608 #[cfg(feature = "native")]
1618 pub(crate) fn apply_pin_policy(&mut self, policy: &crate::segment::pin::PinPolicy) {
1619 use crate::segment::pin::PinReport;
1620
1621 if !policy.is_enabled() {
1622 return;
1623 }
1624 let mut remaining = policy.budget_bytes;
1625 let mut dense_report = PinReport::default();
1626 let mut sparse_report = PinReport::default();
1627
1628 for index in self.vector_indexes.values_mut() {
1630 index.pin_lookup_directory(policy.mode, &mut remaining, &mut dense_report);
1631 }
1632 for bmp in self.bmp_indexes.values_mut() {
1634 bmp.pin_block_starts(policy.mode, &mut remaining, &mut sparse_report);
1635 }
1636 for sparse in self.sparse_indexes.values_mut() {
1638 sparse.pin_skip_section(policy.mode, &mut remaining, &mut sparse_report);
1639 }
1640 for flat in self.flat_vectors.values_mut() {
1642 flat.pin_doc_ids(policy.mode, &mut remaining, &mut dense_report);
1643 }
1644 for bmp in self.bmp_indexes.values_mut() {
1645 bmp.pin_doc_maps(policy.mode, &mut remaining, &mut sparse_report);
1646 }
1647 for bmp in self.bmp_indexes.values_mut() {
1649 bmp.pin_query_hierarchy(policy.mode, &mut remaining, &mut sparse_report);
1650 }
1651
1652 let report = PinReport {
1653 intended_bytes: dense_report
1654 .intended_bytes
1655 .saturating_add(sparse_report.intended_bytes),
1656 pinned_bytes: dense_report
1657 .pinned_bytes
1658 .saturating_add(sparse_report.pinned_bytes),
1659 skipped_budget_bytes: dense_report
1660 .skipped_budget_bytes
1661 .saturating_add(sparse_report.skipped_budget_bytes),
1662 failed_bytes: dense_report
1663 .failed_bytes
1664 .saturating_add(sparse_report.failed_bytes),
1665 heap_copy_bytes: dense_report
1666 .heap_copy_bytes
1667 .saturating_add(sparse_report.heap_copy_bytes),
1668 };
1669 if report.skipped_budget_bytes > 0 || report.failed_bytes > 0 {
1670 log::warn!(
1671 "[pin] segment {:016x}: pinned {}/{} (budget skipped {}, mlock failed {}) — \
1672 raise HERMES_PIN_METADATA_BUDGET_MB or RLIMIT_MEMLOCK for full coverage",
1673 self.meta.id,
1674 crate::format_bytes(report.pinned_bytes),
1675 crate::format_bytes(report.intended_bytes),
1676 crate::format_bytes(report.skipped_budget_bytes),
1677 crate::format_bytes(report.failed_bytes),
1678 );
1679 } else if report.pinned_bytes > 0 {
1680 log::info!(
1681 "[pin] segment {:016x}: pinned {} of hot metadata ({:?})",
1682 self.meta.id,
1683 crate::format_bytes(report.pinned_bytes),
1684 policy.mode,
1685 );
1686 }
1687 self.dense_pin_report = dense_report;
1688 self.sparse_pin_report = sparse_report;
1689 }
1690
1691 pub fn meta(&self) -> &SegmentMeta {
1697 &self.meta
1698 }
1699
1700 pub fn num_docs(&self) -> u32 {
1701 self.meta.num_docs
1702 }
1703
1704 pub fn avg_field_len(&self, field: Field) -> f32 {
1706 self.meta.avg_field_len(field)
1707 }
1708
1709 pub fn schema(&self) -> &Schema {
1710 &self.schema
1711 }
1712
1713 pub fn sparse_indexes(&self) -> &FxHashMap<u32, SparseIndex> {
1715 &self.sparse_indexes
1716 }
1717
1718 pub fn sparse_index(&self, field: Field) -> Option<&SparseIndex> {
1720 self.sparse_indexes.get(&field.0)
1721 }
1722
1723 pub fn bmp_index(&self, field: Field) -> Option<&BmpIndex> {
1725 self.bmp_indexes.get(&field.0)
1726 }
1727
1728 pub fn bmp_indexes(&self) -> &FxHashMap<u32, BmpIndex> {
1730 &self.bmp_indexes
1731 }
1732
1733 pub fn vector_indexes(&self) -> &FxHashMap<u32, VectorIndex> {
1735 &self.vector_indexes
1736 }
1737
1738 pub fn flat_vectors(&self) -> &FxHashMap<u32, LazyFlatVectorData> {
1740 &self.flat_vectors
1741 }
1742
1743 pub fn fast_field(
1745 &self,
1746 field_id: u32,
1747 ) -> Option<&crate::structures::fast_field::FastFieldReader> {
1748 self.fast_fields.get(&field_id)
1749 }
1750
1751 pub fn fast_fields(&self) -> &FxHashMap<u32, crate::structures::fast_field::FastFieldReader> {
1753 &self.fast_fields
1754 }
1755
1756 pub fn term_dict_stats(&self) -> SSTableStats {
1758 self.term_dict.stats()
1759 }
1760
1761 pub fn memory_stats(&self) -> SegmentMemoryStats {
1763 let term_dict_stats = self.term_dict.stats();
1764
1765 let term_dict_cache_bytes = self.term_dict.cached_bytes();
1769 let store_cache_bytes = self.store.cached_bytes();
1770
1771 let sparse_heap_bytes: usize = self
1774 .sparse_indexes
1775 .values()
1776 .map(|s| s.estimated_heap_bytes())
1777 .sum::<usize>()
1778 + self
1779 .bmp_indexes
1780 .values()
1781 .map(|b| b.estimated_heap_bytes())
1782 .sum::<usize>();
1783
1784 let dense_heap_bytes: usize = self
1787 .vector_indexes
1788 .values()
1789 .map(|v| v.estimated_heap_bytes())
1790 .sum::<usize>()
1791 + self
1792 .flat_vectors
1793 .values()
1794 .map(LazyFlatVectorData::estimated_heap_bytes)
1795 .sum::<usize>();
1796
1797 #[cfg(feature = "native")]
1798 let (sparse_heap_bytes, dense_heap_bytes) = (
1799 sparse_heap_bytes.saturating_add(
1800 usize::try_from(self.sparse_pin_report.heap_copy_bytes).unwrap_or(usize::MAX),
1801 ),
1802 dense_heap_bytes.saturating_add(
1803 usize::try_from(self.dense_pin_report.heap_copy_bytes).unwrap_or(usize::MAX),
1804 ),
1805 );
1806
1807 #[cfg(feature = "native")]
1808 let (
1809 sparse_pinned_metadata_bytes,
1810 sparse_pin_intended_bytes,
1811 dense_pinned_metadata_bytes,
1812 dense_pin_intended_bytes,
1813 ) = (
1814 self.sparse_pin_report.pinned_bytes,
1815 self.sparse_pin_report.intended_bytes,
1816 self.dense_pin_report.pinned_bytes,
1817 self.dense_pin_report.intended_bytes,
1818 );
1819 #[cfg(not(feature = "native"))]
1820 let (
1821 sparse_pinned_metadata_bytes,
1822 sparse_pin_intended_bytes,
1823 dense_pinned_metadata_bytes,
1824 dense_pin_intended_bytes,
1825 ) = (0u64, 0u64, 0u64, 0u64);
1826
1827 let pinned_metadata_bytes =
1828 sparse_pinned_metadata_bytes.saturating_add(dense_pinned_metadata_bytes);
1829 let pin_intended_bytes = sparse_pin_intended_bytes.saturating_add(dense_pin_intended_bytes);
1830
1831 SegmentMemoryStats {
1832 segment_id: self.meta.id,
1833 num_docs: self.meta.num_docs,
1834 term_dict_cache_bytes,
1835 store_cache_bytes,
1836 sparse_heap_bytes,
1837 dense_heap_bytes,
1838 term_bloom_file_bytes: term_dict_stats.bloom_filter_size as u64,
1839 sparse_file_backed_bytes: self.sparse_file_backed_bytes,
1840 dense_file_backed_bytes: self.dense_file_backed_bytes,
1841 pinned_metadata_bytes,
1842 pin_intended_bytes,
1843 sparse_pinned_metadata_bytes,
1844 sparse_pin_intended_bytes,
1845 dense_pinned_metadata_bytes,
1846 dense_pin_intended_bytes,
1847 }
1848 }
1849
1850 pub async fn get_postings(
1855 &self,
1856 field: Field,
1857 term: &[u8],
1858 ) -> Result<Option<BlockPostingList>> {
1859 log::debug!(
1860 "SegmentReader::get_postings field={} term_len={}",
1861 field.0,
1862 term.len()
1863 );
1864
1865 let mut key = Vec::with_capacity(4 + term.len());
1867 key.extend_from_slice(&field.0.to_le_bytes());
1868 key.extend_from_slice(term);
1869
1870 let term_info = match self.term_dict.get(&key).await? {
1872 Some(info) => {
1873 log::debug!("SegmentReader::get_postings found term_info");
1874 info
1875 }
1876 None => {
1877 log::debug!("SegmentReader::get_postings term not found");
1878 return Ok(None);
1879 }
1880 };
1881
1882 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
1884 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
1886 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
1887 posting_list.push(doc_id, tf);
1888 }
1889 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
1890 return Ok(Some(block_list));
1891 }
1892
1893 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
1895 Error::Corruption("TermInfo has neither inline nor external data".to_string())
1896 })?;
1897
1898 let range = checked_file_range(
1899 posting_offset,
1900 posting_len,
1901 self.postings_handle.len(),
1902 "posting",
1903 )?;
1904 let posting_bytes = self.postings_handle.read_bytes_range(range).await?;
1905 let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
1906
1907 Ok(Some(block_list))
1908 }
1909
1910 pub async fn get_prefix_postings(
1912 &self,
1913 field: Field,
1914 prefix: &[u8],
1915 ) -> Result<Vec<BlockPostingList>> {
1916 if prefix.is_empty() {
1917 return Err(Error::Query("prefix must not be empty".into()));
1918 }
1919 let mut key_prefix = Vec::with_capacity(4 + prefix.len());
1921 key_prefix.extend_from_slice(&field.0.to_le_bytes());
1922 key_prefix.extend_from_slice(prefix);
1923
1924 let (entries, truncated) = self
1925 .term_dict
1926 .prefix_scan_limited(&key_prefix, MAX_PREFIX_TERMS)
1927 .await?;
1928 if truncated {
1929 return Err(Error::Query(format!(
1930 "prefix expands to more than {MAX_PREFIX_TERMS} terms"
1931 )));
1932 }
1933 let posting_count: u64 = entries
1934 .iter()
1935 .map(|(_, term_info)| term_info.doc_freq() as u64)
1936 .sum();
1937 if posting_count > MAX_PREFIX_POSTINGS {
1938 return Err(Error::Query(format!(
1939 "prefix expands to {posting_count} postings (maximum {MAX_PREFIX_POSTINGS})"
1940 )));
1941 }
1942 let mut results = Vec::with_capacity(entries.len());
1943
1944 for (_key, term_info) in entries {
1945 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
1946 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
1947 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
1948 posting_list.push(doc_id, tf);
1949 }
1950 results.push(BlockPostingList::from_posting_list(&posting_list)?);
1951 } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
1952 let range = checked_file_range(
1953 posting_offset,
1954 posting_len,
1955 self.postings_handle.len(),
1956 "prefix posting",
1957 )?;
1958 let posting_bytes = self.postings_handle.read_bytes_range(range).await?;
1959 results.push(BlockPostingList::deserialize_zero_copy(posting_bytes)?);
1960 }
1961 }
1962
1963 Ok(results)
1964 }
1965
1966 pub async fn doc(&self, local_doc_id: DocId) -> Result<Option<Document>> {
1971 self.doc_with_fields(local_doc_id, None).await
1972 }
1973
1974 pub async fn doc_with_fields(
1980 &self,
1981 local_doc_id: DocId,
1982 fields: Option<&rustc_hash::FxHashSet<u32>>,
1983 ) -> Result<Option<Document>> {
1984 let mut doc = match fields {
1985 Some(set) => {
1986 let field_ids: Vec<u32> = set.iter().copied().collect();
1987 match self
1988 .store
1989 .get_fields(local_doc_id, &self.schema, &field_ids)
1990 .await
1991 {
1992 Ok(Some(d)) => d,
1993 Ok(None) => return Ok(None),
1994 Err(e) => return Err(Error::from(e)),
1995 }
1996 }
1997 None => match self.store.get(local_doc_id, &self.schema).await {
1998 Ok(Some(d)) => d,
1999 Ok(None) => return Ok(None),
2000 Err(e) => return Err(Error::from(e)),
2001 },
2002 };
2003
2004 for (&field_id, lazy_flat) in &self.flat_vectors {
2006 if let Some(set) = fields
2008 && !set.contains(&field_id)
2009 {
2010 continue;
2011 }
2012
2013 let is_binary = lazy_flat.quantization == DenseVectorQuantization::Binary;
2014 let (start, entries) = lazy_flat.flat_indexes_for_doc(local_doc_id);
2015 for (j, &(_doc_id, _ordinal)) in entries.iter().enumerate() {
2016 let flat_idx = start + j;
2017 if is_binary {
2018 let vbs = lazy_flat.vector_byte_size();
2019 let mut raw = vec![0u8; vbs];
2020 match lazy_flat.read_vector_raw_into(flat_idx, &mut raw).await {
2021 Ok(()) => {
2022 doc.add_binary_dense_vector(Field(field_id), raw);
2023 }
2024 Err(e) => {
2025 log::warn!(
2026 "Failed to hydrate binary dense vector field {}: {}",
2027 field_id,
2028 e
2029 );
2030 }
2031 }
2032 } else {
2033 match lazy_flat.get_vector(flat_idx).await {
2034 Ok(vec) => {
2035 doc.add_dense_vector(Field(field_id), vec);
2036 }
2037 Err(e) => {
2038 log::warn!("Failed to hydrate dense vector field {}: {}", field_id, e);
2039 }
2040 }
2041 }
2042 }
2043 }
2044
2045 Ok(Some(doc))
2046 }
2047
2048 pub async fn prefetch_terms(
2050 &self,
2051 field: Field,
2052 start_term: &[u8],
2053 end_term: &[u8],
2054 ) -> Result<()> {
2055 let mut start_key = Vec::with_capacity(4 + start_term.len());
2056 start_key.extend_from_slice(&field.0.to_le_bytes());
2057 start_key.extend_from_slice(start_term);
2058
2059 let mut end_key = Vec::with_capacity(4 + end_term.len());
2060 end_key.extend_from_slice(&field.0.to_le_bytes());
2061 end_key.extend_from_slice(end_term);
2062
2063 self.term_dict.prefetch_range(&start_key, &end_key).await?;
2064 Ok(())
2065 }
2066
2067 pub fn store_has_dict(&self) -> bool {
2069 self.store.has_dict()
2070 }
2071
2072 pub fn store(&self) -> &super::store::AsyncStoreReader {
2074 &self.store
2075 }
2076
2077 pub fn store_raw_blocks(&self) -> Vec<RawStoreBlock> {
2079 self.store.raw_blocks()
2080 }
2081
2082 pub fn store_data_slice(&self) -> &FileHandle {
2084 self.store.data_slice()
2085 }
2086
2087 pub async fn all_terms(&self) -> Result<Vec<(Vec<u8>, TermInfo)>> {
2089 self.term_dict.all_entries().await.map_err(Error::from)
2090 }
2091
2092 pub async fn all_terms_with_stats(&self) -> Result<Vec<(Field, String, u32)>> {
2097 let entries = self.term_dict.all_entries().await?;
2098 let mut result = Vec::with_capacity(entries.len());
2099
2100 for (key, term_info) in entries {
2101 if key.len() > 4 {
2103 let field_id = u32::from_le_bytes([key[0], key[1], key[2], key[3]]);
2104 let term_bytes = &key[4..];
2105 if let Ok(term_str) = std::str::from_utf8(term_bytes) {
2106 result.push((Field(field_id), term_str.to_string(), term_info.doc_freq()));
2107 }
2108 }
2109 }
2110
2111 Ok(result)
2112 }
2113
2114 pub fn term_dict_iter(&self) -> crate::structures::AsyncSSTableIterator<'_, TermInfo> {
2116 self.term_dict.iter()
2117 }
2118
2119 pub async fn prefetch_term_dict(&self) -> crate::Result<()> {
2123 self.term_dict
2124 .prefetch_all_data_bulk()
2125 .await
2126 .map_err(crate::Error::from)
2127 }
2128
2129 pub async fn read_postings(&self, offset: u64, len: u64) -> Result<Vec<u8>> {
2131 let range = checked_file_range(offset, len, self.postings_handle.len(), "posting")?;
2132 let bytes = self.postings_handle.read_bytes_range(range).await?;
2133 Ok(bytes.to_vec())
2134 }
2135
2136 pub async fn read_position_bytes(&self, offset: u64, len: u64) -> Result<Option<Vec<u8>>> {
2138 let handle = match &self.positions_handle {
2139 Some(h) => h,
2140 None => return Ok(None),
2141 };
2142 let range = checked_file_range(offset, len, handle.len(), "position")?;
2143 let bytes = handle.read_bytes_range(range).await?;
2144 Ok(Some(bytes.to_vec()))
2145 }
2146
2147 pub fn has_positions_file(&self) -> bool {
2149 self.positions_handle.is_some()
2150 }
2151
2152 fn validate_dense_search_request(
2156 &self,
2157 field: Field,
2158 query: &[f32],
2159 nprobe: usize,
2160 rerank_factor: f32,
2161 combiner: crate::query::MultiValueCombiner,
2162 ) -> Result<DenseSearchParams> {
2163 let entry = self
2164 .schema
2165 .get_field_entry(field)
2166 .ok_or_else(|| Error::FieldNotFound(field.0.to_string()))?;
2167 if entry.field_type != crate::dsl::FieldType::DenseVector {
2168 return Err(Error::InvalidFieldType {
2169 expected: "dense_vector".to_string(),
2170 got: format!("{:?}", entry.field_type),
2171 });
2172 }
2173 let config = entry.dense_vector_config.as_ref().ok_or_else(|| {
2174 Error::Schema(format!(
2175 "dense vector field '{}' has no dense vector configuration",
2176 entry.name
2177 ))
2178 })?;
2179
2180 if query.is_empty() {
2181 return Err(Error::Query(format!(
2182 "dense query vector for field '{}' must not be empty",
2183 entry.name
2184 )));
2185 }
2186 if query.len() != config.dim {
2187 return Err(Error::Query(format!(
2188 "dense query vector dimension {} does not match field '{}' dimension {}",
2189 query.len(),
2190 entry.name,
2191 config.dim
2192 )));
2193 }
2194 if let Some((index, value)) = query
2195 .iter()
2196 .enumerate()
2197 .find(|(_, value)| !value.is_finite())
2198 {
2199 return Err(Error::Query(format!(
2200 "dense query vector for field '{}' contains non-finite value {value} at index {index}",
2201 entry.name
2202 )));
2203 }
2204
2205 let nprobe = match (nprobe, config.nprobe) {
2208 (0, 0) => 32,
2209 (0, schema_nprobe) => schema_nprobe,
2210 (query_nprobe, _) => query_nprobe,
2211 };
2212 if nprobe > MAX_DENSE_NPROBE {
2213 return Err(Error::Query(format!(
2214 "dense nprobe must be at most {MAX_DENSE_NPROBE}, got {nprobe}"
2215 )));
2216 }
2217
2218 checked_dense_fetch_k(0, rerank_factor)?;
2221 combiner.validate().map_err(Error::Query)?;
2222
2223 Ok(DenseSearchParams {
2224 dim: config.dim,
2225 nprobe,
2226 unit_norm: config.unit_norm,
2227 })
2228 }
2229
2230 fn validate_binary_search_request(&self, field: Field, query: &[u8]) -> Result<usize> {
2231 let entry = self
2232 .schema
2233 .get_field_entry(field)
2234 .ok_or_else(|| Error::FieldNotFound(field.0.to_string()))?;
2235 if entry.field_type != crate::dsl::FieldType::BinaryDenseVector {
2236 return Err(Error::InvalidFieldType {
2237 expected: "binary_dense_vector".to_string(),
2238 got: format!("{:?}", entry.field_type),
2239 });
2240 }
2241 let config = entry.binary_dense_vector_config.as_ref().ok_or_else(|| {
2242 Error::Schema(format!(
2243 "binary dense vector field '{}' has no configuration",
2244 entry.name
2245 ))
2246 })?;
2247 if config.dim == 0 || !config.dim.is_multiple_of(8) {
2248 return Err(Error::Schema(format!(
2249 "binary dense vector field '{}' has invalid dimension {}",
2250 entry.name, config.dim
2251 )));
2252 }
2253 if query.len() != config.byte_len() {
2254 return Err(Error::Query(format!(
2255 "binary query byte length {} does not match field '{}' byte length {}",
2256 query.len(),
2257 entry.name,
2258 config.byte_len()
2259 )));
2260 }
2261 Ok(config.dim)
2262 }
2263
2264 #[cfg(test)]
2266 fn score_quantized_batch_legacy(
2267 query: &[f32],
2268 raw: &[u8],
2269 quant: crate::dsl::DenseVectorQuantization,
2270 dim: usize,
2271 scores: &mut [f32],
2272 unit_norm: bool,
2273 ) -> Result<()> {
2274 use crate::dsl::DenseVectorQuantization;
2275 use crate::structures::simd;
2276
2277 if query.len() != dim {
2278 return Err(Error::Query(format!(
2279 "dense SIMD query dimension {} does not match vector dimension {dim}",
2280 query.len()
2281 )));
2282 }
2283 let element_size = match quant {
2284 DenseVectorQuantization::F32 => std::mem::size_of::<f32>(),
2285 DenseVectorQuantization::F16 => std::mem::size_of::<u16>(),
2286 DenseVectorQuantization::UInt8 => 1,
2287 DenseVectorQuantization::Binary => {
2288 return Err(Error::InvalidFieldType {
2289 expected: "non-binary dense vector".to_string(),
2290 got: "binary dense vector".to_string(),
2291 });
2292 }
2293 };
2294 let required_bytes = scores
2295 .len()
2296 .checked_mul(dim)
2297 .and_then(|elements| elements.checked_mul(element_size))
2298 .ok_or_else(|| Error::Corruption("dense vector batch byte length overflow".into()))?;
2299 if raw.len() < required_bytes {
2300 return Err(Error::Corruption(format!(
2301 "dense vector batch is truncated: need {required_bytes} bytes, got {}",
2302 raw.len()
2303 )));
2304 }
2305 if quant == DenseVectorQuantization::F16
2306 && required_bytes > 0
2307 && !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>())
2308 {
2309 return Err(Error::Corruption(
2310 "f16 vector data is not 2-byte aligned".to_string(),
2311 ));
2312 }
2313
2314 match (quant, unit_norm) {
2315 (DenseVectorQuantization::F32, false) => {
2316 let num_floats = scores.len() * dim;
2317 if !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()) {
2318 return Err(Error::Corruption(
2319 "f32 vector data is not 4-byte aligned".to_string(),
2320 ));
2321 }
2322 let vectors: &[f32] =
2323 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
2324 simd::batch_cosine_scores(query, vectors, dim, scores);
2325 }
2326 (DenseVectorQuantization::F32, true) => {
2327 let num_floats = scores.len() * dim;
2328 if !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()) {
2329 return Err(Error::Corruption(
2330 "f32 vector data is not 4-byte aligned".to_string(),
2331 ));
2332 }
2333 let vectors: &[f32] =
2334 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
2335 simd::batch_dot_scores(query, vectors, dim, scores);
2336 }
2337 (DenseVectorQuantization::F16, false) => {
2338 simd::batch_cosine_scores_f16(query, raw, dim, scores);
2339 }
2340 (DenseVectorQuantization::F16, true) => {
2341 simd::batch_dot_scores_f16(query, raw, dim, scores);
2342 }
2343 (DenseVectorQuantization::UInt8, false) => {
2344 simd::batch_cosine_scores_u8(query, raw, dim, scores);
2345 }
2346 (DenseVectorQuantization::UInt8, true) => {
2347 simd::batch_dot_scores_u8(query, raw, dim, scores);
2348 }
2349 (DenseVectorQuantization::Binary, _) => unreachable!("validated above"),
2350 }
2351 Ok(())
2352 }
2353
2354 pub async fn search_dense_vector(
2360 &self,
2361 field: Field,
2362 query: &[f32],
2363 k: usize,
2364 nprobe: usize,
2365 rerank_factor: f32,
2366 combiner: crate::query::MultiValueCombiner,
2367 ) -> Result<Vec<VectorSearchResult>> {
2368 self.search_dense_vector_impl(field, query, k, nprobe, rerank_factor, combiner, None)
2369 .await
2370 }
2371
2372 #[allow(clippy::too_many_arguments)]
2373 pub(crate) async fn search_dense_vector_with_probe_cache(
2374 &self,
2375 field: Field,
2376 query: &[f32],
2377 k: usize,
2378 nprobe: usize,
2379 rerank_factor: f32,
2380 combiner: crate::query::MultiValueCombiner,
2381 plan_cache: &DensePlanCache,
2382 ) -> Result<Vec<VectorSearchResult>> {
2383 self.search_dense_vector_impl(
2384 field,
2385 query,
2386 k,
2387 nprobe,
2388 rerank_factor,
2389 combiner,
2390 Some(plan_cache),
2391 )
2392 .await
2393 }
2394
2395 #[allow(clippy::too_many_arguments)]
2396 async fn search_dense_vector_impl(
2397 &self,
2398 field: Field,
2399 query: &[f32],
2400 k: usize,
2401 nprobe: usize,
2402 rerank_factor: f32,
2403 combiner: crate::query::MultiValueCombiner,
2404 plan_cache: Option<&DensePlanCache>,
2405 ) -> Result<Vec<VectorSearchResult>> {
2406 let params =
2407 self.validate_dense_search_request(field, query, nprobe, rerank_factor, combiner)?;
2408 let fetch_k = checked_dense_fetch_k(k, rerank_factor)?;
2409 if k == 0 {
2410 return Ok(Vec::new());
2411 }
2412
2413 let configured_ann_index = self.vector_indexes.get(&field.0);
2414 let lazy_flat = self.flat_vectors.get(&field.0);
2415 if configured_ann_index.is_none() && lazy_flat.is_none() {
2417 return Ok(Vec::new());
2418 }
2419
2420 if configured_ann_index.is_some() && lazy_flat.is_none() {
2421 return Err(Error::Corruption(format!(
2422 "dense ANN field {} is missing flat vector storage",
2423 field.0
2424 )));
2425 }
2426
2427 if let Some(flat) = lazy_flat
2428 && flat.dim != params.dim
2429 {
2430 return Err(Error::Corruption(format!(
2431 "dense vector field {} has schema dimension {} but flat storage dimension {}",
2432 field.0, params.dim, flat.dim
2433 )));
2434 }
2435
2436 let needs_document_aggregation = lazy_flat.is_some_and(|flat| {
2437 flat.num_vectors != flat.num_docs_with_vectors()
2438 && !matches!(combiner, crate::query::MultiValueCombiner::Max)
2439 });
2440 let ann_index = configured_ann_index;
2444
2445 let t0 = std::time::Instant::now();
2447 let mut flat_results = None;
2448 let results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
2449 match index {
2451 VectorIndex::Tq { index: lazy, codec } => {
2452 let flat = lazy_flat.expect("ANN/flat pairing validated above");
2453 search_tq_segment(
2455 lazy.get(),
2456 codec,
2457 query,
2458 fetch_k.min(flat.num_docs_with_vectors()),
2459 needs_document_aggregation.then_some(combiner),
2460 field,
2461 params.dim,
2462 plan_cache.map(|cache| &cache.tq),
2463 )?
2464 }
2465 VectorIndex::IvfTq { index: lazy, codec } => {
2466 let index = lazy.get();
2467 let centroids =
2468 self.trained_vectors
2469 .centroids
2470 .get(&field.0)
2471 .ok_or_else(|| {
2472 Error::Schema(format!(
2473 "IVF-TQ index requires coarse centroids for field {}",
2474 field.0
2475 ))
2476 })?;
2477 validate_coarse_centroids(centroids, params.dim)?;
2478 let routing = self
2479 .schema
2480 .get_field_entry(field)
2481 .and_then(|entry| entry.dense_vector_config.as_ref())
2482 .map_or(crate::dsl::IvfRoutingMode::Auto, |config| {
2483 config.ivf_routing
2484 });
2485 validate_ivf_tq_ann(index, centroids, codec, params.dim, routing, field)?;
2486 let flat = lazy_flat.expect("ANN/flat pairing validated above");
2487 search_ivf_tq_segment(
2488 index,
2489 centroids,
2490 codec,
2491 query,
2492 fetch_k.min(flat.num_docs_with_vectors()),
2493 needs_document_aggregation.then_some(combiner),
2494 field,
2495 params.nprobe,
2496 routing,
2497 plan_cache.map(|cache| &cache.ivf_tq),
2498 )?
2499 }
2500 VectorIndex::BinaryIvf(_) => {
2501 Vec::new()
2503 }
2504 }
2505 } else if let Some(lazy_flat) = lazy_flat {
2506 log::debug!(
2510 "[dense_vector_search] field {}: brute-force on {} vectors (dim={}, quant={:?})",
2511 field.0,
2512 lazy_flat.num_vectors,
2513 lazy_flat.dim,
2514 lazy_flat.quantization
2515 );
2516 let dim = lazy_flat.dim;
2517 let n = lazy_flat.num_vectors;
2518 let quant = lazy_flat.quantization;
2519 let batch_len =
2520 bounded_vector_score_batch(lazy_flat.vector_byte_size(), DENSE_SCORE_BATCH);
2521 let mut collector = FlatDocumentCollector::new(fetch_k.min(n), combiner);
2522 let mut scores = vec![0f32; batch_len];
2523 let prepared_query = PreparedDenseScoreQuery::new(query, quant, dim, params.unit_norm)?;
2524
2525 for batch_start in (0..n).step_by(batch_len) {
2526 let batch_count = batch_len.min(n - batch_start);
2527 let batch_bytes = lazy_flat
2528 .read_vectors_batch(batch_start, batch_count)
2529 .await
2530 .map_err(crate::Error::Io)?;
2531 let raw = batch_bytes.as_slice();
2532
2533 prepared_query.score_batch(raw, &mut scores[..batch_count])?;
2534
2535 for (i, &score) in scores.iter().enumerate().take(batch_count) {
2536 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
2537 collector.push(doc_id, ordinal, score);
2538 }
2539 }
2540
2541 flat_results = Some(collector.into_results());
2542 Vec::new()
2543 } else {
2544 return Ok(Vec::new());
2545 };
2546 let l1_elapsed = t0.elapsed();
2547 {
2548 let kind = match ann_index {
2549 Some(VectorIndex::BinaryIvf(_)) => "binary_ivf",
2550 Some(VectorIndex::Tq { .. }) => "tq_flat",
2551 Some(VectorIndex::IvfTq { .. }) => "ivf_tq",
2552 None => "flat",
2553 };
2554 crate::observe::dense_l1(
2555 self.schema.index_label(),
2556 self.schema.get_field_name(field).unwrap_or("?"),
2557 kind,
2558 l1_elapsed.as_secs_f64(),
2559 flat_results.as_ref().map_or(results.len(), Vec::len),
2560 );
2561 }
2562 log::debug!(
2563 "[dense_vector_search] field {}: L1 returned {} candidates in {:.1}ms",
2564 field.0,
2565 flat_results.as_ref().map_or(results.len(), Vec::len),
2566 l1_elapsed.as_secs_f64() * 1000.0
2567 );
2568
2569 if let Some(results) = flat_results {
2570 return Ok(results);
2571 }
2572
2573 if ann_index.is_some()
2576 && !results.is_empty()
2577 && let Some(lazy_flat) = lazy_flat
2578 {
2579 let t_rerank = std::time::Instant::now();
2580 let vbs = lazy_flat.vector_byte_size();
2581 let (reranked, stats) = exact_score_dense_candidate_documents(
2582 &results,
2583 lazy_flat,
2584 query,
2585 params.unit_norm,
2586 combiner,
2587 k,
2588 )
2589 .await?;
2590
2591 crate::observe::dense_rerank(
2592 self.schema.index_label(),
2593 self.schema.get_field_name(field).unwrap_or("?"),
2594 t_rerank.elapsed().as_secs_f64(),
2595 stats.resolve_elapsed.as_secs_f64(),
2596 stats.read_elapsed.as_secs_f64(),
2597 stats.vector_count,
2598 );
2599 log::debug!(
2600 "[dense_vector_search] field {}: rerank {} vectors (dim={}, quant={:?}, bytes_per_vector={}): resolve={:.1}ms read={:.1}ms score={:.1}ms",
2601 field.0,
2602 stats.vector_count,
2603 lazy_flat.dim,
2604 lazy_flat.quantization,
2605 vbs,
2606 stats.resolve_elapsed.as_secs_f64() * 1000.0,
2607 stats.read_elapsed.as_secs_f64() * 1000.0,
2608 stats.score_elapsed.as_secs_f64() * 1000.0,
2609 );
2610
2611 log::debug!(
2612 "[dense_vector_search] field {}: rerank total={:.1}ms",
2613 field.0,
2614 t_rerank.elapsed().as_secs_f64() * 1000.0
2615 );
2616 return Ok(reranked);
2617 }
2618
2619 Ok(combine_grouped_ordinal_results(results, combiner, k))
2620 }
2621
2622 async fn search_binary_dense_vector_impl(
2627 &self,
2628 field: Field,
2629 query: &[u8],
2630 k: usize,
2631 combiner: crate::query::MultiValueCombiner,
2632 probe_cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
2633 ) -> Result<Vec<VectorSearchResult>> {
2634 let schema_dim = self.validate_binary_search_request(field, query)?;
2635 combiner.validate().map_err(Error::Query)?;
2636 if k == 0 {
2637 return Ok(Vec::new());
2638 }
2639 let t0 = crate::observe::Timer::start();
2640 if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0) {
2641 let ivf = lazy.get();
2642 let config = self
2643 .schema
2644 .get_field_entry(field)
2645 .and_then(|entry| entry.binary_dense_vector_config.as_ref())
2646 .ok_or_else(|| {
2647 Error::Schema(format!(
2648 "binary IVF field {} has no schema configuration",
2649 field.0
2650 ))
2651 })?;
2652 let quantizer = self
2653 .trained_vectors
2654 .binary_quantizers
2655 .get(&field.0)
2656 .ok_or_else(|| {
2657 Error::Schema(format!(
2658 "global binary IVF field {} has no loaded quantizer",
2659 field.0
2660 ))
2661 })?;
2662 validate_binary_ann(ivf, quantizer, config, schema_dim, field)?;
2663 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
2664 Error::Corruption(format!(
2665 "global binary IVF field {} is missing flat vector storage",
2666 field.0
2667 ))
2668 })?;
2669 let single_valued = flat.num_vectors == flat.num_docs_with_vectors();
2670 let clusters = binary_probe_clusters(
2671 quantizer,
2672 query,
2673 config.nprobe,
2674 config.ivf_routing,
2675 probe_cache,
2676 )?;
2677 let results = if !single_valued
2678 && !matches!(combiner, crate::query::MultiValueCombiner::Max)
2679 {
2680 let candidate_limit =
2681 checked_binary_combined_fetch_k(k)?.min(flat.num_docs_with_vectors());
2682 let candidate_documents = ivf
2683 .search_binary_combined_documents(candidate_limit, query, &clusters, combiner)
2684 .map_err(|error| {
2685 Error::Corruption(format!(
2686 "invalid binary IVF payload for field {}: {error}",
2687 field.0,
2688 ))
2689 })?;
2690 exact_score_binary_candidate_document_ids(
2691 candidate_documents
2692 .into_iter()
2693 .map(|candidate| candidate.doc_id)
2694 .collect(),
2695 flat,
2696 query,
2697 schema_dim,
2698 combiner,
2699 k,
2700 )
2701 .await?
2702 } else {
2703 let candidate_docs = if single_valued {
2704 k
2705 } else {
2706 checked_binary_combined_fetch_k(k)?
2711 }
2712 .min(flat.num_docs_with_vectors());
2713 let ann_results = if single_valued {
2714 ivf.search_binary_clusters::<false>(query, candidate_docs, &clusters)
2715 } else {
2716 ivf.search_binary_clusters::<true>(query, candidate_docs, &clusters)
2717 }
2718 .map_err(|error| {
2719 Error::Corruption(format!(
2720 "invalid binary IVF payload for field {}: {error}",
2721 field.0,
2722 ))
2723 })?;
2724 if single_valued {
2727 let ann_results = validate_binary_single_value_ann_results(ann_results, flat)?;
2728 combine_ordinal_results(ann_results, combiner, k)
2729 } else {
2730 exact_score_binary_candidate_documents(
2731 &ann_results,
2732 flat,
2733 query,
2734 schema_dim,
2735 combiner,
2736 k,
2737 )
2738 .await?
2739 }
2740 };
2741 crate::observe::dense_l1(
2742 self.schema.index_label(),
2743 self.schema.get_field_name(field).unwrap_or("?"),
2744 "global_binary_ivf",
2745 t0.secs(),
2746 results.len(),
2747 );
2748 return Ok(results);
2749 }
2750 let lazy_flat = match self.flat_vectors.get(&field.0) {
2751 Some(f) => f,
2752 None => return Ok(Vec::new()),
2753 };
2754
2755 let dim_bits = lazy_flat.dim;
2756 let byte_len = lazy_flat.vector_byte_size();
2757 let n = lazy_flat.num_vectors;
2758
2759 if dim_bits != schema_dim {
2760 return Err(Error::Corruption(format!(
2761 "binary vector field {} has schema dimension {} but flat storage dimension {}",
2762 field.0, schema_dim, dim_bits
2763 )));
2764 }
2765
2766 if byte_len != query.len() {
2767 return Err(Error::Schema(format!(
2768 "Binary query vector byte length {} != field byte length {}",
2769 query.len(),
2770 byte_len
2771 )));
2772 }
2773
2774 let batch_len = bounded_vector_score_batch(byte_len, BINARY_SCORE_BATCH);
2775 let mut collector = FlatDocumentCollector::new(k, combiner);
2776 let mut scores = vec![0f32; batch_len];
2777
2778 for batch_start in (0..n).step_by(batch_len) {
2779 let batch_count = batch_len.min(n - batch_start);
2780 let batch_bytes = lazy_flat
2781 .read_vectors_batch(batch_start, batch_count)
2782 .await
2783 .map_err(crate::Error::Io)?;
2784 let raw = batch_bytes.as_slice();
2785
2786 crate::structures::simd::batch_hamming_scores(
2787 query,
2788 raw,
2789 byte_len,
2790 dim_bits,
2791 &mut scores[..batch_count],
2792 );
2793
2794 for (i, &score) in scores.iter().enumerate().take(batch_count) {
2795 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
2796 collector.push(doc_id, ordinal, score);
2797 }
2798 }
2799
2800 let results = collector.into_results();
2801
2802 crate::observe::dense_l1(
2803 self.schema.index_label(),
2804 self.schema.get_field_name(field).unwrap_or("?"),
2805 "binary_flat",
2806 t0.secs(),
2807 results.len(),
2808 );
2809 Ok(results)
2810 }
2811
2812 pub async fn search_binary_dense_vector(
2813 &self,
2814 field: Field,
2815 query: &[u8],
2816 k: usize,
2817 combiner: crate::query::MultiValueCombiner,
2818 ) -> Result<Vec<VectorSearchResult>> {
2819 self.search_binary_dense_vector_impl(field, query, k, combiner, None)
2820 .await
2821 }
2822
2823 pub(crate) async fn search_binary_dense_vector_with_probe_cache(
2824 &self,
2825 field: Field,
2826 query: &[u8],
2827 k: usize,
2828 combiner: crate::query::MultiValueCombiner,
2829 probe_cache: &std::sync::Mutex<Option<crate::structures::IvfProbePlan>>,
2830 ) -> Result<Vec<VectorSearchResult>> {
2831 self.search_binary_dense_vector_impl(field, query, k, combiner, Some(probe_cache))
2832 .await
2833 }
2834
2835 pub fn coarse_centroids(&self, field_id: u32) -> Option<&Arc<CoarseCentroids>> {
2837 self.trained_vectors.centroids.get(&field_id)
2838 }
2839
2840 pub fn set_trained_vectors(
2841 &mut self,
2842 trained_vectors: Arc<crate::segment::TrainedVectorStructures>,
2843 ) {
2844 self.trained_vectors = trained_vectors;
2845 }
2846
2847 pub fn get_vector_index(&self, field: Field) -> Option<&VectorIndex> {
2849 self.vector_indexes.get(&field.0)
2850 }
2851
2852 pub async fn get_positions(
2857 &self,
2858 field: Field,
2859 term: &[u8],
2860 ) -> Result<Option<crate::structures::PositionPostingList>> {
2861 let handle = match &self.positions_handle {
2863 Some(h) => h,
2864 None => return Ok(None),
2865 };
2866
2867 let mut key = Vec::with_capacity(4 + term.len());
2869 key.extend_from_slice(&field.0.to_le_bytes());
2870 key.extend_from_slice(term);
2871
2872 let term_info = match self.term_dict.get(&key).await? {
2874 Some(info) => info,
2875 None => return Ok(None),
2876 };
2877
2878 let (offset, length) = match term_info.position_info() {
2880 Some((o, l)) => (o, l),
2881 None => return Ok(None),
2882 };
2883
2884 let range = checked_file_range(offset, length, handle.len(), "position list")?;
2888 let slice = handle.slice(range);
2889 let data = slice.read_bytes().await?;
2890
2891 let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
2893
2894 Ok(Some(pos_list))
2895 }
2896
2897 pub fn has_positions(&self, field: Field) -> bool {
2899 if let Some(entry) = self.schema.get_field_entry(field) {
2901 entry.positions.is_some()
2902 } else {
2903 false
2904 }
2905 }
2906}
2907
2908#[cfg(feature = "sync")]
2910impl SegmentReader {
2911 pub fn get_postings_sync(&self, field: Field, term: &[u8]) -> Result<Option<BlockPostingList>> {
2913 let mut key = Vec::with_capacity(4 + term.len());
2915 key.extend_from_slice(&field.0.to_le_bytes());
2916 key.extend_from_slice(term);
2917
2918 let term_info = match self.term_dict.get_sync(&key)? {
2920 Some(info) => info,
2921 None => return Ok(None),
2922 };
2923
2924 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
2926 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
2927 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
2928 posting_list.push(doc_id, tf);
2929 }
2930 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
2931 return Ok(Some(block_list));
2932 }
2933
2934 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
2936 Error::Corruption("TermInfo has neither inline nor external data".to_string())
2937 })?;
2938
2939 let range = checked_file_range(
2940 posting_offset,
2941 posting_len,
2942 self.postings_handle.len(),
2943 "posting",
2944 )?;
2945 let posting_bytes = self.postings_handle.read_bytes_range_sync(range)?;
2946 let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
2947
2948 Ok(Some(block_list))
2949 }
2950
2951 pub fn get_prefix_postings_sync(
2953 &self,
2954 field: Field,
2955 prefix: &[u8],
2956 ) -> Result<Vec<BlockPostingList>> {
2957 if prefix.is_empty() {
2958 return Err(Error::Query("prefix must not be empty".into()));
2959 }
2960 let mut key_prefix = Vec::with_capacity(4 + prefix.len());
2961 key_prefix.extend_from_slice(&field.0.to_le_bytes());
2962 key_prefix.extend_from_slice(prefix);
2963
2964 let (entries, truncated) = self
2965 .term_dict
2966 .prefix_scan_limited_sync(&key_prefix, MAX_PREFIX_TERMS)?;
2967 if truncated {
2968 return Err(Error::Query(format!(
2969 "prefix expands to more than {MAX_PREFIX_TERMS} terms"
2970 )));
2971 }
2972 let posting_count: u64 = entries
2973 .iter()
2974 .map(|(_, term_info)| term_info.doc_freq() as u64)
2975 .sum();
2976 if posting_count > MAX_PREFIX_POSTINGS {
2977 return Err(Error::Query(format!(
2978 "prefix expands to {posting_count} postings (maximum {MAX_PREFIX_POSTINGS})"
2979 )));
2980 }
2981 let mut results = Vec::with_capacity(entries.len());
2982
2983 for (_key, term_info) in entries {
2984 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
2985 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
2986 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
2987 posting_list.push(doc_id, tf);
2988 }
2989 results.push(BlockPostingList::from_posting_list(&posting_list)?);
2990 } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
2991 let range = checked_file_range(
2992 posting_offset,
2993 posting_len,
2994 self.postings_handle.len(),
2995 "prefix posting",
2996 )?;
2997 let posting_bytes = self.postings_handle.read_bytes_range_sync(range)?;
2998 results.push(BlockPostingList::deserialize_zero_copy(posting_bytes)?);
2999 }
3000 }
3001
3002 Ok(results)
3003 }
3004
3005 pub fn get_positions_sync(
3007 &self,
3008 field: Field,
3009 term: &[u8],
3010 ) -> Result<Option<crate::structures::PositionPostingList>> {
3011 let handle = match &self.positions_handle {
3012 Some(h) => h,
3013 None => return Ok(None),
3014 };
3015
3016 let mut key = Vec::with_capacity(4 + term.len());
3018 key.extend_from_slice(&field.0.to_le_bytes());
3019 key.extend_from_slice(term);
3020
3021 let term_info = match self.term_dict.get_sync(&key)? {
3023 Some(info) => info,
3024 None => return Ok(None),
3025 };
3026
3027 let (offset, length) = match term_info.position_info() {
3028 Some((o, l)) => (o, l),
3029 None => return Ok(None),
3030 };
3031
3032 let range = checked_file_range(offset, length, handle.len(), "position list")?;
3033 let slice = handle.slice(range);
3034 let data = slice.read_bytes_sync()?;
3035
3036 let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
3037 Ok(Some(pos_list))
3038 }
3039
3040 pub fn search_dense_vector_sync(
3043 &self,
3044 field: Field,
3045 query: &[f32],
3046 k: usize,
3047 nprobe: usize,
3048 rerank_factor: f32,
3049 combiner: crate::query::MultiValueCombiner,
3050 ) -> Result<Vec<VectorSearchResult>> {
3051 self.search_dense_vector_sync_impl(field, query, k, nprobe, rerank_factor, combiner, None)
3052 }
3053
3054 #[cfg(feature = "sync")]
3055 #[allow(clippy::too_many_arguments)]
3056 pub(crate) fn search_dense_vector_sync_with_probe_cache(
3057 &self,
3058 field: Field,
3059 query: &[f32],
3060 k: usize,
3061 nprobe: usize,
3062 rerank_factor: f32,
3063 combiner: crate::query::MultiValueCombiner,
3064 plan_cache: &DensePlanCache,
3065 ) -> Result<Vec<VectorSearchResult>> {
3066 self.search_dense_vector_sync_impl(
3067 field,
3068 query,
3069 k,
3070 nprobe,
3071 rerank_factor,
3072 combiner,
3073 Some(plan_cache),
3074 )
3075 }
3076
3077 #[cfg(feature = "sync")]
3078 #[allow(clippy::too_many_arguments)]
3079 fn search_dense_vector_sync_impl(
3080 &self,
3081 field: Field,
3082 query: &[f32],
3083 k: usize,
3084 nprobe: usize,
3085 rerank_factor: f32,
3086 combiner: crate::query::MultiValueCombiner,
3087 plan_cache: Option<&DensePlanCache>,
3088 ) -> Result<Vec<VectorSearchResult>> {
3089 let params =
3090 self.validate_dense_search_request(field, query, nprobe, rerank_factor, combiner)?;
3091 let fetch_k = checked_dense_fetch_k(k, rerank_factor)?;
3092 if k == 0 {
3093 return Ok(Vec::new());
3094 }
3095
3096 let configured_ann_index = self.vector_indexes.get(&field.0);
3097 let lazy_flat = self.flat_vectors.get(&field.0);
3098 if configured_ann_index.is_none() && lazy_flat.is_none() {
3099 return Ok(Vec::new());
3100 }
3101
3102 if configured_ann_index.is_some() && lazy_flat.is_none() {
3103 return Err(Error::Corruption(format!(
3104 "dense ANN field {} is missing flat vector storage",
3105 field.0
3106 )));
3107 }
3108
3109 if let Some(flat) = lazy_flat
3110 && flat.dim != params.dim
3111 {
3112 return Err(Error::Corruption(format!(
3113 "dense vector field {} has schema dimension {} but flat storage dimension {}",
3114 field.0, params.dim, flat.dim
3115 )));
3116 }
3117
3118 let needs_document_aggregation = lazy_flat.is_some_and(|flat| {
3119 flat.num_vectors != flat.num_docs_with_vectors()
3120 && !matches!(combiner, crate::query::MultiValueCombiner::Max)
3121 });
3122 let ann_index = configured_ann_index;
3125
3126 let results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
3127 match index {
3129 VectorIndex::Tq { index: lazy, codec } => {
3130 let flat = lazy_flat.expect("ANN/flat pairing validated above");
3131 search_tq_segment(
3132 lazy.get(),
3133 codec,
3134 query,
3135 fetch_k.min(flat.num_docs_with_vectors()),
3136 needs_document_aggregation.then_some(combiner),
3137 field,
3138 params.dim,
3139 plan_cache.map(|cache| &cache.tq),
3140 )?
3141 }
3142 VectorIndex::IvfTq { index: lazy, codec } => {
3143 let index = lazy.get();
3144 let centroids =
3145 self.trained_vectors
3146 .centroids
3147 .get(&field.0)
3148 .ok_or_else(|| {
3149 Error::Schema(format!(
3150 "IVF-TQ index requires coarse centroids for field {}",
3151 field.0
3152 ))
3153 })?;
3154 validate_coarse_centroids(centroids, params.dim)?;
3155 let routing = self
3156 .schema
3157 .get_field_entry(field)
3158 .and_then(|entry| entry.dense_vector_config.as_ref())
3159 .map_or(crate::dsl::IvfRoutingMode::Auto, |config| {
3160 config.ivf_routing
3161 });
3162 validate_ivf_tq_ann(index, centroids, codec, params.dim, routing, field)?;
3163 let flat = lazy_flat.expect("ANN/flat pairing validated above");
3164 search_ivf_tq_segment(
3165 index,
3166 centroids,
3167 codec,
3168 query,
3169 fetch_k.min(flat.num_docs_with_vectors()),
3170 needs_document_aggregation.then_some(combiner),
3171 field,
3172 params.nprobe,
3173 routing,
3174 plan_cache.map(|cache| &cache.ivf_tq),
3175 )?
3176 }
3177 VectorIndex::BinaryIvf(_) => {
3178 Vec::new()
3180 }
3181 }
3182 } else if let Some(lazy_flat) = lazy_flat {
3183 let dim = lazy_flat.dim;
3185 let n = lazy_flat.num_vectors;
3186 let quant = lazy_flat.quantization;
3187 let batch_len =
3188 bounded_vector_score_batch(lazy_flat.vector_byte_size(), DENSE_SCORE_BATCH);
3189 let mut collector = FlatDocumentCollector::new(fetch_k.min(n), combiner);
3190 let mut scores = vec![0f32; batch_len];
3191 let prepared_query = PreparedDenseScoreQuery::new(query, quant, dim, params.unit_norm)?;
3192
3193 for batch_start in (0..n).step_by(batch_len) {
3194 let batch_count = batch_len.min(n - batch_start);
3195 let batch_bytes = lazy_flat
3196 .read_vectors_batch_sync(batch_start, batch_count)
3197 .map_err(crate::Error::Io)?;
3198 let raw = batch_bytes.as_slice();
3199
3200 prepared_query.score_batch(raw, &mut scores[..batch_count])?;
3201
3202 for (i, &score) in scores.iter().enumerate().take(batch_count) {
3203 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
3204 collector.push(doc_id, ordinal, score);
3205 }
3206 }
3207
3208 return Ok(collector.into_results());
3209 } else {
3210 return Ok(Vec::new());
3211 };
3212
3213 if ann_index.is_some()
3215 && !results.is_empty()
3216 && let Some(lazy_flat) = lazy_flat
3217 {
3218 return exact_score_dense_candidate_documents_sync(
3219 &results,
3220 lazy_flat,
3221 query,
3222 params.unit_norm,
3223 combiner,
3224 k,
3225 );
3226 }
3227
3228 Ok(combine_grouped_ordinal_results(results, combiner, k))
3229 }
3230
3231 #[cfg(feature = "sync")]
3236 fn search_binary_dense_vector_sync_impl(
3237 &self,
3238 field: Field,
3239 query: &[u8],
3240 k: usize,
3241 combiner: crate::query::MultiValueCombiner,
3242 probe_cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
3243 ) -> Result<Vec<VectorSearchResult>> {
3244 let schema_dim = self.validate_binary_search_request(field, query)?;
3245 combiner.validate().map_err(Error::Query)?;
3246 if k == 0 {
3247 return Ok(Vec::new());
3248 }
3249 let t0 = crate::observe::Timer::start();
3250 if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0) {
3251 let ivf = lazy.get();
3252 let config = self
3253 .schema
3254 .get_field_entry(field)
3255 .and_then(|entry| entry.binary_dense_vector_config.as_ref())
3256 .ok_or_else(|| {
3257 Error::Schema(format!(
3258 "binary IVF field {} has no schema configuration",
3259 field.0
3260 ))
3261 })?;
3262 let quantizer = self
3263 .trained_vectors
3264 .binary_quantizers
3265 .get(&field.0)
3266 .ok_or_else(|| {
3267 Error::Schema(format!(
3268 "global binary IVF field {} has no loaded quantizer",
3269 field.0
3270 ))
3271 })?;
3272 validate_binary_ann(ivf, quantizer, config, schema_dim, field)?;
3273 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
3274 Error::Corruption(format!(
3275 "global binary IVF field {} is missing flat vector storage",
3276 field.0
3277 ))
3278 })?;
3279 let single_valued = flat.num_vectors == flat.num_docs_with_vectors();
3280 let clusters = binary_probe_clusters(
3281 quantizer,
3282 query,
3283 config.nprobe,
3284 config.ivf_routing,
3285 probe_cache,
3286 )?;
3287 let results = if !single_valued
3288 && !matches!(combiner, crate::query::MultiValueCombiner::Max)
3289 {
3290 let candidate_limit =
3291 checked_binary_combined_fetch_k(k)?.min(flat.num_docs_with_vectors());
3292 let candidate_documents = ivf
3293 .search_binary_combined_documents(candidate_limit, query, &clusters, combiner)
3294 .map_err(|error| {
3295 Error::Corruption(format!(
3296 "invalid binary IVF payload for field {}: {error}",
3297 field.0,
3298 ))
3299 })?;
3300 exact_score_binary_candidate_document_ids_sync(
3301 candidate_documents
3302 .into_iter()
3303 .map(|candidate| candidate.doc_id)
3304 .collect(),
3305 flat,
3306 query,
3307 schema_dim,
3308 combiner,
3309 k,
3310 )?
3311 } else {
3312 let candidate_docs = if single_valued {
3313 k
3314 } else {
3315 checked_binary_combined_fetch_k(k)?
3316 }
3317 .min(flat.num_docs_with_vectors());
3318 let ann_results = if single_valued {
3319 ivf.search_binary_clusters::<false>(query, candidate_docs, &clusters)
3320 } else {
3321 ivf.search_binary_clusters::<true>(query, candidate_docs, &clusters)
3322 }
3323 .map_err(|error| {
3324 Error::Corruption(format!(
3325 "invalid binary IVF payload for field {}: {error}",
3326 field.0,
3327 ))
3328 })?;
3329 if single_valued {
3330 let ann_results = validate_binary_single_value_ann_results(ann_results, flat)?;
3331 combine_ordinal_results(ann_results, combiner, k)
3332 } else {
3333 exact_score_binary_candidate_documents_sync(
3334 &ann_results,
3335 flat,
3336 query,
3337 schema_dim,
3338 combiner,
3339 k,
3340 )?
3341 }
3342 };
3343 crate::observe::dense_l1(
3344 self.schema.index_label(),
3345 self.schema.get_field_name(field).unwrap_or("?"),
3346 "global_binary_ivf",
3347 t0.secs(),
3348 results.len(),
3349 );
3350 return Ok(results);
3351 }
3352 let lazy_flat = match self.flat_vectors.get(&field.0) {
3353 Some(f) => f,
3354 None => return Ok(Vec::new()),
3355 };
3356
3357 let dim_bits = lazy_flat.dim;
3358 let byte_len = lazy_flat.vector_byte_size();
3359 let n = lazy_flat.num_vectors;
3360
3361 if dim_bits != schema_dim {
3362 return Err(Error::Corruption(format!(
3363 "binary vector field {} has schema dimension {} but flat storage dimension {}",
3364 field.0, schema_dim, dim_bits
3365 )));
3366 }
3367
3368 if byte_len != query.len() {
3369 return Err(Error::Schema(format!(
3370 "Binary query vector byte length {} != field byte length {}",
3371 query.len(),
3372 byte_len
3373 )));
3374 }
3375
3376 let batch_len = bounded_vector_score_batch(byte_len, BINARY_SCORE_BATCH);
3377 let mut collector = FlatDocumentCollector::new(k, combiner);
3378 let mut scores = vec![0f32; batch_len];
3379
3380 for batch_start in (0..n).step_by(batch_len) {
3381 let batch_count = batch_len.min(n - batch_start);
3382 let batch_bytes = lazy_flat
3383 .read_vectors_batch_sync(batch_start, batch_count)
3384 .map_err(crate::Error::Io)?;
3385 let raw = batch_bytes.as_slice();
3386
3387 crate::structures::simd::batch_hamming_scores(
3388 query,
3389 raw,
3390 byte_len,
3391 dim_bits,
3392 &mut scores[..batch_count],
3393 );
3394
3395 for (i, &score) in scores.iter().enumerate().take(batch_count) {
3396 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
3397 collector.push(doc_id, ordinal, score);
3398 }
3399 }
3400
3401 let results = collector.into_results();
3402
3403 crate::observe::dense_l1(
3404 self.schema.index_label(),
3405 self.schema.get_field_name(field).unwrap_or("?"),
3406 "binary_flat",
3407 t0.secs(),
3408 results.len(),
3409 );
3410 Ok(results)
3411 }
3412
3413 #[cfg(feature = "sync")]
3414 pub fn search_binary_dense_vector_sync(
3415 &self,
3416 field: Field,
3417 query: &[u8],
3418 k: usize,
3419 combiner: crate::query::MultiValueCombiner,
3420 ) -> Result<Vec<VectorSearchResult>> {
3421 self.search_binary_dense_vector_sync_impl(field, query, k, combiner, None)
3422 }
3423
3424 #[cfg(feature = "sync")]
3425 pub(crate) fn search_binary_dense_vector_sync_with_probe_cache(
3426 &self,
3427 field: Field,
3428 query: &[u8],
3429 k: usize,
3430 combiner: crate::query::MultiValueCombiner,
3431 probe_cache: &std::sync::Mutex<Option<crate::structures::IvfProbePlan>>,
3432 ) -> Result<Vec<VectorSearchResult>> {
3433 self.search_binary_dense_vector_sync_impl(field, query, k, combiner, Some(probe_cache))
3434 }
3435}
3436
3437#[cfg(test)]
3438mod dense_search_safety_tests {
3439 use super::*;
3440
3441 #[test]
3442 fn dense_fetch_count_rejects_non_finite_and_unbounded_factors() {
3443 for factor in [
3444 f32::NAN,
3445 f32::INFINITY,
3446 f32::NEG_INFINITY,
3447 0.0,
3448 0.5,
3449 2.01,
3450 MAX_DENSE_RERANK_FACTOR + 1.0,
3451 ] {
3452 assert!(
3453 checked_dense_fetch_k(10, factor).is_err(),
3454 "factor={factor}"
3455 );
3456 }
3457 }
3458
3459 fn values_as_bytes<T>(values: &[T]) -> &[u8] {
3460 unsafe {
3461 std::slice::from_raw_parts(values.as_ptr() as *const u8, std::mem::size_of_val(values))
3462 }
3463 }
3464
3465 fn assert_prepared_dense_scores_match_legacy(
3466 quantization: DenseVectorQuantization,
3467 raw: &[u8],
3468 unit_norm: bool,
3469 ) {
3470 const DIM: usize = 4;
3471 const VECTOR_COUNT: usize = 4;
3472 let query = [0.25, -0.5, 0.75, 1.0];
3473 let mut expected = [0.0; VECTOR_COUNT];
3474 SegmentReader::score_quantized_batch_legacy(
3475 &query,
3476 raw,
3477 quantization,
3478 DIM,
3479 &mut expected,
3480 unit_norm,
3481 )
3482 .unwrap();
3483
3484 let prepared = PreparedDenseScoreQuery::new(&query, quantization, DIM, unit_norm).unwrap();
3485 let vector_bytes = DIM
3486 * match quantization {
3487 DenseVectorQuantization::F32 => std::mem::size_of::<f32>(),
3488 DenseVectorQuantization::F16 => std::mem::size_of::<u16>(),
3489 DenseVectorQuantization::UInt8 => 1,
3490 DenseVectorQuantization::Binary => unreachable!(),
3491 };
3492 let split = 2 * vector_bytes;
3493 let mut actual = [0.0; VECTOR_COUNT];
3494 prepared
3495 .score_batch(&raw[..split], &mut actual[..2])
3496 .unwrap();
3497 prepared
3498 .score_batch(&raw[split..], &mut actual[2..])
3499 .unwrap();
3500
3501 assert_eq!(
3502 actual.map(f32::to_bits),
3503 expected.map(f32::to_bits),
3504 "quantization={quantization:?}, unit_norm={unit_norm}"
3505 );
3506 }
3507
3508 #[test]
3509 fn prepared_dense_query_matches_legacy_scoring_across_batches() {
3510 let vectors_f32 = [
3511 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,
3512 -0.25,
3513 ];
3514 let vectors_f16: Vec<u16> = vectors_f32
3515 .iter()
3516 .map(|&value| crate::structures::simd::f32_to_f16(value))
3517 .collect();
3518 let vectors_u8 = [
3519 255, 96, 224, 160, 0, 192, 144, 128, 128, 128, 128, 128, 224, 192, 64, 96,
3520 ];
3521
3522 for unit_norm in [false, true] {
3523 assert_prepared_dense_scores_match_legacy(
3524 DenseVectorQuantization::F32,
3525 values_as_bytes(&vectors_f32),
3526 unit_norm,
3527 );
3528 assert_prepared_dense_scores_match_legacy(
3529 DenseVectorQuantization::F16,
3530 values_as_bytes(&vectors_f16),
3531 unit_norm,
3532 );
3533 assert_prepared_dense_scores_match_legacy(
3534 DenseVectorQuantization::UInt8,
3535 &vectors_u8,
3536 unit_norm,
3537 );
3538 }
3539 }
3540
3541 #[test]
3542 fn prepared_dense_query_preserves_scoring_validation_errors() {
3543 assert!(matches!(
3544 PreparedDenseScoreQuery::new(&[1.0], DenseVectorQuantization::F32, 2, false).err(),
3545 Some(Error::Query(_))
3546 ));
3547 assert!(matches!(
3548 PreparedDenseScoreQuery::new(&[1.0], DenseVectorQuantization::Binary, 1, false).err(),
3549 Some(Error::InvalidFieldType { .. })
3550 ));
3551
3552 let query = [1.0, 2.0];
3553 let prepared =
3554 PreparedDenseScoreQuery::new(&query, DenseVectorQuantization::F32, 2, false).unwrap();
3555 let mut scores = [0.0];
3556 assert!(matches!(
3557 prepared.score_batch(&[0; 7], &mut scores),
3558 Err(Error::Corruption(_))
3559 ));
3560 }
3561
3562 #[test]
3563 fn flat_document_collector_does_not_let_one_multivalue_doc_crowd_out_others() {
3564 let mut collector = FlatDocumentCollector::new(2, crate::query::MultiValueCombiner::Max);
3565 collector.push(1, 0, 1.0);
3566 collector.push(1, 1, 0.9);
3567 collector.push(2, 0, 0.8);
3568
3569 let results = collector.into_results();
3570 assert_eq!(
3571 results
3572 .iter()
3573 .map(|result| result.doc_id)
3574 .collect::<Vec<_>>(),
3575 vec![1, 2]
3576 );
3577 assert_eq!(results[0].ordinals.len(), 2);
3578 }
3579
3580 #[test]
3581 fn flat_document_collector_evicts_by_score_then_doc_id() {
3582 let mut collector = FlatDocumentCollector::new(2, crate::query::MultiValueCombiner::Max);
3583 collector.push(1, 0, 0.5);
3584 collector.push(3, 0, 0.8);
3585 collector.push(2, 0, 0.9);
3586 let results = collector.into_results();
3587 assert_eq!(
3588 results
3589 .iter()
3590 .map(|result| result.doc_id)
3591 .collect::<Vec<_>>(),
3592 vec![2, 3]
3593 );
3594
3595 let mut tied = FlatDocumentCollector::new(1, crate::query::MultiValueCombiner::Max);
3596 tied.push(2, 0, 1.0);
3597 tied.push(1, 0, 1.0);
3598 let results = tied.into_results();
3599 assert_eq!(results[0].doc_id, 1);
3600 }
3601
3602 #[test]
3603 fn dense_fetch_count_rounds_up_and_detects_overflow() {
3604 assert_eq!(checked_dense_fetch_k(3, 1.5).unwrap(), 5);
3605 assert_eq!(checked_dense_fetch_k(10_000, 2.0).unwrap(), 20_000);
3606 assert!(checked_dense_fetch_k(10_001, 2.0).is_err());
3607 assert!(checked_dense_fetch_k(usize::MAX, 2.0).is_err());
3608 }
3609
3610 #[test]
3611 fn binary_combined_fetch_count_uses_shared_bounded_oversampling() {
3612 assert_eq!(checked_binary_combined_fetch_k(3).unwrap(), 6);
3613 assert_eq!(checked_binary_combined_fetch_k(10_000).unwrap(), 20_000);
3614 assert_eq!(checked_binary_combined_fetch_k(10_001).unwrap(), 20_000);
3615 assert_eq!(checked_binary_combined_fetch_k(20_000).unwrap(), 20_000);
3616 assert!(checked_binary_combined_fetch_k(20_001).is_err());
3617 assert!(checked_binary_combined_fetch_k(usize::MAX).is_err());
3618 }
3619
3620 #[cfg(feature = "native")]
3621 #[test]
3622 fn legacy_ivf_tq_generation_is_rejected_while_opening() {
3623 use crate::directories::OwnedBytes;
3624 use crate::dsl::IvfRoutingMode;
3625 use crate::segment::ann_disk::{AnnDiskIndex, AnnKind};
3626
3627 let centroids = CoarseCentroids {
3628 num_clusters: 1,
3629 dim: 2,
3630 centroids: vec![1.0, 0.0],
3631 version: 7,
3632 soar_config: None,
3633 routing_index: None,
3634 };
3635 let mut build_centroids = centroids.clone();
3636 build_centroids.version =
3637 crate::structures::mark_ivf_tq_cosine_generation(build_centroids.version);
3638 let mut bytes = crate::segment::ann_build::build_ivf_tq(
3639 2,
3640 IvfRoutingMode::Flat,
3641 &build_centroids,
3642 &[(0, 0)],
3643 &[1.0, 0.0],
3644 )
3645 .unwrap();
3646 bytes[24..32].copy_from_slice(¢roids.version.to_le_bytes());
3649 let error = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::IvfTq, 1)
3650 .err()
3651 .expect("legacy IVF-TQ payload must fail while opening")
3652 .to_string();
3653 assert!(error.contains("unsupported legacy generation"), "{error}");
3654 }
3655
3656 #[test]
3657 fn rerank_batch_is_capped_by_actual_candidate_vectors() {
3658 assert_eq!(bounded_rerank_batch(3_072, DENSE_SCORE_BATCH, 20), 20);
3659 assert_eq!(
3660 bounded_rerank_batch(3_072, DENSE_SCORE_BATCH, 10_000),
3661 MAX_VECTOR_SCORE_BATCH_BYTES / 3_072
3662 );
3663 assert_eq!(bounded_rerank_batch(3_072, DENSE_SCORE_BATCH, 0), 1);
3664 }
3665
3666 #[test]
3667 fn file_ranges_reject_overflow_and_truncation() {
3668 assert_eq!(checked_file_range(4, 3, 7, "test").unwrap(), 4..7);
3669 assert!(checked_file_range(u64::MAX, 1, u64::MAX, "test").is_err());
3670 assert!(checked_file_range(5, 3, 7, "test").is_err());
3671 }
3672
3673 #[test]
3674 fn shared_tq_plan_cache_rebuilds_for_divergent_query_clones() {
3675 let codec = crate::structures::TqCodec::new(4);
3676 let cache = std::sync::Mutex::new(None);
3677 let original_query = vec![1.0, 2.0, 3.0, 4.0];
3678
3679 let original =
3680 cached_tq_query_plan(&codec, &original_query, Some(&cache)).expect("build plan");
3681 let reused =
3682 cached_tq_query_plan(&codec, &original_query, Some(&cache)).expect("reuse plan");
3683 assert!(
3684 std::sync::Arc::ptr_eq(&original, &reused),
3685 "unchanged queries must share their plan across segments"
3686 );
3687
3688 let mut divergent_clone = original_query.clone();
3689 divergent_clone[0] = -1.0;
3690 let rebuilt =
3691 cached_tq_query_plan(&codec, &divergent_clone, Some(&cache)).expect("rebuild plan");
3692 assert!(
3693 !std::sync::Arc::ptr_eq(&original, &rebuilt),
3694 "a clone with a mutated vector must not reuse stale LUTs"
3695 );
3696 assert!(rebuilt.matches_query(&divergent_clone));
3697 assert!(!rebuilt.matches_query(&original_query));
3698 }
3699
3700 #[test]
3701 fn candidate_vector_reads_coalesce_contiguous_values() {
3702 let mut runs = Vec::new();
3703 plan_vector_read_runs(&[3, 4, 5, 9, 12, 13], &mut runs).unwrap();
3704 assert_eq!(runs.len(), 3);
3705 assert!(matches!(
3706 runs.as_slice(),
3707 [
3708 VectorReadRun {
3709 buffer_start: 0,
3710 flat_start: 3,
3711 count: 3,
3712 },
3713 VectorReadRun {
3714 buffer_start: 3,
3715 flat_start: 9,
3716 count: 1,
3717 },
3718 VectorReadRun {
3719 buffer_start: 4,
3720 flat_start: 12,
3721 count: 2,
3722 },
3723 ]
3724 ));
3725 assert!(plan_vector_read_runs(&[3, 3], &mut runs).is_err());
3726 }
3727
3728 #[tokio::test]
3729 async fn binary_single_value_ann_fast_path_validates_and_deduplicates() {
3730 use crate::directories::{FileHandle, OwnedBytes};
3731 use crate::segment::FlatVectorData;
3732
3733 let mut encoded = Vec::new();
3734 FlatVectorData::serialize_binary_from_bits_streaming(
3735 8,
3736 &[0x0f, 0xf0],
3737 &[(1, 0), (3, 2)],
3738 &mut encoded,
3739 )
3740 .unwrap();
3741 let flat = LazyFlatVectorData::open_with_doc_limit(
3742 FileHandle::from_bytes(OwnedBytes::new(encoded)),
3743 Some(4),
3744 )
3745 .await
3746 .unwrap();
3747 assert_eq!(flat.num_vectors, flat.num_docs_with_vectors());
3748
3749 let validated = validate_binary_single_value_ann_results(
3750 vec![(3, 2, 0.9), (1, 0, 0.8), (3, 2, 0.7)],
3751 &flat,
3752 )
3753 .unwrap();
3754 assert_eq!(validated, vec![(3, 2, 0.9), (1, 0, 0.8)]);
3755
3756 assert!(matches!(
3757 validate_binary_single_value_ann_results(vec![(2, 0, 1.0)], &flat),
3758 Err(Error::Corruption(_))
3759 ));
3760 assert!(matches!(
3761 validate_binary_single_value_ann_results(vec![(3, 0, 1.0)], &flat),
3762 Err(Error::Corruption(_))
3763 ));
3764 }
3765
3766 #[tokio::test]
3767 async fn multivalue_ann_rerank_streams_past_document_candidate_cap() {
3768 use crate::directories::{FileHandle, OwnedBytes};
3769 use crate::segment::FlatVectorData;
3770
3771 const VALUES: usize = MAX_DENSE_CANDIDATES_PER_SEGMENT + 1;
3772 let mut encoded = Vec::new();
3773 let vectors = vec![1.0f32; VALUES];
3774 let doc_ids: Vec<_> = (0..VALUES).map(|ordinal| (0, ordinal as u16)).collect();
3775 FlatVectorData::serialize_binary_from_flat_streaming(
3776 1,
3777 &vectors,
3778 &doc_ids,
3779 DenseVectorQuantization::F32,
3780 &mut encoded,
3781 )
3782 .unwrap();
3783 let flat = LazyFlatVectorData::open_with_doc_limit(
3784 FileHandle::from_bytes(OwnedBytes::new(encoded)),
3785 Some(1),
3786 )
3787 .await
3788 .unwrap();
3789
3790 let (results, stats) = exact_score_dense_candidate_documents(
3791 &[(0, 0, 0.0)],
3792 &flat,
3793 &[1.0],
3794 false,
3795 crate::query::MultiValueCombiner::Max,
3796 1,
3797 )
3798 .await
3799 .unwrap();
3800 assert_eq!(stats.vector_count, VALUES);
3801 assert_eq!(results.len(), 1);
3802 assert_eq!(results[0].ordinals.len(), VALUES);
3803 assert!((results[0].score - 1.0).abs() < 1e-5);
3804 }
3805}