diff --git a/Cargo.lock b/Cargo.lock index a92a0f5be59..d519a17da89 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10325,6 +10325,7 @@ dependencies = [ "async-stream", "async-trait", "bit-vec", + "codspeed-divan-compat", "flatbuffers", "futures", "insta", diff --git a/vortex-array/src/scalar_fn/internal/row_count.rs b/vortex-array/src/scalar_fn/internal/row_count.rs index 290378c30a7..d69728c169b 100644 --- a/vortex-array/src/scalar_fn/internal/row_count.rs +++ b/vortex-array/src/scalar_fn/internal/row_count.rs @@ -5,9 +5,6 @@ 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; @@ -20,28 +17,25 @@ 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. +/// 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. /// -/// This expression *MUST* be replaced with a concrete array before evaluation. -/// Currently, the rewrite only happens in the context of stats pruning. +/// 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. /// -/// `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. +/// 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. /// -/// 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; @@ -92,66 +86,6 @@ impl ScalarFnVTable for RowCount { } } -/// 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; diff --git a/vortex-array/src/stats/bind.rs b/vortex-array/src/stats/bind.rs index e07a588de94..8535bf8ad33 100644 --- a/vortex-array/src/stats/bind.rs +++ b/vortex-array/src/stats/bind.rs @@ -12,6 +12,16 @@ //! 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 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; @@ -22,7 +32,10 @@ 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; /// A target that can bind abstract statistics to concrete expressions. /// @@ -42,6 +55,16 @@ pub trait StatBinder { stat_dtype: &DType, ) -> VortexResult>; + /// Bind the number of rows covered by each row of the stats scope. + /// + /// 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, 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. /// /// The default is a nullable null literal, which preserves three-valued @@ -61,19 +84,50 @@ pub fn bind_stats( binder: &B, ) -> VortexResult { Ok(predicate - .transform_down(|expr| { - if !expr.is::() { - return Ok(Transformed::no(expr)); - } - - match bind_stat_fn(&expr, binder)? { - 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), @@ -95,17 +149,23 @@ 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::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::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; @@ -144,6 +204,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()), + RowCountFn.new_expr(EmptyOptions, []), + ) + .bind(&self.stats_scope)?, + )); + } + let Some(stat) = Stat::from_aggregate_fn(aggregate_fn) else { return Ok(None); }; @@ -156,6 +229,10 @@ mod tests { Ok(None) } } + + fn bind_row_count(&self) -> VortexResult> { + lit(10u64).bind(&self.stats_scope).map(Some) + } } #[test] @@ -181,6 +258,58 @@ mod tests { Ok(()) } + #[test] + 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)?; + + assert_eq!( + bound, + eq(col("f_nan_count"), lit(10u64)).bind(&binder.stats_scope)? + ); + 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-array/src/stats/rewrite/builtins.rs b/vortex-array/src/stats/rewrite/builtins.rs index 571c8b5ff84..3cb5fdb06df 100644 --- a/vortex-array/src/stats/rewrite/builtins.rs +++ b/vortex-array/src/stats/rewrite/builtins.rs @@ -168,10 +168,19 @@ fn binary_falsify( let rhs_falsifier = ctx.falsify(rhs)?; or_collect(lhs_falsifier.into_iter().chain(rhs_falsifier)) } - Operator::Or => match (ctx.falsify(lhs)?, ctx.falsify(rhs)?) { - (Some(lhs), Some(rhs)) if P::EMIT_UNGUARDED_REWRITES => Some(and(lhs, rhs)), - _ => None, - }, + Operator::Or => { + // Check before recursing: falsifying the children first would repeat the + // whole recursive rewrite once per registered `Binary` rule, which is + // exponential in `Or`-nesting depth for chains like `a = 1 OR a = 2 OR ...`. + if !P::EMIT_UNGUARDED_REWRITES { + return Ok(None); + } + + match (ctx.falsify(lhs)?, ctx.falsify(rhs)?) { + (Some(lhs), Some(rhs)) => Some(and(lhs, rhs)), + _ => None, + } + } Operator::Add | Operator::Sub | Operator::Mul | Operator::Div => None, }) } @@ -704,6 +713,8 @@ fn stat_fn(expr: BoundExpression, aggregate_fn: AggregateFnRef) -> BoundExpressi mod tests { use std::sync::Arc; use std::sync::LazyLock; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; use vortex_error::VortexResult; use vortex_session::VortexSession; @@ -712,6 +723,7 @@ mod tests { use crate::aggregate_fn::AggregateFnVTableExt; use crate::aggregate_fn::EmptyOptions as AggregateEmptyOptions; use crate::aggregate_fn::fns::all_non_nan::AllNonNan; + use crate::array_session; use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; @@ -737,17 +749,23 @@ mod tests { 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; use crate::scalar_fn::fns::between::BetweenOptions; use crate::scalar_fn::fns::between::StrictComparison; + use crate::scalar_fn::fns::binary::Binary; 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; + use crate::stats::rewrite::StatsRewriteRule; + use crate::stats::session::StatsSessionExt; - static SESSION: LazyLock = LazyLock::new(crate::array_session); + static SESSION: LazyLock = LazyLock::new(array_session); fn stat(expr: Expression, stat: Stat) -> Expression { let aggregate_fn = stat.aggregate_fn().expect("stat should have aggregate fn"); @@ -856,6 +874,53 @@ mod tests { gt_eq(stat(col("a"), Stat::Min), lit(50)), )) ); + + let expr = or(gt(col("a"), lit(10)), lt(col("a"), lit(5))); + assert_rewrite_eq!( + falsify(&expr)?, + Some(and( + lt_eq(stat(col("a"), Stat::Max), lit(10)), + gt_eq(stat(col("a"), Stat::Min), lit(5)), + )) + ); + Ok(()) + } + + /// Counts how many times the stats rewrite visits a `Binary` node. + #[derive(Debug)] + struct BinaryVisitCounter(Arc); + + impl StatsRewriteRule for BinaryVisitCounter { + fn scalar_fn_id(&self) -> ScalarFnId { + Binary.id() + } + + fn falsify( + &self, + _expr: &BoundExpression, + _ctx: &StatsRewriteCtx<'_>, + ) -> VortexResult> { + self.0.fetch_add(1, Ordering::Relaxed); + Ok(None) + } + } + + #[test] + fn or_chain_falsify_visits_each_node_once() -> VortexResult<()> { + let session = array_session(); + let visits = Arc::new(AtomicUsize::new(0)); + session + .stats() + .register_rewrite(BinaryVisitCounter(Arc::clone(&visits))); + + let expr = (1..16).fold(eq(col("a"), lit(0)), |chain, i| { + or(chain, eq(col("a"), lit(i))) + }); + let falsifier = expr.bind(&test_scope())?.falsify(&session)?; + assert!(falsifier.is_some()); + + // One visit per `Binary` node: 16 comparisons plus 15 `or`s. + assert_eq!(visits.load(Ordering::Relaxed), 31); Ok(()) } diff --git a/vortex-file/src/pruning.rs b/vortex-file/src/pruning.rs index df327638d00..f3ff618661a 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 the scope's single row 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..49ecf4cabcb 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 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", + 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/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 c84c0b443dd..10b3e8984de 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -21,6 +21,8 @@ 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::StructFields; use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::expr::eq; @@ -30,9 +32,8 @@ 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::fns::literal::Literal; 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; @@ -59,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 { @@ -91,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, } } @@ -131,42 +130,89 @@ impl ZoneMap { /// [`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. + /// 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, 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)?; - - if !contains_row_count(&applied) { - return applied.null_as_false().execute(&mut ctx); + // 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 = self.scope.len(); + return Ok(match scalar.as_bool().value() { + Some(true) => Mask::new_true(len), + Some(false) | None => Mask::new_false(len), + }); } - 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) + self.scope + .clone() + .into_array() + .apply_bound(&predicate)? + .null_as_false() + .execute(&mut ctx) } fn lower_stats(&self, predicate: BoundExpression) -> VortexResult { - let binder = ZoneMapStatsBinder { zone_map: self }; + let binder = ZoneMapStatsBinder { + zone_map: self, + 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"; + +/// 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, + scope_dtype: &'a DType, } impl StatBinder for ZoneMapStatsBinder<'_> { @@ -232,21 +278,36 @@ 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.zone_map.array.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()), )); } @@ -259,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 => {} } @@ -283,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 @@ -307,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 { @@ -331,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)] @@ -368,15 +435,19 @@ 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; 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; @@ -386,6 +457,8 @@ 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::STATS_FIELD; use crate::layouts::zoned::zone_map::ZoneMap; use crate::test::SESSION; @@ -731,6 +804,102 @@ 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 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_dtype = zone_map.scope.dtype(); + + let pruning_expr = falsify(&is_not_null(root()), PType::U64.into()); + 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) + .unwrap(); + assert_eq!(lowered, expected); + } + #[test] fn unavailable_stat_fn_lowers_to_unknown_mask() { let zone_map = ZoneMap::try_new(