Add BlockedFoR: a frame of reference per 1024 values

encodings/fastlanes/goldenfiles/blockedfor.metadata

@@ -0,0 +1 @@
+�

encodings/fastlanes/src/blocked_for/array/blocked_for_compress.rs

@@ -0,0 +1,170 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use num_traits::PrimInt;
+use num_traits::WrappingSub;
+use vortex_array::ArrayView;
+use vortex_array::ExecutionCtx;
+use vortex_array::IntoArray;
+use vortex_array::arrays::Primitive;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::dtype::NativePType;
+use vortex_array::match_each_integer_ptype;
+use vortex_array::validity::Validity;
+use vortex_buffer::BitBuffer;
+use vortex_buffer::Buffer;
+use vortex_buffer::BufferMut;
+use vortex_error::VortexResult;
+use vortex_mask::AllOr;
+
+use crate::BlockedFoR;
+use crate::BlockedFoRArray;
+use crate::BlockedFoRData;
+use crate::blocked_for::array::BLOCK_SIZE;
+
+impl BlockedFoRData {
+    /// Encode a primitive array, subtracting a per-block minimum from every block of
+    /// [`BLOCK_SIZE`] values.
+    pub fn encode(array: PrimitiveArray, ctx: &mut ExecutionCtx) -> VortexResult<BlockedFoRArray> {
+        let validity = array.validity()?;
+        let (encoded, references) = match_each_integer_ptype!(array.ptype(), |T| {
+            let (residuals, references) = compress_primitive::<T>(&array, ctx)?;
+            (
+                PrimitiveArray::new(residuals, validity).into_array(),
+                PrimitiveArray::new(references, Validity::NonNullable).into_array(),
+            )
+        });
+        BlockedFoR::try_new(encoded, references, 0)
+    }
+}
+
+/// Split `array` into [`BLOCK_SIZE`] blocks, returning the block-local residuals and the
+/// per-block reference (minimum) values.
+///
+/// Null values encode as a zero residual, exactly as [`crate::FoR`] does, so that they decode
+/// back to their block's reference rather than wrapping out of the primitive range.
+fn compress_primitive<T: NativePType + WrappingSub + PrimInt>(
+    array: &PrimitiveArray,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<(Buffer<T>, Buffer<T>)> {
+    let len = array.len();
+    let values = array.as_slice::<T>();
+    let num_blocks = len.div_ceil(BLOCK_SIZE);
+
+    let mut references = BufferMut::<T>::with_capacity(num_blocks);
+    let mut residuals = BufferMut::<T>::with_capacity(len);
+
+    match array.validity()?.execute_mask(len, ctx)?.bit_buffer() {
+        AllOr::All => {
+            for block in values.chunks(BLOCK_SIZE) {
+                let min = block.iter().copied().min().unwrap_or_else(T::zero);
+                references.push(min);
+                residuals.extend(block.iter().map(|v| v.wrapping_sub(&min)));
+            }
+        }
+        AllOr::None => {
+            references.extend(std::iter::repeat_n(T::zero(), num_blocks));
+            residuals.extend(std::iter::repeat_n(T::zero(), len));
+        }
+        AllOr::Some(bits) => {
+            // Materialize the validity bits once: the per-block loops below need random access
+            // to them, and `BitBuffer::value` is far too costly to call per element.
+            let valid = bits.iter().collect::<Vec<_>>();
+            for (block, valid) in values.chunks(BLOCK_SIZE).zip(valid.chunks(BLOCK_SIZE)) {
+                let min = block
+                    .iter()
+                    .zip(valid)
+                    .filter_map(|(v, valid)| valid.then_some(*v))
+                    .min()
+                    .unwrap_or_else(T::zero);
+                references.push(min);
+                residuals.extend(block.iter().zip(valid).map(|(v, valid)| {
+                    if *valid {
+                        v.wrapping_sub(&min)
+                    } else {
+                        T::zero()
+                    }
+                }));
+            }
+        }
+    }
+
+    Ok((residuals.freeze(), references.freeze()))
+}
+
+/// Per-block summary used to estimate how well blocked FoR will pack.
+pub struct BlockSummary {
+    /// The widest `max - min` over the valid values of any block.
+    pub max_range: u128,
+    /// Whether every block's minimum is already zero, making the encoding a no-op.
+    pub all_minima_zero: bool,
+}
+
+/// Summarize `array` a block at a time, or `None` if it holds no valid values.
+pub fn block_summary(
+    array: ArrayView<'_, Primitive>,
+    exec_ctx: &mut ExecutionCtx,
+) -> Option<BlockSummary> {
+    let len = array.len();
+    let validity = array.validity().ok()?.execute_mask(len, exec_ctx).ok()?;
+
+    match_each_integer_ptype!(array.ptype(), |T| {
+        block_summary_typed::<T>(array.as_slice::<T>(), validity.bit_buffer())
+    })
+}
+
+fn block_summary_typed<T: NativePType + PrimInt>(
+    values: &[T],
+    validity: AllOr<&BitBuffer>,
+) -> Option<BlockSummary> {
+    // Materialize the validity bits once rather than probing them per element.
+    let valid = match validity {
+        AllOr::All => None,
+        AllOr::None => return None,
+        AllOr::Some(bits) => Some(bits.iter().collect::<Vec<_>>()),
+    };
+
+    let mut max_range = 0u128;
+    let mut all_minima_zero = true;
+    let mut any_valid = false;
+
+    for (index, block) in values.chunks(BLOCK_SIZE).enumerate() {
+        let (min, max) = match &valid {
+            None => (block.iter().copied().min(), block.iter().copied().max()),
+            Some(valid) => {
+                let valid = &valid[index * BLOCK_SIZE..][..block.len()];
+                let mut it = block
+                    .iter()
+                    .zip(valid)
+                    .filter_map(|(v, valid)| valid.then_some(*v));
+                match it.next() {
+                    None => (None, None),
+                    Some(first) => {
+                        let (min, max) =
+                            it.fold((first, first), |(min, max), v| (min.min(v), max.max(v)));
+                        (Some(min), Some(max))
+                    }
+                }
+            }
+        };
+
+        let (Some(min), Some(max)) = (min, max) else {
+            // An all-null block encodes to a zero reference and zero residuals.
+            continue;
+        };
+        any_valid = true;
+        all_minima_zero &= min.is_zero();
+        let range = max.to_i128()? - min.to_i128()?;
+        #[expect(
+            clippy::cast_sign_loss,
+            reason = "max >= min, so the range is non-negative"
+        )]
+        let range = range as u128;
+        max_range = max_range.max(range);
+    }
+
+    any_valid.then_some(BlockSummary {
+        max_range,
+        all_minima_zero,
+    })
+}

