Skip to main content

hermes_core/query/vector/
sparse.rs

1//! Sparse vector queries for similarity search (MaxScore-based)
2
3use crate::dsl::Field;
4use crate::segment::SegmentReader;
5use crate::{DocId, Score, TERMINATED};
6
7use super::combiner::MultiValueCombiner;
8use crate::query::ScoredPosition;
9use crate::query::traits::{CountFuture, MatchedPositions, Query, Scorer, ScorerFuture};
10
11const DEFAULT_SPARSE_OVER_FETCH_FACTOR: f32 = crate::query::MAX_CANDIDATE_OVERSUBSCRIPTION as f32;
12
13enum SparseQueryInfos {
14    Local(Vec<crate::query::SparseTermQueryInfo>),
15    Shared(std::sync::Arc<[crate::query::SparseTermQueryInfo]>),
16}
17
18impl SparseQueryInfos {
19    fn as_slice(&self) -> &[crate::query::SparseTermQueryInfo] {
20        match self {
21            Self::Local(infos) => infos,
22            Self::Shared(infos) => infos,
23        }
24    }
25}
26
27/// Sparse vector query for similarity search
28#[derive(Debug, Clone)]
29pub struct SparseVectorQuery {
30    /// Field containing the sparse vectors
31    pub field: Field,
32    /// Query vector as (dimension_id, weight) pairs
33    pub vector: Vec<(u32, f32)>,
34    /// How to combine scores for multi-valued documents
35    pub combiner: MultiValueCombiner,
36    /// Approximate search factor (1.0 = exact, lower values = faster but approximate)
37    /// Controls MaxScore pruning aggressiveness in block-max scoring
38    pub heap_factor: f32,
39    /// Minimum abs(weight) for query dimensions (0.0 = no filtering)
40    /// Dimensions below this threshold are dropped from candidate generation.
41    /// BMP still uses the bounded full query when scoring visited candidates.
42    pub weight_threshold: f32,
43    /// Maximum candidate-generation dimensions (None = implementation cap).
44    /// Keeps only the top-k dimensions by abs(weight); BMP final scoring uses
45    /// up to `MAX_QUERY_TERMS` dimensions from the full query.
46    pub max_query_dims: Option<usize>,
47    /// Fraction of query dimensions to keep (0.0-1.0), same semantics as
48    /// indexing-time `pruning`: sort by abs(weight) descending,
49    /// keep top fraction. BMP applies it to candidate generation and scores
50    /// visited candidates with the bounded full query. None or 1.0 = no pruning.
51    pub pruning: Option<f32>,
52    /// Minimum number of query dimensions before pruning and weight_threshold
53    /// filtering are applied. Protects short queries from losing signal.
54    /// Default: 4. Set to 0 to always apply.
55    pub min_query_dims: usize,
56    /// Multiplier on executor limit for ordinal deduplication (1.0 = no over-fetch)
57    pub over_fetch_factor: f32,
58    /// LSP/0 γ. None is depth-derived; Some(0) is exhaustive.
59    pub lsp_gamma: Option<usize>,
60    /// Cached pruned vector; None = use `vector` as-is (no pruning applied)
61    pruned: Option<Vec<(u32, f32)>>,
62}
63
64impl std::fmt::Display for SparseVectorQuery {
65    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
66        let dims = self.pruned_dims();
67        write!(f, "Sparse({}, dims={}", self.field.0, dims.len())?;
68        if self.heap_factor < 1.0 {
69            write!(f, ", heap={}", self.heap_factor)?;
70        }
71        if self.vector.len() != dims.len() {
72            write!(f, ", orig={}", self.vector.len())?;
73        }
74        write!(f, ")")
75    }
76}
77
78impl SparseVectorQuery {
79    /// Create a new sparse vector query
80    ///
81    /// Default combiner is `LogSumExp { temperature: 0.7 }` which provides
82    /// saturation for documents with many sparse vectors (e.g., 100+ ordinals).
83    /// This prevents over-weighting from multiple matches while still allowing
84    /// additional matches to contribute to the score.
85    pub fn new(field: Field, vector: Vec<(u32, f32)>) -> Self {
86        let mut q = Self {
87            field,
88            vector,
89            combiner: MultiValueCombiner::LogSumExp { temperature: 0.7 },
90            heap_factor: 1.0,
91            weight_threshold: 0.0,
92            max_query_dims: Some(crate::query::MAX_QUERY_TERMS),
93            pruning: None,
94            min_query_dims: 4,
95            over_fetch_factor: DEFAULT_SPARSE_OVER_FETCH_FACTOR,
96            lsp_gamma: None,
97            pruned: None,
98        };
99        q.pruned = q.compute_pruned_vector();
100        q
101    }
102
103    /// Effective query dimensions after pruning. Returns `vector` if no pruning is configured.
104    pub(crate) fn pruned_dims(&self) -> &[(u32, f32)] {
105        self.pruned.as_deref().unwrap_or(&self.vector)
106    }
107
108    fn validate(&self, reader: &SegmentReader) -> crate::Result<()> {
109        let entry = reader
110            .schema()
111            .get_field_entry(self.field)
112            .ok_or_else(|| crate::Error::FieldNotFound(self.field.0.to_string()))?;
113        if entry.field_type != crate::dsl::FieldType::SparseVector {
114            return Err(crate::Error::InvalidFieldType {
115                expected: "sparse_vector".to_string(),
116                got: format!("{:?}", entry.field_type),
117            });
118        }
119        if self.vector.iter().any(|(_, weight)| !weight.is_finite()) {
120            return Err(crate::Error::Query(
121                "sparse query contains a non-finite weight".to_string(),
122            ));
123        }
124        if self.pruned_dims().len() > crate::query::MAX_QUERY_TERMS {
125            return Err(crate::Error::Query(format!(
126                "sparse query contains more than {} effective dimensions",
127                crate::query::MAX_QUERY_TERMS
128            )));
129        }
130        if !self.heap_factor.is_finite() || !(0.0..=1.0).contains(&self.heap_factor) {
131            return Err(crate::Error::Query(format!(
132                "sparse heap_factor must be finite and in [0, 1], got {}",
133                self.heap_factor
134            )));
135        }
136        if !self.over_fetch_factor.is_finite()
137            || !(1.0..=DEFAULT_SPARSE_OVER_FETCH_FACTOR).contains(&self.over_fetch_factor)
138        {
139            return Err(crate::Error::Query(format!(
140                "sparse over_fetch_factor must be finite and in [1, {DEFAULT_SPARSE_OVER_FETCH_FACTOR}], got {}",
141                self.over_fetch_factor
142            )));
143        }
144        self.combiner.validate().map_err(crate::Error::Query)
145    }
146
147    /// Set the multi-value score combiner
148    pub fn with_combiner(mut self, combiner: MultiValueCombiner) -> Self {
149        self.combiner = combiner;
150        self
151    }
152
153    /// Set executor over-fetch factor for multi-valued fields.
154    /// After MaxScore execution, ordinal combining may reduce result count;
155    /// this multiplier compensates by fetching more from the executor.
156    /// (1.0 = no over-fetch, 2.0 = fetch 2x then combine down)
157    pub fn with_over_fetch_factor(mut self, factor: f32) -> Self {
158        self.over_fetch_factor = factor.clamp(1.0, DEFAULT_SPARSE_OVER_FETCH_FACTOR);
159        self
160    }
161
162    /// Set the heap factor for approximate search
163    ///
164    /// Controls the trade-off between speed and recall:
165    /// - 1.0 = exact search (default)
166    /// - 0.8-0.9 = ~20-40% faster with minimal recall loss
167    /// - Lower values = more aggressive pruning, faster but lower recall
168    pub fn with_heap_factor(mut self, heap_factor: f32) -> Self {
169        self.heap_factor = heap_factor.clamp(0.0, 1.0);
170        self
171    }
172
173    /// Set minimum weight threshold for query dimensions
174    /// Dimensions with abs(weight) below this are dropped before search.
175    pub fn with_weight_threshold(mut self, threshold: f32) -> Self {
176        self.weight_threshold = threshold;
177        self.pruned = self.compute_pruned_vector();
178        self
179    }
180
181    /// Set maximum number of query dimensions (top-k by weight)
182    pub fn with_max_query_dims(mut self, max_dims: usize) -> Self {
183        // MaxScore and BMP use a u64 query-term mask in their hot paths.  Keep
184        // this invariant here even when an SDL or RPC override asks for more.
185        self.max_query_dims = Some(max_dims.min(crate::query::MAX_QUERY_TERMS));
186        self.pruned = self.compute_pruned_vector();
187        self
188    }
189
190    /// Set pruning fraction (0.0-1.0): keep top fraction of query dims by weight.
191    /// Same semantics as indexing-time `pruning`.
192    pub fn with_pruning(mut self, fraction: f32) -> Self {
193        self.pruning = Some(fraction.clamp(0.0, 1.0));
194        self.pruned = self.compute_pruned_vector();
195        self
196    }
197
198    /// Set minimum query dimensions before pruning/filtering are applied.
199    /// Queries with fewer dimensions than this skip weight_threshold and pruning.
200    pub fn with_min_query_dims(mut self, min_dims: usize) -> Self {
201        self.min_query_dims = min_dims;
202        self.pruned = self.compute_pruned_vector();
203        self
204    }
205
206    /// Select at most the top-γ superblocks by SBMax using LSP/0.
207    /// Zero retains exhaustive SBMax-ordered traversal.
208    pub fn with_lsp_gamma(mut self, gamma: usize) -> Self {
209        self.lsp_gamma = Some(gamma);
210        self
211    }
212
213    /// Apply weight_threshold, pruning, and max_query_dims. `None` aliases the
214    /// original query vector and avoids a second allocation on the default
215    /// unpruned path.
216    fn compute_pruned_vector(&self) -> Option<Vec<(u32, f32)>> {
217        let original_len = self.vector.len();
218        let max_dims = self
219            .max_query_dims
220            .unwrap_or(crate::query::MAX_QUERY_TERMS)
221            .min(crate::query::MAX_QUERY_TERMS);
222        let filtering_enabled = self.weight_threshold > 0.0 && original_len > self.min_query_dims;
223        let pruning_enabled = self
224            .pruning
225            .is_some_and(|fraction| fraction < 1.0 && original_len > self.min_query_dims);
226        if !filtering_enabled && !pruning_enabled && original_len <= max_dims {
227            return None;
228        }
229
230        // Step 1: weight_threshold — drop dimensions below minimum weight
231        // Skip when query has fewer than min_query_dims dimensions
232        let mut v: Vec<(u32, f32)> = if filtering_enabled {
233            self.vector
234                .iter()
235                .copied()
236                .filter(|(_, w)| w.abs() >= self.weight_threshold)
237                .collect()
238        } else {
239            self.vector.clone()
240        };
241        let after_threshold = v.len();
242
243        // Step 2: pruning — keep top fraction by abs(weight), same as indexing
244        // Skip when query has fewer than min_query_dims dimensions
245        let mut sorted_by_weight = false;
246        if let Some(fraction) = self.pruning
247            && fraction < 1.0
248            && v.len() > self.min_query_dims
249        {
250            v.sort_unstable_by(|a, b| b.1.abs().total_cmp(&a.1.abs()).then_with(|| a.0.cmp(&b.0)));
251            sorted_by_weight = true;
252            let keep = ((v.len() as f64 * fraction as f64).ceil() as usize).max(1);
253            v.truncate(keep);
254        }
255        let after_pruning = v.len();
256
257        // Step 3: max_query_dims — absolute cap on dimensions.  The hard
258        // MAX_QUERY_TERMS bound is a correctness requirement, not merely a
259        // tuning default: both sparse executors represent query terms in u64.
260        if v.len() > max_dims {
261            if !sorted_by_weight {
262                v.sort_unstable_by(|a, b| {
263                    b.1.abs().total_cmp(&a.1.abs()).then_with(|| a.0.cmp(&b.0))
264                });
265            }
266            v.truncate(max_dims);
267        }
268
269        if v.len() < original_len && log::log_enabled!(log::Level::Debug) {
270            let src: Vec<_> = self
271                .vector
272                .iter()
273                .map(|(d, w)| format!("({},{:.4})", d, w))
274                .collect();
275            let pruned_fmt: Vec<_> = v.iter().map(|(d, w)| format!("({},{:.4})", d, w)).collect();
276            log::debug!(
277                "[sparse query] field={}: pruned {}->{} dims \
278                 (threshold: {}->{}, pruning: {}->{}, max_dims: {}->{}), \
279                 source=[{}], pruned=[{}]",
280                self.field.0,
281                original_len,
282                v.len(),
283                original_len,
284                after_threshold,
285                after_threshold,
286                after_pruning,
287                after_pruning,
288                v.len(),
289                src.join(", "),
290                pruned_fmt.join(", "),
291            );
292        }
293
294        Some(v)
295    }
296
297    /// Create from separate indices and weights vectors
298    pub fn from_indices_weights(field: Field, indices: Vec<u32>, weights: Vec<f32>) -> Self {
299        let vector: Vec<(u32, f32)> = indices.into_iter().zip(weights).collect();
300        Self::new(field, vector)
301    }
302
303    /// Create from raw text using a HuggingFace tokenizer (single segment)
304    ///
305    /// This method tokenizes the text and creates a sparse vector query.
306    /// For multi-segment indexes, use `from_text_with_stats` instead.
307    ///
308    /// # Arguments
309    /// * `field` - The sparse vector field to search
310    /// * `text` - Raw text to tokenize
311    /// * `tokenizer_name` - HuggingFace tokenizer path (e.g., "bert-base-uncased")
312    /// * `weighting` - Weighting strategy for tokens
313    /// * `sparse_index` - Optional sparse index for IDF lookup (required for IDF weighting)
314    #[cfg(feature = "native")]
315    pub fn from_text(
316        field: Field,
317        text: &str,
318        tokenizer_name: &str,
319        weighting: crate::structures::QueryWeighting,
320        sparse_index: Option<&crate::segment::SparseIndex>,
321    ) -> crate::Result<Self> {
322        use crate::structures::QueryWeighting;
323        use crate::tokenizer::tokenizer_cache;
324
325        let tokenizer = tokenizer_cache().get_or_load(tokenizer_name)?;
326        let token_ids = tokenizer.tokenize_unique(text)?;
327
328        let weights: Vec<f32> = match weighting {
329            QueryWeighting::One => vec![1.0f32; token_ids.len()],
330            QueryWeighting::Idf => {
331                if let Some(index) = sparse_index {
332                    index.idf_weights(&token_ids)
333                } else {
334                    vec![1.0f32; token_ids.len()]
335                }
336            }
337            QueryWeighting::IdfFile => {
338                use crate::tokenizer::idf_weights_cache;
339                if let Some(idf) = idf_weights_cache().get_or_load(tokenizer_name, None) {
340                    token_ids.iter().map(|&id| idf.get(id)).collect()
341                } else {
342                    vec![1.0f32; token_ids.len()]
343                }
344            }
345        };
346
347        let vector: Vec<(u32, f32)> = token_ids.into_iter().zip(weights).collect();
348        Ok(Self::new(field, vector))
349    }
350
351    /// Create from raw text using global statistics (multi-segment)
352    ///
353    /// This is the recommended method for multi-segment indexes as it uses
354    /// aggregated IDF values across all segments for consistent ranking.
355    ///
356    /// # Arguments
357    /// * `field` - The sparse vector field to search
358    /// * `text` - Raw text to tokenize
359    /// * `tokenizer` - Pre-loaded HuggingFace tokenizer
360    /// * `weighting` - Weighting strategy for tokens
361    /// * `global_stats` - Global statistics for IDF computation
362    #[cfg(feature = "native")]
363    pub fn from_text_with_stats(
364        field: Field,
365        text: &str,
366        tokenizer: &crate::tokenizer::HfTokenizer,
367        weighting: crate::structures::QueryWeighting,
368        global_stats: Option<&crate::query::GlobalStats>,
369    ) -> crate::Result<Self> {
370        use crate::structures::QueryWeighting;
371
372        let token_ids = tokenizer.tokenize_unique(text)?;
373
374        let weights: Vec<f32> = match weighting {
375            QueryWeighting::One => vec![1.0f32; token_ids.len()],
376            QueryWeighting::Idf => {
377                if let Some(stats) = global_stats {
378                    // Clamp to zero: negative weights don't make sense for IDF
379                    stats
380                        .sparse_idf_weights(field, &token_ids)
381                        .into_iter()
382                        .map(|w| w.max(0.0))
383                        .collect()
384                } else {
385                    vec![1.0f32; token_ids.len()]
386                }
387            }
388            QueryWeighting::IdfFile => {
389                // IdfFile requires a tokenizer name for HF model lookup;
390                // this code path doesn't have one, so fall back to 1.0
391                vec![1.0f32; token_ids.len()]
392            }
393        };
394
395        let vector: Vec<(u32, f32)> = token_ids.into_iter().zip(weights).collect();
396        Ok(Self::new(field, vector))
397    }
398
399    /// Create from raw text, loading tokenizer from index directory
400    ///
401    /// This method supports the `index://` prefix for tokenizer paths,
402    /// loading tokenizer.json from the index directory.
403    ///
404    /// # Arguments
405    /// * `field` - The sparse vector field to search
406    /// * `text` - Raw text to tokenize
407    /// * `tokenizer_bytes` - Tokenizer JSON bytes (pre-loaded from directory)
408    /// * `weighting` - Weighting strategy for tokens
409    /// * `global_stats` - Global statistics for IDF computation
410    #[cfg(feature = "native")]
411    pub fn from_text_with_tokenizer_bytes(
412        field: Field,
413        text: &str,
414        tokenizer_bytes: &[u8],
415        weighting: crate::structures::QueryWeighting,
416        global_stats: Option<&crate::query::GlobalStats>,
417    ) -> crate::Result<Self> {
418        use crate::structures::QueryWeighting;
419        use crate::tokenizer::HfTokenizer;
420
421        let tokenizer = HfTokenizer::from_bytes(tokenizer_bytes)?;
422        let token_ids = tokenizer.tokenize_unique(text)?;
423
424        let weights: Vec<f32> = match weighting {
425            QueryWeighting::One => vec![1.0f32; token_ids.len()],
426            QueryWeighting::Idf => {
427                if let Some(stats) = global_stats {
428                    // Clamp to zero: negative weights don't make sense for IDF
429                    stats
430                        .sparse_idf_weights(field, &token_ids)
431                        .into_iter()
432                        .map(|w| w.max(0.0))
433                        .collect()
434                } else {
435                    vec![1.0f32; token_ids.len()]
436                }
437            }
438            QueryWeighting::IdfFile => {
439                // IdfFile requires a tokenizer name for HF model lookup;
440                // this code path doesn't have one, so fall back to 1.0
441                vec![1.0f32; token_ids.len()]
442            }
443        };
444
445        let vector: Vec<(u32, f32)> = token_ids.into_iter().zip(weights).collect();
446        Ok(Self::new(field, vector))
447    }
448}
449
450impl SparseVectorQuery {
451    /// Build a bounded full-query decomposition and mark the pruned terms used
452    /// for candidate generation. LSP/BMP scores visited documents with the
453    /// full list; MaxScore continues to consume only marked candidate terms.
454    fn sparse_infos(&self) -> Vec<crate::query::SparseTermQueryInfo> {
455        let candidate_dims: Option<rustc_hash::FxHashSet<u32>> = self
456            .pruned
457            .as_ref()
458            .map(|dimensions| dimensions.iter().map(|&(dimension, _)| dimension).collect());
459        let make_info = |(dim_id, weight)| crate::query::SparseTermQueryInfo {
460            field: self.field,
461            dim_id,
462            weight,
463            candidate: candidate_dims
464                .as_ref()
465                .is_none_or(|dimensions| dimensions.contains(&dim_id)),
466            heap_factor: self.heap_factor,
467            combiner: self.combiner,
468            over_fetch_factor: self.over_fetch_factor,
469            lsp_gamma: self.lsp_gamma,
470        };
471        if self.vector.len() <= crate::query::MAX_QUERY_TERMS {
472            return self.vector.iter().copied().map(make_info).collect();
473        }
474
475        let mut scoring_dims = self.vector.clone();
476        scoring_dims.sort_unstable_by(|left, right| {
477            right
478                .1
479                .abs()
480                .total_cmp(&left.1.abs())
481                .then_with(|| left.0.cmp(&right.0))
482        });
483        scoring_dims.truncate(crate::query::MAX_QUERY_TERMS);
484        scoring_dims.into_iter().map(make_info).collect()
485    }
486
487    fn sparse_infos_for_plan(
488        &self,
489        plan: Option<&std::sync::Arc<crate::query::bmp::LspSegmentPlan>>,
490    ) -> SparseQueryInfos {
491        match plan {
492            Some(plan) => SparseQueryInfos::Shared(std::sync::Arc::clone(&plan.infos)),
493            None => SparseQueryInfos::Local(self.sparse_infos()),
494        }
495    }
496}
497
498impl Query for SparseVectorQuery {
499    fn scorer<'a>(&self, reader: &'a SegmentReader, limit: usize) -> ScorerFuture<'a> {
500        self.scorer_with_options(reader, limit, crate::query::ScorerOptions::with_positions())
501    }
502
503    fn scorer_with_options<'a>(
504        &self,
505        reader: &'a SegmentReader,
506        limit: usize,
507        options: crate::query::ScorerOptions,
508    ) -> ScorerFuture<'a> {
509        let validation = self.validate(reader);
510        let infos = self.sparse_infos_for_plan(options.lsp_plan.as_ref());
511
512        Box::pin(async move {
513            validation?;
514            let infos = infos.as_slice();
515            if infos.is_empty() {
516                return Ok(Box::new(crate::query::EmptyScorer) as Box<dyn Scorer>);
517            }
518
519            // Auto-detect: try BMP executor first (coupled to index format)
520            if let Some((raw, info)) =
521                crate::query::planner::build_sparse_bmp_results(infos, reader, limit, &options)?
522            {
523                return Ok(crate::query::planner::combine_sparse_results(
524                    raw,
525                    info.combiner,
526                    info.field,
527                    limit,
528                ));
529            }
530
531            // Fall back to MaxScore execution
532            if let Some((executor, info)) =
533                crate::query::planner::build_sparse_maxscore_executor(infos, reader, limit, None)
534            {
535                let raw = executor.execute().await?;
536                return Ok(crate::query::planner::combine_sparse_results(
537                    raw,
538                    info.combiner,
539                    info.field,
540                    limit,
541                ));
542            }
543
544            Ok(Box::new(crate::query::EmptyScorer) as Box<dyn Scorer>)
545        })
546    }
547
548    #[cfg(feature = "sync")]
549    fn scorer_sync<'a>(
550        &self,
551        reader: &'a SegmentReader,
552        limit: usize,
553    ) -> crate::Result<Box<dyn Scorer + 'a>> {
554        self.scorer_sync_with_options(reader, limit, crate::query::ScorerOptions::with_positions())
555    }
556
557    #[cfg(feature = "sync")]
558    fn scorer_sync_with_options<'a>(
559        &self,
560        reader: &'a SegmentReader,
561        limit: usize,
562        options: crate::query::ScorerOptions,
563    ) -> crate::Result<Box<dyn Scorer + 'a>> {
564        self.validate(reader)?;
565        let infos = self.sparse_infos_for_plan(options.lsp_plan.as_ref());
566        let infos = infos.as_slice();
567        if infos.is_empty() {
568            return Ok(Box::new(crate::query::EmptyScorer) as Box<dyn Scorer + 'a>);
569        }
570
571        // Auto-detect: try BMP executor first (coupled to index format)
572        if let Some((raw, info)) =
573            crate::query::planner::build_sparse_bmp_results(infos, reader, limit, &options)?
574        {
575            return Ok(crate::query::planner::combine_sparse_results(
576                raw,
577                info.combiner,
578                info.field,
579                limit,
580            ));
581        }
582
583        // Fall back to MaxScore execution
584        if let Some((executor, info)) =
585            crate::query::planner::build_sparse_maxscore_executor(infos, reader, limit, None)
586        {
587            let raw = executor.execute_sync()?;
588            return Ok(crate::query::planner::combine_sparse_results(
589                raw,
590                info.combiner,
591                info.field,
592                limit,
593            ));
594        }
595
596        Ok(Box::new(crate::query::EmptyScorer) as Box<dyn Scorer + 'a>)
597    }
598
599    fn count_estimate<'a>(&self, _reader: &'a SegmentReader) -> CountFuture<'a> {
600        Box::pin(async move { Ok(u32::MAX) })
601    }
602
603    fn decompose(&self) -> crate::query::QueryDecomposition {
604        let infos = self.sparse_infos();
605        if infos.is_empty() {
606            crate::query::QueryDecomposition::Opaque
607        } else {
608            crate::query::QueryDecomposition::SparseTerms(infos)
609        }
610    }
611}
612
613// ── SparseTermQuery: single sparse dimension query (like TermQuery for text) ──
614
615/// Query for a single sparse vector dimension.
616///
617/// Analogous to `TermQuery` for text: searches one dimension's posting list
618/// with a given weight. Multiple `SparseTermQuery` instances are combined as
619/// `BooleanQuery` SHOULD clauses to form a full sparse vector search.
620#[derive(Debug, Clone)]
621pub struct SparseTermQuery {
622    pub field: Field,
623    pub dim_id: u32,
624    pub weight: f32,
625    /// MaxScore heap factor (1.0 = exact, lower = approximate)
626    pub heap_factor: f32,
627    /// Multi-value combiner for ordinal deduplication
628    pub combiner: MultiValueCombiner,
629    /// Multiplier on executor limit to compensate for ordinal deduplication
630    pub over_fetch_factor: f32,
631}
632
633impl std::fmt::Display for SparseTermQuery {
634    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
635        write!(
636            f,
637            "SparseTerm({}, dim={}, w={:.3})",
638            self.field.0, self.dim_id, self.weight
639        )
640    }
641}
642
643impl SparseTermQuery {
644    pub fn new(field: Field, dim_id: u32, weight: f32) -> Self {
645        Self {
646            field,
647            dim_id,
648            weight,
649            heap_factor: 1.0,
650            combiner: MultiValueCombiner::default(),
651            over_fetch_factor: DEFAULT_SPARSE_OVER_FETCH_FACTOR,
652        }
653    }
654
655    pub fn with_heap_factor(mut self, heap_factor: f32) -> Self {
656        self.heap_factor = heap_factor;
657        self
658    }
659
660    pub fn with_combiner(mut self, combiner: MultiValueCombiner) -> Self {
661        self.combiner = combiner;
662        self
663    }
664
665    pub fn with_over_fetch_factor(mut self, factor: f32) -> Self {
666        self.over_fetch_factor = factor.clamp(1.0, DEFAULT_SPARSE_OVER_FETCH_FACTOR);
667        self
668    }
669
670    fn validate(&self, reader: &SegmentReader) -> crate::Result<()> {
671        let entry = reader
672            .schema()
673            .get_field_entry(self.field)
674            .ok_or_else(|| crate::Error::FieldNotFound(self.field.0.to_string()))?;
675        if entry.field_type != crate::dsl::FieldType::SparseVector {
676            return Err(crate::Error::InvalidFieldType {
677                expected: "sparse_vector".to_string(),
678                got: format!("{:?}", entry.field_type),
679            });
680        }
681        if !self.weight.is_finite() {
682            return Err(crate::Error::Query(
683                "sparse term query weight must be finite".to_string(),
684            ));
685        }
686        if !self.heap_factor.is_finite() || !(0.0..=1.0).contains(&self.heap_factor) {
687            return Err(crate::Error::Query(format!(
688                "sparse heap_factor must be finite and in [0, 1], got {}",
689                self.heap_factor
690            )));
691        }
692        if !self.over_fetch_factor.is_finite()
693            || !(1.0..=DEFAULT_SPARSE_OVER_FETCH_FACTOR).contains(&self.over_fetch_factor)
694        {
695            return Err(crate::Error::Query(format!(
696                "sparse over_fetch_factor must be finite and in [1, {DEFAULT_SPARSE_OVER_FETCH_FACTOR}], got {}",
697                self.over_fetch_factor
698            )));
699        }
700        self.combiner.validate().map_err(crate::Error::Query)
701    }
702
703    /// BMP fallback: execute BMP for this single dimension and wrap in a TopK scorer.
704    fn bmp_fallback_scorer<'a>(
705        &self,
706        reader: &'a SegmentReader,
707        limit: usize,
708        options: &crate::query::ScorerOptions,
709    ) -> crate::Result<Box<dyn Scorer + 'a>> {
710        let infos = [crate::query::SparseTermQueryInfo {
711            field: self.field,
712            dim_id: self.dim_id,
713            weight: self.weight,
714            candidate: true,
715            heap_factor: self.heap_factor,
716            combiner: self.combiner,
717            over_fetch_factor: self.over_fetch_factor,
718            lsp_gamma: None,
719        }];
720        if let Some((raw, info)) =
721            crate::query::planner::build_sparse_bmp_results(&infos, reader, limit, options)?
722        {
723            return Ok(crate::query::planner::combine_sparse_results(
724                raw,
725                info.combiner,
726                info.field,
727                limit,
728            ));
729        }
730        Ok(Box::new(crate::query::EmptyScorer))
731    }
732
733    /// Create a SparseTermScorer from this query's config against a segment.
734    /// Returns EmptyScorer if the dimension doesn't exist.
735    fn make_scorer<'a>(
736        &self,
737        reader: &'a SegmentReader,
738    ) -> crate::Result<Option<SparseTermScorer<'a>>> {
739        let si = match reader.sparse_index(self.field) {
740            Some(si) => si,
741            None => return Ok(None),
742        };
743        let (skip_start, skip_count, global_max, block_data_offset) =
744            match si.get_skip_range_full(self.dim_id) {
745                Some(v) => v,
746                None => return Ok(None),
747            };
748        let cursor = crate::query::TermCursor::sparse(
749            si,
750            self.weight,
751            skip_start,
752            skip_count,
753            global_max,
754            block_data_offset,
755        );
756        Ok(Some(SparseTermScorer {
757            cursor,
758            field_id: self.field.0,
759        }))
760    }
761}
762
763impl Query for SparseTermQuery {
764    fn scorer<'a>(&self, reader: &'a SegmentReader, limit: usize) -> ScorerFuture<'a> {
765        self.scorer_with_options(reader, limit, crate::query::ScorerOptions::with_positions())
766    }
767
768    fn scorer_with_options<'a>(
769        &self,
770        reader: &'a SegmentReader,
771        limit: usize,
772        options: crate::query::ScorerOptions,
773    ) -> ScorerFuture<'a> {
774        let query = self.clone();
775        Box::pin(async move {
776            query.validate(reader)?;
777            let mut scorer = match query.make_scorer(reader)? {
778                Some(s) => s,
779                None => return query.bmp_fallback_scorer(reader, limit, &options),
780            };
781            scorer.cursor.ensure_block_loaded().await.ok();
782            Ok(Box::new(scorer) as Box<dyn Scorer + 'a>)
783        })
784    }
785
786    #[cfg(feature = "sync")]
787    fn scorer_sync<'a>(
788        &self,
789        reader: &'a SegmentReader,
790        limit: usize,
791    ) -> crate::Result<Box<dyn Scorer + 'a>> {
792        self.scorer_sync_with_options(reader, limit, crate::query::ScorerOptions::with_positions())
793    }
794
795    #[cfg(feature = "sync")]
796    fn scorer_sync_with_options<'a>(
797        &self,
798        reader: &'a SegmentReader,
799        limit: usize,
800        options: crate::query::ScorerOptions,
801    ) -> crate::Result<Box<dyn Scorer + 'a>> {
802        self.validate(reader)?;
803        let mut scorer = match self.make_scorer(reader)? {
804            Some(s) => s,
805            None => return self.bmp_fallback_scorer(reader, limit, &options),
806        };
807        scorer.cursor.ensure_block_loaded_sync().ok();
808        Ok(Box::new(scorer) as Box<dyn Scorer + 'a>)
809    }
810
811    fn count_estimate<'a>(&self, reader: &'a SegmentReader) -> CountFuture<'a> {
812        let field = self.field;
813        let dim_id = self.dim_id;
814        Box::pin(async move {
815            let si = match reader.sparse_index(field) {
816                Some(si) => si,
817                None => return Ok(0),
818            };
819            match si.get_skip_range_full(dim_id) {
820                Some((_, skip_count, _, _)) => Ok((skip_count * 256) as u32),
821                None => Ok(0),
822            }
823        })
824    }
825
826    fn decompose(&self) -> crate::query::QueryDecomposition {
827        crate::query::QueryDecomposition::SparseTerms(vec![crate::query::SparseTermQueryInfo {
828            field: self.field,
829            dim_id: self.dim_id,
830            weight: self.weight,
831            candidate: true,
832            heap_factor: self.heap_factor,
833            combiner: self.combiner,
834            over_fetch_factor: self.over_fetch_factor,
835            lsp_gamma: None,
836        }])
837    }
838}
839
840/// Lazy scorer for a single sparse dimension, backed by `TermCursor::Sparse`.
841///
842/// Iterates through the posting list block-by-block using sync I/O.
843/// Score for each doc = `query_weight * quantized_stored_weight`.
844struct SparseTermScorer<'a> {
845    cursor: crate::query::TermCursor<'a>,
846    field_id: u32,
847}
848
849impl crate::query::docset::DocSet for SparseTermScorer<'_> {
850    fn doc(&self) -> DocId {
851        let d = self.cursor.doc();
852        if d == u32::MAX { TERMINATED } else { d }
853    }
854
855    fn advance(&mut self) -> DocId {
856        match self.cursor.advance_sync() {
857            Ok(d) if d == u32::MAX => TERMINATED,
858            Ok(d) => d,
859            Err(_) => TERMINATED,
860        }
861    }
862
863    fn seek(&mut self, target: DocId) -> DocId {
864        match self.cursor.seek_sync(target) {
865            Ok(d) if d == u32::MAX => TERMINATED,
866            Ok(d) => d,
867            Err(_) => TERMINATED,
868        }
869    }
870
871    fn size_hint(&self) -> u32 {
872        0
873    }
874}
875
876impl Scorer for SparseTermScorer<'_> {
877    fn score(&self) -> Score {
878        self.cursor.score()
879    }
880
881    fn matched_positions(&self) -> Option<MatchedPositions> {
882        let ordinal = self.cursor.ordinal();
883        let score = self.cursor.score();
884        if score == 0.0 {
885            return None;
886        }
887        Some(vec![(
888            self.field_id,
889            vec![ScoredPosition::new(ordinal as u32, score)],
890        )])
891    }
892}
893
894#[cfg(test)]
895mod tests {
896    use super::*;
897    use crate::dsl::Field;
898
899    #[test]
900    fn test_sparse_vector_query_new() {
901        let sparse = vec![(1, 0.5), (5, 0.3), (10, 0.2)];
902        let query = SparseVectorQuery::new(Field(0), sparse.clone());
903
904        assert_eq!(query.field, Field(0));
905        assert_eq!(query.vector, sparse);
906        assert!(
907            query.pruned.is_none(),
908            "the default path must alias the source vector instead of cloning it"
909        );
910    }
911
912    #[test]
913    fn test_sparse_vector_query_from_indices_weights() {
914        let query =
915            SparseVectorQuery::from_indices_weights(Field(0), vec![1, 5, 10], vec![0.5, 0.3, 0.2]);
916
917        assert_eq!(query.vector, vec![(1, 0.5), (5, 0.3), (10, 0.2)]);
918    }
919
920    #[test]
921    fn max_query_dims_cannot_exceed_executor_mask_width() {
922        let vector: Vec<(u32, f32)> = (0..100).map(|dim| (dim, dim as f32 + 1.0)).collect();
923        let query = SparseVectorQuery::new(Field(0), vector).with_max_query_dims(usize::MAX);
924
925        assert_eq!(query.pruned_dims().len(), crate::query::MAX_QUERY_TERMS);
926        // Pruning retains the dimensions with the largest absolute weights.
927        assert!(query.pruned_dims().iter().all(|(dim, _)| *dim >= 36));
928    }
929
930    #[test]
931    fn decomposition_keeps_full_scores_and_marks_pruned_candidates() {
932        let query = SparseVectorQuery::new(Field(0), vec![(3, 1.0), (7, 0.8), (11, 0.2)])
933            .with_min_query_dims(0)
934            .with_pruning(0.34);
935        let infos = query.sparse_infos();
936
937        assert_eq!(infos.len(), 3);
938        assert_eq!(
939            infos
940                .iter()
941                .filter(|info| info.candidate)
942                .map(|info| info.dim_id)
943                .collect::<Vec<_>>(),
944            vec![3, 7]
945        );
946        assert_eq!(
947            infos
948                .iter()
949                .map(|info| (info.dim_id, info.weight))
950                .collect::<Vec<_>>(),
951            vec![(3, 1.0), (7, 0.8), (11, 0.2)]
952        );
953    }
954
955    #[test]
956    fn sparse_over_fetch_factor_uses_shared_candidate_bound() {
957        let query = SparseVectorQuery::new(Field(0), vec![]).with_over_fetch_factor(99.0);
958        let term = SparseTermQuery::new(Field(0), 1, 1.0).with_over_fetch_factor(99.0);
959
960        assert_eq!(
961            query.over_fetch_factor,
962            crate::query::MAX_CANDIDATE_OVERSUBSCRIPTION as f32
963        );
964        assert_eq!(
965            term.over_fetch_factor,
966            crate::query::MAX_CANDIDATE_OVERSUBSCRIPTION as f32
967        );
968    }
969}