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
486fn prepare_vector_read_runs(
490 flat: &LazyFlatVectorData,
491 indexes: &[usize],
492 runs: &mut Vec<VectorReadRun>,
493) -> Result<()> {
494 plan_vector_read_runs(indexes, runs)?;
495 #[cfg(feature = "native")]
496 flat.prefetch_vectors(indexes.iter().copied());
497 #[cfg(not(feature = "native"))]
498 let _ = flat;
499 Ok(())
500}
501
502async fn read_vector_runs(
503 flat: &LazyFlatVectorData,
504 indexes: &[usize],
505 runs: &mut Vec<VectorReadRun>,
506 output: &mut [u8],
507) -> Result<()> {
508 prepare_vector_read_runs(flat, indexes, runs)?;
509 let vector_byte_size = flat.vector_byte_size();
510 for run in runs {
511 let bytes = flat
512 .read_vectors_batch(run.flat_start, run.count)
513 .await
514 .map_err(Error::Io)?;
515 let start = run
516 .buffer_start
517 .checked_mul(vector_byte_size)
518 .ok_or_else(|| Error::Query("dense rerank buffer offset overflow".into()))?;
519 let end = start
520 .checked_add(bytes.len())
521 .ok_or_else(|| Error::Query("dense rerank buffer range overflow".into()))?;
522 let destination = output
523 .get_mut(start..end)
524 .ok_or_else(|| Error::Corruption("dense rerank buffer is too short".into()))?;
525 destination.copy_from_slice(bytes.as_slice());
526 }
527 Ok(())
528}
529
530#[cfg(feature = "sync")]
531fn read_vector_runs_sync(
532 flat: &LazyFlatVectorData,
533 indexes: &[usize],
534 runs: &mut Vec<VectorReadRun>,
535 output: &mut [u8],
536) -> Result<()> {
537 prepare_vector_read_runs(flat, indexes, runs)?;
538 let vector_byte_size = flat.vector_byte_size();
539 for run in runs {
540 let bytes = flat
541 .read_vectors_batch_sync(run.flat_start, run.count)
542 .map_err(Error::Io)?;
543 let start = run
544 .buffer_start
545 .checked_mul(vector_byte_size)
546 .ok_or_else(|| Error::Query("dense rerank buffer offset overflow".into()))?;
547 let end = start
548 .checked_add(bytes.len())
549 .ok_or_else(|| Error::Query("dense rerank buffer range overflow".into()))?;
550 let destination = output
551 .get_mut(start..end)
552 .ok_or_else(|| Error::Corruption("dense rerank buffer is too short".into()))?;
553 destination.copy_from_slice(bytes.as_slice());
554 }
555 Ok(())
556}
557
558#[derive(Default)]
559struct DenseRerankStats {
560 vector_count: usize,
561 resolve_elapsed: std::time::Duration,
562 read_elapsed: std::time::Duration,
563 score_elapsed: std::time::Duration,
564}
565
566async fn exact_score_dense_candidate_documents(
567 ann_results: &[RawVectorCandidate],
568 flat: &LazyFlatVectorData,
569 query: &[f32],
570 unit_norm: bool,
571 combiner: crate::query::MultiValueCombiner,
572 limit: usize,
573) -> Result<(Vec<VectorSearchResult>, DenseRerankStats)> {
574 let resolve_started = std::time::Instant::now();
575 let documents = ann_candidate_document_ranges(ann_results, flat)?;
576 let mut stats = DenseRerankStats {
577 vector_count: documents.vector_count,
578 resolve_elapsed: resolve_started.elapsed(),
579 ..Default::default()
580 };
581 let vector_byte_size = flat.vector_byte_size();
582 let batch_len = bounded_vector_score_batch(vector_byte_size, DENSE_SCORE_BATCH);
583 let raw_capacity = batch_len
584 .checked_mul(vector_byte_size)
585 .ok_or_else(|| Error::Query("dense rerank buffer size overflow".to_string()))?;
586 let mut raw = vec![0u8; raw_capacity];
587 let mut scores = vec![0.0f32; batch_len];
588 let mut batch = Vec::with_capacity(batch_len);
589 let mut flat_indexes = Vec::with_capacity(batch_len);
590 let mut read_runs = Vec::new();
591 let mut cursor = CandidateVectorCursor::new(&documents.ranges);
592 let mut collector = FlatDocumentCollector::new(limit, combiner);
593 let mut scored = 0usize;
594
595 while cursor.fill_batch(flat, &mut batch, batch_len)? {
596 flat_indexes.clear();
597 flat_indexes.extend(batch.iter().map(|&(_, _, flat_index)| flat_index));
598 let raw_len = batch
599 .len()
600 .checked_mul(vector_byte_size)
601 .ok_or_else(|| Error::Query("dense rerank buffer size overflow".to_string()))?;
602 let raw = &mut raw[..raw_len];
603
604 let read_started = std::time::Instant::now();
605 read_vector_runs(flat, &flat_indexes, &mut read_runs, raw).await?;
606 stats.read_elapsed += read_started.elapsed();
607
608 let score_started = std::time::Instant::now();
609 SegmentReader::score_quantized_batch(
610 query,
611 raw,
612 flat.quantization,
613 flat.dim,
614 &mut scores[..batch.len()],
615 unit_norm,
616 )?;
617 stats.score_elapsed += score_started.elapsed();
618 for (buffer_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
619 collector.push(doc_id, ordinal, scores[buffer_index]);
620 }
621 scored += batch.len();
622 }
623 debug_assert_eq!(scored, documents.vector_count);
624 Ok((collector.into_results(), stats))
625}
626
627#[cfg(feature = "sync")]
628fn exact_score_dense_candidate_documents_sync(
629 ann_results: &[RawVectorCandidate],
630 flat: &LazyFlatVectorData,
631 query: &[f32],
632 unit_norm: bool,
633 combiner: crate::query::MultiValueCombiner,
634 limit: usize,
635) -> Result<Vec<VectorSearchResult>> {
636 let documents = ann_candidate_document_ranges(ann_results, flat)?;
637 let vector_byte_size = flat.vector_byte_size();
638 let batch_len = bounded_vector_score_batch(vector_byte_size, DENSE_SCORE_BATCH);
639 let raw_capacity = batch_len
640 .checked_mul(vector_byte_size)
641 .ok_or_else(|| Error::Query("dense rerank buffer size overflow".to_string()))?;
642 let mut raw = vec![0u8; raw_capacity];
643 let mut scores = vec![0.0f32; batch_len];
644 let mut batch = Vec::with_capacity(batch_len);
645 let mut flat_indexes = Vec::with_capacity(batch_len);
646 let mut read_runs = Vec::new();
647 let mut cursor = CandidateVectorCursor::new(&documents.ranges);
648 let mut collector = FlatDocumentCollector::new(limit, combiner);
649 let mut scored = 0usize;
650
651 while cursor.fill_batch(flat, &mut batch, batch_len)? {
652 flat_indexes.clear();
653 flat_indexes.extend(batch.iter().map(|&(_, _, flat_index)| flat_index));
654 let raw_len = batch
655 .len()
656 .checked_mul(vector_byte_size)
657 .ok_or_else(|| Error::Query("dense rerank buffer size overflow".to_string()))?;
658 let raw = &mut raw[..raw_len];
659 read_vector_runs_sync(flat, &flat_indexes, &mut read_runs, raw)?;
660 SegmentReader::score_quantized_batch(
661 query,
662 raw,
663 flat.quantization,
664 flat.dim,
665 &mut scores[..batch.len()],
666 unit_norm,
667 )?;
668 for (buffer_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
669 collector.push(doc_id, ordinal, scores[buffer_index]);
670 }
671 scored += batch.len();
672 }
673 debug_assert_eq!(scored, documents.vector_count);
674 Ok(collector.into_results())
675}
676
677async fn exact_score_binary_candidate_documents(
678 ann_results: &[RawVectorCandidate],
679 flat: &LazyFlatVectorData,
680 query: &[u8],
681 dim_bits: usize,
682 combiner: crate::query::MultiValueCombiner,
683 limit: usize,
684) -> Result<Vec<VectorSearchResult>> {
685 let documents = ann_candidate_document_ranges(ann_results, flat)?;
686 let probe_scores: FxHashMap<(DocId, u16), f32> = ann_results
687 .iter()
688 .map(|&(doc_id, ordinal, score)| ((doc_id, ordinal), score))
689 .collect();
690 let vector_byte_size = flat.vector_byte_size();
691 let batch_len = bounded_vector_score_batch(vector_byte_size, BINARY_SCORE_BATCH);
692 let raw_capacity = batch_len
693 .checked_mul(vector_byte_size)
694 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
695 let mut raw = vec![0u8; raw_capacity];
696 let mut scores = vec![0.0f32; batch_len];
697 let mut batch_scores = vec![0.0f32; batch_len];
698 let mut batch = Vec::with_capacity(batch_len);
699 let mut unresolved = Vec::with_capacity(batch_len);
700 let mut unresolved_flat_indexes = Vec::with_capacity(batch_len);
701 let mut read_runs = Vec::new();
702 let mut cursor = CandidateVectorCursor::new(&documents.ranges);
703 let mut collector = FlatDocumentCollector::new(limit, combiner);
704 let mut scored = 0usize;
705
706 while cursor.fill_batch(flat, &mut batch, batch_len)? {
707 unresolved.clear();
708 for (batch_index, &(doc_id, ordinal, flat_index)) in batch.iter().enumerate() {
709 if let Some(&score) = probe_scores.get(&(doc_id, ordinal)) {
710 batch_scores[batch_index] = score;
711 } else {
712 unresolved.push((batch_index, flat_index));
713 }
714 }
715 unresolved_flat_indexes.clear();
716 unresolved_flat_indexes.extend(unresolved.iter().map(|&(_, flat_index)| flat_index));
717 let raw_len = unresolved
718 .len()
719 .checked_mul(vector_byte_size)
720 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
721 let raw = &mut raw[..raw_len];
722 read_vector_runs(flat, &unresolved_flat_indexes, &mut read_runs, raw).await?;
723 crate::structures::simd::batch_hamming_scores(
724 query,
725 raw,
726 vector_byte_size,
727 dim_bits,
728 &mut scores[..unresolved.len()],
729 );
730 for (buffer_index, &(batch_index, _)) in unresolved.iter().enumerate() {
731 batch_scores[batch_index] = scores[buffer_index];
732 }
733 for (batch_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
734 collector.push(doc_id, ordinal, batch_scores[batch_index]);
735 }
736 scored += batch.len();
737 }
738 debug_assert_eq!(scored, documents.vector_count);
739 Ok(collector.into_results())
740}
741
742#[cfg(feature = "sync")]
743fn exact_score_binary_candidate_documents_sync(
744 ann_results: &[RawVectorCandidate],
745 flat: &LazyFlatVectorData,
746 query: &[u8],
747 dim_bits: usize,
748 combiner: crate::query::MultiValueCombiner,
749 limit: usize,
750) -> Result<Vec<VectorSearchResult>> {
751 let documents = ann_candidate_document_ranges(ann_results, flat)?;
752 let probe_scores: FxHashMap<(DocId, u16), f32> = ann_results
753 .iter()
754 .map(|&(doc_id, ordinal, score)| ((doc_id, ordinal), score))
755 .collect();
756 let vector_byte_size = flat.vector_byte_size();
757 let batch_len = bounded_vector_score_batch(vector_byte_size, BINARY_SCORE_BATCH);
758 let raw_capacity = batch_len
759 .checked_mul(vector_byte_size)
760 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
761 let mut raw = vec![0u8; raw_capacity];
762 let mut scores = vec![0.0f32; batch_len];
763 let mut batch_scores = vec![0.0f32; batch_len];
764 let mut batch = Vec::with_capacity(batch_len);
765 let mut unresolved = Vec::with_capacity(batch_len);
766 let mut unresolved_flat_indexes = Vec::with_capacity(batch_len);
767 let mut read_runs = Vec::new();
768 let mut cursor = CandidateVectorCursor::new(&documents.ranges);
769 let mut collector = FlatDocumentCollector::new(limit, combiner);
770 let mut scored = 0usize;
771
772 while cursor.fill_batch(flat, &mut batch, batch_len)? {
773 unresolved.clear();
774 for (batch_index, &(doc_id, ordinal, flat_index)) in batch.iter().enumerate() {
775 if let Some(&score) = probe_scores.get(&(doc_id, ordinal)) {
776 batch_scores[batch_index] = score;
777 } else {
778 unresolved.push((batch_index, flat_index));
779 }
780 }
781 unresolved_flat_indexes.clear();
782 unresolved_flat_indexes.extend(unresolved.iter().map(|&(_, flat_index)| flat_index));
783 let raw_len = unresolved
784 .len()
785 .checked_mul(vector_byte_size)
786 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
787 let raw = &mut raw[..raw_len];
788 read_vector_runs_sync(flat, &unresolved_flat_indexes, &mut read_runs, raw)?;
789 crate::structures::simd::batch_hamming_scores(
790 query,
791 raw,
792 vector_byte_size,
793 dim_bits,
794 &mut scores[..unresolved.len()],
795 );
796 for (buffer_index, &(batch_index, _)) in unresolved.iter().enumerate() {
797 batch_scores[batch_index] = scores[buffer_index];
798 }
799 for (batch_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
800 collector.push(doc_id, ordinal, batch_scores[batch_index]);
801 }
802 scored += batch.len();
803 }
804 debug_assert_eq!(scored, documents.vector_count);
805 Ok(collector.into_results())
806}
807
808fn validate_coarse_centroids(centroids: &CoarseCentroids, dim: usize) -> Result<()> {
809 let expected = (centroids.num_clusters as usize)
810 .checked_mul(dim)
811 .ok_or_else(|| Error::Corruption("coarse centroid size overflow".into()))?;
812 if centroids.num_clusters == 0
813 || centroids.dim != dim
814 || centroids.centroids.len() != expected
815 || centroids.centroids.iter().any(|value| !value.is_finite())
816 {
817 return Err(Error::Corruption(format!(
818 "invalid coarse centroids: clusters={}, dim={}, values={} (expected dim={dim}, values={expected})",
819 centroids.num_clusters,
820 centroids.dim,
821 centroids.centroids.len()
822 )));
823 }
824 Ok(())
825}
826
827fn validate_ivf_pq_ann(
828 index: &crate::segment::ann_disk::AnnDiskIndex,
829 centroids: &CoarseCentroids,
830 codebook: &crate::structures::PQCodebook,
831 dim: usize,
832 routing: crate::dsl::IvfRoutingMode,
833) -> Result<()> {
834 let header = index.header();
835 if header.dim != dim
836 || codebook.config.dim != dim
837 || header.code_size != codebook.config.num_subspaces
838 || header.num_clusters != centroids.num_clusters
839 || header.quantizer_version != centroids.version
840 || header.codebook_version != codebook.version
841 || header.routing != routing
842 {
843 return Err(Error::Corruption(format!(
844 "IVF-PQ payload/codebook/centroid metadata does not match schema dimension {dim}"
845 )));
846 }
847 Ok(())
848}
849
850#[derive(Debug, Default)]
855pub struct DensePlanCache {
856 pub(crate) ivf_pq: std::sync::Mutex<Option<std::sync::Arc<crate::structures::IvfPqQueryPlan>>>,
857 pub(crate) tq: std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqQueryPlan>>>,
858 pub(crate) ivf_tq: std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqIvfQueryPlan>>>,
859}
860
861fn search_tq_segment(
866 index: &crate::segment::ann_disk::AnnDiskIndex,
867 codec: &crate::structures::TqCodec,
868 query: &[f32],
869 fetch_k: usize,
870 field: Field,
871 dim: usize,
872 plan_cache: Option<&std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqQueryPlan>>>>,
873) -> Result<Vec<RawVectorCandidate>> {
874 validate_tq_ann(index, codec, dim, field)?;
875 let plan = match plan_cache {
876 Some(cache) => {
877 let mut cached = cache
878 .lock()
879 .map_err(|_| Error::Internal("TQ plan cache is poisoned".into()))?;
880 match cached.as_ref() {
881 Some(plan) if plan.fingerprint() == codec.fingerprint() => {
882 std::sync::Arc::clone(plan)
883 }
884 _ => {
885 let plan =
886 std::sync::Arc::new(crate::structures::TqQueryPlan::build(codec, query));
887 *cached = Some(std::sync::Arc::clone(&plan));
888 plan
889 }
890 }
891 }
892 None => std::sync::Arc::new(crate::structures::TqQueryPlan::build(codec, query)),
893 };
894 index.search_tq_distinct(fetch_k, &plan).map_err(|error| {
895 Error::Corruption(format!("invalid TQ payload for field {}: {error}", field.0))
896 })
897}
898
899fn validate_tq_ann(
900 index: &crate::segment::ann_disk::AnnDiskIndex,
901 codec: &crate::structures::TqCodec,
902 dim: usize,
903 field: Field,
904) -> Result<()> {
905 let header = index.header();
906 if header.dim != dim
907 || codec.dim() != dim
908 || header.code_size != codec.code_size()
909 || header.quantizer_version != codec.fingerprint()
910 || header.codebook_version != 0
911 || header.num_clusters != 1
912 {
913 return Err(Error::Corruption(format!(
914 "TQ payload for field {} does not match the codec derived from schema dimension {dim}",
915 field.0,
916 )));
917 }
918 Ok(())
919}
920
921#[allow(clippy::too_many_arguments)]
925fn search_ivf_tq_segment(
926 index: &crate::segment::ann_disk::AnnDiskIndex,
927 centroids: &CoarseCentroids,
928 codec: &crate::structures::TqCodec,
929 query: &[f32],
930 fetch_k: usize,
931 field: Field,
932 nprobe: usize,
933 routing: crate::dsl::IvfRoutingMode,
934 plan_cache: Option<
935 &std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqIvfQueryPlan>>>,
936 >,
937) -> Result<Vec<RawVectorCandidate>> {
938 let effective_nprobe = nprobe.clamp(1, centroids.num_clusters as usize);
939 let request_fingerprint = crate::structures::vector::ivf::routing::float_probe_fingerprint(
940 query,
941 effective_nprobe,
942 routing,
943 );
944 let build = || {
945 std::sync::Arc::new(crate::structures::TqIvfQueryPlan::build(
946 centroids,
947 codec,
948 query,
949 effective_nprobe,
950 routing,
951 ))
952 };
953 let plan = match plan_cache {
954 Some(cache) => {
955 let mut cached = cache
956 .lock()
957 .map_err(|_| Error::Internal("IVF-TQ plan cache is poisoned".into()))?;
958 match cached.as_ref() {
959 Some(plan)
960 if plan.quantizer_version == centroids.version
961 && plan.fingerprint == codec.fingerprint()
962 && plan.request_fingerprint == request_fingerprint
963 && plan.cluster_ids.len() == effective_nprobe =>
964 {
965 std::sync::Arc::clone(plan)
966 }
967 _ => {
968 let plan = build();
969 *cached = Some(std::sync::Arc::clone(&plan));
970 plan
971 }
972 }
973 }
974 None => build(),
975 };
976 index
977 .search_ivf_tq_distinct(fetch_k, &plan)
978 .map_err(|error| {
979 Error::Corruption(format!(
980 "invalid IVF-TQ payload for field {}: {error}",
981 field.0
982 ))
983 })
984}
985
986fn validate_ivf_tq_ann(
987 index: &crate::segment::ann_disk::AnnDiskIndex,
988 centroids: &CoarseCentroids,
989 codec: &crate::structures::TqCodec,
990 dim: usize,
991 routing: crate::dsl::IvfRoutingMode,
992 field: Field,
993) -> Result<()> {
994 let header = index.header();
995 if header.dim != dim
996 || codec.dim() != dim
997 || header.code_size != codec.code_size()
998 || header.num_clusters != centroids.num_clusters
999 || header.quantizer_version != centroids.version
1000 || header.codebook_version != codec.fingerprint()
1001 || header.routing != routing
1002 {
1003 return Err(Error::Corruption(format!(
1004 "IVF-TQ payload for field {} does not match its quantizer/codec generation",
1005 field.0,
1006 )));
1007 }
1008 Ok(())
1009}
1010
1011fn validate_binary_ann(
1012 index: &crate::segment::ann_disk::AnnDiskIndex,
1013 quantizer: &crate::structures::BinaryCoarseQuantizer,
1014 config: &crate::dsl::BinaryDenseVectorConfig,
1015 dim: usize,
1016 field: Field,
1017) -> Result<()> {
1018 let header = index.header();
1019 if header.dim != dim
1020 || header.code_size != config.byte_len()
1021 || header.num_clusters != quantizer.num_clusters
1022 || header.quantizer_version != quantizer.version
1023 || header.codebook_version != 0
1024 || header.routing != config.ivf_routing
1025 || quantizer.dim_bits != dim
1026 {
1027 return Err(Error::Corruption(format!(
1028 "binary IVF field {} does not match its quantizer/schema generation",
1029 field.0,
1030 )));
1031 }
1032 Ok(())
1033}
1034
1035fn float_query_plan(
1036 centroids: &CoarseCentroids,
1037 codebook: &crate::structures::PQCodebook,
1038 query: &[f32],
1039 nprobe: usize,
1040 routing: crate::dsl::IvfRoutingMode,
1041 cache: Option<&std::sync::Mutex<Option<std::sync::Arc<crate::structures::IvfPqQueryPlan>>>>,
1042) -> Result<std::sync::Arc<crate::structures::IvfPqQueryPlan>> {
1043 let effective_nprobe = nprobe.clamp(1, centroids.num_clusters as usize);
1044 let request_fingerprint = crate::structures::vector::ivf::routing::float_probe_fingerprint(
1045 query,
1046 effective_nprobe,
1047 routing,
1048 );
1049 if let Some(cache) = cache {
1050 let mut cached = cache
1051 .lock()
1052 .map_err(|_| Error::Internal("dense IVF probe cache is poisoned".into()))?;
1053 if let Some(plan) = cached.as_ref()
1054 && plan.quantizer_version == centroids.version
1055 && plan.codebook_version == codebook.version
1056 && plan.request_fingerprint == request_fingerprint
1057 && plan.cluster_ids.len() == effective_nprobe
1058 {
1059 return Ok(std::sync::Arc::clone(plan));
1060 }
1061 let plan = std::sync::Arc::new(crate::structures::IvfPqQueryPlan::build(
1062 centroids,
1063 codebook,
1064 query,
1065 effective_nprobe,
1066 routing,
1067 ));
1068 *cached = Some(std::sync::Arc::clone(&plan));
1069 return Ok(plan);
1070 }
1071 Ok(std::sync::Arc::new(
1072 crate::structures::IvfPqQueryPlan::build(
1073 centroids,
1074 codebook,
1075 query,
1076 effective_nprobe,
1077 routing,
1078 ),
1079 ))
1080}
1081
1082fn binary_probe_clusters(
1083 quantizer: &crate::structures::BinaryCoarseQuantizer,
1084 query: &[u8],
1085 nprobe: usize,
1086 routing: crate::dsl::IvfRoutingMode,
1087 cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
1088) -> Result<std::sync::Arc<[u32]>> {
1089 let effective_nprobe = nprobe.clamp(1, quantizer.num_clusters as usize);
1090 let request_fingerprint = crate::structures::vector::ivf::routing::binary_probe_fingerprint(
1091 query,
1092 effective_nprobe,
1093 routing,
1094 );
1095 if let Some(cache) = cache {
1096 let mut cached = cache
1097 .lock()
1098 .map_err(|_| Error::Internal("binary IVF probe cache is poisoned".into()))?;
1099 if let Some(plan) = cached.as_ref()
1100 && plan.quantizer_version == quantizer.version
1101 && plan.request_fingerprint == request_fingerprint
1102 && plan.cluster_ids.len() == effective_nprobe
1103 {
1104 return Ok(std::sync::Arc::clone(&plan.cluster_ids));
1105 }
1106 let plan = quantizer.probe(query, effective_nprobe, routing);
1107 let clusters = std::sync::Arc::clone(&plan.cluster_ids);
1108 *cached = Some(plan);
1109 return Ok(clusters);
1110 }
1111 Ok(quantizer
1112 .probe(query, effective_nprobe, routing)
1113 .cluster_ids)
1114}
1115
1116pub struct SegmentReader {
1122 meta: SegmentMeta,
1123 term_dict: Arc<AsyncSSTableReader<TermInfo>>,
1125 postings_handle: FileHandle,
1127 store: Arc<AsyncStoreReader>,
1129 schema: Arc<Schema>,
1130 vector_indexes: FxHashMap<u32, VectorIndex>,
1132 flat_vectors: FxHashMap<u32, LazyFlatVectorData>,
1134 dense_file_backed_bytes: u64,
1136 trained_vectors: Arc<crate::segment::TrainedVectorStructures>,
1138 sparse_indexes: FxHashMap<u32, SparseIndex>,
1140 bmp_indexes: FxHashMap<u32, BmpIndex>,
1142 sparse_file_backed_bytes: u64,
1144 positions_handle: Option<FileHandle>,
1146 fast_fields: FxHashMap<u32, crate::structures::fast_field::FastFieldReader>,
1148 #[cfg(feature = "native")]
1150 dense_pin_report: crate::segment::pin::PinReport,
1151 #[cfg(feature = "native")]
1153 sparse_pin_report: crate::segment::pin::PinReport,
1154}
1155
1156impl SegmentReader {
1157 pub async fn open<D: Directory>(
1159 dir: &D,
1160 segment_id: SegmentId,
1161 schema: Arc<Schema>,
1162 cache_blocks: usize,
1163 ) -> Result<Self> {
1164 Self::open_with_cache_blocks(dir, segment_id, schema, cache_blocks, cache_blocks).await
1165 }
1166
1167 pub async fn open_with_cache_blocks<D: Directory>(
1173 dir: &D,
1174 segment_id: SegmentId,
1175 schema: Arc<Schema>,
1176 term_cache_blocks: usize,
1177 store_cache_blocks: usize,
1178 ) -> Result<Self> {
1179 let files = SegmentFiles::new(segment_id.0);
1180
1181 let meta_slice = dir.open_read(&files.meta).await?;
1183 let meta_bytes = meta_slice.read_bytes().await?;
1184 let meta = SegmentMeta::deserialize(meta_bytes.as_slice())?;
1185 debug_assert_eq!(meta.id, segment_id.0);
1186
1187 let term_dict_handle = dir.open_lazy(&files.term_dict).await?;
1189 let term_dict = AsyncSSTableReader::open(term_dict_handle, term_cache_blocks).await?;
1190
1191 let postings_handle = dir.open_lazy(&files.postings).await?;
1193
1194 let store_handle = dir.open_lazy(&files.store).await?;
1196 let store = AsyncStoreReader::open(store_handle, store_cache_blocks).await?;
1197
1198 let vectors_data = loader::load_vectors_file(dir, &files, &schema, meta.num_docs).await?;
1200 let dense_file_backed_bytes = vectors_data.file_backed_bytes;
1201 let vector_indexes = vectors_data.indexes;
1202 let flat_vectors = vectors_data.flat_vectors;
1203
1204 #[cfg(feature = "native")]
1209 for (field_id, lazy_flat) in &flat_vectors {
1210 if vector_indexes.contains_key(field_id) {
1211 lazy_flat.advise_random_access();
1212 }
1213 }
1214
1215 let sparse_data = loader::load_sparse_file(dir, &files, meta.num_docs, &schema).await?;
1217 let sparse_file_backed_bytes = sparse_data.file_backed_bytes;
1218 let sparse_indexes = sparse_data.maxscore_indexes;
1219 let bmp_indexes = sparse_data.bmp_indexes;
1220
1221 let positions_handle = loader::open_positions_file(dir, &files, &schema).await?;
1223
1224 let fast_fields = loader::load_fast_fields_file(dir, &files, &schema).await?;
1226
1227 {
1229 let mut parts = vec![format!(
1230 "[segment] loaded {:016x}: docs={}",
1231 segment_id.0, meta.num_docs
1232 )];
1233 if !vector_indexes.is_empty() || !flat_vectors.is_empty() {
1234 parts.push(format!(
1235 "dense vectors: {} ANN + {} flat fields",
1236 vector_indexes.len(),
1237 flat_vectors.len()
1238 ));
1239 }
1240 for (field_id, idx) in &sparse_indexes {
1241 parts.push(format!(
1242 "sparse vector field {}: {} dims, ~{}",
1243 field_id,
1244 idx.num_dimensions(),
1245 crate::format_bytes(idx.num_dimensions() as u64 * 24)
1246 ));
1247 }
1248 for (field_id, idx) in &bmp_indexes {
1249 parts.push(format!(
1250 "bmp field {}: {} dims, {} blocks",
1251 field_id,
1252 idx.dims(),
1253 idx.num_blocks
1254 ));
1255 }
1256 if !fast_fields.is_empty() {
1257 parts.push(format!("fast: {} fields", fast_fields.len()));
1258 }
1259 log::debug!("{}", parts.join(", "));
1260 }
1261
1262 #[allow(unused_mut)]
1263 let mut reader = Self {
1264 meta,
1265 term_dict: Arc::new(term_dict),
1266 postings_handle,
1267 store: Arc::new(store),
1268 schema,
1269 vector_indexes,
1270 flat_vectors,
1271 dense_file_backed_bytes,
1272 trained_vectors: Arc::new(crate::segment::TrainedVectorStructures::default()),
1273 sparse_indexes,
1274 bmp_indexes,
1275 sparse_file_backed_bytes,
1276 positions_handle,
1277 fast_fields,
1278 #[cfg(feature = "native")]
1279 dense_pin_report: Default::default(),
1280 #[cfg(feature = "native")]
1281 sparse_pin_report: Default::default(),
1282 };
1283
1284 #[cfg(feature = "native")]
1286 reader.apply_pin_policy(&crate::segment::pin::pin_policy().to_owned());
1287
1288 Ok(reader)
1289 }
1290
1291 #[cfg(feature = "native")]
1301 pub(crate) fn apply_pin_policy(&mut self, policy: &crate::segment::pin::PinPolicy) {
1302 use crate::segment::pin::PinReport;
1303
1304 if !policy.is_enabled() {
1305 return;
1306 }
1307 let mut remaining = policy.budget_bytes;
1308 let mut dense_report = PinReport::default();
1309 let mut sparse_report = PinReport::default();
1310
1311 for index in self.vector_indexes.values_mut() {
1313 index.pin_lookup_directory(policy.mode, &mut remaining, &mut dense_report);
1314 }
1315 for bmp in self.bmp_indexes.values_mut() {
1317 bmp.pin_block_starts(policy.mode, &mut remaining, &mut sparse_report);
1318 }
1319 for sparse in self.sparse_indexes.values_mut() {
1321 sparse.pin_skip_section(policy.mode, &mut remaining, &mut sparse_report);
1322 }
1323 for flat in self.flat_vectors.values_mut() {
1325 flat.pin_doc_ids(policy.mode, &mut remaining, &mut dense_report);
1326 }
1327 for bmp in self.bmp_indexes.values_mut() {
1328 bmp.pin_doc_maps(policy.mode, &mut remaining, &mut sparse_report);
1329 }
1330 for bmp in self.bmp_indexes.values_mut() {
1332 bmp.pin_sb_grid(policy.mode, &mut remaining, &mut sparse_report);
1333 }
1334
1335 let report = PinReport {
1336 intended_bytes: dense_report
1337 .intended_bytes
1338 .saturating_add(sparse_report.intended_bytes),
1339 pinned_bytes: dense_report
1340 .pinned_bytes
1341 .saturating_add(sparse_report.pinned_bytes),
1342 skipped_budget_bytes: dense_report
1343 .skipped_budget_bytes
1344 .saturating_add(sparse_report.skipped_budget_bytes),
1345 failed_bytes: dense_report
1346 .failed_bytes
1347 .saturating_add(sparse_report.failed_bytes),
1348 heap_copy_bytes: dense_report
1349 .heap_copy_bytes
1350 .saturating_add(sparse_report.heap_copy_bytes),
1351 };
1352 if report.skipped_budget_bytes > 0 || report.failed_bytes > 0 {
1353 log::warn!(
1354 "[pin] segment {:016x}: pinned {}/{} (budget skipped {}, mlock failed {}) — \
1355 raise HERMES_PIN_METADATA_BUDGET_MB or RLIMIT_MEMLOCK for full coverage",
1356 self.meta.id,
1357 crate::format_bytes(report.pinned_bytes),
1358 crate::format_bytes(report.intended_bytes),
1359 crate::format_bytes(report.skipped_budget_bytes),
1360 crate::format_bytes(report.failed_bytes),
1361 );
1362 } else if report.pinned_bytes > 0 {
1363 log::info!(
1364 "[pin] segment {:016x}: pinned {} of hot metadata ({:?})",
1365 self.meta.id,
1366 crate::format_bytes(report.pinned_bytes),
1367 policy.mode,
1368 );
1369 }
1370 self.dense_pin_report = dense_report;
1371 self.sparse_pin_report = sparse_report;
1372 }
1373
1374 pub fn meta(&self) -> &SegmentMeta {
1380 &self.meta
1381 }
1382
1383 pub fn num_docs(&self) -> u32 {
1384 self.meta.num_docs
1385 }
1386
1387 pub fn avg_field_len(&self, field: Field) -> f32 {
1389 self.meta.avg_field_len(field)
1390 }
1391
1392 pub fn schema(&self) -> &Schema {
1393 &self.schema
1394 }
1395
1396 pub fn sparse_indexes(&self) -> &FxHashMap<u32, SparseIndex> {
1398 &self.sparse_indexes
1399 }
1400
1401 pub fn sparse_index(&self, field: Field) -> Option<&SparseIndex> {
1403 self.sparse_indexes.get(&field.0)
1404 }
1405
1406 pub fn bmp_index(&self, field: Field) -> Option<&BmpIndex> {
1408 self.bmp_indexes.get(&field.0)
1409 }
1410
1411 pub fn bmp_indexes(&self) -> &FxHashMap<u32, BmpIndex> {
1413 &self.bmp_indexes
1414 }
1415
1416 pub fn vector_indexes(&self) -> &FxHashMap<u32, VectorIndex> {
1418 &self.vector_indexes
1419 }
1420
1421 pub fn flat_vectors(&self) -> &FxHashMap<u32, LazyFlatVectorData> {
1423 &self.flat_vectors
1424 }
1425
1426 pub fn fast_field(
1428 &self,
1429 field_id: u32,
1430 ) -> Option<&crate::structures::fast_field::FastFieldReader> {
1431 self.fast_fields.get(&field_id)
1432 }
1433
1434 pub fn fast_fields(&self) -> &FxHashMap<u32, crate::structures::fast_field::FastFieldReader> {
1436 &self.fast_fields
1437 }
1438
1439 pub fn term_dict_stats(&self) -> SSTableStats {
1441 self.term_dict.stats()
1442 }
1443
1444 pub fn memory_stats(&self) -> SegmentMemoryStats {
1446 let term_dict_stats = self.term_dict.stats();
1447
1448 let term_dict_cache_bytes = self.term_dict.cached_bytes();
1452 let store_cache_bytes = self.store.cached_bytes();
1453
1454 let sparse_heap_bytes: usize = self
1457 .sparse_indexes
1458 .values()
1459 .map(|s| s.estimated_heap_bytes())
1460 .sum::<usize>()
1461 + self
1462 .bmp_indexes
1463 .values()
1464 .map(|b| b.estimated_heap_bytes())
1465 .sum::<usize>();
1466
1467 let dense_heap_bytes: usize = self
1470 .vector_indexes
1471 .values()
1472 .map(|v| v.estimated_heap_bytes())
1473 .sum::<usize>()
1474 + self
1475 .flat_vectors
1476 .values()
1477 .map(LazyFlatVectorData::estimated_heap_bytes)
1478 .sum::<usize>();
1479
1480 #[cfg(feature = "native")]
1481 let (sparse_heap_bytes, dense_heap_bytes) = (
1482 sparse_heap_bytes.saturating_add(
1483 usize::try_from(self.sparse_pin_report.heap_copy_bytes).unwrap_or(usize::MAX),
1484 ),
1485 dense_heap_bytes.saturating_add(
1486 usize::try_from(self.dense_pin_report.heap_copy_bytes).unwrap_or(usize::MAX),
1487 ),
1488 );
1489
1490 #[cfg(feature = "native")]
1491 let (
1492 sparse_pinned_metadata_bytes,
1493 sparse_pin_intended_bytes,
1494 dense_pinned_metadata_bytes,
1495 dense_pin_intended_bytes,
1496 ) = (
1497 self.sparse_pin_report.pinned_bytes,
1498 self.sparse_pin_report.intended_bytes,
1499 self.dense_pin_report.pinned_bytes,
1500 self.dense_pin_report.intended_bytes,
1501 );
1502 #[cfg(not(feature = "native"))]
1503 let (
1504 sparse_pinned_metadata_bytes,
1505 sparse_pin_intended_bytes,
1506 dense_pinned_metadata_bytes,
1507 dense_pin_intended_bytes,
1508 ) = (0u64, 0u64, 0u64, 0u64);
1509
1510 let pinned_metadata_bytes =
1511 sparse_pinned_metadata_bytes.saturating_add(dense_pinned_metadata_bytes);
1512 let pin_intended_bytes = sparse_pin_intended_bytes.saturating_add(dense_pin_intended_bytes);
1513
1514 SegmentMemoryStats {
1515 segment_id: self.meta.id,
1516 num_docs: self.meta.num_docs,
1517 term_dict_cache_bytes,
1518 store_cache_bytes,
1519 sparse_heap_bytes,
1520 dense_heap_bytes,
1521 term_bloom_file_bytes: term_dict_stats.bloom_filter_size as u64,
1522 sparse_file_backed_bytes: self.sparse_file_backed_bytes,
1523 dense_file_backed_bytes: self.dense_file_backed_bytes,
1524 pinned_metadata_bytes,
1525 pin_intended_bytes,
1526 sparse_pinned_metadata_bytes,
1527 sparse_pin_intended_bytes,
1528 dense_pinned_metadata_bytes,
1529 dense_pin_intended_bytes,
1530 }
1531 }
1532
1533 pub async fn get_postings(
1538 &self,
1539 field: Field,
1540 term: &[u8],
1541 ) -> Result<Option<BlockPostingList>> {
1542 log::debug!(
1543 "SegmentReader::get_postings field={} term_len={}",
1544 field.0,
1545 term.len()
1546 );
1547
1548 let mut key = Vec::with_capacity(4 + term.len());
1550 key.extend_from_slice(&field.0.to_le_bytes());
1551 key.extend_from_slice(term);
1552
1553 let term_info = match self.term_dict.get(&key).await? {
1555 Some(info) => {
1556 log::debug!("SegmentReader::get_postings found term_info");
1557 info
1558 }
1559 None => {
1560 log::debug!("SegmentReader::get_postings term not found");
1561 return Ok(None);
1562 }
1563 };
1564
1565 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
1567 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
1569 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
1570 posting_list.push(doc_id, tf);
1571 }
1572 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
1573 return Ok(Some(block_list));
1574 }
1575
1576 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
1578 Error::Corruption("TermInfo has neither inline nor external data".to_string())
1579 })?;
1580
1581 let range = checked_file_range(
1582 posting_offset,
1583 posting_len,
1584 self.postings_handle.len(),
1585 "posting",
1586 )?;
1587 let posting_bytes = self.postings_handle.read_bytes_range(range).await?;
1588 let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
1589
1590 Ok(Some(block_list))
1591 }
1592
1593 pub async fn get_prefix_postings(
1595 &self,
1596 field: Field,
1597 prefix: &[u8],
1598 ) -> Result<Vec<BlockPostingList>> {
1599 if prefix.is_empty() {
1600 return Err(Error::Query("prefix must not be empty".into()));
1601 }
1602 let mut key_prefix = Vec::with_capacity(4 + prefix.len());
1604 key_prefix.extend_from_slice(&field.0.to_le_bytes());
1605 key_prefix.extend_from_slice(prefix);
1606
1607 let (entries, truncated) = self
1608 .term_dict
1609 .prefix_scan_limited(&key_prefix, MAX_PREFIX_TERMS)
1610 .await?;
1611 if truncated {
1612 return Err(Error::Query(format!(
1613 "prefix expands to more than {MAX_PREFIX_TERMS} terms"
1614 )));
1615 }
1616 let posting_count: u64 = entries
1617 .iter()
1618 .map(|(_, term_info)| term_info.doc_freq() as u64)
1619 .sum();
1620 if posting_count > MAX_PREFIX_POSTINGS {
1621 return Err(Error::Query(format!(
1622 "prefix expands to {posting_count} postings (maximum {MAX_PREFIX_POSTINGS})"
1623 )));
1624 }
1625 let mut results = Vec::with_capacity(entries.len());
1626
1627 for (_key, term_info) in entries {
1628 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
1629 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
1630 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
1631 posting_list.push(doc_id, tf);
1632 }
1633 results.push(BlockPostingList::from_posting_list(&posting_list)?);
1634 } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
1635 let range = checked_file_range(
1636 posting_offset,
1637 posting_len,
1638 self.postings_handle.len(),
1639 "prefix posting",
1640 )?;
1641 let posting_bytes = self.postings_handle.read_bytes_range(range).await?;
1642 results.push(BlockPostingList::deserialize_zero_copy(posting_bytes)?);
1643 }
1644 }
1645
1646 Ok(results)
1647 }
1648
1649 pub async fn doc(&self, local_doc_id: DocId) -> Result<Option<Document>> {
1654 self.doc_with_fields(local_doc_id, None).await
1655 }
1656
1657 pub async fn doc_with_fields(
1663 &self,
1664 local_doc_id: DocId,
1665 fields: Option<&rustc_hash::FxHashSet<u32>>,
1666 ) -> Result<Option<Document>> {
1667 let mut doc = match fields {
1668 Some(set) => {
1669 let field_ids: Vec<u32> = set.iter().copied().collect();
1670 match self
1671 .store
1672 .get_fields(local_doc_id, &self.schema, &field_ids)
1673 .await
1674 {
1675 Ok(Some(d)) => d,
1676 Ok(None) => return Ok(None),
1677 Err(e) => return Err(Error::from(e)),
1678 }
1679 }
1680 None => match self.store.get(local_doc_id, &self.schema).await {
1681 Ok(Some(d)) => d,
1682 Ok(None) => return Ok(None),
1683 Err(e) => return Err(Error::from(e)),
1684 },
1685 };
1686
1687 for (&field_id, lazy_flat) in &self.flat_vectors {
1689 if let Some(set) = fields
1691 && !set.contains(&field_id)
1692 {
1693 continue;
1694 }
1695
1696 let is_binary = lazy_flat.quantization == DenseVectorQuantization::Binary;
1697 let (start, entries) = lazy_flat.flat_indexes_for_doc(local_doc_id);
1698 for (j, &(_doc_id, _ordinal)) in entries.iter().enumerate() {
1699 let flat_idx = start + j;
1700 if is_binary {
1701 let vbs = lazy_flat.vector_byte_size();
1702 let mut raw = vec![0u8; vbs];
1703 match lazy_flat.read_vector_raw_into(flat_idx, &mut raw).await {
1704 Ok(()) => {
1705 doc.add_binary_dense_vector(Field(field_id), raw);
1706 }
1707 Err(e) => {
1708 log::warn!(
1709 "Failed to hydrate binary dense vector field {}: {}",
1710 field_id,
1711 e
1712 );
1713 }
1714 }
1715 } else {
1716 match lazy_flat.get_vector(flat_idx).await {
1717 Ok(vec) => {
1718 doc.add_dense_vector(Field(field_id), vec);
1719 }
1720 Err(e) => {
1721 log::warn!("Failed to hydrate dense vector field {}: {}", field_id, e);
1722 }
1723 }
1724 }
1725 }
1726 }
1727
1728 Ok(Some(doc))
1729 }
1730
1731 pub async fn prefetch_terms(
1733 &self,
1734 field: Field,
1735 start_term: &[u8],
1736 end_term: &[u8],
1737 ) -> Result<()> {
1738 let mut start_key = Vec::with_capacity(4 + start_term.len());
1739 start_key.extend_from_slice(&field.0.to_le_bytes());
1740 start_key.extend_from_slice(start_term);
1741
1742 let mut end_key = Vec::with_capacity(4 + end_term.len());
1743 end_key.extend_from_slice(&field.0.to_le_bytes());
1744 end_key.extend_from_slice(end_term);
1745
1746 self.term_dict.prefetch_range(&start_key, &end_key).await?;
1747 Ok(())
1748 }
1749
1750 pub fn store_has_dict(&self) -> bool {
1752 self.store.has_dict()
1753 }
1754
1755 pub fn store(&self) -> &super::store::AsyncStoreReader {
1757 &self.store
1758 }
1759
1760 pub fn store_raw_blocks(&self) -> Vec<RawStoreBlock> {
1762 self.store.raw_blocks()
1763 }
1764
1765 pub fn store_data_slice(&self) -> &FileHandle {
1767 self.store.data_slice()
1768 }
1769
1770 pub async fn all_terms(&self) -> Result<Vec<(Vec<u8>, TermInfo)>> {
1772 self.term_dict.all_entries().await.map_err(Error::from)
1773 }
1774
1775 pub async fn all_terms_with_stats(&self) -> Result<Vec<(Field, String, u32)>> {
1780 let entries = self.term_dict.all_entries().await?;
1781 let mut result = Vec::with_capacity(entries.len());
1782
1783 for (key, term_info) in entries {
1784 if key.len() > 4 {
1786 let field_id = u32::from_le_bytes([key[0], key[1], key[2], key[3]]);
1787 let term_bytes = &key[4..];
1788 if let Ok(term_str) = std::str::from_utf8(term_bytes) {
1789 result.push((Field(field_id), term_str.to_string(), term_info.doc_freq()));
1790 }
1791 }
1792 }
1793
1794 Ok(result)
1795 }
1796
1797 pub fn term_dict_iter(&self) -> crate::structures::AsyncSSTableIterator<'_, TermInfo> {
1799 self.term_dict.iter()
1800 }
1801
1802 pub async fn prefetch_term_dict(&self) -> crate::Result<()> {
1806 self.term_dict
1807 .prefetch_all_data_bulk()
1808 .await
1809 .map_err(crate::Error::from)
1810 }
1811
1812 pub async fn read_postings(&self, offset: u64, len: u64) -> Result<Vec<u8>> {
1814 let range = checked_file_range(offset, len, self.postings_handle.len(), "posting")?;
1815 let bytes = self.postings_handle.read_bytes_range(range).await?;
1816 Ok(bytes.to_vec())
1817 }
1818
1819 pub async fn read_position_bytes(&self, offset: u64, len: u64) -> Result<Option<Vec<u8>>> {
1821 let handle = match &self.positions_handle {
1822 Some(h) => h,
1823 None => return Ok(None),
1824 };
1825 let range = checked_file_range(offset, len, handle.len(), "position")?;
1826 let bytes = handle.read_bytes_range(range).await?;
1827 Ok(Some(bytes.to_vec()))
1828 }
1829
1830 pub fn has_positions_file(&self) -> bool {
1832 self.positions_handle.is_some()
1833 }
1834
1835 fn validate_dense_search_request(
1839 &self,
1840 field: Field,
1841 query: &[f32],
1842 nprobe: usize,
1843 rerank_factor: f32,
1844 combiner: crate::query::MultiValueCombiner,
1845 ) -> Result<DenseSearchParams> {
1846 let entry = self
1847 .schema
1848 .get_field_entry(field)
1849 .ok_or_else(|| Error::FieldNotFound(field.0.to_string()))?;
1850 if entry.field_type != crate::dsl::FieldType::DenseVector {
1851 return Err(Error::InvalidFieldType {
1852 expected: "dense_vector".to_string(),
1853 got: format!("{:?}", entry.field_type),
1854 });
1855 }
1856 let config = entry.dense_vector_config.as_ref().ok_or_else(|| {
1857 Error::Schema(format!(
1858 "dense vector field '{}' has no dense vector configuration",
1859 entry.name
1860 ))
1861 })?;
1862
1863 if query.is_empty() {
1864 return Err(Error::Query(format!(
1865 "dense query vector for field '{}' must not be empty",
1866 entry.name
1867 )));
1868 }
1869 if query.len() != config.dim {
1870 return Err(Error::Query(format!(
1871 "dense query vector dimension {} does not match field '{}' dimension {}",
1872 query.len(),
1873 entry.name,
1874 config.dim
1875 )));
1876 }
1877 if let Some((index, value)) = query
1878 .iter()
1879 .enumerate()
1880 .find(|(_, value)| !value.is_finite())
1881 {
1882 return Err(Error::Query(format!(
1883 "dense query vector for field '{}' contains non-finite value {value} at index {index}",
1884 entry.name
1885 )));
1886 }
1887
1888 let nprobe = match (nprobe, config.nprobe) {
1891 (0, 0) => 32,
1892 (0, schema_nprobe) => schema_nprobe,
1893 (query_nprobe, _) => query_nprobe,
1894 };
1895 if nprobe > MAX_DENSE_NPROBE {
1896 return Err(Error::Query(format!(
1897 "dense nprobe must be at most {MAX_DENSE_NPROBE}, got {nprobe}"
1898 )));
1899 }
1900
1901 checked_dense_fetch_k(0, rerank_factor)?;
1904 combiner.validate().map_err(Error::Query)?;
1905
1906 Ok(DenseSearchParams {
1907 dim: config.dim,
1908 nprobe,
1909 unit_norm: config.unit_norm,
1910 })
1911 }
1912
1913 fn validate_binary_search_request(&self, field: Field, query: &[u8]) -> Result<usize> {
1914 let entry = self
1915 .schema
1916 .get_field_entry(field)
1917 .ok_or_else(|| Error::FieldNotFound(field.0.to_string()))?;
1918 if entry.field_type != crate::dsl::FieldType::BinaryDenseVector {
1919 return Err(Error::InvalidFieldType {
1920 expected: "binary_dense_vector".to_string(),
1921 got: format!("{:?}", entry.field_type),
1922 });
1923 }
1924 let config = entry.binary_dense_vector_config.as_ref().ok_or_else(|| {
1925 Error::Schema(format!(
1926 "binary dense vector field '{}' has no configuration",
1927 entry.name
1928 ))
1929 })?;
1930 if config.dim == 0 || !config.dim.is_multiple_of(8) {
1931 return Err(Error::Schema(format!(
1932 "binary dense vector field '{}' has invalid dimension {}",
1933 entry.name, config.dim
1934 )));
1935 }
1936 if query.len() != config.byte_len() {
1937 return Err(Error::Query(format!(
1938 "binary query byte length {} does not match field '{}' byte length {}",
1939 query.len(),
1940 entry.name,
1941 config.byte_len()
1942 )));
1943 }
1944 Ok(config.dim)
1945 }
1946
1947 fn score_quantized_batch(
1953 query: &[f32],
1954 raw: &[u8],
1955 quant: crate::dsl::DenseVectorQuantization,
1956 dim: usize,
1957 scores: &mut [f32],
1958 unit_norm: bool,
1959 ) -> Result<()> {
1960 use crate::dsl::DenseVectorQuantization;
1961 use crate::structures::simd;
1962
1963 if query.len() != dim {
1964 return Err(Error::Query(format!(
1965 "dense SIMD query dimension {} does not match vector dimension {dim}",
1966 query.len()
1967 )));
1968 }
1969 let element_size = match quant {
1970 DenseVectorQuantization::F32 => std::mem::size_of::<f32>(),
1971 DenseVectorQuantization::F16 => std::mem::size_of::<u16>(),
1972 DenseVectorQuantization::UInt8 => 1,
1973 DenseVectorQuantization::Binary => {
1974 return Err(Error::InvalidFieldType {
1975 expected: "non-binary dense vector".to_string(),
1976 got: "binary dense vector".to_string(),
1977 });
1978 }
1979 };
1980 let required_bytes = scores
1981 .len()
1982 .checked_mul(dim)
1983 .and_then(|elements| elements.checked_mul(element_size))
1984 .ok_or_else(|| Error::Corruption("dense vector batch byte length overflow".into()))?;
1985 if raw.len() < required_bytes {
1986 return Err(Error::Corruption(format!(
1987 "dense vector batch is truncated: need {required_bytes} bytes, got {}",
1988 raw.len()
1989 )));
1990 }
1991 if quant == DenseVectorQuantization::F16
1992 && required_bytes > 0
1993 && !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>())
1994 {
1995 return Err(Error::Corruption(
1996 "f16 vector data is not 2-byte aligned".to_string(),
1997 ));
1998 }
1999
2000 match (quant, unit_norm) {
2001 (DenseVectorQuantization::F32, false) => {
2002 let num_floats = scores.len() * dim;
2003 if !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()) {
2004 return Err(Error::Corruption(
2005 "f32 vector data is not 4-byte aligned".to_string(),
2006 ));
2007 }
2008 let vectors: &[f32] =
2009 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
2010 simd::batch_cosine_scores(query, vectors, dim, scores);
2011 }
2012 (DenseVectorQuantization::F32, true) => {
2013 let num_floats = scores.len() * dim;
2014 if !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()) {
2015 return Err(Error::Corruption(
2016 "f32 vector data is not 4-byte aligned".to_string(),
2017 ));
2018 }
2019 let vectors: &[f32] =
2020 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
2021 simd::batch_dot_scores(query, vectors, dim, scores);
2022 }
2023 (DenseVectorQuantization::F16, false) => {
2024 simd::batch_cosine_scores_f16(query, raw, dim, scores);
2025 }
2026 (DenseVectorQuantization::F16, true) => {
2027 simd::batch_dot_scores_f16(query, raw, dim, scores);
2028 }
2029 (DenseVectorQuantization::UInt8, false) => {
2030 simd::batch_cosine_scores_u8(query, raw, dim, scores);
2031 }
2032 (DenseVectorQuantization::UInt8, true) => {
2033 simd::batch_dot_scores_u8(query, raw, dim, scores);
2034 }
2035 (DenseVectorQuantization::Binary, _) => unreachable!("validated above"),
2036 }
2037 Ok(())
2038 }
2039
2040 pub async fn search_dense_vector(
2046 &self,
2047 field: Field,
2048 query: &[f32],
2049 k: usize,
2050 nprobe: usize,
2051 rerank_factor: f32,
2052 combiner: crate::query::MultiValueCombiner,
2053 ) -> Result<Vec<VectorSearchResult>> {
2054 self.search_dense_vector_impl(field, query, k, nprobe, rerank_factor, combiner, None)
2055 .await
2056 }
2057
2058 #[allow(clippy::too_many_arguments)]
2059 pub(crate) async fn search_dense_vector_with_probe_cache(
2060 &self,
2061 field: Field,
2062 query: &[f32],
2063 k: usize,
2064 nprobe: usize,
2065 rerank_factor: f32,
2066 combiner: crate::query::MultiValueCombiner,
2067 plan_cache: &DensePlanCache,
2068 ) -> Result<Vec<VectorSearchResult>> {
2069 self.search_dense_vector_impl(
2070 field,
2071 query,
2072 k,
2073 nprobe,
2074 rerank_factor,
2075 combiner,
2076 Some(plan_cache),
2077 )
2078 .await
2079 }
2080
2081 #[allow(clippy::too_many_arguments)]
2082 async fn search_dense_vector_impl(
2083 &self,
2084 field: Field,
2085 query: &[f32],
2086 k: usize,
2087 nprobe: usize,
2088 rerank_factor: f32,
2089 combiner: crate::query::MultiValueCombiner,
2090 plan_cache: Option<&DensePlanCache>,
2091 ) -> Result<Vec<VectorSearchResult>> {
2092 let params =
2093 self.validate_dense_search_request(field, query, nprobe, rerank_factor, combiner)?;
2094 let fetch_k = checked_dense_fetch_k(k, rerank_factor)?;
2095 if k == 0 {
2096 return Ok(Vec::new());
2097 }
2098
2099 let ann_index = self.vector_indexes.get(&field.0);
2100 let lazy_flat = self.flat_vectors.get(&field.0);
2101 if ann_index.is_none() && lazy_flat.is_none() {
2103 return Ok(Vec::new());
2104 }
2105
2106 if ann_index.is_some() && lazy_flat.is_none() {
2107 return Err(Error::Corruption(format!(
2108 "dense ANN field {} is missing flat vector storage",
2109 field.0
2110 )));
2111 }
2112
2113 if let Some(flat) = lazy_flat
2114 && flat.dim != params.dim
2115 {
2116 return Err(Error::Corruption(format!(
2117 "dense vector field {} has schema dimension {} but flat storage dimension {}",
2118 field.0, params.dim, flat.dim
2119 )));
2120 }
2121
2122 let t0 = std::time::Instant::now();
2124 let mut flat_results = None;
2125 let results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
2126 match index {
2128 VectorIndex::IvfPq(lazy) => {
2129 let index = lazy.get();
2130 let codebook =
2131 self.trained_vectors
2132 .codebooks
2133 .get(&field.0)
2134 .ok_or_else(|| {
2135 Error::Schema(format!(
2136 "IVF-PQ index requires a global codebook for field {}",
2137 field.0
2138 ))
2139 })?;
2140 let centroids =
2141 self.trained_vectors
2142 .centroids
2143 .get(&field.0)
2144 .ok_or_else(|| {
2145 Error::Schema(format!(
2146 "IVF-PQ index requires coarse centroids for field {}",
2147 field.0
2148 ))
2149 })?;
2150 validate_coarse_centroids(centroids, params.dim)?;
2151 let routing = self
2152 .schema
2153 .get_field_entry(field)
2154 .and_then(|entry| entry.dense_vector_config.as_ref())
2155 .map_or(crate::dsl::IvfRoutingMode::Auto, |config| {
2156 config.ivf_routing
2157 });
2158 validate_ivf_pq_ann(index, centroids, codebook, params.dim, routing)?;
2159 let query_plan = float_query_plan(
2160 centroids,
2161 codebook,
2162 query,
2163 params.nprobe,
2164 routing,
2165 plan_cache.map(|cache| &cache.ivf_pq),
2166 )?;
2167 let flat = lazy_flat.expect("ANN/flat pairing validated above");
2168 index
2169 .search_ivf_pq_distinct(
2170 fetch_k.min(flat.num_docs_with_vectors()),
2171 &query_plan,
2172 )
2173 .map_err(|error| {
2174 Error::Corruption(format!(
2175 "invalid IVF-PQ payload for field {}: {error}",
2176 field.0,
2177 ))
2178 })?
2179 .into_iter()
2180 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
2181 .collect()
2182 }
2183 VectorIndex::Tq { index: lazy, codec } => {
2184 let flat = lazy_flat.expect("ANN/flat pairing validated above");
2185 search_tq_segment(
2187 lazy.get(),
2188 codec,
2189 query,
2190 fetch_k.min(flat.num_docs_with_vectors()),
2191 field,
2192 params.dim,
2193 plan_cache.map(|cache| &cache.tq),
2194 )?
2195 }
2196 VectorIndex::IvfTq { index: lazy, codec } => {
2197 let index = lazy.get();
2198 let centroids =
2199 self.trained_vectors
2200 .centroids
2201 .get(&field.0)
2202 .ok_or_else(|| {
2203 Error::Schema(format!(
2204 "IVF-TQ index requires coarse centroids for field {}",
2205 field.0
2206 ))
2207 })?;
2208 validate_coarse_centroids(centroids, params.dim)?;
2209 let routing = self
2210 .schema
2211 .get_field_entry(field)
2212 .and_then(|entry| entry.dense_vector_config.as_ref())
2213 .map_or(crate::dsl::IvfRoutingMode::Auto, |config| {
2214 config.ivf_routing
2215 });
2216 validate_ivf_tq_ann(index, centroids, codec, params.dim, routing, field)?;
2217 let flat = lazy_flat.expect("ANN/flat pairing validated above");
2218 search_ivf_tq_segment(
2219 index,
2220 centroids,
2221 codec,
2222 query,
2223 fetch_k.min(flat.num_docs_with_vectors()),
2224 field,
2225 params.nprobe,
2226 routing,
2227 plan_cache.map(|cache| &cache.ivf_tq),
2228 )?
2229 }
2230 VectorIndex::BinaryIvf(_) => {
2231 Vec::new()
2233 }
2234 }
2235 } else if let Some(lazy_flat) = lazy_flat {
2236 log::debug!(
2240 "[dense_vector_search] field {}: brute-force on {} vectors (dim={}, quant={:?})",
2241 field.0,
2242 lazy_flat.num_vectors,
2243 lazy_flat.dim,
2244 lazy_flat.quantization
2245 );
2246 let dim = lazy_flat.dim;
2247 let n = lazy_flat.num_vectors;
2248 let quant = lazy_flat.quantization;
2249 let batch_len =
2250 bounded_vector_score_batch(lazy_flat.vector_byte_size(), DENSE_SCORE_BATCH);
2251 let mut collector = FlatDocumentCollector::new(fetch_k.min(n), combiner);
2252 let mut scores = vec![0f32; batch_len];
2253
2254 for batch_start in (0..n).step_by(batch_len) {
2255 let batch_count = batch_len.min(n - batch_start);
2256 let batch_bytes = lazy_flat
2257 .read_vectors_batch(batch_start, batch_count)
2258 .await
2259 .map_err(crate::Error::Io)?;
2260 let raw = batch_bytes.as_slice();
2261
2262 Self::score_quantized_batch(
2263 query,
2264 raw,
2265 quant,
2266 dim,
2267 &mut scores[..batch_count],
2268 params.unit_norm,
2269 )?;
2270
2271 for (i, &score) in scores.iter().enumerate().take(batch_count) {
2272 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
2273 collector.push(doc_id, ordinal, score);
2274 }
2275 }
2276
2277 flat_results = Some(collector.into_results());
2278 Vec::new()
2279 } else {
2280 return Ok(Vec::new());
2281 };
2282 let l1_elapsed = t0.elapsed();
2283 {
2284 let kind = match ann_index {
2285 Some(VectorIndex::IvfPq(_)) => "ivf_pq",
2286 Some(VectorIndex::BinaryIvf(_)) => "binary_ivf",
2287 Some(VectorIndex::Tq { .. }) => "tq_flat",
2288 Some(VectorIndex::IvfTq { .. }) => "ivf_tq",
2289 None => "flat",
2290 };
2291 crate::observe::dense_l1(
2292 self.schema.index_label(),
2293 self.schema.get_field_name(field).unwrap_or("?"),
2294 kind,
2295 l1_elapsed.as_secs_f64(),
2296 flat_results.as_ref().map_or(results.len(), Vec::len),
2297 );
2298 }
2299 log::debug!(
2300 "[dense_vector_search] field {}: L1 returned {} candidates in {:.1}ms",
2301 field.0,
2302 flat_results.as_ref().map_or(results.len(), Vec::len),
2303 l1_elapsed.as_secs_f64() * 1000.0
2304 );
2305
2306 if let Some(results) = flat_results {
2307 return Ok(results);
2308 }
2309
2310 if ann_index.is_some()
2313 && !results.is_empty()
2314 && let Some(lazy_flat) = lazy_flat
2315 {
2316 let t_rerank = std::time::Instant::now();
2317 let vbs = lazy_flat.vector_byte_size();
2318 let (reranked, stats) = exact_score_dense_candidate_documents(
2319 &results,
2320 lazy_flat,
2321 query,
2322 params.unit_norm,
2323 combiner,
2324 k,
2325 )
2326 .await?;
2327
2328 crate::observe::dense_rerank(
2329 self.schema.index_label(),
2330 self.schema.get_field_name(field).unwrap_or("?"),
2331 t_rerank.elapsed().as_secs_f64(),
2332 stats.resolve_elapsed.as_secs_f64(),
2333 stats.read_elapsed.as_secs_f64(),
2334 stats.vector_count,
2335 );
2336 log::debug!(
2337 "[dense_vector_search] field {}: rerank {} vectors (dim={}, quant={:?}, bytes_per_vector={}): resolve={:.1}ms read={:.1}ms score={:.1}ms",
2338 field.0,
2339 stats.vector_count,
2340 lazy_flat.dim,
2341 lazy_flat.quantization,
2342 vbs,
2343 stats.resolve_elapsed.as_secs_f64() * 1000.0,
2344 stats.read_elapsed.as_secs_f64() * 1000.0,
2345 stats.score_elapsed.as_secs_f64() * 1000.0,
2346 );
2347
2348 log::debug!(
2349 "[dense_vector_search] field {}: rerank total={:.1}ms",
2350 field.0,
2351 t_rerank.elapsed().as_secs_f64() * 1000.0
2352 );
2353 return Ok(reranked);
2354 }
2355
2356 Ok(combine_grouped_ordinal_results(results, combiner, k))
2357 }
2358
2359 async fn search_binary_dense_vector_impl(
2364 &self,
2365 field: Field,
2366 query: &[u8],
2367 k: usize,
2368 combiner: crate::query::MultiValueCombiner,
2369 probe_cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
2370 ) -> Result<Vec<VectorSearchResult>> {
2371 let schema_dim = self.validate_binary_search_request(field, query)?;
2372 combiner.validate().map_err(Error::Query)?;
2373 if k == 0 {
2374 return Ok(Vec::new());
2375 }
2376 let t0 = crate::observe::Timer::start();
2377 if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0) {
2378 let ivf = lazy.get();
2379 let config = self
2380 .schema
2381 .get_field_entry(field)
2382 .and_then(|entry| entry.binary_dense_vector_config.as_ref())
2383 .ok_or_else(|| {
2384 Error::Schema(format!(
2385 "binary IVF field {} has no schema configuration",
2386 field.0
2387 ))
2388 })?;
2389 let quantizer = self
2390 .trained_vectors
2391 .binary_quantizers
2392 .get(&field.0)
2393 .ok_or_else(|| {
2394 Error::Schema(format!(
2395 "global binary IVF field {} has no loaded quantizer",
2396 field.0
2397 ))
2398 })?;
2399 validate_binary_ann(ivf, quantizer, config, schema_dim, field)?;
2400 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
2401 Error::Corruption(format!(
2402 "global binary IVF field {} is missing flat vector storage",
2403 field.0
2404 ))
2405 })?;
2406 let clusters = binary_probe_clusters(
2407 quantizer,
2408 query,
2409 config.nprobe,
2410 config.ivf_routing,
2411 probe_cache,
2412 )?;
2413 let single_valued = flat.num_vectors == flat.num_docs_with_vectors();
2414 let candidate_docs = k.min(flat.num_docs_with_vectors());
2415 let ann_results = if single_valued {
2416 ivf.search_binary_clusters::<false>(query, candidate_docs, &clusters)
2417 } else {
2418 ivf.search_binary_clusters::<true>(query, candidate_docs, &clusters)
2419 }
2420 .map_err(|error| {
2421 Error::Corruption(format!(
2422 "invalid binary IVF payload for field {}: {error}",
2423 field.0,
2424 ))
2425 })?;
2426 let results = exact_score_binary_candidate_documents(
2427 &ann_results,
2428 flat,
2429 query,
2430 schema_dim,
2431 combiner,
2432 k,
2433 )
2434 .await?;
2435 crate::observe::dense_l1(
2436 self.schema.index_label(),
2437 self.schema.get_field_name(field).unwrap_or("?"),
2438 "global_binary_ivf",
2439 t0.secs(),
2440 results.len(),
2441 );
2442 return Ok(results);
2443 }
2444 let lazy_flat = match self.flat_vectors.get(&field.0) {
2445 Some(f) => f,
2446 None => return Ok(Vec::new()),
2447 };
2448
2449 let dim_bits = lazy_flat.dim;
2450 let byte_len = lazy_flat.vector_byte_size();
2451 let n = lazy_flat.num_vectors;
2452
2453 if dim_bits != schema_dim {
2454 return Err(Error::Corruption(format!(
2455 "binary vector field {} has schema dimension {} but flat storage dimension {}",
2456 field.0, schema_dim, dim_bits
2457 )));
2458 }
2459
2460 if byte_len != query.len() {
2461 return Err(Error::Schema(format!(
2462 "Binary query vector byte length {} != field byte length {}",
2463 query.len(),
2464 byte_len
2465 )));
2466 }
2467
2468 let batch_len = bounded_vector_score_batch(byte_len, BINARY_SCORE_BATCH);
2469 let mut collector = FlatDocumentCollector::new(k, combiner);
2470 let mut scores = vec![0f32; batch_len];
2471
2472 for batch_start in (0..n).step_by(batch_len) {
2473 let batch_count = batch_len.min(n - batch_start);
2474 let batch_bytes = lazy_flat
2475 .read_vectors_batch(batch_start, batch_count)
2476 .await
2477 .map_err(crate::Error::Io)?;
2478 let raw = batch_bytes.as_slice();
2479
2480 crate::structures::simd::batch_hamming_scores(
2481 query,
2482 raw,
2483 byte_len,
2484 dim_bits,
2485 &mut scores[..batch_count],
2486 );
2487
2488 for (i, &score) in scores.iter().enumerate().take(batch_count) {
2489 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
2490 collector.push(doc_id, ordinal, score);
2491 }
2492 }
2493
2494 let results = collector.into_results();
2495
2496 crate::observe::dense_l1(
2497 self.schema.index_label(),
2498 self.schema.get_field_name(field).unwrap_or("?"),
2499 "binary_flat",
2500 t0.secs(),
2501 results.len(),
2502 );
2503 Ok(results)
2504 }
2505
2506 pub async fn search_binary_dense_vector(
2507 &self,
2508 field: Field,
2509 query: &[u8],
2510 k: usize,
2511 combiner: crate::query::MultiValueCombiner,
2512 ) -> Result<Vec<VectorSearchResult>> {
2513 self.search_binary_dense_vector_impl(field, query, k, combiner, None)
2514 .await
2515 }
2516
2517 pub(crate) async fn search_binary_dense_vector_with_probe_cache(
2518 &self,
2519 field: Field,
2520 query: &[u8],
2521 k: usize,
2522 combiner: crate::query::MultiValueCombiner,
2523 probe_cache: &std::sync::Mutex<Option<crate::structures::IvfProbePlan>>,
2524 ) -> Result<Vec<VectorSearchResult>> {
2525 self.search_binary_dense_vector_impl(field, query, k, combiner, Some(probe_cache))
2526 .await
2527 }
2528
2529 pub fn coarse_centroids(&self, field_id: u32) -> Option<&Arc<CoarseCentroids>> {
2531 self.trained_vectors.centroids.get(&field_id)
2532 }
2533
2534 pub fn set_trained_vectors(
2535 &mut self,
2536 trained_vectors: Arc<crate::segment::TrainedVectorStructures>,
2537 ) {
2538 self.trained_vectors = trained_vectors;
2539 }
2540
2541 pub fn get_vector_index(&self, field: Field) -> Option<&VectorIndex> {
2543 self.vector_indexes.get(&field.0)
2544 }
2545
2546 pub async fn get_positions(
2551 &self,
2552 field: Field,
2553 term: &[u8],
2554 ) -> Result<Option<crate::structures::PositionPostingList>> {
2555 let handle = match &self.positions_handle {
2557 Some(h) => h,
2558 None => return Ok(None),
2559 };
2560
2561 let mut key = Vec::with_capacity(4 + term.len());
2563 key.extend_from_slice(&field.0.to_le_bytes());
2564 key.extend_from_slice(term);
2565
2566 let term_info = match self.term_dict.get(&key).await? {
2568 Some(info) => info,
2569 None => return Ok(None),
2570 };
2571
2572 let (offset, length) = match term_info.position_info() {
2574 Some((o, l)) => (o, l),
2575 None => return Ok(None),
2576 };
2577
2578 let range = checked_file_range(offset, length, handle.len(), "position list")?;
2582 let slice = handle.slice(range);
2583 let data = slice.read_bytes().await?;
2584
2585 let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
2587
2588 Ok(Some(pos_list))
2589 }
2590
2591 pub fn has_positions(&self, field: Field) -> bool {
2593 if let Some(entry) = self.schema.get_field_entry(field) {
2595 entry.positions.is_some()
2596 } else {
2597 false
2598 }
2599 }
2600}
2601
2602#[cfg(feature = "sync")]
2604impl SegmentReader {
2605 pub fn get_postings_sync(&self, field: Field, term: &[u8]) -> Result<Option<BlockPostingList>> {
2607 let mut key = Vec::with_capacity(4 + term.len());
2609 key.extend_from_slice(&field.0.to_le_bytes());
2610 key.extend_from_slice(term);
2611
2612 let term_info = match self.term_dict.get_sync(&key)? {
2614 Some(info) => info,
2615 None => return Ok(None),
2616 };
2617
2618 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
2620 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
2621 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
2622 posting_list.push(doc_id, tf);
2623 }
2624 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
2625 return Ok(Some(block_list));
2626 }
2627
2628 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
2630 Error::Corruption("TermInfo has neither inline nor external data".to_string())
2631 })?;
2632
2633 let range = checked_file_range(
2634 posting_offset,
2635 posting_len,
2636 self.postings_handle.len(),
2637 "posting",
2638 )?;
2639 let posting_bytes = self.postings_handle.read_bytes_range_sync(range)?;
2640 let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
2641
2642 Ok(Some(block_list))
2643 }
2644
2645 pub fn get_prefix_postings_sync(
2647 &self,
2648 field: Field,
2649 prefix: &[u8],
2650 ) -> Result<Vec<BlockPostingList>> {
2651 if prefix.is_empty() {
2652 return Err(Error::Query("prefix must not be empty".into()));
2653 }
2654 let mut key_prefix = Vec::with_capacity(4 + prefix.len());
2655 key_prefix.extend_from_slice(&field.0.to_le_bytes());
2656 key_prefix.extend_from_slice(prefix);
2657
2658 let (entries, truncated) = self
2659 .term_dict
2660 .prefix_scan_limited_sync(&key_prefix, MAX_PREFIX_TERMS)?;
2661 if truncated {
2662 return Err(Error::Query(format!(
2663 "prefix expands to more than {MAX_PREFIX_TERMS} terms"
2664 )));
2665 }
2666 let posting_count: u64 = entries
2667 .iter()
2668 .map(|(_, term_info)| term_info.doc_freq() as u64)
2669 .sum();
2670 if posting_count > MAX_PREFIX_POSTINGS {
2671 return Err(Error::Query(format!(
2672 "prefix expands to {posting_count} postings (maximum {MAX_PREFIX_POSTINGS})"
2673 )));
2674 }
2675 let mut results = Vec::with_capacity(entries.len());
2676
2677 for (_key, term_info) in entries {
2678 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
2679 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
2680 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
2681 posting_list.push(doc_id, tf);
2682 }
2683 results.push(BlockPostingList::from_posting_list(&posting_list)?);
2684 } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
2685 let range = checked_file_range(
2686 posting_offset,
2687 posting_len,
2688 self.postings_handle.len(),
2689 "prefix posting",
2690 )?;
2691 let posting_bytes = self.postings_handle.read_bytes_range_sync(range)?;
2692 results.push(BlockPostingList::deserialize_zero_copy(posting_bytes)?);
2693 }
2694 }
2695
2696 Ok(results)
2697 }
2698
2699 pub fn get_positions_sync(
2701 &self,
2702 field: Field,
2703 term: &[u8],
2704 ) -> Result<Option<crate::structures::PositionPostingList>> {
2705 let handle = match &self.positions_handle {
2706 Some(h) => h,
2707 None => return Ok(None),
2708 };
2709
2710 let mut key = Vec::with_capacity(4 + term.len());
2712 key.extend_from_slice(&field.0.to_le_bytes());
2713 key.extend_from_slice(term);
2714
2715 let term_info = match self.term_dict.get_sync(&key)? {
2717 Some(info) => info,
2718 None => return Ok(None),
2719 };
2720
2721 let (offset, length) = match term_info.position_info() {
2722 Some((o, l)) => (o, l),
2723 None => return Ok(None),
2724 };
2725
2726 let range = checked_file_range(offset, length, handle.len(), "position list")?;
2727 let slice = handle.slice(range);
2728 let data = slice.read_bytes_sync()?;
2729
2730 let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
2731 Ok(Some(pos_list))
2732 }
2733
2734 pub fn search_dense_vector_sync(
2737 &self,
2738 field: Field,
2739 query: &[f32],
2740 k: usize,
2741 nprobe: usize,
2742 rerank_factor: f32,
2743 combiner: crate::query::MultiValueCombiner,
2744 ) -> Result<Vec<VectorSearchResult>> {
2745 self.search_dense_vector_sync_impl(field, query, k, nprobe, rerank_factor, combiner, None)
2746 }
2747
2748 #[cfg(feature = "sync")]
2749 #[allow(clippy::too_many_arguments)]
2750 pub(crate) fn search_dense_vector_sync_with_probe_cache(
2751 &self,
2752 field: Field,
2753 query: &[f32],
2754 k: usize,
2755 nprobe: usize,
2756 rerank_factor: f32,
2757 combiner: crate::query::MultiValueCombiner,
2758 plan_cache: &DensePlanCache,
2759 ) -> Result<Vec<VectorSearchResult>> {
2760 self.search_dense_vector_sync_impl(
2761 field,
2762 query,
2763 k,
2764 nprobe,
2765 rerank_factor,
2766 combiner,
2767 Some(plan_cache),
2768 )
2769 }
2770
2771 #[cfg(feature = "sync")]
2772 #[allow(clippy::too_many_arguments)]
2773 fn search_dense_vector_sync_impl(
2774 &self,
2775 field: Field,
2776 query: &[f32],
2777 k: usize,
2778 nprobe: usize,
2779 rerank_factor: f32,
2780 combiner: crate::query::MultiValueCombiner,
2781 plan_cache: Option<&DensePlanCache>,
2782 ) -> Result<Vec<VectorSearchResult>> {
2783 let params =
2784 self.validate_dense_search_request(field, query, nprobe, rerank_factor, combiner)?;
2785 let fetch_k = checked_dense_fetch_k(k, rerank_factor)?;
2786 if k == 0 {
2787 return Ok(Vec::new());
2788 }
2789
2790 let ann_index = self.vector_indexes.get(&field.0);
2791 let lazy_flat = self.flat_vectors.get(&field.0);
2792 if ann_index.is_none() && lazy_flat.is_none() {
2793 return Ok(Vec::new());
2794 }
2795
2796 if ann_index.is_some() && lazy_flat.is_none() {
2797 return Err(Error::Corruption(format!(
2798 "dense ANN field {} is missing flat vector storage",
2799 field.0
2800 )));
2801 }
2802
2803 if let Some(flat) = lazy_flat
2804 && flat.dim != params.dim
2805 {
2806 return Err(Error::Corruption(format!(
2807 "dense vector field {} has schema dimension {} but flat storage dimension {}",
2808 field.0, params.dim, flat.dim
2809 )));
2810 }
2811
2812 let results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
2813 match index {
2815 VectorIndex::IvfPq(lazy) => {
2816 let index = lazy.get();
2817 let codebook =
2818 self.trained_vectors
2819 .codebooks
2820 .get(&field.0)
2821 .ok_or_else(|| {
2822 Error::Schema(format!(
2823 "IVF-PQ index requires a global codebook for field {}",
2824 field.0
2825 ))
2826 })?;
2827 let centroids =
2828 self.trained_vectors
2829 .centroids
2830 .get(&field.0)
2831 .ok_or_else(|| {
2832 Error::Schema(format!(
2833 "IVF-PQ index requires coarse centroids for field {}",
2834 field.0
2835 ))
2836 })?;
2837 validate_coarse_centroids(centroids, params.dim)?;
2838 let routing = self
2839 .schema
2840 .get_field_entry(field)
2841 .and_then(|entry| entry.dense_vector_config.as_ref())
2842 .map_or(crate::dsl::IvfRoutingMode::Auto, |config| {
2843 config.ivf_routing
2844 });
2845 validate_ivf_pq_ann(index, centroids, codebook, params.dim, routing)?;
2846 let query_plan = float_query_plan(
2847 centroids,
2848 codebook,
2849 query,
2850 params.nprobe,
2851 routing,
2852 plan_cache.map(|cache| &cache.ivf_pq),
2853 )?;
2854 let flat = lazy_flat.expect("ANN/flat pairing validated above");
2855 index
2856 .search_ivf_pq_distinct(
2857 fetch_k.min(flat.num_docs_with_vectors()),
2858 &query_plan,
2859 )
2860 .map_err(|error| {
2861 Error::Corruption(format!(
2862 "invalid IVF-PQ payload for field {}: {error}",
2863 field.0,
2864 ))
2865 })?
2866 .into_iter()
2867 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
2868 .collect()
2869 }
2870 VectorIndex::Tq { index: lazy, codec } => {
2871 let flat = lazy_flat.expect("ANN/flat pairing validated above");
2872 search_tq_segment(
2873 lazy.get(),
2874 codec,
2875 query,
2876 fetch_k.min(flat.num_docs_with_vectors()),
2877 field,
2878 params.dim,
2879 plan_cache.map(|cache| &cache.tq),
2880 )?
2881 }
2882 VectorIndex::IvfTq { index: lazy, codec } => {
2883 let index = lazy.get();
2884 let centroids =
2885 self.trained_vectors
2886 .centroids
2887 .get(&field.0)
2888 .ok_or_else(|| {
2889 Error::Schema(format!(
2890 "IVF-TQ index requires coarse centroids for field {}",
2891 field.0
2892 ))
2893 })?;
2894 validate_coarse_centroids(centroids, params.dim)?;
2895 let routing = self
2896 .schema
2897 .get_field_entry(field)
2898 .and_then(|entry| entry.dense_vector_config.as_ref())
2899 .map_or(crate::dsl::IvfRoutingMode::Auto, |config| {
2900 config.ivf_routing
2901 });
2902 validate_ivf_tq_ann(index, centroids, codec, params.dim, routing, field)?;
2903 let flat = lazy_flat.expect("ANN/flat pairing validated above");
2904 search_ivf_tq_segment(
2905 index,
2906 centroids,
2907 codec,
2908 query,
2909 fetch_k.min(flat.num_docs_with_vectors()),
2910 field,
2911 params.nprobe,
2912 routing,
2913 plan_cache.map(|cache| &cache.ivf_tq),
2914 )?
2915 }
2916 VectorIndex::BinaryIvf(_) => {
2917 Vec::new()
2919 }
2920 }
2921 } else if let Some(lazy_flat) = lazy_flat {
2922 let dim = lazy_flat.dim;
2924 let n = lazy_flat.num_vectors;
2925 let quant = lazy_flat.quantization;
2926 let batch_len =
2927 bounded_vector_score_batch(lazy_flat.vector_byte_size(), DENSE_SCORE_BATCH);
2928 let mut collector = FlatDocumentCollector::new(fetch_k.min(n), combiner);
2929 let mut scores = vec![0f32; batch_len];
2930
2931 for batch_start in (0..n).step_by(batch_len) {
2932 let batch_count = batch_len.min(n - batch_start);
2933 let batch_bytes = lazy_flat
2934 .read_vectors_batch_sync(batch_start, batch_count)
2935 .map_err(crate::Error::Io)?;
2936 let raw = batch_bytes.as_slice();
2937
2938 Self::score_quantized_batch(
2939 query,
2940 raw,
2941 quant,
2942 dim,
2943 &mut scores[..batch_count],
2944 params.unit_norm,
2945 )?;
2946
2947 for (i, &score) in scores.iter().enumerate().take(batch_count) {
2948 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
2949 collector.push(doc_id, ordinal, score);
2950 }
2951 }
2952
2953 return Ok(collector.into_results());
2954 } else {
2955 return Ok(Vec::new());
2956 };
2957
2958 if ann_index.is_some()
2960 && !results.is_empty()
2961 && let Some(lazy_flat) = lazy_flat
2962 {
2963 return exact_score_dense_candidate_documents_sync(
2964 &results,
2965 lazy_flat,
2966 query,
2967 params.unit_norm,
2968 combiner,
2969 k,
2970 );
2971 }
2972
2973 Ok(combine_grouped_ordinal_results(results, combiner, k))
2974 }
2975
2976 #[cfg(feature = "sync")]
2981 fn search_binary_dense_vector_sync_impl(
2982 &self,
2983 field: Field,
2984 query: &[u8],
2985 k: usize,
2986 combiner: crate::query::MultiValueCombiner,
2987 probe_cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
2988 ) -> Result<Vec<VectorSearchResult>> {
2989 let schema_dim = self.validate_binary_search_request(field, query)?;
2990 combiner.validate().map_err(Error::Query)?;
2991 if k == 0 {
2992 return Ok(Vec::new());
2993 }
2994 let t0 = crate::observe::Timer::start();
2995 if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0) {
2996 let ivf = lazy.get();
2997 let config = self
2998 .schema
2999 .get_field_entry(field)
3000 .and_then(|entry| entry.binary_dense_vector_config.as_ref())
3001 .ok_or_else(|| {
3002 Error::Schema(format!(
3003 "binary IVF field {} has no schema configuration",
3004 field.0
3005 ))
3006 })?;
3007 let quantizer = self
3008 .trained_vectors
3009 .binary_quantizers
3010 .get(&field.0)
3011 .ok_or_else(|| {
3012 Error::Schema(format!(
3013 "global binary IVF field {} has no loaded quantizer",
3014 field.0
3015 ))
3016 })?;
3017 validate_binary_ann(ivf, quantizer, config, schema_dim, field)?;
3018 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
3019 Error::Corruption(format!(
3020 "global binary IVF field {} is missing flat vector storage",
3021 field.0
3022 ))
3023 })?;
3024 let clusters = binary_probe_clusters(
3025 quantizer,
3026 query,
3027 config.nprobe,
3028 config.ivf_routing,
3029 probe_cache,
3030 )?;
3031 let candidate_docs = k.min(flat.num_docs_with_vectors());
3032 let ann_results = if flat.num_vectors == flat.num_docs_with_vectors() {
3033 ivf.search_binary_clusters::<false>(query, candidate_docs, &clusters)
3034 } else {
3035 ivf.search_binary_clusters::<true>(query, candidate_docs, &clusters)
3036 }
3037 .map_err(|error| {
3038 Error::Corruption(format!(
3039 "invalid binary IVF payload for field {}: {error}",
3040 field.0,
3041 ))
3042 })?;
3043 let results = exact_score_binary_candidate_documents_sync(
3044 &ann_results,
3045 flat,
3046 query,
3047 schema_dim,
3048 combiner,
3049 k,
3050 )?;
3051 crate::observe::dense_l1(
3052 self.schema.index_label(),
3053 self.schema.get_field_name(field).unwrap_or("?"),
3054 "global_binary_ivf",
3055 t0.secs(),
3056 results.len(),
3057 );
3058 return Ok(results);
3059 }
3060 let lazy_flat = match self.flat_vectors.get(&field.0) {
3061 Some(f) => f,
3062 None => return Ok(Vec::new()),
3063 };
3064
3065 let dim_bits = lazy_flat.dim;
3066 let byte_len = lazy_flat.vector_byte_size();
3067 let n = lazy_flat.num_vectors;
3068
3069 if dim_bits != schema_dim {
3070 return Err(Error::Corruption(format!(
3071 "binary vector field {} has schema dimension {} but flat storage dimension {}",
3072 field.0, schema_dim, dim_bits
3073 )));
3074 }
3075
3076 if byte_len != query.len() {
3077 return Err(Error::Schema(format!(
3078 "Binary query vector byte length {} != field byte length {}",
3079 query.len(),
3080 byte_len
3081 )));
3082 }
3083
3084 let batch_len = bounded_vector_score_batch(byte_len, BINARY_SCORE_BATCH);
3085 let mut collector = FlatDocumentCollector::new(k, combiner);
3086 let mut scores = vec![0f32; batch_len];
3087
3088 for batch_start in (0..n).step_by(batch_len) {
3089 let batch_count = batch_len.min(n - batch_start);
3090 let batch_bytes = lazy_flat
3091 .read_vectors_batch_sync(batch_start, batch_count)
3092 .map_err(crate::Error::Io)?;
3093 let raw = batch_bytes.as_slice();
3094
3095 crate::structures::simd::batch_hamming_scores(
3096 query,
3097 raw,
3098 byte_len,
3099 dim_bits,
3100 &mut scores[..batch_count],
3101 );
3102
3103 for (i, &score) in scores.iter().enumerate().take(batch_count) {
3104 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
3105 collector.push(doc_id, ordinal, score);
3106 }
3107 }
3108
3109 let results = collector.into_results();
3110
3111 crate::observe::dense_l1(
3112 self.schema.index_label(),
3113 self.schema.get_field_name(field).unwrap_or("?"),
3114 "binary_flat",
3115 t0.secs(),
3116 results.len(),
3117 );
3118 Ok(results)
3119 }
3120
3121 #[cfg(feature = "sync")]
3122 pub fn search_binary_dense_vector_sync(
3123 &self,
3124 field: Field,
3125 query: &[u8],
3126 k: usize,
3127 combiner: crate::query::MultiValueCombiner,
3128 ) -> Result<Vec<VectorSearchResult>> {
3129 self.search_binary_dense_vector_sync_impl(field, query, k, combiner, None)
3130 }
3131
3132 #[cfg(feature = "sync")]
3133 pub(crate) fn search_binary_dense_vector_sync_with_probe_cache(
3134 &self,
3135 field: Field,
3136 query: &[u8],
3137 k: usize,
3138 combiner: crate::query::MultiValueCombiner,
3139 probe_cache: &std::sync::Mutex<Option<crate::structures::IvfProbePlan>>,
3140 ) -> Result<Vec<VectorSearchResult>> {
3141 self.search_binary_dense_vector_sync_impl(field, query, k, combiner, Some(probe_cache))
3142 }
3143}
3144
3145#[cfg(test)]
3146mod dense_search_safety_tests {
3147 use super::*;
3148
3149 #[test]
3150 fn dense_fetch_count_rejects_non_finite_and_unbounded_factors() {
3151 for factor in [
3152 f32::NAN,
3153 f32::INFINITY,
3154 f32::NEG_INFINITY,
3155 0.0,
3156 0.5,
3157 2.01,
3158 MAX_DENSE_RERANK_FACTOR + 1.0,
3159 ] {
3160 assert!(
3161 checked_dense_fetch_k(10, factor).is_err(),
3162 "factor={factor}"
3163 );
3164 }
3165 }
3166
3167 #[test]
3168 fn flat_document_collector_does_not_let_one_multivalue_doc_crowd_out_others() {
3169 let mut collector = FlatDocumentCollector::new(2, crate::query::MultiValueCombiner::Max);
3170 collector.push(1, 0, 1.0);
3171 collector.push(1, 1, 0.9);
3172 collector.push(2, 0, 0.8);
3173
3174 let results = collector.into_results();
3175 assert_eq!(
3176 results
3177 .iter()
3178 .map(|result| result.doc_id)
3179 .collect::<Vec<_>>(),
3180 vec![1, 2]
3181 );
3182 assert_eq!(results[0].ordinals.len(), 2);
3183 }
3184
3185 #[test]
3186 fn flat_document_collector_evicts_by_score_then_doc_id() {
3187 let mut collector = FlatDocumentCollector::new(2, crate::query::MultiValueCombiner::Max);
3188 collector.push(1, 0, 0.5);
3189 collector.push(3, 0, 0.8);
3190 collector.push(2, 0, 0.9);
3191 let results = collector.into_results();
3192 assert_eq!(
3193 results
3194 .iter()
3195 .map(|result| result.doc_id)
3196 .collect::<Vec<_>>(),
3197 vec![2, 3]
3198 );
3199
3200 let mut tied = FlatDocumentCollector::new(1, crate::query::MultiValueCombiner::Max);
3201 tied.push(2, 0, 1.0);
3202 tied.push(1, 0, 1.0);
3203 let results = tied.into_results();
3204 assert_eq!(results[0].doc_id, 1);
3205 }
3206
3207 #[test]
3208 fn dense_fetch_count_rounds_up_and_detects_overflow() {
3209 assert_eq!(checked_dense_fetch_k(3, 1.5).unwrap(), 5);
3210 assert_eq!(checked_dense_fetch_k(10_000, 2.0).unwrap(), 20_000);
3211 assert!(checked_dense_fetch_k(10_001, 2.0).is_err());
3212 assert!(checked_dense_fetch_k(usize::MAX, 2.0).is_err());
3213 }
3214
3215 #[test]
3216 fn file_ranges_reject_overflow_and_truncation() {
3217 assert_eq!(checked_file_range(4, 3, 7, "test").unwrap(), 4..7);
3218 assert!(checked_file_range(u64::MAX, 1, u64::MAX, "test").is_err());
3219 assert!(checked_file_range(5, 3, 7, "test").is_err());
3220 }
3221
3222 #[test]
3223 fn candidate_vector_reads_coalesce_contiguous_values() {
3224 let mut runs = Vec::new();
3225 plan_vector_read_runs(&[3, 4, 5, 9, 12, 13], &mut runs).unwrap();
3226 assert_eq!(runs.len(), 3);
3227 assert!(matches!(
3228 runs.as_slice(),
3229 [
3230 VectorReadRun {
3231 buffer_start: 0,
3232 flat_start: 3,
3233 count: 3,
3234 },
3235 VectorReadRun {
3236 buffer_start: 3,
3237 flat_start: 9,
3238 count: 1,
3239 },
3240 VectorReadRun {
3241 buffer_start: 4,
3242 flat_start: 12,
3243 count: 2,
3244 },
3245 ]
3246 ));
3247 assert!(plan_vector_read_runs(&[3, 3], &mut runs).is_err());
3248 }
3249
3250 #[tokio::test]
3251 async fn multivalue_ann_rerank_streams_past_document_candidate_cap() {
3252 use crate::directories::{FileHandle, OwnedBytes};
3253 use crate::segment::FlatVectorData;
3254
3255 const VALUES: usize = MAX_DENSE_CANDIDATES_PER_SEGMENT + 1;
3256 let mut encoded = Vec::new();
3257 let vectors = vec![1.0f32; VALUES];
3258 let doc_ids: Vec<_> = (0..VALUES).map(|ordinal| (0, ordinal as u16)).collect();
3259 FlatVectorData::serialize_binary_from_flat_streaming(
3260 1,
3261 &vectors,
3262 &doc_ids,
3263 DenseVectorQuantization::F32,
3264 &mut encoded,
3265 )
3266 .unwrap();
3267 let flat = LazyFlatVectorData::open_with_doc_limit(
3268 FileHandle::from_bytes(OwnedBytes::new(encoded)),
3269 Some(1),
3270 )
3271 .await
3272 .unwrap();
3273
3274 let (results, stats) = exact_score_dense_candidate_documents(
3275 &[(0, 0, 0.0)],
3276 &flat,
3277 &[1.0],
3278 false,
3279 crate::query::MultiValueCombiner::Max,
3280 1,
3281 )
3282 .await
3283 .unwrap();
3284 assert_eq!(stats.vector_count, VALUES);
3285 assert_eq!(results.len(), 1);
3286 assert_eq!(results[0].ordinals.len(), VALUES);
3287 assert!((results[0].score - 1.0).abs() < 1e-5);
3288 }
3289}