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 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
54/// Async segment reader with lazy loading
55///
56/// - Term dictionary: only index loaded, blocks loaded on-demand
57/// - Postings: loaded on-demand per term via HTTP range requests
58/// - Document store: only index loaded, blocks loaded on-demand via HTTP range requests
59pub struct AsyncSegmentReader {
60    meta: SegmentMeta,
61    /// Term dictionary with lazy block loading
62    term_dict: Arc<AsyncSSTableReader<TermInfo>>,
63    /// Postings file handle - fetches ranges on demand
64    postings_handle: LazyFileHandle,
65    /// Document store with lazy block loading
66    store: Arc<AsyncStoreReader>,
67    schema: Arc<Schema>,
68    /// Base doc_id offset for this segment
69    doc_id_offset: DocId,
70    /// Dense vector indexes per field (RaBitQ or IVF-RaBitQ) — for search
71    vector_indexes: FxHashMap<u32, VectorIndex>,
72    /// Lazy flat vectors per field — for reranking and merge (doc_ids in memory, vectors via mmap)
73    flat_vectors: FxHashMap<u32, LazyFlatVectorData>,
74    /// Per-field coarse centroids for IVF/ScaNN search
75    coarse_centroids: FxHashMap<u32, Arc<CoarseCentroids>>,
76    /// Sparse vector indexes per field
77    sparse_indexes: FxHashMap<u32, SparseIndex>,
78    /// Position file handle for phrase queries (lazy loading)
79    positions_handle: Option<LazyFileHandle>,
80}
81
82impl AsyncSegmentReader {
83    /// Open a segment with lazy loading
84    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        // Read metadata (small, always loaded)
94        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        // Open term dictionary with lazy loading (fetches ranges on demand)
100        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        // Get postings file handle (lazy - fetches ranges on demand)
104        let postings_handle = dir.open_lazy(&files.postings).await?;
105
106        // Open store with lazy loading
107        let store_handle = dir.open_lazy(&files.store).await?;
108        let store = AsyncStoreReader::open(store_handle, cache_blocks).await?;
109
110        // Load dense vector indexes from unified .vectors file
111        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        // Load sparse vector indexes from .sparse file
116        let sparse_indexes = loader::load_sparse_file(dir, &files, meta.num_docs, &schema).await?;
117
118        // Open positions file handle (if exists) - offsets are now in TermInfo
119        let positions_handle = loader::open_positions_file(dir, &files, &schema).await?;
120
121        // Log segment loading stats
122        {
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    /// Get average field length for BM25F scoring
169    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    /// Set the doc_id_offset (used for parallel segment loading)
178    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    /// Get sparse indexes for all fields
187    pub fn sparse_indexes(&self) -> &FxHashMap<u32, SparseIndex> {
188        &self.sparse_indexes
189    }
190
191    /// Get vector indexes for all fields
192    pub fn vector_indexes(&self) -> &FxHashMap<u32, VectorIndex> {
193        &self.vector_indexes
194    }
195
196    /// Get lazy flat vectors for all fields (for reranking and merge)
197    pub fn flat_vectors(&self) -> &FxHashMap<u32, LazyFlatVectorData> {
198        &self.flat_vectors
199    }
200
201    /// Get term dictionary stats for debugging
202    pub fn term_dict_stats(&self) -> SSTableStats {
203        self.term_dict.stats()
204    }
205
206    /// Estimate memory usage of this segment reader
207    pub fn memory_stats(&self) -> SegmentMemoryStats {
208        let term_dict_stats = self.term_dict.stats();
209
210        // Term dict cache: num_blocks * avg_block_size (estimate 4KB per cached block)
211        let term_dict_cache_bytes = self.term_dict.cached_blocks() * 4096;
212
213        // Store cache: similar estimate
214        let store_cache_bytes = self.store.cached_blocks() * 4096;
215
216        // Sparse index: SoA dim table + OwnedBytes skip section
217        let sparse_index_bytes: usize = self
218            .sparse_indexes
219            .values()
220            .map(|s| s.estimated_memory_bytes())
221            .sum();
222
223        // Dense index: vectors are memory-mapped, but we track index structures
224        // RaBitQ/IVF indexes have cluster assignments in memory
225        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    /// Get posting list for a term (async - loads on demand)
243    ///
244    /// For small posting lists (1-3 docs), the data is inlined in the term dictionary
245    /// and no additional I/O is needed. For larger lists, reads from .post file.
246    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        // Build key: field_id + term
258        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        // Look up in term dictionary
263        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        // Check if posting list is inlined
275        if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
276            // Build BlockPostingList from inline data (no I/O needed!)
277            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        // External posting list - read from postings file handle (lazy - HTTP range request)
286        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    /// Get document by local doc_id (async - loads on demand).
306    ///
307    /// Dense vector fields are hydrated from LazyFlatVectorData (not stored in .store).
308    /// Uses binary search on sorted doc_ids for O(log N) lookup.
309    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    /// Get document by local doc_id, hydrating only the specified fields.
314    ///
315    /// If `fields` is `None`, all fields (including dense vectors) are hydrated.
316    /// If `fields` is `Some(set)`, only dense vector fields in the set are hydrated,
317    /// skipping expensive mmap reads + dequantization for unrequested vector fields.
318    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        // Hydrate dense vector fields from flat vector data
330        for (&field_id, lazy_flat) in &self.flat_vectors {
331            // Skip vector fields not in the requested set
332            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    /// Prefetch term dictionary blocks for a key range
356    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    /// Check if store uses dictionary compression (incompatible with raw merging)
375    pub fn store_has_dict(&self) -> bool {
376        self.store.has_dict()
377    }
378
379    /// Get store reference for merge operations
380    pub fn store(&self) -> &super::store::AsyncStoreReader {
381        &self.store
382    }
383
384    /// Get raw store blocks for optimized merging
385    pub fn store_raw_blocks(&self) -> Vec<RawStoreBlock> {
386        self.store.raw_blocks()
387    }
388
389    /// Get store data slice for raw block access
390    pub fn store_data_slice(&self) -> &LazyFileSlice {
391        self.store.data_slice()
392    }
393
394    /// Get all terms from this segment (for merge)
395    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    /// Get all terms with parsed field and term string (for statistics aggregation)
400    ///
401    /// Returns (field, term_string, doc_freq) for each term in the dictionary.
402    /// Skips terms that aren't valid UTF-8.
403    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            // Key format: field_id (4 bytes little-endian) + term bytes
409            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    /// Get streaming iterator over term dictionary (for memory-efficient merge)
422    pub fn term_dict_iter(&self) -> crate::structures::AsyncSSTableIterator<'_, TermInfo> {
423        self.term_dict.iter()
424    }
425
426    /// Prefetch all term dictionary blocks in a single bulk I/O call.
427    ///
428    /// Call before merge iteration to eliminate per-block cache misses.
429    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    /// Read raw posting bytes at offset
437    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    /// Read raw position bytes at offset (for merge)
445    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    /// Check if this segment has a positions file
457    pub fn has_positions_file(&self) -> bool {
458        self.positions_handle.is_some()
459    }
460
461    /// Batch cosine scoring on raw quantized bytes.
462    ///
463    /// Dispatches to the appropriate SIMD scorer based on quantization type.
464    /// Vectors file uses data-first layout (offset 0) with 8-byte padding between
465    /// fields, so mmap slices are always properly aligned for f32/f16/u8 access.
466    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    /// Search dense vectors using RaBitQ
513    ///
514    /// Returns VectorSearchResult with ordinal tracking for multi-value fields.
515    /// The doc_ids are adjusted by doc_id_offset for this segment.
516    /// For multi-valued documents, scores are combined using the specified combiner.
517    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        // No vectors at all for this field
530        if ann_index.is_none() && lazy_flat.is_none() {
531            return Ok(Vec::new());
532        }
533
534        // Check if vectors are pre-normalized (skip per-vector norm in scoring)
535        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        /// Batch size for brute-force scoring (4096 vectors × 768 dims × 4 bytes ≈ 12MB)
542        const BRUTE_FORCE_BATCH: usize = 4096;
543
544        // Results are (doc_id, ordinal, score) where score = similarity (higher = better)
545        let t0 = std::time::Instant::now();
546        let mut results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
547            // ANN search (RaBitQ, IVF, ScaNN)
548            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            // Batched brute-force from lazy flat vectors (native-precision SIMD scoring)
599            // Uses a top-k heap to avoid collecting and sorting all N candidates.
600            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        // Rerank ANN candidates using raw vectors from lazy flat (binary search lookup)
654        // Uses native-precision SIMD scoring on quantized bytes — no dequantization overhead.
655        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            // Resolve flat indexes for each candidate via binary search
665            let mut resolved: Vec<(usize, usize)> = Vec::new(); // (result_idx, flat_idx)
666            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                // Sort by flat_idx for sequential mmap access (better page locality)
679                resolved.sort_unstable_by_key(|&(_, flat_idx)| flat_idx);
680
681                // Batch-read raw quantized bytes into contiguous buffer
682                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                // Native-precision batch SIMD cosine scoring
696                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                // Write scores back to results
702                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        // Track ordinals with individual scores for each doc_id
729        // Note: doc_id_offset is NOT applied here - the collector applies it uniformly
730        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        // Combine scores and build results with ordinal tracking
738        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        // Sort by score descending and take top k
747        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    /// Check if this segment has dense vectors for the given field
758    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    /// Get the dense vector index for a field (if available)
763    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    /// Get the IVF vector index for a field (if available)
771    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    /// Get coarse centroids for a field
782    pub fn coarse_centroids(&self, field_id: u32) -> Option<&Arc<CoarseCentroids>> {
783        self.coarse_centroids.get(&field_id)
784    }
785
786    /// Set per-field coarse centroids from index-level trained structures
787    pub fn set_coarse_centroids(&mut self, centroids: FxHashMap<u32, Arc<CoarseCentroids>>) {
788        self.coarse_centroids = centroids;
789    }
790
791    /// Get the ScaNN vector index for a field (if available)
792    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    /// Get the vector index type for a field
803    pub fn get_vector_index(&self, field: Field) -> Option<&VectorIndex> {
804        self.vector_indexes.get(&field.0)
805    }
806
807    /// Search for similar sparse vectors using dedicated sparse posting lists
808    ///
809    /// Uses shared `WandExecutor` with `SparseTermScorer` for efficient top-k retrieval.
810    /// Optimizations (via WandExecutor):
811    /// 1. **MaxScore pruning**: Dimensions sorted by max contribution
812    /// 2. **Block-Max WAND**: Skips blocks where max contribution < threshold
813    /// 3. **Top-K heap**: Efficient score collection
814    ///
815    /// Returns VectorSearchResult with ordinal tracking for multi-value fields.
816    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        // Get sparse index for this field
829        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        // Filter query terms to only those present in the index
843        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        // Select executor based on number of query terms:
868        // - 12+ terms: BMP (block-at-a-time, lazy block loading, best for SPLADE)
869        // - 1-11 terms: LazyBlockMaxScoreExecutor (cursor-based traversal + lazy block loading)
870        let num_terms = matched_terms.len();
871        let over_fetch = limit * 2; // Over-fetch for multi-value combining
872        let raw_results = if num_terms > 12 {
873            // BMP: lazy block loading — only skip entries in memory, blocks loaded on-demand
874            BmpExecutor::new(sparse_index, matched_terms, over_fetch, heap_factor)
875                .execute()
876                .await?
877        } else {
878            // Lazy BlockMaxScore: skip entries drive navigation, blocks loaded on-demand.
879            // No upfront get_posting() — blocks only loaded when cursor actually visits them.
880            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        // Track ordinals with individual scores for each doc_id
908        // Now using real ordinals from the posting lists
909        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        // Combine scores and build results with ordinal tracking
917        // Note: doc_id_offset is NOT applied here - the collector applies it uniformly
918        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        // Sort by score descending and take top limit
927        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    /// Get positions for a term (for phrase queries)
938    ///
939    /// Position offsets are now embedded in TermInfo, so we first look up
940    /// the term to get its TermInfo, then use position_info() to get the offset.
941    pub async fn get_positions(
942        &self,
943        field: Field,
944        term: &[u8],
945    ) -> Result<Option<crate::structures::PositionPostingList>> {
946        // Get positions handle
947        let handle = match &self.positions_handle {
948            Some(h) => h,
949            None => return Ok(None),
950        };
951
952        // Build key: field_id + term
953        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        // Look up term in dictionary to get TermInfo with position offset
958        let term_info = match self.term_dict.get(&key).await? {
959            Some(info) => info,
960            None => return Ok(None),
961        };
962
963        // Get position offset from TermInfo
964        let (offset, length) = match term_info.position_info() {
965            Some((o, l)) => (o, l),
966            None => return Ok(None),
967        };
968
969        // Read the position data
970        let slice = handle.slice(offset..offset + length as u64);
971        let data = slice.read_bytes().await?;
972
973        // Deserialize
974        let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
975
976        Ok(Some(pos_list))
977    }
978
979    /// Check if positions are available for a field
980    pub fn has_positions(&self, field: Field) -> bool {
981        // Check schema for position mode on this field
982        if let Some(entry) = self.schema.get_field_entry(field) {
983            entry.positions.is_some()
984        } else {
985            false
986        }
987    }
988}
989
990/// Alias for AsyncSegmentReader
991pub type SegmentReader = AsyncSegmentReader;