Optimize run-end decode kernel with chunked splat stores

encodings/runend/Cargo.toml

@@ -57,3 +57,11 @@ harness = false
 [[bench]]
 name = "run_end_filter"
 harness = false
+
+[[bench]]
+name = "run_end_decode_ablation"
+harness = false
+
+[[bench]]
+name = "run_end_decode_distribution"
+harness = false

encodings/runend/benches/run_end_decode.rs

@@ -7,12 +7,17 @@ use std::fmt;
 use std::sync::LazyLock;
 
 use divan::Bencher;
+use rand::RngExt;
+use rand::SeedableRng;
+use rand::rngs::StdRng;
 use vortex_array::VortexSessionExecute;
 use vortex_array::arrays::BoolArray;
 use vortex_array::arrays::PrimitiveArray;
+use vortex_array::dtype::NativePType;
 use vortex_array::validity::Validity;
 use vortex_buffer::BitBuffer;
 use vortex_buffer::BufferMut;
+use vortex_runend::compress::runend_decode_primitive;
 use vortex_runend::decompress_bool::runend_decode_bools;
 use vortex_session::VortexSession;
 
@@ -369,6 +374,111 @@ const NULLABLE_BOOL_ARGS: &[NullableBoolBenchArgs] = &[
     },
 ];
 
