lb_tantivy/query/range_query/
range_query.rs

1use std::io;
2use std::ops::Bound;
3
4use common::bounds::{map_bound, BoundsRange};
5use common::BitSet;
6
7use super::range_query_fastfield::FastFieldRangeWeight;
8use crate::index::SegmentReader;
9use crate::query::explanation::does_not_match;
10use crate::query::range_query::is_type_valid_for_fastfield_range_query;
11use crate::query::{BitSetDocSet, ConstScorer, EnableScoring, Explanation, Query, Scorer, Weight};
12use crate::schema::{Field, IndexRecordOption, Term, Type};
13use crate::termdict::{TermDictionary, TermStreamer};
14use crate::{DocId, Score};
15
16/// `RangeQuery` matches all documents that have at least one term within a defined range.
17///
18/// Matched document will all get a constant `Score` of one.
19///
20/// # Implementation
21///
22/// ## Default
23/// The default implementation collects all documents _upfront_ into a `BitSet`.
24/// This is done by iterating over the terms within the range and loading all docs for each
25/// `TermInfo` from the inverted index (posting list) and put them into a `BitSet`.
26/// Depending on the number of terms matched, this is a potentially expensive operation.
27///
28/// ## IP fast field
29/// For IP fast fields a custom variant is used, by scanning the fast field. Unlike the default
30/// variant we can walk in a lazy fashion over it, since the fastfield is implicit orderered by
31/// DocId.
32///
33///
34/// # Example
35///
36/// ```rust
37/// use tantivy::collector::Count;
38/// use tantivy::query::RangeQuery;
39/// use tantivy::Term;
40/// use tantivy::schema::{Schema, INDEXED};
41/// use tantivy::{doc, Index, IndexWriter};
42/// use std::ops::Bound;
43/// # fn test() -> tantivy::Result<()> {
44/// let mut schema_builder = Schema::builder();
45/// let year_field = schema_builder.add_u64_field("year", INDEXED);
46/// let schema = schema_builder.build();
47///
48/// let index = Index::create_in_ram(schema);
49/// let mut index_writer: IndexWriter = index.writer_with_num_threads(1, 20_000_000)?;
50/// for year in 1950u64..2017u64 {
51///     let num_docs_within_year = 10 + (year - 1950) * (year - 1950);
52///     for _ in 0..num_docs_within_year {
53///       index_writer.add_document(doc!(year_field => year))?;
54///     }
55/// }
56/// index_writer.commit()?;
57///
58/// let reader = index.reader()?;
59/// let searcher = reader.searcher();
60/// let docs_in_the_sixties = RangeQuery::new(
61///     Bound::Included(Term::from_field_u64(year_field, 1960)),
62///     Bound::Excluded(Term::from_field_u64(year_field, 1970)),
63/// );
64/// let num_60s_books = searcher.search(&docs_in_the_sixties, &Count)?;
65/// assert_eq!(num_60s_books, 2285);
66/// Ok(())
67/// # }
68/// # assert!(test().is_ok());
69/// ```
70#[derive(Clone, Debug)]
71pub struct RangeQuery {
72    bounds: BoundsRange<Term>,
73}
74
75impl RangeQuery {
76    /// Creates a new `RangeQuery` from bounded start and end terms.
77    ///
78    /// If the value type is not correct, something may go terribly wrong when
79    /// the `Weight` object is created.
80    pub fn new(lower_bound: Bound<Term>, upper_bound: Bound<Term>) -> RangeQuery {
81        RangeQuery {
82            bounds: BoundsRange::new(lower_bound, upper_bound),
83        }
84    }
85
86    /// Field to search over
87    pub fn field(&self) -> Field {
88        self.get_term().field()
89    }
90
91    /// The value type of the field
92    pub fn value_type(&self) -> Type {
93        self.get_term().typ()
94    }
95
96    pub(crate) fn get_term(&self) -> &Term {
97        self.bounds
98            .get_inner()
99            .expect("At least one bound must be set")
100    }
101}
102
103impl Query for RangeQuery {
104    fn weight(&self, enable_scoring: EnableScoring<'_>) -> crate::Result<Box<dyn Weight>> {
105        let schema = enable_scoring.schema();
106        let field_type = schema.get_field_entry(self.field()).field_type();
107
108        if field_type.is_fast() && is_type_valid_for_fastfield_range_query(self.value_type()) {
109            Ok(Box::new(FastFieldRangeWeight::new(self.bounds.clone())))
110        } else {
111            if field_type.is_json() {
112                return Err(crate::TantivyError::InvalidArgument(
113                    "RangeQuery on JSON is only supported for fast fields currently".to_string(),
114                ));
115            }
116            Ok(Box::new(InvertedIndexRangeWeight::new(
117                self.field(),
118                &self.bounds.lower_bound,
119                &self.bounds.upper_bound,
120                None,
121            )))
122        }
123    }
124}
125
126#[derive(Clone, Debug)]
127/// `InvertedIndexRangeQuery` is the same as [RangeQuery] but only uses the inverted index
128pub struct InvertedIndexRangeQuery {
129    bounds: BoundsRange<Term>,
130    limit: Option<u64>,
131}
132impl InvertedIndexRangeQuery {
133    /// Create new `InvertedIndexRangeQuery`
134    pub fn new(lower_bound: Bound<Term>, upper_bound: Bound<Term>) -> InvertedIndexRangeQuery {
135        InvertedIndexRangeQuery {
136            bounds: BoundsRange::new(lower_bound, upper_bound),
137            limit: None,
138        }
139    }
140    /// Limit the number of term the `RangeQuery` will go through.
141    ///
142    /// This does not limit the number of matching document, only the number of
143    /// different terms that get matched.
144    pub fn limit(&mut self, limit: u64) {
145        self.limit = Some(limit);
146    }
147}
148
149impl Query for InvertedIndexRangeQuery {
150    fn weight(&self, _enable_scoring: EnableScoring<'_>) -> crate::Result<Box<dyn Weight>> {
151        let field = self
152            .bounds
153            .get_inner()
154            .expect("At least one bound must be set")
155            .field();
156
157        Ok(Box::new(InvertedIndexRangeWeight::new(
158            field,
159            &self.bounds.lower_bound,
160            &self.bounds.upper_bound,
161            self.limit,
162        )))
163    }
164}
165
166/// Range weight on the inverted index
167pub struct InvertedIndexRangeWeight {
168    field: Field,
169    lower_bound: Bound<Vec<u8>>,
170    upper_bound: Bound<Vec<u8>>,
171    limit: Option<u64>,
172}
173
174impl InvertedIndexRangeWeight {
175    /// Creates a new RangeWeight
176    ///
177    /// Note: The limit is only enabled with the quickwit feature flag.
178    pub fn new(
179        field: Field,
180        lower_bound: &Bound<Term>,
181        upper_bound: &Bound<Term>,
182        limit: Option<u64>,
183    ) -> Self {
184        let verify_and_unwrap_term = |val: &Term| val.serialized_value_bytes().to_owned();
185        Self {
186            field,
187            lower_bound: map_bound(lower_bound, verify_and_unwrap_term),
188            upper_bound: map_bound(upper_bound, verify_and_unwrap_term),
189            limit,
190        }
191    }
192
193    fn term_range<'a>(&self, term_dict: &'a TermDictionary) -> io::Result<TermStreamer<'a>> {
194        use std::ops::Bound::*;
195        let mut term_stream_builder = term_dict.range();
196        term_stream_builder = match self.lower_bound {
197            Included(ref term_val) => term_stream_builder.ge(term_val),
198            Excluded(ref term_val) => term_stream_builder.gt(term_val),
199            Unbounded => term_stream_builder,
200        };
201        term_stream_builder = match self.upper_bound {
202            Included(ref term_val) => term_stream_builder.le(term_val),
203            Excluded(ref term_val) => term_stream_builder.lt(term_val),
204            Unbounded => term_stream_builder,
205        };
206        #[cfg(feature = "quickwit")]
207        if let Some(limit) = self.limit {
208            term_stream_builder = term_stream_builder.limit(limit);
209        }
210        term_stream_builder.into_stream()
211    }
212}
213
214impl Weight for InvertedIndexRangeWeight {
215    fn scorer(&self, reader: &SegmentReader, boost: Score) -> crate::Result<Box<dyn Scorer>> {
216        let max_doc = reader.max_doc();
217        let mut doc_bitset = BitSet::with_max_value(max_doc);
218
219        let inverted_index = reader.inverted_index(self.field)?;
220        let term_dict = inverted_index.terms();
221        let mut term_range = self.term_range(term_dict)?;
222        let mut processed_count = 0;
223        while term_range.advance() {
224            if let Some(limit) = self.limit {
225                if limit <= processed_count {
226                    break;
227                }
228            }
229            processed_count += 1;
230            let term_info = term_range.value();
231            let mut block_segment_postings = inverted_index
232                .read_block_postings_from_terminfo(term_info, IndexRecordOption::Basic)?;
233            loop {
234                let docs = block_segment_postings.docs();
235                if docs.is_empty() {
236                    break;
237                }
238                for &doc in block_segment_postings.docs() {
239                    doc_bitset.insert(doc);
240                }
241                block_segment_postings.advance();
242            }
243        }
244        let doc_bitset = BitSetDocSet::from(doc_bitset);
245        Ok(Box::new(ConstScorer::new(doc_bitset, boost)))
246    }
247
248    fn explain(&self, reader: &SegmentReader, doc: DocId) -> crate::Result<Explanation> {
249        let mut scorer = self.scorer(reader, 1.0)?;
250        if scorer.seek(doc) != doc {
251            return Err(does_not_match(doc));
252        }
253        Ok(Explanation::new("RangeQuery", 1.0))
254    }
255}
256
257#[cfg(test)]
258mod tests {
259
260    use std::net::IpAddr;
261    use std::ops::Bound;
262    use std::str::FromStr;
263
264    use rand::seq::SliceRandom;
265
266    use super::RangeQuery;
267    use crate::collector::{Count, TopDocs};
268    use crate::indexer::NoMergePolicy;
269    use crate::query::range_query::range_query::InvertedIndexRangeQuery;
270    use crate::query::QueryParser;
271    use crate::schema::{
272        Field, IntoIpv6Addr, Schema, TantivyDocument, FAST, INDEXED, STORED, TEXT,
273    };
274    use crate::{Index, IndexWriter, Term};
275
276    #[test]
277    fn test_range_query_simple() -> crate::Result<()> {
278        let mut schema_builder = Schema::builder();
279        let year_field = schema_builder.add_u64_field("year", INDEXED);
280        let schema = schema_builder.build();
281
282        let index = Index::create_in_ram(schema);
283        {
284            let mut index_writer = index.writer_for_tests()?;
285            for year in 1950u64..2017u64 {
286                let num_docs_within_year = 10 + (year - 1950) * (year - 1950);
287                for _ in 0..num_docs_within_year {
288                    index_writer.add_document(doc!(year_field => year))?;
289                }
290            }
291            index_writer.commit()?;
292        }
293        let reader = index.reader()?;
294        let searcher = reader.searcher();
295
296        let docs_in_the_sixties = InvertedIndexRangeQuery::new(
297            Bound::Included(Term::from_field_u64(year_field, 1960)),
298            Bound::Excluded(Term::from_field_u64(year_field, 1970)),
299        );
300
301        // ... or `1960..=1969` if inclusive range is enabled.
302        let count = searcher.search(&docs_in_the_sixties, &Count)?;
303        assert_eq!(count, 2285);
304        Ok(())
305    }
306
307    #[test]
308    fn test_range_query_with_limit() -> crate::Result<()> {
309        let mut schema_builder = Schema::builder();
310        let year_field = schema_builder.add_u64_field("year", INDEXED);
311        let schema = schema_builder.build();
312
313        let index = Index::create_in_ram(schema);
314        {
315            let mut index_writer = index.writer_for_tests()?;
316            for year in 1950u64..2017u64 {
317                if year == 1963 {
318                    continue;
319                }
320                let num_docs_within_year = 10 + (year - 1950) * (year - 1950);
321                for _ in 0..num_docs_within_year {
322                    index_writer.add_document(doc!(year_field => year))?;
323                }
324            }
325            index_writer.commit()?;
326        }
327        let reader = index.reader()?;
328        let searcher = reader.searcher();
329
330        let mut docs_in_the_sixties = InvertedIndexRangeQuery::new(
331            Bound::Included(Term::from_field_u64(year_field, 1960)),
332            Bound::Excluded(Term::from_field_u64(year_field, 1970)),
333        );
334        docs_in_the_sixties.limit(5);
335
336        // due to the limit and no docs in 1963, it's really only 1960..=1965
337        let count = searcher.search(&docs_in_the_sixties, &Count)?;
338        assert_eq!(count, 836);
339        Ok(())
340    }
341
342    #[test]
343    fn test_range_query() -> crate::Result<()> {
344        let int_field: Field;
345        let schema = {
346            let mut schema_builder = Schema::builder();
347            int_field = schema_builder.add_i64_field("intfield", INDEXED);
348            schema_builder.build()
349        };
350
351        let index = Index::create_in_ram(schema);
352        {
353            let mut index_writer = index.writer_with_num_threads(1, 60_000_000)?;
354            index_writer.set_merge_policy(Box::new(NoMergePolicy));
355
356            for i in 1..100 {
357                let mut doc = TantivyDocument::new();
358                for j in 1..100 {
359                    if i % j == 0 {
360                        doc.add_i64(int_field, j as i64);
361                    }
362                }
363                index_writer.add_document(doc)?;
364                if i == 10 {
365                    index_writer.commit()?;
366                }
367            }
368
369            index_writer.commit()?;
370        }
371        let reader = index.reader().unwrap();
372        let searcher = reader.searcher();
373        assert_eq!(searcher.segment_readers().len(), 2);
374        let count_multiples =
375            |range_query: RangeQuery| searcher.search(&range_query, &Count).unwrap();
376
377        assert_eq!(
378            count_multiples(RangeQuery::new(
379                Bound::Included(Term::from_field_i64(int_field, 10)),
380                Bound::Excluded(Term::from_field_i64(int_field, 11)),
381            )),
382            9
383        );
384        assert_eq!(
385            count_multiples(RangeQuery::new(
386                Bound::Included(Term::from_field_i64(int_field, 10)),
387                Bound::Included(Term::from_field_i64(int_field, 11)),
388            )),
389            18
390        );
391        assert_eq!(
392            count_multiples(RangeQuery::new(
393                Bound::Excluded(Term::from_field_i64(int_field, 9)),
394                Bound::Included(Term::from_field_i64(int_field, 10)),
395            )),
396            9
397        );
398        assert_eq!(
399            count_multiples(RangeQuery::new(
400                Bound::Included(Term::from_field_i64(int_field, 9)),
401                Bound::Unbounded
402            )),
403            91
404        );
405        Ok(())
406    }
407
408    #[test]
409    fn test_range_float() -> crate::Result<()> {
410        let float_field: Field;
411        let schema = {
412            let mut schema_builder = Schema::builder();
413            float_field = schema_builder.add_f64_field("floatfield", INDEXED);
414            schema_builder.build()
415        };
416
417        let index = Index::create_in_ram(schema);
418        {
419            let mut index_writer = index.writer_with_num_threads(1, 60_000_000).unwrap();
420            let mut docs = Vec::new();
421            for i in 1..100 {
422                let mut doc = TantivyDocument::new();
423                for j in 1..100 {
424                    if i % j == 0 {
425                        doc.add_f64(float_field, j as f64);
426                    }
427                }
428                docs.push(doc);
429            }
430
431            docs.shuffle(&mut rand::thread_rng());
432            let mut docs_it = docs.into_iter();
433            for doc in (&mut docs_it).take(50) {
434                index_writer.add_document(doc)?;
435            }
436            index_writer.commit()?;
437            for doc in docs_it {
438                index_writer.add_document(doc)?;
439            }
440            index_writer.commit()?;
441        }
442        let reader = index.reader()?;
443        let searcher = reader.searcher();
444        assert_eq!(searcher.segment_readers().len(), 2);
445        let count_multiples =
446            |range_query: RangeQuery| searcher.search(&range_query, &Count).unwrap();
447
448        assert_eq!(
449            count_multiples(RangeQuery::new(
450                Bound::Included(Term::from_field_f64(float_field, 10.0)),
451                Bound::Excluded(Term::from_field_f64(float_field, 11.0)),
452            )),
453            9
454        );
455        assert_eq!(
456            count_multiples(RangeQuery::new(
457                Bound::Included(Term::from_field_f64(float_field, 10.0)),
458                Bound::Included(Term::from_field_f64(float_field, 11.0)),
459            )),
460            18
461        );
462        assert_eq!(
463            count_multiples(RangeQuery::new(
464                Bound::Excluded(Term::from_field_f64(float_field, 9.0)),
465                Bound::Included(Term::from_field_f64(float_field, 10.0)),
466            )),
467            9
468        );
469        assert_eq!(
470            count_multiples(RangeQuery::new(
471                Bound::Included(Term::from_field_f64(float_field, 9.0)),
472                Bound::Unbounded
473            )),
474            91
475        );
476        Ok(())
477    }
478
479    #[test]
480    fn test_bug_reproduce_range_query() -> crate::Result<()> {
481        let mut schema_builder = Schema::builder();
482        schema_builder.add_text_field("title", TEXT);
483        schema_builder.add_i64_field("year", INDEXED);
484        let schema = schema_builder.build();
485        let index = Index::create_in_ram(schema.clone());
486        let mut index_writer = index.writer_for_tests()?;
487        let title = schema.get_field("title").unwrap();
488        let year = schema.get_field("year").unwrap();
489        index_writer.add_document(doc!(
490          title => "hemoglobin blood",
491          year => 1990_i64
492        ))?;
493        index_writer.commit()?;
494        let reader = index.reader()?;
495        let searcher = reader.searcher();
496        let query_parser = QueryParser::for_index(&index, vec![title]);
497        let query = query_parser.parse_query("hemoglobin AND year:[1970 TO 1990]")?;
498        let top_docs = searcher.search(&query, &TopDocs::with_limit(10))?;
499        assert_eq!(top_docs.len(), 1);
500        Ok(())
501    }
502
503    #[test]
504    fn search_ip_range_test_posting_list() {
505        search_ip_range_test_opt(false);
506    }
507
508    #[test]
509    fn search_ip_range_test() {
510        search_ip_range_test_opt(true);
511    }
512
513    fn search_ip_range_test_opt(with_fast_field: bool) {
514        let mut schema_builder = Schema::builder();
515        let ip_field = if with_fast_field {
516            schema_builder.add_ip_addr_field("ip", INDEXED | STORED | FAST)
517        } else {
518            schema_builder.add_ip_addr_field("ip", INDEXED | STORED)
519        };
520        let text_field = schema_builder.add_text_field("text", TEXT | STORED);
521        let schema = schema_builder.build();
522        let index = Index::create_in_ram(schema);
523        let ip_addr_1 = IpAddr::from_str("127.0.0.10").unwrap().into_ipv6_addr();
524        let ip_addr_2 = IpAddr::from_str("127.0.0.20").unwrap().into_ipv6_addr();
525
526        {
527            let mut index_writer: IndexWriter = index.writer_for_tests().unwrap();
528            for _ in 0..1_000 {
529                index_writer
530                    .add_document(doc!(
531                        ip_field => ip_addr_1,
532                        text_field => "BLUBBER"
533                    ))
534                    .unwrap();
535            }
536            for _ in 0..1_000 {
537                index_writer
538                    .add_document(doc!(
539                        ip_field => ip_addr_2,
540                        text_field => "BLOBBER"
541                    ))
542                    .unwrap();
543            }
544            index_writer.commit().unwrap();
545        }
546        let reader = index.reader().unwrap();
547        let searcher = reader.searcher();
548        assert_eq!(searcher.segment_readers().len(), 1);
549
550        let get_num_hits = |query| {
551            let (_top_docs, count) = searcher
552                .search(&query, &(TopDocs::with_limit(10), Count))
553                .unwrap();
554            count
555        };
556        let query_from_text = |text: &str| {
557            QueryParser::for_index(&index, vec![])
558                .parse_query(text)
559                .unwrap()
560        };
561
562        // Inclusive range
563        assert_eq!(
564            get_num_hits(query_from_text("ip:[127.0.0.1 TO 127.0.0.20]")),
565            2000
566        );
567
568        assert_eq!(
569            get_num_hits(query_from_text("ip:[127.0.0.10 TO 127.0.0.20]")),
570            2000
571        );
572
573        assert_eq!(
574            get_num_hits(query_from_text("ip:[127.0.0.11 TO 127.0.0.20]")),
575            1000
576        );
577
578        assert_eq!(
579            get_num_hits(query_from_text("ip:[127.0.0.11 TO 127.0.0.19]")),
580            0
581        );
582
583        assert_eq!(get_num_hits(query_from_text("ip:[127.0.0.11 TO *]")), 1000);
584        assert_eq!(get_num_hits(query_from_text("ip:[127.0.0.21 TO *]")), 0);
585        assert_eq!(get_num_hits(query_from_text("ip:[* TO 127.0.0.9]")), 0);
586        assert_eq!(get_num_hits(query_from_text("ip:[* TO 127.0.0.10]")), 1000);
587
588        // Exclusive range
589        assert_eq!(
590            get_num_hits(query_from_text("ip:{127.0.0.1 TO 127.0.0.20}")),
591            1000
592        );
593
594        assert_eq!(
595            get_num_hits(query_from_text("ip:{127.0.0.1 TO 127.0.0.21}")),
596            2000
597        );
598
599        assert_eq!(
600            get_num_hits(query_from_text("ip:{127.0.0.10 TO 127.0.0.20}")),
601            0
602        );
603
604        assert_eq!(
605            get_num_hits(query_from_text("ip:{127.0.0.11 TO 127.0.0.20}")),
606            0
607        );
608
609        assert_eq!(
610            get_num_hits(query_from_text("ip:{127.0.0.11 TO 127.0.0.19}")),
611            0
612        );
613
614        assert_eq!(get_num_hits(query_from_text("ip:{127.0.0.11 TO *}")), 1000);
615        assert_eq!(get_num_hits(query_from_text("ip:{127.0.0.10 TO *}")), 1000);
616        assert_eq!(get_num_hits(query_from_text("ip:{127.0.0.21 TO *}")), 0);
617        assert_eq!(get_num_hits(query_from_text("ip:{127.0.0.20 TO *}")), 0);
618        assert_eq!(get_num_hits(query_from_text("ip:{127.0.0.19 TO *}")), 1000);
619        assert_eq!(get_num_hits(query_from_text("ip:{* TO 127.0.0.9}")), 0);
620        assert_eq!(get_num_hits(query_from_text("ip:{* TO 127.0.0.10}")), 0);
621        assert_eq!(get_num_hits(query_from_text("ip:{* TO 127.0.0.11}")), 1000);
622
623        // Inclusive/Exclusive range
624        assert_eq!(
625            get_num_hits(query_from_text("ip:[127.0.0.1 TO 127.0.0.20}")),
626            1000
627        );
628
629        assert_eq!(
630            get_num_hits(query_from_text("ip:{127.0.0.1 TO 127.0.0.20]")),
631            2000
632        );
633
634        // Intersection
635        assert_eq!(
636            get_num_hits(query_from_text(
637                "text:BLUBBER AND ip:[127.0.0.10 TO 127.0.0.10]"
638            )),
639            1000
640        );
641
642        assert_eq!(
643            get_num_hits(query_from_text(
644                "text:BLOBBER AND ip:[127.0.0.10 TO 127.0.0.10]"
645            )),
646            0
647        );
648
649        assert_eq!(
650            get_num_hits(query_from_text(
651                "text:BLOBBER AND ip:[127.0.0.20 TO 127.0.0.20]"
652            )),
653            1000
654        );
655
656        assert_eq!(
657            get_num_hits(query_from_text(
658                "text:BLUBBER AND ip:[127.0.0.20 TO 127.0.0.20]"
659            )),
660            0
661        );
662    }
663}
lb_tantivy/query/range_query/range_query.rs

lb_tantivy/query/range_query/
range_query.rs
