diff --git a/encodings/runend/Cargo.toml b/encodings/runend/Cargo.toml index 7ec149b9dd9..6b80c1b1124 100644 --- a/encodings/runend/Cargo.toml +++ b/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 diff --git a/encodings/runend/benches/run_end_decode.rs b/encodings/runend/benches/run_end_decode.rs index 9256a4b3a53..80dc082dc82 100644 --- a/encodings/runend/benches/run_end_decode.rs +++ b/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>( + 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::::with_capacity(num_runs); + let mut values = BufferMut::::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(>::from(rng.random::())); + 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>(bencher: Bencher, args: PrimitiveBenchArgs) { + let PrimitiveBenchArgs { + total_length, + avg_run_length, + } = args; + let (ends, values) = create_primitive_test_data::(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>( + bencher: Bencher, + args: PrimitiveBenchArgs, +) { + let PrimitiveBenchArgs { + total_length, + avg_run_length, + } = args; + let (ends, values) = create_primitive_test_data::(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 { diff --git a/encodings/runend/benches/run_end_decode_ablation.rs b/encodings/runend/benches/run_end_decode_ablation.rs new file mode 100644 index 00000000000..a0ad78df016 --- /dev/null +++ b/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>( + args: Args, +) -> (Buffer, Buffer, vortex_buffer::BitBuffer) { + // Uniform 1..=(2*avg-1) has mean `avg`; matches the distribution bench's convention. + decode_variants::make_data::( + SEED, + args.total_length, + 2 * args.avg_run_length - 1, + DENSITY, + ) +} + +#[divan::bench(types = [u8, u32, u64], args = ARGS)] +fn v0_original>(bencher: Bencher, args: Args) { + let (ends, values, _) = data::(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>(bencher: Bencher, args: Args) { + let (ends, values, _) = data::(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>(bencher: Bencher, args: Args) { + let (ends, values, _) = data::(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>(bencher: Bencher, args: Args) { + let (ends, values, _) = data::(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>(bencher: Bencher, args: Args) { + let (ends, values, run_validity) = data::(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>(bencher: Bencher, args: Args) { + let (ends, values, run_validity) = data::(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>(bencher: Bencher, args: Args) { + let (ends, values, run_validity) = data::(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, + ) + }); +} diff --git a/encodings/runend/benches/run_end_decode_distribution.rs b/encodings/runend/benches/run_end_decode_distribution.rs new file mode 100644 index 00000000000..07a8dfe81af --- /dev/null +++ b/encodings/runend/benches/run_end_decode_distribution.rs @@ -0,0 +1,373 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Distribution-sweep ablation: attributes each run-end decode change to the data +//! distribution it is sensitive to. Every stage kernel is shared with +//! `run_end_decode_ablation` via `shared/decode_variants.rs`, so the two benchmarks measure +//! byte-identical code. +//! +//! Throughput is reported in decoded elements/sec (`ItemsCount`), which normalizes across run +//! lengths: a per-run fixed cost shows up as throughput that climbs with run length, while a +//! per-element cost shows up as throughput flat across run length. +//! +//! Groups: +//! - `nonnull_{v0,v1,v2,v3}` — run-length sweep. Isolates the non-null structural wins +//! (v0->v1 iterator/bookkeeping, v1->v2 scalar-fill -> chunk stores) against run length and +//! element width. +//! - `nonnull_v3_{elem_fill,byte_splat}` — alternative fill strategies for the shipped kernel. +//! - `zeros_v3{,_prev}` / `rand_v3_prev` — byte-uniform versus arbitrary values, isolating the +//! `memset` fast path. +//! - `nullable_{n0,n2,n3}` — run-length sweep at 90% valid. Isolates the nullable structural +//! wins and the majority-prefill validity (n2->n3), including the long-run fill dispatch. +//! - `density_{n0,n2,n3}` — validity-density sweep, u32 at run length 8. Isolates the +//! prefill validity win (n2->n3) against validity skew: prefill rewrites only the minority +//! runs, so its advantage grows as validity moves away from 50/50. +//! +//! # Reading these numbers +//! +//! `nonnull_v3` and `zeros_v3` call the real `runend_decode_typed_primitive` across a crate +//! boundary; every other variant is a copy compiled into this binary and inlines more +//! aggressively. That difference is worth up to ~40% at short runs *on identical algorithms*, +//! and it consistently favours the bench-local copies. So: +//! +//! - Comparing a bench-local variant against another bench-local variant is apples to apples. +//! - Comparing `nonnull_v3`/`zeros_v3` against a bench-local variant understates the shipped +//! kernel. A win measured that way is a lower bound; a small loss may be the boundary alone. +//! - To isolate a change cleanly, prefer a *within-function* contrast (the same bench over two +//! data distributions, as `zeros_v3` versus `nonnull_v3`), which cancels the effect. + +#![expect(clippy::cast_possible_truncation)] + +use std::fmt; + +use divan::Bencher; +use divan::counter::ItemsCount; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; +use vortex_buffer::BitBuffer; +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; +use decode_variants::decode_v3_byte_splat; +use decode_variants::decode_v3_elem_fill; +use decode_variants::decode_v3_no_memset; +use decode_variants::make_data; +use decode_variants::make_data_values; + +fn main() { + divan::main(); +} + +const SEED: u64 = 0x5eed; +const TOTAL_LENGTH: usize = 65_536; + +/// Average run length. Uniform run lengths in `1..=(2*avg-1)` have this mean. +/// +/// Extends past 1024 so the 2 KiB doubling-fill threshold is crossed for every element width +/// (u16 needs 1024+ elements, u32 512+, u64 256+). +const RUN_LENGTHS: &[usize] = &[2, 4, 8, 16, 32, 64, 256, 1024, 4096]; + +/// Fixed short run length for the density sweep, where per-run validity work is visible. +const DENSITY_RUN_LENGTH: usize = 8; + +fn max_run_len(avg: usize) -> usize { + 2 * avg - 1 +} + +// ---- Group A: non-nullable run-length sweep ---- + +fn nonnull_data>(avg: usize) -> (Buffer, Buffer) { + let (ends, values, _) = make_data::(SEED, TOTAL_LENGTH, max_run_len(avg), 1.0); + (ends, values) +} + +#[divan::bench(types = [u8, u32, u64], args = RUN_LENGTHS)] +fn nonnull_v0>(bencher: Bencher, avg: usize) { + let (ends, values) = nonnull_data::(avg); + bencher + .counter(ItemsCount::new(TOTAL_LENGTH)) + .with_inputs(|| (ends.clone(), values.clone())) + .bench_refs(|(ends, values)| { + let (buf, validity) = decode_v0( + trimmed_ends_iter(ends.as_slice(), 0, TOTAL_LENGTH), + values.as_slice(), + Mask::new_true(values.len()), + TOTAL_LENGTH, + ); + PrimitiveArray::new(buf, validity) + }); +} + +#[divan::bench(types = [u8, u32, u64], args = RUN_LENGTHS)] +fn nonnull_v1>(bencher: Bencher, avg: usize) { + let (ends, values) = nonnull_data::(avg); + bencher + .counter(ItemsCount::new(TOTAL_LENGTH)) + .with_inputs(|| (ends.clone(), values.clone())) + .bench_refs(|(ends, values)| { + let (buf, validity) = decode_v1(ends.as_slice(), values.as_slice(), TOTAL_LENGTH); + PrimitiveArray::new(buf, validity) + }); +} + +#[divan::bench(types = [u8, u32, u64], args = RUN_LENGTHS)] +fn nonnull_v2>(bencher: Bencher, avg: usize) { + let (ends, values) = nonnull_data::(avg); + bencher + .counter(ItemsCount::new(TOTAL_LENGTH)) + .with_inputs(|| (ends.clone(), values.clone())) + .bench_refs(|(ends, values)| { + let (buf, validity) = decode_v2(ends.as_slice(), values.as_slice(), TOTAL_LENGTH); + PrimitiveArray::new(buf, validity) + }); +} + +#[divan::bench(types = [u8, u16, u32, u64], args = RUN_LENGTHS)] +fn nonnull_v3>(bencher: Bencher, avg: usize) { + let (ends, values) = nonnull_data::(avg); + bencher + .counter(ItemsCount::new(TOTAL_LENGTH)) + .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, + TOTAL_LENGTH, + ) + }); +} + +/// The previous long-run fill (element loop, baseline SSE2 width) instead of the shipped +/// doubling `memcpy`. Identical to `nonnull_v3` below the 2 KiB doubling threshold; above it +/// the shipped kernel should win. +#[divan::bench(types = [u8, u16, u32, u64], args = RUN_LENGTHS)] +fn nonnull_v3_elem_fill>(bencher: Bencher, avg: usize) { + let (ends, values) = nonnull_data::(avg); + bencher + .counter(ItemsCount::new(TOTAL_LENGTH)) + .with_inputs(|| (ends.clone(), values.clone())) + .bench_refs(|(ends, values)| { + let (buf, validity) = + decode_v3_elem_fill(ends.as_slice(), values.as_slice(), TOTAL_LENGTH); + PrimitiveArray::new(buf, validity) + }); +} + +/// The rejected byte word-splat kernel. For widths > 1 this is identical to `nonnull_v3`; +/// only u8 differs, and there the shipped `nonnull_v3` (generic path) is faster across all run +/// lengths — which is why the shipped kernel carries no byte special case. +#[divan::bench(types = [u8, u16, u32, u64], args = RUN_LENGTHS)] +fn nonnull_v3_byte_splat>(bencher: Bencher, avg: usize) { + let (ends, values) = nonnull_data::(avg); + bencher + .counter(ItemsCount::new(TOTAL_LENGTH)) + .with_inputs(|| (ends.clone(), values.clone())) + .bench_refs(|(ends, values)| { + let (buf, validity) = + decode_v3_byte_splat(ends.as_slice(), values.as_slice(), TOTAL_LENGTH); + PrimitiveArray::new(buf, validity) + }); +} + +// ---- Group A2: byte-uniform values (zeros) vs arbitrary values ---- +// +// A byte-uniform value fills with `memset`; an arbitrary one takes the doubling path. +// `zeros_v3` vs `zeros_v3_prev` isolates that fast path, and `rand_v3_prev` vs `nonnull_v3` +// confirms arbitrary values are unaffected by it. + +fn zero_data>(avg: usize) -> (Buffer, Buffer) { + let (ends, values, _) = make_data_values::(SEED, TOTAL_LENGTH, max_run_len(avg), 1.0, true); + (ends, values) +} + +#[divan::bench(types = [u32, u64], args = RUN_LENGTHS)] +fn zeros_v3>(bencher: Bencher, avg: usize) { + let (ends, values) = zero_data::(avg); + bencher + .counter(ItemsCount::new(TOTAL_LENGTH)) + .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, + TOTAL_LENGTH, + ) + }); +} + +#[divan::bench(types = [u32, u64], args = RUN_LENGTHS)] +fn zeros_v3_prev>(bencher: Bencher, avg: usize) { + let (ends, values) = zero_data::(avg); + bencher + .counter(ItemsCount::new(TOTAL_LENGTH)) + .with_inputs(|| (ends.clone(), values.clone())) + .bench_refs(|(ends, values)| { + let (buf, validity) = + decode_v3_no_memset(ends.as_slice(), values.as_slice(), TOTAL_LENGTH); + PrimitiveArray::new(buf, validity) + }); +} + +/// Arbitrary (not byte-uniform) values through the pre-`memset` kernel; compare against +/// `nonnull_v3` to confirm the fast path costs nothing when it does not apply. +#[divan::bench(types = [u32, u64], args = RUN_LENGTHS)] +fn rand_v3_prev>(bencher: Bencher, avg: usize) { + let (ends, values) = nonnull_data::(avg); + bencher + .counter(ItemsCount::new(TOTAL_LENGTH)) + .with_inputs(|| (ends.clone(), values.clone())) + .bench_refs(|(ends, values)| { + let (buf, validity) = + decode_v3_no_memset(ends.as_slice(), values.as_slice(), TOTAL_LENGTH); + PrimitiveArray::new(buf, validity) + }); +} + +// ---- Group B: nullable run-length sweep at 90% valid ---- + +fn nullable_data>(avg: usize) -> (Buffer, Buffer, BitBuffer) { + make_data::(SEED, TOTAL_LENGTH, max_run_len(avg), 0.9) +} + +#[divan::bench(types = [u8, u32, u64], args = RUN_LENGTHS)] +fn nullable_n0>(bencher: Bencher, avg: usize) { + let (ends, values, validity) = nullable_data::(avg); + bencher + .counter(ItemsCount::new(TOTAL_LENGTH)) + .with_inputs(|| (ends.clone(), values.clone(), validity.clone())) + .bench_refs(|(ends, values, validity)| { + let (buf, decoded_validity) = decode_v0( + trimmed_ends_iter(ends.as_slice(), 0, TOTAL_LENGTH), + values.as_slice(), + Mask::from_buffer(validity.clone()), + TOTAL_LENGTH, + ); + PrimitiveArray::new(buf, decoded_validity) + }); +} + +#[divan::bench(types = [u8, u32, u64], args = RUN_LENGTHS)] +fn nullable_n2>(bencher: Bencher, avg: usize) { + let (ends, values, validity) = nullable_data::(avg); + bencher + .counter(ItemsCount::new(TOTAL_LENGTH)) + .with_inputs(|| (ends.clone(), values.clone(), validity.clone())) + .bench_refs(|(ends, values, validity)| { + let (buf, decoded_validity) = + decode_n2(ends.as_slice(), values.as_slice(), validity, TOTAL_LENGTH); + PrimitiveArray::new(buf, decoded_validity) + }); +} + +#[divan::bench(types = [u8, u32, u64], args = RUN_LENGTHS)] +fn nullable_n3>(bencher: Bencher, avg: usize) { + let (ends, values, validity) = nullable_data::(avg); + bencher + .counter(ItemsCount::new(TOTAL_LENGTH)) + .with_inputs(|| (ends.clone(), values.clone(), validity.clone())) + .bench_refs(|(ends, values, validity)| { + runend_decode_typed_primitive( + ends.as_slice(), + 0, + values.as_slice(), + Mask::from_buffer(validity.clone()), + Nullability::Nullable, + TOTAL_LENGTH, + ) + }); +} + +// ---- Group C: validity-density sweep (u32, run length 8) ---- + +fn density_data(density_pct: u32) -> (Buffer, Buffer, BitBuffer) { + make_data::( + SEED, + TOTAL_LENGTH, + max_run_len(DENSITY_RUN_LENGTH), + f64::from(density_pct) / 100.0, + ) +} + +#[derive(Clone, Copy)] +struct Density(u32); + +impl fmt::Display for Density { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "valid{:02}pct", self.0) + } +} + +/// Fraction-valid densities (percent). Spans balanced (50) to strongly skewed (99) validity, +/// plus a null-heavy point (10) to show the prefill win is symmetric about 50%. +const DENSITIES: &[Density] = &[ + Density(10), + Density(50), + Density(70), + Density(90), + Density(95), + Density(99), +]; + +#[divan::bench(args = DENSITIES)] +fn density_n0(bencher: Bencher, density: Density) { + let (ends, values, validity) = density_data(density.0); + bencher + .counter(ItemsCount::new(TOTAL_LENGTH)) + .with_inputs(|| (ends.clone(), values.clone(), validity.clone())) + .bench_refs(|(ends, values, validity)| { + let (buf, decoded_validity) = decode_v0( + trimmed_ends_iter(ends.as_slice(), 0, TOTAL_LENGTH), + values.as_slice(), + Mask::from_buffer(validity.clone()), + TOTAL_LENGTH, + ); + PrimitiveArray::new(buf, decoded_validity) + }); +} + +#[divan::bench(args = DENSITIES)] +fn density_n2(bencher: Bencher, density: Density) { + let (ends, values, validity) = density_data(density.0); + bencher + .counter(ItemsCount::new(TOTAL_LENGTH)) + .with_inputs(|| (ends.clone(), values.clone(), validity.clone())) + .bench_refs(|(ends, values, validity)| { + let (buf, decoded_validity) = + decode_n2(ends.as_slice(), values.as_slice(), validity, TOTAL_LENGTH); + PrimitiveArray::new(buf, decoded_validity) + }); +} + +#[divan::bench(args = DENSITIES)] +fn density_n3(bencher: Bencher, density: Density) { + let (ends, values, validity) = density_data(density.0); + bencher + .counter(ItemsCount::new(TOTAL_LENGTH)) + .with_inputs(|| (ends.clone(), values.clone(), validity.clone())) + .bench_refs(|(ends, values, validity)| { + runend_decode_typed_primitive( + ends.as_slice(), + 0, + values.as_slice(), + Mask::from_buffer(validity.clone()), + Nullability::Nullable, + TOTAL_LENGTH, + ) + }); +} diff --git a/encodings/runend/benches/shared/decode_variants.rs b/encodings/runend/benches/shared/decode_variants.rs new file mode 100644 index 00000000000..1bc9b0ebc7d --- /dev/null +++ b/encodings/runend/benches/shared/decode_variants.rs @@ -0,0 +1,591 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Shared run-end decode kernel variants for the ablation benchmarks. +//! +//! Each intermediate stage of the decode optimization lives here exactly once so that +//! `run_end_decode_ablation` and `run_end_decode_distribution` measure byte-identical code. +//! An ablation only proves a change if every benchmark that names a stage runs the same +//! implementation of it. +//! +//! Stages (non-nullable): +//! - [`decode_v0`]: original iterator chain (`trimmed_ends_iter` + `zip_eq`) + `push_n_unchecked` +//! - [`decode_v1`]: slice iteration + inline trim, still `push_n_unchecked` per run +//! - [`decode_v2`]: `decode_v1` + unconditional chunked element splat stores +//! - shipped: `runend_decode_typed_primitive` (adds the byte word/memset + long-run fill paths) +//! +//! Stages (nullable): +//! - `decode_v0` with a `Mask::Values`: Option-zip + `append_n` validity + `push_n_unchecked` +//! - [`decode_n2`]: slice loop + chunk stores, validity still built with per-run `append_n` +//! - shipped: `runend_decode_typed_primitive` (majority prefill + `fill_range_unchecked`) + +// Included via `#[path]` into each bench binary, which uses only a subset of these stages. +#![allow(dead_code)] +// The `E::from_usize(length)` conversions cannot fail for benchmark lengths. +#![allow(clippy::unwrap_used)] + +use std::cmp::min; +use std::mem::MaybeUninit; + +use itertools::Itertools; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex_array::dtype::IntegerPType; +use vortex_array::dtype::NativePType; +use vortex_array::dtype::Nullability; +use vortex_array::validity::Validity; +use vortex_buffer::BitBuffer; +use vortex_buffer::BitBufferMut; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_mask::Mask; + +/// Build run-end test data with uniformly-random run lengths (average `(max_run_len + 1) / 2`), +/// random values, and random run validity at the requested fraction-valid density. +/// +/// Randomized so the branch-heavy decode stages are measured against unpredictable inputs +/// rather than a periodic pattern the predictor can learn. +pub fn make_data>( + seed: u64, + total_length: usize, + max_run_len: usize, + validity_density: f64, +) -> (Buffer, Buffer, BitBuffer) { + make_data_values(seed, total_length, max_run_len, validity_density, false) +} + +/// As [`make_data`], but `zero_values` forces every run value to zero. +/// +/// Zero is byte-uniform, which the decode kernel fills with `memset` instead of the general +/// path. Real run-end columns carry both kinds of value, so both are benchmarked. +pub fn make_data_values>( + seed: u64, + total_length: usize, + max_run_len: usize, + validity_density: f64, + zero_values: bool, +) -> (Buffer, Buffer, BitBuffer) { + let mut rng = StdRng::seed_from_u64(seed); + let mut ends = BufferMut::::empty(); + let mut values = BufferMut::::empty(); + let mut validity = Vec::new(); + let max_run_len = max_run_len.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); + let byte = if zero_values { 0 } else { rng.random::() }; + values.push(>::from(byte)); + validity.push(rng.random_bool(validity_density)); + } + (ends.freeze(), values.freeze(), BitBuffer::from(validity)) +} + +/// Number of runs produced by [`make_data`] for the given parameters, for per-run normalization. +pub fn run_count(ends: &Buffer) -> usize { + ends.len() +} + +// ---- helpers shared by v1/v2/n2 (mirror the shipped private helpers) ---- + +pub const DECODE_CHUNK_BYTES: usize = 32; + +pub const fn decode_chunk_len() -> usize { + let n = DECODE_CHUNK_BYTES / size_of::(); + if n == 0 { 1 } else { n } +} + +#[inline(always)] +fn trim_end(end: E, offset: E, length: E) -> usize { + assert!(end >= offset, "run end must be >= offset"); + min(end - offset, length).as_() +} + +/// The v2 element-splat store (the shipped wide-element path, applied to all widths). +/// +/// # Safety +/// +/// The allocation behind `base` must have room for at least +/// `max(pos, end) + decode_chunk_len::()` elements. +#[inline(always)] +unsafe fn splat_run_elements(base: *mut MaybeUninit, pos: usize, end: usize, value: T) { + let chunk = const { decode_chunk_len::() }; + // SAFETY: caller guarantees one chunk of slack past max(pos, end). + unsafe { + let mut p = base.add(pos); + let stop = base.add(end); + loop { + for i in 0..chunk { + p.add(i).write(MaybeUninit::new(value)); + } + p = p.add(chunk); + if p >= stop { + break; + } + } + } +} + +// ---- REJECTED byte word-splat variant ---- +// +// Reconstructs the byte special case that an earlier revision of the shipped kernel carried: +// for `u8`, splat a replicated `0x0101..` word (whose runtime multiply is opaque to LLVM's +// loop-idiom pass) instead of the generic element chunk stores. It existed to stop the byte +// path collapsing into a `memset` call per short run. +// +// It was removed because the unconditional-first-chunk restructuring already prevents that +// collapse for short runs, making the word splat redundant *and* measurably slower for u8 +// (the `wrapping_mul` + unaligned `u64` stores lose to the generic 16-byte vector stores). +// `nonnull_v3` (shipped) vs `nonnull_v3_byte_splat` in `run_end_decode_distribution` +// reproduces that: for widths > 1 they are identical; for `u8` the shipped path wins. + +const LONG_RUN_FILL_BYTES: usize = 256; + +#[inline(never)] +unsafe fn fill_run(dst: *mut MaybeUninit, len: usize, value: T) { + unsafe { + if size_of::() == 1 { + let byte: u8 = std::mem::transmute_copy(&value); + dst.cast::().write_bytes(byte, len); + } else { + for i in 0..len { + dst.add(i).write(MaybeUninit::new(value)); + } + } + } +} + +/// # Safety +/// +/// The allocation behind `base` must have room for at least +/// `max(pos, end) + decode_chunk_len::()` elements. +#[inline(always)] +unsafe fn splat_run_byte_splat( + base: *mut MaybeUninit, + pos: usize, + end: usize, + value: T, +) { + if size_of::() == 1 { + // SAFETY: size_of::() == 1, and the caller guarantees a chunk of slack. + unsafe { + let byte: u8 = std::mem::transmute_copy(&value); + let word = u64::from(byte).wrapping_mul(0x0101_0101_0101_0101); + let mut p = base.add(pos).cast::(); + let stop = base.add(end).cast::(); + for i in 0..DECODE_CHUNK_BYTES / 8 { + p.add(i * 8).cast::().write_unaligned(word); + } + p = p.add(DECODE_CHUNK_BYTES); + if p >= stop { + return; + } + let len = end - pos; + if len >= LONG_RUN_FILL_BYTES { + base.add(pos).cast::().write_bytes(byte, len); + return; + } + loop { + for i in 0..DECODE_CHUNK_BYTES / 8 { + p.add(i * 8).cast::().write_unaligned(word); + } + p = p.add(DECODE_CHUNK_BYTES); + if p >= stop { + break; + } + } + } + } else { + let chunk = const { decode_chunk_len::() }; + // SAFETY: caller guarantees one chunk of slack past max(pos, end). + unsafe { + let mut p = base.add(pos); + let stop = base.add(end); + for i in 0..chunk { + p.add(i).write(MaybeUninit::new(value)); + } + p = p.add(chunk); + if p >= stop { + return; + } + let len = end - pos; + if len * size_of::() >= LONG_RUN_FILL_BYTES { + fill_run(base.add(pos), len, value); + return; + } + loop { + for i in 0..chunk { + p.add(i).write(MaybeUninit::new(value)); + } + p = p.add(chunk); + if p >= stop { + break; + } + } + } + } +} + +// ---- PREVIOUS fill: doubling memcpy, but no byte-uniform `memset` fast path ---- +// +// Reconstructs the kernel as it stood before `repeated_byte` was added: byte-sized elements +// fill with `memset`, every wider element takes the doubling/element path regardless of +// whether its bytes happen to be identical. `zeros_v3` vs `zeros_v3_prev` isolates the +// `memset` fast path on byte-uniform values (zero); `rand_v3_prev` vs `nonnull_v3` confirms +// arbitrary values are unaffected. + +/// # Safety +/// +/// `dst` must be valid for writes of `len` elements. +#[inline(never)] +unsafe fn fill_run_no_memset(dst: *mut MaybeUninit, len: usize, value: T) { + unsafe { + if size_of::() == 1 { + let byte: u8 = std::mem::transmute_copy(&value); + dst.cast::().write_bytes(byte, len); + return; + } + if len * size_of::() >= DOUBLING_FILL_BYTES { + let seed = (64 / size_of::()).max(1); + for i in 0..seed { + dst.add(i).write(MaybeUninit::new(value)); + } + let mut filled = seed; + while filled < len { + let n = filled.min(len - filled); + std::ptr::copy_nonoverlapping(dst, dst.add(filled), n); + filled += n; + } + return; + } + for i in 0..len { + dst.add(i).write(MaybeUninit::new(value)); + } + } +} + +/// # Safety +/// +/// The allocation behind `base` must have room for at least +/// `max(pos, end) + decode_chunk_len::()` elements. +#[inline(always)] +unsafe fn splat_run_no_memset( + base: *mut MaybeUninit, + pos: usize, + end: usize, + value: T, +) { + let chunk = const { decode_chunk_len::() }; + // SAFETY: caller guarantees one chunk of slack past max(pos, end). + unsafe { + let mut p = base.add(pos); + let stop = base.add(end); + for i in 0..chunk { + p.add(i).write(MaybeUninit::new(value)); + } + p = p.add(chunk); + if p >= stop { + return; + } + let len = end - pos; + if len * size_of::() >= LONG_RUN_FILL_BYTES { + fill_run_no_memset(base.add(pos), len, value); + return; + } + loop { + for i in 0..chunk { + p.add(i).write(MaybeUninit::new(value)); + } + p = p.add(chunk); + if p >= stop { + break; + } + } + } +} + +/// Non-nullable decode without the byte-uniform `memset` fast path (see above). +pub fn decode_v3_no_memset( + run_ends: &[E], + values: &[T], + length: usize, +) -> (Buffer, Validity) { + let offset_e = E::zero(); + let length_e = E::from_usize(length).unwrap(); + let mut decoded = BufferMut::::with_capacity(length + decode_chunk_len::()); + let base = decoded.spare_capacity_mut().as_mut_ptr(); + let mut pos = 0usize; + for (&end, &value) in run_ends.iter().zip(values) { + let end = trim_end(end, offset_e, length_e); + assert!( + end >= pos, + "Runend ends must be monotonic, got {end} after {pos}" + ); + // SAFETY: pos <= end <= length and the buffer has one chunk of padding. + unsafe { splat_run_no_memset(base, pos, end, value) }; + pos = end; + } + // SAFETY: every element in 0..pos was initialized above. + unsafe { decoded.set_len(pos) }; + (decoded.into(), Nullability::NonNullable.into()) +} + +/// The doubling threshold, mirrored from the shipped kernel. +const DOUBLING_FILL_BYTES: usize = 2048; + +// ---- PREVIOUS long-run fill: element loop instead of doubling memcpy ---- +// +// The shipped `fill_run` grows long fills by doubling `copy_nonoverlapping`, reaching the +// libc's runtime-dispatched (AVX2/ERMS) store width. This variant keeps the structure but +// fills with the plain element loop, which the compiler emits at baseline SSE2 width. +// `nonnull_v3` (shipped) vs `nonnull_v3_elem_fill` isolates that change: identical below the +// 2 KiB doubling threshold, shipped pulls ahead above it. + +/// # Safety +/// +/// `dst` must be valid for writes of `len` elements. +#[inline(never)] +unsafe fn fill_run_elems(dst: *mut MaybeUninit, len: usize, value: T) { + unsafe { + if size_of::() == 1 { + let byte: u8 = std::mem::transmute_copy(&value); + dst.cast::().write_bytes(byte, len); + } else { + for i in 0..len { + dst.add(i).write(MaybeUninit::new(value)); + } + } + } +} + +/// # Safety +/// +/// The allocation behind `base` must have room for at least +/// `max(pos, end) + decode_chunk_len::()` elements. +#[inline(always)] +unsafe fn splat_run_elem_fill( + base: *mut MaybeUninit, + pos: usize, + end: usize, + value: T, +) { + let chunk = const { decode_chunk_len::() }; + // SAFETY: caller guarantees one chunk of slack past max(pos, end). + unsafe { + let mut p = base.add(pos); + let stop = base.add(end); + for i in 0..chunk { + p.add(i).write(MaybeUninit::new(value)); + } + p = p.add(chunk); + if p >= stop { + return; + } + let len = end - pos; + if len * size_of::() >= LONG_RUN_FILL_BYTES { + fill_run_elems(base.add(pos), len, value); + return; + } + loop { + for i in 0..chunk { + p.add(i).write(MaybeUninit::new(value)); + } + p = p.add(chunk); + if p >= stop { + break; + } + } + } +} + +/// Non-nullable decode whose long-run fill uses the element loop (see above). +pub fn decode_v3_elem_fill( + run_ends: &[E], + values: &[T], + length: usize, +) -> (Buffer, Validity) { + let offset_e = E::zero(); + let length_e = E::from_usize(length).unwrap(); + let mut decoded = BufferMut::::with_capacity(length + decode_chunk_len::()); + let base = decoded.spare_capacity_mut().as_mut_ptr(); + let mut pos = 0usize; + for (&end, &value) in run_ends.iter().zip(values) { + let end = trim_end(end, offset_e, length_e); + assert!( + end >= pos, + "Runend ends must be monotonic, got {end} after {pos}" + ); + // SAFETY: pos <= end <= length and the buffer has one chunk of padding. + unsafe { splat_run_elem_fill(base, pos, end, value) }; + pos = end; + } + // SAFETY: every element in 0..pos was initialized above. + unsafe { decoded.set_len(pos) }; + (decoded.into(), Nullability::NonNullable.into()) +} + +/// Non-nullable decode using the rejected byte word-splat kernel (see above). +pub fn decode_v3_byte_splat( + run_ends: &[E], + values: &[T], + length: usize, +) -> (Buffer, Validity) { + let offset_e = E::zero(); + let length_e = E::from_usize(length).unwrap(); + let mut decoded = BufferMut::::with_capacity(length + decode_chunk_len::()); + let base = decoded.spare_capacity_mut().as_mut_ptr(); + let mut pos = 0usize; + for (&end, &value) in run_ends.iter().zip(values) { + let end = trim_end(end, offset_e, length_e); + assert!( + end >= pos, + "Runend ends must be monotonic, got {end} after {pos}" + ); + // SAFETY: pos <= end <= length and the buffer has one chunk of padding. + unsafe { splat_run_byte_splat(base, pos, end, value) }; + pos = end; + } + // SAFETY: every element in 0..pos was initialized above. + unsafe { decoded.set_len(pos) }; + (decoded.into(), Nullability::NonNullable.into()) +} + +// ---- v0: original implementation (verbatim from the pre-change kernel) ---- + +pub fn decode_v0( + run_ends: impl Iterator, + values: &[T], + values_validity: Mask, + length: usize, +) -> (Buffer, Validity) { + match values_validity { + Mask::AllTrue(_) => { + let mut decoded: BufferMut = BufferMut::with_capacity(length); + for (end, value) in run_ends.zip_eq(values) { + assert!( + end >= decoded.len(), + "Runend ends must be monotonic, got {end} after {}", + decoded.len() + ); + assert!(end <= length, "Runend end must be less than overall length"); + // SAFETY: we preallocate enough capacity because we know the total length + unsafe { decoded.push_n_unchecked(*value, end - decoded.len()) }; + } + (decoded.into(), Nullability::NonNullable.into()) + } + Mask::AllFalse(_) => (Buffer::::zeroed(length), Validity::AllInvalid), + Mask::Values(mask) => { + let mut decoded = BufferMut::with_capacity(length); + let mut decoded_validity = BitBufferMut::with_capacity(length); + for (end, value) in run_ends.zip_eq( + values + .iter() + .zip(mask.bit_buffer().iter()) + .map(|(&v, is_valid)| is_valid.then_some(v)), + ) { + assert!( + end >= decoded.len(), + "Runend ends must be monotonic, got {end} after {}", + decoded.len() + ); + assert!(end <= length, "Runend end must be less than overall length"); + match value { + None => { + decoded_validity.append_n(false, end - decoded.len()); + // SAFETY: we preallocate enough capacity + unsafe { decoded.push_n_unchecked(T::default(), end - decoded.len()) }; + } + Some(value) => { + decoded_validity.append_n(true, end - decoded.len()); + // SAFETY: we preallocate enough capacity + unsafe { decoded.push_n_unchecked(value, end - decoded.len()) }; + } + } + } + (decoded.into(), Validity::from(decoded_validity.freeze())) + } + } +} + +// ---- v1: slice loop, still push_n_unchecked ---- + +pub fn decode_v1( + run_ends: &[E], + values: &[T], + length: usize, +) -> (Buffer, Validity) { + let offset_e = E::zero(); + let length_e = E::from_usize(length).unwrap(); + let mut decoded: BufferMut = BufferMut::with_capacity(length); + for (&end, &value) in run_ends.iter().zip(values) { + let end = trim_end(end, offset_e, length_e); + assert!( + end >= decoded.len(), + "Runend ends must be monotonic, got {end} after {}", + decoded.len() + ); + // SAFETY: we preallocate enough capacity because we know the total length + unsafe { decoded.push_n_unchecked(value, end - decoded.len()) }; + } + (decoded.into(), Nullability::NonNullable.into()) +} + +// ---- v2: slice loop + chunked element splat stores ---- + +pub fn decode_v2( + run_ends: &[E], + values: &[T], + length: usize, +) -> (Buffer, Validity) { + let offset_e = E::zero(); + let length_e = E::from_usize(length).unwrap(); + let mut decoded = BufferMut::::with_capacity(length + decode_chunk_len::()); + let base = decoded.spare_capacity_mut().as_mut_ptr(); + let mut pos = 0usize; + for (&end, &value) in run_ends.iter().zip(values) { + let end = trim_end(end, offset_e, length_e); + assert!( + end >= pos, + "Runend ends must be monotonic, got {end} after {pos}" + ); + // SAFETY: pos <= end <= length and the buffer has one chunk of padding. + unsafe { splat_run_elements(base, pos, end, value) }; + pos = end; + } + // SAFETY: every element in 0..pos was initialized above. + unsafe { decoded.set_len(pos) }; + (decoded.into(), Nullability::NonNullable.into()) +} + +// ---- n2: nullable with chunk stores but validity still via per-run append_n ---- + +pub fn decode_n2( + run_ends: &[E], + values: &[T], + run_validity: &BitBuffer, + length: usize, +) -> (Buffer, Validity) { + let offset_e = E::zero(); + let length_e = E::from_usize(length).unwrap(); + let mut decoded = BufferMut::::with_capacity(length + decode_chunk_len::()); + let mut decoded_validity = BitBufferMut::with_capacity(length); + let base = decoded.spare_capacity_mut().as_mut_ptr(); + let mut pos = 0usize; + for ((&end, &value), is_valid) in run_ends.iter().zip(values).zip(run_validity.iter()) { + let end = trim_end(end, offset_e, length_e); + assert!( + end >= pos, + "Runend ends must be monotonic, got {end} after {pos}" + ); + decoded_validity.append_n(is_valid, end - pos); + let value = if is_valid { value } else { T::default() }; + // SAFETY: pos <= end <= length and the buffer has one chunk of padding. + unsafe { splat_run_elements(base, pos, end, value) }; + pos = end; + } + // SAFETY: every element in 0..pos was initialized above. + unsafe { decoded.set_len(pos) }; + (decoded.into(), Validity::from(decoded_validity.freeze())) +} diff --git a/encodings/runend/src/compress.rs b/encodings/runend/src/compress.rs index 544b677dfa3..9314faff14f 100644 --- a/encodings/runend/src/compress.rs +++ b/encodings/runend/src/compress.rs @@ -1,7 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use itertools::Itertools; +use std::cmp::min; +use std::mem::MaybeUninit; + use vortex_array::ArrayRef; use vortex_array::ArrayView; use vortex_array::ExecutionCtx; @@ -14,6 +16,7 @@ use vortex_array::arrays::VarBinViewArray; use vortex_array::arrays::bool::BoolArrayExt; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::buffer::BufferHandle; +use vortex_array::dtype::IntegerPType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::expr::stats::Precision; @@ -29,10 +32,9 @@ use vortex_buffer::BufferMut; use vortex_buffer::buffer; use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_error::vortex_panic; use vortex_mask::Mask; -use crate::iter::trimmed_ends_iter; - /// Run-end encode a `PrimitiveArray`, returning a tuple of `(ends, values)`. pub fn runend_encode( array: ArrayView, @@ -195,7 +197,8 @@ pub fn runend_decode_primitive( Ok(match_each_native_ptype!(values.ptype(), |P| { match_each_unsigned_integer_ptype!(ends.ptype(), |E| { runend_decode_typed_primitive( - trimmed_ends_iter(ends.as_slice::(), offset, length), + ends.as_slice::(), + offset, values.as_slice::

(), validity_mask, values.dtype().nullability(), @@ -205,70 +208,271 @@ pub fn runend_decode_primitive( })) } +/// Number of bytes written by each unconditional splat store in the decode loop. +/// +/// Each run is expanded chunk-at-a-time: short runs complete in a single store into the +/// buffer's padding, so they never pay for a per-element tail loop. +const DECODE_CHUNK_BYTES: usize = 32; + +/// Number of elements of `T` written per splat store (at least one). +const fn decode_chunk_len() -> usize { + let n = DECODE_CHUNK_BYTES / size_of::(); + if n == 0 { 1 } else { n } +} + +/// Trim a raw run end down by `offset` and clamp it to `length`, panicking if the end +/// precedes the offset. +#[inline(always)] +fn trim_end(end: E, offset: E, length: E) -> usize { + if end < offset { + vortex_panic!("run end {end} must be >= offset {offset}"); + } + min(end - offset, length).as_() +} + +/// Runs at least this many bytes long are filled by [`fill_run`] instead of the inline +/// chunked stores. The out-of-line fill keeps its own clean codegen (a memset for bytes, an +/// aligned vector loop otherwise) regardless of register pressure in the calling decode +/// loop, and its call cost is amortized over the run. +const LONG_RUN_FILL_BYTES: usize = 256; + +/// Fills of at least this many bytes grow by doubling `copy_nonoverlapping` instead of an +/// element loop. +/// +/// The element loop is compiled against the baseline target features (16-byte SSE2 stores), +/// whereas `memcpy` is resolved by the libc at runtime to the widest implementation the host +/// supports (AVX2, or `rep movsb` on ERMS parts). Doubling therefore reaches a store width the +/// compiled loop cannot, but each `memcpy` costs a call, so it only pays once the run is long +/// enough to amortize it. Measured crossover is ~2 KiB across u16/u32/u64; below it doubling +/// is up to 2x slower, at and above it is 1.1-1.4x faster. See `run_end_decode_distribution`. +const DOUBLING_FILL_BYTES: usize = 2048; + +/// The repeated byte of `value` when all of its bytes are equal, otherwise `None`. +/// +/// Such a value can be filled with a single `memset`, which writes without reading anything +/// back and so beats the doubling fill below. The common cases are `0` and all-ones (`-1` for +/// signed integers), both frequent in real columns. +/// +/// The comparison is arithmetic on a same-sized integer rather than a walk over `value`'s +/// bytes, so it never reads padding. Only power-of-two widths up to 8 are recognised; wider +/// elements (such as binary views) simply take the general path. +#[inline(always)] +fn repeated_byte(value: T) -> Option { + // SAFETY: each branch transmutes `T` to an unsigned integer of exactly the same size. + // Every type reaching this function is a primitive or a binary view, none of which have + // padding, so all bytes of the source are initialized. + unsafe { + match size_of::() { + 1 => Some(std::mem::transmute_copy::(&value)), + 2 => { + let v = std::mem::transmute_copy::(&value); + (v == (v & 0xFF) * 0x0101).then(|| v.to_le_bytes()[0]) + } + 4 => { + let v = std::mem::transmute_copy::(&value); + (v == (v & 0xFF) * 0x0101_0101).then(|| v.to_le_bytes()[0]) + } + 8 => { + let v = std::mem::transmute_copy::(&value); + (v == (v & 0xFF) * 0x0101_0101_0101_0101).then(|| v.to_le_bytes()[0]) + } + _ => None, + } + } +} + +/// Fill `dst[..len]` with `value`, used for long runs. +/// +/// Deliberately not inlined: inlining a fill loop into the decode loop leaves its codegen at +/// the mercy of the surrounding register pressure — inside the nullable decode arm the inline +/// loop otherwise loses its wide vector stores. +/// +/// # Safety +/// +/// `dst` must be valid for writes of `len` elements. +#[inline(never)] +unsafe fn fill_run(dst: *mut MaybeUninit, len: usize, value: T) { + unsafe { + // `memset` needs no read traffic and the libc dispatches it well at every length, so + // it wins outright whenever the value is byte-uniform. Byte-sized elements always are. + if let Some(byte) = repeated_byte(value) { + dst.cast::().write_bytes(byte, len * size_of::()); + return; + } + if len * size_of::() >= DOUBLING_FILL_BYTES { + // Seed one cache line by hand, then repeatedly copy the filled prefix onto the + // tail. `seed <= len` holds because `seed * size_of::()` is at most 64 while + // `len * size_of::()` is at least `DOUBLING_FILL_BYTES`. + let seed = (64 / size_of::()).max(1); + for i in 0..seed { + dst.add(i).write(MaybeUninit::new(value)); + } + let mut filled = seed; + while filled < len { + // `n <= filled` keeps the source `dst[..n]` disjoint from the destination + // `dst[filled..filled + n]`, and `filled + n <= len` stays in bounds. + let n = filled.min(len - filled); + std::ptr::copy_nonoverlapping(dst, dst.add(filled), n); + filled += n; + } + return; + } + for i in 0..len { + dst.add(i).write(MaybeUninit::new(value)); + } + } +} + +/// Splat `value` into `base[pos..end]`. +/// +/// Short runs use unconditional chunk-wide stores: they always write at least one full chunk +/// (rounding the run up to a multiple of the chunk length), so the overshoot past `end` is +/// written and later either overwritten by the next run or discarded by the final `set_len`. +/// Long runs dispatch to [`fill_run`] and write exactly `end - pos` elements. +/// +/// The unconditional first chunk keeps byte-element codegen fast without a special case: a run +/// no longer than one chunk finishes in the inline, fixed-length vector stores below and never +/// reaches the trailing loop that LLVM's loop-idiom pass would collapse into a per-run +/// `memset` call. (An explicit replicated-word store for bytes was measurably slower here.) +/// +/// # Safety +/// +/// The allocation behind `base` must have room for at least +/// `max(pos, end) + decode_chunk_len::()` elements. +#[inline(always)] +unsafe fn splat_run(base: *mut MaybeUninit, pos: usize, end: usize, value: T) { + let chunk = const { decode_chunk_len::() }; + // SAFETY: the caller guarantees the allocation extends one chunk past max(pos, end), + // so every store below lands inside it. + unsafe { + let mut p = base.add(pos); + let stop = base.add(end); + // The first chunk is unconditional: runs up to one chunk finish on one branch. + for i in 0..chunk { + p.add(i).write(MaybeUninit::new(value)); + } + p = p.add(chunk); + if p >= stop { + return; + } + // Only multi-chunk runs pay for the long-run check; the out-of-line fill keeps wide + // vector stores regardless of register pressure in the calling loop and uses memset + // for byte fills. + let len = end - pos; + if len * size_of::() >= LONG_RUN_FILL_BYTES { + // Refilling from `pos` overlaps the chunk written above, which is harmless. + fill_run(base.add(pos), len, value); + return; + } + loop { + for i in 0..chunk { + p.add(i).write(MaybeUninit::new(value)); + } + p = p.add(chunk); + if p >= stop { + break; + } + } + } +} + /// Decode a run-end encoded slice of values into a flat `Buffer` and `Validity`. /// -/// This is the core decode loop shared by primitive and varbinview run-end decoding. -fn runend_decode_slice( - run_ends: impl Iterator, +/// This is the core decode loop shared by primitive and varbinview run-end decoding. Run +/// ends are adjusted by `offset` and clamped to `length` while decoding. +fn runend_decode_slice( + run_ends: &[E], + offset: usize, values: &[T], values_validity: Mask, values_nullability: Nullability, length: usize, ) -> (Buffer, Validity) { + assert_eq!( + run_ends.len(), + values.len(), + "runend ends and values must have equal lengths" + ); + let offset_e = E::from_usize(offset).unwrap_or_else(|| { + vortex_panic!( + "offset {} cannot be converted to {}", + offset, + std::any::type_name::() + ) + }); + let length_e = E::from_usize(length).unwrap_or_else(|| { + vortex_panic!( + "length {} cannot be converted to {}", + length, + std::any::type_name::() + ) + }); + match values_validity { Mask::AllTrue(_) => { - let mut decoded: BufferMut = BufferMut::with_capacity(length); - for (end, value) in run_ends.zip_eq(values) { + let mut decoded = BufferMut::::with_capacity(length + decode_chunk_len::()); + let base = decoded.spare_capacity_mut().as_mut_ptr(); + let mut pos = 0usize; + for (&end, &value) in run_ends.iter().zip(values) { + let end = trim_end(end, offset_e, length_e); assert!( - end >= decoded.len(), - "Runend ends must be monotonic, got {end} after {}", - decoded.len() + end >= pos, + "Runend ends must be monotonic, got {end} after {pos}" ); - assert!(end <= length, "Runend end must be less than overall length"); - // SAFETY: - // We preallocate enough capacity because we know the total length - unsafe { decoded.push_n_unchecked(*value, end - decoded.len()) }; + // SAFETY: pos <= end <= length and the buffer was allocated with one chunk + // of padding beyond `length`. + unsafe { splat_run(base, pos, end, value) }; + pos = end; } + // SAFETY: every element in 0..pos was initialized by the loop above. + unsafe { decoded.set_len(pos) }; (decoded.into(), values_nullability.into()) } Mask::AllFalse(_) => (Buffer::::zeroed(length), Validity::AllInvalid), Mask::Values(mask) => { - let mut decoded = BufferMut::with_capacity(length); - let mut decoded_validity = BitBufferMut::with_capacity(length); - for (end, value) in run_ends.zip_eq( - values - .iter() - .zip(mask.bit_buffer().iter()) - .map(|(&v, is_valid)| is_valid.then_some(v)), - ) { + let run_validity = mask.bit_buffer(); + assert_eq!( + run_ends.len(), + run_validity.len(), + "runend ends and validity must have equal lengths" + ); + // Prefill the element validity with the majority run validity; only minority + // runs then need their bit range rewritten. + let prefill = 2 * run_validity.true_count() > run_validity.len(); + let mut decoded_validity = BitBufferMut::full(prefill, length); + let mut decoded = BufferMut::::with_capacity(length + decode_chunk_len::()); + let base = decoded.spare_capacity_mut().as_mut_ptr(); + let mut pos = 0usize; + for ((&end, &value), is_valid) in run_ends.iter().zip(values).zip(run_validity.iter()) { + let end = trim_end(end, offset_e, length_e); assert!( - end >= decoded.len(), - "Runend ends must be monotonic, got {end} after {}", - decoded.len() + end >= pos, + "Runend ends must be monotonic, got {end} after {pos}" ); - assert!(end <= length, "Runend end must be less than overall length"); - match value { - None => { - decoded_validity.append_n(false, end - decoded.len()); - // SAFETY: - // We preallocate enough capacity because we know the total length - unsafe { decoded.push_n_unchecked(T::default(), end - decoded.len()) }; - } - Some(value) => { - decoded_validity.append_n(true, end - decoded.len()); - // SAFETY: - // We preallocate enough capacity because we know the total length - unsafe { decoded.push_n_unchecked(value, end - decoded.len()) }; - } + if is_valid != prefill && end > pos { + // SAFETY: pos <= end <= length == decoded_validity.len() + unsafe { decoded_validity.fill_range_unchecked(pos, end, is_valid) }; } + let value = if is_valid { value } else { T::default() }; + // SAFETY: pos <= end <= length and the buffer was allocated with one chunk + // of padding beyond `length`. + unsafe { splat_run(base, pos, end, value) }; + pos = end; } + // SAFETY: every element in 0..pos was initialized by the loop above. + unsafe { decoded.set_len(pos) }; + decoded_validity.truncate(pos); (decoded.into(), Validity::from(decoded_validity.freeze())) } } } -pub fn runend_decode_typed_primitive( - run_ends: impl Iterator, +/// Decode run-end encoded `values` with the given `run_ends` into a flat [`PrimitiveArray`]. +/// +/// Run ends are adjusted by `offset` and clamped to `length` while decoding. +pub fn runend_decode_typed_primitive( + run_ends: &[E], + offset: usize, values: &[T], values_validity: Mask, values_nullability: Nullability, @@ -276,6 +480,7 @@ pub fn runend_decode_typed_primitive( ) -> PrimitiveArray { let (decoded, validity) = runend_decode_slice( run_ends, + offset, values, values_validity, values_nullability, @@ -300,7 +505,8 @@ pub fn runend_decode_varbinview( let (decoded_views, validity) = match_each_unsigned_integer_ptype!(ends.ptype(), |E| { runend_decode_slice( - trimmed_ends_iter(ends.as_slice::(), offset, length), + ends.as_slice::(), + offset, views, validity_mask, values.dtype().nullability(), @@ -322,12 +528,18 @@ pub fn runend_decode_varbinview( mod tests { use std::sync::LazyLock; + use rand::RngExt; + use rand::SeedableRng; + use rand::rngs::StdRng; use vortex_array::VortexSessionExecute; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; + use vortex_array::dtype::NativePType; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; + use vortex_buffer::Buffer; use vortex_buffer::buffer; + use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -402,4 +614,159 @@ mod tests { assert_arrays_eq!(decoded, expected, &mut ctx); Ok(()) } + + #[test] + fn decode_with_offset() -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let ends = PrimitiveArray::from_iter([2u32, 5, 10]); + let values = PrimitiveArray::from_iter([1i32, 2, 3]); + let decoded = runend_decode_primitive(ends, values, 2, 6, &mut ctx)?; + + let expected = PrimitiveArray::from_iter(vec![2i32, 2, 2, 3, 3, 3]); + assert_arrays_eq!(decoded, expected, &mut ctx); + Ok(()) + } + + #[rstest::rstest] + #[case::mostly_valid(vec![Some(1i32), None, Some(3)])] + #[case::mostly_null(vec![None, Some(2i32), None])] + fn decode_nullable(#[case] run_values: Vec>) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let ends = PrimitiveArray::from_iter([2u32, 5, 10]); + let values = PrimitiveArray::from_option_iter(run_values.clone()); + let decoded = runend_decode_primitive(ends, values, 0, 10, &mut ctx)?; + + let mut expanded = Vec::new(); + let mut prev = 0usize; + for (&end, value) in [2usize, 5, 10].iter().zip(run_values) { + expanded.extend(std::iter::repeat_n(value, end - prev)); + prev = end; + } + let expected = PrimitiveArray::from_option_iter(expanded); + assert_arrays_eq!(decoded, expected, &mut ctx); + Ok(()) + } + + #[test] + fn decode_long_runs() -> VortexResult<()> { + // Runs longer than the splat chunk exercise the chunked-store loop's iteration path. + let mut ctx = SESSION.create_execution_ctx(); + let ends = PrimitiveArray::from_iter([100u32, 101, 300]); + let values = PrimitiveArray::from_iter([1i64, 2, 3]); + let decoded = runend_decode_primitive(ends, values, 0, 300, &mut ctx)?; + + let mut expanded = vec![1i64; 100]; + expanded.push(2); + expanded.extend(std::iter::repeat_n(3i64, 199)); + let expected = PrimitiveArray::from_iter(expanded); + assert_arrays_eq!(decoded, expected, &mut ctx); + Ok(()) + } + + /// Decode random runs with `runend_decode_primitive` and compare against a naive + /// element-by-element reference expansion. + /// + /// Random run lengths straddle every length-dependent branch in the kernel: the inline + /// chunk stores, the out-of-line long-run fill, and the doubling `memcpy` fill (2 KiB, + /// which even the widest tested element type only crosses past ~1024-element runs). Also + /// covers zero-length runs, offsets that partially trim the first run, and a final run + /// overshooting the logical length; validity covers non-nullable, mixed, and all-invalid. + fn check_decode_matches_reference>(seed: u64) -> VortexResult<()> { + let mut ctx = SESSION.create_execution_ctx(); + let mut rng = StdRng::seed_from_u64(seed); + + for &max_run_len in &[3usize, 40, 700, 3000] { + for nullable in [false, true] { + let length = rng.random_range(1..=5000); + let offset = if rng.random_bool(0.5) { + rng.random_range(0..50) + } else { + 0 + }; + let total = offset + length; + + let mut ends = Vec::new(); + let mut run_values: Vec = Vec::new(); + let mut run_valid = Vec::new(); + let mut pos = 0usize; + while pos < total { + let run_len = if !ends.is_empty() && rng.random_bool(0.05) { + // Zero-length runs exercise the empty-run edge of the splat loop. + 0 + } else if ends.is_empty() { + // The first run must reach past the offset. + rng.random_range(offset + 1..=offset + max_run_len) + } else { + rng.random_range(1..=max_run_len) + }; + pos += run_len; + ends.push(u32::try_from(pos).vortex_expect("test lengths fit in u32")); + // A quarter of runs take value zero, which is byte-uniform and so routes + // through the `memset` fill for every element width rather than the + // general path taken by arbitrary values. + let byte = if rng.random_bool(0.25) { + 0 + } else { + rng.random::() + }; + run_values.push(>::from(byte)); + run_valid.push(rng.random_bool(0.9)); + } + + let all_invalid = nullable && rng.random_bool(0.1); + let validity = if all_invalid { + Validity::AllInvalid + } else if nullable { + Validity::from(BitBuffer::from(run_valid.clone())) + } else { + Validity::NonNullable + }; + + let decoded = runend_decode_primitive( + PrimitiveArray::from_iter(ends.clone()), + PrimitiveArray::new(Buffer::from(run_values.clone()), validity), + offset, + length, + &mut ctx, + )?; + + let mut expected: Vec> = Vec::with_capacity(length); + for (&end, (value, valid)) in + ends.iter().zip(run_values.iter().zip(run_valid.iter())) + { + let end = (end as usize - offset).min(length); + let value = match (nullable, all_invalid, valid) { + (false, ..) => Some(*value), + (true, true, _) => None, + (true, false, valid) => valid.then_some(*value), + }; + expected.resize(end, value); + } + + if nullable { + let expected = PrimitiveArray::from_option_iter(expected); + assert_arrays_eq!(decoded, expected, &mut ctx); + } else { + let expected = PrimitiveArray::from_iter( + expected.into_iter().map(|v| v.vortex_expect("non-null")), + ); + assert_arrays_eq!(decoded, expected, &mut ctx); + } + } + } + Ok(()) + } + + #[rstest::rstest] + #[case(0)] + #[case(1)] + #[case(0x5eed)] + fn decode_matches_reference(#[case] seed: u64) -> VortexResult<()> { + check_decode_matches_reference::(seed)?; + check_decode_matches_reference::(seed)?; + check_decode_matches_reference::(seed)?; + check_decode_matches_reference::(seed)?; + check_decode_matches_reference::(seed)?; + Ok(()) + } }