Add PatchesV2 with chunk-local patch indices

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

@@ -2,6 +2,10 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 use std::mem::MaybeUninit;
+use std::sync::LazyLock;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::AtomicU64;
+use std::sync::atomic::Ordering;
 
 use fastlanes::BitPacking;
 use itertools::Itertools;
@@ -16,6 +20,7 @@ use vortex_array::dtype::NativePType;
 use vortex_array::match_each_integer_ptype;
 use vortex_array::match_each_unsigned_integer_ptype;
 use vortex_array::patches::Patches;
+use vortex_array::patches_v2::PatchesV2;
 use vortex_array::scalar::Scalar;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
@@ -142,11 +147,20 @@ pub(crate) fn apply_patches_to_uninit_range<S: NativePType, T: NativePType, F: F
 ) -> VortexResult<()> {
     assert_eq!(patches.array_len(), dst.len());
 
-    let indices = patches.indices().clone().execute::<PrimitiveArray>(ctx)?;
     let values = patches.values().clone().execute::<PrimitiveArray>(ctx)?;
     assert!(values.all_valid(ctx)?, "Patch values must be all valid");
     let values = values.as_slice::<S>();
 
+    // When enabled, chunked patch sets scatter through the chunk-local PatchesV2 form to
+    // exercise it on the real decompression path. Converting per decompression costs a pass and
+    // allocations over the patch set, so this stays opt-in until the stored layout is
+    // chunk-local; the default path below is unchanged and the branch stays out of line to keep
+    // it out of the hot scatter loop's codegen.
+    if use_patches_v2_scatter() && patches.chunk_offsets().is_some() {
+        return apply_patches_v2(dst, patches, values, ctx, f);
+    }
+
+    let indices = patches.indices().clone().execute::<PrimitiveArray>(ctx)?;
     match_each_unsigned_integer_ptype!(indices.ptype(), |P| {
         for (index, &value) in indices.as_slice::<P>().iter().zip_eq(values) {
             dst.set_value(
@@ -158,6 +172,50 @@ pub(crate) fn apply_patches_to_uninit_range<S: NativePType, T: NativePType, F: F
     Ok(())
 }
 
+/// Scatter patch values through the chunk-local [`PatchesV2`] form.
+#[cold]
+#[inline(never)]
+fn apply_patches_v2<S: NativePType, T: NativePType, F: Fn(S) -> T>(
+    dst: &mut UninitRange<T>,
+    patches: &Patches,
+    values: &[S],
+    ctx: &mut ExecutionCtx,
+    f: F,
+) -> VortexResult<()> {
+    let v2 = PatchesV2::from_patches(patches, ctx)?;
+    v2.apply_each(ctx, |logical, ordinal| {
+        dst.set_value(logical, f(values[ordinal]));
+    })?;
+    PATCHES_V2_APPLIES.fetch_add(1, Ordering::Relaxed);
+    Ok(())
+}
+
+static PATCHES_V2_APPLIES: AtomicU64 = AtomicU64::new(0);
+static PATCHES_V2_SCATTER: AtomicBool = AtomicBool::new(false);
+
+/// Returns whether decompression scatters chunked patches through [`PatchesV2`].
+///
+/// Enabled by [`force_patches_v2_scatter`] or the `VORTEX_PATCHES_V2_SCATTER=1` environment
+/// variable.
+pub fn use_patches_v2_scatter() -> bool {
+    static FROM_ENV: LazyLock<bool> =
+        LazyLock::new(|| std::env::var("VORTEX_PATCHES_V2_SCATTER").is_ok_and(|v| v == "1"));
+    PATCHES_V2_SCATTER.load(Ordering::Relaxed) || *FROM_ENV
+}
+
+/// Force the chunk-local patch scatter on or off for this process.
+pub fn force_patches_v2_scatter(enabled: bool) {
+    PATCHES_V2_SCATTER.store(enabled, Ordering::Relaxed);
+}
+
+/// The number of decompressions that scattered patches through [`PatchesV2`].
+///
+/// This instruments the chunk-local patch path so integration tests can assert real read paths
+/// exercise it.
+pub fn patches_v2_apply_count() -> u64 {
+    PATCHES_V2_APPLIES.load(Ordering::Relaxed)
+}
+
 pub fn unpack_single(array: ArrayView<'_, BitPacked>, index: usize) -> Scalar {
     let bit_width = array.bit_width() as usize;
     let ptype = array.dtype().as_ptype();

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

@@ -385,11 +385,17 @@ mod test {
 
         let packed_with_patches = BitPackedData::encode(&parray, 9, &mut ctx).unwrap();
         assert!(packed_with_patches.patches().is_some());
+        crate::bitpack_decompress::force_patches_v2_scatter(true);
+        let applies_before = crate::bitpack_decompress::patches_v2_apply_count();
         let packed_primitive = packed_with_patches
             .as_array()
             .clone()
             .execute::<PrimitiveArray>(&mut ctx)
             .unwrap();
+        assert!(
+            crate::bitpack_decompress::patches_v2_apply_count() > applies_before,
+            "expected the chunk-local patch path"
+        );
         assert_arrays_eq!(
             packed_primitive,
             PrimitiveArray::new(values, vortex_array::validity::Validity::NonNullable),

vortex-array/benches/patches_lookup.rs

@@ -9,8 +9,11 @@ use rand::RngExt;
 use rand::SeedableRng;
 use rand::rngs::StdRng;
 use vortex_array::IntoArray;
+use vortex_array::VortexSessionExecute;
+use vortex_array::array_session;
 use vortex_array::patches::PATCH_CHUNK_SIZE;
 use vortex_array::patches::Patches;
+use vortex_array::patches_v2::PatchesV2;
 use vortex_buffer::Buffer;
 
 fn main() {
@@ -165,3 +168,116 @@ fn search_index_full_range_random(bencher: Bencher) {
 fn search_index_full_range_random_chunked(bencher: Bencher) {
     bench_search_index(bencher, full_range_patches(true), queries_full_range());
 }
+
+fn patches_v2_from(patches: &Patches) -> PatchesV2 {
+    let mut ctx = array_session().create_execution_ctx();
+    PatchesV2::from_patches(patches, &mut ctx).unwrap()
+}
+
+fn bench_search_index_v2(bencher: Bencher, patches: PatchesV2, queries: Vec<usize>) {
+    let mut ctx = array_session().create_execution_ctx();
+    bencher
+        .with_inputs(|| (&patches, &queries))
+        .bench_local_refs(|(patches, queries)| {
+            for &q in queries.iter() {
+                divan::black_box(patches.search_index(q, &mut ctx).unwrap());
+            }
+        });
+}
+
+#[divan::bench]
+fn search_index_below_min_v2(bencher: Bencher) {
+    bench_search_index_v2(
+        bencher,
+        patches_v2_from(&narrow_band_patches(false)),
+        queries_below_min(),
+    );
+}
+
+#[divan::bench]
+fn search_index_above_max_v2(bencher: Bencher) {
+    bench_search_index_v2(
+        bencher,
+        patches_v2_from(&narrow_band_patches(false)),
+        queries_above_max(),
+    );
+}
+
+#[divan::bench]
+fn search_index_mixed_out_of_range_v2(bencher: Bencher) {
+    bench_search_index_v2(
+        bencher,
+        patches_v2_from(&narrow_band_patches(false)),
+        queries_mixed_out_of_range(),
+    );
+}
+
+#[divan::bench]
+fn search_index_in_range_v2(bencher: Bencher) {
+    bench_search_index_v2(
+        bencher,
+        patches_v2_from(&narrow_band_patches(false)),
+        queries_in_range(),
+    );
+}
+
+#[divan::bench]
+fn search_index_full_range_random_v2(bencher: Bencher) {
+    bench_search_index_v2(
+        bencher,
+        patches_v2_from(&full_range_patches(false)),
+        queries_full_range(),
+    );
+}
+
+fn bench_apply_v1(bencher: Bencher, patches: Patches) {
+    let mut ctx = array_session().create_execution_ctx();
+    let indices = patches
+        .indices()
+        .clone()
+        .execute::<vortex_array::arrays::PrimitiveArray>(&mut ctx)
+        .unwrap();
+    bencher
+        .with_inputs(|| vec![0i64; ARRAY_LEN])
+        .bench_local_values(|mut dst| {
+            for &index in indices.as_slice::<u64>() {
+                dst[index as usize] = 1;
+            }
+            divan::black_box(dst);
+        });
+}
+
+fn bench_apply_v2(bencher: Bencher, patches: PatchesV2) {
+    let mut ctx = array_session().create_execution_ctx();
+    bencher
+        .with_inputs(|| vec![0i64; ARRAY_LEN])
+        .bench_local_values(|mut dst| {
+            patches
+                .apply_each(&mut ctx, |logical, _ordinal| dst[logical] = 1)
+                .unwrap();
+            divan::black_box(dst);
+        });
+}
+
+#[divan::bench]
+fn apply_full_range(bencher: Bencher) {
+    bench_apply_v1(bencher, full_range_patches(false));
+}
+
+#[divan::bench]
+fn apply_full_range_v2(bencher: Bencher) {
+    bench_apply_v2(bencher, patches_v2_from(&full_range_patches(false)));
+}
+
+#[divan::bench]
+fn slice_unaligned(bencher: Bencher) {
+    let patches = full_range_patches(true);
+    bencher.bench(|| divan::black_box(patches.slice(1_000..900_000).unwrap()));
+}
+
+#[divan::bench]
+fn slice_unaligned_v2(bencher: Bencher) {
+    let patches = patches_v2_from(&full_range_patches(true));
+    let mut ctx = array_session().create_execution_ctx();
+    bencher.bench_local(|| divan::black_box(patches.slice(1_000..900_000, &mut ctx).unwrap()));
+}

vortex-array/src/lib.rs

@@ -133,6 +133,7 @@ pub mod normalize;
 pub mod optimizer;
 mod partial_ord;
 pub mod patches;
+pub mod patches_v2;
 pub mod scalar;
 pub mod scalar_fn;
 pub mod search_sorted;