hermes_core/query/
proximity.rs

1//! Sequential-dependence proximity rescoring for BM25 text queries.
2//!
3//! A second stage over the top candidates of a text MaxScore pass: for every
4//! pair of adjacent query terms the document's positions are scanned for
5//! ordered windows (the second term directly after the first) and unordered
6//! windows (both within `window` positions). Each count is saturated with the
7//! field's BM25 parameters and length, weighted by the pair's mean idf, and
8//! added to the BM25 score scaled by `weight` (Metzler & Croft, "A Markov
9//! random field model for term dependencies", SIGIR 2005; ordered windows
10//! count full weight, unordered windows half).
11//!
12//! The stage is approximate on purpose: the MaxScore pass over-fetches
13//! `PROXIMITY_OVER_FETCH` times the requested limit, the bonus is added to
14//! those candidates only, and cross-segment threshold seeding is disabled
15//! for the pass because the bonus lifts scores above the BM25 floor.
16
17#[cfg(feature = "sync")]
18use crate::DocId;
19use crate::dsl::Field;
20use crate::segment::SegmentReader;
21
22/// Proximity rescoring of a text query (`MatchQuery.proximity_weight`).
23#[derive(Debug, Clone, Copy, PartialEq)]
24pub struct ProximityConfig {
25    /// Multiplier of the window bonus; 0 disables the stage.
26    pub weight: f32,
27    /// Maximum distance of an unordered window.
28    pub window: u32,
29}
30
31impl ProximityConfig {
32    /// Default unordered window (Metzler & Croft used 8).
33    pub const DEFAULT_WINDOW: u32 = 8;
34
35    pub fn new(weight: f32, window: u32) -> Self {
36        Self {
37            weight,
38            window: if window == 0 {
39                Self::DEFAULT_WINDOW
40            } else {
41                window
42            },
43        }
44    }
45
46    pub fn is_active(&self) -> bool {
47        self.weight > 0.0
48    }
49}
50
51/// Candidates fetched per requested hit before rescoring.
52pub(crate) const PROXIMITY_OVER_FETCH: usize = 4;
53
54/// Ordered (`b` right after `a`) and unordered (within `window`) windows
55/// between two ascending position lists.
56#[cfg_attr(not(feature = "sync"), allow(dead_code))]
57pub(crate) fn count_windows(a: &[u32], b: &[u32], window: u32) -> (u32, u32) {
58    let mut ordered = 0u32;
59    let mut unordered = 0u32;
60    let (mut lo, mut hi, mut adj) = (0usize, 0usize, 0usize);
61    for &pa in a {
62        let low = pa.saturating_sub(window);
63        let high = pa.saturating_add(window);
64        while lo < b.len() && b[lo] < low {
65            lo += 1;
66        }
67        if hi < lo {
68            hi = lo;
69        }
70        while hi < b.len() && b[hi] <= high {
71            hi += 1;
72        }
73        unordered += (hi - lo) as u32;
74        while adj < b.len() && b[adj] < pa + 1 {
75            adj += 1;
76        }
77        if adj < b.len() && b[adj] == pa + 1 {
78            ordered += 1;
79        }
80    }
81    (ordered, unordered)
82}
83
84/// Add the proximity bonus to `hits` (documents or chunk ids of `field`).
85/// `terms` are the query terms in query order with their idf.
86#[cfg(feature = "sync")]
87#[allow(clippy::too_many_arguments)]
88pub(crate) fn rescore_sync(
89    reader: &SegmentReader,
90    field: Field,
91    terms: &[(Vec<u8>, f32)],
92    params: super::Bm25Params,
93    lengths: Option<super::LengthSource<'_>>,
94    avg_len: f32,
95    config: ProximityConfig,
96    hits: &mut [super::ScoredDoc],
97) -> crate::Result<()> {
98    use crate::structures::TERMINATED;
99
100    if !config.is_active() || terms.len() < 2 || hits.is_empty() {
101        return Ok(());
102    }
103    let mut cursors = Vec::with_capacity(terms.len());
104    for (term, _) in terms {
105        let list = reader.get_postings_sync(field, term)?;
106        let positions = reader.get_positions_sync(field, term)?;
107        cursors.push(match (list, positions) {
108            (Some(list), Some(positions)) => Some((list.into_iterator(), positions.into_cursor())),
109            _ => None,
110        });
111    }
112    let mut order: Vec<usize> = (0..hits.len()).collect();
113    order.sort_unstable_by_key(|&i| hits[i].doc_id);
114    let mut bufs: Vec<Vec<u32>> = vec![Vec::new(); terms.len()];
115    let avg = avg_len.max(1.0);
116    for i in order {
117        let doc: DocId = hits[i].doc_id;
118        for (t, cursor) in cursors.iter_mut().enumerate() {
119            bufs[t].clear();
120            if let Some((it, positions)) = cursor
121                && it.doc() != TERMINATED
122                && it.seek(doc) == doc
123            {
124                positions.read_into(doc, it.position_cursor(), it.term_freq(), &mut bufs[t]);
125            }
126        }
127        let len = lengths
128            .map(|source| source.length(doc) as f32)
129            .filter(|len| *len > 0.0)
130            .unwrap_or(avg);
131        let mut bonus = 0.0f32;
132        for pair in 0..terms.len() - 1 {
133            let (a, b) = (&bufs[pair], &bufs[pair + 1]);
134            if a.is_empty() || b.is_empty() {
135                continue;
136            }
137            let (ordered, unordered) = count_windows(a, b, config.window);
138            let idf = (terms[pair].1 + terms[pair + 1].1) * 0.5;
139            if ordered > 0 {
140                bonus += params.score(ordered as f32, idf, len, avg);
141            }
142            if unordered > 0 {
143                bonus += 0.5 * params.score(unordered as f32, idf, len, avg);
144            }
145        }
146        hits[i].score += config.weight * bonus;
147    }
148    Ok(())
149}
150
151/// Without synchronous file handles the stage cannot read positions per
152/// candidate; the BM25 ranking is returned unchanged.
153#[cfg(not(feature = "sync"))]
154#[allow(clippy::too_many_arguments)]
155pub(crate) fn rescore_sync(
156    _reader: &SegmentReader,
157    _field: Field,
158    _terms: &[(Vec<u8>, f32)],
159    _params: super::Bm25Params,
160    _lengths: Option<super::LengthSource<'_>>,
161    _avg_len: f32,
162    _config: ProximityConfig,
163    _hits: &mut [super::ScoredDoc],
164) -> crate::Result<()> {
165    log::debug!("proximity rescoring needs the `sync` feature; returning BM25 order");
166    Ok(())
167}
168
169#[cfg(test)]
170mod tests {
171    use super::*;
172
173    #[test]
174    fn windows_count_ordered_and_unordered_pairs() {
175        // a at 0, 10; b at 1, 3, 12: a@0 sees b@1 (ordered) and b@3 within 8,
176        // a@10 sees b@3 and b@12 within 8 (b@12 is not adjacent).
177        assert_eq!(count_windows(&[0, 10], &[1, 3, 12], 8), (1, 4));
178        assert_eq!(count_windows(&[0, 10], &[1, 3, 12], 1), (1, 1));
179        assert_eq!(count_windows(&[5], &[4], 2), (0, 1));
180        assert_eq!(count_windows(&[5], &[], 2), (0, 0));
181        // Window 0 counts co-located positions only; ordered pairs are the
182        // three adjacent ones.
183        assert_eq!(count_windows(&[0, 1, 2], &[1, 2, 3], 0), (3, 2));
184    }
185
186    #[test]
187    fn config_defaults_window() {
188        assert_eq!(ProximityConfig::new(1.0, 0).window, 8);
189        assert!(!ProximityConfig::new(0.0, 4).is_active());
190    }
191}
hermes_core/query/proximity.rs

hermes_core/query/
proximity.rs
