lb_tantivy/query/
boost_query.rs

1use std::fmt;
2
3use crate::docset::COLLECT_BLOCK_BUFFER_LEN;
4use crate::fastfield::AliveBitSet;
5use crate::query::{EnableScoring, Explanation, Query, Scorer, Weight};
6use crate::{DocId, DocSet, Score, SegmentReader, Term};
7
8/// `BoostQuery` is a wrapper over a query used to boost its score.
9///
10/// The document set matched by the `BoostQuery` is strictly the same as the underlying query.
11/// The score of each document, is the score of the underlying query multiplied by the `boost`
12/// factor.
13pub struct BoostQuery {
14    query: Box<dyn Query>,
15    boost: Score,
16}
17
18impl BoostQuery {
19    /// Builds a boost query.
20    pub fn new(query: Box<dyn Query>, boost: Score) -> BoostQuery {
21        BoostQuery { query, boost }
22    }
23}
24
25impl Clone for BoostQuery {
26    fn clone(&self) -> Self {
27        BoostQuery {
28            query: self.query.box_clone(),
29            boost: self.boost,
30        }
31    }
32}
33
34impl fmt::Debug for BoostQuery {
35    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
36        write!(f, "Boost(query={:?}, boost={})", self.query, self.boost)
37    }
38}
39
40impl Query for BoostQuery {
41    fn weight(&self, enable_scoring: EnableScoring<'_>) -> crate::Result<Box<dyn Weight>> {
42        let weight_without_boost = self.query.weight(enable_scoring)?;
43        let boosted_weight = if enable_scoring.is_scoring_enabled() {
44            Box::new(BoostWeight::new(weight_without_boost, self.boost))
45        } else {
46            weight_without_boost
47        };
48        Ok(boosted_weight)
49    }
50
51    fn query_terms<'a>(&'a self, visitor: &mut dyn FnMut(&'a Term, bool)) {
52        self.query.query_terms(visitor)
53    }
54}
55
56/// Weight associated to the BoostQuery.
57pub struct BoostWeight {
58    weight: Box<dyn Weight>,
59    boost: Score,
60}
61
62impl BoostWeight {
63    /// Creates a new BoostWeight.
64    pub fn new(weight: Box<dyn Weight>, boost: Score) -> Self {
65        BoostWeight { weight, boost }
66    }
67}
68
69impl Weight for BoostWeight {
70    fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result<Box<dyn Scorer>> {
71        self.weight.scorer(reader, boost * self.boost)
72    }
73
74    fn explain(&self, reader: &SegmentReader, doc: u32) -> crate::Result<Explanation> {
75        let underlying_explanation = self.weight.explain(reader, doc)?;
76        let score = underlying_explanation.value() * self.boost;
77        let mut explanation =
78            Explanation::new_with_string(format!("Boost x{} of ...", self.boost), score);
79        explanation.add_detail(underlying_explanation);
80        Ok(explanation)
81    }
82
83    fn count(&self, reader: &SegmentReader) -> crate::Result<u32> {
84        self.weight.count(reader)
85    }
86}
87
88pub(crate) struct BoostScorer<S: Scorer> {
89    underlying: S,
90    boost: Score,
91}
92
93impl<S: Scorer> BoostScorer<S> {
94    pub fn new(underlying: S, boost: Score) -> BoostScorer<S> {
95        BoostScorer { underlying, boost }
96    }
97}
98
99impl<S: Scorer> DocSet for BoostScorer<S> {
100    fn advance(&mut self) -> DocId {
101        self.underlying.advance()
102    }
103
104    fn seek(&mut self, target: DocId) -> DocId {
105        self.underlying.seek(target)
106    }
107
108    fn fill_buffer(&mut self, buffer: &mut [DocId; COLLECT_BLOCK_BUFFER_LEN]) -> usize {
109        self.underlying.fill_buffer(buffer)
110    }
111
112    fn doc(&self) -> u32 {
113        self.underlying.doc()
114    }
115
116    fn size_hint(&self) -> u32 {
117        self.underlying.size_hint()
118    }
119
120    fn count(&mut self, alive_bitset: &AliveBitSet) -> u32 {
121        self.underlying.count(alive_bitset)
122    }
123
124    fn count_including_deleted(&mut self) -> u32 {
125        self.underlying.count_including_deleted()
126    }
127}
128
129impl<S: Scorer> Scorer for BoostScorer<S> {
130    fn score(&mut self) -> Score {
131        self.underlying.score() * self.boost
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use super::BoostQuery;
138    use crate::query::{AllQuery, Query};
139    use crate::schema::Schema;
140    use crate::{DocAddress, Index, IndexWriter, TantivyDocument};
141
142    #[test]
143    fn test_boost_query_explain() -> crate::Result<()> {
144        let schema = Schema::builder().build();
145        let index = Index::create_in_ram(schema);
146        let mut index_writer: IndexWriter = index.writer_for_tests()?;
147        index_writer.add_document(TantivyDocument::new())?;
148        index_writer.commit()?;
149        let reader = index.reader()?;
150        let searcher = reader.searcher();
151        let query = BoostQuery::new(Box::new(AllQuery), 0.2);
152        let explanation = query.explain(&searcher, DocAddress::new(0, 0u32)).unwrap();
153        assert_eq!(
154            explanation.to_pretty_json(),
155            "{\n  \"value\": 0.2,\n  \"description\": \"Boost x0.2 of ...\",\n  \"details\": [\n    {\n      \"value\": 1.0,\n      \"description\": \"AllQuery\"\n    }\n  ]\n}"
156        );
157        Ok(())
158    }
159}
lb_tantivy/query/boost_query.rs

lb_tantivy/query/
boost_query.rs
