Merge pull request #1505 from quickwit-oss/refact-fast-field

Refact fast field

Author: PSeitz

fastfield_codecs/benches/bench.rs

@@ -4,14 +4,17 @@ extern crate test;
 
 #[cfg(test)]
 mod tests {
+    use std::sync::Arc;
+
     use fastfield_codecs::bitpacked::BitpackedCodec;
     use fastfield_codecs::blockwise_linear::BlockwiseLinearCodec;
     use fastfield_codecs::linear::LinearCodec;
     use fastfield_codecs::*;
 
     fn get_data() -> Vec<u64> {
+        let mut rng = StdRng::seed_from_u64(2u64);
         let mut data: Vec<_> = (100..55000_u64)
-            .map(|num| num + rand::random::<u8>() as u64)
+            .map(|num| num + rng.gen::<u8>() as u64)
             .collect();
         data.push(99_000);
         data.insert(1000, 2000);
@@ -22,32 +25,59 @@ mod tests {
         data
     }
 
+    #[inline(never)]
     fn value_iter() -> impl Iterator<Item = u64> {
         0..20_000
     }
+    fn get_reader_for_bench<Codec: FastFieldCodec>(data: &[u64]) -> Codec::Reader {
+        let mut bytes = Vec::new();
+        let col = VecColumn::from(&data);
+        let normalized_header = fastfield_codecs::NormalizedHeader {
+            num_vals: col.num_vals(),
+            max_value: col.max_value(),
+        };
+        Codec::serialize(&VecColumn::from(data), &mut bytes).unwrap();
+        Codec::open_from_bytes(OwnedBytes::new(bytes), normalized_header).unwrap()
+    }
     fn bench_get<Codec: FastFieldCodec>(b: &mut Bencher, data: &[u64]) {
-        let mut bytes = vec![];
-        Codec::serialize(&mut bytes, &VecColumn::from(data)).unwrap();
-        let reader = Codec::open_from_bytes(OwnedBytes::new(bytes)).unwrap();
+        let col = get_reader_for_bench::<Codec>(data);
+        b.iter(|| {
+            let mut sum = 0u64;
+            for pos in value_iter() {
+                let val = col.get_val(pos as u64);
+                sum = sum.wrapping_add(val);
+            }
+            sum
+        });
+    }
+
+    #[inline(never)]
+    fn bench_get_dynamic_helper(b: &mut Bencher, col: Arc<dyn Column>) {
         b.iter(|| {
             let mut sum = 0u64;
             for pos in value_iter() {
-                let val = reader.get_val(pos as u64);
-                debug_assert_eq!(data[pos as usize], val);
+                let val = col.get_val(pos as u64);
                 sum = sum.wrapping_add(val);
             }
             sum
         });
     }
+
+    fn bench_get_dynamic<Codec: FastFieldCodec>(b: &mut Bencher, data: &[u64]) {
+        let col = Arc::new(get_reader_for_bench::<Codec>(data));
+        bench_get_dynamic_helper(b, col);
+    }
     fn bench_create<Codec: FastFieldCodec>(b: &mut Bencher, data: &[u64]) {
         let mut bytes = Vec::new();
         b.iter(|| {
             bytes.clear();
-            Codec::serialize(&mut bytes, &VecColumn::from(data)).unwrap();
+            Codec::serialize(&VecColumn::from(data), &mut bytes).unwrap();
         });
     }
 
     use ownedbytes::OwnedBytes;
+    use rand::rngs::StdRng;
+    use rand::{Rng, SeedableRng};
     use test::Bencher;
     #[bench]
     fn bench_fastfield_bitpack_create(b: &mut Bencher) {
@@ -70,22 +100,28 @@ mod tests {
         bench_get::<BitpackedCodec>(b, &data);
     }
     #[bench]
+    fn bench_fastfield_bitpack_get_dynamic(b: &mut Bencher) {
+        let data: Vec<_> = get_data();
+        bench_get_dynamic::<BitpackedCodec>(b, &data);
+    }
+    #[bench]
     fn bench_fastfield_linearinterpol_get(b: &mut Bencher) {
         let data: Vec<_> = get_data();
         bench_get::<LinearCodec>(b, &data);
     }
     #[bench]
+    fn bench_fastfield_linearinterpol_get_dynamic(b: &mut Bencher) {
+        let data: Vec<_> = get_data();
+        bench_get_dynamic::<LinearCodec>(b, &data);
+    }
+    #[bench]
     fn bench_fastfield_multilinearinterpol_get(b: &mut Bencher) {
         let data: Vec<_> = get_data();
         bench_get::<BlockwiseLinearCodec>(b, &data);
     }
-    pub fn stats_from_vec(data: &[u64]) -> FastFieldStats {
-        let min_value = data.iter().cloned().min().unwrap_or(0);
-        let max_value = data.iter().cloned().max().unwrap_or(0);
-        FastFieldStats {
-            min_value,
-            max_value,
-            num_vals: data.len() as u64,
-        }
+    #[bench]
+    fn bench_fastfield_multilinearinterpol_get_dynamic(b: &mut Bencher) {
+        let data: Vec<_> = get_data();
+        bench_get_dynamic::<BlockwiseLinearCodec>(b, &data);
     }
 }

fastfield_codecs/src/bitpacked.rs

@@ -1,9 +1,9 @@
 use std::io::{self, Write};
 
-use common::BinarySerializable;
 use ownedbytes::OwnedBytes;
 use tantivy_bitpacker::{compute_num_bits, BitPacker, BitUnpacker};
 
+use crate::serialize::NormalizedHeader;
 use crate::{Column, FastFieldCodec, FastFieldCodecType};
 
 /// Depending on the field type, a different
@@ -12,80 +12,26 @@ use crate::{Column, FastFieldCodec, FastFieldCodecType};
 pub struct BitpackedReader {
     data: OwnedBytes,
     bit_unpacker: BitUnpacker,
-    min_value_u64: u64,
-    max_value_u64: u64,
-    num_vals: u64,
+    normalized_header: NormalizedHeader,
 }
 
 impl Column for BitpackedReader {
     #[inline]
     fn get_val(&self, doc: u64) -> u64 {
-        self.min_value_u64 + self.bit_unpacker.get(doc, &self.data)
+        self.bit_unpacker.get(doc, &self.data)
     }
     #[inline]
     fn min_value(&self) -> u64 {
-        self.min_value_u64
+        // The BitpackedReader assumes a normalized vector.
+        0
     }
     #[inline]
     fn max_value(&self) -> u64 {
-        self.max_value_u64
+        self.normalized_header.max_value
     }
     #[inline]
     fn num_vals(&self) -> u64 {
-        self.num_vals
-    }
-}
-pub struct BitpackedSerializerLegacy<'a, W: 'a + Write> {
-    bit_packer: BitPacker,
-    write: &'a mut W,
-    min_value: u64,
-    num_vals: u64,
-    amplitude: u64,
-    num_bits: u8,
-}
-
-impl<'a, W: Write> BitpackedSerializerLegacy<'a, W> {
-    /// Creates a new fast field serializer.
-    ///
-    /// The serializer in fact encode the values by bitpacking
-    /// `(val - min_value)`.
-    ///
-    /// It requires a `min_value` and a `max_value` to compute
-    /// compute the minimum number of bits required to encode
-    /// values.
-    pub fn open(
-        write: &'a mut W,
-        min_value: u64,
-        max_value: u64,
-    ) -> io::Result<BitpackedSerializerLegacy<'a, W>> {
-        assert!(min_value <= max_value);
-        let amplitude = max_value - min_value;
-        let num_bits = compute_num_bits(amplitude);
-        let bit_packer = BitPacker::new();
-        Ok(BitpackedSerializerLegacy {
-            bit_packer,
-            write,
-            min_value,
-            num_vals: 0,
-            amplitude,
-            num_bits,
-        })
-    }
-    /// Pushes a new value to the currently open u64 fast field.
-    #[inline]
-    pub fn add_val(&mut self, val: u64) -> io::Result<()> {
-        let val_to_write: u64 = val - self.min_value;
-        self.bit_packer
-            .write(val_to_write, self.num_bits, &mut self.write)?;
-        self.num_vals += 1;
-        Ok(())
-    }
-    pub fn close_field(mut self) -> io::Result<()> {
-        self.bit_packer.close(&mut self.write)?;
-        self.min_value.serialize(&mut self.write)?;
-        self.amplitude.serialize(&mut self.write)?;
-        self.num_vals.serialize(&mut self.write)?;
-        Ok(())
+        self.normalized_header.num_vals
     }
 }
 
@@ -98,50 +44,39 @@ impl FastFieldCodec for BitpackedCodec {
     type Reader = BitpackedReader;
 
     /// Opens a fast field given a file.
-    fn open_from_bytes(bytes: OwnedBytes) -> io::Result<Self::Reader> {
-        let footer_offset = bytes.len() - 24;
-        let (data, mut footer) = bytes.split(footer_offset);
-        let min_value = u64::deserialize(&mut footer)?;
-        let amplitude = u64::deserialize(&mut footer)?;
-        let num_vals = u64::deserialize(&mut footer)?;
-        let max_value = min_value + amplitude;
-        let num_bits = compute_num_bits(amplitude);
+    fn open_from_bytes(
+        data: OwnedBytes,
+        normalized_header: NormalizedHeader,
+    ) -> io::Result<Self::Reader> {
+        let num_bits = compute_num_bits(normalized_header.max_value);
         let bit_unpacker = BitUnpacker::new(num_bits);
         Ok(BitpackedReader {
             data,
             bit_unpacker,
-            min_value_u64: min_value,
-            max_value_u64: max_value,
-            num_vals,
+            normalized_header,
         })
     }
 
     /// Serializes data with the BitpackedFastFieldSerializer.
     ///
-    /// The serializer in fact encode the values by bitpacking
-    /// `(val - min_value)`.
+    /// The bitpacker assumes that the column has been normalized.
+    /// i.e. It has already been shifted by its minimum value, so that its
+    /// current minimum value is 0.
     ///
-    /// It requires a `min_value` and a `max_value` to compute
-    /// compute the minimum number of bits required to encode
-    /// values.
-    fn serialize(write: &mut impl Write, fastfield_accessor: &dyn Column) -> io::Result<()> {
-        let mut serializer = BitpackedSerializerLegacy::open(
-            write,
-            fastfield_accessor.min_value(),
-            fastfield_accessor.max_value(),
-        )?;
-
-        for val in fastfield_accessor.iter() {
-            serializer.add_val(val)?;
+    /// Ideally, we made a shift upstream on the column so that `col.min_value() == 0`.
+    fn serialize(column: &dyn Column, write: &mut impl Write) -> io::Result<()> {
+        assert_eq!(column.min_value(), 0u64);
+        let num_bits = compute_num_bits(column.max_value());
+        let mut bit_packer = BitPacker::new();
+        for val in column.iter() {
+            bit_packer.write(val, num_bits, write)?;
         }
-        serializer.close_field()?;
-
+        bit_packer.close(write)?;
         Ok(())
     }
 
-    fn estimate(fastfield_accessor: &impl Column) -> Option<f32> {
-        let amplitude = fastfield_accessor.max_value() - fastfield_accessor.min_value();
-        let num_bits = compute_num_bits(amplitude);
+    fn estimate(column: &impl Column) -> Option<f32> {
+        let num_bits = compute_num_bits(column.max_value());
         let num_bits_uncompressed = 64;
         Some(num_bits as f32 / num_bits_uncompressed as f32)
     }