perf: Lazily load in BitpackedCodec (#56)

We would like to be able to lazily load `BitpackedCodec` columns (similar to what 020bdff did for `BlockwiseLinearCodec`), because in the context of `pg_search`, immediately constructing `OwnedBytes` means copying the entire content of the column into memory.

To do so, we expose some (slightly overlapped) block boundaries from `BitUnpacker`, and then lazily load each block when it is requested. Only the `get_val` function uses the cache: `get_row_ids_for_value_range` does not (yet), because it would be necessary to partition the row ids by block, and most of the time consumers using it are already loading reasonably large ranges anyway.

See paradedb/paradedb#2894 for usage. There are a few 2x speedups in the benchmark suite, as well as a 1.8x speedup on a representative customer query. Unfortunately there are also some 13-19% slowdowns on aggregates: it looks like that is because aggregates use `get_vals`, for which the default implementation is to just call `get_val` in a loop.

(cherry picked from commit 57d2ef8)

Author: stuhood

bitpacker/src/bitpacker.rs

@@ -69,6 +69,12 @@ pub struct BitUnpacker {
     mask: u64,
 }
 
+pub type BlockNumber = usize;
+
+// 16k
+const BLOCK_SIZE_MIN_POW: u8 = 14;
+const BLOCK_SIZE_MIN: usize = 2 << BLOCK_SIZE_MIN_POW;
+
 impl BitUnpacker {
     /// Creates a bit unpacker, that assumes the same bitwidth for all values.
     ///
@@ -82,6 +88,7 @@ impl BitUnpacker {
         } else {
             (1u64 << num_bits) - 1u64
         };
+
         BitUnpacker {
             num_bits: u32::from(num_bits),
             mask,
@@ -92,10 +99,63 @@ impl BitUnpacker {
         self.num_bits as u8
     }
 
+    /// Calculates a block number for the given `idx`.
+    #[inline]
+    pub fn block_num(&self, idx: u32) -> BlockNumber {
+        // Find the address in bits of the index.
+        let addr_in_bits = (idx * self.num_bits) as usize;
+
+        // Then round down to the nearest byte.
+        let addr_in_bytes = addr_in_bits >> 3;
+
+        // And compute the containing BlockNumber.
+        addr_in_bytes >> (BLOCK_SIZE_MIN_POW + 1)
+    }
+
+    /// Given a block number and dataset length, calculates a data Range for the block.
+    pub fn block(&self, block: BlockNumber, data_len: usize) -> Range<usize> {
+        let block_addr = block << (BLOCK_SIZE_MIN_POW + 1);
+        // We extend the end of the block by a constant factor, so that it overlaps the next
+        // block. That ensures that we never need to read on a block boundary.
+        block_addr..(std::cmp::min(block_addr + BLOCK_SIZE_MIN + 8, data_len))
+    }
+
+    /// Calculates the number of blocks for the given data_len.
+    ///
+    /// Usually only called at startup to pre-allocate structures.
+    pub fn block_count(&self, data_len: usize) -> usize {
+        let block_count = data_len / (BLOCK_SIZE_MIN as usize);
+        if data_len % (BLOCK_SIZE_MIN as usize) == 0 {
+            block_count
+        } else {
+            block_count + 1
+        }
+    }
+
+    /// Returns a range within the data which covers the given id_range.
+    ///
+    /// NOTE: This method is used for batch reads which bypass blocks to avoid dealing with block
+    /// boundaries.
+    #[inline]
+    pub fn block_oblivious_range(&self, id_range: Range<u32>, data_len: usize) -> Range<usize> {
+        let start_in_bits = id_range.start * self.num_bits;
+        let start = (start_in_bits >> 3) as usize;
+        let end_in_bits = id_range.end * self.num_bits;
+        let end = (end_in_bits >> 3) as usize;
+        // TODO: We fetch more than we need and then truncate.
+        start..(std::cmp::min(end + 8, data_len))
+    }
+
     #[inline]
     pub fn get(&self, idx: u32, data: &[u8]) -> u64 {
+        self.get_from_subset(idx, 0, data)
+    }
+
+    /// Get the value at the given idx, which must exist within the given subset of the data.
+    #[inline]
+    pub fn get_from_subset(&self, idx: u32, data_offset: usize, data: &[u8]) -> u64 {
         let addr_in_bits = idx * self.num_bits;
-        let addr = (addr_in_bits >> 3) as usize;
+        let addr = (addr_in_bits >> 3) as usize - data_offset;
         if addr + 8 > data.len() {
             if self.num_bits == 0 {
                 return 0;
@@ -129,7 +189,7 @@ impl BitUnpacker {
     // #Panics
     //
     // This methods panics if `num_bits` is > 32.
-    fn get_batch_u32s(&self, start_idx: u32, data: &[u8], output: &mut [u32]) {
+    fn get_batch_u32s(&self, start_idx: u32, data_offset: usize, data: &[u8], output: &mut [u32]) {
         assert!(
             self.bit_width() <= 32,
             "Bitwidth must be <= 32 to use this method."
@@ -140,14 +200,14 @@ impl BitUnpacker {
         let end_bit_read = end_idx * self.num_bits;
         let end_byte_read = (end_bit_read + 7) / 8;
         assert!(
-            end_byte_read as usize <= data.len(),
+            end_byte_read as usize <= data_offset + data.len(),
             "Requested index is out of bounds."
         );
 
         // Simple slow implementation of get_batch_u32s, to deal with our ramps.
         let get_batch_ramp = |start_idx: u32, output: &mut [u32]| {
             for (out, idx) in output.iter_mut().zip(start_idx..) {
-                *out = self.get(idx, data) as u32;
+                *out = self.get_from_subset(idx, data_offset, data) as u32;
             }
         };
 
@@ -177,7 +237,7 @@ impl BitUnpacker {
         get_batch_ramp(start_idx, &mut output[..entrance_ramp_len as usize]);
 
         // Highway
-        let mut offset = (highway_start * self.num_bits) as usize / 8;
+        let mut offset = ((highway_start * self.num_bits) as usize / 8) - data_offset;
         let mut output_cursor = (highway_start - start_idx) as usize;
         for _ in 0..num_blocks {
             offset += BitPacker1x.decompress(
@@ -199,31 +259,43 @@ impl BitUnpacker {
         id_range: Range<u32>,
         data: &[u8],
         positions: &mut Vec<u32>,
+    ) {
+        self.get_ids_for_value_range_from_subset(range, id_range, 0, data, positions)
+    }
+
+    pub fn get_ids_for_value_range_from_subset(
+        &self,
+        range: RangeInclusive<u64>,
+        id_range: Range<u32>,
+        data_offset: usize,
+        data: &[u8],
+        positions: &mut Vec<u32>,
     ) {
         if self.bit_width() > 32 {
-            self.get_ids_for_value_range_slow(range, id_range, data, positions)
+            self.get_ids_for_value_range_slow(range, id_range, data_offset, data, positions)
         } else {
             if *range.start() > u32::MAX as u64 {
                 positions.clear();
                 return;
             }
             let range_u32 = (*range.start() as u32)..=(*range.end()).min(u32::MAX as u64) as u32;
-            self.get_ids_for_value_range_fast(range_u32, id_range, data, positions)
+            self.get_ids_for_value_range_fast(range_u32, id_range, data_offset, data, positions)
         }
     }
 
     fn get_ids_for_value_range_slow(
         &self,
         range: RangeInclusive<u64>,
         id_range: Range<u32>,
+        data_offset: usize,
         data: &[u8],
         positions: &mut Vec<u32>,
     ) {
         positions.clear();
         for i in id_range {
             // If we cared we could make this branchless, but the slow implementation should rarely
             // kick in.
-            let val = self.get(i, data);
+            let val = self.get_from_subset(i, data_offset, data);
             if range.contains(&val) {
                 positions.push(i);
             }
@@ -234,11 +306,12 @@ impl BitUnpacker {
         &self,
         value_range: RangeInclusive<u32>,
         id_range: Range<u32>,
+        data_offset: usize,
         data: &[u8],
         positions: &mut Vec<u32>,
     ) {
         positions.resize(id_range.len(), 0u32);
-        self.get_batch_u32s(id_range.start, data, positions);
+        self.get_batch_u32s(id_range.start, data_offset, data, positions);
         crate::filter_vec::filter_vec_in_place(value_range, id_range.start, positions)
     }
 }
@@ -329,14 +402,14 @@ mod test {
     fn test_get_batch_panics_over_32_bits() {
         let bitunpacker = BitUnpacker::new(33);
         let mut output: [u32; 1] = [0u32];
-        bitunpacker.get_batch_u32s(0, &[0, 0, 0, 0, 0, 0, 0, 0], &mut output[..]);
+        bitunpacker.get_batch_u32s(0, 0, &[0, 0, 0, 0, 0, 0, 0, 0], &mut output[..]);
     }
 
     #[test]
     fn test_get_batch_limit() {
         let bitunpacker = BitUnpacker::new(1);
         let mut output: [u32; 3] = [0u32, 0u32, 0u32];
-        bitunpacker.get_batch_u32s(8 * 4 - 3, &[0u8, 0u8, 0u8, 0u8], &mut output[..]);
+        bitunpacker.get_batch_u32s(8 * 4 - 3, 0, &[0u8, 0u8, 0u8, 0u8], &mut output[..]);
     }
 
     #[test]
@@ -345,7 +418,7 @@ mod test {
         let bitunpacker = BitUnpacker::new(1);
         let mut output: [u32; 3] = [0u32, 0u32, 0u32];
         // We are missing exactly one bit.
-        bitunpacker.get_batch_u32s(8 * 4 - 2, &[0u8, 0u8, 0u8, 0u8], &mut output[..]);
+        bitunpacker.get_batch_u32s(8 * 4 - 2, 0, &[0u8, 0u8, 0u8, 0u8], &mut output[..]);
     }
 
     proptest::proptest! {
@@ -368,7 +441,7 @@ mod test {
             for len in [0, 1, 2, 32, 33, 34, 64] {
                 for start_idx in 0u32..32u32 {
                     output.resize(len, 0);
-                    bitunpacker.get_batch_u32s(start_idx, &buffer, &mut output);
+                    bitunpacker.get_batch_u32s(start_idx, 0, &buffer, &mut output);
                     for (i, output_byte) in output.iter().enumerate() {
                         let expected = (start_idx + i as u32) & mask;
                         assert_eq!(*output_byte, expected);

columnar/src/column_values/stats.rs

@@ -2,7 +2,8 @@ use std::io;
 use std::io::Write;
 use std::num::NonZeroU64;
 
-use common::{BinarySerializable, VInt};
+use common::file_slice::FileSlice;
+use common::{BinarySerializable, HasLen, VInt};
 
 use crate::RowId;
 
@@ -28,9 +29,21 @@ impl ColumnStats {
 }
 
 impl ColumnStats {
+    /// Deserialize from the tail of the given FileSlice, and return the stats and remaining prefix
+    /// FileSlice.
+    pub fn deserialize_from_tail(file_slice: FileSlice) -> io::Result<(Self, FileSlice)> {
+        // [`deserialize_with_size`] deserializes 4 variable-width encoded u64s, which
+        // could end up being, in the worst case, 9 bytes each. this is where the 36 comes from
+        let (stats, _) = file_slice.clone().split(36.min(file_slice.len())); // hope that's enough bytes
+        let mut stats = stats.read_bytes()?;
+        let (stats, stats_nbytes) = ColumnStats::deserialize_with_size(&mut stats)?;
+        let (_, remainder) = file_slice.split(stats_nbytes);
+        Ok((stats, remainder))
+    }
+
     /// Same as [`BinarySeerializable::deserialize`] but also returns the number of bytes
     /// consumed from the reader `R`
-    pub fn deserialize_with_size<R: io::Read>(reader: &mut R) -> io::Result<(Self, usize)> {
+    fn deserialize_with_size<R: io::Read>(reader: &mut R) -> io::Result<(Self, usize)> {
         let mut nbytes = 0;
 
         let (min_value, len) = VInt::deserialize_with_size(reader)?;

columnar/src/column_values/u64_based/bitpacked.rs

@@ -1,9 +1,10 @@
 use std::io::{self, Write};
 use std::num::NonZeroU64;
 use std::ops::{Range, RangeInclusive};
+use std::sync::{Arc, OnceLock};
 
 use common::file_slice::FileSlice;
-use common::{BinarySerializable, OwnedBytes};
+use common::{BinarySerializable, HasLen, OwnedBytes};
 use fastdivide::DividerU64;
 use tantivy_bitpacker::{BitPacker, BitUnpacker, compute_num_bits};
 
@@ -14,9 +15,40 @@ use crate::{ColumnValues, RowId};
 /// fast field is required.
 #[derive(Clone)]
 pub struct BitpackedReader {
-    data: OwnedBytes,
+    data: FileSlice,
     bit_unpacker: BitUnpacker,
     stats: ColumnStats,
+    blocks: Arc<[OnceLock<Block>]>,
+}
+
+impl BitpackedReader {
+    #[inline(always)]
+    fn unpack_val(&self, doc: u32) -> u64 {
+        let block_num = self.bit_unpacker.block_num(doc);
+
+        if block_num == 0 && self.blocks.len() == 0 {
+            return 0;
+        }
+
+        let block = self.blocks[block_num].get_or_init(|| {
+            let block_range = self.bit_unpacker.block(block_num, self.data.len());
+            let offset = block_range.start;
+            let data = self
+                .data
+                .slice(block_range)
+                .read_bytes()
+                .expect("Failed to read column values.");
+            Block { offset, data }
+        });
+
+        self.bit_unpacker
+            .get_from_subset(doc, block.offset, &block.data)
+    }
+}
+
+struct Block {
+    offset: usize,
+    data: OwnedBytes,
 }
 
 #[inline(always)]
@@ -52,8 +84,9 @@ fn transform_range_before_linear_transformation(
 impl ColumnValues for BitpackedReader {
     #[inline(always)]
     fn get_val(&self, doc: u32) -> u64 {
-        self.stats.min_value + self.stats.gcd.get() * self.bit_unpacker.get(doc, &self.data)
+        self.stats.min_value + self.stats.gcd.get() * self.unpack_val(doc)
     }
+
     #[inline]
     fn min_value(&self) -> u64 {
         self.stats.min_value
@@ -79,10 +112,23 @@ impl ColumnValues for BitpackedReader {
             positions.clear();
             return;
         };
-        self.bit_unpacker.get_ids_for_value_range(
+        // TODO: This does not use the `self.blocks` cache, because callers are usually already
+        // doing sequential, and fairly dense reads. Fix it to iterate over blocks if that
+        // assumption turns out to be incorrect!
+        let data_range = self
+            .bit_unpacker
+            .block_oblivious_range(doc_id_range.clone(), self.data.len());
+        let data_offset = data_range.start;
+        let data_subset = self
+            .data
+            .slice(data_range)
+            .read_bytes()
+            .expect("Failed to read column values.");
+        self.bit_unpacker.get_ids_for_value_range_from_subset(
             transformed_range,
             doc_id_range,
-            &self.data,
+            data_offset,
+            &data_subset,
             positions,
         );
     }
@@ -129,14 +175,19 @@ impl ColumnCodec for BitpackedCodec {
 
     /// Opens a fast field given a file.
     fn load(file_slice: FileSlice) -> io::Result<Self::ColumnValues> {
-        let mut data = file_slice.read_bytes()?;
-        let stats = ColumnStats::deserialize(&mut data)?;
+        let (stats, data) = ColumnStats::deserialize_from_tail(file_slice)?;
+
         let num_bits = num_bits(&stats);
         let bit_unpacker = BitUnpacker::new(num_bits);
+        let block_count = bit_unpacker.block_count(data.len());
         Ok(BitpackedReader {
             data,
             bit_unpacker,
             stats,
+            blocks: (0..block_count)
+                .into_iter()
+                .map(|_| OnceLock::new())
+                .collect(),
         })
     }
 }

columnar/src/column_values/u64_based/blockwise_linear.rs

@@ -175,13 +175,7 @@ impl ColumnCodec<u64> for BlockwiseLinearCodec {
     type Estimator = BlockwiseLinearEstimator;
 
     fn load(file_slice: FileSlice) -> io::Result<Self::ColumnValues> {
-        // [`ColumnStats::deserialize_with_size`] deserializes 4 variable-width encoded u64s, which
-        // could end up being, in the worst case, 9 bytes each.  this is where the 36 comes from
-        let (stats, _) = file_slice.clone().split(36.min(file_slice.len())); // hope that's enough bytes
-        let mut stats = stats.read_bytes()?;
-        let (stats, stats_nbytes) = ColumnStats::deserialize_with_size(&mut stats)?;
-
-        let (_, body) = file_slice.split(stats_nbytes);
+        let (stats, body) = ColumnStats::deserialize_from_tail(file_slice)?;
 
         let (_, footer) = body.clone().split_from_end(4);
 

src/fastfield/mod.rs

@@ -395,7 +395,7 @@ mod tests {
             .unwrap()
             .first_or_default_col(0);
         for a in 0..n {
-            assert_eq!(col.get_val(a as u32), permutation[a]);
+            assert_eq!(col.get_val(a as u32), permutation[a], "for doc {a}");
         }
     }