1mod loader;
4mod types;
5
6pub use types::{SparseIndex, VectorIndex, VectorSearchResult};
7
8#[derive(Debug, Clone, Default)]
10pub struct SegmentMemoryStats {
11 pub segment_id: u128,
13 pub num_docs: u32,
15 pub term_dict_cache_bytes: usize,
17 pub store_cache_bytes: usize,
19 pub sparse_index_bytes: usize,
21 pub dense_index_bytes: usize,
23 pub bloom_filter_bytes: usize,
25}
26
27impl SegmentMemoryStats {
28 pub fn total_bytes(&self) -> usize {
30 self.term_dict_cache_bytes
31 + self.store_cache_bytes
32 + self.sparse_index_bytes
33 + self.dense_index_bytes
34 + self.bloom_filter_bytes
35 }
36}
37
38use crate::structures::BlockSparsePostingList;
39
40use std::sync::Arc;
41
42use rustc_hash::FxHashMap;
43
44use super::vector_data::LazyFlatVectorData;
45use crate::directories::{AsyncFileRead, Directory, LazyFileHandle, LazyFileSlice};
46use crate::dsl::{Document, Field, Schema};
47use crate::structures::{
48 AsyncSSTableReader, BlockPostingList, CoarseCentroids, IVFPQIndex, IVFRaBitQIndex, PQCodebook,
49 RaBitQIndex, SSTableStats, TermInfo,
50};
51use crate::{DocId, Error, Result};
52
53use super::store::{AsyncStoreReader, RawStoreBlock};
54use super::types::{SegmentFiles, SegmentId, SegmentMeta};
55
56pub struct AsyncSegmentReader {
62 meta: SegmentMeta,
63 term_dict: Arc<AsyncSSTableReader<TermInfo>>,
65 postings_handle: LazyFileHandle,
67 store: Arc<AsyncStoreReader>,
69 schema: Arc<Schema>,
70 doc_id_offset: DocId,
72 vector_indexes: FxHashMap<u32, VectorIndex>,
74 flat_vectors: FxHashMap<u32, LazyFlatVectorData>,
76 coarse_centroids: FxHashMap<u32, Arc<CoarseCentroids>>,
78 sparse_indexes: FxHashMap<u32, SparseIndex>,
80 positions_handle: Option<LazyFileHandle>,
82}
83
84impl AsyncSegmentReader {
85 pub async fn open<D: Directory>(
87 dir: &D,
88 segment_id: SegmentId,
89 schema: Arc<Schema>,
90 doc_id_offset: DocId,
91 cache_blocks: usize,
92 ) -> Result<Self> {
93 let files = SegmentFiles::new(segment_id.0);
94
95 let meta_slice = dir.open_read(&files.meta).await?;
97 let meta_bytes = meta_slice.read_bytes().await?;
98 let meta = SegmentMeta::deserialize(meta_bytes.as_slice())?;
99 debug_assert_eq!(meta.id, segment_id.0);
100
101 let term_dict_handle = dir.open_lazy(&files.term_dict).await?;
103 let term_dict = AsyncSSTableReader::open(term_dict_handle, cache_blocks).await?;
104
105 let postings_handle = dir.open_lazy(&files.postings).await?;
107
108 let store_handle = dir.open_lazy(&files.store).await?;
110 let store = AsyncStoreReader::open(store_handle, cache_blocks).await?;
111
112 let vectors_data = loader::load_vectors_file(dir, &files, &schema).await?;
114 let vector_indexes = vectors_data.indexes;
115 let flat_vectors = vectors_data.flat_vectors;
116
117 let sparse_indexes = loader::load_sparse_file(dir, &files, meta.num_docs, &schema).await?;
119
120 let positions_handle = loader::open_positions_file(dir, &files, &schema).await?;
122
123 {
125 let mut parts = vec![format!(
126 "[segment] loaded {:016x}: docs={}",
127 segment_id.0, meta.num_docs
128 )];
129 if !vector_indexes.is_empty() || !flat_vectors.is_empty() {
130 parts.push(format!(
131 "dense: {} ann + {} flat fields",
132 vector_indexes.len(),
133 flat_vectors.len()
134 ));
135 }
136 for (field_id, idx) in &sparse_indexes {
137 parts.push(format!(
138 "sparse field {}: {} dims, ~{:.1} KB",
139 field_id,
140 idx.num_dimensions(),
141 idx.num_dimensions() as f64 * 24.0 / 1024.0
142 ));
143 }
144 log::debug!("{}", parts.join(", "));
145 }
146
147 Ok(Self {
148 meta,
149 term_dict: Arc::new(term_dict),
150 postings_handle,
151 store: Arc::new(store),
152 schema,
153 doc_id_offset,
154 vector_indexes,
155 flat_vectors,
156 coarse_centroids: FxHashMap::default(),
157 sparse_indexes,
158 positions_handle,
159 })
160 }
161
162 pub fn meta(&self) -> &SegmentMeta {
163 &self.meta
164 }
165
166 pub fn num_docs(&self) -> u32 {
167 self.meta.num_docs
168 }
169
170 pub fn avg_field_len(&self, field: Field) -> f32 {
172 self.meta.avg_field_len(field)
173 }
174
175 pub fn doc_id_offset(&self) -> DocId {
176 self.doc_id_offset
177 }
178
179 pub fn set_doc_id_offset(&mut self, offset: DocId) {
181 self.doc_id_offset = offset;
182 }
183
184 pub fn schema(&self) -> &Schema {
185 &self.schema
186 }
187
188 pub fn sparse_indexes(&self) -> &FxHashMap<u32, SparseIndex> {
190 &self.sparse_indexes
191 }
192
193 pub fn vector_indexes(&self) -> &FxHashMap<u32, VectorIndex> {
195 &self.vector_indexes
196 }
197
198 pub fn flat_vectors(&self) -> &FxHashMap<u32, LazyFlatVectorData> {
200 &self.flat_vectors
201 }
202
203 pub fn term_dict_stats(&self) -> SSTableStats {
205 self.term_dict.stats()
206 }
207
208 pub fn memory_stats(&self) -> SegmentMemoryStats {
210 let term_dict_stats = self.term_dict.stats();
211
212 let term_dict_cache_bytes = self.term_dict.cached_blocks() * 4096;
214
215 let store_cache_bytes = self.store.cached_blocks() * 4096;
217
218 let sparse_index_bytes: usize = self
220 .sparse_indexes
221 .values()
222 .map(|s| s.estimated_memory_bytes())
223 .sum();
224
225 let dense_index_bytes: usize = self
228 .vector_indexes
229 .values()
230 .map(|v| v.estimated_memory_bytes())
231 .sum();
232
233 SegmentMemoryStats {
234 segment_id: self.meta.id,
235 num_docs: self.meta.num_docs,
236 term_dict_cache_bytes,
237 store_cache_bytes,
238 sparse_index_bytes,
239 dense_index_bytes,
240 bloom_filter_bytes: term_dict_stats.bloom_filter_size,
241 }
242 }
243
244 pub async fn get_postings(
249 &self,
250 field: Field,
251 term: &[u8],
252 ) -> Result<Option<BlockPostingList>> {
253 log::debug!(
254 "SegmentReader::get_postings field={} term_len={}",
255 field.0,
256 term.len()
257 );
258
259 let mut key = Vec::with_capacity(4 + term.len());
261 key.extend_from_slice(&field.0.to_le_bytes());
262 key.extend_from_slice(term);
263
264 let term_info = match self.term_dict.get(&key).await? {
266 Some(info) => {
267 log::debug!("SegmentReader::get_postings found term_info");
268 info
269 }
270 None => {
271 log::debug!("SegmentReader::get_postings term not found");
272 return Ok(None);
273 }
274 };
275
276 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
278 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
280 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs.into_iter()) {
281 posting_list.push(doc_id, tf);
282 }
283 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
284 return Ok(Some(block_list));
285 }
286
287 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
289 Error::Corruption("TermInfo has neither inline nor external data".to_string())
290 })?;
291
292 let start = posting_offset;
293 let end = start + posting_len as u64;
294
295 if end > self.postings_handle.len() {
296 return Err(Error::Corruption(
297 "Posting offset out of bounds".to_string(),
298 ));
299 }
300
301 let posting_bytes = self.postings_handle.read_bytes_range(start..end).await?;
302 let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
303
304 Ok(Some(block_list))
305 }
306
307 pub async fn doc(&self, local_doc_id: DocId) -> Result<Option<Document>> {
312 self.doc_with_fields(local_doc_id, None).await
313 }
314
315 pub async fn doc_with_fields(
321 &self,
322 local_doc_id: DocId,
323 fields: Option<&rustc_hash::FxHashSet<u32>>,
324 ) -> Result<Option<Document>> {
325 let mut doc = match self.store.get(local_doc_id, &self.schema).await {
326 Ok(Some(d)) => d,
327 Ok(None) => return Ok(None),
328 Err(e) => return Err(Error::from(e)),
329 };
330
331 for (&field_id, lazy_flat) in &self.flat_vectors {
333 if let Some(set) = fields
335 && !set.contains(&field_id)
336 {
337 continue;
338 }
339
340 let (start, entries) = lazy_flat.flat_indexes_for_doc(local_doc_id);
341 for (j, &(_doc_id, _ordinal)) in entries.iter().enumerate() {
342 let flat_idx = start + j;
343 match lazy_flat.get_vector(flat_idx).await {
344 Ok(vec) => {
345 doc.add_dense_vector(Field(field_id), vec);
346 }
347 Err(e) => {
348 log::warn!("Failed to hydrate vector field {}: {}", field_id, e);
349 }
350 }
351 }
352 }
353
354 Ok(Some(doc))
355 }
356
357 pub async fn prefetch_terms(
359 &self,
360 field: Field,
361 start_term: &[u8],
362 end_term: &[u8],
363 ) -> Result<()> {
364 let mut start_key = Vec::with_capacity(4 + start_term.len());
365 start_key.extend_from_slice(&field.0.to_le_bytes());
366 start_key.extend_from_slice(start_term);
367
368 let mut end_key = Vec::with_capacity(4 + end_term.len());
369 end_key.extend_from_slice(&field.0.to_le_bytes());
370 end_key.extend_from_slice(end_term);
371
372 self.term_dict.prefetch_range(&start_key, &end_key).await?;
373 Ok(())
374 }
375
376 pub fn store_has_dict(&self) -> bool {
378 self.store.has_dict()
379 }
380
381 pub fn store(&self) -> &super::store::AsyncStoreReader {
383 &self.store
384 }
385
386 pub fn store_raw_blocks(&self) -> Vec<RawStoreBlock> {
388 self.store.raw_blocks()
389 }
390
391 pub fn store_data_slice(&self) -> &LazyFileSlice {
393 self.store.data_slice()
394 }
395
396 pub async fn all_terms(&self) -> Result<Vec<(Vec<u8>, TermInfo)>> {
398 self.term_dict.all_entries().await.map_err(Error::from)
399 }
400
401 pub async fn all_terms_with_stats(&self) -> Result<Vec<(Field, String, u32)>> {
406 let entries = self.term_dict.all_entries().await?;
407 let mut result = Vec::with_capacity(entries.len());
408
409 for (key, term_info) in entries {
410 if key.len() > 4 {
412 let field_id = u32::from_le_bytes([key[0], key[1], key[2], key[3]]);
413 let term_bytes = &key[4..];
414 if let Ok(term_str) = std::str::from_utf8(term_bytes) {
415 result.push((Field(field_id), term_str.to_string(), term_info.doc_freq()));
416 }
417 }
418 }
419
420 Ok(result)
421 }
422
423 pub fn term_dict_iter(&self) -> crate::structures::AsyncSSTableIterator<'_, TermInfo> {
425 self.term_dict.iter()
426 }
427
428 pub async fn prefetch_term_dict(&self) -> crate::Result<()> {
432 self.term_dict
433 .prefetch_all_data_bulk()
434 .await
435 .map_err(crate::Error::from)
436 }
437
438 pub async fn read_postings(&self, offset: u64, len: u32) -> Result<Vec<u8>> {
440 let start = offset;
441 let end = start + len as u64;
442 let bytes = self.postings_handle.read_bytes_range(start..end).await?;
443 Ok(bytes.to_vec())
444 }
445
446 pub async fn read_position_bytes(&self, offset: u64, len: u32) -> Result<Option<Vec<u8>>> {
448 let handle = match &self.positions_handle {
449 Some(h) => h,
450 None => return Ok(None),
451 };
452 let start = offset;
453 let end = start + len as u64;
454 let bytes = handle.read_bytes_range(start..end).await?;
455 Ok(Some(bytes.to_vec()))
456 }
457
458 pub fn has_positions_file(&self) -> bool {
460 self.positions_handle.is_some()
461 }
462
463 fn score_quantized_batch(
469 query: &[f32],
470 raw: &[u8],
471 quant: crate::dsl::DenseVectorQuantization,
472 dim: usize,
473 scores: &mut [f32],
474 ) {
475 match quant {
476 crate::dsl::DenseVectorQuantization::F32 => {
477 let num_floats = scores.len() * dim;
478 debug_assert!(
479 (raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()),
480 "f32 vector data not 4-byte aligned — vectors file may use legacy format"
481 );
482 let vectors: &[f32] =
483 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
484 crate::structures::simd::batch_cosine_scores(query, vectors, dim, scores);
485 }
486 crate::dsl::DenseVectorQuantization::F16 => {
487 crate::structures::simd::batch_cosine_scores_f16(query, raw, dim, scores);
488 }
489 crate::dsl::DenseVectorQuantization::UInt8 => {
490 crate::structures::simd::batch_cosine_scores_u8(query, raw, dim, scores);
491 }
492 }
493 }
494
495 pub async fn search_dense_vector(
501 &self,
502 field: Field,
503 query: &[f32],
504 k: usize,
505 nprobe: usize,
506 rerank_factor: usize,
507 combiner: crate::query::MultiValueCombiner,
508 ) -> Result<Vec<VectorSearchResult>> {
509 let ann_index = self.vector_indexes.get(&field.0);
510 let lazy_flat = self.flat_vectors.get(&field.0);
511
512 if ann_index.is_none() && lazy_flat.is_none() {
514 return Ok(Vec::new());
515 }
516
517 const BRUTE_FORCE_BATCH: usize = 4096;
519
520 let t0 = std::time::Instant::now();
522 let mut results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
523 match index {
525 VectorIndex::RaBitQ(lazy) => {
526 let rabitq = lazy.get().ok_or_else(|| {
527 Error::Schema("RaBitQ index deserialization failed".to_string())
528 })?;
529 let fetch_k = k * rerank_factor.max(1);
530 rabitq
531 .search(query, fetch_k, rerank_factor)
532 .into_iter()
533 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
534 .collect()
535 }
536 VectorIndex::IVF(lazy) => {
537 let (index, codebook) = lazy.get().ok_or_else(|| {
538 Error::Schema("IVF index deserialization failed".to_string())
539 })?;
540 let centroids = self.coarse_centroids.get(&field.0).ok_or_else(|| {
541 Error::Schema(format!(
542 "IVF index requires coarse centroids for field {}",
543 field.0
544 ))
545 })?;
546 let effective_nprobe = if nprobe > 0 { nprobe } else { 32 };
547 let fetch_k = k * rerank_factor.max(1);
548 index
549 .search(centroids, codebook, query, fetch_k, Some(effective_nprobe))
550 .into_iter()
551 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
552 .collect()
553 }
554 VectorIndex::ScaNN(lazy) => {
555 let (index, codebook) = lazy.get().ok_or_else(|| {
556 Error::Schema("ScaNN index deserialization failed".to_string())
557 })?;
558 let centroids = self.coarse_centroids.get(&field.0).ok_or_else(|| {
559 Error::Schema(format!(
560 "ScaNN index requires coarse centroids for field {}",
561 field.0
562 ))
563 })?;
564 let effective_nprobe = if nprobe > 0 { nprobe } else { 32 };
565 let fetch_k = k * rerank_factor.max(1);
566 index
567 .search(centroids, codebook, query, fetch_k, Some(effective_nprobe))
568 .into_iter()
569 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
570 .collect()
571 }
572 }
573 } else if let Some(lazy_flat) = lazy_flat {
574 log::debug!(
577 "[search_dense] field {}: brute-force on {} vectors (dim={}, quant={:?})",
578 field.0,
579 lazy_flat.num_vectors,
580 lazy_flat.dim,
581 lazy_flat.quantization
582 );
583 let dim = lazy_flat.dim;
584 let n = lazy_flat.num_vectors;
585 let quant = lazy_flat.quantization;
586 let fetch_k = k * rerank_factor.max(1);
587 let mut collector = crate::query::ScoreCollector::new(fetch_k);
588 let mut scores = vec![0f32; BRUTE_FORCE_BATCH];
589
590 for batch_start in (0..n).step_by(BRUTE_FORCE_BATCH) {
591 let batch_count = BRUTE_FORCE_BATCH.min(n - batch_start);
592 let batch_bytes = lazy_flat
593 .read_vectors_batch(batch_start, batch_count)
594 .await
595 .map_err(crate::Error::Io)?;
596 let raw = batch_bytes.as_slice();
597
598 Self::score_quantized_batch(query, raw, quant, dim, &mut scores[..batch_count]);
599
600 for (i, &score) in scores.iter().enumerate().take(batch_count) {
601 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
602 collector.insert_with_ordinal(doc_id, score, ordinal);
603 }
604 }
605
606 collector
607 .into_sorted_results()
608 .into_iter()
609 .map(|(doc_id, score, ordinal)| (doc_id, ordinal, score))
610 .collect()
611 } else {
612 return Ok(Vec::new());
613 };
614 let l1_elapsed = t0.elapsed();
615 log::debug!(
616 "[search_dense] field {}: L1 returned {} candidates in {:.1}ms",
617 field.0,
618 results.len(),
619 l1_elapsed.as_secs_f64() * 1000.0
620 );
621
622 if ann_index.is_some()
625 && !results.is_empty()
626 && let Some(lazy_flat) = lazy_flat
627 {
628 let t_rerank = std::time::Instant::now();
629 let dim = lazy_flat.dim;
630 let quant = lazy_flat.quantization;
631 let vbs = lazy_flat.vector_byte_size();
632
633 let mut resolved: Vec<(usize, usize)> = Vec::new(); for (ri, c) in results.iter().enumerate() {
636 let (start, entries) = lazy_flat.flat_indexes_for_doc(c.0);
637 for (j, &(_, ord)) in entries.iter().enumerate() {
638 if ord == c.1 {
639 resolved.push((ri, start + j));
640 break;
641 }
642 }
643 }
644
645 let t_resolve = t_rerank.elapsed();
646 if !resolved.is_empty() {
647 resolved.sort_unstable_by_key(|&(_, flat_idx)| flat_idx);
649
650 let t_read = std::time::Instant::now();
652 let mut raw_buf = vec![0u8; resolved.len() * vbs];
653 for (buf_idx, &(_, flat_idx)) in resolved.iter().enumerate() {
654 let _ = lazy_flat
655 .read_vector_raw_into(
656 flat_idx,
657 &mut raw_buf[buf_idx * vbs..(buf_idx + 1) * vbs],
658 )
659 .await;
660 }
661
662 let read_elapsed = t_read.elapsed();
663
664 let t_score = std::time::Instant::now();
666 let mut scores = vec![0f32; resolved.len()];
667 Self::score_quantized_batch(query, &raw_buf, quant, dim, &mut scores);
668 let score_elapsed = t_score.elapsed();
669
670 for (buf_idx, &(ri, _)) in resolved.iter().enumerate() {
672 results[ri].2 = scores[buf_idx];
673 }
674
675 log::debug!(
676 "[search_dense] field {}: rerank {} vectors (dim={}, quant={:?}, {}B/vec): resolve={:.1}ms read={:.1}ms score={:.1}ms",
677 field.0,
678 resolved.len(),
679 dim,
680 quant,
681 vbs,
682 t_resolve.as_secs_f64() * 1000.0,
683 read_elapsed.as_secs_f64() * 1000.0,
684 score_elapsed.as_secs_f64() * 1000.0,
685 );
686 }
687
688 results.sort_by(|a, b| b.2.partial_cmp(&a.2).unwrap_or(std::cmp::Ordering::Equal));
689 results.truncate(k * rerank_factor.max(1));
690 log::debug!(
691 "[search_dense] field {}: rerank total={:.1}ms",
692 field.0,
693 t_rerank.elapsed().as_secs_f64() * 1000.0
694 );
695 }
696
697 let mut doc_ordinals: rustc_hash::FxHashMap<DocId, Vec<(u32, f32)>> =
700 rustc_hash::FxHashMap::default();
701 for (doc_id, ordinal, score) in results {
702 let ordinals = doc_ordinals.entry(doc_id as DocId).or_default();
703 ordinals.push((ordinal as u32, score));
704 }
705
706 let mut final_results: Vec<VectorSearchResult> = doc_ordinals
708 .into_iter()
709 .map(|(doc_id, ordinals)| {
710 let combined_score = combiner.combine(&ordinals);
711 VectorSearchResult::new(doc_id, combined_score, ordinals)
712 })
713 .collect();
714
715 final_results.sort_by(|a, b| {
717 b.score
718 .partial_cmp(&a.score)
719 .unwrap_or(std::cmp::Ordering::Equal)
720 });
721 final_results.truncate(k);
722
723 Ok(final_results)
724 }
725
726 pub fn has_dense_vector_index(&self, field: Field) -> bool {
728 self.vector_indexes.contains_key(&field.0) || self.flat_vectors.contains_key(&field.0)
729 }
730
731 pub fn get_dense_vector_index(&self, field: Field) -> Option<Arc<RaBitQIndex>> {
733 match self.vector_indexes.get(&field.0) {
734 Some(VectorIndex::RaBitQ(lazy)) => lazy.get().cloned(),
735 _ => None,
736 }
737 }
738
739 pub fn get_ivf_vector_index(
741 &self,
742 field: Field,
743 ) -> Option<(Arc<IVFRaBitQIndex>, Arc<crate::structures::RaBitQCodebook>)> {
744 match self.vector_indexes.get(&field.0) {
745 Some(VectorIndex::IVF(lazy)) => lazy.get().map(|(i, c)| (i.clone(), c.clone())),
746 _ => None,
747 }
748 }
749
750 pub fn coarse_centroids(&self, field_id: u32) -> Option<&Arc<CoarseCentroids>> {
752 self.coarse_centroids.get(&field_id)
753 }
754
755 pub fn set_coarse_centroids(&mut self, centroids: FxHashMap<u32, Arc<CoarseCentroids>>) {
757 self.coarse_centroids = centroids;
758 }
759
760 pub fn get_scann_vector_index(
762 &self,
763 field: Field,
764 ) -> Option<(Arc<IVFPQIndex>, Arc<PQCodebook>)> {
765 match self.vector_indexes.get(&field.0) {
766 Some(VectorIndex::ScaNN(lazy)) => lazy.get().map(|(i, c)| (i.clone(), c.clone())),
767 _ => None,
768 }
769 }
770
771 pub fn get_vector_index(&self, field: Field) -> Option<&VectorIndex> {
773 self.vector_indexes.get(&field.0)
774 }
775
776 pub async fn search_sparse_vector(
786 &self,
787 field: Field,
788 vector: &[(u32, f32)],
789 limit: usize,
790 combiner: crate::query::MultiValueCombiner,
791 heap_factor: f32,
792 ) -> Result<Vec<VectorSearchResult>> {
793 use crate::query::{BlockMaxScoreExecutor, BmpExecutor, SparseTermScorer};
794
795 let query_tokens = vector.len();
796
797 let sparse_index = match self.sparse_indexes.get(&field.0) {
799 Some(idx) => idx,
800 None => {
801 log::debug!(
802 "Sparse vector search: no index for field {}, returning empty",
803 field.0
804 );
805 return Ok(Vec::new());
806 }
807 };
808
809 let index_dimensions = sparse_index.num_dimensions();
810
811 let mut matched_terms: Vec<(u32, f32)> = Vec::with_capacity(vector.len());
813 let mut missing_count = 0usize;
814
815 for &(dim_id, query_weight) in vector {
816 if sparse_index.has_dimension(dim_id) {
817 matched_terms.push((dim_id, query_weight));
818 } else {
819 missing_count += 1;
820 }
821 }
822
823 log::debug!(
824 "Sparse vector search: query_tokens={}, matched={}, missing={}, index_dimensions={}",
825 query_tokens,
826 matched_terms.len(),
827 missing_count,
828 index_dimensions
829 );
830
831 if matched_terms.is_empty() {
832 log::debug!("Sparse vector search: no matching tokens, returning empty");
833 return Ok(Vec::new());
834 }
835
836 let num_terms = matched_terms.len();
840 let over_fetch = limit * 2; let raw_results = if num_terms > 12 {
842 BmpExecutor::new(sparse_index, matched_terms, over_fetch, heap_factor)
844 .execute()
845 .await?
846 } else {
847 let mut posting_lists: Vec<(u32, f32, Arc<BlockSparsePostingList>)> =
849 Vec::with_capacity(num_terms);
850 for &(dim_id, query_weight) in &matched_terms {
851 if let Some(pl) = sparse_index.get_posting(dim_id).await? {
852 posting_lists.push((dim_id, query_weight, pl));
853 }
854 }
855 let scorers: Vec<SparseTermScorer> = posting_lists
856 .iter()
857 .map(|(_, query_weight, pl)| SparseTermScorer::from_arc(pl, *query_weight))
858 .collect();
859 if scorers.is_empty() {
860 return Ok(Vec::new());
861 }
862 BlockMaxScoreExecutor::with_heap_factor(scorers, over_fetch, heap_factor).execute()
863 };
864
865 log::trace!(
866 "Sparse WAND returned {} raw results for segment (doc_id_offset={})",
867 raw_results.len(),
868 self.doc_id_offset
869 );
870 if log::log_enabled!(log::Level::Trace) && !raw_results.is_empty() {
871 for r in raw_results.iter().take(5) {
872 log::trace!(
873 " Raw result: doc_id={} (global={}), score={:.4}, ordinal={}",
874 r.doc_id,
875 r.doc_id + self.doc_id_offset,
876 r.score,
877 r.ordinal
878 );
879 }
880 }
881
882 let mut doc_ordinals: rustc_hash::FxHashMap<u32, Vec<(u32, f32)>> =
885 rustc_hash::FxHashMap::default();
886 for r in raw_results {
887 let ordinals = doc_ordinals.entry(r.doc_id).or_default();
888 ordinals.push((r.ordinal as u32, r.score));
889 }
890
891 let mut results: Vec<VectorSearchResult> = doc_ordinals
894 .into_iter()
895 .map(|(doc_id, ordinals)| {
896 let combined_score = combiner.combine(&ordinals);
897 VectorSearchResult::new(doc_id, combined_score, ordinals)
898 })
899 .collect();
900
901 results.sort_by(|a, b| {
903 b.score
904 .partial_cmp(&a.score)
905 .unwrap_or(std::cmp::Ordering::Equal)
906 });
907 results.truncate(limit);
908
909 Ok(results)
910 }
911
912 pub async fn get_positions(
917 &self,
918 field: Field,
919 term: &[u8],
920 ) -> Result<Option<crate::structures::PositionPostingList>> {
921 let handle = match &self.positions_handle {
923 Some(h) => h,
924 None => return Ok(None),
925 };
926
927 let mut key = Vec::with_capacity(4 + term.len());
929 key.extend_from_slice(&field.0.to_le_bytes());
930 key.extend_from_slice(term);
931
932 let term_info = match self.term_dict.get(&key).await? {
934 Some(info) => info,
935 None => return Ok(None),
936 };
937
938 let (offset, length) = match term_info.position_info() {
940 Some((o, l)) => (o, l),
941 None => return Ok(None),
942 };
943
944 let slice = handle.slice(offset..offset + length as u64);
946 let data = slice.read_bytes().await?;
947
948 let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
950
951 Ok(Some(pos_list))
952 }
953
954 pub fn has_positions(&self, field: Field) -> bool {
956 if let Some(entry) = self.schema.get_field_entry(field) {
958 entry.positions.is_some()
959 } else {
960 false
961 }
962 }
963}
964
965pub type SegmentReader = AsyncSegmentReader;