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    /// Shared coarse centroids for IVF search (loaded once)
77    coarse_centroids: Option<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        let coarse_centroids = vectors_data.coarse_centroids;
117
118        // Load sparse vector indexes from .sparse file
119        let sparse_indexes = loader::load_sparse_file(dir, &files, meta.num_docs, &schema).await?;
120
121        // Open positions file handle (if exists) - offsets are now in TermInfo
122        let positions_handle = loader::open_positions_file(dir, &files, &schema).await?;
123
124        // Log segment loading stats (compact format: ~24 bytes per active dim in hashmap)
125        let sparse_dims: usize = sparse_indexes.values().map(|s| s.num_dimensions()).sum();
126        let sparse_mem = sparse_dims * 24; // HashMap entry overhead
127        log::debug!(
128            "[segment] loaded {:016x}: docs={}, sparse_dims={}, sparse_mem={:.2} KB, vectors={}",
129            segment_id.0,
130            meta.num_docs,
131            sparse_dims,
132            sparse_mem as f64 / 1024.0,
133            vector_indexes.len()
134        );
135
136        Ok(Self {
137            meta,
138            term_dict: Arc::new(term_dict),
139            postings_handle,
140            store: Arc::new(store),
141            schema,
142            doc_id_offset,
143            vector_indexes,
144            flat_vectors,
145            coarse_centroids,
146            sparse_indexes,
147            positions_handle,
148        })
149    }
150
151    pub fn meta(&self) -> &SegmentMeta {
152        &self.meta
153    }
154
155    pub fn num_docs(&self) -> u32 {
156        self.meta.num_docs
157    }
158
159    /// Get average field length for BM25F scoring
160    pub fn avg_field_len(&self, field: Field) -> f32 {
161        self.meta.avg_field_len(field)
162    }
163
164    pub fn doc_id_offset(&self) -> DocId {
165        self.doc_id_offset
166    }
167
168    /// Set the doc_id_offset (used for parallel segment loading)
169    pub fn set_doc_id_offset(&mut self, offset: DocId) {
170        self.doc_id_offset = offset;
171    }
172
173    pub fn schema(&self) -> &Schema {
174        &self.schema
175    }
176
177    /// Get sparse indexes for all fields
178    pub fn sparse_indexes(&self) -> &FxHashMap<u32, SparseIndex> {
179        &self.sparse_indexes
180    }
181
182    /// Get vector indexes for all fields
183    pub fn vector_indexes(&self) -> &FxHashMap<u32, VectorIndex> {
184        &self.vector_indexes
185    }
186
187    /// Get lazy flat vectors for all fields (for reranking and merge)
188    pub fn flat_vectors(&self) -> &FxHashMap<u32, LazyFlatVectorData> {
189        &self.flat_vectors
190    }
191
192    /// Get term dictionary stats for debugging
193    pub fn term_dict_stats(&self) -> SSTableStats {
194        self.term_dict.stats()
195    }
196
197    /// Estimate memory usage of this segment reader
198    pub fn memory_stats(&self) -> SegmentMemoryStats {
199        let term_dict_stats = self.term_dict.stats();
200
201        // Term dict cache: num_blocks * avg_block_size (estimate 4KB per cached block)
202        let term_dict_cache_bytes = self.term_dict.cached_blocks() * 4096;
203
204        // Store cache: similar estimate
205        let store_cache_bytes = self.store.cached_blocks() * 4096;
206
207        // Sparse index: each dimension has a posting list in memory
208        // Estimate: ~24 bytes per active dimension (HashMap entry overhead)
209        let sparse_index_bytes: usize = self
210            .sparse_indexes
211            .values()
212            .map(|s| s.num_dimensions() * 24)
213            .sum();
214
215        // Dense index: vectors are memory-mapped, but we track index structures
216        // RaBitQ/IVF indexes have cluster assignments in memory
217        let dense_index_bytes: usize = self
218            .vector_indexes
219            .values()
220            .map(|v| v.estimated_memory_bytes())
221            .sum();
222
223        SegmentMemoryStats {
224            segment_id: self.meta.id,
225            num_docs: self.meta.num_docs,
226            term_dict_cache_bytes,
227            store_cache_bytes,
228            sparse_index_bytes,
229            dense_index_bytes,
230            bloom_filter_bytes: term_dict_stats.bloom_filter_size,
231        }
232    }
233
234    /// Get posting list for a term (async - loads on demand)
235    ///
236    /// For small posting lists (1-3 docs), the data is inlined in the term dictionary
237    /// and no additional I/O is needed. For larger lists, reads from .post file.
238    pub async fn get_postings(
239        &self,
240        field: Field,
241        term: &[u8],
242    ) -> Result<Option<BlockPostingList>> {
243        log::debug!(
244            "SegmentReader::get_postings field={} term_len={}",
245            field.0,
246            term.len()
247        );
248
249        // Build key: field_id + term
250        let mut key = Vec::with_capacity(4 + term.len());
251        key.extend_from_slice(&field.0.to_le_bytes());
252        key.extend_from_slice(term);
253
254        // Look up in term dictionary
255        let term_info = match self.term_dict.get(&key).await? {
256            Some(info) => {
257                log::debug!("SegmentReader::get_postings found term_info");
258                info
259            }
260            None => {
261                log::debug!("SegmentReader::get_postings term not found");
262                return Ok(None);
263            }
264        };
265
266        // Check if posting list is inlined
267        if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
268            // Build BlockPostingList from inline data (no I/O needed!)
269            let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
270            for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs.into_iter()) {
271                posting_list.push(doc_id, tf);
272            }
273            let block_list = BlockPostingList::from_posting_list(&posting_list)?;
274            return Ok(Some(block_list));
275        }
276
277        // External posting list - read from postings file handle (lazy - HTTP range request)
278        let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
279            Error::Corruption("TermInfo has neither inline nor external data".to_string())
280        })?;
281
282        let start = posting_offset;
283        let end = start + posting_len as u64;
284
285        if end > self.postings_handle.len() {
286            return Err(Error::Corruption(
287                "Posting offset out of bounds".to_string(),
288            ));
289        }
290
291        let posting_bytes = self.postings_handle.read_bytes_range(start..end).await?;
292        let block_list = BlockPostingList::deserialize(&mut posting_bytes.as_slice())?;
293
294        Ok(Some(block_list))
295    }
296
297    /// Get document by local doc_id (async - loads on demand).
298    ///
299    /// Dense vector fields are hydrated from LazyFlatVectorData (not stored in .store).
300    /// Uses binary search on sorted doc_ids for O(log N) lookup.
301    pub async fn doc(&self, local_doc_id: DocId) -> Result<Option<Document>> {
302        let mut doc = match self.store.get(local_doc_id, &self.schema).await {
303            Ok(Some(d)) => d,
304            Ok(None) => return Ok(None),
305            Err(e) => return Err(Error::from(e)),
306        };
307
308        // Hydrate dense vector fields from flat vector data
309        for (&field_id, lazy_flat) in &self.flat_vectors {
310            let (start, entries) = lazy_flat.flat_indexes_for_doc(local_doc_id);
311            for (j, &(_doc_id, _ordinal)) in entries.iter().enumerate() {
312                let flat_idx = start + j;
313                match lazy_flat.get_vector(flat_idx).await {
314                    Ok(vec) => {
315                        doc.add_dense_vector(Field(field_id), vec);
316                    }
317                    Err(e) => {
318                        log::warn!("Failed to hydrate vector field {}: {}", field_id, e);
319                    }
320                }
321            }
322        }
323
324        Ok(Some(doc))
325    }
326
327    /// Prefetch term dictionary blocks for a key range
328    pub async fn prefetch_terms(
329        &self,
330        field: Field,
331        start_term: &[u8],
332        end_term: &[u8],
333    ) -> Result<()> {
334        let mut start_key = Vec::with_capacity(4 + start_term.len());
335        start_key.extend_from_slice(&field.0.to_le_bytes());
336        start_key.extend_from_slice(start_term);
337
338        let mut end_key = Vec::with_capacity(4 + end_term.len());
339        end_key.extend_from_slice(&field.0.to_le_bytes());
340        end_key.extend_from_slice(end_term);
341
342        self.term_dict.prefetch_range(&start_key, &end_key).await?;
343        Ok(())
344    }
345
346    /// Check if store uses dictionary compression (incompatible with raw merging)
347    pub fn store_has_dict(&self) -> bool {
348        self.store.has_dict()
349    }
350
351    /// Get store reference for merge operations
352    pub fn store(&self) -> &super::store::AsyncStoreReader {
353        &self.store
354    }
355
356    /// Get raw store blocks for optimized merging
357    pub fn store_raw_blocks(&self) -> Vec<RawStoreBlock> {
358        self.store.raw_blocks()
359    }
360
361    /// Get store data slice for raw block access
362    pub fn store_data_slice(&self) -> &LazyFileSlice {
363        self.store.data_slice()
364    }
365
366    /// Get all terms from this segment (for merge)
367    pub async fn all_terms(&self) -> Result<Vec<(Vec<u8>, TermInfo)>> {
368        self.term_dict.all_entries().await.map_err(Error::from)
369    }
370
371    /// Get all terms with parsed field and term string (for statistics aggregation)
372    ///
373    /// Returns (field, term_string, doc_freq) for each term in the dictionary.
374    /// Skips terms that aren't valid UTF-8.
375    pub async fn all_terms_with_stats(&self) -> Result<Vec<(Field, String, u32)>> {
376        let entries = self.term_dict.all_entries().await?;
377        let mut result = Vec::with_capacity(entries.len());
378
379        for (key, term_info) in entries {
380            // Key format: field_id (4 bytes little-endian) + term bytes
381            if key.len() > 4 {
382                let field_id = u32::from_le_bytes([key[0], key[1], key[2], key[3]]);
383                let term_bytes = &key[4..];
384                if let Ok(term_str) = std::str::from_utf8(term_bytes) {
385                    result.push((Field(field_id), term_str.to_string(), term_info.doc_freq()));
386                }
387            }
388        }
389
390        Ok(result)
391    }
392
393    /// Get streaming iterator over term dictionary (for memory-efficient merge)
394    pub fn term_dict_iter(&self) -> crate::structures::AsyncSSTableIterator<'_, TermInfo> {
395        self.term_dict.iter()
396    }
397
398    /// Prefetch all term dictionary blocks in a single bulk I/O call.
399    ///
400    /// Call before merge iteration to eliminate per-block cache misses.
401    pub async fn prefetch_term_dict(&self) -> crate::Result<()> {
402        self.term_dict
403            .prefetch_all_data_bulk()
404            .await
405            .map_err(crate::Error::from)
406    }
407
408    /// Read raw posting bytes at offset
409    pub async fn read_postings(&self, offset: u64, len: u32) -> Result<Vec<u8>> {
410        let start = offset;
411        let end = start + len as u64;
412        let bytes = self.postings_handle.read_bytes_range(start..end).await?;
413        Ok(bytes.to_vec())
414    }
415
416    /// Read raw position bytes at offset (for merge)
417    pub async fn read_position_bytes(&self, offset: u64, len: u32) -> Result<Option<Vec<u8>>> {
418        let handle = match &self.positions_handle {
419            Some(h) => h,
420            None => return Ok(None),
421        };
422        let start = offset;
423        let end = start + len as u64;
424        let bytes = handle.read_bytes_range(start..end).await?;
425        Ok(Some(bytes.to_vec()))
426    }
427
428    /// Check if this segment has a positions file
429    pub fn has_positions_file(&self) -> bool {
430        self.positions_handle.is_some()
431    }
432
433    /// Search dense vectors using RaBitQ
434    ///
435    /// Returns VectorSearchResult with ordinal tracking for multi-value fields.
436    /// The doc_ids are adjusted by doc_id_offset for this segment.
437    /// For multi-valued documents, scores are combined using the specified combiner.
438    pub async fn search_dense_vector(
439        &self,
440        field: Field,
441        query: &[f32],
442        k: usize,
443        nprobe: usize,
444        rerank_factor: usize,
445        combiner: crate::query::MultiValueCombiner,
446    ) -> Result<Vec<VectorSearchResult>> {
447        let ann_index = self.vector_indexes.get(&field.0);
448        let lazy_flat = self.flat_vectors.get(&field.0);
449
450        // No vectors at all for this field
451        if ann_index.is_none() && lazy_flat.is_none() {
452            return Ok(Vec::new());
453        }
454
455        /// Batch size for brute-force scoring (4096 vectors × 768 dims × 4 bytes ≈ 12MB)
456        const BRUTE_FORCE_BATCH: usize = 4096;
457
458        // Results are (doc_id, ordinal, score) where score = similarity (higher = better)
459        let mut results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
460            // ANN search (RaBitQ, IVF, ScaNN)
461            match index {
462                VectorIndex::RaBitQ(rabitq) => {
463                    let fetch_k = k * rerank_factor.max(1);
464                    rabitq
465                        .search(query, fetch_k, rerank_factor)
466                        .into_iter()
467                        .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
468                        .collect()
469                }
470                VectorIndex::IVF { index, codebook } => {
471                    let centroids = self.coarse_centroids.as_ref().ok_or_else(|| {
472                        Error::Schema("IVF index requires coarse centroids".to_string())
473                    })?;
474                    let effective_nprobe = if nprobe > 0 { nprobe } else { 32 };
475                    let fetch_k = k * rerank_factor.max(1);
476                    index
477                        .search(centroids, codebook, query, fetch_k, Some(effective_nprobe))
478                        .into_iter()
479                        .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
480                        .collect()
481                }
482                VectorIndex::ScaNN { index, codebook } => {
483                    let centroids = self.coarse_centroids.as_ref().ok_or_else(|| {
484                        Error::Schema("ScaNN index requires coarse centroids".to_string())
485                    })?;
486                    let effective_nprobe = if nprobe > 0 { nprobe } else { 32 };
487                    let fetch_k = k * rerank_factor.max(1);
488                    index
489                        .search(centroids, codebook, query, fetch_k, Some(effective_nprobe))
490                        .into_iter()
491                        .map(|(doc_id, ordinal, dist)| (doc_id, ordinal, 1.0 / (1.0 + dist)))
492                        .collect()
493                }
494            }
495        } else if let Some(lazy_flat) = lazy_flat {
496            // Batched brute-force from lazy flat vectors (native-precision SIMD scoring)
497            let dim = lazy_flat.dim;
498            let n = lazy_flat.num_vectors;
499            let quant = lazy_flat.quantization;
500            let fetch_k = k * rerank_factor.max(1);
501            let mut candidates: Vec<(u32, u16, f32)> = Vec::new();
502
503            for batch_start in (0..n).step_by(BRUTE_FORCE_BATCH) {
504                let batch_count = BRUTE_FORCE_BATCH.min(n - batch_start);
505                let batch_bytes = lazy_flat
506                    .read_vectors_batch(batch_start, batch_count)
507                    .await
508                    .map_err(crate::Error::Io)?;
509                let raw = batch_bytes.as_slice();
510
511                let mut scores = vec![0f32; batch_count];
512
513                match quant {
514                    crate::dsl::DenseVectorQuantization::F32 => {
515                        let batch_floats = batch_count * dim;
516                        let mut aligned_buf: Vec<f32> = Vec::new();
517                        let vectors: &[f32] = if (raw.as_ptr() as usize)
518                            .is_multiple_of(std::mem::align_of::<f32>())
519                        {
520                            unsafe {
521                                std::slice::from_raw_parts(raw.as_ptr() as *const f32, batch_floats)
522                            }
523                        } else {
524                            aligned_buf.resize(batch_floats, 0.0);
525                            unsafe {
526                                std::ptr::copy_nonoverlapping(
527                                    raw.as_ptr(),
528                                    aligned_buf.as_mut_ptr() as *mut u8,
529                                    batch_floats * std::mem::size_of::<f32>(),
530                                );
531                            }
532                            &aligned_buf
533                        };
534                        crate::structures::simd::batch_cosine_scores(
535                            query,
536                            vectors,
537                            dim,
538                            &mut scores,
539                        );
540                    }
541                    crate::dsl::DenseVectorQuantization::F16 => {
542                        crate::structures::simd::batch_cosine_scores_f16(
543                            query,
544                            raw,
545                            dim,
546                            &mut scores,
547                        );
548                    }
549                    crate::dsl::DenseVectorQuantization::UInt8 => {
550                        crate::structures::simd::batch_cosine_scores_u8(
551                            query,
552                            raw,
553                            dim,
554                            &mut scores,
555                        );
556                    }
557                }
558
559                for (i, &score) in scores.iter().enumerate().take(batch_count) {
560                    let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
561                    candidates.push((doc_id, ordinal, score));
562                }
563            }
564
565            candidates.sort_by(|a, b| b.2.partial_cmp(&a.2).unwrap_or(std::cmp::Ordering::Equal));
566            candidates.truncate(fetch_k);
567            candidates
568        } else {
569            return Ok(Vec::new());
570        };
571
572        // Rerank ANN candidates using raw vectors from lazy flat (binary search lookup)
573        // Uses native-precision SIMD scoring on quantized bytes — no dequantization overhead.
574        if ann_index.is_some()
575            && !results.is_empty()
576            && let Some(lazy_flat) = lazy_flat
577        {
578            let dim = lazy_flat.dim;
579            let quant = lazy_flat.quantization;
580            let vbs = lazy_flat.vector_byte_size();
581
582            // Resolve flat indexes for each candidate via binary search
583            let mut resolved: Vec<(usize, usize)> = Vec::new(); // (result_idx, flat_idx)
584            for (ri, c) in results.iter().enumerate() {
585                let (start, entries) = lazy_flat.flat_indexes_for_doc(c.0);
586                for (j, &(_, ord)) in entries.iter().enumerate() {
587                    if ord == c.1 {
588                        resolved.push((ri, start + j));
589                        break;
590                    }
591                }
592            }
593
594            if !resolved.is_empty() {
595                // Batch-read raw quantized bytes into contiguous buffer
596                let mut raw_buf = vec![0u8; resolved.len() * vbs];
597                for (buf_idx, &(_, flat_idx)) in resolved.iter().enumerate() {
598                    let _ = lazy_flat
599                        .read_vector_raw_into(
600                            flat_idx,
601                            &mut raw_buf[buf_idx * vbs..(buf_idx + 1) * vbs],
602                        )
603                        .await;
604                }
605
606                // Native-precision batch SIMD cosine scoring
607                let mut scores = vec![0f32; resolved.len()];
608                match quant {
609                    crate::dsl::DenseVectorQuantization::F32 => {
610                        let floats = resolved.len() * dim;
611                        let mut aligned_buf: Vec<f32> = Vec::new();
612                        let vectors: &[f32] = if (raw_buf.as_ptr() as usize)
613                            .is_multiple_of(std::mem::align_of::<f32>())
614                        {
615                            unsafe {
616                                std::slice::from_raw_parts(raw_buf.as_ptr() as *const f32, floats)
617                            }
618                        } else {
619                            aligned_buf.resize(floats, 0.0);
620                            unsafe {
621                                std::ptr::copy_nonoverlapping(
622                                    raw_buf.as_ptr(),
623                                    aligned_buf.as_mut_ptr() as *mut u8,
624                                    floats * std::mem::size_of::<f32>(),
625                                );
626                            }
627                            &aligned_buf
628                        };
629                        crate::structures::simd::batch_cosine_scores(
630                            query,
631                            vectors,
632                            dim,
633                            &mut scores,
634                        );
635                    }
636                    crate::dsl::DenseVectorQuantization::F16 => {
637                        crate::structures::simd::batch_cosine_scores_f16(
638                            query,
639                            &raw_buf,
640                            dim,
641                            &mut scores,
642                        );
643                    }
644                    crate::dsl::DenseVectorQuantization::UInt8 => {
645                        crate::structures::simd::batch_cosine_scores_u8(
646                            query,
647                            &raw_buf,
648                            dim,
649                            &mut scores,
650                        );
651                    }
652                }
653
654                // Write scores back to results
655                for (buf_idx, &(ri, _)) in resolved.iter().enumerate() {
656                    results[ri].2 = scores[buf_idx];
657                }
658            }
659
660            results.sort_by(|a, b| b.2.partial_cmp(&a.2).unwrap_or(std::cmp::Ordering::Equal));
661            results.truncate(k * rerank_factor.max(1));
662        }
663
664        // Track ordinals with individual scores for each doc_id
665        // Note: doc_id_offset is NOT applied here - the collector applies it uniformly
666        let mut doc_ordinals: rustc_hash::FxHashMap<DocId, Vec<(u32, f32)>> =
667            rustc_hash::FxHashMap::default();
668        for (doc_id, ordinal, score) in results {
669            let ordinals = doc_ordinals.entry(doc_id as DocId).or_default();
670            ordinals.push((ordinal as u32, score));
671        }
672
673        // Combine scores and build results with ordinal tracking
674        let mut final_results: Vec<VectorSearchResult> = doc_ordinals
675            .into_iter()
676            .map(|(doc_id, ordinals)| {
677                let combined_score = combiner.combine(&ordinals);
678                VectorSearchResult::new(doc_id, combined_score, ordinals)
679            })
680            .collect();
681
682        // Sort by score descending and take top k
683        final_results.sort_by(|a, b| {
684            b.score
685                .partial_cmp(&a.score)
686                .unwrap_or(std::cmp::Ordering::Equal)
687        });
688        final_results.truncate(k);
689
690        Ok(final_results)
691    }
692
693    /// Check if this segment has dense vectors for the given field
694    pub fn has_dense_vector_index(&self, field: Field) -> bool {
695        self.vector_indexes.contains_key(&field.0) || self.flat_vectors.contains_key(&field.0)
696    }
697
698    /// Get the dense vector index for a field (if available)
699    pub fn get_dense_vector_index(&self, field: Field) -> Option<Arc<RaBitQIndex>> {
700        match self.vector_indexes.get(&field.0) {
701            Some(VectorIndex::RaBitQ(idx)) => Some(idx.clone()),
702            _ => None,
703        }
704    }
705
706    /// Get the IVF vector index for a field (if available)
707    pub fn get_ivf_vector_index(
708        &self,
709        field: Field,
710    ) -> Option<(Arc<IVFRaBitQIndex>, Arc<crate::structures::RaBitQCodebook>)> {
711        match self.vector_indexes.get(&field.0) {
712            Some(VectorIndex::IVF { index, codebook }) => Some((index.clone(), codebook.clone())),
713            _ => None,
714        }
715    }
716
717    /// Get coarse centroids (shared across IVF/ScaNN indexes)
718    pub fn coarse_centroids(&self) -> Option<&Arc<CoarseCentroids>> {
719        self.coarse_centroids.as_ref()
720    }
721
722    /// Get the ScaNN vector index for a field (if available)
723    pub fn get_scann_vector_index(
724        &self,
725        field: Field,
726    ) -> Option<(Arc<IVFPQIndex>, Arc<PQCodebook>)> {
727        match self.vector_indexes.get(&field.0) {
728            Some(VectorIndex::ScaNN { index, codebook }) => Some((index.clone(), codebook.clone())),
729            _ => None,
730        }
731    }
732
733    /// Get the vector index type for a field
734    pub fn get_vector_index(&self, field: Field) -> Option<&VectorIndex> {
735        self.vector_indexes.get(&field.0)
736    }
737
738    /// Search for similar sparse vectors using dedicated sparse posting lists
739    ///
740    /// Uses shared `WandExecutor` with `SparseTermScorer` for efficient top-k retrieval.
741    /// Optimizations (via WandExecutor):
742    /// 1. **MaxScore pruning**: Dimensions sorted by max contribution
743    /// 2. **Block-Max WAND**: Skips blocks where max contribution < threshold
744    /// 3. **Top-K heap**: Efficient score collection
745    ///
746    /// Returns VectorSearchResult with ordinal tracking for multi-value fields.
747    pub async fn search_sparse_vector(
748        &self,
749        field: Field,
750        vector: &[(u32, f32)],
751        limit: usize,
752        combiner: crate::query::MultiValueCombiner,
753        heap_factor: f32,
754    ) -> Result<Vec<VectorSearchResult>> {
755        use crate::query::{BlockMaxScoreExecutor, BmpExecutor, SparseTermScorer};
756
757        let query_tokens = vector.len();
758
759        // Get sparse index for this field
760        let sparse_index = match self.sparse_indexes.get(&field.0) {
761            Some(idx) => idx,
762            None => {
763                log::debug!(
764                    "Sparse vector search: no index for field {}, returning empty",
765                    field.0
766                );
767                return Ok(Vec::new());
768            }
769        };
770
771        let index_dimensions = sparse_index.num_dimensions();
772
773        // Build scorers for each dimension that exists in the index
774        // Load posting lists on-demand (lazy loading via mmap)
775        // Keep Arc references alive for the duration of scoring
776        let mut matched_tokens = Vec::new();
777        let mut missing_tokens = Vec::new();
778        let mut posting_lists: Vec<(u32, f32, Arc<BlockSparsePostingList>)> =
779            Vec::with_capacity(vector.len());
780
781        for &(dim_id, query_weight) in vector {
782            // Check if dimension exists before loading
783            if !sparse_index.has_dimension(dim_id) {
784                missing_tokens.push(dim_id);
785                continue;
786            }
787
788            // Load posting list on-demand (async, uses mmap)
789            match sparse_index.get_posting(dim_id).await? {
790                Some(pl) => {
791                    matched_tokens.push(dim_id);
792                    posting_lists.push((dim_id, query_weight, pl));
793                }
794                None => {
795                    missing_tokens.push(dim_id);
796                }
797            }
798        }
799
800        // Create scorers from the loaded posting lists (borrows from posting_lists)
801        let scorers: Vec<SparseTermScorer> = posting_lists
802            .iter()
803            .map(|(_, query_weight, pl)| SparseTermScorer::from_arc(pl, *query_weight))
804            .collect();
805
806        log::debug!(
807            "Sparse vector search: query_tokens={}, matched={}, missing={}, index_dimensions={}",
808            query_tokens,
809            matched_tokens.len(),
810            missing_tokens.len(),
811            index_dimensions
812        );
813
814        // Log query tokens with their IDs and weights
815        if log::log_enabled!(log::Level::Debug) {
816            let query_details: Vec<_> = vector
817                .iter()
818                .take(30)
819                .map(|(id, w)| format!("{}:{:.3}", id, w))
820                .collect();
821            log::debug!("Query tokens (id:weight): [{}]", query_details.join(", "));
822        }
823
824        if !matched_tokens.is_empty() {
825            log::debug!(
826                "Matched token IDs: {:?}",
827                matched_tokens.iter().take(20).collect::<Vec<_>>()
828            );
829        }
830
831        if !missing_tokens.is_empty() {
832            log::debug!(
833                "Missing token IDs (not in index): {:?}",
834                missing_tokens.iter().take(20).collect::<Vec<_>>()
835            );
836        }
837
838        if scorers.is_empty() {
839            log::debug!("Sparse vector search: no matching tokens, returning empty");
840            return Ok(Vec::new());
841        }
842
843        // Select executor based on number of query terms:
844        // - 12+ terms: BMP (block-at-a-time, best for SPLADE expansions)
845        // - 1-11 terms: BlockMaxScoreExecutor (unified MaxScore + block-max + conjunction)
846        let num_terms = scorers.len();
847        let over_fetch = limit * 2; // Over-fetch for multi-value combining
848        let raw_results = if num_terms > 12 {
849            // BMP: use posting lists directly (not iterators)
850            let pl_refs: Vec<_> = posting_lists
851                .iter()
852                .map(|(_, _, pl)| Arc::clone(pl))
853                .collect();
854            let weights: Vec<_> = posting_lists.iter().map(|(_, qw, _)| *qw).collect();
855            drop(scorers); // Release borrowing iterators before using posting_lists
856            BmpExecutor::new(pl_refs, weights, over_fetch, heap_factor).execute()
857        } else {
858            BlockMaxScoreExecutor::with_heap_factor(scorers, over_fetch, heap_factor).execute()
859        };
860
861        log::trace!(
862            "Sparse WAND returned {} raw results for segment (doc_id_offset={})",
863            raw_results.len(),
864            self.doc_id_offset
865        );
866        if log::log_enabled!(log::Level::Trace) && !raw_results.is_empty() {
867            for r in raw_results.iter().take(5) {
868                log::trace!(
869                    "  Raw result: doc_id={} (global={}), score={:.4}, ordinal={}",
870                    r.doc_id,
871                    r.doc_id + self.doc_id_offset,
872                    r.score,
873                    r.ordinal
874                );
875            }
876        }
877
878        // Track ordinals with individual scores for each doc_id
879        // Now using real ordinals from the posting lists
880        let mut doc_ordinals: rustc_hash::FxHashMap<u32, Vec<(u32, f32)>> =
881            rustc_hash::FxHashMap::default();
882        for r in raw_results {
883            let ordinals = doc_ordinals.entry(r.doc_id).or_default();
884            ordinals.push((r.ordinal as u32, r.score));
885        }
886
887        // Combine scores and build results with ordinal tracking
888        // Note: doc_id_offset is NOT applied here - the collector applies it uniformly
889        let mut results: Vec<VectorSearchResult> = doc_ordinals
890            .into_iter()
891            .map(|(doc_id, ordinals)| {
892                let combined_score = combiner.combine(&ordinals);
893                VectorSearchResult::new(doc_id, combined_score, ordinals)
894            })
895            .collect();
896
897        // Sort by score descending and take top limit
898        results.sort_by(|a, b| {
899            b.score
900                .partial_cmp(&a.score)
901                .unwrap_or(std::cmp::Ordering::Equal)
902        });
903        results.truncate(limit);
904
905        Ok(results)
906    }
907
908    /// Get positions for a term (for phrase queries)
909    ///
910    /// Position offsets are now embedded in TermInfo, so we first look up
911    /// the term to get its TermInfo, then use position_info() to get the offset.
912    pub async fn get_positions(
913        &self,
914        field: Field,
915        term: &[u8],
916    ) -> Result<Option<crate::structures::PositionPostingList>> {
917        use std::io::Cursor;
918
919        // Get positions handle
920        let handle = match &self.positions_handle {
921            Some(h) => h,
922            None => return Ok(None),
923        };
924
925        // Build key: field_id + term
926        let mut key = Vec::with_capacity(4 + term.len());
927        key.extend_from_slice(&field.0.to_le_bytes());
928        key.extend_from_slice(term);
929
930        // Look up term in dictionary to get TermInfo with position offset
931        let term_info = match self.term_dict.get(&key).await? {
932            Some(info) => info,
933            None => return Ok(None),
934        };
935
936        // Get position offset from TermInfo
937        let (offset, length) = match term_info.position_info() {
938            Some((o, l)) => (o, l),
939            None => return Ok(None),
940        };
941
942        // Read the position data
943        let slice = handle.slice(offset..offset + length as u64);
944        let data = slice.read_bytes().await?;
945
946        // Deserialize
947        let mut cursor = Cursor::new(data.as_slice());
948        let pos_list = crate::structures::PositionPostingList::deserialize(&mut cursor)?;
949
950        Ok(Some(pos_list))
951    }
952
953    /// Check if positions are available for a field
954    pub fn has_positions(&self, field: Field) -> bool {
955        // Check schema for position mode on this field
956        if let Some(entry) = self.schema.get_field_entry(field) {
957            entry.positions.is_some()
958        } else {
959            false
960        }
961    }
962}
963
964/// Alias for AsyncSegmentReader
965pub type SegmentReader = AsyncSegmentReader;