1use crate::dsl::Field;
4use crate::segment::SegmentReader;
5use crate::{DocId, Score, TERMINATED};
6
7use super::combiner::MultiValueCombiner;
8use crate::query::ScoredPosition;
9use crate::query::traits::{CountFuture, MatchedPositions, Query, Scorer, ScorerFuture};
10
11const DEFAULT_SPARSE_OVER_FETCH_FACTOR: f32 = crate::query::MAX_CANDIDATE_OVERSUBSCRIPTION as f32;
12
13enum SparseQueryInfos {
14 Local(Vec<crate::query::SparseTermQueryInfo>),
15 Shared(std::sync::Arc<[crate::query::SparseTermQueryInfo]>),
16}
17
18impl SparseQueryInfos {
19 fn as_slice(&self) -> &[crate::query::SparseTermQueryInfo] {
20 match self {
21 Self::Local(infos) => infos,
22 Self::Shared(infos) => infos,
23 }
24 }
25}
26
27#[derive(Debug, Clone)]
29pub struct SparseVectorQuery {
30 pub field: Field,
32 pub vector: Vec<(u32, f32)>,
34 pub combiner: MultiValueCombiner,
36 pub heap_factor: f32,
39 pub weight_threshold: f32,
43 pub max_query_dims: Option<usize>,
47 pub pruning: Option<f32>,
52 pub min_query_dims: usize,
56 pub over_fetch_factor: f32,
58 pub lsp_gamma: Option<usize>,
60 pruned: Option<Vec<(u32, f32)>>,
62}
63
64impl std::fmt::Display for SparseVectorQuery {
65 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
66 let dims = self.pruned_dims();
67 write!(f, "Sparse({}, dims={}", self.field.0, dims.len())?;
68 if self.heap_factor < 1.0 {
69 write!(f, ", heap={}", self.heap_factor)?;
70 }
71 if self.vector.len() != dims.len() {
72 write!(f, ", orig={}", self.vector.len())?;
73 }
74 write!(f, ")")
75 }
76}
77
78impl SparseVectorQuery {
79 pub fn new(field: Field, vector: Vec<(u32, f32)>) -> Self {
86 let mut q = Self {
87 field,
88 vector,
89 combiner: MultiValueCombiner::LogSumExp { temperature: 0.7 },
90 heap_factor: 1.0,
91 weight_threshold: 0.0,
92 max_query_dims: Some(crate::query::MAX_QUERY_TERMS),
93 pruning: None,
94 min_query_dims: 4,
95 over_fetch_factor: DEFAULT_SPARSE_OVER_FETCH_FACTOR,
96 lsp_gamma: None,
97 pruned: None,
98 };
99 q.pruned = q.compute_pruned_vector();
100 q
101 }
102
103 pub(crate) fn pruned_dims(&self) -> &[(u32, f32)] {
105 self.pruned.as_deref().unwrap_or(&self.vector)
106 }
107
108 fn validate(&self, reader: &SegmentReader) -> crate::Result<()> {
109 let entry = reader
110 .schema()
111 .get_field_entry(self.field)
112 .ok_or_else(|| crate::Error::FieldNotFound(self.field.0.to_string()))?;
113 if entry.field_type != crate::dsl::FieldType::SparseVector {
114 return Err(crate::Error::InvalidFieldType {
115 expected: "sparse_vector".to_string(),
116 got: format!("{:?}", entry.field_type),
117 });
118 }
119 if self.vector.iter().any(|(_, weight)| !weight.is_finite()) {
120 return Err(crate::Error::Query(
121 "sparse query contains a non-finite weight".to_string(),
122 ));
123 }
124 if self.pruned_dims().len() > crate::query::MAX_QUERY_TERMS {
125 return Err(crate::Error::Query(format!(
126 "sparse query contains more than {} effective dimensions",
127 crate::query::MAX_QUERY_TERMS
128 )));
129 }
130 if !self.heap_factor.is_finite() || !(0.0..=1.0).contains(&self.heap_factor) {
131 return Err(crate::Error::Query(format!(
132 "sparse heap_factor must be finite and in [0, 1], got {}",
133 self.heap_factor
134 )));
135 }
136 if !self.over_fetch_factor.is_finite()
137 || !(1.0..=DEFAULT_SPARSE_OVER_FETCH_FACTOR).contains(&self.over_fetch_factor)
138 {
139 return Err(crate::Error::Query(format!(
140 "sparse over_fetch_factor must be finite and in [1, {DEFAULT_SPARSE_OVER_FETCH_FACTOR}], got {}",
141 self.over_fetch_factor
142 )));
143 }
144 self.combiner.validate().map_err(crate::Error::Query)
145 }
146
147 pub fn with_combiner(mut self, combiner: MultiValueCombiner) -> Self {
149 self.combiner = combiner;
150 self
151 }
152
153 pub fn with_over_fetch_factor(mut self, factor: f32) -> Self {
158 self.over_fetch_factor = factor.clamp(1.0, DEFAULT_SPARSE_OVER_FETCH_FACTOR);
159 self
160 }
161
162 pub fn with_heap_factor(mut self, heap_factor: f32) -> Self {
169 self.heap_factor = heap_factor.clamp(0.0, 1.0);
170 self
171 }
172
173 pub fn with_weight_threshold(mut self, threshold: f32) -> Self {
176 self.weight_threshold = threshold;
177 self.pruned = self.compute_pruned_vector();
178 self
179 }
180
181 pub fn with_max_query_dims(mut self, max_dims: usize) -> Self {
183 self.max_query_dims = Some(max_dims.min(crate::query::MAX_QUERY_TERMS));
186 self.pruned = self.compute_pruned_vector();
187 self
188 }
189
190 pub fn with_pruning(mut self, fraction: f32) -> Self {
193 self.pruning = Some(fraction.clamp(0.0, 1.0));
194 self.pruned = self.compute_pruned_vector();
195 self
196 }
197
198 pub fn with_min_query_dims(mut self, min_dims: usize) -> Self {
201 self.min_query_dims = min_dims;
202 self.pruned = self.compute_pruned_vector();
203 self
204 }
205
206 pub fn with_lsp_gamma(mut self, gamma: usize) -> Self {
209 self.lsp_gamma = Some(gamma);
210 self
211 }
212
213 fn compute_pruned_vector(&self) -> Option<Vec<(u32, f32)>> {
217 let original_len = self.vector.len();
218 let max_dims = self
219 .max_query_dims
220 .unwrap_or(crate::query::MAX_QUERY_TERMS)
221 .min(crate::query::MAX_QUERY_TERMS);
222 let filtering_enabled = self.weight_threshold > 0.0 && original_len > self.min_query_dims;
223 let pruning_enabled = self
224 .pruning
225 .is_some_and(|fraction| fraction < 1.0 && original_len > self.min_query_dims);
226 if !filtering_enabled && !pruning_enabled && original_len <= max_dims {
227 return None;
228 }
229
230 let mut v: Vec<(u32, f32)> = if filtering_enabled {
233 self.vector
234 .iter()
235 .copied()
236 .filter(|(_, w)| w.abs() >= self.weight_threshold)
237 .collect()
238 } else {
239 self.vector.clone()
240 };
241 let after_threshold = v.len();
242
243 let mut sorted_by_weight = false;
246 if let Some(fraction) = self.pruning
247 && fraction < 1.0
248 && v.len() > self.min_query_dims
249 {
250 v.sort_unstable_by(|a, b| b.1.abs().total_cmp(&a.1.abs()).then_with(|| a.0.cmp(&b.0)));
251 sorted_by_weight = true;
252 let keep = ((v.len() as f64 * fraction as f64).ceil() as usize).max(1);
253 v.truncate(keep);
254 }
255 let after_pruning = v.len();
256
257 if v.len() > max_dims {
261 if !sorted_by_weight {
262 v.sort_unstable_by(|a, b| {
263 b.1.abs().total_cmp(&a.1.abs()).then_with(|| a.0.cmp(&b.0))
264 });
265 }
266 v.truncate(max_dims);
267 }
268
269 if v.len() < original_len && log::log_enabled!(log::Level::Debug) {
270 let src: Vec<_> = self
271 .vector
272 .iter()
273 .map(|(d, w)| format!("({},{:.4})", d, w))
274 .collect();
275 let pruned_fmt: Vec<_> = v.iter().map(|(d, w)| format!("({},{:.4})", d, w)).collect();
276 log::debug!(
277 "[sparse query] field={}: pruned {}->{} dims \
278 (threshold: {}->{}, pruning: {}->{}, max_dims: {}->{}), \
279 source=[{}], pruned=[{}]",
280 self.field.0,
281 original_len,
282 v.len(),
283 original_len,
284 after_threshold,
285 after_threshold,
286 after_pruning,
287 after_pruning,
288 v.len(),
289 src.join(", "),
290 pruned_fmt.join(", "),
291 );
292 }
293
294 Some(v)
295 }
296
297 pub fn from_indices_weights(field: Field, indices: Vec<u32>, weights: Vec<f32>) -> Self {
299 let vector: Vec<(u32, f32)> = indices.into_iter().zip(weights).collect();
300 Self::new(field, vector)
301 }
302
303 #[cfg(feature = "native")]
315 pub fn from_text(
316 field: Field,
317 text: &str,
318 tokenizer_name: &str,
319 weighting: crate::structures::QueryWeighting,
320 sparse_index: Option<&crate::segment::SparseIndex>,
321 ) -> crate::Result<Self> {
322 use crate::structures::QueryWeighting;
323 use crate::tokenizer::tokenizer_cache;
324
325 let tokenizer = tokenizer_cache().get_or_load(tokenizer_name)?;
326 let token_ids = tokenizer.tokenize_unique(text)?;
327
328 let weights: Vec<f32> = match weighting {
329 QueryWeighting::One => vec![1.0f32; token_ids.len()],
330 QueryWeighting::Idf => {
331 if let Some(index) = sparse_index {
332 index.idf_weights(&token_ids)
333 } else {
334 vec![1.0f32; token_ids.len()]
335 }
336 }
337 QueryWeighting::IdfFile => {
338 use crate::tokenizer::idf_weights_cache;
339 if let Some(idf) = idf_weights_cache().get_or_load(tokenizer_name, None) {
340 token_ids.iter().map(|&id| idf.get(id)).collect()
341 } else {
342 vec![1.0f32; token_ids.len()]
343 }
344 }
345 };
346
347 let vector: Vec<(u32, f32)> = token_ids.into_iter().zip(weights).collect();
348 Ok(Self::new(field, vector))
349 }
350
351 #[cfg(feature = "native")]
363 pub fn from_text_with_stats(
364 field: Field,
365 text: &str,
366 tokenizer: &crate::tokenizer::HfTokenizer,
367 weighting: crate::structures::QueryWeighting,
368 global_stats: Option<&crate::query::GlobalStats>,
369 ) -> crate::Result<Self> {
370 use crate::structures::QueryWeighting;
371
372 let token_ids = tokenizer.tokenize_unique(text)?;
373
374 let weights: Vec<f32> = match weighting {
375 QueryWeighting::One => vec![1.0f32; token_ids.len()],
376 QueryWeighting::Idf => {
377 if let Some(stats) = global_stats {
378 stats
380 .sparse_idf_weights(field, &token_ids)
381 .into_iter()
382 .map(|w| w.max(0.0))
383 .collect()
384 } else {
385 vec![1.0f32; token_ids.len()]
386 }
387 }
388 QueryWeighting::IdfFile => {
389 vec![1.0f32; token_ids.len()]
392 }
393 };
394
395 let vector: Vec<(u32, f32)> = token_ids.into_iter().zip(weights).collect();
396 Ok(Self::new(field, vector))
397 }
398
399 #[cfg(feature = "native")]
411 pub fn from_text_with_tokenizer_bytes(
412 field: Field,
413 text: &str,
414 tokenizer_bytes: &[u8],
415 weighting: crate::structures::QueryWeighting,
416 global_stats: Option<&crate::query::GlobalStats>,
417 ) -> crate::Result<Self> {
418 use crate::structures::QueryWeighting;
419 use crate::tokenizer::HfTokenizer;
420
421 let tokenizer = HfTokenizer::from_bytes(tokenizer_bytes)?;
422 let token_ids = tokenizer.tokenize_unique(text)?;
423
424 let weights: Vec<f32> = match weighting {
425 QueryWeighting::One => vec![1.0f32; token_ids.len()],
426 QueryWeighting::Idf => {
427 if let Some(stats) = global_stats {
428 stats
430 .sparse_idf_weights(field, &token_ids)
431 .into_iter()
432 .map(|w| w.max(0.0))
433 .collect()
434 } else {
435 vec![1.0f32; token_ids.len()]
436 }
437 }
438 QueryWeighting::IdfFile => {
439 vec![1.0f32; token_ids.len()]
442 }
443 };
444
445 let vector: Vec<(u32, f32)> = token_ids.into_iter().zip(weights).collect();
446 Ok(Self::new(field, vector))
447 }
448}
449
450impl SparseVectorQuery {
451 fn sparse_infos(&self) -> Vec<crate::query::SparseTermQueryInfo> {
455 let candidate_dims: Option<rustc_hash::FxHashSet<u32>> = self
456 .pruned
457 .as_ref()
458 .map(|dimensions| dimensions.iter().map(|&(dimension, _)| dimension).collect());
459 let make_info = |(dim_id, weight)| crate::query::SparseTermQueryInfo {
460 field: self.field,
461 dim_id,
462 weight,
463 candidate: candidate_dims
464 .as_ref()
465 .is_none_or(|dimensions| dimensions.contains(&dim_id)),
466 heap_factor: self.heap_factor,
467 combiner: self.combiner,
468 over_fetch_factor: self.over_fetch_factor,
469 lsp_gamma: self.lsp_gamma,
470 };
471 if self.vector.len() <= crate::query::MAX_QUERY_TERMS {
472 return self.vector.iter().copied().map(make_info).collect();
473 }
474
475 let mut scoring_dims = self.vector.clone();
476 scoring_dims.sort_unstable_by(|left, right| {
477 right
478 .1
479 .abs()
480 .total_cmp(&left.1.abs())
481 .then_with(|| left.0.cmp(&right.0))
482 });
483 scoring_dims.truncate(crate::query::MAX_QUERY_TERMS);
484 scoring_dims.into_iter().map(make_info).collect()
485 }
486
487 fn sparse_infos_for_plan(
488 &self,
489 plan: Option<&std::sync::Arc<crate::query::bmp::LspSegmentPlan>>,
490 ) -> SparseQueryInfos {
491 match plan {
492 Some(plan) => SparseQueryInfos::Shared(std::sync::Arc::clone(&plan.infos)),
493 None => SparseQueryInfos::Local(self.sparse_infos()),
494 }
495 }
496}
497
498impl Query for SparseVectorQuery {
499 fn scorer<'a>(&self, reader: &'a SegmentReader, limit: usize) -> ScorerFuture<'a> {
500 self.scorer_with_options(reader, limit, crate::query::ScorerOptions::with_positions())
501 }
502
503 fn scorer_with_options<'a>(
504 &self,
505 reader: &'a SegmentReader,
506 limit: usize,
507 options: crate::query::ScorerOptions,
508 ) -> ScorerFuture<'a> {
509 let validation = self.validate(reader);
510 let infos = self.sparse_infos_for_plan(options.lsp_plan.as_ref());
511
512 Box::pin(async move {
513 validation?;
514 let infos = infos.as_slice();
515 if infos.is_empty() {
516 return Ok(Box::new(crate::query::EmptyScorer) as Box<dyn Scorer>);
517 }
518
519 if let Some((raw, info)) =
521 crate::query::planner::build_sparse_bmp_results(infos, reader, limit, &options)?
522 {
523 return Ok(crate::query::planner::combine_sparse_results(
524 raw,
525 info.combiner,
526 info.field,
527 limit,
528 ));
529 }
530
531 if let Some((executor, info)) =
533 crate::query::planner::build_sparse_maxscore_executor(infos, reader, limit, None)
534 {
535 let raw = executor.execute().await?;
536 return Ok(crate::query::planner::combine_sparse_results(
537 raw,
538 info.combiner,
539 info.field,
540 limit,
541 ));
542 }
543
544 Ok(Box::new(crate::query::EmptyScorer) as Box<dyn Scorer>)
545 })
546 }
547
548 #[cfg(feature = "sync")]
549 fn scorer_sync<'a>(
550 &self,
551 reader: &'a SegmentReader,
552 limit: usize,
553 ) -> crate::Result<Box<dyn Scorer + 'a>> {
554 self.scorer_sync_with_options(reader, limit, crate::query::ScorerOptions::with_positions())
555 }
556
557 #[cfg(feature = "sync")]
558 fn scorer_sync_with_options<'a>(
559 &self,
560 reader: &'a SegmentReader,
561 limit: usize,
562 options: crate::query::ScorerOptions,
563 ) -> crate::Result<Box<dyn Scorer + 'a>> {
564 self.validate(reader)?;
565 let infos = self.sparse_infos_for_plan(options.lsp_plan.as_ref());
566 let infos = infos.as_slice();
567 if infos.is_empty() {
568 return Ok(Box::new(crate::query::EmptyScorer) as Box<dyn Scorer + 'a>);
569 }
570
571 if let Some((raw, info)) =
573 crate::query::planner::build_sparse_bmp_results(infos, reader, limit, &options)?
574 {
575 return Ok(crate::query::planner::combine_sparse_results(
576 raw,
577 info.combiner,
578 info.field,
579 limit,
580 ));
581 }
582
583 if let Some((executor, info)) =
585 crate::query::planner::build_sparse_maxscore_executor(infos, reader, limit, None)
586 {
587 let raw = executor.execute_sync()?;
588 return Ok(crate::query::planner::combine_sparse_results(
589 raw,
590 info.combiner,
591 info.field,
592 limit,
593 ));
594 }
595
596 Ok(Box::new(crate::query::EmptyScorer) as Box<dyn Scorer + 'a>)
597 }
598
599 fn count_estimate<'a>(&self, _reader: &'a SegmentReader) -> CountFuture<'a> {
600 Box::pin(async move { Ok(u32::MAX) })
601 }
602
603 fn decompose(&self) -> crate::query::QueryDecomposition {
604 let infos = self.sparse_infos();
605 if infos.is_empty() {
606 crate::query::QueryDecomposition::Opaque
607 } else {
608 crate::query::QueryDecomposition::SparseTerms(infos)
609 }
610 }
611}
612
613#[derive(Debug, Clone)]
621pub struct SparseTermQuery {
622 pub field: Field,
623 pub dim_id: u32,
624 pub weight: f32,
625 pub heap_factor: f32,
627 pub combiner: MultiValueCombiner,
629 pub over_fetch_factor: f32,
631}
632
633impl std::fmt::Display for SparseTermQuery {
634 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
635 write!(
636 f,
637 "SparseTerm({}, dim={}, w={:.3})",
638 self.field.0, self.dim_id, self.weight
639 )
640 }
641}
642
643impl SparseTermQuery {
644 pub fn new(field: Field, dim_id: u32, weight: f32) -> Self {
645 Self {
646 field,
647 dim_id,
648 weight,
649 heap_factor: 1.0,
650 combiner: MultiValueCombiner::default(),
651 over_fetch_factor: DEFAULT_SPARSE_OVER_FETCH_FACTOR,
652 }
653 }
654
655 pub fn with_heap_factor(mut self, heap_factor: f32) -> Self {
656 self.heap_factor = heap_factor;
657 self
658 }
659
660 pub fn with_combiner(mut self, combiner: MultiValueCombiner) -> Self {
661 self.combiner = combiner;
662 self
663 }
664
665 pub fn with_over_fetch_factor(mut self, factor: f32) -> Self {
666 self.over_fetch_factor = factor.clamp(1.0, DEFAULT_SPARSE_OVER_FETCH_FACTOR);
667 self
668 }
669
670 fn validate(&self, reader: &SegmentReader) -> crate::Result<()> {
671 let entry = reader
672 .schema()
673 .get_field_entry(self.field)
674 .ok_or_else(|| crate::Error::FieldNotFound(self.field.0.to_string()))?;
675 if entry.field_type != crate::dsl::FieldType::SparseVector {
676 return Err(crate::Error::InvalidFieldType {
677 expected: "sparse_vector".to_string(),
678 got: format!("{:?}", entry.field_type),
679 });
680 }
681 if !self.weight.is_finite() {
682 return Err(crate::Error::Query(
683 "sparse term query weight must be finite".to_string(),
684 ));
685 }
686 if !self.heap_factor.is_finite() || !(0.0..=1.0).contains(&self.heap_factor) {
687 return Err(crate::Error::Query(format!(
688 "sparse heap_factor must be finite and in [0, 1], got {}",
689 self.heap_factor
690 )));
691 }
692 if !self.over_fetch_factor.is_finite()
693 || !(1.0..=DEFAULT_SPARSE_OVER_FETCH_FACTOR).contains(&self.over_fetch_factor)
694 {
695 return Err(crate::Error::Query(format!(
696 "sparse over_fetch_factor must be finite and in [1, {DEFAULT_SPARSE_OVER_FETCH_FACTOR}], got {}",
697 self.over_fetch_factor
698 )));
699 }
700 self.combiner.validate().map_err(crate::Error::Query)
701 }
702
703 fn bmp_fallback_scorer<'a>(
705 &self,
706 reader: &'a SegmentReader,
707 limit: usize,
708 options: &crate::query::ScorerOptions,
709 ) -> crate::Result<Box<dyn Scorer + 'a>> {
710 let infos = [crate::query::SparseTermQueryInfo {
711 field: self.field,
712 dim_id: self.dim_id,
713 weight: self.weight,
714 candidate: true,
715 heap_factor: self.heap_factor,
716 combiner: self.combiner,
717 over_fetch_factor: self.over_fetch_factor,
718 lsp_gamma: None,
719 }];
720 if let Some((raw, info)) =
721 crate::query::planner::build_sparse_bmp_results(&infos, reader, limit, options)?
722 {
723 return Ok(crate::query::planner::combine_sparse_results(
724 raw,
725 info.combiner,
726 info.field,
727 limit,
728 ));
729 }
730 Ok(Box::new(crate::query::EmptyScorer))
731 }
732
733 fn make_scorer<'a>(
736 &self,
737 reader: &'a SegmentReader,
738 ) -> crate::Result<Option<SparseTermScorer<'a>>> {
739 let si = match reader.sparse_index(self.field) {
740 Some(si) => si,
741 None => return Ok(None),
742 };
743 let (skip_start, skip_count, global_max, block_data_offset) =
744 match si.get_skip_range_full(self.dim_id) {
745 Some(v) => v,
746 None => return Ok(None),
747 };
748 let cursor = crate::query::TermCursor::sparse(
749 si,
750 self.weight,
751 skip_start,
752 skip_count,
753 global_max,
754 block_data_offset,
755 );
756 Ok(Some(SparseTermScorer {
757 cursor,
758 field_id: self.field.0,
759 }))
760 }
761}
762
763impl Query for SparseTermQuery {
764 fn scorer<'a>(&self, reader: &'a SegmentReader, limit: usize) -> ScorerFuture<'a> {
765 self.scorer_with_options(reader, limit, crate::query::ScorerOptions::with_positions())
766 }
767
768 fn scorer_with_options<'a>(
769 &self,
770 reader: &'a SegmentReader,
771 limit: usize,
772 options: crate::query::ScorerOptions,
773 ) -> ScorerFuture<'a> {
774 let query = self.clone();
775 Box::pin(async move {
776 query.validate(reader)?;
777 let mut scorer = match query.make_scorer(reader)? {
778 Some(s) => s,
779 None => return query.bmp_fallback_scorer(reader, limit, &options),
780 };
781 scorer.cursor.ensure_block_loaded().await.ok();
782 Ok(Box::new(scorer) as Box<dyn Scorer + 'a>)
783 })
784 }
785
786 #[cfg(feature = "sync")]
787 fn scorer_sync<'a>(
788 &self,
789 reader: &'a SegmentReader,
790 limit: usize,
791 ) -> crate::Result<Box<dyn Scorer + 'a>> {
792 self.scorer_sync_with_options(reader, limit, crate::query::ScorerOptions::with_positions())
793 }
794
795 #[cfg(feature = "sync")]
796 fn scorer_sync_with_options<'a>(
797 &self,
798 reader: &'a SegmentReader,
799 limit: usize,
800 options: crate::query::ScorerOptions,
801 ) -> crate::Result<Box<dyn Scorer + 'a>> {
802 self.validate(reader)?;
803 let mut scorer = match self.make_scorer(reader)? {
804 Some(s) => s,
805 None => return self.bmp_fallback_scorer(reader, limit, &options),
806 };
807 scorer.cursor.ensure_block_loaded_sync().ok();
808 Ok(Box::new(scorer) as Box<dyn Scorer + 'a>)
809 }
810
811 fn count_estimate<'a>(&self, reader: &'a SegmentReader) -> CountFuture<'a> {
812 let field = self.field;
813 let dim_id = self.dim_id;
814 Box::pin(async move {
815 let si = match reader.sparse_index(field) {
816 Some(si) => si,
817 None => return Ok(0),
818 };
819 match si.get_skip_range_full(dim_id) {
820 Some((_, skip_count, _, _)) => Ok((skip_count * 256) as u32),
821 None => Ok(0),
822 }
823 })
824 }
825
826 fn decompose(&self) -> crate::query::QueryDecomposition {
827 crate::query::QueryDecomposition::SparseTerms(vec![crate::query::SparseTermQueryInfo {
828 field: self.field,
829 dim_id: self.dim_id,
830 weight: self.weight,
831 candidate: true,
832 heap_factor: self.heap_factor,
833 combiner: self.combiner,
834 over_fetch_factor: self.over_fetch_factor,
835 lsp_gamma: None,
836 }])
837 }
838}
839
840struct SparseTermScorer<'a> {
845 cursor: crate::query::TermCursor<'a>,
846 field_id: u32,
847}
848
849impl crate::query::docset::DocSet for SparseTermScorer<'_> {
850 fn doc(&self) -> DocId {
851 let d = self.cursor.doc();
852 if d == u32::MAX { TERMINATED } else { d }
853 }
854
855 fn advance(&mut self) -> DocId {
856 match self.cursor.advance_sync() {
857 Ok(d) if d == u32::MAX => TERMINATED,
858 Ok(d) => d,
859 Err(_) => TERMINATED,
860 }
861 }
862
863 fn seek(&mut self, target: DocId) -> DocId {
864 match self.cursor.seek_sync(target) {
865 Ok(d) if d == u32::MAX => TERMINATED,
866 Ok(d) => d,
867 Err(_) => TERMINATED,
868 }
869 }
870
871 fn size_hint(&self) -> u32 {
872 0
873 }
874}
875
876impl Scorer for SparseTermScorer<'_> {
877 fn score(&self) -> Score {
878 self.cursor.score()
879 }
880
881 fn matched_positions(&self) -> Option<MatchedPositions> {
882 let ordinal = self.cursor.ordinal();
883 let score = self.cursor.score();
884 if score == 0.0 {
885 return None;
886 }
887 Some(vec![(
888 self.field_id,
889 vec![ScoredPosition::new(ordinal as u32, score)],
890 )])
891 }
892}
893
894#[cfg(test)]
895mod tests {
896 use super::*;
897 use crate::dsl::Field;
898
899 #[test]
900 fn test_sparse_vector_query_new() {
901 let sparse = vec![(1, 0.5), (5, 0.3), (10, 0.2)];
902 let query = SparseVectorQuery::new(Field(0), sparse.clone());
903
904 assert_eq!(query.field, Field(0));
905 assert_eq!(query.vector, sparse);
906 assert!(
907 query.pruned.is_none(),
908 "the default path must alias the source vector instead of cloning it"
909 );
910 }
911
912 #[test]
913 fn test_sparse_vector_query_from_indices_weights() {
914 let query =
915 SparseVectorQuery::from_indices_weights(Field(0), vec![1, 5, 10], vec![0.5, 0.3, 0.2]);
916
917 assert_eq!(query.vector, vec![(1, 0.5), (5, 0.3), (10, 0.2)]);
918 }
919
920 #[test]
921 fn max_query_dims_cannot_exceed_executor_mask_width() {
922 let vector: Vec<(u32, f32)> = (0..100).map(|dim| (dim, dim as f32 + 1.0)).collect();
923 let query = SparseVectorQuery::new(Field(0), vector).with_max_query_dims(usize::MAX);
924
925 assert_eq!(query.pruned_dims().len(), crate::query::MAX_QUERY_TERMS);
926 assert!(query.pruned_dims().iter().all(|(dim, _)| *dim >= 36));
928 }
929
930 #[test]
931 fn decomposition_keeps_full_scores_and_marks_pruned_candidates() {
932 let query = SparseVectorQuery::new(Field(0), vec![(3, 1.0), (7, 0.8), (11, 0.2)])
933 .with_min_query_dims(0)
934 .with_pruning(0.34);
935 let infos = query.sparse_infos();
936
937 assert_eq!(infos.len(), 3);
938 assert_eq!(
939 infos
940 .iter()
941 .filter(|info| info.candidate)
942 .map(|info| info.dim_id)
943 .collect::<Vec<_>>(),
944 vec![3, 7]
945 );
946 assert_eq!(
947 infos
948 .iter()
949 .map(|info| (info.dim_id, info.weight))
950 .collect::<Vec<_>>(),
951 vec![(3, 1.0), (7, 0.8), (11, 0.2)]
952 );
953 }
954
955 #[test]
956 fn sparse_over_fetch_factor_uses_shared_candidate_bound() {
957 let query = SparseVectorQuery::new(Field(0), vec![]).with_over_fetch_factor(99.0);
958 let term = SparseTermQuery::new(Field(0), 1, 1.0).with_over_fetch_factor(99.0);
959
960 assert_eq!(
961 query.over_fetch_factor,
962 crate::query::MAX_CANDIDATE_OVERSUBSCRIPTION as f32
963 );
964 assert_eq!(
965 term.over_fetch_factor,
966 crate::query::MAX_CANDIDATE_OVERSUBSCRIPTION as f32
967 );
968 }
969}