feat(array): migrate BitPacked encoding to scalar probes

Signed-off-by: Gunther Xing <jiengup@gmail.com>

Author: jiengup

encodings/fastlanes/Cargo.toml

@@ -59,6 +59,10 @@ harness = false
 name = "bitpack_blocked_decompress"
 harness = false
 
+[[bench]]
+name = "bitpacked_probe"
+harness = false
+
 [[bench]]
 name = "canonicalize_bench"
 harness = false

encodings/fastlanes/benches/bitpacked_probe.rs

@@ -0,0 +1,160 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+#![expect(clippy::unwrap_used)]
+
+use std::sync::LazyLock;
+
+use divan::Bencher;
+use mimalloc::MiMalloc;
+use vortex_array::ArrayRef;
+use vortex_array::IntoArray;
+use vortex_array::VortexSessionExecute;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::scalar::Scalar;
+use vortex_fastlanes::BitPacked;
+use vortex_fastlanes::BitPackedArrayExt;
+use vortex_fastlanes::BitPackedArraySlotsExt;
+use vortex_fastlanes::FoR;
+use vortex_fastlanes::bitpack_compress::bitpack_encode;
+use vortex_fastlanes::bitpack_compress::bitpack_encode_blocked;
+use vortex_session::VortexSession;
+
+#[global_allocator]
+static GLOBAL: MiMalloc = MiMalloc;
+
+fn main() {
+    LazyLock::force(&SESSION);
+    divan::main();
+}
+
+const LEN: usize = 65_536;
+const ACCESSES: usize = 256;
+const FIXTURES: &[(&str, bool)] = &[
+    ("global", false),
+    ("global", true),
+    ("blocked_primitive", false),
+    ("blocked_primitive", true),
+    ("blocked_for", false),
+    ("blocked_for", true),
+];
+const PATCHED_FIXTURES: &[&str] = &["global", "blocked_primitive", "blocked_for"];
+
+static SESSION: LazyLock<VortexSession> = LazyLock::new(|| {
+    let session = vortex_array::array_session();
+    vortex_fastlanes::initialize(&session);
+    session
+});
+
+fn fixture(mode: &str, patched: bool) -> ArrayRef {
+    let widths: Vec<u8> = (0..LEN / 1024).map(|block| 4 + (block % 4) as u8).collect();
+    let values = PrimitiveArray::from_iter((0..LEN).map(|index| {
+        let width = widths[index / 1024];
+        if patched && index % 64 == 0 {
+            65_535u32
+        } else {
+            ((index * 7919) as u32) & ((1u32 << width) - 1)
+        }
+    }));
+    let mut ctx = SESSION.create_execution_ctx();
+    let encoded = if mode == "global" {
+        bitpack_encode(&values, 7, None, &mut ctx).unwrap()
+    } else {
+        bitpack_encode_blocked(&values, &widths, None, &mut ctx).unwrap()
+    };
+    assert_eq!(encoded.patches().is_some(), patched);
+    if mode == "blocked_for" {
+        let offsets = encoded.block_offsets().unwrap().clone();
+        let reference = Scalar::zero_value(offsets.dtype());
+        let offsets = FoR::try_new(offsets, reference).unwrap().into_array();
+        BitPacked::try_new_with_block_offsets(
+            encoded.packed().clone(),
+            encoded.dtype().as_ptype(),
+            encoded.validity().unwrap(),
+            encoded.patches(),
+            offsets,
+            encoded.len(),
+            encoded.offset(),
+        )
+        .unwrap()
+        .into_array()
+    } else {
+        encoded.into_array()
+    }
+}
+
+fn random_index(read: usize) -> usize {
+    (read * 7919) % LEN
+}
+
+#[divan::bench(args = FIXTURES)]
+fn random_once(bencher: Bencher, (mode, patched): (&str, bool)) {
+    let array = fixture(mode, patched);
+    bencher
+        .with_inputs(|| SESSION.create_execution_ctx())
+        .bench_refs(|ctx| {
+            for read in 0..ACCESSES {
+                divan::black_box(array.execute_scalar(random_index(read), ctx).unwrap());
+            }
+        });
+}
+
+#[divan::bench(args = FIXTURES)]
+fn random_repeated(bencher: Bencher, (mode, patched): (&str, bool)) {
+    let array = fixture(mode, patched);
+    bencher
+        .with_inputs(|| (SESSION.create_execution_ctx(), array.repeated_probe()))
+        .bench_refs(|(ctx, probe)| {
+            for read in 0..ACCESSES {
+                divan::black_box(probe.execute_scalar(random_index(read), ctx).unwrap());
+            }
+        });
+}
+
+#[divan::bench(args = FIXTURES)]
+fn same_block_once(bencher: Bencher, (mode, patched): (&str, bool)) {
+    let array = fixture(mode, patched);
+    bencher
+        .with_inputs(|| SESSION.create_execution_ctx())
+        .bench_refs(|ctx| {
+            for read in 0..ACCESSES {
+                divan::black_box(array.execute_scalar(1024 + read, ctx).unwrap());
+            }
+        });
+}
+
+#[divan::bench(args = FIXTURES)]
+fn same_block_repeated(bencher: Bencher, (mode, patched): (&str, bool)) {
+    let array = fixture(mode, patched);
+    bencher
+        .with_inputs(|| (SESSION.create_execution_ctx(), array.repeated_probe()))
+        .bench_refs(|(ctx, probe)| {
+            for read in 0..ACCESSES {
+                divan::black_box(probe.execute_scalar(1024 + read, ctx).unwrap());
+            }
+        });
+}
+
+#[divan::bench(args = PATCHED_FIXTURES)]
+fn patch_hit_repeated(bencher: Bencher, mode: &str) {
+    let array = fixture(mode, true);
+    bencher
+        .with_inputs(|| (SESSION.create_execution_ctx(), array.repeated_probe()))
+        .bench_refs(|(ctx, probe)| {
+            for _ in 0..ACCESSES {
+                divan::black_box(probe.execute_scalar(1024, ctx).unwrap());
+            }
+        });
+}
+
+#[divan::bench(args = PATCHED_FIXTURES)]
+fn patch_miss_repeated(bencher: Bencher, mode: &str) {
+    let array = fixture(mode, true);
+    bencher
+        .with_inputs(|| (SESSION.create_execution_ctx(), array.repeated_probe()))
+        .bench_refs(|(ctx, probe)| {
+            for _ in 0..ACCESSES {
+                divan::black_box(probe.execute_scalar(1025, ctx).unwrap());
+            }
+        });
+}

