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

hermes_core/query/vector/
sparse.rs
