From 7ddb2b2eea6632c36b4f948b738f5e0120577d84 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 14 Aug 2026 12:03:46 +0100 Subject: [PATCH 1/5] simpler stats rewrites Signed-off-by: Robert Kruszewski --- vortex-array/src/scalar_fn/fns/is_not_null.rs | 14 +- vortex-array/src/scalar_fn/fns/is_null.rs | 6 +- vortex-array/src/scalar_fn/internal/mod.rs | 4 - .../src/scalar_fn/internal/row_count.rs | 174 ------------------ vortex-array/src/scalar_fn/mod.rs | 1 - vortex-array/src/stats/bind.rs | 157 +++++++++++++++- vortex-array/src/stats/rewrite/builtins.rs | 125 +++---------- vortex-file/src/pruning.rs | 11 +- vortex-file/src/tests.rs | 43 +++++ vortex-layout/src/layouts/zoned/zone_map.rs | 129 ++++++------- 10 files changed, 290 insertions(+), 374 deletions(-) delete mode 100644 vortex-array/src/scalar_fn/internal/mod.rs delete mode 100644 vortex-array/src/scalar_fn/internal/row_count.rs diff --git a/vortex-array/src/scalar_fn/fns/is_not_null.rs b/vortex-array/src/scalar_fn/fns/is_not_null.rs index 49bb9937b13..68a30771b3d 100644 --- a/vortex-array/src/scalar_fn/fns/is_not_null.rs +++ b/vortex-array/src/scalar_fn/fns/is_not_null.rs @@ -116,19 +116,13 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; use crate::expr::col; - use crate::expr::eq; use crate::expr::get_item; use crate::expr::is_not_null; - use crate::expr::or; use crate::expr::root; use crate::expr::test_harness; use crate::scalar::Scalar; - use crate::scalar_fn::EmptyOptions; - use crate::scalar_fn::ScalarFnVTableExt; - use crate::scalar_fn::internal::row_count::RowCount; use crate::stats::StatsSession; use crate::stats::all_null; - use crate::stats::null_count; static STATS_SESSION: LazyLock = LazyLock::new(|| VortexSession::empty().with::()); @@ -261,13 +255,7 @@ mod tests { assert_eq!( expr.bind(&dtype)?.falsify(&STATS_SESSION)?, - Some( - or( - eq(null_count(col("a")), RowCount.new_expr(EmptyOptions, []),), - all_null(col("a")), - ) - .bind(&dtype)? - ) + Some(all_null(col("a")).bind(&dtype)?) ); Ok(()) } diff --git a/vortex-array/src/scalar_fn/fns/is_null.rs b/vortex-array/src/scalar_fn/fns/is_null.rs index 87aca6748e5..253b527e9e5 100644 --- a/vortex-array/src/scalar_fn/fns/is_null.rs +++ b/vortex-array/src/scalar_fn/fns/is_null.rs @@ -106,17 +106,13 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; use crate::expr::col; - use crate::expr::eq; use crate::expr::get_item; use crate::expr::is_null; - use crate::expr::lit; - use crate::expr::or; use crate::expr::root; use crate::expr::test_harness; use crate::scalar::Scalar; use crate::stats::StatsSession; use crate::stats::all_non_null; - use crate::stats::null_count; static STATS_SESSION: LazyLock = LazyLock::new(|| VortexSession::empty().with::()); @@ -242,7 +238,7 @@ mod tests { assert_eq!( expr.bind(&dtype)?.falsify(&STATS_SESSION)?, - Some(or(eq(null_count(col("a")), lit(0u64)), all_non_null(col("a")),).bind(&dtype)?) + Some(all_non_null(col("a")).bind(&dtype)?) ); Ok(()) } diff --git a/vortex-array/src/scalar_fn/internal/mod.rs b/vortex-array/src/scalar_fn/internal/mod.rs deleted file mode 100644 index c91a7541c8a..00000000000 --- a/vortex-array/src/scalar_fn/internal/mod.rs +++ /dev/null @@ -1,4 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -pub mod row_count; diff --git a/vortex-array/src/scalar_fn/internal/row_count.rs b/vortex-array/src/scalar_fn/internal/row_count.rs deleted file mode 100644 index 290378c30a7..00000000000 --- a/vortex-array/src/scalar_fn/internal/row_count.rs +++ /dev/null @@ -1,174 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::fmt::Formatter; - -use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; -use vortex_array::arrays::ScalarFn; -use vortex_array::arrays::scalar_fn::ExactScalarFn; -use vortex_array::arrays::scalar_fn::ScalarFnArrayExt; -use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::dtype::PType; -use vortex_array::expr::display::ExprDisplay; -use vortex_array::scalar_fn::Arity; -use vortex_array::scalar_fn::ChildName; -use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ExecutionArgs; -use vortex_array::scalar_fn::ScalarFnId; -use vortex_array::scalar_fn::ScalarFnVTable; -use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_error::vortex_ensure; -use vortex_session::registry::CachedId; - -/// Zero-argument placeholder for the row count of the current evaluation scope. -/// -/// This is a legacy pruning hack for readers that only have a `null_count` -/// stat and need to support `is_not_null` pruning. It is currently substituted -/// by the zoned/file stats pruning paths before execution. New stats rewrites -/// should prefer boolean `all_null` and `all_non_null` aggregates instead of -/// depending on this scope-level placeholder. -/// -/// This expression *MUST* be replaced with a concrete array before evaluation. -/// Currently, the rewrite only happens in the context of stats pruning. -/// -/// `RowCount` is emitted while building pruning predicates that need a -/// scope-level value which is not stored as a regular stats column, such as the -/// row count of the current file or zone. The layer that owns that scope must -/// replace each placeholder with a concrete array via [`substitute_row_count`] -/// before evaluation. -/// -/// Calling [`ScalarFnVTable::execute`] directly returns an error because this -/// node is only a marker in a lazy expression tree. -#[derive(Clone)] -pub struct RowCount; - -impl ScalarFnVTable for RowCount { - type Options = EmptyOptions; - - fn id(&self) -> ScalarFnId { - static ID: CachedId = CachedId::new("vortex.row_count"); - *ID - } - - fn arity(&self, _options: &Self::Options) -> Arity { - Arity::Exact(0) - } - - fn child_name(&self, _options: &Self::Options, _child_idx: usize) -> ChildName { - unreachable!("RowCount has arity 0") - } - - fn fmt_sql( - &self, - _options: &Self::Options, - _expr: &dyn ExprDisplay, - f: &mut Formatter<'_>, - ) -> std::fmt::Result { - write!(f, "row_count()") - } - - fn return_dtype(&self, _options: &Self::Options, _args: &[DType]) -> VortexResult { - Ok(DType::Primitive(PType::U64, Nullability::NonNullable)) - } - - fn execute( - &self, - _options: &Self::Options, - _args: &dyn ExecutionArgs, - _ctx: &mut ExecutionCtx, - ) -> VortexResult { - vortex_bail!("RowCount must be substituted before evaluation") - } - - fn is_strict(&self, _options: &Self::Options) -> bool { - true - } - - fn is_fallible(&self, _options: &Self::Options) -> bool { - false - } -} - -/// Returns whether `array` contains a [`RowCount`] placeholder. -/// -/// Traversal is limited to lazy [`ScalarFnArray`] nodes produced by -/// [`ArrayRef::apply`][crate::ArrayRef::apply]. Other arrays are evaluation -/// leaves and cannot contain unevaluated placeholders. -/// -/// [`ScalarFnArray`]: vortex_array::arrays::ScalarFnArray -pub fn contains_row_count(array: &ArrayRef) -> bool { - if array.is::>() { - return true; - } - match array.as_opt::() { - Some(view) => view.iter_children().any(contains_row_count), - None => false, - } -} - -/// Replaces every [`RowCount`] placeholder with `replacement`. -/// -/// The replacement must have the same dtype and length as each placeholder. -/// Lazy [`ScalarFnArray`] ancestors are rewritten through slot take/put so -/// unaffected children are preserved, while non-[`ScalarFn`] arrays are returned -/// unchanged. -/// -/// [`ScalarFnArray`]: vortex_array::arrays::ScalarFnArray -pub fn substitute_row_count(array: ArrayRef, replacement: &ArrayRef) -> VortexResult { - if array.is::>() { - vortex_ensure!( - replacement.len() == array.len(), - "RowCount replacement length {} does not match scope length {}", - replacement.len(), - array.len(), - ); - vortex_ensure!( - replacement.dtype() == array.dtype(), - "RowCount replacement dtype {} does not match scope dtype {}", - replacement.dtype(), - array.dtype(), - ); - return Ok(replacement.clone()); - } - - if !array.is::() { - return Ok(array); - } - - let nchildren = array.nchildren(); - let mut array = array; - for slot_idx in 0..nchildren { - // SAFETY: `substitute_row_count` always returns an array with the same dtype and - // length as its input — `RowCount` placeholders are replaced with a checked - // replacement (same dtype and length), and `ScalarFn` recursion preserves both by - // operating on each slot in place. - let (taken, child) = unsafe { array.take_slot_unchecked(slot_idx)? }; - let new_child = substitute_row_count(child, replacement)?; - array = unsafe { taken.put_slot_unchecked(slot_idx, new_child)? }; - } - Ok(array) -} - -#[cfg(test)] -mod tests { - use vortex_array::dtype::DType; - use vortex_array::dtype::Nullability; - use vortex_array::dtype::PType; - - use crate::scalar_fn::EmptyOptions; - use crate::scalar_fn::internal::row_count::RowCount; - use crate::scalar_fn::vtable::ScalarFnVTableExt; - - #[test] - fn row_count_helper_dtype() { - let expr = RowCount.new_expr(EmptyOptions, []); - assert_eq!( - expr.return_dtype(&DType::Primitive(PType::I32, Nullability::Nullable)) - .unwrap(), - DType::Primitive(PType::U64, Nullability::NonNullable), - ); - } -} diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index 6be34ce1f34..423b8d14bc5 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -48,7 +48,6 @@ pub mod unstable; pub(crate) mod unstable; pub mod fns; -pub mod internal; pub mod session; /// A unique identifier for a scalar function. diff --git a/vortex-array/src/stats/bind.rs b/vortex-array/src/stats/bind.rs index e07a588de94..7c36ee9ca02 100644 --- a/vortex-array/src/stats/bind.rs +++ b/vortex-array/src/stats/bind.rs @@ -12,16 +12,31 @@ //! by a caller: zone-map field references, file-level stat literals, or typed nulls for missing //! stats. This lets all callers share the same falsification rules while keeping layout-specific //! stat storage behind [`StatBinder`]. +//! +//! Binding is also where scope-level quantities enter the predicate. The boolean aggregates +//! `all_null`, `all_non_null`, `all_nan`, and `all_non_nan` are derived from the corresponding +//! count statistic when a binder does not store them directly, which for the "all" variants needs +//! the number of rows the scope covers — see [`StatBinder::bind_row_count`]. use vortex_error::VortexResult; use crate::aggregate_fn::AggregateFnRef; +use crate::aggregate_fn::AggregateFnVTableExt; +use crate::aggregate_fn::EmptyOptions; +use crate::aggregate_fn::fns::all_nan::AllNan; +use crate::aggregate_fn::fns::all_non_nan::AllNonNan; +use crate::aggregate_fn::fns::all_non_null::AllNonNull; +use crate::aggregate_fn::fns::all_null::AllNull; +use crate::aggregate_fn::fns::nan_count::NanCount; +use crate::aggregate_fn::fns::null_count::NullCount; use crate::dtype::DType; use crate::expr::BoundExpression; +use crate::expr::bound::eq; use crate::expr::bound::lit; use crate::expr::traversal::NodeExt; use crate::expr::traversal::Transformed; use crate::scalar::Scalar; +use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::stat::StatFn; /// A target that can bind abstract statistics to concrete expressions. @@ -42,6 +57,15 @@ pub trait StatBinder { stat_dtype: &DType, ) -> VortexResult>; + /// Bind the number of rows covered by each row of the stats scope. + /// + /// This backs the derivation of `all_null` and `all_nan` from their count statistics. It is an + /// expression rather than a scalar because a scope may cover a different number of rows per + /// row of its stats table: a zone map's final zone is often shorter than the rest. + /// + /// Implementations return `Ok(None)` when the row count is unknown. + fn bind_row_count(&self) -> VortexResult>; + /// Expression to use when a stat is unavailable. /// /// The default is a nullable null literal, which preserves three-valued @@ -83,7 +107,64 @@ fn bind_stat_fn( // `StatFn` has exactly one child: the expression the aggregate statistic is computed over. let input = expr.child(0); - binder.bind_aggregate(input, aggregate_fn, expr.dtype()) + if let Some(bound) = binder.bind_aggregate(input, aggregate_fn, expr.dtype())? { + return Ok(Some(bound)); + } + + derive_from_count(input, aggregate_fn, binder) +} + +/// The count statistic a boolean "all" aggregate is derived from, and whether the derivation +/// compares it against the row count (`true`) or against zero (`false`). +fn count_derivation(aggregate_fn: &AggregateFnRef) -> Option<(AggregateFnRef, bool)> { + if aggregate_fn.is::() { + Some((NullCount.bind(EmptyOptions), true)) + } else if aggregate_fn.is::() { + Some((NullCount.bind(EmptyOptions), false)) + } else if aggregate_fn.is::() { + Some((NanCount.bind(EmptyOptions), true)) + } else if aggregate_fn.is::() { + Some((NanCount.bind(EmptyOptions), false)) + } else { + None + } +} + +/// Derive a boolean "all" aggregate from the count statistic a binder does store. +/// +/// `all_null` holds exactly when every row is null, so it is `null_count == row_count`, and +/// `all_non_null` is `null_count == 0`; the NaN variants are the same shape over `nan_count`. +/// Binders that store the boolean aggregate directly answer from [`StatBinder::bind_aggregate`] +/// and never reach here. +fn derive_from_count( + input: &BoundExpression, + aggregate_fn: &AggregateFnRef, + binder: &(impl StatBinder + ?Sized), +) -> VortexResult> { + // A cast can change how many values are null or NaN, so a count over the cast input proves + // nothing about the cast output. The rewrite rules refuse to push counts through a cast for + // the same reason. + if input.is::() { + return Ok(None); + } + + let Some((count_fn, against_row_count)) = count_derivation(aggregate_fn) else { + return Ok(None); + }; + let Some(count_dtype) = count_fn.state_dtype(input.dtype()) else { + return Ok(None); + }; + let Some(count) = binder.bind_aggregate(input, &count_fn, &count_dtype.as_nullable())? else { + return Ok(None); + }; + + if !against_row_count { + return Ok(Some(eq(count, lit(0u64)))); + } + + Ok(binder + .bind_row_count()? + .map(|row_count| eq(count, row_count))) } fn null_expr(dtype: DType) -> VortexResult { @@ -99,13 +180,16 @@ mod tests { use crate::dtype::PType; use crate::dtype::StructFields; use crate::expr::and; + use crate::expr::cast; use crate::expr::col; + use crate::expr::eq; use crate::expr::get_item; use crate::expr::is_null; use crate::expr::lit; use crate::expr::or; use crate::expr::root; use crate::expr::stats::Stat; + use crate::stats::all_nan; use crate::stats::all_non_nan; use crate::stats::nan_count; @@ -113,6 +197,7 @@ mod tests { input_scope: DType, stats_scope: DType, bind_nan_count: bool, + row_count: Option, } impl TestBinder { @@ -133,8 +218,14 @@ mod tests { Nullability::NonNullable, ), bind_nan_count, + row_count: Some(10), } } + + fn without_row_count(mut self) -> Self { + self.row_count = None; + self + } } impl StatBinder for TestBinder { @@ -156,6 +247,12 @@ mod tests { Ok(None) } } + + fn bind_row_count(&self) -> VortexResult> { + self.row_count + .map(|row_count| lit(row_count).bind(&self.stats_scope)) + .transpose() + } } #[test] @@ -169,9 +266,65 @@ mod tests { } #[test] - fn all_non_nan_does_not_derive_from_nan_count() -> VortexResult<()> { + fn all_non_nan_derives_from_nan_count() -> VortexResult<()> { + let binder = TestBinder::new(true); + + let bound = bind_stats(all_non_nan(col("f")).bind(&binder.input_scope)?, &binder)?; + + assert_eq!( + bound, + eq(col("f_nan_count"), lit(0u64)).bind(&binder.stats_scope)? + ); + Ok(()) + } + + #[test] + fn all_nan_derives_from_nan_count_and_row_count() -> VortexResult<()> { let binder = TestBinder::new(true); + let bound = bind_stats(all_nan(col("f")).bind(&binder.input_scope)?, &binder)?; + + assert_eq!( + bound, + eq(col("f_nan_count"), lit(10u64)).bind(&binder.stats_scope)? + ); + Ok(()) + } + + #[test] + fn all_nan_is_missing_without_a_row_count() -> VortexResult<()> { + let binder = TestBinder::new(true).without_row_count(); + + let bound = bind_stats(all_nan(col("f")).bind(&binder.input_scope)?, &binder)?; + + assert_eq!( + bound, + lit(Scalar::null(DType::Bool(Nullability::Nullable))).bind(&binder.stats_scope)? + ); + Ok(()) + } + + #[test] + fn all_non_nan_does_not_derive_through_a_cast() -> VortexResult<()> { + let binder = TestBinder::new(true); + let cast_dtype = DType::Primitive(PType::F64, Nullability::NonNullable); + + let bound = bind_stats( + all_non_nan(cast(col("f"), cast_dtype)).bind(&binder.input_scope)?, + &binder, + )?; + + assert_eq!( + bound, + lit(Scalar::null(DType::Bool(Nullability::Nullable))).bind(&binder.stats_scope)? + ); + Ok(()) + } + + #[test] + fn all_non_nan_is_missing_without_a_nan_count() -> VortexResult<()> { + let binder = TestBinder::new(false); + let bound = bind_stats(all_non_nan(col("f")).bind(&binder.input_scope)?, &binder)?; assert_eq!( diff --git a/vortex-array/src/stats/rewrite/builtins.rs b/vortex-array/src/stats/rewrite/builtins.rs index 3cb5fdb06df..d2d23150e56 100644 --- a/vortex-array/src/stats/rewrite/builtins.rs +++ b/vortex-array/src/stats/rewrite/builtins.rs @@ -3,7 +3,6 @@ use std::sync::Arc; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::aggregate_fn::AggregateFnRef; @@ -29,10 +28,8 @@ use crate::expr::bound::or; use crate::expr::bound::or_collect; use crate::expr::stats::Stat; use crate::scalar::StringLike; -use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::ScalarFnVTableExt; use crate::scalar_fn::fns::between::Between; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::cast::Cast; @@ -46,7 +43,6 @@ use crate::scalar_fn::fns::list_contains::ListContains; use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::fns::operators::Operator; -use crate::scalar_fn::internal::row_count::RowCount; use crate::stats::bound::stat; use crate::stats::rewrite::StatsRewriteCtx; use crate::stats::rewrite::StatsRewriteRule; @@ -57,10 +53,8 @@ pub(crate) fn register_builtins(session: &StatsSession) { session.register_rewrite(BinaryNanCountStatsRewrite); session.register_rewrite(BinaryAllNonNanStatsRewrite); session.register_rewrite(BetweenStatsRewrite); - session.register_rewrite(IsNullNullCountStatsRewrite); session.register_rewrite(IsNullAllNonNullStatsRewrite); session.register_rewrite(IsNullAllNullStatsRewrite); - session.register_rewrite(IsNotNullNullCountStatsRewrite); session.register_rewrite(IsNotNullAllNullStatsRewrite); session.register_rewrite(IsNotNullAllNonNullStatsRewrite); session.register_rewrite(LikeStatsRewrite); @@ -70,12 +64,6 @@ pub(crate) fn register_builtins(session: &StatsSession) { session.register_rewrite(DynamicComparisonAllNonNanStatsRewrite); } -fn row_count() -> BoundExpression { - RowCount - .try_new_bound_expr(EmptyOptions, []) - .vortex_expect("row-count expressions are always well-typed") -} - #[derive(Debug)] struct BinaryNanCountStatsRewrite; @@ -209,31 +197,6 @@ impl StatsRewriteRule for BetweenStatsRewrite { } } -#[derive(Debug)] -struct IsNullNullCountStatsRewrite; - -impl StatsRewriteRule for IsNullNullCountStatsRewrite { - fn scalar_fn_id(&self) -> ScalarFnId { - IsNull.id() - } - - fn falsify( - &self, - expr: &BoundExpression, - ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { - Ok(null_count(expr.child(0), ctx).map(|null_count| eq(null_count, lit(0u64)))) - } - - fn satisfy( - &self, - expr: &BoundExpression, - ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { - Ok(null_count(expr.child(0), ctx).map(|null_count| eq(null_count, row_count()))) - } -} - #[derive(Debug)] struct IsNullAllNonNullStatsRewrite; @@ -268,31 +231,6 @@ impl StatsRewriteRule for IsNullAllNullStatsRewrite { } } -#[derive(Debug)] -struct IsNotNullNullCountStatsRewrite; - -impl StatsRewriteRule for IsNotNullNullCountStatsRewrite { - fn scalar_fn_id(&self) -> ScalarFnId { - IsNotNull.id() - } - - fn falsify( - &self, - expr: &BoundExpression, - ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { - Ok(null_count(expr.child(0), ctx).map(|null_count| eq(null_count, row_count()))) - } - - fn satisfy( - &self, - expr: &BoundExpression, - ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { - Ok(null_count(expr.child(0), ctx).map(|null_count| eq(null_count, lit(0u64)))) - } -} - #[derive(Debug)] struct IsNotNullAllNullStatsRewrite; @@ -526,16 +464,18 @@ fn max(expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>) -> Option) -> Option { - stat_expr(expr, Stat::NullCount, ctx) -} - fn all_null(expr: &BoundExpression) -> BoundExpression { - stat_fn(expr.clone(), AllNull.bind(AggregateEmptyOptions)) + match expr.as_opt::() { + Some(scalar) => lit(scalar.is_null()), + None => stat_fn(expr.clone(), AllNull.bind(AggregateEmptyOptions)), + } } fn all_non_null(expr: &BoundExpression) -> BoundExpression { - stat_fn(expr.clone(), AllNonNull.bind(AggregateEmptyOptions)) + match expr.as_opt::() { + Some(scalar) => lit(!scalar.is_null()), + None => stat_fn(expr.clone(), AllNonNull.bind(AggregateEmptyOptions)), + } } enum NanCheck { @@ -748,7 +688,6 @@ mod tests { use crate::expr::or; use crate::expr::stats::Stat; use crate::scalar::Scalar; - use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; @@ -758,7 +697,6 @@ mod tests { use crate::scalar_fn::fns::dynamic::DynamicComparison; use crate::scalar_fn::fns::dynamic::DynamicComparisonExpr; use crate::scalar_fn::fns::operators::CompareOperator; - use crate::scalar_fn::internal::row_count::RowCount; use crate::stats::expr::StatFn; use crate::stats::expr::StatOptions; use crate::stats::rewrite::StatsRewriteCtx; @@ -948,47 +886,32 @@ mod tests { #[test] fn rewrites_null_falsifiers() -> VortexResult<()> { - assert_rewrite_eq!( - falsify(&is_null(col("a")))?, - Some(or( - eq(stat(col("a"), Stat::NullCount), lit(0u64)), - all_non_null(&col("a")), - )) - ); - - assert_rewrite_eq!( - falsify(&is_not_null(col("a")))?, - Some(or( - eq( - stat(col("a"), Stat::NullCount), - RowCount.new_expr(EmptyOptions, []), - ), - all_null(&col("a")), - )) - ); + assert_rewrite_eq!(falsify(&is_null(col("a")))?, Some(all_non_null(&col("a")))); + assert_rewrite_eq!(falsify(&is_not_null(col("a")))?, Some(all_null(&col("a")))); Ok(()) } #[test] fn rewrites_null_satisfiers() -> VortexResult<()> { + assert_rewrite_eq!(satisfy(&is_null(col("a")))?, Some(all_null(&col("a")))); assert_rewrite_eq!( - satisfy(&is_null(col("a")))?, - Some(or( - eq( - stat(col("a"), Stat::NullCount), - RowCount.new_expr(EmptyOptions, []), - ), - all_null(&col("a")), - )) + satisfy(&is_not_null(col("a")))?, + Some(all_non_null(&col("a"))) ); + Ok(()) + } + #[test] + fn null_rewrites_fold_over_literals() -> VortexResult<()> { + assert_rewrite_eq!(falsify(&is_null(lit(1i32)))?, Some(lit(true))); assert_rewrite_eq!( - satisfy(&is_not_null(col("a")))?, - Some(or( - eq(stat(col("a"), Stat::NullCount), lit(0u64)), - all_non_null(&col("a")), - )) + falsify(&is_null(lit(Scalar::null(DType::Primitive( + PType::I32, + Nullability::Nullable + )))))?, + Some(lit(false)) ); + assert_rewrite_eq!(falsify(&is_not_null(lit(1i32)))?, Some(lit(false))); Ok(()) } diff --git a/vortex-file/src/pruning.rs b/vortex-file/src/pruning.rs index df327638d00..d53e9ca8e5f 100644 --- a/vortex-file/src/pruning.rs +++ b/vortex-file/src/pruning.rs @@ -5,7 +5,6 @@ use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnRef; -use vortex_array::arrays::ConstantArray; use vortex_array::arrays::NullArray; use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; @@ -17,7 +16,6 @@ use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::cast::Cast; use vortex_array::scalar_fn::fns::get_item::GetItem; use vortex_array::scalar_fn::fns::literal::Literal; -use vortex_array::scalar_fn::internal::row_count::substitute_row_count; use vortex_array::stats::bind::StatBinder; use vortex_array::stats::bind::bind_stats; use vortex_error::VortexResult; @@ -39,6 +37,7 @@ pub(crate) fn can_prune_file_stats( let binder = FileStatsBinder { file_stats, struct_fields, + row_count, }; let pruning_expr = bind_stats(pruning_expr, &binder)?; @@ -47,8 +46,6 @@ pub(crate) fn can_prune_file_stats( } let pruning = NullArray::new(1).into_array().apply_bound(&pruning_expr)?; - let row_count_replacement = ConstantArray::new(row_count, pruning.len()).into_array(); - let pruning = substitute_row_count(pruning, &row_count_replacement)?; let mut ctx = session.create_execution_ctx(); let result = pruning @@ -63,6 +60,7 @@ pub(crate) fn can_prune_file_stats( struct FileStatsBinder<'a> { file_stats: &'a FileStatistics, struct_fields: &'a StructFields, + row_count: u64, } impl StatBinder for FileStatsBinder<'_> { @@ -80,6 +78,11 @@ impl StatBinder for FileStatsBinder<'_> { }; Ok(self.stat_ref(&field_path, stat)) } + + /// File statistics cover the whole file, so every row of this scope covers `row_count` rows. + fn bind_row_count(&self) -> VortexResult> { + Ok(Some(lit(self.row_count))) + } } impl FileStatsBinder<'_> { diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index f5c177c9cdf..67322842354 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -47,6 +47,8 @@ use vortex_array::expr::eq; use vortex_array::expr::get_item; use vortex_array::expr::gt; use vortex_array::expr::gt_eq; +use vortex_array::expr::is_not_null; +use vortex_array::expr::is_null; use vortex_array::expr::lit; use vortex_array::expr::lt; use vortex_array::expr::lt_eq; @@ -2721,6 +2723,47 @@ async fn test_can_prune_composite_predicates() -> VortexResult<()> { Ok(()) } +#[tokio::test] +#[cfg_attr(miri, ignore)] +async fn test_can_prune_null_predicates() -> VortexResult<()> { + // File stats only store `null_count`, so `all_null` and `all_non_null` must be derived from + // it during stat binding. `all_null` additionally needs the file's row count. + let st = StructArray::from_fields(&[ + ( + "never_null", + PrimitiveArray::from_option_iter([Some(1i32), Some(2), Some(3)]).into_array(), + ), + ( + "always_null", + PrimitiveArray::from_option_iter::([None, None, None]).into_array(), + ), + ( + "sometimes_null", + PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array(), + ), + ])?; + let mut buf = ByteBufferMut::empty(); + SESSION + .write_options() + .write(&mut buf, st.into_array().to_array_stream()) + .await?; + let file = SESSION.open_options().open_buffer(buf)?; + + // `null_count == 0` proves no row is null. + assert!(file.can_prune(&is_null(col("never_null")))?); + assert!(!file.can_prune(&is_not_null(col("never_null")))?); + + // `null_count == row_count` proves every row is null. + assert!(file.can_prune(&is_not_null(col("always_null")))?); + assert!(!file.can_prune(&is_null(col("always_null")))?); + + // Mixed nullability proves nothing either way. + assert!(!file.can_prune(&is_null(col("sometimes_null")))?); + assert!(!file.can_prune(&is_not_null(col("sometimes_null")))?); + + Ok(()) +} + #[tokio::test] #[cfg_attr(miri, ignore)] async fn repro_8166_binary_gt_all_ff_max() -> VortexResult<()> { diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index c84c0b443dd..020b749b83b 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -10,10 +10,6 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnSatisfaction; -use vortex_array::aggregate_fn::fns::all_nan::AllNan; -use vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan; -use vortex_array::aggregate_fn::fns::all_non_null::AllNonNull; -use vortex_array::aggregate_fn::fns::all_null::AllNull; use vortex_array::aggregate_fn::fns::bounded_max::BOUNDED_MAX_BOUND; use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax; use vortex_array::arrays::ConstantArray; @@ -23,16 +19,9 @@ use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; -use vortex_array::expr::eq; use vortex_array::expr::get_item; -use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::expr::stats::Stat; -use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ScalarFnVTableExt; -use vortex_array::scalar_fn::internal::row_count::RowCount; -use vortex_array::scalar_fn::internal::row_count::contains_row_count; -use vortex_array::scalar_fn::internal::row_count::substitute_row_count; use vortex_array::stats::bind::StatBinder; use vortex_array::stats::bind::bind_stats; use vortex_array::validity::Validity; @@ -130,43 +119,72 @@ impl ZoneMap { /// `predicate` should be a stats rewrite expression such as the result of /// [`BoundExpression::falsify`]. The returned mask has one value per zone, where /// `true` means the zone cannot contain matching rows and can be skipped. - /// - /// If the predicate contains [`row_count`][vortex_array::scalar_fn::internal::row_count] - /// placeholders, they are replaced after [`ArrayRef::apply_bound`] with per-zone - /// counts derived from `zone_len` and `row_count`. Uniform zones use a - /// [`ConstantArray`]; a short final zone uses a run-end encoded array. - /// `row_count` is a layout property rather than a stored stats field, and the - /// final zone may be shorter than the nominal zone length, so it is materialized - /// only after the predicate has been lowered to the zone-map table. pub fn prune( &self, predicate: &BoundExpression, session: &VortexSession, ) -> VortexResult { let mut ctx = session.create_execution_ctx(); - let num_zones = self.array.len(); - let predicate = self.lower_stats(predicate.clone())?; - - let array = self.array.clone().into_array(); - let applied = array.apply_bound(&predicate)?; + let scope = self.pruning_scope()?; + let predicate = self.lower_stats(predicate.clone(), scope.dtype())?; - if !contains_row_count(&applied) { - return applied.null_as_false().execute(&mut ctx); - } + scope + .into_array() + .apply_bound(&predicate)? + .null_as_false() + .execute(&mut ctx) + } - let row_count_array = row_count_array(self.zone_len, self.row_count, num_zones)?; - let substituted = substitute_row_count(applied, &row_count_array)?; - substituted.null_as_false().execute(&mut ctx) + /// The scope that a lowered pruning predicate is evaluated against. + /// + /// This is the stored zone-map table plus a synthetic [`ROW_COUNT_FIELD`] column holding the + /// number of rows in each zone. The row count is a layout property rather than a stored stat, + /// and the final zone may be shorter than the nominal zone length, so it cannot be bound to a + /// literal. Materializing it costs nothing: uniform zones use a [`ConstantArray`] and a short + /// final zone uses a two-run run-end encoded array. + fn pruning_scope(&self) -> VortexResult { + let num_zones = self.array.len(); + let names = self + .array + .names() + .iter() + .cloned() + .chain([ROW_COUNT_FIELD.into()]); + let fields = self + .array + .iter_unmasked_fields() + .cloned() + .chain([row_count_array(self.zone_len, self.row_count, num_zones)?]); + + StructArray::try_new( + names.collect(), + fields, + num_zones, + self.array.struct_validity(), + ) } - fn lower_stats(&self, predicate: BoundExpression) -> VortexResult { - let binder = ZoneMapStatsBinder { zone_map: self }; + fn lower_stats( + &self, + predicate: BoundExpression, + scope_dtype: &DType, + ) -> VortexResult { + let binder = ZoneMapStatsBinder { + zone_map: self, + scope_dtype, + }; bind_stats(predicate, &binder) } } +/// Name of the synthetic per-zone row-count column added by [`ZoneMap::pruning_scope`]. +/// +/// The `$` prefix cannot collide with a stat name or an aggregate function's display name. +const ROW_COUNT_FIELD: &str = "$row_count"; + struct ZoneMapStatsBinder<'a> { zone_map: &'a ZoneMap, + scope_dtype: &'a DType, } impl StatBinder for ZoneMapStatsBinder<'_> { @@ -190,38 +208,8 @@ impl StatBinder for ZoneMapStatsBinder<'_> { return Ok(Some(self.bind_target(stat_expr)?)); } - if aggregate_fn.is::() { - return self - .zone_map - .stat_field_expr(Stat::NullCount) - .map(|null_count| self.bind_target(eq(null_count, row_count_expr()))) - .transpose(); - } - - if aggregate_fn.is::() { - return self - .zone_map - .stat_field_expr(Stat::NullCount) - .map(|null_count| self.bind_target(eq(null_count, lit(0u64)))) - .transpose(); - } - - if aggregate_fn.is::() { - return self - .zone_map - .stat_field_expr(Stat::NaNCount) - .map(|nan_count| self.bind_target(eq(nan_count, row_count_expr()))) - .transpose(); - } - - if aggregate_fn.is::() { - return self - .zone_map - .stat_field_expr(Stat::NaNCount) - .map(|nan_count| self.bind_target(eq(nan_count, lit(0u64)))) - .transpose(); - } - + // The boolean `all_null` / `all_nan` family is derived from the count stats by + // `bind_stats`, using `bind_row_count` below. if let Some(stat) = Stat::from_aggregate_fn(aggregate_fn) { return self .zone_map @@ -232,11 +220,16 @@ impl StatBinder for ZoneMapStatsBinder<'_> { Ok(None) } + + fn bind_row_count(&self) -> VortexResult> { + self.bind_target(get_item(ROW_COUNT_FIELD, root())) + .map(Some) + } } impl ZoneMapStatsBinder<'_> { fn bind_target(&self, expr: Expression) -> VortexResult { - expr.bind(self.zone_map.array.dtype()) + expr.bind(self.scope_dtype) } } @@ -298,10 +291,6 @@ fn aggregate_result_expr(stored: &AggregateFnRef, state_expr: Expression) -> Exp } } -fn row_count_expr() -> Expression { - RowCount.new_expr(EmptyOptions, []) -} - /// Build per-zone row counts for a zone map. /// /// `zone_len` is the nominal zone size; only the final zone may be shorter. The @@ -560,7 +549,7 @@ mod tests { } #[test] - fn row_count_substitution_handles_empty_zone_map() { + fn pruning_handles_empty_zone_map() { let zone_map = ZoneMap::try_new_legacy( PType::U64.into(), StructArray::from_fields(&[( From 867e3a93f48c223f21d5607a8e52f5ed3106f71b Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 14 Aug 2026 14:07:47 +0100 Subject: [PATCH 2/5] less Signed-off-by: Robert Kruszewski --- vortex-array/src/scalar_fn/fns/is_not_null.rs | 14 +- vortex-array/src/scalar_fn/fns/is_null.rs | 6 +- vortex-array/src/scalar_fn/internal/mod.rs | 4 + .../src/scalar_fn/internal/row_count.rs | 108 ++++++++++ vortex-array/src/scalar_fn/mod.rs | 1 + vortex-array/src/stats/bind.rs | 176 ++++------------ vortex-array/src/stats/rewrite/builtins.rs | 125 +++++++++--- vortex-file/src/pruning.rs | 2 +- vortex-file/src/tests.rs | 4 +- vortex-layout/src/layouts/zoned/zone_map.rs | 192 +++++++++++++++--- 10 files changed, 437 insertions(+), 195 deletions(-) create mode 100644 vortex-array/src/scalar_fn/internal/mod.rs create mode 100644 vortex-array/src/scalar_fn/internal/row_count.rs diff --git a/vortex-array/src/scalar_fn/fns/is_not_null.rs b/vortex-array/src/scalar_fn/fns/is_not_null.rs index 68a30771b3d..49bb9937b13 100644 --- a/vortex-array/src/scalar_fn/fns/is_not_null.rs +++ b/vortex-array/src/scalar_fn/fns/is_not_null.rs @@ -116,13 +116,19 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; use crate::expr::col; + use crate::expr::eq; use crate::expr::get_item; use crate::expr::is_not_null; + use crate::expr::or; use crate::expr::root; use crate::expr::test_harness; use crate::scalar::Scalar; + use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::ScalarFnVTableExt; + use crate::scalar_fn::internal::row_count::RowCount; use crate::stats::StatsSession; use crate::stats::all_null; + use crate::stats::null_count; static STATS_SESSION: LazyLock = LazyLock::new(|| VortexSession::empty().with::()); @@ -255,7 +261,13 @@ mod tests { assert_eq!( expr.bind(&dtype)?.falsify(&STATS_SESSION)?, - Some(all_null(col("a")).bind(&dtype)?) + Some( + or( + eq(null_count(col("a")), RowCount.new_expr(EmptyOptions, []),), + all_null(col("a")), + ) + .bind(&dtype)? + ) ); Ok(()) } diff --git a/vortex-array/src/scalar_fn/fns/is_null.rs b/vortex-array/src/scalar_fn/fns/is_null.rs index 253b527e9e5..87aca6748e5 100644 --- a/vortex-array/src/scalar_fn/fns/is_null.rs +++ b/vortex-array/src/scalar_fn/fns/is_null.rs @@ -106,13 +106,17 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; use crate::expr::col; + use crate::expr::eq; use crate::expr::get_item; use crate::expr::is_null; + use crate::expr::lit; + use crate::expr::or; use crate::expr::root; use crate::expr::test_harness; use crate::scalar::Scalar; use crate::stats::StatsSession; use crate::stats::all_non_null; + use crate::stats::null_count; static STATS_SESSION: LazyLock = LazyLock::new(|| VortexSession::empty().with::()); @@ -238,7 +242,7 @@ mod tests { assert_eq!( expr.bind(&dtype)?.falsify(&STATS_SESSION)?, - Some(all_non_null(col("a")).bind(&dtype)?) + Some(or(eq(null_count(col("a")), lit(0u64)), all_non_null(col("a")),).bind(&dtype)?) ); Ok(()) } diff --git a/vortex-array/src/scalar_fn/internal/mod.rs b/vortex-array/src/scalar_fn/internal/mod.rs new file mode 100644 index 00000000000..c91a7541c8a --- /dev/null +++ b/vortex-array/src/scalar_fn/internal/mod.rs @@ -0,0 +1,4 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +pub mod row_count; diff --git a/vortex-array/src/scalar_fn/internal/row_count.rs b/vortex-array/src/scalar_fn/internal/row_count.rs new file mode 100644 index 00000000000..d69728c169b --- /dev/null +++ b/vortex-array/src/scalar_fn/internal/row_count.rs @@ -0,0 +1,108 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Formatter; + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::expr::display::ExprDisplay; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_session::registry::CachedId; + +/// Zero-argument placeholder for the row count of the current evaluation scope. +/// +/// Stats rewrite rules emit `RowCount` when a proof needs a scope-level value that is not stored +/// as a regular stats column — `is_not_null` is falsified by `null_count == row_count`, for +/// example. Keeping it as a placeholder lets a rewrite rule name the row count without knowing +/// anything about where the stats it sits beside are stored. +/// +/// It is resolved during stat binding, by [`bind_stats`], which asks the [`StatBinder`] for the +/// row count of its scope. Binding is a single top-down pass that recurses into the expressions +/// it substitutes, so a binder may itself emit `RowCount` and have it resolved in the same pass. +/// +/// This expression *MUST* be replaced before evaluation; calling +/// [`ScalarFnVTable::execute`] directly returns an error because this node is only a marker in a +/// lazy expression tree. +/// +/// [`bind_stats`]: crate::stats::bind::bind_stats +/// [`StatBinder`]: crate::stats::bind::StatBinder +#[derive(Clone)] +pub struct RowCount; + +impl ScalarFnVTable for RowCount { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.row_count"); + *ID + } + + fn arity(&self, _options: &Self::Options) -> Arity { + Arity::Exact(0) + } + + fn child_name(&self, _options: &Self::Options, _child_idx: usize) -> ChildName { + unreachable!("RowCount has arity 0") + } + + fn fmt_sql( + &self, + _options: &Self::Options, + _expr: &dyn ExprDisplay, + f: &mut Formatter<'_>, + ) -> std::fmt::Result { + write!(f, "row_count()") + } + + fn return_dtype(&self, _options: &Self::Options, _args: &[DType]) -> VortexResult { + Ok(DType::Primitive(PType::U64, Nullability::NonNullable)) + } + + fn execute( + &self, + _options: &Self::Options, + _args: &dyn ExecutionArgs, + _ctx: &mut ExecutionCtx, + ) -> VortexResult { + vortex_bail!("RowCount must be substituted before evaluation") + } + + fn is_strict(&self, _options: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _options: &Self::Options) -> bool { + false + } +} + +#[cfg(test)] +mod tests { + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + + use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::internal::row_count::RowCount; + use crate::scalar_fn::vtable::ScalarFnVTableExt; + + #[test] + fn row_count_helper_dtype() { + let expr = RowCount.new_expr(EmptyOptions, []); + assert_eq!( + expr.return_dtype(&DType::Primitive(PType::I32, Nullability::Nullable)) + .unwrap(), + DType::Primitive(PType::U64, Nullability::NonNullable), + ); + } +} diff --git a/vortex-array/src/scalar_fn/mod.rs b/vortex-array/src/scalar_fn/mod.rs index 423b8d14bc5..6be34ce1f34 100644 --- a/vortex-array/src/scalar_fn/mod.rs +++ b/vortex-array/src/scalar_fn/mod.rs @@ -48,6 +48,7 @@ pub mod unstable; pub(crate) mod unstable; pub mod fns; +pub mod internal; pub mod session; /// A unique identifier for a scalar function. diff --git a/vortex-array/src/stats/bind.rs b/vortex-array/src/stats/bind.rs index 7c36ee9ca02..b0d75af2735 100644 --- a/vortex-array/src/stats/bind.rs +++ b/vortex-array/src/stats/bind.rs @@ -13,31 +13,20 @@ //! stats. This lets all callers share the same falsification rules while keeping layout-specific //! stat storage behind [`StatBinder`]. //! -//! Binding is also where scope-level quantities enter the predicate. The boolean aggregates -//! `all_null`, `all_non_null`, `all_nan`, and `all_non_nan` are derived from the corresponding -//! count statistic when a binder does not store them directly, which for the "all" variants needs -//! the number of rows the scope covers — see [`StatBinder::bind_row_count`]. +//! Binding also resolves [`RowCount`] placeholders, which rewrite rules emit when a proof needs +//! the number of rows the scope covers rather than a stored statistic. use vortex_error::VortexResult; use crate::aggregate_fn::AggregateFnRef; -use crate::aggregate_fn::AggregateFnVTableExt; -use crate::aggregate_fn::EmptyOptions; -use crate::aggregate_fn::fns::all_nan::AllNan; -use crate::aggregate_fn::fns::all_non_nan::AllNonNan; -use crate::aggregate_fn::fns::all_non_null::AllNonNull; -use crate::aggregate_fn::fns::all_null::AllNull; -use crate::aggregate_fn::fns::nan_count::NanCount; -use crate::aggregate_fn::fns::null_count::NullCount; use crate::dtype::DType; use crate::expr::BoundExpression; -use crate::expr::bound::eq; use crate::expr::bound::lit; use crate::expr::traversal::NodeExt; use crate::expr::traversal::Transformed; use crate::scalar::Scalar; -use crate::scalar_fn::fns::cast::Cast; use crate::scalar_fn::fns::stat::StatFn; +use crate::scalar_fn::internal::row_count::RowCount; /// A target that can bind abstract statistics to concrete expressions. /// @@ -59,11 +48,12 @@ pub trait StatBinder { /// Bind the number of rows covered by each row of the stats scope. /// - /// This backs the derivation of `all_null` and `all_nan` from their count statistics. It is an - /// expression rather than a scalar because a scope may cover a different number of rows per - /// row of its stats table: a zone map's final zone is often shorter than the rest. + /// This resolves the [`RowCount`] placeholders that rewrite rules emit. It is an expression + /// rather than a scalar because a scope may cover a different number of rows per row of its + /// stats table: a zone map's final zone is often shorter than the rest. /// - /// Implementations return `Ok(None)` when the row count is unknown. + /// Implementations return `Ok(None)` when the row count is unknown, and must not return an + /// expression that itself contains a [`RowCount`]. fn bind_row_count(&self) -> VortexResult>; /// Expression to use when a stat is unavailable. @@ -86,11 +76,19 @@ pub fn bind_stats( ) -> VortexResult { Ok(predicate .transform_down(|expr| { - if !expr.is::() { + // `transform_down` recurses into whatever it substitutes, so a binder may answer with + // an expression that itself contains placeholders and have them resolved in this same + // pass. That is what lets a stats source express `all_null` as + // `null_count == row_count` without a second traversal. + let bound = if expr.is::() { + bind_stat_fn(&expr, binder)? + } else if expr.is::() { + binder.bind_row_count()? + } else { return Ok(Transformed::no(expr)); - } + }; - match bind_stat_fn(&expr, binder)? { + match bound { Some(bound) => Ok(Transformed::yes(bound)), None => Ok(Transformed::yes(binder.missing_stat(expr.dtype().clone())?)), } @@ -107,64 +105,7 @@ fn bind_stat_fn( // `StatFn` has exactly one child: the expression the aggregate statistic is computed over. let input = expr.child(0); - if let Some(bound) = binder.bind_aggregate(input, aggregate_fn, expr.dtype())? { - return Ok(Some(bound)); - } - - derive_from_count(input, aggregate_fn, binder) -} - -/// The count statistic a boolean "all" aggregate is derived from, and whether the derivation -/// compares it against the row count (`true`) or against zero (`false`). -fn count_derivation(aggregate_fn: &AggregateFnRef) -> Option<(AggregateFnRef, bool)> { - if aggregate_fn.is::() { - Some((NullCount.bind(EmptyOptions), true)) - } else if aggregate_fn.is::() { - Some((NullCount.bind(EmptyOptions), false)) - } else if aggregate_fn.is::() { - Some((NanCount.bind(EmptyOptions), true)) - } else if aggregate_fn.is::() { - Some((NanCount.bind(EmptyOptions), false)) - } else { - None - } -} - -/// Derive a boolean "all" aggregate from the count statistic a binder does store. -/// -/// `all_null` holds exactly when every row is null, so it is `null_count == row_count`, and -/// `all_non_null` is `null_count == 0`; the NaN variants are the same shape over `nan_count`. -/// Binders that store the boolean aggregate directly answer from [`StatBinder::bind_aggregate`] -/// and never reach here. -fn derive_from_count( - input: &BoundExpression, - aggregate_fn: &AggregateFnRef, - binder: &(impl StatBinder + ?Sized), -) -> VortexResult> { - // A cast can change how many values are null or NaN, so a count over the cast input proves - // nothing about the cast output. The rewrite rules refuse to push counts through a cast for - // the same reason. - if input.is::() { - return Ok(None); - } - - let Some((count_fn, against_row_count)) = count_derivation(aggregate_fn) else { - return Ok(None); - }; - let Some(count_dtype) = count_fn.state_dtype(input.dtype()) else { - return Ok(None); - }; - let Some(count) = binder.bind_aggregate(input, &count_fn, &count_dtype.as_nullable())? else { - return Ok(None); - }; - - if !against_row_count { - return Ok(Some(eq(count, lit(0u64)))); - } - - Ok(binder - .bind_row_count()? - .map(|row_count| eq(count, row_count))) + binder.bind_aggregate(input, aggregate_fn, expr.dtype()) } fn null_expr(dtype: DType) -> VortexResult { @@ -176,11 +117,11 @@ mod tests { use vortex_error::VortexResult; use super::*; + use crate::aggregate_fn::fns::all_nan::AllNan; use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; use crate::expr::and; - use crate::expr::cast; use crate::expr::col; use crate::expr::eq; use crate::expr::get_item; @@ -189,6 +130,8 @@ mod tests { use crate::expr::or; use crate::expr::root; use crate::expr::stats::Stat; + use crate::scalar_fn::EmptyOptions; + use crate::scalar_fn::ScalarFnVTableExt; use crate::stats::all_nan; use crate::stats::all_non_nan; use crate::stats::nan_count; @@ -197,7 +140,6 @@ mod tests { input_scope: DType, stats_scope: DType, bind_nan_count: bool, - row_count: Option, } impl TestBinder { @@ -218,14 +160,8 @@ mod tests { Nullability::NonNullable, ), bind_nan_count, - row_count: Some(10), } } - - fn without_row_count(mut self) -> Self { - self.row_count = None; - self - } } impl StatBinder for TestBinder { @@ -235,6 +171,19 @@ mod tests { aggregate_fn: &AggregateFnRef, _stat_dtype: &DType, ) -> VortexResult> { + // `all_nan` is not stored, but it is `nan_count == row_count`. Answering with an + // expression that still contains a `RowCount` exercises the binding pass recursing + // into what it substitutes. + if aggregate_fn.is::() && self.bind_nan_count { + return Ok(Some( + eq( + get_item("f_nan_count", root()), + RowCount.new_expr(EmptyOptions, []), + ) + .bind(&self.stats_scope)?, + )); + } + let Some(stat) = Stat::from_aggregate_fn(aggregate_fn) else { return Ok(None); }; @@ -249,9 +198,7 @@ mod tests { } fn bind_row_count(&self) -> VortexResult> { - self.row_count - .map(|row_count| lit(row_count).bind(&self.stats_scope)) - .transpose() + lit(10u64).bind(&self.stats_scope).map(Some) } } @@ -266,20 +213,20 @@ mod tests { } #[test] - fn all_non_nan_derives_from_nan_count() -> VortexResult<()> { + fn all_non_nan_does_not_derive_from_nan_count() -> VortexResult<()> { let binder = TestBinder::new(true); let bound = bind_stats(all_non_nan(col("f")).bind(&binder.input_scope)?, &binder)?; assert_eq!( bound, - eq(col("f_nan_count"), lit(0u64)).bind(&binder.stats_scope)? + lit(Scalar::null(DType::Bool(Nullability::Nullable))).bind(&binder.stats_scope)? ); Ok(()) } #[test] - fn all_nan_derives_from_nan_count_and_row_count() -> VortexResult<()> { + fn binder_emitted_row_count_resolves_in_the_same_pass() -> VortexResult<()> { let binder = TestBinder::new(true); let bound = bind_stats(all_nan(col("f")).bind(&binder.input_scope)?, &binder)?; @@ -291,49 +238,6 @@ mod tests { Ok(()) } - #[test] - fn all_nan_is_missing_without_a_row_count() -> VortexResult<()> { - let binder = TestBinder::new(true).without_row_count(); - - let bound = bind_stats(all_nan(col("f")).bind(&binder.input_scope)?, &binder)?; - - assert_eq!( - bound, - lit(Scalar::null(DType::Bool(Nullability::Nullable))).bind(&binder.stats_scope)? - ); - Ok(()) - } - - #[test] - fn all_non_nan_does_not_derive_through_a_cast() -> VortexResult<()> { - let binder = TestBinder::new(true); - let cast_dtype = DType::Primitive(PType::F64, Nullability::NonNullable); - - let bound = bind_stats( - all_non_nan(cast(col("f"), cast_dtype)).bind(&binder.input_scope)?, - &binder, - )?; - - assert_eq!( - bound, - lit(Scalar::null(DType::Bool(Nullability::Nullable))).bind(&binder.stats_scope)? - ); - Ok(()) - } - - #[test] - fn all_non_nan_is_missing_without_a_nan_count() -> VortexResult<()> { - let binder = TestBinder::new(false); - - let bound = bind_stats(all_non_nan(col("f")).bind(&binder.input_scope)?, &binder)?; - - assert_eq!( - bound, - lit(Scalar::null(DType::Bool(Nullability::Nullable))).bind(&binder.stats_scope)? - ); - Ok(()) - } - #[test] fn missing_stats_bind_to_null_without_reducing() -> VortexResult<()> { let binder = TestBinder::new(false); diff --git a/vortex-array/src/stats/rewrite/builtins.rs b/vortex-array/src/stats/rewrite/builtins.rs index d2d23150e56..3cb5fdb06df 100644 --- a/vortex-array/src/stats/rewrite/builtins.rs +++ b/vortex-array/src/stats/rewrite/builtins.rs @@ -3,6 +3,7 @@ use std::sync::Arc; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::aggregate_fn::AggregateFnRef; @@ -28,8 +29,10 @@ use crate::expr::bound::or; use crate::expr::bound::or_collect; use crate::expr::stats::Stat; use crate::scalar::StringLike; +use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::ScalarFnVTableExt; use crate::scalar_fn::fns::between::Between; use crate::scalar_fn::fns::binary::Binary; use crate::scalar_fn::fns::cast::Cast; @@ -43,6 +46,7 @@ use crate::scalar_fn::fns::list_contains::ListContains; use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::fns::operators::Operator; +use crate::scalar_fn::internal::row_count::RowCount; use crate::stats::bound::stat; use crate::stats::rewrite::StatsRewriteCtx; use crate::stats::rewrite::StatsRewriteRule; @@ -53,8 +57,10 @@ pub(crate) fn register_builtins(session: &StatsSession) { session.register_rewrite(BinaryNanCountStatsRewrite); session.register_rewrite(BinaryAllNonNanStatsRewrite); session.register_rewrite(BetweenStatsRewrite); + session.register_rewrite(IsNullNullCountStatsRewrite); session.register_rewrite(IsNullAllNonNullStatsRewrite); session.register_rewrite(IsNullAllNullStatsRewrite); + session.register_rewrite(IsNotNullNullCountStatsRewrite); session.register_rewrite(IsNotNullAllNullStatsRewrite); session.register_rewrite(IsNotNullAllNonNullStatsRewrite); session.register_rewrite(LikeStatsRewrite); @@ -64,6 +70,12 @@ pub(crate) fn register_builtins(session: &StatsSession) { session.register_rewrite(DynamicComparisonAllNonNanStatsRewrite); } +fn row_count() -> BoundExpression { + RowCount + .try_new_bound_expr(EmptyOptions, []) + .vortex_expect("row-count expressions are always well-typed") +} + #[derive(Debug)] struct BinaryNanCountStatsRewrite; @@ -197,6 +209,31 @@ impl StatsRewriteRule for BetweenStatsRewrite { } } +#[derive(Debug)] +struct IsNullNullCountStatsRewrite; + +impl StatsRewriteRule for IsNullNullCountStatsRewrite { + fn scalar_fn_id(&self) -> ScalarFnId { + IsNull.id() + } + + fn falsify( + &self, + expr: &BoundExpression, + ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + Ok(null_count(expr.child(0), ctx).map(|null_count| eq(null_count, lit(0u64)))) + } + + fn satisfy( + &self, + expr: &BoundExpression, + ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + Ok(null_count(expr.child(0), ctx).map(|null_count| eq(null_count, row_count()))) + } +} + #[derive(Debug)] struct IsNullAllNonNullStatsRewrite; @@ -231,6 +268,31 @@ impl StatsRewriteRule for IsNullAllNullStatsRewrite { } } +#[derive(Debug)] +struct IsNotNullNullCountStatsRewrite; + +impl StatsRewriteRule for IsNotNullNullCountStatsRewrite { + fn scalar_fn_id(&self) -> ScalarFnId { + IsNotNull.id() + } + + fn falsify( + &self, + expr: &BoundExpression, + ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + Ok(null_count(expr.child(0), ctx).map(|null_count| eq(null_count, row_count()))) + } + + fn satisfy( + &self, + expr: &BoundExpression, + ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + Ok(null_count(expr.child(0), ctx).map(|null_count| eq(null_count, lit(0u64)))) + } +} + #[derive(Debug)] struct IsNotNullAllNullStatsRewrite; @@ -464,18 +526,16 @@ fn max(expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>) -> Option) -> Option { + stat_expr(expr, Stat::NullCount, ctx) +} + fn all_null(expr: &BoundExpression) -> BoundExpression { - match expr.as_opt::() { - Some(scalar) => lit(scalar.is_null()), - None => stat_fn(expr.clone(), AllNull.bind(AggregateEmptyOptions)), - } + stat_fn(expr.clone(), AllNull.bind(AggregateEmptyOptions)) } fn all_non_null(expr: &BoundExpression) -> BoundExpression { - match expr.as_opt::() { - Some(scalar) => lit(!scalar.is_null()), - None => stat_fn(expr.clone(), AllNonNull.bind(AggregateEmptyOptions)), - } + stat_fn(expr.clone(), AllNonNull.bind(AggregateEmptyOptions)) } enum NanCheck { @@ -688,6 +748,7 @@ mod tests { use crate::expr::or; use crate::expr::stats::Stat; use crate::scalar::Scalar; + use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; @@ -697,6 +758,7 @@ mod tests { use crate::scalar_fn::fns::dynamic::DynamicComparison; use crate::scalar_fn::fns::dynamic::DynamicComparisonExpr; use crate::scalar_fn::fns::operators::CompareOperator; + use crate::scalar_fn::internal::row_count::RowCount; use crate::stats::expr::StatFn; use crate::stats::expr::StatOptions; use crate::stats::rewrite::StatsRewriteCtx; @@ -886,32 +948,47 @@ mod tests { #[test] fn rewrites_null_falsifiers() -> VortexResult<()> { - assert_rewrite_eq!(falsify(&is_null(col("a")))?, Some(all_non_null(&col("a")))); - assert_rewrite_eq!(falsify(&is_not_null(col("a")))?, Some(all_null(&col("a")))); + assert_rewrite_eq!( + falsify(&is_null(col("a")))?, + Some(or( + eq(stat(col("a"), Stat::NullCount), lit(0u64)), + all_non_null(&col("a")), + )) + ); + + assert_rewrite_eq!( + falsify(&is_not_null(col("a")))?, + Some(or( + eq( + stat(col("a"), Stat::NullCount), + RowCount.new_expr(EmptyOptions, []), + ), + all_null(&col("a")), + )) + ); Ok(()) } #[test] fn rewrites_null_satisfiers() -> VortexResult<()> { - assert_rewrite_eq!(satisfy(&is_null(col("a")))?, Some(all_null(&col("a")))); assert_rewrite_eq!( - satisfy(&is_not_null(col("a")))?, - Some(all_non_null(&col("a"))) + satisfy(&is_null(col("a")))?, + Some(or( + eq( + stat(col("a"), Stat::NullCount), + RowCount.new_expr(EmptyOptions, []), + ), + all_null(&col("a")), + )) ); - Ok(()) - } - #[test] - fn null_rewrites_fold_over_literals() -> VortexResult<()> { - assert_rewrite_eq!(falsify(&is_null(lit(1i32)))?, Some(lit(true))); assert_rewrite_eq!( - falsify(&is_null(lit(Scalar::null(DType::Primitive( - PType::I32, - Nullability::Nullable - )))))?, - Some(lit(false)) + satisfy(&is_not_null(col("a")))?, + Some(or( + eq(stat(col("a"), Stat::NullCount), lit(0u64)), + all_non_null(&col("a")), + )) ); - assert_rewrite_eq!(falsify(&is_not_null(lit(1i32)))?, Some(lit(false))); Ok(()) } diff --git a/vortex-file/src/pruning.rs b/vortex-file/src/pruning.rs index d53e9ca8e5f..f3ff618661a 100644 --- a/vortex-file/src/pruning.rs +++ b/vortex-file/src/pruning.rs @@ -79,7 +79,7 @@ impl StatBinder for FileStatsBinder<'_> { Ok(self.stat_ref(&field_path, stat)) } - /// File statistics cover the whole file, so every row of this scope covers `row_count` rows. + /// File statistics cover the whole file, so the scope's single row covers `row_count` rows. fn bind_row_count(&self) -> VortexResult> { Ok(Some(lit(self.row_count))) } diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 67322842354..49ecf4cabcb 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -2726,8 +2726,8 @@ async fn test_can_prune_composite_predicates() -> VortexResult<()> { #[tokio::test] #[cfg_attr(miri, ignore)] async fn test_can_prune_null_predicates() -> VortexResult<()> { - // File stats only store `null_count`, so `all_null` and `all_non_null` must be derived from - // it during stat binding. `all_null` additionally needs the file's row count. + // File stats store `null_count` but not the row count, so `is_not_null` falsification depends + // on the `RowCount` placeholder being resolved against the file's own row count. let st = StructArray::from_fields(&[ ( "never_null", diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index 020b749b83b..c54bd8f971d 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -10,6 +10,10 @@ use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::aggregate_fn::AggregateFnSatisfaction; +use vortex_array::aggregate_fn::fns::all_nan::AllNan; +use vortex_array::aggregate_fn::fns::all_non_nan::AllNonNan; +use vortex_array::aggregate_fn::fns::all_non_null::AllNonNull; +use vortex_array::aggregate_fn::fns::all_null::AllNull; use vortex_array::aggregate_fn::fns::bounded_max::BOUNDED_MAX_BOUND; use vortex_array::aggregate_fn::fns::bounded_max::BoundedMax; use vortex_array::arrays::ConstantArray; @@ -17,11 +21,19 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::DType; +use vortex_array::dtype::FieldNames; use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; +use vortex_array::expr::eq; use vortex_array::expr::get_item; +use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::expr::stats::Stat; +use vortex_array::expr::transform::replace; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::scalar_fn::fns::literal::Literal; +use vortex_array::scalar_fn::internal::row_count::RowCount; use vortex_array::stats::bind::StatBinder; use vortex_array::stats::bind::bind_stats; use vortex_array::validity::Validity; @@ -119,6 +131,9 @@ impl ZoneMap { /// `predicate` should be a stats rewrite expression such as the result of /// [`BoundExpression::falsify`]. The returned mask has one value per zone, where /// `true` means the zone cannot contain matching rows and can be skipped. + /// + /// Row-count placeholders are resolved during lowering against a per-zone column that + /// [`ZoneMap::pruning_scope`] materializes, so the lowered predicate is directly evaluable. pub fn prune( &self, predicate: &BoundExpression, @@ -128,6 +143,17 @@ impl ZoneMap { let scope = self.pruning_scope()?; let predicate = self.lower_stats(predicate.clone(), scope.dtype())?; + // A predicate that binds to a constant needs no per-zone evaluation. This is the common + // shape when the zone map is missing the stats a proof needs: every placeholder binds to a + // typed null, which collapses the whole tree to a null literal. + if let Some(scalar) = predicate.as_opt::() { + let len = scope.len(); + return Ok(match scalar.as_bool().value() { + Some(true) => Mask::new_true(len), + Some(false) | None => Mask::new_false(len), + }); + } + scope .into_array() .apply_bound(&predicate)? @@ -135,32 +161,32 @@ impl ZoneMap { .execute(&mut ctx) } - /// The scope that a lowered pruning predicate is evaluated against. + /// Build the scope that a lowered pruning predicate is evaluated against. + /// + /// The scope is a two-field struct: [`STATS_FIELD`] nests the stored zone-map table, and + /// [`ROW_COUNT_FIELD`] holds the number of rows in each zone. The row count is a layout + /// property rather than a stored stat, and the final zone may be shorter than the nominal zone + /// length, so it cannot be resolved to a literal. Materializing it costs nothing: uniform zones + /// use a [`ConstantArray`] and a short final zone uses a two-run run-end encoded array. /// - /// This is the stored zone-map table plus a synthetic [`ROW_COUNT_FIELD`] column holding the - /// number of rows in each zone. The row count is a layout property rather than a stored stat, - /// and the final zone may be shorter than the nominal zone length, so it cannot be bound to a - /// literal. Materializing it costs nothing: uniform zones use a [`ConstantArray`] and a short - /// final zone uses a two-run run-end encoded array. + /// Nesting is what keeps the row count addressable. Appending it beside the stat columns would + /// put a name this module chooses into a namespace that aggregate display names and legacy stat + /// names also write to, and a collision would not be loud: [`StructFields`] permits duplicate + /// names and resolves lookups to the *first* match, so a colliding stat column would silently + /// shadow the row count and corrupt pruning. One level down, stat names cannot reach the two + /// names this module owns. + /// + /// [`StructFields`]: vortex_array::dtype::StructFields fn pruning_scope(&self) -> VortexResult { let num_zones = self.array.len(); - let names = self - .array - .names() - .iter() - .cloned() - .chain([ROW_COUNT_FIELD.into()]); - let fields = self - .array - .iter_unmasked_fields() - .cloned() - .chain([row_count_array(self.zone_len, self.row_count, num_zones)?]); - StructArray::try_new( - names.collect(), - fields, + FieldNames::from([STATS_FIELD, ROW_COUNT_FIELD]), + [ + self.array.clone().into_array(), + row_count_array(self.zone_len, self.row_count, num_zones)?, + ], num_zones, - self.array.struct_validity(), + Validity::NonNullable, ) } @@ -177,10 +203,11 @@ impl ZoneMap { } } -/// Name of the synthetic per-zone row-count column added by [`ZoneMap::pruning_scope`]. -/// -/// The `$` prefix cannot collide with a stat name or an aggregate function's display name. -const ROW_COUNT_FIELD: &str = "$row_count"; +/// Field of the pruning scope nesting the stored zone-map table. +const STATS_FIELD: &str = "stats"; + +/// Field of the pruning scope holding the number of rows in each zone. +const ROW_COUNT_FIELD: &str = "row_count"; struct ZoneMapStatsBinder<'a> { zone_map: &'a ZoneMap, @@ -208,8 +235,38 @@ impl StatBinder for ZoneMapStatsBinder<'_> { return Ok(Some(self.bind_target(stat_expr)?)); } - // The boolean `all_null` / `all_nan` family is derived from the count stats by - // `bind_stats`, using `bind_row_count` below. + if aggregate_fn.is::() { + return self + .zone_map + .stat_field_expr(Stat::NullCount) + .map(|null_count| self.bind_target(eq(null_count, row_count_expr()))) + .transpose(); + } + + if aggregate_fn.is::() { + return self + .zone_map + .stat_field_expr(Stat::NullCount) + .map(|null_count| self.bind_target(eq(null_count, lit(0u64)))) + .transpose(); + } + + if aggregate_fn.is::() { + return self + .zone_map + .stat_field_expr(Stat::NaNCount) + .map(|nan_count| self.bind_target(eq(nan_count, row_count_expr()))) + .transpose(); + } + + if aggregate_fn.is::() { + return self + .zone_map + .stat_field_expr(Stat::NaNCount) + .map(|nan_count| self.bind_target(eq(nan_count, lit(0u64)))) + .transpose(); + } + if let Some(stat) = Stat::from_aggregate_fn(aggregate_fn) { return self .zone_map @@ -222,14 +279,17 @@ impl StatBinder for ZoneMapStatsBinder<'_> { } fn bind_row_count(&self) -> VortexResult> { - self.bind_target(get_item(ROW_COUNT_FIELD, root())) + get_item(ROW_COUNT_FIELD, root()) + .bind(self.scope_dtype) .map(Some) } } impl ZoneMapStatsBinder<'_> { + /// Bind a stat expression, which is built against the stored zone-map table, by rebasing it + /// onto the nested [`STATS_FIELD`] of the pruning scope. fn bind_target(&self, expr: Expression) -> VortexResult { - expr.bind(self.scope_dtype) + replace(expr, &root(), get_item(STATS_FIELD, root())).bind(self.scope_dtype) } } @@ -291,6 +351,10 @@ fn aggregate_result_expr(stored: &AggregateFnRef, state_expr: Expression) -> Exp } } +fn row_count_expr() -> Expression { + RowCount.new_expr(EmptyOptions, []) +} + /// Build per-zone row counts for a zone map. /// /// `zone_len` is the nominal zone size; only the final zone may be shorter. The @@ -361,11 +425,13 @@ mod tests { use vortex_array::expr::gt_eq; use vortex_array::expr::is_not_null; use vortex_array::expr::is_null; + use vortex_array::expr::list_contains; use vortex_array::expr::lit; use vortex_array::expr::lt; use vortex_array::expr::not_eq; use vortex_array::expr::root; use vortex_array::expr::stats::Stat; + use vortex_array::scalar::Scalar; use vortex_array::stats::all_nan; use vortex_array::stats::all_non_nan; use vortex_array::stats::all_non_null; @@ -375,6 +441,7 @@ mod tests { use vortex_error::VortexResult; use vortex_mask::Mask; + use crate::layouts::zoned::zone_map::ROW_COUNT_FIELD; use crate::layouts::zoned::zone_map::ZoneMap; use crate::test::SESSION; @@ -549,7 +616,7 @@ mod tests { } #[test] - fn pruning_handles_empty_zone_map() { + fn row_count_substitution_handles_empty_zone_map() { let zone_map = ZoneMap::try_new_legacy( PType::U64.into(), StructArray::from_fields(&[( @@ -720,6 +787,71 @@ mod tests { } } + #[test] + fn stat_column_cannot_shadow_the_row_count_field() { + // A stats column named `row_count` sits one level below the pruning scope's own + // `row_count` field, so it cannot shadow it. Were the two flattened into one namespace, + // this column would win every lookup: `StructFields` resolves duplicates to the first + // match, and pruning would silently read `999` as each zone's row count. + let zone_map = unsafe { + ZoneMap::new_unchecked( + PType::U64.into(), + StructArray::from_fields(&[ + ( + "null_count", + PrimitiveArray::new(buffer![0u64, 4, 2], Validity::AllValid).into_array(), + ), + ( + ROW_COUNT_FIELD, + PrimitiveArray::new(buffer![999u64, 999, 999], Validity::AllValid) + .into_array(), + ), + ]) + .unwrap(), + Arc::new([]), + 4, + 10, + ) + }; + + // Zones hold 4, 4 and 2 rows, so the last two are entirely null. + let expr = is_not_null(root()); + let pruning_expr = falsify(&expr, PType::U64.into()); + let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap(); + assert_arrays_eq!( + mask.into_array(), + BoolArray::from_iter([false, true, true]), + &mut SESSION.create_execution_ctx() + ); + } + + #[test] + fn constant_predicate_skips_per_zone_evaluation() { + let zone_map = ZoneMap::try_new( + PType::U64.into(), + StructArray::try_new(FieldNames::empty(), vec![], 3, Validity::NonNullable).unwrap(), + Arc::new([]), + 4, + 10, + ) + .unwrap(); + + // An empty list contains nothing, so the falsifier is `true` for every zone. + let empty_list = Scalar::list( + Arc::new(DType::Primitive(PType::U64, Nullability::NonNullable)), + vec![], + Nullability::NonNullable, + ); + let expr = list_contains(lit(empty_list), root()); + let pruning_expr = falsify(&expr, PType::U64.into()); + let mask = zone_map.prune(&pruning_expr, &SESSION).unwrap(); + assert_arrays_eq!( + mask.into_array(), + BoolArray::from_iter([true, true, true]), + &mut SESSION.create_execution_ctx() + ); + } + #[test] fn unavailable_stat_fn_lowers_to_unknown_mask() { let zone_map = ZoneMap::try_new( From 5aa0f07698aad51fdf28767f8d8de50946fab6f3 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 14 Aug 2026 14:43:17 +0100 Subject: [PATCH 3/5] Collapse duplicate expressions on the way up from stat lowering Signed-off-by: Robert Kruszewski --- vortex-array/src/stats/bind.rs | 110 ++++++++++++++++---- vortex-layout/src/layouts/zoned/zone_map.rs | 34 ++++++ 2 files changed, 125 insertions(+), 19 deletions(-) diff --git a/vortex-array/src/stats/bind.rs b/vortex-array/src/stats/bind.rs index b0d75af2735..8535bf8ad33 100644 --- a/vortex-array/src/stats/bind.rs +++ b/vortex-array/src/stats/bind.rs @@ -15,6 +15,13 @@ //! //! Binding also resolves [`RowCount`] placeholders, which rewrite rules emit when a proof needs //! the number of rows the scope covers rather than a stored statistic. +//! +//! Rewrite rules are independent and are combined with `or`, so several of them may prove the same +//! thing through different statistics — `is_not_null` is falsified both by +//! `null_count == row_count` and by `all_null`. Which of those a stats source can actually answer +//! is only known here, and a source that answers both from the same column lowers them to the same +//! expression. Binding collapses those duplicates on the way back up, so the predicate is not +//! evaluated twice per row. use vortex_error::VortexResult; @@ -25,6 +32,8 @@ use crate::expr::bound::lit; use crate::expr::traversal::NodeExt; use crate::expr::traversal::Transformed; use crate::scalar::Scalar; +use crate::scalar_fn::fns::binary::Binary; +use crate::scalar_fn::fns::operators::Operator; use crate::scalar_fn::fns::stat::StatFn; use crate::scalar_fn::internal::row_count::RowCount; @@ -75,27 +84,50 @@ pub fn bind_stats( binder: &B, ) -> VortexResult { Ok(predicate - .transform_down(|expr| { - // `transform_down` recurses into whatever it substitutes, so a binder may answer with - // an expression that itself contains placeholders and have them resolved in this same - // pass. That is what lets a stats source express `all_null` as - // `null_count == row_count` without a second traversal. - let bound = if expr.is::() { - bind_stat_fn(&expr, binder)? - } else if expr.is::() { - binder.bind_row_count()? - } else { - return Ok(Transformed::no(expr)); - }; - - match bound { - Some(bound) => Ok(Transformed::yes(bound)), - None => Ok(Transformed::yes(binder.missing_stat(expr.dtype().clone())?)), - } - })? + .transform(bind_placeholder(binder), collapse_duplicate_operand)? .into_inner()) } +/// Substitute a `vortex.stat` or `vortex.row_count` placeholder with the binder's representation. +fn bind_placeholder( + binder: &B, +) -> impl FnMut(BoundExpression) -> VortexResult> + '_ { + move |expr| { + // The traversal recurses into whatever it substitutes, so a binder may answer with an + // expression that itself contains placeholders and have them resolved in this same pass. + // That is what lets a stats source express `all_null` as `null_count == row_count` + // without a second traversal. + let bound = if expr.is::() { + bind_stat_fn(&expr, binder)? + } else if expr.is::() { + binder.bind_row_count()? + } else { + return Ok(Transformed::no(expr)); + }; + + match bound { + Some(bound) => Ok(Transformed::yes(bound)), + None => Ok(Transformed::yes(binder.missing_stat(expr.dtype().clone())?)), + } + } +} + +/// Collapse `a or a` and `a and a` to `a`. +/// +/// Both are idempotent under the three-valued logic pruning uses — `null or null` is `null`, just +/// as `null` is — so this only removes work, never changes the proof. +fn collapse_duplicate_operand(expr: BoundExpression) -> VortexResult> { + let is_duplicate = expr + .as_opt::() + .is_some_and(|operator| matches!(operator, Operator::Or | Operator::And)) + && expr.child(0) == expr.child(1); + + if is_duplicate { + return Ok(Transformed::yes(expr.child(0).clone())); + } + Ok(Transformed::no(expr)) +} + fn bind_stat_fn( expr: &BoundExpression, binder: &(impl StatBinder + ?Sized), @@ -132,6 +164,7 @@ mod tests { use crate::expr::stats::Stat; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnVTableExt; + use crate::scalar_fn::internal::row_count::RowCount as RowCountFn; use crate::stats::all_nan; use crate::stats::all_non_nan; use crate::stats::nan_count; @@ -178,7 +211,7 @@ mod tests { return Ok(Some( eq( get_item("f_nan_count", root()), - RowCount.new_expr(EmptyOptions, []), + RowCountFn.new_expr(EmptyOptions, []), ) .bind(&self.stats_scope)?, )); @@ -238,6 +271,45 @@ mod tests { Ok(()) } + #[test] + fn duplicate_proofs_collapse() -> VortexResult<()> { + let binder = TestBinder::new(true); + + // Two independent proofs of the same fact, reached through different placeholders: one + // states `nan_count == row_count` directly, the other asks for `all_nan`, which this + // binder answers the same way. + let predicate = or( + eq(nan_count(col("f")), RowCountFn.new_expr(EmptyOptions, [])), + all_nan(col("f")), + ); + let bound = bind_stats(predicate.bind(&binder.input_scope)?, &binder)?; + + assert_eq!( + bound, + eq(col("f_nan_count"), lit(10u64)).bind(&binder.stats_scope)? + ); + Ok(()) + } + + #[test] + fn distinct_proofs_are_both_kept() -> VortexResult<()> { + let binder = TestBinder::new(true); + + // Only collapse operands that are actually equal. + let predicate = or(eq(nan_count(col("f")), lit(0u64)), all_nan(col("f"))); + let bound = bind_stats(predicate.bind(&binder.input_scope)?, &binder)?; + + assert_eq!( + bound, + or( + eq(col("f_nan_count"), lit(0u64)), + eq(col("f_nan_count"), lit(10u64)), + ) + .bind(&binder.stats_scope)? + ); + Ok(()) + } + #[test] fn missing_stats_bind_to_null_without_reducing() -> VortexResult<()> { let binder = TestBinder::new(false); diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index c54bd8f971d..65a08629350 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -421,6 +421,8 @@ mod tests { use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::expr::cast; + use vortex_array::expr::eq; + use vortex_array::expr::get_item; use vortex_array::expr::gt; use vortex_array::expr::gt_eq; use vortex_array::expr::is_not_null; @@ -442,6 +444,7 @@ mod tests { use vortex_mask::Mask; use crate::layouts::zoned::zone_map::ROW_COUNT_FIELD; + use crate::layouts::zoned::zone_map::STATS_FIELD; use crate::layouts::zoned::zone_map::ZoneMap; use crate::test::SESSION; @@ -852,6 +855,37 @@ mod tests { ); } + #[test] + fn duplicate_proofs_collapse_during_lowering() { + // `is_not_null` is falsified both by `null_count == row_count` and by `all_null`. This + // zone map answers both from its `null_count` column, so the two disjuncts lower to the + // same expression and must not be evaluated twice per zone. + let zone_map = ZoneMap::try_new_legacy( + PType::U64.into(), + StructArray::from_fields(&[( + "null_count", + PrimitiveArray::new(buffer![0u64, 4, 2], Validity::AllValid).into_array(), + )]) + .unwrap(), + Arc::new([Stat::NullCount]), + 4, + 10, + ) + .unwrap(); + let scope = zone_map.pruning_scope().unwrap(); + + let pruning_expr = falsify(&is_not_null(root()), PType::U64.into()); + let lowered = zone_map.lower_stats(pruning_expr, scope.dtype()).unwrap(); + + let expected = eq( + get_item("null_count", get_item(STATS_FIELD, root())), + get_item(ROW_COUNT_FIELD, root()), + ) + .bind(scope.dtype()) + .unwrap(); + assert_eq!(lowered, expected); + } + #[test] fn unavailable_stat_fn_lowers_to_unknown_mask() { let zone_map = ZoneMap::try_new( From ab93872eeaf05a9bebfbfb616b5376e394394215 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Fri, 14 Aug 2026 17:08:05 +0100 Subject: [PATCH 4/5] rewrites Signed-off-by: Robert Kruszewski --- Cargo.lock | 1 + vortex-layout/Cargo.toml | 5 + vortex-layout/benches/zone_map_prune.rs | 294 ++++++++++++++++++++ vortex-layout/src/layouts/zoned/zone_map.rs | 140 +++++----- 4 files changed, 377 insertions(+), 63 deletions(-) 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(), + ); +} diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index 65a08629350..10b3e8984de 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -22,6 +22,7 @@ use vortex_array::arrays::StructArray; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::FieldNames; +use vortex_array::dtype::StructFields; use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::expr::eq; @@ -29,7 +30,6 @@ use vortex_array::expr::get_item; use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::expr::stats::Stat; -use vortex_array::expr::transform::replace; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::fns::literal::Literal; @@ -60,10 +60,8 @@ pub struct ZoneMap { array: StructArray, // Aggregate functions stored in the zone map, ordered by their stats-table fields. aggregate_fns: Arc<[AggregateFnRef]>, - // The length of each zone in the zone map. - zone_len: u64, - // Number of rows that the zone map covers - row_count: u64, + // Scope that lowered pruning predicates are evaluated against. See [`pruning_scope`]. + scope: StructArray, } impl ZoneMap { @@ -92,12 +90,12 @@ impl ZoneMap { zone_len: u64, row_count: u64, ) -> Self { + let scope = pruning_scope(&array, zone_len, row_count); Self { column_dtype, array, aggregate_fns, - zone_len, - row_count, + scope, } } @@ -140,69 +138,72 @@ impl ZoneMap { session: &VortexSession, ) -> VortexResult { let mut ctx = session.create_execution_ctx(); - let scope = self.pruning_scope()?; - let predicate = self.lower_stats(predicate.clone(), scope.dtype())?; + let predicate = self.lower_stats(predicate.clone())?; - // A predicate that binds to a constant needs no per-zone evaluation. This is the common - // shape when the zone map is missing the stats a proof needs: every placeholder binds to a - // typed null, which collapses the whole tree to a null literal. + // A rewrite rule that proves its case from the predicate alone lowers to a constant, which + // needs no per-zone evaluation. if let Some(scalar) = predicate.as_opt::() { - let len = scope.len(); + let len = self.scope.len(); return Ok(match scalar.as_bool().value() { Some(true) => Mask::new_true(len), Some(false) | None => Mask::new_false(len), }); } - scope + self.scope + .clone() .into_array() .apply_bound(&predicate)? .null_as_false() .execute(&mut ctx) } - /// Build the scope that a lowered pruning predicate is evaluated against. - /// - /// The scope is a two-field struct: [`STATS_FIELD`] nests the stored zone-map table, and - /// [`ROW_COUNT_FIELD`] holds the number of rows in each zone. The row count is a layout - /// property rather than a stored stat, and the final zone may be shorter than the nominal zone - /// length, so it cannot be resolved to a literal. Materializing it costs nothing: uniform zones - /// use a [`ConstantArray`] and a short final zone uses a two-run run-end encoded array. - /// - /// Nesting is what keeps the row count addressable. Appending it beside the stat columns would - /// put a name this module chooses into a namespace that aggregate display names and legacy stat - /// names also write to, and a collision would not be loud: [`StructFields`] permits duplicate - /// names and resolves lookups to the *first* match, so a colliding stat column would silently - /// shadow the row count and corrupt pruning. One level down, stat names cannot reach the two - /// names this module owns. - /// - /// [`StructFields`]: vortex_array::dtype::StructFields - fn pruning_scope(&self) -> VortexResult { - let num_zones = self.array.len(); - StructArray::try_new( - FieldNames::from([STATS_FIELD, ROW_COUNT_FIELD]), - [ - self.array.clone().into_array(), - row_count_array(self.zone_len, self.row_count, num_zones)?, - ], - num_zones, - Validity::NonNullable, - ) - } - - fn lower_stats( - &self, - predicate: BoundExpression, - scope_dtype: &DType, - ) -> VortexResult { + fn lower_stats(&self, predicate: BoundExpression) -> VortexResult { let binder = ZoneMapStatsBinder { zone_map: self, - scope_dtype, + scope_dtype: self.scope.dtype(), }; bind_stats(predicate, &binder) } } +/// Build the scope that a lowered pruning predicate is evaluated against. +/// +/// The scope is a two-field struct: [`STATS_FIELD`] nests the stored zone-map table, and +/// [`ROW_COUNT_FIELD`] holds the number of rows in each zone. The row count is a layout property +/// rather than a stored stat, and the final zone may be shorter than the nominal zone length, so it +/// cannot be resolved to a literal. Materializing it costs nothing: uniform zones use a +/// [`ConstantArray`] and a short final zone uses a two-run run-end encoded array. +/// +/// Nesting is what keeps the row count addressable. Appending it beside the stat columns would put +/// a name this module chooses into a namespace that aggregate display names and legacy stat names +/// also write to, and a collision would not be loud: [`StructFields`] permits duplicate names and +/// resolves lookups to the *first* match, so a colliding stat column would silently shadow the row +/// count and corrupt pruning. One level down, stat names cannot reach the two names this module +/// owns. +/// +/// The scope is built once per zone map rather than per predicate, because its dtype and its +/// row-count column depend only on the stats table and the layout's zone geometry. +fn pruning_scope(array: &StructArray, zone_len: u64, row_count: u64) -> StructArray { + let num_zones = array.len(); + let row_counts = row_count_array(zone_len, row_count, num_zones); + let fields = StructFields::new( + FieldNames::from([STATS_FIELD, ROW_COUNT_FIELD]), + vec![array.dtype().clone(), row_counts.dtype().clone()], + ); + + // SAFETY: both fields are `num_zones` long and their dtypes are taken from the arrays + // themselves, so they match the struct dtype by construction. + unsafe { + StructArray::new_unchecked( + [array.clone().into_array(), row_counts], + fields, + num_zones, + Validity::NonNullable, + ) + } +} + /// Field of the pruning scope nesting the stored zone-map table. const STATS_FIELD: &str = "stats"; @@ -286,20 +287,27 @@ impl StatBinder for ZoneMapStatsBinder<'_> { } impl ZoneMapStatsBinder<'_> { - /// Bind a stat expression, which is built against the stored zone-map table, by rebasing it - /// onto the nested [`STATS_FIELD`] of the pruning scope. + /// Bind a stat expression against the pruning scope it was built for. fn bind_target(&self, expr: Expression) -> VortexResult { - replace(expr, &root(), get_item(STATS_FIELD, root())).bind(self.scope_dtype) + expr.bind(self.scope_dtype) } } +/// Root of the stored zone-map table within the pruning scope. +/// +/// Stat expressions are built against this rather than against the scope root, so that binding a +/// stat is a single walk rather than a build-then-rebase. +fn stats_root() -> Expression { + get_item(STATS_FIELD, root()) +} + impl ZoneMap { fn aggregate_field_expr(&self, requested: &AggregateFnRef) -> Option { let field_name = requested.to_string(); if self.array.unmasked_field_by_name_opt(&field_name).is_some() { return Some(aggregate_result_expr( requested, - get_item(field_name, root()), + get_item(field_name, stats_root()), )); } @@ -312,10 +320,16 @@ impl ZoneMap { match stored.can_satisfy(requested) { AggregateFnSatisfaction::Exact => { - return Some(aggregate_result_expr(stored, get_item(field_name, root()))); + return Some(aggregate_result_expr( + stored, + get_item(field_name, stats_root()), + )); } AggregateFnSatisfaction::Approximate => { - approximate = Some(aggregate_result_expr(stored, get_item(field_name, root()))); + approximate = Some(aggregate_result_expr( + stored, + get_item(field_name, stats_root()), + )); } AggregateFnSatisfaction::No => {} } @@ -336,7 +350,7 @@ impl ZoneMap { fn legacy_stat_field_expr(&self, stat: Stat) -> Option { if self.array.unmasked_field_by_name_opt(stat.name()).is_some() { - return Some(get_item(stat.name(), root())); + return Some(get_item(stat.name(), stats_root())); } None @@ -360,14 +374,14 @@ fn row_count_expr() -> Expression { /// `zone_len` is the nominal zone size; only the final zone may be shorter. The /// result is a [`ConstantArray`] for uniform zone sizes, otherwise a two-run /// run-end encoded array whose trailing run carries the final zone length. -fn row_count_array(zone_len: u64, row_count: u64, num_zones: usize) -> VortexResult { +fn row_count_array(zone_len: u64, row_count: u64, num_zones: usize) -> ArrayRef { if num_zones == 0 { - return Ok(ConstantArray::new(0u64, 0).into_array()); + return ConstantArray::new(0u64, 0).into_array(); } let last_zone_len = row_count - zone_len.saturating_mul((num_zones as u64) - 1); if num_zones == 1 || last_zone_len == zone_len { - return Ok(ConstantArray::new(last_zone_len, num_zones).into_array()); + return ConstantArray::new(last_zone_len, num_zones).into_array(); } let ends = unsafe { @@ -384,7 +398,7 @@ fn row_count_array(zone_len: u64, row_count: u64, num_zones: usize) -> VortexRes // SAFETY: `ends` are strictly increasing, terminate at `num_zones`, and align one-to-one // with the non-null run values. - Ok(unsafe { RunEnd::new_unchecked(ends, values, 0, num_zones) }.into_array()) + unsafe { RunEnd::new_unchecked(ends, values, 0, num_zones) }.into_array() } #[cfg(test)] @@ -872,16 +886,16 @@ mod tests { 10, ) .unwrap(); - let scope = zone_map.pruning_scope().unwrap(); + let scope_dtype = zone_map.scope.dtype(); let pruning_expr = falsify(&is_not_null(root()), PType::U64.into()); - let lowered = zone_map.lower_stats(pruning_expr, scope.dtype()).unwrap(); + let lowered = zone_map.lower_stats(pruning_expr).unwrap(); let expected = eq( get_item("null_count", get_item(STATS_FIELD, root())), get_item(ROW_COUNT_FIELD, root()), ) - .bind(scope.dtype()) + .bind(scope_dtype) .unwrap(); assert_eq!(lowered, expected); } From 84fdfb71a35d82fe3c6b05350f928c0002f96191 Mon Sep 17 00:00:00 2001 From: Robert Kruszewski Date: Sun, 16 Aug 2026 19:46:57 +0100 Subject: [PATCH 5/5] simplify Signed-off-by: Robert Kruszewski --- vortex-array/src/expr/bound_expression.rs | 24 ++- vortex-array/src/expr/traversal/mod.rs | 44 ++-- vortex-layout/src/layouts/zoned/zone_map.rs | 215 +++++++++++++------- 3 files changed, 197 insertions(+), 86 deletions(-) diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index ad95ae61c9d..4ea0ce28afd 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -143,7 +143,12 @@ impl BoundExpression { children: impl IntoIterator, ) -> VortexResult { let children = Vec::from_iter(children); - let BoundExpression::Scalar { scalar_fn, .. } = &self else { + let BoundExpression::Scalar { + dtype, + scalar_fn, + children: old_children, + } = &self + else { vortex_ensure!( children.is_empty(), "Root expression cannot have {} children", @@ -152,6 +157,23 @@ impl BoundExpression { return Ok(self); }; + // A return dtype is a function of the argument dtypes alone, so replacing children that + // type the same cannot change it. Rewrites usually substitute deep inside a tree and leave + // every dtype on the path to the root untouched, and recomputing one means a vector of + // cloned dtypes and a virtual call per node on that path. + if children.len() == old_children.len() + && children + .iter() + .zip(old_children.iter()) + .all(|(new, old)| new.dtype() == old.dtype()) + { + return Ok(Self::Scalar { + dtype: dtype.clone(), + scalar_fn: scalar_fn.clone(), + children: children.into(), + }); + } + Self::try_new(scalar_fn.clone(), children) } diff --git a/vortex-array/src/expr/traversal/mod.rs b/vortex-array/src/expr/traversal/mod.rs index 952a73f2657..c50e4e2c978 100644 --- a/vortex-array/src/expr/traversal/mod.rs +++ b/vortex-array/src/expr/traversal/mod.rs @@ -556,28 +556,38 @@ impl Node for BoundExpression { }; let mut order = TraversalOrder::Continue; - let mut changed = false; - let children = children - .iter() - .cloned() - .map(|child| match order { - TraversalOrder::Continue | TraversalOrder::Skip => f(child).map(|result| { + // Stays `None` until a child actually changes. Most nodes of a rewritten tree are + // untouched, and collecting their children into a vector only to discard it is the + // dominant cost of a rewrite over a large predicate. + let mut rewritten: Option> = None; + + for (index, child) in children.iter().enumerate() { + let value = match order { + TraversalOrder::Continue | TraversalOrder::Skip => { + let result = f(child.clone())?; order = result.order; - changed |= result.changed; + if result.changed && rewritten.is_none() { + let mut prefix = Vec::with_capacity(children.len()); + prefix.extend_from_slice(&children[..index]); + rewritten = Some(prefix); + } result.value - }), - TraversalOrder::Stop => Ok(child), - }) - .collect::>>()?; + } + TraversalOrder::Stop => child.clone(), + }; - if changed { - Ok(Transformed { - value: self.with_children(children)?, + if let Some(rewritten) = &mut rewritten { + rewritten.push(value); + } + } + + match rewritten { + Some(rewritten) => Ok(Transformed { + value: self.with_children(rewritten)?, order, changed: true, - }) - } else { - Ok(Transformed::no(self)) + }), + None => Ok(Transformed::no(self)), } } diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index 10b3e8984de..ee811cb50ce 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -3,6 +3,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::cell::RefCell; use std::sync::Arc; use vortex_array::ArrayRef; @@ -21,7 +22,7 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::DType; -use vortex_array::dtype::FieldNames; +use vortex_array::dtype::FieldName; use vortex_array::dtype::StructFields; use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; @@ -38,6 +39,7 @@ use vortex_array::stats::bind::StatBinder; use vortex_array::stats::bind::bind_stats; use vortex_array::validity::Validity; use vortex_buffer::buffer; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -60,8 +62,14 @@ pub struct ZoneMap { array: StructArray, // Aggregate functions stored in the zone map, ordered by their stats-table fields. aggregate_fns: Arc<[AggregateFnRef]>, + // Stats-table column backing each entry of `aggregate_fns`, or `None` when the table has no + // column for it. Resolving one means formatting the aggregate's display name and scanning the + // table, which is too much to repeat for every stat a predicate asks for. + aggregate_field_names: Arc<[Option]>, // Scope that lowered pruning predicates are evaluated against. See [`pruning_scope`]. scope: StructArray, + // Name of the row-count column within `scope`. See [`row_count_field_name`]. + row_count_field: FieldName, } impl ZoneMap { @@ -90,12 +98,25 @@ impl ZoneMap { zone_len: u64, row_count: u64, ) -> Self { - let scope = pruning_scope(&array, zone_len, row_count); + let (scope, row_count_field) = pruning_scope(&array, zone_len, row_count); + let aggregate_field_names = aggregate_fns + .iter() + .map(|aggregate_fn| { + let field_name = FieldName::from(aggregate_fn.to_string()); + array + .unmasked_field_by_name_opt(&field_name) + .is_some() + .then_some(field_name) + }) + .collect(); + Self { column_dtype, array, aggregate_fns, + aggregate_field_names, scope, + row_count_field, } } @@ -131,13 +152,12 @@ impl ZoneMap { /// `true` means the zone cannot contain matching rows and can be skipped. /// /// Row-count placeholders are resolved during lowering against a per-zone column that - /// [`ZoneMap::pruning_scope`] materializes, so the lowered predicate is directly evaluable. + /// [`pruning_scope`] materializes, so the lowered predicate is directly evaluable. pub fn prune( &self, predicate: &BoundExpression, session: &VortexSession, ) -> VortexResult { - let mut ctx = session.create_execution_ctx(); let predicate = self.lower_stats(predicate.clone())?; // A rewrite rule that proves its case from the predicate alone lowers to a constant, which @@ -150,6 +170,7 @@ impl ZoneMap { }); } + let mut ctx = session.create_execution_ctx(); self.scope .clone() .into_array() @@ -162,6 +183,7 @@ impl ZoneMap { let binder = ZoneMapStatsBinder { zone_map: self, scope_dtype: self.scope.dtype(), + bound: RefCell::default(), }; bind_stats(predicate, &binder) } @@ -169,50 +191,59 @@ impl ZoneMap { /// Build the scope that a lowered pruning predicate is evaluated against. /// -/// The scope is a two-field struct: [`STATS_FIELD`] nests the stored zone-map table, and -/// [`ROW_COUNT_FIELD`] holds the number of rows in each zone. The row count is a layout property -/// rather than a stored stat, and the final zone may be shorter than the nominal zone length, so it -/// cannot be resolved to a literal. Materializing it costs nothing: uniform zones use a -/// [`ConstantArray`] and a short final zone uses a two-run run-end encoded array. +/// The scope is the stored zone-map table with one extra column appended, holding the number of +/// rows in each zone. The row count is a layout property rather than a stored stat, and the final +/// zone may be shorter than the nominal zone length, so it cannot be resolved to a literal. +/// Materializing it costs nothing: uniform zones use a [`ConstantArray`] and a short final zone +/// uses a two-run run-end encoded array. /// -/// Nesting is what keeps the row count addressable. Appending it beside the stat columns would put -/// a name this module chooses into a namespace that aggregate display names and legacy stat names -/// also write to, and a collision would not be loud: [`StructFields`] permits duplicate names and -/// resolves lookups to the *first* match, so a colliding stat column would silently shadow the row -/// count and corrupt pruning. One level down, stat names cannot reach the two names this module -/// owns. +/// Keeping the row count in the same namespace as the stat columns is what makes a stat reference +/// a single [`get_item`]. Every reference is evaluated per zone, and a lowered predicate can hold +/// dozens of them, so the extra hop a nested scope would add is paid over and over. The name +/// collision that flatness risks is handled once, at construction, by [`row_count_field_name`]. /// /// The scope is built once per zone map rather than per predicate, because its dtype and its /// row-count column depend only on the stats table and the layout's zone geometry. -fn pruning_scope(array: &StructArray, zone_len: u64, row_count: u64) -> StructArray { - let num_zones = array.len(); - let row_counts = row_count_array(zone_len, row_count, num_zones); - let fields = StructFields::new( - FieldNames::from([STATS_FIELD, ROW_COUNT_FIELD]), - vec![array.dtype().clone(), row_counts.dtype().clone()], - ); - - // SAFETY: both fields are `num_zones` long and their dtypes are taken from the arrays - // themselves, so they match the struct dtype by construction. - unsafe { - StructArray::new_unchecked( - [array.clone().into_array(), row_counts], - fields, - num_zones, - Validity::NonNullable, - ) - } +fn pruning_scope(array: &StructArray, zone_len: u64, row_count: u64) -> (StructArray, FieldName) { + let row_counts = row_count_array(zone_len, row_count, array.len()); + let row_count_field = row_count_field_name(array.struct_fields()); + let scope = array + .with_column(row_count_field.clone(), row_counts) + .vortex_expect("row-count column matches the stats table length"); + + (scope, row_count_field) } -/// Field of the pruning scope nesting the stored zone-map table. -const STATS_FIELD: &str = "stats"; +/// Pick a name for the pruning scope's row-count column that no stat column already uses. +/// +/// A collision would not be loud: [`StructFields`] permits duplicate names and resolves lookups to +/// the *first* match, so a stats column named `row_count` would silently shadow the row count and +/// corrupt pruning. Stat column names come from aggregate display names and legacy stat names, so +/// this module cannot reserve a name up front; it appends underscores until the name is free. +fn row_count_field_name(stats_fields: &StructFields) -> FieldName { + if stats_fields.find(ROW_COUNT_FIELD).is_none() { + return FieldName::from(ROW_COUNT_FIELD); + } + + let mut name = String::from(ROW_COUNT_FIELD); + while stats_fields.find(&name).is_some() { + name.push('_'); + } + FieldName::from(name) +} -/// Field of the pruning scope holding the number of rows in each zone. +/// Preferred name of the pruning scope's row-count column. const ROW_COUNT_FIELD: &str = "row_count"; struct ZoneMapStatsBinder<'a> { zone_map: &'a ZoneMap, scope_dtype: &'a DType, + // Aggregates already lowered during this pass, keyed by the requested aggregate function. + // + // A falsifier repeats the same stat many times — a sixteen-term `IN` list asks for `min` and + // `max` once per term — and resolving one means scanning the stats table and binding a fresh + // expression against the scope. + bound: RefCell)>>, } impl StatBinder for ZoneMapStatsBinder<'_> { @@ -232,6 +263,34 @@ impl StatBinder for ZoneMapStatsBinder<'_> { self.zone_map.column_dtype ); + if let Some((_, cached)) = self + .bound + .borrow() + .iter() + .find(|(cached_fn, _)| cached_fn == aggregate_fn) + { + return Ok(cached.clone()); + } + + let bound = self.bind_aggregate_uncached(aggregate_fn)?; + self.bound + .borrow_mut() + .push((aggregate_fn.clone(), bound.clone())); + Ok(bound) + } + + fn bind_row_count(&self) -> VortexResult> { + get_item(self.zone_map.row_count_field.clone(), root()) + .bind(self.scope_dtype) + .map(Some) + } +} + +impl ZoneMapStatsBinder<'_> { + fn bind_aggregate_uncached( + &self, + aggregate_fn: &AggregateFnRef, + ) -> VortexResult> { if let Some(stat_expr) = self.zone_map.aggregate_field_expr(aggregate_fn) { return Ok(Some(self.bind_target(stat_expr)?)); } @@ -279,56 +338,46 @@ impl StatBinder for ZoneMapStatsBinder<'_> { Ok(None) } - fn bind_row_count(&self) -> VortexResult> { - get_item(ROW_COUNT_FIELD, root()) - .bind(self.scope_dtype) - .map(Some) - } -} - -impl ZoneMapStatsBinder<'_> { /// Bind a stat expression against the pruning scope it was built for. fn bind_target(&self, expr: Expression) -> VortexResult { expr.bind(self.scope_dtype) } } -/// Root of the stored zone-map table within the pruning scope. -/// -/// Stat expressions are built against this rather than against the scope root, so that binding a -/// stat is a single walk rather than a build-then-rebase. -fn stats_root() -> Expression { - get_item(STATS_FIELD, root()) -} - impl ZoneMap { fn aggregate_field_expr(&self, requested: &AggregateFnRef) -> Option { + // A stat the zone map stores verbatim is the overwhelmingly common case, and it needs no + // name at all: stats-table columns are named after the aggregates in `aggregate_fns`, so a + // stored aggregate equal to the requested one names the same column the display-name + // lookup below would find. + if let Some(field_name) = self.stored_field_name(requested) { + return Some(aggregate_result_expr( + requested, + get_item(field_name.clone(), root()), + )); + } + let field_name = requested.to_string(); if self.array.unmasked_field_by_name_opt(&field_name).is_some() { return Some(aggregate_result_expr( requested, - get_item(field_name, stats_root()), + get_item(field_name, root()), )); } let mut approximate = None; - for stored in self.aggregate_fns.iter() { - let field_name = stored.to_string(); - if self.array.unmasked_field_by_name_opt(&field_name).is_none() { - continue; - } - + for (stored, field_name) in self.stored_fields() { match stored.can_satisfy(requested) { AggregateFnSatisfaction::Exact => { return Some(aggregate_result_expr( stored, - get_item(field_name, stats_root()), + get_item(field_name.clone(), root()), )); } AggregateFnSatisfaction::Approximate => { approximate = Some(aggregate_result_expr( stored, - get_item(field_name, stats_root()), + get_item(field_name.clone(), root()), )); } AggregateFnSatisfaction::No => {} @@ -338,6 +387,21 @@ impl ZoneMap { approximate } + /// The stats-table column storing `requested` verbatim, if the zone map has one. + fn stored_field_name(&self, requested: &AggregateFnRef) -> Option<&FieldName> { + self.stored_fields() + .find(|(stored, _)| *stored == requested) + .map(|(_, field_name)| field_name) + } + + /// The stored aggregates that have a column in the stats table, paired with that column. + fn stored_fields(&self) -> impl Iterator { + self.aggregate_fns + .iter() + .zip(self.aggregate_field_names.iter()) + .filter_map(|(stored, field_name)| Some((stored, field_name.as_ref()?))) + } + fn stat_field_expr(&self, stat: Stat) -> Option { if let Some(aggregate_fn) = stat.aggregate_fn() && let Some(expr) = self.aggregate_field_expr(&aggregate_fn) @@ -350,7 +414,7 @@ impl ZoneMap { fn legacy_stat_field_expr(&self, stat: Stat) -> Option { if self.array.unmasked_field_by_name_opt(stat.name()).is_some() { - return Some(get_item(stat.name(), stats_root())); + return Some(get_item(stat.name(), root())); } None @@ -432,6 +496,7 @@ mod tests { use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; + use vortex_array::dtype::StructFields; use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::expr::cast; @@ -458,8 +523,8 @@ mod tests { use vortex_mask::Mask; use crate::layouts::zoned::zone_map::ROW_COUNT_FIELD; - use crate::layouts::zoned::zone_map::STATS_FIELD; use crate::layouts::zoned::zone_map::ZoneMap; + use crate::layouts::zoned::zone_map::row_count_field_name; use crate::test::SESSION; fn falsify(expr: &Expression, dtype: DType) -> BoundExpression { @@ -806,10 +871,10 @@ mod tests { #[test] fn stat_column_cannot_shadow_the_row_count_field() { - // A stats column named `row_count` sits one level below the pruning scope's own - // `row_count` field, so it cannot shadow it. Were the two flattened into one namespace, - // this column would win every lookup: `StructFields` resolves duplicates to the first - // match, and pruning would silently read `999` as each zone's row count. + // A stats column already named `row_count` forces the scope's own row-count column to a + // different name. Without that, this column would win every lookup: `StructFields` + // resolves duplicate names to the first match, and pruning would silently read `999` as + // each zone's row count. let zone_map = unsafe { ZoneMap::new_unchecked( PType::U64.into(), @@ -831,6 +896,8 @@ mod tests { ) }; + assert_ne!(zone_map.row_count_field.as_ref(), ROW_COUNT_FIELD); + // Zones hold 4, 4 and 2 rows, so the last two are entirely null. let expr = is_not_null(root()); let pruning_expr = falsify(&expr, PType::U64.into()); @@ -842,6 +909,18 @@ mod tests { ); } + #[test] + fn row_count_field_name_keeps_looking_until_it_is_free() { + let taken = StructFields::new( + FieldNames::from(["row_count", "row_count_"]), + vec![PType::U64.into(), PType::U64.into()], + ); + assert_eq!(row_count_field_name(&taken).as_ref(), "row_count__"); + + let free = StructFields::new(FieldNames::from(["max"]), vec![PType::U64.into()]); + assert_eq!(row_count_field_name(&free).as_ref(), ROW_COUNT_FIELD); + } + #[test] fn constant_predicate_skips_per_zone_evaluation() { let zone_map = ZoneMap::try_new( @@ -892,7 +971,7 @@ mod tests { let lowered = zone_map.lower_stats(pruning_expr).unwrap(); let expected = eq( - get_item("null_count", get_item(STATS_FIELD, root())), + get_item("null_count", root()), get_item(ROW_COUNT_FIELD, root()), ) .bind(scope_dtype)