diff --git a/encodings/runend/Cargo.toml b/encodings/runend/Cargo.toml index 01a5b8d7a3e..268aa6ac3ca 100644 --- a/encodings/runend/Cargo.toml +++ b/encodings/runend/Cargo.toml @@ -52,3 +52,7 @@ harness = false [[bench]] name = "run_end_decode" harness = false + +[[bench]] +name = "run_end_take" +harness = false diff --git a/encodings/runend/benches/run_end_take.rs b/encodings/runend/benches/run_end_take.rs new file mode 100644 index 00000000000..5c5c75a712a --- /dev/null +++ b/encodings/runend/benches/run_end_take.rs @@ -0,0 +1,168 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![expect(clippy::unwrap_used)] + +use std::fmt; +use std::sync::LazyLock; + +use divan::Bencher; +use rand::RngExt; +use rand::SeedableRng; +use rand::rngs::StdRng; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::RecursiveCanonical; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::session::ArraySession; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_runend::RunEnd; +use vortex_session::VortexSession; + +fn main() { + divan::main(); +} + +static SESSION: LazyLock = + LazyLock::new(|| VortexSession::empty().with::()); + +#[derive(Clone, Copy)] +enum IndexPattern { + SortedEven, + ReverseDense, + Random, +} + +#[derive(Clone, Copy)] +enum IndexValidity { + NonNullable, + EveryFourthNull, + AllNull, +} + +#[derive(Clone, Copy)] +struct TakeBenchArgs { + name: &'static str, + array_len: usize, + run_step: usize, + take_len: usize, + pattern: IndexPattern, + validity: IndexValidity, +} + +impl fmt::Display for TakeBenchArgs { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "{}_len{}_run{}_take{}", + self.name, self.array_len, self.run_step, self.take_len + ) + } +} + +const BENCH_ARGS: &[TakeBenchArgs] = &[ + // Sorted sparse takes should use the single-pass linear scan over run ends. + TakeBenchArgs { + name: "sorted_linear", + array_len: 16_384, + run_step: 8, + take_len: 512, + pattern: IndexPattern::SortedEven, + validity: IndexValidity::NonNullable, + }, + // Dense unsorted takes should build the logical-position-to-run table. + TakeBenchArgs { + name: "dense_table", + array_len: 8_192, + run_step: 1, + take_len: 2_048, + pattern: IndexPattern::ReverseDense, + validity: IndexValidity::NonNullable, + }, + // Sparse unsorted takes below the large-run threshold should stay on binary search. + TakeBenchArgs { + name: "binary_sparse", + array_len: 65_536, + run_step: 4, + take_len: 512, + pattern: IndexPattern::Random, + validity: IndexValidity::NonNullable, + }, + // Nullable indices exercise masked stats and table lookup. + TakeBenchArgs { + name: "nullable_dense_table", + array_len: 8_192, + run_step: 1, + take_len: 2_048, + pattern: IndexPattern::ReverseDense, + validity: IndexValidity::EveryFourthNull, + }, + // All-null indices should return before touching run ends or values. + TakeBenchArgs { + name: "all_null", + array_len: 65_536, + run_step: 4, + take_len: 2_048, + pattern: IndexPattern::Random, + validity: IndexValidity::AllNull, + }, +]; + +#[divan::bench(args = BENCH_ARGS)] +fn take(bencher: Bencher, args: TakeBenchArgs) { + let array = run_end_array(args.array_len, args.run_step); + let indices = take_indices(args); + + bencher + .with_inputs(|| (&array, &indices, SESSION.create_execution_ctx())) + .bench_refs(|(array, indices, execution_ctx)| { + array + .take(indices.clone()) + .unwrap() + .execute::(execution_ctx) + .unwrap() + }); +} + +fn run_end_array(len: usize, run_step: usize) -> ArrayRef { + let num_runs = len.div_ceil(run_step); + let ends = (0..num_runs) + .map(|run_idx| ((run_idx + 1) * run_step).min(len) as u64) + .collect::>() + .into_array(); + let values = PrimitiveArray::from_iter((0..num_runs).map(|idx| idx as u64)).into_array(); + + RunEnd::new(ends, values, &mut SESSION.create_execution_ctx()).into_array() +} + +fn take_indices(args: TakeBenchArgs) -> ArrayRef { + let values = index_values(args); + let validity = match args.validity { + IndexValidity::NonNullable => Validity::NonNullable, + IndexValidity::EveryFourthNull => { + Validity::from_iter((0..args.take_len).map(|i| !i.is_multiple_of(4))) + } + IndexValidity::AllNull => Validity::AllInvalid, + }; + + PrimitiveArray::new(values, validity).into_array() +} + +fn index_values(args: TakeBenchArgs) -> Buffer { + let values = match args.pattern { + IndexPattern::SortedEven => (0..args.take_len) + .map(|idx| ((idx * args.array_len) / args.take_len) as u64) + .collect::>(), + IndexPattern::ReverseDense => (0..args.take_len).rev().map(|idx| idx as u64).collect(), + IndexPattern::Random => { + let mut rng = StdRng::seed_from_u64(0); + (0..args.take_len) + .map(|_| rng.random_range(0..args.array_len) as u64) + .collect() + } + }; + + Buffer::from(values) +}