Skip to main content

hermes_core/segment/reader/
mod.rs

1//! Async segment reader with lazy loading
2
3mod loader;
4mod types;
5
6pub use types::{SparseIndex, VectorIndex, VectorSearchResult};
7
8/// Memory statistics for a single segment
9#[derive(Debug, Clone, Default)]
10pub struct SegmentMemoryStats {
11    /// Segment ID
12    pub segment_id: u128,
13    /// Number of documents in segment
14    pub num_docs: u32,
15    /// Term dictionary block cache bytes
16    pub term_dict_cache_bytes: usize,
17    /// Document store block cache bytes
18    pub store_cache_bytes: usize,
19    /// Sparse vector index bytes (in-memory posting lists)
20    pub sparse_index_bytes: usize,
21    /// Dense vector index bytes (cluster assignments, quantized codes)
22    pub dense_index_bytes: usize,
23    /// Bloom filter bytes
24    pub bloom_filter_bytes: usize,
25}
26
27impl SegmentMemoryStats {
28    /// Total estimated memory for this segment
29    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 crate::directories::{AsyncFileRead, Directory, LazyFileHandle, LazyFileSlice};
45use crate::dsl::{Document, Field, Schema};
46use crate::structures::{
47    AsyncSSTableReader, BlockPostingList, CoarseCentroids, IVFPQIndex, IVFRaBitQIndex, PQCodebook,
48    RaBitQIndex, SSTableStats, TermInfo,
49};
50use crate::{DocId, Error, Result};
51
52use super::store::{AsyncStoreReader, RawStoreBlock};
53use super::types::{SegmentFiles, SegmentId, SegmentMeta};
54
55/// Async segment reader with lazy loading
56///
57/// - Term dictionary: only index loaded, blocks loaded on-demand
58/// - Postings: loaded on-demand per term via HTTP range requests
59/// - Document store: only index loaded, blocks loaded on-demand via HTTP range requests
60pub struct AsyncSegmentReader {
61    meta: SegmentMeta,
62    /// Term dictionary with lazy block loading
63    term_dict: Arc<AsyncSSTableReader<TermInfo>>,
64    /// Postings file handle - fetches ranges on demand
65    postings_handle: LazyFileHandle,
66    /// Document store with lazy block loading
67    store: Arc<AsyncStoreReader>,
68    schema: Arc<Schema>,
69    /// Base doc_id offset for this segment
70    doc_id_offset: DocId,
71    /// Dense vector indexes per field (RaBitQ or IVF-RaBitQ)
72    vector_indexes: FxHashMap<u32, VectorIndex>,
73    /// Shared coarse centroids for IVF search (loaded once)
74    coarse_centroids: Option<Arc<CoarseCentroids>>,
75    /// Sparse vector indexes per field
76    sparse_indexes: FxHashMap<u32, SparseIndex>,
77    /// Position file handle for phrase queries (lazy loading)
78    positions_handle: Option<LazyFileHandle>,
79}
80
81impl AsyncSegmentReader {
82    /// Open a segment with lazy loading
83    pub async fn open<D: Directory>(
84        dir: &D,
85        segment_id: SegmentId,
86        schema: Arc<Schema>,
87        doc_id_offset: DocId,
88        cache_blocks: usize,
89    ) -> Result<Self> {
90        let files = SegmentFiles::new(segment_id.0);
91
92        // Read metadata (small, always loaded)
93        let meta_slice = dir.open_read(&files.meta).await?;
94        let meta_bytes = meta_slice.read_bytes().await?;
95        let meta = SegmentMeta::deserialize(meta_bytes.as_slice())?;
96        debug_assert_eq!(meta.id, segment_id.0);
97
98        // Open term dictionary with lazy loading (fetches ranges on demand)
99        let term_dict_handle = dir.open_lazy(&files.term_dict).await?;
100        let term_dict = AsyncSSTableReader::open(term_dict_handle, cache_blocks).await?;
101
102        // Get postings file handle (lazy - fetches ranges on demand)
103        let postings_handle = dir.open_lazy(&files.postings).await?;
104
105        // Open store with lazy loading
106        let store_handle = dir.open_lazy(&files.store).await?;
107        let store = AsyncStoreReader::open(store_handle, cache_blocks).await?;
108
109        // Load dense vector indexes from unified .vectors file
110        let (vector_indexes, coarse_centroids) =
111            loader::load_vectors_file(dir, &files, &schema).await?;
112
113        // Load sparse vector indexes from .sparse file
114        let sparse_indexes = loader::load_sparse_file(dir, &files, meta.num_docs, &schema).await?;
115
116        // Open positions file handle (if exists) - offsets are now in TermInfo
117        let positions_handle = loader::open_positions_file(dir, &files, &schema).await?;
118
119        // Log segment loading stats (compact format: ~24 bytes per active dim in hashmap)
120        let sparse_dims: usize = sparse_indexes.values().map(|s| s.num_dimensions()).sum();
121        let sparse_mem = sparse_dims * 24; // HashMap entry overhead
122        log::debug!(
123            "[segment] loaded {:016x}: docs={}, sparse_dims={}, sparse_mem={:.2} KB, vectors={}",
124            segment_id.0,
125            meta.num_docs,
126            sparse_dims,
127            sparse_mem as f64 / 1024.0,
128            vector_indexes.len()
129        );
130
131        Ok(Self {
132            meta,
133            term_dict: Arc::new(term_dict),
134            postings_handle,
135            store: Arc::new(store),
136            schema,
137            doc_id_offset,
138            vector_indexes,
139            coarse_centroids,
140            sparse_indexes,
141            positions_handle,
142        })
143    }
144
145    pub fn meta(&self) -> &SegmentMeta {
146        &self.meta
147    }
148
149    pub fn num_docs(&self) -> u32 {
150        self.meta.num_docs
151    }
152
153    /// Get average field length for BM25F scoring
154    pub fn avg_field_len(&self, field: Field) -> f32 {
155        self.meta.avg_field_len(field)
156    }
157
158    pub fn doc_id_offset(&self) -> DocId {
159        self.doc_id_offset
160    }
161
162    /// Set the doc_id_offset (used for parallel segment loading)
163    pub fn set_doc_id_offset(&mut self, offset: DocId) {
164        self.doc_id_offset = offset;
165    }
166
167    pub fn schema(&self) -> &Schema {
168        &self.schema
169    }
170
171    /// Get sparse indexes for all fields
172    pub fn sparse_indexes(&self) -> &FxHashMap<u32, SparseIndex> {
173        &self.sparse_indexes
174    }
175
176    /// Get vector indexes for all fields
177    pub fn vector_indexes(&self) -> &FxHashMap<u32, VectorIndex> {
178        &self.vector_indexes
179    }
180
181    /// Get term dictionary stats for debugging
182    pub fn term_dict_stats(&self) -> SSTableStats {
183        self.term_dict.stats()
184    }
185
186    /// Estimate memory usage of this segment reader
187    pub fn memory_stats(&self) -> SegmentMemoryStats {
188        let term_dict_stats = self.term_dict.stats();
189
190        // Term dict cache: num_blocks * avg_block_size (estimate 4KB per cached block)
191        let term_dict_cache_bytes = self.term_dict.cached_blocks() * 4096;
192
193        // Store cache: similar estimate
194        let store_cache_bytes = self.store.cached_blocks() * 4096;
195
196        // Sparse index: each dimension has a posting list in memory
197        // Estimate: ~24 bytes per active dimension (HashMap entry overhead)
198        let sparse_index_bytes: usize = self
199            .sparse_indexes
200            .values()
201            .map(|s| s.num_dimensions() * 24)
202            .sum();
203
204        // Dense index: vectors are memory-mapped, but we track index structures
205        // RaBitQ/IVF indexes have cluster assignments in memory
206        let dense_index_bytes: usize = self
207            .vector_indexes
208            .values()
209            .map(|v| v.estimated_memory_bytes())
210            .sum();
211
212        SegmentMemoryStats {
213            segment_id: self.meta.id,
214            num_docs: self.meta.num_docs,
215            term_dict_cache_bytes,
216            store_cache_bytes,
217            sparse_index_bytes,
218            dense_index_bytes,
219            bloom_filter_bytes: term_dict_stats.bloom_filter_size,
220        }
221    }
222
223    /// Get posting list for a term (async - loads on demand)
224    ///
225    /// For small posting lists (1-3 docs), the data is inlined in the term dictionary
226    /// and no additional I/O is needed. For larger lists, reads from .post file.
227    pub async fn get_postings(
228        &self,
229        field: Field,
230        term: &[u8],
231    ) -> Result<Option<BlockPostingList>> {
232        log::debug!(
233            "SegmentReader::get_postings field={} term_len={}",
234            field.0,
235            term.len()
236        );
237
238        // Build key: field_id + term
239        let mut key = Vec::with_capacity(4 + term.len());
240        key.extend_from_slice(&field.0.to_le_bytes());
241        key.extend_from_slice(term);
242
243        // Look up in term dictionary
244        let term_info = match self.term_dict.get(&key).await? {
245            Some(info) => {
246                log::debug!("SegmentReader::get_postings found term_info");
247                info
248            }
249            None => {
250                log::debug!("SegmentReader::get_postings term not found");
251                return Ok(None);
252            }
253        };
254
255        // Check if posting list is inlined
256        if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
257            // Build BlockPostingList from inline data (no I/O needed!)
258            let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
259            for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs.into_iter()) {
260                posting_list.push(doc_id, tf);
261            }
262            let block_list = BlockPostingList::from_posting_list(&posting_list)?;
263            return Ok(Some(block_list));
264        }
265
266        // External posting list - read from postings file handle (lazy - HTTP range request)
267        let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
268            Error::Corruption("TermInfo has neither inline nor external data".to_string())
269        })?;
270
271        let start = posting_offset;
272        let end = start + posting_len as u64;
273
274        if end > self.postings_handle.len() {
275            return Err(Error::Corruption(
276                "Posting offset out of bounds".to_string(),
277            ));
278        }
279
280        let posting_bytes = self.postings_handle.read_bytes_range(start..end).await?;
281        let block_list = BlockPostingList::deserialize(&mut posting_bytes.as_slice())?;
282
283        Ok(Some(block_list))
284    }
285
286    /// Get document by local doc_id (async - loads on demand)
287    pub async fn doc(&self, local_doc_id: DocId) -> Result<Option<Document>> {
288        self.store
289            .get(local_doc_id, &self.schema)
290            .await
291            .map_err(Error::from)
292    }
293
294    /// Prefetch term dictionary blocks for a key range
295    pub async fn prefetch_terms(
296        &self,
297        field: Field,
298        start_term: &[u8],
299        end_term: &[u8],
300    ) -> Result<()> {
301        let mut start_key = Vec::with_capacity(4 + start_term.len());
302        start_key.extend_from_slice(&field.0.to_le_bytes());
303        start_key.extend_from_slice(start_term);
304
305        let mut end_key = Vec::with_capacity(4 + end_term.len());
306        end_key.extend_from_slice(&field.0.to_le_bytes());
307        end_key.extend_from_slice(end_term);
308
309        self.term_dict.prefetch_range(&start_key, &end_key).await?;
310        Ok(())
311    }
312
313    /// Check if store uses dictionary compression (incompatible with raw merging)
314    pub fn store_has_dict(&self) -> bool {
315        self.store.has_dict()
316    }
317
318    /// Get store reference for merge operations
319    pub fn store(&self) -> &super::store::AsyncStoreReader {
320        &self.store
321    }
322
323    /// Get raw store blocks for optimized merging
324    pub fn store_raw_blocks(&self) -> Vec<RawStoreBlock> {
325        self.store.raw_blocks()
326    }
327
328    /// Get store data slice for raw block access
329    pub fn store_data_slice(&self) -> &LazyFileSlice {
330        self.store.data_slice()
331    }
332
333    /// Get all terms from this segment (for merge)
334    pub async fn all_terms(&self) -> Result<Vec<(Vec<u8>, TermInfo)>> {
335        self.term_dict.all_entries().await.map_err(Error::from)
336    }
337
338    /// Get all terms with parsed field and term string (for statistics aggregation)
339    ///
340    /// Returns (field, term_string, doc_freq) for each term in the dictionary.
341    /// Skips terms that aren't valid UTF-8.
342    pub async fn all_terms_with_stats(&self) -> Result<Vec<(Field, String, u32)>> {
343        let entries = self.term_dict.all_entries().await?;
344        let mut result = Vec::with_capacity(entries.len());
345
346        for (key, term_info) in entries {
347            // Key format: field_id (4 bytes little-endian) + term bytes
348            if key.len() > 4 {
349                let field_id = u32::from_le_bytes([key[0], key[1], key[2], key[3]]);
350                let term_bytes = &key[4..];
351                if let Ok(term_str) = std::str::from_utf8(term_bytes) {
352                    result.push((Field(field_id), term_str.to_string(), term_info.doc_freq()));
353                }
354            }
355        }
356
357        Ok(result)
358    }
359
360    /// Get streaming iterator over term dictionary (for memory-efficient merge)
361    pub fn term_dict_iter(&self) -> crate::structures::AsyncSSTableIterator<'_, TermInfo> {
362        self.term_dict.iter()
363    }
364
365    /// Prefetch all term dictionary blocks in a single bulk I/O call.
366    ///
367    /// Call before merge iteration to eliminate per-block cache misses.
368    pub async fn prefetch_term_dict(&self) -> crate::Result<()> {
369        self.term_dict
370            .prefetch_all_data_bulk()
371            .await
372            .map_err(crate::Error::from)
373    }
374
375    /// Read raw posting bytes at offset
376    pub async fn read_postings(&self, offset: u64, len: u32) -> Result<Vec<u8>> {
377        let start = offset;
378        let end = start + len as u64;
379        let bytes = self.postings_handle.read_bytes_range(start..end).await?;
380        Ok(bytes.to_vec())
381    }
382
383    /// Read raw position bytes at offset (for merge)
384    pub async fn read_position_bytes(&self, offset: u64, len: u32) -> Result<Option<Vec<u8>>> {
385        let handle = match &self.positions_handle {
386            Some(h) => h,
387            None => return Ok(None),
388        };
389        let start = offset;
390        let end = start + len as u64;
391        let bytes = handle.read_bytes_range(start..end).await?;
392        Ok(Some(bytes.to_vec()))
393    }
394
395    /// Check if this segment has a positions file
396    pub fn has_positions_file(&self) -> bool {
397        self.positions_handle.is_some()
398    }
399
400    /// Search dense vectors using RaBitQ
401    ///
402    /// Returns VectorSearchResult with ordinal tracking for multi-value fields.
403    /// The doc_ids are adjusted by doc_id_offset for this segment.
404    /// If mrl_dim is configured, the query vector is automatically trimmed.
405    /// For multi-valued documents, scores are combined using the specified combiner.
406    pub fn search_dense_vector(
407        &self,
408        field: Field,
409        query: &[f32],
410        k: usize,
411        rerank_factor: usize,
412        combiner: crate::query::MultiValueCombiner,
413    ) -> Result<Vec<VectorSearchResult>> {
414        let index = self
415            .vector_indexes
416            .get(&field.0)
417            .ok_or_else(|| Error::Schema(format!("No dense vector index for field {}", field.0)))?;
418
419        // Get mrl_dim from config to trim query vector if needed
420        let mrl_dim = self
421            .schema
422            .get_field_entry(field)
423            .and_then(|e| e.dense_vector_config.as_ref())
424            .and_then(|c| c.mrl_dim);
425
426        // Trim query vector if mrl_dim is set
427        let query_vec: Vec<f32>;
428        let effective_query = if let Some(trim_dim) = mrl_dim {
429            if trim_dim < query.len() {
430                query_vec = query[..trim_dim].to_vec();
431                query_vec.as_slice()
432            } else {
433                query
434            }
435        } else {
436            query
437        };
438
439        // Results include (doc_id, ordinal, distance)
440        let results: Vec<(u32, u16, f32)> = match index {
441            VectorIndex::Flat(flat_data) => {
442                // Brute-force search over raw vectors using SIMD-accelerated distance
443                use crate::structures::simd::squared_euclidean_distance;
444
445                let mut candidates: Vec<(u32, u16, f32)> = flat_data
446                    .vectors
447                    .iter()
448                    .zip(flat_data.doc_ids.iter())
449                    .map(|(vec, &(doc_id, ordinal))| {
450                        let dist = squared_euclidean_distance(effective_query, vec);
451                        (doc_id, ordinal, dist)
452                    })
453                    .collect();
454                candidates
455                    .sort_by(|a, b| a.2.partial_cmp(&b.2).unwrap_or(std::cmp::Ordering::Equal));
456                candidates.truncate(k);
457                candidates
458            }
459            VectorIndex::RaBitQ(rabitq) => rabitq.search(effective_query, k, rerank_factor),
460            VectorIndex::IVF { index, codebook } => {
461                let centroids = self.coarse_centroids.as_ref().ok_or_else(|| {
462                    Error::Schema("IVF index requires coarse centroids".to_string())
463                })?;
464                let nprobe = rerank_factor.max(32); // Use rerank_factor as nprobe hint
465                index
466                    .search(centroids, codebook, effective_query, k, Some(nprobe))
467                    .into_iter()
468                    .map(|(doc_id, dist)| (doc_id, 0u16, dist)) // IVF doesn't track ordinals yet
469                    .collect()
470            }
471            VectorIndex::ScaNN { index, codebook } => {
472                let centroids = self.coarse_centroids.as_ref().ok_or_else(|| {
473                    Error::Schema("ScaNN index requires coarse centroids".to_string())
474                })?;
475                let nprobe = rerank_factor.max(32);
476                index
477                    .search(centroids, codebook, effective_query, k, Some(nprobe))
478                    .into_iter()
479                    .map(|(doc_id, dist)| (doc_id, 0u16, dist)) // ScaNN doesn't track ordinals yet
480                    .collect()
481            }
482        };
483
484        // Convert distance to score (smaller distance = higher score)
485        // Note: doc_id_offset is NOT applied here - the collector applies it uniformly
486        // Track ordinals with individual scores for each doc_id
487        let mut doc_ordinals: rustc_hash::FxHashMap<DocId, Vec<(u32, f32)>> =
488            rustc_hash::FxHashMap::default();
489        for (doc_id, ordinal, dist) in results {
490            let score = 1.0 / (1.0 + dist); // Convert distance to similarity score
491            let ordinals = doc_ordinals.entry(doc_id as DocId).or_default();
492            ordinals.push((ordinal as u32, score));
493        }
494
495        // Combine scores and build results with ordinal tracking
496        let mut final_results: Vec<VectorSearchResult> = doc_ordinals
497            .into_iter()
498            .map(|(doc_id, ordinals)| {
499                let combined_score = combiner.combine(&ordinals);
500                VectorSearchResult::new(doc_id, combined_score, ordinals)
501            })
502            .collect();
503
504        // Sort by score descending and take top k
505        final_results.sort_by(|a, b| {
506            b.score
507                .partial_cmp(&a.score)
508                .unwrap_or(std::cmp::Ordering::Equal)
509        });
510        final_results.truncate(k);
511
512        Ok(final_results)
513    }
514
515    /// Check if this segment has a dense vector index for the given field
516    pub fn has_dense_vector_index(&self, field: Field) -> bool {
517        self.vector_indexes.contains_key(&field.0)
518    }
519
520    /// Get the dense vector index for a field (if available)
521    pub fn get_dense_vector_index(&self, field: Field) -> Option<Arc<RaBitQIndex>> {
522        match self.vector_indexes.get(&field.0) {
523            Some(VectorIndex::RaBitQ(idx)) => Some(idx.clone()),
524            _ => None,
525        }
526    }
527
528    /// Get the IVF vector index for a field (if available)
529    pub fn get_ivf_vector_index(&self, field: Field) -> Option<Arc<IVFRaBitQIndex>> {
530        match self.vector_indexes.get(&field.0) {
531            Some(VectorIndex::IVF { index, .. }) => Some(index.clone()),
532            _ => None,
533        }
534    }
535
536    /// Get the ScaNN vector index for a field (if available)
537    pub fn get_scann_vector_index(
538        &self,
539        field: Field,
540    ) -> Option<(Arc<IVFPQIndex>, Arc<PQCodebook>)> {
541        match self.vector_indexes.get(&field.0) {
542            Some(VectorIndex::ScaNN { index, codebook }) => Some((index.clone(), codebook.clone())),
543            _ => None,
544        }
545    }
546
547    /// Get the vector index type for a field
548    pub fn get_vector_index(&self, field: Field) -> Option<&VectorIndex> {
549        self.vector_indexes.get(&field.0)
550    }
551
552    /// Search for similar sparse vectors using dedicated sparse posting lists
553    ///
554    /// Uses shared `WandExecutor` with `SparseTermScorer` for efficient top-k retrieval.
555    /// Optimizations (via WandExecutor):
556    /// 1. **MaxScore pruning**: Dimensions sorted by max contribution
557    /// 2. **Block-Max WAND**: Skips blocks where max contribution < threshold
558    /// 3. **Top-K heap**: Efficient score collection
559    ///
560    /// Returns VectorSearchResult with ordinal tracking for multi-value fields.
561    pub async fn search_sparse_vector(
562        &self,
563        field: Field,
564        vector: &[(u32, f32)],
565        limit: usize,
566        combiner: crate::query::MultiValueCombiner,
567        heap_factor: f32,
568    ) -> Result<Vec<VectorSearchResult>> {
569        use crate::query::{BlockMaxScoreExecutor, BmpExecutor, SparseTermScorer};
570
571        let query_tokens = vector.len();
572
573        // Get sparse index for this field
574        let sparse_index = match self.sparse_indexes.get(&field.0) {
575            Some(idx) => idx,
576            None => {
577                log::debug!(
578                    "Sparse vector search: no index for field {}, returning empty",
579                    field.0
580                );
581                return Ok(Vec::new());
582            }
583        };
584
585        let index_dimensions = sparse_index.num_dimensions();
586
587        // Build scorers for each dimension that exists in the index
588        // Load posting lists on-demand (lazy loading via mmap)
589        // Keep Arc references alive for the duration of scoring
590        let mut matched_tokens = Vec::new();
591        let mut missing_tokens = Vec::new();
592        let mut posting_lists: Vec<(u32, f32, Arc<BlockSparsePostingList>)> =
593            Vec::with_capacity(vector.len());
594
595        for &(dim_id, query_weight) in vector {
596            // Check if dimension exists before loading
597            if !sparse_index.has_dimension(dim_id) {
598                missing_tokens.push(dim_id);
599                continue;
600            }
601
602            // Load posting list on-demand (async, uses mmap)
603            match sparse_index.get_posting(dim_id).await? {
604                Some(pl) => {
605                    matched_tokens.push(dim_id);
606                    posting_lists.push((dim_id, query_weight, pl));
607                }
608                None => {
609                    missing_tokens.push(dim_id);
610                }
611            }
612        }
613
614        // Create scorers from the loaded posting lists (borrows from posting_lists)
615        let scorers: Vec<SparseTermScorer> = posting_lists
616            .iter()
617            .map(|(_, query_weight, pl)| SparseTermScorer::from_arc(pl, *query_weight))
618            .collect();
619
620        log::debug!(
621            "Sparse vector search: query_tokens={}, matched={}, missing={}, index_dimensions={}",
622            query_tokens,
623            matched_tokens.len(),
624            missing_tokens.len(),
625            index_dimensions
626        );
627
628        // Log query tokens with their IDs and weights
629        if log::log_enabled!(log::Level::Debug) {
630            let query_details: Vec<_> = vector
631                .iter()
632                .take(30)
633                .map(|(id, w)| format!("{}:{:.3}", id, w))
634                .collect();
635            log::debug!("Query tokens (id:weight): [{}]", query_details.join(", "));
636        }
637
638        if !matched_tokens.is_empty() {
639            log::debug!(
640                "Matched token IDs: {:?}",
641                matched_tokens.iter().take(20).collect::<Vec<_>>()
642            );
643        }
644
645        if !missing_tokens.is_empty() {
646            log::debug!(
647                "Missing token IDs (not in index): {:?}",
648                missing_tokens.iter().take(20).collect::<Vec<_>>()
649            );
650        }
651
652        if scorers.is_empty() {
653            log::debug!("Sparse vector search: no matching tokens, returning empty");
654            return Ok(Vec::new());
655        }
656
657        // Select executor based on number of query terms:
658        // - 12+ terms: BMP (block-at-a-time, best for SPLADE expansions)
659        // - 1-11 terms: BlockMaxScoreExecutor (unified MaxScore + block-max + conjunction)
660        let num_terms = scorers.len();
661        let over_fetch = limit * 2; // Over-fetch for multi-value combining
662        let raw_results = if num_terms > 12 {
663            // BMP: use posting lists directly (not iterators)
664            let pl_refs: Vec<_> = posting_lists
665                .iter()
666                .map(|(_, _, pl)| Arc::clone(pl))
667                .collect();
668            let weights: Vec<_> = posting_lists.iter().map(|(_, qw, _)| *qw).collect();
669            drop(scorers); // Release borrowing iterators before using posting_lists
670            BmpExecutor::new(pl_refs, weights, over_fetch, heap_factor).execute()
671        } else {
672            BlockMaxScoreExecutor::with_heap_factor(scorers, over_fetch, heap_factor).execute()
673        };
674
675        log::trace!(
676            "Sparse WAND returned {} raw results for segment (doc_id_offset={})",
677            raw_results.len(),
678            self.doc_id_offset
679        );
680        if log::log_enabled!(log::Level::Trace) && !raw_results.is_empty() {
681            for r in raw_results.iter().take(5) {
682                log::trace!(
683                    "  Raw result: doc_id={} (global={}), score={:.4}, ordinal={}",
684                    r.doc_id,
685                    r.doc_id + self.doc_id_offset,
686                    r.score,
687                    r.ordinal
688                );
689            }
690        }
691
692        // Track ordinals with individual scores for each doc_id
693        // Now using real ordinals from the posting lists
694        let mut doc_ordinals: rustc_hash::FxHashMap<u32, Vec<(u32, f32)>> =
695            rustc_hash::FxHashMap::default();
696        for r in raw_results {
697            let ordinals = doc_ordinals.entry(r.doc_id).or_default();
698            ordinals.push((r.ordinal as u32, r.score));
699        }
700
701        // Combine scores and build results with ordinal tracking
702        // Note: doc_id_offset is NOT applied here - the collector applies it uniformly
703        let mut results: Vec<VectorSearchResult> = doc_ordinals
704            .into_iter()
705            .map(|(doc_id, ordinals)| {
706                let combined_score = combiner.combine(&ordinals);
707                VectorSearchResult::new(doc_id, combined_score, ordinals)
708            })
709            .collect();
710
711        // Sort by score descending and take top limit
712        results.sort_by(|a, b| {
713            b.score
714                .partial_cmp(&a.score)
715                .unwrap_or(std::cmp::Ordering::Equal)
716        });
717        results.truncate(limit);
718
719        Ok(results)
720    }
721
722    /// Get positions for a term (for phrase queries)
723    ///
724    /// Position offsets are now embedded in TermInfo, so we first look up
725    /// the term to get its TermInfo, then use position_info() to get the offset.
726    pub async fn get_positions(
727        &self,
728        field: Field,
729        term: &[u8],
730    ) -> Result<Option<crate::structures::PositionPostingList>> {
731        use std::io::Cursor;
732
733        // Get positions handle
734        let handle = match &self.positions_handle {
735            Some(h) => h,
736            None => return Ok(None),
737        };
738
739        // Build key: field_id + term
740        let mut key = Vec::with_capacity(4 + term.len());
741        key.extend_from_slice(&field.0.to_le_bytes());
742        key.extend_from_slice(term);
743
744        // Look up term in dictionary to get TermInfo with position offset
745        let term_info = match self.term_dict.get(&key).await? {
746            Some(info) => info,
747            None => return Ok(None),
748        };
749
750        // Get position offset from TermInfo
751        let (offset, length) = match term_info.position_info() {
752            Some((o, l)) => (o, l),
753            None => return Ok(None),
754        };
755
756        // Read the position data
757        let slice = handle.slice(offset..offset + length as u64);
758        let data = slice.read_bytes().await?;
759
760        // Deserialize
761        let mut cursor = Cursor::new(data.as_slice());
762        let pos_list = crate::structures::PositionPostingList::deserialize(&mut cursor)?;
763
764        Ok(Some(pos_list))
765    }
766
767    /// Check if positions are available for a field
768    pub fn has_positions(&self, field: Field) -> bool {
769        // Check schema for position mode on this field
770        if let Some(entry) = self.schema.get_field_entry(field) {
771            entry.positions.is_some()
772        } else {
773            false
774        }
775    }
776}
777
778/// Alias for AsyncSegmentReader
779pub type SegmentReader = AsyncSegmentReader;