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 = 200_000;
20const MAX_ANN_ORDINAL_OVERFETCH: usize = 32;
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)]
28pub struct SegmentMemoryStats {
29 pub segment_id: u128,
31 pub num_docs: u32,
33 pub term_dict_cache_bytes: usize,
35 pub store_cache_bytes: usize,
37 pub sparse_index_bytes: usize,
39 pub dense_index_bytes: usize,
41 pub bloom_filter_bytes: usize,
43 pub pinned_metadata_bytes: u64,
45 pub pin_intended_bytes: u64,
48}
49
50impl SegmentMemoryStats {
51 pub fn total_bytes(&self) -> usize {
53 self.term_dict_cache_bytes
54 + self.store_cache_bytes
55 + self.sparse_index_bytes
56 + self.dense_index_bytes
57 + self.bloom_filter_bytes
58 }
59}
60
61use std::cmp::Ordering;
62use std::collections::BinaryHeap;
63use std::sync::Arc;
64
65use rustc_hash::FxHashMap;
66
67use super::vector_data::LazyFlatVectorData;
68use crate::directories::{Directory, FileHandle};
69use crate::dsl::{DenseVectorQuantization, Document, Field, Schema};
70use crate::query::{MAX_DENSE_NPROBE, MAX_DENSE_RERANK_FACTOR};
71use crate::structures::{
72 AsyncSSTableReader, BlockPostingList, CoarseCentroids, IVFPQIndex, IVFRaBitQIndex, PQCodebook,
73 RaBitQIndex, SSTableStats, TermInfo,
74};
75use crate::{DocId, Error, Result};
76
77use super::store::{AsyncStoreReader, RawStoreBlock};
78use super::types::{SegmentFiles, SegmentId, SegmentMeta};
79
80pub(crate) fn combine_ordinal_results(
86 raw: impl IntoIterator<Item = (u32, u16, f32)>,
87 combiner: crate::query::MultiValueCombiner,
88 limit: usize,
89) -> Vec<VectorSearchResult> {
90 let collected: Vec<(u32, u16, f32)> = raw.into_iter().collect();
91
92 let num_raw = collected.len();
93 if log::log_enabled!(log::Level::Debug) {
94 let mut ids: Vec<u32> = collected.iter().map(|(d, _, _)| *d).collect();
95 ids.sort_unstable();
96 ids.dedup();
97 log::debug!(
98 "combine_ordinal_results: {} raw entries, {} unique docs, combiner={:?}, limit={}",
99 num_raw,
100 ids.len(),
101 combiner,
102 limit
103 );
104 }
105
106 let all_single = collected.iter().all(|&(_, ord, _)| ord == 0);
108 if all_single {
109 let mut results: Vec<VectorSearchResult> = collected
110 .into_iter()
111 .map(|(doc_id, _, score)| VectorSearchResult::new(doc_id, score, vec![(0, score)]))
112 .collect();
113 results.sort_unstable_by(|a, b| {
114 b.score
115 .total_cmp(&a.score)
116 .then_with(|| a.doc_id.cmp(&b.doc_id))
117 });
118 results.truncate(limit);
119 return results;
120 }
121
122 let mut doc_ordinals: rustc_hash::FxHashMap<DocId, Vec<(u32, f32)>> =
124 rustc_hash::FxHashMap::default();
125 for (doc_id, ordinal, score) in collected {
126 doc_ordinals
127 .entry(doc_id as DocId)
128 .or_default()
129 .push((ordinal as u32, score));
130 }
131 let mut results: Vec<VectorSearchResult> = doc_ordinals
132 .into_iter()
133 .map(|(doc_id, ordinals)| {
134 let combined_score = combiner.combine(&ordinals);
135 VectorSearchResult::new(doc_id, combined_score, ordinals)
136 })
137 .collect();
138 results.sort_unstable_by(|a, b| {
139 b.score
140 .total_cmp(&a.score)
141 .then_with(|| a.doc_id.cmp(&b.doc_id))
142 });
143 results.truncate(limit);
144 results
145}
146
147struct HeapVectorResult(VectorSearchResult);
152
153impl PartialEq for HeapVectorResult {
154 fn eq(&self, other: &Self) -> bool {
155 self.0.score.to_bits() == other.0.score.to_bits() && self.0.doc_id == other.0.doc_id
156 }
157}
158
159impl Eq for HeapVectorResult {}
160
161impl Ord for HeapVectorResult {
162 fn cmp(&self, other: &Self) -> Ordering {
163 other
166 .0
167 .score
168 .total_cmp(&self.0.score)
169 .then_with(|| self.0.doc_id.cmp(&other.0.doc_id))
170 }
171}
172
173impl PartialOrd for HeapVectorResult {
174 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
175 Some(self.cmp(other))
176 }
177}
178
179struct FlatDocumentCollector {
183 heap: BinaryHeap<HeapVectorResult>,
184 limit: usize,
185 combiner: crate::query::MultiValueCombiner,
186 current_doc: Option<DocId>,
187 current_ordinals: Vec<(u32, f32)>,
188}
189
190impl FlatDocumentCollector {
191 fn new(limit: usize, combiner: crate::query::MultiValueCombiner) -> Self {
192 Self {
193 heap: BinaryHeap::with_capacity(limit.min(8 * 1024)),
194 limit,
195 combiner,
196 current_doc: None,
197 current_ordinals: Vec::new(),
198 }
199 }
200
201 fn push(&mut self, doc_id: DocId, ordinal: u16, score: f32) {
202 if self.current_doc.is_some_and(|current| current != doc_id) {
203 self.finish_current();
204 }
205 self.current_doc = Some(doc_id);
206 self.current_ordinals.push((ordinal as u32, score));
207 }
208
209 fn finish_current(&mut self) {
210 let Some(doc_id) = self.current_doc.take() else {
211 return;
212 };
213 let score = self.combiner.combine(&self.current_ordinals);
214 let should_retain = self.heap.len() < self.limit
215 || self.heap.peek().is_some_and(|worst| {
216 HeapVectorResult(VectorSearchResult::new(doc_id, score, Vec::new()))
217 .cmp(worst)
218 .is_lt()
219 });
220
221 if !should_retain {
222 self.current_ordinals.clear();
226 return;
227 }
228
229 let ordinals = std::mem::take(&mut self.current_ordinals);
230 let entry = HeapVectorResult(VectorSearchResult::new(doc_id, score, ordinals));
231 if self.heap.len() < self.limit {
232 self.heap.push(entry);
233 } else if let Some(mut worst) = self.heap.peek_mut() {
234 let mut evicted = std::mem::replace(&mut worst.0, entry.0);
237 evicted.ordinals.clear();
238 self.current_ordinals = evicted.ordinals;
239 }
240 }
241
242 fn into_results(mut self) -> Vec<VectorSearchResult> {
243 self.finish_current();
244 let mut results: Vec<_> = self.heap.into_iter().map(|entry| entry.0).collect();
245 results.sort_unstable_by(|a, b| {
246 b.score
247 .total_cmp(&a.score)
248 .then_with(|| a.doc_id.cmp(&b.doc_id))
249 });
250 results
251 }
252}
253
254fn combine_grouped_ordinal_results(
257 raw: impl IntoIterator<Item = RawVectorCandidate>,
258 combiner: crate::query::MultiValueCombiner,
259 limit: usize,
260) -> Vec<VectorSearchResult> {
261 let mut collector = FlatDocumentCollector::new(limit, combiner);
262 for (doc_id, ordinal, score) in raw {
263 collector.push(doc_id, ordinal, score);
264 }
265 collector.into_results()
266}
267
268#[derive(Clone, Copy)]
269struct DenseSearchParams {
270 dim: usize,
271 nprobe: usize,
272 unit_norm: bool,
273}
274
275fn checked_dense_fetch_k(k: usize, rerank_factor: f32) -> Result<usize> {
277 if !rerank_factor.is_finite() || !(1.0..=MAX_DENSE_RERANK_FACTOR).contains(&rerank_factor) {
278 return Err(Error::Query(format!(
279 "dense rerank_factor must be finite and in [1, {MAX_DENSE_RERANK_FACTOR}], got {rerank_factor}"
280 )));
281 }
282
283 let fetch = (k as f64) * (rerank_factor as f64);
284 if !fetch.is_finite()
285 || fetch > usize::MAX as f64
286 || fetch > MAX_DENSE_CANDIDATES_PER_SEGMENT as f64
287 {
288 return Err(Error::Query(format!(
289 "dense candidate count exceeds the per-segment maximum of \
290 {MAX_DENSE_CANDIDATES_PER_SEGMENT}: k={k}, rerank_factor={rerank_factor}"
291 )));
292 }
293 Ok(fetch.ceil() as usize)
294}
295
296fn ann_ordinal_fetch_k(fetch_k: usize, num_vectors: usize, num_docs: usize) -> usize {
297 if num_vectors == 0 || num_docs == 0 {
298 return 0;
299 }
300 let average_values_per_doc = num_vectors
301 .div_ceil(num_docs)
302 .clamp(1, MAX_ANN_ORDINAL_OVERFETCH);
303 fetch_k
304 .saturating_mul(average_values_per_doc)
305 .min(MAX_DENSE_CANDIDATES_PER_SEGMENT)
306 .min(num_vectors)
307}
308
309fn progressive_ann_search<F>(
314 target_docs: usize,
315 initial_fetch: usize,
316 max_vectors: usize,
317 mut search: F,
318) -> Result<Vec<RawVectorCandidate>>
319where
320 F: FnMut(usize) -> Vec<RawVectorCandidate>,
321{
322 let max_fetch = max_vectors.min(MAX_DENSE_CANDIDATES_PER_SEGMENT);
323 if target_docs == 0 || max_fetch == 0 {
324 return Ok(Vec::new());
325 }
326
327 let target_docs = target_docs.min(max_fetch);
328 let mut fetch = initial_fetch.max(target_docs).min(max_fetch);
329 loop {
330 let results = search(fetch);
331 let mut docs = rustc_hash::FxHashSet::with_capacity_and_hasher(
332 target_docs.min(results.len()),
333 Default::default(),
334 );
335 for &(doc_id, _, _) in &results {
336 docs.insert(doc_id);
337 if docs.len() >= target_docs {
338 return Ok(results);
339 }
340 }
341 if results.len() < fetch || fetch == max_vectors {
342 return Ok(results);
343 }
344 if fetch == max_fetch {
345 return Err(Error::Query(format!(
346 "ANN search reached the per-segment candidate limit of \
347 {MAX_DENSE_CANDIDATES_PER_SEGMENT} with only {} of {target_docs} requested \
348 documents; reduce k or the number of vector values per document",
349 docs.len()
350 )));
351 }
352
353 let next_fetch = fetch.saturating_mul(2).min(max_fetch);
354 if next_fetch == fetch {
355 return Ok(results);
356 }
357 fetch = next_fetch;
358 }
359}
360
361#[inline]
362fn bounded_vector_score_batch(vector_byte_size: usize, preferred: usize) -> usize {
363 preferred.min((MAX_VECTOR_SCORE_BATCH_BYTES / vector_byte_size.max(1)).max(1))
364}
365
366fn checked_file_range(
367 offset: u64,
368 length: u64,
369 file_length: u64,
370 description: &str,
371) -> Result<std::ops::Range<u64>> {
372 let end = offset
373 .checked_add(length)
374 .ok_or_else(|| Error::Corruption(format!("{description} byte range overflows u64")))?;
375 if end > file_length {
376 return Err(Error::Corruption(format!(
377 "{description} byte range {offset}..{end} exceeds file length {file_length}"
378 )));
379 }
380 Ok(offset..end)
381}
382
383type RawVectorCandidate = (u32, u16, f32);
384type ResolvedVectorCandidate = (usize, usize); fn expand_ann_candidate_documents(
390 ann_results: &[RawVectorCandidate],
391 flat: &LazyFlatVectorData,
392) -> Result<(Vec<RawVectorCandidate>, Vec<ResolvedVectorCandidate>)> {
393 let mut candidate_docs: Vec<DocId> = ann_results.iter().map(|candidate| candidate.0).collect();
394 candidate_docs.sort_unstable();
395 candidate_docs.dedup();
396
397 let mut expanded = Vec::new();
398 let mut resolved = Vec::new();
399 for doc_id in candidate_docs {
400 let (start, count) = flat.flat_indexes_for_doc_range(doc_id);
401 if count == 0 {
402 return Err(Error::Corruption(format!(
403 "ANN candidate document {doc_id} is missing from flat vector storage"
404 )));
405 }
406 let next_len = expanded
407 .len()
408 .checked_add(count)
409 .ok_or_else(|| Error::Query("ANN candidate vector expansion overflow".to_string()))?;
410 if next_len > MAX_DENSE_CANDIDATES_PER_SEGMENT {
411 return Err(Error::Query(format!(
412 "ANN candidate documents expand to more than \
413 {MAX_DENSE_CANDIDATES_PER_SEGMENT} vectors in one segment"
414 )));
415 }
416 expanded.reserve(count);
417 resolved.reserve(count);
418 let end = start
419 .checked_add(count)
420 .ok_or_else(|| Error::Corruption("flat vector range overflow".to_string()))?;
421 for flat_index in start..end {
422 let (stored_doc_id, ordinal) = flat.get_doc_id(flat_index);
423 if stored_doc_id != doc_id {
424 return Err(Error::Corruption(format!(
425 "flat vector doc map is not contiguous for document {doc_id}"
426 )));
427 }
428 let result_index = expanded.len();
429 expanded.push((doc_id, ordinal, 0.0));
430 resolved.push((result_index, flat_index));
431 }
432 }
433 Ok((expanded, resolved))
434}
435
436async fn exact_score_binary_candidate_documents(
437 ann_results: &[RawVectorCandidate],
438 flat: &LazyFlatVectorData,
439 query: &[u8],
440 dim_bits: usize,
441) -> Result<Vec<RawVectorCandidate>> {
442 let (mut expanded, resolved) = expand_ann_candidate_documents(ann_results, flat)?;
443 let vector_byte_size = flat.vector_byte_size();
444 let batch_len = bounded_vector_score_batch(vector_byte_size, BINARY_SCORE_BATCH);
445 let raw_capacity = batch_len
446 .checked_mul(vector_byte_size)
447 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
448 let mut raw = vec![0u8; raw_capacity];
449 let mut scores = vec![0.0f32; batch_len];
450
451 for chunk in resolved.chunks(batch_len) {
452 #[cfg(feature = "native")]
453 flat.prefetch_vectors(chunk.iter().map(|&(_, flat_index)| flat_index));
454 let raw_len = chunk
455 .len()
456 .checked_mul(vector_byte_size)
457 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
458 let raw = &mut raw[..raw_len];
459 for (buffer_index, &(_, flat_index)) in chunk.iter().enumerate() {
460 flat.read_vector_raw_into(
461 flat_index,
462 &mut raw[buffer_index * vector_byte_size..(buffer_index + 1) * vector_byte_size],
463 )
464 .await
465 .map_err(Error::Io)?;
466 }
467 crate::structures::simd::batch_hamming_scores(
468 query,
469 raw,
470 vector_byte_size,
471 dim_bits,
472 &mut scores[..chunk.len()],
473 );
474 for (buffer_index, &(result_index, _)) in chunk.iter().enumerate() {
475 expanded[result_index].2 = scores[buffer_index];
476 }
477 }
478 Ok(expanded)
479}
480
481#[cfg(feature = "sync")]
482fn exact_score_binary_candidate_documents_sync(
483 ann_results: &[RawVectorCandidate],
484 flat: &LazyFlatVectorData,
485 query: &[u8],
486 dim_bits: usize,
487) -> Result<Vec<RawVectorCandidate>> {
488 let (mut expanded, resolved) = expand_ann_candidate_documents(ann_results, flat)?;
489 let vector_byte_size = flat.vector_byte_size();
490 let batch_len = bounded_vector_score_batch(vector_byte_size, BINARY_SCORE_BATCH);
491 let raw_capacity = batch_len
492 .checked_mul(vector_byte_size)
493 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
494 let mut raw = vec![0u8; raw_capacity];
495 let mut scores = vec![0.0f32; batch_len];
496
497 for chunk in resolved.chunks(batch_len) {
498 let raw_len = chunk
499 .len()
500 .checked_mul(vector_byte_size)
501 .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
502 let raw = &mut raw[..raw_len];
503 for (buffer_index, &(_, flat_index)) in chunk.iter().enumerate() {
504 flat.read_vector_raw_into_sync(
505 flat_index,
506 &mut raw[buffer_index * vector_byte_size..(buffer_index + 1) * vector_byte_size],
507 )
508 .map_err(Error::Io)?;
509 }
510 crate::structures::simd::batch_hamming_scores(
511 query,
512 raw,
513 vector_byte_size,
514 dim_bits,
515 &mut scores[..chunk.len()],
516 );
517 for (buffer_index, &(result_index, _)) in chunk.iter().enumerate() {
518 expanded[result_index].2 = scores[buffer_index];
519 }
520 }
521 Ok(expanded)
522}
523
524fn validate_coarse_centroids(centroids: &CoarseCentroids, dim: usize) -> Result<()> {
525 let expected = (centroids.num_clusters as usize)
526 .checked_mul(dim)
527 .ok_or_else(|| Error::Corruption("coarse centroid size overflow".into()))?;
528 if centroids.num_clusters == 0
529 || centroids.dim != dim
530 || centroids.centroids.len() != expected
531 || centroids.centroids.iter().any(|value| !value.is_finite())
532 {
533 return Err(Error::Corruption(format!(
534 "invalid coarse centroids: clusters={}, dim={}, values={} (expected dim={dim}, values={expected})",
535 centroids.num_clusters,
536 centroids.dim,
537 centroids.centroids.len()
538 )));
539 }
540 Ok(())
541}
542
543pub struct SegmentReader {
549 meta: SegmentMeta,
550 term_dict: Arc<AsyncSSTableReader<TermInfo>>,
552 postings_handle: FileHandle,
554 store: Arc<AsyncStoreReader>,
556 schema: Arc<Schema>,
557 vector_indexes: FxHashMap<u32, VectorIndex>,
559 flat_vectors: FxHashMap<u32, LazyFlatVectorData>,
561 coarse_centroids: FxHashMap<u32, Arc<CoarseCentroids>>,
563 sparse_indexes: FxHashMap<u32, SparseIndex>,
565 bmp_indexes: FxHashMap<u32, BmpIndex>,
567 positions_handle: Option<FileHandle>,
569 fast_fields: FxHashMap<u32, crate::structures::fast_field::FastFieldReader>,
571 #[cfg(feature = "native")]
576 pin_report: crate::segment::pin::PinReport,
577}
578
579impl SegmentReader {
580 pub async fn open<D: Directory>(
582 dir: &D,
583 segment_id: SegmentId,
584 schema: Arc<Schema>,
585 cache_blocks: usize,
586 ) -> Result<Self> {
587 Self::open_with_cache_blocks(dir, segment_id, schema, cache_blocks, cache_blocks).await
588 }
589
590 pub async fn open_with_cache_blocks<D: Directory>(
596 dir: &D,
597 segment_id: SegmentId,
598 schema: Arc<Schema>,
599 term_cache_blocks: usize,
600 store_cache_blocks: usize,
601 ) -> Result<Self> {
602 let files = SegmentFiles::new(segment_id.0);
603
604 let meta_slice = dir.open_read(&files.meta).await?;
606 let meta_bytes = meta_slice.read_bytes().await?;
607 let meta = SegmentMeta::deserialize(meta_bytes.as_slice())?;
608 debug_assert_eq!(meta.id, segment_id.0);
609
610 let term_dict_handle = dir.open_lazy(&files.term_dict).await?;
612 let term_dict = AsyncSSTableReader::open(term_dict_handle, term_cache_blocks).await?;
613
614 let postings_handle = dir.open_lazy(&files.postings).await?;
616
617 let store_handle = dir.open_lazy(&files.store).await?;
619 let store = AsyncStoreReader::open(store_handle, store_cache_blocks).await?;
620
621 let vectors_data = loader::load_vectors_file(dir, &files, &schema, meta.num_docs).await?;
623 let vector_indexes = vectors_data.indexes;
624 let flat_vectors = vectors_data.flat_vectors;
625
626 #[cfg(feature = "native")]
631 for (field_id, lazy_flat) in &flat_vectors {
632 if vector_indexes.contains_key(field_id) {
633 lazy_flat.advise_random_access();
634 }
635 }
636
637 let sparse_data = loader::load_sparse_file(dir, &files, meta.num_docs, &schema).await?;
639 let sparse_indexes = sparse_data.maxscore_indexes;
640 let bmp_indexes = sparse_data.bmp_indexes;
641
642 let positions_handle = loader::open_positions_file(dir, &files, &schema).await?;
644
645 let fast_fields = loader::load_fast_fields_file(dir, &files, &schema).await?;
647
648 {
650 let mut parts = vec![format!(
651 "[segment] loaded {:016x}: docs={}",
652 segment_id.0, meta.num_docs
653 )];
654 if !vector_indexes.is_empty() || !flat_vectors.is_empty() {
655 parts.push(format!(
656 "dense: {} ann + {} flat fields",
657 vector_indexes.len(),
658 flat_vectors.len()
659 ));
660 }
661 for (field_id, idx) in &sparse_indexes {
662 parts.push(format!(
663 "sparse field {}: {} dims, ~{:.1} KB",
664 field_id,
665 idx.num_dimensions(),
666 idx.num_dimensions() as f64 * 24.0 / 1024.0
667 ));
668 }
669 for (field_id, idx) in &bmp_indexes {
670 parts.push(format!(
671 "bmp field {}: {} dims, {} blocks",
672 field_id,
673 idx.dims(),
674 idx.num_blocks
675 ));
676 }
677 if !fast_fields.is_empty() {
678 parts.push(format!("fast: {} fields", fast_fields.len()));
679 }
680 log::debug!("{}", parts.join(", "));
681 }
682
683 #[allow(unused_mut)]
684 let mut reader = Self {
685 meta,
686 term_dict: Arc::new(term_dict),
687 postings_handle,
688 store: Arc::new(store),
689 schema,
690 vector_indexes,
691 flat_vectors,
692 coarse_centroids: FxHashMap::default(),
693 sparse_indexes,
694 bmp_indexes,
695 positions_handle,
696 fast_fields,
697 #[cfg(feature = "native")]
698 pin_report: Default::default(),
699 };
700
701 #[cfg(feature = "native")]
703 reader.apply_pin_policy(&crate::segment::pin::pin_policy().to_owned());
704
705 Ok(reader)
706 }
707
708 #[cfg(feature = "native")]
717 pub(crate) fn apply_pin_policy(&mut self, policy: &crate::segment::pin::PinPolicy) {
718 use crate::segment::pin::PinReport;
719
720 if !policy.is_enabled() {
721 return;
722 }
723 let mut remaining = policy.budget_bytes;
724 let mut report = PinReport::default();
725
726 for bmp in self.bmp_indexes.values_mut() {
728 bmp.pin_block_starts(policy.mode, &mut remaining, &mut report);
729 }
730 for sparse in self.sparse_indexes.values_mut() {
732 sparse.pin_skip_section(policy.mode, &mut remaining, &mut report);
733 }
734 for flat in self.flat_vectors.values_mut() {
736 flat.pin_doc_ids(policy.mode, &mut remaining, &mut report);
737 }
738 for bmp in self.bmp_indexes.values_mut() {
739 bmp.pin_doc_maps(policy.mode, &mut remaining, &mut report);
740 }
741 for bmp in self.bmp_indexes.values_mut() {
743 bmp.pin_sb_grid(policy.mode, &mut remaining, &mut report);
744 }
745
746 if report.skipped_budget_bytes > 0 || report.failed_bytes > 0 {
747 log::warn!(
748 "[pin] segment {:016x}: pinned {}/{} bytes (budget skipped {}, mlock failed {}) — raise HERMES_PIN_METADATA_BUDGET_MB or RLIMIT_MEMLOCK for full coverage",
749 self.meta.id,
750 report.pinned_bytes,
751 report.intended_bytes,
752 report.skipped_budget_bytes,
753 report.failed_bytes,
754 );
755 } else if report.pinned_bytes > 0 {
756 log::info!(
757 "[pin] segment {:016x}: pinned {} bytes of hot metadata ({:?})",
758 self.meta.id,
759 report.pinned_bytes,
760 policy.mode,
761 );
762 }
763 self.pin_report = report;
764 }
765
766 pub fn meta(&self) -> &SegmentMeta {
772 &self.meta
773 }
774
775 pub fn num_docs(&self) -> u32 {
776 self.meta.num_docs
777 }
778
779 pub fn avg_field_len(&self, field: Field) -> f32 {
781 self.meta.avg_field_len(field)
782 }
783
784 pub fn schema(&self) -> &Schema {
785 &self.schema
786 }
787
788 pub fn sparse_indexes(&self) -> &FxHashMap<u32, SparseIndex> {
790 &self.sparse_indexes
791 }
792
793 pub fn sparse_index(&self, field: Field) -> Option<&SparseIndex> {
795 self.sparse_indexes.get(&field.0)
796 }
797
798 pub fn bmp_index(&self, field: Field) -> Option<&BmpIndex> {
800 self.bmp_indexes.get(&field.0)
801 }
802
803 pub fn bmp_indexes(&self) -> &FxHashMap<u32, BmpIndex> {
805 &self.bmp_indexes
806 }
807
808 pub fn vector_indexes(&self) -> &FxHashMap<u32, VectorIndex> {
810 &self.vector_indexes
811 }
812
813 pub fn flat_vectors(&self) -> &FxHashMap<u32, LazyFlatVectorData> {
815 &self.flat_vectors
816 }
817
818 pub fn fast_field(
820 &self,
821 field_id: u32,
822 ) -> Option<&crate::structures::fast_field::FastFieldReader> {
823 self.fast_fields.get(&field_id)
824 }
825
826 pub fn fast_fields(&self) -> &FxHashMap<u32, crate::structures::fast_field::FastFieldReader> {
828 &self.fast_fields
829 }
830
831 pub fn term_dict_stats(&self) -> SSTableStats {
833 self.term_dict.stats()
834 }
835
836 pub fn memory_stats(&self) -> SegmentMemoryStats {
838 let term_dict_stats = self.term_dict.stats();
839
840 let term_dict_cache_bytes = self.term_dict.cached_bytes();
844 let store_cache_bytes = self.store.cached_bytes();
845
846 let sparse_index_bytes: usize = self
848 .sparse_indexes
849 .values()
850 .map(|s| s.estimated_memory_bytes())
851 .sum::<usize>()
852 + self
853 .bmp_indexes
854 .values()
855 .map(|b| b.estimated_memory_bytes())
856 .sum::<usize>();
857
858 let dense_index_bytes: usize = self
861 .vector_indexes
862 .values()
863 .map(|v| v.estimated_memory_bytes())
864 .sum();
865
866 #[cfg(feature = "native")]
867 let (pinned_metadata_bytes, pin_intended_bytes) =
868 (self.pin_report.pinned_bytes, self.pin_report.intended_bytes);
869 #[cfg(not(feature = "native"))]
870 let (pinned_metadata_bytes, pin_intended_bytes) = (0u64, 0u64);
871
872 SegmentMemoryStats {
873 segment_id: self.meta.id,
874 num_docs: self.meta.num_docs,
875 term_dict_cache_bytes,
876 store_cache_bytes,
877 sparse_index_bytes,
878 dense_index_bytes,
879 bloom_filter_bytes: term_dict_stats.bloom_filter_size,
880 pinned_metadata_bytes,
881 pin_intended_bytes,
882 }
883 }
884
885 pub async fn get_postings(
890 &self,
891 field: Field,
892 term: &[u8],
893 ) -> Result<Option<BlockPostingList>> {
894 log::debug!(
895 "SegmentReader::get_postings field={} term_len={}",
896 field.0,
897 term.len()
898 );
899
900 let mut key = Vec::with_capacity(4 + term.len());
902 key.extend_from_slice(&field.0.to_le_bytes());
903 key.extend_from_slice(term);
904
905 let term_info = match self.term_dict.get(&key).await? {
907 Some(info) => {
908 log::debug!("SegmentReader::get_postings found term_info");
909 info
910 }
911 None => {
912 log::debug!("SegmentReader::get_postings term not found");
913 return Ok(None);
914 }
915 };
916
917 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
919 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
921 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
922 posting_list.push(doc_id, tf);
923 }
924 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
925 return Ok(Some(block_list));
926 }
927
928 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
930 Error::Corruption("TermInfo has neither inline nor external data".to_string())
931 })?;
932
933 let range = checked_file_range(
934 posting_offset,
935 posting_len,
936 self.postings_handle.len(),
937 "posting",
938 )?;
939 let posting_bytes = self.postings_handle.read_bytes_range(range).await?;
940 let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
941
942 Ok(Some(block_list))
943 }
944
945 pub async fn get_prefix_postings(
947 &self,
948 field: Field,
949 prefix: &[u8],
950 ) -> Result<Vec<BlockPostingList>> {
951 if prefix.is_empty() {
952 return Err(Error::Query("prefix must not be empty".into()));
953 }
954 let mut key_prefix = Vec::with_capacity(4 + prefix.len());
956 key_prefix.extend_from_slice(&field.0.to_le_bytes());
957 key_prefix.extend_from_slice(prefix);
958
959 let (entries, truncated) = self
960 .term_dict
961 .prefix_scan_limited(&key_prefix, MAX_PREFIX_TERMS)
962 .await?;
963 if truncated {
964 return Err(Error::Query(format!(
965 "prefix expands to more than {MAX_PREFIX_TERMS} terms"
966 )));
967 }
968 let posting_count: u64 = entries
969 .iter()
970 .map(|(_, term_info)| term_info.doc_freq() as u64)
971 .sum();
972 if posting_count > MAX_PREFIX_POSTINGS {
973 return Err(Error::Query(format!(
974 "prefix expands to {posting_count} postings (maximum {MAX_PREFIX_POSTINGS})"
975 )));
976 }
977 let mut results = Vec::with_capacity(entries.len());
978
979 for (_key, term_info) in entries {
980 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
981 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
982 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
983 posting_list.push(doc_id, tf);
984 }
985 results.push(BlockPostingList::from_posting_list(&posting_list)?);
986 } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
987 let range = checked_file_range(
988 posting_offset,
989 posting_len,
990 self.postings_handle.len(),
991 "prefix posting",
992 )?;
993 let posting_bytes = self.postings_handle.read_bytes_range(range).await?;
994 results.push(BlockPostingList::deserialize_zero_copy(posting_bytes)?);
995 }
996 }
997
998 Ok(results)
999 }
1000
1001 pub async fn doc(&self, local_doc_id: DocId) -> Result<Option<Document>> {
1006 self.doc_with_fields(local_doc_id, None).await
1007 }
1008
1009 pub async fn doc_with_fields(
1015 &self,
1016 local_doc_id: DocId,
1017 fields: Option<&rustc_hash::FxHashSet<u32>>,
1018 ) -> Result<Option<Document>> {
1019 let mut doc = match fields {
1020 Some(set) => {
1021 let field_ids: Vec<u32> = set.iter().copied().collect();
1022 match self
1023 .store
1024 .get_fields(local_doc_id, &self.schema, &field_ids)
1025 .await
1026 {
1027 Ok(Some(d)) => d,
1028 Ok(None) => return Ok(None),
1029 Err(e) => return Err(Error::from(e)),
1030 }
1031 }
1032 None => match self.store.get(local_doc_id, &self.schema).await {
1033 Ok(Some(d)) => d,
1034 Ok(None) => return Ok(None),
1035 Err(e) => return Err(Error::from(e)),
1036 },
1037 };
1038
1039 for (&field_id, lazy_flat) in &self.flat_vectors {
1041 if let Some(set) = fields
1043 && !set.contains(&field_id)
1044 {
1045 continue;
1046 }
1047
1048 let is_binary = lazy_flat.quantization == DenseVectorQuantization::Binary;
1049 let (start, entries) = lazy_flat.flat_indexes_for_doc(local_doc_id);
1050 for (j, &(_doc_id, _ordinal)) in entries.iter().enumerate() {
1051 let flat_idx = start + j;
1052 if is_binary {
1053 let vbs = lazy_flat.vector_byte_size();
1054 let mut raw = vec![0u8; vbs];
1055 match lazy_flat.read_vector_raw_into(flat_idx, &mut raw).await {
1056 Ok(()) => {
1057 doc.add_binary_dense_vector(Field(field_id), raw);
1058 }
1059 Err(e) => {
1060 log::warn!("Failed to hydrate binary vector field {}: {}", field_id, e);
1061 }
1062 }
1063 } else {
1064 match lazy_flat.get_vector(flat_idx).await {
1065 Ok(vec) => {
1066 doc.add_dense_vector(Field(field_id), vec);
1067 }
1068 Err(e) => {
1069 log::warn!("Failed to hydrate vector field {}: {}", field_id, e);
1070 }
1071 }
1072 }
1073 }
1074 }
1075
1076 Ok(Some(doc))
1077 }
1078
1079 pub async fn prefetch_terms(
1081 &self,
1082 field: Field,
1083 start_term: &[u8],
1084 end_term: &[u8],
1085 ) -> Result<()> {
1086 let mut start_key = Vec::with_capacity(4 + start_term.len());
1087 start_key.extend_from_slice(&field.0.to_le_bytes());
1088 start_key.extend_from_slice(start_term);
1089
1090 let mut end_key = Vec::with_capacity(4 + end_term.len());
1091 end_key.extend_from_slice(&field.0.to_le_bytes());
1092 end_key.extend_from_slice(end_term);
1093
1094 self.term_dict.prefetch_range(&start_key, &end_key).await?;
1095 Ok(())
1096 }
1097
1098 pub fn store_has_dict(&self) -> bool {
1100 self.store.has_dict()
1101 }
1102
1103 pub fn store(&self) -> &super::store::AsyncStoreReader {
1105 &self.store
1106 }
1107
1108 pub fn store_raw_blocks(&self) -> Vec<RawStoreBlock> {
1110 self.store.raw_blocks()
1111 }
1112
1113 pub fn store_data_slice(&self) -> &FileHandle {
1115 self.store.data_slice()
1116 }
1117
1118 pub async fn all_terms(&self) -> Result<Vec<(Vec<u8>, TermInfo)>> {
1120 self.term_dict.all_entries().await.map_err(Error::from)
1121 }
1122
1123 pub async fn all_terms_with_stats(&self) -> Result<Vec<(Field, String, u32)>> {
1128 let entries = self.term_dict.all_entries().await?;
1129 let mut result = Vec::with_capacity(entries.len());
1130
1131 for (key, term_info) in entries {
1132 if key.len() > 4 {
1134 let field_id = u32::from_le_bytes([key[0], key[1], key[2], key[3]]);
1135 let term_bytes = &key[4..];
1136 if let Ok(term_str) = std::str::from_utf8(term_bytes) {
1137 result.push((Field(field_id), term_str.to_string(), term_info.doc_freq()));
1138 }
1139 }
1140 }
1141
1142 Ok(result)
1143 }
1144
1145 pub fn term_dict_iter(&self) -> crate::structures::AsyncSSTableIterator<'_, TermInfo> {
1147 self.term_dict.iter()
1148 }
1149
1150 pub async fn prefetch_term_dict(&self) -> crate::Result<()> {
1154 self.term_dict
1155 .prefetch_all_data_bulk()
1156 .await
1157 .map_err(crate::Error::from)
1158 }
1159
1160 pub async fn read_postings(&self, offset: u64, len: u64) -> Result<Vec<u8>> {
1162 let range = checked_file_range(offset, len, self.postings_handle.len(), "posting")?;
1163 let bytes = self.postings_handle.read_bytes_range(range).await?;
1164 Ok(bytes.to_vec())
1165 }
1166
1167 pub async fn read_position_bytes(&self, offset: u64, len: u64) -> Result<Option<Vec<u8>>> {
1169 let handle = match &self.positions_handle {
1170 Some(h) => h,
1171 None => return Ok(None),
1172 };
1173 let range = checked_file_range(offset, len, handle.len(), "position")?;
1174 let bytes = handle.read_bytes_range(range).await?;
1175 Ok(Some(bytes.to_vec()))
1176 }
1177
1178 pub fn has_positions_file(&self) -> bool {
1180 self.positions_handle.is_some()
1181 }
1182
1183 fn validate_dense_search_request(
1187 &self,
1188 field: Field,
1189 query: &[f32],
1190 nprobe: usize,
1191 rerank_factor: f32,
1192 combiner: crate::query::MultiValueCombiner,
1193 ) -> Result<DenseSearchParams> {
1194 let entry = self
1195 .schema
1196 .get_field_entry(field)
1197 .ok_or_else(|| Error::FieldNotFound(field.0.to_string()))?;
1198 if entry.field_type != crate::dsl::FieldType::DenseVector {
1199 return Err(Error::InvalidFieldType {
1200 expected: "dense_vector".to_string(),
1201 got: format!("{:?}", entry.field_type),
1202 });
1203 }
1204 let config = entry.dense_vector_config.as_ref().ok_or_else(|| {
1205 Error::Schema(format!(
1206 "dense vector field '{}' has no dense vector configuration",
1207 entry.name
1208 ))
1209 })?;
1210
1211 if query.is_empty() {
1212 return Err(Error::Query(format!(
1213 "dense query vector for field '{}' must not be empty",
1214 entry.name
1215 )));
1216 }
1217 if query.len() != config.dim {
1218 return Err(Error::Query(format!(
1219 "dense query vector dimension {} does not match field '{}' dimension {}",
1220 query.len(),
1221 entry.name,
1222 config.dim
1223 )));
1224 }
1225 if let Some((index, value)) = query
1226 .iter()
1227 .enumerate()
1228 .find(|(_, value)| !value.is_finite())
1229 {
1230 return Err(Error::Query(format!(
1231 "dense query vector for field '{}' contains non-finite value {value} at index {index}",
1232 entry.name
1233 )));
1234 }
1235
1236 let nprobe = match (nprobe, config.nprobe) {
1239 (0, 0) => 32,
1240 (0, schema_nprobe) => schema_nprobe,
1241 (query_nprobe, _) => query_nprobe,
1242 };
1243 if nprobe > MAX_DENSE_NPROBE {
1244 return Err(Error::Query(format!(
1245 "dense nprobe must be at most {MAX_DENSE_NPROBE}, got {nprobe}"
1246 )));
1247 }
1248
1249 checked_dense_fetch_k(0, rerank_factor)?;
1252 combiner.validate().map_err(Error::Query)?;
1253
1254 Ok(DenseSearchParams {
1255 dim: config.dim,
1256 nprobe,
1257 unit_norm: config.unit_norm,
1258 })
1259 }
1260
1261 fn validate_binary_search_request(&self, field: Field, query: &[u8]) -> Result<usize> {
1262 let entry = self
1263 .schema
1264 .get_field_entry(field)
1265 .ok_or_else(|| Error::FieldNotFound(field.0.to_string()))?;
1266 if entry.field_type != crate::dsl::FieldType::BinaryDenseVector {
1267 return Err(Error::InvalidFieldType {
1268 expected: "binary_dense_vector".to_string(),
1269 got: format!("{:?}", entry.field_type),
1270 });
1271 }
1272 let config = entry.binary_dense_vector_config.as_ref().ok_or_else(|| {
1273 Error::Schema(format!(
1274 "binary dense vector field '{}' has no configuration",
1275 entry.name
1276 ))
1277 })?;
1278 if config.dim == 0 || !config.dim.is_multiple_of(8) {
1279 return Err(Error::Schema(format!(
1280 "binary dense vector field '{}' has invalid dimension {}",
1281 entry.name, config.dim
1282 )));
1283 }
1284 if query.len() != config.byte_len() {
1285 return Err(Error::Query(format!(
1286 "binary query byte length {} does not match field '{}' byte length {}",
1287 query.len(),
1288 entry.name,
1289 config.byte_len()
1290 )));
1291 }
1292 Ok(config.dim)
1293 }
1294
1295 fn score_quantized_batch(
1301 query: &[f32],
1302 raw: &[u8],
1303 quant: crate::dsl::DenseVectorQuantization,
1304 dim: usize,
1305 scores: &mut [f32],
1306 unit_norm: bool,
1307 ) -> Result<()> {
1308 use crate::dsl::DenseVectorQuantization;
1309 use crate::structures::simd;
1310
1311 if query.len() != dim {
1312 return Err(Error::Query(format!(
1313 "dense SIMD query dimension {} does not match vector dimension {dim}",
1314 query.len()
1315 )));
1316 }
1317 let element_size = match quant {
1318 DenseVectorQuantization::F32 => std::mem::size_of::<f32>(),
1319 DenseVectorQuantization::F16 => std::mem::size_of::<u16>(),
1320 DenseVectorQuantization::UInt8 => 1,
1321 DenseVectorQuantization::Binary => {
1322 return Err(Error::InvalidFieldType {
1323 expected: "non-binary dense vector".to_string(),
1324 got: "binary dense vector".to_string(),
1325 });
1326 }
1327 };
1328 let required_bytes = scores
1329 .len()
1330 .checked_mul(dim)
1331 .and_then(|elements| elements.checked_mul(element_size))
1332 .ok_or_else(|| Error::Corruption("dense vector batch byte length overflow".into()))?;
1333 if raw.len() < required_bytes {
1334 return Err(Error::Corruption(format!(
1335 "dense vector batch is truncated: need {required_bytes} bytes, got {}",
1336 raw.len()
1337 )));
1338 }
1339 if quant == DenseVectorQuantization::F16
1340 && required_bytes > 0
1341 && !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>())
1342 {
1343 return Err(Error::Corruption(
1344 "f16 vector data is not 2-byte aligned".to_string(),
1345 ));
1346 }
1347
1348 match (quant, unit_norm) {
1349 (DenseVectorQuantization::F32, false) => {
1350 let num_floats = scores.len() * dim;
1351 if !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()) {
1352 return Err(Error::Corruption(
1353 "f32 vector data is not 4-byte aligned".to_string(),
1354 ));
1355 }
1356 let vectors: &[f32] =
1357 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
1358 simd::batch_cosine_scores(query, vectors, dim, scores);
1359 }
1360 (DenseVectorQuantization::F32, true) => {
1361 let num_floats = scores.len() * dim;
1362 if !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()) {
1363 return Err(Error::Corruption(
1364 "f32 vector data is not 4-byte aligned".to_string(),
1365 ));
1366 }
1367 let vectors: &[f32] =
1368 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
1369 simd::batch_dot_scores(query, vectors, dim, scores);
1370 }
1371 (DenseVectorQuantization::F16, false) => {
1372 simd::batch_cosine_scores_f16(query, raw, dim, scores);
1373 }
1374 (DenseVectorQuantization::F16, true) => {
1375 simd::batch_dot_scores_f16(query, raw, dim, scores);
1376 }
1377 (DenseVectorQuantization::UInt8, false) => {
1378 simd::batch_cosine_scores_u8(query, raw, dim, scores);
1379 }
1380 (DenseVectorQuantization::UInt8, true) => {
1381 simd::batch_dot_scores_u8(query, raw, dim, scores);
1382 }
1383 (DenseVectorQuantization::Binary, _) => unreachable!("validated above"),
1384 }
1385 Ok(())
1386 }
1387
1388 pub async fn search_dense_vector(
1394 &self,
1395 field: Field,
1396 query: &[f32],
1397 k: usize,
1398 nprobe: usize,
1399 rerank_factor: f32,
1400 combiner: crate::query::MultiValueCombiner,
1401 ) -> Result<Vec<VectorSearchResult>> {
1402 let params =
1403 self.validate_dense_search_request(field, query, nprobe, rerank_factor, combiner)?;
1404 let fetch_k = checked_dense_fetch_k(k, rerank_factor)?;
1405 if k == 0 {
1406 return Ok(Vec::new());
1407 }
1408
1409 let ann_index = self.vector_indexes.get(&field.0);
1410 let lazy_flat = self.flat_vectors.get(&field.0);
1411 let ann_fetch_k = lazy_flat.map_or(fetch_k, |flat| {
1412 ann_ordinal_fetch_k(fetch_k, flat.num_vectors, flat.num_docs_with_vectors())
1413 });
1414
1415 if ann_index.is_none() && lazy_flat.is_none() {
1417 return Ok(Vec::new());
1418 }
1419
1420 if ann_index.is_some() && lazy_flat.is_none() {
1421 return Err(Error::Corruption(format!(
1422 "dense ANN field {} is missing flat vector storage",
1423 field.0
1424 )));
1425 }
1426
1427 if let Some(flat) = lazy_flat
1428 && flat.dim != params.dim
1429 {
1430 return Err(Error::Corruption(format!(
1431 "dense vector field {} has schema dimension {} but flat storage dimension {}",
1432 field.0, params.dim, flat.dim
1433 )));
1434 }
1435
1436 let t0 = std::time::Instant::now();
1438 let mut flat_results = None;
1439 let mut results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
1440 match index {
1442 VectorIndex::RaBitQ(lazy) => {
1443 let rabitq = lazy.get().ok_or_else(|| {
1444 Error::Schema("RaBitQ index deserialization failed".to_string())
1445 })?;
1446 if rabitq.codebook.config.dim != params.dim {
1447 return Err(Error::Corruption(format!(
1448 "RaBitQ index dimension {} does not match schema dimension {}",
1449 rabitq.codebook.config.dim, params.dim
1450 )));
1451 }
1452 let flat = lazy_flat.expect("ANN/flat pairing validated above");
1453 progressive_ann_search(
1454 fetch_k.min(flat.num_docs_with_vectors()),
1455 ann_fetch_k,
1456 rabitq.len().min(flat.num_vectors),
1457 |candidate_k| {
1458 rabitq
1459 .search(query, candidate_k)
1460 .into_iter()
1461 .map(|(doc_id, ordinal, dist)| {
1462 (doc_id, ordinal, 1.0 / (1.0 + dist))
1463 })
1464 .collect()
1465 },
1466 )?
1467 }
1468 VectorIndex::IVF(lazy) => {
1469 let (index, codebook) = lazy.get().ok_or_else(|| {
1470 Error::Schema("IVF index deserialization failed".to_string())
1471 })?;
1472 let centroids = self.coarse_centroids.get(&field.0).ok_or_else(|| {
1473 Error::Schema(format!(
1474 "IVF index requires coarse centroids for field {}",
1475 field.0
1476 ))
1477 })?;
1478 validate_coarse_centroids(centroids, params.dim)?;
1479 if index.config.dim != params.dim
1480 || codebook.config.dim != params.dim
1481 || index.centroids_version != centroids.version
1482 || index.codebook_version != codebook.version
1483 || index
1484 .clusters
1485 .iter()
1486 .any(|(cluster_id, _)| cluster_id >= centroids.num_clusters)
1487 {
1488 return Err(Error::Corruption(format!(
1489 "IVF index/codebook/centroid metadata does not match schema dimension {}",
1490 params.dim
1491 )));
1492 }
1493 let effective_nprobe = params.nprobe.min(centroids.num_clusters as usize);
1494 let flat = lazy_flat.expect("ANN/flat pairing validated above");
1495 progressive_ann_search(
1496 fetch_k.min(flat.num_docs_with_vectors()),
1497 ann_fetch_k,
1498 index.len().min(flat.num_vectors),
1499 |candidate_k| {
1500 index
1501 .search(
1502 centroids,
1503 codebook,
1504 query,
1505 candidate_k,
1506 Some(effective_nprobe),
1507 )
1508 .into_iter()
1509 .map(|(doc_id, ordinal, dist)| {
1510 (doc_id, ordinal, 1.0 / (1.0 + dist))
1511 })
1512 .collect()
1513 },
1514 )?
1515 }
1516 VectorIndex::ScaNN(lazy) => {
1517 let (index, codebook) = lazy.get().ok_or_else(|| {
1518 Error::Schema("ScaNN index deserialization failed".to_string())
1519 })?;
1520 let centroids = self.coarse_centroids.get(&field.0).ok_or_else(|| {
1521 Error::Schema(format!(
1522 "ScaNN index requires coarse centroids for field {}",
1523 field.0
1524 ))
1525 })?;
1526 validate_coarse_centroids(centroids, params.dim)?;
1527 if index.config.dim != params.dim
1528 || codebook.config.dim != params.dim
1529 || index.centroids_version != centroids.version
1530 || index.codebook_version != codebook.version
1531 || index
1532 .clusters
1533 .iter()
1534 .any(|(cluster_id, _)| cluster_id >= centroids.num_clusters)
1535 {
1536 return Err(Error::Corruption(format!(
1537 "ScaNN index/codebook/centroid metadata does not match schema dimension {}",
1538 params.dim
1539 )));
1540 }
1541 let effective_nprobe = params.nprobe.min(centroids.num_clusters as usize);
1542 let flat = lazy_flat.expect("ANN/flat pairing validated above");
1543 progressive_ann_search(
1544 fetch_k.min(flat.num_docs_with_vectors()),
1545 ann_fetch_k,
1546 index.len().min(flat.num_vectors),
1547 |candidate_k| {
1548 index
1549 .search(
1550 centroids,
1551 codebook,
1552 query,
1553 candidate_k,
1554 Some(effective_nprobe),
1555 )
1556 .into_iter()
1557 .map(|(doc_id, ordinal, dist)| {
1558 (doc_id, ordinal, 1.0 / (1.0 + dist))
1559 })
1560 .collect()
1561 },
1562 )?
1563 }
1564 VectorIndex::BinaryIvf(_) => {
1565 Vec::new()
1567 }
1568 }
1569 } else if let Some(lazy_flat) = lazy_flat {
1570 log::debug!(
1574 "[search_dense] field {}: brute-force on {} vectors (dim={}, quant={:?})",
1575 field.0,
1576 lazy_flat.num_vectors,
1577 lazy_flat.dim,
1578 lazy_flat.quantization
1579 );
1580 let dim = lazy_flat.dim;
1581 let n = lazy_flat.num_vectors;
1582 let quant = lazy_flat.quantization;
1583 let batch_len =
1584 bounded_vector_score_batch(lazy_flat.vector_byte_size(), DENSE_SCORE_BATCH);
1585 let mut collector = FlatDocumentCollector::new(fetch_k.min(n), combiner);
1586 let mut scores = vec![0f32; batch_len];
1587
1588 for batch_start in (0..n).step_by(batch_len) {
1589 let batch_count = batch_len.min(n - batch_start);
1590 let batch_bytes = lazy_flat
1591 .read_vectors_batch(batch_start, batch_count)
1592 .await
1593 .map_err(crate::Error::Io)?;
1594 let raw = batch_bytes.as_slice();
1595
1596 Self::score_quantized_batch(
1597 query,
1598 raw,
1599 quant,
1600 dim,
1601 &mut scores[..batch_count],
1602 params.unit_norm,
1603 )?;
1604
1605 for (i, &score) in scores.iter().enumerate().take(batch_count) {
1606 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
1607 collector.push(doc_id, ordinal, score);
1608 }
1609 }
1610
1611 flat_results = Some(collector.into_results());
1612 Vec::new()
1613 } else {
1614 return Ok(Vec::new());
1615 };
1616 let l1_elapsed = t0.elapsed();
1617 {
1618 let kind = match ann_index {
1619 Some(VectorIndex::RaBitQ(_)) => "rabitq",
1620 Some(VectorIndex::IVF(_)) => "ivf_rabitq",
1621 Some(VectorIndex::ScaNN(_)) => "scann",
1622 Some(VectorIndex::BinaryIvf(_)) => "binary_ivf",
1623 None => "flat",
1624 };
1625 crate::observe::dense_l1(
1626 self.schema.index_label(),
1627 self.schema.get_field_name(field).unwrap_or("?"),
1628 kind,
1629 l1_elapsed.as_secs_f64(),
1630 flat_results.as_ref().map_or(results.len(), Vec::len),
1631 );
1632 }
1633 log::debug!(
1634 "[search_dense] field {}: L1 returned {} candidates in {:.1}ms",
1635 field.0,
1636 flat_results.as_ref().map_or(results.len(), Vec::len),
1637 l1_elapsed.as_secs_f64() * 1000.0
1638 );
1639
1640 if let Some(results) = flat_results {
1641 return Ok(results);
1642 }
1643
1644 if ann_index.is_some()
1647 && !results.is_empty()
1648 && let Some(lazy_flat) = lazy_flat
1649 {
1650 let t_rerank = std::time::Instant::now();
1651 let dim = lazy_flat.dim;
1652 let quant = lazy_flat.quantization;
1653 let vbs = lazy_flat.vector_byte_size();
1654
1655 let (expanded, mut resolved) = expand_ann_candidate_documents(&results, lazy_flat)?;
1659 results = expanded;
1660
1661 let t_resolve = t_rerank.elapsed();
1662 if !resolved.is_empty() {
1663 resolved.sort_unstable_by_key(|&(_, flat_idx)| flat_idx);
1665
1666 let batch_len = bounded_vector_score_batch(vbs, DENSE_SCORE_BATCH);
1670 let max_batch = batch_len.min(resolved.len());
1671 let max_raw_len = max_batch
1672 .checked_mul(vbs)
1673 .ok_or_else(|| Error::Query("dense rerank buffer size overflow".into()))?;
1674 let mut raw_buf = vec![0u8; max_raw_len];
1675 let mut scores = vec![0f32; max_batch];
1676 let mut read_elapsed = std::time::Duration::ZERO;
1677 let mut score_elapsed = std::time::Duration::ZERO;
1678
1679 for chunk in resolved.chunks(batch_len) {
1680 #[cfg(feature = "native")]
1684 lazy_flat.prefetch_vectors(chunk.iter().map(|&(_, flat_idx)| flat_idx));
1685 let raw_len = chunk
1686 .len()
1687 .checked_mul(vbs)
1688 .ok_or_else(|| Error::Query("dense rerank buffer size overflow".into()))?;
1689 let raw = &mut raw_buf[..raw_len];
1690
1691 let t_read = std::time::Instant::now();
1692 for (buf_idx, &(_, flat_idx)) in chunk.iter().enumerate() {
1693 lazy_flat
1694 .read_vector_raw_into(
1695 flat_idx,
1696 &mut raw[buf_idx * vbs..(buf_idx + 1) * vbs],
1697 )
1698 .await
1699 .map_err(crate::Error::Io)?;
1700 }
1701 read_elapsed += t_read.elapsed();
1702
1703 let t_score = std::time::Instant::now();
1704 Self::score_quantized_batch(
1705 query,
1706 raw,
1707 quant,
1708 dim,
1709 &mut scores[..chunk.len()],
1710 params.unit_norm,
1711 )?;
1712 score_elapsed += t_score.elapsed();
1713
1714 for (buf_idx, &(ri, _)) in chunk.iter().enumerate() {
1715 results[ri].2 = scores[buf_idx];
1716 }
1717 }
1718
1719 crate::observe::dense_rerank(
1720 self.schema.index_label(),
1721 self.schema.get_field_name(field).unwrap_or("?"),
1722 t_rerank.elapsed().as_secs_f64(),
1723 t_resolve.as_secs_f64(),
1724 read_elapsed.as_secs_f64(),
1725 resolved.len(),
1726 );
1727 log::debug!(
1728 "[search_dense] field {}: rerank {} vectors (dim={}, quant={:?}, {}B/vec): resolve={:.1}ms read={:.1}ms score={:.1}ms",
1729 field.0,
1730 resolved.len(),
1731 dim,
1732 quant,
1733 vbs,
1734 t_resolve.as_secs_f64() * 1000.0,
1735 read_elapsed.as_secs_f64() * 1000.0,
1736 score_elapsed.as_secs_f64() * 1000.0,
1737 );
1738 }
1739
1740 log::debug!(
1741 "[search_dense] field {}: rerank total={:.1}ms",
1742 field.0,
1743 t_rerank.elapsed().as_secs_f64() * 1000.0
1744 );
1745 }
1746
1747 Ok(combine_grouped_ordinal_results(results, combiner, k))
1748 }
1749
1750 fn binary_ivf_nprobe(&self, field: Field) -> Option<usize> {
1752 self.schema
1753 .get_field_entry(field)
1754 .and_then(|e| e.binary_dense_vector_config.as_ref())
1755 .map(|c| c.nprobe)
1756 .filter(|&n| n > 0)
1757 }
1758
1759 pub async fn search_binary_dense_vector(
1763 &self,
1764 field: Field,
1765 query: &[u8],
1766 k: usize,
1767 combiner: crate::query::MultiValueCombiner,
1768 ) -> Result<Vec<VectorSearchResult>> {
1769 let schema_dim = self.validate_binary_search_request(field, query)?;
1770 combiner.validate().map_err(Error::Query)?;
1771 if k == 0 {
1772 return Ok(Vec::new());
1773 }
1774 let t0 = crate::observe::Timer::start();
1775 if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0)
1779 && let Some(ivf) = lazy.get()
1780 {
1781 if ivf.config.dim_bits != schema_dim {
1782 return Err(Error::Corruption(format!(
1783 "binary IVF field {} has schema dimension {} but index dimension {}",
1784 field.0, schema_dim, ivf.config.dim_bits
1785 )));
1786 }
1787 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
1788 Error::Corruption(format!(
1789 "binary IVF field {} is missing flat vector storage",
1790 field.0
1791 ))
1792 })?;
1793 if flat.dim != schema_dim
1794 || flat.quantization != crate::dsl::DenseVectorQuantization::Binary
1795 {
1796 return Err(Error::Corruption(format!(
1797 "binary IVF field {} has inconsistent flat vector metadata",
1798 field.0
1799 )));
1800 }
1801 let nprobe = self.binary_ivf_nprobe(field);
1802 let initial_fetch =
1803 ann_ordinal_fetch_k(k, flat.num_vectors, flat.num_docs_with_vectors());
1804 let ann_results = progressive_ann_search(
1805 k.min(flat.num_docs_with_vectors()),
1806 initial_fetch,
1807 ivf.len().min(flat.num_vectors),
1808 |candidate_k| ivf.search(query, candidate_k, nprobe),
1809 )?;
1810 let results =
1811 exact_score_binary_candidate_documents(&ann_results, flat, query, schema_dim)
1812 .await?;
1813 crate::observe::dense_l1(
1814 self.schema.index_label(),
1815 self.schema.get_field_name(field).unwrap_or("?"),
1816 "binary_ivf",
1817 t0.secs(),
1818 results.len(),
1819 );
1820 return Ok(combine_grouped_ordinal_results(results, combiner, k));
1821 }
1822
1823 let lazy_flat = match self.flat_vectors.get(&field.0) {
1824 Some(f) => f,
1825 None => return Ok(Vec::new()),
1826 };
1827
1828 let dim_bits = lazy_flat.dim;
1829 let byte_len = lazy_flat.vector_byte_size();
1830 let n = lazy_flat.num_vectors;
1831
1832 if dim_bits != schema_dim {
1833 return Err(Error::Corruption(format!(
1834 "binary vector field {} has schema dimension {} but flat storage dimension {}",
1835 field.0, schema_dim, dim_bits
1836 )));
1837 }
1838
1839 if byte_len != query.len() {
1840 return Err(Error::Schema(format!(
1841 "Binary query vector byte length {} != field byte length {}",
1842 query.len(),
1843 byte_len
1844 )));
1845 }
1846
1847 let batch_len = bounded_vector_score_batch(byte_len, BINARY_SCORE_BATCH);
1848 let mut collector = FlatDocumentCollector::new(k, combiner);
1849 let mut scores = vec![0f32; batch_len];
1850
1851 for batch_start in (0..n).step_by(batch_len) {
1852 let batch_count = batch_len.min(n - batch_start);
1853 let batch_bytes = lazy_flat
1854 .read_vectors_batch(batch_start, batch_count)
1855 .await
1856 .map_err(crate::Error::Io)?;
1857 let raw = batch_bytes.as_slice();
1858
1859 crate::structures::simd::batch_hamming_scores(
1860 query,
1861 raw,
1862 byte_len,
1863 dim_bits,
1864 &mut scores[..batch_count],
1865 );
1866
1867 for (i, &score) in scores.iter().enumerate().take(batch_count) {
1868 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
1869 collector.push(doc_id, ordinal, score);
1870 }
1871 }
1872
1873 let results = collector.into_results();
1874
1875 crate::observe::dense_l1(
1876 self.schema.index_label(),
1877 self.schema.get_field_name(field).unwrap_or("?"),
1878 "binary_flat",
1879 t0.secs(),
1880 results.len(),
1881 );
1882 Ok(results)
1883 }
1884
1885 pub fn has_dense_vector_index(&self, field: Field) -> bool {
1887 self.vector_indexes.contains_key(&field.0) || self.flat_vectors.contains_key(&field.0)
1888 }
1889
1890 pub fn get_dense_vector_index(&self, field: Field) -> Option<Arc<RaBitQIndex>> {
1892 match self.vector_indexes.get(&field.0) {
1893 Some(VectorIndex::RaBitQ(lazy)) => lazy.get().cloned(),
1894 _ => None,
1895 }
1896 }
1897
1898 pub fn get_ivf_vector_index(
1900 &self,
1901 field: Field,
1902 ) -> Option<(Arc<IVFRaBitQIndex>, Arc<crate::structures::RaBitQCodebook>)> {
1903 match self.vector_indexes.get(&field.0) {
1904 Some(VectorIndex::IVF(lazy)) => lazy.get().map(|(i, c)| (i.clone(), c.clone())),
1905 _ => None,
1906 }
1907 }
1908
1909 pub fn coarse_centroids(&self, field_id: u32) -> Option<&Arc<CoarseCentroids>> {
1911 self.coarse_centroids.get(&field_id)
1912 }
1913
1914 pub fn set_coarse_centroids(&mut self, centroids: FxHashMap<u32, Arc<CoarseCentroids>>) {
1916 self.coarse_centroids = centroids;
1917 }
1918
1919 pub fn get_scann_vector_index(
1921 &self,
1922 field: Field,
1923 ) -> Option<(Arc<IVFPQIndex>, Arc<PQCodebook>)> {
1924 match self.vector_indexes.get(&field.0) {
1925 Some(VectorIndex::ScaNN(lazy)) => lazy.get().map(|(i, c)| (i.clone(), c.clone())),
1926 _ => None,
1927 }
1928 }
1929
1930 pub fn get_vector_index(&self, field: Field) -> Option<&VectorIndex> {
1932 self.vector_indexes.get(&field.0)
1933 }
1934
1935 pub async fn get_positions(
1940 &self,
1941 field: Field,
1942 term: &[u8],
1943 ) -> Result<Option<crate::structures::PositionPostingList>> {
1944 let handle = match &self.positions_handle {
1946 Some(h) => h,
1947 None => return Ok(None),
1948 };
1949
1950 let mut key = Vec::with_capacity(4 + term.len());
1952 key.extend_from_slice(&field.0.to_le_bytes());
1953 key.extend_from_slice(term);
1954
1955 let term_info = match self.term_dict.get(&key).await? {
1957 Some(info) => info,
1958 None => return Ok(None),
1959 };
1960
1961 let (offset, length) = match term_info.position_info() {
1963 Some((o, l)) => (o, l),
1964 None => return Ok(None),
1965 };
1966
1967 let range = checked_file_range(offset, length, handle.len(), "position list")?;
1971 let slice = handle.slice(range);
1972 let data = slice.read_bytes().await?;
1973
1974 let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
1976
1977 Ok(Some(pos_list))
1978 }
1979
1980 pub fn has_positions(&self, field: Field) -> bool {
1982 if let Some(entry) = self.schema.get_field_entry(field) {
1984 entry.positions.is_some()
1985 } else {
1986 false
1987 }
1988 }
1989}
1990
1991#[cfg(feature = "sync")]
1993impl SegmentReader {
1994 pub fn get_postings_sync(&self, field: Field, term: &[u8]) -> Result<Option<BlockPostingList>> {
1996 let mut key = Vec::with_capacity(4 + term.len());
1998 key.extend_from_slice(&field.0.to_le_bytes());
1999 key.extend_from_slice(term);
2000
2001 let term_info = match self.term_dict.get_sync(&key)? {
2003 Some(info) => info,
2004 None => return Ok(None),
2005 };
2006
2007 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
2009 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
2010 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
2011 posting_list.push(doc_id, tf);
2012 }
2013 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
2014 return Ok(Some(block_list));
2015 }
2016
2017 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
2019 Error::Corruption("TermInfo has neither inline nor external data".to_string())
2020 })?;
2021
2022 let range = checked_file_range(
2023 posting_offset,
2024 posting_len,
2025 self.postings_handle.len(),
2026 "posting",
2027 )?;
2028 let posting_bytes = self.postings_handle.read_bytes_range_sync(range)?;
2029 let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
2030
2031 Ok(Some(block_list))
2032 }
2033
2034 pub fn get_prefix_postings_sync(
2036 &self,
2037 field: Field,
2038 prefix: &[u8],
2039 ) -> Result<Vec<BlockPostingList>> {
2040 if prefix.is_empty() {
2041 return Err(Error::Query("prefix must not be empty".into()));
2042 }
2043 let mut key_prefix = Vec::with_capacity(4 + prefix.len());
2044 key_prefix.extend_from_slice(&field.0.to_le_bytes());
2045 key_prefix.extend_from_slice(prefix);
2046
2047 let (entries, truncated) = self
2048 .term_dict
2049 .prefix_scan_limited_sync(&key_prefix, MAX_PREFIX_TERMS)?;
2050 if truncated {
2051 return Err(Error::Query(format!(
2052 "prefix expands to more than {MAX_PREFIX_TERMS} terms"
2053 )));
2054 }
2055 let posting_count: u64 = entries
2056 .iter()
2057 .map(|(_, term_info)| term_info.doc_freq() as u64)
2058 .sum();
2059 if posting_count > MAX_PREFIX_POSTINGS {
2060 return Err(Error::Query(format!(
2061 "prefix expands to {posting_count} postings (maximum {MAX_PREFIX_POSTINGS})"
2062 )));
2063 }
2064 let mut results = Vec::with_capacity(entries.len());
2065
2066 for (_key, term_info) in entries {
2067 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
2068 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
2069 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
2070 posting_list.push(doc_id, tf);
2071 }
2072 results.push(BlockPostingList::from_posting_list(&posting_list)?);
2073 } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
2074 let range = checked_file_range(
2075 posting_offset,
2076 posting_len,
2077 self.postings_handle.len(),
2078 "prefix posting",
2079 )?;
2080 let posting_bytes = self.postings_handle.read_bytes_range_sync(range)?;
2081 results.push(BlockPostingList::deserialize_zero_copy(posting_bytes)?);
2082 }
2083 }
2084
2085 Ok(results)
2086 }
2087
2088 pub fn get_positions_sync(
2090 &self,
2091 field: Field,
2092 term: &[u8],
2093 ) -> Result<Option<crate::structures::PositionPostingList>> {
2094 let handle = match &self.positions_handle {
2095 Some(h) => h,
2096 None => return Ok(None),
2097 };
2098
2099 let mut key = Vec::with_capacity(4 + term.len());
2101 key.extend_from_slice(&field.0.to_le_bytes());
2102 key.extend_from_slice(term);
2103
2104 let term_info = match self.term_dict.get_sync(&key)? {
2106 Some(info) => info,
2107 None => return Ok(None),
2108 };
2109
2110 let (offset, length) = match term_info.position_info() {
2111 Some((o, l)) => (o, l),
2112 None => return Ok(None),
2113 };
2114
2115 let range = checked_file_range(offset, length, handle.len(), "position list")?;
2116 let slice = handle.slice(range);
2117 let data = slice.read_bytes_sync()?;
2118
2119 let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
2120 Ok(Some(pos_list))
2121 }
2122
2123 pub fn search_dense_vector_sync(
2126 &self,
2127 field: Field,
2128 query: &[f32],
2129 k: usize,
2130 nprobe: usize,
2131 rerank_factor: f32,
2132 combiner: crate::query::MultiValueCombiner,
2133 ) -> Result<Vec<VectorSearchResult>> {
2134 let params =
2135 self.validate_dense_search_request(field, query, nprobe, rerank_factor, combiner)?;
2136 let fetch_k = checked_dense_fetch_k(k, rerank_factor)?;
2137 if k == 0 {
2138 return Ok(Vec::new());
2139 }
2140
2141 let ann_index = self.vector_indexes.get(&field.0);
2142 let lazy_flat = self.flat_vectors.get(&field.0);
2143 let ann_fetch_k = lazy_flat.map_or(fetch_k, |flat| {
2144 ann_ordinal_fetch_k(fetch_k, flat.num_vectors, flat.num_docs_with_vectors())
2145 });
2146
2147 if ann_index.is_none() && lazy_flat.is_none() {
2148 return Ok(Vec::new());
2149 }
2150
2151 if ann_index.is_some() && lazy_flat.is_none() {
2152 return Err(Error::Corruption(format!(
2153 "dense ANN field {} is missing flat vector storage",
2154 field.0
2155 )));
2156 }
2157
2158 if let Some(flat) = lazy_flat
2159 && flat.dim != params.dim
2160 {
2161 return Err(Error::Corruption(format!(
2162 "dense vector field {} has schema dimension {} but flat storage dimension {}",
2163 field.0, params.dim, flat.dim
2164 )));
2165 }
2166
2167 let mut results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
2168 match index {
2170 VectorIndex::RaBitQ(lazy) => {
2171 let rabitq = lazy.get().ok_or_else(|| {
2172 Error::Schema("RaBitQ index deserialization failed".to_string())
2173 })?;
2174 if rabitq.codebook.config.dim != params.dim {
2175 return Err(Error::Corruption(format!(
2176 "RaBitQ index dimension {} does not match schema dimension {}",
2177 rabitq.codebook.config.dim, params.dim
2178 )));
2179 }
2180 let flat = lazy_flat.expect("ANN/flat pairing validated above");
2181 progressive_ann_search(
2182 fetch_k.min(flat.num_docs_with_vectors()),
2183 ann_fetch_k,
2184 rabitq.len().min(flat.num_vectors),
2185 |candidate_k| {
2186 rabitq
2187 .search(query, candidate_k)
2188 .into_iter()
2189 .map(|(doc_id, ordinal, dist)| {
2190 (doc_id, ordinal, 1.0 / (1.0 + dist))
2191 })
2192 .collect()
2193 },
2194 )?
2195 }
2196 VectorIndex::IVF(lazy) => {
2197 let (index, codebook) = lazy.get().ok_or_else(|| {
2198 Error::Schema("IVF index deserialization failed".to_string())
2199 })?;
2200 let centroids = self.coarse_centroids.get(&field.0).ok_or_else(|| {
2201 Error::Schema(format!(
2202 "IVF index requires coarse centroids for field {}",
2203 field.0
2204 ))
2205 })?;
2206 validate_coarse_centroids(centroids, params.dim)?;
2207 if index.config.dim != params.dim
2208 || codebook.config.dim != params.dim
2209 || index.centroids_version != centroids.version
2210 || index.codebook_version != codebook.version
2211 || index
2212 .clusters
2213 .iter()
2214 .any(|(cluster_id, _)| cluster_id >= centroids.num_clusters)
2215 {
2216 return Err(Error::Corruption(format!(
2217 "IVF index/codebook/centroid metadata does not match schema dimension {}",
2218 params.dim
2219 )));
2220 }
2221 let effective_nprobe = params.nprobe.min(centroids.num_clusters as usize);
2222 let flat = lazy_flat.expect("ANN/flat pairing validated above");
2223 progressive_ann_search(
2224 fetch_k.min(flat.num_docs_with_vectors()),
2225 ann_fetch_k,
2226 index.len().min(flat.num_vectors),
2227 |candidate_k| {
2228 index
2229 .search(
2230 centroids,
2231 codebook,
2232 query,
2233 candidate_k,
2234 Some(effective_nprobe),
2235 )
2236 .into_iter()
2237 .map(|(doc_id, ordinal, dist)| {
2238 (doc_id, ordinal, 1.0 / (1.0 + dist))
2239 })
2240 .collect()
2241 },
2242 )?
2243 }
2244 VectorIndex::ScaNN(lazy) => {
2245 let (index, codebook) = lazy.get().ok_or_else(|| {
2246 Error::Schema("ScaNN index deserialization failed".to_string())
2247 })?;
2248 let centroids = self.coarse_centroids.get(&field.0).ok_or_else(|| {
2249 Error::Schema(format!(
2250 "ScaNN index requires coarse centroids for field {}",
2251 field.0
2252 ))
2253 })?;
2254 validate_coarse_centroids(centroids, params.dim)?;
2255 if index.config.dim != params.dim
2256 || codebook.config.dim != params.dim
2257 || index.centroids_version != centroids.version
2258 || index.codebook_version != codebook.version
2259 || index
2260 .clusters
2261 .iter()
2262 .any(|(cluster_id, _)| cluster_id >= centroids.num_clusters)
2263 {
2264 return Err(Error::Corruption(format!(
2265 "ScaNN index/codebook/centroid metadata does not match schema dimension {}",
2266 params.dim
2267 )));
2268 }
2269 let effective_nprobe = params.nprobe.min(centroids.num_clusters as usize);
2270 let flat = lazy_flat.expect("ANN/flat pairing validated above");
2271 progressive_ann_search(
2272 fetch_k.min(flat.num_docs_with_vectors()),
2273 ann_fetch_k,
2274 index.len().min(flat.num_vectors),
2275 |candidate_k| {
2276 index
2277 .search(
2278 centroids,
2279 codebook,
2280 query,
2281 candidate_k,
2282 Some(effective_nprobe),
2283 )
2284 .into_iter()
2285 .map(|(doc_id, ordinal, dist)| {
2286 (doc_id, ordinal, 1.0 / (1.0 + dist))
2287 })
2288 .collect()
2289 },
2290 )?
2291 }
2292 VectorIndex::BinaryIvf(_) => {
2293 Vec::new()
2295 }
2296 }
2297 } else if let Some(lazy_flat) = lazy_flat {
2298 let dim = lazy_flat.dim;
2300 let n = lazy_flat.num_vectors;
2301 let quant = lazy_flat.quantization;
2302 let batch_len =
2303 bounded_vector_score_batch(lazy_flat.vector_byte_size(), DENSE_SCORE_BATCH);
2304 let mut collector = FlatDocumentCollector::new(fetch_k.min(n), combiner);
2305 let mut scores = vec![0f32; batch_len];
2306
2307 for batch_start in (0..n).step_by(batch_len) {
2308 let batch_count = batch_len.min(n - batch_start);
2309 let batch_bytes = lazy_flat
2310 .read_vectors_batch_sync(batch_start, batch_count)
2311 .map_err(crate::Error::Io)?;
2312 let raw = batch_bytes.as_slice();
2313
2314 Self::score_quantized_batch(
2315 query,
2316 raw,
2317 quant,
2318 dim,
2319 &mut scores[..batch_count],
2320 params.unit_norm,
2321 )?;
2322
2323 for (i, &score) in scores.iter().enumerate().take(batch_count) {
2324 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
2325 collector.push(doc_id, ordinal, score);
2326 }
2327 }
2328
2329 return Ok(collector.into_results());
2330 } else {
2331 return Ok(Vec::new());
2332 };
2333
2334 if ann_index.is_some()
2336 && !results.is_empty()
2337 && let Some(lazy_flat) = lazy_flat
2338 {
2339 let dim = lazy_flat.dim;
2340 let quant = lazy_flat.quantization;
2341 let vbs = lazy_flat.vector_byte_size();
2342
2343 let (expanded, mut resolved) = expand_ann_candidate_documents(&results, lazy_flat)?;
2344 results = expanded;
2345
2346 if !resolved.is_empty() {
2347 resolved.sort_unstable_by_key(|&(_, flat_idx)| flat_idx);
2348 let batch_len = bounded_vector_score_batch(vbs, DENSE_SCORE_BATCH);
2349 let max_batch = batch_len.min(resolved.len());
2350 let max_raw_len = max_batch
2351 .checked_mul(vbs)
2352 .ok_or_else(|| Error::Query("dense rerank buffer size overflow".into()))?;
2353 let mut raw_buf = vec![0u8; max_raw_len];
2354 let mut scores = vec![0f32; max_batch];
2355
2356 for chunk in resolved.chunks(batch_len) {
2357 let raw_len = chunk
2358 .len()
2359 .checked_mul(vbs)
2360 .ok_or_else(|| Error::Query("dense rerank buffer size overflow".into()))?;
2361 let raw = &mut raw_buf[..raw_len];
2362 for (buf_idx, &(_, flat_idx)) in chunk.iter().enumerate() {
2363 lazy_flat
2364 .read_vector_raw_into_sync(
2365 flat_idx,
2366 &mut raw[buf_idx * vbs..(buf_idx + 1) * vbs],
2367 )
2368 .map_err(crate::Error::Io)?;
2369 }
2370
2371 Self::score_quantized_batch(
2372 query,
2373 raw,
2374 quant,
2375 dim,
2376 &mut scores[..chunk.len()],
2377 params.unit_norm,
2378 )?;
2379
2380 for (buf_idx, &(ri, _)) in chunk.iter().enumerate() {
2381 results[ri].2 = scores[buf_idx];
2382 }
2383 }
2384 }
2385 }
2386
2387 Ok(combine_grouped_ordinal_results(results, combiner, k))
2388 }
2389
2390 #[cfg(feature = "sync")]
2395 pub fn search_binary_dense_vector_sync(
2396 &self,
2397 field: Field,
2398 query: &[u8],
2399 k: usize,
2400 combiner: crate::query::MultiValueCombiner,
2401 ) -> Result<Vec<VectorSearchResult>> {
2402 let schema_dim = self.validate_binary_search_request(field, query)?;
2403 combiner.validate().map_err(Error::Query)?;
2404 if k == 0 {
2405 return Ok(Vec::new());
2406 }
2407 let t0 = crate::observe::Timer::start();
2408 if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0)
2410 && let Some(ivf) = lazy.get()
2411 {
2412 if ivf.config.dim_bits != schema_dim {
2413 return Err(Error::Corruption(format!(
2414 "binary IVF field {} has schema dimension {} but index dimension {}",
2415 field.0, schema_dim, ivf.config.dim_bits
2416 )));
2417 }
2418 let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
2419 Error::Corruption(format!(
2420 "binary IVF field {} is missing flat vector storage",
2421 field.0
2422 ))
2423 })?;
2424 if flat.dim != schema_dim
2425 || flat.quantization != crate::dsl::DenseVectorQuantization::Binary
2426 {
2427 return Err(Error::Corruption(format!(
2428 "binary IVF field {} has inconsistent flat vector metadata",
2429 field.0
2430 )));
2431 }
2432 let nprobe = self.binary_ivf_nprobe(field);
2433 let initial_fetch =
2434 ann_ordinal_fetch_k(k, flat.num_vectors, flat.num_docs_with_vectors());
2435 let ann_results = progressive_ann_search(
2436 k.min(flat.num_docs_with_vectors()),
2437 initial_fetch,
2438 ivf.len().min(flat.num_vectors),
2439 |candidate_k| ivf.search(query, candidate_k, nprobe),
2440 )?;
2441 let results =
2442 exact_score_binary_candidate_documents_sync(&ann_results, flat, query, schema_dim)?;
2443 crate::observe::dense_l1(
2444 self.schema.index_label(),
2445 self.schema.get_field_name(field).unwrap_or("?"),
2446 "binary_ivf",
2447 t0.secs(),
2448 results.len(),
2449 );
2450 return Ok(combine_grouped_ordinal_results(results, combiner, k));
2451 }
2452
2453 let lazy_flat = match self.flat_vectors.get(&field.0) {
2454 Some(f) => f,
2455 None => return Ok(Vec::new()),
2456 };
2457
2458 let dim_bits = lazy_flat.dim;
2459 let byte_len = lazy_flat.vector_byte_size();
2460 let n = lazy_flat.num_vectors;
2461
2462 if dim_bits != schema_dim {
2463 return Err(Error::Corruption(format!(
2464 "binary vector field {} has schema dimension {} but flat storage dimension {}",
2465 field.0, schema_dim, dim_bits
2466 )));
2467 }
2468
2469 if byte_len != query.len() {
2470 return Err(Error::Schema(format!(
2471 "Binary query vector byte length {} != field byte length {}",
2472 query.len(),
2473 byte_len
2474 )));
2475 }
2476
2477 let batch_len = bounded_vector_score_batch(byte_len, BINARY_SCORE_BATCH);
2478 let mut collector = FlatDocumentCollector::new(k, combiner);
2479 let mut scores = vec![0f32; batch_len];
2480
2481 for batch_start in (0..n).step_by(batch_len) {
2482 let batch_count = batch_len.min(n - batch_start);
2483 let batch_bytes = lazy_flat
2484 .read_vectors_batch_sync(batch_start, batch_count)
2485 .map_err(crate::Error::Io)?;
2486 let raw = batch_bytes.as_slice();
2487
2488 crate::structures::simd::batch_hamming_scores(
2489 query,
2490 raw,
2491 byte_len,
2492 dim_bits,
2493 &mut scores[..batch_count],
2494 );
2495
2496 for (i, &score) in scores.iter().enumerate().take(batch_count) {
2497 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
2498 collector.push(doc_id, ordinal, score);
2499 }
2500 }
2501
2502 let results = collector.into_results();
2503
2504 crate::observe::dense_l1(
2505 self.schema.index_label(),
2506 self.schema.get_field_name(field).unwrap_or("?"),
2507 "binary_flat",
2508 t0.secs(),
2509 results.len(),
2510 );
2511 Ok(results)
2512 }
2513}
2514
2515#[cfg(test)]
2516mod dense_search_safety_tests {
2517 use super::*;
2518
2519 #[test]
2520 fn dense_fetch_count_rejects_non_finite_and_unbounded_factors() {
2521 for factor in [
2522 f32::NAN,
2523 f32::INFINITY,
2524 f32::NEG_INFINITY,
2525 0.0,
2526 0.5,
2527 MAX_DENSE_RERANK_FACTOR + 1.0,
2528 ] {
2529 assert!(
2530 checked_dense_fetch_k(10, factor).is_err(),
2531 "factor={factor}"
2532 );
2533 }
2534 }
2535
2536 #[test]
2537 fn flat_document_collector_does_not_let_one_multivalue_doc_crowd_out_others() {
2538 let mut collector = FlatDocumentCollector::new(2, crate::query::MultiValueCombiner::Max);
2539 collector.push(1, 0, 1.0);
2540 collector.push(1, 1, 0.9);
2541 collector.push(2, 0, 0.8);
2542
2543 let results = collector.into_results();
2544 assert_eq!(
2545 results
2546 .iter()
2547 .map(|result| result.doc_id)
2548 .collect::<Vec<_>>(),
2549 vec![1, 2]
2550 );
2551 assert_eq!(results[0].ordinals.len(), 2);
2552 }
2553
2554 #[test]
2555 fn flat_document_collector_evicts_by_score_then_doc_id() {
2556 let mut collector = FlatDocumentCollector::new(2, crate::query::MultiValueCombiner::Max);
2557 collector.push(1, 0, 0.5);
2558 collector.push(3, 0, 0.8);
2559 collector.push(2, 0, 0.9);
2560 let results = collector.into_results();
2561 assert_eq!(
2562 results
2563 .iter()
2564 .map(|result| result.doc_id)
2565 .collect::<Vec<_>>(),
2566 vec![2, 3]
2567 );
2568
2569 let mut tied = FlatDocumentCollector::new(1, crate::query::MultiValueCombiner::Max);
2570 tied.push(2, 0, 1.0);
2571 tied.push(1, 0, 1.0);
2572 let results = tied.into_results();
2573 assert_eq!(results[0].doc_id, 1);
2574 }
2575
2576 #[test]
2577 fn dense_fetch_count_rounds_up_and_detects_overflow() {
2578 assert_eq!(checked_dense_fetch_k(3, 1.5).unwrap(), 5);
2579 assert_eq!(checked_dense_fetch_k(50_000, 3.0).unwrap(), 150_000);
2580 assert!(checked_dense_fetch_k(50_000, 32.0).is_err());
2581 assert!(checked_dense_fetch_k(usize::MAX, 2.0).is_err());
2582 }
2583
2584 #[test]
2585 fn ann_fetch_depth_accounts_for_multivalue_density_and_stays_bounded() {
2586 assert_eq!(ann_ordinal_fetch_k(100, 1_000, 100), 1_000);
2587 assert_eq!(
2588 ann_ordinal_fetch_k(100_000, usize::MAX, 1),
2589 MAX_DENSE_CANDIDATES_PER_SEGMENT
2590 );
2591 assert_eq!(ann_ordinal_fetch_k(100, 0, 10), 0);
2592 }
2593
2594 #[test]
2595 fn ann_search_deepens_until_skewed_results_contain_enough_documents() {
2596 let ranked = [
2597 (1, 0, 1.0),
2598 (1, 1, 0.99),
2599 (1, 2, 0.98),
2600 (1, 3, 0.97),
2601 (1, 4, 0.96),
2602 (1, 5, 0.95),
2603 (2, 0, 0.9),
2604 (3, 0, 0.8),
2605 ];
2606 let mut fetches = Vec::new();
2607 let results = progressive_ann_search(2, 2, ranked.len(), |fetch| {
2608 fetches.push(fetch);
2609 ranked.iter().copied().take(fetch).collect()
2610 })
2611 .unwrap();
2612
2613 assert_eq!(fetches, vec![2, 4, 8]);
2614 assert!(results.iter().any(|&(doc_id, _, _)| doc_id == 2));
2615 }
2616
2617 #[test]
2618 fn ann_search_stops_when_probed_population_is_exhausted() {
2619 let ranked = [(1, 0, 1.0), (1, 1, 0.9)];
2620 let mut calls = 0;
2621 let results = progressive_ann_search(3, 4, 100, |fetch| {
2622 calls += 1;
2623 ranked.iter().copied().take(fetch).collect()
2624 })
2625 .unwrap();
2626
2627 assert_eq!(calls, 1);
2628 assert_eq!(results.len(), 2);
2629 }
2630
2631 #[test]
2632 fn file_ranges_reject_overflow_and_truncation() {
2633 assert_eq!(checked_file_range(4, 3, 7, "test").unwrap(), 4..7);
2634 assert!(checked_file_range(u64::MAX, 1, u64::MAX, "test").is_err());
2635 assert!(checked_file_range(5, 3, 7, "test").is_err());
2636 }
2637}