enable range query on fast field for u64 compatible types (#1762)

PSeitz · web-flow · commit 7c6cc818ae26 · 2023-01-10T04:08:26.000+01:00
* enable range query on fast field for u64 compatible types

* rename, update benches
diff --git a/fastfield_codecs/benches/bench.rs b/fastfield_codecs/benches/bench.rs
@@ -4,7 +4,7 @@ extern crate test;
 
 #[cfg(test)]
 mod tests {
-    use std::iter;
+    use std::ops::RangeInclusive;
     use std::sync::Arc;
 
     use common::OwnedBytes;
@@ -71,27 +71,24 @@ mod tests {
         });
     }
 
-    fn get_exp_data() -> Vec<u64> {
+    const FIFTY_PERCENT_RANGE: RangeInclusive<u64> = 1..=50;
+    const SINGLE_ITEM: u64 = 90;
+    const SINGLE_ITEM_RANGE: RangeInclusive<u64> = 90..=90;
+    const ONE_PERCENT_ITEM_RANGE: RangeInclusive<u64> = 49..=49;
+    fn get_data_50percent_item() -> Vec<u128> {
+        let mut rng = StdRng::from_seed([1u8; 32]);
+
         let mut data = vec![];
-        for i in 0..100 {
-            let num = i * i;
-            data.extend(iter::repeat(i as u64).take(num));
+        for _ in 0..300_000 {
+            let val = rng.gen_range(1..=100);
+            data.push(val);
         }
-        data.shuffle(&mut StdRng::from_seed([1u8; 32]));
+        data.push(SINGLE_ITEM);
 
-        // lengt = 328350
+        data.shuffle(&mut rng);
+        let data = data.iter().map(|el| *el as u128).collect::<Vec<_>>();
         data
     }
-
-    fn get_data_50percent_item() -> (u128, u128, Vec<u128>) {
-        let mut permutation = get_exp_data();
-        let major_item = 20;
-        let minor_item = 10;
-        permutation.extend(iter::repeat(major_item).take(permutation.len()));
-        permutation.shuffle(&mut StdRng::from_seed([1u8; 32]));
-        let permutation = permutation.iter().map(|el| *el as u128).collect::<Vec<_>>();
-        (major_item as u128, minor_item as u128, permutation)
-    }
     fn get_u128_column_random() -> Arc<dyn Column<u128>> {
         let permutation = generate_random();
         let permutation = permutation.iter().map(|el| *el as u128).collect::<Vec<_>>();
@@ -106,15 +103,82 @@ mod tests {
         open_u128::<u128>(out).unwrap()
     }
 
+    // U64 RANGE START
+    #[bench]
+    fn bench_intfastfield_getrange_u64_50percent_hit(b: &mut Bencher) {
+        let data = get_data_50percent_item();
+        let data = data.iter().map(|el| *el as u64).collect::<Vec<_>>();
+        let column: Arc<dyn Column<u64>> = serialize_and_load(&data);
+
+        b.iter(|| {
+            let mut positions = Vec::new();
+            column.get_docids_for_value_range(
+                FIFTY_PERCENT_RANGE,
+                0..data.len() as u32,
+                &mut positions,
+            );
+            positions
+        });
+    }
+
+    #[bench]
+    fn bench_intfastfield_getrange_u64_1percent_hit(b: &mut Bencher) {
+        let data = get_data_50percent_item();
+        let data = data.iter().map(|el| *el as u64).collect::<Vec<_>>();
+        let column: Arc<dyn Column<u64>> = serialize_and_load(&data);
+
+        b.iter(|| {
+            let mut positions = Vec::new();
+            column.get_docids_for_value_range(
+                ONE_PERCENT_ITEM_RANGE,
+                0..data.len() as u32,
+                &mut positions,
+            );
+            positions
+        });
+    }
+
+    #[bench]
+    fn bench_intfastfield_getrange_u64_single_hit(b: &mut Bencher) {
+        let data = get_data_50percent_item();
+        let data = data.iter().map(|el| *el as u64).collect::<Vec<_>>();
+        let column: Arc<dyn Column<u64>> = serialize_and_load(&data);
+
+        b.iter(|| {
+            let mut positions = Vec::new();
+            column.get_docids_for_value_range(
+                SINGLE_ITEM_RANGE,
+                0..data.len() as u32,
+                &mut positions,
+            );
+            positions
+        });
+    }
+
+    #[bench]
+    fn bench_intfastfield_getrange_u64_hit_all(b: &mut Bencher) {
+        let data = get_data_50percent_item();
+        let data = data.iter().map(|el| *el as u64).collect::<Vec<_>>();
+        let column: Arc<dyn Column<u64>> = serialize_and_load(&data);
+
+        b.iter(|| {
+            let mut positions = Vec::new();
+            column.get_docids_for_value_range(0..=u64::MAX, 0..data.len() as u32, &mut positions);
+            positions
+        });
+    }
+    // U64 RANGE END
+
+    // U128 RANGE START
     #[bench]
     fn bench_intfastfield_getrange_u128_50percent_hit(b: &mut Bencher) {
-        let (major_item, _minor_item, data) = get_data_50percent_item();
+        let data = get_data_50percent_item();
         let column = get_u128_column_from_data(&data);
 
         b.iter(|| {
             let mut positions = Vec::new();
             column.get_docids_for_value_range(
-                major_item..=major_item,
+                *FIFTY_PERCENT_RANGE.start() as u128..=*FIFTY_PERCENT_RANGE.end() as u128,
                 0..data.len() as u32,
                 &mut positions,
             );
@@ -124,13 +188,13 @@ mod tests {
 
     #[bench]
     fn bench_intfastfield_getrange_u128_single_hit(b: &mut Bencher) {
-        let (_major_item, minor_item, data) = get_data_50percent_item();
+        let data = get_data_50percent_item();
         let column = get_u128_column_from_data(&data);
 
         b.iter(|| {
             let mut positions = Vec::new();
             column.get_docids_for_value_range(
-                minor_item..=minor_item,
+                *SINGLE_ITEM_RANGE.start() as u128..=*SINGLE_ITEM_RANGE.end() as u128,
                 0..data.len() as u32,
                 &mut positions,
             );
@@ -140,7 +204,7 @@ mod tests {
 
     #[bench]
     fn bench_intfastfield_getrange_u128_hit_all(b: &mut Bencher) {
-        let (_major_item, _minor_item, data) = get_data_50percent_item();
+        let data = get_data_50percent_item();
         let column = get_u128_column_from_data(&data);
 
         b.iter(|| {
@@ -149,6 +213,7 @@ mod tests {
             positions
         });
     }
+    // U128 RANGE END
 
     #[bench]
     fn bench_intfastfield_scan_all_fflookup_u128(b: &mut Bencher) {
diff --git a/src/fastfield/multivalued/reader.rs b/src/fastfield/multivalued/reader.rs
@@ -122,8 +122,78 @@ impl<T: PartialOrd + MakeZero + Clone> MultiValuedFastFieldReader<T> {
 #[cfg(test)]
 mod tests {
 
+    use time::{Duration, OffsetDateTime};
+
+    use crate::collector::Count;
     use crate::core::Index;
+    use crate::query::RangeQuery;
     use crate::schema::{Cardinality, Facet, FacetOptions, NumericOptions, Schema};
+    use crate::{DateOptions, DateTime};
+
+    #[test]
+    fn test_multivalued_date_docids_for_value_range() -> crate::Result<()> {
+        let mut schema_builder = Schema::builder();
+        let date_field = schema_builder.add_date_field(
+            "multi_date_field",
+            DateOptions::default()
+                .set_fast(Cardinality::MultiValues)
+                .set_indexed()
+                .set_fieldnorm()
+                .set_stored(),
+        );
+        let schema = schema_builder.build();
+        let index = Index::create_in_ram(schema);
+        let mut index_writer = index.writer_for_tests()?;
+        let first_time_stamp = OffsetDateTime::now_utc();
+        index_writer.add_document(doc!(
+            date_field => DateTime::from_utc(first_time_stamp),
+            date_field => DateTime::from_utc(first_time_stamp),
+        ))?;
+        index_writer.add_document(doc!())?;
+        // add one second
+        index_writer.add_document(doc!(
+            date_field => DateTime::from_utc(first_time_stamp + Duration::seconds(1)),
+        ))?;
+        // add another second
+        let two_secs_ahead = first_time_stamp + Duration::seconds(2);
+        index_writer.add_document(doc!(
+            date_field => DateTime::from_utc(two_secs_ahead),
+            date_field => DateTime::from_utc(two_secs_ahead),
+            date_field => DateTime::from_utc(two_secs_ahead),
+        ))?;
+        // add three seconds
+        index_writer.add_document(doc!(
+            date_field => DateTime::from_utc(first_time_stamp + Duration::seconds(3)),
+        ))?;
+        index_writer.commit()?;
+
+        let reader = index.reader()?;
+        let searcher = reader.searcher();
+        let reader = searcher.segment_reader(0);
+        assert_eq!(reader.num_docs(), 5);
+
+        let date_ff_reader = reader.fast_fields().dates(date_field).unwrap();
+        let mut docids = vec![];
+        date_ff_reader.get_docids_for_value_range(
+            DateTime::from_utc(first_time_stamp)..=DateTime::from_utc(two_secs_ahead),
+            0..5,
+            &mut docids,
+        );
+        assert_eq!(docids, vec![0, 2, 3]);
+
+        let count_multiples =
+            |range_query: RangeQuery| searcher.search(&range_query, &Count).unwrap();
+
+        assert_eq!(
+            count_multiples(RangeQuery::new_date(
+                date_field,
+                DateTime::from_utc(first_time_stamp)..DateTime::from_utc(two_secs_ahead)
+            )),
+            2
+        );
+
+        Ok(())
+    }
 
     #[test]
     fn test_multifastfield_reader() -> crate::Result<()> {
diff --git a/src/query/range_query/mod.rs b/src/query/range_query/mod.rs
@@ -1,5 +1,6 @@
 mod fast_field_range_query;
 mod range_query;
 mod range_query_ip_fastfield;
+mod range_query_u64_fastfield;
 
 pub use self::range_query::RangeQuery;
diff --git a/src/query/range_query/range_query.rs b/src/query/range_query/range_query.rs
@@ -1,16 +1,17 @@
 use std::io;
 use std::ops::{Bound, Range};
 
-use common::BitSet;
+use common::{BinarySerializable, BitSet};
 
+use super::range_query_u64_fastfield::FastFieldRangeWeight;
 use crate::core::SegmentReader;
 use crate::error::TantivyError;
 use crate::query::explanation::does_not_match;
 use crate::query::range_query::range_query_ip_fastfield::IPFastFieldRangeWeight;
 use crate::query::{BitSetDocSet, ConstScorer, EnableScoring, Explanation, Query, Scorer, Weight};
 use crate::schema::{Field, IndexRecordOption, Term, Type};
 use crate::termdict::{TermDictionary, TermStreamer};
-use crate::{DocId, Score};
+use crate::{DateTime, DocId, Score};
 
 pub(crate) fn map_bound<TFrom, TTo, Transform: Fn(&TFrom) -> TTo>(
     bound: &Bound<TFrom>,
@@ -203,6 +204,40 @@ impl RangeQuery {
         )
     }
 
+    /// Create a new `RangeQuery` over a `date` field.
+    ///
+    /// The two `Bound` arguments make it possible to create more complex
+    /// ranges than semi-inclusive range.
+    ///
+    /// If the field is not of the type `date`, tantivy
+    /// will panic when the `Weight` object is created.
+    pub fn new_date_bounds(
+        field: Field,
+        left_bound: Bound<DateTime>,
+        right_bound: Bound<DateTime>,
+    ) -> RangeQuery {
+        let make_term_val =
+            |val: &DateTime| Term::from_field_date(field, *val).value_bytes().to_owned();
+        RangeQuery {
+            field,
+            value_type: Type::Date,
+            left_bound: map_bound(&left_bound, &make_term_val),
+            right_bound: map_bound(&right_bound, &make_term_val),
+        }
+    }
+
+    /// Create a new `RangeQuery` over a `date` field.
+    ///
+    /// If the field is not of the type `date`, tantivy
+    /// will panic when the `Weight` object is created.
+    pub fn new_date(field: Field, range: Range<DateTime>) -> RangeQuery {
+        RangeQuery::new_date_bounds(
+            field,
+            Bound::Included(range.start),
+            Bound::Excluded(range.end),
+        )
+    }
+
     /// Create a new `RangeQuery` over a `Str` field.
     ///
     /// The two `Bound` arguments make it possible to create more complex
@@ -252,6 +287,23 @@ impl RangeQuery {
     }
 }
 
+fn is_type_valid_for_fastfield_range_query(typ: Type) -> bool {
+    match typ {
+        Type::U64 | Type::I64 | Type::F64 | Type::Bool | Type::Date => true,
+        Type::IpAddr => true,
+        Type::Str | Type::Facet | Type::Bytes | Type::Json => false,
+    }
+}
+
+/// Returns true if the type maps to a u64 fast field
+pub(crate) fn maps_to_u64_fastfield(typ: Type) -> bool {
+    match typ {
+        Type::U64 | Type::I64 | Type::F64 | Type::Bool | Type::Date => true,
+        Type::IpAddr => false,
+        Type::Str | Type::Facet | Type::Bytes | Type::Json => false,
+    }
+}
+
 impl Query for RangeQuery {
     fn weight(&self, enable_scoring: EnableScoring<'_>) -> crate::Result<Box<dyn Weight>> {
         let schema = enable_scoring.schema();
@@ -265,12 +317,29 @@ impl Query for RangeQuery {
             return Err(TantivyError::SchemaError(err_msg));
         }
 
-        if field_type.is_ip_addr() && field_type.is_fast() {
-            Ok(Box::new(IPFastFieldRangeWeight::new(
-                self.field,
-                &self.left_bound,
-                &self.right_bound,
-            )))
+        if field_type.is_fast() && is_type_valid_for_fastfield_range_query(self.value_type) {
+            if field_type.is_ip_addr() {
+                Ok(Box::new(IPFastFieldRangeWeight::new(
+                    self.field,
+                    &self.left_bound,
+                    &self.right_bound,
+                )))
+            } else {
+                // We run the range query on u64 value space for performance reasons and simpicity
+                // assert the type maps to u64
+                assert!(maps_to_u64_fastfield(self.value_type));
+                let parse_from_bytes = |data: &Vec<u8>| {
+                    u64::from_be(BinarySerializable::deserialize(&mut &data[..]).unwrap())
+                };
+
+                let left_bound = map_bound(&self.left_bound, &parse_from_bytes);
+                let right_bound = map_bound(&self.right_bound, &parse_from_bytes);
+                Ok(Box::new(FastFieldRangeWeight::new(
+                    self.field,
+                    left_bound,
+                    right_bound,
+                )))
+            }
         } else {
             Ok(Box::new(RangeWeight {
                 field: self.field,
diff --git a/src/query/range_query/range_query_ip_fastfield.rs b/src/query/range_query/range_query_ip_fastfield.rs
@@ -23,13 +23,13 @@ pub struct IPFastFieldRangeWeight {
 
 impl IPFastFieldRangeWeight {
     pub fn new(field: Field, left_bound: &Bound<Vec<u8>>, right_bound: &Bound<Vec<u8>>) -> Self {
-        let ip_from_bound_raw_data = |data: &Vec<u8>| {
-            let left_ip_u128: u128 =
+        let parse_ip_from_bytes = |data: &Vec<u8>| {
+            let ip_u128: u128 =
                 u128::from_be(BinarySerializable::deserialize(&mut &data[..]).unwrap());
-            Ipv6Addr::from_u128(left_ip_u128)
+            Ipv6Addr::from_u128(ip_u128)
         };
-        let left_bound = map_bound(left_bound, &ip_from_bound_raw_data);
-        let right_bound = map_bound(right_bound, &ip_from_bound_raw_data);
+        let left_bound = map_bound(left_bound, &parse_ip_from_bytes);
+        let right_bound = map_bound(right_bound, &parse_ip_from_bytes);
         Self {
             field,
             left_bound,
diff --git a/src/query/range_query/range_query_u64_fastfield.rs b/src/query/range_query/range_query_u64_fastfield.rs
diff --git a/src/termdict/fst_termdict/merger.rs b/src/termdict/fst_termdict/merger.rs
