Blocked FoR with one reference per 1024-element chunk (#10105)

Stacked on #10104.

## Summary

The in-memory `FoR` array now allows one reference per 1024-element
chunk instead of a single global reference per array.

This is supported by adding a `references` child that records each
reference. Element `i` decodes as `encoded[i] + references[(offset + i)
/ 1024]` with wrapping arithmetic.

This also requires adding an `offset` pointer. This is the position of
the first element within the first chunk and only matters for sliced
arrays.

The wire format is unchanged. `fastlanes.for` reads its single reference
as a `ConstantArray` child, and `FoRPlugin` writes `fastlanes.for` only
when the references are constant.

Note that FoR encoding always uses a constant reference per array still
and serialization refuses arrays that have non-constant references, so
this change will not lead to breaks.

## Example

### References

A 3000-row array spans three chunks, each with its own reference.
`encoded` holds each value minus its chunk's reference.

```text
rows         0 ─────────── 1023 │ 1024 ────────── 2047 │ 2048 ──────── 2999
references   [    1_000_000     │      5_000_000       │     9_000_000     ]
encoded      [  small values    │    small values      │   small values    ]

value[1500] = encoded[1500] + references[1500 / 1024]
            = encoded[1500] + references[1]
            = encoded[1500] + 5_000_000
```

### Offset

Slicing is zero-copy, so a slice can start partway through a chunk.
`offset` records how far into its first chunk the slice starts.

```text
             │◄──── chunk 0 ────►│◄──── chunk 1 ────►│◄──── chunk 2 ────►│
rows         0                 1024      1500      2048      2500      3000
slice                                     [═══════════════════)
                                   │◄ 476 ►│

slice(1500..2500):
  encoded    = encoded.slice(1500..2500)          1000 rows
  references = references.slice(1..3)             [5_000_000, 9_000_000]
  offset     = 1500 % 1024                        476

slice row i uses references[(476 + i) / 1024]:
  i = 0..=547    →  references[0] = 5_000_000     (original chunk 1)
  i = 548..=999  →  references[1] = 9_000_000     (original chunk 2)
```

The new offset is always below 1024. Slicing a slice adds the offsets
and slices the references again.

### Constant references (existing files)

```text
fastlanes.for on disk      metadata = reference 42, children = [encoded]
        │ read                                  ▲ write (only when references are constant)
        ▼                                       │
in memory                  references = ConstantArray(42, len = num_chunks), offset = 0
```

## API

- `FoRSlots` gains `references`.
- `FoRData` holds only `offset: u16`. The stored reference scalar is
gone.
- `FoRArrayExt::reference_scalar() -> &Scalar` is replaced by
`constant_reference() -> Option<Scalar>`, which is `Some` when
`references` is a `ConstantArray`. `offset()` and `ptype()` move onto
the extension trait.
- `FoR::try_new(encoded, reference)` keeps its signature and builds the
constant child. `FoR::try_new_chunked(encoded, references, offset)`
accepts any references.

## Kernels

With constant references, every kernel takes its existing path. That
includes Eq/NotEq compare pushdown, is_constant, is_sorted, take,
filter, the fused BitPacked decode, and CUDA FFOR and dynamic dispatch.

With varying references:
- Decoding adds each chunk's reference in place. `scalar_at` looks up
the element's chunk reference. `slice` slices the references by chunk
and carries the new offset. `cast` casts the references.
- Compare, is_constant, is_sorted, take and filter return `None` and
fall back to decoding.
- CUDA decodes these arrays on the CPU and never fuses them into a
dispatch plan.

Callers that rewrap a FoR's encoded child (the btrblocks FoR scheme,
benches, CUDA tests) now pass through `references()` and `offset()` via
`try_new_chunked`.

---------

Signed-off-by: Matt Katz <mhkatz97@gmail.com>

Author: mhk197

encodings/fastlanes/src/for/array/for_compress.rs

@@ -83,7 +83,10 @@ mod test {
         let mut ctx = SESSION.create_execution_ctx();
         let array = PrimitiveArray::new((1i32..10).collect::<Buffer<_>>(), Validity::NonNullable);
         let compressed = FoRData::encode(array.clone(), &mut ctx).unwrap();
-        assert_eq!(i32::try_from(compressed.reference_scalar()).unwrap(), 1);
+        assert_eq!(
+            i32::try_from(&compressed.constant_reference().unwrap()).unwrap(),
+            1
+        );
 
         assert_arrays_eq!(compressed, array, &mut ctx);
     }
