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

hermes_core/query/vector/
sparse.rs
