add method to fetch block of first vals in columnar (#2330)

* add method to fetch block of first vals in columnar

add method to fetch block of first vals in columnar (this is way faster
than single calls for full columns)
add benchmark
fix import warnings

```
test bench_get_block_first_on_full_column                  ... bench:          56 ns/iter (+/- 26)
test bench_get_block_first_on_full_column_single_calls     ... bench:         311 ns/iter (+/- 6)
test bench_get_block_first_on_multi_column                 ... bench:         378 ns/iter (+/- 15)
test bench_get_block_first_on_multi_column_single_calls    ... bench:         546 ns/iter (+/- 13)
test bench_get_block_first_on_optional_column              ... bench:         291 ns/iter (+/- 6)
test bench_get_block_first_on_optional_column_single_calls ... bench:         362 ns/iter (+/- 8)
```

* use remainder

columnar/src/column/mod.rs

@@ -13,7 +13,7 @@ pub use serialize::{
     open_column_u64, serialize_column_mappable_to_u128, serialize_column_mappable_to_u64,
 };
 
-use crate::column_index::ColumnIndex;
+use crate::column_index::{ColumnIndex, Set};
 use crate::column_values::monotonic_mapping::StrictlyMonotonicMappingToInternal;
 use crate::column_values::{monotonic_map_column, ColumnValues};
 use crate::{Cardinality, DocId, EmptyColumnValues, MonotonicallyMappableToU64, RowId};
@@ -83,10 +83,36 @@ impl<T: PartialOrd + Copy + Debug + Send + Sync + 'static> Column<T> {
         self.values.max_value()
     }
 
+    #[inline]
     pub fn first(&self, row_id: RowId) -> Option<T> {
         self.values_for_doc(row_id).next()
     }
 
+    /// Load the first value for each docid in the provided slice.
+    #[inline]
+    pub fn first_vals(&self, docids: &[DocId], output: &mut [Option<T>]) {
+        match &self.index {
+            ColumnIndex::Empty { .. } => {}
+            ColumnIndex::Full => self.values.get_vals_opt(docids, output),
+            ColumnIndex::Optional(optional_index) => {
+                for (i, docid) in docids.iter().enumerate() {
+                    output[i] = optional_index
+                        .rank_if_exists(*docid)
+                        .map(|rowid| self.values.get_val(rowid));
+                }
+            }
+            ColumnIndex::Multivalued(multivalued_index) => {
+                for (i, docid) in docids.iter().enumerate() {
+                    let range = multivalued_index.range(*docid);
+                    let is_empty = range.start == range.end;
+                    if !is_empty {
+                        output[i] = Some(self.values.get_val(range.start));
+                    }
+                }
+            }
+        }
+    }
+
     /// Translates a block of docis to row_ids.
     ///
     /// returns the row_ids and the matching docids on the same index

columnar/src/column_index/optional_index/set_block/dense.rs

@@ -1,4 +1,3 @@
-use std::convert::TryInto;
 use std::io::{self, Write};
 
 use common::BinarySerializable;

columnar/src/column_index/optional_index/tests.rs

@@ -1,5 +1,4 @@
-use proptest::prelude::{any, prop, *};
-use proptest::strategy::Strategy;
+use proptest::prelude::*;
 use proptest::{prop_oneof, proptest};
 
 use super::*;

columnar/src/column_values/mod.rs

@@ -72,11 +72,40 @@ pub trait ColumnValues<T: PartialOrd = u64>: Send + Sync + DowncastSync {
             out_x4[3] = self.get_val(idx_x4[3]);
         }
 
-        let step_size = 4;
-        let cutoff = indexes.len() - indexes.len() % step_size;
+        let out_and_idx_chunks = output
+            .chunks_exact_mut(4)
+            .into_remainder()
+            .into_iter()
+            .zip(indexes.chunks_exact(4).remainder());
+        for (out, idx) in out_and_idx_chunks {
+            *out = self.get_val(*idx);
+        }
+    }
 
-        for idx in cutoff..indexes.len() {
-            output[idx] = self.get_val(indexes[idx]);
+    /// Allows to push down multiple fetch calls, to avoid dynamic dispatch overhead.
+    /// The slightly weird `Option<T>` in output allows pushdown to full columns.
+    ///
+    /// idx and output should have the same length
+    ///
+    /// # Panics
+    ///
+    /// May panic if `idx` is greater than the column length.
+    fn get_vals_opt(&self, indexes: &[u32], output: &mut [Option<T>]) {
+        assert!(indexes.len() == output.len());
+        let out_and_idx_chunks = output.chunks_exact_mut(4).zip(indexes.chunks_exact(4));
+        for (out_x4, idx_x4) in out_and_idx_chunks {
+            out_x4[0] = Some(self.get_val(idx_x4[0]));
+            out_x4[1] = Some(self.get_val(idx_x4[1]));
+            out_x4[2] = Some(self.get_val(idx_x4[2]));
+            out_x4[3] = Some(self.get_val(idx_x4[3]));
+        }
+        let out_and_idx_chunks = output
+            .chunks_exact_mut(4)
+            .into_remainder()
+            .into_iter()
+            .zip(indexes.chunks_exact(4).remainder());
+        for (out, idx) in out_and_idx_chunks {
+            *out = Some(self.get_val(*idx));
         }
     }
 
@@ -172,6 +201,11 @@ impl<T: Copy + PartialOrd + Debug + 'static> ColumnValues<T> for Arc<dyn ColumnV
         self.as_ref().get_val(idx)
     }
 
+    #[inline(always)]
+    fn get_vals_opt(&self, indexes: &[u32], output: &mut [Option<T>]) {
+        self.as_ref().get_vals_opt(indexes, output)
+    }
+
     #[inline(always)]
     fn min_value(&self) -> T {
         self.as_ref().min_value()

columnar/src/column_values/u128_based/compact_space/mod.rs

@@ -22,7 +22,7 @@ mod build_compact_space;
 
 use build_compact_space::get_compact_space;
 use common::{BinarySerializable, CountingWriter, OwnedBytes, VInt, VIntU128};
-use tantivy_bitpacker::{self, BitPacker, BitUnpacker};
+use tantivy_bitpacker::{BitPacker, BitUnpacker};
 
 use crate::column_values::ColumnValues;
 use crate::RowId;

columnar/src/column_values/u64_based/tests.rs

@@ -1,5 +1,4 @@
 use proptest::prelude::*;
-use proptest::strategy::Strategy;
 use proptest::{prop_oneof, proptest};
 
 #[test]

columnar/src/columnar/merge/tests.rs

@@ -1,7 +1,3 @@
-use std::collections::BTreeMap;
-
-use itertools::Itertools;
-
 use super::*;
 use crate::{Cardinality, ColumnarWriter, HasAssociatedColumnType, RowId};
 

columnar/src/columnar/writer/serializer.rs

@@ -96,7 +96,6 @@ impl<'a, W: io::Write> io::Write for ColumnSerializer<'a, W> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::columnar::column_type::ColumnType;
 
     #[test]
     fn test_prepare_key_bytes() {