From 457d57b4c4473bc79c42173bd4bab8f2ac58fafe Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 17 Aug 2026 10:46:10 +0100 Subject: [PATCH 1/4] bench: add zone map pruning benchmark Measures `ZoneMap::prune` end-to-end for the overhead beyond expression rewriting Signed-off-by: Robert Kruszewski Signed-off-by: Robert Kruszewski --- Cargo.lock | 1 + vortex-layout/Cargo.toml | 5 + vortex-layout/benches/zone_map_prune.rs | 294 ++++++++++++++++++++++++ 3 files changed, 300 insertions(+) create mode 100644 vortex-layout/benches/zone_map_prune.rs diff --git a/Cargo.lock b/Cargo.lock index b0f98336b0b..f602495d363 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10326,6 +10326,7 @@ dependencies = [ "async-stream", "async-trait", "bit-vec", + "codspeed-divan-compat", "flatbuffers", "futures", "insta", diff --git a/vortex-layout/Cargo.toml b/vortex-layout/Cargo.toml index f772b9ab639..a9a9f08c5cd 100644 --- a/vortex-layout/Cargo.toml +++ b/vortex-layout/Cargo.toml @@ -55,6 +55,7 @@ vortex-session = { workspace = true } vortex-utils = { workspace = true, features = ["dashmap"] } [dev-dependencies] +divan = { workspace = true } futures = { workspace = true, features = ["executor"] } insta = { workspace = true } rstest = { workspace = true } @@ -67,6 +68,10 @@ vortex-io = { path = "../vortex-io", features = ["tokio"] } _test-harness = [] tokio = ["dep:tokio", "vortex-error/tokio"] +[[bench]] +name = "zone_map_prune" +harness = false + [lints] workspace = true diff --git a/vortex-layout/benches/zone_map_prune.rs b/vortex-layout/benches/zone_map_prune.rs new file mode 100644 index 00000000000..3b6c5d3f4b7 --- /dev/null +++ b/vortex-layout/benches/zone_map_prune.rs @@ -0,0 +1,294 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Microbenchmarks for [`ZoneMap::prune`]. +//! +//! Each case pre-falsifies its predicate so the timed region covers exactly what `prune` does: +//! lowering the stats placeholders against the zone map, then evaluating the result per zone. + +#![expect(clippy::unwrap_used)] + +use std::sync::Arc; +use std::sync::LazyLock; + +use divan::Bencher; +use parking_lot::Mutex; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::aggregate_fn::AggregateFnRef; +use vortex_array::aggregate_fn::AggregateFnVTableExt; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::aggregate_fn::NumericalAggregateOpts; +use vortex_array::aggregate_fn::fns::max::Max; +use vortex_array::aggregate_fn::fns::min::Min; +use vortex_array::aggregate_fn::fns::nan_count::NanCount; +use vortex_array::aggregate_fn::fns::null_count::NullCount; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::Expression; +use vortex_array::expr::eq; +use vortex_array::expr::gt; +use vortex_array::expr::is_not_null; +use vortex_array::expr::lit; +use vortex_array::expr::or; +use vortex_array::expr::root; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_layout::layouts::zoned::zone_map::ZoneMap; +use vortex_layout::session::LayoutSession; +use vortex_session::VortexSession; +use vortex_utils::aliases::hash_map::HashMap; + +fn main() { + divan::main(); +} + +static SESSION: LazyLock = + LazyLock::new(|| vortex_array::array_session().with::()); + +/// Zone counts to sweep. The small case exposes per-call fixed cost, the large case exposes +/// per-zone evaluation cost. +const ZONE_COUNTS: &[usize] = &[16, 1024, 65536]; + +const ZONE_LEN: u64 = 8192; + +/// Deterministic pseudo-random values, so both branches benchmark identical data. +fn pseudo_random(len: usize, seed: u64) -> impl Iterator { + let mut state = seed | 1; + (0..len).map(move |_| { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + state + }) +} + +fn i32_stats(num_zones: usize) -> (Vec, Vec) { + let mins: Vec = pseudo_random(num_zones, 0x5eed) + .map(|v| (v % 10_000) as i32) + .collect(); + let maxs = mins.iter().map(|min| min + 100).collect(); + (mins, maxs) +} + +fn counts(num_zones: usize, seed: u64, modulus: u64) -> Buffer { + pseudo_random(num_zones, seed) + .map(|v| v % modulus) + .collect() +} + +/// The aggregates the zoned writer stores by default for a numeric column. `nan_count` only has a +/// state dtype for floats, so it is omitted for integers. +fn min_max_fields(column_dtype: &DType, num_zones: usize) -> Vec<(String, ArrayRef)> { + let (mins, maxs) = i32_stats(num_zones); + let max = Max.bind(NumericalAggregateOpts::skip_nans()); + let min = Min.bind(NumericalAggregateOpts::skip_nans()); + + let (min_array, max_array) = if column_dtype.is_float() { + ( + PrimitiveArray::new( + mins.iter().map(|v| f64::from(*v)).collect::>(), + Validity::AllValid, + ) + .into_array(), + PrimitiveArray::new( + maxs.iter().map(|v| f64::from(*v)).collect::>(), + Validity::AllValid, + ) + .into_array(), + ) + } else { + ( + PrimitiveArray::new( + mins.iter().copied().collect::>(), + Validity::AllValid, + ) + .into_array(), + PrimitiveArray::new( + maxs.iter().copied().collect::>(), + Validity::AllValid, + ) + .into_array(), + ) + }; + + vec![(max.to_string(), max_array), (min.to_string(), min_array)] +} + +fn count_fields(column_dtype: &DType, num_zones: usize) -> Vec<(String, ArrayRef)> { + let mut fields = Vec::new(); + if column_dtype.is_float() { + fields.push(( + NanCount.bind(EmptyOptions).to_string(), + PrimitiveArray::new(counts(num_zones, 0xfeed, 4), Validity::AllValid).into_array(), + )); + } + fields.push(( + NullCount.bind(EmptyOptions).to_string(), + PrimitiveArray::new( + counts(num_zones, 0xc0ffee, ZONE_LEN + 1), + Validity::AllValid, + ) + .into_array(), + )); + fields +} + +/// Key identifying a cached zone map: column dtype, number of zones, and whether min/max are +/// omitted. +type ZoneMapKey = (DType, usize, bool); + +/// Divan calls a benchmark function once per sample, so zone maps are cached to keep construction +/// out of both the reported times and the profile. +static ZONE_MAPS: LazyLock>> = LazyLock::new(Mutex::default); + +fn zone_map(column_dtype: DType, num_zones: usize, counts_only: bool) -> ZoneMap { + ZONE_MAPS + .lock() + .entry((column_dtype.clone(), num_zones, counts_only)) + .or_insert_with(|| { + let mut fields = if counts_only { + Vec::new() + } else { + min_max_fields(&column_dtype, num_zones) + }; + fields.extend(count_fields(&column_dtype, num_zones)); + build(column_dtype.clone(), fields, num_zones) + }) + .clone() +} + +/// A zone map carrying every aggregate the zoned writer stores by default. +fn numeric_zone_map(column_dtype: DType, num_zones: usize) -> ZoneMap { + zone_map(column_dtype, num_zones, false) +} + +/// A zone map carrying only the count aggregates, so min/max proofs find no stat to bind. +fn counts_only_zone_map(column_dtype: DType, num_zones: usize) -> ZoneMap { + zone_map(column_dtype, num_zones, true) +} + +fn build(column_dtype: DType, fields: Vec<(String, ArrayRef)>, num_zones: usize) -> ZoneMap { + let aggregate_fns: Arc<[AggregateFnRef]> = [ + Max.bind(NumericalAggregateOpts::skip_nans()), + Min.bind(NumericalAggregateOpts::skip_nans()), + NanCount.bind(EmptyOptions), + NullCount.bind(EmptyOptions), + ] + .into_iter() + .filter(|aggregate_fn| { + fields + .iter() + .any(|(name, _)| name == &aggregate_fn.to_string()) + }) + .collect(); + + let stats = StructArray::from_fields( + &fields + .iter() + .map(|(name, array)| (name.as_str(), array.clone())) + .collect::>(), + ) + .unwrap(); + + // A trailing short zone, which is the common shape and forces the run-end row-count array. + let row_count = ZONE_LEN * (num_zones as u64 - 1) + ZONE_LEN / 2; + ZoneMap::try_new(column_dtype, stats, aggregate_fns, ZONE_LEN, row_count).unwrap() +} + +fn i32_dtype() -> DType { + DType::Primitive(PType::I32, Nullability::Nullable) +} + +fn f64_dtype() -> DType { + DType::Primitive(PType::F64, Nullability::Nullable) +} + +fn falsify(expr: Expression, column_dtype: &DType) -> BoundExpression { + expr.bind(column_dtype) + .unwrap() + .falsify(&SESSION) + .unwrap() + .unwrap() +} + +fn run(bencher: Bencher, zone_map: ZoneMap, predicate: BoundExpression) { + bencher.bench(|| { + divan::black_box( + zone_map + .prune(divan::black_box(&predicate), &SESSION) + .unwrap(), + ) + }); +} + +/// Integer range predicate: binds to `max` only, no row count, no NaN guard. +#[divan::bench(args = ZONE_COUNTS)] +fn int_gt(bencher: Bencher, num_zones: usize) { + static PREDICATE: LazyLock = + LazyLock::new(|| falsify(gt(root(), lit(5_000i32)), &i32_dtype())); + run( + bencher, + numeric_zone_map(i32_dtype(), num_zones), + PREDICATE.clone(), + ); +} + +/// Float range predicate: the NaN-guarded and unguarded rules both fire, and on a zone map that +/// stores `nan_count` they lower to the same expression. +#[divan::bench(args = ZONE_COUNTS)] +fn float_gt(bencher: Bencher, num_zones: usize) { + static PREDICATE: LazyLock = + LazyLock::new(|| falsify(gt(root(), lit(5_000f64)), &f64_dtype())); + run( + bencher, + numeric_zone_map(f64_dtype(), num_zones), + PREDICATE.clone(), + ); +} + +/// Null predicate: lowers to `null_count == row_count`, exercising the row-count path. +#[divan::bench(args = ZONE_COUNTS)] +fn is_not_null_pred(bencher: Bencher, num_zones: usize) { + static PREDICATE: LazyLock = + LazyLock::new(|| falsify(is_not_null(root()), &i32_dtype())); + run( + bencher, + numeric_zone_map(i32_dtype(), num_zones), + PREDICATE.clone(), + ); +} + +/// A 16-term `OR` chain, which is where lowering cost grows relative to evaluation cost. +#[divan::bench(args = ZONE_COUNTS)] +fn or_chain(bencher: Bencher, num_zones: usize) { + static PREDICATE: LazyLock = LazyLock::new(|| { + let expr = (0..16i32) + .map(|i| eq(root(), lit(i * 500))) + .reduce(or) + .unwrap(); + falsify(expr, &i32_dtype()) + }); + run( + bencher, + numeric_zone_map(i32_dtype(), num_zones), + PREDICATE.clone(), + ); +} + +/// The zone map lacks min/max, so every proof binds to a null literal and the lowered predicate is +/// constant. +#[divan::bench(args = ZONE_COUNTS)] +fn missing_stats(bencher: Bencher, num_zones: usize) { + static PREDICATE: LazyLock = + LazyLock::new(|| falsify(gt(root(), lit(5_000i32)), &i32_dtype())); + run( + bencher, + counts_only_zone_map(i32_dtype(), num_zones), + PREDICATE.clone(), + ); +} From 97c6412ae9aec8b23057b6fa6cbd9160759b1759 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Mon, 17 Aug 2026 19:10:05 +0100 Subject: [PATCH 2/4] fixes Signed-off-by: Robert Kruszewski --- .github/workflows/codspeed.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index 13393072b8c..b61df171865 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -55,7 +55,7 @@ jobs: - { shard: 4, name: "Encodings 1", packages: "vortex-alp vortex-bytebool vortex-datetime-parts" } - { shard: 5, name: "Encodings 2", packages: "vortex-decimal-byte-parts vortex-fastlanes vortex-fsst", features: "--features _test-harness" } - { shard: 6, name: "Encodings 3", packages: "vortex-pco vortex-runend vortex-sequence" } - - { shard: 7, name: "Encodings 4", packages: "vortex-sparse vortex-zigzag vortex-zstd" } + - { shard: 7, name: "Encodings 4 & layout", packages: "vortex-sparse vortex-zigzag vortex-zstd vortex-layout" } - { shard: 8, name: "Storage formats & row encoding", packages: "vortex-flatbuffers vortex-proto vortex-btrblocks vortex-row" } - { shard: 9, name: "Tensor & spatial", packages: "vortex-tensor vortex-spatial" } name: "Benchmark with Codspeed (Shard #${{ matrix.shard }})" From 802ceb7865462cb5cc1241ab73ba0d78fc1530a1 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Tue, 18 Aug 2026 13:50:30 +0100 Subject: [PATCH 3/4] less- --- vortex-layout/benches/zone_map_prune.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vortex-layout/benches/zone_map_prune.rs b/vortex-layout/benches/zone_map_prune.rs index 3b6c5d3f4b7..a441d6b436b 100644 --- a/vortex-layout/benches/zone_map_prune.rs +++ b/vortex-layout/benches/zone_map_prune.rs @@ -52,7 +52,7 @@ static SESSION: LazyLock = /// Zone counts to sweep. The small case exposes per-call fixed cost, the large case exposes /// per-zone evaluation cost. -const ZONE_COUNTS: &[usize] = &[16, 1024, 65536]; +const ZONE_COUNTS: &[usize] = &[16, 1024, 8192]; const ZONE_LEN: u64 = 8192; From 5edb54d63f0a6b7cd7b314fd1765b026c1a5b3c5 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Wed, 19 Aug 2026 10:42:46 +0100 Subject: [PATCH 4/4] bench Signed-off-by: Robert Kruszewski --- vortex-layout/benches/zone_map_prune.rs | 86 ++++++++++++++++++------- 1 file changed, 63 insertions(+), 23 deletions(-) diff --git a/vortex-layout/benches/zone_map_prune.rs b/vortex-layout/benches/zone_map_prune.rs index a441d6b436b..87eb21fa9e4 100644 --- a/vortex-layout/benches/zone_map_prune.rs +++ b/vortex-layout/benches/zone_map_prune.rs @@ -75,62 +75,101 @@ fn i32_stats(num_zones: usize) -> (Vec, Vec) { (mins, maxs) } -fn counts(num_zones: usize, seed: u64, modulus: u64) -> Buffer { - pseudo_random(num_zones, seed) - .map(|v| v % modulus) - .collect() +struct ZoneCounts { + null: Vec, + nan: Option>, + has_min_max: Vec, +} + +fn zone_counts(column_dtype: &DType, num_zones: usize) -> ZoneCounts { + let row_counts = (0..num_zones) + .map(|zone| { + if zone + 1 == num_zones { + ZONE_LEN / 2 + } else { + ZONE_LEN + } + }) + .collect::>(); + let null = pseudo_random(num_zones, 0xc0ffee) + .zip(&row_counts) + .map(|(value, row_count)| value % (row_count + 1)) + .collect::>(); + let nan = column_dtype.is_float().then(|| { + pseudo_random(num_zones, 0xfeed) + .zip(row_counts.iter().zip(&null)) + .map(|(value, (row_count, null_count))| value % (row_count - null_count + 1)) + .collect::>() + }); + let has_min_max = row_counts + .iter() + .enumerate() + .map(|(zone, row_count)| { + null[zone] + nan.as_ref().map_or(0, |counts| counts[zone]) < *row_count + }) + .collect(); + + ZoneCounts { + null, + nan, + has_min_max, + } } /// The aggregates the zoned writer stores by default for a numeric column. `nan_count` only has a /// state dtype for floats, so it is omitted for integers. -fn min_max_fields(column_dtype: &DType, num_zones: usize) -> Vec<(String, ArrayRef)> { +fn min_max_fields( + column_dtype: &DType, + num_zones: usize, + has_min_max: &[bool], +) -> Vec<(String, ArrayRef)> { let (mins, maxs) = i32_stats(num_zones); let max = Max.bind(NumericalAggregateOpts::skip_nans()); let min = Min.bind(NumericalAggregateOpts::skip_nans()); + let min_validity = Validity::from_iter(has_min_max.iter().copied()); + let max_validity = Validity::from_iter(has_min_max.iter().copied()); let (min_array, max_array) = if column_dtype.is_float() { ( PrimitiveArray::new( mins.iter().map(|v| f64::from(*v)).collect::>(), - Validity::AllValid, + min_validity, ) .into_array(), PrimitiveArray::new( maxs.iter().map(|v| f64::from(*v)).collect::>(), - Validity::AllValid, + max_validity, ) .into_array(), ) } else { ( - PrimitiveArray::new( - mins.iter().copied().collect::>(), - Validity::AllValid, - ) - .into_array(), - PrimitiveArray::new( - maxs.iter().copied().collect::>(), - Validity::AllValid, - ) - .into_array(), + PrimitiveArray::new(mins.iter().copied().collect::>(), min_validity) + .into_array(), + PrimitiveArray::new(maxs.iter().copied().collect::>(), max_validity) + .into_array(), ) }; vec![(max.to_string(), max_array), (min.to_string(), min_array)] } -fn count_fields(column_dtype: &DType, num_zones: usize) -> Vec<(String, ArrayRef)> { +fn count_fields(counts: &ZoneCounts) -> Vec<(String, ArrayRef)> { let mut fields = Vec::new(); - if column_dtype.is_float() { + if let Some(nan) = &counts.nan { fields.push(( NanCount.bind(EmptyOptions).to_string(), - PrimitiveArray::new(counts(num_zones, 0xfeed, 4), Validity::AllValid).into_array(), + PrimitiveArray::new( + nan.iter().copied().collect::>(), + Validity::AllValid, + ) + .into_array(), )); } fields.push(( NullCount.bind(EmptyOptions).to_string(), PrimitiveArray::new( - counts(num_zones, 0xc0ffee, ZONE_LEN + 1), + counts.null.iter().copied().collect::>(), Validity::AllValid, ) .into_array(), @@ -151,12 +190,13 @@ fn zone_map(column_dtype: DType, num_zones: usize, counts_only: bool) -> ZoneMap .lock() .entry((column_dtype.clone(), num_zones, counts_only)) .or_insert_with(|| { + let counts = zone_counts(&column_dtype, num_zones); let mut fields = if counts_only { Vec::new() } else { - min_max_fields(&column_dtype, num_zones) + min_max_fields(&column_dtype, num_zones, &counts.has_min_max) }; - fields.extend(count_fields(&column_dtype, num_zones)); + fields.extend(count_fields(&counts)); build(column_dtype.clone(), fields, num_zones) }) .clone()