encodings/fastlanes/src/blocked_for/array/blocked_for_decompress.rs

@@ -0,0 +1,67 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use num_traits::PrimInt;
+use num_traits::WrappingAdd;
+use vortex_array::ExecutionCtx;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::dtype::NativePType;
+use vortex_array::match_each_integer_ptype;
+use vortex_buffer::Buffer;
+use vortex_error::VortexResult;
+
+use crate::BlockedFoRArray;
+use crate::blocked_for::array::BLOCK_SIZE;
+use crate::blocked_for::array::BlockedFoRArrayExt;
+use crate::blocked_for::array::BlockedFoRArraySlotsExt;
+
+pub fn decompress(array: &BlockedFoRArray, ctx: &mut ExecutionCtx) -> VortexResult<PrimitiveArray> {
+    let ptype = array.as_ref().dtype().as_ptype();
+    let offset = array.offset() as usize;
+
+    let references = array.references().clone().execute::<PrimitiveArray>(ctx)?;
+    let encoded = array.encoded().clone().execute::<PrimitiveArray>(ctx)?;
+    let validity = encoded.validity()?;
+
+    Ok(match_each_integer_ptype!(ptype, |T| {
+        PrimitiveArray::new(
+            decompress_primitive::<T>(
+                encoded.into_buffer::<T>(),
+                references.as_slice::<T>(),
+                offset,
+            ),
+            validity,
+        )
+    }))
+}
+
+/// Add each block's reference back onto its residuals.
+///
+/// `offset` is the position of the first value within its block, so the first (possibly short)
+/// block covers `BLOCK_SIZE - offset` values and every subsequent block a full `BLOCK_SIZE`.
+fn decompress_primitive<T: NativePType + WrappingAdd + PrimInt>(
+    values: Buffer<T>,
+    references: &[T],
+    offset: usize,
+) -> Buffer<T> {
+    let len = values.len();
+    let mut values = values.into_mut();
+
+    let mut pos = 0;
+    for (block, reference) in references.iter().enumerate() {
+        let block_len = if block == 0 {
+            (BLOCK_SIZE - offset).min(len)
+        } else {
+            BLOCK_SIZE.min(len - pos)
+        };
+        if !reference.is_zero() {
+            for v in &mut values[pos..pos + block_len] {
+                *v = v.wrapping_add(reference);
+            }
+        }
+        pos += block_len;
+    }
+    debug_assert_eq!(pos, len);
+
+    values.freeze()
+}

