Skip per-split zone-mask expansion when covering zones are uniform

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"]

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

vortex-layout/src/layouts/zoned/reader.rs

@@ -6,16 +6,15 @@ 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;
 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<dyn SegmentSource>, LayoutRef),
+        #[case] row_range: Range<u64>,
+        #[case] expected: Vec<bool>,
+    ) -> 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();