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 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
56/// Async segment reader with lazy loading
57///
58/// - Term dictionary: only index loaded, blocks loaded on-demand
59/// - Postings: loaded on-demand per term via HTTP range requests
60/// - Document store: only index loaded, blocks loaded on-demand via HTTP range requests
61pub struct AsyncSegmentReader {
62    meta: SegmentMeta,
63    /// Term dictionary with lazy block loading
64    term_dict: Arc<AsyncSSTableReader<TermInfo>>,
65    /// Postings file handle - fetches ranges on demand
66    postings_handle: LazyFileHandle,
67    /// Document store with lazy block loading
68    store: Arc<AsyncStoreReader>,
69    schema: Arc<Schema>,
70    /// Base doc_id offset for this segment
71    doc_id_offset: DocId,
72    /// Dense vector indexes per field (RaBitQ or IVF-RaBitQ) — for search
73    vector_indexes: FxHashMap<u32, VectorIndex>,
74    /// Lazy flat vectors per field — for reranking and merge (doc_ids in memory, vectors via mmap)
75    flat_vectors: FxHashMap<u32, LazyFlatVectorData>,
76    /// Per-field coarse centroids for IVF/ScaNN search
77    coarse_centroids: FxHashMap<u32, Arc<CoarseCentroids>>,
78    /// Sparse vector indexes per field
79    sparse_indexes: FxHashMap<u32, SparseIndex>,
80    /// Position file handle for phrase queries (lazy loading)
81    positions_handle: Option<LazyFileHandle>,
82}
83
84impl AsyncSegmentReader {
85    /// Open a segment with lazy loading
86    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        // Read metadata (small, always loaded)
96        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        // Open term dictionary with lazy loading (fetches ranges on demand)
102        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        // Get postings file handle (lazy - fetches ranges on demand)
106        let postings_handle = dir.open_lazy(&files.postings).await?;
107
108        // Open store with lazy loading
109        let store_handle = dir.open_lazy(&files.store).await?;
110        let store = AsyncStoreReader::open(store_handle, cache_blocks).await?;
111
112        // Load dense vector indexes from unified .vectors file
113        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        // Load sparse vector indexes from .sparse file
118        let sparse_indexes = loader::load_sparse_file(dir, &files, meta.num_docs, &schema).await?;
119
120        // Open positions file handle (if exists) - offsets are now in TermInfo
121        let positions_handle = loader::open_positions_file(dir, &files, &schema).await?;
122
123        // Log segment loading stats
124        {
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    /// Get average field length for BM25F scoring
171    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    /// Set the doc_id_offset (used for parallel segment loading)
180    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    /// Get sparse indexes for all fields
189    pub fn sparse_indexes(&self) -> &FxHashMap<u32, SparseIndex> {
190        &self.sparse_indexes
191    }
192
193    /// Get vector indexes for all fields
194    pub fn vector_indexes(&self) -> &FxHashMap<u32, VectorIndex> {
195        &self.vector_indexes
196    }
197
198    /// Get lazy flat vectors for all fields (for reranking and merge)
199    pub fn flat_vectors(&self) -> &FxHashMap<u32, LazyFlatVectorData> {
200        &self.flat_vectors
201    }
202
203    /// Get term dictionary stats for debugging
204    pub fn term_dict_stats(&self) -> SSTableStats {
205        self.term_dict.stats()
206    }
207
208    /// Estimate memory usage of this segment reader
209    pub fn memory_stats(&self) -> SegmentMemoryStats {
210        let term_dict_stats = self.term_dict.stats();
211
212        // Term dict cache: num_blocks * avg_block_size (estimate 4KB per cached block)
213        let term_dict_cache_bytes = self.term_dict.cached_blocks() * 4096;
214
215        // Store cache: similar estimate
216        let store_cache_bytes = self.store.cached_blocks() * 4096;
217
218        // Sparse index: SoA dim table + OwnedBytes skip section
219        let sparse_index_bytes: usize = self
220            .sparse_indexes
221            .values()
222            .map(|s| s.estimated_memory_bytes())
223            .sum();
224
225        // Dense index: vectors are memory-mapped, but we track index structures
226        // RaBitQ/IVF indexes have cluster assignments in memory
227        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    /// Get posting list for a term (async - loads on demand)
245    ///
246    /// For small posting lists (1-3 docs), the data is inlined in the term dictionary
247    /// and no additional I/O is needed. For larger lists, reads from .post file.
248    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        // Build key: field_id + term
260        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        // Look up in term dictionary
265        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        // Check if posting list is inlined
277        if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
278            // Build BlockPostingList from inline data (no I/O needed!)
279            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        // External posting list - read from postings file handle (lazy - HTTP range request)
288        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    /// Get document by local doc_id (async - loads on demand).
308    ///
309    /// Dense vector fields are hydrated from LazyFlatVectorData (not stored in .store).
310    /// Uses binary search on sorted doc_ids for O(log N) lookup.
311    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    /// Get document by local doc_id, hydrating only the specified fields.
316    ///
317    /// If `fields` is `None`, all fields (including dense vectors) are hydrated.
318    /// If `fields` is `Some(set)`, only dense vector fields in the set are hydrated,
319    /// skipping expensive mmap reads + dequantization for unrequested vector fields.
320    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        // Hydrate dense vector fields from flat vector data
332        for (&field_id, lazy_flat) in &self.flat_vectors {
333            // Skip vector fields not in the requested set
334            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    /// Prefetch term dictionary blocks for a key range
358    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    /// Check if store uses dictionary compression (incompatible with raw merging)
377    pub fn store_has_dict(&self) -> bool {
378        self.store.has_dict()
379    }
380
381    /// Get store reference for merge operations
382    pub fn store(&self) -> &super::store::AsyncStoreReader {
383        &self.store
384    }
385
386    /// Get raw store blocks for optimized merging
387    pub fn store_raw_blocks(&self) -> Vec<RawStoreBlock> {
388        self.store.raw_blocks()
389    }
390
391    /// Get store data slice for raw block access
392    pub fn store_data_slice(&self) -> &LazyFileSlice {
393        self.store.data_slice()
394    }
395
396    /// Get all terms from this segment (for merge)
397    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    /// Get all terms with parsed field and term string (for statistics aggregation)
402    ///
403    /// Returns (field, term_string, doc_freq) for each term in the dictionary.
404    /// Skips terms that aren't valid UTF-8.
405    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            // Key format: field_id (4 bytes little-endian) + term bytes
411            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    /// Get streaming iterator over term dictionary (for memory-efficient merge)
424    pub fn term_dict_iter(&self) -> crate::structures::AsyncSSTableIterator<'_, TermInfo> {
425        self.term_dict.iter()
426    }
427
428    /// Prefetch all term dictionary blocks in a single bulk I/O call.
429    ///
430    /// Call before merge iteration to eliminate per-block cache misses.
431    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    /// Read raw posting bytes at offset
439    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    /// Read raw position bytes at offset (for merge)
447    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    /// Check if this segment has a positions file
459    pub fn has_positions_file(&self) -> bool {
460        self.positions_handle.is_some()
461    }
462
463    /// Batch cosine scoring on raw quantized bytes.
464    ///
465    /// Dispatches to the appropriate SIMD scorer based on quantization type.
466    /// Vectors file uses data-first layout (offset 0) with 8-byte padding between
467    /// fields, so mmap slices are always properly aligned for f32/f16/u8 access.
468    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    /// Search dense vectors using RaBitQ
496    ///
497    /// Returns VectorSearchResult with ordinal tracking for multi-value fields.
498    /// The doc_ids are adjusted by doc_id_offset for this segment.
499    /// For multi-valued documents, scores are combined using the specified combiner.
500    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        // No vectors at all for this field
513        if ann_index.is_none() && lazy_flat.is_none() {
514            return Ok(Vec::new());
515        }
516
517        /// Batch size for brute-force scoring (4096 vectors × 768 dims × 4 bytes ≈ 12MB)
518        const BRUTE_FORCE_BATCH: usize = 4096;
519
520        // Results are (doc_id, ordinal, score) where score = similarity (higher = better)
521        let t0 = std::time::Instant::now();
522        let mut results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
523            // ANN search (RaBitQ, IVF, ScaNN)
524            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            // Batched brute-force from lazy flat vectors (native-precision SIMD scoring)
575            // Uses a top-k heap to avoid collecting and sorting all N candidates.
576            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        // Rerank ANN candidates using raw vectors from lazy flat (binary search lookup)
623        // Uses native-precision SIMD scoring on quantized bytes — no dequantization overhead.
624        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            // Resolve flat indexes for each candidate via binary search
634            let mut resolved: Vec<(usize, usize)> = Vec::new(); // (result_idx, flat_idx)
635            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                // Sort by flat_idx for sequential mmap access (better page locality)
648                resolved.sort_unstable_by_key(|&(_, flat_idx)| flat_idx);
649
650                // Batch-read raw quantized bytes into contiguous buffer
651                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                // Native-precision batch SIMD cosine scoring
665                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                // Write scores back to results
671                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        // Track ordinals with individual scores for each doc_id
698        // Note: doc_id_offset is NOT applied here - the collector applies it uniformly
699        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        // Combine scores and build results with ordinal tracking
707        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        // Sort by score descending and take top k
716        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    /// Check if this segment has dense vectors for the given field
727    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    /// Get the dense vector index for a field (if available)
732    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    /// Get the IVF vector index for a field (if available)
740    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    /// Get coarse centroids for a field
751    pub fn coarse_centroids(&self, field_id: u32) -> Option<&Arc<CoarseCentroids>> {
752        self.coarse_centroids.get(&field_id)
753    }
754
755    /// Set per-field coarse centroids from index-level trained structures
756    pub fn set_coarse_centroids(&mut self, centroids: FxHashMap<u32, Arc<CoarseCentroids>>) {
757        self.coarse_centroids = centroids;
758    }
759
760    /// Get the ScaNN vector index for a field (if available)
761    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    /// Get the vector index type for a field
772    pub fn get_vector_index(&self, field: Field) -> Option<&VectorIndex> {
773        self.vector_indexes.get(&field.0)
774    }
775
776    /// Search for similar sparse vectors using dedicated sparse posting lists
777    ///
778    /// Uses shared `WandExecutor` with `SparseTermScorer` for efficient top-k retrieval.
779    /// Optimizations (via WandExecutor):
780    /// 1. **MaxScore pruning**: Dimensions sorted by max contribution
781    /// 2. **Block-Max WAND**: Skips blocks where max contribution < threshold
782    /// 3. **Top-K heap**: Efficient score collection
783    ///
784    /// Returns VectorSearchResult with ordinal tracking for multi-value fields.
785    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        // Get sparse index for this field
798        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        // Filter query terms to only those present in the index
812        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        // Select executor based on number of query terms:
837        // - 12+ terms: BMP (block-at-a-time, lazy block loading, best for SPLADE)
838        // - 1-11 terms: BlockMaxScoreExecutor (unified MaxScore + block-max + conjunction)
839        let num_terms = matched_terms.len();
840        let over_fetch = limit * 2; // Over-fetch for multi-value combining
841        let raw_results = if num_terms > 12 {
842            // BMP: lazy block loading — only skip entries in memory, blocks loaded on-demand
843            BmpExecutor::new(sparse_index, matched_terms, over_fetch, heap_factor)
844                .execute()
845                .await?
846        } else {
847            // Load posting lists only for the few terms (1-11) used by BlockMaxScore
848            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        // Track ordinals with individual scores for each doc_id
883        // Now using real ordinals from the posting lists
884        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        // Combine scores and build results with ordinal tracking
892        // Note: doc_id_offset is NOT applied here - the collector applies it uniformly
893        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        // Sort by score descending and take top limit
902        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    /// Get positions for a term (for phrase queries)
913    ///
914    /// Position offsets are now embedded in TermInfo, so we first look up
915    /// the term to get its TermInfo, then use position_info() to get the offset.
916    pub async fn get_positions(
917        &self,
918        field: Field,
919        term: &[u8],
920    ) -> Result<Option<crate::structures::PositionPostingList>> {
921        // Get positions handle
922        let handle = match &self.positions_handle {
923            Some(h) => h,
924            None => return Ok(None),
925        };
926
927        // Build key: field_id + term
928        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        // Look up term in dictionary to get TermInfo with position offset
933        let term_info = match self.term_dict.get(&key).await? {
934            Some(info) => info,
935            None => return Ok(None),
936        };
937
938        // Get position offset from TermInfo
939        let (offset, length) = match term_info.position_info() {
940            Some((o, l)) => (o, l),
941            None => return Ok(None),
942        };
943
944        // Read the position data
945        let slice = handle.slice(offset..offset + length as u64);
946        let data = slice.read_bytes().await?;
947
948        // Deserialize
949        let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
950
951        Ok(Some(pos_list))
952    }
953
954    /// Check if positions are available for a field
955    pub fn has_positions(&self, field: Field) -> bool {
956        // Check schema for position mode on this field
957        if let Some(entry) = self.schema.get_field_entry(field) {
958            entry.positions.is_some()
959        } else {
960            false
961        }
962    }
963}
964
965/// Alias for AsyncSegmentReader
966pub type SegmentReader = AsyncSegmentReader;