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;
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
298fn checked_dense_fetch_k(k: usize, rerank_factor: f32) -> Result<usize> {
300 if !rerank_factor.is_finite() || !(1.0..=MAX_DENSE_RERANK_FACTOR).contains(&rerank_factor) {
301 return Err(Error::Query(format!(
302 "dense rerank_factor must be finite and in [1, {MAX_DENSE_RERANK_FACTOR}], got {rerank_factor}"
303 )));
304 }
305
306 let fetch = (k as f64) * (rerank_factor as f64);
307 if !fetch.is_finite()
308 || fetch > usize::MAX as f64
309 || fetch > MAX_DENSE_CANDIDATES_PER_SEGMENT as f64
310 {
311 return Err(Error::Query(format!(
312 "dense candidate count exceeds the per-segment maximum of \
313 {MAX_DENSE_CANDIDATES_PER_SEGMENT}: k={k}, rerank_factor={rerank_factor}"
314 )));
315 }
316 Ok(fetch.ceil() as usize)
317}
318
319#[inline]
320fn bounded_vector_score_batch(vector_byte_size: usize, preferred: usize) -> usize {
321 preferred.min((MAX_VECTOR_SCORE_BATCH_BYTES / vector_byte_size.max(1)).max(1))
322}
323
324fn checked_file_range(
325 offset: u64,
326 length: u64,
327 file_length: u64,
328 description: &str,
329) -> Result<std::ops::Range<u64>> {
330 let end = offset
331 .checked_add(length)
332 .ok_or_else(|| Error::Corruption(format!("{description} byte range overflows u64")))?;
333 if end > file_length {
334 return Err(Error::Corruption(format!(
335 "{description} byte range {offset}..{end} exceeds file length {file_length}"
336 )));
337 }
338 Ok(offset..end)
339}
340
341type RawVectorCandidate = (u32, u16, f32);
342type CandidateVectorRef = (DocId, u16, usize); #[derive(Clone, Copy)]
345struct CandidateDocumentRange {
346 doc_id: DocId,
347 start: usize,
348 end: usize,
349}
350
351struct AnnCandidateDocuments {
352 ranges: Vec<CandidateDocumentRange>,
353 vector_count: usize,
354}
355
356fn ann_candidate_document_ranges(
365 ann_results: &[RawVectorCandidate],
366 flat: &LazyFlatVectorData,
367) -> Result<AnnCandidateDocuments> {
368 let mut candidate_docs: Vec<DocId> = ann_results.iter().map(|candidate| candidate.0).collect();
369 candidate_docs.sort_unstable();
370 candidate_docs.dedup();
371
372 let mut ranges = Vec::with_capacity(candidate_docs.len());
373 let mut vector_count = 0usize;
374 for doc_id in candidate_docs {
375 let (start, count) = flat.flat_indexes_for_doc_range(doc_id);
376 if count == 0 {
377 return Err(Error::Corruption(format!(
378 "ANN candidate document {doc_id} is missing from flat vector storage"
379 )));
380 }
381 vector_count = vector_count
382 .checked_add(count)
383 .ok_or_else(|| Error::Query("ANN candidate vector expansion overflow".to_string()))?;
384 let end = start
385 .checked_add(count)
386 .ok_or_else(|| Error::Corruption("flat vector range overflow".to_string()))?;
387 if end > flat.num_vectors {
388 return Err(Error::Corruption(format!(
389 "flat vector range {start}..{end} for document {doc_id} exceeds {} vectors",
390 flat.num_vectors
391 )));
392 }
393 ranges.push(CandidateDocumentRange { doc_id, start, end });
394 }
395 Ok(AnnCandidateDocuments {
396 ranges,
397 vector_count,
398 })
399}
400
401struct CandidateVectorCursor<'a> {
402 ranges: &'a [CandidateDocumentRange],
403 range_index: usize,
404 flat_index: usize,
405}
406
407impl<'a> CandidateVectorCursor<'a> {
408 fn new(ranges: &'a [CandidateDocumentRange]) -> Self {
409 Self {
410 ranges,
411 range_index: 0,
412 flat_index: ranges.first().map_or(0, |range| range.start),
413 }
414 }
415
416 fn fill_batch(
420 &mut self,
421 flat: &LazyFlatVectorData,
422 batch: &mut Vec<CandidateVectorRef>,
423 limit: usize,
424 ) -> Result<bool> {
425 batch.clear();
426 while batch.len() < limit && self.range_index < self.ranges.len() {
427 let range = self.ranges[self.range_index];
428 if self.flat_index == range.end {
429 self.range_index += 1;
430 if let Some(next) = self.ranges.get(self.range_index) {
431 self.flat_index = next.start;
432 }
433 continue;
434 }
435 let (stored_doc_id, ordinal) = flat.get_doc_id(self.flat_index);
436 if stored_doc_id != range.doc_id {
437 return Err(Error::Corruption(format!(
438 "flat vector doc map is not contiguous for document {}",
439 range.doc_id
440 )));
441 }
442 batch.push((range.doc_id, ordinal, self.flat_index));
443 self.flat_index += 1;
444 }
445 Ok(!batch.is_empty())
446 }
447}
448
449#[derive(Clone, Copy)]
450struct VectorReadRun {
451 buffer_start: usize,
452 flat_start: usize,
453 count: usize,
454}
455
456fn plan_vector_read_runs(indexes: &[usize], runs: &mut Vec<VectorReadRun>) -> Result<()> {
461 runs.clear();
462 for (buffer_index, &flat_index) in indexes.iter().enumerate() {
463 if let Some(run) = runs.last_mut()
464 && run
465 .flat_start
466 .checked_add(run.count)
467 .is_some_and(|next| next == flat_index)
468 {
469 run.count += 1;
470 continue;
471 }
472 if buffer_index > 0 && flat_index <= indexes[buffer_index - 1] {
473 return Err(Error::Corruption(
474 "candidate flat-vector indexes are not strictly ordered".into(),
475 ));
476 }
477 runs.push(VectorReadRun {
478 buffer_start: buffer_index,
479 flat_start: flat_index,
480 count: 1,
481 });
482 }
483 Ok(())
484}
485
486async fn read_vector_runs(
487 flat: &LazyFlatVectorData,
488 indexes: &[usize],
489 runs: &mut Vec<VectorReadRun>,
490 output: &mut [u8],
491) -> Result<()> {
492 plan_vector_read_runs(indexes, runs)?;
493 let vector_byte_size = flat.vector_byte_size();
494 for run in runs {
495 let bytes = flat
496 .read_vectors_batch(run.flat_start, run.count)
497 .await
498 .map_err(Error::Io)?;
499 let start = run
500 .buffer_start
501 .checked_mul(vector_byte_size)
502 .ok_or_else(|| Error::Query("dense rerank buffer offset overflow".into()))?;
503 let end = start
504 .checked_add(bytes.len())
505 .ok_or_else(|| Error::Query("dense rerank buffer range overflow".into()))?;
506 let destination = output
507 .get_mut(start..end)
508 .ok_or_else(|| Error::Corruption("dense rerank buffer is too short".into()))?;
509 destination.copy_from_slice(bytes.as_slice());
510 }
511 Ok(())
512}
513
514#[cfg(feature = "sync")]
515fn read_vector_runs_sync(
516 flat: &LazyFlatVectorData,
517 indexes: &[usize],
518 runs: &mut Vec<VectorReadRun>,
519 output: &mut [u8],
520) -> Result<()> {
521 plan_vector_read_runs(indexes, runs)?;
522 let vector_byte_size = flat.vector_byte_size();
523 for run in runs {
524 let bytes = flat
525 .read_vectors_batch_sync(run.flat_start, run.count)
526 .map_err(Error::Io)?;
527 let start = run
528 .buffer_start
529 .checked_mul(vector_byte_size)
530 .ok_or_else(|| Error::Query("dense rerank buffer offset overflow".into()))?;
531 let end = start
532 .checked_add(bytes.len())
533 .ok_or_else(|| Error::Query("dense rerank buffer range overflow".into()))?;
534 let destination = output
535 .get_mut(start..end)
536 .ok_or_else(|| Error::Corruption("dense rerank buffer is too short".into()))?;
537 destination.copy_from_slice(bytes.as_slice());
538 }
539 Ok(())
540}
541
542#[derive(Default)]
543struct DenseRerankStats {
544 vector_count: usize,
545 resolve_elapsed: std::time::Duration,
546 read_elapsed: std::time::Duration,
547 score_elapsed: std::time::Duration,
548}
549
550async fn exact_score_dense_candidate_documents(
551 ann_results: &[RawVectorCandidate],
552 flat: &LazyFlatVectorData,
553 query: &[f32],
554 unit_norm: bool,
555 combiner: crate::query::MultiValueCombiner,
556 limit: usize,
557) -> Result<(Vec<VectorSearchResult>, DenseRerankStats)> {
558 let resolve_started = std::time::Instant::now();
559 let documents = ann_candidate_document_ranges(ann_results, flat)?;
560 let mut stats = DenseRerankStats {
561 vector_count: documents.vector_count,
562 resolve_elapsed: resolve_started.elapsed(),
563 ..Default::default()
564 };
565 let vector_byte_size = flat.vector_byte_size();
566 let batch_len = bounded_vector_score_batch(vector_byte_size, DENSE_SCORE_BATCH);
567 let raw_capacity = batch_len
568 .checked_mul(vector_byte_size)
569 .ok_or_else(|| Error::Query("dense rerank buffer size overflow".to_string()))?;
570 let mut raw = vec![0u8; raw_capacity];
571 let mut scores = vec![0.0f32; batch_len];
572 let mut batch = Vec::with_capacity(batch_len);
573 let mut flat_indexes = Vec::with_capacity(batch_len);
574 let mut read_runs = Vec::new();
575 let mut cursor = CandidateVectorCursor::new(&documents.ranges);
576 let mut collector = FlatDocumentCollector::new(limit, combiner);
577 let mut scored = 0usize;
578
579 while cursor.fill_batch(flat, &mut batch, batch_len)? {
580 flat_indexes.clear();
581 flat_indexes.extend(batch.iter().map(|&(_, _, flat_index)| flat_index));
582 #[cfg(feature = "native")]
583 flat.prefetch_vectors(flat_indexes.iter().copied());
584 let raw_len = batch
585 .len()
586 .checked_mul(vector_byte_size)
587 .ok_or_else(|| Error::Query("dense rerank buffer size overflow".to_string()))?;
588 let raw = &mut raw[..raw_len];
589
590 let read_started = std::time::Instant::now();
591 read_vector_runs(flat, &flat_indexes, &mut read_runs, raw).await?;
592 stats.read_elapsed += read_started.elapsed();
593
594 let score_started = std::time::Instant::now();
595 SegmentReader::score_quantized_batch(
596 query,
597 raw,
598 flat.quantization,
599 flat.dim,
600 &mut scores[..batch.len()],
601 unit_norm,
602 )?;
603 stats.score_elapsed += score_started.elapsed();
604 for (buffer_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
605 collector.push(doc_id, ordinal, scores[buffer_index]);
606 }
607 scored += batch.len();
608 }
609 debug_assert_eq!(scored, documents.vector_count);
610 Ok((collector.into_results(), stats))
611}
612
613#[cfg(feature = "sync")]
614fn exact_score_dense_candidate_documents_sync(
615 ann_results: &[RawVectorCandidate],
616 flat: &LazyFlatVectorData,
617 query: &[f32],
618 unit_norm: bool,
619 combiner: crate::query::MultiValueCombiner,
620 limit: usize,
621) -> Result<Vec<VectorSearchResult>> {
622 let documents = ann_candidate_document_ranges(ann_results, flat)?;
623 let vector_byte_size = flat.vector_byte_size();
624 let batch_len = bounded_vector_score_batch(vector_byte_size, DENSE_SCORE_BATCH);
625 let raw_capacity = batch_len
626 .checked_mul(vector_byte_size)
627 .ok_or_else(|| Error::Query("dense rerank buffer size overflow".to_string()))?;
628 let mut raw = vec![0u8; raw_capacity];
629 let mut scores = vec![0.0f32; batch_len];
630 let mut batch = Vec::with_capacity(batch_len);
631 let mut flat_indexes = Vec::with_capacity(batch_len);
632 let mut read_runs = Vec::new();
633 let mut cursor = CandidateVectorCursor::new(&documents.ranges);
634 let mut collector = FlatDocumentCollector::new(limit, combiner);
635 let mut scored = 0usize;
636
637 while cursor.fill_batch(flat, &mut batch, batch_len)? {
638 flat_indexes.clear();
639 flat_indexes.extend(batch.iter().map(|&(_, _, flat_index)| flat_index));
640 let raw_len = batch
641 .len()
642 .checked_mul(vector_byte_size)
643 .ok_or_else(|| Error::Query("dense rerank buffer size overflow".to_string()))?;
644 let raw = &mut raw[..raw_len];
645 read_vector_runs_sync(flat, &flat_indexes, &mut read_runs, raw)?;
646 SegmentReader::score_quantized_batch(
647 query,
648 raw,
649 flat.quantization,
650 flat.dim,
651 &mut scores[..batch.len()],
652 unit_norm,
653 )?;
654 for (buffer_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
655 collector.push(doc_id, ordinal, scores[buffer_index]);
656 }
657 scored += batch.len();
658 }
659 debug_assert_eq!(scored, documents.vector_count);
660 Ok(collector.into_results())
661}
662
663async fn exact_score_binary_candidate_documents(
664 ann_results: &[RawVectorCandidate],
665 flat: &LazyFlatVectorData,
666 query: &[u8],
667 dim_bits: usize,
668 combiner: crate::query::MultiValueCombiner,
669 limit: usize,
670) -> Result<Vec<VectorSearchResult>> {
671 let documents = ann_candidate_document_ranges(ann_results, flat)?;
672 let probe_scores: FxHashMap<(DocId, u16), f32> = ann_results
673 .iter()
674 .map(|&(doc_id, ordinal, score)| ((doc_id, ordinal), score))
675 .collect();
676 let vector_byte_size = flat.vector_byte_size();
677 let batch_len = bounded_vector_score_batch(vector_byte_size, BINARY_SCORE_BATCH);
678 let raw_capacity = batch_len
679 .checked_mul(vector_byte_size)
680 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
681 let mut raw = vec![0u8; raw_capacity];
682 let mut scores = vec![0.0f32; batch_len];
683 let mut batch_scores = vec![0.0f32; batch_len];
684 let mut batch = Vec::with_capacity(batch_len);
685 let mut unresolved = Vec::with_capacity(batch_len);
686 let mut unresolved_flat_indexes = Vec::with_capacity(batch_len);
687 let mut read_runs = Vec::new();
688 let mut cursor = CandidateVectorCursor::new(&documents.ranges);
689 let mut collector = FlatDocumentCollector::new(limit, combiner);
690 let mut scored = 0usize;
691
692 while cursor.fill_batch(flat, &mut batch, batch_len)? {
693 unresolved.clear();
694 for (batch_index, &(doc_id, ordinal, flat_index)) in batch.iter().enumerate() {
695 if let Some(&score) = probe_scores.get(&(doc_id, ordinal)) {
696 batch_scores[batch_index] = score;
697 } else {
698 unresolved.push((batch_index, flat_index));
699 }
700 }
701 unresolved_flat_indexes.clear();
702 unresolved_flat_indexes.extend(unresolved.iter().map(|&(_, flat_index)| flat_index));
703 #[cfg(feature = "native")]
704 flat.prefetch_vectors(unresolved_flat_indexes.iter().copied());
705 let raw_len = unresolved
706 .len()
707 .checked_mul(vector_byte_size)
708 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
709 let raw = &mut raw[..raw_len];
710 read_vector_runs(flat, &unresolved_flat_indexes, &mut read_runs, raw).await?;
711 crate::structures::simd::batch_hamming_scores(
712 query,
713 raw,
714 vector_byte_size,
715 dim_bits,
716 &mut scores[..unresolved.len()],
717 );
718 for (buffer_index, &(batch_index, _)) in unresolved.iter().enumerate() {
719 batch_scores[batch_index] = scores[buffer_index];
720 }
721 for (batch_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
722 collector.push(doc_id, ordinal, batch_scores[batch_index]);
723 }
724 scored += batch.len();
725 }
726 debug_assert_eq!(scored, documents.vector_count);
727 Ok(collector.into_results())
728}
729
730#[cfg(feature = "sync")]
731fn exact_score_binary_candidate_documents_sync(
732 ann_results: &[RawVectorCandidate],
733 flat: &LazyFlatVectorData,
734 query: &[u8],
735 dim_bits: usize,
736 combiner: crate::query::MultiValueCombiner,
737 limit: usize,
738) -> Result<Vec<VectorSearchResult>> {
739 let documents = ann_candidate_document_ranges(ann_results, flat)?;
740 let probe_scores: FxHashMap<(DocId, u16), f32> = ann_results
741 .iter()
742 .map(|&(doc_id, ordinal, score)| ((doc_id, ordinal), score))
743 .collect();
744 let vector_byte_size = flat.vector_byte_size();
745 let batch_len = bounded_vector_score_batch(vector_byte_size, BINARY_SCORE_BATCH);
746 let raw_capacity = batch_len
747 .checked_mul(vector_byte_size)
748 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
749 let mut raw = vec![0u8; raw_capacity];
750 let mut scores = vec![0.0f32; batch_len];
751 let mut batch_scores = vec![0.0f32; batch_len];
752 let mut batch = Vec::with_capacity(batch_len);
753 let mut unresolved = Vec::with_capacity(batch_len);
754 let mut unresolved_flat_indexes = Vec::with_capacity(batch_len);
755 let mut read_runs = Vec::new();
756 let mut cursor = CandidateVectorCursor::new(&documents.ranges);
757 let mut collector = FlatDocumentCollector::new(limit, combiner);
758 let mut scored = 0usize;
759
760 while cursor.fill_batch(flat, &mut batch, batch_len)? {
761 unresolved.clear();
762 for (batch_index, &(doc_id, ordinal, flat_index)) in batch.iter().enumerate() {
763 if let Some(&score) = probe_scores.get(&(doc_id, ordinal)) {
764 batch_scores[batch_index] = score;
765 } else {
766 unresolved.push((batch_index, flat_index));
767 }
768 }
769 unresolved_flat_indexes.clear();
770 unresolved_flat_indexes.extend(unresolved.iter().map(|&(_, flat_index)| flat_index));
771 let raw_len = unresolved
772 .len()
773 .checked_mul(vector_byte_size)
774 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
775 let raw = &mut raw[..raw_len];
776 read_vector_runs_sync(flat, &unresolved_flat_indexes, &mut read_runs, raw)?;
777 crate::structures::simd::batch_hamming_scores(
778 query,
779 raw,
780 vector_byte_size,
781 dim_bits,
782 &mut scores[..unresolved.len()],
783 );
784 for (buffer_index, &(batch_index, _)) in unresolved.iter().enumerate() {
785 batch_scores[batch_index] = scores[buffer_index];
786 }
787 for (batch_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
788 collector.push(doc_id, ordinal, batch_scores[batch_index]);
789 }
790 scored += batch.len();
791 }
792 debug_assert_eq!(scored, documents.vector_count);
793 Ok(collector.into_results())
794}
795
796fn validate_coarse_centroids(centroids: &CoarseCentroids, dim: usize) -> Result<()> {
797 let expected = (centroids.num_clusters as usize)
798 .checked_mul(dim)
799 .ok_or_else(|| Error::Corruption("coarse centroid size overflow".into()))?;
800 if centroids.num_clusters == 0
801 || centroids.dim != dim
802 || centroids.centroids.len() != expected
803 || centroids.centroids.iter().any(|value| !value.is_finite())
804 {
805 return Err(Error::Corruption(format!(
806 "invalid coarse centroids: clusters={}, dim={}, values={} (expected dim={dim}, values={expected})",
807 centroids.num_clusters,
808 centroids.dim,
809 centroids.centroids.len()
810 )));
811 }
812 Ok(())
813}
814
815fn validate_ivf_pq_ann(
816 index: &crate::segment::ann_disk::AnnDiskIndex,
817 centroids: &CoarseCentroids,
818 codebook: &crate::structures::PQCodebook,
819 dim: usize,
820 routing: crate::dsl::IvfRoutingMode,
821) -> Result<()> {
822 let header = index.header();
823 if header.dim != dim
824 || codebook.config.dim != dim
825 || header.code_size != codebook.config.num_subspaces
826 || header.num_clusters != centroids.num_clusters
827 || header.quantizer_version != centroids.version
828 || header.codebook_version != codebook.version
829 || header.routing != routing
830 {
831 return Err(Error::Corruption(format!(
832 "IVF-PQ payload/codebook/centroid metadata does not match schema dimension {dim}"
833 )));
834 }
835 Ok(())
836}
837
838fn validate_binary_ann(
839 index: &crate::segment::ann_disk::AnnDiskIndex,
840 quantizer: &crate::structures::BinaryCoarseQuantizer,
841 config: &crate::dsl::BinaryDenseVectorConfig,
842 dim: usize,
843 field: Field,
844) -> Result<()> {
845 let header = index.header();
846 if header.dim != dim
847 || header.code_size != config.byte_len()
848 || header.num_clusters != quantizer.num_clusters
849 || header.quantizer_version != quantizer.version
850 || header.codebook_version != 0
851 || header.routing != config.ivf_routing
852 || quantizer.dim_bits != dim
853 {
854 return Err(Error::Corruption(format!(
855 "binary IVF field {} does not match its quantizer/schema generation",
856 field.0,
857 )));
858 }
859 Ok(())
860}
861
862fn float_query_plan(
863 centroids: &CoarseCentroids,
864 codebook: &crate::structures::PQCodebook,
865 query: &[f32],
866 nprobe: usize,
867 routing: crate::dsl::IvfRoutingMode,
868 cache: Option<&std::sync::Mutex<Option<std::sync::Arc<crate::structures::IvfPqQueryPlan>>>>,
869) -> Result<std::sync::Arc<crate::structures::IvfPqQueryPlan>> {
870 let effective_nprobe = nprobe.clamp(1, centroids.num_clusters as usize);
871 let request_fingerprint = crate::structures::vector::ivf::routing::float_probe_fingerprint(
872 query,
873 effective_nprobe,
874 routing,
875 );
876 if let Some(cache) = cache {
877 let mut cached = cache
878 .lock()
879 .map_err(|_| Error::Internal("dense IVF probe cache is poisoned".into()))?;
880 if let Some(plan) = cached.as_ref()
881 && plan.quantizer_version == centroids.version
882 && plan.codebook_version == codebook.version
883 && plan.request_fingerprint == request_fingerprint
884 && plan.cluster_ids.len() == effective_nprobe
885 {
886 return Ok(std::sync::Arc::clone(plan));
887 }
888 let plan = std::sync::Arc::new(crate::structures::IvfPqQueryPlan::build(
889 centroids,
890 codebook,
891 query,
892 effective_nprobe,
893 routing,
894 ));
895 *cached = Some(std::sync::Arc::clone(&plan));
896 return Ok(plan);
897 }
898 Ok(std::sync::Arc::new(
899 crate::structures::IvfPqQueryPlan::build(
900 centroids,
901 codebook,
902 query,
903 effective_nprobe,
904 routing,
905 ),
906 ))
907}
908
909fn binary_probe_clusters(
910 quantizer: &crate::structures::BinaryCoarseQuantizer,
911 query: &[u8],
912 nprobe: usize,
913 routing: crate::dsl::IvfRoutingMode,
914 cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
915) -> Result<std::sync::Arc<[u32]>> {
916 let effective_nprobe = nprobe.clamp(1, quantizer.num_clusters as usize);
917 let request_fingerprint = crate::structures::vector::ivf::routing::binary_probe_fingerprint(
918 query,
919 effective_nprobe,
920 routing,
921 );
922 if let Some(cache) = cache {
923 let mut cached = cache
924 .lock()
925 .map_err(|_| Error::Internal("binary IVF probe cache is poisoned".into()))?;
926 if let Some(plan) = cached.as_ref()
927 && plan.quantizer_version == quantizer.version
928 && plan.request_fingerprint == request_fingerprint
929 && plan.cluster_ids.len() == effective_nprobe
930 {
931 return Ok(std::sync::Arc::clone(&plan.cluster_ids));
932 }
933 let plan = quantizer.probe(query, effective_nprobe, routing);
934 let clusters = std::sync::Arc::clone(&plan.cluster_ids);
935 *cached = Some(plan);
936 return Ok(clusters);
937 }
938 Ok(quantizer
939 .probe(query, effective_nprobe, routing)
940 .cluster_ids)
941}
942
943pub struct SegmentReader {
949 meta: SegmentMeta,
950 term_dict: Arc<AsyncSSTableReader<TermInfo>>,
952 postings_handle: FileHandle,
954 store: Arc<AsyncStoreReader>,
956 schema: Arc<Schema>,
957 vector_indexes: FxHashMap<u32, VectorIndex>,
959 flat_vectors: FxHashMap<u32, LazyFlatVectorData>,
961 dense_file_backed_bytes: u64,
963 trained_vectors: Arc<crate::segment::TrainedVectorStructures>,
965 sparse_indexes: FxHashMap<u32, SparseIndex>,
967 bmp_indexes: FxHashMap<u32, BmpIndex>,
969 sparse_file_backed_bytes: u64,
971 positions_handle: Option<FileHandle>,
973 fast_fields: FxHashMap<u32, crate::structures::fast_field::FastFieldReader>,
975 #[cfg(feature = "native")]
977 dense_pin_report: crate::segment::pin::PinReport,
978 #[cfg(feature = "native")]
980 sparse_pin_report: crate::segment::pin::PinReport,
981}
982
983impl SegmentReader {
984 pub async fn open<D: Directory>(
986 dir: &D,
987 segment_id: SegmentId,
988 schema: Arc<Schema>,
989 cache_blocks: usize,
990 ) -> Result<Self> {
991 Self::open_with_cache_blocks(dir, segment_id, schema, cache_blocks, cache_blocks).await
992 }
993
994 pub async fn open_with_cache_blocks<D: Directory>(
1000 dir: &D,
1001 segment_id: SegmentId,
1002 schema: Arc<Schema>,
1003 term_cache_blocks: usize,
1004 store_cache_blocks: usize,
1005 ) -> Result<Self> {
1006 let files = SegmentFiles::new(segment_id.0);
1007
1008 let meta_slice = dir.open_read(&files.meta).await?;
1010 let meta_bytes = meta_slice.read_bytes().await?;
1011 let meta = SegmentMeta::deserialize(meta_bytes.as_slice())?;
1012 debug_assert_eq!(meta.id, segment_id.0);
1013
1014 let term_dict_handle = dir.open_lazy(&files.term_dict).await?;
1016 let term_dict = AsyncSSTableReader::open(term_dict_handle, term_cache_blocks).await?;
1017
1018 let postings_handle = dir.open_lazy(&files.postings).await?;
1020
1021 let store_handle = dir.open_lazy(&files.store).await?;
1023 let store = AsyncStoreReader::open(store_handle, store_cache_blocks).await?;
1024
1025 let vectors_data = loader::load_vectors_file(dir, &files, &schema, meta.num_docs).await?;
1027 let dense_file_backed_bytes = vectors_data.file_backed_bytes;
1028 let vector_indexes = vectors_data.indexes;
1029 let flat_vectors = vectors_data.flat_vectors;
1030
1031 #[cfg(feature = "native")]
1036 for (field_id, lazy_flat) in &flat_vectors {
1037 if vector_indexes.contains_key(field_id) {
1038 lazy_flat.advise_random_access();
1039 }
1040 }
1041
1042 let sparse_data = loader::load_sparse_file(dir, &files, meta.num_docs, &schema).await?;
1044 let sparse_file_backed_bytes = sparse_data.file_backed_bytes;
1045 let sparse_indexes = sparse_data.maxscore_indexes;
1046 let bmp_indexes = sparse_data.bmp_indexes;
1047
1048 let positions_handle = loader::open_positions_file(dir, &files, &schema).await?;
1050
1051 let fast_fields = loader::load_fast_fields_file(dir, &files, &schema).await?;
1053
1054 {
1056 let mut parts = vec![format!(
1057 "[segment] loaded {:016x}: docs={}",
1058 segment_id.0, meta.num_docs
1059 )];
1060 if !vector_indexes.is_empty() || !flat_vectors.is_empty() {
1061 parts.push(format!(
1062 "dense vectors: {} ANN + {} flat fields",
1063 vector_indexes.len(),
1064 flat_vectors.len()
1065 ));
1066 }
1067 for (field_id, idx) in &sparse_indexes {
1068 parts.push(format!(
1069 "sparse vector field {}: {} dims, ~{}",
1070 field_id,
1071 idx.num_dimensions(),
1072 crate::format_bytes(idx.num_dimensions() as u64 * 24)
1073 ));
1074 }
1075 for (field_id, idx) in &bmp_indexes {
1076 parts.push(format!(
1077 "bmp field {}: {} dims, {} blocks",
1078 field_id,
1079 idx.dims(),
1080 idx.num_blocks
1081 ));
1082 }
1083 if !fast_fields.is_empty() {
1084 parts.push(format!("fast: {} fields", fast_fields.len()));
1085 }
1086 log::debug!("{}", parts.join(", "));
1087 }
1088
1089 #[allow(unused_mut)]
1090 let mut reader = Self {
1091 meta,
1092 term_dict: Arc::new(term_dict),
1093 postings_handle,
1094 store: Arc::new(store),
1095 schema,
1096 vector_indexes,
1097 flat_vectors,
1098 dense_file_backed_bytes,
1099 trained_vectors: Arc::new(crate::segment::TrainedVectorStructures::default()),
1100 sparse_indexes,
1101 bmp_indexes,
1102 sparse_file_backed_bytes,
1103 positions_handle,
1104 fast_fields,
1105 #[cfg(feature = "native")]
1106 dense_pin_report: Default::default(),
1107 #[cfg(feature = "native")]
1108 sparse_pin_report: Default::default(),
1109 };
1110
1111 #[cfg(feature = "native")]
1113 reader.apply_pin_policy(&crate::segment::pin::pin_policy().to_owned());
1114
1115 Ok(reader)
1116 }
1117
1118 #[cfg(feature = "native")]
1128 pub(crate) fn apply_pin_policy(&mut self, policy: &crate::segment::pin::PinPolicy) {
1129 use crate::segment::pin::PinReport;
1130
1131 if !policy.is_enabled() {
1132 return;
1133 }
1134 let mut remaining = policy.budget_bytes;
1135 let mut dense_report = PinReport::default();
1136 let mut sparse_report = PinReport::default();
1137
1138 for index in self.vector_indexes.values_mut() {
1140 index.pin_lookup_directory(policy.mode, &mut remaining, &mut dense_report);
1141 }
1142 for bmp in self.bmp_indexes.values_mut() {
1144 bmp.pin_block_starts(policy.mode, &mut remaining, &mut sparse_report);
1145 }
1146 for sparse in self.sparse_indexes.values_mut() {
1148 sparse.pin_skip_section(policy.mode, &mut remaining, &mut sparse_report);
1149 }
1150 for flat in self.flat_vectors.values_mut() {
1152 flat.pin_doc_ids(policy.mode, &mut remaining, &mut dense_report);
1153 }
1154 for bmp in self.bmp_indexes.values_mut() {
1155 bmp.pin_doc_maps(policy.mode, &mut remaining, &mut sparse_report);
1156 }
1157 for bmp in self.bmp_indexes.values_mut() {
1159 bmp.pin_sb_grid(policy.mode, &mut remaining, &mut sparse_report);
1160 }
1161
1162 let report = PinReport {
1163 intended_bytes: dense_report
1164 .intended_bytes
1165 .saturating_add(sparse_report.intended_bytes),
1166 pinned_bytes: dense_report
1167 .pinned_bytes
1168 .saturating_add(sparse_report.pinned_bytes),
1169 skipped_budget_bytes: dense_report
1170 .skipped_budget_bytes
1171 .saturating_add(sparse_report.skipped_budget_bytes),
1172 failed_bytes: dense_report
1173 .failed_bytes
1174 .saturating_add(sparse_report.failed_bytes),
1175 heap_copy_bytes: dense_report
1176 .heap_copy_bytes
1177 .saturating_add(sparse_report.heap_copy_bytes),
1178 };
1179 if report.skipped_budget_bytes > 0 || report.failed_bytes > 0 {
1180 log::warn!(
1181 "[pin] segment {:016x}: pinned {}/{} (budget skipped {}, mlock failed {}) — \
1182 raise HERMES_PIN_METADATA_BUDGET_MB or RLIMIT_MEMLOCK for full coverage",
1183 self.meta.id,
1184 crate::format_bytes(report.pinned_bytes),
1185 crate::format_bytes(report.intended_bytes),
1186 crate::format_bytes(report.skipped_budget_bytes),
1187 crate::format_bytes(report.failed_bytes),
1188 );
1189 } else if report.pinned_bytes > 0 {
1190 log::info!(
1191 "[pin] segment {:016x}: pinned {} of hot metadata ({:?})",
1192 self.meta.id,
1193 crate::format_bytes(report.pinned_bytes),
1194 policy.mode,
1195 );
1196 }
1197 self.dense_pin_report = dense_report;
1198 self.sparse_pin_report = sparse_report;
1199 }
1200
1201 pub fn meta(&self) -> &SegmentMeta {
1207 &self.meta
1208 }
1209
1210 pub fn num_docs(&self) -> u32 {
1211 self.meta.num_docs
1212 }
1213
1214 pub fn avg_field_len(&self, field: Field) -> f32 {
1216 self.meta.avg_field_len(field)
1217 }
1218
1219 pub fn schema(&self) -> &Schema {
1220 &self.schema
1221 }
1222
1223 pub fn sparse_indexes(&self) -> &FxHashMap<u32, SparseIndex> {
1225 &self.sparse_indexes
1226 }
1227
1228 pub fn sparse_index(&self, field: Field) -> Option<&SparseIndex> {
1230 self.sparse_indexes.get(&field.0)
1231 }
1232
1233 pub fn bmp_index(&self, field: Field) -> Option<&BmpIndex> {
1235 self.bmp_indexes.get(&field.0)
1236 }
1237
1238 pub fn bmp_indexes(&self) -> &FxHashMap<u32, BmpIndex> {
1240 &self.bmp_indexes
1241 }
1242
1243 pub fn vector_indexes(&self) -> &FxHashMap<u32, VectorIndex> {
1245 &self.vector_indexes
1246 }
1247
1248 pub fn flat_vectors(&self) -> &FxHashMap<u32, LazyFlatVectorData> {
1250 &self.flat_vectors
1251 }
1252
1253 pub fn fast_field(
1255 &self,
1256 field_id: u32,
1257 ) -> Option<&crate::structures::fast_field::FastFieldReader> {
1258 self.fast_fields.get(&field_id)
1259 }
1260
1261 pub fn fast_fields(&self) -> &FxHashMap<u32, crate::structures::fast_field::FastFieldReader> {
1263 &self.fast_fields
1264 }
1265
1266 pub fn term_dict_stats(&self) -> SSTableStats {
1268 self.term_dict.stats()
1269 }
1270
1271 pub fn memory_stats(&self) -> SegmentMemoryStats {
1273 let term_dict_stats = self.term_dict.stats();
1274
1275 let term_dict_cache_bytes = self.term_dict.cached_bytes();
1279 let store_cache_bytes = self.store.cached_bytes();
1280
1281 let sparse_heap_bytes: usize = self
1284 .sparse_indexes
1285 .values()
1286 .map(|s| s.estimated_heap_bytes())
1287 .sum::<usize>()
1288 + self
1289 .bmp_indexes
1290 .values()
1291 .map(|b| b.estimated_heap_bytes())
1292 .sum::<usize>();
1293
1294 let dense_heap_bytes: usize = self
1297 .vector_indexes
1298 .values()
1299 .map(|v| v.estimated_heap_bytes())
1300 .sum::<usize>()
1301 + self
1302 .flat_vectors
1303 .values()
1304 .map(LazyFlatVectorData::estimated_heap_bytes)
1305 .sum::<usize>();
1306
1307 #[cfg(feature = "native")]
1308 let (sparse_heap_bytes, dense_heap_bytes) = (
1309 sparse_heap_bytes.saturating_add(
1310 usize::try_from(self.sparse_pin_report.heap_copy_bytes).unwrap_or(usize::MAX),
1311 ),
1312 dense_heap_bytes.saturating_add(
1313 usize::try_from(self.dense_pin_report.heap_copy_bytes).unwrap_or(usize::MAX),
1314 ),
1315 );
1316
1317 #[cfg(feature = "native")]
1318 let (
1319 sparse_pinned_metadata_bytes,
1320 sparse_pin_intended_bytes,
1321 dense_pinned_metadata_bytes,
1322 dense_pin_intended_bytes,
1323 ) = (
1324 self.sparse_pin_report.pinned_bytes,
1325 self.sparse_pin_report.intended_bytes,
1326 self.dense_pin_report.pinned_bytes,
1327 self.dense_pin_report.intended_bytes,
1328 );
1329 #[cfg(not(feature = "native"))]
1330 let (
1331 sparse_pinned_metadata_bytes,
1332 sparse_pin_intended_bytes,
1333 dense_pinned_metadata_bytes,
1334 dense_pin_intended_bytes,
1335 ) = (0u64, 0u64, 0u64, 0u64);
1336
1337 let pinned_metadata_bytes =
1338 sparse_pinned_metadata_bytes.saturating_add(dense_pinned_metadata_bytes);
1339 let pin_intended_bytes = sparse_pin_intended_bytes.saturating_add(dense_pin_intended_bytes);
1340
1341 SegmentMemoryStats {
1342 segment_id: self.meta.id,
1343 num_docs: self.meta.num_docs,
1344 term_dict_cache_bytes,
1345 store_cache_bytes,
1346 sparse_heap_bytes,
1347 dense_heap_bytes,
1348 term_bloom_file_bytes: term_dict_stats.bloom_filter_size as u64,
1349 sparse_file_backed_bytes: self.sparse_file_backed_bytes,
1350 dense_file_backed_bytes: self.dense_file_backed_bytes,
1351 pinned_metadata_bytes,
1352 pin_intended_bytes,
1353 sparse_pinned_metadata_bytes,
1354 sparse_pin_intended_bytes,
1355 dense_pinned_metadata_bytes,
1356 dense_pin_intended_bytes,
1357 }
1358 }
1359
1360 pub async fn get_postings(
1365 &self,
1366 field: Field,
1367 term: &[u8],
1368 ) -> Result<Option<BlockPostingList>> {
1369 log::debug!(
1370 "SegmentReader::get_postings field={} term_len={}",
1371 field.0,
1372 term.len()
1373 );
1374
1375 let mut key = Vec::with_capacity(4 + term.len());
1377 key.extend_from_slice(&field.0.to_le_bytes());
1378 key.extend_from_slice(term);
1379
1380 let term_info = match self.term_dict.get(&key).await? {
1382 Some(info) => {
1383 log::debug!("SegmentReader::get_postings found term_info");
1384 info
1385 }
1386 None => {
1387 log::debug!("SegmentReader::get_postings term not found");
1388 return Ok(None);
1389 }
1390 };
1391
1392 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
1394 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
1396 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
1397 posting_list.push(doc_id, tf);
1398 }
1399 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
1400 return Ok(Some(block_list));
1401 }
1402
1403 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
1405 Error::Corruption("TermInfo has neither inline nor external data".to_string())
1406 })?;
1407
1408 let range = checked_file_range(
1409 posting_offset,
1410 posting_len,
1411 self.postings_handle.len(),
1412 "posting",
1413 )?;
1414 let posting_bytes = self.postings_handle.read_bytes_range(range).await?;
1415 let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
1416
1417 Ok(Some(block_list))
1418 }
1419
1420 pub async fn get_prefix_postings(
1422 &self,
1423 field: Field,
1424 prefix: &[u8],
1425 ) -> Result<Vec<BlockPostingList>> {
1426 if prefix.is_empty() {
1427 return Err(Error::Query("prefix must not be empty".into()));
1428 }
1429 let mut key_prefix = Vec::with_capacity(4 + prefix.len());
1431 key_prefix.extend_from_slice(&field.0.to_le_bytes());
1432 key_prefix.extend_from_slice(prefix);
1433
1434 let (entries, truncated) = self
1435 .term_dict
1436 .prefix_scan_limited(&key_prefix, MAX_PREFIX_TERMS)
1437 .await?;
1438 if truncated {
1439 return Err(Error::Query(format!(
1440 "prefix expands to more than {MAX_PREFIX_TERMS} terms"
1441 )));
1442 }
1443 let posting_count: u64 = entries
1444 .iter()
1445 .map(|(_, term_info)| term_info.doc_freq() as u64)
1446 .sum();
1447 if posting_count > MAX_PREFIX_POSTINGS {
1448 return Err(Error::Query(format!(
1449 "prefix expands to {posting_count} postings (maximum {MAX_PREFIX_POSTINGS})"
1450 )));
1451 }
1452 let mut results = Vec::with_capacity(entries.len());
1453
1454 for (_key, term_info) in entries {
1455 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
1456 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
1457 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
1458 posting_list.push(doc_id, tf);
1459 }
1460 results.push(BlockPostingList::from_posting_list(&posting_list)?);
1461 } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
1462 let range = checked_file_range(
1463 posting_offset,
1464 posting_len,
1465 self.postings_handle.len(),
1466 "prefix posting",
1467 )?;
1468 let posting_bytes = self.postings_handle.read_bytes_range(range).await?;
1469 results.push(BlockPostingList::deserialize_zero_copy(posting_bytes)?);
1470 }
1471 }
1472
1473 Ok(results)
1474 }
1475
1476 pub async fn doc(&self, local_doc_id: DocId) -> Result<Option<Document>> {
1481 self.doc_with_fields(local_doc_id, None).await
1482 }
1483
1484 pub async fn doc_with_fields(
1490 &self,
1491 local_doc_id: DocId,
1492 fields: Option<&rustc_hash::FxHashSet<u32>>,
1493 ) -> Result<Option<Document>> {
1494 let mut doc = match fields {
1495 Some(set) => {
1496 let field_ids: Vec<u32> = set.iter().copied().collect();
1497 match self
1498 .store
1499 .get_fields(local_doc_id, &self.schema, &field_ids)
1500 .await
1501 {
1502 Ok(Some(d)) => d,
1503 Ok(None) => return Ok(None),
1504 Err(e) => return Err(Error::from(e)),
1505 }
1506 }
1507 None => match self.store.get(local_doc_id, &self.schema).await {
1508 Ok(Some(d)) => d,
1509 Ok(None) => return Ok(None),
1510 Err(e) => return Err(Error::from(e)),
1511 },
1512 };
1513
1514 for (&field_id, lazy_flat) in &self.flat_vectors {
1516 if let Some(set) = fields
1518 && !set.contains(&field_id)
1519 {
1520 continue;
1521 }
1522
1523 let is_binary = lazy_flat.quantization == DenseVectorQuantization::Binary;
1524 let (start, entries) = lazy_flat.flat_indexes_for_doc(local_doc_id);
1525 for (j, &(_doc_id, _ordinal)) in entries.iter().enumerate() {
1526 let flat_idx = start + j;
1527 if is_binary {
1528 let vbs = lazy_flat.vector_byte_size();
1529 let mut raw = vec![0u8; vbs];
1530 match lazy_flat.read_vector_raw_into(flat_idx, &mut raw).await {
1531 Ok(()) => {
1532 doc.add_binary_dense_vector(Field(field_id), raw);
1533 }
1534 Err(e) => {
1535 log::warn!(
1536 "Failed to hydrate binary dense vector field {}: {}",
1537 field_id,
1538 e
1539 );
1540 }
1541 }
1542 } else {
1543 match lazy_flat.get_vector(flat_idx).await {
1544 Ok(vec) => {
1545 doc.add_dense_vector(Field(field_id), vec);
1546 }
1547 Err(e) => {
1548 log::warn!("Failed to hydrate dense vector field {}: {}", field_id, e);
1549 }
1550 }
1551 }
1552 }
1553 }
1554
1555 Ok(Some(doc))
1556 }
1557
1558 pub async fn prefetch_terms(
1560 &self,
1561 field: Field,
1562 start_term: &[u8],
1563 end_term: &[u8],
1564 ) -> Result<()> {
1565 let mut start_key = Vec::with_capacity(4 + start_term.len());
1566 start_key.extend_from_slice(&field.0.to_le_bytes());
1567 start_key.extend_from_slice(start_term);
1568
1569 let mut end_key = Vec::with_capacity(4 + end_term.len());
1570 end_key.extend_from_slice(&field.0.to_le_bytes());
1571 end_key.extend_from_slice(end_term);
1572
1573 self.term_dict.prefetch_range(&start_key, &end_key).await?;
1574 Ok(())
1575 }
1576
1577 pub fn store_has_dict(&self) -> bool {
1579 self.store.has_dict()
1580 }
1581
1582 pub fn store(&self) -> &super::store::AsyncStoreReader {
1584 &self.store
1585 }
1586
1587 pub fn store_raw_blocks(&self) -> Vec<RawStoreBlock> {
1589 self.store.raw_blocks()
1590 }
1591
1592 pub fn store_data_slice(&self) -> &FileHandle {
1594 self.store.data_slice()
1595 }
1596
1597 pub async fn all_terms(&self) -> Result<Vec<(Vec<u8>, TermInfo)>> {
1599 self.term_dict.all_entries().await.map_err(Error::from)
1600 }
1601
1602 pub async fn all_terms_with_stats(&self) -> Result<Vec<(Field, String, u32)>> {
1607 let entries = self.term_dict.all_entries().await?;
1608 let mut result = Vec::with_capacity(entries.len());
1609
1610 for (key, term_info) in entries {
1611 if key.len() > 4 {
1613 let field_id = u32::from_le_bytes([key[0], key[1], key[2], key[3]]);
1614 let term_bytes = &key[4..];
1615 if let Ok(term_str) = std::str::from_utf8(term_bytes) {
1616 result.push((Field(field_id), term_str.to_string(), term_info.doc_freq()));
1617 }
1618 }
1619 }
1620
1621 Ok(result)
1622 }
1623
1624 pub fn term_dict_iter(&self) -> crate::structures::AsyncSSTableIterator<'_, TermInfo> {
1626 self.term_dict.iter()
1627 }
1628
1629 pub async fn prefetch_term_dict(&self) -> crate::Result<()> {
1633 self.term_dict
1634 .prefetch_all_data_bulk()
1635 .await
1636 .map_err(crate::Error::from)
1637 }
1638
1639 pub async fn read_postings(&self, offset: u64, len: u64) -> Result<Vec<u8>> {
1641 let range = checked_file_range(offset, len, self.postings_handle.len(), "posting")?;
1642 let bytes = self.postings_handle.read_bytes_range(range).await?;
1643 Ok(bytes.to_vec())
1644 }
1645
1646 pub async fn read_position_bytes(&self, offset: u64, len: u64) -> Result<Option<Vec<u8>>> {
1648 let handle = match &self.positions_handle {
1649 Some(h) => h,
1650 None => return Ok(None),
1651 };
1652 let range = checked_file_range(offset, len, handle.len(), "position")?;
1653 let bytes = handle.read_bytes_range(range).await?;
1654 Ok(Some(bytes.to_vec()))
1655 }
1656
1657 pub fn has_positions_file(&self) -> bool {
1659 self.positions_handle.is_some()
1660 }
1661
1662 fn validate_dense_search_request(
1666 &self,
1667 field: Field,
1668 query: &[f32],
1669 nprobe: usize,
1670 rerank_factor: f32,
1671 combiner: crate::query::MultiValueCombiner,
1672 ) -> Result<DenseSearchParams> {
1673 let entry = self
1674 .schema
1675 .get_field_entry(field)
1676 .ok_or_else(|| Error::FieldNotFound(field.0.to_string()))?;
1677 if entry.field_type != crate::dsl::FieldType::DenseVector {
1678 return Err(Error::InvalidFieldType {
1679 expected: "dense_vector".to_string(),
1680 got: format!("{:?}", entry.field_type),
1681 });
1682 }
1683 let config = entry.dense_vector_config.as_ref().ok_or_else(|| {
1684 Error::Schema(format!(
1685 "dense vector field '{}' has no dense vector configuration",
1686 entry.name
1687 ))
1688 })?;
1689
1690 if query.is_empty() {
1691 return Err(Error::Query(format!(
1692 "dense query vector for field '{}' must not be empty",
1693 entry.name
1694 )));
1695 }
1696 if query.len() != config.dim {
1697 return Err(Error::Query(format!(
1698 "dense query vector dimension {} does not match field '{}' dimension {}",
1699 query.len(),
1700 entry.name,
1701 config.dim
1702 )));
1703 }
1704 if let Some((index, value)) = query
1705 .iter()
1706 .enumerate()
1707 .find(|(_, value)| !value.is_finite())
1708 {
1709 return Err(Error::Query(format!(
1710 "dense query vector for field '{}' contains non-finite value {value} at index {index}",
1711 entry.name
1712 )));
1713 }
1714
1715 let nprobe = match (nprobe, config.nprobe) {
1718 (0, 0) => 32,
1719 (0, schema_nprobe) => schema_nprobe,
1720 (query_nprobe, _) => query_nprobe,
1721 };
1722 if nprobe > MAX_DENSE_NPROBE {
1723 return Err(Error::Query(format!(
1724 "dense nprobe must be at most {MAX_DENSE_NPROBE}, got {nprobe}"
1725 )));
1726 }
1727
1728 checked_dense_fetch_k(0, rerank_factor)?;
1731 combiner.validate().map_err(Error::Query)?;
1732
1733 Ok(DenseSearchParams {
1734 dim: config.dim,
1735 nprobe,
1736 unit_norm: config.unit_norm,
1737 })
1738 }
1739
1740 fn validate_binary_search_request(&self, field: Field, query: &[u8]) -> Result<usize> {
1741 let entry = self
1742 .schema
1743 .get_field_entry(field)
1744 .ok_or_else(|| Error::FieldNotFound(field.0.to_string()))?;
1745 if entry.field_type != crate::dsl::FieldType::BinaryDenseVector {
1746 return Err(Error::InvalidFieldType {
1747 expected: "binary_dense_vector".to_string(),
1748 got: format!("{:?}", entry.field_type),
1749 });
1750 }
1751 let config = entry.binary_dense_vector_config.as_ref().ok_or_else(|| {
1752 Error::Schema(format!(
1753 "binary dense vector field '{}' has no configuration",
1754 entry.name
1755 ))
1756 })?;
1757 if config.dim == 0 || !config.dim.is_multiple_of(8) {
1758 return Err(Error::Schema(format!(
1759 "binary dense vector field '{}' has invalid dimension {}",
1760 entry.name, config.dim
1761 )));
1762 }
1763 if query.len() != config.byte_len() {
1764 return Err(Error::Query(format!(
1765 "binary query byte length {} does not match field '{}' byte length {}",
1766 query.len(),
1767 entry.name,
1768 config.byte_len()
1769 )));
1770 }
1771 Ok(config.dim)
1772 }
1773
1774 fn score_quantized_batch(
1780 query: &[f32],
1781 raw: &[u8],
1782 quant: crate::dsl::DenseVectorQuantization,
1783 dim: usize,
1784 scores: &mut [f32],
1785 unit_norm: bool,
1786 ) -> Result<()> {
1787 use crate::dsl::DenseVectorQuantization;
1788 use crate::structures::simd;
1789
1790 if query.len() != dim {
1791 return Err(Error::Query(format!(
1792 "dense SIMD query dimension {} does not match vector dimension {dim}",
1793 query.len()
1794 )));
1795 }
1796 let element_size = match quant {
1797 DenseVectorQuantization::F32 => std::mem::size_of::<f32>(),
1798 DenseVectorQuantization::F16 => std::mem::size_of::<u16>(),
1799 DenseVectorQuantization::UInt8 => 1,
1800 DenseVectorQuantization::Binary => {
1801 return Err(Error::InvalidFieldType {
1802 expected: "non-binary dense vector".to_string(),
1803 got: "binary dense vector".to_string(),
1804 });
1805 }
1806 };
1807 let required_bytes = scores
1808 .len()
1809 .checked_mul(dim)
1810 .and_then(|elements| elements.checked_mul(element_size))
1811 .ok_or_else(|| Error::Corruption("dense vector batch byte length overflow".into()))?;
1812 if raw.len() < required_bytes {
1813 return Err(Error::Corruption(format!(
1814 "dense vector batch is truncated: need {required_bytes} bytes, got {}",
1815 raw.len()
1816 )));
1817 }
1818 if quant == DenseVectorQuantization::F16
1819 && required_bytes > 0
1820 && !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>())
1821 {
1822 return Err(Error::Corruption(
1823 "f16 vector data is not 2-byte aligned".to_string(),
1824 ));
1825 }
1826
1827 match (quant, unit_norm) {
1828 (DenseVectorQuantization::F32, false) => {
1829 let num_floats = scores.len() * dim;
1830 if !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()) {
1831 return Err(Error::Corruption(
1832 "f32 vector data is not 4-byte aligned".to_string(),
1833 ));
1834 }
1835 let vectors: &[f32] =
1836 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
1837 simd::batch_cosine_scores(query, vectors, dim, scores);
1838 }
1839 (DenseVectorQuantization::F32, true) => {
1840 let num_floats = scores.len() * dim;
1841 if !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()) {
1842 return Err(Error::Corruption(
1843 "f32 vector data is not 4-byte aligned".to_string(),
1844 ));
1845 }
1846 let vectors: &[f32] =
1847 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
1848 simd::batch_dot_scores(query, vectors, dim, scores);
1849 }
1850 (DenseVectorQuantization::F16, false) => {
1851 simd::batch_cosine_scores_f16(query, raw, dim, scores);
1852 }
1853 (DenseVectorQuantization::F16, true) => {
1854 simd::batch_dot_scores_f16(query, raw, dim, scores);
1855 }
1856 (DenseVectorQuantization::UInt8, false) => {
1857 simd::batch_cosine_scores_u8(query, raw, dim, scores);
1858 }
1859 (DenseVectorQuantization::UInt8, true) => {
1860 simd::batch_dot_scores_u8(query, raw, dim, scores);
1861 }
1862 (DenseVectorQuantization::Binary, _) => unreachable!("validated above"),
1863 }
1864 Ok(())
1865 }
1866
1867 pub async fn search_dense_vector(
1873 &self,
1874 field: Field,
1875 query: &[f32],
1876 k: usize,
1877 nprobe: usize,
1878 rerank_factor: f32,
1879 combiner: crate::query::MultiValueCombiner,
1880 ) -> Result<Vec<VectorSearchResult>> {
1881 self.search_dense_vector_impl(field, query, k, nprobe, rerank_factor, combiner, None)
1882 .await
1883 }
1884
1885 #[allow(clippy::too_many_arguments)]
1886 pub(crate) async fn search_dense_vector_with_probe_cache(
1887 &self,
1888 field: Field,
1889 query: &[f32],
1890 k: usize,
1891 nprobe: usize,
1892 rerank_factor: f32,
1893 combiner: crate::query::MultiValueCombiner,
1894 probe_cache: &std::sync::Mutex<Option<std::sync::Arc<crate::structures::IvfPqQueryPlan>>>,
1895 ) -> Result<Vec<VectorSearchResult>> {
1896 self.search_dense_vector_impl(
1897 field,
1898 query,
1899 k,
1900 nprobe,
1901 rerank_factor,
1902 combiner,
1903 Some(probe_cache),
1904 )
1905 .await
1906 }
1907
1908 #[allow(clippy::too_many_arguments)]
1909 async fn search_dense_vector_impl(
1910 &self,
1911 field: Field,
1912 query: &[f32],
1913 k: usize,
1914 nprobe: usize,
1915 rerank_factor: f32,
1916 combiner: crate::query::MultiValueCombiner,
1917 probe_cache: Option<
1918 &std::sync::Mutex<Option<std::sync::Arc<crate::structures::IvfPqQueryPlan>>>,
1919 >,
1920 ) -> Result<Vec<VectorSearchResult>> {
1921 let params =
1922 self.validate_dense_search_request(field, query, nprobe, rerank_factor, combiner)?;
1923 let fetch_k = checked_dense_fetch_k(k, rerank_factor)?;
1924 if k == 0 {
1925 return Ok(Vec::new());
1926 }
1927
1928 let ann_index = self.vector_indexes.get(&field.0);
1929 let lazy_flat = self.flat_vectors.get(&field.0);
1930 if ann_index.is_none() && lazy_flat.is_none() {
1932 return Ok(Vec::new());
1933 }
1934
1935 if ann_index.is_some() && lazy_flat.is_none() {
1936 return Err(Error::Corruption(format!(
1937 "dense ANN field {} is missing flat vector storage",
1938 field.0
1939 )));
1940 }
1941
1942 if let Some(flat) = lazy_flat
1943 && flat.dim != params.dim
1944 {
1945 return Err(Error::Corruption(format!(
1946 "dense vector field {} has schema dimension {} but flat storage dimension {}",
1947 field.0, params.dim, flat.dim
1948 )));
1949 }
1950
1951 let t0 = std::time::Instant::now();
1953 let mut flat_results = None;
1954 let results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
1955 match index {
1957 VectorIndex::IvfPq(lazy) => {
1958 let index = lazy.get();
1959 let codebook =
1960 self.trained_vectors
1961 .codebooks
1962 .get(&field.0)
1963 .ok_or_else(|| {
1964 Error::Schema(format!(
1965 "IVF-PQ index requires a global codebook for field {}",
1966 field.0
1967 ))
1968 })?;
1969 let centroids =
1970 self.trained_vectors
1971 .centroids
1972 .get(&field.0)
1973 .ok_or_else(|| {
1974 Error::Schema(format!(
1975 "IVF-PQ index requires coarse centroids for field {}",
1976 field.0
1977 ))
1978 })?;
1979 validate_coarse_centroids(centroids, params.dim)?;
1980 let routing = self
1981 .schema
1982 .get_field_entry(field)
1983 .and_then(|entry| entry.dense_vector_config.as_ref())
1984 .map_or(crate::dsl::IvfRoutingMode::Auto, |config| {
1985 config.ivf_routing
1986 });
1987 validate_ivf_pq_ann(index, centroids, codebook, params.dim, routing)?;
1988 let query_plan = float_query_plan(
1989 centroids,
1990 codebook,
1991 query,
1992 params.nprobe,
1993 routing,
1994 probe_cache,
1995 )?;
1996 let flat = lazy_flat.expect("ANN/flat pairing validated above");
1997 index
1998 .search_ivf_pq_distinct(
1999 fetch_k.min(flat.num_docs_with_vectors()),
2000 &query_plan,
2001 )
2002 .map_err(|error| {
2003 Error::Corruption(format!(
2004 "invalid IVF-PQ payload for field {}: {error}",
2005 field.0,
2006 ))
2007 })?
2008 .into_iter()
2009 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
2010 .collect()
2011 }
2012 VectorIndex::BinaryIvf(_) => {
2013 Vec::new()
2015 }
2016 }
2017 } else if let Some(lazy_flat) = lazy_flat {
2018 log::debug!(
2022 "[dense_vector_search] field {}: brute-force on {} vectors (dim={}, quant={:?})",
2023 field.0,
2024 lazy_flat.num_vectors,
2025 lazy_flat.dim,
2026 lazy_flat.quantization
2027 );
2028 let dim = lazy_flat.dim;
2029 let n = lazy_flat.num_vectors;
2030 let quant = lazy_flat.quantization;
2031 let batch_len =
2032 bounded_vector_score_batch(lazy_flat.vector_byte_size(), DENSE_SCORE_BATCH);
2033 let mut collector = FlatDocumentCollector::new(fetch_k.min(n), combiner);
2034 let mut scores = vec![0f32; batch_len];
2035
2036 for batch_start in (0..n).step_by(batch_len) {
2037 let batch_count = batch_len.min(n - batch_start);
2038 let batch_bytes = lazy_flat
2039 .read_vectors_batch(batch_start, batch_count)
2040 .await
2041 .map_err(crate::Error::Io)?;
2042 let raw = batch_bytes.as_slice();
2043
2044 Self::score_quantized_batch(
2045 query,
2046 raw,
2047 quant,
2048 dim,
2049 &mut scores[..batch_count],
2050 params.unit_norm,
2051 )?;
2052
2053 for (i, &score) in scores.iter().enumerate().take(batch_count) {
2054 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
2055 collector.push(doc_id, ordinal, score);
2056 }
2057 }
2058
2059 flat_results = Some(collector.into_results());
2060 Vec::new()
2061 } else {
2062 return Ok(Vec::new());
2063 };
2064 let l1_elapsed = t0.elapsed();
2065 {
2066 let kind = match ann_index {
2067 Some(VectorIndex::IvfPq(_)) => "ivf_pq",
2068 Some(VectorIndex::BinaryIvf(_)) => "binary_ivf",
2069 None => "flat",
2070 };
2071 crate::observe::dense_l1(
2072 self.schema.index_label(),
2073 self.schema.get_field_name(field).unwrap_or("?"),
2074 kind,
2075 l1_elapsed.as_secs_f64(),
2076 flat_results.as_ref().map_or(results.len(), Vec::len),
2077 );
2078 }
2079 log::debug!(
2080 "[dense_vector_search] field {}: L1 returned {} candidates in {:.1}ms",
2081 field.0,
2082 flat_results.as_ref().map_or(results.len(), Vec::len),
2083 l1_elapsed.as_secs_f64() * 1000.0
2084 );
2085
2086 if let Some(results) = flat_results {
2087 return Ok(results);
2088 }
2089
2090 if ann_index.is_some()
2093 && !results.is_empty()
2094 && let Some(lazy_flat) = lazy_flat
2095 {
2096 let t_rerank = std::time::Instant::now();
2097 let vbs = lazy_flat.vector_byte_size();
2098 let (reranked, stats) = exact_score_dense_candidate_documents(
2099 &results,
2100 lazy_flat,
2101 query,
2102 params.unit_norm,
2103 combiner,
2104 k,
2105 )
2106 .await?;
2107
2108 crate::observe::dense_rerank(
2109 self.schema.index_label(),
2110 self.schema.get_field_name(field).unwrap_or("?"),
2111 t_rerank.elapsed().as_secs_f64(),
2112 stats.resolve_elapsed.as_secs_f64(),
2113 stats.read_elapsed.as_secs_f64(),
2114 stats.vector_count,
2115 );
2116 log::debug!(
2117 "[dense_vector_search] field {}: rerank {} vectors (dim={}, quant={:?}, bytes_per_vector={}): resolve={:.1}ms read={:.1}ms score={:.1}ms",
2118 field.0,
2119 stats.vector_count,
2120 lazy_flat.dim,
2121 lazy_flat.quantization,
2122 vbs,
2123 stats.resolve_elapsed.as_secs_f64() * 1000.0,
2124 stats.read_elapsed.as_secs_f64() * 1000.0,
2125 stats.score_elapsed.as_secs_f64() * 1000.0,
2126 );
2127
2128 log::debug!(
2129 "[dense_vector_search] field {}: rerank total={:.1}ms",
2130 field.0,
2131 t_rerank.elapsed().as_secs_f64() * 1000.0
2132 );
2133 return Ok(reranked);
2134 }
2135
2136 Ok(combine_grouped_ordinal_results(results, combiner, k))
2137 }
2138
2139 async fn search_binary_dense_vector_impl(
2144 &self,
2145 field: Field,
2146 query: &[u8],
2147 k: usize,
2148 combiner: crate::query::MultiValueCombiner,
2149 probe_cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
2150 ) -> Result<Vec<VectorSearchResult>> {
2151 let schema_dim = self.validate_binary_search_request(field, query)?;
2152 combiner.validate().map_err(Error::Query)?;
2153 if k == 0 {
2154 return Ok(Vec::new());
2155 }
2156 let t0 = crate::observe::Timer::start();
2157 if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0) {
2158 let ivf = lazy.get();
2159 let config = self
2160 .schema
2161 .get_field_entry(field)
2162 .and_then(|entry| entry.binary_dense_vector_config.as_ref())
2163 .ok_or_else(|| {
2164 Error::Schema(format!(
2165 "binary IVF field {} has no schema configuration",
2166 field.0
2167 ))
2168 })?;
2169 let quantizer = self
2170 .trained_vectors
2171 .binary_quantizers
2172 .get(&field.0)
2173 .ok_or_else(|| {
2174 Error::Schema(format!(
2175 "global binary IVF field {} has no loaded quantizer",
2176 field.0
2177 ))
2178 })?;
2179 validate_binary_ann(ivf, quantizer, config, schema_dim, field)?;
2180 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
2181 Error::Corruption(format!(
2182 "global binary IVF field {} is missing flat vector storage",
2183 field.0
2184 ))
2185 })?;
2186 let clusters = binary_probe_clusters(
2187 quantizer,
2188 query,
2189 config.nprobe,
2190 config.ivf_routing,
2191 probe_cache,
2192 )?;
2193 let single_valued = flat.num_vectors == flat.num_docs_with_vectors();
2194 let candidate_docs = k.min(flat.num_docs_with_vectors());
2195 let ann_results = if single_valued {
2196 ivf.search_binary_clusters::<false>(query, candidate_docs, &clusters)
2197 } else {
2198 ivf.search_binary_clusters::<true>(query, candidate_docs, &clusters)
2199 }
2200 .map_err(|error| {
2201 Error::Corruption(format!(
2202 "invalid binary IVF payload for field {}: {error}",
2203 field.0,
2204 ))
2205 })?;
2206 let results = exact_score_binary_candidate_documents(
2207 &ann_results,
2208 flat,
2209 query,
2210 schema_dim,
2211 combiner,
2212 k,
2213 )
2214 .await?;
2215 crate::observe::dense_l1(
2216 self.schema.index_label(),
2217 self.schema.get_field_name(field).unwrap_or("?"),
2218 "global_binary_ivf",
2219 t0.secs(),
2220 results.len(),
2221 );
2222 return Ok(results);
2223 }
2224 let lazy_flat = match self.flat_vectors.get(&field.0) {
2225 Some(f) => f,
2226 None => return Ok(Vec::new()),
2227 };
2228
2229 let dim_bits = lazy_flat.dim;
2230 let byte_len = lazy_flat.vector_byte_size();
2231 let n = lazy_flat.num_vectors;
2232
2233 if dim_bits != schema_dim {
2234 return Err(Error::Corruption(format!(
2235 "binary vector field {} has schema dimension {} but flat storage dimension {}",
2236 field.0, schema_dim, dim_bits
2237 )));
2238 }
2239
2240 if byte_len != query.len() {
2241 return Err(Error::Schema(format!(
2242 "Binary query vector byte length {} != field byte length {}",
2243 query.len(),
2244 byte_len
2245 )));
2246 }
2247
2248 let batch_len = bounded_vector_score_batch(byte_len, BINARY_SCORE_BATCH);
2249 let mut collector = FlatDocumentCollector::new(k, combiner);
2250 let mut scores = vec![0f32; batch_len];
2251
2252 for batch_start in (0..n).step_by(batch_len) {
2253 let batch_count = batch_len.min(n - batch_start);
2254 let batch_bytes = lazy_flat
2255 .read_vectors_batch(batch_start, batch_count)
2256 .await
2257 .map_err(crate::Error::Io)?;
2258 let raw = batch_bytes.as_slice();
2259
2260 crate::structures::simd::batch_hamming_scores(
2261 query,
2262 raw,
2263 byte_len,
2264 dim_bits,
2265 &mut scores[..batch_count],
2266 );
2267
2268 for (i, &score) in scores.iter().enumerate().take(batch_count) {
2269 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
2270 collector.push(doc_id, ordinal, score);
2271 }
2272 }
2273
2274 let results = collector.into_results();
2275
2276 crate::observe::dense_l1(
2277 self.schema.index_label(),
2278 self.schema.get_field_name(field).unwrap_or("?"),
2279 "binary_flat",
2280 t0.secs(),
2281 results.len(),
2282 );
2283 Ok(results)
2284 }
2285
2286 pub async fn search_binary_dense_vector(
2287 &self,
2288 field: Field,
2289 query: &[u8],
2290 k: usize,
2291 combiner: crate::query::MultiValueCombiner,
2292 ) -> Result<Vec<VectorSearchResult>> {
2293 self.search_binary_dense_vector_impl(field, query, k, combiner, None)
2294 .await
2295 }
2296
2297 pub(crate) async fn search_binary_dense_vector_with_probe_cache(
2298 &self,
2299 field: Field,
2300 query: &[u8],
2301 k: usize,
2302 combiner: crate::query::MultiValueCombiner,
2303 probe_cache: &std::sync::Mutex<Option<crate::structures::IvfProbePlan>>,
2304 ) -> Result<Vec<VectorSearchResult>> {
2305 self.search_binary_dense_vector_impl(field, query, k, combiner, Some(probe_cache))
2306 .await
2307 }
2308
2309 pub fn coarse_centroids(&self, field_id: u32) -> Option<&Arc<CoarseCentroids>> {
2311 self.trained_vectors.centroids.get(&field_id)
2312 }
2313
2314 pub fn set_trained_vectors(
2315 &mut self,
2316 trained_vectors: Arc<crate::segment::TrainedVectorStructures>,
2317 ) {
2318 self.trained_vectors = trained_vectors;
2319 }
2320
2321 pub fn get_vector_index(&self, field: Field) -> Option<&VectorIndex> {
2323 self.vector_indexes.get(&field.0)
2324 }
2325
2326 pub async fn get_positions(
2331 &self,
2332 field: Field,
2333 term: &[u8],
2334 ) -> Result<Option<crate::structures::PositionPostingList>> {
2335 let handle = match &self.positions_handle {
2337 Some(h) => h,
2338 None => return Ok(None),
2339 };
2340
2341 let mut key = Vec::with_capacity(4 + term.len());
2343 key.extend_from_slice(&field.0.to_le_bytes());
2344 key.extend_from_slice(term);
2345
2346 let term_info = match self.term_dict.get(&key).await? {
2348 Some(info) => info,
2349 None => return Ok(None),
2350 };
2351
2352 let (offset, length) = match term_info.position_info() {
2354 Some((o, l)) => (o, l),
2355 None => return Ok(None),
2356 };
2357
2358 let range = checked_file_range(offset, length, handle.len(), "position list")?;
2362 let slice = handle.slice(range);
2363 let data = slice.read_bytes().await?;
2364
2365 let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
2367
2368 Ok(Some(pos_list))
2369 }
2370
2371 pub fn has_positions(&self, field: Field) -> bool {
2373 if let Some(entry) = self.schema.get_field_entry(field) {
2375 entry.positions.is_some()
2376 } else {
2377 false
2378 }
2379 }
2380}
2381
2382#[cfg(feature = "sync")]
2384impl SegmentReader {
2385 pub fn get_postings_sync(&self, field: Field, term: &[u8]) -> Result<Option<BlockPostingList>> {
2387 let mut key = Vec::with_capacity(4 + term.len());
2389 key.extend_from_slice(&field.0.to_le_bytes());
2390 key.extend_from_slice(term);
2391
2392 let term_info = match self.term_dict.get_sync(&key)? {
2394 Some(info) => info,
2395 None => return Ok(None),
2396 };
2397
2398 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
2400 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
2401 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
2402 posting_list.push(doc_id, tf);
2403 }
2404 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
2405 return Ok(Some(block_list));
2406 }
2407
2408 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
2410 Error::Corruption("TermInfo has neither inline nor external data".to_string())
2411 })?;
2412
2413 let range = checked_file_range(
2414 posting_offset,
2415 posting_len,
2416 self.postings_handle.len(),
2417 "posting",
2418 )?;
2419 let posting_bytes = self.postings_handle.read_bytes_range_sync(range)?;
2420 let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
2421
2422 Ok(Some(block_list))
2423 }
2424
2425 pub fn get_prefix_postings_sync(
2427 &self,
2428 field: Field,
2429 prefix: &[u8],
2430 ) -> Result<Vec<BlockPostingList>> {
2431 if prefix.is_empty() {
2432 return Err(Error::Query("prefix must not be empty".into()));
2433 }
2434 let mut key_prefix = Vec::with_capacity(4 + prefix.len());
2435 key_prefix.extend_from_slice(&field.0.to_le_bytes());
2436 key_prefix.extend_from_slice(prefix);
2437
2438 let (entries, truncated) = self
2439 .term_dict
2440 .prefix_scan_limited_sync(&key_prefix, MAX_PREFIX_TERMS)?;
2441 if truncated {
2442 return Err(Error::Query(format!(
2443 "prefix expands to more than {MAX_PREFIX_TERMS} terms"
2444 )));
2445 }
2446 let posting_count: u64 = entries
2447 .iter()
2448 .map(|(_, term_info)| term_info.doc_freq() as u64)
2449 .sum();
2450 if posting_count > MAX_PREFIX_POSTINGS {
2451 return Err(Error::Query(format!(
2452 "prefix expands to {posting_count} postings (maximum {MAX_PREFIX_POSTINGS})"
2453 )));
2454 }
2455 let mut results = Vec::with_capacity(entries.len());
2456
2457 for (_key, term_info) in entries {
2458 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
2459 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
2460 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
2461 posting_list.push(doc_id, tf);
2462 }
2463 results.push(BlockPostingList::from_posting_list(&posting_list)?);
2464 } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
2465 let range = checked_file_range(
2466 posting_offset,
2467 posting_len,
2468 self.postings_handle.len(),
2469 "prefix posting",
2470 )?;
2471 let posting_bytes = self.postings_handle.read_bytes_range_sync(range)?;
2472 results.push(BlockPostingList::deserialize_zero_copy(posting_bytes)?);
2473 }
2474 }
2475
2476 Ok(results)
2477 }
2478
2479 pub fn get_positions_sync(
2481 &self,
2482 field: Field,
2483 term: &[u8],
2484 ) -> Result<Option<crate::structures::PositionPostingList>> {
2485 let handle = match &self.positions_handle {
2486 Some(h) => h,
2487 None => return Ok(None),
2488 };
2489
2490 let mut key = Vec::with_capacity(4 + term.len());
2492 key.extend_from_slice(&field.0.to_le_bytes());
2493 key.extend_from_slice(term);
2494
2495 let term_info = match self.term_dict.get_sync(&key)? {
2497 Some(info) => info,
2498 None => return Ok(None),
2499 };
2500
2501 let (offset, length) = match term_info.position_info() {
2502 Some((o, l)) => (o, l),
2503 None => return Ok(None),
2504 };
2505
2506 let range = checked_file_range(offset, length, handle.len(), "position list")?;
2507 let slice = handle.slice(range);
2508 let data = slice.read_bytes_sync()?;
2509
2510 let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
2511 Ok(Some(pos_list))
2512 }
2513
2514 pub fn search_dense_vector_sync(
2517 &self,
2518 field: Field,
2519 query: &[f32],
2520 k: usize,
2521 nprobe: usize,
2522 rerank_factor: f32,
2523 combiner: crate::query::MultiValueCombiner,
2524 ) -> Result<Vec<VectorSearchResult>> {
2525 self.search_dense_vector_sync_impl(field, query, k, nprobe, rerank_factor, combiner, None)
2526 }
2527
2528 #[cfg(feature = "sync")]
2529 #[allow(clippy::too_many_arguments)]
2530 pub(crate) fn search_dense_vector_sync_with_probe_cache(
2531 &self,
2532 field: Field,
2533 query: &[f32],
2534 k: usize,
2535 nprobe: usize,
2536 rerank_factor: f32,
2537 combiner: crate::query::MultiValueCombiner,
2538 probe_cache: &std::sync::Mutex<Option<std::sync::Arc<crate::structures::IvfPqQueryPlan>>>,
2539 ) -> Result<Vec<VectorSearchResult>> {
2540 self.search_dense_vector_sync_impl(
2541 field,
2542 query,
2543 k,
2544 nprobe,
2545 rerank_factor,
2546 combiner,
2547 Some(probe_cache),
2548 )
2549 }
2550
2551 #[cfg(feature = "sync")]
2552 #[allow(clippy::too_many_arguments)]
2553 fn search_dense_vector_sync_impl(
2554 &self,
2555 field: Field,
2556 query: &[f32],
2557 k: usize,
2558 nprobe: usize,
2559 rerank_factor: f32,
2560 combiner: crate::query::MultiValueCombiner,
2561 probe_cache: Option<
2562 &std::sync::Mutex<Option<std::sync::Arc<crate::structures::IvfPqQueryPlan>>>,
2563 >,
2564 ) -> Result<Vec<VectorSearchResult>> {
2565 let params =
2566 self.validate_dense_search_request(field, query, nprobe, rerank_factor, combiner)?;
2567 let fetch_k = checked_dense_fetch_k(k, rerank_factor)?;
2568 if k == 0 {
2569 return Ok(Vec::new());
2570 }
2571
2572 let ann_index = self.vector_indexes.get(&field.0);
2573 let lazy_flat = self.flat_vectors.get(&field.0);
2574 if ann_index.is_none() && lazy_flat.is_none() {
2575 return Ok(Vec::new());
2576 }
2577
2578 if ann_index.is_some() && lazy_flat.is_none() {
2579 return Err(Error::Corruption(format!(
2580 "dense ANN field {} is missing flat vector storage",
2581 field.0
2582 )));
2583 }
2584
2585 if let Some(flat) = lazy_flat
2586 && flat.dim != params.dim
2587 {
2588 return Err(Error::Corruption(format!(
2589 "dense vector field {} has schema dimension {} but flat storage dimension {}",
2590 field.0, params.dim, flat.dim
2591 )));
2592 }
2593
2594 let results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
2595 match index {
2597 VectorIndex::IvfPq(lazy) => {
2598 let index = lazy.get();
2599 let codebook =
2600 self.trained_vectors
2601 .codebooks
2602 .get(&field.0)
2603 .ok_or_else(|| {
2604 Error::Schema(format!(
2605 "IVF-PQ index requires a global codebook for field {}",
2606 field.0
2607 ))
2608 })?;
2609 let centroids =
2610 self.trained_vectors
2611 .centroids
2612 .get(&field.0)
2613 .ok_or_else(|| {
2614 Error::Schema(format!(
2615 "IVF-PQ index requires coarse centroids for field {}",
2616 field.0
2617 ))
2618 })?;
2619 validate_coarse_centroids(centroids, params.dim)?;
2620 let routing = self
2621 .schema
2622 .get_field_entry(field)
2623 .and_then(|entry| entry.dense_vector_config.as_ref())
2624 .map_or(crate::dsl::IvfRoutingMode::Auto, |config| {
2625 config.ivf_routing
2626 });
2627 validate_ivf_pq_ann(index, centroids, codebook, params.dim, routing)?;
2628 let query_plan = float_query_plan(
2629 centroids,
2630 codebook,
2631 query,
2632 params.nprobe,
2633 routing,
2634 probe_cache,
2635 )?;
2636 let flat = lazy_flat.expect("ANN/flat pairing validated above");
2637 index
2638 .search_ivf_pq_distinct(
2639 fetch_k.min(flat.num_docs_with_vectors()),
2640 &query_plan,
2641 )
2642 .map_err(|error| {
2643 Error::Corruption(format!(
2644 "invalid IVF-PQ payload for field {}: {error}",
2645 field.0,
2646 ))
2647 })?
2648 .into_iter()
2649 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
2650 .collect()
2651 }
2652 VectorIndex::BinaryIvf(_) => {
2653 Vec::new()
2655 }
2656 }
2657 } else if let Some(lazy_flat) = lazy_flat {
2658 let dim = lazy_flat.dim;
2660 let n = lazy_flat.num_vectors;
2661 let quant = lazy_flat.quantization;
2662 let batch_len =
2663 bounded_vector_score_batch(lazy_flat.vector_byte_size(), DENSE_SCORE_BATCH);
2664 let mut collector = FlatDocumentCollector::new(fetch_k.min(n), combiner);
2665 let mut scores = vec![0f32; batch_len];
2666
2667 for batch_start in (0..n).step_by(batch_len) {
2668 let batch_count = batch_len.min(n - batch_start);
2669 let batch_bytes = lazy_flat
2670 .read_vectors_batch_sync(batch_start, batch_count)
2671 .map_err(crate::Error::Io)?;
2672 let raw = batch_bytes.as_slice();
2673
2674 Self::score_quantized_batch(
2675 query,
2676 raw,
2677 quant,
2678 dim,
2679 &mut scores[..batch_count],
2680 params.unit_norm,
2681 )?;
2682
2683 for (i, &score) in scores.iter().enumerate().take(batch_count) {
2684 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
2685 collector.push(doc_id, ordinal, score);
2686 }
2687 }
2688
2689 return Ok(collector.into_results());
2690 } else {
2691 return Ok(Vec::new());
2692 };
2693
2694 if ann_index.is_some()
2696 && !results.is_empty()
2697 && let Some(lazy_flat) = lazy_flat
2698 {
2699 return exact_score_dense_candidate_documents_sync(
2700 &results,
2701 lazy_flat,
2702 query,
2703 params.unit_norm,
2704 combiner,
2705 k,
2706 );
2707 }
2708
2709 Ok(combine_grouped_ordinal_results(results, combiner, k))
2710 }
2711
2712 #[cfg(feature = "sync")]
2717 fn search_binary_dense_vector_sync_impl(
2718 &self,
2719 field: Field,
2720 query: &[u8],
2721 k: usize,
2722 combiner: crate::query::MultiValueCombiner,
2723 probe_cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
2724 ) -> Result<Vec<VectorSearchResult>> {
2725 let schema_dim = self.validate_binary_search_request(field, query)?;
2726 combiner.validate().map_err(Error::Query)?;
2727 if k == 0 {
2728 return Ok(Vec::new());
2729 }
2730 let t0 = crate::observe::Timer::start();
2731 if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0) {
2732 let ivf = lazy.get();
2733 let config = self
2734 .schema
2735 .get_field_entry(field)
2736 .and_then(|entry| entry.binary_dense_vector_config.as_ref())
2737 .ok_or_else(|| {
2738 Error::Schema(format!(
2739 "binary IVF field {} has no schema configuration",
2740 field.0
2741 ))
2742 })?;
2743 let quantizer = self
2744 .trained_vectors
2745 .binary_quantizers
2746 .get(&field.0)
2747 .ok_or_else(|| {
2748 Error::Schema(format!(
2749 "global binary IVF field {} has no loaded quantizer",
2750 field.0
2751 ))
2752 })?;
2753 validate_binary_ann(ivf, quantizer, config, schema_dim, field)?;
2754 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
2755 Error::Corruption(format!(
2756 "global binary IVF field {} is missing flat vector storage",
2757 field.0
2758 ))
2759 })?;
2760 let clusters = binary_probe_clusters(
2761 quantizer,
2762 query,
2763 config.nprobe,
2764 config.ivf_routing,
2765 probe_cache,
2766 )?;
2767 let candidate_docs = k.min(flat.num_docs_with_vectors());
2768 let ann_results = if flat.num_vectors == flat.num_docs_with_vectors() {
2769 ivf.search_binary_clusters::<false>(query, candidate_docs, &clusters)
2770 } else {
2771 ivf.search_binary_clusters::<true>(query, candidate_docs, &clusters)
2772 }
2773 .map_err(|error| {
2774 Error::Corruption(format!(
2775 "invalid binary IVF payload for field {}: {error}",
2776 field.0,
2777 ))
2778 })?;
2779 let results = exact_score_binary_candidate_documents_sync(
2780 &ann_results,
2781 flat,
2782 query,
2783 schema_dim,
2784 combiner,
2785 k,
2786 )?;
2787 crate::observe::dense_l1(
2788 self.schema.index_label(),
2789 self.schema.get_field_name(field).unwrap_or("?"),
2790 "global_binary_ivf",
2791 t0.secs(),
2792 results.len(),
2793 );
2794 return Ok(results);
2795 }
2796 let lazy_flat = match self.flat_vectors.get(&field.0) {
2797 Some(f) => f,
2798 None => return Ok(Vec::new()),
2799 };
2800
2801 let dim_bits = lazy_flat.dim;
2802 let byte_len = lazy_flat.vector_byte_size();
2803 let n = lazy_flat.num_vectors;
2804
2805 if dim_bits != schema_dim {
2806 return Err(Error::Corruption(format!(
2807 "binary vector field {} has schema dimension {} but flat storage dimension {}",
2808 field.0, schema_dim, dim_bits
2809 )));
2810 }
2811
2812 if byte_len != query.len() {
2813 return Err(Error::Schema(format!(
2814 "Binary query vector byte length {} != field byte length {}",
2815 query.len(),
2816 byte_len
2817 )));
2818 }
2819
2820 let batch_len = bounded_vector_score_batch(byte_len, BINARY_SCORE_BATCH);
2821 let mut collector = FlatDocumentCollector::new(k, combiner);
2822 let mut scores = vec![0f32; batch_len];
2823
2824 for batch_start in (0..n).step_by(batch_len) {
2825 let batch_count = batch_len.min(n - batch_start);
2826 let batch_bytes = lazy_flat
2827 .read_vectors_batch_sync(batch_start, batch_count)
2828 .map_err(crate::Error::Io)?;
2829 let raw = batch_bytes.as_slice();
2830
2831 crate::structures::simd::batch_hamming_scores(
2832 query,
2833 raw,
2834 byte_len,
2835 dim_bits,
2836 &mut scores[..batch_count],
2837 );
2838
2839 for (i, &score) in scores.iter().enumerate().take(batch_count) {
2840 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
2841 collector.push(doc_id, ordinal, score);
2842 }
2843 }
2844
2845 let results = collector.into_results();
2846
2847 crate::observe::dense_l1(
2848 self.schema.index_label(),
2849 self.schema.get_field_name(field).unwrap_or("?"),
2850 "binary_flat",
2851 t0.secs(),
2852 results.len(),
2853 );
2854 Ok(results)
2855 }
2856
2857 #[cfg(feature = "sync")]
2858 pub fn search_binary_dense_vector_sync(
2859 &self,
2860 field: Field,
2861 query: &[u8],
2862 k: usize,
2863 combiner: crate::query::MultiValueCombiner,
2864 ) -> Result<Vec<VectorSearchResult>> {
2865 self.search_binary_dense_vector_sync_impl(field, query, k, combiner, None)
2866 }
2867
2868 #[cfg(feature = "sync")]
2869 pub(crate) fn search_binary_dense_vector_sync_with_probe_cache(
2870 &self,
2871 field: Field,
2872 query: &[u8],
2873 k: usize,
2874 combiner: crate::query::MultiValueCombiner,
2875 probe_cache: &std::sync::Mutex<Option<crate::structures::IvfProbePlan>>,
2876 ) -> Result<Vec<VectorSearchResult>> {
2877 self.search_binary_dense_vector_sync_impl(field, query, k, combiner, Some(probe_cache))
2878 }
2879}
2880
2881#[cfg(test)]
2882mod dense_search_safety_tests {
2883 use super::*;
2884
2885 #[test]
2886 fn dense_fetch_count_rejects_non_finite_and_unbounded_factors() {
2887 for factor in [
2888 f32::NAN,
2889 f32::INFINITY,
2890 f32::NEG_INFINITY,
2891 0.0,
2892 0.5,
2893 2.01,
2894 MAX_DENSE_RERANK_FACTOR + 1.0,
2895 ] {
2896 assert!(
2897 checked_dense_fetch_k(10, factor).is_err(),
2898 "factor={factor}"
2899 );
2900 }
2901 }
2902
2903 #[test]
2904 fn flat_document_collector_does_not_let_one_multivalue_doc_crowd_out_others() {
2905 let mut collector = FlatDocumentCollector::new(2, crate::query::MultiValueCombiner::Max);
2906 collector.push(1, 0, 1.0);
2907 collector.push(1, 1, 0.9);
2908 collector.push(2, 0, 0.8);
2909
2910 let results = collector.into_results();
2911 assert_eq!(
2912 results
2913 .iter()
2914 .map(|result| result.doc_id)
2915 .collect::<Vec<_>>(),
2916 vec![1, 2]
2917 );
2918 assert_eq!(results[0].ordinals.len(), 2);
2919 }
2920
2921 #[test]
2922 fn flat_document_collector_evicts_by_score_then_doc_id() {
2923 let mut collector = FlatDocumentCollector::new(2, crate::query::MultiValueCombiner::Max);
2924 collector.push(1, 0, 0.5);
2925 collector.push(3, 0, 0.8);
2926 collector.push(2, 0, 0.9);
2927 let results = collector.into_results();
2928 assert_eq!(
2929 results
2930 .iter()
2931 .map(|result| result.doc_id)
2932 .collect::<Vec<_>>(),
2933 vec![2, 3]
2934 );
2935
2936 let mut tied = FlatDocumentCollector::new(1, crate::query::MultiValueCombiner::Max);
2937 tied.push(2, 0, 1.0);
2938 tied.push(1, 0, 1.0);
2939 let results = tied.into_results();
2940 assert_eq!(results[0].doc_id, 1);
2941 }
2942
2943 #[test]
2944 fn dense_fetch_count_rounds_up_and_detects_overflow() {
2945 assert_eq!(checked_dense_fetch_k(3, 1.5).unwrap(), 5);
2946 assert_eq!(checked_dense_fetch_k(10_000, 2.0).unwrap(), 20_000);
2947 assert!(checked_dense_fetch_k(10_001, 2.0).is_err());
2948 assert!(checked_dense_fetch_k(usize::MAX, 2.0).is_err());
2949 }
2950
2951 #[test]
2952 fn file_ranges_reject_overflow_and_truncation() {
2953 assert_eq!(checked_file_range(4, 3, 7, "test").unwrap(), 4..7);
2954 assert!(checked_file_range(u64::MAX, 1, u64::MAX, "test").is_err());
2955 assert!(checked_file_range(5, 3, 7, "test").is_err());
2956 }
2957
2958 #[test]
2959 fn candidate_vector_reads_coalesce_contiguous_values() {
2960 let mut runs = Vec::new();
2961 plan_vector_read_runs(&[3, 4, 5, 9, 12, 13], &mut runs).unwrap();
2962 assert_eq!(runs.len(), 3);
2963 assert!(matches!(
2964 runs.as_slice(),
2965 [
2966 VectorReadRun {
2967 buffer_start: 0,
2968 flat_start: 3,
2969 count: 3,
2970 },
2971 VectorReadRun {
2972 buffer_start: 3,
2973 flat_start: 9,
2974 count: 1,
2975 },
2976 VectorReadRun {
2977 buffer_start: 4,
2978 flat_start: 12,
2979 count: 2,
2980 },
2981 ]
2982 ));
2983 assert!(plan_vector_read_runs(&[3, 3], &mut runs).is_err());
2984 }
2985
2986 #[tokio::test]
2987 async fn multivalue_ann_rerank_streams_past_document_candidate_cap() {
2988 use crate::directories::{FileHandle, OwnedBytes};
2989 use crate::segment::FlatVectorData;
2990
2991 const VALUES: usize = MAX_DENSE_CANDIDATES_PER_SEGMENT + 1;
2992 let mut encoded = Vec::new();
2993 let vectors = vec![1.0f32; VALUES];
2994 let doc_ids: Vec<_> = (0..VALUES).map(|ordinal| (0, ordinal as u16)).collect();
2995 FlatVectorData::serialize_binary_from_flat_streaming(
2996 1,
2997 &vectors,
2998 &doc_ids,
2999 DenseVectorQuantization::F32,
3000 &mut encoded,
3001 )
3002 .unwrap();
3003 let flat = LazyFlatVectorData::open_with_doc_limit(
3004 FileHandle::from_bytes(OwnedBytes::new(encoded)),
3005 Some(1),
3006 )
3007 .await
3008 .unwrap();
3009
3010 let (results, stats) = exact_score_dense_candidate_documents(
3011 &[(0, 0, 0.0)],
3012 &flat,
3013 &[1.0],
3014 false,
3015 crate::query::MultiValueCombiner::Max,
3016 1,
3017 )
3018 .await
3019 .unwrap();
3020 assert_eq!(stats.vector_count, VALUES);
3021 assert_eq!(results.len(), 1);
3022 assert_eq!(results[0].ordinals.len(), VALUES);
3023 assert!((results[0].score - 1.0).abs() < 1e-5);
3024 }
3025}