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 std::sync::Arc;
39
40use rustc_hash::FxHashMap;
41
42use super::vector_data::LazyFlatVectorData;
43use crate::directories::{AsyncFileRead, Directory, LazyFileHandle, LazyFileSlice};
44use crate::dsl::{Document, Field, Schema};
45use crate::structures::{
46 AsyncSSTableReader, BlockPostingList, CoarseCentroids, IVFPQIndex, IVFRaBitQIndex, PQCodebook,
47 RaBitQIndex, SSTableStats, TermInfo,
48};
49use crate::{DocId, Error, Result};
50
51use super::store::{AsyncStoreReader, RawStoreBlock};
52use super::types::{SegmentFiles, SegmentId, SegmentMeta};
53
54pub struct AsyncSegmentReader {
60 meta: SegmentMeta,
61 term_dict: Arc<AsyncSSTableReader<TermInfo>>,
63 postings_handle: LazyFileHandle,
65 store: Arc<AsyncStoreReader>,
67 schema: Arc<Schema>,
68 doc_id_offset: DocId,
70 vector_indexes: FxHashMap<u32, VectorIndex>,
72 flat_vectors: FxHashMap<u32, LazyFlatVectorData>,
74 coarse_centroids: FxHashMap<u32, Arc<CoarseCentroids>>,
76 sparse_indexes: FxHashMap<u32, SparseIndex>,
78 positions_handle: Option<LazyFileHandle>,
80}
81
82impl AsyncSegmentReader {
83 pub async fn open<D: Directory>(
85 dir: &D,
86 segment_id: SegmentId,
87 schema: Arc<Schema>,
88 doc_id_offset: DocId,
89 cache_blocks: usize,
90 ) -> Result<Self> {
91 let files = SegmentFiles::new(segment_id.0);
92
93 let meta_slice = dir.open_read(&files.meta).await?;
95 let meta_bytes = meta_slice.read_bytes().await?;
96 let meta = SegmentMeta::deserialize(meta_bytes.as_slice())?;
97 debug_assert_eq!(meta.id, segment_id.0);
98
99 let term_dict_handle = dir.open_lazy(&files.term_dict).await?;
101 let term_dict = AsyncSSTableReader::open(term_dict_handle, cache_blocks).await?;
102
103 let postings_handle = dir.open_lazy(&files.postings).await?;
105
106 let store_handle = dir.open_lazy(&files.store).await?;
108 let store = AsyncStoreReader::open(store_handle, cache_blocks).await?;
109
110 let vectors_data = loader::load_vectors_file(dir, &files, &schema).await?;
112 let vector_indexes = vectors_data.indexes;
113 let flat_vectors = vectors_data.flat_vectors;
114
115 let sparse_indexes = loader::load_sparse_file(dir, &files, meta.num_docs, &schema).await?;
117
118 let positions_handle = loader::open_positions_file(dir, &files, &schema).await?;
120
121 {
123 let mut parts = vec![format!(
124 "[segment] loaded {:016x}: docs={}",
125 segment_id.0, meta.num_docs
126 )];
127 if !vector_indexes.is_empty() || !flat_vectors.is_empty() {
128 parts.push(format!(
129 "dense: {} ann + {} flat fields",
130 vector_indexes.len(),
131 flat_vectors.len()
132 ));
133 }
134 for (field_id, idx) in &sparse_indexes {
135 parts.push(format!(
136 "sparse field {}: {} dims, ~{:.1} KB",
137 field_id,
138 idx.num_dimensions(),
139 idx.num_dimensions() as f64 * 24.0 / 1024.0
140 ));
141 }
142 log::debug!("{}", parts.join(", "));
143 }
144
145 Ok(Self {
146 meta,
147 term_dict: Arc::new(term_dict),
148 postings_handle,
149 store: Arc::new(store),
150 schema,
151 doc_id_offset,
152 vector_indexes,
153 flat_vectors,
154 coarse_centroids: FxHashMap::default(),
155 sparse_indexes,
156 positions_handle,
157 })
158 }
159
160 pub fn meta(&self) -> &SegmentMeta {
161 &self.meta
162 }
163
164 pub fn num_docs(&self) -> u32 {
165 self.meta.num_docs
166 }
167
168 pub fn avg_field_len(&self, field: Field) -> f32 {
170 self.meta.avg_field_len(field)
171 }
172
173 pub fn doc_id_offset(&self) -> DocId {
174 self.doc_id_offset
175 }
176
177 pub fn set_doc_id_offset(&mut self, offset: DocId) {
179 self.doc_id_offset = offset;
180 }
181
182 pub fn schema(&self) -> &Schema {
183 &self.schema
184 }
185
186 pub fn sparse_indexes(&self) -> &FxHashMap<u32, SparseIndex> {
188 &self.sparse_indexes
189 }
190
191 pub fn vector_indexes(&self) -> &FxHashMap<u32, VectorIndex> {
193 &self.vector_indexes
194 }
195
196 pub fn flat_vectors(&self) -> &FxHashMap<u32, LazyFlatVectorData> {
198 &self.flat_vectors
199 }
200
201 pub fn term_dict_stats(&self) -> SSTableStats {
203 self.term_dict.stats()
204 }
205
206 pub fn memory_stats(&self) -> SegmentMemoryStats {
208 let term_dict_stats = self.term_dict.stats();
209
210 let term_dict_cache_bytes = self.term_dict.cached_blocks() * 4096;
212
213 let store_cache_bytes = self.store.cached_blocks() * 4096;
215
216 let sparse_index_bytes: usize = self
218 .sparse_indexes
219 .values()
220 .map(|s| s.estimated_memory_bytes())
221 .sum();
222
223 let dense_index_bytes: usize = self
226 .vector_indexes
227 .values()
228 .map(|v| v.estimated_memory_bytes())
229 .sum();
230
231 SegmentMemoryStats {
232 segment_id: self.meta.id,
233 num_docs: self.meta.num_docs,
234 term_dict_cache_bytes,
235 store_cache_bytes,
236 sparse_index_bytes,
237 dense_index_bytes,
238 bloom_filter_bytes: term_dict_stats.bloom_filter_size,
239 }
240 }
241
242 pub async fn get_postings(
247 &self,
248 field: Field,
249 term: &[u8],
250 ) -> Result<Option<BlockPostingList>> {
251 log::debug!(
252 "SegmentReader::get_postings field={} term_len={}",
253 field.0,
254 term.len()
255 );
256
257 let mut key = Vec::with_capacity(4 + term.len());
259 key.extend_from_slice(&field.0.to_le_bytes());
260 key.extend_from_slice(term);
261
262 let term_info = match self.term_dict.get(&key).await? {
264 Some(info) => {
265 log::debug!("SegmentReader::get_postings found term_info");
266 info
267 }
268 None => {
269 log::debug!("SegmentReader::get_postings term not found");
270 return Ok(None);
271 }
272 };
273
274 if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
276 let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
278 for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs.into_iter()) {
279 posting_list.push(doc_id, tf);
280 }
281 let block_list = BlockPostingList::from_posting_list(&posting_list)?;
282 return Ok(Some(block_list));
283 }
284
285 let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
287 Error::Corruption("TermInfo has neither inline nor external data".to_string())
288 })?;
289
290 let start = posting_offset;
291 let end = start + posting_len as u64;
292
293 if end > self.postings_handle.len() {
294 return Err(Error::Corruption(
295 "Posting offset out of bounds".to_string(),
296 ));
297 }
298
299 let posting_bytes = self.postings_handle.read_bytes_range(start..end).await?;
300 let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
301
302 Ok(Some(block_list))
303 }
304
305 pub async fn doc(&self, local_doc_id: DocId) -> Result<Option<Document>> {
310 self.doc_with_fields(local_doc_id, None).await
311 }
312
313 pub async fn doc_with_fields(
319 &self,
320 local_doc_id: DocId,
321 fields: Option<&rustc_hash::FxHashSet<u32>>,
322 ) -> Result<Option<Document>> {
323 let mut doc = match self.store.get(local_doc_id, &self.schema).await {
324 Ok(Some(d)) => d,
325 Ok(None) => return Ok(None),
326 Err(e) => return Err(Error::from(e)),
327 };
328
329 for (&field_id, lazy_flat) in &self.flat_vectors {
331 if let Some(set) = fields
333 && !set.contains(&field_id)
334 {
335 continue;
336 }
337
338 let (start, entries) = lazy_flat.flat_indexes_for_doc(local_doc_id);
339 for (j, &(_doc_id, _ordinal)) in entries.iter().enumerate() {
340 let flat_idx = start + j;
341 match lazy_flat.get_vector(flat_idx).await {
342 Ok(vec) => {
343 doc.add_dense_vector(Field(field_id), vec);
344 }
345 Err(e) => {
346 log::warn!("Failed to hydrate vector field {}: {}", field_id, e);
347 }
348 }
349 }
350 }
351
352 Ok(Some(doc))
353 }
354
355 pub async fn prefetch_terms(
357 &self,
358 field: Field,
359 start_term: &[u8],
360 end_term: &[u8],
361 ) -> Result<()> {
362 let mut start_key = Vec::with_capacity(4 + start_term.len());
363 start_key.extend_from_slice(&field.0.to_le_bytes());
364 start_key.extend_from_slice(start_term);
365
366 let mut end_key = Vec::with_capacity(4 + end_term.len());
367 end_key.extend_from_slice(&field.0.to_le_bytes());
368 end_key.extend_from_slice(end_term);
369
370 self.term_dict.prefetch_range(&start_key, &end_key).await?;
371 Ok(())
372 }
373
374 pub fn store_has_dict(&self) -> bool {
376 self.store.has_dict()
377 }
378
379 pub fn store(&self) -> &super::store::AsyncStoreReader {
381 &self.store
382 }
383
384 pub fn store_raw_blocks(&self) -> Vec<RawStoreBlock> {
386 self.store.raw_blocks()
387 }
388
389 pub fn store_data_slice(&self) -> &LazyFileSlice {
391 self.store.data_slice()
392 }
393
394 pub async fn all_terms(&self) -> Result<Vec<(Vec<u8>, TermInfo)>> {
396 self.term_dict.all_entries().await.map_err(Error::from)
397 }
398
399 pub async fn all_terms_with_stats(&self) -> Result<Vec<(Field, String, u32)>> {
404 let entries = self.term_dict.all_entries().await?;
405 let mut result = Vec::with_capacity(entries.len());
406
407 for (key, term_info) in entries {
408 if key.len() > 4 {
410 let field_id = u32::from_le_bytes([key[0], key[1], key[2], key[3]]);
411 let term_bytes = &key[4..];
412 if let Ok(term_str) = std::str::from_utf8(term_bytes) {
413 result.push((Field(field_id), term_str.to_string(), term_info.doc_freq()));
414 }
415 }
416 }
417
418 Ok(result)
419 }
420
421 pub fn term_dict_iter(&self) -> crate::structures::AsyncSSTableIterator<'_, TermInfo> {
423 self.term_dict.iter()
424 }
425
426 pub async fn prefetch_term_dict(&self) -> crate::Result<()> {
430 self.term_dict
431 .prefetch_all_data_bulk()
432 .await
433 .map_err(crate::Error::from)
434 }
435
436 pub async fn read_postings(&self, offset: u64, len: u32) -> Result<Vec<u8>> {
438 let start = offset;
439 let end = start + len as u64;
440 let bytes = self.postings_handle.read_bytes_range(start..end).await?;
441 Ok(bytes.to_vec())
442 }
443
444 pub async fn read_position_bytes(&self, offset: u64, len: u32) -> Result<Option<Vec<u8>>> {
446 let handle = match &self.positions_handle {
447 Some(h) => h,
448 None => return Ok(None),
449 };
450 let start = offset;
451 let end = start + len as u64;
452 let bytes = handle.read_bytes_range(start..end).await?;
453 Ok(Some(bytes.to_vec()))
454 }
455
456 pub fn has_positions_file(&self) -> bool {
458 self.positions_handle.is_some()
459 }
460
461 fn score_quantized_batch(
467 query: &[f32],
468 raw: &[u8],
469 quant: crate::dsl::DenseVectorQuantization,
470 dim: usize,
471 scores: &mut [f32],
472 unit_norm: bool,
473 ) {
474 use crate::dsl::DenseVectorQuantization;
475 use crate::structures::simd;
476 match (quant, unit_norm) {
477 (DenseVectorQuantization::F32, false) => {
478 let num_floats = scores.len() * dim;
479 debug_assert!(
480 (raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()),
481 "f32 vector data not 4-byte aligned — vectors file may use legacy format"
482 );
483 let vectors: &[f32] =
484 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
485 simd::batch_cosine_scores(query, vectors, dim, scores);
486 }
487 (DenseVectorQuantization::F32, true) => {
488 let num_floats = scores.len() * dim;
489 debug_assert!(
490 (raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()),
491 "f32 vector data not 4-byte aligned"
492 );
493 let vectors: &[f32] =
494 unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
495 simd::batch_dot_scores(query, vectors, dim, scores);
496 }
497 (DenseVectorQuantization::F16, false) => {
498 simd::batch_cosine_scores_f16(query, raw, dim, scores);
499 }
500 (DenseVectorQuantization::F16, true) => {
501 simd::batch_dot_scores_f16(query, raw, dim, scores);
502 }
503 (DenseVectorQuantization::UInt8, false) => {
504 simd::batch_cosine_scores_u8(query, raw, dim, scores);
505 }
506 (DenseVectorQuantization::UInt8, true) => {
507 simd::batch_dot_scores_u8(query, raw, dim, scores);
508 }
509 }
510 }
511
512 pub async fn search_dense_vector(
518 &self,
519 field: Field,
520 query: &[f32],
521 k: usize,
522 nprobe: usize,
523 rerank_factor: usize,
524 combiner: crate::query::MultiValueCombiner,
525 ) -> Result<Vec<VectorSearchResult>> {
526 let ann_index = self.vector_indexes.get(&field.0);
527 let lazy_flat = self.flat_vectors.get(&field.0);
528
529 if ann_index.is_none() && lazy_flat.is_none() {
531 return Ok(Vec::new());
532 }
533
534 let unit_norm = self
536 .schema
537 .get_field_entry(field)
538 .and_then(|e| e.dense_vector_config.as_ref())
539 .is_some_and(|c| c.unit_norm);
540
541 const BRUTE_FORCE_BATCH: usize = 4096;
543
544 let t0 = std::time::Instant::now();
546 let mut results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
547 match index {
549 VectorIndex::RaBitQ(lazy) => {
550 let rabitq = lazy.get().ok_or_else(|| {
551 Error::Schema("RaBitQ index deserialization failed".to_string())
552 })?;
553 let fetch_k = k * rerank_factor.max(1);
554 rabitq
555 .search(query, fetch_k, rerank_factor)
556 .into_iter()
557 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
558 .collect()
559 }
560 VectorIndex::IVF(lazy) => {
561 let (index, codebook) = lazy.get().ok_or_else(|| {
562 Error::Schema("IVF index deserialization failed".to_string())
563 })?;
564 let centroids = self.coarse_centroids.get(&field.0).ok_or_else(|| {
565 Error::Schema(format!(
566 "IVF index requires coarse centroids for field {}",
567 field.0
568 ))
569 })?;
570 let effective_nprobe = if nprobe > 0 { nprobe } else { 32 };
571 let fetch_k = k * rerank_factor.max(1);
572 index
573 .search(centroids, codebook, query, fetch_k, Some(effective_nprobe))
574 .into_iter()
575 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
576 .collect()
577 }
578 VectorIndex::ScaNN(lazy) => {
579 let (index, codebook) = lazy.get().ok_or_else(|| {
580 Error::Schema("ScaNN index deserialization failed".to_string())
581 })?;
582 let centroids = self.coarse_centroids.get(&field.0).ok_or_else(|| {
583 Error::Schema(format!(
584 "ScaNN index requires coarse centroids for field {}",
585 field.0
586 ))
587 })?;
588 let effective_nprobe = if nprobe > 0 { nprobe } else { 32 };
589 let fetch_k = k * rerank_factor.max(1);
590 index
591 .search(centroids, codebook, query, fetch_k, Some(effective_nprobe))
592 .into_iter()
593 .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
594 .collect()
595 }
596 }
597 } else if let Some(lazy_flat) = lazy_flat {
598 log::debug!(
601 "[search_dense] field {}: brute-force on {} vectors (dim={}, quant={:?})",
602 field.0,
603 lazy_flat.num_vectors,
604 lazy_flat.dim,
605 lazy_flat.quantization
606 );
607 let dim = lazy_flat.dim;
608 let n = lazy_flat.num_vectors;
609 let quant = lazy_flat.quantization;
610 let fetch_k = k * rerank_factor.max(1);
611 let mut collector = crate::query::ScoreCollector::new(fetch_k);
612 let mut scores = vec![0f32; BRUTE_FORCE_BATCH];
613
614 for batch_start in (0..n).step_by(BRUTE_FORCE_BATCH) {
615 let batch_count = BRUTE_FORCE_BATCH.min(n - batch_start);
616 let batch_bytes = lazy_flat
617 .read_vectors_batch(batch_start, batch_count)
618 .await
619 .map_err(crate::Error::Io)?;
620 let raw = batch_bytes.as_slice();
621
622 Self::score_quantized_batch(
623 query,
624 raw,
625 quant,
626 dim,
627 &mut scores[..batch_count],
628 unit_norm,
629 );
630
631 for (i, &score) in scores.iter().enumerate().take(batch_count) {
632 let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
633 collector.insert_with_ordinal(doc_id, score, ordinal);
634 }
635 }
636
637 collector
638 .into_sorted_results()
639 .into_iter()
640 .map(|(doc_id, score, ordinal)| (doc_id, ordinal, score))
641 .collect()
642 } else {
643 return Ok(Vec::new());
644 };
645 let l1_elapsed = t0.elapsed();
646 log::debug!(
647 "[search_dense] field {}: L1 returned {} candidates in {:.1}ms",
648 field.0,
649 results.len(),
650 l1_elapsed.as_secs_f64() * 1000.0
651 );
652
653 if ann_index.is_some()
656 && !results.is_empty()
657 && let Some(lazy_flat) = lazy_flat
658 {
659 let t_rerank = std::time::Instant::now();
660 let dim = lazy_flat.dim;
661 let quant = lazy_flat.quantization;
662 let vbs = lazy_flat.vector_byte_size();
663
664 let mut resolved: Vec<(usize, usize)> = Vec::new(); for (ri, c) in results.iter().enumerate() {
667 let (start, entries) = lazy_flat.flat_indexes_for_doc(c.0);
668 for (j, &(_, ord)) in entries.iter().enumerate() {
669 if ord == c.1 {
670 resolved.push((ri, start + j));
671 break;
672 }
673 }
674 }
675
676 let t_resolve = t_rerank.elapsed();
677 if !resolved.is_empty() {
678 resolved.sort_unstable_by_key(|&(_, flat_idx)| flat_idx);
680
681 let t_read = std::time::Instant::now();
683 let mut raw_buf = vec![0u8; resolved.len() * vbs];
684 for (buf_idx, &(_, flat_idx)) in resolved.iter().enumerate() {
685 let _ = lazy_flat
686 .read_vector_raw_into(
687 flat_idx,
688 &mut raw_buf[buf_idx * vbs..(buf_idx + 1) * vbs],
689 )
690 .await;
691 }
692
693 let read_elapsed = t_read.elapsed();
694
695 let t_score = std::time::Instant::now();
697 let mut scores = vec![0f32; resolved.len()];
698 Self::score_quantized_batch(query, &raw_buf, quant, dim, &mut scores, unit_norm);
699 let score_elapsed = t_score.elapsed();
700
701 for (buf_idx, &(ri, _)) in resolved.iter().enumerate() {
703 results[ri].2 = scores[buf_idx];
704 }
705
706 log::debug!(
707 "[search_dense] field {}: rerank {} vectors (dim={}, quant={:?}, {}B/vec): resolve={:.1}ms read={:.1}ms score={:.1}ms",
708 field.0,
709 resolved.len(),
710 dim,
711 quant,
712 vbs,
713 t_resolve.as_secs_f64() * 1000.0,
714 read_elapsed.as_secs_f64() * 1000.0,
715 score_elapsed.as_secs_f64() * 1000.0,
716 );
717 }
718
719 results.sort_by(|a, b| b.2.partial_cmp(&a.2).unwrap_or(std::cmp::Ordering::Equal));
720 results.truncate(k * rerank_factor.max(1));
721 log::debug!(
722 "[search_dense] field {}: rerank total={:.1}ms",
723 field.0,
724 t_rerank.elapsed().as_secs_f64() * 1000.0
725 );
726 }
727
728 let mut doc_ordinals: rustc_hash::FxHashMap<DocId, Vec<(u32, f32)>> =
731 rustc_hash::FxHashMap::default();
732 for (doc_id, ordinal, score) in results {
733 let ordinals = doc_ordinals.entry(doc_id as DocId).or_default();
734 ordinals.push((ordinal as u32, score));
735 }
736
737 let mut final_results: Vec<VectorSearchResult> = doc_ordinals
739 .into_iter()
740 .map(|(doc_id, ordinals)| {
741 let combined_score = combiner.combine(&ordinals);
742 VectorSearchResult::new(doc_id, combined_score, ordinals)
743 })
744 .collect();
745
746 final_results.sort_by(|a, b| {
748 b.score
749 .partial_cmp(&a.score)
750 .unwrap_or(std::cmp::Ordering::Equal)
751 });
752 final_results.truncate(k);
753
754 Ok(final_results)
755 }
756
757 pub fn has_dense_vector_index(&self, field: Field) -> bool {
759 self.vector_indexes.contains_key(&field.0) || self.flat_vectors.contains_key(&field.0)
760 }
761
762 pub fn get_dense_vector_index(&self, field: Field) -> Option<Arc<RaBitQIndex>> {
764 match self.vector_indexes.get(&field.0) {
765 Some(VectorIndex::RaBitQ(lazy)) => lazy.get().cloned(),
766 _ => None,
767 }
768 }
769
770 pub fn get_ivf_vector_index(
772 &self,
773 field: Field,
774 ) -> Option<(Arc<IVFRaBitQIndex>, Arc<crate::structures::RaBitQCodebook>)> {
775 match self.vector_indexes.get(&field.0) {
776 Some(VectorIndex::IVF(lazy)) => lazy.get().map(|(i, c)| (i.clone(), c.clone())),
777 _ => None,
778 }
779 }
780
781 pub fn coarse_centroids(&self, field_id: u32) -> Option<&Arc<CoarseCentroids>> {
783 self.coarse_centroids.get(&field_id)
784 }
785
786 pub fn set_coarse_centroids(&mut self, centroids: FxHashMap<u32, Arc<CoarseCentroids>>) {
788 self.coarse_centroids = centroids;
789 }
790
791 pub fn get_scann_vector_index(
793 &self,
794 field: Field,
795 ) -> Option<(Arc<IVFPQIndex>, Arc<PQCodebook>)> {
796 match self.vector_indexes.get(&field.0) {
797 Some(VectorIndex::ScaNN(lazy)) => lazy.get().map(|(i, c)| (i.clone(), c.clone())),
798 _ => None,
799 }
800 }
801
802 pub fn get_vector_index(&self, field: Field) -> Option<&VectorIndex> {
804 self.vector_indexes.get(&field.0)
805 }
806
807 pub async fn search_sparse_vector(
817 &self,
818 field: Field,
819 vector: &[(u32, f32)],
820 limit: usize,
821 combiner: crate::query::MultiValueCombiner,
822 heap_factor: f32,
823 ) -> Result<Vec<VectorSearchResult>> {
824 use crate::query::BmpExecutor;
825
826 let query_tokens = vector.len();
827
828 let sparse_index = match self.sparse_indexes.get(&field.0) {
830 Some(idx) => idx,
831 None => {
832 log::debug!(
833 "Sparse vector search: no index for field {}, returning empty",
834 field.0
835 );
836 return Ok(Vec::new());
837 }
838 };
839
840 let index_dimensions = sparse_index.num_dimensions();
841
842 let mut matched_terms: Vec<(u32, f32)> = Vec::with_capacity(vector.len());
844 let mut missing_count = 0usize;
845
846 for &(dim_id, query_weight) in vector {
847 if sparse_index.has_dimension(dim_id) {
848 matched_terms.push((dim_id, query_weight));
849 } else {
850 missing_count += 1;
851 }
852 }
853
854 log::debug!(
855 "Sparse vector search: query_tokens={}, matched={}, missing={}, index_dimensions={}",
856 query_tokens,
857 matched_terms.len(),
858 missing_count,
859 index_dimensions
860 );
861
862 if matched_terms.is_empty() {
863 log::debug!("Sparse vector search: no matching tokens, returning empty");
864 return Ok(Vec::new());
865 }
866
867 let num_terms = matched_terms.len();
871 let over_fetch = limit * 2; let raw_results = if num_terms > 12 {
873 BmpExecutor::new(sparse_index, matched_terms, over_fetch, heap_factor)
875 .execute()
876 .await?
877 } else {
878 crate::query::LazyBlockMaxScoreExecutor::new(
881 sparse_index,
882 matched_terms,
883 over_fetch,
884 heap_factor,
885 )
886 .execute()
887 .await?
888 };
889
890 log::trace!(
891 "Sparse WAND returned {} raw results for segment (doc_id_offset={})",
892 raw_results.len(),
893 self.doc_id_offset
894 );
895 if log::log_enabled!(log::Level::Trace) && !raw_results.is_empty() {
896 for r in raw_results.iter().take(5) {
897 log::trace!(
898 " Raw result: doc_id={} (global={}), score={:.4}, ordinal={}",
899 r.doc_id,
900 r.doc_id + self.doc_id_offset,
901 r.score,
902 r.ordinal
903 );
904 }
905 }
906
907 let mut doc_ordinals: rustc_hash::FxHashMap<u32, Vec<(u32, f32)>> =
910 rustc_hash::FxHashMap::default();
911 for r in raw_results {
912 let ordinals = doc_ordinals.entry(r.doc_id).or_default();
913 ordinals.push((r.ordinal as u32, r.score));
914 }
915
916 let mut results: Vec<VectorSearchResult> = doc_ordinals
919 .into_iter()
920 .map(|(doc_id, ordinals)| {
921 let combined_score = combiner.combine(&ordinals);
922 VectorSearchResult::new(doc_id, combined_score, ordinals)
923 })
924 .collect();
925
926 results.sort_by(|a, b| {
928 b.score
929 .partial_cmp(&a.score)
930 .unwrap_or(std::cmp::Ordering::Equal)
931 });
932 results.truncate(limit);
933
934 Ok(results)
935 }
936
937 pub async fn get_positions(
942 &self,
943 field: Field,
944 term: &[u8],
945 ) -> Result<Option<crate::structures::PositionPostingList>> {
946 let handle = match &self.positions_handle {
948 Some(h) => h,
949 None => return Ok(None),
950 };
951
952 let mut key = Vec::with_capacity(4 + term.len());
954 key.extend_from_slice(&field.0.to_le_bytes());
955 key.extend_from_slice(term);
956
957 let term_info = match self.term_dict.get(&key).await? {
959 Some(info) => info,
960 None => return Ok(None),
961 };
962
963 let (offset, length) = match term_info.position_info() {
965 Some((o, l)) => (o, l),
966 None => return Ok(None),
967 };
968
969 let slice = handle.slice(offset..offset + length as u64);
971 let data = slice.read_bytes().await?;
972
973 let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
975
976 Ok(Some(pos_list))
977 }
978
979 pub fn has_positions(&self, field: Field) -> bool {
981 if let Some(entry) = self.schema.get_field_entry(field) {
983 entry.positions.is_some()
984 } else {
985 false
986 }
987 }
988}
989
990pub type SegmentReader = AsyncSegmentReader;