1use crate::core::{DocId, NO_MORE_DOCS, Result, ScoreMode, Scorer, TwoPhaseIterator};
9
10use crate::query::{BoundQuery, Query, ScorerSupplier};
11use crate::search::searcher::Searcher;
12use crate::segment::reader::SegmentReader;
13
14pub struct RangeQuery {
16 pub field: String,
17 pub gte: Option<f64>,
18 pub gt: Option<f64>,
19 pub lte: Option<f64>,
20 pub lt: Option<f64>,
21}
22
23impl Query for RangeQuery {
24 fn bind(&self, _searcher: &Searcher, _score_mode: ScoreMode) -> Result<Box<dyn BoundQuery>> {
25 Ok(Box::new(BoundRangeQuery {
26 field: self.field.clone(),
27 gte: self.gte,
28 gt: self.gt,
29 lte: self.lte,
30 lt: self.lt,
31 }))
32 }
33}
34
35struct BoundRangeQuery {
36 field: String,
37 gte: Option<f64>,
38 gt: Option<f64>,
39 lte: Option<f64>,
40 lt: Option<f64>,
41}
42
43impl BoundRangeQuery {
44 fn matches(&self, value: f64) -> bool {
46 if let Some(gte) = self.gte {
47 if value < gte {
48 return false;
49 }
50 }
51 if let Some(gt) = self.gt {
52 if value <= gt {
53 return false;
54 }
55 }
56 if let Some(lte) = self.lte {
57 if value > lte {
58 return false;
59 }
60 }
61 if let Some(lt) = self.lt {
62 if value >= lt {
63 return false;
64 }
65 }
66 true
67 }
68}
69
70impl BoundQuery for BoundRangeQuery {
71 fn scorer_supplier(&self, reader: &SegmentReader) -> Result<Option<Box<dyn ScorerSupplier>>> {
72 let field_id = match reader
73 .header()
74 .fields
75 .iter()
76 .find(|f| f.field_name == self.field)
77 .map(|f| f.field_id)
78 {
79 Some(id) => id,
80 None => return Ok(None),
81 };
82
83 let col = match reader.column(field_id) {
84 Some(c) => c,
85 None => return Ok(None),
86 };
87
88 if let Some(stats) = col.stats() {
90 let segment_min = stats.min;
91 let segment_max = stats.max;
92
93 if let Some(gte) = self.gte {
95 if gte > segment_max {
96 return Ok(None);
97 }
98 }
99 if let Some(gt) = self.gt {
100 if gt >= segment_max {
101 return Ok(None);
102 }
103 }
104 if let Some(lte) = self.lte {
106 if lte < segment_min {
107 return Ok(None);
108 }
109 }
110 if let Some(lt) = self.lt {
111 if lt <= segment_min {
112 return Ok(None);
113 }
114 }
115 }
116
117 let doc_count = col.doc_count();
120 let mut matching_docs: Vec<u32> = Vec::new();
121 for i in 0..doc_count {
122 if let Some(v) = col.numeric_value(i) {
123 if self.matches(v) {
124 matching_docs.push(i);
125 }
126 }
127 }
128
129 if matching_docs.is_empty() {
130 return Ok(None);
131 }
132
133 Ok(Some(Box::new(RangeScorerSupplier { matching_docs })))
134 }
135}
136
137struct RangeScorerSupplier {
138 matching_docs: Vec<u32>,
139}
140
141impl ScorerSupplier for RangeScorerSupplier {
142 fn cost(&self) -> u64 {
143 self.matching_docs.len() as u64
144 }
145
146 fn scorer(self: Box<Self>) -> Result<Box<dyn Scorer>> {
147 Ok(Box::new(RangeScorer {
148 matching_docs: self.matching_docs,
149 pos: 0,
150 }))
151 }
152}
153
154struct RangeScorer {
155 matching_docs: Vec<u32>,
156 pos: usize,
157}
158
159impl Scorer for RangeScorer {
160 fn doc_id(&self) -> DocId {
161 if self.pos < self.matching_docs.len() {
162 DocId::new(self.matching_docs[self.pos])
163 } else {
164 NO_MORE_DOCS
165 }
166 }
167
168 fn next(&mut self) -> DocId {
169 self.pos += 1;
170 self.doc_id()
171 }
172
173 fn advance(&mut self, target: DocId) -> DocId {
174 let target_u32 = target.as_u32();
175 while self.pos < self.matching_docs.len() && self.matching_docs[self.pos] < target_u32 {
176 self.pos += 1;
177 }
178 self.doc_id()
179 }
180
181 fn score(&mut self) -> f32 {
182 1.0 }
184
185 fn two_phase(&mut self) -> Option<&mut dyn TwoPhaseIterator> {
186 None
187 }
188}
189
190#[cfg(test)]
191mod tests {
192 use super::*;
193 use crate::columnar::writer::ColumnValue;
194 use crate::core::SegmentId;
195 use crate::mapping::{FieldType, Mapping};
196 use crate::segment::builder::SegmentBuilder;
197
198 fn build_numeric_segment(values: &[f64]) -> SegmentReader {
199 let schema = Mapping::builder().field("price", FieldType::Float).build();
200 let mut builder = SegmentBuilder::new(SegmentId::new(1), &schema);
201 let field_id = schema.field_id("price").unwrap();
202
203 for &v in values {
204 builder.add_document(&[], b"{}");
205 builder.add_column_value(field_id, ColumnValue::F64(v));
206 }
207
208 SegmentReader::open(builder.build()).unwrap()
209 }
210
211 #[test]
212 fn range_gte_lte() {
213 let reader = build_numeric_segment(&[1.0, 5.0, 10.0, 15.0, 20.0]);
214 let store = crate::search::segment_store::SegmentStore::new(
215 vec![reader],
216 crate::analysis::AnalyzerRegistry::new(),
217 None,
218 None,
219 );
220 let searcher = Searcher::new(&store);
221
222 let query = RangeQuery {
223 field: "price".into(),
224 gte: Some(5.0),
225 gt: None,
226 lte: Some(15.0),
227 lt: None,
228 };
229
230 let results = searcher.search_query(&query, 10, 0).unwrap();
231 assert_eq!(results.total_hits.value, 3); }
233
234 #[test]
235 fn range_gt_lt() {
236 let reader = build_numeric_segment(&[1.0, 5.0, 10.0, 15.0, 20.0]);
237 let store = crate::search::segment_store::SegmentStore::new(
238 vec![reader],
239 crate::analysis::AnalyzerRegistry::new(),
240 None,
241 None,
242 );
243 let searcher = Searcher::new(&store);
244
245 let query = RangeQuery {
246 field: "price".into(),
247 gte: None,
248 gt: Some(5.0),
249 lte: None,
250 lt: Some(15.0),
251 };
252
253 let results = searcher.search_query(&query, 10, 0).unwrap();
254 assert_eq!(results.total_hits.value, 1); }
256
257 #[test]
258 fn range_no_matches() {
259 let reader = build_numeric_segment(&[1.0, 2.0, 3.0]);
260 let store = crate::search::segment_store::SegmentStore::new(
261 vec![reader],
262 crate::analysis::AnalyzerRegistry::new(),
263 None,
264 None,
265 );
266 let searcher = Searcher::new(&store);
267
268 let query = RangeQuery {
269 field: "price".into(),
270 gte: Some(100.0),
271 gt: None,
272 lte: None,
273 lt: None,
274 };
275
276 let results = searcher.search_query(&query, 10, 0).unwrap();
277 assert_eq!(results.total_hits.value, 0);
278 }
279
280 #[test]
281 fn range_all_match() {
282 let reader = build_numeric_segment(&[5.0, 10.0, 15.0]);
283 let store = crate::search::segment_store::SegmentStore::new(
284 vec![reader],
285 crate::analysis::AnalyzerRegistry::new(),
286 None,
287 None,
288 );
289 let searcher = Searcher::new(&store);
290
291 let query = RangeQuery {
292 field: "price".into(),
293 gte: Some(0.0),
294 gt: None,
295 lte: Some(100.0),
296 lt: None,
297 };
298
299 let results = searcher.search_query(&query, 10, 0).unwrap();
300 assert_eq!(results.total_hits.value, 3);
301 }
302
303 #[test]
304 fn range_open_ended() {
305 let reader = build_numeric_segment(&[1.0, 5.0, 10.0, 15.0, 20.0]);
306 let store = crate::search::segment_store::SegmentStore::new(
307 vec![reader],
308 crate::analysis::AnalyzerRegistry::new(),
309 None,
310 None,
311 );
312 let searcher = Searcher::new(&store);
313
314 let query = RangeQuery {
316 field: "price".into(),
317 gte: Some(10.0),
318 gt: None,
319 lte: None,
320 lt: None,
321 };
322
323 let results = searcher.search_query(&query, 10, 0).unwrap();
324 assert_eq!(results.total_hits.value, 3); }
326
327 #[test]
328 fn range_missing_field() {
329 let reader = build_numeric_segment(&[1.0, 5.0]);
330 let store = crate::search::segment_store::SegmentStore::new(
331 vec![reader],
332 crate::analysis::AnalyzerRegistry::new(),
333 None,
334 None,
335 );
336 let searcher = Searcher::new(&store);
337
338 let query = RangeQuery {
339 field: "nonexistent".into(),
340 gte: Some(0.0),
341 gt: None,
342 lte: None,
343 lt: None,
344 };
345
346 let results = searcher.search_query(&query, 10, 0).unwrap();
347 assert_eq!(results.total_hits.value, 0);
348 }
349}