encodings/fastlanes/src/bitpacking/array/bitpack_decompress/blocked.rs

@@ -12,6 +12,7 @@ use num_traits::AsPrimitive;
 use vortex_array::ArrayRef;
 use vortex_array::ArrayView;
 use vortex_array::ExecutionCtx;
+use vortex_array::ProbeChildren;
 use vortex_array::arrays::Primitive;
 use vortex_array::arrays::PrimitiveArray;
 use vortex_array::builders::ArrayBuilder;
@@ -21,12 +22,14 @@ 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_ensure;
 
 use super::global::apply_patches_to_uninit_range;
 use crate::BitPacked;
 use crate::BitPackedArrayExt;
+use crate::BitPackedSlots;
 use crate::FL_CHUNK_SIZE;
 use crate::bitpacking::array::block_range;
 use crate::bitpacking::array::validate_primitive_offsets;
@@ -223,17 +226,8 @@ pub fn unpack_single_blocked(
         block + 1 < offsets.len(),
         "BitPacked index {index} has no block boundaries"
     );
-    let (base, start, end) = if let Some(primitive) = offsets.as_opt::<Primitive>()
-        && primitive.buffer_handle().is_on_host()
-    {
-        match_each_unsigned_integer_ptype!(primitive.ptype(), |I| {
-            let offsets = primitive.as_slice::<I>();
-            (
-                AsPrimitive::<u64>::as_(offsets[0]),
-                AsPrimitive::<u64>::as_(offsets[block]),
-                AsPrimitive::<u64>::as_(offsets[block + 1]),
-            )
-        })
+    let (base, start, end) = if let Some(bounds) = primitive_block_bounds(offsets, block) {
+        bounds
     } else {
         let base = u64::try_from(&offsets.execute_scalar(0, ctx)?)?;
         let start = if block == 0 {
@@ -251,6 +245,104 @@ pub fn unpack_single_blocked(
         array.dtype().as_ptype().bit_width() as u64,
         array.packed().len(),
     )?;
+    unpack_single_in_range(array, index_in_encoded, range)
+}
+
+/// Read only the selected block's boundaries through the parent's stable child slot.
+pub(crate) fn unpack_single_blocked_with_probes(
+    array: ArrayView<'_, BitPacked>,
+    offsets: &ArrayRef,
+    index: usize,
+    children: &mut ProbeChildren<'_, '_, BitPacked>,
+    mut cache: Option<&mut BlockProbeCache>,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<Scalar> {
+    let index_in_encoded = index + array.offset() as usize;
+    let block = index_in_encoded / FL_CHUNK_SIZE;
+    if let Some(cached) = cache.as_ref().and_then(|entry| entry.last)
+        && cached.block == block
+    {
+        return unpack_single_in_range(array, index_in_encoded, cached.start..cached.end);
+    }
+    vortex_ensure!(
+        block + 1 < offsets.len(),
+        "BitPacked index {index} has no block boundaries"
+    );
+    let (base, start, end) = if let Some(bounds) = primitive_block_bounds(offsets, block) {
+        bounds
+    } else {
+        let mut probe = children
+            .slot(BitPackedSlots::BLOCK_OFFSETS)?
+            .vortex_expect("blocked BitPacked must have block offsets");
+        let base = if let Some(base) = cache.as_ref().and_then(|entry| entry.base) {
+            base
+        } else {
+            let base = u64::try_from(&probe.execute_scalar(0, ctx)?)?;
+            if let Some(cache) = cache.as_deref_mut() {
+                cache.base = Some(base);
+            }
+            base
+        };
+        let start = if block == 0 {
+            base
+        } else {
+            u64::try_from(&probe.execute_scalar(block, ctx)?)?
+        };
+        let end = u64::try_from(&probe.execute_scalar(block + 1, ctx)?)?;
+        (base, start, end)
+    };
+    let range = block_range(
+        base,
+        start,
+        end,
+        array.dtype().as_ptype().bit_width() as u64,
+        array.packed().len(),
+    )?;
+    if let Some(cache) = cache {
+        cache.last = Some(ValidatedBlock {
+            block,
+            start: range.start,
+            end: range.end,
+        });
+    }
+    unpack_single_in_range(array, index_in_encoded, range)
+}
+
+/// Preparation retained for one immutable BitPacked array.
+#[derive(Default)]
+pub(crate) struct BlockProbeCache {
+    base: Option<u64>,
+    last: Option<ValidatedBlock>,
+}
+
+/// A validated byte range for one block of an immutable BitPacked array.
+#[derive(Clone, Copy)]
+struct ValidatedBlock {
+    block: usize,
+    start: usize,
+    end: usize,
+}
+
+fn primitive_block_bounds(offsets: &ArrayRef, block: usize) -> Option<(u64, u64, u64)> {
+    let primitive = offsets.as_opt::<Primitive>()?;
+    if !primitive.buffer_handle().is_on_host() {
+        return None;
+    }
+    Some(match_each_unsigned_integer_ptype!(primitive.ptype(), |I| {
+        let offsets = primitive.as_slice::<I>();
+        (
+            AsPrimitive::<u64>::as_(offsets[0]),
+            AsPrimitive::<u64>::as_(offsets[block]),
+            AsPrimitive::<u64>::as_(offsets[block + 1]),
+        )
+    }))
+}
+
+fn unpack_single_in_range(
+    array: ArrayView<'_, BitPacked>,
+    index_in_encoded: usize,
+    range: Range<usize>,
+) -> VortexResult<Scalar> {
     let bit_width = range.len() / 128;
     match_each_integer_ptype!(array.dtype().as_ptype(), |P| {
         let packed = array.packed_slice::<<P as PhysicalPType>::Physical>();

encodings/fastlanes/src/bitpacking/array/bitpack_decompress/mod.rs

@@ -9,6 +9,8 @@ pub(crate) use blocked::for_each_packed_block;
 pub use blocked::unpack_array_blocked;
 pub(crate) use blocked::unpack_into_primitive_builder_blocked;
 pub use blocked::unpack_single_blocked;
+pub(crate) use blocked::unpack_single_blocked_with_probes;
+pub(crate) use blocked::BlockProbeCache;
 pub use global::count_exceptions;
 pub use global::unpack_array;
 pub(crate) use global::unpack_into_primitive_builder;