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: Option<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 let coarse_centroids = vectors_data.coarse_centroids;
117
118 let sparse_indexes = loader::load_sparse_file(dir, &files, meta.num_docs, &schema).await?;
120
121 let positions_handle = loader::open_positions_file(dir, &files, &schema).await?;
123
124 let sparse_dims: usize = sparse_indexes.values().map(|s| s.num_dimensions()).sum();
126 let sparse_mem = sparse_dims * 24; log::debug!(
128 "[segment] loaded {:016x}: docs={}, sparse_dims={}, sparse_mem={:.2} KB, vectors={}",
129 segment_id.0,
130 meta.num_docs,
131 sparse_dims,
132 sparse_mem as f64 / 1024.0,
133 vector_indexes.len()
134 );
135
136 Ok(Self {
137 meta,
138 term_dict: Arc::new(term_dict),
139 postings_handle,
140 store: Arc::new(store),
141 schema,
142 doc_id_offset,
143 vector_indexes,
144 flat_vectors,
145 coarse_centroids,
146 sparse_indexes,
147 positions_handle,
148 })
149 }
150
151 pub fn meta(&self) -> &SegmentMeta {
152 &self.meta
153 }
154
155 pub fn num_docs(&self) -> u32 {
156 self.meta.num_docs
157 }
158
159 pub fn avg_field_len(&self, field: Field) -> f32 {
161 self.meta.avg_field_len(field)
162 }
163
164 pub fn doc_id_offset(&self) -> DocId {
165 self.doc_id_offset
166 }
167
168 pub fn set_doc_id_offset(&mut self, offset: DocId) {
170 self.doc_id_offset = offset;
171 }
172
173 pub fn schema(&self) -> &Schema {
174 &self.schema
175 }
176
177 pub fn sparse_indexes(&self) -> &FxHashMap<u32, SparseIndex> {
179 &self.sparse_indexes
180 }
181
182 pub fn vector_indexes(&self) -> &FxHashMap<u32, VectorIndex> {
184 &self.vector_indexes
185 }
186
187 pub fn flat_vectors(&self) -> &FxHashMap<u32, LazyFlatVectorData> {
189 &self.flat_vectors
190 }
191
192 pub fn term_dict_stats(&self) -> SSTableStats {
194 self.term_dict.stats()
195 }
196
197 pub fn memory_stats(&self) -> SegmentMemoryStats {
199 let term_dict_stats = self.term_dict.stats();
200
201 let term_dict_cache_bytes = self.term_dict.cached_blocks() * 4096;
203
204 let store_cache_bytes = self.store.cached_blocks() * 4096;
206
207 let sparse_index_bytes: usize = self
210 .sparse_indexes
211 .values()
212 .map(|s| s.num_dimensions() * 24)
213 .sum();
214
215 let dense_index_bytes: usize = self
218 .vector_indexes
219 .values()
220 .map(|v| v.estimated_memory_bytes())
221 .sum();
222
223 SegmentMemoryStats {
224 segment_id: self.meta.id,
225 num_docs: self.meta.num_docs,
226 term_dict_cache_bytes,
227 store_cache_bytes,
228 sparse_index_bytes,
229 dense_index_bytes,
230 bloom_filter_bytes: term_dict_stats.bloom_filter_size,
231 }
232 }
233
234 pub async fn get_postings(
239 &self,
240 field: Field,
241 term: &[u8],
242 ) -> Result<Option<BlockPostingList>> {
243 log::debug!(
244 "SegmentReader::get_postings field={} term_len={}",
245 field.0,
246 term.len()
247 );
248
249 let mut key = Vec::with_capacity(4 + term.len());
251 key.extend_from_slice(&field.0.to_le_bytes());
252 key.extend_from_slice(term);
253
254 let term_info = match self.term_dict.get(&key).await? {
256 Some(info) => {
257 log::debug!("SegmentReader::get_postings found term_info");
258 info
259 }
260 None => {
261 log::debug!("SegmentReader::get_postings term not found");
262 return Ok(None);
263 }
264 };
265
266 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
268 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
270 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs.into_iter()) {
271 posting_list.push(doc_id, tf);
272 }
273 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
274 return Ok(Some(block_list));
275 }
276
277 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
279 Error::Corruption("TermInfo has neither inline nor external data".to_string())
280 })?;
281
282 let start = posting_offset;
283 let end = start + posting_len as u64;
284
285 if end > self.postings_handle.len() {
286 return Err(Error::Corruption(
287 "Posting offset out of bounds".to_string(),
288 ));
289 }
290
291 let posting_bytes = self.postings_handle.read_bytes_range(start..end).await?;
292 let block_list = BlockPostingList::deserialize(&mut posting_bytes.as_slice())?;
293
294 Ok(Some(block_list))
295 }
296
297 pub async fn doc(&self, local_doc_id: DocId) -> Result<Option<Document>> {
302 let mut doc = match self.store.get(local_doc_id, &self.schema).await {
303 Ok(Some(d)) => d,
304 Ok(None) => return Ok(None),
305 Err(e) => return Err(Error::from(e)),
306 };
307
308 for (&field_id, lazy_flat) in &self.flat_vectors {
310 let (start, entries) = lazy_flat.flat_indexes_for_doc(local_doc_id);
311 for (j, &(_doc_id, _ordinal)) in entries.iter().enumerate() {
312 let flat_idx = start + j;
313 match lazy_flat.get_vector(flat_idx).await {
314 Ok(vec) => {
315 doc.add_dense_vector(Field(field_id), vec);
316 }
317 Err(e) => {
318 log::warn!("Failed to hydrate vector field {}: {}", field_id, e);
319 }
320 }
321 }
322 }
323
324 Ok(Some(doc))
325 }
326
327 pub async fn prefetch_terms(
329 &self,
330 field: Field,
331 start_term: &[u8],
332 end_term: &[u8],
333 ) -> Result<()> {
334 let mut start_key = Vec::with_capacity(4 + start_term.len());
335 start_key.extend_from_slice(&field.0.to_le_bytes());
336 start_key.extend_from_slice(start_term);
337
338 let mut end_key = Vec::with_capacity(4 + end_term.len());
339 end_key.extend_from_slice(&field.0.to_le_bytes());
340 end_key.extend_from_slice(end_term);
341
342 self.term_dict.prefetch_range(&start_key, &end_key).await?;
343 Ok(())
344 }
345
346 pub fn store_has_dict(&self) -> bool {
348 self.store.has_dict()
349 }
350
351 pub fn store(&self) -> &super::store::AsyncStoreReader {
353 &self.store
354 }
355
356 pub fn store_raw_blocks(&self) -> Vec<RawStoreBlock> {
358 self.store.raw_blocks()
359 }
360
361 pub fn store_data_slice(&self) -> &LazyFileSlice {
363 self.store.data_slice()
364 }
365
366 pub async fn all_terms(&self) -> Result<Vec<(Vec<u8>, TermInfo)>> {
368 self.term_dict.all_entries().await.map_err(Error::from)
369 }
370
371 pub async fn all_terms_with_stats(&self) -> Result<Vec<(Field, String, u32)>> {
376 let entries = self.term_dict.all_entries().await?;
377 let mut result = Vec::with_capacity(entries.len());
378
379 for (key, term_info) in entries {
380 if key.len() > 4 {
382 let field_id = u32::from_le_bytes([key[0], key[1], key[2], key[3]]);
383 let term_bytes = &key[4..];
384 if let Ok(term_str) = std::str::from_utf8(term_bytes) {
385 result.push((Field(field_id), term_str.to_string(), term_info.doc_freq()));
386 }
387 }
388 }
389
390 Ok(result)
391 }
392
393 pub fn term_dict_iter(&self) -> crate::structures::AsyncSSTableIterator<'_, TermInfo> {
395 self.term_dict.iter()
396 }
397
398 pub async fn prefetch_term_dict(&self) -> crate::Result<()> {
402 self.term_dict
403 .prefetch_all_data_bulk()
404 .await
405 .map_err(crate::Error::from)
406 }
407
408 pub async fn read_postings(&self, offset: u64, len: u32) -> Result<Vec<u8>> {
410 let start = offset;
411 let end = start + len as u64;
412 let bytes = self.postings_handle.read_bytes_range(start..end).await?;
413 Ok(bytes.to_vec())
414 }
415
416 pub async fn read_position_bytes(&self, offset: u64, len: u32) -> Result<Option<Vec<u8>>> {
418 let handle = match &self.positions_handle {
419 Some(h) => h,
420 None => return Ok(None),
421 };
422 let start = offset;
423 let end = start + len as u64;
424 let bytes = handle.read_bytes_range(start..end).await?;
425 Ok(Some(bytes.to_vec()))
426 }
427
428 pub fn has_positions_file(&self) -> bool {
430 self.positions_handle.is_some()
431 }
432
433 pub async fn search_dense_vector(
439 &self,
440 field: Field,
441 query: &[f32],
442 k: usize,
443 nprobe: usize,
444 rerank_factor: usize,
445 combiner: crate::query::MultiValueCombiner,
446 ) -> Result<Vec<VectorSearchResult>> {
447 let ann_index = self.vector_indexes.get(&field.0);
448 let lazy_flat = self.flat_vectors.get(&field.0);
449
450 if ann_index.is_none() && lazy_flat.is_none() {
452 return Ok(Vec::new());
453 }
454
455 const BRUTE_FORCE_BATCH: usize = 4096;
457
458 let mut results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
460 match index {
462 VectorIndex::RaBitQ(rabitq) => {
463 let fetch_k = k * rerank_factor.max(1);
464 rabitq
465 .search(query, fetch_k, rerank_factor)
466 .into_iter()
467 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
468 .collect()
469 }
470 VectorIndex::IVF { index, codebook } => {
471 let centroids = self.coarse_centroids.as_ref().ok_or_else(|| {
472 Error::Schema("IVF index requires coarse centroids".to_string())
473 })?;
474 let effective_nprobe = if nprobe > 0 { nprobe } else { 32 };
475 let fetch_k = k * rerank_factor.max(1);
476 index
477 .search(centroids, codebook, query, fetch_k, Some(effective_nprobe))
478 .into_iter()
479 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
480 .collect()
481 }
482 VectorIndex::ScaNN { index, codebook } => {
483 let centroids = self.coarse_centroids.as_ref().ok_or_else(|| {
484 Error::Schema("ScaNN index requires coarse centroids".to_string())
485 })?;
486 let effective_nprobe = if nprobe > 0 { nprobe } else { 32 };
487 let fetch_k = k * rerank_factor.max(1);
488 index
489 .search(centroids, codebook, query, fetch_k, Some(effective_nprobe))
490 .into_iter()
491 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
492 .collect()
493 }
494 }
495 } else if let Some(lazy_flat) = lazy_flat {
496 let dim = lazy_flat.dim;
498 let n = lazy_flat.num_vectors;
499 let quant = lazy_flat.quantization;
500 let fetch_k = k * rerank_factor.max(1);
501 let mut candidates: Vec<(u32, u16, f32)> = Vec::new();
502
503 for batch_start in (0..n).step_by(BRUTE_FORCE_BATCH) {
504 let batch_count = BRUTE_FORCE_BATCH.min(n - batch_start);
505 let batch_bytes = lazy_flat
506 .read_vectors_batch(batch_start, batch_count)
507 .await
508 .map_err(crate::Error::Io)?;
509 let raw = batch_bytes.as_slice();
510
511 let mut scores = vec![0f32; batch_count];
512
513 match quant {
514 crate::dsl::DenseVectorQuantization::F32 => {
515 let batch_floats = batch_count * dim;
516 let mut aligned_buf: Vec<f32> = Vec::new();
517 let vectors: &[f32] = if (raw.as_ptr() as usize)
518 .is_multiple_of(std::mem::align_of::<f32>())
519 {
520 unsafe {
521 std::slice::from_raw_parts(raw.as_ptr() as *const f32, batch_floats)
522 }
523 } else {
524 aligned_buf.resize(batch_floats, 0.0);
525 unsafe {
526 std::ptr::copy_nonoverlapping(
527 raw.as_ptr(),
528 aligned_buf.as_mut_ptr() as *mut u8,
529 batch_floats * std::mem::size_of::<f32>(),
530 );
531 }
532 &aligned_buf
533 };
534 crate::structures::simd::batch_cosine_scores(
535 query,
536 vectors,
537 dim,
538 &mut scores,
539 );
540 }
541 crate::dsl::DenseVectorQuantization::F16 => {
542 crate::structures::simd::batch_cosine_scores_f16(
543 query,
544 raw,
545 dim,
546 &mut scores,
547 );
548 }
549 crate::dsl::DenseVectorQuantization::UInt8 => {
550 crate::structures::simd::batch_cosine_scores_u8(
551 query,
552 raw,
553 dim,
554 &mut scores,
555 );
556 }
557 }
558
559 for (i, &score) in scores.iter().enumerate().take(batch_count) {
560 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
561 candidates.push((doc_id, ordinal, score));
562 }
563 }
564
565 candidates.sort_by(|a, b| b.2.partial_cmp(&a.2).unwrap_or(std::cmp::Ordering::Equal));
566 candidates.truncate(fetch_k);
567 candidates
568 } else {
569 return Ok(Vec::new());
570 };
571
572 if ann_index.is_some()
575 && !results.is_empty()
576 && let Some(lazy_flat) = lazy_flat
577 {
578 let dim = lazy_flat.dim;
579 let quant = lazy_flat.quantization;
580 let vbs = lazy_flat.vector_byte_size();
581
582 let mut resolved: Vec<(usize, usize)> = Vec::new(); for (ri, c) in results.iter().enumerate() {
585 let (start, entries) = lazy_flat.flat_indexes_for_doc(c.0);
586 for (j, &(_, ord)) in entries.iter().enumerate() {
587 if ord == c.1 {
588 resolved.push((ri, start + j));
589 break;
590 }
591 }
592 }
593
594 if !resolved.is_empty() {
595 let mut raw_buf = vec![0u8; resolved.len() * vbs];
597 for (buf_idx, &(_, flat_idx)) in resolved.iter().enumerate() {
598 let _ = lazy_flat
599 .read_vector_raw_into(
600 flat_idx,
601 &mut raw_buf[buf_idx * vbs..(buf_idx + 1) * vbs],
602 )
603 .await;
604 }
605
606 let mut scores = vec![0f32; resolved.len()];
608 match quant {
609 crate::dsl::DenseVectorQuantization::F32 => {
610 let floats = resolved.len() * dim;
611 let mut aligned_buf: Vec<f32> = Vec::new();
612 let vectors: &[f32] = if (raw_buf.as_ptr() as usize)
613 .is_multiple_of(std::mem::align_of::<f32>())
614 {
615 unsafe {
616 std::slice::from_raw_parts(raw_buf.as_ptr() as *const f32, floats)
617 }
618 } else {
619 aligned_buf.resize(floats, 0.0);
620 unsafe {
621 std::ptr::copy_nonoverlapping(
622 raw_buf.as_ptr(),
623 aligned_buf.as_mut_ptr() as *mut u8,
624 floats * std::mem::size_of::<f32>(),
625 );
626 }
627 &aligned_buf
628 };
629 crate::structures::simd::batch_cosine_scores(
630 query,
631 vectors,
632 dim,
633 &mut scores,
634 );
635 }
636 crate::dsl::DenseVectorQuantization::F16 => {
637 crate::structures::simd::batch_cosine_scores_f16(
638 query,
639 &raw_buf,
640 dim,
641 &mut scores,
642 );
643 }
644 crate::dsl::DenseVectorQuantization::UInt8 => {
645 crate::structures::simd::batch_cosine_scores_u8(
646 query,
647 &raw_buf,
648 dim,
649 &mut scores,
650 );
651 }
652 }
653
654 for (buf_idx, &(ri, _)) in resolved.iter().enumerate() {
656 results[ri].2 = scores[buf_idx];
657 }
658 }
659
660 results.sort_by(|a, b| b.2.partial_cmp(&a.2).unwrap_or(std::cmp::Ordering::Equal));
661 results.truncate(k * rerank_factor.max(1));
662 }
663
664 let mut doc_ordinals: rustc_hash::FxHashMap<DocId, Vec<(u32, f32)>> =
667 rustc_hash::FxHashMap::default();
668 for (doc_id, ordinal, score) in results {
669 let ordinals = doc_ordinals.entry(doc_id as DocId).or_default();
670 ordinals.push((ordinal as u32, score));
671 }
672
673 let mut final_results: Vec<VectorSearchResult> = doc_ordinals
675 .into_iter()
676 .map(|(doc_id, ordinals)| {
677 let combined_score = combiner.combine(&ordinals);
678 VectorSearchResult::new(doc_id, combined_score, ordinals)
679 })
680 .collect();
681
682 final_results.sort_by(|a, b| {
684 b.score
685 .partial_cmp(&a.score)
686 .unwrap_or(std::cmp::Ordering::Equal)
687 });
688 final_results.truncate(k);
689
690 Ok(final_results)
691 }
692
693 pub fn has_dense_vector_index(&self, field: Field) -> bool {
695 self.vector_indexes.contains_key(&field.0) || self.flat_vectors.contains_key(&field.0)
696 }
697
698 pub fn get_dense_vector_index(&self, field: Field) -> Option<Arc<RaBitQIndex>> {
700 match self.vector_indexes.get(&field.0) {
701 Some(VectorIndex::RaBitQ(idx)) => Some(idx.clone()),
702 _ => None,
703 }
704 }
705
706 pub fn get_ivf_vector_index(
708 &self,
709 field: Field,
710 ) -> Option<(Arc<IVFRaBitQIndex>, Arc<crate::structures::RaBitQCodebook>)> {
711 match self.vector_indexes.get(&field.0) {
712 Some(VectorIndex::IVF { index, codebook }) => Some((index.clone(), codebook.clone())),
713 _ => None,
714 }
715 }
716
717 pub fn coarse_centroids(&self) -> Option<&Arc<CoarseCentroids>> {
719 self.coarse_centroids.as_ref()
720 }
721
722 pub fn get_scann_vector_index(
724 &self,
725 field: Field,
726 ) -> Option<(Arc<IVFPQIndex>, Arc<PQCodebook>)> {
727 match self.vector_indexes.get(&field.0) {
728 Some(VectorIndex::ScaNN { index, codebook }) => Some((index.clone(), codebook.clone())),
729 _ => None,
730 }
731 }
732
733 pub fn get_vector_index(&self, field: Field) -> Option<&VectorIndex> {
735 self.vector_indexes.get(&field.0)
736 }
737
738 pub async fn search_sparse_vector(
748 &self,
749 field: Field,
750 vector: &[(u32, f32)],
751 limit: usize,
752 combiner: crate::query::MultiValueCombiner,
753 heap_factor: f32,
754 ) -> Result<Vec<VectorSearchResult>> {
755 use crate::query::{BlockMaxScoreExecutor, BmpExecutor, SparseTermScorer};
756
757 let query_tokens = vector.len();
758
759 let sparse_index = match self.sparse_indexes.get(&field.0) {
761 Some(idx) => idx,
762 None => {
763 log::debug!(
764 "Sparse vector search: no index for field {}, returning empty",
765 field.0
766 );
767 return Ok(Vec::new());
768 }
769 };
770
771 let index_dimensions = sparse_index.num_dimensions();
772
773 let mut matched_terms: Vec<(u32, f32)> = Vec::with_capacity(vector.len());
775 let mut missing_tokens = Vec::new();
776
777 for &(dim_id, query_weight) in vector {
778 if sparse_index.has_dimension(dim_id) {
779 matched_terms.push((dim_id, query_weight));
780 } else {
781 missing_tokens.push(dim_id);
782 }
783 }
784
785 log::debug!(
786 "Sparse vector search: query_tokens={}, matched={}, missing={}, index_dimensions={}",
787 query_tokens,
788 matched_terms.len(),
789 missing_tokens.len(),
790 index_dimensions
791 );
792
793 if log::log_enabled!(log::Level::Debug) {
794 let query_details: Vec<_> = vector
795 .iter()
796 .take(30)
797 .map(|(id, w)| format!("{}:{:.3}", id, w))
798 .collect();
799 log::debug!("Query tokens (id:weight): [{}]", query_details.join(", "));
800 }
801
802 if !missing_tokens.is_empty() {
803 log::debug!(
804 "Missing token IDs (not in index): {:?}",
805 missing_tokens.iter().take(20).collect::<Vec<_>>()
806 );
807 }
808
809 if matched_terms.is_empty() {
810 log::debug!("Sparse vector search: no matching tokens, returning empty");
811 return Ok(Vec::new());
812 }
813
814 let num_terms = matched_terms.len();
818 let over_fetch = limit * 2; let raw_results = if num_terms > 12 {
820 BmpExecutor::new(sparse_index, matched_terms, over_fetch, heap_factor)
822 .execute()
823 .await?
824 } else {
825 let mut posting_lists: Vec<(u32, f32, Arc<BlockSparsePostingList>)> =
827 Vec::with_capacity(num_terms);
828 for &(dim_id, query_weight) in &matched_terms {
829 if let Some(pl) = sparse_index.get_posting(dim_id).await? {
830 posting_lists.push((dim_id, query_weight, pl));
831 }
832 }
833 let scorers: Vec<SparseTermScorer> = posting_lists
834 .iter()
835 .map(|(_, query_weight, pl)| SparseTermScorer::from_arc(pl, *query_weight))
836 .collect();
837 if scorers.is_empty() {
838 return Ok(Vec::new());
839 }
840 BlockMaxScoreExecutor::with_heap_factor(scorers, over_fetch, heap_factor).execute()
841 };
842
843 log::trace!(
844 "Sparse WAND returned {} raw results for segment (doc_id_offset={})",
845 raw_results.len(),
846 self.doc_id_offset
847 );
848 if log::log_enabled!(log::Level::Trace) && !raw_results.is_empty() {
849 for r in raw_results.iter().take(5) {
850 log::trace!(
851 " Raw result: doc_id={} (global={}), score={:.4}, ordinal={}",
852 r.doc_id,
853 r.doc_id + self.doc_id_offset,
854 r.score,
855 r.ordinal
856 );
857 }
858 }
859
860 let mut doc_ordinals: rustc_hash::FxHashMap<u32, Vec<(u32, f32)>> =
863 rustc_hash::FxHashMap::default();
864 for r in raw_results {
865 let ordinals = doc_ordinals.entry(r.doc_id).or_default();
866 ordinals.push((r.ordinal as u32, r.score));
867 }
868
869 let mut results: Vec<VectorSearchResult> = doc_ordinals
872 .into_iter()
873 .map(|(doc_id, ordinals)| {
874 let combined_score = combiner.combine(&ordinals);
875 VectorSearchResult::new(doc_id, combined_score, ordinals)
876 })
877 .collect();
878
879 results.sort_by(|a, b| {
881 b.score
882 .partial_cmp(&a.score)
883 .unwrap_or(std::cmp::Ordering::Equal)
884 });
885 results.truncate(limit);
886
887 Ok(results)
888 }
889
890 pub async fn get_positions(
895 &self,
896 field: Field,
897 term: &[u8],
898 ) -> Result<Option<crate::structures::PositionPostingList>> {
899 use std::io::Cursor;
900
901 let handle = match &self.positions_handle {
903 Some(h) => h,
904 None => return Ok(None),
905 };
906
907 let mut key = Vec::with_capacity(4 + term.len());
909 key.extend_from_slice(&field.0.to_le_bytes());
910 key.extend_from_slice(term);
911
912 let term_info = match self.term_dict.get(&key).await? {
914 Some(info) => info,
915 None => return Ok(None),
916 };
917
918 let (offset, length) = match term_info.position_info() {
920 Some((o, l)) => (o, l),
921 None => return Ok(None),
922 };
923
924 let slice = handle.slice(offset..offset + length as u64);
926 let data = slice.read_bytes().await?;
927
928 let mut cursor = Cursor::new(data.as_slice());
930 let pos_list = crate::structures::PositionPostingList::deserialize(&mut cursor)?;
931
932 Ok(Some(pos_list))
933 }
934
935 pub fn has_positions(&self, field: Field) -> bool {
937 if let Some(entry) = self.schema.get_field_entry(field) {
939 entry.positions.is_some()
940 } else {
941 false
942 }
943 }
944}
945
946pub type SegmentReader = AsyncSegmentReader;