Skip to main content

hermes_core/segment/reader/
mod.rs

1//! Async segment reader with lazy loading
2
3pub(crate) mod bmp;
4pub(crate) mod loader;
5mod types;
6
7pub use bmp::{BmpDimStats, BmpIndex};
8#[cfg(feature = "native")]
9pub(crate) use types::DimRawData;
10pub use types::{SparseIndex, VectorIndex, VectorSearchResult};
11
12/// Bound vocabulary and posting expansion before a prefix query starts loading
13/// posting payloads. These are per-segment limits; callers should use exact-term
14/// or a more selective prefix when they are exceeded.
15const MAX_PREFIX_TERMS: usize = 1_024;
16const MAX_PREFIX_POSTINGS: u64 = 5_000_000;
17/// Hard guard for explicitly requested dense candidate documents. Values of
18/// those documents are exact-scored through bounded streaming batches, so a
19/// valid multi-valued document is not rejected merely for owning many values.
20const MAX_DENSE_CANDIDATES_PER_SEGMENT: usize = 20_000;
21/// Preferred vector count; wide vectors reduce it to stay under the byte cap.
22const DENSE_SCORE_BATCH: usize = 4_096;
23const BINARY_SCORE_BATCH: usize = 8_192;
24const MAX_VECTOR_SCORE_BATCH_BYTES: usize = 8 * 1024 * 1024;
25
26/// Runtime memory accounting for a single segment.
27///
28/// Heap, file-backed address space, and pinned residency are deliberately
29/// separate: file-backed bytes are not resident merely because they are
30/// mapped, and pinned bytes are a subset rather than an additive allocation.
31#[derive(Debug, Clone, Default)]
32pub struct SegmentMemoryStats {
33    /// Segment ID
34    pub segment_id: u128,
35    /// Number of documents in segment
36    pub num_docs: u32,
37    /// Term dictionary block cache bytes
38    pub term_dict_cache_bytes: usize,
39    /// Document store block cache bytes
40    pub store_cache_bytes: usize,
41    /// Sparse-vector lookup structures retained on the heap.
42    pub sparse_heap_bytes: usize,
43    /// Dense-vector ANN lookup structures retained on the heap.
44    pub dense_heap_bytes: usize,
45    /// File-backed term-dictionary bloom-filter bytes.
46    pub term_bloom_file_bytes: u64,
47    /// Logical `.sparse` file bytes retained by the reader.
48    pub sparse_file_backed_bytes: u64,
49    /// Logical `.vectors` file bytes retained by the reader.
50    pub dense_file_backed_bytes: u64,
51    /// Hot metadata bytes actually pinned (mlock/heap-copy) at open
52    pub pinned_metadata_bytes: u64,
53    /// Hot metadata bytes eligible for pinning (gap vs pinned = budget
54    /// exhausted or mlock failures — operator-visible)
55    pub pin_intended_bytes: u64,
56    /// Sparse-vector subset of `pinned_metadata_bytes`.
57    pub sparse_pinned_metadata_bytes: u64,
58    /// Sparse-vector bytes eligible for pinning.
59    pub sparse_pin_intended_bytes: u64,
60    /// Dense-vector subset of `pinned_metadata_bytes`.
61    pub dense_pinned_metadata_bytes: u64,
62    /// Dense-vector bytes eligible for pinning.
63    pub dense_pin_intended_bytes: u64,
64}
65
66impl SegmentMemoryStats {
67    /// Total estimated heap retained by this segment reader.
68    pub fn estimated_heap_bytes(&self) -> usize {
69        self.term_dict_cache_bytes
70            + self.store_cache_bytes
71            + self.sparse_heap_bytes
72            + self.dense_heap_bytes
73    }
74
75    /// Total logical bytes in the explicitly accounted file-backed sections.
76    ///
77    /// This is mapped address space for `MmapDirectory`, not resident memory.
78    pub fn file_backed_bytes(&self) -> u64 {
79        self.term_bloom_file_bytes
80            .saturating_add(self.sparse_file_backed_bytes)
81            .saturating_add(self.dense_file_backed_bytes)
82    }
83}
84
85use std::cmp::Ordering;
86use std::collections::BinaryHeap;
87use std::sync::Arc;
88
89use rustc_hash::{FxHashMap, FxHashSet};
90
91use super::vector_data::LazyFlatVectorData;
92use crate::directories::{Directory, FileHandle};
93use crate::dsl::{DenseVectorQuantization, Document, Field, Schema};
94use crate::query::{MAX_DENSE_NPROBE, MAX_DENSE_RERANK_FACTOR};
95use crate::structures::{
96    AsyncSSTableReader, BlockPostingList, CoarseCentroids, SSTableStats, TermInfo,
97};
98use crate::{DocId, Error, Result};
99
100use super::store::{AsyncStoreReader, RawStoreBlock};
101use super::types::{SegmentFiles, SegmentId, SegmentMeta};
102
103/// Combine per-ordinal (doc_id, ordinal, score) triples into VectorSearchResults,
104/// applying the multi-value combiner, sorting by score desc, and truncating to `limit`.
105///
106/// Fast path: when all ordinals are 0 (single-valued field), skips the HashMap
107/// grouping entirely and just sorts + truncates the raw results.
108pub(crate) fn combine_ordinal_results(
109    raw: impl IntoIterator<Item = (u32, u16, f32)>,
110    combiner: crate::query::MultiValueCombiner,
111    limit: usize,
112) -> Vec<VectorSearchResult> {
113    let collected: Vec<(u32, u16, f32)> = raw.into_iter().collect();
114
115    let num_raw = collected.len();
116    if log::log_enabled!(log::Level::Debug) {
117        let mut ids: Vec<u32> = collected.iter().map(|(d, _, _)| *d).collect();
118        ids.sort_unstable();
119        ids.dedup();
120        log::debug!(
121            "combine_ordinal_results: {} raw entries, {} unique docs, combiner={:?}, limit={}",
122            num_raw,
123            ids.len(),
124            combiner,
125            limit
126        );
127    }
128
129    // Fast path: all ordinals are 0 → no grouping needed, skip HashMap
130    let all_single = collected.iter().all(|&(_, ord, _)| ord == 0);
131    if all_single {
132        let mut results: Vec<VectorSearchResult> = collected
133            .into_iter()
134            .map(|(doc_id, _, score)| VectorSearchResult::new(doc_id, score, vec![(0, score)]))
135            .collect();
136        results.sort_unstable_by(|a, b| {
137            b.score
138                .total_cmp(&a.score)
139                .then_with(|| a.doc_id.cmp(&b.doc_id))
140        });
141        results.truncate(limit);
142        return results;
143    }
144
145    // Slow path: multi-valued field — group by doc_id, apply combiner
146    let mut doc_ordinals: rustc_hash::FxHashMap<DocId, Vec<(u32, f32)>> =
147        rustc_hash::FxHashMap::default();
148    for (doc_id, ordinal, score) in collected {
149        doc_ordinals
150            .entry(doc_id as DocId)
151            .or_default()
152            .push((ordinal as u32, score));
153    }
154    let mut results: Vec<VectorSearchResult> = doc_ordinals
155        .into_iter()
156        .map(|(doc_id, ordinals)| {
157            let combined_score = combiner.combine(&ordinals);
158            VectorSearchResult::new(doc_id, combined_score, ordinals)
159        })
160        .collect();
161    results.sort_unstable_by(|a, b| {
162        b.score
163            .total_cmp(&a.score)
164            .then_with(|| a.doc_id.cmp(&b.doc_id))
165    });
166    results.truncate(limit);
167    results
168}
169
170/// Heap entry used by exact flat-vector search after all values belonging to
171/// one document have been combined. Keeping the heap at document granularity
172/// prevents several strong values from one document from crowding other
173/// documents out of the raw vector top-k.
174struct HeapVectorResult(VectorSearchResult);
175
176impl PartialEq for HeapVectorResult {
177    fn eq(&self, other: &Self) -> bool {
178        self.0.score.to_bits() == other.0.score.to_bits() && self.0.doc_id == other.0.doc_id
179    }
180}
181
182impl Eq for HeapVectorResult {}
183
184impl Ord for HeapVectorResult {
185    fn cmp(&self, other: &Self) -> Ordering {
186        // BinaryHeap top is the worst retained document: lower score, then
187        // larger doc ID for deterministic equal-score eviction.
188        other
189            .0
190            .score
191            .total_cmp(&self.0.score)
192            .then_with(|| self.0.doc_id.cmp(&other.0.doc_id))
193    }
194}
195
196impl PartialOrd for HeapVectorResult {
197    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
198        Some(self.cmp(other))
199    }
200}
201
202/// Incrementally combine a flat vector stream sorted by `(doc_id, ordinal)`
203/// and retain only the best `limit` documents. Scratch is O(values in the
204/// current document + retained output), independent of the segment size.
205struct FlatDocumentCollector {
206    heap: BinaryHeap<HeapVectorResult>,
207    limit: usize,
208    combiner: crate::query::MultiValueCombiner,
209    current_doc: Option<DocId>,
210    current_ordinals: Vec<(u32, f32)>,
211}
212
213impl FlatDocumentCollector {
214    fn new(limit: usize, combiner: crate::query::MultiValueCombiner) -> Self {
215        Self {
216            heap: BinaryHeap::with_capacity(limit.min(8 * 1024)),
217            limit,
218            combiner,
219            current_doc: None,
220            current_ordinals: Vec::new(),
221        }
222    }
223
224    fn push(&mut self, doc_id: DocId, ordinal: u16, score: f32) {
225        if self.current_doc.is_some_and(|current| current != doc_id) {
226            self.finish_current();
227        }
228        self.current_doc = Some(doc_id);
229        self.current_ordinals.push((ordinal as u32, score));
230    }
231
232    fn finish_current(&mut self) {
233        let Some(doc_id) = self.current_doc.take() else {
234            return;
235        };
236        let score = self.combiner.combine(&self.current_ordinals);
237        let should_retain = self.heap.len() < self.limit
238            || self.heap.peek().is_some_and(|worst| {
239                HeapVectorResult(VectorSearchResult::new(doc_id, score, Vec::new()))
240                    .cmp(worst)
241                    .is_lt()
242            });
243
244        if !should_retain {
245            // The overwhelmingly common path once the heap is full. Reuse
246            // the ordinal scratch instead of allocating a fresh Vec for
247            // every rejected document in a flat scan.
248            self.current_ordinals.clear();
249            return;
250        }
251
252        let ordinals = std::mem::take(&mut self.current_ordinals);
253        let entry = HeapVectorResult(VectorSearchResult::new(doc_id, score, ordinals));
254        if self.heap.len() < self.limit {
255            self.heap.push(entry);
256        } else if let Some(mut worst) = self.heap.peek_mut() {
257            // Recycle the evicted result's allocation as the next document's
258            // scratch. PeekMut restores heap order when it is dropped.
259            let mut evicted = std::mem::replace(&mut worst.0, entry.0);
260            evicted.ordinals.clear();
261            self.current_ordinals = evicted.ordinals;
262        }
263    }
264
265    fn into_results(mut self) -> Vec<VectorSearchResult> {
266        self.finish_current();
267        let mut results: Vec<_> = self.heap.into_iter().map(|entry| entry.0).collect();
268        results.sort_unstable_by(|a, b| {
269            b.score
270                .total_cmp(&a.score)
271                .then_with(|| a.doc_id.cmp(&b.doc_id))
272        });
273        results
274    }
275}
276
277/// Collect a stream already grouped by document (the layout produced by flat
278/// storage expansion) without rebuilding a hash table for every candidate.
279fn combine_grouped_ordinal_results(
280    raw: impl IntoIterator<Item = RawVectorCandidate>,
281    combiner: crate::query::MultiValueCombiner,
282    limit: usize,
283) -> Vec<VectorSearchResult> {
284    let mut collector = FlatDocumentCollector::new(limit, combiner);
285    for (doc_id, ordinal, score) in raw {
286        collector.push(doc_id, ordinal, score);
287    }
288    collector.into_results()
289}
290
291#[derive(Clone, Copy)]
292struct DenseSearchParams {
293    dim: usize,
294    nprobe: usize,
295    unit_norm: bool,
296}
297
298/// Query-derived state shared by every native-precision scoring batch in one
299/// flat scan or exact rerank operation.
300///
301/// Computing the query norm is O(dim), and f16 scoring additionally quantizes
302/// the query. Keeping both here avoids repeating that work for every bounded
303/// vector batch.
304struct PreparedDenseScoreQuery<'a> {
305    query: &'a [f32],
306    query_f16: Vec<u16>,
307    inv_norm_q: f32,
308    quantization: DenseVectorQuantization,
309    dim: usize,
310    unit_norm: bool,
311}
312
313impl<'a> PreparedDenseScoreQuery<'a> {
314    fn new(
315        query: &'a [f32],
316        quantization: DenseVectorQuantization,
317        dim: usize,
318        unit_norm: bool,
319    ) -> Result<Self> {
320        use crate::structures::simd;
321
322        if query.len() != dim {
323            return Err(Error::Query(format!(
324                "dense SIMD query dimension {} does not match vector dimension {dim}",
325                query.len()
326            )));
327        }
328        if quantization == DenseVectorQuantization::Binary {
329            return Err(Error::InvalidFieldType {
330                expected: "non-binary dense vector".to_string(),
331                got: "binary dense vector".to_string(),
332            });
333        }
334
335        let norm_q_sq = simd::dot_product_f32(query, query, dim);
336        let inv_norm_q = if norm_q_sq < f32::EPSILON {
337            0.0
338        } else {
339            simd::fast_inv_sqrt(norm_q_sq)
340        };
341        let query_f16 = if quantization == DenseVectorQuantization::F16 {
342            query.iter().map(|&value| simd::f32_to_f16(value)).collect()
343        } else {
344            Vec::new()
345        };
346
347        Ok(Self {
348            query,
349            query_f16,
350            inv_norm_q,
351            quantization,
352            dim,
353            unit_norm,
354        })
355    }
356
357    fn score_batch(&self, raw: &[u8], scores: &mut [f32]) -> Result<()> {
358        use crate::structures::simd;
359
360        let element_size = match self.quantization {
361            DenseVectorQuantization::F32 => std::mem::size_of::<f32>(),
362            DenseVectorQuantization::F16 => std::mem::size_of::<u16>(),
363            DenseVectorQuantization::UInt8 => 1,
364            DenseVectorQuantization::Binary => {
365                return Err(Error::InvalidFieldType {
366                    expected: "non-binary dense vector".to_string(),
367                    got: "binary dense vector".to_string(),
368                });
369            }
370        };
371        let required_bytes = scores
372            .len()
373            .checked_mul(self.dim)
374            .and_then(|elements| elements.checked_mul(element_size))
375            .ok_or_else(|| Error::Corruption("dense vector batch byte length overflow".into()))?;
376        if raw.len() < required_bytes {
377            return Err(Error::Corruption(format!(
378                "dense vector batch is truncated: need {required_bytes} bytes, got {}",
379                raw.len()
380            )));
381        }
382        if self.quantization == DenseVectorQuantization::F16
383            && required_bytes > 0
384            && !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>())
385        {
386            return Err(Error::Corruption(
387                "f16 vector data is not 2-byte aligned".to_string(),
388            ));
389        }
390        if self.quantization == DenseVectorQuantization::F32
391            && !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>())
392        {
393            return Err(Error::Corruption(
394                "f32 vector data is not 4-byte aligned".to_string(),
395            ));
396        }
397
398        // The legacy batch scorers leave the destination untouched for empty
399        // dimensions or batches. Retain that boundary behavior before calling
400        // the precomputed kernels.
401        if self.dim == 0 || scores.is_empty() {
402            return Ok(());
403        }
404
405        match (self.quantization, self.unit_norm) {
406            (DenseVectorQuantization::F32, false) => {
407                let num_floats = scores.len() * self.dim;
408                let vectors: &[f32] =
409                    unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
410                simd::batch_cosine_scores_precomp(
411                    self.query,
412                    vectors,
413                    self.dim,
414                    scores,
415                    self.inv_norm_q,
416                );
417            }
418            (DenseVectorQuantization::F32, true) => {
419                let num_floats = scores.len() * self.dim;
420                let vectors: &[f32] =
421                    unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
422                simd::batch_dot_scores_precomp(
423                    self.query,
424                    vectors,
425                    self.dim,
426                    scores,
427                    self.inv_norm_q,
428                );
429            }
430            (DenseVectorQuantization::F16, false) => {
431                simd::batch_cosine_scores_f16_precomp(
432                    &self.query_f16,
433                    raw,
434                    self.dim,
435                    scores,
436                    self.inv_norm_q,
437                );
438            }
439            (DenseVectorQuantization::F16, true) => {
440                simd::batch_dot_scores_f16_precomp(
441                    &self.query_f16,
442                    raw,
443                    self.dim,
444                    scores,
445                    self.inv_norm_q,
446                );
447            }
448            (DenseVectorQuantization::UInt8, false) => {
449                simd::batch_cosine_scores_u8_precomp(
450                    self.query,
451                    raw,
452                    self.dim,
453                    scores,
454                    self.inv_norm_q,
455                );
456            }
457            (DenseVectorQuantization::UInt8, true) => {
458                simd::batch_dot_scores_u8_precomp(
459                    self.query,
460                    raw,
461                    self.dim,
462                    scores,
463                    self.inv_norm_q,
464                );
465            }
466            (DenseVectorQuantization::Binary, _) => unreachable!("validated during preparation"),
467        }
468        Ok(())
469    }
470}
471
472/// Compute the ANN candidate count without relying on saturating float casts.
473fn checked_dense_fetch_k(k: usize, rerank_factor: f32) -> Result<usize> {
474    if !rerank_factor.is_finite() || !(1.0..=MAX_DENSE_RERANK_FACTOR).contains(&rerank_factor) {
475        return Err(Error::Query(format!(
476            "dense rerank_factor must be finite and in [1, {MAX_DENSE_RERANK_FACTOR}], got {rerank_factor}"
477        )));
478    }
479
480    let fetch = (k as f64) * (rerank_factor as f64);
481    if !fetch.is_finite()
482        || fetch > usize::MAX as f64
483        || fetch > MAX_DENSE_CANDIDATES_PER_SEGMENT as f64
484    {
485        return Err(Error::Query(format!(
486            "dense candidate count exceeds the per-segment maximum of \
487             {MAX_DENSE_CANDIDATES_PER_SEGMENT}: k={k}, rerank_factor={rerank_factor}"
488        )));
489    }
490    Ok(fetch.ceil() as usize)
491}
492
493/// Binary queries do not expose a configurable rerank factor. Use the shared
494/// query-level oversubscription policy while retaining the same hard
495/// per-segment candidate bound as float-vector reranking. Reject a result
496/// window larger than that bound instead of silently returning fewer than
497/// requested; candidate oversampling itself may safely clamp at the bound.
498#[inline]
499fn checked_binary_combined_fetch_k(k: usize) -> Result<usize> {
500    if k > MAX_DENSE_CANDIDATES_PER_SEGMENT {
501        return Err(Error::Query(format!(
502            "binary dense result count exceeds the per-segment maximum of \
503             {MAX_DENSE_CANDIDATES_PER_SEGMENT}: k={k}"
504        )));
505    }
506    Ok(crate::query::max_candidate_limit(k).min(MAX_DENSE_CANDIDATES_PER_SEGMENT))
507}
508
509#[inline]
510fn bounded_vector_score_batch(vector_byte_size: usize, preferred: usize) -> usize {
511    preferred.min((MAX_VECTOR_SCORE_BATCH_BYTES / vector_byte_size.max(1)).max(1))
512}
513
514#[inline]
515fn bounded_rerank_batch(vector_byte_size: usize, preferred: usize, vector_count: usize) -> usize {
516    bounded_vector_score_batch(vector_byte_size, preferred).min(vector_count.max(1))
517}
518
519fn checked_file_range(
520    offset: u64,
521    length: u64,
522    file_length: u64,
523    description: &str,
524) -> Result<std::ops::Range<u64>> {
525    let end = offset
526        .checked_add(length)
527        .ok_or_else(|| Error::Corruption(format!("{description} byte range overflows u64")))?;
528    if end > file_length {
529        return Err(Error::Corruption(format!(
530            "{description} byte range {offset}..{end} exceeds file length {file_length}"
531        )));
532    }
533    Ok(offset..end)
534}
535
536type RawVectorCandidate = (u32, u16, f32);
537type CandidateVectorRef = (DocId, u16, usize); // (doc ID, ordinal, flat-vector index)
538
539#[derive(Clone, Copy)]
540struct CandidateDocumentRange {
541    doc_id: DocId,
542    start: usize,
543    end: usize,
544}
545
546struct AnnCandidateDocuments {
547    ranges: Vec<CandidateDocumentRange>,
548    vector_count: usize,
549}
550
551/// Resolve the document union returned by ANN to compact flat-vector ranges.
552///
553/// The document union is bounded by ANN document top-k, while the number of
554/// values those documents own is intentionally not capped. A valid
555/// multi-valued document may have many ordinals;
556/// materializing one result and one flat-index entry per ordinal used to turn
557/// that into a spurious query error at 20,000 vectors. Callers stream these
558/// ranges through a fixed-size score buffer instead.
559fn ann_candidate_document_ranges(
560    ann_results: &[RawVectorCandidate],
561    flat: &LazyFlatVectorData,
562) -> Result<AnnCandidateDocuments> {
563    ann_candidate_document_ranges_from_ids(ann_results.iter().map(|candidate| candidate.0), flat)
564}
565
566fn ann_candidate_document_ranges_from_ids(
567    doc_ids: impl IntoIterator<Item = DocId>,
568    flat: &LazyFlatVectorData,
569) -> Result<AnnCandidateDocuments> {
570    let mut candidate_docs: Vec<DocId> = doc_ids.into_iter().collect();
571    candidate_docs.sort_unstable();
572    candidate_docs.dedup();
573
574    let mut ranges = Vec::with_capacity(candidate_docs.len());
575    let mut vector_count = 0usize;
576    for doc_id in candidate_docs {
577        let (start, count) = flat.flat_indexes_for_doc_range(doc_id);
578        if count == 0 {
579            return Err(Error::Corruption(format!(
580                "ANN candidate document {doc_id} is missing from flat vector storage"
581            )));
582        }
583        vector_count = vector_count
584            .checked_add(count)
585            .ok_or_else(|| Error::Query("ANN candidate vector expansion overflow".to_string()))?;
586        let end = start
587            .checked_add(count)
588            .ok_or_else(|| Error::Corruption("flat vector range overflow".to_string()))?;
589        if end > flat.num_vectors {
590            return Err(Error::Corruption(format!(
591                "flat vector range {start}..{end} for document {doc_id} exceeds {} vectors",
592                flat.num_vectors
593            )));
594        }
595        ranges.push(CandidateDocumentRange { doc_id, start, end });
596    }
597    Ok(AnnCandidateDocuments {
598        ranges,
599        vector_count,
600    })
601}
602
603/// Validate the no-rerank binary IVF fast path against exact flat metadata.
604///
605/// Binary IVF stores the original packed codes, so a single-valued field does
606/// not need vector-data I/O to recompute scores. It still needs the same
607/// ANN/flat consistency checks the rerank path provided: every candidate must
608/// name the field's sole stored ordinal. Deduplicate by document as well so a
609/// malformed ANN payload cannot surface the same document more than once.
610fn validate_binary_single_value_ann_results(
611    ann_results: Vec<RawVectorCandidate>,
612    flat: &LazyFlatVectorData,
613) -> Result<Vec<RawVectorCandidate>> {
614    let mut seen_docs = FxHashSet::default();
615    let mut validated = Vec::with_capacity(ann_results.len());
616    for (doc_id, ordinal, score) in ann_results {
617        let (start, count) = flat.flat_indexes_for_doc_range(doc_id);
618        if count == 0 {
619            return Err(Error::Corruption(format!(
620                "ANN candidate document {doc_id} is missing from flat vector storage"
621            )));
622        }
623        if count != 1 {
624            return Err(Error::Corruption(format!(
625                "binary ANN single-valued candidate document {doc_id} has {count} flat vectors"
626            )));
627        }
628        let (stored_doc_id, stored_ordinal) = flat.get_doc_id(start);
629        if stored_doc_id != doc_id {
630            return Err(Error::Corruption(format!(
631                "flat vector doc map is not contiguous for document {doc_id}"
632            )));
633        }
634        if stored_ordinal != ordinal {
635            return Err(Error::Corruption(format!(
636                "binary ANN candidate document {doc_id} ordinal {ordinal} is missing from flat vector storage"
637            )));
638        }
639        if seen_docs.insert(doc_id) {
640            validated.push((doc_id, ordinal, score));
641        }
642    }
643    Ok(validated)
644}
645
646struct CandidateVectorCursor<'a> {
647    ranges: &'a [CandidateDocumentRange],
648    range_index: usize,
649    flat_index: usize,
650}
651
652impl<'a> CandidateVectorCursor<'a> {
653    fn new(ranges: &'a [CandidateDocumentRange]) -> Self {
654        Self {
655            ranges,
656            range_index: 0,
657            flat_index: ranges.first().map_or(0, |range| range.start),
658        }
659    }
660
661    /// Fill `batch` in `(doc_id, ordinal)` order. The cursor validates the
662    /// contiguity promise made by the flat doc map while it streams, avoiding
663    /// an O(all candidate ordinals) validation allocation.
664    fn fill_batch(
665        &mut self,
666        flat: &LazyFlatVectorData,
667        batch: &mut Vec<CandidateVectorRef>,
668        limit: usize,
669    ) -> Result<bool> {
670        batch.clear();
671        while batch.len() < limit && self.range_index < self.ranges.len() {
672            let range = self.ranges[self.range_index];
673            if self.flat_index == range.end {
674                self.range_index += 1;
675                if let Some(next) = self.ranges.get(self.range_index) {
676                    self.flat_index = next.start;
677                }
678                continue;
679            }
680            let (stored_doc_id, ordinal) = flat.get_doc_id(self.flat_index);
681            if stored_doc_id != range.doc_id {
682                return Err(Error::Corruption(format!(
683                    "flat vector doc map is not contiguous for document {}",
684                    range.doc_id
685                )));
686            }
687            batch.push((range.doc_id, ordinal, self.flat_index));
688            self.flat_index += 1;
689        }
690        Ok(!batch.is_empty())
691    }
692}
693
694#[derive(Clone, Copy)]
695struct VectorReadRun {
696    buffer_start: usize,
697    flat_start: usize,
698    count: usize,
699}
700
701/// Coalesce an ordered set of selected flat indexes into contiguous reads.
702/// Multi-valued document bodies are stored consecutively, so this turns the
703/// common case from one range lookup per value into one lookup per bounded
704/// run while retaining a packed score buffer.
705fn plan_vector_read_runs(indexes: &[usize], runs: &mut Vec<VectorReadRun>) -> Result<()> {
706    runs.clear();
707    for (buffer_index, &flat_index) in indexes.iter().enumerate() {
708        if let Some(run) = runs.last_mut()
709            && run
710                .flat_start
711                .checked_add(run.count)
712                .is_some_and(|next| next == flat_index)
713        {
714            run.count += 1;
715            continue;
716        }
717        if buffer_index > 0 && flat_index <= indexes[buffer_index - 1] {
718            return Err(Error::Corruption(
719                "candidate flat-vector indexes are not strictly ordered".into(),
720            ));
721        }
722        runs.push(VectorReadRun {
723            buffer_start: buffer_index,
724            flat_start: flat_index,
725            count: 1,
726        });
727    }
728    Ok(())
729}
730
731/// Plan contiguous raw-vector reads and initiate page-in before either the
732/// synchronous or asynchronous reader starts copying. Keeping prefetch here
733/// prevents the two execution paths from drifting.
734fn prepare_vector_read_runs(
735    flat: &LazyFlatVectorData,
736    indexes: &[usize],
737    runs: &mut Vec<VectorReadRun>,
738) -> Result<()> {
739    plan_vector_read_runs(indexes, runs)?;
740    #[cfg(feature = "native")]
741    flat.prefetch_vectors(indexes.iter().copied());
742    #[cfg(not(feature = "native"))]
743    let _ = flat;
744    Ok(())
745}
746
747async fn read_vector_runs(
748    flat: &LazyFlatVectorData,
749    indexes: &[usize],
750    runs: &mut Vec<VectorReadRun>,
751    output: &mut [u8],
752) -> Result<()> {
753    prepare_vector_read_runs(flat, indexes, runs)?;
754    let vector_byte_size = flat.vector_byte_size();
755    for run in runs {
756        let bytes = flat
757            .read_vectors_batch(run.flat_start, run.count)
758            .await
759            .map_err(Error::Io)?;
760        let start = run
761            .buffer_start
762            .checked_mul(vector_byte_size)
763            .ok_or_else(|| Error::Query("dense rerank buffer offset overflow".into()))?;
764        let end = start
765            .checked_add(bytes.len())
766            .ok_or_else(|| Error::Query("dense rerank buffer range overflow".into()))?;
767        let destination = output
768            .get_mut(start..end)
769            .ok_or_else(|| Error::Corruption("dense rerank buffer is too short".into()))?;
770        destination.copy_from_slice(bytes.as_slice());
771    }
772    Ok(())
773}
774
775#[cfg(feature = "sync")]
776fn read_vector_runs_sync(
777    flat: &LazyFlatVectorData,
778    indexes: &[usize],
779    runs: &mut Vec<VectorReadRun>,
780    output: &mut [u8],
781) -> Result<()> {
782    prepare_vector_read_runs(flat, indexes, runs)?;
783    let vector_byte_size = flat.vector_byte_size();
784    for run in runs {
785        let bytes = flat
786            .read_vectors_batch_sync(run.flat_start, run.count)
787            .map_err(Error::Io)?;
788        let start = run
789            .buffer_start
790            .checked_mul(vector_byte_size)
791            .ok_or_else(|| Error::Query("dense rerank buffer offset overflow".into()))?;
792        let end = start
793            .checked_add(bytes.len())
794            .ok_or_else(|| Error::Query("dense rerank buffer range overflow".into()))?;
795        let destination = output
796            .get_mut(start..end)
797            .ok_or_else(|| Error::Corruption("dense rerank buffer is too short".into()))?;
798        destination.copy_from_slice(bytes.as_slice());
799    }
800    Ok(())
801}
802
803#[derive(Default)]
804struct DenseRerankStats {
805    vector_count: usize,
806    resolve_elapsed: std::time::Duration,
807    read_elapsed: std::time::Duration,
808    score_elapsed: std::time::Duration,
809}
810
811async fn exact_score_dense_candidate_documents(
812    ann_results: &[RawVectorCandidate],
813    flat: &LazyFlatVectorData,
814    query: &[f32],
815    unit_norm: bool,
816    combiner: crate::query::MultiValueCombiner,
817    limit: usize,
818) -> Result<(Vec<VectorSearchResult>, DenseRerankStats)> {
819    let resolve_started = std::time::Instant::now();
820    let documents = ann_candidate_document_ranges(ann_results, flat)?;
821    let mut stats = DenseRerankStats {
822        vector_count: documents.vector_count,
823        resolve_elapsed: resolve_started.elapsed(),
824        ..Default::default()
825    };
826    let vector_byte_size = flat.vector_byte_size();
827    let batch_len =
828        bounded_rerank_batch(vector_byte_size, DENSE_SCORE_BATCH, documents.vector_count);
829    let raw_capacity = batch_len
830        .checked_mul(vector_byte_size)
831        .ok_or_else(|| Error::Query("dense rerank buffer size overflow".to_string()))?;
832    let mut raw = vec![0u8; raw_capacity];
833    let mut scores = vec![0.0f32; batch_len];
834    let mut batch = Vec::with_capacity(batch_len);
835    let mut flat_indexes = Vec::with_capacity(batch_len);
836    let mut read_runs = Vec::new();
837    let mut cursor = CandidateVectorCursor::new(&documents.ranges);
838    let mut collector = FlatDocumentCollector::new(limit, combiner);
839    let mut scored = 0usize;
840    let prepared_query =
841        PreparedDenseScoreQuery::new(query, flat.quantization, flat.dim, unit_norm)?;
842
843    while cursor.fill_batch(flat, &mut batch, batch_len)? {
844        flat_indexes.clear();
845        flat_indexes.extend(batch.iter().map(|&(_, _, flat_index)| flat_index));
846        let raw_len = batch
847            .len()
848            .checked_mul(vector_byte_size)
849            .ok_or_else(|| Error::Query("dense rerank buffer size overflow".to_string()))?;
850        let raw = &mut raw[..raw_len];
851
852        let read_started = std::time::Instant::now();
853        read_vector_runs(flat, &flat_indexes, &mut read_runs, raw).await?;
854        stats.read_elapsed += read_started.elapsed();
855
856        let score_started = std::time::Instant::now();
857        prepared_query.score_batch(raw, &mut scores[..batch.len()])?;
858        stats.score_elapsed += score_started.elapsed();
859        for (buffer_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
860            collector.push(doc_id, ordinal, scores[buffer_index]);
861        }
862        scored += batch.len();
863    }
864    debug_assert_eq!(scored, documents.vector_count);
865    Ok((collector.into_results(), stats))
866}
867
868#[cfg(feature = "sync")]
869fn exact_score_dense_candidate_documents_sync(
870    ann_results: &[RawVectorCandidate],
871    flat: &LazyFlatVectorData,
872    query: &[f32],
873    unit_norm: bool,
874    combiner: crate::query::MultiValueCombiner,
875    limit: usize,
876) -> Result<Vec<VectorSearchResult>> {
877    let documents = ann_candidate_document_ranges(ann_results, flat)?;
878    let vector_byte_size = flat.vector_byte_size();
879    let batch_len =
880        bounded_rerank_batch(vector_byte_size, DENSE_SCORE_BATCH, documents.vector_count);
881    let raw_capacity = batch_len
882        .checked_mul(vector_byte_size)
883        .ok_or_else(|| Error::Query("dense rerank buffer size overflow".to_string()))?;
884    let mut raw = vec![0u8; raw_capacity];
885    let mut scores = vec![0.0f32; batch_len];
886    let mut batch = Vec::with_capacity(batch_len);
887    let mut flat_indexes = Vec::with_capacity(batch_len);
888    let mut read_runs = Vec::new();
889    let mut cursor = CandidateVectorCursor::new(&documents.ranges);
890    let mut collector = FlatDocumentCollector::new(limit, combiner);
891    let mut scored = 0usize;
892    let prepared_query =
893        PreparedDenseScoreQuery::new(query, flat.quantization, flat.dim, unit_norm)?;
894
895    while cursor.fill_batch(flat, &mut batch, batch_len)? {
896        flat_indexes.clear();
897        flat_indexes.extend(batch.iter().map(|&(_, _, flat_index)| flat_index));
898        let raw_len = batch
899            .len()
900            .checked_mul(vector_byte_size)
901            .ok_or_else(|| Error::Query("dense rerank buffer size overflow".to_string()))?;
902        let raw = &mut raw[..raw_len];
903        read_vector_runs_sync(flat, &flat_indexes, &mut read_runs, raw)?;
904        prepared_query.score_batch(raw, &mut scores[..batch.len()])?;
905        for (buffer_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
906            collector.push(doc_id, ordinal, scores[buffer_index]);
907        }
908        scored += batch.len();
909    }
910    debug_assert_eq!(scored, documents.vector_count);
911    Ok(collector.into_results())
912}
913
914async fn exact_score_binary_candidate_documents(
915    ann_results: &[RawVectorCandidate],
916    flat: &LazyFlatVectorData,
917    query: &[u8],
918    dim_bits: usize,
919    combiner: crate::query::MultiValueCombiner,
920    limit: usize,
921) -> Result<Vec<VectorSearchResult>> {
922    let documents = ann_candidate_document_ranges(ann_results, flat)?;
923    let probe_scores: FxHashMap<(DocId, u16), f32> = ann_results
924        .iter()
925        .map(|&(doc_id, ordinal, score)| ((doc_id, ordinal), score))
926        .collect();
927    exact_score_binary_resolved_documents(
928        documents,
929        &probe_scores,
930        flat,
931        query,
932        dim_bits,
933        combiner,
934        limit,
935    )
936    .await
937}
938
939async fn exact_score_binary_candidate_document_ids(
940    candidate_doc_ids: Vec<DocId>,
941    probed_ordinal_scores: &[(u32, u16, f32)],
942    flat: &LazyFlatVectorData,
943    query: &[u8],
944    dim_bits: usize,
945    combiner: crate::query::MultiValueCombiner,
946    limit: usize,
947) -> Result<Vec<VectorSearchResult>> {
948    let documents = ann_candidate_document_ranges_from_ids(candidate_doc_ids, flat)?;
949    // Binary leaves hold the original packed codes, so probed ordinals already
950    // have exact scores; only ordinals outside the probed leaves are read back.
951    let probe_scores = binary_probe_score_map(probed_ordinal_scores);
952    exact_score_binary_resolved_documents(
953        documents,
954        &probe_scores,
955        flat,
956        query,
957        dim_bits,
958        combiner,
959        limit,
960    )
961    .await
962}
963
964fn binary_probe_score_map(
965    probed_ordinal_scores: &[(u32, u16, f32)],
966) -> FxHashMap<(DocId, u16), f32> {
967    probed_ordinal_scores
968        .iter()
969        .map(|&(doc_id, ordinal, score)| ((doc_id, ordinal), score))
970        .collect()
971}
972
973async fn exact_score_binary_resolved_documents(
974    documents: AnnCandidateDocuments,
975    probe_scores: &FxHashMap<(DocId, u16), f32>,
976    flat: &LazyFlatVectorData,
977    query: &[u8],
978    dim_bits: usize,
979    combiner: crate::query::MultiValueCombiner,
980    limit: usize,
981) -> Result<Vec<VectorSearchResult>> {
982    let vector_byte_size = flat.vector_byte_size();
983    let batch_len =
984        bounded_rerank_batch(vector_byte_size, BINARY_SCORE_BATCH, documents.vector_count);
985    let raw_capacity = batch_len
986        .checked_mul(vector_byte_size)
987        .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
988    let mut raw = vec![0u8; raw_capacity];
989    let mut scores = vec![0.0f32; batch_len];
990    let mut batch_scores = vec![0.0f32; batch_len];
991    let mut batch = Vec::with_capacity(batch_len);
992    let mut unresolved = Vec::with_capacity(batch_len);
993    let mut unresolved_flat_indexes = Vec::with_capacity(batch_len);
994    let mut read_runs = Vec::new();
995    let mut cursor = CandidateVectorCursor::new(&documents.ranges);
996    let mut collector = FlatDocumentCollector::new(limit, combiner);
997    let mut scored = 0usize;
998
999    while cursor.fill_batch(flat, &mut batch, batch_len)? {
1000        unresolved.clear();
1001        for (batch_index, &(doc_id, ordinal, flat_index)) in batch.iter().enumerate() {
1002            if let Some(&score) = probe_scores.get(&(doc_id, ordinal)) {
1003                batch_scores[batch_index] = score;
1004            } else {
1005                unresolved.push((batch_index, flat_index));
1006            }
1007        }
1008        unresolved_flat_indexes.clear();
1009        unresolved_flat_indexes.extend(unresolved.iter().map(|&(_, flat_index)| flat_index));
1010        let raw_len = unresolved
1011            .len()
1012            .checked_mul(vector_byte_size)
1013            .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
1014        let raw = &mut raw[..raw_len];
1015        read_vector_runs(flat, &unresolved_flat_indexes, &mut read_runs, raw).await?;
1016        crate::structures::simd::batch_hamming_scores(
1017            query,
1018            raw,
1019            vector_byte_size,
1020            dim_bits,
1021            &mut scores[..unresolved.len()],
1022        );
1023        for (buffer_index, &(batch_index, _)) in unresolved.iter().enumerate() {
1024            batch_scores[batch_index] = scores[buffer_index];
1025        }
1026        for (batch_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
1027            collector.push(doc_id, ordinal, batch_scores[batch_index]);
1028        }
1029        scored += batch.len();
1030    }
1031    debug_assert_eq!(scored, documents.vector_count);
1032    Ok(collector.into_results())
1033}
1034
1035#[cfg(feature = "sync")]
1036fn exact_score_binary_candidate_documents_sync(
1037    ann_results: &[RawVectorCandidate],
1038    flat: &LazyFlatVectorData,
1039    query: &[u8],
1040    dim_bits: usize,
1041    combiner: crate::query::MultiValueCombiner,
1042    limit: usize,
1043) -> Result<Vec<VectorSearchResult>> {
1044    let documents = ann_candidate_document_ranges(ann_results, flat)?;
1045    let probe_scores: FxHashMap<(DocId, u16), f32> = ann_results
1046        .iter()
1047        .map(|&(doc_id, ordinal, score)| ((doc_id, ordinal), score))
1048        .collect();
1049    exact_score_binary_resolved_documents_sync(
1050        documents,
1051        &probe_scores,
1052        flat,
1053        query,
1054        dim_bits,
1055        combiner,
1056        limit,
1057    )
1058}
1059
1060#[cfg(feature = "sync")]
1061fn exact_score_binary_candidate_document_ids_sync(
1062    candidate_doc_ids: Vec<DocId>,
1063    probed_ordinal_scores: &[(u32, u16, f32)],
1064    flat: &LazyFlatVectorData,
1065    query: &[u8],
1066    dim_bits: usize,
1067    combiner: crate::query::MultiValueCombiner,
1068    limit: usize,
1069) -> Result<Vec<VectorSearchResult>> {
1070    let documents = ann_candidate_document_ranges_from_ids(candidate_doc_ids, flat)?;
1071    let probe_scores = binary_probe_score_map(probed_ordinal_scores);
1072    exact_score_binary_resolved_documents_sync(
1073        documents,
1074        &probe_scores,
1075        flat,
1076        query,
1077        dim_bits,
1078        combiner,
1079        limit,
1080    )
1081}
1082
1083#[cfg(feature = "sync")]
1084fn exact_score_binary_resolved_documents_sync(
1085    documents: AnnCandidateDocuments,
1086    probe_scores: &FxHashMap<(DocId, u16), f32>,
1087    flat: &LazyFlatVectorData,
1088    query: &[u8],
1089    dim_bits: usize,
1090    combiner: crate::query::MultiValueCombiner,
1091    limit: usize,
1092) -> Result<Vec<VectorSearchResult>> {
1093    let vector_byte_size = flat.vector_byte_size();
1094    let batch_len =
1095        bounded_rerank_batch(vector_byte_size, BINARY_SCORE_BATCH, documents.vector_count);
1096    let raw_capacity = batch_len
1097        .checked_mul(vector_byte_size)
1098        .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
1099    let mut raw = vec![0u8; raw_capacity];
1100    let mut scores = vec![0.0f32; batch_len];
1101    let mut batch_scores = vec![0.0f32; batch_len];
1102    let mut batch = Vec::with_capacity(batch_len);
1103    let mut unresolved = Vec::with_capacity(batch_len);
1104    let mut unresolved_flat_indexes = Vec::with_capacity(batch_len);
1105    let mut read_runs = Vec::new();
1106    let mut cursor = CandidateVectorCursor::new(&documents.ranges);
1107    let mut collector = FlatDocumentCollector::new(limit, combiner);
1108    let mut scored = 0usize;
1109
1110    while cursor.fill_batch(flat, &mut batch, batch_len)? {
1111        unresolved.clear();
1112        for (batch_index, &(doc_id, ordinal, flat_index)) in batch.iter().enumerate() {
1113            if let Some(&score) = probe_scores.get(&(doc_id, ordinal)) {
1114                batch_scores[batch_index] = score;
1115            } else {
1116                unresolved.push((batch_index, flat_index));
1117            }
1118        }
1119        unresolved_flat_indexes.clear();
1120        unresolved_flat_indexes.extend(unresolved.iter().map(|&(_, flat_index)| flat_index));
1121        let raw_len = unresolved
1122            .len()
1123            .checked_mul(vector_byte_size)
1124            .ok_or_else(|| Error::Query("binary candidate buffer size overflow".to_string()))?;
1125        let raw = &mut raw[..raw_len];
1126        read_vector_runs_sync(flat, &unresolved_flat_indexes, &mut read_runs, raw)?;
1127        crate::structures::simd::batch_hamming_scores(
1128            query,
1129            raw,
1130            vector_byte_size,
1131            dim_bits,
1132            &mut scores[..unresolved.len()],
1133        );
1134        for (buffer_index, &(batch_index, _)) in unresolved.iter().enumerate() {
1135            batch_scores[batch_index] = scores[buffer_index];
1136        }
1137        for (batch_index, &(doc_id, ordinal, _)) in batch.iter().enumerate() {
1138            collector.push(doc_id, ordinal, batch_scores[batch_index]);
1139        }
1140        scored += batch.len();
1141    }
1142    debug_assert_eq!(scored, documents.vector_count);
1143    Ok(collector.into_results())
1144}
1145
1146fn validate_coarse_centroids(centroids: &CoarseCentroids, dim: usize) -> Result<()> {
1147    let expected = (centroids.num_clusters as usize)
1148        .checked_mul(dim)
1149        .ok_or_else(|| Error::Corruption("coarse centroid size overflow".into()))?;
1150    if centroids.num_clusters == 0
1151        || centroids.dim != dim
1152        || centroids.centroids.len() != expected
1153        || centroids.centroids.iter().any(|value| !value.is_finite())
1154    {
1155        return Err(Error::Corruption(format!(
1156            "invalid coarse centroids: clusters={}, dim={}, values={} (expected dim={dim}, values={expected})",
1157            centroids.num_clusters,
1158            centroids.dim,
1159            centroids.centroids.len()
1160        )));
1161    }
1162    Ok(())
1163}
1164
1165/// Per-query dense plan caches, shared by every segment scorer the query
1166/// spawns. Both members are query-global: the IVF-TQ probe route and its
1167/// LUTs depend only on the query and index-level artifacts, and the TQ
1168/// LUTs depend only on the query and the schema dimension.
1169#[derive(Debug, Default)]
1170pub struct DensePlanCache {
1171    pub(crate) tq: std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqQueryPlan>>>,
1172    pub(crate) ivf_tq: std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqIvfQueryPlan>>>,
1173    scann: std::sync::Mutex<Option<ScannPlanCacheEntry>>,
1174}
1175
1176#[derive(Debug)]
1177struct ScannPlanCacheEntry {
1178    artifact_id: u64,
1179    probes: usize,
1180    query_bits: Vec<u32>,
1181    plan: std::sync::Arc<crate::structures::vector::scann::FloatScannQuery>,
1182}
1183
1184/// Search one segment's TQ payload, reusing the per-query plan across
1185/// segments: the codec is a pure function of the schema dimension, so the
1186/// LUTs are identical for every segment of the field (mirrors the IVF-PQ
1187/// `probe_cache` hot-path rule — no repeated per-segment plan allocation).
1188#[allow(clippy::too_many_arguments)]
1189fn search_tq_segment(
1190    index: &crate::segment::ann_disk::AnnDiskIndex,
1191    codec: &crate::structures::TqCodec,
1192    query: &[f32],
1193    fetch_k: usize,
1194    document_combiner: Option<crate::query::MultiValueCombiner>,
1195    field: Field,
1196    dim: usize,
1197    plan_cache: Option<&std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqQueryPlan>>>>,
1198) -> Result<Vec<RawVectorCandidate>> {
1199    validate_tq_ann(index, codec, dim, field)?;
1200    let plan = cached_tq_query_plan(codec, query, plan_cache)?;
1201    match document_combiner {
1202        Some(combiner) => index
1203            .search_tq_combined_documents(fetch_k, &plan, combiner)
1204            .map(|candidates| {
1205                candidates
1206                    .into_iter()
1207                    // Exact dense reranking consumes only the document ID.
1208                    // Use a zero placeholder so the document aggregate can
1209                    // never be mistaken for an ordinal score.
1210                    .map(|candidate| (candidate.doc_id, 0, 0.0))
1211                    .collect()
1212            }),
1213        None => index.search_tq_distinct(fetch_k, &plan),
1214    }
1215    .map_err(|error| {
1216        Error::Corruption(format!("invalid TQ payload for field {}: {error}", field.0))
1217    })
1218}
1219
1220/// Return the query-global flat-TQ plan, rebuilding it whenever either the
1221/// codec generation or the exact query bits differ.
1222fn cached_tq_query_plan(
1223    codec: &crate::structures::TqCodec,
1224    query: &[f32],
1225    plan_cache: Option<&std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqQueryPlan>>>>,
1226) -> Result<std::sync::Arc<crate::structures::TqQueryPlan>> {
1227    Ok(match plan_cache {
1228        Some(cache) => {
1229            let mut cached = cache
1230                .lock()
1231                .map_err(|_| Error::Internal("TQ plan cache is poisoned".into()))?;
1232            match cached.as_ref() {
1233                Some(plan)
1234                    if plan.fingerprint() == codec.fingerprint() && plan.matches_query(query) =>
1235                {
1236                    std::sync::Arc::clone(plan)
1237                }
1238                _ => {
1239                    let plan =
1240                        std::sync::Arc::new(crate::structures::TqQueryPlan::build(codec, query));
1241                    *cached = Some(std::sync::Arc::clone(&plan));
1242                    plan
1243                }
1244            }
1245        }
1246        None => std::sync::Arc::new(crate::structures::TqQueryPlan::build(codec, query)),
1247    })
1248}
1249
1250fn validate_tq_ann(
1251    index: &crate::segment::ann_disk::AnnDiskIndex,
1252    codec: &crate::structures::TqCodec,
1253    dim: usize,
1254    field: Field,
1255) -> Result<()> {
1256    let header = index.header();
1257    if header.dim != dim
1258        || codec.dim() != dim
1259        || header.code_size != codec.code_size()
1260        || header.quantizer_version != codec.fingerprint()
1261        || header.codebook_version != 0
1262        || header.num_clusters != 1
1263    {
1264        return Err(Error::Corruption(format!(
1265            "TQ payload for field {} does not match the codec derived from schema dimension {dim}",
1266            field.0,
1267        )));
1268    }
1269    Ok(())
1270}
1271
1272/// Search one segment's IVF-TQ payload. The probe route, the `⟨q̂,c⟩`
1273/// scalars, and the TQ LUTs are all query-global, so the plan is cached and
1274/// shared across every segment of the field.
1275#[allow(clippy::too_many_arguments)]
1276fn search_ivf_tq_segment(
1277    index: &crate::segment::ann_disk::AnnDiskIndex,
1278    centroids: &CoarseCentroids,
1279    codec: &crate::structures::TqCodec,
1280    query: &[f32],
1281    fetch_k: usize,
1282    document_combiner: Option<crate::query::MultiValueCombiner>,
1283    field: Field,
1284    nprobe: usize,
1285    routing: crate::dsl::IvfRoutingMode,
1286    plan_cache: Option<
1287        &std::sync::Mutex<Option<std::sync::Arc<crate::structures::TqIvfQueryPlan>>>,
1288    >,
1289) -> Result<Vec<RawVectorCandidate>> {
1290    let effective_nprobe = nprobe.clamp(1, centroids.num_clusters as usize);
1291    let request_fingerprint = crate::structures::TqIvfQueryPlan::request_fingerprint_for(
1292        centroids,
1293        query,
1294        effective_nprobe,
1295        routing,
1296    );
1297    let build = || {
1298        std::sync::Arc::new(crate::structures::TqIvfQueryPlan::build(
1299            centroids,
1300            codec,
1301            query,
1302            effective_nprobe,
1303            routing,
1304        ))
1305    };
1306    let plan = match plan_cache {
1307        Some(cache) => {
1308            let mut cached = cache
1309                .lock()
1310                .map_err(|_| Error::Internal("IVF-TQ plan cache is poisoned".into()))?;
1311            match cached.as_ref() {
1312                Some(plan)
1313                    if plan.quantizer_version == centroids.version
1314                        && plan.fingerprint == codec.fingerprint()
1315                        && plan.request_fingerprint == request_fingerprint
1316                        && plan.cluster_ids.len() == effective_nprobe =>
1317                {
1318                    std::sync::Arc::clone(plan)
1319                }
1320                _ => {
1321                    let plan = build();
1322                    *cached = Some(std::sync::Arc::clone(&plan));
1323                    plan
1324                }
1325            }
1326        }
1327        None => build(),
1328    };
1329    let candidates = match document_combiner {
1330        Some(combiner) => index
1331            .search_ivf_tq_combined_documents(fetch_k, &plan, combiner)
1332            .map(|documents| {
1333                documents
1334                    .into_iter()
1335                    // The compressed score aggregates a whole document. Exact
1336                    // dense reranking consumes only its ID; a zero placeholder
1337                    // prevents accidental reuse as an ordinal score.
1338                    .map(|candidate| (candidate.doc_id, 0, 0.0))
1339                    .collect()
1340            }),
1341        None => index.search_ivf_tq_distinct(fetch_k, &plan),
1342    };
1343    candidates.map_err(|error| {
1344        Error::Corruption(format!(
1345            "invalid IVF-TQ payload for field {}: {error}",
1346            field.0
1347        ))
1348    })
1349}
1350
1351#[allow(clippy::too_many_arguments)]
1352fn search_scann_ah_segment(
1353    index: &crate::segment::ann_disk::AnnDiskIndex,
1354    artifact: &crate::segment::ScannTrainedArtifactBytes,
1355    query: &[f32],
1356    fetch_k: usize,
1357    combiner: crate::query::MultiValueCombiner,
1358    field: Field,
1359    nprobe: usize,
1360    plan_cache: Option<&std::sync::Mutex<Option<ScannPlanCacheEntry>>>,
1361) -> Result<Vec<RawVectorCandidate>> {
1362    index
1363        .validate_scann_generation(
1364            artifact.config(),
1365            artifact.generation(),
1366            artifact.artifact_id(),
1367        )
1368        .map_err(|error| {
1369            Error::Corruption(format!(
1370                "ScaNN generation mismatch for field {}: {error}",
1371                field.0
1372            ))
1373        })?;
1374    let probes = nprobe.clamp(1, artifact.config().num_leaves as usize);
1375    let query_bits = query
1376        .iter()
1377        .map(|value| value.to_bits())
1378        .collect::<Vec<_>>();
1379    let build = || {
1380        let mut normalized_query = query.to_vec();
1381        crate::structures::vector::ivf::routing::normalize_cosine_in_place(&mut normalized_query);
1382        artifact
1383            .float_model()
1384            .map_err(Error::Io)?
1385            .prepare_query(&normalized_query, probes)
1386            .map(std::sync::Arc::new)
1387            .map_err(|error| Error::Query(format!("invalid ScaNN query: {error}")))
1388    };
1389    let plan = match plan_cache {
1390        Some(cache) => {
1391            let mut cached = cache
1392                .lock()
1393                .map_err(|_| Error::Internal("ScaNN plan cache is poisoned".into()))?;
1394            match cached.as_ref() {
1395                Some(entry)
1396                    if entry.artifact_id == artifact.artifact_id()
1397                        && entry.probes == probes
1398                        && entry.query_bits == query_bits =>
1399                {
1400                    std::sync::Arc::clone(&entry.plan)
1401                }
1402                _ => {
1403                    let plan = build()?;
1404                    *cached = Some(ScannPlanCacheEntry {
1405                        artifact_id: artifact.artifact_id(),
1406                        probes,
1407                        query_bits,
1408                        plan: std::sync::Arc::clone(&plan),
1409                    });
1410                    plan
1411                }
1412            }
1413        }
1414        None => build()?,
1415    };
1416    index
1417        .search_scann_ah_combined_documents(fetch_k, &plan, combiner)
1418        .map(|documents| {
1419            documents
1420                .into_iter()
1421                .map(|candidate| (candidate.doc_id, 0, 0.0))
1422                .collect()
1423        })
1424        .map_err(|error| {
1425            Error::Corruption(format!(
1426                "invalid ScaNN AH payload for field {}: {error}",
1427                field.0
1428            ))
1429        })
1430}
1431
1432fn validate_ivf_tq_ann(
1433    index: &crate::segment::ann_disk::AnnDiskIndex,
1434    centroids: &CoarseCentroids,
1435    codec: &crate::structures::TqCodec,
1436    dim: usize,
1437    routing: crate::dsl::IvfRoutingMode,
1438    field: Field,
1439) -> Result<()> {
1440    let header = index.header();
1441    if !crate::structures::is_ivf_tq_cosine_generation(centroids.version)
1442        || !crate::structures::is_ivf_tq_cosine_generation(header.quantizer_version)
1443    {
1444        return Err(Error::Corruption(format!(
1445            "IVF-TQ field {} uses a legacy unmarked raw-vector generation that cannot \
1446             preserve cosine candidate semantics; rebuild the index with a current \
1447             Hermes version",
1448            field.0,
1449        )));
1450    }
1451    if header.dim != dim
1452        || codec.dim() != dim
1453        || header.code_size != codec.code_size()
1454        || header.num_clusters != centroids.num_clusters
1455        || header.quantizer_version != centroids.version
1456        || header.codebook_version != codec.fingerprint()
1457        || header.routing != routing
1458    {
1459        return Err(Error::Corruption(format!(
1460            "IVF-TQ payload for field {} does not match its quantizer/codec generation",
1461            field.0,
1462        )));
1463    }
1464    Ok(())
1465}
1466
1467fn validate_binary_ann(
1468    index: &crate::segment::ann_disk::AnnDiskIndex,
1469    quantizer: &crate::structures::BinaryCoarseQuantizer,
1470    config: &crate::dsl::BinaryDenseVectorConfig,
1471    dim: usize,
1472    field: Field,
1473) -> Result<()> {
1474    let header = index.header();
1475    if header.dim != dim
1476        || header.code_size != config.byte_len()
1477        || header.num_clusters != quantizer.num_clusters
1478        || header.quantizer_version != quantizer.version
1479        || header.codebook_version != 0
1480        || header.routing != config.ivf_routing
1481        || quantizer.dim_bits != dim
1482    {
1483        return Err(Error::Corruption(format!(
1484            "binary IVF field {} does not match its quantizer/schema generation",
1485            field.0,
1486        )));
1487    }
1488    Ok(())
1489}
1490
1491fn binary_probe_clusters(
1492    quantizer: &crate::structures::BinaryCoarseQuantizer,
1493    query: &[u8],
1494    nprobe: usize,
1495    routing: crate::dsl::IvfRoutingMode,
1496    cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
1497) -> Result<std::sync::Arc<[u32]>> {
1498    let effective_nprobe = nprobe.clamp(1, quantizer.num_clusters as usize);
1499    let request_fingerprint = quantizer.request_fingerprint(query, effective_nprobe, routing);
1500    if let Some(cache) = cache {
1501        let mut cached = cache
1502            .lock()
1503            .map_err(|_| Error::Internal("binary IVF probe cache is poisoned".into()))?;
1504        if let Some(plan) = cached.as_ref()
1505            && plan.quantizer_version == quantizer.version
1506            && plan.request_fingerprint == request_fingerprint
1507            && plan.cluster_ids.len() == effective_nprobe
1508        {
1509            return Ok(std::sync::Arc::clone(&plan.cluster_ids));
1510        }
1511        let plan = quantizer.probe(query, effective_nprobe, routing);
1512        let clusters = std::sync::Arc::clone(&plan.cluster_ids);
1513        *cached = Some(plan);
1514        return Ok(clusters);
1515    }
1516    Ok(quantizer
1517        .probe(query, effective_nprobe, routing)
1518        .cluster_ids)
1519}
1520
1521fn binary_scann_probe_clusters(
1522    model: &crate::structures::vector::scann::QuantizedBinaryScannModelView<'_>,
1523    query: &[u8],
1524    nprobe: usize,
1525    cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
1526) -> Result<std::sync::Arc<[u32]>> {
1527    let probes = nprobe.clamp(1, model.num_leaves() as usize);
1528    // Start with a bounded recall beam. The model widens intermediate levels
1529    // when necessary so every requested terminal leaf remains reachable.
1530    let routing_beam = probes.min(64);
1531    let mut fingerprint = 0xcbf2_9ce4_8422_2325u64;
1532    for byte in query.iter().copied().chain((probes as u64).to_le_bytes()) {
1533        fingerprint ^= u64::from(byte);
1534        fingerprint = fingerprint.wrapping_mul(0x0000_0100_0000_01b3);
1535    }
1536    if let Some(cache) = cache {
1537        let mut cached = cache
1538            .lock()
1539            .map_err(|_| Error::Internal("binary ScaNN probe cache is poisoned".into()))?;
1540        if let Some(plan) = cached.as_ref()
1541            && plan.quantizer_version == model.fingerprint()
1542            && plan.request_fingerprint == fingerprint
1543            && plan.cluster_ids.len() == probes
1544        {
1545            return Ok(std::sync::Arc::clone(&plan.cluster_ids));
1546        }
1547        let mut scratch = crate::structures::vector::scann::BinaryScannSearchScratch::default();
1548        let routed = model
1549            .probe(query, probes, routing_beam, &mut scratch)
1550            .map_err(|error| Error::Query(format!("binary ScaNN routing failed: {error}")))?;
1551        let plan =
1552            crate::structures::IvfProbePlan::new(model.fingerprint(), fingerprint, routed.leaf_ids);
1553        let leaves = std::sync::Arc::clone(&plan.cluster_ids);
1554        *cached = Some(plan);
1555        return Ok(leaves);
1556    }
1557    let mut scratch = crate::structures::vector::scann::BinaryScannSearchScratch::default();
1558    model
1559        .probe(query, probes, routing_beam, &mut scratch)
1560        .map(|plan| plan.leaf_ids.into())
1561        .map_err(|error| Error::Query(format!("binary ScaNN routing failed: {error}")))
1562}
1563
1564/// Async segment reader with lazy loading
1565///
1566/// - Term dictionary: only index loaded, blocks loaded on-demand
1567/// - Postings: loaded on-demand per term via HTTP range requests
1568/// - Document store: only index loaded, blocks loaded on-demand via HTTP range requests
1569pub struct SegmentReader {
1570    meta: SegmentMeta,
1571    /// Term dictionary with lazy block loading
1572    term_dict: Arc<AsyncSSTableReader<TermInfo>>,
1573    /// Postings file handle - fetches ranges on demand
1574    postings_handle: FileHandle,
1575    /// Document store with lazy block loading
1576    store: Arc<AsyncStoreReader>,
1577    schema: Arc<Schema>,
1578    /// Per-segment ANN payloads.
1579    vector_indexes: FxHashMap<u32, VectorIndex>,
1580    /// Lazy flat vectors per field — document maps and vectors stay file-backed.
1581    flat_vectors: FxHashMap<u32, LazyFlatVectorData>,
1582    /// Logical size of the retained `.vectors` file handle.
1583    dense_file_backed_bytes: u64,
1584    /// One immutable generation of all index-global ANN artifacts.
1585    trained_vectors: Arc<crate::segment::TrainedVectorStructures>,
1586    /// Sparse vector indexes per field (MaxScore format)
1587    sparse_indexes: FxHashMap<u32, SparseIndex>,
1588    /// BMP sparse vector indexes per field (BMP format)
1589    bmp_indexes: FxHashMap<u32, BmpIndex>,
1590    /// Logical size of the retained `.sparse` file handle.
1591    sparse_file_backed_bytes: u64,
1592    /// Position file handle for phrase queries (lazy loading)
1593    positions_handle: Option<FileHandle>,
1594    /// Fast-field columnar readers per field_id
1595    fast_fields: FxHashMap<u32, crate::structures::fast_field::FastFieldReader>,
1596    /// Virtual-id maps of chunked text fields per field_id
1597    chunk_maps: FxHashMap<u32, super::chunk_map::ChunkMap>,
1598    /// Dense-vector hot-metadata pin accounting (see `segment::pin`).
1599    #[cfg(feature = "native")]
1600    dense_pin_report: crate::segment::pin::PinReport,
1601    /// Sparse-vector hot-metadata pin accounting (see `segment::pin`).
1602    #[cfg(feature = "native")]
1603    sparse_pin_report: crate::segment::pin::PinReport,
1604}
1605
1606impl SegmentReader {
1607    /// Open a segment with lazy loading
1608    pub async fn open<D: Directory>(
1609        dir: &D,
1610        segment_id: SegmentId,
1611        schema: Arc<Schema>,
1612        term_cache_blocks: usize,
1613    ) -> Result<Self> {
1614        Self::open_with_store_cache(
1615            dir,
1616            segment_id,
1617            schema,
1618            term_cache_blocks,
1619            dir as *const D as usize,
1620            Arc::new(super::SharedStoreCache::new(0)),
1621        )
1622        .await
1623    }
1624
1625    /// Open a search segment against the process-wide document-store cache.
1626    pub(crate) async fn open_with_store_cache<D: Directory>(
1627        dir: &D,
1628        segment_id: SegmentId,
1629        schema: Arc<Schema>,
1630        term_cache_blocks: usize,
1631        store_cache_directory_namespace: usize,
1632        store_cache: Arc<super::SharedStoreCache>,
1633    ) -> Result<Self> {
1634        let files = SegmentFiles::new(segment_id.0);
1635
1636        // Read metadata (small, always loaded)
1637        let meta_slice = dir.open_read(&files.meta).await?;
1638        let meta_bytes = meta_slice.read_bytes().await?;
1639        let meta = SegmentMeta::deserialize(meta_bytes.as_slice())?;
1640        debug_assert_eq!(meta.id, segment_id.0);
1641
1642        // Open term dictionary with lazy loading (fetches ranges on demand)
1643        let term_dict_handle = dir.open_lazy(&files.term_dict).await?;
1644        let term_dict = AsyncSSTableReader::open(term_dict_handle, term_cache_blocks).await?;
1645
1646        // Get postings file handle (lazy - fetches ranges on demand)
1647        let postings_handle = dir.open_lazy(&files.postings).await?;
1648
1649        // Open store with lazy loading
1650        let store_handle = dir.open_lazy(&files.store).await?;
1651        let store = AsyncStoreReader::open(
1652            store_handle,
1653            store_cache_directory_namespace,
1654            segment_id.0,
1655            store_cache,
1656        )
1657        .await?;
1658
1659        // Load dense vector indexes from unified .vectors file
1660        let vectors_data = loader::load_vectors_file(dir, &files, &schema, meta.num_docs).await?;
1661        let dense_file_backed_bytes = vectors_data.file_backed_bytes;
1662        let vector_indexes = vectors_data.indexes;
1663        let flat_vectors = vectors_data.flat_vectors;
1664
1665        // Fields served by an ANN index only touch flat vectors for scattered
1666        // rerank reads — disable readahead for them once at open. Flat-only
1667        // fields keep default advice: brute-force scans them sequentially.
1668        // Advice is sticky on the mapping, so per-query re-advising is wasted.
1669        #[cfg(feature = "native")]
1670        for (field_id, lazy_flat) in &flat_vectors {
1671            if vector_indexes.contains_key(field_id) {
1672                lazy_flat.advise_random_access();
1673            }
1674        }
1675
1676        // Load sparse vector indexes from .sparse file (MaxScore + BMP)
1677        let sparse_data = loader::load_sparse_file(dir, &files, meta.num_docs, &schema).await?;
1678        let sparse_file_backed_bytes = sparse_data.file_backed_bytes;
1679        let sparse_indexes = sparse_data.maxscore_indexes;
1680        let bmp_indexes = sparse_data.bmp_indexes;
1681
1682        // Open positions file handle (if exists) - offsets are now in TermInfo
1683        let positions_handle = loader::open_positions_file(dir, &files, &schema).await?;
1684
1685        // Load fast-field columns from .fast file
1686        let fast_fields = loader::load_fast_fields_file(dir, &files, &schema).await?;
1687
1688        // Load chunk maps of chunked text fields from .chunks file
1689        let chunk_maps = loader::load_chunk_maps_file(dir, &files, &schema).await?;
1690
1691        // Log segment loading stats
1692        {
1693            let mut parts = vec![format!(
1694                "[segment] loaded {:016x}: docs={}",
1695                segment_id.0, meta.num_docs
1696            )];
1697            if !vector_indexes.is_empty() || !flat_vectors.is_empty() {
1698                parts.push(format!(
1699                    "dense vectors: {} ANN + {} flat fields",
1700                    vector_indexes.len(),
1701                    flat_vectors.len()
1702                ));
1703            }
1704            for (field_id, idx) in &sparse_indexes {
1705                parts.push(format!(
1706                    "sparse vector field {}: {} dims, ~{}",
1707                    field_id,
1708                    idx.num_dimensions(),
1709                    crate::format_bytes(idx.num_dimensions() as u64 * 24)
1710                ));
1711            }
1712            for (field_id, idx) in &bmp_indexes {
1713                parts.push(format!(
1714                    "bmp field {}: {} dims, {} blocks",
1715                    field_id,
1716                    idx.dims(),
1717                    idx.num_blocks
1718                ));
1719            }
1720            if !fast_fields.is_empty() {
1721                parts.push(format!("fast: {} fields", fast_fields.len()));
1722            }
1723            for (field_id, map) in &chunk_maps {
1724                parts.push(format!(
1725                    "chunked text field {}: {} chunks",
1726                    field_id,
1727                    map.num_chunks()
1728                ));
1729            }
1730            log::debug!("{}", parts.join(", "));
1731        }
1732
1733        #[allow(unused_mut)]
1734        let mut reader = Self {
1735            meta,
1736            term_dict: Arc::new(term_dict),
1737            postings_handle,
1738            store: Arc::new(store),
1739            schema,
1740            vector_indexes,
1741            flat_vectors,
1742            dense_file_backed_bytes,
1743            trained_vectors: Arc::new(crate::segment::TrainedVectorStructures::default()),
1744            sparse_indexes,
1745            bmp_indexes,
1746            sparse_file_backed_bytes,
1747            positions_handle,
1748            fast_fields,
1749            chunk_maps,
1750            #[cfg(feature = "native")]
1751            dense_pin_report: Default::default(),
1752            #[cfg(feature = "native")]
1753            sparse_pin_report: Default::default(),
1754        };
1755
1756        // Pin hot metadata per the process-wide policy (no-op when disabled)
1757        #[cfg(feature = "native")]
1758        reader.apply_pin_policy(&crate::segment::pin::pin_policy().to_owned());
1759
1760        // Structural ANN health from the already-parsed run directories —
1761        // O(runs) per field, no payload reads. This is the passive tier of
1762        // `docs/diagnostics.md`: leaf collapse and extent fragmentation warn
1763        // here instead of surfacing as unexplained latency.
1764        for (&field_id, vector_index) in &reader.vector_indexes {
1765            match vector_index {
1766                VectorIndex::BinaryIvf(index)
1767                | VectorIndex::IvfTq { index, .. }
1768                | VectorIndex::ScannAh(index)
1769                | VectorIndex::ScannBinary(index) => {
1770                    index.get().report_health(
1771                        reader.schema.index_label(),
1772                        field_id,
1773                        reader.meta.id,
1774                    );
1775                }
1776                // TQ flat payloads have no cluster structure; skew and
1777                // fragmentation metrics would be meaningless there.
1778                VectorIndex::Tq { .. } => {}
1779            }
1780        }
1781
1782        Ok(reader)
1783    }
1784
1785    /// Structural health of one field's IVF payload, if it has one.
1786    ///
1787    /// Cheap (O(runs) over in-memory data); exposed for `hermes-tool diagnose`.
1788    pub fn ann_health(&self, field: Field) -> Option<crate::segment::ann_disk::AnnHealth> {
1789        match self.vector_indexes.get(&field.0)? {
1790            VectorIndex::BinaryIvf(index)
1791            | VectorIndex::IvfTq { index, .. }
1792            | VectorIndex::ScannAh(index)
1793            | VectorIndex::ScannBinary(index) => Some(index.get().health()),
1794            VectorIndex::Tq { .. } => None,
1795        }
1796    }
1797
1798    /// Pin per-query-mandatory metadata sections in priority order until the
1799    /// budget is exhausted (see `segment::pin` and docs/hot-metadata-pinning.md).
1800    ///
1801    /// Priority: ANN run directories → BMP block-offset tables → sparse skip
1802    /// sections → doc-id maps → BMP E offsets + coarse H. Bulk data (ANN codes,
1803    /// D/E grid payloads, block data, raw vectors) is never pinned. Fail-loud: budget
1804    /// exhaustion and mlock failures are
1805    /// logged and visible via `SegmentMemoryStats::{pin_intended_bytes,
1806    /// pinned_metadata_bytes}`.
1807    #[cfg(feature = "native")]
1808    pub(crate) fn apply_pin_policy(&mut self, policy: &crate::segment::pin::PinPolicy) {
1809        use crate::segment::pin::PinReport;
1810
1811        if !policy.is_enabled() {
1812            return;
1813        }
1814        let mut remaining = policy.budget_bytes;
1815        let mut dense_report = PinReport::default();
1816        let mut sparse_report = PinReport::default();
1817
1818        // Priority 1: compact ANN lookup directories
1819        for index in self.vector_indexes.values_mut() {
1820            index.pin_lookup_directory(policy.mode, &mut remaining, &mut dense_report);
1821        }
1822        // Priority 2: BMP block-offset tables
1823        for bmp in self.bmp_indexes.values_mut() {
1824            bmp.pin_block_starts(policy.mode, &mut remaining, &mut sparse_report);
1825        }
1826        // Priority 3: sparse skip sections
1827        for sparse in self.sparse_indexes.values_mut() {
1828            sparse.pin_skip_section(policy.mode, &mut remaining, &mut sparse_report);
1829        }
1830        // Priority 4: doc-id maps
1831        for flat in self.flat_vectors.values_mut() {
1832            flat.pin_doc_ids(policy.mode, &mut remaining, &mut dense_report);
1833        }
1834        for bmp in self.bmp_indexes.values_mut() {
1835            bmp.pin_doc_maps(policy.mode, &mut remaining, &mut sparse_report);
1836        }
1837        // Priority 5: BMP E offsets and coarse H
1838        for bmp in self.bmp_indexes.values_mut() {
1839            bmp.pin_query_hierarchy(policy.mode, &mut remaining, &mut sparse_report);
1840        }
1841
1842        let report = PinReport {
1843            intended_bytes: dense_report
1844                .intended_bytes
1845                .saturating_add(sparse_report.intended_bytes),
1846            pinned_bytes: dense_report
1847                .pinned_bytes
1848                .saturating_add(sparse_report.pinned_bytes),
1849            skipped_budget_bytes: dense_report
1850                .skipped_budget_bytes
1851                .saturating_add(sparse_report.skipped_budget_bytes),
1852            failed_bytes: dense_report
1853                .failed_bytes
1854                .saturating_add(sparse_report.failed_bytes),
1855            heap_copy_bytes: dense_report
1856                .heap_copy_bytes
1857                .saturating_add(sparse_report.heap_copy_bytes),
1858        };
1859        if report.skipped_budget_bytes > 0 || report.failed_bytes > 0 {
1860            log::warn!(
1861                "[pin] index={} segment {:016x}: pinned {}/{} (budget skipped {}, mlock failed {}) — \
1862                 raise HERMES_PIN_METADATA_BUDGET_MB or RLIMIT_MEMLOCK for full coverage",
1863                self.schema.index_label(),
1864                self.meta.id,
1865                crate::format_bytes(report.pinned_bytes),
1866                crate::format_bytes(report.intended_bytes),
1867                crate::format_bytes(report.skipped_budget_bytes),
1868                crate::format_bytes(report.failed_bytes),
1869            );
1870        } else if report.pinned_bytes > 0 {
1871            log::info!(
1872                "[pin] index={} segment {:016x}: pinned {} of hot metadata ({:?})",
1873                self.schema.index_label(),
1874                self.meta.id,
1875                crate::format_bytes(report.pinned_bytes),
1876                policy.mode,
1877            );
1878        }
1879        self.dense_pin_report = dense_report;
1880        self.sparse_pin_report = sparse_report;
1881    }
1882
1883    // NOTE: cross-group MaxScore threshold seeding is query-execution-local
1884    // (a Cell in the boolean planner) — it must never live on the shared
1885    // SegmentReader, where concurrent queries would leak thresholds into
1886    // each other and wrongly prune results.
1887
1888    pub fn meta(&self) -> &SegmentMeta {
1889        &self.meta
1890    }
1891
1892    pub fn num_docs(&self) -> u32 {
1893        self.meta.num_docs
1894    }
1895
1896    /// Get average field length for BM25F scoring
1897    pub fn avg_field_len(&self, field: Field) -> f32 {
1898        self.meta.avg_field_len(field)
1899    }
1900
1901    pub fn schema(&self) -> &Schema {
1902        &self.schema
1903    }
1904
1905    /// Get sparse indexes for all fields
1906    pub fn sparse_indexes(&self) -> &FxHashMap<u32, SparseIndex> {
1907        &self.sparse_indexes
1908    }
1909
1910    /// Get sparse index for a specific field (MaxScore format)
1911    pub fn sparse_index(&self, field: Field) -> Option<&SparseIndex> {
1912        self.sparse_indexes.get(&field.0)
1913    }
1914
1915    /// Get BMP index for a specific field
1916    pub fn bmp_index(&self, field: Field) -> Option<&BmpIndex> {
1917        self.bmp_indexes.get(&field.0)
1918    }
1919
1920    /// Get all BMP indexes
1921    pub fn bmp_indexes(&self) -> &FxHashMap<u32, BmpIndex> {
1922        &self.bmp_indexes
1923    }
1924
1925    /// Get vector indexes for all fields
1926    pub fn vector_indexes(&self) -> &FxHashMap<u32, VectorIndex> {
1927        &self.vector_indexes
1928    }
1929
1930    /// Get lazy flat vectors for all fields (for reranking and merge)
1931    pub fn flat_vectors(&self) -> &FxHashMap<u32, LazyFlatVectorData> {
1932        &self.flat_vectors
1933    }
1934
1935    /// Get a fast-field reader for a specific field.
1936    pub fn fast_field(
1937        &self,
1938        field_id: u32,
1939    ) -> Option<&crate::structures::fast_field::FastFieldReader> {
1940        self.fast_fields.get(&field_id)
1941    }
1942
1943    /// Get all fast-field readers.
1944    pub fn fast_fields(&self) -> &FxHashMap<u32, crate::structures::fast_field::FastFieldReader> {
1945        &self.fast_fields
1946    }
1947
1948    /// Virtual-id map of a chunked text field, when the field is chunked and
1949    /// this segment indexed at least one chunk of it.
1950    pub fn chunk_map(&self, field: Field) -> Option<&super::chunk_map::ChunkMap> {
1951        self.chunk_maps.get(&field.0)
1952    }
1953
1954    /// All chunk maps of this segment.
1955    pub fn chunk_maps(&self) -> &FxHashMap<u32, super::chunk_map::ChunkMap> {
1956        &self.chunk_maps
1957    }
1958
1959    /// Whether `field` is declared chunked in the schema (its postings are
1960    /// keyed by virtual chunk ids, never by document ids).
1961    pub fn is_chunked_field(&self, field: Field) -> bool {
1962        self.schema
1963            .get_field_entry(field)
1964            .is_some_and(|entry| entry.chunked)
1965    }
1966
1967    /// Number of chunks a chunked field holds in this segment (0 when none).
1968    pub fn num_chunks(&self, field: Field) -> u32 {
1969        self.chunk_maps
1970            .get(&field.0)
1971            .map_or(0, |map| map.num_chunks())
1972    }
1973
1974    /// BM25 corpus size for `field`: chunks for a chunked field, documents
1975    /// otherwise.
1976    pub fn text_corpus_size(&self, field: Field) -> f32 {
1977        if self.is_chunked_field(field) {
1978            self.num_chunks(field) as f32
1979        } else {
1980            self.meta.num_docs as f32
1981        }
1982    }
1983
1984    /// Whether this segment carries a `.chunks` file.
1985    pub fn has_chunks_file(&self) -> bool {
1986        !self.chunk_maps.is_empty()
1987    }
1988
1989    /// Get term dictionary stats for debugging
1990    pub fn term_dict_stats(&self) -> SSTableStats {
1991        self.term_dict.stats()
1992    }
1993
1994    /// Account for heap, file-backed, and pinned bytes separately.
1995    pub fn memory_stats(&self) -> SegmentMemoryStats {
1996        let term_dict_stats = self.term_dict.stats();
1997
1998        // Report actual decompressed heap retention. Both caches use variable
1999        // boundary blocks, so multiplying a block count by a guessed size can
2000        // materially under-report resident memory.
2001        let term_dict_cache_bytes = self.term_dict.cached_bytes();
2002        let store_cache_bytes = self.store.cached_bytes();
2003
2004        // Sparse heap: SoA dimension tables and small reader objects. Posting
2005        // payloads, BMP grids, and document maps remain file-backed.
2006        let sparse_heap_bytes: usize = self
2007            .sparse_indexes
2008            .values()
2009            .map(|s| s.estimated_heap_bytes())
2010            .sum::<usize>()
2011            + self
2012                .bmp_indexes
2013                .values()
2014                .map(|b| b.estimated_heap_bytes())
2015                .sum::<usize>();
2016
2017        // Dense corpus columns are file-backed. Only compact ANN run
2018        // directories and flat-reader objects count as heap here.
2019        let dense_heap_bytes: usize = self
2020            .vector_indexes
2021            .values()
2022            .map(|v| v.estimated_heap_bytes())
2023            .sum::<usize>()
2024            + self
2025                .flat_vectors
2026                .values()
2027                .map(LazyFlatVectorData::estimated_heap_bytes)
2028                .sum::<usize>();
2029
2030        #[cfg(feature = "native")]
2031        let (sparse_heap_bytes, dense_heap_bytes) = (
2032            sparse_heap_bytes.saturating_add(
2033                usize::try_from(self.sparse_pin_report.heap_copy_bytes).unwrap_or(usize::MAX),
2034            ),
2035            dense_heap_bytes.saturating_add(
2036                usize::try_from(self.dense_pin_report.heap_copy_bytes).unwrap_or(usize::MAX),
2037            ),
2038        );
2039
2040        #[cfg(feature = "native")]
2041        let (
2042            sparse_pinned_metadata_bytes,
2043            sparse_pin_intended_bytes,
2044            dense_pinned_metadata_bytes,
2045            dense_pin_intended_bytes,
2046        ) = (
2047            self.sparse_pin_report.pinned_bytes,
2048            self.sparse_pin_report.intended_bytes,
2049            self.dense_pin_report.pinned_bytes,
2050            self.dense_pin_report.intended_bytes,
2051        );
2052        #[cfg(not(feature = "native"))]
2053        let (
2054            sparse_pinned_metadata_bytes,
2055            sparse_pin_intended_bytes,
2056            dense_pinned_metadata_bytes,
2057            dense_pin_intended_bytes,
2058        ) = (0u64, 0u64, 0u64, 0u64);
2059
2060        let pinned_metadata_bytes =
2061            sparse_pinned_metadata_bytes.saturating_add(dense_pinned_metadata_bytes);
2062        let pin_intended_bytes = sparse_pin_intended_bytes.saturating_add(dense_pin_intended_bytes);
2063
2064        SegmentMemoryStats {
2065            segment_id: self.meta.id,
2066            num_docs: self.meta.num_docs,
2067            term_dict_cache_bytes,
2068            store_cache_bytes,
2069            sparse_heap_bytes,
2070            dense_heap_bytes,
2071            term_bloom_file_bytes: term_dict_stats.bloom_filter_size as u64,
2072            sparse_file_backed_bytes: self.sparse_file_backed_bytes,
2073            dense_file_backed_bytes: self.dense_file_backed_bytes,
2074            pinned_metadata_bytes,
2075            pin_intended_bytes,
2076            sparse_pinned_metadata_bytes,
2077            sparse_pin_intended_bytes,
2078            dense_pinned_metadata_bytes,
2079            dense_pin_intended_bytes,
2080        }
2081    }
2082
2083    /// Get posting list for a term (async - loads on demand)
2084    ///
2085    /// For small posting lists (1-3 docs), the data is inlined in the term dictionary
2086    /// and no additional I/O is needed. For larger lists, reads from .post file.
2087    pub async fn get_postings(
2088        &self,
2089        field: Field,
2090        term: &[u8],
2091    ) -> Result<Option<BlockPostingList>> {
2092        log::debug!(
2093            "SegmentReader::get_postings field={} term_len={}",
2094            field.0,
2095            term.len()
2096        );
2097
2098        // Build key: field_id + term
2099        let mut key = Vec::with_capacity(4 + term.len());
2100        key.extend_from_slice(&field.0.to_le_bytes());
2101        key.extend_from_slice(term);
2102
2103        // Look up in term dictionary
2104        let term_info = match self.term_dict.get(&key).await? {
2105            Some(info) => {
2106                log::debug!("SegmentReader::get_postings found term_info");
2107                info
2108            }
2109            None => {
2110                log::debug!("SegmentReader::get_postings term not found");
2111                return Ok(None);
2112            }
2113        };
2114
2115        // Check if posting list is inlined
2116        if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
2117            // Build BlockPostingList from inline data (no I/O needed!)
2118            let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
2119            for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
2120                posting_list.push(doc_id, tf);
2121            }
2122            let block_list = BlockPostingList::from_posting_list(&posting_list)?;
2123            return Ok(Some(block_list));
2124        }
2125
2126        // External posting list - read from postings file handle (lazy - HTTP range request)
2127        let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
2128            Error::Corruption("TermInfo has neither inline nor external data".to_string())
2129        })?;
2130
2131        let range = checked_file_range(
2132            posting_offset,
2133            posting_len,
2134            self.postings_handle.len(),
2135            "posting",
2136        )?;
2137        let posting_bytes = self.postings_handle.read_bytes_range(range).await?;
2138        let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
2139
2140        Ok(Some(block_list))
2141    }
2142
2143    /// Get all posting lists for terms that start with `prefix` in the given field.
2144    pub async fn get_prefix_postings(
2145        &self,
2146        field: Field,
2147        prefix: &[u8],
2148    ) -> Result<Vec<BlockPostingList>> {
2149        if prefix.is_empty() {
2150            return Err(Error::Query("prefix must not be empty".into()));
2151        }
2152        // Build composite key prefix: field_id ++ prefix
2153        let mut key_prefix = Vec::with_capacity(4 + prefix.len());
2154        key_prefix.extend_from_slice(&field.0.to_le_bytes());
2155        key_prefix.extend_from_slice(prefix);
2156
2157        let (entries, truncated) = self
2158            .term_dict
2159            .prefix_scan_limited(&key_prefix, MAX_PREFIX_TERMS)
2160            .await?;
2161        if truncated {
2162            return Err(Error::Query(format!(
2163                "prefix expands to more than {MAX_PREFIX_TERMS} terms"
2164            )));
2165        }
2166        let posting_count: u64 = entries
2167            .iter()
2168            .map(|(_, term_info)| term_info.doc_freq() as u64)
2169            .sum();
2170        if posting_count > MAX_PREFIX_POSTINGS {
2171            return Err(Error::Query(format!(
2172                "prefix expands to {posting_count} postings (maximum {MAX_PREFIX_POSTINGS})"
2173            )));
2174        }
2175        let mut results = Vec::with_capacity(entries.len());
2176
2177        for (_key, term_info) in entries {
2178            if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
2179                let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
2180                for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
2181                    posting_list.push(doc_id, tf);
2182                }
2183                results.push(BlockPostingList::from_posting_list(&posting_list)?);
2184            } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
2185                let range = checked_file_range(
2186                    posting_offset,
2187                    posting_len,
2188                    self.postings_handle.len(),
2189                    "prefix posting",
2190                )?;
2191                let posting_bytes = self.postings_handle.read_bytes_range(range).await?;
2192                results.push(BlockPostingList::deserialize_zero_copy(posting_bytes)?);
2193            }
2194        }
2195
2196        Ok(results)
2197    }
2198
2199    /// Get document by local doc_id (async - loads on demand).
2200    ///
2201    /// Dense vector fields are hydrated from LazyFlatVectorData (not stored in .store).
2202    /// Uses binary search on sorted doc_ids for O(log N) lookup.
2203    pub async fn doc(&self, local_doc_id: DocId) -> Result<Option<Document>> {
2204        self.doc_with_fields(local_doc_id, None).await
2205    }
2206
2207    /// Get document by local doc_id, hydrating only the specified fields.
2208    ///
2209    /// If `fields` is `None`, all fields (including dense vectors) are hydrated.
2210    /// If `fields` is `Some(set)`, only dense vector fields in the set are hydrated,
2211    /// skipping expensive mmap reads + dequantization for unrequested vector fields.
2212    pub async fn doc_with_fields(
2213        &self,
2214        local_doc_id: DocId,
2215        fields: Option<&rustc_hash::FxHashSet<u32>>,
2216    ) -> Result<Option<Document>> {
2217        let mut doc = match fields {
2218            Some(set) => {
2219                let field_ids: Vec<u32> = set.iter().copied().collect();
2220                match self
2221                    .store
2222                    .get_fields(local_doc_id, &self.schema, &field_ids)
2223                    .await
2224                {
2225                    Ok(Some(d)) => d,
2226                    Ok(None) => return Ok(None),
2227                    Err(e) => return Err(Error::from(e)),
2228                }
2229            }
2230            None => match self.store.get(local_doc_id, &self.schema).await {
2231                Ok(Some(d)) => d,
2232                Ok(None) => return Ok(None),
2233                Err(e) => return Err(Error::from(e)),
2234            },
2235        };
2236
2237        // Hydrate dense vector fields from flat vector data
2238        for (&field_id, lazy_flat) in &self.flat_vectors {
2239            // Skip vector fields not in the requested set
2240            if let Some(set) = fields
2241                && !set.contains(&field_id)
2242            {
2243                continue;
2244            }
2245
2246            let is_binary = lazy_flat.quantization == DenseVectorQuantization::Binary;
2247            let (start, entries) = lazy_flat.flat_indexes_for_doc(local_doc_id);
2248            for (j, &(_doc_id, _ordinal)) in entries.iter().enumerate() {
2249                let flat_idx = start + j;
2250                if is_binary {
2251                    let vbs = lazy_flat.vector_byte_size();
2252                    let mut raw = vec![0u8; vbs];
2253                    match lazy_flat.read_vector_raw_into(flat_idx, &mut raw).await {
2254                        Ok(()) => {
2255                            doc.add_binary_dense_vector(Field(field_id), raw);
2256                        }
2257                        Err(e) => {
2258                            log::warn!(
2259                                "Failed to hydrate binary dense vector field {}: {}",
2260                                field_id,
2261                                e
2262                            );
2263                        }
2264                    }
2265                } else {
2266                    match lazy_flat.get_vector(flat_idx).await {
2267                        Ok(vec) => {
2268                            doc.add_dense_vector(Field(field_id), vec);
2269                        }
2270                        Err(e) => {
2271                            log::warn!("Failed to hydrate dense vector field {}: {}", field_id, e);
2272                        }
2273                    }
2274                }
2275            }
2276        }
2277
2278        Ok(Some(doc))
2279    }
2280
2281    /// Prefetch term dictionary blocks for a key range
2282    pub async fn prefetch_terms(
2283        &self,
2284        field: Field,
2285        start_term: &[u8],
2286        end_term: &[u8],
2287    ) -> Result<()> {
2288        let mut start_key = Vec::with_capacity(4 + start_term.len());
2289        start_key.extend_from_slice(&field.0.to_le_bytes());
2290        start_key.extend_from_slice(start_term);
2291
2292        let mut end_key = Vec::with_capacity(4 + end_term.len());
2293        end_key.extend_from_slice(&field.0.to_le_bytes());
2294        end_key.extend_from_slice(end_term);
2295
2296        self.term_dict.prefetch_range(&start_key, &end_key).await?;
2297        Ok(())
2298    }
2299
2300    /// Check if store uses dictionary compression (incompatible with raw merging)
2301    pub fn store_has_dict(&self) -> bool {
2302        self.store.has_dict()
2303    }
2304
2305    /// Get store reference for merge operations
2306    pub fn store(&self) -> &super::store::AsyncStoreReader {
2307        &self.store
2308    }
2309
2310    /// Get raw store blocks for optimized merging
2311    pub fn store_raw_blocks(&self) -> Vec<RawStoreBlock> {
2312        self.store.raw_blocks()
2313    }
2314
2315    /// Get store data slice for raw block access
2316    pub fn store_data_slice(&self) -> &FileHandle {
2317        self.store.data_slice()
2318    }
2319
2320    /// Get all terms from this segment (for merge)
2321    pub async fn all_terms(&self) -> Result<Vec<(Vec<u8>, TermInfo)>> {
2322        self.term_dict.all_entries().await.map_err(Error::from)
2323    }
2324
2325    /// Get all terms with parsed field and term string (for statistics aggregation)
2326    ///
2327    /// Returns (field, term_string, doc_freq) for each term in the dictionary.
2328    /// Skips terms that aren't valid UTF-8.
2329    pub async fn all_terms_with_stats(&self) -> Result<Vec<(Field, String, u32)>> {
2330        let entries = self.term_dict.all_entries().await?;
2331        let mut result = Vec::with_capacity(entries.len());
2332
2333        for (key, term_info) in entries {
2334            // Key format: field_id (4 bytes little-endian) + term bytes
2335            if key.len() > 4 {
2336                let field_id = u32::from_le_bytes([key[0], key[1], key[2], key[3]]);
2337                let term_bytes = &key[4..];
2338                if let Ok(term_str) = std::str::from_utf8(term_bytes) {
2339                    result.push((Field(field_id), term_str.to_string(), term_info.doc_freq()));
2340                }
2341            }
2342        }
2343
2344        Ok(result)
2345    }
2346
2347    /// Get streaming iterator over term dictionary (for memory-efficient merge)
2348    pub fn term_dict_iter(&self) -> crate::structures::AsyncSSTableIterator<'_, TermInfo> {
2349        self.term_dict.iter()
2350    }
2351
2352    /// Prefetch all term dictionary blocks in a single bulk I/O call.
2353    ///
2354    /// Call before merge iteration to eliminate per-block cache misses.
2355    pub async fn prefetch_term_dict(&self) -> crate::Result<()> {
2356        self.term_dict
2357            .prefetch_all_data_bulk()
2358            .await
2359            .map_err(crate::Error::from)
2360    }
2361
2362    /// Read raw posting bytes at offset
2363    pub async fn read_postings(&self, offset: u64, len: u64) -> Result<Vec<u8>> {
2364        let range = checked_file_range(offset, len, self.postings_handle.len(), "posting")?;
2365        let bytes = self.postings_handle.read_bytes_range(range).await?;
2366        Ok(bytes.to_vec())
2367    }
2368
2369    /// Read raw position bytes at offset (for merge)
2370    pub async fn read_position_bytes(&self, offset: u64, len: u64) -> Result<Option<Vec<u8>>> {
2371        let handle = match &self.positions_handle {
2372            Some(h) => h,
2373            None => return Ok(None),
2374        };
2375        let range = checked_file_range(offset, len, handle.len(), "position")?;
2376        let bytes = handle.read_bytes_range(range).await?;
2377        Ok(Some(bytes.to_vec()))
2378    }
2379
2380    /// Check if this segment has a positions file
2381    pub fn has_positions_file(&self) -> bool {
2382        self.positions_handle.is_some()
2383    }
2384
2385    /// Validate all caller-controlled dense-search inputs before touching ANN
2386    /// structures or entering SIMD code. This is deliberately repeated at the
2387    /// segment boundary so non-server users receive the same safety guarantees.
2388    fn validate_dense_search_request(
2389        &self,
2390        field: Field,
2391        query: &[f32],
2392        nprobe: usize,
2393        rerank_factor: f32,
2394        combiner: crate::query::MultiValueCombiner,
2395    ) -> Result<DenseSearchParams> {
2396        let entry = self
2397            .schema
2398            .get_field_entry(field)
2399            .ok_or_else(|| Error::FieldNotFound(field.0.to_string()))?;
2400        if entry.field_type != crate::dsl::FieldType::DenseVector {
2401            return Err(Error::InvalidFieldType {
2402                expected: "dense_vector".to_string(),
2403                got: format!("{:?}", entry.field_type),
2404            });
2405        }
2406        let config = entry.dense_vector_config.as_ref().ok_or_else(|| {
2407            Error::Schema(format!(
2408                "dense vector field '{}' has no dense vector configuration",
2409                entry.name
2410            ))
2411        })?;
2412
2413        if query.is_empty() {
2414            return Err(Error::Query(format!(
2415                "dense query vector for field '{}' must not be empty",
2416                entry.name
2417            )));
2418        }
2419        if query.len() != config.dim {
2420            return Err(Error::Query(format!(
2421                "dense query vector dimension {} does not match field '{}' dimension {}",
2422                query.len(),
2423                entry.name,
2424                config.dim
2425            )));
2426        }
2427        if let Some((index, value)) = query
2428            .iter()
2429            .enumerate()
2430            .find(|(_, value)| !value.is_finite())
2431        {
2432            return Err(Error::Query(format!(
2433                "dense query vector for field '{}' contains non-finite value {value} at index {index}",
2434                entry.name
2435            )));
2436        }
2437
2438        // A zero query override means "use the schema". Legacy schemas may
2439        // contain zero for flat fields, so retain 32 as a final ANN fallback.
2440        let nprobe = match (nprobe, config.nprobe) {
2441            (0, 0) => 32,
2442            (0, schema_nprobe) => schema_nprobe,
2443            (query_nprobe, _) => query_nprobe,
2444        };
2445        if nprobe > MAX_DENSE_NPROBE {
2446            return Err(Error::Query(format!(
2447                "dense nprobe must be at most {MAX_DENSE_NPROBE}, got {nprobe}"
2448            )));
2449        }
2450
2451        // Validate the factor here even for empty segments. Otherwise malformed
2452        // requests would succeed or fail depending on segment contents.
2453        checked_dense_fetch_k(0, rerank_factor)?;
2454        combiner.validate().map_err(Error::Query)?;
2455
2456        Ok(DenseSearchParams {
2457            dim: config.dim,
2458            nprobe,
2459            unit_norm: config.unit_norm,
2460        })
2461    }
2462
2463    fn validate_binary_search_request(&self, field: Field, query: &[u8]) -> Result<usize> {
2464        let entry = self
2465            .schema
2466            .get_field_entry(field)
2467            .ok_or_else(|| Error::FieldNotFound(field.0.to_string()))?;
2468        if entry.field_type != crate::dsl::FieldType::BinaryDenseVector {
2469            return Err(Error::InvalidFieldType {
2470                expected: "binary_dense_vector".to_string(),
2471                got: format!("{:?}", entry.field_type),
2472            });
2473        }
2474        let config = entry.binary_dense_vector_config.as_ref().ok_or_else(|| {
2475            Error::Schema(format!(
2476                "binary dense vector field '{}' has no configuration",
2477                entry.name
2478            ))
2479        })?;
2480        if config.dim == 0 || !config.dim.is_multiple_of(8) {
2481            return Err(Error::Schema(format!(
2482                "binary dense vector field '{}' has invalid dimension {}",
2483                entry.name, config.dim
2484            )));
2485        }
2486        if query.len() != config.byte_len() {
2487            return Err(Error::Query(format!(
2488                "binary query byte length {} does not match field '{}' byte length {}",
2489                query.len(),
2490                entry.name,
2491                config.byte_len()
2492            )));
2493        }
2494        Ok(config.dim)
2495    }
2496
2497    /// Previous per-batch preparation path retained as an equivalence oracle.
2498    #[cfg(test)]
2499    fn score_quantized_batch_legacy(
2500        query: &[f32],
2501        raw: &[u8],
2502        quant: crate::dsl::DenseVectorQuantization,
2503        dim: usize,
2504        scores: &mut [f32],
2505        unit_norm: bool,
2506    ) -> Result<()> {
2507        use crate::dsl::DenseVectorQuantization;
2508        use crate::structures::simd;
2509
2510        if query.len() != dim {
2511            return Err(Error::Query(format!(
2512                "dense SIMD query dimension {} does not match vector dimension {dim}",
2513                query.len()
2514            )));
2515        }
2516        let element_size = match quant {
2517            DenseVectorQuantization::F32 => std::mem::size_of::<f32>(),
2518            DenseVectorQuantization::F16 => std::mem::size_of::<u16>(),
2519            DenseVectorQuantization::UInt8 => 1,
2520            DenseVectorQuantization::Binary => {
2521                return Err(Error::InvalidFieldType {
2522                    expected: "non-binary dense vector".to_string(),
2523                    got: "binary dense vector".to_string(),
2524                });
2525            }
2526        };
2527        let required_bytes = scores
2528            .len()
2529            .checked_mul(dim)
2530            .and_then(|elements| elements.checked_mul(element_size))
2531            .ok_or_else(|| Error::Corruption("dense vector batch byte length overflow".into()))?;
2532        if raw.len() < required_bytes {
2533            return Err(Error::Corruption(format!(
2534                "dense vector batch is truncated: need {required_bytes} bytes, got {}",
2535                raw.len()
2536            )));
2537        }
2538        if quant == DenseVectorQuantization::F16
2539            && required_bytes > 0
2540            && !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<u16>())
2541        {
2542            return Err(Error::Corruption(
2543                "f16 vector data is not 2-byte aligned".to_string(),
2544            ));
2545        }
2546
2547        match (quant, unit_norm) {
2548            (DenseVectorQuantization::F32, false) => {
2549                let num_floats = scores.len() * dim;
2550                if !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()) {
2551                    return Err(Error::Corruption(
2552                        "f32 vector data is not 4-byte aligned".to_string(),
2553                    ));
2554                }
2555                let vectors: &[f32] =
2556                    unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
2557                simd::batch_cosine_scores(query, vectors, dim, scores);
2558            }
2559            (DenseVectorQuantization::F32, true) => {
2560                let num_floats = scores.len() * dim;
2561                if !(raw.as_ptr() as usize).is_multiple_of(std::mem::align_of::<f32>()) {
2562                    return Err(Error::Corruption(
2563                        "f32 vector data is not 4-byte aligned".to_string(),
2564                    ));
2565                }
2566                let vectors: &[f32] =
2567                    unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const f32, num_floats) };
2568                simd::batch_dot_scores(query, vectors, dim, scores);
2569            }
2570            (DenseVectorQuantization::F16, false) => {
2571                simd::batch_cosine_scores_f16(query, raw, dim, scores);
2572            }
2573            (DenseVectorQuantization::F16, true) => {
2574                simd::batch_dot_scores_f16(query, raw, dim, scores);
2575            }
2576            (DenseVectorQuantization::UInt8, false) => {
2577                simd::batch_cosine_scores_u8(query, raw, dim, scores);
2578            }
2579            (DenseVectorQuantization::UInt8, true) => {
2580                simd::batch_dot_scores_u8(query, raw, dim, scores);
2581            }
2582            (DenseVectorQuantization::Binary, _) => unreachable!("validated above"),
2583        }
2584        Ok(())
2585    }
2586
2587    /// Search dense vectors through the production IVF-PQ index.
2588    ///
2589    /// Returns VectorSearchResult with ordinal tracking for multi-value fields.
2590    /// Doc IDs are segment-local.
2591    /// For multi-valued documents, scores are combined using the specified combiner.
2592    pub async fn search_dense_vector(
2593        &self,
2594        field: Field,
2595        query: &[f32],
2596        k: usize,
2597        nprobe: usize,
2598        rerank_factor: f32,
2599        combiner: crate::query::MultiValueCombiner,
2600    ) -> Result<Vec<VectorSearchResult>> {
2601        self.search_dense_vector_impl(field, query, k, nprobe, rerank_factor, combiner, None)
2602            .await
2603    }
2604
2605    #[allow(clippy::too_many_arguments)]
2606    pub(crate) async fn search_dense_vector_with_probe_cache(
2607        &self,
2608        field: Field,
2609        query: &[f32],
2610        k: usize,
2611        nprobe: usize,
2612        rerank_factor: f32,
2613        combiner: crate::query::MultiValueCombiner,
2614        plan_cache: &DensePlanCache,
2615    ) -> Result<Vec<VectorSearchResult>> {
2616        self.search_dense_vector_impl(
2617            field,
2618            query,
2619            k,
2620            nprobe,
2621            rerank_factor,
2622            combiner,
2623            Some(plan_cache),
2624        )
2625        .await
2626    }
2627
2628    #[allow(clippy::too_many_arguments)]
2629    async fn search_dense_vector_impl(
2630        &self,
2631        field: Field,
2632        query: &[f32],
2633        k: usize,
2634        nprobe: usize,
2635        rerank_factor: f32,
2636        combiner: crate::query::MultiValueCombiner,
2637        plan_cache: Option<&DensePlanCache>,
2638    ) -> Result<Vec<VectorSearchResult>> {
2639        let params =
2640            self.validate_dense_search_request(field, query, nprobe, rerank_factor, combiner)?;
2641        let fetch_k = checked_dense_fetch_k(k, rerank_factor)?;
2642        if k == 0 {
2643            return Ok(Vec::new());
2644        }
2645
2646        let configured_ann_index = self.vector_indexes.get(&field.0);
2647        let lazy_flat = self.flat_vectors.get(&field.0);
2648        // No vectors at all for this field
2649        if configured_ann_index.is_none() && lazy_flat.is_none() {
2650            return Ok(Vec::new());
2651        }
2652
2653        if configured_ann_index.is_some() && lazy_flat.is_none() {
2654            return Err(Error::Corruption(format!(
2655                "dense ANN field {} is missing flat vector storage",
2656                field.0
2657            )));
2658        }
2659
2660        if let Some(flat) = lazy_flat
2661            && flat.dim != params.dim
2662        {
2663            return Err(Error::Corruption(format!(
2664                "dense vector field {} has schema dimension {} but flat storage dimension {}",
2665                field.0, params.dim, flat.dim
2666            )));
2667        }
2668
2669        let needs_document_aggregation = lazy_flat.is_some_and(|flat| {
2670            flat.num_vectors != flat.num_docs_with_vectors()
2671                && !matches!(combiner, crate::query::MultiValueCombiner::Max)
2672        });
2673        // Keep every configured ANN index active. Multi-value semantics are
2674        // handled by bounded combiner-aware scans; IVF-TQ accepts only the
2675        // cosine-normalized generation validated below.
2676        let ann_index = configured_ann_index;
2677
2678        // Results are (doc_id, ordinal, score) where score = similarity (higher = better)
2679        let t0 = std::time::Instant::now();
2680        let mut flat_results = None;
2681        let results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
2682            // ANN search through the segment's ANN payload.
2683            match index {
2684                VectorIndex::Tq { index: lazy, codec } => {
2685                    let flat = lazy_flat.expect("ANN/flat pairing validated above");
2686                    // Estimated similarities feed the shared exact re-rank.
2687                    search_tq_segment(
2688                        lazy.get(),
2689                        codec,
2690                        query,
2691                        fetch_k.min(flat.num_docs_with_vectors()),
2692                        needs_document_aggregation.then_some(combiner),
2693                        field,
2694                        params.dim,
2695                        plan_cache.map(|cache| &cache.tq),
2696                    )?
2697                }
2698                VectorIndex::IvfTq { index: lazy, codec } => {
2699                    let index = lazy.get();
2700                    let centroids =
2701                        self.trained_vectors
2702                            .centroids
2703                            .get(&field.0)
2704                            .ok_or_else(|| {
2705                                Error::Schema(format!(
2706                                    "IVF-TQ index requires coarse centroids for field {}",
2707                                    field.0
2708                                ))
2709                            })?;
2710                    validate_coarse_centroids(centroids, params.dim)?;
2711                    let routing = self
2712                        .schema
2713                        .get_field_entry(field)
2714                        .and_then(|entry| entry.dense_vector_config.as_ref())
2715                        .map_or(crate::dsl::IvfRoutingMode::Auto, |config| {
2716                            config.ivf_routing
2717                        });
2718                    validate_ivf_tq_ann(index, centroids, codec, params.dim, routing, field)?;
2719                    let flat = lazy_flat.expect("ANN/flat pairing validated above");
2720                    search_ivf_tq_segment(
2721                        index,
2722                        centroids,
2723                        codec,
2724                        query,
2725                        fetch_k.min(flat.num_docs_with_vectors()),
2726                        needs_document_aggregation.then_some(combiner),
2727                        field,
2728                        params.nprobe,
2729                        routing,
2730                        plan_cache.map(|cache| &cache.ivf_tq),
2731                    )?
2732                }
2733                VectorIndex::BinaryIvf(_) => {
2734                    // Binary IVF serves Hamming queries only (BinaryDenseVectorQuery)
2735                    Vec::new()
2736                }
2737                VectorIndex::ScannAh(lazy) => {
2738                    let artifact = self
2739                        .trained_vectors
2740                        .scann_artifacts
2741                        .get(&field.0)
2742                        .ok_or_else(|| {
2743                            Error::Schema(format!(
2744                                "ScaNN field {} has no loaded global artifact",
2745                                field.0
2746                            ))
2747                        })?;
2748                    let flat = lazy_flat.expect("ANN/flat pairing validated above");
2749                    search_scann_ah_segment(
2750                        lazy.get(),
2751                        artifact,
2752                        query,
2753                        fetch_k.min(flat.num_docs_with_vectors()),
2754                        combiner,
2755                        field,
2756                        params.nprobe,
2757                        plan_cache.map(|cache| &cache.scann),
2758                    )?
2759                }
2760                VectorIndex::ScannBinary(_) => {
2761                    return Err(Error::Corruption(format!(
2762                        "binary ScaNN payload was attached to float field {}",
2763                        field.0
2764                    )));
2765                }
2766            }
2767        } else if let Some(lazy_flat) = lazy_flat {
2768            // Batched brute-force from lazy flat vectors (native-precision SIMD scoring).
2769            // Combine every value of a document before document-level top-k;
2770            // vector-level top-k loses documents on multi-valued fields.
2771            log::debug!(
2772                "[dense_vector_search] index={} field {}: brute-force on {} vectors (dim={}, quant={:?})",
2773                self.schema.index_label(),
2774                field.0,
2775                lazy_flat.num_vectors,
2776                lazy_flat.dim,
2777                lazy_flat.quantization
2778            );
2779            let dim = lazy_flat.dim;
2780            let n = lazy_flat.num_vectors;
2781            let quant = lazy_flat.quantization;
2782            let batch_len =
2783                bounded_vector_score_batch(lazy_flat.vector_byte_size(), DENSE_SCORE_BATCH);
2784            let mut collector = FlatDocumentCollector::new(fetch_k.min(n), combiner);
2785            let mut scores = vec![0f32; batch_len];
2786            let prepared_query = PreparedDenseScoreQuery::new(query, quant, dim, params.unit_norm)?;
2787
2788            for batch_start in (0..n).step_by(batch_len) {
2789                let batch_count = batch_len.min(n - batch_start);
2790                let batch_bytes = lazy_flat
2791                    .read_vectors_batch(batch_start, batch_count)
2792                    .await
2793                    .map_err(crate::Error::Io)?;
2794                let raw = batch_bytes.as_slice();
2795
2796                prepared_query.score_batch(raw, &mut scores[..batch_count])?;
2797
2798                for (i, &score) in scores.iter().enumerate().take(batch_count) {
2799                    let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
2800                    collector.push(doc_id, ordinal, score);
2801                }
2802            }
2803
2804            flat_results = Some(collector.into_results());
2805            Vec::new()
2806        } else {
2807            return Ok(Vec::new());
2808        };
2809        let l1_elapsed = t0.elapsed();
2810        {
2811            let kind = match ann_index {
2812                Some(VectorIndex::BinaryIvf(_)) => "binary_ivf",
2813                Some(VectorIndex::Tq { .. }) => "tq_flat",
2814                Some(VectorIndex::IvfTq { .. }) => "ivf_tq",
2815                Some(VectorIndex::ScannAh(_)) => "scann_ah",
2816                Some(VectorIndex::ScannBinary(_)) => "scann_binary",
2817                None => "flat",
2818            };
2819            crate::observe::dense_l1(
2820                self.schema.index_label(),
2821                self.schema.get_field_name(field).unwrap_or("?"),
2822                kind,
2823                l1_elapsed.as_secs_f64(),
2824                flat_results.as_ref().map_or(results.len(), Vec::len),
2825            );
2826        }
2827        log::debug!(
2828            "[dense_vector_search] index={} field {}: L1 returned {} candidates in {:.1}ms",
2829            self.schema.index_label(),
2830            field.0,
2831            flat_results.as_ref().map_or(results.len(), Vec::len),
2832            l1_elapsed.as_secs_f64() * 1000.0
2833        );
2834
2835        if let Some(results) = flat_results {
2836            return Ok(results);
2837        }
2838
2839        // Rerank ANN candidates using raw vectors from lazy flat (binary search lookup)
2840        // Uses native-precision SIMD scoring on quantized bytes — no dequantization overhead.
2841        if ann_index.is_some()
2842            && !results.is_empty()
2843            && let Some(lazy_flat) = lazy_flat
2844        {
2845            let t_rerank = std::time::Instant::now();
2846            let vbs = lazy_flat.vector_byte_size();
2847            let (reranked, stats) = exact_score_dense_candidate_documents(
2848                &results,
2849                lazy_flat,
2850                query,
2851                params.unit_norm,
2852                combiner,
2853                k,
2854            )
2855            .await?;
2856
2857            crate::observe::dense_rerank(
2858                self.schema.index_label(),
2859                self.schema.get_field_name(field).unwrap_or("?"),
2860                t_rerank.elapsed().as_secs_f64(),
2861                stats.resolve_elapsed.as_secs_f64(),
2862                stats.read_elapsed.as_secs_f64(),
2863                stats.vector_count,
2864            );
2865            log::debug!(
2866                "[dense_vector_search] index={} field {}: rerank {} vectors (dim={}, quant={:?}, bytes_per_vector={}): resolve={:.1}ms read={:.1}ms score={:.1}ms",
2867                self.schema.index_label(),
2868                field.0,
2869                stats.vector_count,
2870                lazy_flat.dim,
2871                lazy_flat.quantization,
2872                vbs,
2873                stats.resolve_elapsed.as_secs_f64() * 1000.0,
2874                stats.read_elapsed.as_secs_f64() * 1000.0,
2875                stats.score_elapsed.as_secs_f64() * 1000.0,
2876            );
2877
2878            log::debug!(
2879                "[dense_vector_search] index={} field {}: rerank total={:.1}ms",
2880                self.schema.index_label(),
2881                field.0,
2882                t_rerank.elapsed().as_secs_f64() * 1000.0
2883            );
2884            return Ok(reranked);
2885        }
2886
2887        Ok(combine_grouped_ordinal_results(results, combiner, k))
2888    }
2889
2890    /// Search binary dense vectors using IVF when available, otherwise
2891    /// brute-force Hamming distance.
2892    ///
2893    /// Returns VectorSearchResult with ordinal tracking.
2894    async fn search_binary_dense_vector_impl(
2895        &self,
2896        field: Field,
2897        query: &[u8],
2898        k: usize,
2899        combiner: crate::query::MultiValueCombiner,
2900        probe_cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
2901    ) -> Result<Vec<VectorSearchResult>> {
2902        let schema_dim = self.validate_binary_search_request(field, query)?;
2903        combiner.validate().map_err(Error::Query)?;
2904        if k == 0 {
2905            return Ok(Vec::new());
2906        }
2907        let t0 = crate::observe::Timer::start();
2908        if let Some(VectorIndex::ScannBinary(lazy)) = self.vector_indexes.get(&field.0) {
2909            let artifact = self
2910                .trained_vectors
2911                .scann_artifacts
2912                .get(&field.0)
2913                .ok_or_else(|| {
2914                    Error::Schema(format!(
2915                        "binary ScaNN field {} has no loaded global artifact",
2916                        field.0
2917                    ))
2918                })?;
2919            lazy.get()
2920                .validate_scann_generation(
2921                    artifact.config(),
2922                    artifact.generation(),
2923                    artifact.artifact_id(),
2924                )
2925                .map_err(|error| {
2926                    Error::Corruption(format!(
2927                        "binary ScaNN generation mismatch for field {}: {error}",
2928                        field.0
2929                    ))
2930                })?;
2931            let config = self
2932                .schema
2933                .get_field_entry(field)
2934                .and_then(|entry| entry.binary_dense_vector_config.as_ref())
2935                .ok_or_else(|| {
2936                    Error::Schema(format!(
2937                        "binary ScaNN field {} has no schema configuration",
2938                        field.0
2939                    ))
2940                })?;
2941            let model = artifact.binary_model().map_err(Error::Io)?;
2942            let clusters = binary_scann_probe_clusters(&model, query, config.nprobe, probe_cache)?;
2943            let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
2944                Error::Corruption(format!(
2945                    "binary ScaNN field {} is missing flat vectors",
2946                    field.0
2947                ))
2948            })?;
2949            let candidate_limit =
2950                checked_binary_combined_fetch_k(k)?.min(flat.num_docs_with_vectors());
2951            let (documents, ordinal_scores) = lazy
2952                .get()
2953                .search_binary_combined_documents(candidate_limit, query, &clusters, combiner)
2954                .map_err(|error| {
2955                    Error::Corruption(format!(
2956                        "invalid binary ScaNN payload for field {}: {error}",
2957                        field.0
2958                    ))
2959                })?;
2960            let results = exact_score_binary_candidate_document_ids(
2961                documents
2962                    .into_iter()
2963                    .map(|candidate| candidate.doc_id)
2964                    .collect(),
2965                &ordinal_scores,
2966                flat,
2967                query,
2968                schema_dim,
2969                combiner,
2970                k,
2971            )
2972            .await?;
2973            crate::observe::dense_l1(
2974                self.schema.index_label(),
2975                self.schema.get_field_name(field).unwrap_or("?"),
2976                "binary_scann",
2977                t0.secs(),
2978                results.len(),
2979            );
2980            return Ok(results);
2981        }
2982        if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0) {
2983            let ivf = lazy.get();
2984            let config = self
2985                .schema
2986                .get_field_entry(field)
2987                .and_then(|entry| entry.binary_dense_vector_config.as_ref())
2988                .ok_or_else(|| {
2989                    Error::Schema(format!(
2990                        "binary IVF field {} has no schema configuration",
2991                        field.0
2992                    ))
2993                })?;
2994            let quantizer = self
2995                .trained_vectors
2996                .binary_quantizers
2997                .get(&field.0)
2998                .ok_or_else(|| {
2999                    Error::Schema(format!(
3000                        "global binary IVF field {} has no loaded quantizer",
3001                        field.0
3002                    ))
3003                })?;
3004            validate_binary_ann(ivf, quantizer, config, schema_dim, field)?;
3005            let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
3006                Error::Corruption(format!(
3007                    "global binary IVF field {} is missing flat vector storage",
3008                    field.0
3009                ))
3010            })?;
3011            let single_valued = flat.num_vectors == flat.num_docs_with_vectors();
3012            let clusters = binary_probe_clusters(
3013                quantizer,
3014                query,
3015                config.nprobe,
3016                config.ivf_routing,
3017                probe_cache,
3018            )?;
3019            let results = if !single_valued
3020                && !matches!(combiner, crate::query::MultiValueCombiner::Max)
3021            {
3022                let candidate_limit =
3023                    checked_binary_combined_fetch_k(k)?.min(flat.num_docs_with_vectors());
3024                let (candidate_documents, probed_ordinal_scores) = ivf
3025                    .search_binary_combined_documents(candidate_limit, query, &clusters, combiner)
3026                    .map_err(|error| {
3027                        Error::Corruption(format!(
3028                            "invalid binary IVF payload for field {}: {error}",
3029                            field.0,
3030                        ))
3031                    })?;
3032                exact_score_binary_candidate_document_ids(
3033                    candidate_documents
3034                        .into_iter()
3035                        .map(|candidate| candidate.doc_id)
3036                        .collect(),
3037                    &probed_ordinal_scores,
3038                    flat,
3039                    query,
3040                    schema_dim,
3041                    combiner,
3042                    k,
3043                )
3044                .await?
3045            } else {
3046                let candidate_docs = if single_valued {
3047                    k
3048                } else {
3049                    // Completing the selected documents from flat storage can
3050                    // reorder a multi-value Max result when another ordinal
3051                    // lives outside the probed leaves. Keep the same bounded
3052                    // oversubscription used by combined binary reranking.
3053                    checked_binary_combined_fetch_k(k)?
3054                }
3055                .min(flat.num_docs_with_vectors());
3056                let ann_results = if single_valued {
3057                    ivf.search_binary_clusters::<false>(query, candidate_docs, &clusters)
3058                } else {
3059                    ivf.search_binary_clusters::<true>(query, candidate_docs, &clusters)
3060                }
3061                .map_err(|error| {
3062                    Error::Corruption(format!(
3063                        "invalid binary IVF payload for field {}: {error}",
3064                        field.0,
3065                    ))
3066                })?;
3067                // Binary IVF stores the original packed codes, so its leaf
3068                // scores are already exact for a single-valued field.
3069                if single_valued {
3070                    let ann_results = validate_binary_single_value_ann_results(ann_results, flat)?;
3071                    combine_ordinal_results(ann_results, combiner, k)
3072                } else {
3073                    exact_score_binary_candidate_documents(
3074                        &ann_results,
3075                        flat,
3076                        query,
3077                        schema_dim,
3078                        combiner,
3079                        k,
3080                    )
3081                    .await?
3082                }
3083            };
3084            crate::observe::dense_l1(
3085                self.schema.index_label(),
3086                self.schema.get_field_name(field).unwrap_or("?"),
3087                "global_binary_ivf",
3088                t0.secs(),
3089                results.len(),
3090            );
3091            return Ok(results);
3092        }
3093        let lazy_flat = match self.flat_vectors.get(&field.0) {
3094            Some(f) => f,
3095            None => return Ok(Vec::new()),
3096        };
3097
3098        let dim_bits = lazy_flat.dim;
3099        let byte_len = lazy_flat.vector_byte_size();
3100        let n = lazy_flat.num_vectors;
3101
3102        if dim_bits != schema_dim {
3103            return Err(Error::Corruption(format!(
3104                "binary vector field {} has schema dimension {} but flat storage dimension {}",
3105                field.0, schema_dim, dim_bits
3106            )));
3107        }
3108
3109        if byte_len != query.len() {
3110            return Err(Error::Schema(format!(
3111                "Binary query vector byte length {} != field byte length {}",
3112                query.len(),
3113                byte_len
3114            )));
3115        }
3116
3117        let batch_len = bounded_vector_score_batch(byte_len, BINARY_SCORE_BATCH);
3118        let mut collector = FlatDocumentCollector::new(k, combiner);
3119        let mut scores = vec![0f32; batch_len];
3120
3121        for batch_start in (0..n).step_by(batch_len) {
3122            let batch_count = batch_len.min(n - batch_start);
3123            let batch_bytes = lazy_flat
3124                .read_vectors_batch(batch_start, batch_count)
3125                .await
3126                .map_err(crate::Error::Io)?;
3127            let raw = batch_bytes.as_slice();
3128
3129            crate::structures::simd::batch_hamming_scores(
3130                query,
3131                raw,
3132                byte_len,
3133                dim_bits,
3134                &mut scores[..batch_count],
3135            );
3136
3137            for (i, &score) in scores.iter().enumerate().take(batch_count) {
3138                let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
3139                collector.push(doc_id, ordinal, score);
3140            }
3141        }
3142
3143        let results = collector.into_results();
3144
3145        crate::observe::dense_l1(
3146            self.schema.index_label(),
3147            self.schema.get_field_name(field).unwrap_or("?"),
3148            "binary_flat",
3149            t0.secs(),
3150            results.len(),
3151        );
3152        Ok(results)
3153    }
3154
3155    pub async fn search_binary_dense_vector(
3156        &self,
3157        field: Field,
3158        query: &[u8],
3159        k: usize,
3160        combiner: crate::query::MultiValueCombiner,
3161    ) -> Result<Vec<VectorSearchResult>> {
3162        self.search_binary_dense_vector_impl(field, query, k, combiner, None)
3163            .await
3164    }
3165
3166    pub(crate) async fn search_binary_dense_vector_with_probe_cache(
3167        &self,
3168        field: Field,
3169        query: &[u8],
3170        k: usize,
3171        combiner: crate::query::MultiValueCombiner,
3172        probe_cache: &std::sync::Mutex<Option<crate::structures::IvfProbePlan>>,
3173    ) -> Result<Vec<VectorSearchResult>> {
3174        self.search_binary_dense_vector_impl(field, query, k, combiner, Some(probe_cache))
3175            .await
3176    }
3177
3178    /// Get coarse centroids for a field.
3179    pub fn coarse_centroids(&self, field_id: u32) -> Option<&Arc<CoarseCentroids>> {
3180        self.trained_vectors.centroids.get(&field_id)
3181    }
3182
3183    pub fn set_trained_vectors(
3184        &mut self,
3185        trained_vectors: Arc<crate::segment::TrainedVectorStructures>,
3186    ) {
3187        self.trained_vectors = trained_vectors;
3188    }
3189
3190    /// Get the vector index type for a field
3191    pub fn get_vector_index(&self, field: Field) -> Option<&VectorIndex> {
3192        self.vector_indexes.get(&field.0)
3193    }
3194
3195    /// Get positions for a term (for phrase queries)
3196    ///
3197    /// Position offsets are now embedded in TermInfo, so we first look up
3198    /// the term to get its TermInfo, then use position_info() to get the offset.
3199    pub async fn get_positions(
3200        &self,
3201        field: Field,
3202        term: &[u8],
3203    ) -> Result<Option<crate::structures::PositionPostingList>> {
3204        // Get positions handle
3205        let handle = match &self.positions_handle {
3206            Some(h) => h,
3207            None => return Ok(None),
3208        };
3209
3210        // Build key: field_id + term
3211        let mut key = Vec::with_capacity(4 + term.len());
3212        key.extend_from_slice(&field.0.to_le_bytes());
3213        key.extend_from_slice(term);
3214
3215        // Look up term in dictionary to get TermInfo with position offset
3216        let term_info = match self.term_dict.get(&key).await? {
3217            Some(info) => info,
3218            None => return Ok(None),
3219        };
3220
3221        // Get position offset from TermInfo
3222        let (offset, length) = match term_info.position_info() {
3223            Some((o, l)) => (o, l),
3224            None => return Ok(None),
3225        };
3226
3227        // Read the position data only after validating untrusted offsets from
3228        // the term dictionary. Direct `offset + length` can wrap in release
3229        // builds and alias an unrelated range.
3230        let range = checked_file_range(offset, length, handle.len(), "position list")?;
3231        let slice = handle.slice(range);
3232        let data = slice.read_bytes().await?;
3233
3234        // Deserialize
3235        let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
3236
3237        Ok(Some(pos_list))
3238    }
3239
3240    /// Check if positions are available for a field
3241    pub fn has_positions(&self, field: Field) -> bool {
3242        // Check schema for position mode on this field
3243        if let Some(entry) = self.schema.get_field_entry(field) {
3244            entry.positions.is_some()
3245        } else {
3246            false
3247        }
3248    }
3249}
3250
3251// ── Synchronous search methods (mmap/RAM only) ─────────────────────────────
3252#[cfg(feature = "sync")]
3253impl SegmentReader {
3254    /// Synchronous posting list lookup — requires Inline (mmap/RAM) file handles.
3255    pub fn get_postings_sync(&self, field: Field, term: &[u8]) -> Result<Option<BlockPostingList>> {
3256        // Build key: field_id + term
3257        let mut key = Vec::with_capacity(4 + term.len());
3258        key.extend_from_slice(&field.0.to_le_bytes());
3259        key.extend_from_slice(term);
3260
3261        // Look up in term dictionary (sync)
3262        let term_info = match self.term_dict.get_sync(&key)? {
3263            Some(info) => info,
3264            None => return Ok(None),
3265        };
3266
3267        // Check if posting list is inlined
3268        if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
3269            let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
3270            for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
3271                posting_list.push(doc_id, tf);
3272            }
3273            let block_list = BlockPostingList::from_posting_list(&posting_list)?;
3274            return Ok(Some(block_list));
3275        }
3276
3277        // External posting list — sync range read
3278        let (posting_offset, posting_len) = term_info.external_info().ok_or_else(|| {
3279            Error::Corruption("TermInfo has neither inline nor external data".to_string())
3280        })?;
3281
3282        let range = checked_file_range(
3283            posting_offset,
3284            posting_len,
3285            self.postings_handle.len(),
3286            "posting",
3287        )?;
3288        let posting_bytes = self.postings_handle.read_bytes_range_sync(range)?;
3289        let block_list = BlockPostingList::deserialize_zero_copy(posting_bytes)?;
3290
3291        Ok(Some(block_list))
3292    }
3293
3294    /// Synchronous prefix posting list lookup — requires Inline (mmap/RAM) file handles.
3295    pub fn get_prefix_postings_sync(
3296        &self,
3297        field: Field,
3298        prefix: &[u8],
3299    ) -> Result<Vec<BlockPostingList>> {
3300        if prefix.is_empty() {
3301            return Err(Error::Query("prefix must not be empty".into()));
3302        }
3303        let mut key_prefix = Vec::with_capacity(4 + prefix.len());
3304        key_prefix.extend_from_slice(&field.0.to_le_bytes());
3305        key_prefix.extend_from_slice(prefix);
3306
3307        let (entries, truncated) = self
3308            .term_dict
3309            .prefix_scan_limited_sync(&key_prefix, MAX_PREFIX_TERMS)?;
3310        if truncated {
3311            return Err(Error::Query(format!(
3312                "prefix expands to more than {MAX_PREFIX_TERMS} terms"
3313            )));
3314        }
3315        let posting_count: u64 = entries
3316            .iter()
3317            .map(|(_, term_info)| term_info.doc_freq() as u64)
3318            .sum();
3319        if posting_count > MAX_PREFIX_POSTINGS {
3320            return Err(Error::Query(format!(
3321                "prefix expands to {posting_count} postings (maximum {MAX_PREFIX_POSTINGS})"
3322            )));
3323        }
3324        let mut results = Vec::with_capacity(entries.len());
3325
3326        for (_key, term_info) in entries {
3327            if let Some((doc_ids, term_freqs)) = term_info.decode_inline() {
3328                let mut posting_list = crate::structures::PostingList::with_capacity(doc_ids.len());
3329                for (doc_id, tf) in doc_ids.into_iter().zip(term_freqs) {
3330                    posting_list.push(doc_id, tf);
3331                }
3332                results.push(BlockPostingList::from_posting_list(&posting_list)?);
3333            } else if let Some((posting_offset, posting_len)) = term_info.external_info() {
3334                let range = checked_file_range(
3335                    posting_offset,
3336                    posting_len,
3337                    self.postings_handle.len(),
3338                    "prefix posting",
3339                )?;
3340                let posting_bytes = self.postings_handle.read_bytes_range_sync(range)?;
3341                results.push(BlockPostingList::deserialize_zero_copy(posting_bytes)?);
3342            }
3343        }
3344
3345        Ok(results)
3346    }
3347
3348    /// Synchronous position list lookup — requires Inline (mmap/RAM) file handles.
3349    pub fn get_positions_sync(
3350        &self,
3351        field: Field,
3352        term: &[u8],
3353    ) -> Result<Option<crate::structures::PositionPostingList>> {
3354        let handle = match &self.positions_handle {
3355            Some(h) => h,
3356            None => return Ok(None),
3357        };
3358
3359        // Build key: field_id + term
3360        let mut key = Vec::with_capacity(4 + term.len());
3361        key.extend_from_slice(&field.0.to_le_bytes());
3362        key.extend_from_slice(term);
3363
3364        // Look up term in dictionary (sync)
3365        let term_info = match self.term_dict.get_sync(&key)? {
3366            Some(info) => info,
3367            None => return Ok(None),
3368        };
3369
3370        let (offset, length) = match term_info.position_info() {
3371            Some((o, l)) => (o, l),
3372            None => return Ok(None),
3373        };
3374
3375        let range = checked_file_range(offset, length, handle.len(), "position list")?;
3376        let slice = handle.slice(range);
3377        let data = slice.read_bytes_sync()?;
3378
3379        let pos_list = crate::structures::PositionPostingList::deserialize(data.as_slice())?;
3380        Ok(Some(pos_list))
3381    }
3382
3383    /// Synchronous dense vector search — ANN indexes are already sync,
3384    /// brute-force uses sync mmap reads.
3385    pub fn search_dense_vector_sync(
3386        &self,
3387        field: Field,
3388        query: &[f32],
3389        k: usize,
3390        nprobe: usize,
3391        rerank_factor: f32,
3392        combiner: crate::query::MultiValueCombiner,
3393    ) -> Result<Vec<VectorSearchResult>> {
3394        self.search_dense_vector_sync_impl(field, query, k, nprobe, rerank_factor, combiner, None)
3395    }
3396
3397    #[cfg(feature = "sync")]
3398    #[allow(clippy::too_many_arguments)]
3399    pub(crate) fn search_dense_vector_sync_with_probe_cache(
3400        &self,
3401        field: Field,
3402        query: &[f32],
3403        k: usize,
3404        nprobe: usize,
3405        rerank_factor: f32,
3406        combiner: crate::query::MultiValueCombiner,
3407        plan_cache: &DensePlanCache,
3408    ) -> Result<Vec<VectorSearchResult>> {
3409        self.search_dense_vector_sync_impl(
3410            field,
3411            query,
3412            k,
3413            nprobe,
3414            rerank_factor,
3415            combiner,
3416            Some(plan_cache),
3417        )
3418    }
3419
3420    #[cfg(feature = "sync")]
3421    #[allow(clippy::too_many_arguments)]
3422    fn search_dense_vector_sync_impl(
3423        &self,
3424        field: Field,
3425        query: &[f32],
3426        k: usize,
3427        nprobe: usize,
3428        rerank_factor: f32,
3429        combiner: crate::query::MultiValueCombiner,
3430        plan_cache: Option<&DensePlanCache>,
3431    ) -> Result<Vec<VectorSearchResult>> {
3432        let params =
3433            self.validate_dense_search_request(field, query, nprobe, rerank_factor, combiner)?;
3434        let fetch_k = checked_dense_fetch_k(k, rerank_factor)?;
3435        if k == 0 {
3436            return Ok(Vec::new());
3437        }
3438
3439        let configured_ann_index = self.vector_indexes.get(&field.0);
3440        let lazy_flat = self.flat_vectors.get(&field.0);
3441        if configured_ann_index.is_none() && lazy_flat.is_none() {
3442            return Ok(Vec::new());
3443        }
3444
3445        if configured_ann_index.is_some() && lazy_flat.is_none() {
3446            return Err(Error::Corruption(format!(
3447                "dense ANN field {} is missing flat vector storage",
3448                field.0
3449            )));
3450        }
3451
3452        if let Some(flat) = lazy_flat
3453            && flat.dim != params.dim
3454        {
3455            return Err(Error::Corruption(format!(
3456                "dense vector field {} has schema dimension {} but flat storage dimension {}",
3457                field.0, params.dim, flat.dim
3458            )));
3459        }
3460
3461        let needs_document_aggregation = lazy_flat.is_some_and(|flat| {
3462            flat.num_vectors != flat.num_docs_with_vectors()
3463                && !matches!(combiner, crate::query::MultiValueCombiner::Max)
3464        });
3465        // Sync and async search share the same ANN candidate modes; neither
3466        // silently substitutes a raw flat scan for an indexed field.
3467        let ann_index = configured_ann_index;
3468
3469        let results: Vec<(u32, u16, f32)> = if let Some(index) = ann_index {
3470            // ANN search (already sync)
3471            match index {
3472                VectorIndex::Tq { index: lazy, codec } => {
3473                    let flat = lazy_flat.expect("ANN/flat pairing validated above");
3474                    search_tq_segment(
3475                        lazy.get(),
3476                        codec,
3477                        query,
3478                        fetch_k.min(flat.num_docs_with_vectors()),
3479                        needs_document_aggregation.then_some(combiner),
3480                        field,
3481                        params.dim,
3482                        plan_cache.map(|cache| &cache.tq),
3483                    )?
3484                }
3485                VectorIndex::IvfTq { index: lazy, codec } => {
3486                    let index = lazy.get();
3487                    let centroids =
3488                        self.trained_vectors
3489                            .centroids
3490                            .get(&field.0)
3491                            .ok_or_else(|| {
3492                                Error::Schema(format!(
3493                                    "IVF-TQ index requires coarse centroids for field {}",
3494                                    field.0
3495                                ))
3496                            })?;
3497                    validate_coarse_centroids(centroids, params.dim)?;
3498                    let routing = self
3499                        .schema
3500                        .get_field_entry(field)
3501                        .and_then(|entry| entry.dense_vector_config.as_ref())
3502                        .map_or(crate::dsl::IvfRoutingMode::Auto, |config| {
3503                            config.ivf_routing
3504                        });
3505                    validate_ivf_tq_ann(index, centroids, codec, params.dim, routing, field)?;
3506                    let flat = lazy_flat.expect("ANN/flat pairing validated above");
3507                    search_ivf_tq_segment(
3508                        index,
3509                        centroids,
3510                        codec,
3511                        query,
3512                        fetch_k.min(flat.num_docs_with_vectors()),
3513                        needs_document_aggregation.then_some(combiner),
3514                        field,
3515                        params.nprobe,
3516                        routing,
3517                        plan_cache.map(|cache| &cache.ivf_tq),
3518                    )?
3519                }
3520                VectorIndex::BinaryIvf(_) => {
3521                    // Binary IVF serves Hamming queries only (BinaryDenseVectorQuery)
3522                    Vec::new()
3523                }
3524                VectorIndex::ScannAh(lazy) => {
3525                    let artifact = self
3526                        .trained_vectors
3527                        .scann_artifacts
3528                        .get(&field.0)
3529                        .ok_or_else(|| {
3530                            Error::Schema(format!(
3531                                "ScaNN field {} has no loaded global artifact",
3532                                field.0
3533                            ))
3534                        })?;
3535                    let flat = lazy_flat.expect("ANN/flat pairing validated above");
3536                    search_scann_ah_segment(
3537                        lazy.get(),
3538                        artifact,
3539                        query,
3540                        fetch_k.min(flat.num_docs_with_vectors()),
3541                        combiner,
3542                        field,
3543                        params.nprobe,
3544                        plan_cache.map(|cache| &cache.scann),
3545                    )?
3546                }
3547                VectorIndex::ScannBinary(_) => {
3548                    return Err(Error::Corruption(format!(
3549                        "binary ScaNN payload was attached to float field {}",
3550                        field.0
3551                    )));
3552                }
3553            }
3554        } else if let Some(lazy_flat) = lazy_flat {
3555            // Batched brute-force (sync mmap reads)
3556            let dim = lazy_flat.dim;
3557            let n = lazy_flat.num_vectors;
3558            let quant = lazy_flat.quantization;
3559            let batch_len =
3560                bounded_vector_score_batch(lazy_flat.vector_byte_size(), DENSE_SCORE_BATCH);
3561            let mut collector = FlatDocumentCollector::new(fetch_k.min(n), combiner);
3562            let mut scores = vec![0f32; batch_len];
3563            let prepared_query = PreparedDenseScoreQuery::new(query, quant, dim, params.unit_norm)?;
3564
3565            for batch_start in (0..n).step_by(batch_len) {
3566                let batch_count = batch_len.min(n - batch_start);
3567                let batch_bytes = lazy_flat
3568                    .read_vectors_batch_sync(batch_start, batch_count)
3569                    .map_err(crate::Error::Io)?;
3570                let raw = batch_bytes.as_slice();
3571
3572                prepared_query.score_batch(raw, &mut scores[..batch_count])?;
3573
3574                for (i, &score) in scores.iter().enumerate().take(batch_count) {
3575                    let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
3576                    collector.push(doc_id, ordinal, score);
3577                }
3578            }
3579
3580            return Ok(collector.into_results());
3581        } else {
3582            return Ok(Vec::new());
3583        };
3584
3585        // Rerank ANN candidates using raw vectors (sync)
3586        if ann_index.is_some()
3587            && !results.is_empty()
3588            && let Some(lazy_flat) = lazy_flat
3589        {
3590            return exact_score_dense_candidate_documents_sync(
3591                &results,
3592                lazy_flat,
3593                query,
3594                params.unit_norm,
3595                combiner,
3596                k,
3597            );
3598        }
3599
3600        Ok(combine_grouped_ordinal_results(results, combiner, k))
3601    }
3602
3603    /// Synchronous binary dense vector search (mmap/RAM only).
3604    ///
3605    /// Mirrors [`Self::search_binary_dense_vector`] for the rayon-parallel
3606    /// sync scorer path used by multi-threaded runtimes.
3607    #[cfg(feature = "sync")]
3608    fn search_binary_dense_vector_sync_impl(
3609        &self,
3610        field: Field,
3611        query: &[u8],
3612        k: usize,
3613        combiner: crate::query::MultiValueCombiner,
3614        probe_cache: Option<&std::sync::Mutex<Option<crate::structures::IvfProbePlan>>>,
3615    ) -> Result<Vec<VectorSearchResult>> {
3616        let schema_dim = self.validate_binary_search_request(field, query)?;
3617        combiner.validate().map_err(Error::Query)?;
3618        if k == 0 {
3619            return Ok(Vec::new());
3620        }
3621        let t0 = crate::observe::Timer::start();
3622        if let Some(VectorIndex::ScannBinary(lazy)) = self.vector_indexes.get(&field.0) {
3623            let artifact = self
3624                .trained_vectors
3625                .scann_artifacts
3626                .get(&field.0)
3627                .ok_or_else(|| {
3628                    Error::Schema(format!(
3629                        "binary ScaNN field {} has no loaded global artifact",
3630                        field.0
3631                    ))
3632                })?;
3633            lazy.get()
3634                .validate_scann_generation(
3635                    artifact.config(),
3636                    artifact.generation(),
3637                    artifact.artifact_id(),
3638                )
3639                .map_err(|error| {
3640                    Error::Corruption(format!(
3641                        "binary ScaNN generation mismatch for field {}: {error}",
3642                        field.0
3643                    ))
3644                })?;
3645            let config = self
3646                .schema
3647                .get_field_entry(field)
3648                .and_then(|entry| entry.binary_dense_vector_config.as_ref())
3649                .ok_or_else(|| {
3650                    Error::Schema(format!(
3651                        "binary ScaNN field {} has no schema configuration",
3652                        field.0
3653                    ))
3654                })?;
3655            let model = artifact.binary_model().map_err(Error::Io)?;
3656            let clusters = binary_scann_probe_clusters(&model, query, config.nprobe, probe_cache)?;
3657            let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
3658                Error::Corruption(format!(
3659                    "binary ScaNN field {} is missing flat vectors",
3660                    field.0
3661                ))
3662            })?;
3663            let candidate_limit =
3664                checked_binary_combined_fetch_k(k)?.min(flat.num_docs_with_vectors());
3665            let (documents, ordinal_scores) = lazy
3666                .get()
3667                .search_binary_combined_documents(candidate_limit, query, &clusters, combiner)
3668                .map_err(|error| {
3669                    Error::Corruption(format!(
3670                        "invalid binary ScaNN payload for field {}: {error}",
3671                        field.0
3672                    ))
3673                })?;
3674            let results = exact_score_binary_candidate_document_ids_sync(
3675                documents
3676                    .into_iter()
3677                    .map(|candidate| candidate.doc_id)
3678                    .collect(),
3679                &ordinal_scores,
3680                flat,
3681                query,
3682                schema_dim,
3683                combiner,
3684                k,
3685            )?;
3686            crate::observe::dense_l1(
3687                self.schema.index_label(),
3688                self.schema.get_field_name(field).unwrap_or("?"),
3689                "binary_scann",
3690                t0.secs(),
3691                results.len(),
3692            );
3693            return Ok(results);
3694        }
3695        if let Some(VectorIndex::BinaryIvf(lazy)) = self.vector_indexes.get(&field.0) {
3696            let ivf = lazy.get();
3697            let config = self
3698                .schema
3699                .get_field_entry(field)
3700                .and_then(|entry| entry.binary_dense_vector_config.as_ref())
3701                .ok_or_else(|| {
3702                    Error::Schema(format!(
3703                        "binary IVF field {} has no schema configuration",
3704                        field.0
3705                    ))
3706                })?;
3707            let quantizer = self
3708                .trained_vectors
3709                .binary_quantizers
3710                .get(&field.0)
3711                .ok_or_else(|| {
3712                    Error::Schema(format!(
3713                        "global binary IVF field {} has no loaded quantizer",
3714                        field.0
3715                    ))
3716                })?;
3717            validate_binary_ann(ivf, quantizer, config, schema_dim, field)?;
3718            let flat = self.flat_vectors.get(&field.0).ok_or_else(|| {
3719                Error::Corruption(format!(
3720                    "global binary IVF field {} is missing flat vector storage",
3721                    field.0
3722                ))
3723            })?;
3724            let single_valued = flat.num_vectors == flat.num_docs_with_vectors();
3725            let clusters = binary_probe_clusters(
3726                quantizer,
3727                query,
3728                config.nprobe,
3729                config.ivf_routing,
3730                probe_cache,
3731            )?;
3732            let results = if !single_valued
3733                && !matches!(combiner, crate::query::MultiValueCombiner::Max)
3734            {
3735                let candidate_limit =
3736                    checked_binary_combined_fetch_k(k)?.min(flat.num_docs_with_vectors());
3737                let (candidate_documents, probed_ordinal_scores) = ivf
3738                    .search_binary_combined_documents(candidate_limit, query, &clusters, combiner)
3739                    .map_err(|error| {
3740                        Error::Corruption(format!(
3741                            "invalid binary IVF payload for field {}: {error}",
3742                            field.0,
3743                        ))
3744                    })?;
3745                exact_score_binary_candidate_document_ids_sync(
3746                    candidate_documents
3747                        .into_iter()
3748                        .map(|candidate| candidate.doc_id)
3749                        .collect(),
3750                    &probed_ordinal_scores,
3751                    flat,
3752                    query,
3753                    schema_dim,
3754                    combiner,
3755                    k,
3756                )?
3757            } else {
3758                let candidate_docs = if single_valued {
3759                    k
3760                } else {
3761                    checked_binary_combined_fetch_k(k)?
3762                }
3763                .min(flat.num_docs_with_vectors());
3764                let ann_results = if single_valued {
3765                    ivf.search_binary_clusters::<false>(query, candidate_docs, &clusters)
3766                } else {
3767                    ivf.search_binary_clusters::<true>(query, candidate_docs, &clusters)
3768                }
3769                .map_err(|error| {
3770                    Error::Corruption(format!(
3771                        "invalid binary IVF payload for field {}: {error}",
3772                        field.0,
3773                    ))
3774                })?;
3775                if single_valued {
3776                    let ann_results = validate_binary_single_value_ann_results(ann_results, flat)?;
3777                    combine_ordinal_results(ann_results, combiner, k)
3778                } else {
3779                    exact_score_binary_candidate_documents_sync(
3780                        &ann_results,
3781                        flat,
3782                        query,
3783                        schema_dim,
3784                        combiner,
3785                        k,
3786                    )?
3787                }
3788            };
3789            crate::observe::dense_l1(
3790                self.schema.index_label(),
3791                self.schema.get_field_name(field).unwrap_or("?"),
3792                "global_binary_ivf",
3793                t0.secs(),
3794                results.len(),
3795            );
3796            return Ok(results);
3797        }
3798        let lazy_flat = match self.flat_vectors.get(&field.0) {
3799            Some(f) => f,
3800            None => return Ok(Vec::new()),
3801        };
3802
3803        let dim_bits = lazy_flat.dim;
3804        let byte_len = lazy_flat.vector_byte_size();
3805        let n = lazy_flat.num_vectors;
3806
3807        if dim_bits != schema_dim {
3808            return Err(Error::Corruption(format!(
3809                "binary vector field {} has schema dimension {} but flat storage dimension {}",
3810                field.0, schema_dim, dim_bits
3811            )));
3812        }
3813
3814        if byte_len != query.len() {
3815            return Err(Error::Schema(format!(
3816                "Binary query vector byte length {} != field byte length {}",
3817                query.len(),
3818                byte_len
3819            )));
3820        }
3821
3822        let batch_len = bounded_vector_score_batch(byte_len, BINARY_SCORE_BATCH);
3823        let mut collector = FlatDocumentCollector::new(k, combiner);
3824        let mut scores = vec![0f32; batch_len];
3825
3826        for batch_start in (0..n).step_by(batch_len) {
3827            let batch_count = batch_len.min(n - batch_start);
3828            let batch_bytes = lazy_flat
3829                .read_vectors_batch_sync(batch_start, batch_count)
3830                .map_err(crate::Error::Io)?;
3831            let raw = batch_bytes.as_slice();
3832
3833            crate::structures::simd::batch_hamming_scores(
3834                query,
3835                raw,
3836                byte_len,
3837                dim_bits,
3838                &mut scores[..batch_count],
3839            );
3840
3841            for (i, &score) in scores.iter().enumerate().take(batch_count) {
3842                let (doc_id, ordinal) = lazy_flat.get_doc_id(batch_start + i);
3843                collector.push(doc_id, ordinal, score);
3844            }
3845        }
3846
3847        let results = collector.into_results();
3848
3849        crate::observe::dense_l1(
3850            self.schema.index_label(),
3851            self.schema.get_field_name(field).unwrap_or("?"),
3852            "binary_flat",
3853            t0.secs(),
3854            results.len(),
3855        );
3856        Ok(results)
3857    }
3858
3859    #[cfg(feature = "sync")]
3860    pub fn search_binary_dense_vector_sync(
3861        &self,
3862        field: Field,
3863        query: &[u8],
3864        k: usize,
3865        combiner: crate::query::MultiValueCombiner,
3866    ) -> Result<Vec<VectorSearchResult>> {
3867        self.search_binary_dense_vector_sync_impl(field, query, k, combiner, None)
3868    }
3869
3870    #[cfg(feature = "sync")]
3871    pub(crate) fn search_binary_dense_vector_sync_with_probe_cache(
3872        &self,
3873        field: Field,
3874        query: &[u8],
3875        k: usize,
3876        combiner: crate::query::MultiValueCombiner,
3877        probe_cache: &std::sync::Mutex<Option<crate::structures::IvfProbePlan>>,
3878    ) -> Result<Vec<VectorSearchResult>> {
3879        self.search_binary_dense_vector_sync_impl(field, query, k, combiner, Some(probe_cache))
3880    }
3881}
3882
3883#[cfg(test)]
3884mod dense_search_safety_tests {
3885    use super::*;
3886
3887    #[test]
3888    fn dense_fetch_count_rejects_non_finite_and_unbounded_factors() {
3889        for factor in [
3890            f32::NAN,
3891            f32::INFINITY,
3892            f32::NEG_INFINITY,
3893            0.0,
3894            0.5,
3895            2.01,
3896            MAX_DENSE_RERANK_FACTOR + 1.0,
3897        ] {
3898            assert!(
3899                checked_dense_fetch_k(10, factor).is_err(),
3900                "factor={factor}"
3901            );
3902        }
3903    }
3904
3905    fn values_as_bytes<T>(values: &[T]) -> &[u8] {
3906        unsafe {
3907            std::slice::from_raw_parts(values.as_ptr() as *const u8, std::mem::size_of_val(values))
3908        }
3909    }
3910
3911    fn assert_prepared_dense_scores_match_legacy(
3912        quantization: DenseVectorQuantization,
3913        raw: &[u8],
3914        unit_norm: bool,
3915    ) {
3916        const DIM: usize = 4;
3917        const VECTOR_COUNT: usize = 4;
3918        let query = [0.25, -0.5, 0.75, 1.0];
3919        let mut expected = [0.0; VECTOR_COUNT];
3920        SegmentReader::score_quantized_batch_legacy(
3921            &query,
3922            raw,
3923            quantization,
3924            DIM,
3925            &mut expected,
3926            unit_norm,
3927        )
3928        .unwrap();
3929
3930        let prepared = PreparedDenseScoreQuery::new(&query, quantization, DIM, unit_norm).unwrap();
3931        let vector_bytes = DIM
3932            * match quantization {
3933                DenseVectorQuantization::F32 => std::mem::size_of::<f32>(),
3934                DenseVectorQuantization::F16 => std::mem::size_of::<u16>(),
3935                DenseVectorQuantization::UInt8 => 1,
3936                DenseVectorQuantization::Binary => unreachable!(),
3937            };
3938        let split = 2 * vector_bytes;
3939        let mut actual = [0.0; VECTOR_COUNT];
3940        prepared
3941            .score_batch(&raw[..split], &mut actual[..2])
3942            .unwrap();
3943        prepared
3944            .score_batch(&raw[split..], &mut actual[2..])
3945            .unwrap();
3946
3947        assert_eq!(
3948            actual.map(f32::to_bits),
3949            expected.map(f32::to_bits),
3950            "quantization={quantization:?}, unit_norm={unit_norm}"
3951        );
3952    }
3953
3954    #[test]
3955    fn prepared_dense_query_matches_legacy_scoring_across_batches() {
3956        let vectors_f32 = [
3957            0.5, -0.25, 0.75, 1.0, -1.0, 0.5, 0.25, 0.125, 0.0, 0.0, 0.0, 0.0, 0.75, 0.5, -0.5,
3958            -0.25,
3959        ];
3960        let vectors_f16: Vec<u16> = vectors_f32
3961            .iter()
3962            .map(|&value| crate::structures::simd::f32_to_f16(value))
3963            .collect();
3964        let vectors_u8 = [
3965            255, 96, 224, 160, 0, 192, 144, 128, 128, 128, 128, 128, 224, 192, 64, 96,
3966        ];
3967
3968        for unit_norm in [false, true] {
3969            assert_prepared_dense_scores_match_legacy(
3970                DenseVectorQuantization::F32,
3971                values_as_bytes(&vectors_f32),
3972                unit_norm,
3973            );
3974            assert_prepared_dense_scores_match_legacy(
3975                DenseVectorQuantization::F16,
3976                values_as_bytes(&vectors_f16),
3977                unit_norm,
3978            );
3979            assert_prepared_dense_scores_match_legacy(
3980                DenseVectorQuantization::UInt8,
3981                &vectors_u8,
3982                unit_norm,
3983            );
3984        }
3985    }
3986
3987    #[test]
3988    fn prepared_dense_query_preserves_scoring_validation_errors() {
3989        assert!(matches!(
3990            PreparedDenseScoreQuery::new(&[1.0], DenseVectorQuantization::F32, 2, false).err(),
3991            Some(Error::Query(_))
3992        ));
3993        assert!(matches!(
3994            PreparedDenseScoreQuery::new(&[1.0], DenseVectorQuantization::Binary, 1, false).err(),
3995            Some(Error::InvalidFieldType { .. })
3996        ));
3997
3998        let query = [1.0, 2.0];
3999        let prepared =
4000            PreparedDenseScoreQuery::new(&query, DenseVectorQuantization::F32, 2, false).unwrap();
4001        let mut scores = [0.0];
4002        assert!(matches!(
4003            prepared.score_batch(&[0; 7], &mut scores),
4004            Err(Error::Corruption(_))
4005        ));
4006    }
4007
4008    #[test]
4009    fn flat_document_collector_does_not_let_one_multivalue_doc_crowd_out_others() {
4010        let mut collector = FlatDocumentCollector::new(2, crate::query::MultiValueCombiner::Max);
4011        collector.push(1, 0, 1.0);
4012        collector.push(1, 1, 0.9);
4013        collector.push(2, 0, 0.8);
4014
4015        let results = collector.into_results();
4016        assert_eq!(
4017            results
4018                .iter()
4019                .map(|result| result.doc_id)
4020                .collect::<Vec<_>>(),
4021            vec![1, 2]
4022        );
4023        assert_eq!(results[0].ordinals.len(), 2);
4024    }
4025
4026    #[test]
4027    fn flat_document_collector_evicts_by_score_then_doc_id() {
4028        let mut collector = FlatDocumentCollector::new(2, crate::query::MultiValueCombiner::Max);
4029        collector.push(1, 0, 0.5);
4030        collector.push(3, 0, 0.8);
4031        collector.push(2, 0, 0.9);
4032        let results = collector.into_results();
4033        assert_eq!(
4034            results
4035                .iter()
4036                .map(|result| result.doc_id)
4037                .collect::<Vec<_>>(),
4038            vec![2, 3]
4039        );
4040
4041        let mut tied = FlatDocumentCollector::new(1, crate::query::MultiValueCombiner::Max);
4042        tied.push(2, 0, 1.0);
4043        tied.push(1, 0, 1.0);
4044        let results = tied.into_results();
4045        assert_eq!(results[0].doc_id, 1);
4046    }
4047
4048    #[test]
4049    fn dense_fetch_count_rounds_up_and_detects_overflow() {
4050        assert_eq!(checked_dense_fetch_k(3, 1.5).unwrap(), 5);
4051        assert_eq!(checked_dense_fetch_k(10_000, 2.0).unwrap(), 20_000);
4052        assert!(checked_dense_fetch_k(10_001, 2.0).is_err());
4053        assert!(checked_dense_fetch_k(usize::MAX, 2.0).is_err());
4054    }
4055
4056    #[test]
4057    fn binary_combined_fetch_count_uses_shared_bounded_oversampling() {
4058        assert_eq!(checked_binary_combined_fetch_k(3).unwrap(), 6);
4059        assert_eq!(checked_binary_combined_fetch_k(10_000).unwrap(), 20_000);
4060        assert_eq!(checked_binary_combined_fetch_k(10_001).unwrap(), 20_000);
4061        assert_eq!(checked_binary_combined_fetch_k(20_000).unwrap(), 20_000);
4062        assert!(checked_binary_combined_fetch_k(20_001).is_err());
4063        assert!(checked_binary_combined_fetch_k(usize::MAX).is_err());
4064    }
4065
4066    #[cfg(feature = "native")]
4067    #[test]
4068    fn legacy_ivf_tq_generation_is_rejected_while_opening() {
4069        use crate::directories::OwnedBytes;
4070        use crate::dsl::IvfRoutingMode;
4071        use crate::segment::ann_disk::{AnnDiskIndex, AnnKind};
4072
4073        let centroids = CoarseCentroids {
4074            num_clusters: 1,
4075            dim: 2,
4076            centroids: vec![1.0, 0.0],
4077            version: 7,
4078            soar_config: None,
4079            routing_index: None,
4080        };
4081        let mut build_centroids = centroids.clone();
4082        build_centroids.version =
4083            crate::structures::mark_ivf_tq_cosine_generation(build_centroids.version);
4084        let mut bytes = crate::segment::ann_build::build_ivf_tq(
4085            2,
4086            IvfRoutingMode::Flat,
4087            &build_centroids,
4088            &[(0, 0)],
4089            &[1.0, 0.0],
4090        )
4091        .unwrap();
4092        // Rewrite only the in-band centroid generation in the header to model
4093        // a persisted pre-cosine artifact.
4094        bytes[24..32].copy_from_slice(&centroids.version.to_le_bytes());
4095        let error = AnnDiskIndex::open(OwnedBytes::new(bytes), AnnKind::IvfTq, 1)
4096            .err()
4097            .expect("legacy IVF-TQ payload must fail while opening")
4098            .to_string();
4099        assert!(error.contains("unsupported legacy generation"), "{error}");
4100    }
4101
4102    #[test]
4103    fn rerank_batch_is_capped_by_actual_candidate_vectors() {
4104        assert_eq!(bounded_rerank_batch(3_072, DENSE_SCORE_BATCH, 20), 20);
4105        assert_eq!(
4106            bounded_rerank_batch(3_072, DENSE_SCORE_BATCH, 10_000),
4107            MAX_VECTOR_SCORE_BATCH_BYTES / 3_072
4108        );
4109        assert_eq!(bounded_rerank_batch(3_072, DENSE_SCORE_BATCH, 0), 1);
4110    }
4111
4112    #[test]
4113    fn file_ranges_reject_overflow_and_truncation() {
4114        assert_eq!(checked_file_range(4, 3, 7, "test").unwrap(), 4..7);
4115        assert!(checked_file_range(u64::MAX, 1, u64::MAX, "test").is_err());
4116        assert!(checked_file_range(5, 3, 7, "test").is_err());
4117    }
4118
4119    #[test]
4120    fn shared_tq_plan_cache_rebuilds_for_divergent_query_clones() {
4121        let codec = crate::structures::TqCodec::new(4);
4122        let cache = std::sync::Mutex::new(None);
4123        let original_query = vec![1.0, 2.0, 3.0, 4.0];
4124
4125        let original =
4126            cached_tq_query_plan(&codec, &original_query, Some(&cache)).expect("build plan");
4127        let reused =
4128            cached_tq_query_plan(&codec, &original_query, Some(&cache)).expect("reuse plan");
4129        assert!(
4130            std::sync::Arc::ptr_eq(&original, &reused),
4131            "unchanged queries must share their plan across segments"
4132        );
4133
4134        let mut divergent_clone = original_query.clone();
4135        divergent_clone[0] = -1.0;
4136        let rebuilt =
4137            cached_tq_query_plan(&codec, &divergent_clone, Some(&cache)).expect("rebuild plan");
4138        assert!(
4139            !std::sync::Arc::ptr_eq(&original, &rebuilt),
4140            "a clone with a mutated vector must not reuse stale LUTs"
4141        );
4142        assert!(rebuilt.matches_query(&divergent_clone));
4143        assert!(!rebuilt.matches_query(&original_query));
4144    }
4145
4146    #[test]
4147    fn candidate_vector_reads_coalesce_contiguous_values() {
4148        let mut runs = Vec::new();
4149        plan_vector_read_runs(&[3, 4, 5, 9, 12, 13], &mut runs).unwrap();
4150        assert_eq!(runs.len(), 3);
4151        assert!(matches!(
4152            runs.as_slice(),
4153            [
4154                VectorReadRun {
4155                    buffer_start: 0,
4156                    flat_start: 3,
4157                    count: 3,
4158                },
4159                VectorReadRun {
4160                    buffer_start: 3,
4161                    flat_start: 9,
4162                    count: 1,
4163                },
4164                VectorReadRun {
4165                    buffer_start: 4,
4166                    flat_start: 12,
4167                    count: 2,
4168                },
4169            ]
4170        ));
4171        assert!(plan_vector_read_runs(&[3, 3], &mut runs).is_err());
4172    }
4173
4174    #[tokio::test]
4175    async fn binary_single_value_ann_fast_path_validates_and_deduplicates() {
4176        use crate::directories::{FileHandle, OwnedBytes};
4177        use crate::segment::FlatVectorData;
4178
4179        let mut encoded = Vec::new();
4180        FlatVectorData::serialize_binary_from_bits_streaming(
4181            8,
4182            &[0x0f, 0xf0],
4183            &[(1, 0), (3, 2)],
4184            &mut encoded,
4185        )
4186        .unwrap();
4187        let flat = LazyFlatVectorData::open_with_doc_limit(
4188            FileHandle::from_bytes(OwnedBytes::new(encoded)),
4189            Some(4),
4190        )
4191        .await
4192        .unwrap();
4193        assert_eq!(flat.num_vectors, flat.num_docs_with_vectors());
4194
4195        let validated = validate_binary_single_value_ann_results(
4196            vec![(3, 2, 0.9), (1, 0, 0.8), (3, 2, 0.7)],
4197            &flat,
4198        )
4199        .unwrap();
4200        assert_eq!(validated, vec![(3, 2, 0.9), (1, 0, 0.8)]);
4201
4202        assert!(matches!(
4203            validate_binary_single_value_ann_results(vec![(2, 0, 1.0)], &flat),
4204            Err(Error::Corruption(_))
4205        ));
4206        assert!(matches!(
4207            validate_binary_single_value_ann_results(vec![(3, 0, 1.0)], &flat),
4208            Err(Error::Corruption(_))
4209        ));
4210    }
4211
4212    #[tokio::test]
4213    async fn multivalue_ann_rerank_streams_past_document_candidate_cap() {
4214        use crate::directories::{FileHandle, OwnedBytes};
4215        use crate::segment::FlatVectorData;
4216
4217        const VALUES: usize = MAX_DENSE_CANDIDATES_PER_SEGMENT + 1;
4218        let mut encoded = Vec::new();
4219        let vectors = vec![1.0f32; VALUES];
4220        let doc_ids: Vec<_> = (0..VALUES).map(|ordinal| (0, ordinal as u16)).collect();
4221        FlatVectorData::serialize_binary_from_flat_streaming(
4222            1,
4223            &vectors,
4224            &doc_ids,
4225            DenseVectorQuantization::F32,
4226            &mut encoded,
4227        )
4228        .unwrap();
4229        let flat = LazyFlatVectorData::open_with_doc_limit(
4230            FileHandle::from_bytes(OwnedBytes::new(encoded)),
4231            Some(1),
4232        )
4233        .await
4234        .unwrap();
4235
4236        let (results, stats) = exact_score_dense_candidate_documents(
4237            &[(0, 0, 0.0)],
4238            &flat,
4239            &[1.0],
4240            false,
4241            crate::query::MultiValueCombiner::Max,
4242            1,
4243        )
4244        .await
4245        .unwrap();
4246        assert_eq!(stats.vector_count, VALUES);
4247        assert_eq!(results.len(), 1);
4248        assert_eq!(results[0].ordinals.len(), VALUES);
4249        assert!((results[0].score - 1.0).abs() < 1e-5);
4250    }
4251}