1use 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#[derive(Debug, Clone)]
27pub enum RangeBound {
28 U64 { min: Option<u64>, max: Option<u64> },
30 I64 { min: Option<i64>, max: Option<i64> },
32 F64 { min: Option<f64>, max: Option<f64> },
34}
35
36impl RangeBound {
37 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 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 fn is_i64(&self) -> bool {
71 matches!(self, RangeBound::I64 { .. })
72 }
73
74 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#[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 pub fn u64(field: Field, min: Option<u64>, max: Option<u64>) -> Self {
130 Self::new(field, RangeBound::U64 { min, max })
131 }
132
133 pub fn i64(field: Field, min: Option<i64>, max: Option<i64>) -> Self {
135 Self::new(field, RangeBound::I64 { min, max })
136 }
137
138 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 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 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 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
225struct RangeScorer<'a> {
233 fast_field: &'a crate::structures::fast_field::FastFieldReader,
235 raw_lo: u64,
237 raw_hi: u64,
238 i64_lo: i64,
240 i64_hi: i64,
241 use_i64: bool,
243 current: u32,
245 num_docs: u32,
246}
247
248struct 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 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 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 self.current = target - 1;
334 self.scan_forward();
335 self.doc()
336 }
337
338 fn size_hint(&self) -> u32 {
339 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}