encodings/fastlanes/src/blocked_for/array/mod.rs

@@ -0,0 +1,94 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use std::fmt::Display;
+use std::fmt::Formatter;
+
+use vortex_array::ArrayRef;
+use vortex_array::TypedArrayRef;
+use vortex_array::array_slots;
+use vortex_array::dtype::DType;
+use vortex_array::dtype::PType;
+use vortex_error::VortexResult;
+use vortex_error::vortex_ensure;
+
+pub mod blocked_for_compress;
+pub mod blocked_for_decompress;
+
+/// Number of values covered by a single reference value.
+///
+/// This matches [`crate::FL_CHUNK_SIZE`] so that a block lines up exactly with a FastLanes
+/// bit-packing chunk, which lets decompression fuse the reference addition into unpacking.
+pub const BLOCK_SIZE: usize = crate::FL_CHUNK_SIZE;
+
+#[array_slots(crate::BlockedFoR)]
+pub struct BlockedFoRSlots {
+    /// The encoded array with the block-local frame-of-reference subtracted.
+    #[slot(0)]
+    pub encoded: ArrayRef,
+    /// One reference (minimum) value per [`BLOCK_SIZE`] values of `encoded`.
+    #[slot(1)]
+    pub references: ArrayRef,
+}
+
+/// Block-wise Frame of Reference (FoR) encoded array.
+///
+/// Where [`crate::FoR`] subtracts a single reference from the whole array, this encoding
+/// subtracts a separate reference from each [`BLOCK_SIZE`]-value block. Locally clustered data
+/// whose values drift over the array — timestamps, sorted keys, counters — then produce much
+/// smaller residuals, so the single bit width that [`crate::BitPacked`] picks for the whole
+/// array can be far narrower.
+#[derive(Clone, Debug)]
+pub struct BlockedFoRData {
+    /// The offset within the first block, created by slicing. `0 <= offset < BLOCK_SIZE`.
+    pub(super) offset: u16,
+}
+
+impl Display for BlockedFoRData {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        write!(f, "offset: {}", self.offset)
+    }
+}
+
+impl BlockedFoRData {
+    pub(crate) fn try_new(offset: u16) -> VortexResult<Self> {
+        vortex_ensure!(
+            (offset as usize) < BLOCK_SIZE,
+            "Offset must be less than the block size {BLOCK_SIZE}, got {offset}"
+        );
+        Ok(Self { offset })
+    }
+
+    /// The offset of the first value within its block.
+    #[inline]
+    pub fn offset(&self) -> u16 {
+        self.offset
+    }
+
+    #[inline]
+    pub fn ptype(&self, dtype: &DType) -> PType {
+        dtype.as_ptype()
+    }
+}
+
+/// The number of reference values needed to cover `len` values starting at `offset`.
+#[inline]
+pub(crate) fn num_blocks(len: usize, offset: u16) -> usize {
+    (len + offset as usize).div_ceil(BLOCK_SIZE)
+}
+
+pub trait BlockedFoRArrayExt: BlockedFoRArraySlotsExt {
+    /// The offset of the first value within its block, `0 <= offset < BLOCK_SIZE`.
+    #[inline]
+    fn offset(&self) -> u16 {
+        BlockedFoRData::offset(self)
+    }
+
+    /// The number of reference values held by this array.
+    #[inline]
+    fn num_blocks(&self) -> usize {
+        num_blocks(self.as_ref().len(), self.offset())
+    }
+}
+
+impl<T: TypedArrayRef<crate::BlockedFoR>> BlockedFoRArrayExt for T {}

encodings/fastlanes/src/blocked_for/mod.rs

@@ -0,0 +1,18 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+mod array;
+pub use array::BLOCK_SIZE;
+pub use array::BlockedFoRArrayExt;
+pub use array::BlockedFoRArraySlotsExt;
+pub use array::BlockedFoRData;
+pub use array::BlockedFoRSlots;
+pub use array::blocked_for_compress::BlockSummary;
+pub use array::blocked_for_compress::block_summary;
+
+#[cfg(test)]
+mod tests;
+
+mod vtable;
+pub use vtable::BlockedFoR;
+pub use vtable::BlockedFoRArray;