hermes_core/query/
boost.rs

1//! Boost query - multiplies the score of the inner query
2
3use std::sync::Arc;
4
5use crate::segment::SegmentReader;
6use crate::{DocId, Score};
7
8use super::{CountFuture, Query, Scorer, ScorerFuture};
9
10/// Boost query - multiplies the score of the inner query
11#[derive(Clone)]
12pub struct BoostQuery {
13    pub inner: Arc<dyn Query>,
14    pub boost: f32,
15}
16
17impl std::fmt::Debug for BoostQuery {
18    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
19        f.debug_struct("BoostQuery")
20            .field("boost", &self.boost)
21            .finish()
22    }
23}
24
25impl std::fmt::Display for BoostQuery {
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27        write!(f, "{}^{}", self.inner, self.boost)
28    }
29}
30
31impl BoostQuery {
32    pub fn new(query: impl Query + 'static, boost: f32) -> Self {
33        Self {
34            inner: Arc::new(query),
35            boost,
36        }
37    }
38}
39
40impl Query for BoostQuery {
41    fn scorer<'a>(&self, reader: &'a SegmentReader, limit: usize) -> ScorerFuture<'a> {
42        self.scorer_with_options(reader, limit, super::ScorerOptions::with_positions())
43    }
44
45    fn scorer_with_options<'a>(
46        &self,
47        reader: &'a SegmentReader,
48        limit: usize,
49        options: super::ScorerOptions,
50    ) -> ScorerFuture<'a> {
51        let inner = self.inner.clone();
52        let boost = self.boost;
53        Box::pin(async move {
54            if !boost.is_finite() {
55                return Err(crate::Error::Query(
56                    "boost must be a finite number".to_string(),
57                ));
58            }
59            let inner_options = if boost == 1.0 {
60                options
61            } else {
62                options.without_threshold()
63            };
64            let inner_scorer = inner
65                .scorer_with_options(reader, limit, inner_options)
66                .await?;
67            Ok(Box::new(BoostScorer {
68                inner: inner_scorer,
69                boost,
70            }) as Box<dyn Scorer + 'a>)
71        })
72    }
73
74    #[cfg(feature = "sync")]
75    fn scorer_sync<'a>(
76        &self,
77        reader: &'a SegmentReader,
78        limit: usize,
79    ) -> crate::Result<Box<dyn Scorer + 'a>> {
80        self.scorer_sync_with_options(reader, limit, super::ScorerOptions::with_positions())
81    }
82
83    #[cfg(feature = "sync")]
84    fn scorer_sync_with_options<'a>(
85        &self,
86        reader: &'a SegmentReader,
87        limit: usize,
88        options: super::ScorerOptions,
89    ) -> crate::Result<Box<dyn Scorer + 'a>> {
90        if !self.boost.is_finite() {
91            return Err(crate::Error::Query(
92                "boost must be a finite number".to_string(),
93            ));
94        }
95        let inner_options = if self.boost == 1.0 {
96            options
97        } else {
98            options.without_threshold()
99        };
100        let inner_scorer = self
101            .inner
102            .scorer_sync_with_options(reader, limit, inner_options)?;
103        Ok(Box::new(BoostScorer {
104            inner: inner_scorer,
105            boost: self.boost,
106        }) as Box<dyn Scorer + 'a>)
107    }
108
109    fn count_estimate<'a>(&self, reader: &'a SegmentReader) -> CountFuture<'a> {
110        let inner = self.inner.clone();
111        Box::pin(async move { inner.count_estimate(reader).await })
112    }
113
114    fn is_filter(&self) -> bool {
115        self.boost == 1.0 && self.inner.is_filter()
116    }
117
118    fn as_doc_predicate<'a>(&self, reader: &'a SegmentReader) -> Option<super::DocPredicate<'a>> {
119        (self.boost == 1.0)
120            .then(|| self.inner.as_doc_predicate(reader))
121            .flatten()
122    }
123
124    fn text_terms(&self, out: &mut Vec<(crate::dsl::Field, Vec<u8>)>) {
125        self.inner.text_terms(out);
126    }
127
128    fn decompose(&self) -> super::QueryDecomposition {
129        match self.inner.decompose() {
130            // A boosted text term stays on the text MaxScore path: the
131            // weight scales its idf, so scores and bounds scale together.
132            super::QueryDecomposition::TextTerm(mut info) => {
133                info.weight *= self.boost;
134                super::QueryDecomposition::TextTerm(info)
135            }
136            other if self.boost == 1.0 => other,
137            _ => super::QueryDecomposition::Opaque,
138        }
139    }
140}
141
142struct BoostScorer<'a> {
143    inner: Box<dyn Scorer + 'a>,
144    boost: f32,
145}
146
147impl super::docset::DocSet for BoostScorer<'_> {
148    fn doc(&self) -> DocId {
149        self.inner.doc()
150    }
151
152    fn advance(&mut self) -> DocId {
153        self.inner.advance()
154    }
155
156    fn seek(&mut self, target: DocId) -> DocId {
157        self.inner.seek(target)
158    }
159
160    fn size_hint(&self) -> u32 {
161        self.inner.size_hint()
162    }
163}
164
165impl Scorer for BoostScorer<'_> {
166    fn score(&self) -> Score {
167        self.inner.score() * self.boost
168    }
169
170    fn matched_positions(&self) -> Option<super::MatchedPositions> {
171        let mut positions = self.inner.matched_positions()?;
172        for (_, scored_positions) in &mut positions {
173            for position in scored_positions {
174                position.score *= self.boost;
175            }
176        }
177        Some(positions)
178    }
179}
hermes_core/query/boost.rs

hermes_core/query/
boost.rs
