From 3062ebbbbbccf06e01c7a254656008275783b0f7 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Fri, 7 Aug 2026 11:35:41 +0000 Subject: [PATCH 1/2] Skip per-split zone-mask expansion when covering zones are uniform `ZonedReader::pruning_evaluation` expanded the cached zone-level pruning mask into a row-aligned bit buffer for every split, then intersected it with the incoming mask. That cost a `Vec` of zone lengths (built eagerly, before the future was polled), a `BitBufferMut` of the full split length, a popcount over it, and a bitand, on every split of every scan. For most splits the zones covering that split are uniform: either none are pruned or all of them are. Both collapse to a constant stats mask, so counting the covered zone bits first - a handful of bit reads via `BitBuffer::count_range` - lets us skip the expansion entirely: - no covered zone pruned: the stats mask is all-true, so forward the incoming mask unchanged, - every covered zone pruned: return `Mask::new_false` directly, - otherwise: fall through to the existing expansion. The zone-length computation now lives on the non-uniform path, so uniform splits allocate nothing. Results and masks are unchanged. The equivalence is exact, not just bit-equal. `Mask::from_buffer` canonicalises an all-ones buffer to `AllTrue` and an all-zeros buffer to `AllFalse`, and owned-left `bitand` short circuits on both (`(_, AllOr::All) => self` and `(_, AllOr::None) => new_false`). So the old code already returned the incoming mask unchanged in the no-prune case, after an allocation and two O(n) passes to rediscover that. The new fast paths return the same `Mask` variant, not merely the same bits. No path in this change does more work than before; the only addition is a `count_range` over the covering zones on the non-uniform path. This is a simplification, not a demonstrated speedup. Two full benchmark sweeps were run over the same commit. Every suite returned "No clear signal", and four of the six suites common to both sweeps flipped sign: statpopgen -6.0% then +1.8%, clickbench-sorted -2.8% then +1.4%, tpcds -1.4% then +0.7%, polarsignals +1.7% then -2.0%. The apparent statpopgen win came from an anomalously slow baseline: q10 measured 4.26s base / 2.94s HEAD in the first sweep and 2.88s base / 3.12s HEAD in the second, so HEAD was stable and the base was not. A tight Parquet control does not rule this out - it bounds host drift, not per-path measurement variance. What holds up is the absence of a regression. The lowest-noise suites all read zero: random-access -0.1% (controls 0.98-1.04), appian +0.2% (all rows 0.98-1.02), string encoding +0.1% (sizes byte-identical), tpch sf=10 -0.3% (controls 0.96-1.03). CodSpeed reports no change across 1934 benchmarks. No suite in either sweep showed a credible regression. Signed-off-by: Joe Isaacs --- vortex-layout/src/layouts/zoned/reader.rs | 120 +++++++++++++++------- 1 file changed, 85 insertions(+), 35 deletions(-) diff --git a/vortex-layout/src/layouts/zoned/reader.rs b/vortex-layout/src/layouts/zoned/reader.rs index 52a57c1b92b..5a289cc8010 100644 --- a/vortex-layout/src/layouts/zoned/reader.rs +++ b/vortex-layout/src/layouts/zoned/reader.rs @@ -6,7 +6,6 @@ use std::ops::Range; use std::sync::Arc; use futures::future::BoxFuture; -use itertools::Itertools; use tracing::trace; use vortex_array::ArrayRef; use vortex_array::MaskFuture; @@ -14,8 +13,8 @@ use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::expr::BoundExpression; use vortex_buffer::BitBufferMut; -use vortex_error::VortexError; use vortex_error::VortexResult; +use vortex_mask::AllOr; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -97,13 +96,13 @@ impl ZonedReader { let zone_end = row_range.end.div_ceil(zone_len_u64); zone_start..zone_end } +} - /// Get the row index for the first row in a zone with the given `zone_index`. - pub(crate) fn first_row_offset(&self, zone_idx: u64) -> u64 { - zone_idx - .saturating_mul(self.zone_len as u64) - .min(self.row_count) - } +/// Get the row index for the first row in a zone with the given `zone_idx`. +/// +/// Free function so that it can be used from a `'static` future without capturing the reader. +fn first_row_offset(zone_idx: u64, zone_len: u64, row_count: u64) -> u64 { + zone_idx.saturating_mul(zone_len).min(row_count) } impl LayoutReader for ZonedReader { @@ -149,44 +148,64 @@ impl LayoutReader for ZonedReader { return Ok(data_eval); }; - let row_count = row_range.end - row_range.start; + let split_row_count = row_range.end - row_range.start; + let row_start = row_range.start; let zone_range = self.zone_range(row_range); - let zone_lengths: Vec<_> = zone_range - .clone() - .map(|zone_idx| { - // Figure out the range in the mask that corresponds to the zone - let start = usize::try_from( - self.first_row_offset(zone_idx) - .saturating_sub(row_range.start), - )?; - let end = usize::try_from( - self.first_row_offset(zone_idx + 1) - .saturating_sub(row_range.start) - .min(row_count), - )?; - Ok::<_, VortexError>(end - start) - }) - .try_collect()?; + let zone_start = usize::try_from(zone_range.start)?; + let zone_end = usize::try_from(zone_range.end)?; + let covered_zones = zone_end - zone_start; + let zone_len = self.zone_len as u64; + let layout_row_count = self.row_count; let name = Arc::clone(&self.name); let expr = expr.clone(); + let mask_len = mask.len(); - Ok(MaskFuture::new(mask.len(), async move { + Ok(MaskFuture::new(mask_len, async move { trace!("Invoking stats pruning evaluation {}: {}", name, expr); let pruning_mask = pruning_mask_future.await?.mask()?; - let mut builder = BitBufferMut::with_capacity(mask.len()); - for (zone_idx, &zone_length) in zone_range.clone().zip_eq(&zone_lengths) { - builder.append_n(!pruning_mask.value(usize::try_from(zone_idx)?), zone_length); - } + // Only the zones covering this row range matter. Counting their pruned bits is a + // handful of bit reads, and the overwhelming majority of splits are uniform: either + // no covered zone is pruned, or all of them are. Both collapse to a constant stats + // mask, so we can skip expanding zones into a row-aligned buffer entirely. + let pruned_zones = match pruning_mask.bit_buffer() { + AllOr::All => covered_zones, + AllOr::None => 0, + AllOr::Some(buffer) => buffer.count_range(zone_start, zone_end), + }; - let stats_mask = Mask::from(builder.freeze()); - assert_eq!(stats_mask.len(), mask.len(), "Mask length mismatch"); - - // Intersect the masks. let mask_density = mask.density(); - let mut stats_mask = mask.bitand(&stats_mask); + let mut stats_mask = if pruned_zones == 0 { + // The stats mask would be all-true, so intersecting it is a no-op. + mask + } else if pruned_zones == covered_zones { + // The stats mask would be all-false. + Mask::new_false(mask_len) + } else { + let mut builder = BitBufferMut::with_capacity(mask_len); + for zone_idx in zone_start..zone_end { + // Figure out the range in the mask that corresponds to the zone + let zone_idx_u64 = zone_idx as u64; + let start = usize::try_from( + first_row_offset(zone_idx_u64, zone_len, layout_row_count) + .saturating_sub(row_start), + )?; + let end = usize::try_from( + first_row_offset(zone_idx_u64 + 1, zone_len, layout_row_count) + .saturating_sub(row_start) + .min(split_row_count), + )?; + builder.append_n(!pruning_mask.value(zone_idx), end - start); + } + + let stats_mask = Mask::from(builder.freeze()); + assert_eq!(stats_mask.len(), mask_len, "Mask length mismatch"); + + // Intersect the masks. + mask.bitand(&stats_mask) + }; // Forward to data child for further pruning. if !stats_mask.all_false() { @@ -231,6 +250,7 @@ impl LayoutReader for ZonedReader { #[cfg(test)] mod test { use std::num::NonZeroUsize; + use std::ops::Range; use std::sync::Arc; use rstest::fixture; @@ -375,6 +395,36 @@ mod test { }) } + /// The zoned reader takes a uniform fast path when every zone covering the requested row + /// range agrees, so exercise all-pruned, all-kept and mixed ranges. + #[rstest] + #[case::only_pruned_zones(0..6, vec![false; 6])] + #[case::only_kept_zones(6..9, vec![true; 3])] + #[case::single_pruned_zone(3..6, vec![false; 3])] + #[case::partial_zones_mixed(1..8, vec![false, false, false, false, false, true, true])] + #[case::empty_range(4..4, vec![])] + fn test_stats_pruning_mask_zone_ranges( + #[from(stats_layout)] (segments, layout): (Arc, LayoutRef), + #[case] row_range: Range, + #[case] expected: Vec, + ) -> VortexResult<()> { + block_on(|handle| async { + let session = session_with_handle(handle); + let reader = layout.new_reader("".into(), segments, &session, &Default::default())?; + + // Values are 1..=9 in zones of 3, so `> 7` prunes zones 0 and 1 and keeps zone 2. + let expr = gt(root(), lit(7)).bind(reader.dtype())?; + let len = usize::try_from(row_range.end - row_range.start)?; + + let result = reader + .pruning_evaluation(&row_range, &expr, Mask::new_true(len))? + .await?; + + assert_eq!(result, Mask::from_iter(expected)); + Ok(()) + }) + } + #[test] fn test_default_zoned_null_count_pruning_mask() { let ctx = ArrayContext::empty(); From adb75ed9abd2b592b97b8916a651397440f05a5d Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 7 Aug 2026 16:00:28 +0000 Subject: [PATCH 2/2] Add a benchmark for the per-split zone-mask expansion The uniform fast path in `ZonedReader::pruning_evaluation` had no measurement behind it - the SQL suites cannot resolve a per-split cost against per-split decode, so two full sweeps disagreed with each other and neither confirmed nor refuted anything. This benchmarks the changed code directly. `expand` reproduces the old behaviour (build a row-aligned buffer, popcount it, intersect) and `uniform` is the fast path that counts the covering zone bits instead. The incoming mask is a `Values` mask, not `AllTrue`, so `bitand` takes the `(AllOr::Some, AllOr::All) => self` branch that the real code hits rather than the all-true short circuit. Medians on a c6id-class host, one zone per split at the default 8192-row zone and block length: no zone pruned 7.4 ns vs 147 ns expanded (~20x) all zones pruned 14.5 ns vs 151 ns expanded (~10x) The gap widens with split length - at 65536 rows over 8 zones the no-prune case is 7.4 ns against 420 ns. The absolute saving is roughly 140-400 ns per split per pruning expression, which is why it does not surface in a query-level benchmark: a split then decodes 8192 rows across every projected column. `mixed_zones` covers the case where the covering zones genuinely disagree and the expansion still has to run. It needs at least two zones to be meaningful, so it takes its own argument set - a one-zone split is uniform by construction and would silently measure a fast path. There the added `count_range` shows a small overhead, within a few percent and at the edge of what this environment resolves, so it is bounded rather than claimed to be free. Signed-off-by: Joe Isaacs --- Cargo.lock | 1 + vortex-layout/Cargo.toml | 5 + vortex-layout/benches/zone_mask_expansion.rs | 166 +++++++++++++++++++ 3 files changed, 172 insertions(+) create mode 100644 vortex-layout/benches/zone_mask_expansion.rs diff --git a/Cargo.lock b/Cargo.lock index 27cf5957fd4..7b74c38115a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10330,6 +10330,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..cfd1574baf7 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 } @@ -63,6 +64,10 @@ tokio = { workspace = true, features = ["rt", "macros"] } vortex-array = { path = "../vortex-array", features = ["_test-harness"] } vortex-io = { path = "../vortex-io", features = ["tokio"] } +[[bench]] +name = "zone_mask_expansion" +harness = false + [features] _test-harness = [] tokio = ["dep:tokio", "vortex-error/tokio"] diff --git a/vortex-layout/benches/zone_mask_expansion.rs b/vortex-layout/benches/zone_mask_expansion.rs new file mode 100644 index 00000000000..6440fd4a66d --- /dev/null +++ b/vortex-layout/benches/zone_mask_expansion.rs @@ -0,0 +1,166 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Benchmarks for the per-split zone-mask expansion in `ZonedReader::pruning_evaluation`. +//! +//! For every split, the reader turns the cached zone-level pruning mask into a row-aligned +//! stats mask and intersects it with the incoming mask. When every zone covering the split +//! agrees - none pruned, or all pruned - that expansion collapses to a constant, because +//! `Mask::from_buffer` canonicalises an all-ones buffer to `AllTrue` and an all-zeros buffer +//! to `AllFalse`, and owned-left `bitand` short circuits on both. The expanded buffer is +//! built, counted, and thrown away. +//! +//! `expand` is the old behaviour and `uniform` is the fast path that reads the covering zone +//! bits directly. `mixed` covers the case where the zones genuinely disagree and the +//! expansion still has to run, bounding what the added `count_range` costs when it cannot +//! short circuit. + +#![allow(clippy::cast_possible_truncation)] + +use std::hint::black_box; +use std::ops::BitAnd; + +use divan::Bencher; +use vortex_buffer::BitBuffer; +use vortex_buffer::BitBufferMut; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +fn main() { + divan::main(); +} + +/// `(split_len, zones_per_split)`. The default zone and block length are both 8192 rows, so +/// one zone per split is the common shape; the wider cases cover multi-zone splits. +const SPLITS: &[(usize, usize)] = &[(8_192, 1), (8_192, 4), (65_536, 8)]; + +/// Splits with at least two covering zones, so the zones can actually disagree. A one-zone +/// split is uniform by construction and would silently measure a fast path instead. +const MIXED_SPLITS: &[(usize, usize)] = &[(8_192, 4), (65_536, 8)]; + +/// A realistic incoming scan mask: `Values`, not `AllTrue`, so `bitand` takes the +/// `(AllOr::Some, AllOr::All) => self` branch rather than the all-true short circuit. +fn incoming(len: usize) -> Mask { + Mask::from_buffer(BitBuffer::from_iter((0..len).map(|i| i % 3 != 0))) +} + +fn zone_lengths(split_len: usize, zones: usize) -> Vec { + let per_zone = split_len / zones; + (0..zones) + .map(|z| { + if z == zones - 1 { + split_len - per_zone * (zones - 1) + } else { + per_zone + } + }) + .collect() +} + +/// The zone-level pruning mask, restricted to the zones covering one split. +fn pruning_mask(zones: usize, pruned: impl Fn(usize) -> bool) -> Mask { + Mask::from_buffer(BitBuffer::from_iter((0..zones).map(pruned))) +} + +/// Expand the covering zones into a row-aligned buffer, then intersect. This is what ran for +/// every split before the uniform fast path. +fn expand(mask: Mask, pruning_mask: &Mask, zone_lengths: &[usize]) -> Mask { + let mut builder = BitBufferMut::with_capacity(mask.len()); + for (zone_idx, &zone_length) in zone_lengths.iter().enumerate() { + builder.append_n(!pruning_mask.value(zone_idx), zone_length); + } + let stats_mask = Mask::from(builder.freeze()); + mask.bitand(&stats_mask) +} + +/// Count the pruned bits among the covering zones first, and skip the expansion when they +/// all agree. +fn uniform(mask: Mask, pruning_mask: &Mask, zone_lengths: &[usize]) -> Mask { + let covered = zone_lengths.len(); + let pruned = match pruning_mask.bit_buffer() { + AllOr::All => covered, + AllOr::None => 0, + AllOr::Some(buffer) => buffer.count_range(0, covered), + }; + + if pruned == 0 { + mask + } else if pruned == covered { + Mask::new_false(mask.len()) + } else { + expand(mask, pruning_mask, zone_lengths) + } +} + +/// No covering zone is pruned: the stats mask is all-true and the intersection is a no-op. +#[divan::bench(args = SPLITS)] +fn no_zone_pruned(bencher: Bencher, (split_len, zones): (usize, usize)) { + let lengths = zone_lengths(split_len, zones); + let pruning = pruning_mask(zones, |_| false); + let mask = incoming(split_len); + + bencher + .with_inputs(|| mask.clone()) + .bench_values(|mask| black_box(uniform(mask, &pruning, &lengths))); +} + +/// The same split through the old expansion, for comparison. +#[divan::bench(args = SPLITS)] +fn no_zone_pruned_expanded(bencher: Bencher, (split_len, zones): (usize, usize)) { + let lengths = zone_lengths(split_len, zones); + let pruning = pruning_mask(zones, |_| false); + let mask = incoming(split_len); + + bencher + .with_inputs(|| mask.clone()) + .bench_values(|mask| black_box(expand(mask, &pruning, &lengths))); +} + +/// Every covering zone is pruned: the stats mask is all-false and the split drops out. +#[divan::bench(args = SPLITS)] +fn all_zones_pruned(bencher: Bencher, (split_len, zones): (usize, usize)) { + let lengths = zone_lengths(split_len, zones); + let pruning = pruning_mask(zones, |_| true); + let mask = incoming(split_len); + + bencher + .with_inputs(|| mask.clone()) + .bench_values(|mask| black_box(uniform(mask, &pruning, &lengths))); +} + +/// The same split through the old expansion, for comparison. +#[divan::bench(args = SPLITS)] +fn all_zones_pruned_expanded(bencher: Bencher, (split_len, zones): (usize, usize)) { + let lengths = zone_lengths(split_len, zones); + let pruning = pruning_mask(zones, |_| true); + let mask = incoming(split_len); + + bencher + .with_inputs(|| mask.clone()) + .bench_values(|mask| black_box(expand(mask, &pruning, &lengths))); +} + +/// The covering zones disagree, so the expansion still runs and the added `count_range` is +/// the whole overhead. +#[divan::bench(args = MIXED_SPLITS)] +fn mixed_zones(bencher: Bencher, (split_len, zones): (usize, usize)) { + let lengths = zone_lengths(split_len, zones); + let pruning = pruning_mask(zones, |z| z % 2 == 0); + let mask = incoming(split_len); + + bencher + .with_inputs(|| mask.clone()) + .bench_values(|mask| black_box(uniform(mask, &pruning, &lengths))); +} + +/// The same split through the old expansion, for comparison. +#[divan::bench(args = MIXED_SPLITS)] +fn mixed_zones_expanded(bencher: Bencher, (split_len, zones): (usize, usize)) { + let lengths = zone_lengths(split_len, zones); + let pruning = pruning_mask(zones, |z| z % 2 == 0); + let mask = incoming(split_len); + + bencher + .with_inputs(|| mask.clone()) + .bench_values(|mask| black_box(expand(mask, &pruning, &lengths))); +}