@@ -98,7 +101,7 @@ mod test {
         );
         let compressed = FoRData::encode(array, &mut ctx).unwrap();
         assert_eq!(
-            u32::try_from(compressed.reference_scalar()).unwrap(),
+            u32::try_from(&compressed.constant_reference().unwrap()).unwrap(),
             1_000_000u32
         );
     }
@@ -111,8 +114,14 @@ mod test {
 
         let dtype = array.dtype().clone();
         let compressed = FoRData::encode(array, &mut ctx).unwrap();
-        assert_eq!(compressed.reference_scalar().dtype(), &dtype);
-        assert!(compressed.reference_scalar().dtype().is_signed_int());
+        assert_eq!(compressed.constant_reference().unwrap().dtype(), &dtype);
+        assert!(
+            compressed
+                .constant_reference()
+                .unwrap()
+                .dtype()
+                .is_signed_int()
+        );
         assert!(compressed.encoded().dtype().is_signed_int());
 
         let encoded = compressed.encoded().execute_scalar(0, &mut ctx).unwrap();
@@ -160,7 +169,8 @@ mod test {
         assert_eq!(
             i8::MIN,
             compressed
-                .reference_scalar()
+                .constant_reference()
+                .unwrap()
                 .as_primitive()
                 .typed_value::<i8>()
                 .unwrap()

encodings/fastlanes/src/for/array/for_decompress.rs

@@ -1,9 +1,11 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
+use std::iter;
 use std::mem::MaybeUninit;
 
 use fastlanes::FoR;
+use itertools::Itertools;
 use num_traits::PrimInt;
 use num_traits::WrappingAdd;
 use vortex_array::ArrayView;
@@ -15,9 +17,11 @@ use vortex_array::dtype::PhysicalPType;
 use vortex_array::dtype::UnsignedPType;
 use vortex_array::match_each_integer_ptype;
 use vortex_array::match_each_unsigned_integer_ptype;
+use vortex_array::scalar::Scalar;
 use vortex_buffer::Buffer;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
+use vortex_error::vortex_err;
 
 use crate::BitPacked;
 use crate::BitPackedArrayExt;
@@ -51,13 +55,27 @@ impl<T: PhysicalPType<Physical = T> + FoR> UnpackStrategy<T> for FoRStrategy<T>
 }
 
 pub fn decompress(array: &FoRArray, ctx: &mut ExecutionCtx) -> VortexResult<PrimitiveArray> {
+    match array.constant_reference() {
+        Some(reference) => decompress_one_ref(array, &reference, ctx),
+        None => {
+            match_each_integer_ptype!(array.ptype(), |T| { decompress_many_refs::<T>(array, ctx) })
+        }
+    }
+}
+
+/// Decompress an array whose chunks all share `reference`.
+fn decompress_one_ref(
+    array: &FoRArray,
+    reference: &Scalar,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<PrimitiveArray> {
     let ptype = array.ptype();
 
     // Try to do fused unpack.
-    if array.reference_scalar().dtype().is_unsigned_int()
+    if ptype.is_unsigned_int()
         && let Some(bp) = array.encoded().as_opt::<BitPacked>()
     {
-        return match_each_unsigned_integer_ptype!(array.ptype(), |T| {
+        return match_each_unsigned_integer_ptype!(ptype, |T| {
             fused_decompress::<T>(array, bp, ctx)
         });
     }
@@ -67,8 +85,7 @@ pub fn decompress(array: &FoRArray, ctx: &mut ExecutionCtx) -> VortexResult<Prim
     let validity = encoded.validity()?;
 
     Ok(match_each_integer_ptype!(ptype, |T| {
-        let min = array
-            .reference_scalar()
+        let min = reference
             .as_primitive()
             .typed_value::<T>()
             .vortex_expect("reference must be non-null");
@@ -83,6 +100,32 @@ pub fn decompress(array: &FoRArray, ctx: &mut ExecutionCtx) -> VortexResult<Prim
     }))
 }
 
+/// Decompress an array whose chunks have different references.
+fn decompress_many_refs<T: NativePType + WrappingAdd + PrimInt>(
+    array: &FoRArray,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<PrimitiveArray> {
+    let encoded = array.encoded().clone().execute::<PrimitiveArray>(ctx)?;
+    if encoded.is_empty() {
+        return Ok(encoded);
+    }
+    let validity = encoded.validity()?;
+    let references = array.references().clone().execute::<PrimitiveArray>(ctx)?;
+    let references = references.as_slice::<T>();
+
+    // The first chunk may be partial when the array was sliced.
+    let first_len = (FL_CHUNK_SIZE - usize::from(array.offset())).min(array.len());
+    let mut values = encoded.into_buffer_mut::<T>();
+    let (first, rest) = values.as_mut_slice().split_at_mut(first_len);
+    let chunks = iter::once(first).chain(rest.chunks_mut(FL_CHUNK_SIZE));
+    for (chunk, &reference) in chunks.zip_eq(references) {
+        for value in chunk {
+            *value = value.wrapping_add(&reference);
+        }
+    }
+    Ok(PrimitiveArray::new(values.freeze(), validity))
+}
+
 pub(crate) fn fused_decompress<
     T: PhysicalPType<Physical = T> + UnsignedPType + FoR + WrappingAdd,
 >(
@@ -91,7 +134,8 @@ pub(crate) fn fused_decompress<
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<PrimitiveArray> {
     let ref_ = for_
-        .reference_scalar()
+        .constant_reference()
+        .ok_or_else(|| vortex_err!("fused FoR decompression requires a constant reference"))?
         .as_primitive()
         .as_::<T>()
         .vortex_expect("cannot be null");
@@ -110,7 +154,7 @@ pub(crate) fn fused_decompress<
     )?;
 
     let mut builder = PrimitiveBuilder::<T>::with_capacity_in(
-        for_.reference_scalar().dtype().nullability(),
+        for_.dtype().nullability(),
         bp.len(),
         ctx.allocator(),
     );

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

@@ -12,52 +12,71 @@ use vortex_array::scalar::Scalar;
 use vortex_error::VortexResult;
 use vortex_error::vortex_ensure;
 
+use crate::FL_CHUNK_SIZE;
+
 pub mod for_compress;
 pub mod for_decompress;
 
 #[array_slots(crate::FoR)]
 pub struct FoRSlots {
-    /// The encoded array with the frame-of-reference (minimum value) subtracted.
+    /// The encoded array with each chunk's reference subtracted.
     #[slot(0)]
     pub encoded: ArrayRef,
+    /// One reference per [`FL_CHUNK_SIZE`]-element chunk.
+    #[slot(1)]
+    pub references: ArrayRef,
 }
 
 /// Frame of Reference (FoR) encoded array.
 ///
 /// This encoding stores values as offsets from a reference value, which can significantly reduce
 /// storage requirements when values are clustered around a specific point.
+///
+/// Every [`FL_CHUNK_SIZE`]-element chunk has its own reference: element `i` decodes as
+/// `encoded[i] + references[(offset + i) / FL_CHUNK_SIZE]` with wrapping arithmetic, where
+/// `offset` is the position of the first element within the first chunk. Arrays with a single
+/// reference store a constant `references` child.
 #[derive(Clone, Debug)]
 pub struct FoRData {
-    pub(super) reference: Scalar,
+    pub(super) offset: u16,
 }
 
 pub trait FoRArrayExt: FoRArraySlotsExt {
-    fn reference_scalar(&self) -> &Scalar {
-        &self.reference
+    /// The reference shared by every chunk, if the references are constant.
+    fn constant_reference(&self) -> Option<Scalar> {
+        self.references().as_constant()
+    }
+
+    /// The position of the first element within the first chunk of `references`.
+    fn offset(&self) -> u16 {
+        self.offset
+    }
+
+    #[inline]
+    fn ptype(&self) -> PType {
+        self.as_ref().dtype().as_ptype()
     }
 }
 
 impl<T: TypedArrayRef<crate::FoR>> FoRArrayExt for T {}
 
 impl Display for FoRData {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "reference: {}", self.reference)
+        write!(f, "offset: {}", self.offset)
     }
 }
 
 impl FoRData {
-    pub(crate) fn try_new(reference: Scalar) -> VortexResult<Self> {
-        vortex_ensure!(!reference.is_null(), "Reference value cannot be null");
+    pub(crate) fn try_new(offset: u16) -> VortexResult<Self> {
         vortex_ensure!(
-            reference.dtype().is_int(),
-            "FoR requires an integer reference dtype, got {}",
-            reference.dtype()
+            usize::from(offset) < FL_CHUNK_SIZE,
+            "FoR offset must be less than {FL_CHUNK_SIZE}, got {offset}"
         );
-        Ok(Self { reference })
+        Ok(Self { offset })
     }
+}
 
-    #[inline]
-    pub fn ptype(&self) -> PType {
-        self.reference.dtype().as_ptype()
-    }
+/// The number of chunks spanned by `len` elements starting at `offset` within the first chunk.
+pub(crate) fn num_chunks(offset: u16, len: usize) -> usize {
+    (usize::from(offset) + len).div_ceil(FL_CHUNK_SIZE)
 }

encodings/fastlanes/src/for/compute/cast.rs

@@ -21,10 +21,11 @@ impl CastReduce for FoR {
 
         // For type changes between integers, cast the components
         let casted_child = array.encoded().cast(dtype.clone())?;
-        let casted_reference = array.reference_scalar().cast(dtype)?;
+        // References are always non-nullable.
+        let casted_references = array.references().cast(dtype.as_nonnullable())?;
 
         Ok(Some(
-            FoR::try_new(casted_child, casted_reference)?.into_array(),
+            FoR::try_new_chunked(casted_child, casted_references, array.offset())?.into_array(),
         ))
     }
 }

encodings/fastlanes/src/for/compute/compare.rs

@@ -69,7 +69,11 @@ where
         return Ok(None);
     }
 
-    let reference = lhs.reference_scalar();
+    // TODO(mk): support many references.
+
+    let Some(reference) = lhs.constant_reference() else {
+        return Ok(None);
+    };
     let reference = reference.as_primitive().typed_value::<T>();
 
     // We encode the RHS into the FoR domain.

encodings/fastlanes/src/for/compute/is_constant.rs

@@ -11,6 +11,7 @@ use vortex_array::scalar::Scalar;
 use vortex_error::VortexResult;
 
 use crate::FoR;
+use crate::r#for::array::FoRArrayExt;
 use crate::r#for::array::FoRArraySlotsExt;
 
 /// FoR-specific is_constant kernel.
@@ -33,6 +34,10 @@ impl DynAggregateKernel for FoRIsConstantKernel {
         let Some(array) = batch.as_opt::<FoR>() else {
             return Ok(None);
         };
+        // TODO(mk): support many references.
+        if array.constant_reference().is_none() {
+            return Ok(None);
+        }
 
         let result = is_constant(array.encoded(), ctx)?;
         Ok(Some(IsConstant::make_partial(batch, result, ctx)?))

encodings/fastlanes/src/for/compute/is_sorted.rs

@@ -14,6 +14,7 @@ use vortex_array::scalar::Scalar;
 use vortex_error::VortexResult;
 
 use crate::FoR;
+use crate::r#for::array::FoRArrayExt;
 use crate::r#for::array::FoRArraySlotsExt;
 
 #[derive(Debug)]
@@ -33,6 +34,10 @@ impl DynAggregateKernel for FoRIsSortedKernel {
         let Some(array) = batch.as_opt::<FoR>() else {
             return Ok(None);
         };
+        // TODO(mk): support many references.
+        if array.constant_reference().is_none() {
+            return Ok(None);
+        }
 
         let encoded = array.encoded().clone().execute::<PrimitiveArray>(ctx)?;
         let unsigned_array = PrimitiveArray::from_buffer_handle(

encodings/fastlanes/src/for/compute/mod.rs

@@ -25,23 +25,23 @@ impl TakeExecute for FoR {
         indices: &ArrayRef,
         _ctx: &mut ExecutionCtx,
     ) -> VortexResult<Option<ArrayRef>> {
+        // TODO(mk): support many references.
+        let Some(reference) = array.constant_reference() else {
+            return Ok(None);
+        };
         Ok(Some(
-            FoR::try_new(
-                array.encoded().take(indices.clone())?,
-                array.reference_scalar().clone(),
-            )?
-            .into_array(),
+            FoR::try_new(array.encoded().take(indices.clone())?, reference)?.into_array(),
         ))
     }
 }
 
 impl FilterReduce for FoR {
     fn filter(array: ArrayView<'_, Self>, mask: &Mask) -> VortexResult<Option<ArrayRef>> {
-        FoR::try_new(
-            array.encoded().filter(mask.clone())?,
-            array.reference_scalar().clone(),
-        )
-        .map(|a| Some(a.into_array()))
+        // TODO(mk): support many references.
+        let Some(reference) = array.constant_reference() else {
+            return Ok(None);
+        };
+        FoR::try_new(array.encoded().filter(mask.clone())?, reference).map(|a| Some(a.into_array()))
     }
 }
 

encodings/fastlanes/src/for/mod.rs

@@ -9,6 +9,9 @@ pub use array::FoRSlots;
 
 pub(crate) mod compute;
 
+#[cfg(test)]
+mod tests;
+
 mod plugin;
 pub use plugin::FoRPlugin;