hermes_core/query/
prefix.rs

1//! Prefix query — matches all documents containing any term that starts with a
2//! given prefix. Materializes the union of matching posting lists into a sorted
3//! doc ID set bounded by the segment's document count. Score is always 1.0
4//! (filter-style, like `RangeQuery`).
5
6use std::sync::Arc;
7
8use crate::dsl::Field;
9use crate::segment::SegmentReader;
10use crate::structures::{BlockPostingList, TERMINATED};
11use crate::{DocId, Score};
12
13use super::docset::{DocSet, SortedVecDocSet};
14use super::traits::{CountFuture, EmptyScorer, Query, Scorer, ScorerFuture};
15
16/// Prefix query — matches documents containing any term starting with `prefix`.
17#[derive(Debug, Clone)]
18pub struct PrefixQuery {
19    pub field: Field,
20    pub prefix: Vec<u8>,
21}
22
23impl std::fmt::Display for PrefixQuery {
24    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25        write!(
26            f,
27            "Prefix({}:\"{}*\")",
28            self.field.0,
29            String::from_utf8_lossy(&self.prefix)
30        )
31    }
32}
33
34impl PrefixQuery {
35    /// Create from raw bytes.
36    pub fn new(field: Field, prefix: impl Into<Vec<u8>>) -> Self {
37        Self {
38            field,
39            prefix: prefix.into(),
40        }
41    }
42
43    /// Create from text — lowercased to match default tokenization.
44    pub fn text(field: Field, text: &str) -> Self {
45        Self {
46            field,
47            prefix: text.to_lowercase().into_bytes(),
48        }
49    }
50}
51
52impl Query for PrefixQuery {
53    fn scorer<'a>(&self, reader: &'a SegmentReader, _limit: usize) -> ScorerFuture<'a> {
54        let field = self.field;
55        let prefix = self.prefix.clone();
56        Box::pin(async move {
57            let postings = reader.get_prefix_postings(field, &prefix).await?;
58            if postings.is_empty() {
59                return Ok(Box::new(EmptyScorer) as Box<dyn Scorer>);
60            }
61            let docs = materialize_union(&postings, reader.num_docs());
62            if docs.is_empty() {
63                return Ok(Box::new(EmptyScorer) as Box<dyn Scorer>);
64            }
65            Ok(Box::new(PrefixScorer::new(docs)) as Box<dyn Scorer>)
66        })
67    }
68
69    #[cfg(feature = "sync")]
70    fn scorer_sync<'a>(
71        &self,
72        reader: &'a SegmentReader,
73        _limit: usize,
74    ) -> crate::Result<Box<dyn Scorer + 'a>> {
75        let postings = reader.get_prefix_postings_sync(self.field, &self.prefix)?;
76        if postings.is_empty() {
77            return Ok(Box::new(EmptyScorer) as Box<dyn Scorer>);
78        }
79        let docs = materialize_union(&postings, reader.num_docs());
80        if docs.is_empty() {
81            return Ok(Box::new(EmptyScorer) as Box<dyn Scorer>);
82        }
83        Ok(Box::new(PrefixScorer::new(docs)) as Box<dyn Scorer>)
84    }
85
86    fn count_estimate<'a>(&self, reader: &'a SegmentReader) -> CountFuture<'a> {
87        let field = self.field;
88        let prefix = self.prefix.clone();
89        Box::pin(async move {
90            let postings = reader.get_prefix_postings(field, &prefix).await?;
91            Ok(postings
92                .iter()
93                .fold(0u32, |sum, posting| sum.saturating_add(posting.doc_count()))
94                .min(reader.num_docs()))
95        })
96    }
97
98    fn is_filter(&self) -> bool {
99        true
100    }
101
102    #[cfg(feature = "sync")]
103    fn as_doc_predicate<'a>(&self, reader: &'a SegmentReader) -> Option<super::DocPredicate<'a>> {
104        let bitset = self.as_doc_bitset(reader)?;
105        Some(Box::new(move |doc_id: DocId| bitset.contains(doc_id)))
106    }
107
108    #[cfg(feature = "sync")]
109    fn as_doc_bitset(&self, reader: &SegmentReader) -> Option<super::DocBitset> {
110        let postings = reader
111            .get_prefix_postings_sync(self.field, &self.prefix)
112            .ok()?;
113        let mut bitset = super::DocBitset::new(reader.num_docs());
114        for posting in &postings {
115            let mut iter = posting.iterator();
116            loop {
117                let d = iter.doc();
118                if d == TERMINATED {
119                    break;
120                }
121                bitset.set(d);
122                iter.advance();
123            }
124        }
125        Some(bitset)
126    }
127}
128
129// ── PrefixScorer ────────────────────────────────────────────────────────
130
131/// Scorer backed by a pre-materialized sorted doc ID set.
132struct PrefixScorer {
133    inner: SortedVecDocSet,
134}
135
136impl PrefixScorer {
137    fn new(docs: Vec<u32>) -> Self {
138        Self {
139            inner: SortedVecDocSet::new(Arc::new(docs)),
140        }
141    }
142}
143
144impl DocSet for PrefixScorer {
145    #[inline]
146    fn doc(&self) -> DocId {
147        self.inner.doc()
148    }
149
150    #[inline]
151    fn advance(&mut self) -> DocId {
152        self.inner.advance()
153    }
154
155    fn seek(&mut self, target: DocId) -> DocId {
156        self.inner.seek(target)
157    }
158
159    fn size_hint(&self) -> u32 {
160        self.inner.size_hint()
161    }
162}
163
164impl Scorer for PrefixScorer {
165    fn score(&self) -> Score {
166        1.0
167    }
168}
169
170// ── Helpers ─────────────────────────────────────────────────────────────
171
172/// Materialize a posting union using the smaller of two bounded scratch forms.
173/// Narrow prefixes append/sort doc IDs; broad, overlapping prefixes use a
174/// segment-sized bitset so duplicate postings cannot multiply memory.
175fn materialize_union(postings: &[BlockPostingList], num_docs: u32) -> Vec<u32> {
176    let posting_count = postings.iter().fold(0usize, |sum, posting| {
177        sum.saturating_add(posting.doc_count() as usize)
178    });
179    let posting_bytes = posting_count.saturating_mul(std::mem::size_of::<u32>());
180    let bitset_bytes = (num_docs as usize)
181        .div_ceil(64)
182        .saturating_mul(std::mem::size_of::<u64>());
183
184    if posting_bytes <= bitset_bytes {
185        let mut docs = Vec::with_capacity(posting_count);
186        for posting in postings {
187            let mut iter = posting.iterator();
188            loop {
189                let d = iter.doc();
190                if d == TERMINATED {
191                    break;
192                }
193                docs.push(d);
194                iter.advance();
195            }
196        }
197        docs.sort_unstable();
198        docs.dedup();
199        return docs;
200    }
201
202    let mut bitset = super::DocBitset::new(num_docs);
203    for posting in postings {
204        let mut iter = posting.iterator();
205        loop {
206            let d = iter.doc();
207            if d == TERMINATED {
208                break;
209            }
210            bitset.set(d);
211            iter.advance();
212        }
213    }
214
215    let mut docs = Vec::with_capacity(bitset.count() as usize);
216    for (word_idx, &word) in bitset.bits.iter().enumerate() {
217        let mut remaining = word;
218        while remaining != 0 {
219            let bit = remaining.trailing_zeros() as usize;
220            docs.push((word_idx * 64 + bit) as u32);
221            remaining &= remaining - 1;
222        }
223    }
224    docs
225}
226
227#[cfg(test)]
228mod tests {
229    use super::*;
230
231    #[test]
232    fn test_materialize_union_empty() {
233        let docs = materialize_union(&[], 0);
234        assert!(docs.is_empty());
235    }
236
237    #[test]
238    fn test_materialize_union_deduplicates() {
239        let mut left = crate::structures::PostingList::new();
240        left.push(1, 1);
241        left.push(5, 1);
242        left.push(9, 1);
243        let mut right = crate::structures::PostingList::new();
244        right.push(2, 1);
245        right.push(5, 1);
246        right.push(10, 1);
247        let postings = vec![
248            BlockPostingList::from_posting_list(&left).unwrap(),
249            BlockPostingList::from_posting_list(&right).unwrap(),
250        ];
251
252        assert_eq!(materialize_union(&postings, 11), vec![1, 2, 5, 9, 10]);
253        // A huge segment with a narrow prefix takes the posting-vector path;
254        // it must not allocate a num_docs-sized bitset.
255        assert_eq!(
256            materialize_union(&postings[..1], 1_000_000_000),
257            vec![1, 5, 9]
258        );
259    }
260
261    #[test]
262    fn test_prefix_scorer_basic() {
263        let mut scorer = PrefixScorer::new(vec![1, 5, 10, 20]);
264        assert_eq!(scorer.doc(), 1);
265        assert_eq!(scorer.score(), 1.0);
266        assert_eq!(scorer.advance(), 5);
267        assert_eq!(scorer.seek(10), 10);
268        assert_eq!(scorer.advance(), 20);
269        assert_eq!(scorer.advance(), TERMINATED);
270    }
271
272    #[test]
273    fn test_prefix_scorer_seek_past() {
274        let mut scorer = PrefixScorer::new(vec![1, 5, 10, 20]);
275        assert_eq!(scorer.seek(7), 10);
276        assert_eq!(scorer.seek(100), TERMINATED);
277    }
278
279    #[test]
280    fn test_prefix_query_display() {
281        let q = PrefixQuery::text(Field(0), "abc");
282        assert_eq!(format!("{}", q), "Prefix(0:\"abc*\")");
283    }
284
285    #[test]
286    fn test_prefix_query_is_filter() {
287        let q = PrefixQuery::text(Field(0), "test");
288        assert!(q.is_filter());
289    }
290}
hermes_core/query/prefix.rs

hermes_core/query/
prefix.rs
