use group_by in translate_codec_idx_to_original_id

common/src/group_by.rs

@@ -0,0 +1,166 @@
+use std::cell::RefCell;
+use std::iter::Peekable;
+use std::rc::Rc;
+
+pub trait GroupByIteratorExtended: Iterator {
+    /// Return an `Iterator` that groups iterator elements. Consecutive elements that map to the
+    /// same key are assigned to the same group.
+    ///
+    /// The returned Iterator item is `(K, impl Iterator)`, where Iterator are the items of the
+    /// group.
+    ///
+    /// ```
+    /// use tantivy_common::GroupByIteratorExtended;
+    ///
+    /// // group data into blocks of larger than zero or not.
+    /// let data: Vec<i32> = vec![1, 3, -2, -2, 1, 0, 1, 2];
+    /// // groups:               |---->|------>|--------->|
+    ///
+    /// let mut data_grouped = Vec::new();
+    /// // Note: group is an iterator
+    /// for (key, group) in data.into_iter().group_by(|val| *val >= 0) {
+    ///     data_grouped.push((key, group.collect()));
+    /// }
+    /// assert_eq!(data_grouped, vec![(true, vec![1, 3]), (false, vec![-2, -2]), (true, vec![1, 0, 1, 2])]);
+    /// ```
+    fn group_by<K, F>(self, key: F) -> GroupByIterator<Self, F, K>
+    where
+        Self: Sized,
+        F: FnMut(&Self::Item) -> K,
+        K: PartialEq + Copy,
+        Self::Item: Copy,
+    {
+        GroupByIterator::new(self, key)
+    }
+}
+impl<I: Iterator> GroupByIteratorExtended for I {}
+
+pub struct GroupByIterator<I, F, K: Copy>
+where
+    I: Iterator,
+    F: FnMut(&I::Item) -> K,
+{
+    // I really would like to avoid the Rc<RefCell>, but the Iterator is shared between
+    // `GroupByIterator` and `GroupIter`. In practice they are used consecutive and
+    // `GroupByIter` is finished before calling next on `GroupByIterator`. I'm not sure there
+    // is a solution with lifetimes for that, because we would need to enforce it in the usage
+    // somehow.
+    //
+    // One potential solution would be to replace the iterator approach with something similar.
+    inner: Rc<RefCell<GroupByShared<I, F, K>>>,
+}
+
+struct GroupByShared<I, F, K: Copy>
+where
+    I: Iterator,
+    F: FnMut(&I::Item) -> K,
+{
+    iter: Peekable<I>,
+    group_by_fn: F,
+}
+
+impl<I, F, K> GroupByIterator<I, F, K>
+where
+    I: Iterator,
+    F: FnMut(&I::Item) -> K,
+    K: Copy,
+{
+    fn new(inner: I, group_by_fn: F) -> Self {
+        let inner = GroupByShared {
+            iter: inner.peekable(),
+            group_by_fn,
+        };
+
+        Self {
+            inner: Rc::new(RefCell::new(inner)),
+        }
+    }
+}
+
+impl<I, F, K> Iterator for GroupByIterator<I, F, K>
+where
+    I: Iterator,
+    I::Item: Copy,
+    F: FnMut(&I::Item) -> K,
+    K: Copy,
+{
+    type Item = (K, GroupIterator<I, F, K>);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let mut inner = self.inner.borrow_mut();
+        let value = *inner.iter.peek()?;
+        let key = (inner.group_by_fn)(&value);
+
+        let inner = self.inner.clone();
+
+        let group_iter = GroupIterator {
+            inner,
+            group_key: key,
+        };
+        Some((key, group_iter))
+    }
+}
+
+pub struct GroupIterator<I, F, K: Copy>
+where
+    I: Iterator,
+    F: FnMut(&I::Item) -> K,
+{
+    inner: Rc<RefCell<GroupByShared<I, F, K>>>,
+    group_key: K,
+}
+
+impl<I, F, K: PartialEq + Copy> Iterator for GroupIterator<I, F, K>
+where
+    I: Iterator,
+    I::Item: Copy,
+    F: FnMut(&I::Item) -> K,
+{
+    type Item = I::Item;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let mut inner = self.inner.borrow_mut();
+        // peek if next value is in group
+        let peek_val = *inner.iter.peek()?;
+        if (inner.group_by_fn)(&peek_val) == self.group_key {
+            inner.iter.next()
+        } else {
+            None
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn group_by_collect<I: Iterator<Item = u32>>(iter: I) -> Vec<(I::Item, Vec<I::Item>)> {
+        iter.group_by(|val| val / 10)
+            .map(|(el, iter)| (el, iter.collect::<Vec<_>>()))
+            .collect::<Vec<_>>()
+    }
+
+    #[test]
+    fn group_by_two_groups() {
+        let vals = vec![1u32, 4, 15];
+        let grouped_vals = group_by_collect(vals.into_iter());
+        assert_eq!(grouped_vals, vec![(0, vec![1, 4]), (1, vec![15])]);
+    }
+
+    #[test]
+    fn group_by_test_empty() {
+        let vals = vec![];
+        let grouped_vals = group_by_collect(vals.into_iter());
+        assert_eq!(grouped_vals, vec![]);
+    }
+
+    #[test]
+    fn group_by_three_groups() {
+        let vals = vec![1u32, 4, 15, 1];
+        let grouped_vals = group_by_collect(vals.into_iter());
+        assert_eq!(
+            grouped_vals,
+            vec![(0, vec![1, 4]), (1, vec![15]), (0, vec![1])]
+        );
+    }
+}

common/src/lib.rs

@@ -6,10 +6,12 @@ pub use byteorder::LittleEndian as Endianness;
 
 mod bitset;
 pub mod file_slice;
+mod group_by;
 mod serialize;
 mod vint;
 mod writer;
 pub use bitset::*;
+pub use group_by::GroupByIteratorExtended;
 pub use ownedbytes::{OwnedBytes, StableDeref};
 pub use serialize::{BinarySerializable, DeserializeFrom, FixedSize};
 pub use vint::{

fastfield_codecs/src/null_index/dense.rs

@@ -347,11 +347,16 @@ mod bench {
         codec
     }
 
-    fn random_range_iterator(start: u32, end: u32, step_size: u32) -> impl Iterator<Item = u32> {
+    fn random_range_iterator(
+        start: u32,
+        end: u32,
+        avg_step_size: u32,
+        avg_deviation: u32,
+    ) -> impl Iterator<Item = u32> {
         let mut rng: StdRng = StdRng::from_seed([1u8; 32]);
         let mut current = start;
         std::iter::from_fn(move || {
-            current += rng.gen_range(1..step_size + 1);
+            current += rng.gen_range(avg_step_size - avg_deviation..=avg_step_size + avg_deviation);
             if current >= end {
                 None
             } else {
@@ -360,10 +365,17 @@ mod bench {
         })
     }
 
-    fn walk_over_data(codec: &DenseCodec, max_step_size: u32) -> Option<u32> {
+    fn n_percent_step_iterator(percent: f32, num_values: u32) -> impl Iterator<Item = u32> {
+        let ratio = percent as f32 / 100.0;
+        let step_size = (1f32 / ratio) as u32;
+        let deviation = step_size - 1;
+        random_range_iterator(0, num_values, step_size, deviation)
+    }
+
+    fn walk_over_data(codec: &DenseCodec, avg_step_size: u32) -> Option<u32> {
         walk_over_data_from_positions(
             codec,
-            random_range_iterator(0, TOTAL_NUM_VALUES, max_step_size),
+            random_range_iterator(0, TOTAL_NUM_VALUES, avg_step_size, 0),
         )
     }
 
@@ -379,69 +391,105 @@ mod bench {
     }
 
     #[bench]
-    fn bench_dense_codec_translate_orig_to_codec_90percent_filled_random_stride(
-        bench: &mut Bencher,
-    ) {
-        let codec = gen_bools(0.9f64);
+    fn bench_translate_orig_to_codec_1percent_filled_10percent_hit(bench: &mut Bencher) {
+        let codec = gen_bools(0.01f64);
         bench.iter(|| walk_over_data(&codec, 100));
     }
 
     #[bench]
-    fn bench_dense_codec_translate_orig_to_codec_50percent_filled_random_stride(
-        bench: &mut Bencher,
-    ) {
-        let codec = gen_bools(0.5f64);
+    fn bench_translate_orig_to_codec_5percent_filled_10percent_hit(bench: &mut Bencher) {
+        let codec = gen_bools(0.05f64);
         bench.iter(|| walk_over_data(&codec, 100));
     }
 
     #[bench]
-    fn bench_dense_codec_translate_orig_to_codec_full_scan_10percent(bench: &mut Bencher) {
+    fn bench_translate_orig_to_codec_5percent_filled_1percent_hit(bench: &mut Bencher) {
+        let codec = gen_bools(0.05f64);
+        bench.iter(|| walk_over_data(&codec, 1000));
+    }
+
+    #[bench]
+    fn bench_translate_orig_to_codec_full_scan_1percent_filled(bench: &mut Bencher) {
+        let codec = gen_bools(0.01f64);
+        bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES));
+    }
+
+    #[bench]
+    fn bench_translate_orig_to_codec_full_scan_10percent_filled(bench: &mut Bencher) {
         let codec = gen_bools(0.1f64);
         bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES));
     }
 
     #[bench]
-    fn bench_dense_codec_translate_orig_to_codec_full_scan_90percent(bench: &mut Bencher) {
+    fn bench_translate_orig_to_codec_full_scan_90percent_filled(bench: &mut Bencher) {
         let codec = gen_bools(0.9f64);
         bench.iter(|| walk_over_data_from_positions(&codec, 0..TOTAL_NUM_VALUES));
     }
 
     #[bench]
-    fn bench_dense_codec_translate_orig_to_codec_10percent_filled_random_stride(
-        bench: &mut Bencher,
-    ) {
+    fn bench_translate_orig_to_codec_10percent_filled_1percent_hit(bench: &mut Bencher) {
         let codec = gen_bools(0.1f64);
         bench.iter(|| walk_over_data(&codec, 100));
     }
 
     #[bench]
-    fn bench_dense_codec_translate_codec_to_orig_90percent_filled_random_stride_big_step(
-        bench: &mut Bencher,
-    ) {
+    fn bench_translate_orig_to_codec_50percent_filled_1percent_hit(bench: &mut Bencher) {
+        let codec = gen_bools(0.5f64);
+        bench.iter(|| walk_over_data(&codec, 100));
+    }
+
+    #[bench]
+    fn bench_translate_orig_to_codec_90percent_filled_1percent_hit(bench: &mut Bencher) {
         let codec = gen_bools(0.9f64);
-        let num_vals = codec.num_non_nulls();
+        bench.iter(|| walk_over_data(&codec, 100));
+    }
+
+    #[bench]
+    fn bench_translate_codec_to_orig_1percent_filled_0comma005percent_hit(bench: &mut Bencher) {
+        let codec = gen_bools(0.01f64);
+        let num_non_nulls = codec.num_non_nulls();
         bench.iter(|| {
             codec
-                .translate_codec_idx_to_original_idx(random_range_iterator(0, num_vals, 50_000))
+                .translate_codec_idx_to_original_idx(n_percent_step_iterator(0.005, num_non_nulls))
                 .last()
         });
     }
 
     #[bench]
-    fn bench_dense_codec_translate_codec_to_orig_90percent_filled_random_stride(
-        bench: &mut Bencher,
-    ) {
-        let codec = gen_bools(0.9f64);
+    fn bench_translate_codec_to_orig_1percent_filled_10percent_hit(bench: &mut Bencher) {
+        let codec = gen_bools(0.01f64);
+        let num_non_nulls = codec.num_non_nulls();
+        bench.iter(|| {
+            codec
+                .translate_codec_idx_to_original_idx(n_percent_step_iterator(10.0, num_non_nulls))
+                .last()
+        });
+    }
+
+    #[bench]
+    fn bench_translate_codec_to_orig_1percent_filled_full_scan(bench: &mut Bencher) {
+        let codec = gen_bools(0.01f64);
         let num_vals = codec.num_non_nulls();
         bench.iter(|| {
             codec
-                .translate_codec_idx_to_original_idx(random_range_iterator(0, num_vals, 100))
+                .translate_codec_idx_to_original_idx(0..num_vals)
+                .last()
+        });
+    }
+
+    #[bench]
+    fn bench_translate_codec_to_orig_90percent_filled_0comma005percent_hit(bench: &mut Bencher) {
+        let codec = gen_bools(0.90f64);
+        let num_non_nulls = codec.num_non_nulls();
+        bench.iter(|| {
+            codec
+                .translate_codec_idx_to_original_idx(n_percent_step_iterator(0.005, num_non_nulls))
                 .last()
         });
     }
 
     #[bench]
-    fn bench_dense_codec_translate_codec_to_orig_90percent_filled_full_scan(bench: &mut Bencher) {
+    fn bench_translate_codec_to_orig_90percent_filled_full_scan(bench: &mut Bencher) {
         let codec = gen_bools(0.9f64);
         let num_vals = codec.num_non_nulls();
         bench.iter(|| {