+#[derive(Clone, Copy)]
+struct PrimitiveBenchArgs {
+    total_length: usize,
+    avg_run_length: usize,
+}
+
+impl fmt::Display for PrimitiveBenchArgs {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}_{}", self.total_length, self.avg_run_length)
+    }
+}
+
+/// Creates primitive test data with random run lengths (uniform with the requested average),
+/// random values, and random 90%-valid run validity, so branch-heavy decode strategies are
+/// measured against unpredictable inputs rather than a fixed pattern.
+fn create_primitive_test_data<T: NativePType + From<u8>>(
+    total_length: usize,
+    avg_run_length: usize,
+    nullable: bool,
+) -> (PrimitiveArray, PrimitiveArray) {
+    let mut rng = StdRng::seed_from_u64(0x5eed);
+    let num_runs = total_length / avg_run_length + 1;
+    let mut ends = BufferMut::<u32>::with_capacity(num_runs);
+    let mut values = BufferMut::<T>::with_capacity(num_runs);
+    let mut validity_bits = Vec::with_capacity(num_runs);
+
+    let max_run_len = (2 * avg_run_length).saturating_sub(1).max(1);
+    let mut pos = 0usize;
+    while pos < total_length {
+        let run_len = rng.random_range(1..=max_run_len).min(total_length - pos);
+        pos += run_len;
+        ends.push(pos as u32);
+        values.push(<T as From<u8>>::from(rng.random::<u8>()));
+        validity_bits.push(rng.random_bool(0.9));
+    }
+
+    let validity = if nullable {
+        Validity::from(BitBuffer::from(validity_bits))
+    } else {
+        Validity::NonNullable
+    };
+    (
+        PrimitiveArray::new(ends.freeze(), Validity::NonNullable),
+        PrimitiveArray::new(values.freeze(), validity),
+    )
+}
+
+const PRIMITIVE_ARGS: &[PrimitiveBenchArgs] = &[
+    PrimitiveBenchArgs {
+        total_length: 65_536,
+        avg_run_length: 2,
+    },
+    PrimitiveBenchArgs {
+        total_length: 65_536,
+        avg_run_length: 4,
+    },
+    PrimitiveBenchArgs {
+        total_length: 65_536,
+        avg_run_length: 8,
+    },
+    PrimitiveBenchArgs {
+        total_length: 65_536,
+        avg_run_length: 16,
+    },
+    PrimitiveBenchArgs {
+        total_length: 65_536,
+        avg_run_length: 64,
+    },
+    PrimitiveBenchArgs {
+        total_length: 65_536,
+        avg_run_length: 1024,
+    },
+];
+
+#[divan::bench(types = [u8, u16, u32, u64], args = PRIMITIVE_ARGS)]
+fn decode_primitive<T: NativePType + From<u8>>(bencher: Bencher, args: PrimitiveBenchArgs) {
+    let PrimitiveBenchArgs {
+        total_length,
+        avg_run_length,
+    } = args;
+    let (ends, values) = create_primitive_test_data::<T>(total_length, avg_run_length, false);
+    bencher
+        .with_inputs(|| (ends.clone(), values.clone(), SESSION.create_execution_ctx()))
+        .bench_refs(|(ends, values, ctx)| {
+            runend_decode_primitive(ends.clone(), values.clone(), 0, total_length, ctx)
+        });
+}
+
+#[divan::bench(types = [u8, u32, u64], args = PRIMITIVE_ARGS)]
+fn decode_primitive_nullable<T: NativePType + From<u8>>(
+    bencher: Bencher,
+    args: PrimitiveBenchArgs,
+) {
+    let PrimitiveBenchArgs {
+        total_length,
+        avg_run_length,
+    } = args;
+    let (ends, values) = create_primitive_test_data::<T>(total_length, avg_run_length, true);
+    bencher
+        .with_inputs(|| (ends.clone(), values.clone(), SESSION.create_execution_ctx()))
+        .bench_refs(|(ends, values, ctx)| {
+            runend_decode_primitive(ends.clone(), values.clone(), 0, total_length, ctx)
+        });
+}
+
 #[divan::bench(args = NULLABLE_BOOL_ARGS)]
 fn decode_bool_nullable(bencher: Bencher, args: NullableBoolBenchArgs) {
     let NullableBoolBenchArgs {

encodings/runend/benches/run_end_decode_ablation.rs

@@ -0,0 +1,186 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Ablation benchmark: isolates each change to the run-end decode kernel so every change is
+//! justified by its own numbers, all measured in one binary run over identical randomized
+//! inputs at a spread of average run lengths.
+//!
+//! The intermediate kernels live in `shared/decode_variants.rs`; see that file for what each
+//! stage (`v0`..`v3`, `n0`..`n3`) contains. For a sweep of the *data distributions* each
+//! change is sensitive to (run length, element width, validity density) see
+//! `run_end_decode_distribution`.
+
+#![expect(clippy::cast_possible_truncation)]
+
+use std::fmt;
+
+use divan::Bencher;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::dtype::NativePType;
+use vortex_array::dtype::Nullability;
+use vortex_buffer::Buffer;
+use vortex_mask::Mask;
+use vortex_runend::compress::runend_decode_typed_primitive;
+use vortex_runend::trimmed_ends_iter;
+
+#[path = "shared/decode_variants.rs"]
+mod decode_variants;
+
+use decode_variants::decode_n2;
+use decode_variants::decode_v0;
+use decode_variants::decode_v1;
+use decode_variants::decode_v2;
+
+fn main() {
+    divan::main();
+}
+
+const SEED: u64 = 0x5eed;
+const DENSITY: f64 = 0.9;
+
+#[derive(Clone, Copy)]
+struct Args {
+    total_length: usize,
+    avg_run_length: usize,
+}
+
+impl fmt::Display for Args {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}_{}", self.total_length, self.avg_run_length)
+    }
+}
+
+const ARGS: &[Args] = &[
+    Args {
+        total_length: 65_536,
+        avg_run_length: 2,
+    },
+    Args {
+        total_length: 65_536,
+        avg_run_length: 8,
+    },
+    Args {
+        total_length: 65_536,
+        avg_run_length: 64,
+    },
+    Args {
+        total_length: 65_536,
+        avg_run_length: 1024,
+    },
+];
+
+fn data<T: NativePType + From<u8>>(
+    args: Args,
+) -> (Buffer<u32>, Buffer<T>, vortex_buffer::BitBuffer) {
+    // Uniform 1..=(2*avg-1) has mean `avg`; matches the distribution bench's convention.
+    decode_variants::make_data::<T>(
+        SEED,
+        args.total_length,
+        2 * args.avg_run_length - 1,
+        DENSITY,
+    )
+}
+
+#[divan::bench(types = [u8, u32, u64], args = ARGS)]
+fn v0_original<T: NativePType + From<u8>>(bencher: Bencher, args: Args) {
+    let (ends, values, _) = data::<T>(args);
+    bencher
+        .with_inputs(|| (ends.clone(), values.clone()))
+        .bench_refs(|(ends, values)| {
+            let (buf, validity) = decode_v0(
+                trimmed_ends_iter(ends.as_slice(), 0, args.total_length),
+                values.as_slice(),
+                Mask::new_true(values.len()),
+                args.total_length,
+            );
+            PrimitiveArray::new(buf, validity)
+        });
+}
+
+#[divan::bench(types = [u8, u32, u64], args = ARGS)]
+fn v1_slice_loop<T: NativePType + From<u8>>(bencher: Bencher, args: Args) {
+    let (ends, values, _) = data::<T>(args);
+    bencher
+        .with_inputs(|| (ends.clone(), values.clone()))
+        .bench_refs(|(ends, values)| {
+            let (buf, validity) = decode_v1(ends.as_slice(), values.as_slice(), args.total_length);
+            PrimitiveArray::new(buf, validity)
+        });
+}
+
+#[divan::bench(types = [u8, u32, u64], args = ARGS)]
+fn v2_chunk_stores<T: NativePType + From<u8>>(bencher: Bencher, args: Args) {
+    let (ends, values, _) = data::<T>(args);
+    bencher
+        .with_inputs(|| (ends.clone(), values.clone()))
+        .bench_refs(|(ends, values)| {
+            let (buf, validity) = decode_v2(ends.as_slice(), values.as_slice(), args.total_length);
+            PrimitiveArray::new(buf, validity)
+        });
+}
+
+#[divan::bench(types = [u8, u32, u64], args = ARGS)]
+fn v3_shipped<T: NativePType + From<u8>>(bencher: Bencher, args: Args) {
+    let (ends, values, _) = data::<T>(args);
+    bencher
+        .with_inputs(|| (ends.clone(), values.clone()))
+        .bench_refs(|(ends, values)| {
+            runend_decode_typed_primitive(
+                ends.as_slice(),
+                0,
+                values.as_slice(),
+                Mask::new_true(values.len()),
+                Nullability::NonNullable,
+                args.total_length,
+            )
+        });
+}
+
+#[divan::bench(types = [u8, u32, u64], args = ARGS)]
+fn n0_original<T: NativePType + From<u8>>(bencher: Bencher, args: Args) {
+    let (ends, values, run_validity) = data::<T>(args);
+    bencher
+        .with_inputs(|| (ends.clone(), values.clone(), run_validity.clone()))
+        .bench_refs(|(ends, values, run_validity)| {
+            let (buf, validity) = decode_v0(
+                trimmed_ends_iter(ends.as_slice(), 0, args.total_length),
+                values.as_slice(),
+                Mask::from_buffer(run_validity.clone()),
+                args.total_length,
+            );
+            PrimitiveArray::new(buf, validity)
+        });
+}
+
+#[divan::bench(types = [u8, u32, u64], args = ARGS)]
+fn n2_chunk_stores<T: NativePType + From<u8>>(bencher: Bencher, args: Args) {
+    let (ends, values, run_validity) = data::<T>(args);
+    bencher
+        .with_inputs(|| (ends.clone(), values.clone(), run_validity.clone()))
+        .bench_refs(|(ends, values, run_validity)| {
+            let (buf, validity) = decode_n2(
+                ends.as_slice(),
+                values.as_slice(),
+                run_validity,
+                args.total_length,
+            );
+            PrimitiveArray::new(buf, validity)
+        });
+}
+
+#[divan::bench(types = [u8, u32, u64], args = ARGS)]
+fn n3_shipped<T: NativePType + From<u8>>(bencher: Bencher, args: Args) {
+    let (ends, values, run_validity) = data::<T>(args);
+    bencher
+        .with_inputs(|| (ends.clone(), values.clone(), run_validity.clone()))
+        .bench_refs(|(ends, values, run_validity)| {
+            runend_decode_typed_primitive(
+                ends.as_slice(),
+                0,
+                values.as_slice(),
+                Mask::from_buffer(run_validity.clone()),
+                Nullability::Nullable,
+                args.total_length,
+            )
+        });
+}