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