From cd4231c73da01759baac37458b1da770d0f506af Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Tue, 15 Sep 2026 20:52:50 +0100 Subject: [PATCH] bench(pco): per-element scalar reads across and within chunks Two divan benchmarks for `execute_scalar` on a PCO array of 16 chunks with 1024-value pages: one read in every chunk in shuffled order, and 32 reads inside a single chunk. Each read decodes the page holding its value, so both patterns cost about the same per read today; the second is the case a retained scalar probe can amortise. Signed-off-by: Joe Isaacs Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_016CqrLKgPqFYGZK5sjk1qe7 --- Cargo.lock | 2 + encodings/pco/Cargo.toml | 5 ++ encodings/pco/benches/scalar_at.rs | 85 ++++++++++++++++++++++++++++++ 3 files changed, 92 insertions(+) create mode 100644 encodings/pco/benches/scalar_at.rs diff --git a/Cargo.lock b/Cargo.lock index a2dd783356b..758cbf88f81 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11411,8 +11411,10 @@ dependencies = [ name = "vortex-pco" version = "0.1.0" dependencies = [ + "codspeed-divan-compat", "pco", "prost 0.14.4", + "rand 0.10.2", "rstest", "vortex-array", "vortex-arrow", diff --git a/encodings/pco/Cargo.toml b/encodings/pco/Cargo.toml index 1683d4875e9..9035953501c 100644 --- a/encodings/pco/Cargo.toml +++ b/encodings/pco/Cargo.toml @@ -26,7 +26,12 @@ vortex-mask = { workspace = true } vortex-session = { workspace = true } [dev-dependencies] +divan = { workspace = true } +rand = { workspace = true } rstest = { workspace = true } vortex-array = { workspace = true, features = ["_test-harness"] } vortex-arrow = { workspace = true } +[[bench]] +name = "scalar_at" +harness = false diff --git a/encodings/pco/benches/scalar_at.rs b/encodings/pco/benches/scalar_at.rs new file mode 100644 index 00000000000..6c14923f9e3 --- /dev/null +++ b/encodings/pco/benches/scalar_at.rs @@ -0,0 +1,85 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Per-element reads from a PCO array through `execute_scalar`. +//! +//! PCO stores values in chunks of pages, and a scalar read decodes the page holding the value. +//! Two access patterns bracket the cost: one read per chunk, so every read lands on a different +//! chunk and page, and many reads inside one chunk, where reads share pages. + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use rand::RngExt; +use rand::SeedableRng; +use rand::prelude::StdRng; +use rand::seq::SliceRandom; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::PrimitiveArray; +use vortex_pco::Pco; +use vortex_session::VortexSession; + +fn main() { + divan::main(); +} + +/// Values per PCO chunk; `Pco::from_primitive` splits input at this granularity. +const VALUES_PER_CHUNK: usize = pco::DEFAULT_MAX_PAGE_N; +const VALUES_PER_PAGE: usize = 1024; +const CHUNKS: usize = 16; +// Sized to keep the CodSpeed simulation under a millisecond per benchmark. +const READS_IN_CHUNK: usize = 32; + +static SESSION: LazyLock = LazyLock::new(vortex_array::array_session); + +fn pco_array() -> ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); + let len = u32::try_from(CHUNKS * VALUES_PER_CHUNK).unwrap(); + let values = PrimitiveArray::from_iter((0..len).map(|i| i / 7)); + Pco::from_primitive(values.as_view(), 3, VALUES_PER_PAGE, &mut ctx) + .unwrap() + .into_array() +} + +/// One index inside every chunk, visited in a shuffled order. +fn one_index_per_chunk() -> Vec { + let mut rng = StdRng::seed_from_u64(0); + let mut indices: Vec = (0..CHUNKS) + .map(|chunk| chunk * VALUES_PER_CHUNK + rng.random_range(0..VALUES_PER_CHUNK)) + .collect(); + indices.shuffle(&mut rng); + indices +} + +/// Many indices inside one chunk in the middle of the array. +fn indices_in_one_chunk() -> Vec { + let mut rng = StdRng::seed_from_u64(1); + let chunk = CHUNKS / 2; + (0..READS_IN_CHUNK) + .map(|_| chunk * VALUES_PER_CHUNK + rng.random_range(0..VALUES_PER_CHUNK)) + .collect() +} + +fn read_all(bencher: Bencher, array: ArrayRef, indices: Vec) { + bencher + .with_inputs(|| SESSION.create_execution_ctx()) + .bench_refs(|ctx| { + for &index in &indices { + divan::black_box(array.execute_scalar(divan::black_box(index), ctx).unwrap()); + } + }); +} + +#[divan::bench] +fn scalar_at_one_per_chunk(bencher: Bencher) { + read_all(bencher, pco_array(), one_index_per_chunk()); +} + +#[divan::bench] +fn scalar_at_within_one_chunk(bencher: Bencher) { + read_all(bencher, pco_array(), indices_in_one_chunk()); +}