1pub(crate) mod bmp;
4pub(crate) mod loader;
5mod types;
6
7pub use bmp::BmpIndex;
8#[cfg(feature = "diagnostics")]
9pub use types::DimRawData;
10pub use types::{SparseIndex, VectorIndex, VectorSearchResult};
11
12#[derive(Debug, Clone, Default)]
14pub struct SegmentMemoryStats {
15 pub segment_id: u128,
17 pub num_docs: u32,
19 pub term_dict_cache_bytes: usize,
21 pub store_cache_bytes: usize,
23 pub sparse_index_bytes: usize,
25 pub dense_index_bytes: usize,
27 pub bloom_filter_bytes: usize,
29}
30
31impl SegmentMemoryStats {
32 pub fn total_bytes(&self) -> usize {
34 self.term_dict_cache_bytes
35 + self.store_cache_bytes
36 + self.sparse_index_bytes
37 + self.dense_index_bytes
38 + self.bloom_filter_bytes
39 }
40}
41
42use std::sync::Arc;
43
44use rustc_hash::FxHashMap;
45
46use super::vector_data::LazyFlatVectorData;
47use crate::directories::{Directory, FileHandle};
48use crate::dsl::{DenseVectorQuantization, Document, Field, Schema};
49use crate::structures::{
50 AsyncSSTableReader, BlockPostingList, CoarseCentroids, IVFPQIndex, IVFRaBitQIndex, PQCodebook,
51 RaBitQIndex, SSTableStats, TermInfo,
52};
53use crate::{DocId, Error, Result};
54
55use super::store::{AsyncStoreReader, RawStoreBlock};
56use super::types::{SegmentFiles, SegmentId, SegmentMeta};
57
58pub(crate) fn combine_ordinal_results(
64 raw: impl IntoIterator<Item = (u32, u16, f32)>,
65 combiner: crate::query::MultiValueCombiner,
66 limit: usize,
67) -> Vec<VectorSearchResult> {
68 let collected: Vec<(u32, u16, f32)> = raw.into_iter().collect();
69
70 let num_raw = collected.len();
71 if log::log_enabled!(log::Level::Debug) {
72 let mut ids: Vec<u32> = collected.iter().map(|(d, _, _)| *d).collect();
73 ids.sort_unstable();
74 ids.dedup();
75 log::debug!(
76 "combine_ordinal_results: {} raw entries, {} unique docs, combiner={:?}, limit={}",
77 num_raw,
78 ids.len(),
79 combiner,
80 limit
81 );
82 }
83
84 let all_single = collected.iter().all(|&(_, ord, _)| ord == 0);
86 if all_single {
87 let mut results: Vec<VectorSearchResult> = collected
88 .into_iter()
89 .map(|(doc_id, _, score)| VectorSearchResult::new(doc_id, score, vec![(0, score)]))
90 .collect();
91 results.sort_unstable_by(|a, b| {
92 b.score
93 .partial_cmp(&a.score)
94 .unwrap_or(std::cmp::Ordering::Equal)
95 });
96 results.truncate(limit);
97 return results;
98 }
99
100 let mut doc_ordinals: rustc_hash::FxHashMap<DocId, Vec<(u32, f32)>> =
102 rustc_hash::FxHashMap::default();
103 for (doc_id, ordinal, score) in collected {
104 doc_ordinals
105 .entry(doc_id as DocId)
106 .or_default()
107 .push((ordinal as u32, score));
108 }
109 let mut results: Vec<VectorSearchResult> = doc_ordinals
110 .into_iter()
111 .map(|(doc_id, ordinals)| {
112 let combined_score = combiner.combine(&ordinals);
113 VectorSearchResult::new(doc_id, combined_score, ordinals)
114 })
115 .collect();
116 results.sort_unstable_by(|a, b| {
117 b.score
118 .partial_cmp(&a.score)
119 .unwrap_or(std::cmp::Ordering::Equal)
120 });
121 results.truncate(limit);
122 results
123}
124
125pub struct SegmentReader {
131 meta: SegmentMeta,
132 term_dict: Arc<AsyncSSTableReader<TermInfo>>,
134 postings_handle: FileHandle,
136 store: Arc<AsyncStoreReader>,
138 schema: Arc<Schema>,
139 vector_indexes: FxHashMap<u32, VectorIndex>,
141 flat_vectors: FxHashMap<u32, LazyFlatVectorData>,
143 coarse_centroids: FxHashMap<u32, Arc<CoarseCentroids>>,
145 sparse_indexes: FxHashMap<u32, SparseIndex>,
147 bmp_indexes: FxHashMap<u32, BmpIndex>,
149 positions_handle: Option<FileHandle>,
151 fast_fields: FxHashMap<u32, crate::structures::fast_field::FastFieldReader>,
153 shared_threshold: std::sync::atomic::AtomicU32,
157}
158
159impl SegmentReader {
160 pub async fn open<D: Directory>(
162 dir: &D,
163 segment_id: SegmentId,
164 schema: Arc<Schema>,
165 cache_blocks: usize,
166 ) -> Result<Self> {
167 let files = SegmentFiles::new(segment_id.0);
168
169 let meta_slice = dir.open_read(&files.meta).await?;
171 let meta_bytes = meta_slice.read_bytes().await?;
172 let meta = SegmentMeta::deserialize(meta_bytes.as_slice())?;
173 debug_assert_eq!(meta.id, segment_id.0);
174
175 let term_dict_handle = dir.open_lazy(&files.term_dict).await?;
177 let term_dict = AsyncSSTableReader::open(term_dict_handle, cache_blocks).await?;
178
179 let postings_handle = dir.open_lazy(&files.postings).await?;
181
182 let store_handle = dir.open_lazy(&files.store).await?;
184 let store = AsyncStoreReader::open(store_handle, cache_blocks).await?;
185
186 let vectors_data = loader::load_vectors_file(dir, &files, &schema).await?;
188 let vector_indexes = vectors_data.indexes;
189 let flat_vectors = vectors_data.flat_vectors;
190
191 #[cfg(feature = "native")]
196 for (field_id, lazy_flat) in &flat_vectors {
197 if vector_indexes.contains_key(field_id) {
198 lazy_flat.advise_random_access();
199 }
200 }
201
202 let sparse_data = loader::load_sparse_file(dir, &files, meta.num_docs, &schema).await?;
204 let sparse_indexes = sparse_data.maxscore_indexes;
205 let bmp_indexes = sparse_data.bmp_indexes;
206
207 let positions_handle = loader::open_positions_file(dir, &files, &schema).await?;
209
210 let fast_fields = loader::load_fast_fields_file(dir, &files, &schema).await?;
212
213 {
215 let mut parts = vec![format!(
216 "[segment] loaded {:016x}: docs={}",
217 segment_id.0, meta.num_docs
218 )];
219 if !vector_indexes.is_empty() || !flat_vectors.is_empty() {
220 parts.push(format!(
221 "dense: {} ann + {} flat fields",
222 vector_indexes.len(),
223 flat_vectors.len()
224 ));
225 }
226 for (field_id, idx) in &sparse_indexes {
227 parts.push(format!(
228 "sparse field {}: {} dims, ~{:.1} KB",
229 field_id,
230 idx.num_dimensions(),
231 idx.num_dimensions() as f64 * 24.0 / 1024.0
232 ));
233 }
234 for (field_id, idx) in &bmp_indexes {
235 parts.push(format!(
236 "bmp field {}: {} dims, {} blocks",
237 field_id,
238 idx.dims(),
239 idx.num_blocks
240 ));
241 }
242 if !fast_fields.is_empty() {
243 parts.push(format!("fast: {} fields", fast_fields.len()));
244 }
245 log::debug!("{}", parts.join(", "));
246 }
247
248 Ok(Self {
249 meta,
250 term_dict: Arc::new(term_dict),
251 postings_handle,
252 store: Arc::new(store),
253 schema,
254 vector_indexes,
255 flat_vectors,
256 coarse_centroids: FxHashMap::default(),
257 sparse_indexes,
258 bmp_indexes,
259 positions_handle,
260 fast_fields,
261 shared_threshold: std::sync::atomic::AtomicU32::new(0),
262 })
263 }
264
265 #[inline]
267 pub fn reset_shared_threshold(&self) {
268 self.shared_threshold
269 .store(0, std::sync::atomic::Ordering::Relaxed);
270 }
271
272 #[inline]
274 pub fn shared_threshold_f32(&self) -> f32 {
275 f32::from_bits(
276 self.shared_threshold
277 .load(std::sync::atomic::Ordering::Relaxed),
278 )
279 }
280
281 #[inline]
283 pub fn update_shared_threshold(&self, new_threshold: f32) {
284 use std::sync::atomic::Ordering::Relaxed;
285 let new_bits = new_threshold.to_bits();
286 let mut current_bits = self.shared_threshold.load(Relaxed);
287 while new_threshold > f32::from_bits(current_bits) {
288 match self.shared_threshold.compare_exchange_weak(
289 current_bits,
290 new_bits,
291 Relaxed,
292 Relaxed,
293 ) {
294 Ok(_) => return,
295 Err(actual) => current_bits = actual,
296 }
297 }
298 }
299
300 pub fn meta(&self) -> &SegmentMeta {
301 &self.meta
302 }
303
304 pub fn num_docs(&self) -> u32 {
305 self.meta.num_docs
306 }
307
308 pub fn avg_field_len(&self, field: Field) -> f32 {
310 self.meta.avg_field_len(field)
311 }
312
313 pub fn schema(&self) -> &Schema {
314 &self.schema
315 }
316
317 pub fn sparse_indexes(&self) -> &FxHashMap<u32, SparseIndex> {
319 &self.sparse_indexes
320 }
321
322 pub fn sparse_index(&self, field: Field) -> Option<&SparseIndex> {
324 self.sparse_indexes.get(&field.0)
325 }
326
327 pub fn bmp_index(&self, field: Field) -> Option<&BmpIndex> {
329 self.bmp_indexes.get(&field.0)
330 }
331
332 pub fn bmp_indexes(&self) -> &FxHashMap<u32, BmpIndex> {
334 &self.bmp_indexes
335 }
336
337 pub fn vector_indexes(&self) -> &FxHashMap<u32, VectorIndex> {
339 &self.vector_indexes
340 }
341
342 pub fn flat_vectors(&self) -> &FxHashMap<u32, LazyFlatVectorData> {
344 &self.flat_vectors
345 }
346
347 pub fn fast_field(
349 &self,
350 field_id: u32,
351 ) -> Option<&crate::structures::fast_field::FastFieldReader> {
352 self.fast_fields.get(&field_id)
353 }
354
355 pub fn fast_fields(&self) -> &FxHashMap<u32, crate::structures::fast_field::FastFieldReader> {
357 &self.fast_fields
358 }
359
360 pub fn term_dict_stats(&self) -> SSTableStats {
362 self.term_dict.stats()
363 }
364
365 pub fn memory_stats(&self) -> SegmentMemoryStats {
367 let term_dict_stats = self.term_dict.stats();
368
369 let term_dict_cache_bytes = self.term_dict.cached_blocks() * 4096;
371
372 let store_cache_bytes = self.store.cached_blocks() * 4096;
374
375 let sparse_index_bytes: usize = self
377 .sparse_indexes
378 .values()
379 .map(|s| s.estimated_memory_bytes())
380 .sum::<usize>()
381 + self
382 .bmp_indexes
383 .values()
384 .map(|b| b.estimated_memory_bytes())
385 .sum::<usize>();
386
387 let dense_index_bytes: usize = self
390 .vector_indexes
391 .values()
392 .map(|v| v.estimated_memory_bytes())
393 .sum();
394
395 SegmentMemoryStats {
396 segment_id: self.meta.id,
397 num_docs: self.meta.num_docs,
398 term_dict_cache_bytes,
399 store_cache_bytes,
400 sparse_index_bytes,
401 dense_index_bytes,
402 bloom_filter_bytes: term_dict_stats.bloom_filter_size,
403 }
404 }
405
406 pub async fn get_postings(
411 &self,
412 field: Field,
413 term: &[u8],
414 ) -> Result<Option<BlockPostingList>> {
415 log::debug!(
416 "SegmentReader::get_postings field={} term_len={}",
417 field.0,
418 term.len()
419 );
420
421 let mut key = Vec::with_capacity(4 + term.len());
423 key.extend_from_slice(&field.0.to_le_bytes());
424 key.extend_from_slice(term);
425
426 let term_info = match self.term_dict.get(&key).await? {
428 Some(info) => {
429 log::debug!("SegmentReader::get_postings found term_info");
430 info
431 }
432 None => {
433 log::debug!("SegmentReader::get_postings term not found");
434 return Ok(None);
435 }
436 };
437
438 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
440 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
442 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs.into_iter()) {
443 posting_list.push(doc_id, tf);
444 }
445 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
446 return Ok(Some(block_list));
447 }
448
449 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
451 Error::Corruption("TermInfo has neither inline nor external data".to_string())
452 })?;
453
454 let start = posting_offset;
455 let end = start + posting_len;
456
457 if end > self.postings_handle.len() {
458 return Err(Error::Corruption(
459 "Posting offset out of bounds".to_string(),
460 ));
461 }
462
463 let posting_bytes = self.postings_handle.read_bytes_range(start..end).await?;
464 let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
465
466 Ok(Some(block_list))
467 }
468
469 pub async fn get_prefix_postings(
471 &self,
472 field: Field,
473 prefix: &[u8],
474 ) -> Result<Vec<BlockPostingList>> {
475 let mut key_prefix = Vec::with_capacity(4 + prefix.len());
477 key_prefix.extend_from_slice(&field.0.to_le_bytes());
478 key_prefix.extend_from_slice(prefix);
479
480 let entries = self.term_dict.prefix_scan(&key_prefix).await?;
481 let mut results = Vec::with_capacity(entries.len());
482
483 for (_key, term_info) in entries {
484 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
485 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
486 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs.into_iter()) {
487 posting_list.push(doc_id, tf);
488 }
489 results.push(BlockPostingList::from_posting_list(&posting_list)?);
490 } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
491 let start = posting_offset;
492 let end = start + posting_len;
493 if end > self.postings_handle.len() {
494 continue;
495 }
496 let posting_bytes = self.postings_handle.read_bytes_range(start..end).await?;
497 results.push(BlockPostingList::deserialize_zero_copy(posting_bytes)?);
498 }
499 }
500
501 Ok(results)
502 }
503
504 pub async fn doc(&self, local_doc_id: DocId) -> Result<Option<Document>> {
509 self.doc_with_fields(local_doc_id, None).await
510 }
511
512 pub async fn doc_with_fields(
518 &self,
519 local_doc_id: DocId,
520 fields: Option<&rustc_hash::FxHashSet<u32>>,
521 ) -> Result<Option<Document>> {
522 let mut doc = match fields {
523 Some(set) => {
524 let field_ids: Vec<u32> = set.iter().copied().collect();
525 match self
526 .store
527 .get_fields(local_doc_id, &self.schema, &field_ids)
528 .await
529 {
530 Ok(Some(d)) => d,
531 Ok(None) => return Ok(None),
532 Err(e) => return Err(Error::from(e)),
533 }
534 }
535 None => match self.store.get(local_doc_id, &self.schema).await {
536 Ok(Some(d)) => d,
537 Ok(None) => return Ok(None),
538 Err(e) => return Err(Error::from(e)),
539 },
540 };
541
542 for (&field_id, lazy_flat) in &self.flat_vectors {
544 if let Some(set) = fields
546 && !set.contains(&field_id)
547 {
548 continue;
549 }
550
551 let is_binary = lazy_flat.quantization == DenseVectorQuantization::Binary;
552 let (start, entries) = lazy_flat.flat_indexes_for_doc(local_doc_id);
553 for (j, &(_doc_id, _ordinal)) in entries.iter().enumerate() {
554 let flat_idx = start + j;
555 if is_binary {
556 let vbs = lazy_flat.vector_byte_size();
557 let mut raw = vec![0u8; vbs];
558 match lazy_flat.read_vector_raw_into(flat_idx, &mut raw).await {
559 Ok(()) => {
560 doc.add_binary_dense_vector(Field(field_id), raw);
561 }
562 Err(e) => {
563 log::warn!("Failed to hydrate binary vector field {}: {}", field_id, e);
564 }
565 }
566 } else {
567 match lazy_flat.get_vector(flat_idx).await {
568 Ok(vec) => {
569 doc.add_dense_vector(Field(field_id), vec);
570 }
571 Err(e) => {
572 log::warn!("Failed to hydrate vector field {}: {}", field_id, e);
573 }
574 }
575 }
576 }
577 }
578
579 Ok(Some(doc))
580 }
581
582 pub async fn prefetch_terms(
584 &self,
585 field: Field,
586 start_term: &[u8],
587 end_term: &[u8],
588 ) -> Result<()> {
589 let mut start_key = Vec::with_capacity(4 + start_term.len());
590 start_key.extend_from_slice(&field.0.to_le_bytes());
591 start_key.extend_from_slice(start_term);
592
593 let mut end_key = Vec::with_capacity(4 + end_term.len());
594 end_key.extend_from_slice(&field.0.to_le_bytes());
595 end_key.extend_from_slice(end_term);
596
597 self.term_dict.prefetch_range(&start_key, &end_key).await?;
598 Ok(())
599 }
600
601 pub fn store_has_dict(&self) -> bool {
603 self.store.has_dict()
604 }
605
606 pub fn store(&self) -> &super::store::AsyncStoreReader {
608 &self.store
609 }
610
611 pub fn store_raw_blocks(&self) -> Vec<RawStoreBlock> {
613 self.store.raw_blocks()
614 }
615
616 pub fn store_data_slice(&self) -> &FileHandle {
618 self.store.data_slice()
619 }
620
621 pub async fn all_terms(&self) -> Result<Vec<(Vec<u8>, TermInfo)>> {
623 self.term_dict.all_entries().await.map_err(Error::from)
624 }
625
626 pub async fn all_terms_with_stats(&self) -> Result<Vec<(Field, String, u32)>> {
631 let entries = self.term_dict.all_entries().await?;
632 let mut result = Vec::with_capacity(entries.len());
633
634 for (key, term_info) in entries {
635 if key.len() > 4 {
637 let field_id = u32::from_le_bytes([key[0], key[1], key[2], key[3]]);
638 let term_bytes = &key[4..];
639 if let Ok(term_str) = std::str::from_utf8(term_bytes) {
640 result.push((Field(field_id), term_str.to_string(), term_info.doc_freq()));
641 }
642 }
643 }
644
645 Ok(result)
646 }
647
648 pub fn term_dict_iter(&self) -> crate::structures::AsyncSSTableIterator<'_, TermInfo> {
650 self.term_dict.iter()
651 }
652
653 pub async fn prefetch_term_dict(&self) -> crate::Result<()> {
657 self.term_dict
658 .prefetch_all_data_bulk()
659 .await
660 .map_err(crate::Error::from)
661 }
662
663 pub async fn read_postings(&self, offset: u64, len: u64) -> Result<Vec<u8>> {
665 let start = offset;
666 let end = start + len;
667 let bytes = self.postings_handle.read_bytes_range(start..end).await?;
668 Ok(bytes.to_vec())
669 }
670
671 pub async fn read_position_bytes(&self, offset: u64, len: u64) -> Result<Option<Vec<u8>>> {
673 let handle = match &self.positions_handle {
674 Some(h) => h,
675 None => return Ok(None),
676 };
677 let start = offset;
678 let end = start + len;
679 let bytes = handle.read_bytes_range(start..end).await?;
680 Ok(Some(bytes.to_vec()))
681 }
682
683 pub fn has_positions_file(&self) -> bool {
685 self.positions_handle.is_some()
686 }
687
688 fn score_quantized_batch(
694 query: &[f32],
695 raw: &[u8],
696 quant: crate::dsl::DenseVectorQuantization,
697 dim: usize,
698 scores: &mut [f32],
699 unit_norm: bool,
700 ) {
701 use crate::dsl::DenseVectorQuantization;
702 use crate::structures::simd;
703 match (quant, unit_norm) {
704 (DenseVectorQuantization::F32, false) => {
705 let num_floats = scores.len() * dim;
706 debug_assert!(
707 (raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()),
708 "f32 vector data not 4-byte aligned — vectors file may use legacy format"
709 );
710 let vectors: &[f32] =
711 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
712 simd::batch_cosine_scores(query, vectors, dim, scores);
713 }
714 (DenseVectorQuantization::F32, true) => {
715 let num_floats = scores.len() * dim;
716 debug_assert!(
717 (raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()),
718 "f32 vector data not 4-byte aligned"
719 );
720 let vectors: &[f32] =
721 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
722 simd::batch_dot_scores(query, vectors, dim, scores);
723 }
724 (DenseVectorQuantization::F16, false) => {
725 simd::batch_cosine_scores_f16(query, raw, dim, scores);
726 }
727 (DenseVectorQuantization::F16, true) => {
728 simd::batch_dot_scores_f16(query, raw, dim, scores);
729 }
730 (DenseVectorQuantization::UInt8, false) => {
731 simd::batch_cosine_scores_u8(query, raw, dim, scores);
732 }
733 (DenseVectorQuantization::UInt8, true) => {
734 simd::batch_dot_scores_u8(query, raw, dim, scores);
735 }
736 (DenseVectorQuantization::Binary, _) => {
737 unreachable!("Binary quantization should not reach score_quantized_batch");
739 }
740 }
741 }
742
743 pub async fn search_dense_vector(
749 &self,
750 field: Field,
751 query: &[f32],
752 k: usize,
753 nprobe: usize,
754 rerank_factor: f32,
755 combiner: crate::query::MultiValueCombiner,
756 ) -> Result<Vec<VectorSearchResult>> {
757 let ann_index = self.vector_indexes.get(&field.0);
758 let lazy_flat = self.flat_vectors.get(&field.0);
759
760 if ann_index.is_none() && lazy_flat.is_none() {
762 return Ok(Vec::new());
763 }
764
765 let unit_norm = self
767 .schema
768 .get_field_entry(field)
769 .and_then(|e| e.dense_vector_config.as_ref())
770 .is_some_and(|c| c.unit_norm);
771
772 const BRUTE_FORCE_BATCH: usize = 4096;
774
775 let fetch_k = (k as f32 * rerank_factor.max(1.0)).ceil() as usize;
776
777 let t0 = std::time::Instant::now();
779 let mut results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
780 match index {
782 VectorIndex::RaBitQ(lazy) => {
783 let rabitq = lazy.get().ok_or_else(|| {
784 Error::Schema("RaBitQ index deserialization failed".to_string())
785 })?;
786 rabitq
787 .search(query, fetch_k)
788 .into_iter()
789 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
790 .collect()
791 }
792 VectorIndex::IVF(lazy) => {
793 let (index, codebook) = lazy.get().ok_or_else(|| {
794 Error::Schema("IVF index deserialization failed".to_string())
795 })?;
796 let centroids = self.coarse_centroids.get(&field.0).ok_or_else(|| {
797 Error::Schema(format!(
798 "IVF index requires coarse centroids for field {}",
799 field.0
800 ))
801 })?;
802 let effective_nprobe = if nprobe > 0 { nprobe } else { 32 };
803 index
804 .search(centroids, codebook, query, fetch_k, Some(effective_nprobe))
805 .into_iter()
806 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
807 .collect()
808 }
809 VectorIndex::ScaNN(lazy) => {
810 let (index, codebook) = lazy.get().ok_or_else(|| {
811 Error::Schema("ScaNN index deserialization failed".to_string())
812 })?;
813 let centroids = self.coarse_centroids.get(&field.0).ok_or_else(|| {
814 Error::Schema(format!(
815 "ScaNN index requires coarse centroids for field {}",
816 field.0
817 ))
818 })?;
819 let effective_nprobe = if nprobe > 0 { nprobe } else { 32 };
820 index
821 .search(centroids, codebook, query, fetch_k, Some(effective_nprobe))
822 .into_iter()
823 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
824 .collect()
825 }
826 VectorIndex::BinaryIvf(_) => {
827 Vec::new()
829 }
830 }
831 } else if let Some(lazy_flat) = lazy_flat {
832 log::debug!(
835 "[search_dense] field {}: brute-force on {} vectors (dim={}, quant={:?})",
836 field.0,
837 lazy_flat.num_vectors,
838 lazy_flat.dim,
839 lazy_flat.quantization
840 );
841 let dim = lazy_flat.dim;
842 let n = lazy_flat.num_vectors;
843 let quant = lazy_flat.quantization;
844 let mut collector = crate::query::ScoreCollector::new(fetch_k);
845 let mut scores = vec![0f32; BRUTE_FORCE_BATCH];
846
847 for batch_start in (0..n).step_by(BRUTE_FORCE_BATCH) {
848 let batch_count = BRUTE_FORCE_BATCH.min(n - batch_start);
849 let batch_bytes = lazy_flat
850 .read_vectors_batch(batch_start, batch_count)
851 .await
852 .map_err(crate::Error::Io)?;
853 let raw = batch_bytes.as_slice();
854
855 Self::score_quantized_batch(
856 query,
857 raw,
858 quant,
859 dim,
860 &mut scores[..batch_count],
861 unit_norm,
862 );
863
864 for (i, &score) in scores.iter().enumerate().take(batch_count) {
865 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
866 collector.insert_with_ordinal(doc_id, score, ordinal);
867 }
868 }
869
870 collector
871 .into_sorted_results()
872 .into_iter()
873 .map(|(doc_id, score, ordinal)| (doc_id, ordinal, score))
874 .collect()
875 } else {
876 return Ok(Vec::new());
877 };
878 let l1_elapsed = t0.elapsed();
879 log::debug!(
880 "[search_dense] field {}: L1 returned {} candidates in {:.1}ms",
881 field.0,
882 results.len(),
883 l1_elapsed.as_secs_f64() * 1000.0
884 );
885
886 if ann_index.is_some()
889 && !results.is_empty()
890 && let Some(lazy_flat) = lazy_flat
891 {
892 let t_rerank = std::time::Instant::now();
893 let dim = lazy_flat.dim;
894 let quant = lazy_flat.quantization;
895 let vbs = lazy_flat.vector_byte_size();
896
897 let mut resolved: Vec<(usize, usize)> = Vec::new(); for (ri, c) in results.iter().enumerate() {
900 let (start, entries) = lazy_flat.flat_indexes_for_doc(c.0);
901 for (j, &(_, ord)) in entries.iter().enumerate() {
902 if ord == c.1 {
903 resolved.push((ri, start + j));
904 break;
905 }
906 }
907 }
908
909 let t_resolve = t_rerank.elapsed();
910 if !resolved.is_empty() {
911 resolved.sort_unstable_by_key(|&(_, flat_idx)| flat_idx);
913
914 #[cfg(feature = "native")]
919 lazy_flat.prefetch_vectors(resolved.iter().map(|&(_, flat_idx)| flat_idx));
920
921 let t_read = std::time::Instant::now();
923 let mut raw_buf = vec![0u8; resolved.len() * vbs];
924 for (buf_idx, &(_, flat_idx)) in resolved.iter().enumerate() {
925 let _ = lazy_flat
926 .read_vector_raw_into(
927 flat_idx,
928 &mut raw_buf[buf_idx * vbs..(buf_idx + 1) * vbs],
929 )
930 .await;
931 }
932
933 let read_elapsed = t_read.elapsed();
934
935 let t_score = std::time::Instant::now();
937 let mut scores = vec![0f32; resolved.len()];
938 Self::score_quantized_batch(query, &raw_buf, quant, dim, &mut scores, unit_norm);
939 let score_elapsed = t_score.elapsed();
940
941 for (buf_idx, &(ri, _)) in resolved.iter().enumerate() {
943 results[ri].2 = scores[buf_idx];
944 }
945
946 log::debug!(
947 "[search_dense] field {}: rerank {} vectors (dim={}, quant={:?}, {}B/vec): resolve={:.1}ms read={:.1}ms score={:.1}ms",
948 field.0,
949 resolved.len(),
950 dim,
951 quant,
952 vbs,
953 t_resolve.as_secs_f64() * 1000.0,
954 read_elapsed.as_secs_f64() * 1000.0,
955 score_elapsed.as_secs_f64() * 1000.0,
956 );
957 }
958
959 if results.len() > fetch_k {
960 results.select_nth_unstable_by(fetch_k, |a, b| b.2.total_cmp(&a.2));
961 results.truncate(fetch_k);
962 }
963 results.sort_unstable_by(|a, b| b.2.total_cmp(&a.2));
964 log::debug!(
965 "[search_dense] field {}: rerank total={:.1}ms",
966 field.0,
967 t_rerank.elapsed().as_secs_f64() * 1000.0
968 );
969 }
970
971 Ok(combine_ordinal_results(results, combiner, k))
972 }
973
974 fn binary_ivf_nprobe(&self, field: Field) -> Option<usize> {
976 self.schema
977 .get_field_entry(field)
978 .and_then(|e| e.binary_dense_vector_config.as_ref())
979 .map(|c| c.nprobe)
980 .filter(|&n| n > 0)
981 }
982
983 pub async fn search_binary_dense_vector(
987 &self,
988 field: Field,
989 query: &[u8],
990 k: usize,
991 combiner: crate::query::MultiValueCombiner,
992 ) -> Result<Vec<VectorSearchResult>> {
993 if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0)
996 && let Some(ivf) = lazy.get()
997 {
998 let nprobe = self.binary_ivf_nprobe(field);
999 let results = ivf.search(query, k, nprobe);
1000 return Ok(combine_ordinal_results(results, combiner, k));
1001 }
1002
1003 let lazy_flat = match self.flat_vectors.get(&field.0) {
1004 Some(f) => f,
1005 None => return Ok(Vec::new()),
1006 };
1007
1008 const BRUTE_FORCE_BATCH: usize = 8192; let dim_bits = lazy_flat.dim;
1011 let byte_len = lazy_flat.vector_byte_size();
1012 let n = lazy_flat.num_vectors;
1013
1014 if byte_len != query.len() {
1015 return Err(Error::Schema(format!(
1016 "Binary query vector byte length {} != field byte length {}",
1017 query.len(),
1018 byte_len
1019 )));
1020 }
1021
1022 let mut collector = crate::query::ScoreCollector::new(k);
1023 let mut scores = vec![0f32; BRUTE_FORCE_BATCH];
1024
1025 for batch_start in (0..n).step_by(BRUTE_FORCE_BATCH) {
1026 let batch_count = BRUTE_FORCE_BATCH.min(n - batch_start);
1027 let batch_bytes = lazy_flat
1028 .read_vectors_batch(batch_start, batch_count)
1029 .await
1030 .map_err(crate::Error::Io)?;
1031 let raw = batch_bytes.as_slice();
1032
1033 crate::structures::simd::batch_hamming_scores(
1034 query,
1035 raw,
1036 byte_len,
1037 dim_bits,
1038 &mut scores[..batch_count],
1039 );
1040
1041 let threshold = collector.threshold();
1042 for (i, &score) in scores.iter().enumerate().take(batch_count) {
1043 if score > threshold {
1044 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
1045 collector.insert_with_ordinal(doc_id, score, ordinal);
1046 }
1047 }
1048 }
1049
1050 let results: Vec<(u32, u16, f32)> = collector
1051 .into_sorted_results()
1052 .into_iter()
1053 .map(|(doc_id, score, ordinal)| (doc_id, ordinal, score))
1054 .collect();
1055
1056 Ok(combine_ordinal_results(results, combiner, k))
1057 }
1058
1059 pub fn has_dense_vector_index(&self, field: Field) -> bool {
1061 self.vector_indexes.contains_key(&field.0) || self.flat_vectors.contains_key(&field.0)
1062 }
1063
1064 pub fn get_dense_vector_index(&self, field: Field) -> Option<Arc<RaBitQIndex>> {
1066 match self.vector_indexes.get(&field.0) {
1067 Some(VectorIndex::RaBitQ(lazy)) => lazy.get().cloned(),
1068 _ => None,
1069 }
1070 }
1071
1072 pub fn get_ivf_vector_index(
1074 &self,
1075 field: Field,
1076 ) -> Option<(Arc<IVFRaBitQIndex>, Arc<crate::structures::RaBitQCodebook>)> {
1077 match self.vector_indexes.get(&field.0) {
1078 Some(VectorIndex::IVF(lazy)) => lazy.get().map(|(i, c)| (i.clone(), c.clone())),
1079 _ => None,
1080 }
1081 }
1082
1083 pub fn coarse_centroids(&self, field_id: u32) -> Option<&Arc<CoarseCentroids>> {
1085 self.coarse_centroids.get(&field_id)
1086 }
1087
1088 pub fn set_coarse_centroids(&mut self, centroids: FxHashMap<u32, Arc<CoarseCentroids>>) {
1090 self.coarse_centroids = centroids;
1091 }
1092
1093 pub fn get_scann_vector_index(
1095 &self,
1096 field: Field,
1097 ) -> Option<(Arc<IVFPQIndex>, Arc<PQCodebook>)> {
1098 match self.vector_indexes.get(&field.0) {
1099 Some(VectorIndex::ScaNN(lazy)) => lazy.get().map(|(i, c)| (i.clone(), c.clone())),
1100 _ => None,
1101 }
1102 }
1103
1104 pub fn get_vector_index(&self, field: Field) -> Option<&VectorIndex> {
1106 self.vector_indexes.get(&field.0)
1107 }
1108
1109 pub async fn get_positions(
1114 &self,
1115 field: Field,
1116 term: &[u8],
1117 ) -> Result<Option<crate::structures::PositionPostingList>> {
1118 let handle = match &self.positions_handle {
1120 Some(h) => h,
1121 None => return Ok(None),
1122 };
1123
1124 let mut key = Vec::with_capacity(4 + term.len());
1126 key.extend_from_slice(&field.0.to_le_bytes());
1127 key.extend_from_slice(term);
1128
1129 let term_info = match self.term_dict.get(&key).await? {
1131 Some(info) => info,
1132 None => return Ok(None),
1133 };
1134
1135 let (offset, length) = match term_info.position_info() {
1137 Some((o, l)) => (o, l),
1138 None => return Ok(None),
1139 };
1140
1141 let slice = handle.slice(offset..offset + length);
1143 let data = slice.read_bytes().await?;
1144
1145 let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
1147
1148 Ok(Some(pos_list))
1149 }
1150
1151 pub fn has_positions(&self, field: Field) -> bool {
1153 if let Some(entry) = self.schema.get_field_entry(field) {
1155 entry.positions.is_some()
1156 } else {
1157 false
1158 }
1159 }
1160}
1161
1162#[cfg(feature = "sync")]
1164impl SegmentReader {
1165 pub fn get_postings_sync(&self, field: Field, term: &[u8]) -> Result<Option<BlockPostingList>> {
1167 let mut key = Vec::with_capacity(4 + term.len());
1169 key.extend_from_slice(&field.0.to_le_bytes());
1170 key.extend_from_slice(term);
1171
1172 let term_info = match self.term_dict.get_sync(&key)? {
1174 Some(info) => info,
1175 None => return Ok(None),
1176 };
1177
1178 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
1180 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
1181 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs.into_iter()) {
1182 posting_list.push(doc_id, tf);
1183 }
1184 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
1185 return Ok(Some(block_list));
1186 }
1187
1188 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
1190 Error::Corruption("TermInfo has neither inline nor external data".to_string())
1191 })?;
1192
1193 let start = posting_offset;
1194 let end = start + posting_len;
1195
1196 if end > self.postings_handle.len() {
1197 return Err(Error::Corruption(
1198 "Posting offset out of bounds".to_string(),
1199 ));
1200 }
1201
1202 let posting_bytes = self.postings_handle.read_bytes_range_sync(start..end)?;
1203 let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
1204
1205 Ok(Some(block_list))
1206 }
1207
1208 pub fn get_prefix_postings_sync(
1210 &self,
1211 field: Field,
1212 prefix: &[u8],
1213 ) -> Result<Vec<BlockPostingList>> {
1214 let mut key_prefix = Vec::with_capacity(4 + prefix.len());
1215 key_prefix.extend_from_slice(&field.0.to_le_bytes());
1216 key_prefix.extend_from_slice(prefix);
1217
1218 let entries = self.term_dict.prefix_scan_sync(&key_prefix)?;
1219 let mut results = Vec::with_capacity(entries.len());
1220
1221 for (_key, term_info) in entries {
1222 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
1223 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
1224 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs.into_iter()) {
1225 posting_list.push(doc_id, tf);
1226 }
1227 results.push(BlockPostingList::from_posting_list(&posting_list)?);
1228 } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
1229 let start = posting_offset;
1230 let end = start + posting_len;
1231 if end > self.postings_handle.len() {
1232 continue;
1233 }
1234 let posting_bytes = self.postings_handle.read_bytes_range_sync(start..end)?;
1235 results.push(BlockPostingList::deserialize_zero_copy(posting_bytes)?);
1236 }
1237 }
1238
1239 Ok(results)
1240 }
1241
1242 pub fn get_positions_sync(
1244 &self,
1245 field: Field,
1246 term: &[u8],
1247 ) -> Result<Option<crate::structures::PositionPostingList>> {
1248 let handle = match &self.positions_handle {
1249 Some(h) => h,
1250 None => return Ok(None),
1251 };
1252
1253 let mut key = Vec::with_capacity(4 + term.len());
1255 key.extend_from_slice(&field.0.to_le_bytes());
1256 key.extend_from_slice(term);
1257
1258 let term_info = match self.term_dict.get_sync(&key)? {
1260 Some(info) => info,
1261 None => return Ok(None),
1262 };
1263
1264 let (offset, length) = match term_info.position_info() {
1265 Some((o, l)) => (o, l),
1266 None => return Ok(None),
1267 };
1268
1269 let slice = handle.slice(offset..offset + length);
1270 let data = slice.read_bytes_sync()?;
1271
1272 let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
1273 Ok(Some(pos_list))
1274 }
1275
1276 pub fn search_dense_vector_sync(
1279 &self,
1280 field: Field,
1281 query: &[f32],
1282 k: usize,
1283 nprobe: usize,
1284 rerank_factor: f32,
1285 combiner: crate::query::MultiValueCombiner,
1286 ) -> Result<Vec<VectorSearchResult>> {
1287 let ann_index = self.vector_indexes.get(&field.0);
1288 let lazy_flat = self.flat_vectors.get(&field.0);
1289
1290 if ann_index.is_none() && lazy_flat.is_none() {
1291 return Ok(Vec::new());
1292 }
1293
1294 let unit_norm = self
1295 .schema
1296 .get_field_entry(field)
1297 .and_then(|e| e.dense_vector_config.as_ref())
1298 .is_some_and(|c| c.unit_norm);
1299
1300 const BRUTE_FORCE_BATCH: usize = 4096;
1301 let fetch_k = (k as f32 * rerank_factor.max(1.0)).ceil() as usize;
1302
1303 let mut results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
1304 match index {
1306 VectorIndex::RaBitQ(lazy) => {
1307 let rabitq = lazy.get().ok_or_else(|| {
1308 Error::Schema("RaBitQ index deserialization failed".to_string())
1309 })?;
1310 rabitq
1311 .search(query, fetch_k)
1312 .into_iter()
1313 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
1314 .collect()
1315 }
1316 VectorIndex::IVF(lazy) => {
1317 let (index, codebook) = lazy.get().ok_or_else(|| {
1318 Error::Schema("IVF index deserialization failed".to_string())
1319 })?;
1320 let centroids = self.coarse_centroids.get(&field.0).ok_or_else(|| {
1321 Error::Schema(format!(
1322 "IVF index requires coarse centroids for field {}",
1323 field.0
1324 ))
1325 })?;
1326 let effective_nprobe = if nprobe > 0 { nprobe } else { 32 };
1327 index
1328 .search(centroids, codebook, query, fetch_k, Some(effective_nprobe))
1329 .into_iter()
1330 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
1331 .collect()
1332 }
1333 VectorIndex::ScaNN(lazy) => {
1334 let (index, codebook) = lazy.get().ok_or_else(|| {
1335 Error::Schema("ScaNN index deserialization failed".to_string())
1336 })?;
1337 let centroids = self.coarse_centroids.get(&field.0).ok_or_else(|| {
1338 Error::Schema(format!(
1339 "ScaNN index requires coarse centroids for field {}",
1340 field.0
1341 ))
1342 })?;
1343 let effective_nprobe = if nprobe > 0 { nprobe } else { 32 };
1344 index
1345 .search(centroids, codebook, query, fetch_k, Some(effective_nprobe))
1346 .into_iter()
1347 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
1348 .collect()
1349 }
1350 VectorIndex::BinaryIvf(_) => {
1351 Vec::new()
1353 }
1354 }
1355 } else if let Some(lazy_flat) = lazy_flat {
1356 let dim = lazy_flat.dim;
1358 let n = lazy_flat.num_vectors;
1359 let quant = lazy_flat.quantization;
1360 let mut collector = crate::query::ScoreCollector::new(fetch_k);
1361 let mut scores = vec![0f32; BRUTE_FORCE_BATCH];
1362
1363 for batch_start in (0..n).step_by(BRUTE_FORCE_BATCH) {
1364 let batch_count = BRUTE_FORCE_BATCH.min(n - batch_start);
1365 let batch_bytes = lazy_flat
1366 .read_vectors_batch_sync(batch_start, batch_count)
1367 .map_err(crate::Error::Io)?;
1368 let raw = batch_bytes.as_slice();
1369
1370 Self::score_quantized_batch(
1371 query,
1372 raw,
1373 quant,
1374 dim,
1375 &mut scores[..batch_count],
1376 unit_norm,
1377 );
1378
1379 for (i, &score) in scores.iter().enumerate().take(batch_count) {
1380 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
1381 collector.insert_with_ordinal(doc_id, score, ordinal);
1382 }
1383 }
1384
1385 collector
1386 .into_sorted_results()
1387 .into_iter()
1388 .map(|(doc_id, score, ordinal)| (doc_id, ordinal, score))
1389 .collect()
1390 } else {
1391 return Ok(Vec::new());
1392 };
1393
1394 if ann_index.is_some()
1396 && !results.is_empty()
1397 && let Some(lazy_flat) = lazy_flat
1398 {
1399 let dim = lazy_flat.dim;
1400 let quant = lazy_flat.quantization;
1401 let vbs = lazy_flat.vector_byte_size();
1402
1403 let mut resolved: Vec<(usize, usize)> = Vec::new();
1404 for (ri, c) in results.iter().enumerate() {
1405 let (start, entries) = lazy_flat.flat_indexes_for_doc(c.0);
1406 for (j, &(_, ord)) in entries.iter().enumerate() {
1407 if ord == c.1 {
1408 resolved.push((ri, start + j));
1409 break;
1410 }
1411 }
1412 }
1413
1414 if !resolved.is_empty() {
1415 resolved.sort_unstable_by_key(|&(_, flat_idx)| flat_idx);
1416 let mut raw_buf = vec![0u8; resolved.len() * vbs];
1417 for (buf_idx, &(_, flat_idx)) in resolved.iter().enumerate() {
1418 let _ = lazy_flat.read_vector_raw_into_sync(
1419 flat_idx,
1420 &mut raw_buf[buf_idx * vbs..(buf_idx + 1) * vbs],
1421 );
1422 }
1423
1424 let mut scores = vec![0f32; resolved.len()];
1425 Self::score_quantized_batch(query, &raw_buf, quant, dim, &mut scores, unit_norm);
1426
1427 for (buf_idx, &(ri, _)) in resolved.iter().enumerate() {
1428 results[ri].2 = scores[buf_idx];
1429 }
1430 }
1431
1432 if results.len() > fetch_k {
1433 results.select_nth_unstable_by(fetch_k, |a, b| b.2.total_cmp(&a.2));
1434 results.truncate(fetch_k);
1435 }
1436 results.sort_unstable_by(|a, b| b.2.total_cmp(&a.2));
1437 }
1438
1439 Ok(combine_ordinal_results(results, combiner, k))
1440 }
1441
1442 #[cfg(feature = "sync")]
1447 pub fn search_binary_dense_vector_sync(
1448 &self,
1449 field: Field,
1450 query: &[u8],
1451 k: usize,
1452 combiner: crate::query::MultiValueCombiner,
1453 ) -> Result<Vec<VectorSearchResult>> {
1454 if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0)
1457 && let Some(ivf) = lazy.get()
1458 {
1459 let nprobe = self.binary_ivf_nprobe(field);
1460 let results = ivf.search(query, k, nprobe);
1461 return Ok(combine_ordinal_results(results, combiner, k));
1462 }
1463
1464 let lazy_flat = match self.flat_vectors.get(&field.0) {
1465 Some(f) => f,
1466 None => return Ok(Vec::new()),
1467 };
1468
1469 const BRUTE_FORCE_BATCH: usize = 8192; let dim_bits = lazy_flat.dim;
1472 let byte_len = lazy_flat.vector_byte_size();
1473 let n = lazy_flat.num_vectors;
1474
1475 if byte_len != query.len() {
1476 return Err(Error::Schema(format!(
1477 "Binary query vector byte length {} != field byte length {}",
1478 query.len(),
1479 byte_len
1480 )));
1481 }
1482
1483 let mut collector = crate::query::ScoreCollector::new(k);
1484 let mut scores = vec![0f32; BRUTE_FORCE_BATCH];
1485
1486 for batch_start in (0..n).step_by(BRUTE_FORCE_BATCH) {
1487 let batch_count = BRUTE_FORCE_BATCH.min(n - batch_start);
1488 let batch_bytes = lazy_flat
1489 .read_vectors_batch_sync(batch_start, batch_count)
1490 .map_err(crate::Error::Io)?;
1491 let raw = batch_bytes.as_slice();
1492
1493 crate::structures::simd::batch_hamming_scores(
1494 query,
1495 raw,
1496 byte_len,
1497 dim_bits,
1498 &mut scores[..batch_count],
1499 );
1500
1501 let threshold = collector.threshold();
1502 for (i, &score) in scores.iter().enumerate().take(batch_count) {
1503 if score > threshold {
1504 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
1505 collector.insert_with_ordinal(doc_id, score, ordinal);
1506 }
1507 }
1508 }
1509
1510 let results: Vec<(u32, u16, f32)> = collector
1511 .into_sorted_results()
1512 .into_iter()
1513 .map(|(doc_id, score, ordinal)| (doc_id, ordinal, score))
1514 .collect();
1515
1516 Ok(combine_ordinal_results(results, combiner, k))
1517 }
1518}