Skip to main content

hermes_core/query/
range.rs

1//! Range query for fast-field numeric filtering.
2//!
3//! `RangeQuery` produces a `RangeScorer` that scans a fast-field column and
4//! yields documents whose value falls within the specified bounds. Score is
5//! always 1.0 — this is a pure filter query.
6//!
7//! Supports u64, i64, and f64 fields. For i64/f64, bounds are encoded to the
8//! same sortable-u64 representation used by fast fields so that a single raw
9//! u64 comparison covers all types.
10//!
11//! When placed in a `BooleanQuery` MUST clause, the `BooleanScorer`'s
12//! seek-based intersection makes this efficient even on large segments.
13
14use crate::dsl::Field;
15use crate::segment::SegmentReader;
16use crate::structures::TERMINATED;
17use crate::structures::fast_field::{FAST_FIELD_MISSING, f64_to_sortable_u64, zigzag_encode};
18use crate::{DocId, Score};
19
20use super::docset::DocSet;
21use super::traits::{CountFuture, Query, Scorer, ScorerFuture};
22
23// ── Typed range bounds ───────────────────────────────────────────────────
24
25/// Inclusive range bounds in the user's type domain.
26#[derive(Debug, Clone)]
27pub enum RangeBound {
28    /// u64 range — stored raw
29    U64 { min: Option<u64>, max: Option<u64> },
30    /// i64 range — will be zigzag-encoded for comparison
31    I64 { min: Option<i64>, max: Option<i64> },
32    /// f64 range — will be sortable-encoded for comparison
33    F64 { min: Option<f64>, max: Option<f64> },
34}
35
36impl RangeBound {
37    /// Compile to raw u64 inclusive bounds suitable for direct fast-field comparison.
38    ///
39    /// Returns `(low, high)` where both are inclusive. Missing bounds become
40    /// 0 / u64::MAX-1 (reserving u64::MAX for FAST_FIELD_MISSING sentinel).
41    fn compile(&self) -> (u64, u64) {
42        match self {
43            RangeBound::U64 { min, max } => {
44                let lo = min.unwrap_or(0);
45                let hi = max.unwrap_or(u64::MAX - 1);
46                (lo, hi)
47            }
48            RangeBound::I64 { min, max } => {
49                // zigzag encoding preserves magnitude, not order.
50                // For correct range comparison on i64, we use sortable encoding
51                // (same as f64 but cast through bits). However, fast fields store
52                // i64 as zigzag. So we must decode per-doc and compare in i64 domain.
53                // We store the raw i64 bounds and handle comparison in the scorer.
54                //
55                // Sentinel: use a special marker to tell the scorer to use i64 path.
56                // We'll handle this in the scorer directly.
57                let lo = min.map(zigzag_encode).unwrap_or(0);
58                let hi = max.map(zigzag_encode).unwrap_or(u64::MAX - 1);
59                (lo, hi)
60            }
61            RangeBound::F64 { min, max } => {
62                let lo = min.map(f64_to_sortable_u64).unwrap_or(0);
63                let hi = max.map(f64_to_sortable_u64).unwrap_or(u64::MAX - 1);
64                (lo, hi)
65            }
66        }
67    }
68
69    /// Whether this bound requires per-doc i64 decoding (zigzag doesn't preserve order).
70    fn is_i64(&self) -> bool {
71        matches!(self, RangeBound::I64 { .. })
72    }
73
74    /// Get the raw i64 bounds for the i64 path.
75    fn i64_bounds(&self) -> (i64, i64) {
76        match self {
77            RangeBound::I64 { min, max } => (min.unwrap_or(i64::MIN), max.unwrap_or(i64::MAX)),
78            _ => (i64::MIN, i64::MAX),
79        }
80    }
81}
82
83// ── RangeQuery ───────────────────────────────────────────────────────────
84
85/// Fast-field range query.
86///
87/// Scans all documents in a segment and yields those whose fast-field value
88/// falls within `[min, max]` (inclusive). Score is always 1.0.
89#[derive(Debug, Clone)]
90pub struct RangeQuery {
91    pub field: Field,
92    pub bound: RangeBound,
93}
94
95impl std::fmt::Display for RangeQuery {
96    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
97        match &self.bound {
98            RangeBound::U64 { min, max } => write!(
99                f,
100                "Range({}:[{} TO {}])",
101                self.field.0,
102                min.map_or("*".to_string(), |v| v.to_string()),
103                max.map_or("*".to_string(), |v| v.to_string()),
104            ),
105            RangeBound::I64 { min, max } => write!(
106                f,
107                "Range({}:[{} TO {}])",
108                self.field.0,
109                min.map_or("*".to_string(), |v| v.to_string()),
110                max.map_or("*".to_string(), |v| v.to_string()),
111            ),
112            RangeBound::F64 { min, max } => write!(
113                f,
114                "Range({}:[{} TO {}])",
115                self.field.0,
116                min.map_or("*".to_string(), |v| v.to_string()),
117                max.map_or("*".to_string(), |v| v.to_string()),
118            ),
119        }
120    }
121}
122
123impl RangeQuery {
124    pub fn new(field: Field, bound: RangeBound) -> Self {
125        Self { field, bound }
126    }
127
128    /// Convenience: u64 range
129    pub fn u64(field: Field, min: Option<u64>, max: Option<u64>) -> Self {
130        Self::new(field, RangeBound::U64 { min, max })
131    }
132
133    /// Convenience: i64 range
134    pub fn i64(field: Field, min: Option<i64>, max: Option<i64>) -> Self {
135        Self::new(field, RangeBound::I64 { min, max })
136    }
137
138    /// Convenience: f64 range
139    pub fn f64(field: Field, min: Option<f64>, max: Option<f64>) -> Self {
140        Self::new(field, RangeBound::F64 { min, max })
141    }
142}
143
144impl Query for RangeQuery {
145    fn scorer<'a>(&self, reader: &'a SegmentReader, _limit: usize) -> ScorerFuture<'a> {
146        let field = self.field;
147        let bound = self.bound.clone();
148        Box::pin(async move {
149            match RangeScorer::new(reader, field, &bound) {
150                Ok(scorer) => Ok(Box::new(scorer) as Box<dyn Scorer>),
151                Err(_) => Ok(Box::new(EmptyRangeScorer) as Box<dyn Scorer>),
152            }
153        })
154    }
155
156    #[cfg(feature = "sync")]
157    fn scorer_sync<'a>(
158        &self,
159        reader: &'a SegmentReader,
160        _limit: usize,
161    ) -> crate::Result<Box<dyn Scorer + 'a>> {
162        match RangeScorer::new(reader, self.field, &self.bound) {
163            Ok(scorer) => Ok(Box::new(scorer) as Box<dyn Scorer + 'a>),
164            Err(_) => Ok(Box::new(EmptyRangeScorer) as Box<dyn Scorer + 'a>),
165        }
166    }
167
168    fn count_estimate<'a>(&self, reader: &'a SegmentReader) -> CountFuture<'a> {
169        let num_docs = reader.num_docs();
170        // Rough estimate: half the segment (we don't know selectivity)
171        Box::pin(async move { Ok(num_docs / 2) })
172    }
173
174    fn is_filter(&self) -> bool {
175        true
176    }
177
178    fn as_doc_predicate<'a>(&self, reader: &'a SegmentReader) -> Option<super::DocPredicate<'a>> {
179        let fast_field = reader.fast_field(self.field.0)?;
180        let (raw_lo, raw_hi) = self.bound.compile();
181        let use_i64 = self.bound.is_i64();
182        let (i64_lo, i64_hi) = self.bound.i64_bounds();
183
184        Some(Box::new(move |doc_id: DocId| -> bool {
185            let raw = fast_field.get_u64(doc_id);
186            if raw == FAST_FIELD_MISSING {
187                return false;
188            }
189            if use_i64 {
190                let val = crate::structures::fast_field::zigzag_decode(raw);
191                val >= i64_lo && val <= i64_hi
192            } else {
193                raw >= raw_lo && raw <= raw_hi
194            }
195        }))
196    }
197
198    fn as_doc_bitset(&self, reader: &SegmentReader) -> Option<super::DocBitset> {
199        // Build bitset from fast-field scan: O(N). Slower than posting-list-based
200        // bitsets but still faster than per-call predicate in BMP (~2ns lookup vs ~30ns).
201        let pred = self.as_doc_predicate(reader)?;
202        Some(super::DocBitset::from_predicate(reader.num_docs(), &*pred))
203    }
204
205    fn bitset_cardinality_estimate(&self, reader: &SegmentReader) -> Option<u64> {
206        // Sampled: probe ~1k evenly spaced docs with the fast-field predicate.
207        // Works for every encoding (incl. i64 zigzag, where min/max
208        // interpolation would mis-order). Rounded up so a rare-but-present
209        // range never estimates to zero.
210        let pred = self.as_doc_predicate(reader)?;
211        let n = reader.num_docs();
212        if n == 0 {
213            return Some(0);
214        }
215        const SAMPLES: u32 = 1024;
216        if n <= SAMPLES {
217            return Some((0..n).filter(|&d| pred(d)).count() as u64);
218        }
219        let step = n / SAMPLES;
220        let hits = (0..SAMPLES).filter(|&i| pred(i * step)).count() as u64;
221        Some(((hits * n as u64) / SAMPLES as u64).max(1))
222    }
223}
224
225// ── RangeScorer ──────────────────────────────────────────────────────────
226
227/// Scorer that scans a fast-field column and yields matching docs.
228///
229/// For u64 and f64 fields, comparison is done in the raw u64 domain (both
230/// use order-preserving encodings). For i64 fields, zigzag encoding does NOT
231/// preserve order, so we decode each value and compare in i64 domain.
232struct RangeScorer<'a> {
233    /// Cached fast-field reader — avoids HashMap lookup per doc in matches()
234    fast_field: &'a crate::structures::fast_field::FastFieldReader,
235    /// For u64/f64: compiled raw bounds. For i64: unused.
236    raw_lo: u64,
237    raw_hi: u64,
238    /// For i64 only: decoded bounds.
239    i64_lo: i64,
240    i64_hi: i64,
241    /// Whether to use i64 comparison path.
242    use_i64: bool,
243    /// Current document position.
244    current: u32,
245    num_docs: u32,
246}
247
248/// Empty scorer returned when the field has no fast-field data.
249struct EmptyRangeScorer;
250
251impl<'a> RangeScorer<'a> {
252    fn new(
253        reader: &'a SegmentReader,
254        field: Field,
255        bound: &RangeBound,
256    ) -> Result<Self, EmptyRangeScorer> {
257        let fast_field = reader.fast_field(field.0).ok_or(EmptyRangeScorer)?;
258        let num_docs = reader.num_docs();
259        let (raw_lo, raw_hi) = bound.compile();
260        let use_i64 = bound.is_i64();
261        let (i64_lo, i64_hi) = bound.i64_bounds();
262
263        let mut scorer = Self {
264            fast_field,
265            raw_lo,
266            raw_hi,
267            i64_lo,
268            i64_hi,
269            use_i64,
270            current: 0,
271            num_docs,
272        };
273
274        // Position on first matching doc
275        if num_docs > 0 && !scorer.matches(0) {
276            scorer.scan_forward();
277        }
278        Ok(scorer)
279    }
280
281    #[inline]
282    fn matches(&self, doc_id: DocId) -> bool {
283        let raw = self.fast_field.get_u64(doc_id);
284        if raw == FAST_FIELD_MISSING {
285            return false;
286        }
287
288        if self.use_i64 {
289            let val = crate::structures::fast_field::zigzag_decode(raw);
290            val >= self.i64_lo && val <= self.i64_hi
291        } else {
292            raw >= self.raw_lo && raw <= self.raw_hi
293        }
294    }
295
296    /// Advance current past non-matching docs.
297    fn scan_forward(&mut self) {
298        loop {
299            self.current += 1;
300            if self.current >= self.num_docs {
301                self.current = self.num_docs;
302                return;
303            }
304            if self.matches(self.current) {
305                return;
306            }
307        }
308    }
309}
310
311impl DocSet for RangeScorer<'_> {
312    fn doc(&self) -> DocId {
313        if self.current >= self.num_docs {
314            TERMINATED
315        } else {
316            self.current
317        }
318    }
319
320    fn advance(&mut self) -> DocId {
321        self.scan_forward();
322        self.doc()
323    }
324
325    fn seek(&mut self, target: DocId) -> DocId {
326        if self.current >= self.num_docs {
327            return TERMINATED;
328        }
329        if target <= self.current {
330            return self.current;
331        }
332        // Position just before target so scan_forward starts at target
333        self.current = target - 1;
334        self.scan_forward();
335        self.doc()
336    }
337
338    fn size_hint(&self) -> u32 {
339        // Upper bound: remaining docs
340        self.num_docs.saturating_sub(self.current)
341    }
342}
343
344impl Scorer for RangeScorer<'_> {
345    fn score(&self) -> Score {
346        1.0
347    }
348}
349
350impl DocSet for EmptyRangeScorer {
351    fn doc(&self) -> DocId {
352        TERMINATED
353    }
354    fn advance(&mut self) -> DocId {
355        TERMINATED
356    }
357    fn seek(&mut self, _target: DocId) -> DocId {
358        TERMINATED
359    }
360    fn size_hint(&self) -> u32 {
361        0
362    }
363}
364
365impl Scorer for EmptyRangeScorer {
366    fn score(&self) -> Score {
367        0.0
368    }
369}
370
371#[cfg(test)]
372mod tests {
373    use super::*;
374
375    #[test]
376    fn test_range_bound_u64_compile() {
377        let b = RangeBound::U64 {
378            min: Some(10),
379            max: Some(100),
380        };
381        let (lo, hi) = b.compile();
382        assert_eq!(lo, 10);
383        assert_eq!(hi, 100);
384    }
385
386    #[test]
387    fn test_range_bound_f64_compile_preserves_order() {
388        let b1 = RangeBound::F64 {
389            min: Some(-1.0),
390            max: Some(1.0),
391        };
392        let (lo1, hi1) = b1.compile();
393        assert!(lo1 < hi1);
394
395        let b2 = RangeBound::F64 {
396            min: Some(0.0),
397            max: Some(100.0),
398        };
399        let (lo2, hi2) = b2.compile();
400        assert!(lo2 < hi2);
401    }
402
403    #[test]
404    fn test_range_bound_open_bounds() {
405        let b = RangeBound::U64 {
406            min: None,
407            max: None,
408        };
409        let (lo, hi) = b.compile();
410        assert_eq!(lo, 0);
411        assert_eq!(hi, u64::MAX - 1);
412    }
413
414    #[test]
415    fn test_range_query_constructors() {
416        let q = RangeQuery::u64(Field(0), Some(10), Some(100));
417        assert_eq!(q.field, Field(0));
418        assert!(matches!(
419            q.bound,
420            RangeBound::U64 {
421                min: Some(10),
422                max: Some(100)
423            }
424        ));
425
426        let q = RangeQuery::i64(Field(1), Some(-50), Some(50));
427        assert!(matches!(
428            q.bound,
429            RangeBound::I64 {
430                min: Some(-50),
431                max: Some(50)
432            }
433        ));
434
435        let q = RangeQuery::f64(Field(2), Some(0.5), Some(9.5));
436        assert!(matches!(q.bound, RangeBound::F64 { .. }));
437    }
438}