Skip to main content

hermes_core/query/
term.rs

1//! Term query - matches documents containing a specific term
2
3use std::sync::Arc;
4
5use crate::dsl::Field;
6use crate::segment::SegmentReader;
7use crate::structures::BlockPostingList;
8use crate::structures::TERMINATED;
9use crate::{DocId, Score};
10
11use super::{CountFuture, EmptyScorer, GlobalStats, Query, Scorer, ScorerFuture, TermQueryInfo};
12
13/// Term query - matches documents containing a specific term
14#[derive(Clone)]
15pub struct TermQuery {
16    pub field: Field,
17    pub term: Vec<u8>,
18    /// Optional global statistics for cross-segment IDF
19    global_stats: Option<Arc<GlobalStats>>,
20}
21
22impl std::fmt::Debug for TermQuery {
23    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
24        f.debug_struct("TermQuery")
25            .field("field", &self.field)
26            .field("term", &String::from_utf8_lossy(&self.term))
27            .field("has_global_stats", &self.global_stats.is_some())
28            .finish()
29    }
30}
31
32impl std::fmt::Display for TermQuery {
33    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34        write!(
35            f,
36            "Term({}:\"{}\")",
37            self.field.0,
38            String::from_utf8_lossy(&self.term)
39        )
40    }
41}
42
43impl TermQuery {
44    pub fn new(field: Field, term: impl Into<Vec<u8>>) -> Self {
45        Self {
46            field,
47            term: term.into(),
48            global_stats: None,
49        }
50    }
51
52    pub fn text(field: Field, text: &str) -> Self {
53        Self {
54            field,
55            term: text.to_lowercase().into_bytes(),
56            global_stats: None,
57        }
58    }
59
60    /// Create with global statistics for cross-segment IDF
61    pub fn with_global_stats(field: Field, text: &str, stats: Arc<GlobalStats>) -> Self {
62        Self {
63            field,
64            term: text.to_lowercase().into_bytes(),
65            global_stats: Some(stats),
66        }
67    }
68
69    /// Set global statistics for cross-segment IDF
70    pub fn set_global_stats(&mut self, stats: Arc<GlobalStats>) {
71        self.global_stats = Some(stats);
72    }
73}
74
75/// Compute (idf, avg_field_len) from a posting list, using global stats when available.
76fn compute_term_idf(
77    posting_list: &BlockPostingList,
78    field: Field,
79    reader: &SegmentReader,
80    global_stats: Option<&Arc<GlobalStats>>,
81    term: &[u8],
82) -> (f32, f32) {
83    if let Some(stats) = global_stats {
84        let term_str = String::from_utf8_lossy(term);
85        let global_idf = stats.text_idf(field, &term_str);
86        if global_idf > 0.0 {
87            return (global_idf, stats.avg_field_len(field));
88        }
89    }
90    let num_docs = reader.num_docs() as f32;
91    let doc_freq = posting_list.doc_count() as f32;
92    (
93        super::bm25_idf(doc_freq, num_docs),
94        reader.avg_field_len(field),
95    )
96}
97
98// ── Unified term scorer macro ────────────────────────────────────────────
99//
100// Parameterised on:
101//   $get_postings_fn – get_postings | get_postings_sync
102//   $get_positions_fn – get_positions | get_positions_sync
103//   $($aw)*          – .await  (present for async, absent for sync)
104macro_rules! term_plan {
105    ($field:expr, $term:expr, $global_stats:expr, $reader:expr,
106     $get_postings_fn:ident, $get_positions_fn:ident
107     $(, $aw:tt)*) => {{
108        let field: Field = $field;
109        let term: &[u8] = $term;
110        let global_stats: Option<&Arc<GlobalStats>> = $global_stats;
111        let reader: &SegmentReader = $reader;
112
113        // Non-indexed fields → fast-field-only path
114        let is_indexed = reader.schema().get_field_entry(field).is_none_or(|e| e.indexed);
115        if !is_indexed {
116            let term_str = String::from_utf8_lossy(term);
117            if let Some(scorer) = FastFieldTextScorer::try_new(reader, field, &term_str) {
118                return Ok(Box::new(scorer) as Box<dyn Scorer + '_>);
119            }
120            return Ok(Box::new(EmptyScorer) as Box<dyn Scorer + '_>);
121        }
122
123        let postings = reader.$get_postings_fn(field, term) $(. $aw)* ?;
124
125        match postings {
126            Some(posting_list) => {
127                let (idf, avg_field_len) =
128                    compute_term_idf(&posting_list, field, reader, global_stats, term);
129
130                let positions = reader.$get_positions_fn(field, term)
131                    $(. $aw)* .ok().flatten();
132
133                let mut scorer = TermScorer::new(posting_list, idf, avg_field_len, 1.0);
134                if let Some(pos) = positions {
135                    scorer = scorer.with_positions(field.0, pos);
136                }
137                Ok(Box::new(scorer) as Box<dyn Scorer + '_>)
138            }
139            None => {
140                let term_str = String::from_utf8_lossy(term);
141                if let Some(scorer) = FastFieldTextScorer::try_new(reader, field, &term_str) {
142                    Ok(Box::new(scorer) as Box<dyn Scorer + '_>)
143                } else {
144                    Ok(Box::new(EmptyScorer) as Box<dyn Scorer + '_>)
145                }
146            }
147        }
148    }};
149}
150
151impl Query for TermQuery {
152    fn scorer<'a>(&self, reader: &'a SegmentReader, _limit: usize) -> ScorerFuture<'a> {
153        let field = self.field;
154        let term = self.term.clone();
155        let global_stats = self.global_stats.clone();
156        Box::pin(async move {
157            term_plan!(
158                field,
159                &term,
160                global_stats.as_ref(),
161                reader,
162                get_postings,
163                get_positions,
164                await
165            )
166        })
167    }
168
169    fn count_estimate<'a>(&self, reader: &'a SegmentReader) -> CountFuture<'a> {
170        let field = self.field;
171        let term = self.term.clone();
172        Box::pin(async move {
173            match reader.get_postings(field, &term).await? {
174                Some(list) => Ok(list.doc_count()),
175                None => Ok(0),
176            }
177        })
178    }
179
180    #[cfg(feature = "sync")]
181    fn scorer_sync<'a>(
182        &self,
183        reader: &'a SegmentReader,
184        _limit: usize,
185    ) -> crate::Result<Box<dyn Scorer + 'a>> {
186        term_plan!(
187            self.field,
188            &self.term,
189            self.global_stats.as_ref(),
190            reader,
191            get_postings_sync,
192            get_positions_sync
193        )
194    }
195
196    fn as_doc_predicate<'a>(&self, reader: &'a SegmentReader) -> Option<super::DocPredicate<'a>> {
197        let fast_field = reader.fast_field(self.field.0)?;
198        let term_str = String::from_utf8_lossy(&self.term);
199        let target_ordinal = fast_field.text_ordinal(&term_str)?;
200        Some(Box::new(move |doc_id: DocId| -> bool {
201            fast_field.get_u64(doc_id) == target_ordinal
202        }))
203    }
204
205    fn decompose(&self) -> super::QueryDecomposition {
206        super::QueryDecomposition::TextTerm(TermQueryInfo {
207            field: self.field,
208            term: self.term.clone(),
209        })
210    }
211}
212
213struct TermScorer {
214    iterator: crate::structures::BlockPostingIterator<'static>,
215    idf: f32,
216    /// Average field length for this field
217    avg_field_len: f32,
218    /// Field boost/weight for BM25F
219    field_boost: f32,
220    /// Field ID for position reporting
221    field_id: u32,
222    /// Position posting list (if positions are enabled)
223    positions: Option<crate::structures::PositionPostingList>,
224}
225
226impl TermScorer {
227    pub fn new(
228        posting_list: BlockPostingList,
229        idf: f32,
230        avg_field_len: f32,
231        field_boost: f32,
232    ) -> Self {
233        Self {
234            iterator: posting_list.into_iterator(),
235            idf,
236            avg_field_len,
237            field_boost,
238            field_id: 0,
239            positions: None,
240        }
241    }
242
243    pub fn with_positions(
244        mut self,
245        field_id: u32,
246        positions: crate::structures::PositionPostingList,
247    ) -> Self {
248        self.field_id = field_id;
249        self.positions = Some(positions);
250        self
251    }
252}
253
254impl super::docset::DocSet for TermScorer {
255    fn doc(&self) -> DocId {
256        self.iterator.doc()
257    }
258
259    fn advance(&mut self) -> DocId {
260        self.iterator.advance()
261    }
262
263    fn seek(&mut self, target: DocId) -> DocId {
264        self.iterator.seek(target)
265    }
266
267    fn size_hint(&self) -> u32 {
268        0
269    }
270}
271
272// ── Fast field text equality scorer ──────────────────────────────────────
273
274/// Scorer that scans a text fast field for exact string equality.
275/// Used as fallback when a TermQuery targets a fast-only text field (no inverted index).
276/// Returns score 1.0 for matching docs (filter-style, like RangeScorer).
277struct FastFieldTextScorer<'a> {
278    fast_field: &'a crate::structures::fast_field::FastFieldReader,
279    target_ordinal: u64,
280    current: u32,
281    num_docs: u32,
282}
283
284impl<'a> FastFieldTextScorer<'a> {
285    fn try_new(reader: &'a SegmentReader, field: Field, text: &str) -> Option<Self> {
286        let fast_field = reader.fast_field(field.0)?;
287        let target_ordinal = fast_field.text_ordinal(text)?;
288        let num_docs = reader.num_docs();
289        let mut scorer = Self {
290            fast_field,
291            target_ordinal,
292            current: 0,
293            num_docs,
294        };
295        // Position on first matching doc
296        if num_docs > 0 && fast_field.get_u64(0) != target_ordinal {
297            scorer.scan_forward();
298        }
299        Some(scorer)
300    }
301
302    fn scan_forward(&mut self) {
303        loop {
304            self.current += 1;
305            if self.current >= self.num_docs {
306                self.current = self.num_docs;
307                return;
308            }
309            if self.fast_field.get_u64(self.current) == self.target_ordinal {
310                return;
311            }
312        }
313    }
314}
315
316impl super::docset::DocSet for FastFieldTextScorer<'_> {
317    fn doc(&self) -> DocId {
318        if self.current >= self.num_docs {
319            TERMINATED
320        } else {
321            self.current
322        }
323    }
324
325    fn advance(&mut self) -> DocId {
326        self.scan_forward();
327        self.doc()
328    }
329
330    fn seek(&mut self, target: DocId) -> DocId {
331        if target > self.current {
332            self.current = target;
333            if self.current < self.num_docs
334                && self.fast_field.get_u64(self.current) != self.target_ordinal
335            {
336                self.scan_forward();
337            }
338        }
339        self.doc()
340    }
341
342    fn size_hint(&self) -> u32 {
343        0
344    }
345}
346
347impl Scorer for FastFieldTextScorer<'_> {
348    fn score(&self) -> Score {
349        1.0
350    }
351}
352
353impl Scorer for TermScorer {
354    fn score(&self) -> Score {
355        let tf = self.iterator.term_freq() as f32;
356        // Note: Using tf as doc_len proxy since we don't store per-doc field lengths.
357        // This is a common approximation - longer docs tend to have higher TF.
358        super::bm25f_score(tf, self.idf, tf, self.avg_field_len, self.field_boost)
359    }
360
361    fn matched_positions(&self) -> Option<super::MatchedPositions> {
362        let positions = self.positions.as_ref()?;
363        let doc_id = self.iterator.doc();
364        let pos = positions.get_positions(doc_id)?;
365        let score = self.score();
366        // Each position contributes equally to the term score
367        let per_position_score = if pos.is_empty() {
368            0.0
369        } else {
370            score / pos.len() as f32
371        };
372        let scored_positions: Vec<super::ScoredPosition> = pos
373            .iter()
374            .map(|&p| super::ScoredPosition::new(p, per_position_score))
375            .collect();
376        Some(vec![(self.field_id, scored_positions)])
377    }
378}