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
827#[derive(Debug, Default)]
832pub struct DensePlanCache {
833 pub(crate) tq: std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqQueryPlan>>>,
834 pub(crate) ivf_tq: std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqIvfQueryPlan>>>,
835}
836
837fn search_tq_segment(
842 index: &crate::segment::ann_disk::AnnDiskIndex,
843 codec: &crate::structures::TqCodec,
844 query: &[f32],
845 fetch_k: usize,
846 field: Field,
847 dim: usize,
848 plan_cache: Option<&std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqQueryPlan>>>>,
849) -> Result<Vec<RawVectorCandidate>> {
850 validate_tq_ann(index, codec, dim, field)?;
851 let plan = match plan_cache {
852 Some(cache) => {
853 let mut cached = cache
854 .lock()
855 .map_err(|_| Error::Internal("TQ plan cache is poisoned".into()))?;
856 match cached.as_ref() {
857 Some(plan) if plan.fingerprint() == codec.fingerprint() => {
858 std::sync::Arc::clone(plan)
859 }
860 _ => {
861 let plan =
862 std::sync::Arc::new(crate::structures::TqQueryPlan::build(codec, query));
863 *cached = Some(std::sync::Arc::clone(&plan));
864 plan
865 }
866 }
867 }
868 None => std::sync::Arc::new(crate::structures::TqQueryPlan::build(codec, query)),
869 };
870 index.search_tq_distinct(fetch_k, &plan).map_err(|error| {
871 Error::Corruption(format!("invalid TQ payload for field {}: {error}", field.0))
872 })
873}
874
875fn validate_tq_ann(
876 index: &crate::segment::ann_disk::AnnDiskIndex,
877 codec: &crate::structures::TqCodec,
878 dim: usize,
879 field: Field,
880) -> Result<()> {
881 let header = index.header();
882 if header.dim != dim
883 || codec.dim() != dim
884 || header.code_size != codec.code_size()
885 || header.quantizer_version != codec.fingerprint()
886 || header.codebook_version != 0
887 || header.num_clusters != 1
888 {
889 return Err(Error::Corruption(format!(
890 "TQ payload for field {} does not match the codec derived from schema dimension {dim}",
891 field.0,
892 )));
893 }
894 Ok(())
895}
896
897#[allow(clippy::too_many_arguments)]
901fn search_ivf_tq_segment(
902 index: &crate::segment::ann_disk::AnnDiskIndex,
903 centroids: &CoarseCentroids,
904 codec: &crate::structures::TqCodec,
905 query: &[f32],
906 fetch_k: usize,
907 field: Field,
908 nprobe: usize,
909 routing: crate::dsl::IvfRoutingMode,
910 plan_cache: Option<
911 &std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqIvfQueryPlan>>>,
912 >,
913) -> Result<Vec<RawVectorCandidate>> {
914 let effective_nprobe = nprobe.clamp(1, centroids.num_clusters as usize);
915 let request_fingerprint = crate::structures::vector::ivf::routing::float_probe_fingerprint(
916 query,
917 effective_nprobe,
918 routing,
919 );
920 let build = || {
921 std::sync::Arc::new(crate::structures::TqIvfQueryPlan::build(
922 centroids,
923 codec,
924 query,
925 effective_nprobe,
926 routing,
927 ))
928 };
929 let plan = match plan_cache {
930 Some(cache) => {
931 let mut cached = cache
932 .lock()
933 .map_err(|_| Error::Internal("IVF-TQ plan cache is poisoned".into()))?;
934 match cached.as_ref() {
935 Some(plan)
936 if plan.quantizer_version == centroids.version
937 && plan.fingerprint == codec.fingerprint()
938 && plan.request_fingerprint == request_fingerprint
939 && plan.cluster_ids.len() == effective_nprobe =>
940 {
941 std::sync::Arc::clone(plan)
942 }
943 _ => {
944 let plan = build();
945 *cached = Some(std::sync::Arc::clone(&plan));
946 plan
947 }
948 }
949 }
950 None => build(),
951 };
952 index
953 .search_ivf_tq_distinct(fetch_k, &plan)
954 .map_err(|error| {
955 Error::Corruption(format!(
956 "invalid IVF-TQ payload for field {}: {error}",
957 field.0
958 ))
959 })
960}
961
962fn validate_ivf_tq_ann(
963 index: &crate::segment::ann_disk::AnnDiskIndex,
964 centroids: &CoarseCentroids,
965 codec: &crate::structures::TqCodec,
966 dim: usize,
967 routing: crate::dsl::IvfRoutingMode,
968 field: Field,
969) -> Result<()> {
970 let header = index.header();
971 if header.dim != dim
972 || codec.dim() != dim
973 || header.code_size != codec.code_size()
974 || header.num_clusters != centroids.num_clusters
975 || header.quantizer_version != centroids.version
976 || header.codebook_version != codec.fingerprint()
977 || header.routing != routing
978 {
979 return Err(Error::Corruption(format!(
980 "IVF-TQ payload for field {} does not match its quantizer/codec generation",
981 field.0,
982 )));
983 }
984 Ok(())
985}
986
987fn validate_binary_ann(
988 index: &crate::segment::ann_disk::AnnDiskIndex,
989 quantizer: &crate::structures::BinaryCoarseQuantizer,
990 config: &crate::dsl::BinaryDenseVectorConfig,
991 dim: usize,
992 field: Field,
993) -> Result<()> {
994 let header = index.header();
995 if header.dim != dim
996 || header.code_size != config.byte_len()
997 || header.num_clusters != quantizer.num_clusters
998 || header.quantizer_version != quantizer.version
999 || header.codebook_version != 0
1000 || header.routing != config.ivf_routing
1001 || quantizer.dim_bits != dim
1002 {
1003 return Err(Error::Corruption(format!(
1004 "binary IVF field {} does not match its quantizer/schema generation",
1005 field.0,
1006 )));
1007 }
1008 Ok(())
1009}
1010
1011fn binary_probe_clusters(
1012 quantizer: &crate::structures::BinaryCoarseQuantizer,
1013 query: &[u8],
1014 nprobe: usize,
1015 routing: crate::dsl::IvfRoutingMode,
1016 cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
1017) -> Result<std::sync::Arc<[u32]>> {
1018 let effective_nprobe = nprobe.clamp(1, quantizer.num_clusters as usize);
1019 let request_fingerprint = crate::structures::vector::ivf::routing::binary_probe_fingerprint(
1020 query,
1021 effective_nprobe,
1022 routing,
1023 );
1024 if let Some(cache) = cache {
1025 let mut cached = cache
1026 .lock()
1027 .map_err(|_| Error::Internal("binary IVF probe cache is poisoned".into()))?;
1028 if let Some(plan) = cached.as_ref()
1029 && plan.quantizer_version == quantizer.version
1030 && plan.request_fingerprint == request_fingerprint
1031 && plan.cluster_ids.len() == effective_nprobe
1032 {
1033 return Ok(std::sync::Arc::clone(&plan.cluster_ids));
1034 }
1035 let plan = quantizer.probe(query, effective_nprobe, routing);
1036 let clusters = std::sync::Arc::clone(&plan.cluster_ids);
1037 *cached = Some(plan);
1038 return Ok(clusters);
1039 }
1040 Ok(quantizer
1041 .probe(query, effective_nprobe, routing)
1042 .cluster_ids)
1043}
1044
1045pub struct SegmentReader {
1051 meta: SegmentMeta,
1052 term_dict: Arc<AsyncSSTableReader<TermInfo>>,
1054 postings_handle: FileHandle,
1056 store: Arc<AsyncStoreReader>,
1058 schema: Arc<Schema>,
1059 vector_indexes: FxHashMap<u32, VectorIndex>,
1061 flat_vectors: FxHashMap<u32, LazyFlatVectorData>,
1063 dense_file_backed_bytes: u64,
1065 trained_vectors: Arc<crate::segment::TrainedVectorStructures>,
1067 sparse_indexes: FxHashMap<u32, SparseIndex>,
1069 bmp_indexes: FxHashMap<u32, BmpIndex>,
1071 sparse_file_backed_bytes: u64,
1073 positions_handle: Option<FileHandle>,
1075 fast_fields: FxHashMap<u32, crate::structures::fast_field::FastFieldReader>,
1077 #[cfg(feature = "native")]
1079 dense_pin_report: crate::segment::pin::PinReport,
1080 #[cfg(feature = "native")]
1082 sparse_pin_report: crate::segment::pin::PinReport,
1083}
1084
1085impl SegmentReader {
1086 pub async fn open<D: Directory>(
1088 dir: &D,
1089 segment_id: SegmentId,
1090 schema: Arc<Schema>,
1091 term_cache_blocks: usize,
1092 ) -> Result<Self> {
1093 Self::open_with_store_cache(
1094 dir,
1095 segment_id,
1096 schema,
1097 term_cache_blocks,
1098 dir as *const D as usize,
1099 Arc::new(super::SharedStoreCache::new(0)),
1100 )
1101 .await
1102 }
1103
1104 pub(crate) async fn open_with_store_cache<D: Directory>(
1106 dir: &D,
1107 segment_id: SegmentId,
1108 schema: Arc<Schema>,
1109 term_cache_blocks: usize,
1110 store_cache_directory_namespace: usize,
1111 store_cache: Arc<super::SharedStoreCache>,
1112 ) -> Result<Self> {
1113 let files = SegmentFiles::new(segment_id.0);
1114
1115 let meta_slice = dir.open_read(&files.meta).await?;
1117 let meta_bytes = meta_slice.read_bytes().await?;
1118 let meta = SegmentMeta::deserialize(meta_bytes.as_slice())?;
1119 debug_assert_eq!(meta.id, segment_id.0);
1120
1121 let term_dict_handle = dir.open_lazy(&files.term_dict).await?;
1123 let term_dict = AsyncSSTableReader::open(term_dict_handle, term_cache_blocks).await?;
1124
1125 let postings_handle = dir.open_lazy(&files.postings).await?;
1127
1128 let store_handle = dir.open_lazy(&files.store).await?;
1130 let store = AsyncStoreReader::open(
1131 store_handle,
1132 store_cache_directory_namespace,
1133 segment_id.0,
1134 store_cache,
1135 )
1136 .await?;
1137
1138 let vectors_data = loader::load_vectors_file(dir, &files, &schema, meta.num_docs).await?;
1140 let dense_file_backed_bytes = vectors_data.file_backed_bytes;
1141 let vector_indexes = vectors_data.indexes;
1142 let flat_vectors = vectors_data.flat_vectors;
1143
1144 #[cfg(feature = "native")]
1149 for (field_id, lazy_flat) in &flat_vectors {
1150 if vector_indexes.contains_key(field_id) {
1151 lazy_flat.advise_random_access();
1152 }
1153 }
1154
1155 let sparse_data = loader::load_sparse_file(dir, &files, meta.num_docs, &schema).await?;
1157 let sparse_file_backed_bytes = sparse_data.file_backed_bytes;
1158 let sparse_indexes = sparse_data.maxscore_indexes;
1159 let bmp_indexes = sparse_data.bmp_indexes;
1160
1161 let positions_handle = loader::open_positions_file(dir, &files, &schema).await?;
1163
1164 let fast_fields = loader::load_fast_fields_file(dir, &files, &schema).await?;
1166
1167 {
1169 let mut parts = vec![format!(
1170 "[segment] loaded {:016x}: docs={}",
1171 segment_id.0, meta.num_docs
1172 )];
1173 if !vector_indexes.is_empty() || !flat_vectors.is_empty() {
1174 parts.push(format!(
1175 "dense vectors: {} ANN + {} flat fields",
1176 vector_indexes.len(),
1177 flat_vectors.len()
1178 ));
1179 }
1180 for (field_id, idx) in &sparse_indexes {
1181 parts.push(format!(
1182 "sparse vector field {}: {} dims, ~{}",
1183 field_id,
1184 idx.num_dimensions(),
1185 crate::format_bytes(idx.num_dimensions() as u64 * 24)
1186 ));
1187 }
1188 for (field_id, idx) in &bmp_indexes {
1189 parts.push(format!(
1190 "bmp field {}: {} dims, {} blocks",
1191 field_id,
1192 idx.dims(),
1193 idx.num_blocks
1194 ));
1195 }
1196 if !fast_fields.is_empty() {
1197 parts.push(format!("fast: {} fields", fast_fields.len()));
1198 }
1199 log::debug!("{}", parts.join(", "));
1200 }
1201
1202 #[allow(unused_mut)]
1203 let mut reader = Self {
1204 meta,
1205 term_dict: Arc::new(term_dict),
1206 postings_handle,
1207 store: Arc::new(store),
1208 schema,
1209 vector_indexes,
1210 flat_vectors,
1211 dense_file_backed_bytes,
1212 trained_vectors: Arc::new(crate::segment::TrainedVectorStructures::default()),
1213 sparse_indexes,
1214 bmp_indexes,
1215 sparse_file_backed_bytes,
1216 positions_handle,
1217 fast_fields,
1218 #[cfg(feature = "native")]
1219 dense_pin_report: Default::default(),
1220 #[cfg(feature = "native")]
1221 sparse_pin_report: Default::default(),
1222 };
1223
1224 #[cfg(feature = "native")]
1226 reader.apply_pin_policy(&crate::segment::pin::pin_policy().to_owned());
1227
1228 Ok(reader)
1229 }
1230
1231 #[cfg(feature = "native")]
1241 pub(crate) fn apply_pin_policy(&mut self, policy: &crate::segment::pin::PinPolicy) {
1242 use crate::segment::pin::PinReport;
1243
1244 if !policy.is_enabled() {
1245 return;
1246 }
1247 let mut remaining = policy.budget_bytes;
1248 let mut dense_report = PinReport::default();
1249 let mut sparse_report = PinReport::default();
1250
1251 for index in self.vector_indexes.values_mut() {
1253 index.pin_lookup_directory(policy.mode, &mut remaining, &mut dense_report);
1254 }
1255 for bmp in self.bmp_indexes.values_mut() {
1257 bmp.pin_block_starts(policy.mode, &mut remaining, &mut sparse_report);
1258 }
1259 for sparse in self.sparse_indexes.values_mut() {
1261 sparse.pin_skip_section(policy.mode, &mut remaining, &mut sparse_report);
1262 }
1263 for flat in self.flat_vectors.values_mut() {
1265 flat.pin_doc_ids(policy.mode, &mut remaining, &mut dense_report);
1266 }
1267 for bmp in self.bmp_indexes.values_mut() {
1268 bmp.pin_doc_maps(policy.mode, &mut remaining, &mut sparse_report);
1269 }
1270 for bmp in self.bmp_indexes.values_mut() {
1272 bmp.pin_query_hierarchy(policy.mode, &mut remaining, &mut sparse_report);
1273 }
1274
1275 let report = PinReport {
1276 intended_bytes: dense_report
1277 .intended_bytes
1278 .saturating_add(sparse_report.intended_bytes),
1279 pinned_bytes: dense_report
1280 .pinned_bytes
1281 .saturating_add(sparse_report.pinned_bytes),
1282 skipped_budget_bytes: dense_report
1283 .skipped_budget_bytes
1284 .saturating_add(sparse_report.skipped_budget_bytes),
1285 failed_bytes: dense_report
1286 .failed_bytes
1287 .saturating_add(sparse_report.failed_bytes),
1288 heap_copy_bytes: dense_report
1289 .heap_copy_bytes
1290 .saturating_add(sparse_report.heap_copy_bytes),
1291 };
1292 if report.skipped_budget_bytes > 0 || report.failed_bytes > 0 {
1293 log::warn!(
1294 "[pin] segment {:016x}: pinned {}/{} (budget skipped {}, mlock failed {}) — \
1295 raise HERMES_PIN_METADATA_BUDGET_MB or RLIMIT_MEMLOCK for full coverage",
1296 self.meta.id,
1297 crate::format_bytes(report.pinned_bytes),
1298 crate::format_bytes(report.intended_bytes),
1299 crate::format_bytes(report.skipped_budget_bytes),
1300 crate::format_bytes(report.failed_bytes),
1301 );
1302 } else if report.pinned_bytes > 0 {
1303 log::info!(
1304 "[pin] segment {:016x}: pinned {} of hot metadata ({:?})",
1305 self.meta.id,
1306 crate::format_bytes(report.pinned_bytes),
1307 policy.mode,
1308 );
1309 }
1310 self.dense_pin_report = dense_report;
1311 self.sparse_pin_report = sparse_report;
1312 }
1313
1314 pub fn meta(&self) -> &SegmentMeta {
1320 &self.meta
1321 }
1322
1323 pub fn num_docs(&self) -> u32 {
1324 self.meta.num_docs
1325 }
1326
1327 pub fn avg_field_len(&self, field: Field) -> f32 {
1329 self.meta.avg_field_len(field)
1330 }
1331
1332 pub fn schema(&self) -> &Schema {
1333 &self.schema
1334 }
1335
1336 pub fn sparse_indexes(&self) -> &FxHashMap<u32, SparseIndex> {
1338 &self.sparse_indexes
1339 }
1340
1341 pub fn sparse_index(&self, field: Field) -> Option<&SparseIndex> {
1343 self.sparse_indexes.get(&field.0)
1344 }
1345
1346 pub fn bmp_index(&self, field: Field) -> Option<&BmpIndex> {
1348 self.bmp_indexes.get(&field.0)
1349 }
1350
1351 pub fn bmp_indexes(&self) -> &FxHashMap<u32, BmpIndex> {
1353 &self.bmp_indexes
1354 }
1355
1356 pub fn vector_indexes(&self) -> &FxHashMap<u32, VectorIndex> {
1358 &self.vector_indexes
1359 }
1360
1361 pub fn flat_vectors(&self) -> &FxHashMap<u32, LazyFlatVectorData> {
1363 &self.flat_vectors
1364 }
1365
1366 pub fn fast_field(
1368 &self,
1369 field_id: u32,
1370 ) -> Option<&crate::structures::fast_field::FastFieldReader> {
1371 self.fast_fields.get(&field_id)
1372 }
1373
1374 pub fn fast_fields(&self) -> &FxHashMap<u32, crate::structures::fast_field::FastFieldReader> {
1376 &self.fast_fields
1377 }
1378
1379 pub fn term_dict_stats(&self) -> SSTableStats {
1381 self.term_dict.stats()
1382 }
1383
1384 pub fn memory_stats(&self) -> SegmentMemoryStats {
1386 let term_dict_stats = self.term_dict.stats();
1387
1388 let term_dict_cache_bytes = self.term_dict.cached_bytes();
1392 let store_cache_bytes = self.store.cached_bytes();
1393
1394 let sparse_heap_bytes: usize = self
1397 .sparse_indexes
1398 .values()
1399 .map(|s| s.estimated_heap_bytes())
1400 .sum::<usize>()
1401 + self
1402 .bmp_indexes
1403 .values()
1404 .map(|b| b.estimated_heap_bytes())
1405 .sum::<usize>();
1406
1407 let dense_heap_bytes: usize = self
1410 .vector_indexes
1411 .values()
1412 .map(|v| v.estimated_heap_bytes())
1413 .sum::<usize>()
1414 + self
1415 .flat_vectors
1416 .values()
1417 .map(LazyFlatVectorData::estimated_heap_bytes)
1418 .sum::<usize>();
1419
1420 #[cfg(feature = "native")]
1421 let (sparse_heap_bytes, dense_heap_bytes) = (
1422 sparse_heap_bytes.saturating_add(
1423 usize::try_from(self.sparse_pin_report.heap_copy_bytes).unwrap_or(usize::MAX),
1424 ),
1425 dense_heap_bytes.saturating_add(
1426 usize::try_from(self.dense_pin_report.heap_copy_bytes).unwrap_or(usize::MAX),
1427 ),
1428 );
1429
1430 #[cfg(feature = "native")]
1431 let (
1432 sparse_pinned_metadata_bytes,
1433 sparse_pin_intended_bytes,
1434 dense_pinned_metadata_bytes,
1435 dense_pin_intended_bytes,
1436 ) = (
1437 self.sparse_pin_report.pinned_bytes,
1438 self.sparse_pin_report.intended_bytes,
1439 self.dense_pin_report.pinned_bytes,
1440 self.dense_pin_report.intended_bytes,
1441 );
1442 #[cfg(not(feature = "native"))]
1443 let (
1444 sparse_pinned_metadata_bytes,
1445 sparse_pin_intended_bytes,
1446 dense_pinned_metadata_bytes,
1447 dense_pin_intended_bytes,
1448 ) = (0u64, 0u64, 0u64, 0u64);
1449
1450 let pinned_metadata_bytes =
1451 sparse_pinned_metadata_bytes.saturating_add(dense_pinned_metadata_bytes);
1452 let pin_intended_bytes = sparse_pin_intended_bytes.saturating_add(dense_pin_intended_bytes);
1453
1454 SegmentMemoryStats {
1455 segment_id: self.meta.id,
1456 num_docs: self.meta.num_docs,
1457 term_dict_cache_bytes,
1458 store_cache_bytes,
1459 sparse_heap_bytes,
1460 dense_heap_bytes,
1461 term_bloom_file_bytes: term_dict_stats.bloom_filter_size as u64,
1462 sparse_file_backed_bytes: self.sparse_file_backed_bytes,
1463 dense_file_backed_bytes: self.dense_file_backed_bytes,
1464 pinned_metadata_bytes,
1465 pin_intended_bytes,
1466 sparse_pinned_metadata_bytes,
1467 sparse_pin_intended_bytes,
1468 dense_pinned_metadata_bytes,
1469 dense_pin_intended_bytes,
1470 }
1471 }
1472
1473 pub async fn get_postings(
1478 &self,
1479 field: Field,
1480 term: &[u8],
1481 ) -> Result<Option<BlockPostingList>> {
1482 log::debug!(
1483 "SegmentReader::get_postings field={} term_len={}",
1484 field.0,
1485 term.len()
1486 );
1487
1488 let mut key = Vec::with_capacity(4 + term.len());
1490 key.extend_from_slice(&field.0.to_le_bytes());
1491 key.extend_from_slice(term);
1492
1493 let term_info = match self.term_dict.get(&key).await? {
1495 Some(info) => {
1496 log::debug!("SegmentReader::get_postings found term_info");
1497 info
1498 }
1499 None => {
1500 log::debug!("SegmentReader::get_postings term not found");
1501 return Ok(None);
1502 }
1503 };
1504
1505 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
1507 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
1509 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
1510 posting_list.push(doc_id, tf);
1511 }
1512 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
1513 return Ok(Some(block_list));
1514 }
1515
1516 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
1518 Error::Corruption("TermInfo has neither inline nor external data".to_string())
1519 })?;
1520
1521 let range = checked_file_range(
1522 posting_offset,
1523 posting_len,
1524 self.postings_handle.len(),
1525 "posting",
1526 )?;
1527 let posting_bytes = self.postings_handle.read_bytes_range(range).await?;
1528 let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
1529
1530 Ok(Some(block_list))
1531 }
1532
1533 pub async fn get_prefix_postings(
1535 &self,
1536 field: Field,
1537 prefix: &[u8],
1538 ) -> Result<Vec<BlockPostingList>> {
1539 if prefix.is_empty() {
1540 return Err(Error::Query("prefix must not be empty".into()));
1541 }
1542 let mut key_prefix = Vec::with_capacity(4 + prefix.len());
1544 key_prefix.extend_from_slice(&field.0.to_le_bytes());
1545 key_prefix.extend_from_slice(prefix);
1546
1547 let (entries, truncated) = self
1548 .term_dict
1549 .prefix_scan_limited(&key_prefix, MAX_PREFIX_TERMS)
1550 .await?;
1551 if truncated {
1552 return Err(Error::Query(format!(
1553 "prefix expands to more than {MAX_PREFIX_TERMS} terms"
1554 )));
1555 }
1556 let posting_count: u64 = entries
1557 .iter()
1558 .map(|(_, term_info)| term_info.doc_freq() as u64)
1559 .sum();
1560 if posting_count > MAX_PREFIX_POSTINGS {
1561 return Err(Error::Query(format!(
1562 "prefix expands to {posting_count} postings (maximum {MAX_PREFIX_POSTINGS})"
1563 )));
1564 }
1565 let mut results = Vec::with_capacity(entries.len());
1566
1567 for (_key, term_info) in entries {
1568 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
1569 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
1570 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
1571 posting_list.push(doc_id, tf);
1572 }
1573 results.push(BlockPostingList::from_posting_list(&posting_list)?);
1574 } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
1575 let range = checked_file_range(
1576 posting_offset,
1577 posting_len,
1578 self.postings_handle.len(),
1579 "prefix posting",
1580 )?;
1581 let posting_bytes = self.postings_handle.read_bytes_range(range).await?;
1582 results.push(BlockPostingList::deserialize_zero_copy(posting_bytes)?);
1583 }
1584 }
1585
1586 Ok(results)
1587 }
1588
1589 pub async fn doc(&self, local_doc_id: DocId) -> Result<Option<Document>> {
1594 self.doc_with_fields(local_doc_id, None).await
1595 }
1596
1597 pub async fn doc_with_fields(
1603 &self,
1604 local_doc_id: DocId,
1605 fields: Option<&rustc_hash::FxHashSet<u32>>,
1606 ) -> Result<Option<Document>> {
1607 let mut doc = match fields {
1608 Some(set) => {
1609 let field_ids: Vec<u32> = set.iter().copied().collect();
1610 match self
1611 .store
1612 .get_fields(local_doc_id, &self.schema, &field_ids)
1613 .await
1614 {
1615 Ok(Some(d)) => d,
1616 Ok(None) => return Ok(None),
1617 Err(e) => return Err(Error::from(e)),
1618 }
1619 }
1620 None => match self.store.get(local_doc_id, &self.schema).await {
1621 Ok(Some(d)) => d,
1622 Ok(None) => return Ok(None),
1623 Err(e) => return Err(Error::from(e)),
1624 },
1625 };
1626
1627 for (&field_id, lazy_flat) in &self.flat_vectors {
1629 if let Some(set) = fields
1631 && !set.contains(&field_id)
1632 {
1633 continue;
1634 }
1635
1636 let is_binary = lazy_flat.quantization == DenseVectorQuantization::Binary;
1637 let (start, entries) = lazy_flat.flat_indexes_for_doc(local_doc_id);
1638 for (j, &(_doc_id, _ordinal)) in entries.iter().enumerate() {
1639 let flat_idx = start + j;
1640 if is_binary {
1641 let vbs = lazy_flat.vector_byte_size();
1642 let mut raw = vec![0u8; vbs];
1643 match lazy_flat.read_vector_raw_into(flat_idx, &mut raw).await {
1644 Ok(()) => {
1645 doc.add_binary_dense_vector(Field(field_id), raw);
1646 }
1647 Err(e) => {
1648 log::warn!(
1649 "Failed to hydrate binary dense vector field {}: {}",
1650 field_id,
1651 e
1652 );
1653 }
1654 }
1655 } else {
1656 match lazy_flat.get_vector(flat_idx).await {
1657 Ok(vec) => {
1658 doc.add_dense_vector(Field(field_id), vec);
1659 }
1660 Err(e) => {
1661 log::warn!("Failed to hydrate dense vector field {}: {}", field_id, e);
1662 }
1663 }
1664 }
1665 }
1666 }
1667
1668 Ok(Some(doc))
1669 }
1670
1671 pub async fn prefetch_terms(
1673 &self,
1674 field: Field,
1675 start_term: &[u8],
1676 end_term: &[u8],
1677 ) -> Result<()> {
1678 let mut start_key = Vec::with_capacity(4 + start_term.len());
1679 start_key.extend_from_slice(&field.0.to_le_bytes());
1680 start_key.extend_from_slice(start_term);
1681
1682 let mut end_key = Vec::with_capacity(4 + end_term.len());
1683 end_key.extend_from_slice(&field.0.to_le_bytes());
1684 end_key.extend_from_slice(end_term);
1685
1686 self.term_dict.prefetch_range(&start_key, &end_key).await?;
1687 Ok(())
1688 }
1689
1690 pub fn store_has_dict(&self) -> bool {
1692 self.store.has_dict()
1693 }
1694
1695 pub fn store(&self) -> &super::store::AsyncStoreReader {
1697 &self.store
1698 }
1699
1700 pub fn store_raw_blocks(&self) -> Vec<RawStoreBlock> {
1702 self.store.raw_blocks()
1703 }
1704
1705 pub fn store_data_slice(&self) -> &FileHandle {
1707 self.store.data_slice()
1708 }
1709
1710 pub async fn all_terms(&self) -> Result<Vec<(Vec<u8>, TermInfo)>> {
1712 self.term_dict.all_entries().await.map_err(Error::from)
1713 }
1714
1715 pub async fn all_terms_with_stats(&self) -> Result<Vec<(Field, String, u32)>> {
1720 let entries = self.term_dict.all_entries().await?;
1721 let mut result = Vec::with_capacity(entries.len());
1722
1723 for (key, term_info) in entries {
1724 if key.len() > 4 {
1726 let field_id = u32::from_le_bytes([key[0], key[1], key[2], key[3]]);
1727 let term_bytes = &key[4..];
1728 if let Ok(term_str) = std::str::from_utf8(term_bytes) {
1729 result.push((Field(field_id), term_str.to_string(), term_info.doc_freq()));
1730 }
1731 }
1732 }
1733
1734 Ok(result)
1735 }
1736
1737 pub fn term_dict_iter(&self) -> crate::structures::AsyncSSTableIterator<'_, TermInfo> {
1739 self.term_dict.iter()
1740 }
1741
1742 pub async fn prefetch_term_dict(&self) -> crate::Result<()> {
1746 self.term_dict
1747 .prefetch_all_data_bulk()
1748 .await
1749 .map_err(crate::Error::from)
1750 }
1751
1752 pub async fn read_postings(&self, offset: u64, len: u64) -> Result<Vec<u8>> {
1754 let range = checked_file_range(offset, len, self.postings_handle.len(), "posting")?;
1755 let bytes = self.postings_handle.read_bytes_range(range).await?;
1756 Ok(bytes.to_vec())
1757 }
1758
1759 pub async fn read_position_bytes(&self, offset: u64, len: u64) -> Result<Option<Vec<u8>>> {
1761 let handle = match &self.positions_handle {
1762 Some(h) => h,
1763 None => return Ok(None),
1764 };
1765 let range = checked_file_range(offset, len, handle.len(), "position")?;
1766 let bytes = handle.read_bytes_range(range).await?;
1767 Ok(Some(bytes.to_vec()))
1768 }
1769
1770 pub fn has_positions_file(&self) -> bool {
1772 self.positions_handle.is_some()
1773 }
1774
1775 fn validate_dense_search_request(
1779 &self,
1780 field: Field,
1781 query: &[f32],
1782 nprobe: usize,
1783 rerank_factor: f32,
1784 combiner: crate::query::MultiValueCombiner,
1785 ) -> Result<DenseSearchParams> {
1786 let entry = self
1787 .schema
1788 .get_field_entry(field)
1789 .ok_or_else(|| Error::FieldNotFound(field.0.to_string()))?;
1790 if entry.field_type != crate::dsl::FieldType::DenseVector {
1791 return Err(Error::InvalidFieldType {
1792 expected: "dense_vector".to_string(),
1793 got: format!("{:?}", entry.field_type),
1794 });
1795 }
1796 let config = entry.dense_vector_config.as_ref().ok_or_else(|| {
1797 Error::Schema(format!(
1798 "dense vector field '{}' has no dense vector configuration",
1799 entry.name
1800 ))
1801 })?;
1802
1803 if query.is_empty() {
1804 return Err(Error::Query(format!(
1805 "dense query vector for field '{}' must not be empty",
1806 entry.name
1807 )));
1808 }
1809 if query.len() != config.dim {
1810 return Err(Error::Query(format!(
1811 "dense query vector dimension {} does not match field '{}' dimension {}",
1812 query.len(),
1813 entry.name,
1814 config.dim
1815 )));
1816 }
1817 if let Some((index, value)) = query
1818 .iter()
1819 .enumerate()
1820 .find(|(_, value)| !value.is_finite())
1821 {
1822 return Err(Error::Query(format!(
1823 "dense query vector for field '{}' contains non-finite value {value} at index {index}",
1824 entry.name
1825 )));
1826 }
1827
1828 let nprobe = match (nprobe, config.nprobe) {
1831 (0, 0) => 32,
1832 (0, schema_nprobe) => schema_nprobe,
1833 (query_nprobe, _) => query_nprobe,
1834 };
1835 if nprobe > MAX_DENSE_NPROBE {
1836 return Err(Error::Query(format!(
1837 "dense nprobe must be at most {MAX_DENSE_NPROBE}, got {nprobe}"
1838 )));
1839 }
1840
1841 checked_dense_fetch_k(0, rerank_factor)?;
1844 combiner.validate().map_err(Error::Query)?;
1845
1846 Ok(DenseSearchParams {
1847 dim: config.dim,
1848 nprobe,
1849 unit_norm: config.unit_norm,
1850 })
1851 }
1852
1853 fn validate_binary_search_request(&self, field: Field, query: &[u8]) -> Result<usize> {
1854 let entry = self
1855 .schema
1856 .get_field_entry(field)
1857 .ok_or_else(|| Error::FieldNotFound(field.0.to_string()))?;
1858 if entry.field_type != crate::dsl::FieldType::BinaryDenseVector {
1859 return Err(Error::InvalidFieldType {
1860 expected: "binary_dense_vector".to_string(),
1861 got: format!("{:?}", entry.field_type),
1862 });
1863 }
1864 let config = entry.binary_dense_vector_config.as_ref().ok_or_else(|| {
1865 Error::Schema(format!(
1866 "binary dense vector field '{}' has no configuration",
1867 entry.name
1868 ))
1869 })?;
1870 if config.dim == 0 || !config.dim.is_multiple_of(8) {
1871 return Err(Error::Schema(format!(
1872 "binary dense vector field '{}' has invalid dimension {}",
1873 entry.name, config.dim
1874 )));
1875 }
1876 if query.len() != config.byte_len() {
1877 return Err(Error::Query(format!(
1878 "binary query byte length {} does not match field '{}' byte length {}",
1879 query.len(),
1880 entry.name,
1881 config.byte_len()
1882 )));
1883 }
1884 Ok(config.dim)
1885 }
1886
1887 fn score_quantized_batch(
1893 query: &[f32],
1894 raw: &[u8],
1895 quant: crate::dsl::DenseVectorQuantization,
1896 dim: usize,
1897 scores: &mut [f32],
1898 unit_norm: bool,
1899 ) -> Result<()> {
1900 use crate::dsl::DenseVectorQuantization;
1901 use crate::structures::simd;
1902
1903 if query.len() != dim {
1904 return Err(Error::Query(format!(
1905 "dense SIMD query dimension {} does not match vector dimension {dim}",
1906 query.len()
1907 )));
1908 }
1909 let element_size = match quant {
1910 DenseVectorQuantization::F32 => std::mem::size_of::<f32>(),
1911 DenseVectorQuantization::F16 => std::mem::size_of::<u16>(),
1912 DenseVectorQuantization::UInt8 => 1,
1913 DenseVectorQuantization::Binary => {
1914 return Err(Error::InvalidFieldType {
1915 expected: "non-binary dense vector".to_string(),
1916 got: "binary dense vector".to_string(),
1917 });
1918 }
1919 };
1920 let required_bytes = scores
1921 .len()
1922 .checked_mul(dim)
1923 .and_then(|elements| elements.checked_mul(element_size))
1924 .ok_or_else(|| Error::Corruption("dense vector batch byte length overflow".into()))?;
1925 if raw.len() < required_bytes {
1926 return Err(Error::Corruption(format!(
1927 "dense vector batch is truncated: need {required_bytes} bytes, got {}",
1928 raw.len()
1929 )));
1930 }
1931 if quant == DenseVectorQuantization::F16
1932 && required_bytes > 0
1933 && !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>())
1934 {
1935 return Err(Error::Corruption(
1936 "f16 vector data is not 2-byte aligned".to_string(),
1937 ));
1938 }
1939
1940 match (quant, unit_norm) {
1941 (DenseVectorQuantization::F32, false) => {
1942 let num_floats = scores.len() * dim;
1943 if !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()) {
1944 return Err(Error::Corruption(
1945 "f32 vector data is not 4-byte aligned".to_string(),
1946 ));
1947 }
1948 let vectors: &[f32] =
1949 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
1950 simd::batch_cosine_scores(query, vectors, dim, scores);
1951 }
1952 (DenseVectorQuantization::F32, true) => {
1953 let num_floats = scores.len() * dim;
1954 if !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()) {
1955 return Err(Error::Corruption(
1956 "f32 vector data is not 4-byte aligned".to_string(),
1957 ));
1958 }
1959 let vectors: &[f32] =
1960 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
1961 simd::batch_dot_scores(query, vectors, dim, scores);
1962 }
1963 (DenseVectorQuantization::F16, false) => {
1964 simd::batch_cosine_scores_f16(query, raw, dim, scores);
1965 }
1966 (DenseVectorQuantization::F16, true) => {
1967 simd::batch_dot_scores_f16(query, raw, dim, scores);
1968 }
1969 (DenseVectorQuantization::UInt8, false) => {
1970 simd::batch_cosine_scores_u8(query, raw, dim, scores);
1971 }
1972 (DenseVectorQuantization::UInt8, true) => {
1973 simd::batch_dot_scores_u8(query, raw, dim, scores);
1974 }
1975 (DenseVectorQuantization::Binary, _) => unreachable!("validated above"),
1976 }
1977 Ok(())
1978 }
1979
1980 pub async fn search_dense_vector(
1986 &self,
1987 field: Field,
1988 query: &[f32],
1989 k: usize,
1990 nprobe: usize,
1991 rerank_factor: f32,
1992 combiner: crate::query::MultiValueCombiner,
1993 ) -> Result<Vec<VectorSearchResult>> {
1994 self.search_dense_vector_impl(field, query, k, nprobe, rerank_factor, combiner, None)
1995 .await
1996 }
1997
1998 #[allow(clippy::too_many_arguments)]
1999 pub(crate) async fn search_dense_vector_with_probe_cache(
2000 &self,
2001 field: Field,
2002 query: &[f32],
2003 k: usize,
2004 nprobe: usize,
2005 rerank_factor: f32,
2006 combiner: crate::query::MultiValueCombiner,
2007 plan_cache: &DensePlanCache,
2008 ) -> Result<Vec<VectorSearchResult>> {
2009 self.search_dense_vector_impl(
2010 field,
2011 query,
2012 k,
2013 nprobe,
2014 rerank_factor,
2015 combiner,
2016 Some(plan_cache),
2017 )
2018 .await
2019 }
2020
2021 #[allow(clippy::too_many_arguments)]
2022 async fn search_dense_vector_impl(
2023 &self,
2024 field: Field,
2025 query: &[f32],
2026 k: usize,
2027 nprobe: usize,
2028 rerank_factor: f32,
2029 combiner: crate::query::MultiValueCombiner,
2030 plan_cache: Option<&DensePlanCache>,
2031 ) -> Result<Vec<VectorSearchResult>> {
2032 let params =
2033 self.validate_dense_search_request(field, query, nprobe, rerank_factor, combiner)?;
2034 let fetch_k = checked_dense_fetch_k(k, rerank_factor)?;
2035 if k == 0 {
2036 return Ok(Vec::new());
2037 }
2038
2039 let ann_index = self.vector_indexes.get(&field.0);
2040 let lazy_flat = self.flat_vectors.get(&field.0);
2041 if ann_index.is_none() && lazy_flat.is_none() {
2043 return Ok(Vec::new());
2044 }
2045
2046 if ann_index.is_some() && lazy_flat.is_none() {
2047 return Err(Error::Corruption(format!(
2048 "dense ANN field {} is missing flat vector storage",
2049 field.0
2050 )));
2051 }
2052
2053 if let Some(flat) = lazy_flat
2054 && flat.dim != params.dim
2055 {
2056 return Err(Error::Corruption(format!(
2057 "dense vector field {} has schema dimension {} but flat storage dimension {}",
2058 field.0, params.dim, flat.dim
2059 )));
2060 }
2061
2062 let t0 = std::time::Instant::now();
2064 let mut flat_results = None;
2065 let results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
2066 match index {
2068 VectorIndex::Tq { index: lazy, codec } => {
2069 let flat = lazy_flat.expect("ANN/flat pairing validated above");
2070 search_tq_segment(
2072 lazy.get(),
2073 codec,
2074 query,
2075 fetch_k.min(flat.num_docs_with_vectors()),
2076 field,
2077 params.dim,
2078 plan_cache.map(|cache| &cache.tq),
2079 )?
2080 }
2081 VectorIndex::IvfTq { index: lazy, codec } => {
2082 let index = lazy.get();
2083 let centroids =
2084 self.trained_vectors
2085 .centroids
2086 .get(&field.0)
2087 .ok_or_else(|| {
2088 Error::Schema(format!(
2089 "IVF-TQ index requires coarse centroids for field {}",
2090 field.0
2091 ))
2092 })?;
2093 validate_coarse_centroids(centroids, params.dim)?;
2094 let routing = self
2095 .schema
2096 .get_field_entry(field)
2097 .and_then(|entry| entry.dense_vector_config.as_ref())
2098 .map_or(crate::dsl::IvfRoutingMode::Auto, |config| {
2099 config.ivf_routing
2100 });
2101 validate_ivf_tq_ann(index, centroids, codec, params.dim, routing, field)?;
2102 let flat = lazy_flat.expect("ANN/flat pairing validated above");
2103 search_ivf_tq_segment(
2104 index,
2105 centroids,
2106 codec,
2107 query,
2108 fetch_k.min(flat.num_docs_with_vectors()),
2109 field,
2110 params.nprobe,
2111 routing,
2112 plan_cache.map(|cache| &cache.ivf_tq),
2113 )?
2114 }
2115 VectorIndex::BinaryIvf(_) => {
2116 Vec::new()
2118 }
2119 }
2120 } else if let Some(lazy_flat) = lazy_flat {
2121 log::debug!(
2125 "[dense_vector_search] field {}: brute-force on {} vectors (dim={}, quant={:?})",
2126 field.0,
2127 lazy_flat.num_vectors,
2128 lazy_flat.dim,
2129 lazy_flat.quantization
2130 );
2131 let dim = lazy_flat.dim;
2132 let n = lazy_flat.num_vectors;
2133 let quant = lazy_flat.quantization;
2134 let batch_len =
2135 bounded_vector_score_batch(lazy_flat.vector_byte_size(), DENSE_SCORE_BATCH);
2136 let mut collector = FlatDocumentCollector::new(fetch_k.min(n), combiner);
2137 let mut scores = vec![0f32; batch_len];
2138
2139 for batch_start in (0..n).step_by(batch_len) {
2140 let batch_count = batch_len.min(n - batch_start);
2141 let batch_bytes = lazy_flat
2142 .read_vectors_batch(batch_start, batch_count)
2143 .await
2144 .map_err(crate::Error::Io)?;
2145 let raw = batch_bytes.as_slice();
2146
2147 Self::score_quantized_batch(
2148 query,
2149 raw,
2150 quant,
2151 dim,
2152 &mut scores[..batch_count],
2153 params.unit_norm,
2154 )?;
2155
2156 for (i, &score) in scores.iter().enumerate().take(batch_count) {
2157 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
2158 collector.push(doc_id, ordinal, score);
2159 }
2160 }
2161
2162 flat_results = Some(collector.into_results());
2163 Vec::new()
2164 } else {
2165 return Ok(Vec::new());
2166 };
2167 let l1_elapsed = t0.elapsed();
2168 {
2169 let kind = match ann_index {
2170 Some(VectorIndex::BinaryIvf(_)) => "binary_ivf",
2171 Some(VectorIndex::Tq { .. }) => "tq_flat",
2172 Some(VectorIndex::IvfTq { .. }) => "ivf_tq",
2173 None => "flat",
2174 };
2175 crate::observe::dense_l1(
2176 self.schema.index_label(),
2177 self.schema.get_field_name(field).unwrap_or("?"),
2178 kind,
2179 l1_elapsed.as_secs_f64(),
2180 flat_results.as_ref().map_or(results.len(), Vec::len),
2181 );
2182 }
2183 log::debug!(
2184 "[dense_vector_search] field {}: L1 returned {} candidates in {:.1}ms",
2185 field.0,
2186 flat_results.as_ref().map_or(results.len(), Vec::len),
2187 l1_elapsed.as_secs_f64() * 1000.0
2188 );
2189
2190 if let Some(results) = flat_results {
2191 return Ok(results);
2192 }
2193
2194 if ann_index.is_some()
2197 && !results.is_empty()
2198 && let Some(lazy_flat) = lazy_flat
2199 {
2200 let t_rerank = std::time::Instant::now();
2201 let vbs = lazy_flat.vector_byte_size();
2202 let (reranked, stats) = exact_score_dense_candidate_documents(
2203 &results,
2204 lazy_flat,
2205 query,
2206 params.unit_norm,
2207 combiner,
2208 k,
2209 )
2210 .await?;
2211
2212 crate::observe::dense_rerank(
2213 self.schema.index_label(),
2214 self.schema.get_field_name(field).unwrap_or("?"),
2215 t_rerank.elapsed().as_secs_f64(),
2216 stats.resolve_elapsed.as_secs_f64(),
2217 stats.read_elapsed.as_secs_f64(),
2218 stats.vector_count,
2219 );
2220 log::debug!(
2221 "[dense_vector_search] field {}: rerank {} vectors (dim={}, quant={:?}, bytes_per_vector={}): resolve={:.1}ms read={:.1}ms score={:.1}ms",
2222 field.0,
2223 stats.vector_count,
2224 lazy_flat.dim,
2225 lazy_flat.quantization,
2226 vbs,
2227 stats.resolve_elapsed.as_secs_f64() * 1000.0,
2228 stats.read_elapsed.as_secs_f64() * 1000.0,
2229 stats.score_elapsed.as_secs_f64() * 1000.0,
2230 );
2231
2232 log::debug!(
2233 "[dense_vector_search] field {}: rerank total={:.1}ms",
2234 field.0,
2235 t_rerank.elapsed().as_secs_f64() * 1000.0
2236 );
2237 return Ok(reranked);
2238 }
2239
2240 Ok(combine_grouped_ordinal_results(results, combiner, k))
2241 }
2242
2243 async fn search_binary_dense_vector_impl(
2248 &self,
2249 field: Field,
2250 query: &[u8],
2251 k: usize,
2252 combiner: crate::query::MultiValueCombiner,
2253 probe_cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
2254 ) -> Result<Vec<VectorSearchResult>> {
2255 let schema_dim = self.validate_binary_search_request(field, query)?;
2256 combiner.validate().map_err(Error::Query)?;
2257 if k == 0 {
2258 return Ok(Vec::new());
2259 }
2260 let t0 = crate::observe::Timer::start();
2261 if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0) {
2262 let ivf = lazy.get();
2263 let config = self
2264 .schema
2265 .get_field_entry(field)
2266 .and_then(|entry| entry.binary_dense_vector_config.as_ref())
2267 .ok_or_else(|| {
2268 Error::Schema(format!(
2269 "binary IVF field {} has no schema configuration",
2270 field.0
2271 ))
2272 })?;
2273 let quantizer = self
2274 .trained_vectors
2275 .binary_quantizers
2276 .get(&field.0)
2277 .ok_or_else(|| {
2278 Error::Schema(format!(
2279 "global binary IVF field {} has no loaded quantizer",
2280 field.0
2281 ))
2282 })?;
2283 validate_binary_ann(ivf, quantizer, config, schema_dim, field)?;
2284 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
2285 Error::Corruption(format!(
2286 "global binary IVF field {} is missing flat vector storage",
2287 field.0
2288 ))
2289 })?;
2290 let clusters = binary_probe_clusters(
2291 quantizer,
2292 query,
2293 config.nprobe,
2294 config.ivf_routing,
2295 probe_cache,
2296 )?;
2297 let single_valued = flat.num_vectors == flat.num_docs_with_vectors();
2298 let candidate_docs = k.min(flat.num_docs_with_vectors());
2299 let ann_results = if single_valued {
2300 ivf.search_binary_clusters::<false>(query, candidate_docs, &clusters)
2301 } else {
2302 ivf.search_binary_clusters::<true>(query, candidate_docs, &clusters)
2303 }
2304 .map_err(|error| {
2305 Error::Corruption(format!(
2306 "invalid binary IVF payload for field {}: {error}",
2307 field.0,
2308 ))
2309 })?;
2310 let results = exact_score_binary_candidate_documents(
2311 &ann_results,
2312 flat,
2313 query,
2314 schema_dim,
2315 combiner,
2316 k,
2317 )
2318 .await?;
2319 crate::observe::dense_l1(
2320 self.schema.index_label(),
2321 self.schema.get_field_name(field).unwrap_or("?"),
2322 "global_binary_ivf",
2323 t0.secs(),
2324 results.len(),
2325 );
2326 return Ok(results);
2327 }
2328 let lazy_flat = match self.flat_vectors.get(&field.0) {
2329 Some(f) => f,
2330 None => return Ok(Vec::new()),
2331 };
2332
2333 let dim_bits = lazy_flat.dim;
2334 let byte_len = lazy_flat.vector_byte_size();
2335 let n = lazy_flat.num_vectors;
2336
2337 if dim_bits != schema_dim {
2338 return Err(Error::Corruption(format!(
2339 "binary vector field {} has schema dimension {} but flat storage dimension {}",
2340 field.0, schema_dim, dim_bits
2341 )));
2342 }
2343
2344 if byte_len != query.len() {
2345 return Err(Error::Schema(format!(
2346 "Binary query vector byte length {} != field byte length {}",
2347 query.len(),
2348 byte_len
2349 )));
2350 }
2351
2352 let batch_len = bounded_vector_score_batch(byte_len, BINARY_SCORE_BATCH);
2353 let mut collector = FlatDocumentCollector::new(k, combiner);
2354 let mut scores = vec![0f32; batch_len];
2355
2356 for batch_start in (0..n).step_by(batch_len) {
2357 let batch_count = batch_len.min(n - batch_start);
2358 let batch_bytes = lazy_flat
2359 .read_vectors_batch(batch_start, batch_count)
2360 .await
2361 .map_err(crate::Error::Io)?;
2362 let raw = batch_bytes.as_slice();
2363
2364 crate::structures::simd::batch_hamming_scores(
2365 query,
2366 raw,
2367 byte_len,
2368 dim_bits,
2369 &mut scores[..batch_count],
2370 );
2371
2372 for (i, &score) in scores.iter().enumerate().take(batch_count) {
2373 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
2374 collector.push(doc_id, ordinal, score);
2375 }
2376 }
2377
2378 let results = collector.into_results();
2379
2380 crate::observe::dense_l1(
2381 self.schema.index_label(),
2382 self.schema.get_field_name(field).unwrap_or("?"),
2383 "binary_flat",
2384 t0.secs(),
2385 results.len(),
2386 );
2387 Ok(results)
2388 }
2389
2390 pub async fn search_binary_dense_vector(
2391 &self,
2392 field: Field,
2393 query: &[u8],
2394 k: usize,
2395 combiner: crate::query::MultiValueCombiner,
2396 ) -> Result<Vec<VectorSearchResult>> {
2397 self.search_binary_dense_vector_impl(field, query, k, combiner, None)
2398 .await
2399 }
2400
2401 pub(crate) async fn search_binary_dense_vector_with_probe_cache(
2402 &self,
2403 field: Field,
2404 query: &[u8],
2405 k: usize,
2406 combiner: crate::query::MultiValueCombiner,
2407 probe_cache: &std::sync::Mutex<Option<crate::structures::IvfProbePlan>>,
2408 ) -> Result<Vec<VectorSearchResult>> {
2409 self.search_binary_dense_vector_impl(field, query, k, combiner, Some(probe_cache))
2410 .await
2411 }
2412
2413 pub fn coarse_centroids(&self, field_id: u32) -> Option<&Arc<CoarseCentroids>> {
2415 self.trained_vectors.centroids.get(&field_id)
2416 }
2417
2418 pub fn set_trained_vectors(
2419 &mut self,
2420 trained_vectors: Arc<crate::segment::TrainedVectorStructures>,
2421 ) {
2422 self.trained_vectors = trained_vectors;
2423 }
2424
2425 pub fn get_vector_index(&self, field: Field) -> Option<&VectorIndex> {
2427 self.vector_indexes.get(&field.0)
2428 }
2429
2430 pub async fn get_positions(
2435 &self,
2436 field: Field,
2437 term: &[u8],
2438 ) -> Result<Option<crate::structures::PositionPostingList>> {
2439 let handle = match &self.positions_handle {
2441 Some(h) => h,
2442 None => return Ok(None),
2443 };
2444
2445 let mut key = Vec::with_capacity(4 + term.len());
2447 key.extend_from_slice(&field.0.to_le_bytes());
2448 key.extend_from_slice(term);
2449
2450 let term_info = match self.term_dict.get(&key).await? {
2452 Some(info) => info,
2453 None => return Ok(None),
2454 };
2455
2456 let (offset, length) = match term_info.position_info() {
2458 Some((o, l)) => (o, l),
2459 None => return Ok(None),
2460 };
2461
2462 let range = checked_file_range(offset, length, handle.len(), "position list")?;
2466 let slice = handle.slice(range);
2467 let data = slice.read_bytes().await?;
2468
2469 let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
2471
2472 Ok(Some(pos_list))
2473 }
2474
2475 pub fn has_positions(&self, field: Field) -> bool {
2477 if let Some(entry) = self.schema.get_field_entry(field) {
2479 entry.positions.is_some()
2480 } else {
2481 false
2482 }
2483 }
2484}
2485
2486#[cfg(feature = "sync")]
2488impl SegmentReader {
2489 pub fn get_postings_sync(&self, field: Field, term: &[u8]) -> Result<Option<BlockPostingList>> {
2491 let mut key = Vec::with_capacity(4 + term.len());
2493 key.extend_from_slice(&field.0.to_le_bytes());
2494 key.extend_from_slice(term);
2495
2496 let term_info = match self.term_dict.get_sync(&key)? {
2498 Some(info) => info,
2499 None => return Ok(None),
2500 };
2501
2502 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
2504 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
2505 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
2506 posting_list.push(doc_id, tf);
2507 }
2508 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
2509 return Ok(Some(block_list));
2510 }
2511
2512 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
2514 Error::Corruption("TermInfo has neither inline nor external data".to_string())
2515 })?;
2516
2517 let range = checked_file_range(
2518 posting_offset,
2519 posting_len,
2520 self.postings_handle.len(),
2521 "posting",
2522 )?;
2523 let posting_bytes = self.postings_handle.read_bytes_range_sync(range)?;
2524 let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
2525
2526 Ok(Some(block_list))
2527 }
2528
2529 pub fn get_prefix_postings_sync(
2531 &self,
2532 field: Field,
2533 prefix: &[u8],
2534 ) -> Result<Vec<BlockPostingList>> {
2535 if prefix.is_empty() {
2536 return Err(Error::Query("prefix must not be empty".into()));
2537 }
2538 let mut key_prefix = Vec::with_capacity(4 + prefix.len());
2539 key_prefix.extend_from_slice(&field.0.to_le_bytes());
2540 key_prefix.extend_from_slice(prefix);
2541
2542 let (entries, truncated) = self
2543 .term_dict
2544 .prefix_scan_limited_sync(&key_prefix, MAX_PREFIX_TERMS)?;
2545 if truncated {
2546 return Err(Error::Query(format!(
2547 "prefix expands to more than {MAX_PREFIX_TERMS} terms"
2548 )));
2549 }
2550 let posting_count: u64 = entries
2551 .iter()
2552 .map(|(_, term_info)| term_info.doc_freq() as u64)
2553 .sum();
2554 if posting_count > MAX_PREFIX_POSTINGS {
2555 return Err(Error::Query(format!(
2556 "prefix expands to {posting_count} postings (maximum {MAX_PREFIX_POSTINGS})"
2557 )));
2558 }
2559 let mut results = Vec::with_capacity(entries.len());
2560
2561 for (_key, term_info) in entries {
2562 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
2563 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
2564 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
2565 posting_list.push(doc_id, tf);
2566 }
2567 results.push(BlockPostingList::from_posting_list(&posting_list)?);
2568 } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
2569 let range = checked_file_range(
2570 posting_offset,
2571 posting_len,
2572 self.postings_handle.len(),
2573 "prefix posting",
2574 )?;
2575 let posting_bytes = self.postings_handle.read_bytes_range_sync(range)?;
2576 results.push(BlockPostingList::deserialize_zero_copy(posting_bytes)?);
2577 }
2578 }
2579
2580 Ok(results)
2581 }
2582
2583 pub fn get_positions_sync(
2585 &self,
2586 field: Field,
2587 term: &[u8],
2588 ) -> Result<Option<crate::structures::PositionPostingList>> {
2589 let handle = match &self.positions_handle {
2590 Some(h) => h,
2591 None => return Ok(None),
2592 };
2593
2594 let mut key = Vec::with_capacity(4 + term.len());
2596 key.extend_from_slice(&field.0.to_le_bytes());
2597 key.extend_from_slice(term);
2598
2599 let term_info = match self.term_dict.get_sync(&key)? {
2601 Some(info) => info,
2602 None => return Ok(None),
2603 };
2604
2605 let (offset, length) = match term_info.position_info() {
2606 Some((o, l)) => (o, l),
2607 None => return Ok(None),
2608 };
2609
2610 let range = checked_file_range(offset, length, handle.len(), "position list")?;
2611 let slice = handle.slice(range);
2612 let data = slice.read_bytes_sync()?;
2613
2614 let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
2615 Ok(Some(pos_list))
2616 }
2617
2618 pub fn search_dense_vector_sync(
2621 &self,
2622 field: Field,
2623 query: &[f32],
2624 k: usize,
2625 nprobe: usize,
2626 rerank_factor: f32,
2627 combiner: crate::query::MultiValueCombiner,
2628 ) -> Result<Vec<VectorSearchResult>> {
2629 self.search_dense_vector_sync_impl(field, query, k, nprobe, rerank_factor, combiner, None)
2630 }
2631
2632 #[cfg(feature = "sync")]
2633 #[allow(clippy::too_many_arguments)]
2634 pub(crate) fn search_dense_vector_sync_with_probe_cache(
2635 &self,
2636 field: Field,
2637 query: &[f32],
2638 k: usize,
2639 nprobe: usize,
2640 rerank_factor: f32,
2641 combiner: crate::query::MultiValueCombiner,
2642 plan_cache: &DensePlanCache,
2643 ) -> Result<Vec<VectorSearchResult>> {
2644 self.search_dense_vector_sync_impl(
2645 field,
2646 query,
2647 k,
2648 nprobe,
2649 rerank_factor,
2650 combiner,
2651 Some(plan_cache),
2652 )
2653 }
2654
2655 #[cfg(feature = "sync")]
2656 #[allow(clippy::too_many_arguments)]
2657 fn search_dense_vector_sync_impl(
2658 &self,
2659 field: Field,
2660 query: &[f32],
2661 k: usize,
2662 nprobe: usize,
2663 rerank_factor: f32,
2664 combiner: crate::query::MultiValueCombiner,
2665 plan_cache: Option<&DensePlanCache>,
2666 ) -> Result<Vec<VectorSearchResult>> {
2667 let params =
2668 self.validate_dense_search_request(field, query, nprobe, rerank_factor, combiner)?;
2669 let fetch_k = checked_dense_fetch_k(k, rerank_factor)?;
2670 if k == 0 {
2671 return Ok(Vec::new());
2672 }
2673
2674 let ann_index = self.vector_indexes.get(&field.0);
2675 let lazy_flat = self.flat_vectors.get(&field.0);
2676 if ann_index.is_none() && lazy_flat.is_none() {
2677 return Ok(Vec::new());
2678 }
2679
2680 if ann_index.is_some() && lazy_flat.is_none() {
2681 return Err(Error::Corruption(format!(
2682 "dense ANN field {} is missing flat vector storage",
2683 field.0
2684 )));
2685 }
2686
2687 if let Some(flat) = lazy_flat
2688 && flat.dim != params.dim
2689 {
2690 return Err(Error::Corruption(format!(
2691 "dense vector field {} has schema dimension {} but flat storage dimension {}",
2692 field.0, params.dim, flat.dim
2693 )));
2694 }
2695
2696 let results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
2697 match index {
2699 VectorIndex::Tq { index: lazy, codec } => {
2700 let flat = lazy_flat.expect("ANN/flat pairing validated above");
2701 search_tq_segment(
2702 lazy.get(),
2703 codec,
2704 query,
2705 fetch_k.min(flat.num_docs_with_vectors()),
2706 field,
2707 params.dim,
2708 plan_cache.map(|cache| &cache.tq),
2709 )?
2710 }
2711 VectorIndex::IvfTq { index: lazy, codec } => {
2712 let index = lazy.get();
2713 let centroids =
2714 self.trained_vectors
2715 .centroids
2716 .get(&field.0)
2717 .ok_or_else(|| {
2718 Error::Schema(format!(
2719 "IVF-TQ index requires coarse centroids for field {}",
2720 field.0
2721 ))
2722 })?;
2723 validate_coarse_centroids(centroids, params.dim)?;
2724 let routing = self
2725 .schema
2726 .get_field_entry(field)
2727 .and_then(|entry| entry.dense_vector_config.as_ref())
2728 .map_or(crate::dsl::IvfRoutingMode::Auto, |config| {
2729 config.ivf_routing
2730 });
2731 validate_ivf_tq_ann(index, centroids, codec, params.dim, routing, field)?;
2732 let flat = lazy_flat.expect("ANN/flat pairing validated above");
2733 search_ivf_tq_segment(
2734 index,
2735 centroids,
2736 codec,
2737 query,
2738 fetch_k.min(flat.num_docs_with_vectors()),
2739 field,
2740 params.nprobe,
2741 routing,
2742 plan_cache.map(|cache| &cache.ivf_tq),
2743 )?
2744 }
2745 VectorIndex::BinaryIvf(_) => {
2746 Vec::new()
2748 }
2749 }
2750 } else if let Some(lazy_flat) = lazy_flat {
2751 let dim = lazy_flat.dim;
2753 let n = lazy_flat.num_vectors;
2754 let quant = lazy_flat.quantization;
2755 let batch_len =
2756 bounded_vector_score_batch(lazy_flat.vector_byte_size(), DENSE_SCORE_BATCH);
2757 let mut collector = FlatDocumentCollector::new(fetch_k.min(n), combiner);
2758 let mut scores = vec![0f32; batch_len];
2759
2760 for batch_start in (0..n).step_by(batch_len) {
2761 let batch_count = batch_len.min(n - batch_start);
2762 let batch_bytes = lazy_flat
2763 .read_vectors_batch_sync(batch_start, batch_count)
2764 .map_err(crate::Error::Io)?;
2765 let raw = batch_bytes.as_slice();
2766
2767 Self::score_quantized_batch(
2768 query,
2769 raw,
2770 quant,
2771 dim,
2772 &mut scores[..batch_count],
2773 params.unit_norm,
2774 )?;
2775
2776 for (i, &score) in scores.iter().enumerate().take(batch_count) {
2777 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
2778 collector.push(doc_id, ordinal, score);
2779 }
2780 }
2781
2782 return Ok(collector.into_results());
2783 } else {
2784 return Ok(Vec::new());
2785 };
2786
2787 if ann_index.is_some()
2789 && !results.is_empty()
2790 && let Some(lazy_flat) = lazy_flat
2791 {
2792 return exact_score_dense_candidate_documents_sync(
2793 &results,
2794 lazy_flat,
2795 query,
2796 params.unit_norm,
2797 combiner,
2798 k,
2799 );
2800 }
2801
2802 Ok(combine_grouped_ordinal_results(results, combiner, k))
2803 }
2804
2805 #[cfg(feature = "sync")]
2810 fn search_binary_dense_vector_sync_impl(
2811 &self,
2812 field: Field,
2813 query: &[u8],
2814 k: usize,
2815 combiner: crate::query::MultiValueCombiner,
2816 probe_cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
2817 ) -> Result<Vec<VectorSearchResult>> {
2818 let schema_dim = self.validate_binary_search_request(field, query)?;
2819 combiner.validate().map_err(Error::Query)?;
2820 if k == 0 {
2821 return Ok(Vec::new());
2822 }
2823 let t0 = crate::observe::Timer::start();
2824 if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0) {
2825 let ivf = lazy.get();
2826 let config = self
2827 .schema
2828 .get_field_entry(field)
2829 .and_then(|entry| entry.binary_dense_vector_config.as_ref())
2830 .ok_or_else(|| {
2831 Error::Schema(format!(
2832 "binary IVF field {} has no schema configuration",
2833 field.0
2834 ))
2835 })?;
2836 let quantizer = self
2837 .trained_vectors
2838 .binary_quantizers
2839 .get(&field.0)
2840 .ok_or_else(|| {
2841 Error::Schema(format!(
2842 "global binary IVF field {} has no loaded quantizer",
2843 field.0
2844 ))
2845 })?;
2846 validate_binary_ann(ivf, quantizer, config, schema_dim, field)?;
2847 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
2848 Error::Corruption(format!(
2849 "global binary IVF field {} is missing flat vector storage",
2850 field.0
2851 ))
2852 })?;
2853 let clusters = binary_probe_clusters(
2854 quantizer,
2855 query,
2856 config.nprobe,
2857 config.ivf_routing,
2858 probe_cache,
2859 )?;
2860 let candidate_docs = k.min(flat.num_docs_with_vectors());
2861 let ann_results = if flat.num_vectors == flat.num_docs_with_vectors() {
2862 ivf.search_binary_clusters::<false>(query, candidate_docs, &clusters)
2863 } else {
2864 ivf.search_binary_clusters::<true>(query, candidate_docs, &clusters)
2865 }
2866 .map_err(|error| {
2867 Error::Corruption(format!(
2868 "invalid binary IVF payload for field {}: {error}",
2869 field.0,
2870 ))
2871 })?;
2872 let results = exact_score_binary_candidate_documents_sync(
2873 &ann_results,
2874 flat,
2875 query,
2876 schema_dim,
2877 combiner,
2878 k,
2879 )?;
2880 crate::observe::dense_l1(
2881 self.schema.index_label(),
2882 self.schema.get_field_name(field).unwrap_or("?"),
2883 "global_binary_ivf",
2884 t0.secs(),
2885 results.len(),
2886 );
2887 return Ok(results);
2888 }
2889 let lazy_flat = match self.flat_vectors.get(&field.0) {
2890 Some(f) => f,
2891 None => return Ok(Vec::new()),
2892 };
2893
2894 let dim_bits = lazy_flat.dim;
2895 let byte_len = lazy_flat.vector_byte_size();
2896 let n = lazy_flat.num_vectors;
2897
2898 if dim_bits != schema_dim {
2899 return Err(Error::Corruption(format!(
2900 "binary vector field {} has schema dimension {} but flat storage dimension {}",
2901 field.0, schema_dim, dim_bits
2902 )));
2903 }
2904
2905 if byte_len != query.len() {
2906 return Err(Error::Schema(format!(
2907 "Binary query vector byte length {} != field byte length {}",
2908 query.len(),
2909 byte_len
2910 )));
2911 }
2912
2913 let batch_len = bounded_vector_score_batch(byte_len, BINARY_SCORE_BATCH);
2914 let mut collector = FlatDocumentCollector::new(k, combiner);
2915 let mut scores = vec![0f32; batch_len];
2916
2917 for batch_start in (0..n).step_by(batch_len) {
2918 let batch_count = batch_len.min(n - batch_start);
2919 let batch_bytes = lazy_flat
2920 .read_vectors_batch_sync(batch_start, batch_count)
2921 .map_err(crate::Error::Io)?;
2922 let raw = batch_bytes.as_slice();
2923
2924 crate::structures::simd::batch_hamming_scores(
2925 query,
2926 raw,
2927 byte_len,
2928 dim_bits,
2929 &mut scores[..batch_count],
2930 );
2931
2932 for (i, &score) in scores.iter().enumerate().take(batch_count) {
2933 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
2934 collector.push(doc_id, ordinal, score);
2935 }
2936 }
2937
2938 let results = collector.into_results();
2939
2940 crate::observe::dense_l1(
2941 self.schema.index_label(),
2942 self.schema.get_field_name(field).unwrap_or("?"),
2943 "binary_flat",
2944 t0.secs(),
2945 results.len(),
2946 );
2947 Ok(results)
2948 }
2949
2950 #[cfg(feature = "sync")]
2951 pub fn search_binary_dense_vector_sync(
2952 &self,
2953 field: Field,
2954 query: &[u8],
2955 k: usize,
2956 combiner: crate::query::MultiValueCombiner,
2957 ) -> Result<Vec<VectorSearchResult>> {
2958 self.search_binary_dense_vector_sync_impl(field, query, k, combiner, None)
2959 }
2960
2961 #[cfg(feature = "sync")]
2962 pub(crate) fn search_binary_dense_vector_sync_with_probe_cache(
2963 &self,
2964 field: Field,
2965 query: &[u8],
2966 k: usize,
2967 combiner: crate::query::MultiValueCombiner,
2968 probe_cache: &std::sync::Mutex<Option<crate::structures::IvfProbePlan>>,
2969 ) -> Result<Vec<VectorSearchResult>> {
2970 self.search_binary_dense_vector_sync_impl(field, query, k, combiner, Some(probe_cache))
2971 }
2972}
2973
2974#[cfg(test)]
2975mod dense_search_safety_tests {
2976 use super::*;
2977
2978 #[test]
2979 fn dense_fetch_count_rejects_non_finite_and_unbounded_factors() {
2980 for factor in [
2981 f32::NAN,
2982 f32::INFINITY,
2983 f32::NEG_INFINITY,
2984 0.0,
2985 0.5,
2986 2.01,
2987 MAX_DENSE_RERANK_FACTOR + 1.0,
2988 ] {
2989 assert!(
2990 checked_dense_fetch_k(10, factor).is_err(),
2991 "factor={factor}"
2992 );
2993 }
2994 }
2995
2996 #[test]
2997 fn flat_document_collector_does_not_let_one_multivalue_doc_crowd_out_others() {
2998 let mut collector = FlatDocumentCollector::new(2, crate::query::MultiValueCombiner::Max);
2999 collector.push(1, 0, 1.0);
3000 collector.push(1, 1, 0.9);
3001 collector.push(2, 0, 0.8);
3002
3003 let results = collector.into_results();
3004 assert_eq!(
3005 results
3006 .iter()
3007 .map(|result| result.doc_id)
3008 .collect::<Vec<_>>(),
3009 vec![1, 2]
3010 );
3011 assert_eq!(results[0].ordinals.len(), 2);
3012 }
3013
3014 #[test]
3015 fn flat_document_collector_evicts_by_score_then_doc_id() {
3016 let mut collector = FlatDocumentCollector::new(2, crate::query::MultiValueCombiner::Max);
3017 collector.push(1, 0, 0.5);
3018 collector.push(3, 0, 0.8);
3019 collector.push(2, 0, 0.9);
3020 let results = collector.into_results();
3021 assert_eq!(
3022 results
3023 .iter()
3024 .map(|result| result.doc_id)
3025 .collect::<Vec<_>>(),
3026 vec![2, 3]
3027 );
3028
3029 let mut tied = FlatDocumentCollector::new(1, crate::query::MultiValueCombiner::Max);
3030 tied.push(2, 0, 1.0);
3031 tied.push(1, 0, 1.0);
3032 let results = tied.into_results();
3033 assert_eq!(results[0].doc_id, 1);
3034 }
3035
3036 #[test]
3037 fn dense_fetch_count_rounds_up_and_detects_overflow() {
3038 assert_eq!(checked_dense_fetch_k(3, 1.5).unwrap(), 5);
3039 assert_eq!(checked_dense_fetch_k(10_000, 2.0).unwrap(), 20_000);
3040 assert!(checked_dense_fetch_k(10_001, 2.0).is_err());
3041 assert!(checked_dense_fetch_k(usize::MAX, 2.0).is_err());
3042 }
3043
3044 #[test]
3045 fn file_ranges_reject_overflow_and_truncation() {
3046 assert_eq!(checked_file_range(4, 3, 7, "test").unwrap(), 4..7);
3047 assert!(checked_file_range(u64::MAX, 1, u64::MAX, "test").is_err());
3048 assert!(checked_file_range(5, 3, 7, "test").is_err());
3049 }
3050
3051 #[test]
3052 fn candidate_vector_reads_coalesce_contiguous_values() {
3053 let mut runs = Vec::new();
3054 plan_vector_read_runs(&[3, 4, 5, 9, 12, 13], &mut runs).unwrap();
3055 assert_eq!(runs.len(), 3);
3056 assert!(matches!(
3057 runs.as_slice(),
3058 [
3059 VectorReadRun {
3060 buffer_start: 0,
3061 flat_start: 3,
3062 count: 3,
3063 },
3064 VectorReadRun {
3065 buffer_start: 3,
3066 flat_start: 9,
3067 count: 1,
3068 },
3069 VectorReadRun {
3070 buffer_start: 4,
3071 flat_start: 12,
3072 count: 2,
3073 },
3074 ]
3075 ));
3076 assert!(plan_vector_read_runs(&[3, 3], &mut runs).is_err());
3077 }
3078
3079 #[tokio::test]
3080 async fn multivalue_ann_rerank_streams_past_document_candidate_cap() {
3081 use crate::directories::{FileHandle, OwnedBytes};
3082 use crate::segment::FlatVectorData;
3083
3084 const VALUES: usize = MAX_DENSE_CANDIDATES_PER_SEGMENT + 1;
3085 let mut encoded = Vec::new();
3086 let vectors = vec![1.0f32; VALUES];
3087 let doc_ids: Vec<_> = (0..VALUES).map(|ordinal| (0, ordinal as u16)).collect();
3088 FlatVectorData::serialize_binary_from_flat_streaming(
3089 1,
3090 &vectors,
3091 &doc_ids,
3092 DenseVectorQuantization::F32,
3093 &mut encoded,
3094 )
3095 .unwrap();
3096 let flat = LazyFlatVectorData::open_with_doc_limit(
3097 FileHandle::from_bytes(OwnedBytes::new(encoded)),
3098 Some(1),
3099 )
3100 .await
3101 .unwrap();
3102
3103 let (results, stats) = exact_score_dense_candidate_documents(
3104 &[(0, 0, 0.0)],
3105 &flat,
3106 &[1.0],
3107 false,
3108 crate::query::MultiValueCombiner::Max,
3109 1,
3110 )
3111 .await
3112 .unwrap();
3113 assert_eq!(stats.vector_count, VALUES);
3114 assert_eq!(results.len(), 1);
3115 assert_eq!(results[0].ordinals.len(), VALUES);
3116 assert!((results[0].score - 1.0).abs() < 1e-5);
3117 }
3118}