lb_tantivy/query/
const_score_query.rs

1use std::fmt;
2
3use crate::docset::COLLECT_BLOCK_BUFFER_LEN;
4use crate::query::{EnableScoring, Explanation, Query, Scorer, Weight};
5use crate::{DocId, DocSet, Score, SegmentReader, TantivyError, Term};
6
7/// `ConstScoreQuery` is a wrapper over a query to provide a constant score.
8/// It can avoid unnecessary score computation on the wrapped query.
9///
10/// The document set matched by the `ConstScoreQuery` is strictly the same as the underlying query.
11/// The configured score is used for each document.
12pub struct ConstScoreQuery {
13    query: Box<dyn Query>,
14    score: Score,
15}
16
17impl ConstScoreQuery {
18    /// Builds a const score query.
19    pub fn new(query: Box<dyn Query>, score: Score) -> ConstScoreQuery {
20        ConstScoreQuery { query, score }
21    }
22}
23
24impl Clone for ConstScoreQuery {
25    fn clone(&self) -> Self {
26        ConstScoreQuery {
27            query: self.query.box_clone(),
28            score: self.score,
29        }
30    }
31}
32
33impl fmt::Debug for ConstScoreQuery {
34    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
35        write!(f, "Const(score={}, query={:?})", self.score, self.query)
36    }
37}
38
39impl Query for ConstScoreQuery {
40    fn weight(&self, enable_scoring: EnableScoring<'_>) -> crate::Result<Box<dyn Weight>> {
41        let inner_weight = self.query.weight(enable_scoring)?;
42        Ok(if enable_scoring.is_scoring_enabled() {
43            Box::new(ConstWeight::new(inner_weight, self.score))
44        } else {
45            inner_weight
46        })
47    }
48
49    fn query_terms<'a>(&'a self, visitor: &mut dyn FnMut(&'a Term, bool)) {
50        self.query.query_terms(visitor);
51    }
52}
53
54struct ConstWeight {
55    weight: Box<dyn Weight>,
56    score: Score,
57}
58
59impl ConstWeight {
60    pub fn new(weight: Box<dyn Weight>, score: Score) -> Self {
61        ConstWeight { weight, score }
62    }
63}
64
65impl Weight for ConstWeight {
66    fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result<Box<dyn Scorer>> {
67        let inner_scorer = self.weight.scorer(reader, boost)?;
68        Ok(Box::new(ConstScorer::new(inner_scorer, boost * self.score)))
69    }
70
71    fn explain(&self, reader: &SegmentReader, doc: u32) -> crate::Result<Explanation> {
72        let mut scorer = self.scorer(reader, 1.0)?;
73        if scorer.seek(doc) != doc {
74            return Err(TantivyError::InvalidArgument(format!(
75                "Document #({doc}) does not match"
76            )));
77        }
78        let mut explanation = Explanation::new("Const", self.score);
79        let underlying_explanation = self.weight.explain(reader, doc)?;
80        explanation.add_detail(underlying_explanation);
81        Ok(explanation)
82    }
83
84    fn count(&self, reader: &SegmentReader) -> crate::Result<u32> {
85        self.weight.count(reader)
86    }
87}
88
89/// Wraps a `DocSet` and simply returns a constant `Scorer`.
90/// The `ConstScorer` is useful if you have a `DocSet` where
91/// you needed a scorer.
92///
93/// The `ConstScorer`'s constant score can be set
94/// by calling `.set_score(...)`.
95pub struct ConstScorer<TDocSet: DocSet> {
96    docset: TDocSet,
97    score: Score,
98}
99
100impl<TDocSet: DocSet> ConstScorer<TDocSet> {
101    /// Creates a new `ConstScorer`.
102    pub fn new(docset: TDocSet, score: Score) -> ConstScorer<TDocSet> {
103        ConstScorer { docset, score }
104    }
105}
106
107impl<TDocSet: DocSet> From<TDocSet> for ConstScorer<TDocSet> {
108    fn from(docset: TDocSet) -> Self {
109        ConstScorer::new(docset, 1.0)
110    }
111}
112
113impl<TDocSet: DocSet> DocSet for ConstScorer<TDocSet> {
114    fn advance(&mut self) -> DocId {
115        self.docset.advance()
116    }
117
118    fn seek(&mut self, target: DocId) -> DocId {
119        self.docset.seek(target)
120    }
121
122    fn fill_buffer(&mut self, buffer: &mut [DocId; COLLECT_BLOCK_BUFFER_LEN]) -> usize {
123        self.docset.fill_buffer(buffer)
124    }
125
126    fn doc(&self) -> DocId {
127        self.docset.doc()
128    }
129
130    fn size_hint(&self) -> u32 {
131        self.docset.size_hint()
132    }
133}
134
135impl<TDocSet: DocSet + 'static> Scorer for ConstScorer<TDocSet> {
136    fn score(&mut self) -> Score {
137        self.score
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::ConstScoreQuery;
144    use crate::query::{AllQuery, Query};
145    use crate::schema::Schema;
146    use crate::{DocAddress, Index, IndexWriter, TantivyDocument};
147
148    #[test]
149    fn test_const_score_query_explain() -> crate::Result<()> {
150        let schema = Schema::builder().build();
151        let index = Index::create_in_ram(schema);
152        let mut index_writer: IndexWriter = index.writer_for_tests()?;
153        index_writer.add_document(TantivyDocument::new())?;
154        index_writer.commit()?;
155        let reader = index.reader()?;
156        let searcher = reader.searcher();
157        let query = ConstScoreQuery::new(Box::new(AllQuery), 0.42);
158        let explanation = query.explain(&searcher, DocAddress::new(0, 0u32)).unwrap();
159        assert_eq!(
160            explanation.to_pretty_json(),
161            r#"{
162  "value": 0.42,
163  "description": "Const",
164  "details": [
165    {
166      "value": 1.0,
167      "description": "AllQuery"
168    }
169  ]
170}"#
171        );
172        Ok(())
173    }
174}
lb_tantivy/query/const_score_query.rs

lb_tantivy/query/
const_score_query.rs
