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))); +} 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();