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
11#[derive(Debug, Clone)]
13pub struct SparseVectorQuery {
14 pub field: Field,
16 pub vector: Vec<(u32, f32)>,
18 pub combiner: MultiValueCombiner,
20 pub heap_factor: f32,
23 pub weight_threshold: f32,
26 pub max_query_dims: Option<usize>,
29 pub pruning: Option<f32>,
33 pub min_query_dims: usize,
37 pub over_fetch_factor: f32,
39 pub max_superblocks: usize,
41 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 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 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 pub fn with_combiner(mut self, combiner: MultiValueCombiner) -> Self {
128 self.combiner = combiner;
129 self
130 }
131
132 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 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 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 pub fn with_max_query_dims(mut self, max_dims: usize) -> Self {
162 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 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 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 fn compute_pruned_vector(&self) -> Vec<(u32, f32)> {
187 let original_len = self.vector.len();
188
189 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 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 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 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 #[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 #[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 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 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 #[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 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 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 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 let validation = self.validate(reader);
441 let infos = self.sparse_infos();
442
443 Box::pin(async move {
444 validation?;
445 if infos.is_empty() {
446 return Ok(Box::new(crate::query::EmptyScorer) as Box<dyn Scorer>);
447 }
448
449 if let Some((raw, info)) =
451 crate::query::planner::build_sparse_bmp_results(&infos, reader, limit)
452 {
453 return Ok(crate::query::planner::combine_sparse_results(
454 raw,
455 info.combiner,
456 info.field,
457 limit,
458 ));
459 }
460
461 if let Some((executor, info)) =
463 crate::query::planner::build_sparse_maxscore_executor(&infos, reader, limit, None)
464 {
465 let raw = executor.execute().await?;
466 return Ok(crate::query::planner::combine_sparse_results(
467 raw,
468 info.combiner,
469 info.field,
470 limit,
471 ));
472 }
473
474 Ok(Box::new(crate::query::EmptyScorer) as Box<dyn Scorer>)
475 })
476 }
477
478 #[cfg(feature = "sync")]
479 fn scorer_sync<'a>(
480 &self,
481 reader: &'a SegmentReader,
482 limit: usize,
483 ) -> crate::Result<Box<dyn Scorer + 'a>> {
484 self.validate(reader)?;
485 let infos = self.sparse_infos();
486 if infos.is_empty() {
487 return Ok(Box::new(crate::query::EmptyScorer) as Box<dyn Scorer + 'a>);
488 }
489
490 if let Some((raw, info)) =
492 crate::query::planner::build_sparse_bmp_results(&infos, reader, limit)
493 {
494 return Ok(crate::query::planner::combine_sparse_results(
495 raw,
496 info.combiner,
497 info.field,
498 limit,
499 ));
500 }
501
502 if let Some((executor, info)) =
504 crate::query::planner::build_sparse_maxscore_executor(&infos, reader, limit, None)
505 {
506 let raw = executor.execute_sync()?;
507 return Ok(crate::query::planner::combine_sparse_results(
508 raw,
509 info.combiner,
510 info.field,
511 limit,
512 ));
513 }
514
515 Ok(Box::new(crate::query::EmptyScorer) as Box<dyn Scorer + 'a>)
516 }
517
518 fn count_estimate<'a>(&self, _reader: &'a SegmentReader) -> CountFuture<'a> {
519 Box::pin(async move { Ok(u32::MAX) })
520 }
521
522 fn decompose(&self) -> crate::query::QueryDecomposition {
523 let infos = self.sparse_infos();
524 if infos.is_empty() {
525 crate::query::QueryDecomposition::Opaque
526 } else {
527 crate::query::QueryDecomposition::SparseTerms(infos)
528 }
529 }
530}
531
532#[derive(Debug, Clone)]
540pub struct SparseTermQuery {
541 pub field: Field,
542 pub dim_id: u32,
543 pub weight: f32,
544 pub heap_factor: f32,
546 pub combiner: MultiValueCombiner,
548 pub over_fetch_factor: f32,
550}
551
552impl std::fmt::Display for SparseTermQuery {
553 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
554 write!(
555 f,
556 "SparseTerm({}, dim={}, w={:.3})",
557 self.field.0, self.dim_id, self.weight
558 )
559 }
560}
561
562impl SparseTermQuery {
563 pub fn new(field: Field, dim_id: u32, weight: f32) -> Self {
564 Self {
565 field,
566 dim_id,
567 weight,
568 heap_factor: 1.0,
569 combiner: MultiValueCombiner::default(),
570 over_fetch_factor: 2.0,
571 }
572 }
573
574 pub fn with_heap_factor(mut self, heap_factor: f32) -> Self {
575 self.heap_factor = heap_factor;
576 self
577 }
578
579 pub fn with_combiner(mut self, combiner: MultiValueCombiner) -> Self {
580 self.combiner = combiner;
581 self
582 }
583
584 pub fn with_over_fetch_factor(mut self, factor: f32) -> Self {
585 self.over_fetch_factor = factor.max(1.0);
586 self
587 }
588
589 fn validate(&self, reader: &SegmentReader) -> crate::Result<()> {
590 let entry = reader
591 .schema()
592 .get_field_entry(self.field)
593 .ok_or_else(|| crate::Error::FieldNotFound(self.field.0.to_string()))?;
594 if entry.field_type != crate::dsl::FieldType::SparseVector {
595 return Err(crate::Error::InvalidFieldType {
596 expected: "sparse_vector".to_string(),
597 got: format!("{:?}", entry.field_type),
598 });
599 }
600 if !self.weight.is_finite() {
601 return Err(crate::Error::Query(
602 "sparse term query weight must be finite".to_string(),
603 ));
604 }
605 if !self.heap_factor.is_finite() || !(0.0..=1.0).contains(&self.heap_factor) {
606 return Err(crate::Error::Query(format!(
607 "sparse heap_factor must be finite and in [0, 1], got {}",
608 self.heap_factor
609 )));
610 }
611 if !self.over_fetch_factor.is_finite() || self.over_fetch_factor < 1.0 {
612 return Err(crate::Error::Query(format!(
613 "sparse over_fetch_factor must be finite and at least 1, got {}",
614 self.over_fetch_factor
615 )));
616 }
617 self.combiner.validate().map_err(crate::Error::Query)
618 }
619
620 fn bmp_fallback_scorer<'a>(
622 &self,
623 reader: &'a SegmentReader,
624 limit: usize,
625 ) -> crate::Result<Box<dyn Scorer + 'a>> {
626 if let Some(bmp) = reader.bmp_index(self.field) {
627 let executor_limit =
628 crate::query::planner::bounded_sparse_executor_limit(limit, self.over_fetch_factor)
629 .min(bmp.num_virtual_docs as usize);
630 let results = crate::query::bmp::execute_bmp(
631 bmp,
632 reader.schema().index_label(),
633 reader.schema().get_field_name(self.field).unwrap_or("?"),
634 &[(self.dim_id, self.weight)],
635 executor_limit,
636 self.heap_factor,
637 0,
638 )?;
639 let combined = crate::segment::combine_ordinal_results(
640 results.into_iter().map(|r| (r.doc_id, r.ordinal, r.score)),
641 self.combiner,
642 limit,
643 );
644 return Ok(Box::new(
645 crate::query::planner::VectorTopKResultScorer::new(combined, self.field.0),
646 ));
647 }
648 Ok(Box::new(crate::query::EmptyScorer))
649 }
650
651 fn make_scorer<'a>(
654 &self,
655 reader: &'a SegmentReader,
656 ) -> crate::Result<Option<SparseTermScorer<'a>>> {
657 let si = match reader.sparse_index(self.field) {
658 Some(si) => si,
659 None => return Ok(None),
660 };
661 let (skip_start, skip_count, global_max, block_data_offset) =
662 match si.get_skip_range_full(self.dim_id) {
663 Some(v) => v,
664 None => return Ok(None),
665 };
666 let cursor = crate::query::TermCursor::sparse(
667 si,
668 self.weight,
669 skip_start,
670 skip_count,
671 global_max,
672 block_data_offset,
673 );
674 Ok(Some(SparseTermScorer {
675 cursor,
676 field_id: self.field.0,
677 }))
678 }
679}
680
681impl Query for SparseTermQuery {
682 fn scorer<'a>(&self, reader: &'a SegmentReader, limit: usize) -> ScorerFuture<'a> {
683 let query = self.clone();
684 Box::pin(async move {
685 query.validate(reader)?;
686 let mut scorer = match query.make_scorer(reader)? {
687 Some(s) => s,
688 None => return query.bmp_fallback_scorer(reader, limit),
689 };
690 scorer.cursor.ensure_block_loaded().await.ok();
691 Ok(Box::new(scorer) as Box<dyn Scorer + 'a>)
692 })
693 }
694
695 #[cfg(feature = "sync")]
696 fn scorer_sync<'a>(
697 &self,
698 reader: &'a SegmentReader,
699 limit: usize,
700 ) -> crate::Result<Box<dyn Scorer + 'a>> {
701 self.validate(reader)?;
702 let mut scorer = match self.make_scorer(reader)? {
703 Some(s) => s,
704 None => return self.bmp_fallback_scorer(reader, limit),
705 };
706 scorer.cursor.ensure_block_loaded_sync().ok();
707 Ok(Box::new(scorer) as Box<dyn Scorer + 'a>)
708 }
709
710 fn count_estimate<'a>(&self, reader: &'a SegmentReader) -> CountFuture<'a> {
711 let field = self.field;
712 let dim_id = self.dim_id;
713 Box::pin(async move {
714 let si = match reader.sparse_index(field) {
715 Some(si) => si,
716 None => return Ok(0),
717 };
718 match si.get_skip_range_full(dim_id) {
719 Some((_, skip_count, _, _)) => Ok((skip_count * 256) as u32),
720 None => Ok(0),
721 }
722 })
723 }
724
725 fn decompose(&self) -> crate::query::QueryDecomposition {
726 crate::query::QueryDecomposition::SparseTerms(vec![crate::query::SparseTermQueryInfo {
727 field: self.field,
728 dim_id: self.dim_id,
729 weight: self.weight,
730 heap_factor: self.heap_factor,
731 combiner: self.combiner,
732 over_fetch_factor: self.over_fetch_factor,
733 max_superblocks: 0,
734 }])
735 }
736}
737
738struct SparseTermScorer<'a> {
743 cursor: crate::query::TermCursor<'a>,
744 field_id: u32,
745}
746
747impl crate::query::docset::DocSet for SparseTermScorer<'_> {
748 fn doc(&self) -> DocId {
749 let d = self.cursor.doc();
750 if d == u32::MAX { TERMINATED } else { d }
751 }
752
753 fn advance(&mut self) -> DocId {
754 match self.cursor.advance_sync() {
755 Ok(d) if d == u32::MAX => TERMINATED,
756 Ok(d) => d,
757 Err(_) => TERMINATED,
758 }
759 }
760
761 fn seek(&mut self, target: DocId) -> DocId {
762 match self.cursor.seek_sync(target) {
763 Ok(d) if d == u32::MAX => TERMINATED,
764 Ok(d) => d,
765 Err(_) => TERMINATED,
766 }
767 }
768
769 fn size_hint(&self) -> u32 {
770 0
771 }
772}
773
774impl Scorer for SparseTermScorer<'_> {
775 fn score(&self) -> Score {
776 self.cursor.score()
777 }
778
779 fn matched_positions(&self) -> Option<MatchedPositions> {
780 let ordinal = self.cursor.ordinal();
781 let score = self.cursor.score();
782 if score == 0.0 {
783 return None;
784 }
785 Some(vec![(
786 self.field_id,
787 vec![ScoredPosition::new(ordinal as u32, score)],
788 )])
789 }
790}
791
792#[cfg(test)]
793mod tests {
794 use super::*;
795 use crate::dsl::Field;
796
797 #[test]
798 fn test_sparse_vector_query_new() {
799 let sparse = vec![(1, 0.5), (5, 0.3), (10, 0.2)];
800 let query = SparseVectorQuery::new(Field(0), sparse.clone());
801
802 assert_eq!(query.field, Field(0));
803 assert_eq!(query.vector, sparse);
804 }
805
806 #[test]
807 fn test_sparse_vector_query_from_indices_weights() {
808 let query =
809 SparseVectorQuery::from_indices_weights(Field(0), vec![1, 5, 10], vec![0.5, 0.3, 0.2]);
810
811 assert_eq!(query.vector, vec![(1, 0.5), (5, 0.3), (10, 0.2)]);
812 }
813
814 #[test]
815 fn max_query_dims_cannot_exceed_executor_mask_width() {
816 let vector: Vec<(u32, f32)> = (0..100).map(|dim| (dim, dim as f32 + 1.0)).collect();
817 let query = SparseVectorQuery::new(Field(0), vector).with_max_query_dims(usize::MAX);
818
819 assert_eq!(query.pruned_dims().len(), crate::query::MAX_QUERY_TERMS);
820 assert!(query.pruned_dims().iter().all(|(dim, _)| *dim >= 36));
822 }
823}