Add PatchesV2 with chunk-local patch indices

.github/workflows/sql-bench-matrix.yml

@@ -102,6 +102,7 @@ jobs:
     timeout-minutes: 120
     env:
       VORTEX_EXPERIMENTAL_PATCHED_ARRAY: "1"
+      VORTEX_PATCHES_V2_SCATTER: "1"
       FLAT_LAYOUT_INLINE_ARRAY_NODE: "1"
       # Makes python output nicer
       COLUMNS: 120

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

@@ -16,6 +16,8 @@ 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::patches_v2::use_patches_v2_scatter;
 use vortex_array::scalar::Scalar;
 use vortex_error::VortexExpect;
 use vortex_error::VortexResult;
@@ -142,11 +144,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 +169,22 @@ 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]));
+    })
+}
+
 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

@@ -331,6 +331,7 @@ mod test {
     use vortex_array::VortexSessionExecute;
     use vortex_array::arrays::PrimitiveArray;
     use vortex_array::assert_arrays_eq;
+    use vortex_array::patches_v2::force_patches_v2_scatter;
     use vortex_buffer::Buffer;
     use vortex_session::VortexSession;
 
@@ -385,15 +386,25 @@ mod test {
 
         let packed_with_patches = BitPackedData::encode(&parray, 9, &mut ctx).unwrap();
         assert!(packed_with_patches.patches().is_some());
-        let packed_primitive = packed_with_patches
-            .as_array()
-            .clone()
-            .execute::<PrimitiveArray>(&mut ctx)
-            .unwrap();
-        assert_arrays_eq!(
-            packed_primitive,
-            PrimitiveArray::new(values, vortex_array::validity::Validity::NonNullable),
-            &mut ctx
-        );
+
+        // Both scatter paths must decompress to the original values. Toggling the chunk-local
+        // scatter around the decode pins down that the two agree on a real patched array.
+        for chunk_local in [false, true] {
+            force_patches_v2_scatter(chunk_local);
+            let packed_primitive = packed_with_patches
+                .as_array()
+                .clone()
+                .execute::<PrimitiveArray>(&mut ctx)
+                .unwrap();
+            assert_arrays_eq!(
+                packed_primitive,
+                PrimitiveArray::new(
+                    values.clone(),
+                    vortex_array::validity::Validity::NonNullable
+                ),
+                &mut ctx
+            );
+        }
+        force_patches_v2_scatter(false);
     }
 }

encodings/sparse/src/canonical.rs

@@ -53,13 +53,16 @@ use vortex_array::match_each_native_ptype;
 use vortex_array::match_each_unsigned_integer_ptype;
 use vortex_array::match_smallest_list_offset_type;
 use vortex_array::patches::Patches;
+use vortex_array::patches_v2::PatchesV2;
+use vortex_array::patches_v2::use_patches_v2_scatter;
 use vortex_array::scalar::DecimalScalar;
 use vortex_array::scalar::ListScalar;
 use vortex_array::scalar::Scalar;
 use vortex_array::scalar::StructScalar;
 use vortex_array::validity::Validity;
 use vortex_buffer::BitBuffer;
 use vortex_buffer::Buffer;
+use vortex_buffer::BufferMut;
 use vortex_buffer::BufferString;
 use vortex_buffer::ByteBuffer;
 use vortex_buffer::buffer;
@@ -705,6 +708,17 @@ fn execute_sparse_primitives<T: NativePType + for<'a> TryFrom<&'a Scalar, Error
         )
     };
 
+    // A non-nullable result has no validity to patch, so the whole canonicalization is a value
+    // scatter over a constant fill -- exactly what PatchesV2 addresses. Chunked patch sets take
+    // that path when the chunk-local scatter is enabled; everything else falls back to `patch`,
+    // which also handles patching validity.
+    if use_patches_v2_scatter() && matches!(validity, Validity::NonNullable) {
+        let mut out = BufferMut::full(primitive_fill, patches.array_len());
+        let v2 = PatchesV2::from_patches(patches, ctx)?;
+        v2.apply_into(out.as_mut_slice(), ctx)?;
+        return Ok(PrimitiveArray::new(out.freeze(), Validity::NonNullable).into_array());
+    }
+
     let parray = PrimitiveArray::new(buffer![primitive_fill; patches.array_len()], validity);
 
     Ok(parray.patch(patches, ctx)?.into_array())