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)),
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 scratch: Vec<u32> = Vec::new();
115    let mut bufs: Vec<Vec<u32>> = vec![Vec::new(); terms.len()];
116    let avg = avg_len.max(1.0);
117    for i in order {
118        let doc: DocId = hits[i].doc_id;
119        for (t, cursor) in cursors.iter_mut().enumerate() {
120            bufs[t].clear();
121            if let Some((it, positions)) = cursor
122                && it.doc() != TERMINATED
123                && it.seek(doc) == doc
124            {
125                positions.positions_into(
126                    doc,
127                    it.position_cursor(),
128                    it.term_freq(),
129                    &mut scratch,
130                    &mut bufs[t],
131                );
132            }
133        }
134        let len = lengths
135            .map(|source| source.length(doc) as f32)
136            .filter(|len| *len > 0.0)
137            .unwrap_or(avg);
138        let mut bonus = 0.0f32;
139        for pair in 0..terms.len() - 1 {
140            let (a, b) = (&bufs[pair], &bufs[pair + 1]);
141            if a.is_empty() || b.is_empty() {
142                continue;
143            }
144            let (ordered, unordered) = count_windows(a, b, config.window);
145            let idf = (terms[pair].1 + terms[pair + 1].1) * 0.5;
146            if ordered > 0 {
147                bonus += params.score(ordered as f32, idf, len, avg);
148            }
149            if unordered > 0 {
150                bonus += 0.5 * params.score(unordered as f32, idf, len, avg);
151            }
152        }
153        hits[i].score += config.weight * bonus;
154    }
155    Ok(())
156}
157
158/// Without synchronous file handles the stage cannot read positions per
159/// candidate; the BM25 ranking is returned unchanged.
160#[cfg(not(feature = "sync"))]
161#[allow(clippy::too_many_arguments)]
162pub(crate) fn rescore_sync(
163    _reader: &SegmentReader,
164    _field: Field,
165    _terms: &[(Vec<u8>, f32)],
166    _params: super::Bm25Params,
167    _lengths: Option<super::LengthSource<'_>>,
168    _avg_len: f32,
169    _config: ProximityConfig,
170    _hits: &mut [super::ScoredDoc],
171) -> crate::Result<()> {
172    log::debug!("proximity rescoring needs the `sync` feature; returning BM25 order");
173    Ok(())
174}
175
176#[cfg(test)]
177mod tests {
178    use super::*;
179
180    #[test]
181    fn windows_count_ordered_and_unordered_pairs() {
182        // a at 0, 10; b at 1, 3, 12: a@0 sees b@1 (ordered) and b@3 within 8,
183        // a@10 sees b@3 and b@12 within 8 (b@12 is not adjacent).
184        assert_eq!(count_windows(&[0, 10], &[1, 3, 12], 8), (1, 4));
185        assert_eq!(count_windows(&[0, 10], &[1, 3, 12], 1), (1, 1));
186        assert_eq!(count_windows(&[5], &[4], 2), (0, 1));
187        assert_eq!(count_windows(&[5], &[], 2), (0, 0));
188        // Window 0 counts co-located positions only; ordered pairs are the
189        // three adjacent ones.
190        assert_eq!(count_windows(&[0, 1, 2], &[1, 2, 3], 0), (3, 2));
191    }
192
193    #[test]
194    fn config_defaults_window() {
195        assert_eq!(ProximityConfig::new(1.0, 0).window, 8);
196        assert!(!ProximityConfig::new(0.0, 4).is_active());
197    }
198}
hermes_core/query/proximity.rs

hermes_core/query/
proximity.rs
