Skip to main content

hermes_core/segment/reader/
mod.rs

1//! Async segment reader with lazy loading
2
3pub(crate) mod bmp;
4mod loader;
5mod types;
6
7pub use bmp::BmpIndex;
8#[cfg(feature = "diagnostics")]
9pub use types::DimRawData;
10pub use types::{SparseIndex, VectorIndex, VectorSearchResult};
11
12/// Memory statistics for a single segment
13#[derive(Debug, Clone, Default)]
14pub struct SegmentMemoryStats {
15    /// Segment ID
16    pub segment_id: u128,
17    /// Number of documents in segment
18    pub num_docs: u32,
19    /// Term dictionary block cache bytes
20    pub term_dict_cache_bytes: usize,
21    /// Document store block cache bytes
22    pub store_cache_bytes: usize,
23    /// Sparse vector index bytes (in-memory posting lists)
24    pub sparse_index_bytes: usize,
25    /// Dense vector index bytes (cluster assignments, quantized codes)
26    pub dense_index_bytes: usize,
27    /// Bloom filter bytes
28    pub bloom_filter_bytes: usize,
29}
30
31impl SegmentMemoryStats {
32    /// Total estimated memory for this segment
33    pub fn total_bytes(&self) -> usize {
34        self.term_dict_cache_bytes
35            + self.store_cache_bytes
36            + self.sparse_index_bytes
37            + self.dense_index_bytes
38            + self.bloom_filter_bytes
39    }
40}
41
42use std::sync::Arc;
43
44use rustc_hash::FxHashMap;
45
46use super::vector_data::LazyFlatVectorData;
47use crate::directories::{Directory, FileHandle};
48use crate::dsl::{Document, Field, Schema};
49use crate::structures::{
50    AsyncSSTableReader, BlockPostingList, CoarseCentroids, IVFPQIndex, IVFRaBitQIndex, PQCodebook,
51    RaBitQIndex, SSTableStats, TermInfo,
52};
53use crate::{DocId, Error, Result};
54
55use super::store::{AsyncStoreReader, RawStoreBlock};
56use super::types::{SegmentFiles, SegmentId, SegmentMeta};
57
58/// Combine per-ordinal (doc_id, ordinal, score) triples into VectorSearchResults,
59/// applying the multi-value combiner, sorting by score desc, and truncating to `limit`.
60///
61/// Fast path: when all ordinals are 0 (single-valued field), skips the HashMap
62/// grouping entirely and just sorts + truncates the raw results.
63pub(crate) fn combine_ordinal_results(
64    raw: impl IntoIterator<Item = (u32, u16, f32)>,
65    combiner: crate::query::MultiValueCombiner,
66    limit: usize,
67) -> Vec<VectorSearchResult> {
68    let collected: Vec<(u32, u16, f32)> = raw.into_iter().collect();
69
70    let num_raw = collected.len();
71    if log::log_enabled!(log::Level::Debug) {
72        let mut ids: Vec<u32> = collected.iter().map(|(d, _, _)| *d).collect();
73        ids.sort_unstable();
74        ids.dedup();
75        log::debug!(
76            "combine_ordinal_results: {} raw entries, {} unique docs, combiner={:?}, limit={}",
77            num_raw,
78            ids.len(),
79            combiner,
80            limit
81        );
82    }
83
84    // Fast path: all ordinals are 0 → no grouping needed, skip HashMap
85    let all_single = collected.iter().all(|&(_, ord, _)| ord == 0);
86    if all_single {
87        let mut results: Vec<VectorSearchResult> = collected
88            .into_iter()
89            .map(|(doc_id, _, score)| VectorSearchResult::new(doc_id, score, vec![(0, score)]))
90            .collect();
91        results.sort_unstable_by(|a, b| {
92            b.score
93                .partial_cmp(&a.score)
94                .unwrap_or(std::cmp::Ordering::Equal)
95        });
96        results.truncate(limit);
97        return results;
98    }
99
100    // Slow path: multi-valued field — group by doc_id, apply combiner
101    let mut doc_ordinals: rustc_hash::FxHashMap<DocId, Vec<(u32, f32)>> =
102        rustc_hash::FxHashMap::default();
103    for (doc_id, ordinal, score) in collected {
104        doc_ordinals
105            .entry(doc_id as DocId)
106            .or_default()
107            .push((ordinal as u32, score));
108    }
109    let mut results: Vec<VectorSearchResult> = doc_ordinals
110        .into_iter()
111        .map(|(doc_id, ordinals)| {
112            let combined_score = combiner.combine(&ordinals);
113            VectorSearchResult::new(doc_id, combined_score, ordinals)
114        })
115        .collect();
116    results.sort_unstable_by(|a, b| {
117        b.score
118            .partial_cmp(&a.score)
119            .unwrap_or(std::cmp::Ordering::Equal)
120    });
121    results.truncate(limit);
122    results
123}
124
125/// Async segment reader with lazy loading
126///
127/// - Term dictionary: only index loaded, blocks loaded on-demand
128/// - Postings: loaded on-demand per term via HTTP range requests
129/// - Document store: only index loaded, blocks loaded on-demand via HTTP range requests
130pub struct SegmentReader {
131    meta: SegmentMeta,
132    /// Term dictionary with lazy block loading
133    term_dict: Arc<AsyncSSTableReader<TermInfo>>,
134    /// Postings file handle - fetches ranges on demand
135    postings_handle: FileHandle,
136    /// Document store with lazy block loading
137    store: Arc<AsyncStoreReader>,
138    schema: Arc<Schema>,
139    /// Dense vector indexes per field (RaBitQ or IVF-RaBitQ) — for search
140    vector_indexes: FxHashMap<u32, VectorIndex>,
141    /// Lazy flat vectors per field — for reranking and merge (doc_ids in memory, vectors via mmap)
142    flat_vectors: FxHashMap<u32, LazyFlatVectorData>,
143    /// Per-field coarse centroids for IVF/ScaNN search
144    coarse_centroids: FxHashMap<u32, Arc<CoarseCentroids>>,
145    /// Sparse vector indexes per field (MaxScore format)
146    sparse_indexes: FxHashMap<u32, SparseIndex>,
147    /// BMP sparse vector indexes per field (BMP format)
148    bmp_indexes: FxHashMap<u32, BmpIndex>,
149    /// Position file handle for phrase queries (lazy loading)
150    positions_handle: Option<FileHandle>,
151    /// Fast-field columnar readers per field_id
152    fast_fields: FxHashMap<u32, crate::structures::fast_field::FastFieldReader>,
153}
154
155impl SegmentReader {
156    /// Open a segment with lazy loading
157    pub async fn open<D: Directory>(
158        dir: &D,
159        segment_id: SegmentId,
160        schema: Arc<Schema>,
161        cache_blocks: usize,
162    ) -> Result<Self> {
163        let files = SegmentFiles::new(segment_id.0);
164
165        // Read metadata (small, always loaded)
166        let meta_slice = dir.open_read(&files.meta).await?;
167        let meta_bytes = meta_slice.read_bytes().await?;
168        let meta = SegmentMeta::deserialize(meta_bytes.as_slice())?;
169        debug_assert_eq!(meta.id, segment_id.0);
170
171        // Open term dictionary with lazy loading (fetches ranges on demand)
172        let term_dict_handle = dir.open_lazy(&files.term_dict).await?;
173        let term_dict = AsyncSSTableReader::open(term_dict_handle, cache_blocks).await?;
174
175        // Get postings file handle (lazy - fetches ranges on demand)
176        let postings_handle = dir.open_lazy(&files.postings).await?;
177
178        // Open store with lazy loading
179        let store_handle = dir.open_lazy(&files.store).await?;
180        let store = AsyncStoreReader::open(store_handle, cache_blocks).await?;
181
182        // Load dense vector indexes from unified .vectors file
183        let vectors_data = loader::load_vectors_file(dir, &files, &schema).await?;
184        let vector_indexes = vectors_data.indexes;
185        let flat_vectors = vectors_data.flat_vectors;
186
187        // Load sparse vector indexes from .sparse file (MaxScore + BMP)
188        let sparse_data = loader::load_sparse_file(dir, &files, meta.num_docs, &schema).await?;
189        let sparse_indexes = sparse_data.maxscore_indexes;
190        let bmp_indexes = sparse_data.bmp_indexes;
191
192        // Open positions file handle (if exists) - offsets are now in TermInfo
193        let positions_handle = loader::open_positions_file(dir, &files, &schema).await?;
194
195        // Load fast-field columns from .fast file
196        let fast_fields = loader::load_fast_fields_file(dir, &files, &schema).await?;
197
198        // Log segment loading stats
199        {
200            let mut parts = vec![format!(
201                "[segment] loaded {:016x}: docs={}",
202                segment_id.0, meta.num_docs
203            )];
204            if !vector_indexes.is_empty() || !flat_vectors.is_empty() {
205                parts.push(format!(
206                    "dense: {} ann + {} flat fields",
207                    vector_indexes.len(),
208                    flat_vectors.len()
209                ));
210            }
211            for (field_id, idx) in &sparse_indexes {
212                parts.push(format!(
213                    "sparse field {}: {} dims, ~{:.1} KB",
214                    field_id,
215                    idx.num_dimensions(),
216                    idx.num_dimensions() as f64 * 24.0 / 1024.0
217                ));
218            }
219            for (field_id, idx) in &bmp_indexes {
220                parts.push(format!(
221                    "bmp field {}: {} dims, {} blocks",
222                    field_id,
223                    idx.dims(),
224                    idx.num_blocks
225                ));
226            }
227            if !fast_fields.is_empty() {
228                parts.push(format!("fast: {} fields", fast_fields.len()));
229            }
230            log::debug!("{}", parts.join(", "));
231        }
232
233        Ok(Self {
234            meta,
235            term_dict: Arc::new(term_dict),
236            postings_handle,
237            store: Arc::new(store),
238            schema,
239            vector_indexes,
240            flat_vectors,
241            coarse_centroids: FxHashMap::default(),
242            sparse_indexes,
243            bmp_indexes,
244            positions_handle,
245            fast_fields,
246        })
247    }
248
249    pub fn meta(&self) -> &SegmentMeta {
250        &self.meta
251    }
252
253    pub fn num_docs(&self) -> u32 {
254        self.meta.num_docs
255    }
256
257    /// Get average field length for BM25F scoring
258    pub fn avg_field_len(&self, field: Field) -> f32 {
259        self.meta.avg_field_len(field)
260    }
261
262    pub fn schema(&self) -> &Schema {
263        &self.schema
264    }
265
266    /// Get sparse indexes for all fields
267    pub fn sparse_indexes(&self) -> &FxHashMap<u32, SparseIndex> {
268        &self.sparse_indexes
269    }
270
271    /// Get sparse index for a specific field (MaxScore format)
272    pub fn sparse_index(&self, field: Field) -> Option<&SparseIndex> {
273        self.sparse_indexes.get(&field.0)
274    }
275
276    /// Get BMP index for a specific field
277    pub fn bmp_index(&self, field: Field) -> Option<&BmpIndex> {
278        self.bmp_indexes.get(&field.0)
279    }
280
281    /// Get all BMP indexes
282    pub fn bmp_indexes(&self) -> &FxHashMap<u32, BmpIndex> {
283        &self.bmp_indexes
284    }
285
286    /// Get vector indexes for all fields
287    pub fn vector_indexes(&self) -> &FxHashMap<u32, VectorIndex> {
288        &self.vector_indexes
289    }
290
291    /// Get lazy flat vectors for all fields (for reranking and merge)
292    pub fn flat_vectors(&self) -> &FxHashMap<u32, LazyFlatVectorData> {
293        &self.flat_vectors
294    }
295
296    /// Get a fast-field reader for a specific field.
297    pub fn fast_field(
298        &self,
299        field_id: u32,
300    ) -> Option<&crate::structures::fast_field::FastFieldReader> {
301        self.fast_fields.get(&field_id)
302    }
303
304    /// Get all fast-field readers.
305    pub fn fast_fields(&self) -> &FxHashMap<u32, crate::structures::fast_field::FastFieldReader> {
306        &self.fast_fields
307    }
308
309    /// Get term dictionary stats for debugging
310    pub fn term_dict_stats(&self) -> SSTableStats {
311        self.term_dict.stats()
312    }
313
314    /// Estimate memory usage of this segment reader
315    pub fn memory_stats(&self) -> SegmentMemoryStats {
316        let term_dict_stats = self.term_dict.stats();
317
318        // Term dict cache: num_blocks * avg_block_size (estimate 4KB per cached block)
319        let term_dict_cache_bytes = self.term_dict.cached_blocks() * 4096;
320
321        // Store cache: similar estimate
322        let store_cache_bytes = self.store.cached_blocks() * 4096;
323
324        // Sparse index: SoA dim table + OwnedBytes skip section + BMP grids
325        let sparse_index_bytes: usize = self
326            .sparse_indexes
327            .values()
328            .map(|s| s.estimated_memory_bytes())
329            .sum::<usize>()
330            + self
331                .bmp_indexes
332                .values()
333                .map(|b| b.estimated_memory_bytes())
334                .sum::<usize>();
335
336        // Dense index: vectors are memory-mapped, but we track index structures
337        // RaBitQ/IVF indexes have cluster assignments in memory
338        let dense_index_bytes: usize = self
339            .vector_indexes
340            .values()
341            .map(|v| v.estimated_memory_bytes())
342            .sum();
343
344        SegmentMemoryStats {
345            segment_id: self.meta.id,
346            num_docs: self.meta.num_docs,
347            term_dict_cache_bytes,
348            store_cache_bytes,
349            sparse_index_bytes,
350            dense_index_bytes,
351            bloom_filter_bytes: term_dict_stats.bloom_filter_size,
352        }
353    }
354
355    /// Get posting list for a term (async - loads on demand)
356    ///
357    /// For small posting lists (1-3 docs), the data is inlined in the term dictionary
358    /// and no additional I/O is needed. For larger lists, reads from .post file.
359    pub async fn get_postings(
360        &self,
361        field: Field,
362        term: &[u8],
363    ) -> Result<Option<BlockPostingList>> {
364        log::debug!(
365            "SegmentReader::get_postings field={} term_len={}",
366            field.0,
367            term.len()
368        );
369
370        // Build key: field_id + term
371        let mut key = Vec::with_capacity(4 + term.len());
372        key.extend_from_slice(&field.0.to_le_bytes());
373        key.extend_from_slice(term);
374
375        // Look up in term dictionary
376        let term_info = match self.term_dict.get(&key).await? {
377            Some(info) => {
378                log::debug!("SegmentReader::get_postings found term_info");
379                info
380            }
381            None => {
382                log::debug!("SegmentReader::get_postings term not found");
383                return Ok(None);
384            }
385        };
386
387        // Check if posting list is inlined
388        if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
389            // Build BlockPostingList from inline data (no I/O needed!)
390            let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
391            for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs.into_iter()) {
392                posting_list.push(doc_id, tf);
393            }
394            let block_list = BlockPostingList::from_posting_list(&posting_list)?;
395            return Ok(Some(block_list));
396        }
397
398        // External posting list - read from postings file handle (lazy - HTTP range request)
399        let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
400            Error::Corruption("TermInfo has neither inline nor external data".to_string())
401        })?;
402
403        let start = posting_offset;
404        let end = start + posting_len;
405
406        if end > self.postings_handle.len() {
407            return Err(Error::Corruption(
408                "Posting offset out of bounds".to_string(),
409            ));
410        }
411
412        let posting_bytes = self.postings_handle.read_bytes_range(start..end).await?;
413        let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
414
415        Ok(Some(block_list))
416    }
417
418    /// Get document by local doc_id (async - loads on demand).
419    ///
420    /// Dense vector fields are hydrated from LazyFlatVectorData (not stored in .store).
421    /// Uses binary search on sorted doc_ids for O(log N) lookup.
422    pub async fn doc(&self, local_doc_id: DocId) -> Result<Option<Document>> {
423        self.doc_with_fields(local_doc_id, None).await
424    }
425
426    /// Get document by local doc_id, hydrating only the specified fields.
427    ///
428    /// If `fields` is `None`, all fields (including dense vectors) are hydrated.
429    /// If `fields` is `Some(set)`, only dense vector fields in the set are hydrated,
430    /// skipping expensive mmap reads + dequantization for unrequested vector fields.
431    pub async fn doc_with_fields(
432        &self,
433        local_doc_id: DocId,
434        fields: Option<&rustc_hash::FxHashSet<u32>>,
435    ) -> Result<Option<Document>> {
436        let mut doc = match fields {
437            Some(set) => {
438                let field_ids: Vec<u32> = set.iter().copied().collect();
439                match self
440                    .store
441                    .get_fields(local_doc_id, &self.schema, &field_ids)
442                    .await
443                {
444                    Ok(Some(d)) => d,
445                    Ok(None) => return Ok(None),
446                    Err(e) => return Err(Error::from(e)),
447                }
448            }
449            None => match self.store.get(local_doc_id, &self.schema).await {
450                Ok(Some(d)) => d,
451                Ok(None) => return Ok(None),
452                Err(e) => return Err(Error::from(e)),
453            },
454        };
455
456        // Hydrate dense vector fields from flat vector data
457        for (&field_id, lazy_flat) in &self.flat_vectors {
458            // Skip vector fields not in the requested set
459            if let Some(set) = fields
460                && !set.contains(&field_id)
461            {
462                continue;
463            }
464
465            let (start, entries) = lazy_flat.flat_indexes_for_doc(local_doc_id);
466            for (j, &(_doc_id, _ordinal)) in entries.iter().enumerate() {
467                let flat_idx = start + j;
468                match lazy_flat.get_vector(flat_idx).await {
469                    Ok(vec) => {
470                        doc.add_dense_vector(Field(field_id), vec);
471                    }
472                    Err(e) => {
473                        log::warn!("Failed to hydrate vector field {}: {}", field_id, e);
474                    }
475                }
476            }
477        }
478
479        Ok(Some(doc))
480    }
481
482    /// Prefetch term dictionary blocks for a key range
483    pub async fn prefetch_terms(
484        &self,
485        field: Field,
486        start_term: &[u8],
487        end_term: &[u8],
488    ) -> Result<()> {
489        let mut start_key = Vec::with_capacity(4 + start_term.len());
490        start_key.extend_from_slice(&field.0.to_le_bytes());
491        start_key.extend_from_slice(start_term);
492
493        let mut end_key = Vec::with_capacity(4 + end_term.len());
494        end_key.extend_from_slice(&field.0.to_le_bytes());
495        end_key.extend_from_slice(end_term);
496
497        self.term_dict.prefetch_range(&start_key, &end_key).await?;
498        Ok(())
499    }
500
501    /// Check if store uses dictionary compression (incompatible with raw merging)
502    pub fn store_has_dict(&self) -> bool {
503        self.store.has_dict()
504    }
505
506    /// Get store reference for merge operations
507    pub fn store(&self) -> &super::store::AsyncStoreReader {
508        &self.store
509    }
510
511    /// Get raw store blocks for optimized merging
512    pub fn store_raw_blocks(&self) -> Vec<RawStoreBlock> {
513        self.store.raw_blocks()
514    }
515
516    /// Get store data slice for raw block access
517    pub fn store_data_slice(&self) -> &FileHandle {
518        self.store.data_slice()
519    }
520
521    /// Get all terms from this segment (for merge)
522    pub async fn all_terms(&self) -> Result<Vec<(Vec<u8>, TermInfo)>> {
523        self.term_dict.all_entries().await.map_err(Error::from)
524    }
525
526    /// Get all terms with parsed field and term string (for statistics aggregation)
527    ///
528    /// Returns (field, term_string, doc_freq) for each term in the dictionary.
529    /// Skips terms that aren't valid UTF-8.
530    pub async fn all_terms_with_stats(&self) -> Result<Vec<(Field, String, u32)>> {
531        let entries = self.term_dict.all_entries().await?;
532        let mut result = Vec::with_capacity(entries.len());
533
534        for (key, term_info) in entries {
535            // Key format: field_id (4 bytes little-endian) + term bytes
536            if key.len() > 4 {
537                let field_id = u32::from_le_bytes([key[0], key[1], key[2], key[3]]);
538                let term_bytes = &key[4..];
539                if let Ok(term_str) = std::str::from_utf8(term_bytes) {
540                    result.push((Field(field_id), term_str.to_string(), term_info.doc_freq()));
541                }
542            }
543        }
544
545        Ok(result)
546    }
547
548    /// Get streaming iterator over term dictionary (for memory-efficient merge)
549    pub fn term_dict_iter(&self) -> crate::structures::AsyncSSTableIterator<'_, TermInfo> {
550        self.term_dict.iter()
551    }
552
553    /// Prefetch all term dictionary blocks in a single bulk I/O call.
554    ///
555    /// Call before merge iteration to eliminate per-block cache misses.
556    pub async fn prefetch_term_dict(&self) -> crate::Result<()> {
557        self.term_dict
558            .prefetch_all_data_bulk()
559            .await
560            .map_err(crate::Error::from)
561    }
562
563    /// Read raw posting bytes at offset
564    pub async fn read_postings(&self, offset: u64, len: u64) -> Result<Vec<u8>> {
565        let start = offset;
566        let end = start + len;
567        let bytes = self.postings_handle.read_bytes_range(start..end).await?;
568        Ok(bytes.to_vec())
569    }
570
571    /// Read raw position bytes at offset (for merge)
572    pub async fn read_position_bytes(&self, offset: u64, len: u64) -> Result<Option<Vec<u8>>> {
573        let handle = match &self.positions_handle {
574            Some(h) => h,
575            None => return Ok(None),
576        };
577        let start = offset;
578        let end = start + len;
579        let bytes = handle.read_bytes_range(start..end).await?;
580        Ok(Some(bytes.to_vec()))
581    }
582
583    /// Check if this segment has a positions file
584    pub fn has_positions_file(&self) -> bool {
585        self.positions_handle.is_some()
586    }
587
588    /// Batch cosine scoring on raw quantized bytes.
589    ///
590    /// Dispatches to the appropriate SIMD scorer based on quantization type.
591    /// Vectors file uses data-first layout (offset 0) with 8-byte padding between
592    /// fields, so mmap slices are always properly aligned for f32/f16/u8 access.
593    fn score_quantized_batch(
594        query: &[f32],
595        raw: &[u8],
596        quant: crate::dsl::DenseVectorQuantization,
597        dim: usize,
598        scores: &mut [f32],
599        unit_norm: bool,
600    ) {
601        use crate::dsl::DenseVectorQuantization;
602        use crate::structures::simd;
603        match (quant, unit_norm) {
604            (DenseVectorQuantization::F32, false) => {
605                let num_floats = scores.len() * dim;
606                debug_assert!(
607                    (raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()),
608                    "f32 vector data not 4-byte aligned — vectors file may use legacy format"
609                );
610                let vectors: &[f32] =
611                    unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
612                simd::batch_cosine_scores(query, vectors, dim, scores);
613            }
614            (DenseVectorQuantization::F32, true) => {
615                let num_floats = scores.len() * dim;
616                debug_assert!(
617                    (raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()),
618                    "f32 vector data not 4-byte aligned"
619                );
620                let vectors: &[f32] =
621                    unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
622                simd::batch_dot_scores(query, vectors, dim, scores);
623            }
624            (DenseVectorQuantization::F16, false) => {
625                simd::batch_cosine_scores_f16(query, raw, dim, scores);
626            }
627            (DenseVectorQuantization::F16, true) => {
628                simd::batch_dot_scores_f16(query, raw, dim, scores);
629            }
630            (DenseVectorQuantization::UInt8, false) => {
631                simd::batch_cosine_scores_u8(query, raw, dim, scores);
632            }
633            (DenseVectorQuantization::UInt8, true) => {
634                simd::batch_dot_scores_u8(query, raw, dim, scores);
635            }
636        }
637    }
638
639    /// Search dense vectors using RaBitQ
640    ///
641    /// Returns VectorSearchResult with ordinal tracking for multi-value fields.
642    /// Doc IDs are segment-local.
643    /// For multi-valued documents, scores are combined using the specified combiner.
644    pub async fn search_dense_vector(
645        &self,
646        field: Field,
647        query: &[f32],
648        k: usize,
649        nprobe: usize,
650        rerank_factor: f32,
651        combiner: crate::query::MultiValueCombiner,
652    ) -> Result<Vec<VectorSearchResult>> {
653        let ann_index = self.vector_indexes.get(&field.0);
654        let lazy_flat = self.flat_vectors.get(&field.0);
655
656        // No vectors at all for this field
657        if ann_index.is_none() && lazy_flat.is_none() {
658            return Ok(Vec::new());
659        }
660
661        // Check if vectors are pre-normalized (skip per-vector norm in scoring)
662        let unit_norm = self
663            .schema
664            .get_field_entry(field)
665            .and_then(|e| e.dense_vector_config.as_ref())
666            .is_some_and(|c| c.unit_norm);
667
668        /// Batch size for brute-force scoring (4096 vectors × 768 dims × 4 bytes ≈ 12MB)
669        const BRUTE_FORCE_BATCH: usize = 4096;
670
671        let fetch_k = (k as f32 * rerank_factor.max(1.0)).ceil() as usize;
672
673        // Results are (doc_id, ordinal, score) where score = similarity (higher = better)
674        let t0 = std::time::Instant::now();
675        let mut results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
676            // ANN search (RaBitQ, IVF, ScaNN)
677            match index {
678                VectorIndex::RaBitQ(lazy) => {
679                    let rabitq = lazy.get().ok_or_else(|| {
680                        Error::Schema("RaBitQ index deserialization failed".to_string())
681                    })?;
682                    rabitq
683                        .search(query, fetch_k)
684                        .into_iter()
685                        .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
686                        .collect()
687                }
688                VectorIndex::IVF(lazy) => {
689                    let (index, codebook) = lazy.get().ok_or_else(|| {
690                        Error::Schema("IVF index deserialization failed".to_string())
691                    })?;
692                    let centroids = self.coarse_centroids.get(&field.0).ok_or_else(|| {
693                        Error::Schema(format!(
694                            "IVF index requires coarse centroids for field {}",
695                            field.0
696                        ))
697                    })?;
698                    let effective_nprobe = if nprobe > 0 { nprobe } else { 32 };
699                    index
700                        .search(centroids, codebook, query, fetch_k, Some(effective_nprobe))
701                        .into_iter()
702                        .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
703                        .collect()
704                }
705                VectorIndex::ScaNN(lazy) => {
706                    let (index, codebook) = lazy.get().ok_or_else(|| {
707                        Error::Schema("ScaNN index deserialization failed".to_string())
708                    })?;
709                    let centroids = self.coarse_centroids.get(&field.0).ok_or_else(|| {
710                        Error::Schema(format!(
711                            "ScaNN index requires coarse centroids for field {}",
712                            field.0
713                        ))
714                    })?;
715                    let effective_nprobe = if nprobe > 0 { nprobe } else { 32 };
716                    index
717                        .search(centroids, codebook, query, fetch_k, Some(effective_nprobe))
718                        .into_iter()
719                        .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
720                        .collect()
721                }
722            }
723        } else if let Some(lazy_flat) = lazy_flat {
724            // Batched brute-force from lazy flat vectors (native-precision SIMD scoring)
725            // Uses a top-k heap to avoid collecting and sorting all N candidates.
726            log::debug!(
727                "[search_dense] field {}: brute-force on {} vectors (dim={}, quant={:?})",
728                field.0,
729                lazy_flat.num_vectors,
730                lazy_flat.dim,
731                lazy_flat.quantization
732            );
733            let dim = lazy_flat.dim;
734            let n = lazy_flat.num_vectors;
735            let quant = lazy_flat.quantization;
736            let mut collector = crate::query::ScoreCollector::new(fetch_k);
737            let mut scores = vec![0f32; BRUTE_FORCE_BATCH];
738
739            for batch_start in (0..n).step_by(BRUTE_FORCE_BATCH) {
740                let batch_count = BRUTE_FORCE_BATCH.min(n - batch_start);
741                let batch_bytes = lazy_flat
742                    .read_vectors_batch(batch_start, batch_count)
743                    .await
744                    .map_err(crate::Error::Io)?;
745                let raw = batch_bytes.as_slice();
746
747                Self::score_quantized_batch(
748                    query,
749                    raw,
750                    quant,
751                    dim,
752                    &mut scores[..batch_count],
753                    unit_norm,
754                );
755
756                for (i, &score) in scores.iter().enumerate().take(batch_count) {
757                    let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
758                    collector.insert_with_ordinal(doc_id, score, ordinal);
759                }
760            }
761
762            collector
763                .into_sorted_results()
764                .into_iter()
765                .map(|(doc_id, score, ordinal)| (doc_id, ordinal, score))
766                .collect()
767        } else {
768            return Ok(Vec::new());
769        };
770        let l1_elapsed = t0.elapsed();
771        log::debug!(
772            "[search_dense] field {}: L1 returned {} candidates in {:.1}ms",
773            field.0,
774            results.len(),
775            l1_elapsed.as_secs_f64() * 1000.0
776        );
777
778        // Rerank ANN candidates using raw vectors from lazy flat (binary search lookup)
779        // Uses native-precision SIMD scoring on quantized bytes — no dequantization overhead.
780        if ann_index.is_some()
781            && !results.is_empty()
782            && let Some(lazy_flat) = lazy_flat
783        {
784            let t_rerank = std::time::Instant::now();
785            let dim = lazy_flat.dim;
786            let quant = lazy_flat.quantization;
787            let vbs = lazy_flat.vector_byte_size();
788
789            // Resolve flat indexes for each candidate via binary search
790            let mut resolved: Vec<(usize, usize)> = Vec::new(); // (result_idx, flat_idx)
791            for (ri, c) in results.iter().enumerate() {
792                let (start, entries) = lazy_flat.flat_indexes_for_doc(c.0);
793                for (j, &(_, ord)) in entries.iter().enumerate() {
794                    if ord == c.1 {
795                        resolved.push((ri, start + j));
796                        break;
797                    }
798                }
799            }
800
801            let t_resolve = t_rerank.elapsed();
802            if !resolved.is_empty() {
803                // Sort by flat_idx for sequential mmap access (better page locality)
804                resolved.sort_unstable_by_key(|&(_, flat_idx)| flat_idx);
805
806                // Batch-read raw quantized bytes into contiguous buffer
807                let t_read = std::time::Instant::now();
808                let mut raw_buf = vec![0u8; resolved.len() * vbs];
809                for (buf_idx, &(_, flat_idx)) in resolved.iter().enumerate() {
810                    let _ = lazy_flat
811                        .read_vector_raw_into(
812                            flat_idx,
813                            &mut raw_buf[buf_idx * vbs..(buf_idx + 1) * vbs],
814                        )
815                        .await;
816                }
817
818                let read_elapsed = t_read.elapsed();
819
820                // Native-precision batch SIMD cosine scoring
821                let t_score = std::time::Instant::now();
822                let mut scores = vec![0f32; resolved.len()];
823                Self::score_quantized_batch(query, &raw_buf, quant, dim, &mut scores, unit_norm);
824                let score_elapsed = t_score.elapsed();
825
826                // Write scores back to results
827                for (buf_idx, &(ri, _)) in resolved.iter().enumerate() {
828                    results[ri].2 = scores[buf_idx];
829                }
830
831                log::debug!(
832                    "[search_dense] field {}: rerank {} vectors (dim={}, quant={:?}, {}B/vec): resolve={:.1}ms read={:.1}ms score={:.1}ms",
833                    field.0,
834                    resolved.len(),
835                    dim,
836                    quant,
837                    vbs,
838                    t_resolve.as_secs_f64() * 1000.0,
839                    read_elapsed.as_secs_f64() * 1000.0,
840                    score_elapsed.as_secs_f64() * 1000.0,
841                );
842            }
843
844            if results.len() > fetch_k {
845                results.select_nth_unstable_by(fetch_k, |a, b| b.2.total_cmp(&a.2));
846                results.truncate(fetch_k);
847            }
848            results.sort_unstable_by(|a, b| b.2.total_cmp(&a.2));
849            log::debug!(
850                "[search_dense] field {}: rerank total={:.1}ms",
851                field.0,
852                t_rerank.elapsed().as_secs_f64() * 1000.0
853            );
854        }
855
856        Ok(combine_ordinal_results(results, combiner, k))
857    }
858
859    /// Check if this segment has dense vectors for the given field
860    pub fn has_dense_vector_index(&self, field: Field) -> bool {
861        self.vector_indexes.contains_key(&field.0) || self.flat_vectors.contains_key(&field.0)
862    }
863
864    /// Get the dense vector index for a field (if available)
865    pub fn get_dense_vector_index(&self, field: Field) -> Option<Arc<RaBitQIndex>> {
866        match self.vector_indexes.get(&field.0) {
867            Some(VectorIndex::RaBitQ(lazy)) => lazy.get().cloned(),
868            _ => None,
869        }
870    }
871
872    /// Get the IVF vector index for a field (if available)
873    pub fn get_ivf_vector_index(
874        &self,
875        field: Field,
876    ) -> Option<(Arc<IVFRaBitQIndex>, Arc<crate::structures::RaBitQCodebook>)> {
877        match self.vector_indexes.get(&field.0) {
878            Some(VectorIndex::IVF(lazy)) => lazy.get().map(|(i, c)| (i.clone(), c.clone())),
879            _ => None,
880        }
881    }
882
883    /// Get coarse centroids for a field
884    pub fn coarse_centroids(&self, field_id: u32) -> Option<&Arc<CoarseCentroids>> {
885        self.coarse_centroids.get(&field_id)
886    }
887
888    /// Set per-field coarse centroids from index-level trained structures
889    pub fn set_coarse_centroids(&mut self, centroids: FxHashMap<u32, Arc<CoarseCentroids>>) {
890        self.coarse_centroids = centroids;
891    }
892
893    /// Get the ScaNN vector index for a field (if available)
894    pub fn get_scann_vector_index(
895        &self,
896        field: Field,
897    ) -> Option<(Arc<IVFPQIndex>, Arc<PQCodebook>)> {
898        match self.vector_indexes.get(&field.0) {
899            Some(VectorIndex::ScaNN(lazy)) => lazy.get().map(|(i, c)| (i.clone(), c.clone())),
900            _ => None,
901        }
902    }
903
904    /// Get the vector index type for a field
905    pub fn get_vector_index(&self, field: Field) -> Option<&VectorIndex> {
906        self.vector_indexes.get(&field.0)
907    }
908
909    /// Get positions for a term (for phrase queries)
910    ///
911    /// Position offsets are now embedded in TermInfo, so we first look up
912    /// the term to get its TermInfo, then use position_info() to get the offset.
913    pub async fn get_positions(
914        &self,
915        field: Field,
916        term: &[u8],
917    ) -> Result<Option<crate::structures::PositionPostingList>> {
918        // Get positions handle
919        let handle = match &self.positions_handle {
920            Some(h) => h,
921            None => return Ok(None),
922        };
923
924        // Build key: field_id + term
925        let mut key = Vec::with_capacity(4 + term.len());
926        key.extend_from_slice(&field.0.to_le_bytes());
927        key.extend_from_slice(term);
928
929        // Look up term in dictionary to get TermInfo with position offset
930        let term_info = match self.term_dict.get(&key).await? {
931            Some(info) => info,
932            None => return Ok(None),
933        };
934
935        // Get position offset from TermInfo
936        let (offset, length) = match term_info.position_info() {
937            Some((o, l)) => (o, l),
938            None => return Ok(None),
939        };
940
941        // Read the position data
942        let slice = handle.slice(offset..offset + length);
943        let data = slice.read_bytes().await?;
944
945        // Deserialize
946        let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
947
948        Ok(Some(pos_list))
949    }
950
951    /// Check if positions are available for a field
952    pub fn has_positions(&self, field: Field) -> bool {
953        // Check schema for position mode on this field
954        if let Some(entry) = self.schema.get_field_entry(field) {
955            entry.positions.is_some()
956        } else {
957            false
958        }
959    }
960}
961
962// ── Synchronous search methods (mmap/RAM only) ─────────────────────────────
963#[cfg(feature = "sync")]
964impl SegmentReader {
965    /// Synchronous posting list lookup — requires Inline (mmap/RAM) file handles.
966    pub fn get_postings_sync(&self, field: Field, term: &[u8]) -> Result<Option<BlockPostingList>> {
967        // Build key: field_id + term
968        let mut key = Vec::with_capacity(4 + term.len());
969        key.extend_from_slice(&field.0.to_le_bytes());
970        key.extend_from_slice(term);
971
972        // Look up in term dictionary (sync)
973        let term_info = match self.term_dict.get_sync(&key)? {
974            Some(info) => info,
975            None => return Ok(None),
976        };
977
978        // Check if posting list is inlined
979        if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
980            let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
981            for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs.into_iter()) {
982                posting_list.push(doc_id, tf);
983            }
984            let block_list = BlockPostingList::from_posting_list(&posting_list)?;
985            return Ok(Some(block_list));
986        }
987
988        // External posting list — sync range read
989        let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
990            Error::Corruption("TermInfo has neither inline nor external data".to_string())
991        })?;
992
993        let start = posting_offset;
994        let end = start + posting_len;
995
996        if end > self.postings_handle.len() {
997            return Err(Error::Corruption(
998                "Posting offset out of bounds".to_string(),
999            ));
1000        }
1001
1002        let posting_bytes = self.postings_handle.read_bytes_range_sync(start..end)?;
1003        let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
1004
1005        Ok(Some(block_list))
1006    }
1007
1008    /// Synchronous position list lookup — requires Inline (mmap/RAM) file handles.
1009    pub fn get_positions_sync(
1010        &self,
1011        field: Field,
1012        term: &[u8],
1013    ) -> Result<Option<crate::structures::PositionPostingList>> {
1014        let handle = match &self.positions_handle {
1015            Some(h) => h,
1016            None => return Ok(None),
1017        };
1018
1019        // Build key: field_id + term
1020        let mut key = Vec::with_capacity(4 + term.len());
1021        key.extend_from_slice(&field.0.to_le_bytes());
1022        key.extend_from_slice(term);
1023
1024        // Look up term in dictionary (sync)
1025        let term_info = match self.term_dict.get_sync(&key)? {
1026            Some(info) => info,
1027            None => return Ok(None),
1028        };
1029
1030        let (offset, length) = match term_info.position_info() {
1031            Some((o, l)) => (o, l),
1032            None => return Ok(None),
1033        };
1034
1035        let slice = handle.slice(offset..offset + length);
1036        let data = slice.read_bytes_sync()?;
1037
1038        let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
1039        Ok(Some(pos_list))
1040    }
1041
1042    /// Synchronous dense vector search — ANN indexes are already sync,
1043    /// brute-force uses sync mmap reads.
1044    pub fn search_dense_vector_sync(
1045        &self,
1046        field: Field,
1047        query: &[f32],
1048        k: usize,
1049        nprobe: usize,
1050        rerank_factor: f32,
1051        combiner: crate::query::MultiValueCombiner,
1052    ) -> Result<Vec<VectorSearchResult>> {
1053        let ann_index = self.vector_indexes.get(&field.0);
1054        let lazy_flat = self.flat_vectors.get(&field.0);
1055
1056        if ann_index.is_none() && lazy_flat.is_none() {
1057            return Ok(Vec::new());
1058        }
1059
1060        let unit_norm = self
1061            .schema
1062            .get_field_entry(field)
1063            .and_then(|e| e.dense_vector_config.as_ref())
1064            .is_some_and(|c| c.unit_norm);
1065
1066        const BRUTE_FORCE_BATCH: usize = 4096;
1067        let fetch_k = (k as f32 * rerank_factor.max(1.0)).ceil() as usize;
1068
1069        let mut results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
1070            // ANN search (already sync)
1071            match index {
1072                VectorIndex::RaBitQ(lazy) => {
1073                    let rabitq = lazy.get().ok_or_else(|| {
1074                        Error::Schema("RaBitQ index deserialization failed".to_string())
1075                    })?;
1076                    rabitq
1077                        .search(query, fetch_k)
1078                        .into_iter()
1079                        .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
1080                        .collect()
1081                }
1082                VectorIndex::IVF(lazy) => {
1083                    let (index, codebook) = lazy.get().ok_or_else(|| {
1084                        Error::Schema("IVF index deserialization failed".to_string())
1085                    })?;
1086                    let centroids = self.coarse_centroids.get(&field.0).ok_or_else(|| {
1087                        Error::Schema(format!(
1088                            "IVF index requires coarse centroids for field {}",
1089                            field.0
1090                        ))
1091                    })?;
1092                    let effective_nprobe = if nprobe > 0 { nprobe } else { 32 };
1093                    index
1094                        .search(centroids, codebook, query, fetch_k, Some(effective_nprobe))
1095                        .into_iter()
1096                        .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
1097                        .collect()
1098                }
1099                VectorIndex::ScaNN(lazy) => {
1100                    let (index, codebook) = lazy.get().ok_or_else(|| {
1101                        Error::Schema("ScaNN index deserialization failed".to_string())
1102                    })?;
1103                    let centroids = self.coarse_centroids.get(&field.0).ok_or_else(|| {
1104                        Error::Schema(format!(
1105                            "ScaNN index requires coarse centroids for field {}",
1106                            field.0
1107                        ))
1108                    })?;
1109                    let effective_nprobe = if nprobe > 0 { nprobe } else { 32 };
1110                    index
1111                        .search(centroids, codebook, query, fetch_k, Some(effective_nprobe))
1112                        .into_iter()
1113                        .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
1114                        .collect()
1115                }
1116            }
1117        } else if let Some(lazy_flat) = lazy_flat {
1118            // Batched brute-force (sync mmap reads)
1119            let dim = lazy_flat.dim;
1120            let n = lazy_flat.num_vectors;
1121            let quant = lazy_flat.quantization;
1122            let mut collector = crate::query::ScoreCollector::new(fetch_k);
1123            let mut scores = vec![0f32; BRUTE_FORCE_BATCH];
1124
1125            for batch_start in (0..n).step_by(BRUTE_FORCE_BATCH) {
1126                let batch_count = BRUTE_FORCE_BATCH.min(n - batch_start);
1127                let batch_bytes = lazy_flat
1128                    .read_vectors_batch_sync(batch_start, batch_count)
1129                    .map_err(crate::Error::Io)?;
1130                let raw = batch_bytes.as_slice();
1131
1132                Self::score_quantized_batch(
1133                    query,
1134                    raw,
1135                    quant,
1136                    dim,
1137                    &mut scores[..batch_count],
1138                    unit_norm,
1139                );
1140
1141                for (i, &score) in scores.iter().enumerate().take(batch_count) {
1142                    let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
1143                    collector.insert_with_ordinal(doc_id, score, ordinal);
1144                }
1145            }
1146
1147            collector
1148                .into_sorted_results()
1149                .into_iter()
1150                .map(|(doc_id, score, ordinal)| (doc_id, ordinal, score))
1151                .collect()
1152        } else {
1153            return Ok(Vec::new());
1154        };
1155
1156        // Rerank ANN candidates using raw vectors (sync)
1157        if ann_index.is_some()
1158            && !results.is_empty()
1159            && let Some(lazy_flat) = lazy_flat
1160        {
1161            let dim = lazy_flat.dim;
1162            let quant = lazy_flat.quantization;
1163            let vbs = lazy_flat.vector_byte_size();
1164
1165            let mut resolved: Vec<(usize, usize)> = Vec::new();
1166            for (ri, c) in results.iter().enumerate() {
1167                let (start, entries) = lazy_flat.flat_indexes_for_doc(c.0);
1168                for (j, &(_, ord)) in entries.iter().enumerate() {
1169                    if ord == c.1 {
1170                        resolved.push((ri, start + j));
1171                        break;
1172                    }
1173                }
1174            }
1175
1176            if !resolved.is_empty() {
1177                resolved.sort_unstable_by_key(|&(_, flat_idx)| flat_idx);
1178                let mut raw_buf = vec![0u8; resolved.len() * vbs];
1179                for (buf_idx, &(_, flat_idx)) in resolved.iter().enumerate() {
1180                    let _ = lazy_flat.read_vector_raw_into_sync(
1181                        flat_idx,
1182                        &mut raw_buf[buf_idx * vbs..(buf_idx + 1) * vbs],
1183                    );
1184                }
1185
1186                let mut scores = vec![0f32; resolved.len()];
1187                Self::score_quantized_batch(query, &raw_buf, quant, dim, &mut scores, unit_norm);
1188
1189                for (buf_idx, &(ri, _)) in resolved.iter().enumerate() {
1190                    results[ri].2 = scores[buf_idx];
1191                }
1192            }
1193
1194            if results.len() > fetch_k {
1195                results.select_nth_unstable_by(fetch_k, |a, b| b.2.total_cmp(&a.2));
1196                results.truncate(fetch_k);
1197            }
1198            results.sort_unstable_by(|a, b| b.2.total_cmp(&a.2));
1199        }
1200
1201        Ok(combine_ordinal_results(results, combiner, k))
1202    }
1203}