From 2e29e005bdeb2c798d0a3b049384ec10cf4eaf7e Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Wed, 5 Aug 2026 11:14:14 +0100 Subject: [PATCH 1/4] refactor: remove bound expression unbinding Signed-off-by: Joe Isaacs --- .../src/arrays/scalar_fn/vtable/mod.rs | 3 +- vortex-array/src/expr/bound_expression.rs | 87 ++--- vortex-array/src/expr/display.rs | 37 +++ vortex-array/src/expr/expression.rs | 32 +- .../src/expr/transform/bound_partition.rs | 32 +- vortex-array/src/scalar_fn/erased.rs | 7 +- vortex-array/src/scalar_fn/fns/between/mod.rs | 9 +- vortex-array/src/scalar_fn/fns/binary/mod.rs | 7 +- vortex-array/src/scalar_fn/fns/case_when.rs | 9 +- vortex-array/src/scalar_fn/fns/cast/mod.rs | 10 +- vortex-array/src/scalar_fn/fns/dynamic.rs | 17 +- vortex-array/src/scalar_fn/fns/get_item.rs | 5 +- vortex-array/src/scalar_fn/fns/is_not_null.rs | 20 +- vortex-array/src/scalar_fn/fns/is_null.rs | 8 +- vortex-array/src/scalar_fn/fns/like/mod.rs | 7 +- .../src/scalar_fn/fns/list_contains/mod.rs | 27 +- vortex-array/src/scalar_fn/fns/literal.rs | 3 +- vortex-array/src/scalar_fn/fns/pack.rs | 7 +- vortex-array/src/scalar_fn/fns/root.rs | 4 +- vortex-array/src/scalar_fn/fns/select.rs | 5 +- vortex-array/src/scalar_fn/fns/stat.rs | 6 +- .../src/scalar_fn/fns/variant_get/mod.rs | 6 +- vortex-array/src/scalar_fn/fns/zip/mod.rs | 9 +- vortex-array/src/scalar_fn/foreign.rs | 8 +- .../src/scalar_fn/internal/row_count.rs | 4 +- vortex-array/src/scalar_fn/typed.rs | 7 +- vortex-array/src/scalar_fn/vtable.rs | 22 +- vortex-array/src/stats/bind.rs | 93 +++--- vortex-array/src/stats/rewrite.rs | 93 +++--- vortex-array/src/stats/rewrite/builtins.rs | 300 +++++++++++------- vortex-file/src/file.rs | 3 +- vortex-file/src/pruning.rs | 33 +- vortex-file/src/v2/file_stats_reader.rs | 7 +- vortex-geo/src/prune/distance.rs | 66 ++-- vortex-geo/src/prune/intersects.rs | 39 +-- vortex-geo/src/prune/mod.rs | 103 ++++-- vortex-layout/src/layouts/dict/reader.rs | 65 ++-- vortex-layout/src/layouts/row_idx/expr.rs | 3 +- vortex-layout/src/layouts/zoned/pruning.rs | 14 +- vortex-layout/src/layouts/zoned/zone_map.rs | 172 +++++----- vortex-layout/src/scan/filter.rs | 5 +- 41 files changed, 802 insertions(+), 592 deletions(-) diff --git a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs index fa98bdb9683..bb597bc6bb8 100644 --- a/vortex-array/src/arrays/scalar_fn/vtable/mod.rs +++ b/vortex-array/src/arrays/scalar_fn/vtable/mod.rs @@ -38,6 +38,7 @@ use crate::dtype::DType; use crate::executor::ExecutionCtx; use crate::executor::ExecutionResult; use crate::expr::Expression; +use crate::expr::display::ExprDisplay; use crate::matcher::Matcher; use crate::scalar_fn; use crate::scalar_fn::Arity; @@ -309,7 +310,7 @@ impl scalar_fn::ScalarFnVTable for ArrayExpr { fn fmt_sql( &self, options: &Self::Options, - _expr: &Expression, + _expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> std::fmt::Result { write!(f, "{}", options.0.encoding_id()) diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index 3a35757068a..66e3944fd7e 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -12,13 +12,16 @@ use itertools::Itertools; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; +use vortex_session::VortexSession; use crate::dtype::DType; use crate::expr::Expression; use crate::expr::display::DisplayTreeExpr; use crate::expr::scope::Scope; use crate::scalar_fn::ScalarFnRef; +use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::root::Root; +use crate::stats::rewrite::StatsRewriteCtx; /// An [`Expression`] that has been type-checked against a [`Scope`]. /// @@ -171,6 +174,11 @@ impl BoundExpression { } } + /// Return the child at `index`. + pub fn child(&self, index: usize) -> &BoundExpression { + &self.children()[index] + } + /// The scalar function for this node, or `None` if it is the scope root. pub fn as_scalar(&self) -> Option<&ScalarFnRef> { match &self.kind { @@ -179,56 +187,53 @@ impl BoundExpression { } } + /// Return whether this node uses the given scalar-function vtable. + pub fn is(&self) -> bool { + self.as_scalar().is_some_and(ScalarFnRef::is::) + } + + /// Return the typed scalar-function options when this node uses the given vtable. + pub fn as_opt(&self) -> Option<&V::Options> { + self.as_scalar().and_then(ScalarFnRef::as_opt::) + } + + /// Return the typed scalar-function options for this node. + /// + /// # Panics + /// + /// Panics when this node is the scope root or uses a different scalar-function vtable. + pub fn as_(&self) -> &V::Options { + self.as_opt::() + .vortex_expect("Bound expression options type mismatch") + } + /// Whether this node is the scope root. pub fn is_root(&self) -> bool { matches!(self.kind, BoundKind::Root) } - /// Display the bound expression as a formatted tree structure. - pub fn display_tree(&self) -> impl Display { - DisplayTreeExpr(self) + /// Return an expression that proves this predicate is definitely false from statistics. + pub fn falsify(&self, session: &VortexSession) -> VortexResult> { + StatsRewriteCtx::new(session).falsify(self) } - /// Convert this bound tree back into its unbound logical representation. - /// - /// This rebuilds the expression iteratively; the bound representation does not retain a - /// second expression tree. - // TODO: This is temporary artifact of the migration from using `Expression`s to - // `BoundExpression`s - pub fn unbind(&self) -> Expression { - let mut pending = vec![(self, false)]; - let mut expressions = Vec::new(); - - while let Some((node, visited)) = pending.pop() { - match node.kind() { - BoundKind::Root => expressions.push(crate::expr::root()), - BoundKind::Scalar { - scalar_fn, - children, - } if visited => { - let child_start = expressions.len() - children.len(); - let child_expressions = expressions.split_off(child_start); - expressions.push( - Expression::try_new(scalar_fn.clone(), child_expressions) - .vortex_expect("a bound expression always has valid arity"), - ); - } - BoundKind::Scalar { children, .. } => { - pending.push((node, true)); - pending.extend(children.iter().rev().map(|child| (child, false))); - } - } - } + /// Return an expression that proves this predicate is definitely true from statistics. + pub fn satisfy(&self, session: &VortexSession) -> VortexResult> { + StatsRewriteCtx::new(session).satisfy(self) + } - expressions - .pop() - .vortex_expect("binding always produces one expression root") + /// Display the bound expression as a formatted tree structure. + pub fn display_tree(&self) -> impl Display { + DisplayTreeExpr(self) } } impl Display for BoundExpression { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - Display::fmt(&self.unbind(), f) + match self.kind() { + BoundKind::Scalar { scalar_fn, .. } => scalar_fn.fmt_sql(self, f), + BoundKind::Root => f.write_str("$"), + } } } @@ -300,7 +305,7 @@ mod tests { let bound = root().bind_scope(&scope())?; assert!(bound.is_root()); assert_eq!(bound.dtype(), &struct_dtype()); - assert_eq!(bound.unbind(), root()); + assert_eq!(bound, BoundExpression::new_root(struct_dtype())); Ok(()) } @@ -375,11 +380,7 @@ mod tests { assert_eq!(bound, independently_bound); assert_eq!(ExactBoundExpr(bound.clone()), ExactBoundExpr(bound.clone())); - assert_ne!( - ExactBoundExpr(bound.clone()), - ExactBoundExpr(independently_bound) - ); - assert_eq!(bound.unbind(), expr); + assert_ne!(ExactBoundExpr(bound), ExactBoundExpr(independently_bound)); Ok(()) } diff --git a/vortex-array/src/expr/display.rs b/vortex-array/src/expr/display.rs index 685af718134..cc078ff592e 100644 --- a/vortex-array/src/expr/display.rs +++ b/vortex-array/src/expr/display.rs @@ -16,6 +16,43 @@ pub enum DisplayFormat { Tree, } +/// Read-only expression-tree interface used by scalar functions for SQL-style formatting. +/// +/// Both [`Expression`] and [`BoundExpression`] implement this interface, allowing scalar +/// functions to format either representation without converting between them. +pub trait ExprDisplay: Display { + /// Return the child at `index`. + fn display_child(&self, index: usize) -> &dyn ExprDisplay; + + /// Return the number of children in this node. + fn display_children_count(&self) -> usize; + + /// Format the child at `index` using its compact SQL-style representation. + fn fmt_display_child(&self, index: usize, f: &mut Formatter<'_>) -> fmt::Result { + Display::fmt(self.display_child(index), f) + } +} + +impl ExprDisplay for Expression { + fn display_child(&self, index: usize) -> &dyn ExprDisplay { + Expression::child(self, index) + } + + fn display_children_count(&self) -> usize { + self.children().len() + } +} + +impl ExprDisplay for BoundExpression { + fn display_child(&self, index: usize) -> &dyn ExprDisplay { + &self.children()[index] + } + + fn display_children_count(&self) -> usize { + self.children().len() + } +} + trait DisplayTreeNode: Sized { fn tree_children(&self) -> &[Self]; diff --git a/vortex-array/src/expr/expression.rs b/vortex-array/src/expr/expression.rs index d7f85825dbe..853df76e446 100644 --- a/vortex-array/src/expr/expression.rs +++ b/vortex-array/src/expr/expression.rs @@ -12,13 +12,11 @@ use std::sync::Arc; use itertools::Itertools; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use vortex_session::VortexSession; use crate::dtype::DType; use crate::expr::display::DisplayTreeExpr; use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::fns::root::Root; -use crate::stats::rewrite::StatsRewriteCtx; /// A node in a Vortex expression tree. /// @@ -113,40 +111,12 @@ impl Expression { self.scalar_fn.validity(self) } - /// Returns an expression that proves this predicate is definitely false from stats. - /// - /// `scope` is the dtype of the row this expression evaluates over. - /// - /// If the returned expression evaluates to `true` for a stats scope, this expression is - /// guaranteed to be false for every row in that scope. `false` and `null` are unknown. - pub fn falsify( - &self, - scope: &DType, - session: &VortexSession, - ) -> VortexResult> { - StatsRewriteCtx::new(session, scope).falsify(self) - } - - /// Returns an expression that proves this predicate is definitely true from stats. - /// - /// `scope` is the dtype of the row this expression evaluates over. - /// - /// If the returned expression evaluates to `true` for a stats scope, this expression is - /// guaranteed to be true for every row in that scope. `false` and `null` are unknown. - pub fn satisfy( - &self, - scope: &DType, - session: &VortexSession, - ) -> VortexResult> { - StatsRewriteCtx::new(session, scope).satisfy(self) - } - /// Format the expression as a compact string. /// /// Since this is a recursive formatter, it is exposed on the public Expression type. /// See fmt_data that is only implemented on the vtable trait. pub fn fmt_sql(&self, f: &mut Formatter<'_>) -> fmt::Result { - self.scalar_fn().fmt_sql(self, f) + self.scalar_fn.fmt_sql(self, f) } /// Display the expression as a formatted tree structure. diff --git a/vortex-array/src/expr/transform/bound_partition.rs b/vortex-array/src/expr/transform/bound_partition.rs index 95fe3f4a7ea..6488192a831 100644 --- a/vortex-array/src/expr/transform/bound_partition.rs +++ b/vortex-array/src/expr/transform/bound_partition.rs @@ -372,7 +372,7 @@ mod tests { // An un-expanded root expression is annotated by all fields, but since it is a single node assert_eq!(partitioned.partitions.len(), 0); - assert_eq!(partitioned.root.unbind(), root()); + assert_eq!(partitioned.root, root().bind(&dtype).unwrap()); // Instead, callers must expand the root expression themselves. let expr = replace_root_fields(expr, fields); @@ -386,9 +386,13 @@ mod tests { let expr = get_item("y", get_item("a", root())); let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); + let root_dtype = + partition_root_dtype(&partitioned.partition_names, &partitioned.partitions); assert_eq!( - partitioned.root.unbind(), + partitioned.root, get_item("a_0", get_item("a", root())) + .bind(&root_dtype) + .unwrap() ); } @@ -406,14 +410,16 @@ mod tests { let split_a = partitioned.find_partition(&"a".into()).unwrap(); assert_eq!( - split_a.unbind(), - pack( + split_a, + &pack( [ ("a_0", get_item("x", get_item("a", root()))), ("a_1", get_item("y", get_item("a", root()))) ], NonNullable ) + .bind(&dtype) + .unwrap() ); } @@ -441,9 +447,11 @@ mod tests { let partitioned = partition_by_field(expr.bind(&dtype).unwrap(), &dtype).unwrap(); let expected = merge([get_item("a_0", col("a")), get_item("b_0", col("b"))]); + let root_dtype = + partition_root_dtype(&partitioned.partition_names, &partitioned.partitions); assert_eq!( - partitioned.root.unbind(), - expected, + partitioned.root, + expected.bind(&root_dtype).unwrap(), "{} {}", partitioned.root, expected @@ -453,11 +461,19 @@ mod tests { let part_a = partitioned.find_partition(&"a".into()).unwrap(); let expected_a = pack([("a_0", col("a"))], NonNullable); - assert_eq!(part_a.unbind(), expected_a, "{part_a} {expected_a}"); + assert_eq!( + part_a, + &expected_a.bind(&dtype).unwrap(), + "{part_a} {expected_a}" + ); let part_b = partitioned.find_partition(&"b".into()).unwrap(); let expected_b = pack([("b_0", pack([("b", col("b"))], NonNullable))], NonNullable); - assert_eq!(part_b.unbind(), expected_b, "{part_b} {expected_b}"); + assert_eq!( + part_b, + &expected_b.bind(&dtype).unwrap(), + "{part_b} {expected_b}" + ); } #[rstest] diff --git a/vortex-array/src/scalar_fn/erased.rs b/vortex-array/src/scalar_fn/erased.rs index 6e0011c297a..355678973c5 100644 --- a/vortex-array/src/scalar_fn/erased.rs +++ b/vortex-array/src/scalar_fn/erased.rs @@ -20,6 +20,7 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; use crate::expr::Expression; +use crate::expr::display::ExprDisplay; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ReduceCtx; @@ -161,7 +162,11 @@ impl ScalarFnRef { // ------------------------------------------------------------------ /// Format this expression in SQL-style format. - pub(crate) fn fmt_sql(&self, expr: &Expression, f: &mut Formatter<'_>) -> std::fmt::Result { + pub(crate) fn fmt_sql( + &self, + expr: &dyn ExprDisplay, + f: &mut Formatter<'_>, + ) -> std::fmt::Result { self.0.fmt_sql(expr, f) } diff --git a/vortex-array/src/scalar_fn/fns/between/mod.rs b/vortex-array/src/scalar_fn/fns/between/mod.rs index 1a2e038b2fb..c3131dbd342 100644 --- a/vortex-array/src/scalar_fn/fns/between/mod.rs +++ b/vortex-array/src/scalar_fn/fns/between/mod.rs @@ -25,6 +25,7 @@ use crate::arrays::Primitive; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::DType::Bool; +use crate::expr::display::ExprDisplay; use crate::expr::expression::Expression; use crate::scalar::Scalar; use crate::scalar_fn::Arity; @@ -224,7 +225,7 @@ impl ScalarFnVTable for Between { fn fmt_sql( &self, options: &Self::Options, - expr: &Expression, + expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> std::fmt::Result { let lower_op = if options.lower_strict.is_strict() { @@ -240,11 +241,11 @@ impl ScalarFnVTable for Between { write!( f, "({} {} {} {} {})", - expr.child(1), + expr.display_child(1), lower_op, - expr.child(0), + expr.display_child(0), upper_op, - expr.child(2) + expr.display_child(2) ) } diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index 8a986c079a0..b0c6e3c1841 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -19,6 +19,7 @@ use crate::ExecutionCtx; use crate::dtype::DType; use crate::dtype::Nullability; use crate::expr::and; +use crate::expr::display::ExprDisplay; use crate::expr::expression::Expression; use crate::expr::lit; use crate::scalar_fn::Arity; @@ -89,13 +90,13 @@ impl ScalarFnVTable for Binary { fn fmt_sql( &self, operator: &Operator, - expr: &Expression, + expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> std::fmt::Result { write!(f, "(")?; - expr.child(0).fmt_sql(f)?; + expr.fmt_display_child(0, f)?; write!(f, " {} ", operator)?; - expr.child(1).fmt_sql(f)?; + expr.fmt_display_child(1, f)?; write!(f, ")") } diff --git a/vortex-array/src/scalar_fn/fns/case_when.rs b/vortex-array/src/scalar_fn/fns/case_when.rs index 189556c1fdb..8779e0a179c 100644 --- a/vortex-array/src/scalar_fn/fns/case_when.rs +++ b/vortex-array/src/scalar_fn/fns/case_when.rs @@ -35,6 +35,7 @@ use crate::builders::builder_with_capacity; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::expr::Expression; +use crate::expr::display::ExprDisplay; use crate::scalar::Scalar; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -136,7 +137,7 @@ impl ScalarFnVTable for CaseWhen { fn fmt_sql( &self, options: &Self::Options, - expr: &Expression, + expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> fmt::Result { write!(f, "CASE")?; @@ -144,13 +145,13 @@ impl ScalarFnVTable for CaseWhen { write!( f, " WHEN {} THEN {}", - expr.child(i * 2), - expr.child(i * 2 + 1) + expr.display_child(i * 2), + expr.display_child(i * 2 + 1) )?; } if options.has_else { let else_idx = options.num_when_then_pairs as usize * 2; - write!(f, " ELSE {}", expr.child(else_idx))?; + write!(f, " ELSE {}", expr.display_child(else_idx))?; } write!(f, " END") } diff --git a/vortex-array/src/scalar_fn/fns/cast/mod.rs b/vortex-array/src/scalar_fn/fns/cast/mod.rs index 046e19a7c9f..dfc7fab6e10 100644 --- a/vortex-array/src/scalar_fn/fns/cast/mod.rs +++ b/vortex-array/src/scalar_fn/fns/cast/mod.rs @@ -32,6 +32,7 @@ use crate::arrays::VarBinView; use crate::arrays::struct_::compute::cast::struct_cast; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; +use crate::expr::display::ExprDisplay; use crate::expr::expression::Expression; use crate::expr::lit; use crate::scalar_fn::Arity; @@ -90,9 +91,14 @@ impl ScalarFnVTable for Cast { } } - fn fmt_sql(&self, dtype: &DType, expr: &Expression, f: &mut Formatter<'_>) -> std::fmt::Result { + fn fmt_sql( + &self, + dtype: &DType, + expr: &dyn ExprDisplay, + f: &mut Formatter<'_>, + ) -> std::fmt::Result { write!(f, "cast(")?; - expr.children()[0].fmt_sql(f)?; + expr.fmt_display_child(0, f)?; write!(f, " as {}", dtype)?; write!(f, ")") } diff --git a/vortex-array/src/scalar_fn/fns/dynamic.rs b/vortex-array/src/scalar_fn/fns/dynamic.rs index a741fcd9217..bcee437dfef 100644 --- a/vortex-array/src/scalar_fn/fns/dynamic.rs +++ b/vortex-array/src/scalar_fn/fns/dynamic.rs @@ -19,7 +19,8 @@ use crate::ExecutionCtx; use crate::IntoArray; use crate::arrays::ConstantArray; use crate::dtype::DType; -use crate::expr::Expression; +use crate::expr::BoundExpression; +use crate::expr::display::ExprDisplay; use crate::expr::traversal::NodeExt; use crate::expr::traversal::NodeVisitor; use crate::expr::traversal::TraversalOrder; @@ -64,10 +65,10 @@ impl ScalarFnVTable for DynamicComparison { fn fmt_sql( &self, dynamic: &DynamicComparisonExpr, - expr: &Expression, + expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> std::fmt::Result { - expr.child(0).fmt_sql(f)?; + expr.fmt_display_child(0, f)?; write!(f, " {} dynamic(", dynamic.operator)?; match dynamic.scalar() { None => write!(f, "scalar=")?, @@ -204,15 +205,19 @@ pub struct DynamicExprUpdates { } impl DynamicExprUpdates { - pub fn new(expr: &Expression) -> Option { + /// Track dynamic scalar functions contained in a bound expression tree. + pub fn new(expr: &BoundExpression) -> Option { #[derive(Default)] struct Visitor(Vec); impl NodeVisitor<'_> for Visitor { - type NodeTy = Expression; + type NodeTy = BoundExpression; fn visit_down(&mut self, node: &'_ Self::NodeTy) -> VortexResult { - if let Some(dynamic) = node.as_opt::() { + if let Some(dynamic) = node + .as_scalar() + .and_then(|scalar_fn| scalar_fn.as_opt::()) + { self.0.push(dynamic.clone()); } Ok(TraversalOrder::Continue) diff --git a/vortex-array/src/scalar_fn/fns/get_item.rs b/vortex-array/src/scalar_fn/fns/get_item.rs index af319ecf742..d132c441d24 100644 --- a/vortex-array/src/scalar_fn/fns/get_item.rs +++ b/vortex-array/src/scalar_fn/fns/get_item.rs @@ -20,6 +20,7 @@ use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::Nullability; use crate::expr::Expression; +use crate::expr::display::ExprDisplay; use crate::expr::lit; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -78,10 +79,10 @@ impl ScalarFnVTable for GetItem { fn fmt_sql( &self, field_name: &FieldName, - expr: &Expression, + expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> std::fmt::Result { - expr.children()[0].fmt_sql(f)?; + expr.fmt_display_child(0, f)?; write!(f, ".{}", field_name) } 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 acb3b9e0f35..48875083f59 100644 --- a/vortex-array/src/scalar_fn/fns/is_not_null.rs +++ b/vortex-array/src/scalar_fn/fns/is_not_null.rs @@ -13,7 +13,7 @@ use crate::IntoArray; use crate::arrays::ConstantArray; use crate::dtype::DType; use crate::dtype::Nullability; -use crate::expr::Expression; +use crate::expr::display::ExprDisplay; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::EmptyOptions; @@ -60,11 +60,11 @@ impl ScalarFnVTable for IsNotNull { fn fmt_sql( &self, _options: &Self::Options, - expr: &Expression, + expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> std::fmt::Result { write!(f, "is_not_null(")?; - expr.child(0).fmt_sql(f)?; + expr.fmt_display_child(0, f)?; write!(f, ")") } @@ -252,13 +252,17 @@ mod tests { #[test] fn test_is_not_null_falsification() -> VortexResult<()> { let expr = is_not_null(col("a")); + let dtype = test_harness::struct_dtype(); assert_eq!( - expr.falsify(&test_harness::struct_dtype(), &STATS_SESSION)?, - Some(or( - eq(null_count(col("a")), RowCount.new_expr(EmptyOptions, []),), - all_null(col("a")), - )) + expr.bind(&dtype)?.falsify(&STATS_SESSION)?, + 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 118c00ab987..773422455d9 100644 --- a/vortex-array/src/scalar_fn/fns/is_null.rs +++ b/vortex-array/src/scalar_fn/fns/is_null.rs @@ -238,13 +238,11 @@ mod tests { #[test] fn test_is_null_falsification() -> VortexResult<()> { let expr = is_null(col("a")); + let dtype = test_harness::struct_dtype(); assert_eq!( - expr.falsify(&test_harness::struct_dtype(), &STATS_SESSION)?, - Some(or( - eq(null_count(col("a")), lit(0u64)), - all_non_null(col("a")), - )) + expr.bind(&dtype)?.falsify(&STATS_SESSION)?, + 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/fns/like/mod.rs b/vortex-array/src/scalar_fn/fns/like/mod.rs index 3637e0384a3..a57c01ed937 100644 --- a/vortex-array/src/scalar_fn/fns/like/mod.rs +++ b/vortex-array/src/scalar_fn/fns/like/mod.rs @@ -31,6 +31,7 @@ use crate::dtype::DType; use crate::dtype::Nullability; use crate::expr::Expression; use crate::expr::and; +use crate::expr::display::ExprDisplay; use crate::scalar::Scalar; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -107,10 +108,10 @@ impl ScalarFnVTable for Like { fn fmt_sql( &self, options: &Self::Options, - expr: &Expression, + expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> std::fmt::Result { - expr.child(0).fmt_sql(f)?; + expr.fmt_display_child(0, f)?; if options.negated { write!(f, " not")?; } @@ -119,7 +120,7 @@ impl ScalarFnVTable for Like { } else { write!(f, " like ")?; } - expr.child(1).fmt_sql(f) + expr.fmt_display_child(1, f) } fn return_dtype(&self, _options: &Self::Options, arg_dtypes: &[DType]) -> VortexResult { diff --git a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs index 18fce32d701..30d8950da2c 100644 --- a/vortex-array/src/scalar_fn/fns/list_contains/mod.rs +++ b/vortex-array/src/scalar_fn/fns/list_contains/mod.rs @@ -600,23 +600,26 @@ mod tests { ); assert_eq!( - expr.falsify(&scope, &STATS_SESSION)?, - Some(and( + expr.bind(&scope)?.falsify(&STATS_SESSION)?, + Some( and( - or( - lt(stat(col("a"), Stat::Max), lit(1i32)), - gt(stat(col("a"), Stat::Min), lit(1i32)), + and( + or( + lt(stat(col("a"), Stat::Max), lit(1i32)), + gt(stat(col("a"), Stat::Min), lit(1i32)), + ), + or( + lt(stat(col("a"), Stat::Max), lit(2i32)), + gt(stat(col("a"), Stat::Min), lit(2i32)), + ) ), or( - lt(stat(col("a"), Stat::Max), lit(2i32)), - gt(stat(col("a"), Stat::Min), lit(2i32)), + lt(stat(col("a"), Stat::Max), lit(3i32)), + gt(stat(col("a"), Stat::Min), lit(3i32)), ) - ), - or( - lt(stat(col("a"), Stat::Max), lit(3i32)), - gt(stat(col("a"), Stat::Min), lit(3i32)), ) - )) + .bind(&scope)? + ) ); Ok(()) } diff --git a/vortex-array/src/scalar_fn/fns/literal.rs b/vortex-array/src/scalar_fn/fns/literal.rs index 0449ae27e29..4c95fd3172d 100644 --- a/vortex-array/src/scalar_fn/fns/literal.rs +++ b/vortex-array/src/scalar_fn/fns/literal.rs @@ -16,6 +16,7 @@ use crate::IntoArray; use crate::arrays::ConstantArray; use crate::dtype::DType; use crate::expr::Expression; +use crate::expr::display::ExprDisplay; use crate::scalar::Scalar; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -74,7 +75,7 @@ impl ScalarFnVTable for Literal { fn fmt_sql( &self, scalar: &Scalar, - _expr: &Expression, + _expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> std::fmt::Result { write!(f, "{}", scalar) diff --git a/vortex-array/src/scalar_fn/fns/pack.rs b/vortex-array/src/scalar_fn/fns/pack.rs index 84983d9d55e..606c7e07734 100644 --- a/vortex-array/src/scalar_fn/fns/pack.rs +++ b/vortex-array/src/scalar_fn/fns/pack.rs @@ -23,6 +23,7 @@ use crate::dtype::FieldNames; use crate::dtype::Nullability; use crate::dtype::StructFields; use crate::expr::Expression; +use crate::expr::display::ExprDisplay; use crate::expr::lit; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -105,13 +106,13 @@ impl ScalarFnVTable for Pack { fn fmt_sql( &self, options: &Self::Options, - expr: &Expression, + expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> std::fmt::Result { write!(f, "pack(")?; - for (i, (name, child)) in options.names.iter().zip(expr.children().iter()).enumerate() { + for (i, name) in options.names.iter().enumerate() { write!(f, "{}: ", name)?; - child.fmt_sql(f)?; + expr.fmt_display_child(i, f)?; if i + 1 < options.names.len() { write!(f, ", ")?; } diff --git a/vortex-array/src/scalar_fn/fns/root.rs b/vortex-array/src/scalar_fn/fns/root.rs index 646bed19a1b..0831e5b8585 100644 --- a/vortex-array/src/scalar_fn/fns/root.rs +++ b/vortex-array/src/scalar_fn/fns/root.rs @@ -11,7 +11,7 @@ use vortex_session::registry::CachedId; use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; -use crate::expr::expression::Expression; +use crate::expr::display::ExprDisplay; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::EmptyOptions; @@ -58,7 +58,7 @@ impl ScalarFnVTable for Root { fn fmt_sql( &self, _options: &Self::Options, - _expr: &Expression, + _expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> std::fmt::Result { write!(f, "$") diff --git a/vortex-array/src/scalar_fn/fns/select.rs b/vortex-array/src/scalar_fn/fns/select.rs index 2e0bbd02975..c61fc050596 100644 --- a/vortex-array/src/scalar_fn/fns/select.rs +++ b/vortex-array/src/scalar_fn/fns/select.rs @@ -24,6 +24,7 @@ use crate::arrays::struct_::StructArrayExt; use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::FieldNames; +use crate::expr::display::ExprDisplay; use crate::expr::expression::Expression; use crate::expr::field::DisplayFieldNames; use crate::expr::get_item; @@ -104,10 +105,10 @@ impl ScalarFnVTable for Select { fn fmt_sql( &self, selection: &FieldSelection, - expr: &Expression, + expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> std::fmt::Result { - expr.child(0).fmt_sql(f)?; + expr.fmt_display_child(0, f)?; match selection { FieldSelection::Include(fields) => { write!(f, "{{{}}}", DisplayFieldNames(fields)) diff --git a/vortex-array/src/scalar_fn/fns/stat.rs b/vortex-array/src/scalar_fn/fns/stat.rs index 612298424e4..93e06bf0c67 100644 --- a/vortex-array/src/scalar_fn/fns/stat.rs +++ b/vortex-array/src/scalar_fn/fns/stat.rs @@ -20,7 +20,7 @@ use crate::aggregate_fn::fns::all_non_null::AllNonNull; use crate::aggregate_fn::fns::all_null::AllNull; use crate::arrays::ConstantArray; use crate::dtype::DType; -use crate::expr::Expression; +use crate::expr::display::ExprDisplay; use crate::expr::stats::Precision; use crate::expr::stats::Stat; use crate::expr::stats::StatsProvider; @@ -101,11 +101,11 @@ impl ScalarFnVTable for StatFn { fn fmt_sql( &self, options: &Self::Options, - expr: &Expression, + expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> std::fmt::Result { write!(f, "stat(")?; - expr.child(0).fmt_sql(f)?; + expr.fmt_display_child(0, f)?; write!(f, ", {})", options.aggregate_fn()) } diff --git a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs index d939235d24f..6851584f023 100644 --- a/vortex-array/src/scalar_fn/fns/variant_get/mod.rs +++ b/vortex-array/src/scalar_fn/fns/variant_get/mod.rs @@ -25,7 +25,7 @@ use crate::builders::builder_with_capacity_in; use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::Nullability; -use crate::expr::Expression; +use crate::expr::display::ExprDisplay; use crate::scalar::Scalar; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -92,11 +92,11 @@ impl ScalarFnVTable for VariantGet { fn fmt_sql( &self, options: &Self::Options, - expr: &Expression, + expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> fmt::Result { write!(f, "variant_get(")?; - expr.child(0).fmt_sql(f)?; + expr.fmt_display_child(0, f)?; let path = options.path().to_string(); write!(f, ", \"{}\"", StringEscape(&path))?; if let Some(dtype) = options.dtype() { diff --git a/vortex-array/src/scalar_fn/fns/zip/mod.rs b/vortex-array/src/scalar_fn/fns/zip/mod.rs index 5cb84bc737e..65bd41f5c39 100644 --- a/vortex-array/src/scalar_fn/fns/zip/mod.rs +++ b/vortex-array/src/scalar_fn/fns/zip/mod.rs @@ -24,6 +24,7 @@ use crate::builders::builder_with_capacity; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::expr::Expression; +use crate::expr::display::ExprDisplay; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::EmptyOptions; @@ -80,15 +81,15 @@ impl ScalarFnVTable for Zip { fn fmt_sql( &self, _options: &Self::Options, - expr: &Expression, + expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> std::fmt::Result { write!(f, "zip(")?; - expr.child(0).fmt_sql(f)?; + expr.fmt_display_child(0, f)?; write!(f, ", ")?; - expr.child(1).fmt_sql(f)?; + expr.fmt_display_child(1, f)?; write!(f, ", ")?; - expr.child(2).fmt_sql(f)?; + expr.fmt_display_child(2, f)?; write!(f, ")") } diff --git a/vortex-array/src/scalar_fn/foreign.rs b/vortex-array/src/scalar_fn/foreign.rs index 94fd75b08f1..9161f54f289 100644 --- a/vortex-array/src/scalar_fn/foreign.rs +++ b/vortex-array/src/scalar_fn/foreign.rs @@ -13,7 +13,7 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; use crate::dtype::Nullability; -use crate::expr::Expression; +use crate::expr::display::ExprDisplay; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; @@ -93,15 +93,15 @@ impl ScalarFnVTable for ForeignScalarFnVTable { fn fmt_sql( &self, _options: &Self::Options, - expr: &Expression, + expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> fmt::Result { write!(f, "{}(", self.id)?; - for i in 0..expr.children().len() { + for i in 0..expr.display_children_count() { if i > 0 { write!(f, ", ")?; } - expr.child(i).fmt_sql(f)?; + expr.fmt_display_child(i, f)?; } write!(f, ")") } diff --git a/vortex-array/src/scalar_fn/internal/row_count.rs b/vortex-array/src/scalar_fn/internal/row_count.rs index 7f78acbef64..290378c30a7 100644 --- a/vortex-array/src/scalar_fn/internal/row_count.rs +++ b/vortex-array/src/scalar_fn/internal/row_count.rs @@ -11,7 +11,7 @@ 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::Expression; +use vortex_array::expr::display::ExprDisplay; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; @@ -64,7 +64,7 @@ impl ScalarFnVTable for RowCount { fn fmt_sql( &self, _options: &Self::Options, - _expr: &Expression, + _expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> std::fmt::Result { write!(f, "row_count()") diff --git a/vortex-array/src/scalar_fn/typed.rs b/vortex-array/src/scalar_fn/typed.rs index a6620735ff2..f18854394c9 100644 --- a/vortex-array/src/scalar_fn/typed.rs +++ b/vortex-array/src/scalar_fn/typed.rs @@ -24,6 +24,7 @@ use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; use crate::expr::Expression; +use crate::expr::display::ExprDisplay; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; @@ -92,8 +93,8 @@ pub(super) trait DynScalarFn: 'static + Send + Sync + super::sealed::Sealed { fn is_strict(&self) -> bool; fn is_fallible(&self) -> bool; - // Expression methods — take &Expression for tree traversal - fn fmt_sql(&self, expression: &Expression, f: &mut Formatter<'_>) -> fmt::Result; + // Expression-tree methods + fn fmt_sql(&self, expression: &dyn ExprDisplay, f: &mut Formatter<'_>) -> fmt::Result; fn simplify( &self, expression: &Expression, @@ -192,7 +193,7 @@ impl DynScalarFn for TypedScalarFnInstance { V::is_fallible(&self.vtable, &self.options) } - fn fmt_sql(&self, expression: &Expression, f: &mut Formatter<'_>) -> fmt::Result { + fn fmt_sql(&self, expression: &dyn ExprDisplay, f: &mut Formatter<'_>) -> fmt::Result { V::fmt_sql(&self.vtable, &self.options, expression, f) } diff --git a/vortex-array/src/scalar_fn/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index ece52f6c4b4..2c3b33423ba 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -19,8 +19,9 @@ use vortex_session::VortexSession; use crate::ArrayRef; use crate::ExecutionCtx; use crate::dtype::DType; +use crate::expr::BoundExpression; use crate::expr::Expression; -use crate::expr::traversal::Node; +use crate::expr::display::ExprDisplay; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::TypedScalarFnInstance; @@ -70,17 +71,17 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { /// Format this expression in a nice human-readable SQL-style format /// /// The implementation should recursively format child expressions by calling - /// `expr.child(i).fmt_sql(f)`. + /// `expr.fmt_display_child(i, f)`. fn fmt_sql( &self, options: &Self::Options, - expr: &Expression, + expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> fmt::Result { write!(f, "{}(", self.id())?; - let nchildren = expr.children_count(); - for (i, child) in expr.children().iter().enumerate() { - child.fmt_sql(f)?; + let nchildren = expr.display_children_count(); + for i in 0..nchildren { + expr.fmt_display_child(i, f)?; if i + 1 < nchildren { write!(f, ", ")?; } @@ -392,6 +393,15 @@ pub trait ScalarFnVTableExt: ScalarFnVTable { ) -> VortexResult { Expression::try_new(self.bind(options), children) } + + /// Try to create a bound expression with this vtable, the given options, and bound children. + fn try_new_bound_expr( + &self, + options: Self::Options, + children: impl IntoIterator, + ) -> VortexResult { + BoundExpression::try_new(self.bind(options), children) + } } impl ScalarFnVTableExt for V {} diff --git a/vortex-array/src/stats/bind.rs b/vortex-array/src/stats/bind.rs index 6d921fd0ab1..0b709f981a7 100644 --- a/vortex-array/src/stats/bind.rs +++ b/vortex-array/src/stats/bind.rs @@ -17,11 +17,12 @@ use vortex_error::VortexResult; use crate::aggregate_fn::AggregateFnRef; use crate::dtype::DType; -use crate::expr::Expression; -use crate::expr::lit; +use crate::expr::BoundExpression; use crate::expr::traversal::NodeExt; use crate::expr::traversal::Transformed; use crate::scalar::Scalar; +use crate::scalar_fn::ScalarFnVTableExt; +use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::stat::StatFn; /// A target that can bind abstract statistics to concrete expressions. @@ -31,26 +32,23 @@ use crate::scalar_fn::fns::stat::StatFn; /// field reference in the per-zone stats table, while a file-stats binder can translate the same /// placeholder into a literal value from the file footer. pub trait StatBinder { - /// The dtype scope used to type-check expressions before stats are bound. - fn scope(&self) -> &DType; - /// Bind `aggregate_fn(input)` to a concrete expression. /// /// Implementations should return `Ok(None)` when the requested aggregate /// statistic is unavailable in their backing representation. fn bind_aggregate( &self, - input: &Expression, + input: &BoundExpression, aggregate_fn: &AggregateFnRef, stat_dtype: &DType, - ) -> VortexResult>; + ) -> VortexResult>; /// Expression to use when a stat is unavailable. /// /// The default is a nullable null literal, which preserves three-valued /// pruning semantics for stats-table execution. - fn missing_stat(&self, dtype: DType) -> VortexResult { - Ok(null_expr(dtype)) + fn missing_stat(&self, dtype: DType) -> VortexResult { + null_expr(dtype) } } @@ -60,43 +58,37 @@ pub trait StatBinder { /// are responsible for expressing stat semantics; binding maps aggregate-backed /// stat requests to the concrete stats representation supported by the binder. pub fn bind_stats( - predicate: Expression, + predicate: BoundExpression, binder: &B, -) -> VortexResult { - let scope = binder.scope().clone(); +) -> VortexResult { Ok(predicate .transform_down(|expr| { if !expr.is::() { return Ok(Transformed::no(expr)); } - match bind_stat_fn(&expr, &scope, binder)? { + match bind_stat_fn(&expr, binder)? { Some(bound) => Ok(Transformed::yes(bound)), - None => { - let dtype = expr.return_dtype(&scope)?; - Ok(Transformed::yes(binder.missing_stat(dtype)?)) - } + None => Ok(Transformed::yes(binder.missing_stat(expr.dtype().clone())?)), } })? .into_inner()) } fn bind_stat_fn( - expr: &Expression, - scope: &DType, + expr: &BoundExpression, binder: &(impl StatBinder + ?Sized), -) -> VortexResult> { +) -> VortexResult> { let options = expr.as_::(); let aggregate_fn = options.aggregate_fn(); // `StatFn` has exactly one child: the expression the aggregate statistic is computed over. let input = expr.child(0); - let stat_dtype = expr.return_dtype(scope)?; - binder.bind_aggregate(input, aggregate_fn, &stat_dtype) + binder.bind_aggregate(input, aggregate_fn, expr.dtype()) } -fn null_expr(dtype: DType) -> Expression { - lit(Scalar::null(dtype.as_nullable())) +fn null_expr(dtype: DType) -> VortexResult { + Literal.try_new_bound_expr(Scalar::null(dtype.as_nullable()), []) } #[cfg(test)] @@ -111,6 +103,7 @@ mod tests { use crate::expr::col; 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; @@ -119,6 +112,7 @@ mod tests { struct TestBinder { input_scope: DType, + stats_scope: DType, bind_nan_count: bool, } @@ -132,28 +126,33 @@ mod tests { )]), Nullability::NonNullable, ), + stats_scope: DType::Struct( + StructFields::from_iter([( + "f_nan_count", + DType::Primitive(PType::U64, Nullability::NonNullable), + )]), + Nullability::NonNullable, + ), bind_nan_count, } } } impl StatBinder for TestBinder { - fn scope(&self) -> &DType { - &self.input_scope - } - fn bind_aggregate( &self, - _input: &Expression, + _input: &BoundExpression, aggregate_fn: &AggregateFnRef, _stat_dtype: &DType, - ) -> VortexResult> { + ) -> VortexResult> { let Some(stat) = Stat::from_aggregate_fn(aggregate_fn) else { return Ok(None); }; if stat == Stat::NaNCount && self.bind_nan_count { - Ok(Some(get_item("f_nan_count", root()))) + Ok(Some( + get_item("f_nan_count", root()).bind(&self.stats_scope)?, + )) } else { Ok(None) } @@ -164,9 +163,9 @@ mod tests { fn nan_count_binds_to_direct_stat_slot() -> VortexResult<()> { let binder = TestBinder::new(true); - let bound = bind_stats(nan_count(col("f")), &binder)?; + let bound = bind_stats(nan_count(col("f")).bind(&binder.input_scope)?, &binder)?; - assert_eq!(bound, col("f_nan_count")); + assert_eq!(bound, col("f_nan_count").bind(&binder.stats_scope)?); Ok(()) } @@ -174,9 +173,12 @@ mod tests { 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")), &binder)?; + let bound = bind_stats(all_non_nan(col("f")).bind(&binder.input_scope)?, &binder)?; - assert_eq!(bound, lit(Scalar::null(DType::Bool(Nullability::Nullable)))); + assert_eq!( + bound, + lit(Scalar::null(DType::Bool(Nullability::Nullable))).bind(&binder.stats_scope)? + ); Ok(()) } @@ -185,13 +187,22 @@ mod tests { let binder = TestBinder::new(false); let null_bool = lit(Scalar::null(DType::Bool(Nullability::Nullable))); - let bound = bind_stats(and(lit(false), all_non_nan(col("f"))), &binder)?; + let bound = bind_stats( + and(lit(false), all_non_nan(col("f"))).bind(&binder.input_scope)?, + &binder, + )?; - assert_eq!(bound, and(lit(false), null_bool.clone())); + assert_eq!( + bound, + and(lit(false), null_bool.clone()).bind(&binder.stats_scope)? + ); - let bound = bind_stats(or(lit(true), all_non_nan(col("f"))), &binder)?; + let bound = bind_stats( + or(lit(true), all_non_nan(col("f"))).bind(&binder.input_scope)?, + &binder, + )?; - assert_eq!(bound, or(lit(true), null_bool)); + assert_eq!(bound, or(lit(true), null_bool).bind(&binder.stats_scope)?); Ok(()) } @@ -199,9 +210,9 @@ mod tests { fn unrelated_expressions_do_not_request_nan_count() -> VortexResult<()> { let binder = TestBinder::new(false); - let bound = bind_stats(is_null(col("f")), &binder)?; + let bound = bind_stats(is_null(col("f")).bind(&binder.input_scope)?, &binder)?; - assert_eq!(bound, is_null(col("f"))); + assert_eq!(bound, is_null(col("f")).bind(&binder.input_scope)?); Ok(()) } } diff --git a/vortex-array/src/stats/rewrite.rs b/vortex-array/src/stats/rewrite.rs index 2bbeeb00022..45c8f6a4ce9 100644 --- a/vortex-array/src/stats/rewrite.rs +++ b/vortex-array/src/stats/rewrite.rs @@ -9,11 +9,14 @@ use std::sync::Arc; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_session::VortexSession; +use vortex_utils::iter::ReduceBalancedIterExt; use crate::dtype::DType; -use crate::expr::Expression; -use crate::expr::or_collect; +use crate::expr::BoundExpression; use crate::scalar_fn::ScalarFnId; +use crate::scalar_fn::ScalarFnVTableExt; +use crate::scalar_fn::fns::binary::Binary; +use crate::scalar_fn::fns::operators::Operator; use crate::stats::session::StatsSessionExt; mod builtins; @@ -53,9 +56,9 @@ pub trait StatsRewriteRule: Debug + Send + Sync + 'static { /// Returns `Ok(None)` when this rule cannot construct a sound falsity proof for `expr`. fn falsify( &self, - expr: &Expression, + expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { _ = expr; _ = ctx; Ok(None) @@ -73,9 +76,9 @@ pub trait StatsRewriteRule: Debug + Send + Sync + 'static { /// Returns `Ok(None)` when this rule cannot construct a sound truth proof for `expr`. fn satisfy( &self, - expr: &Expression, + expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { _ = expr; _ = ctx; Ok(None) @@ -85,13 +88,12 @@ pub trait StatsRewriteRule: Debug + Send + Sync + 'static { /// Context passed to stats rewrite rules. pub struct StatsRewriteCtx<'a> { session: &'a VortexSession, - scope: &'a DType, } impl<'a> StatsRewriteCtx<'a> { /// Create a rewrite context for `session`. - pub fn new(session: &'a VortexSession, scope: &'a DType) -> Self { - Self { session, scope } + pub fn new(session: &'a VortexSession) -> Self { + Self { session } } /// Returns the session that owns the rewrite registry. @@ -100,23 +102,23 @@ impl<'a> StatsRewriteCtx<'a> { } /// Return the dtype of `expr` within this rewrite scope. - pub fn return_dtype(&self, expr: &Expression) -> VortexResult { - expr.return_dtype(self.scope) + pub fn return_dtype(&self, expr: &BoundExpression) -> VortexResult { + Ok(expr.dtype().clone()) } /// Rewrite `expr` into a stats-backed falsifier. - pub fn falsify(&self, expr: &Expression) -> VortexResult> { + pub fn falsify(&self, expr: &BoundExpression) -> VortexResult> { self.ensure_predicate(expr)?; rewrite(expr, self, StatsRewriteRule::falsify) } /// Rewrite `expr` into a stats-backed satisfier. - pub fn satisfy(&self, expr: &Expression) -> VortexResult> { + pub fn satisfy(&self, expr: &BoundExpression) -> VortexResult> { self.ensure_predicate(expr)?; rewrite(expr, self, StatsRewriteRule::satisfy) } - fn ensure_predicate(&self, expr: &Expression) -> VortexResult<()> { + fn ensure_predicate(&self, expr: &BoundExpression) -> VortexResult<()> { let dtype = self.return_dtype(expr)?; vortex_ensure!( matches!(dtype, DType::Bool(_)), @@ -127,18 +129,18 @@ impl<'a> StatsRewriteCtx<'a> { } fn rewrite( - expr: &Expression, + expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>, apply: fn( &dyn StatsRewriteRule, - &Expression, + &BoundExpression, &StatsRewriteCtx<'_>, - ) -> VortexResult>, -) -> VortexResult> { - let rules = ctx - .session() - .stats() - .rewrite_rules_for(expr.scalar_fn().id()); + ) -> VortexResult>, +) -> VortexResult> { + let Some(scalar_fn) = expr.as_scalar() else { + return Ok(None); + }; + let rules = ctx.session().stats().rewrite_rules_for(scalar_fn.id()); let Some(rules) = rules else { return Ok(None); }; @@ -150,7 +152,9 @@ fn rewrite( } } - Ok(or_collect(rewrites)) + rewrites + .into_iter() + .try_reduce_balanced(|lhs, rhs| Binary.try_new_bound_expr(Operator::Or, [lhs, rhs])) } #[cfg(test)] @@ -162,7 +166,7 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; - use crate::expr::Expression; + use crate::expr::BoundExpression; use crate::expr::lit; use crate::expr::or; use crate::scalar_fn::ScalarFnId; @@ -172,8 +176,8 @@ mod tests { #[derive(Debug)] struct StaticLiteralRule { - falsifier: Option, - satisfier: Option, + falsifier: Option, + satisfier: Option, } impl StatsRewriteRule for StaticLiteralRule { @@ -183,17 +187,17 @@ mod tests { fn falsify( &self, - _expr: &Expression, + _expr: &BoundExpression, _ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(self.falsifier.clone()) } fn satisfy( &self, - _expr: &Expression, + _expr: &BoundExpression, _ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(self.satisfier.clone()) } } @@ -203,17 +207,17 @@ mod tests { let session = crate::array_session(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); session.stats().register_rewrite(StaticLiteralRule { - falsifier: Some(lit(false)), + falsifier: Some(lit(false).bind(&dtype)?), satisfier: None, }); session.stats().register_rewrite(StaticLiteralRule { - falsifier: Some(lit(true)), + falsifier: Some(lit(true).bind(&dtype)?), satisfier: None, }); assert_eq!( - lit(true).falsify(&dtype, &session)?, - Some(or(lit(false), lit(true))) + lit(true).bind(&dtype)?.falsify(&session)?, + Some(or(lit(false), lit(true)).bind(&dtype)?) ); Ok(()) } @@ -224,16 +228,16 @@ mod tests { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); session.stats().register_rewrite(StaticLiteralRule { falsifier: None, - satisfier: Some(lit(false)), + satisfier: Some(lit(false).bind(&dtype)?), }); session.stats().register_rewrite(StaticLiteralRule { falsifier: None, - satisfier: Some(lit(true)), + satisfier: Some(lit(true).bind(&dtype)?), }); assert_eq!( - lit(true).satisfy(&dtype, &session)?, - Some(or(lit(false), lit(true))) + lit(true).bind(&dtype)?.satisfy(&session)?, + Some(or(lit(false), lit(true)).bind(&dtype)?) ); Ok(()) } @@ -243,17 +247,20 @@ mod tests { let session = crate::array_session(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - assert_eq!(lit(true).falsify(&dtype, &session)?, None); - assert_eq!(lit(true).satisfy(&dtype, &session)?, None); + let expr = lit(true).bind(&dtype)?; + assert_eq!(expr.falsify(&session)?, None); + assert_eq!(expr.satisfy(&session)?, None); Ok(()) } #[test] - fn non_predicate_expression_errors() { + fn non_predicate_expression_errors() -> VortexResult<()> { let session = crate::array_session(); let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); - assert!(lit(7).falsify(&dtype, &session).is_err()); - assert!(lit(7).satisfy(&dtype, &session).is_err()); + let expr = lit(7).bind(&dtype)?; + assert!(expr.falsify(&session).is_err()); + assert!(expr.satisfy(&session).is_err()); + Ok(()) } } diff --git a/vortex-array/src/stats/rewrite/builtins.rs b/vortex-array/src/stats/rewrite/builtins.rs index 77d5b31bca9..b68eb579a19 100644 --- a/vortex-array/src/stats/rewrite/builtins.rs +++ b/vortex-array/src/stats/rewrite/builtins.rs @@ -3,7 +3,9 @@ use std::sync::Arc; +use vortex_error::VortexExpect; use vortex_error::VortexResult; +use vortex_utils::iter::ReduceBalancedIterExt; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTableExt; @@ -12,19 +14,9 @@ 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::dtype::DType; -use crate::expr::Expression; -use crate::expr::and; -use crate::expr::and_collect; -use crate::expr::cast; -use crate::expr::eq; -use crate::expr::gt; -use crate::expr::gt_eq; -use crate::expr::lit; -use crate::expr::lt; -use crate::expr::lt_eq; -use crate::expr::or; -use crate::expr::or_collect; +use crate::expr::BoundExpression as BoundExpr; use crate::expr::stats::Stat; +use crate::scalar::Scalar; use crate::scalar::StringLike; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnId; @@ -68,6 +60,65 @@ pub(crate) fn register_builtins(session: &StatsSession) { session.register_rewrite(DynamicComparisonAllNonNanStatsRewrite); } +fn binary(operator: Operator, lhs: BoundExpr, rhs: BoundExpr) -> BoundExpr { + Binary + .try_new_bound_expr(operator, [lhs, rhs]) + .vortex_expect("stats rewrites must construct well-typed binary expressions") +} + +fn and(lhs: BoundExpr, rhs: BoundExpr) -> BoundExpr { + binary(Operator::And, lhs, rhs) +} + +fn or(lhs: BoundExpr, rhs: BoundExpr) -> BoundExpr { + binary(Operator::Or, lhs, rhs) +} + +fn eq(lhs: BoundExpr, rhs: BoundExpr) -> BoundExpr { + binary(Operator::Eq, lhs, rhs) +} + +fn gt(lhs: BoundExpr, rhs: BoundExpr) -> BoundExpr { + binary(Operator::Gt, lhs, rhs) +} + +fn gt_eq(lhs: BoundExpr, rhs: BoundExpr) -> BoundExpr { + binary(Operator::Gte, lhs, rhs) +} + +fn lt(lhs: BoundExpr, rhs: BoundExpr) -> BoundExpr { + binary(Operator::Lt, lhs, rhs) +} + +fn lt_eq(lhs: BoundExpr, rhs: BoundExpr) -> BoundExpr { + binary(Operator::Lte, lhs, rhs) +} + +fn and_collect(exprs: impl IntoIterator) -> Option { + exprs.into_iter().reduce_balanced(and) +} + +fn or_collect(exprs: impl IntoIterator) -> Option { + exprs.into_iter().reduce_balanced(or) +} + +fn lit(value: impl Into) -> BoundExpr { + Literal + .try_new_bound_expr(value.into(), []) + .vortex_expect("literal expressions are always well-typed") +} + +fn cast(expr: BoundExpr, dtype: DType) -> BoundExpr { + Cast.try_new_bound_expr(dtype, [expr]) + .vortex_expect("stats rewrites only preserve casts from a bound predicate") +} + +fn row_count() -> BoundExpr { + RowCount + .try_new_bound_expr(EmptyOptions, []) + .vortex_expect("row-count expressions are always well-typed") +} + #[derive(Debug)] struct BinaryNanCountStatsRewrite; @@ -78,9 +129,9 @@ impl StatsRewriteRule for BinaryNanCountStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpr, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { binary_falsify::(expr, ctx) } } @@ -95,17 +146,17 @@ impl StatsRewriteRule for BinaryAllNonNanStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpr, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { binary_falsify::(expr, ctx) } } fn binary_falsify( - expr: &Expression, + expr: &BoundExpr, ctx: &StatsRewriteCtx<'_>, -) -> VortexResult> { +) -> VortexResult> { let operator = expr.as_::(); let lhs = expr.child(0); let rhs = expr.child(1); @@ -178,16 +229,16 @@ impl StatsRewriteRule for BetweenStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpr, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { let options = expr.as_::(); let arr = expr.child(0).clone(); let lower = expr.child(1).clone(); let upper = expr.child(2).clone(); - let lhs = Binary.new_expr(options.lower_strict.to_operator(), [lower, arr.clone()]); - let rhs = Binary.new_expr(options.upper_strict.to_operator(), [arr, upper]); + let lhs = binary(options.lower_strict.to_operator(), lower, arr.clone()); + let rhs = binary(options.upper_strict.to_operator(), arr, upper); ctx.falsify(&and(lhs, rhs)) } } @@ -202,19 +253,18 @@ impl StatsRewriteRule for IsNullNullCountStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpr, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(null_count(expr.child(0), ctx).map(|null_count| eq(null_count, lit(0u64)))) } fn satisfy( &self, - expr: &Expression, + expr: &BoundExpr, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { - Ok(null_count(expr.child(0), ctx) - .map(|null_count| eq(null_count, RowCount.new_expr(EmptyOptions, [])))) + ) -> VortexResult> { + Ok(null_count(expr.child(0), ctx).map(|null_count| eq(null_count, row_count()))) } } @@ -228,9 +278,9 @@ impl StatsRewriteRule for IsNullAllNonNullStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpr, _ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(Some(all_non_null(expr.child(0)))) } } @@ -245,9 +295,9 @@ impl StatsRewriteRule for IsNullAllNullStatsRewrite { fn satisfy( &self, - expr: &Expression, + expr: &BoundExpr, _ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(Some(all_null(expr.child(0)))) } } @@ -262,18 +312,17 @@ impl StatsRewriteRule for IsNotNullNullCountStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpr, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { - Ok(null_count(expr.child(0), ctx) - .map(|null_count| eq(null_count, RowCount.new_expr(EmptyOptions, [])))) + ) -> VortexResult> { + Ok(null_count(expr.child(0), ctx).map(|null_count| eq(null_count, row_count()))) } fn satisfy( &self, - expr: &Expression, + expr: &BoundExpr, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(null_count(expr.child(0), ctx).map(|null_count| eq(null_count, lit(0u64)))) } } @@ -288,9 +337,9 @@ impl StatsRewriteRule for IsNotNullAllNullStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpr, _ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(Some(all_null(expr.child(0)))) } } @@ -305,9 +354,9 @@ impl StatsRewriteRule for IsNotNullAllNonNullStatsRewrite { fn satisfy( &self, - expr: &Expression, + expr: &BoundExpr, _ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(Some(all_non_null(expr.child(0)))) } } @@ -322,9 +371,9 @@ impl StatsRewriteRule for LikeStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpr, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { let like_options = expr.as_::(); if like_options.negated || like_options.case_insensitive { return Ok(None); @@ -377,9 +426,9 @@ impl StatsRewriteRule for ListContainsNanCountStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpr, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { list_contains_falsify::(expr, ctx) } } @@ -394,17 +443,17 @@ impl StatsRewriteRule for ListContainsAllNonNanStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpr, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { list_contains_falsify::(expr, ctx) } } fn list_contains_falsify( - expr: &Expression, + expr: &BoundExpr, ctx: &StatsRewriteCtx<'_>, -) -> VortexResult> { +) -> VortexResult> { let list = expr.child(0); let needle = expr.child(1); @@ -451,9 +500,9 @@ impl StatsRewriteRule for DynamicComparisonNanCountStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpr, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { dynamic_comparison_falsify::(expr, ctx) } } @@ -468,17 +517,17 @@ impl StatsRewriteRule for DynamicComparisonAllNonNanStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpr, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { dynamic_comparison_falsify::(expr, ctx) } } fn dynamic_comparison_falsify( - expr: &Expression, + expr: &BoundExpr, ctx: &StatsRewriteCtx<'_>, -) -> VortexResult> { +) -> VortexResult> { let dynamic = expr.as_::(); let lhs = expr.child(0); @@ -492,47 +541,49 @@ fn dynamic_comparison_falsify( return Ok(None); }; - let value_predicate = DynamicComparison.new_expr( - DynamicComparisonExpr { - operator, - rhs: Arc::clone(&dynamic.rhs), - default: !dynamic.default, - }, - [lhs_stat], - ); + let value_predicate = DynamicComparison + .try_new_bound_expr( + DynamicComparisonExpr { + operator, + rhs: Arc::clone(&dynamic.rhs), + default: !dynamic.default, + }, + [lhs_stat], + ) + .vortex_expect("a rewritten dynamic comparison preserves its bound input type"); with_non_nan_guards::

(ctx, [lhs], value_predicate) } -fn min(expr: &Expression, ctx: &StatsRewriteCtx<'_>) -> Option { +fn min(expr: &BoundExpr, ctx: &StatsRewriteCtx<'_>) -> Option { stat_expr(expr, Stat::Min, ctx) } -fn max(expr: &Expression, ctx: &StatsRewriteCtx<'_>) -> Option { +fn max(expr: &BoundExpr, ctx: &StatsRewriteCtx<'_>) -> Option { stat_expr(expr, Stat::Max, ctx) } -fn null_count(expr: &Expression, ctx: &StatsRewriteCtx<'_>) -> Option { +fn null_count(expr: &BoundExpr, ctx: &StatsRewriteCtx<'_>) -> Option { stat_expr(expr, Stat::NullCount, ctx) } -fn all_null(expr: &Expression) -> Expression { +fn all_null(expr: &BoundExpr) -> BoundExpr { stat_fn(expr.clone(), AllNull.bind(AggregateEmptyOptions)) } -fn all_non_null(expr: &Expression) -> Expression { +fn all_non_null(expr: &BoundExpr) -> BoundExpr { stat_fn(expr.clone(), AllNonNull.bind(AggregateEmptyOptions)) } enum NanCheck { NotNeeded, - Check(Expression), + Check(BoundExpr), Unavailable, } trait NonNanProof { const EMIT_UNGUARDED_REWRITES: bool; - fn check(ctx: &StatsRewriteCtx<'_>, expr: &Expression) -> VortexResult; + fn check(ctx: &StatsRewriteCtx<'_>, expr: &BoundExpr) -> VortexResult; } struct NanCountProof; @@ -540,7 +591,7 @@ struct NanCountProof; impl NonNanProof for NanCountProof { const EMIT_UNGUARDED_REWRITES: bool = true; - fn check(ctx: &StatsRewriteCtx<'_>, expr: &Expression) -> VortexResult { + fn check(ctx: &StatsRewriteCtx<'_>, expr: &BoundExpr) -> VortexResult { non_nan_check(ctx, expr, |expr| { match stat_expr(expr, Stat::NaNCount, ctx) { Some(nan_count) => NanCheck::Check(eq(nan_count, lit(0u64))), @@ -555,7 +606,7 @@ struct AllNonNanProof; impl NonNanProof for AllNonNanProof { const EMIT_UNGUARDED_REWRITES: bool = false; - fn check(ctx: &StatsRewriteCtx<'_>, expr: &Expression) -> VortexResult { + fn check(ctx: &StatsRewriteCtx<'_>, expr: &BoundExpr) -> VortexResult { non_nan_check(ctx, expr, |expr| { NanCheck::Check(stat_fn(expr.clone(), AllNonNan.bind(AggregateEmptyOptions))) }) @@ -567,8 +618,8 @@ impl NonNanProof for AllNonNanProof { // from float to non-float still needs a proof about the float source values. fn non_nan_check( ctx: &StatsRewriteCtx<'_>, - expr: &Expression, - proof: impl FnOnce(&Expression) -> NanCheck, + expr: &BoundExpr, + proof: impl FnOnce(&BoundExpr) -> NanCheck, ) -> VortexResult { if let Some(scalar) = expr.as_opt::() { if !scalar.dtype().is_float() { @@ -600,7 +651,7 @@ fn has_nans(dtype: &DType) -> bool { dtype.is_float() } -fn stat_expr(expr: &Expression, stat: Stat, ctx: &StatsRewriteCtx<'_>) -> Option { +fn stat_expr(expr: &BoundExpr, stat: Stat, ctx: &StatsRewriteCtx<'_>) -> Option { if let Some(literal) = literal_stat(expr, stat) { return Some(literal); } @@ -629,9 +680,9 @@ fn stat_expr(expr: &Expression, stat: Stat, ctx: &StatsRewriteCtx<'_>) -> Option fn with_non_nan_guards<'a, P: NonNanProof>( ctx: &StatsRewriteCtx<'_>, - exprs: impl IntoIterator, - value_predicate: Expression, -) -> VortexResult> { + exprs: impl IntoIterator, + value_predicate: BoundExpr, +) -> VortexResult> { let mut nan_checks = Vec::new(); for expr in exprs { match P::check(ctx, expr)? { @@ -652,7 +703,7 @@ fn with_non_nan_guards<'a, P: NonNanProof>( }) } -fn literal_stat(expr: &Expression, stat: Stat) -> Option { +fn literal_stat(expr: &BoundExpr, stat: Stat) -> Option { let scalar = expr.as_opt::()?; match stat { Stat::Min | Stat::Max => Some(lit(scalar.clone())), @@ -674,11 +725,11 @@ fn literal_stat(expr: &Expression, stat: Stat) -> Option { } fn cast_stat( - expr: &Expression, + expr: &BoundExpr, dtype: &DType, stat: Stat, ctx: &StatsRewriteCtx<'_>, -) -> Option { +) -> Option { match stat { Stat::Min | Stat::Max => stat_expr(expr, stat, ctx).map(|stat| cast(stat, dtype.clone())), Stat::NaNCount | Stat::Sum | Stat::UncompressedSizeInBytes => stat_expr(expr, stat, ctx), @@ -686,8 +737,10 @@ fn cast_stat( } } -fn stat_fn(expr: Expression, aggregate_fn: AggregateFnRef) -> Expression { - StatFn.new_expr(StatOptions::new(aggregate_fn), [expr]) +fn stat_fn(expr: BoundExpr, aggregate_fn: AggregateFnRef) -> BoundExpr { + StatFn + .try_new_bound_expr(StatOptions::new(aggregate_fn), [expr]) + .vortex_expect("stats rewrites only construct supported aggregate expressions") } #[cfg(test)] @@ -700,8 +753,6 @@ mod tests { use super::StatFn; use super::StatOptions; - use super::all_non_null; - use super::all_null; use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTableExt; use crate::aggregate_fn::EmptyOptions as AggregateEmptyOptions; @@ -710,7 +761,8 @@ mod tests { use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; - use crate::expr::Expression; + use crate::expr::BoundExpression; + use crate::expr::Expression as BoundExpr; use crate::expr::and; use crate::expr::between; use crate::expr::cast; @@ -740,12 +792,12 @@ mod tests { static SESSION: LazyLock = LazyLock::new(crate::array_session); - fn stat(expr: Expression, stat: Stat) -> Expression { + fn stat(expr: BoundExpr, stat: Stat) -> BoundExpr { let aggregate_fn = stat.aggregate_fn().expect("stat should have aggregate fn"); stat_fn(expr, aggregate_fn) } - fn stat_fn(expr: Expression, aggregate_fn: AggregateFnRef) -> Expression { + fn stat_fn(expr: BoundExpr, aggregate_fn: AggregateFnRef) -> BoundExpr { StatFn.new_expr(StatOptions::new(aggregate_fn), [expr]) } @@ -770,15 +822,33 @@ mod tests { ) } - fn falsify(expr: &Expression) -> VortexResult> { - expr.falsify(&test_scope(), &SESSION) + fn falsify(expr: &BoundExpr) -> VortexResult> { + expr.bind(&test_scope())?.falsify(&SESSION) } - fn satisfy(expr: &Expression) -> VortexResult> { - expr.satisfy(&test_scope(), &SESSION) + fn satisfy(expr: &BoundExpr) -> VortexResult> { + expr.bind(&test_scope())?.satisfy(&SESSION) + } + + fn bind_expected(expr: Option) -> VortexResult> { + expr.map(|expr| expr.bind(&test_scope())).transpose() + } + + fn all_null(expr: &BoundExpr) -> BoundExpr { + crate::stats::all_null(expr.clone()) + } + + fn all_non_null(expr: &BoundExpr) -> BoundExpr { + crate::stats::all_non_null(expr.clone()) + } + + macro_rules! assert_rewrite_eq { + ($actual:expr, $expected:expr) => { + assert_eq!($actual, bind_expected($expected)?) + }; } - fn nan_guarded(expr: Expression, value_predicate: Expression) -> Expression { + fn nan_guarded(expr: BoundExpr, value_predicate: BoundExpr) -> BoundExpr { or( and( eq(stat(expr.clone(), Stat::NaNCount), lit(0u64)), @@ -794,13 +864,13 @@ mod tests { #[test] fn rewrites_comparison_falsifier() -> VortexResult<()> { let expr = gt(col("a"), lit(10)); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(lt_eq(stat(col("a"), Stat::Max), lit(10))) ); let expr = eq(col("a"), col("b")); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( gt(stat(col("a"), Stat::Min), stat(col("b"), Stat::Max)), @@ -809,7 +879,7 @@ mod tests { ); let expr = eq(col("s"), col("t")); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( gt(stat(col("s"), Stat::Min), stat(col("t"), Stat::Max)), @@ -822,7 +892,7 @@ mod tests { #[test] fn rewrites_boolean_falsifiers() -> VortexResult<()> { let expr = and(gt(col("a"), lit(10)), lt(col("a"), lit(50))); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( lt_eq(stat(col("a"), Stat::Max), lit(10)), @@ -844,7 +914,7 @@ mod tests { }, ); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( gt(lit(10), stat(col("a"), Stat::Max)), @@ -856,7 +926,7 @@ mod tests { #[test] fn rewrites_null_falsifiers() -> VortexResult<()> { - assert_eq!( + assert_rewrite_eq!( falsify(&is_null(col("a")))?, Some(or( eq(stat(col("a"), Stat::NullCount), lit(0u64)), @@ -864,7 +934,7 @@ mod tests { )) ); - assert_eq!( + assert_rewrite_eq!( falsify(&is_not_null(col("a")))?, Some(or( eq( @@ -879,7 +949,7 @@ mod tests { #[test] fn rewrites_null_satisfiers() -> VortexResult<()> { - assert_eq!( + assert_rewrite_eq!( satisfy(&is_null(col("a")))?, Some(or( eq( @@ -890,7 +960,7 @@ mod tests { )) ); - assert_eq!( + assert_rewrite_eq!( satisfy(&is_not_null(col("a")))?, Some(or( eq(stat(col("a"), Stat::NullCount), lit(0u64)), @@ -909,7 +979,7 @@ mod tests { ); let expr = list_contains(lit(list), col("a")); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(and( and( @@ -934,7 +1004,7 @@ mod tests { #[test] fn rewrites_like_falsifier() -> VortexResult<()> { let expr = like(col("s"), lit("prefix%")); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( gt_eq(stat(col("s"), Stat::Min), lit("prefiy")), @@ -943,7 +1013,7 @@ mod tests { ); let expr = like(col("s"), lit(r"\%%")); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( gt_eq(stat(col("s"), Stat::Min), lit("&")), @@ -952,7 +1022,7 @@ mod tests { ); let expr = like(col("s"), lit("pref%ix%")); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( gt_eq(stat(col("s"), Stat::Min), lit("preg")), @@ -961,7 +1031,7 @@ mod tests { ); let expr = like(col("s"), lit("pref_ix_")); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( gt_eq(stat(col("s"), Stat::Min), lit("preg")), @@ -970,7 +1040,7 @@ mod tests { ); let expr = like(col("s"), lit("exact")); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( gt(stat(col("s"), Stat::Min), lit("exact")), @@ -979,7 +1049,7 @@ mod tests { ); let expr = like(col("s"), lit("%suffix")); - assert_eq!(falsify(&expr)?, None); + assert_rewrite_eq!(falsify(&expr)?, None); Ok(()) } @@ -994,7 +1064,7 @@ mod tests { ); let dynamic = expr.as_::(); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(DynamicComparison.new_expr( DynamicComparisonExpr { @@ -1013,7 +1083,7 @@ mod tests { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); let expr = gt(cast(col("f"), dtype.clone()), lit(5i32)); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(nan_guarded( col("f"), @@ -1031,8 +1101,8 @@ mod tests { nested_struct_dtype(), vec![Scalar::primitive(1.0f32, Nullability::Nullable)], ); - assert_eq!(falsify(<_eq(col("n"), lit(struct_scalar.clone())))?, None); - assert_eq!(falsify(&eq(col("n"), lit(struct_scalar)))?, None); + assert_rewrite_eq!(falsify(<_eq(col("n"), lit(struct_scalar.clone())))?, None); + assert_rewrite_eq!(falsify(&eq(col("n"), lit(struct_scalar)))?, None); Ok(()) } @@ -1041,7 +1111,7 @@ mod tests { let dtype = DType::Primitive(PType::I64, Nullability::NonNullable); let expr = eq(cast(col("a"), dtype.clone()), lit(42i64)); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( gt(cast(stat(col("a"), Stat::Min), dtype.clone()), lit(42i64)), diff --git a/vortex-file/src/file.rs b/vortex-file/src/file.rs index 739d7c9fd8b..a235011b9c0 100644 --- a/vortex-file/src/file.rs +++ b/vortex-file/src/file.rs @@ -236,8 +236,7 @@ impl VortexFile { }; can_prune_file_stats( - filter, - self.footer.dtype(), + &filter.bind(self.footer.dtype())?, self.footer.row_count(), stats, fields, diff --git a/vortex-file/src/pruning.rs b/vortex-file/src/pruning.rs index 559df97d2d0..b115562abe5 100644 --- a/vortex-file/src/pruning.rs +++ b/vortex-file/src/pruning.rs @@ -10,11 +10,10 @@ use vortex_array::arrays::NullArray; use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; use vortex_array::dtype::StructFields; -use vortex_array::expr::Expression; -use vortex_array::expr::is_root; -use vortex_array::expr::lit; +use vortex_array::expr::BoundExpression; use vortex_array::expr::stats::Stat; use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::ScalarFnVTableExt; 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; @@ -27,30 +26,27 @@ use vortex_session::VortexSession; use crate::FileStatistics; pub(crate) fn can_prune_file_stats( - expr: &Expression, - dtype: &DType, + expr: &BoundExpression, row_count: u64, file_stats: &FileStatistics, struct_fields: &StructFields, session: &VortexSession, ) -> VortexResult { - let Some(pruning_expr) = expr.falsify(dtype, session)? else { + let Some(pruning_expr) = expr.falsify(session)? else { return Ok(false); }; let binder = FileStatsBinder { - dtype, file_stats, struct_fields, }; let pruning_expr = bind_stats(pruning_expr, &binder)?; - let simplified = pruning_expr.optimize_recursive(&DType::Null)?; - if let Some(result) = simplified.as_opt::() { + if let Some(result) = pruning_expr.as_opt::() { return Ok(result.as_bool().value() == Some(true)); } - let pruning = NullArray::new(1).into_array().apply(&pruning_expr)?; + 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)?; @@ -65,22 +61,17 @@ pub(crate) fn can_prune_file_stats( } struct FileStatsBinder<'a> { - dtype: &'a DType, file_stats: &'a FileStatistics, struct_fields: &'a StructFields, } impl StatBinder for FileStatsBinder<'_> { - fn scope(&self) -> &DType { - self.dtype - } - fn bind_aggregate( &self, - input: &Expression, + input: &BoundExpression, aggregate_fn: &AggregateFnRef, _stat_dtype: &DType, - ) -> VortexResult> { + ) -> VortexResult> { let Some(stat) = Stat::from_aggregate_fn(aggregate_fn) else { return Ok(None); }; @@ -92,7 +83,7 @@ impl StatBinder for FileStatsBinder<'_> { } impl FileStatsBinder<'_> { - fn stat_ref(&self, field_path: &FieldPath, stat: Stat) -> Option { + fn stat_ref(&self, field_path: &FieldPath, stat: Stat) -> Option { // FileStats currently only holds top-level field statistics. if field_path.parts().len() != 1 { return None; @@ -107,12 +98,12 @@ impl FileStatsBinder<'_> { let stat_dtype = stat.dtype(&field_dtype)?; let stat_scalar = Scalar::try_new(stat_dtype, Some(stat_value)).ok()?; - Some(lit(stat_scalar)) + Literal.try_new_bound_expr(stat_scalar, []).ok() } } -fn direct_field_path(expr: &Expression) -> Option { - if is_root(expr) { +fn direct_field_path(expr: &BoundExpression) -> Option { + if expr.is_root() { return Some(FieldPath::root()); } diff --git a/vortex-file/src/v2/file_stats_reader.rs b/vortex-file/src/v2/file_stats_reader.rs index 20102a2ad66..03ad2ab88d4 100644 --- a/vortex-file/src/v2/file_stats_reader.rs +++ b/vortex-file/src/v2/file_stats_reader.rs @@ -16,7 +16,6 @@ use vortex_array::dtype::FieldMask; use vortex_array::dtype::StructFields; use vortex_array::expr::BoundExpression; use vortex_array::expr::ExactBoundExpr; -use vortex_array::expr::Expression; use vortex_error::VortexResult; use vortex_layout::ArrayFuture; use vortex_layout::LayoutReader; @@ -73,10 +72,9 @@ impl FileStatsLayoutReader { /// /// Row-count placeholders are resolved against the full file row count, /// independent of the requested row range. - fn evaluate_file_stats(&self, expr: &Expression) -> VortexResult { + fn evaluate_file_stats(&self, expr: &BoundExpression) -> VortexResult { can_prune_file_stats( expr, - self.child.dtype(), self.child.row_count(), &self.file_stats, &self.struct_fields, @@ -129,8 +127,7 @@ impl LayoutReader for FileStatsLayoutReader { } // Evaluate and cache. - let expression = expr.unbind(); - let pruned = self.evaluate_file_stats(&expression)?; + let pruned = self.evaluate_file_stats(expr)?; self.prune_cache.insert(key, pruned); if pruned { diff --git a/vortex-geo/src/prune/distance.rs b/vortex-geo/src/prune/distance.rs index 37de001f026..ce771ff3c0b 100644 --- a/vortex-geo/src/prune/distance.rs +++ b/vortex-geo/src/prune/distance.rs @@ -4,12 +4,7 @@ //! `ST_Distance(geom, const) radius` pruning. use geo::Rect as GeoRect; -use vortex_array::expr::Expression; -use vortex_array::expr::gt; -use vortex_array::expr::gt_eq; -use vortex_array::expr::lit; -use vortex_array::expr::lt; -use vortex_array::expr::lt_eq; +use vortex_array::expr::BoundExpression as BoundExpr; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::fns::binary::Binary; @@ -21,6 +16,11 @@ use vortex_error::VortexResult; use super::aabb_stat; use super::geometry_and_constant; +use super::gt; +use super::gt_eq; +use super::lit; +use super::lt; +use super::lt_eq; use super::max_dist_sq; use super::min_dist_sq; use super::query_aabb; @@ -42,9 +42,9 @@ impl StatsRewriteRule for GeoDistancePrune { fn falsify( &self, - expr: &Expression, + expr: &BoundExpr, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { // Only the ordered comparisons prune today. `== r` could prune in the future (a chunk is // provably empty when `r` lies outside its box's [min, max] distance interval), it's just // not implemented. `!= r` cannot: pruning would need every row's distance to equal `r`, @@ -94,11 +94,11 @@ impl StatsRewriteRule for GeoDistancePrune { /// /// A distance is always `>= 0`, which decides the degenerate radii up front. fn distance_prune_proof( - geom: &Expression, + geom: &BoundExpr, query: GeoRect, op: Operator, radius: f64, -) -> Option { +) -> Option { // A distance is always non-negative, so degenerate radii resolve without touching the box. match op { // `<= r` / `< r` with a negative radius (or zero, for `<`) match nothing: prune every chunk. @@ -130,7 +130,7 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::Expression; + use vortex_array::expr::BoundExpression; use vortex_array::expr::gt_eq; use vortex_array::expr::lit; use vortex_array::expr::lt_eq; @@ -158,7 +158,7 @@ mod tests { operator: Operator, geom_first: bool, radius: impl Into, - ) -> VortexResult> { + ) -> VortexResult> { let session = geo_session(); let mut ctx = session.create_execution_ctx(); @@ -170,9 +170,11 @@ mod tests { [lit(origin), root()] }; let distance = GeoDistance.new_expr(EmptyOptions, operands); - let predicate = Binary.new_expr(operator, [distance, lit(radius.into())]); + let predicate = Binary + .new_expr(operator, [distance, lit(radius.into())]) + .bind(&scope)?; - GeoDistancePrune.falsify(&predicate, &StatsRewriteCtx::new(&session, &scope)) + GeoDistancePrune.falsify(&predicate, &StatsRewriteCtx::new(&session)) } /// A null geometry literal (`ST_Distance(geom, NULL) <= r`) declines cleanly instead of @@ -184,9 +186,11 @@ mod tests { let scope = point_column(vec![0.0], vec![0.0])?.dtype().clone(); let null_query = Scalar::null(scope.as_nullable()); let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(null_query)]); - let predicate = Binary.new_expr(Operator::Lte, [distance, lit(0.5f64)]); + let predicate = Binary + .new_expr(Operator::Lte, [distance, lit(0.5f64)]) + .bind(&scope)?; - let ctx = StatsRewriteCtx::new(&session, &scope); + let ctx = StatsRewriteCtx::new(&session); assert!(GeoDistancePrune.falsify(&predicate, &ctx)?.is_none()); Ok(()) } @@ -231,11 +235,10 @@ mod tests { Ok(()) } - /// Filter expressions arrive uncoerced, so `distance <= 10` may carry an integer literal - - /// it casts and prunes like an f64 radius. + /// An uncoerced integer radius is rejected while binding the comparison. #[test] - fn integer_radius_prunes() -> VortexResult<()> { - assert!(falsify_distance(Operator::Lte, true, 10i64)?.is_some()); + fn uncoerced_integer_radius_fails_to_bind() -> VortexResult<()> { + assert!(falsify_distance(Operator::Lte, true, 10i64).is_err()); Ok(()) } @@ -259,8 +262,7 @@ mod tests { Ok(()) } - /// A scope dtype without `GeometryAabb` support gets no proof - the stat reference would - /// fail to bind at prune time. + /// A non-geometry scope is rejected while binding, before stats rewriting. #[test] fn unsupported_scope_is_not_pruned() -> VortexResult<()> { let session = geo_session(); @@ -270,9 +272,7 @@ mod tests { let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); let predicate = lt_eq(distance, lit(0.5f64)); - - let ctx = StatsRewriteCtx::new(&session, &scope); - assert!(GeoDistancePrune.falsify(&predicate, &ctx)?.is_none()); + assert!(predicate.bind(&scope).is_err()); Ok(()) } @@ -282,8 +282,8 @@ mod tests { let session = geo_session(); let scope = point_column(vec![0.0], vec![0.0])?.dtype().clone(); - let predicate = lt_eq(lit(1.0f64), lit(2.0f64)); - let ctx = StatsRewriteCtx::new(&session, &scope); + let predicate = lt_eq(lit(1.0f64), lit(2.0f64)).bind(&scope)?; + let ctx = StatsRewriteCtx::new(&session); assert!(GeoDistancePrune.falsify(&predicate, &ctx)?.is_none()); Ok(()) } @@ -305,7 +305,8 @@ mod tests { let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); let predicate = lt_eq(distance, lit(0.5f64)); let proof = predicate - .falsify(&point_dtype, &session)? + .bind(&point_dtype)? + .falsify(&session)? .expect("distance filter should be falsifiable"); // `true` means the zone is pruned: chunk 0 (near origin) is kept, chunk 1 (far) is skipped. @@ -330,7 +331,8 @@ mod tests { let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); let predicate = lt_eq(distance, lit(1.0f64)); let proof = predicate - .falsify(&point_dtype, &session)? + .bind(&point_dtype)? + .falsify(&session)? .expect("distance filter should be falsifiable"); assert_eq!( @@ -361,7 +363,8 @@ mod tests { let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); let proof = gt_eq(distance, lit(2.0f64)) - .falsify(&point_dtype, &session)? + .bind(&point_dtype)? + .falsify(&session)? .expect("distance filter should be falsifiable"); // Chunk 0 (within 2) is pruned for `>= 2`; chunk 1 (beyond 2) is kept. @@ -383,7 +386,8 @@ mod tests { let origin = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; let distance = GeoDistance.new_expr(EmptyOptions, [root(), lit(origin)]); let proof = lt_eq(distance, lit(0.5f64)) - .falsify(&point_dtype, &session)? + .bind(&point_dtype)? + .falsify(&session)? .expect("distance filter should be falsifiable"); let mask = zone_map.prune(&proof, &session)?; diff --git a/vortex-geo/src/prune/intersects.rs b/vortex-geo/src/prune/intersects.rs index fa5b6489773..c4883cd34ec 100644 --- a/vortex-geo/src/prune/intersects.rs +++ b/vortex-geo/src/prune/intersects.rs @@ -3,9 +3,7 @@ //! `ST_Intersects(geom, const)` pruning. -use vortex_array::expr::Expression; -use vortex_array::expr::gt; -use vortex_array::expr::lit; +use vortex_array::expr::BoundExpression as BoundExpr; use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::stats::rewrite::StatsRewriteCtx; @@ -14,6 +12,8 @@ use vortex_error::VortexResult; use super::aabb_stat; use super::geometry_and_constant; +use super::gt; +use super::lit; use super::min_dist_sq; use super::query_aabb; use crate::scalar_fn::intersects::GeoIntersects; @@ -34,9 +34,9 @@ impl StatsRewriteRule for GeoIntersectsPrune { fn falsify( &self, - expr: &Expression, + expr: &BoundExpr, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { let Some((geom, constant)) = geometry_and_constant(expr, ctx)? else { return Ok(None); }; @@ -56,7 +56,7 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; - use vortex_array::expr::Expression; + use vortex_array::expr::BoundExpression; use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::scalar::Scalar; @@ -75,7 +75,7 @@ mod tests { /// Run the intersects rule against `GeoIntersects(root, point(1.0, 0.5))`, operands swapped /// when `geom_first` is false. - fn falsify_intersects(geom_first: bool) -> VortexResult> { + fn falsify_intersects(geom_first: bool) -> VortexResult> { let session = geo_session(); let mut ctx = session.create_execution_ctx(); @@ -86,8 +86,10 @@ mod tests { } else { [lit(query), root()] }; - let predicate = GeoIntersects.new_expr(EmptyOptions, operands); - GeoIntersectsPrune.falsify(&predicate, &StatsRewriteCtx::new(&session, &scope)) + let predicate = GeoIntersects + .new_expr(EmptyOptions, operands) + .bind(&scope)?; + GeoIntersectsPrune.falsify(&predicate, &StatsRewriteCtx::new(&session)) } /// Intersects is symmetric: both operand orders produce a proof. @@ -99,8 +101,7 @@ mod tests { Ok(()) } - /// A scope dtype without `GeometryAabb` support gets no proof, the stat reference would - /// fail to bind at prune time. + /// A non-geometry scope is rejected while binding, before stats rewriting. #[test] fn unsupported_scope_is_not_pruned() -> VortexResult<()> { let session = geo_session(); @@ -109,9 +110,7 @@ mod tests { let scope = DType::Primitive(PType::F64, Nullability::NonNullable); let query = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; let predicate = GeoIntersects.new_expr(EmptyOptions, [root(), lit(query)]); - - let ctx = StatsRewriteCtx::new(&session, &scope); - assert!(GeoIntersectsPrune.falsify(&predicate, &ctx)?.is_none()); + assert!(predicate.bind(&scope).is_err()); Ok(()) } @@ -123,9 +122,11 @@ mod tests { let scope = point_column(vec![0.0], vec![0.0])?.dtype().clone(); let null_query = Scalar::null(scope.as_nullable()); - let predicate = GeoIntersects.new_expr(EmptyOptions, [root(), lit(null_query)]); + let predicate = GeoIntersects + .new_expr(EmptyOptions, [root(), lit(null_query)]) + .bind(&scope)?; - let ctx = StatsRewriteCtx::new(&session, &scope); + let ctx = StatsRewriteCtx::new(&session); assert!(GeoIntersectsPrune.falsify(&predicate, &ctx)?.is_none()); Ok(()) } @@ -152,7 +153,8 @@ mod tests { let query = point_column(vec![1.0], vec![0.5])?.execute_scalar(0, &mut ctx)?; let predicate = GeoIntersects.new_expr(EmptyOptions, [root(), lit(query)]); let proof = predicate - .falsify(&point_dtype, &session)? + .bind(&point_dtype)? + .falsify(&session)? .expect("intersects filter should be falsifiable"); let mask = zone_map.prune(&proof, &session)?; @@ -172,7 +174,8 @@ mod tests { let query = point_column(vec![0.0], vec![0.0])?.execute_scalar(0, &mut ctx)?; let proof = GeoIntersects .new_expr(EmptyOptions, [root(), lit(query)]) - .falsify(&point_dtype, &session)? + .bind(&point_dtype)? + .falsify(&session)? .expect("intersects filter should be falsifiable"); let mask = zone_map.prune(&proof, &session)?; diff --git a/vortex-geo/src/prune/mod.rs b/vortex-geo/src/prune/mod.rs index 22d54ac0248..83dd892e593 100644 --- a/vortex-geo/src/prune/mod.rs +++ b/vortex-geo/src/prune/mod.rs @@ -22,22 +22,21 @@ pub use intersects::GeoIntersectsPrune; use vortex_array::VortexSessionExecute; use vortex_array::aggregate_fn::AggregateFnVTableExt; use vortex_array::aggregate_fn::EmptyOptions; -use vortex_array::expr::Expression; -use vortex_array::expr::case_when; -use vortex_array::expr::checked_add; -use vortex_array::expr::ext_storage; -use vortex_array::expr::get_item; -use vortex_array::expr::gt; -use vortex_array::expr::is_root; -use vortex_array::expr::lit; -use vortex_array::expr::lt; +use vortex_array::dtype::FieldName; +use vortex_array::expr::BoundExpression as BoundExpr; use vortex_array::scalar::Scalar; +use vortex_array::scalar_fn::EmptyOptions as ScalarEmptyOptions; use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::fns::binary::Binary; +use vortex_array::scalar_fn::fns::case_when::CaseWhen; +use vortex_array::scalar_fn::fns::case_when::CaseWhenOptions; +use vortex_array::scalar_fn::fns::ext_storage::ExtStorage; +use vortex_array::scalar_fn::fns::get_item::GetItem; use vortex_array::scalar_fn::fns::literal::Literal; use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::scalar_fn::fns::stat::StatFn; +use vortex_array::scalar_fn::fns::stat::StatOptions; use vortex_array::stats::rewrite::StatsRewriteCtx; -use vortex_array::stats::stat; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -53,15 +52,15 @@ use crate::extension::single_geometry; /// An asymmetric predicate (e.g. a future contains) must recover which operand is the column /// itself instead of calling this. fn geometry_and_constant<'a>( - expr: &'a Expression, + expr: &'a BoundExpr, ctx: &StatsRewriteCtx<'_>, -) -> VortexResult> { +) -> VortexResult> { // The predicate is symmetric, so the column (scope root) and the constant may be on either // side. let (lhs, rhs) = (expr.child(0), expr.child(1)); - let (geom, constant) = if is_root(lhs) { + let (geom, constant) = if lhs.is_root() { (lhs, rhs) - } else if is_root(rhs) { + } else if rhs.is_root() { (rhs, lhs) } else { return Ok(None); @@ -96,22 +95,78 @@ fn query_aabb(constant: &Scalar, ctx: &StatsRewriteCtx<'_>) -> VortexResult Expression { +fn aabb_stat(geom: &BoundExpr) -> BoundExpr { // `ext_storage` unwraps the native `geoarrow.box` stat value to its backing struct, so // proofs can `get_item` the coordinate fields. ext_storage(stat(geom.clone(), GeometryAabb.bind(EmptyOptions))) } +fn stat(expr: BoundExpr, aggregate_fn: vortex_array::aggregate_fn::AggregateFnRef) -> BoundExpr { + StatFn + .try_new_bound_expr(StatOptions::new(aggregate_fn), [expr]) + .vortex_expect("geometry pruning only requests a supported aggregate") +} + +fn ext_storage(input: BoundExpr) -> BoundExpr { + ExtStorage + .try_new_bound_expr(ScalarEmptyOptions, [input]) + .vortex_expect("the geometry AABB aggregate returns an extension value") +} + +fn get_item(field: impl Into, child: BoundExpr) -> BoundExpr { + GetItem + .try_new_bound_expr(field.into(), [child]) + .vortex_expect("geometry AABB fields are fixed by its storage dtype") +} + +fn lit(value: impl Into) -> BoundExpr { + Literal + .try_new_bound_expr(value.into(), []) + .vortex_expect("literal expressions are always well-typed") +} + +fn case_when(condition: BoundExpr, then_value: BoundExpr, else_value: BoundExpr) -> BoundExpr { + CaseWhen + .try_new_bound_expr( + CaseWhenOptions { + num_when_then_pairs: 1, + has_else: true, + }, + [condition, then_value, else_value], + ) + .vortex_expect("geometry pruning case expressions have matching branch dtypes") +} + +fn gt(lhs: BoundExpr, rhs: BoundExpr) -> BoundExpr { + binop(Operator::Gt, lhs, rhs) +} + +fn gt_eq(lhs: BoundExpr, rhs: BoundExpr) -> BoundExpr { + binop(Operator::Gte, lhs, rhs) +} + +fn lt(lhs: BoundExpr, rhs: BoundExpr) -> BoundExpr { + binop(Operator::Lt, lhs, rhs) +} + +fn lt_eq(lhs: BoundExpr, rhs: BoundExpr) -> BoundExpr { + binop(Operator::Lte, lhs, rhs) +} + +fn checked_add(lhs: BoundExpr, rhs: BoundExpr) -> BoundExpr { + binop(Operator::Add, lhs, rhs) +} + /// Lower bound on every row's squared distance to the query AABB: zero when the boxes overlap or /// touch, positive iff they are strictly separated. /// /// Prunes "near" predicates: `min_dist_sq > r^2` proves every row is farther than `r`. -fn min_dist_sq(aabb: &Expression, query: GeoRect) -> Expression { +fn min_dist_sq(aabb: &BoundExpr, query: GeoRect) -> BoundExpr { let field = |name: &str| get_item(name, aabb.clone()); // Per axis: gap = max(0, q_lo - aabb_hi, aabb_lo - q_hi), positive only when the intervals // are separated. The nearest two points of the boxes are one axis-gap apart per axis, so the // squared distance is gap_x^2 + gap_y^2 (squared throughout to avoid a sqrt). - let gap = |q_lo: f64, q_hi: f64, lo: Expression, hi: Expression| { + let gap = |q_lo: f64, q_hi: f64, lo: BoundExpr, hi: BoundExpr| { maximum( lit(0.0), maximum( @@ -128,12 +183,12 @@ fn min_dist_sq(aabb: &Expression, query: GeoRect) -> Expression { /// Upper bound on every row's squared distance to the query AABB. /// /// Prunes "far" predicates: `max_dist_sq < r^2` proves every row is within `r`. -fn max_dist_sq(aabb: &Expression, query: GeoRect) -> Expression { +fn max_dist_sq(aabb: &BoundExpr, query: GeoRect) -> BoundExpr { let field = |name: &str| get_item(name, aabb.clone()); // Per axis: span = max(q_hi, aabb_hi) - min(q_lo, aabb_lo), the farthest two points of the // boxes can be apart. The nullable AABB field is the second `maximum`/`minimum` argument so // that `case_when`'s else branch carries the nullability - a missing stat propagates null. - let span = |q_lo: f64, q_hi: f64, lo: Expression, hi: Expression| { + let span = |q_lo: f64, q_hi: f64, lo: BoundExpr, hi: BoundExpr| { binop( Operator::Sub, maximum(lit(q_hi), hi), @@ -146,23 +201,23 @@ fn max_dist_sq(aabb: &Expression, query: GeoRect) -> Expression { } /// `a b`. -fn binop(op: Operator, a: Expression, b: Expression) -> Expression { +fn binop(op: Operator, a: BoundExpr, b: BoundExpr) -> BoundExpr { Binary - .try_new_expr(op, [a, b]) + .try_new_bound_expr(op, [a, b]) .vortex_expect("binary expression") } /// `e * e`. -fn square(e: Expression) -> Expression { +fn square(e: BoundExpr) -> BoundExpr { binop(Operator::Mul, e.clone(), e) } /// `max(a, b)`. -fn maximum(a: Expression, b: Expression) -> Expression { +fn maximum(a: BoundExpr, b: BoundExpr) -> BoundExpr { case_when(gt(a.clone(), b.clone()), a, b) } /// `min(a, b)`. -fn minimum(a: Expression, b: Expression) -> Expression { +fn minimum(a: BoundExpr, b: BoundExpr) -> BoundExpr { case_when(lt(a.clone(), b.clone()), a, b) } diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index eefdee6e4c4..5718d7bc985 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -376,6 +376,8 @@ 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::byte_length; use vortex_array::expr::cast; @@ -386,6 +388,9 @@ mod tests { use vortex_array::expr::lit; use vortex_array::expr::pack; use vortex_array::expr::root; + use vortex_array::scalar_fn::ScalarFnVTableExt; + use vortex_array::scalar_fn::fns::pack::Pack; + use vortex_array::scalar_fn::fns::pack::PackOptions; use vortex_array::validity::Validity; use vortex_btrblocks::BtrBlocksCompressor; use vortex_error::VortexExpect; @@ -746,7 +751,11 @@ mod tests { get_item(format!("_{idx}"), get_item("", root())) } - fn test_apply(original: Expression, outer: Expression, inner: Expression) -> VortexResult<()> { + fn test_apply( + original: Expression, + outer: BoundExpression, + inner: BoundExpression, + ) -> VortexResult<()> { let mut ctx = array_session().create_execution_ctx(); let array = VarBinArray::from_iter( [Some("abc"), Some("def"), None], @@ -754,29 +763,39 @@ mod tests { ) .into_array(); - let pushed = array.clone().apply(&pack( - [(PUSHDOWN_ANNOTATION, inner)], - Nullability::NonNullable, - ))?; - let actual = pushed.apply(&outer)?; + let pushed_expr = Pack.try_new_bound_expr( + PackOptions { + names: [FieldName::from(PUSHDOWN_ANNOTATION)].into(), + nullability: Nullability::NonNullable, + }, + [inner], + )?; + let pushed = array.clone().apply_bound(&pushed_expr)?; + let actual = pushed.apply_bound(&outer)?; let expected = array.apply(&original)?; assert_arrays_eq!(actual, expected, &mut ctx); Ok(()) } - fn split_unbound( + fn split_bound( expr: Expression, dtype: &DType, - ) -> VortexResult<(Expression, Option)> { + ) -> VortexResult<(BoundExpression, Option)> { let bound = expr.bind(dtype)?; - let (outer, inner) = split_expression_for_pushdown(&bound)?; - Ok((outer.unbind(), inner.map(|expr| expr.unbind()))) + split_expression_for_pushdown(&bound) + } + + fn pushed_scope(inner: &BoundExpression) -> DType { + DType::Struct( + StructFields::from_iter([(PUSHDOWN_ANNOTATION, inner.dtype().clone())]), + Nullability::NonNullable, + ) } #[test] fn split_expr_root() { - let (outer, inner) = split_unbound(root(), &DType::Null).unwrap(); - assert_eq!(outer, root()); + let (outer, inner) = split_bound(root(), &DType::Null).unwrap(); + assert_eq!(outer, root().bind(&DType::Null).unwrap()); assert_eq!(inner, None); } @@ -785,22 +804,27 @@ mod tests { // cast is fallible, thus not pushed let target = DType::Primitive(PType::I64, Nullability::Nullable); let expr = cast(byte_length(root()), target.clone()); - let (outer, inner) = split_unbound(expr.clone(), &DType::Utf8(false.into()))?; + let dtype = DType::Utf8(false.into()); + let (outer, inner) = split_bound(expr.clone(), &dtype)?; let inner = inner.unwrap(); // [0] = cast([1], dtype) // [1] = byte_length(root) - assert_eq!(outer, cast(pushed_ref(0), target)); - assert_eq!(inner, pushed_inner([byte_length(root())])); + assert_eq!( + outer, + cast(pushed_ref(0), target).bind(&pushed_scope(&inner))? + ); + assert_eq!(inner, pushed_inner([byte_length(root())]).bind(&dtype)?); test_apply(expr, outer, inner) } #[test] fn split_expr_full_pushdown() -> VortexResult<()> { let expr = byte_length(root()); - let (outer, inner) = split_unbound(expr.clone(), &DType::Utf8(false.into()))?; + let dtype = DType::Utf8(false.into()); + let (outer, inner) = split_bound(expr.clone(), &dtype)?; let inner = inner.unwrap(); - assert_eq!(outer, pushed_ref(0)); - assert_eq!(inner, pushed_inner([byte_length(root())])); + assert_eq!(outer, pushed_ref(0).bind(&pushed_scope(&inner))?); + assert_eq!(inner, pushed_inner([byte_length(root())]).bind(&dtype)?); test_apply(expr, outer, inner) } @@ -808,8 +832,9 @@ mod tests { fn split_expr_no_pushdown() { // like is fallible, thus not pushed. lit() does not reference root() let expr = like(root(), lit("abc")); - let (outer, inner) = split_unbound(expr.clone(), &DType::Utf8(true.into())).unwrap(); - assert_eq!(outer, expr); + let dtype = DType::Utf8(true.into()); + let (outer, inner) = split_bound(expr.clone(), &dtype).unwrap(); + assert_eq!(outer, expr.bind(&dtype).unwrap()); assert_eq!(inner, None); } } diff --git a/vortex-layout/src/layouts/row_idx/expr.rs b/vortex-layout/src/layouts/row_idx/expr.rs index c0043542f93..31dbbd4d47e 100644 --- a/vortex-layout/src/layouts/row_idx/expr.rs +++ b/vortex-layout/src/layouts/row_idx/expr.rs @@ -8,6 +8,7 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::expr::Expression; +use vortex_array::expr::display::ExprDisplay; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; @@ -41,7 +42,7 @@ impl ScalarFnVTable for RowIdx { fn fmt_sql( &self, _options: &Self::Options, - _expr: &Expression, + _expr: &dyn ExprDisplay, f: &mut Formatter<'_>, ) -> std::fmt::Result { write!(f, "#row_idx") diff --git a/vortex-layout/src/layouts/zoned/pruning.rs b/vortex-layout/src/layouts/zoned/pruning.rs index 51d59b19cd0..089f2e58e67 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/vortex-layout/src/layouts/zoned/pruning.rs @@ -19,7 +19,6 @@ use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::arrays::StructArray; use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; -use vortex_array::expr::Expression; use vortex_array::expr::root; use vortex_array::scalar_fn::fns::dynamic::DynamicExprUpdates; use vortex_error::SharedVortexResult; @@ -38,7 +37,7 @@ use crate::layouts::zoned::zone_map::ZoneMap; type SharedZoneMap = Shared>>; pub(super) type SharedPruningResult = Shared>>>; -type PredicateCache = Arc>>; +type PredicateCache = Arc>>; pub(super) struct PruningState { zone_count: usize, @@ -50,7 +49,7 @@ pub(super) struct PruningState { session: VortexSession, pruning_result: LazyLock>>, zone_map: OnceLock, - pruning_predicates: LazyLock>>, + pruning_predicates: LazyLock>>, } impl PruningState { @@ -86,7 +85,7 @@ impl PruningState { self.pruning_result .entry(expr.clone()) .or_insert_with(|| { - let expr = expr.unbind(); + let dynamic_updates = DynamicExprUpdates::new(&expr); match self.pruning_predicate(expr.clone()) { None => { trace!(%expr, "no pruning predicate"); @@ -95,7 +94,6 @@ impl PruningState { Some(predicate) => { trace!(%expr, ?predicate, "constructed pruning predicate"); let zone_map = self.zone_map(); - let dynamic_updates = DynamicExprUpdates::new(&expr); let session = self.session.clone(); Some( @@ -125,11 +123,11 @@ impl PruningState { .clone() } - fn pruning_predicate(&self, expr: Expression) -> Option { + fn pruning_predicate(&self, expr: BoundExpression) -> Option { self.pruning_predicates .entry(expr.clone()) .or_default() - .get_or_init(move || match expr.falsify(&self.dtype, &self.session) { + .get_or_init(move || match expr.falsify(&self.session) { Ok(predicate) => predicate, Err(error) => { trace!(%expr, %error, "failed to construct stats rewrite predicate"); @@ -188,7 +186,7 @@ impl PruningState { pub(super) struct PruningResult { zone_map: ZoneMap, - predicate: Expression, + predicate: BoundExpression, dynamic_updates: Option, latest_result: RwLock<(u64, Mask)>, session: VortexSession, diff --git a/vortex-layout/src/layouts/zoned/zone_map.rs b/vortex-layout/src/layouts/zoned/zone_map.rs index 9c7ff87b693..c84c0b443dd 100644 --- a/vortex-layout/src/layouts/zoned/zone_map.rs +++ b/vortex-layout/src/layouts/zoned/zone_map.rs @@ -16,24 +16,20 @@ 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::aggregate_fn::fns::nan_count::NanCount; use vortex_array::arrays::ConstantArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; 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::is_root; use vortex_array::expr::lit; use vortex_array::expr::root; use vortex_array::expr::stats::Stat; -use vortex_array::expr::traversal::NodeExt; -use vortex_array::expr::traversal::Transformed; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ScalarFnVTableExt; -use vortex_array::scalar_fn::fns::stat::StatFn; 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; @@ -43,6 +39,7 @@ use vortex_array::validity::Validity; use vortex_buffer::buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; use vortex_mask::Mask; use vortex_runend::RunEnd; use vortex_session::VortexSession; @@ -131,7 +128,7 @@ impl ZoneMap { /// Apply a pruning predicate to this zone map. /// /// `predicate` should be a stats rewrite expression such as the result of - /// [`Expression::falsify`]. The returned mask has one value per zone, where + /// [`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] @@ -141,13 +138,17 @@ impl ZoneMap { /// `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: &Expression, session: &VortexSession) -> VortexResult { + 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(&predicate)?; + let applied = array.apply_bound(&predicate)?; if !contains_row_count(&applied) { return applied.null_as_false().execute(&mut ctx); @@ -158,42 +159,9 @@ impl ZoneMap { substituted.null_as_false().execute(&mut ctx) } - fn lower_stats(&self, predicate: Expression) -> VortexResult { - let predicate = self.lower_non_float_nan_stats(predicate)?; + fn lower_stats(&self, predicate: BoundExpression) -> VortexResult { let binder = ZoneMapStatsBinder { zone_map: self }; - bind_stats(predicate, &binder)?.optimize_recursive(self.array.dtype()) - } - - fn lower_non_float_nan_stats(&self, predicate: Expression) -> VortexResult { - predicate - .transform_down(|expr| { - if !expr.is::() { - return Ok(Transformed::no(expr)); - } - - let options = expr.as_::(); - let aggregate_fn = options.aggregate_fn(); - let input_dtype = expr.child(0).return_dtype(&self.column_dtype)?; - - if has_nans(&input_dtype) { - return Ok(Transformed::no(expr)); - } - - if aggregate_fn.is::() { - return Ok(Transformed::yes(lit(0u64))); - } - - if aggregate_fn.is::() { - return Ok(Transformed::yes(lit(false))); - } - - if aggregate_fn.is::() { - return Ok(Transformed::yes(lit(true))); - } - - Ok(Transformed::no(expr)) - }) - .map(Transformed::into_inner) + bind_stats(predicate, &binder) } } @@ -202,60 +170,76 @@ struct ZoneMapStatsBinder<'a> { } impl StatBinder for ZoneMapStatsBinder<'_> { - fn scope(&self) -> &DType { - &self.zone_map.column_dtype - } - fn bind_aggregate( &self, - input: &Expression, + input: &BoundExpression, aggregate_fn: &AggregateFnRef, _stat_dtype: &DType, - ) -> VortexResult> { - if !is_root(input) { + ) -> VortexResult> { + if !input.is_root() { return Ok(None); } + vortex_ensure!( + input.dtype() == &self.zone_map.column_dtype, + "Stats predicate root dtype {} does not match zone-map column dtype {}", + input.dtype(), + self.zone_map.column_dtype + ); if let Some(stat_expr) = self.zone_map.aggregate_field_expr(aggregate_fn) { - return Ok(Some(stat_expr)); + return Ok(Some(self.bind_target(stat_expr)?)); } if aggregate_fn.is::() { - return Ok(self + return self .zone_map .stat_field_expr(Stat::NullCount) - .map(|null_count| eq(null_count, row_count_expr()))); + .map(|null_count| self.bind_target(eq(null_count, row_count_expr()))) + .transpose(); } if aggregate_fn.is::() { - return Ok(self + return self .zone_map .stat_field_expr(Stat::NullCount) - .map(|null_count| eq(null_count, lit(0u64)))); + .map(|null_count| self.bind_target(eq(null_count, lit(0u64)))) + .transpose(); } if aggregate_fn.is::() { - return Ok(self + return self .zone_map .stat_field_expr(Stat::NaNCount) - .map(|nan_count| eq(nan_count, row_count_expr()))); + .map(|nan_count| self.bind_target(eq(nan_count, row_count_expr()))) + .transpose(); } if aggregate_fn.is::() { - return Ok(self + return self .zone_map .stat_field_expr(Stat::NaNCount) - .map(|nan_count| eq(nan_count, lit(0u64)))); + .map(|nan_count| self.bind_target(eq(nan_count, lit(0u64)))) + .transpose(); } if let Some(stat) = Stat::from_aggregate_fn(aggregate_fn) { - return Ok(self.zone_map.stat_field_expr(stat)); + return self + .zone_map + .stat_field_expr(stat) + .map(|expr| self.bind_target(expr)) + .transpose(); } Ok(None) } } +impl ZoneMapStatsBinder<'_> { + fn bind_target(&self, expr: Expression) -> VortexResult { + expr.bind(self.zone_map.array.dtype()) + } +} + impl ZoneMap { fn aggregate_field_expr(&self, requested: &AggregateFnRef) -> Option { let field_name = requested.to_string(); @@ -318,10 +302,6 @@ fn row_count_expr() -> Expression { RowCount.new_expr(EmptyOptions, []) } -fn has_nans(dtype: &DType) -> bool { - matches!(dtype, DType::Primitive(ptype, _) if ptype.is_float()) -} - /// Build per-zone row counts for a zone map. /// /// `zone_len` is the nominal zone size; only the final zone may be shorter. The @@ -385,6 +365,7 @@ mod tests { use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; + use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::expr::cast; use vortex_array::expr::gt; @@ -402,12 +383,22 @@ mod tests { use vortex_array::stats::all_null; use vortex_array::validity::Validity; use vortex_buffer::buffer; + use vortex_error::VortexResult; + use vortex_mask::Mask; use crate::layouts::zoned::zone_map::ZoneMap; use crate::test::SESSION; - fn falsify(expr: &Expression, dtype: DType) -> Expression { - expr.falsify(&dtype, &SESSION).unwrap().unwrap() + fn falsify(expr: &Expression, dtype: DType) -> BoundExpression { + expr.bind(&dtype) + .unwrap() + .falsify(&SESSION) + .unwrap() + .unwrap() + } + + fn prune(zone_map: &ZoneMap, predicate: &Expression) -> VortexResult { + zone_map.prune(&predicate.bind(&zone_map.column_dtype)?, &SESSION) } fn default_bounded_stat_max_bytes() -> NonZeroUsize { @@ -631,7 +622,7 @@ mod tests { ) .unwrap(); - let mask = zone_map.prune(&all_null(root()), &SESSION).unwrap(); + let mask = prune(&zone_map, &all_null(root())).unwrap(); assert_arrays_eq!( mask.into_array(), BoolArray::from_iter([false, true, true]), @@ -654,7 +645,7 @@ mod tests { ) .unwrap(); - let mask = zone_map.prune(&all_non_null(root()), &SESSION).unwrap(); + let mask = prune(&zone_map, &all_non_null(root())).unwrap(); assert_arrays_eq!( mask.into_array(), BoolArray::from_iter([true, false, false]), @@ -679,14 +670,14 @@ mod tests { .unwrap(); let ctx = &mut SESSION.create_execution_ctx(); - let mask = zone_map.prune(&all_null(root()), &SESSION).unwrap(); + let mask = prune(&zone_map, &all_null(root())).unwrap(); assert_arrays_eq!( mask.into_array(), BoolArray::from_iter([true, false, true]), ctx ); - let mask = zone_map.prune(&all_non_null(root()), &SESSION).unwrap(); + let mask = prune(&zone_map, &all_non_null(root())).unwrap(); assert_arrays_eq!( mask.into_array(), BoolArray::from_iter([false, true, false]), @@ -711,14 +702,14 @@ mod tests { .unwrap(); let ctx = &mut SESSION.create_execution_ctx(); - let mask = zone_map.prune(&all_nan(root()), &SESSION).unwrap(); + let mask = prune(&zone_map, &all_nan(root())).unwrap(); assert_arrays_eq!( mask.into_array(), BoolArray::from_iter([true, false, true]), ctx ); - let mask = zone_map.prune(&all_non_nan(root()), &SESSION).unwrap(); + let mask = prune(&zone_map, &all_non_nan(root())).unwrap(); assert_arrays_eq!( mask.into_array(), BoolArray::from_iter([false, true, false]), @@ -727,22 +718,17 @@ mod tests { } #[test] - fn non_float_nan_stat_fns_lower_to_constants() { - let zone_map = ZoneMap::try_new( - PType::I32.into(), - StructArray::try_new(FieldNames::empty(), vec![], 2, Validity::NonNullable).unwrap(), - Arc::new([]), - 4, - 8, - ) - .unwrap(); - let ctx = &mut SESSION.create_execution_ctx(); - - let mask = zone_map.prune(&all_nan(root()), &SESSION).unwrap(); - assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([false, false]), ctx); - - let mask = zone_map.prune(&all_non_nan(root()), &SESSION).unwrap(); - assert_arrays_eq!(mask.into_array(), BoolArray::from_iter([true, true]), ctx); + fn non_float_nan_stat_fns_fail_to_bind() { + let dtype = DType::from(PType::I32); + for expr in [all_nan(root()), all_non_nan(root())] { + let error = expr.bind(&dtype).unwrap_err(); + assert!( + error + .to_string() + .contains("does not support input dtype i32"), + "{error}" + ); + } } #[test] @@ -757,7 +743,7 @@ mod tests { .unwrap(); let ctx = &mut SESSION.create_execution_ctx(); - let mask = zone_map.prune(&all_non_null(root()), &SESSION).unwrap(); + let mask = prune(&zone_map, &all_non_null(root())).unwrap(); assert_arrays_eq!( mask.into_array(), BoolArray::from_iter([false, false, false]), @@ -899,7 +885,7 @@ mod tests { let predicate = is_null(vortex_array::stats::stat(root(), max_fn)); // Missing StatFn lowers to a nullable null literal, so `is_null(...)` is true for every zone. - let mask = zone_map.prune(&predicate, &SESSION).unwrap(); + let mask = prune(&zone_map, &predicate).unwrap(); assert_arrays_eq!( mask.into_array(), BoolArray::from_iter([true, true, true]), @@ -922,7 +908,7 @@ mod tests { .aggregate_fn() .expect("max should have an aggregate function"); let predicate = is_null(vortex_array::stats::stat(root(), max_fn)); - let error = zone_map.prune(&predicate, &SESSION).unwrap_err(); + let error = prune(&zone_map, &predicate).unwrap_err(); assert!( error @@ -975,7 +961,7 @@ mod tests { ) .unwrap(); - let mask = zone_map.prune(&all_null(root()), &SESSION).unwrap(); + let mask = prune(&zone_map, &all_null(root())).unwrap(); assert_arrays_eq!( mask.into_array(), BoolArray::from_iter([false, true, true]), @@ -999,7 +985,7 @@ mod tests { ) .unwrap(); - let mask = zone_map.prune(&all_non_null(root()), &SESSION).unwrap(); + let mask = prune(&zone_map, &all_non_null(root())).unwrap(); assert_arrays_eq!( mask.into_array(), BoolArray::from_iter([true, false, false]), diff --git a/vortex-layout/src/scan/filter.rs b/vortex-layout/src/scan/filter.rs index 29b9165741f..fdba393a46f 100644 --- a/vortex-layout/src/scan/filter.rs +++ b/vortex-layout/src/scan/filter.rs @@ -58,10 +58,7 @@ impl FilterExpr { let conjuncts = bound_conjuncts(&expr); let num_conjuncts = conjuncts.len(); - let dynamic_conjuncts = conjuncts - .iter() - .map(|expr| DynamicExprUpdates::new(&expr.unbind())) - .collect_vec(); + let dynamic_conjuncts = conjuncts.iter().map(DynamicExprUpdates::new).collect_vec(); Self { conjuncts, From 688e8479b5c3c2e82397595304e8a35cca7efa45 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Wed, 5 Aug 2026 13:37:14 +0100 Subject: [PATCH 2/4] refactor: bind scan expressions once Signed-off-by: Joe Isaacs --- benchmarks/compress-bench/src/vortex.rs | 8 +- docs/developer-guide/internals/session.md | 3 +- fuzz/fuzz_targets/file_io.rs | 16 +- vortex-bench/src/datasets/tpch_l_comment.rs | 7 +- vortex-datafusion/src/persistent/opener.rs | 13 +- vortex-file/src/tests.rs | 153 ++++++++++++-------- vortex-layout/src/layouts/dict/reader.rs | 8 +- vortex-layout/src/layouts/zoned/pruning.rs | 15 +- vortex-layout/src/lib.rs | 6 +- vortex-layout/src/scan/layout.rs | 29 ++-- vortex-layout/src/scan/multi.rs | 77 ++++++++-- vortex-layout/src/scan/scan_builder.rs | 66 +++++---- vortex-python/src/dataset.rs | 15 +- vortex-python/src/file.rs | 16 +- vortex/src/lib.rs | 28 ++-- 15 files changed, 308 insertions(+), 152 deletions(-) diff --git a/benchmarks/compress-bench/src/vortex.rs b/benchmarks/compress-bench/src/vortex.rs index fae0cc44189..71b9ddfae3d 100644 --- a/benchmarks/compress-bench/src/vortex.rs +++ b/benchmarks/compress-bench/src/vortex.rs @@ -18,6 +18,7 @@ use vortex::expr::root; use vortex::expr::select; use vortex::file::OpenOptionsSessionExt; use vortex::file::WriteOptionsSessionExt; +use vortex::layout::scan::scan_builder::optimize_and_bind; use vortex_arrow::ToArrowType; use vortex_bench::Format; use vortex_bench::SESSION; @@ -64,14 +65,15 @@ impl Compressor for VortexCompressor { let start = Instant::now(); let data = Bytes::from(buf); let mut scan = SESSION.open_options().open_buffer(data)?.scan()?; - let root_columns = scan - .dtype()? + let source_dtype = scan.dtype()?; + let root_columns = source_dtype .as_struct_fields_opt() .map_or(0, |fields| fields.nfields()); if let Some(cols) = read_projection(root_columns) { // Columns are named "0".."num_columns-1"; project the given subset. let names: FieldNames = cols.iter().map(|i| i.to_string()).collect(); - scan = scan.with_projection(select(names, root())); + let projection = optimize_and_bind(select(names, root()), &source_dtype)?; + scan = scan.with_projection(projection); } let schema = Arc::new(scan.dtype()?.to_arrow_schema()?); diff --git a/docs/developer-guide/internals/session.md b/docs/developer-guide/internals/session.md index 0312d496cfe..45952f3a2bc 100644 --- a/docs/developer-guide/internals/session.md +++ b/docs/developer-guide/internals/session.md @@ -78,8 +78,9 @@ session.write_options() .await?; // Scanning a layout +let filter = optimize_and_bind(expr, layout_reader.dtype())?; ScanBuilder::new(session.clone(), layout_reader) - .with_filter(expr) + .with_filter(filter) .into_array_stream()?; ``` diff --git a/fuzz/fuzz_targets/file_io.rs b/fuzz/fuzz_targets/file_io.rs index f62693ccc49..0fc6bb8970c 100644 --- a/fuzz/fuzz_targets/file_io.rs +++ b/fuzz/fuzz_targets/file_io.rs @@ -6,6 +6,7 @@ use itertools::Itertools; use libfuzzer_sys::Corpus; use libfuzzer_sys::fuzz_target; +use vortex::layout::scan::scan_builder::optimize_and_bind; use vortex_array::Canonical; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; @@ -78,14 +79,21 @@ fuzz_target!(|fuzz: FuzzFileAction| -> Corpus { .write(&mut full_buff, array_data.to_array_iterator()) .vortex_expect("file write should succeed in fuzz test"); - let mut output = SESSION + let file = SESSION .open_options() .open_buffer(full_buff) - .vortex_expect("open_buffer should succeed in fuzz test") + .vortex_expect("open_buffer should succeed in fuzz test"); + let projection = optimize_and_bind(projection_expr.unwrap_or_else(root), file.dtype()) + .vortex_expect("projection should bind in fuzz test"); + let filter = filter_expr + .map(|filter| optimize_and_bind(filter, file.dtype())) + .transpose() + .vortex_expect("filter should bind in fuzz test"); + let mut output = file .scan() .vortex_expect("scan should succeed in fuzz test") - .with_projection(projection_expr.unwrap_or_else(root)) - .with_some_filter(filter_expr) + .with_projection(projection) + .with_some_filter(filter) .into_array_iter(&*RUNTIME) .vortex_expect("into_array_iter should succeed in fuzz test") .try_collect::<_, Vec<_>, _>() diff --git a/vortex-bench/src/datasets/tpch_l_comment.rs b/vortex-bench/src/datasets/tpch_l_comment.rs index 55d8497c2b3..064e7ad99f8 100644 --- a/vortex-bench/src/datasets/tpch_l_comment.rs +++ b/vortex-bench/src/datasets/tpch_l_comment.rs @@ -16,6 +16,7 @@ use vortex::dtype::Nullability::NonNullable; use vortex::expr::col; use vortex::expr::pack; use vortex::file::OpenOptionsSessionExt; +use vortex::layout::scan::scan_builder::optimize_and_bind; use crate::Format; use crate::IdempotentPath; @@ -66,9 +67,13 @@ impl Dataset for TPCHLCommentChunked { let path = data_dir.join("lineitem.vortex"); let file = SESSION.open_options().open_path(path).await?; + let projection = optimize_and_bind( + pack(vec![("l_comment", col("l_comment"))], NonNullable), + file.dtype(), + )?; let chunks: Vec<_> = file .scan()? - .with_projection(pack(vec![("l_comment", col("l_comment"))], NonNullable)) + .with_projection(projection) .map({ let ctx = ctx.clone(); move |a| { diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index d0ffb472ebe..d799d1bf01d 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -48,6 +48,7 @@ use vortex::file::OpenOptionsSessionExt; use vortex::io::InstrumentedReadAt; use vortex::layout::LayoutReader; use vortex::layout::scan::scan_builder::ScanBuilder; +use vortex::layout::scan::scan_builder::optimize_and_bind; use vortex::layout::scan::split_by::SplitBy; use vortex::metrics::Label; use vortex::metrics::MetricsRegistry; @@ -302,9 +303,11 @@ impl FileOpener for VortexOpener { // The schema of the stream returned from the vortex scan. // We use a reference schema for types that don't roundtrip (Dictionary, Utf8, etc.). - let scan_dtype = scan_projection.return_dtype(vxf.dtype()).map_err(|_e| { - exec_datafusion_err!("Couldn't get the dtype for the underlying Vortex scan") - })?; + let scan_projection = + optimize_and_bind(scan_projection, vxf.dtype()).map_err(|_e| { + exec_datafusion_err!("Couldn't get the dtype for the underlying Vortex scan") + })?; + let scan_dtype = scan_projection.dtype().clone(); // When projection pushdown is enabled, the scan outputs the projected columns. // When disabled, the scan outputs raw columns and the projection is applied after. @@ -419,6 +422,10 @@ impl FileOpener for VortexOpener { make_vortex_predicate(expr_convertor.as_ref(), &pushed).transpose() }) .transpose()?; + let filter = filter + .map(|filter| optimize_and_bind(filter, vxf.dtype())) + .transpose() + .map_err(|e| exec_datafusion_err!("Couldn't bind Vortex scan filter: {e}"))?; if let Some(limit) = limit && filter.is_none() diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index 06af9e25e8c..d4529e05496 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -36,6 +36,8 @@ use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::dtype::PType::I32; use vortex_array::dtype::StructFields; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::Expression; use vortex_array::expr::and; use vortex_array::expr::cast; use vortex_array::expr::col; @@ -83,6 +85,7 @@ use vortex_layout::layouts::table::TableStrategy; use vortex_layout::layouts::zoned::LegacyStats; use vortex_layout::layouts::zoned::Zoned; use vortex_layout::scan::scan_builder::ScanBuilder; +use vortex_layout::scan::scan_builder::optimize_and_bind; use vortex_layout::scan::split_by::SplitBy; use vortex_layout::session::LayoutSession; use vortex_session::VortexSession; @@ -106,6 +109,10 @@ static SESSION: LazyLock = LazyLock::new(|| { session }); +fn bind_scan_expr(file: &VortexFile, expr: Expression) -> BoundExpression { + optimize_and_bind(expr, file.dtype()).vortex_expect("scan expression should bind") +} + #[tokio::test] async fn test_eof_values() { // this test exists as a reminder to think about whether we should increment the version @@ -313,7 +320,7 @@ async fn test_read_projection() { let array = file .scan() .unwrap() - .with_projection(select(["strings"], root())) + .with_projection(bind_scan_expr(&file, select(["strings"], root()))) .into_array_stream() .unwrap() .read_all() @@ -339,7 +346,7 @@ async fn test_read_projection() { let array = file .scan() .unwrap() - .with_projection(select(["numbers"], root())) + .with_projection(bind_scan_expr(&file, select(["numbers"], root()))) .into_array_stream() .unwrap() .read_all() @@ -510,18 +517,19 @@ async fn issue_5385_filter_casted_column() { .await .unwrap(); - let result = SESSION - .open_options() - .open_buffer(buf) - .unwrap() + let file = SESSION.open_options().open_buffer(buf).unwrap(); + let result = file .scan() .unwrap() - .with_filter(eq( - cast( - get_item("x", root()), - DType::Primitive(PType::U16, Nullability::NonNullable), + .with_filter(bind_scan_expr( + &file, + eq( + cast( + get_item("x", root()), + DType::Primitive(PType::U16, Nullability::NonNullable), + ), + lit(1u16), ), - lit(1u16), )) .into_array_stream() .unwrap() @@ -562,13 +570,14 @@ async fn filter_string() { .await .unwrap(); - let result: Vec<_> = SESSION - .open_options() - .open_buffer(buf) - .unwrap() + let file = SESSION.open_options().open_buffer(buf).unwrap(); + let result: Vec<_> = file .scan() .unwrap() - .with_filter(eq(get_item("name", root()), lit("Joseph"))) + .with_filter(bind_scan_expr( + &file, + eq(get_item("name", root()), lit("Joseph")), + )) .into_array_stream() .unwrap() .try_collect() @@ -622,17 +631,18 @@ async fn filter_or() { .await .unwrap(); - let result: Vec<_> = SESSION - .open_options() - .open_buffer(buf) - .unwrap() + let file = SESSION.open_options().open_buffer(buf).unwrap(); + let result: Vec<_> = file .scan() .unwrap() - .with_filter(or( - eq(get_item("name", root()), lit("Angela")), - and( - gt_eq(get_item("age", root()), lit(20)), - lt_eq(get_item("age", root()), lit(30)), + .with_filter(bind_scan_expr( + &file, + or( + eq(get_item("name", root()), lit("Angela")), + and( + gt_eq(get_item("age", root()), lit(20)), + lt_eq(get_item("age", root()), lit(30)), + ), ), )) .into_array_stream() @@ -690,15 +700,16 @@ async fn filter_and() { .await .unwrap(); - let result: Vec<_> = SESSION - .open_options() - .open_buffer(buf) - .unwrap() + let file = SESSION.open_options().open_buffer(buf).unwrap(); + let result: Vec<_> = file .scan() .unwrap() - .with_filter(and( - gt(get_item("age", root()), lit(21)), - lt_eq(get_item("age", root()), lit(33)), + .with_filter(bind_scan_expr( + &file, + and( + gt(get_item("age", root()), lit(21)), + lt_eq(get_item("age", root()), lit(33)), + ), )) .into_array_stream() .unwrap() @@ -904,7 +915,10 @@ async fn test_with_indices_and_with_row_filter_simple() { let actual_kept_array = file .scan() .unwrap() - .with_filter(gt(get_item("numbers", root()), lit(50_i16))) + .with_filter(bind_scan_expr( + &file, + gt(get_item("numbers", root()), lit(50_i16)), + )) .with_row_indices(Buffer::empty()) .into_array_stream() .unwrap() @@ -922,7 +936,10 @@ async fn test_with_indices_and_with_row_filter_simple() { let actual_kept_array = file .scan() .unwrap() - .with_filter(gt(get_item("numbers", root()), lit(50_i16))) + .with_filter(bind_scan_expr( + &file, + gt(get_item("numbers", root()), lit(50_i16)), + )) .with_row_indices(Buffer::from_iter(kept_indices)) .into_array_stream() .unwrap() @@ -950,7 +967,10 @@ async fn test_with_indices_and_with_row_filter_simple() { let actual_array = file .scan() .unwrap() - .with_filter(gt(get_item("numbers", root()), lit(50_i16))) + .with_filter(bind_scan_expr( + &file, + gt(get_item("numbers", root()), lit(50_i16)), + )) .with_row_indices((0..500).collect::>()) .into_array_stream() .unwrap() @@ -1012,7 +1032,10 @@ async fn filter_string_chunked() { let actual_array = file .scan() .unwrap() - .with_filter(eq(get_item("name", root()), lit("Joseph"))) + .with_filter(bind_scan_expr( + &file, + eq(get_item("name", root()), lit("Joseph")), + )) .into_array_stream() .unwrap() .read_all() @@ -1102,9 +1125,12 @@ async fn test_pruning_with_or() { let actual_array = file .scan() .unwrap() - .with_filter(or( - lt_eq(get_item("letter", root()), lit("J")), - lt(get_item("number", root()), lit(25)), + .with_filter(bind_scan_expr( + &file, + or( + lt_eq(get_item("letter", root()), lit("J")), + lt(get_item("number", root()), lit(25)), + ), )) .into_array_stream() .unwrap() @@ -1177,7 +1203,10 @@ async fn test_repeated_projection() { let actual = file .scan() .unwrap() - .with_projection(select(["strings", "strings"], root())) + .with_projection(bind_scan_expr( + &file, + select(["strings", "strings"], root()), + )) .into_array_stream() .unwrap() .read_all() @@ -1291,7 +1320,7 @@ async fn write_nullable_top_level_struct() { async fn round_trip( array: &ArrayRef, - f: impl Fn(ScanBuilder) -> VortexResult>, + f: impl FnOnce(ScanBuilder) -> VortexResult>, ) -> VortexResult { let mut writer = vec![]; SESSION @@ -1354,15 +1383,19 @@ async fn write_nullable_nested_struct() -> VortexResult<()> { #[tokio::test] async fn scan_empty_fields() -> VortexResult<()> { let array = (0..10000).collect::(); - - let result = round_trip(&array.clone().into_array(), |scan| { - Ok(scan.with_projection(Pack.new_expr( + let projection = optimize_and_bind( + Pack.new_expr( PackOptions { names: Default::default(), nullability: Nullability::Nullable, }, [], - ))) + ), + array.dtype(), + )?; + + let result = round_trip(&array.clone().into_array(), |scan| { + Ok(scan.with_projection(projection)) }) .await?; @@ -2107,13 +2140,9 @@ async fn timestamp_unit_mismatch() -> Result<(), Box> { )), ); - let mut stream = SESSION - .open_options() - .open_buffer(buf)? - .scan()? - .with_filter(filter_expr) - .into_array_stream()?; - + let file = SESSION.open_options().open_buffer(buf)?; + let filter = optimize_and_bind(filter_expr, file.dtype())?; + let mut stream = file.scan()?.with_filter(filter).into_array_stream()?; let result = stream.try_next().await; assert!(result.is_err()); @@ -2160,13 +2189,9 @@ async fn timestamp_unit_mismatch_errors_with_constant_children() )), ); - let stream = SESSION - .open_options() - .open_buffer(buf)? - .scan()? - .with_filter(filter_expr) - .into_array_stream()?; - + let file = SESSION.open_options().open_buffer(buf)?; + let filter = optimize_and_bind(filter_expr, file.dtype())?; + let stream = file.scan()?.with_filter(filter).into_array_stream()?; let results = stream.try_collect::>().await; assert!( @@ -2259,7 +2284,7 @@ async fn test_large_flat_chunk_scan_subdivides_splits() -> VortexResult<()> { // A filtered scan crossing sub-split boundaries selects exactly the matching rows. let result = file .scan()? - .with_filter(gt(root(), lit(0i32))) + .with_filter(bind_scan_expr(&file, gt(root(), lit(0i32)))) .into_array_stream()? .read_all() .await?; @@ -2306,7 +2331,7 @@ async fn test_flat_chunk_scan_with_row_count_splits( let result = file .scan()? .with_split_by(SplitBy::RowCount(rows_per_split)) - .with_filter(gt(root(), lit(0i32))) + .with_filter(bind_scan_expr(&file, gt(root(), lit(0i32)))) .into_array_stream()? .read_all() .await?; @@ -2669,9 +2694,9 @@ async fn repro_8166_binary_gt_all_ff_max() -> VortexResult<()> { )), ); - let result = SESSION - .open_options() - .open_buffer(buf)? + let file = SESSION.open_options().open_buffer(buf)?; + let filter = optimize_and_bind(filter, file.dtype())?; + let result = file .scan()? .with_filter(filter) .into_array_stream()? diff --git a/vortex-layout/src/layouts/dict/reader.rs b/vortex-layout/src/layouts/dict/reader.rs index 5718d7bc985..cc8023ea89d 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -55,7 +55,7 @@ pub struct DictReader { /// Cached dict values array values_array: OnceLock, /// Cache of expression evaluation results on the values array by expression - values_evals: DashMap, + values_evals: DashMap, values: LayoutReaderRef, codes: LayoutReaderRef, @@ -153,13 +153,15 @@ impl DictReader { // shouldn't. // TODO(joe): fixme + let key = ExactBoundExpr(expr.clone()); + // Check cache first with read-only lock - if let Some(fut) = self.values_evals.get(&expr) { + if let Some(fut) = self.values_evals.get(&key) { return fut.clone(); } self.values_evals - .entry(expr.clone()) + .entry(key) .or_insert_with(|| { self.values_array_uncanonical() .map(move |array| { diff --git a/vortex-layout/src/layouts/zoned/pruning.rs b/vortex-layout/src/layouts/zoned/pruning.rs index 089f2e58e67..700aab86ce4 100644 --- a/vortex-layout/src/layouts/zoned/pruning.rs +++ b/vortex-layout/src/layouts/zoned/pruning.rs @@ -19,6 +19,7 @@ use vortex_array::aggregate_fn::AggregateFnRef; use vortex_array::arrays::StructArray; use vortex_array::dtype::DType; use vortex_array::expr::BoundExpression; +use vortex_array::expr::ExactBoundExpr; use vortex_array::expr::root; use vortex_array::scalar_fn::fns::dynamic::DynamicExprUpdates; use vortex_error::SharedVortexResult; @@ -47,9 +48,9 @@ pub(super) struct PruningState { aggregate_fns: Arc<[AggregateFnRef]>, lazy_children: Arc, session: VortexSession, - pruning_result: LazyLock>>, + pruning_result: LazyLock>>, zone_map: OnceLock, - pruning_predicates: LazyLock>>, + pruning_predicates: LazyLock>>, } impl PruningState { @@ -78,12 +79,14 @@ impl PruningState { } pub(super) fn pruning_mask_future(&self, expr: BoundExpression) -> Option { - if let Some(result) = self.pruning_result.get(&expr) { + let key = ExactBoundExpr(expr.clone()); + + if let Some(result) = self.pruning_result.get(&key) { return result.value().clone(); } self.pruning_result - .entry(expr.clone()) + .entry(key) .or_insert_with(|| { let dynamic_updates = DynamicExprUpdates::new(&expr); match self.pruning_predicate(expr.clone()) { @@ -124,8 +127,10 @@ impl PruningState { } fn pruning_predicate(&self, expr: BoundExpression) -> Option { + let key = ExactBoundExpr(expr.clone()); + self.pruning_predicates - .entry(expr.clone()) + .entry(key) .or_default() .get_or_init(move || match expr.falsify(&self.session) { Ok(predicate) => predicate, diff --git a/vortex-layout/src/lib.rs b/vortex-layout/src/lib.rs index 96a34ce75f3..03cd832a280 100644 --- a/vortex-layout/src/lib.rs +++ b/vortex-layout/src/lib.rs @@ -11,9 +11,9 @@ //! Most users enter this crate through file APIs, but extension authors implement [`VTable`] and //! [`LayoutStrategy`] to add new on-disk organizations. //! -//! Scanning is built with [`scan::scan_builder::ScanBuilder`]. It accepts a projection expression, -//! optional filter, optional row range, [`Selection`](vortex_scan::selection::Selection), split -//! strategy, and task concurrency settings, then produces array streams or iterators. +//! Scanning is built with [`scan::scan_builder::ScanBuilder`]. It accepts a bound projection, +//! optional bound filter, optional row range, [`Selection`](vortex_scan::selection::Selection), +//! split strategy, and task concurrency settings, then produces array streams or iterators. pub mod layouts; pub use children::*; diff --git a/vortex-layout/src/scan/layout.rs b/vortex-layout/src/scan/layout.rs index f0e6366a8fc..dadd6661078 100644 --- a/vortex-layout/src/scan/layout.rs +++ b/vortex-layout/src/scan/layout.rs @@ -18,7 +18,7 @@ use vortex_array::arrays::ConstantArray; use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; use vortex_array::dtype::Nullability; -use vortex_array::expr::Expression; +use vortex_array::expr::BoundExpression; use vortex_array::expr::stats::Precision; use vortex_array::scalar::Scalar; use vortex_array::stats::StatsSet; @@ -41,6 +41,7 @@ use vortex_session::VortexSession; use crate::LayoutReaderRef; use crate::scan::scan_builder::ScanBuilder; +use crate::scan::scan_builder::optimize_and_bind; /// An implementation of a [`DataSource`] that reads data from a [`LayoutReaderRef`]. pub struct LayoutReaderDataSource { @@ -115,13 +116,18 @@ impl DataSource for LayoutReaderDataSource { let total_rows = self.reader.row_count(); let row_range = scan_request.row_range.unwrap_or(0..total_rows); - let dtype = scan_request.projection.return_dtype(self.reader.dtype())?; + let projection = optimize_and_bind(scan_request.projection, self.reader.dtype())?; + let filter = scan_request + .filter + .map(|expr| optimize_and_bind(expr, self.reader.dtype())) + .transpose()?; + let dtype = projection.dtype().clone(); // If the dtype is an empty struct, and there is no filter, we can return a special // length-only scan. if let DType::Struct(fields, Nullability::NonNullable) = &dtype && fields.nfields() == 0 - && scan_request.filter.is_none() + && filter.is_none() { // FIXME(ngates): extract out maybe? let row_count = row_range.end - row_range.start; @@ -139,14 +145,13 @@ impl DataSource for LayoutReaderDataSource { // Check file-level pruning: if the filter can be proven false for the entire row range // using file-level statistics (e.g. via FileStatsLayoutReader), skip the scan entirely. - if let Some(filter) = &scan_request.filter { - let filter = filter.bind(self.reader.dtype())?; + if let Some(filter) = &filter { let mask = Mask::new_true( usize::try_from(row_range.end - row_range.start).unwrap_or(usize::MAX), ); let pruning_result = self .reader - .pruning_evaluation(&row_range, &filter, mask)? + .pruning_evaluation(&row_range, filter, mask)? .now_or_never(); if let Some(Ok(result_mask)) = pruning_result && result_mask.all_false() @@ -162,8 +167,8 @@ impl DataSource for LayoutReaderDataSource { reader: Arc::clone(&self.reader), session: self.session.clone(), dtype, - projection: scan_request.projection, - filter: scan_request.filter, + projection, + filter, limit: scan_request.limit, selection: scan_request.selection, ordered: scan_request.ordered, @@ -183,8 +188,8 @@ struct LayoutReaderScan { reader: LayoutReaderRef, session: VortexSession, dtype: DType, - projection: Expression, - filter: Option, + projection: BoundExpression, + filter: Option, limit: Option, ordered: bool, selection: Selection, @@ -276,8 +281,8 @@ impl Stream for LayoutReaderScan { struct LayoutReaderSplit { reader: LayoutReaderRef, session: VortexSession, - projection: Expression, - filter: Option, + projection: BoundExpression, + filter: Option, limit: Option, ordered: bool, row_range: Range, diff --git a/vortex-layout/src/scan/multi.rs b/vortex-layout/src/scan/multi.rs index d251ee15617..f188a5c34d4 100644 --- a/vortex-layout/src/scan/multi.rs +++ b/vortex-layout/src/scan/multi.rs @@ -26,6 +26,7 @@ use std::any::Any; use std::collections::VecDeque; +use std::ops::Range; use std::sync::Arc; use async_trait::async_trait; @@ -36,6 +37,7 @@ use itertools::Itertools; use tracing::Instrument; use vortex_array::dtype::DType; use vortex_array::dtype::FieldPath; +use vortex_array::expr::BoundExpression; use vortex_array::expr::stats::Precision; use vortex_array::stats::StatsSet; use vortex_array::stream::ArrayStreamAdapter; @@ -43,6 +45,7 @@ use vortex_array::stream::ArrayStreamExt; use vortex_array::stream::SendableArrayStream; use vortex_error::VortexResult; use vortex_error::vortex_bail; +use vortex_error::vortex_err; use vortex_io::session::RuntimeSessionExt; use vortex_mask::Mask; use vortex_scan::DataSource; @@ -58,6 +61,7 @@ use vortex_utils::parallelism::get_available_parallelism; use crate::LayoutReaderRef; use crate::scan::scan_builder::ScanBuilder; +use crate::scan::scan_builder::optimize_and_bind; /// Default concurrency for opening deferred readers. const DEFAULT_CONCURRENCY: usize = 8; @@ -299,12 +303,14 @@ impl DataSource for MultiLayoutDataSource { } } - let dtype = scan_request.projection.return_dtype(&self.dtype)?; + let request = BoundScanRequest::try_new(scan_request, &self.dtype)?; + let dtype = request.projection.dtype().clone(); Ok(Box::new(MultiLayoutScan { session: self.session.clone(), + source_dtype: self.dtype.clone(), dtype, - request: scan_request, + request, ready, deferred, handle: self.session.handle(), @@ -317,10 +323,51 @@ impl DataSource for MultiLayoutDataSource { } } +#[derive(Clone)] +struct BoundScanRequest { + projection: BoundExpression, + filter: Option, + row_range: Option>, + selection: Selection, + partition_selection: Selection, + partition_range: Option>, + ordered: bool, + limit: Option, +} + +impl BoundScanRequest { + fn try_new(request: ScanRequest, dtype: &DType) -> VortexResult { + let ScanRequest { + projection, + filter, + row_range, + selection, + partition_selection, + partition_range, + ordered, + limit, + } = request; + + Ok(Self { + projection: optimize_and_bind(projection, dtype)?, + filter: filter + .map(|expr| optimize_and_bind(expr, dtype)) + .transpose()?, + row_range, + selection, + partition_selection, + partition_range, + ordered, + limit, + }) + } +} + struct MultiLayoutScan { session: VortexSession, + source_dtype: DType, dtype: DType, - request: ScanRequest, + request: BoundScanRequest, ready: VecDeque, deferred: VecDeque>, handle: vortex_io::runtime::Handle, @@ -344,6 +391,7 @@ impl DataSourceScan for MultiLayoutScan { fn partitions(self: Box) -> PartitionStream { let Self { session, + source_dtype, dtype: _, request, ready, @@ -400,7 +448,9 @@ impl DataSourceScan for MultiLayoutScan { .chain(deferred_stream) .enumerate() .flat_map(move |(i, reader_result)| match reader_result { - Ok(reader) => reader_partition(i, reader, session.clone(), request.clone()), + Ok(reader) => { + reader_partition(i, reader, session.clone(), &source_dtype, request.clone()) + } Err(e) => stream::once(async move { Err(e) }).boxed(), }) .boxed() @@ -416,8 +466,18 @@ fn reader_partition( partition_idx: usize, reader: LayoutReaderRef, session: VortexSession, - request: ScanRequest, + source_dtype: &DType, + request: BoundScanRequest, ) -> PartitionStream { + if reader.dtype() != source_dtype { + let error = vortex_err!( + "Multi-layout reader dtype mismatch: expected {}, got {}", + source_dtype, + reader.dtype() + ); + return stream::once(async move { Err(error) }).boxed(); + } + let row_count = reader.row_count(); let row_range = request.row_range.clone().unwrap_or(0..row_count); @@ -446,8 +506,7 @@ fn reader_partition( if let Some(filter) = &request.filter { let mask_len = usize::try_from(row_range.end - row_range.start).unwrap_or(usize::MAX); let mask = Mask::new_true(mask_len); - if let Ok(filter) = filter.bind(reader.dtype()) - && let Ok(pruning_future) = reader.pruning_evaluation(&row_range, &filter, mask) + if let Ok(pruning_future) = reader.pruning_evaluation(&row_range, filter, mask) && let Some(Ok(result_mask)) = pruning_future.now_or_never() && result_mask.all_false() { @@ -459,7 +518,7 @@ fn reader_partition( Ok(Box::new(MultiLayoutPartition { reader, session, - request: ScanRequest { + request: BoundScanRequest { row_range: Some(row_range), ..request }, @@ -476,7 +535,7 @@ fn reader_partition( struct MultiLayoutPartition { reader: LayoutReaderRef, session: VortexSession, - request: ScanRequest, + request: BoundScanRequest, index: usize, } diff --git a/vortex-layout/src/scan/scan_builder.rs b/vortex-layout/src/scan/scan_builder.rs index 16a8758c16d..0352c1ad319 100644 --- a/vortex-layout/src/scan/scan_builder.rs +++ b/vortex-layout/src/scan/scan_builder.rs @@ -19,7 +19,6 @@ use vortex_array::dtype::FieldMask; use vortex_array::expr::BoundExpression; use vortex_array::expr::Expression; use vortex_array::expr::analysis::referenced_field_paths; -use vortex_array::expr::root; use vortex_array::iter::ArrayIterator; use vortex_array::iter::ArrayIteratorAdapter; use vortex_array::stats::StatsSet; @@ -46,6 +45,11 @@ use crate::scan::split_by::SplitBy; use crate::scan::splits::Splits; use crate::scan::splits::attempt_split_ranges; +/// Optimize `expr` against `scope`, then bind it into a typed expression tree. +pub fn optimize_and_bind(expr: Expression, scope: &DType) -> VortexResult { + expr.optimize_recursive(scope)?.bind(scope) +} + /// Builder for scanning a [`LayoutReader`] into arrays, streams, iterators, or mapped outputs. /// /// A scan has three independent row restriction mechanisms: @@ -54,14 +58,13 @@ use crate::scan::splits::attempt_split_ranges; /// - [`with_selection`](Self::with_selection) applies a [`Selection`] inside that range. /// - [`with_filter`](Self::with_filter) evaluates an expression predicate during execution. /// -/// Projection and filter expressions are optimized against the reader dtype during -/// [`prepare`](Self::prepare). Work is divided by the configured [`SplitBy`] strategy or by -/// explicit selection ranges. +/// Projection and filter expressions must be bound against the reader dtype. Work is divided by +/// the configured [`SplitBy`] strategy or by explicit selection ranges. pub struct ScanBuilder { session: VortexSession, layout_reader: LayoutReaderRef, - projection: Expression, - filter: Option, + projection: BoundExpression, + filter: Option, /// Whether the scan needs to return splits in the order they appear in the file. ordered: bool, /// Optionally read a subset of the rows in the file. @@ -88,10 +91,11 @@ pub struct ScanBuilder { impl ScanBuilder { /// Create a scan builder over `layout_reader` using `session` for runtime and execution state. pub fn new(session: VortexSession, layout_reader: Arc) -> Self { + let projection = BoundExpression::new_root(layout_reader.dtype().clone()); Self { session, layout_reader, - projection: root(), + projection, filter: None, ordered: true, row_range: None, @@ -132,20 +136,20 @@ impl ScanBuilder { } impl ScanBuilder { - /// Add a filter expression evaluated against the projected row ranges. - pub fn with_filter(mut self, filter: Expression) -> Self { + /// Add a filter expression bound against the reader dtype. + pub fn with_filter(mut self, filter: BoundExpression) -> Self { self.filter = Some(filter); self } - /// Add or clear the filter expression. - pub fn with_some_filter(mut self, filter: Option) -> Self { + /// Add or clear a filter expression bound against the reader dtype. + pub fn with_some_filter(mut self, filter: Option) -> Self { self.filter = filter; self } - /// Set the projection expression for returned rows. - pub fn with_projection(mut self, projection: Expression) -> Self { + /// Set a projection expression bound against the reader dtype. + pub fn with_projection(mut self, projection: BoundExpression) -> Self { self.projection = projection; self } @@ -230,7 +234,7 @@ impl ScanBuilder { /// The [`DType`] returned by the scan, after applying the projection. pub fn dtype(&self) -> VortexResult { - self.projection.return_dtype(self.layout_reader.dtype()) + Ok(self.projection.dtype().clone()) } /// The session used by the scan. @@ -283,19 +287,8 @@ impl ScanBuilder { self.session.clone(), )); - // Normalize and simplify the expressions. - let projection = self.projection.optimize_recursive(layout_reader.dtype())?; - - let filter = self - .filter - .map(|f| f.optimize_recursive(layout_reader.dtype())) - .transpose()?; - - let bound_projection = projection.bind(layout_reader.dtype())?; - let bound_filter = filter - .as_ref() - .map(|expr| expr.bind(layout_reader.dtype())) - .transpose()?; + let bound_projection = self.projection; + let bound_filter = self.filter; // Construct field masks and compute the row splits of the scan. let field_mask = referenced_field_masks(&bound_projection, bound_filter.as_ref())?; @@ -487,6 +480,7 @@ mod test { use vortex_array::dtype::PType; use vortex_array::dtype::StructFields; use vortex_array::expr::BoundExpression; + use vortex_array::expr::ExactBoundExpr; use vortex_array::expr::eq; use vortex_array::expr::get_item; use vortex_array::expr::is_not_null; @@ -526,6 +520,24 @@ mod test { ) } + #[test] + fn bound_setters_preserve_identity() -> VortexResult<()> { + let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let projection = eq(root(), lit(1_i32)).bind(&dtype)?; + let filter = eq(root(), lit(2_i32)).bind(&dtype)?; + let expected_projection = ExactBoundExpr(projection.clone()); + let expected_filter = ExactBoundExpr(filter.clone()); + let reader = Arc::new(CountingLayoutReader::new(Arc::new(AtomicUsize::new(0)))); + + let builder = ScanBuilder::new(SCAN_SESSION.clone(), reader) + .with_projection(projection) + .with_filter(filter); + + assert_eq!(ExactBoundExpr(builder.projection), expected_projection); + assert_eq!(builder.filter.map(ExactBoundExpr), Some(expected_filter)); + Ok(()) + } + #[test] fn root_projection_produces_all_mask() -> VortexResult<()> { let dtype = DType::Primitive(PType::I32, Nullability::NonNullable); diff --git a/vortex-python/src/dataset.rs b/vortex-python/src/dataset.rs index 504ed05aeee..487f6ca9249 100644 --- a/vortex-python/src/dataset.rs +++ b/vortex-python/src/dataset.rs @@ -24,6 +24,7 @@ use vortex::expr::select; use vortex::file::OpenOptionsSessionExt; use vortex::file::VortexFile; use vortex::io::runtime::BlockingRuntime; +use vortex::layout::scan::scan_builder::optimize_and_bind; use vortex::layout::scan::split_by::SplitBy; use vortex_arrow::ToArrowType; @@ -58,6 +59,10 @@ pub fn read_array_from_reader( row_range: Option<(u64, u64)>, ctx: &mut ExecutionCtx, ) -> VortexResult { + let projection = optimize_and_bind(projection, vortex_file.dtype())?; + let filter = filter + .map(|filter| optimize_and_bind(filter, vortex_file.dtype())) + .transpose()?; let mut scan = vortex_file.scan()?.with_projection(projection); if let Some(filter) = filter { @@ -185,6 +190,10 @@ impl PyVortexDataset { let filter = filter_from_python(row_filter); let reader = self_.py().detach(move || { + let projection = optimize_and_bind(projection, vxf.dtype())?; + let filter = filter + .map(|filter| optimize_and_bind(filter, vxf.dtype())) + .transpose()?; let mut scan = vxf .scan()? .with_projection(projection) @@ -224,9 +233,13 @@ impl PyVortexDataset { let vxf = self_.vxf.clone(); let filter = filter_from_python(row_filter); let n_rows: usize = self_.py().detach(move || { + let projection = optimize_and_bind(select(FieldNames::empty(), root()), vxf.dtype())?; + let filter = filter + .map(|filter| optimize_and_bind(filter, vxf.dtype())) + .transpose()?; let mut scan = vxf .scan()? - .with_projection(select(FieldNames::empty(), root())) + .with_projection(projection) .with_some_filter(filter) .with_split_by(split_by.map(SplitBy::RowCount).unwrap_or(SplitBy::Layout)); if let Some((l, r)) = row_range { diff --git a/vortex-python/src/file.rs b/vortex-python/src/file.rs index 5f8a39862b7..ea9ef8e86f2 100644 --- a/vortex-python/src/file.rs +++ b/vortex-python/src/file.rs @@ -26,6 +26,7 @@ use vortex::file::OpenOptionsSessionExt; use vortex::file::VortexFile; use vortex::io::runtime::BlockingRuntime; use vortex::layout::scan::scan_builder::ScanBuilder; +use vortex::layout::scan::scan_builder::optimize_and_bind; use vortex::layout::scan::split_by::SplitBy; use vortex::layout::segments::MokaSegmentCache; use vortex_arrow::ToArrowType; @@ -171,10 +172,15 @@ impl PyVortexFile { .map(Arc::new); let reader = slf.py().detach(|| { + let filter = expr + .map(|e| optimize_and_bind(e.into_inner(), vxf.dtype())) + .transpose()?; + let projection = + optimize_and_bind(projection.map(|p| p.0).unwrap_or_else(root), vxf.dtype())?; let mut builder = vxf .scan()? - .with_some_filter(expr.map(|e| e.into_inner())) - .with_projection(projection.map(|p| p.0).unwrap_or_else(root)); + .with_some_filter(filter) + .with_projection(projection); if let Some(limit) = limit { builder = builder.with_limit(limit); @@ -219,10 +225,14 @@ fn scan_builder( batch_size: Option, ctx: &mut ExecutionCtx, ) -> VortexResult> { + let projection = optimize_and_bind(projection.unwrap_or_else(root), vxf.dtype())?; + let expr = expr + .map(|expr| optimize_and_bind(expr, vxf.dtype())) + .transpose()?; let mut builder = vxf .scan()? .with_some_filter(expr) - .with_projection(projection.unwrap_or_else(root)); + .with_projection(projection); if let Some(limit) = limit { builder = builder.with_limit(limit); diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 5abe779ab01..beba8a99c20 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -70,6 +70,7 @@ //! use vortex::array::validity::Validity; //! use vortex::buffer::{ByteBufferMut, buffer}; //! use vortex::file::{OpenOptionsSessionExt, WriteOptionsSessionExt}; +//! use vortex::layout::scan::scan_builder::optimize_and_bind; //! use vortex::session::VortexSession; //! //! # async fn example() -> vortex::error::VortexResult<()> { @@ -82,11 +83,13 @@ //! .write(&mut bytes, array.into_array().to_array_stream()) //! .await?; //! -//! let filtered = session +//! let file = session //! .open_options() -//! .open_buffer(bytes)? +//! .open_buffer(bytes)?; +//! let filter = optimize_and_bind(gt(root(), lit(2u64)), file.dtype())?; +//! let filtered = file //! .scan()? -//! .with_filter(gt(root(), lit(2u64))) +//! .with_filter(filter) //! .into_array_stream()? //! .read_all() //! .await?; @@ -357,6 +360,7 @@ mod test { use vortex_file::OpenOptionsSessionExt; use vortex_file::WriteOptionsSessionExt; use vortex_file::WriteStrategyBuilder; + use vortex_layout::scan::scan_builder::optimize_and_bind; use vortex_session::VortexSession; use crate as vortex; @@ -438,12 +442,11 @@ mod test { // [write] // [read] - let array = session - .open_options() - .open_path(path.clone()) - .await? + let file = session.open_options().open_path(path.clone()).await?; + let filter = optimize_and_bind(gt(root(), lit(2u64)), file.dtype())?; + let array = file .scan()? - .with_filter(gt(root(), lit(2u64))) + .with_filter(filter) .into_array_stream()? .read_all() .await?; @@ -537,12 +540,11 @@ mod test { .await?; // Read the file back, but project down to just the "value" column. - let projected = session - .open_options() - .open_path(path.clone()) - .await? + let file = session.open_options().open_path(path.clone()).await?; + let projection = optimize_and_bind(select(["value"], root()), file.dtype())?; + let projected = file .scan()? - .with_projection(select(["value"], root())) + .with_projection(projection) .into_array_stream()? .read_all() .await?; From 518643dce315805a4719a4d883c4b2847f032cf2 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Wed, 5 Aug 2026 14:09:38 +0100 Subject: [PATCH 3/4] fix: compare bound conjuncts without unbinding Signed-off-by: Joe Isaacs --- vortex-layout/src/scan/filter.rs | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/vortex-layout/src/scan/filter.rs b/vortex-layout/src/scan/filter.rs index fdba393a46f..2cacfbecf89 100644 --- a/vortex-layout/src/scan/filter.rs +++ b/vortex-layout/src/scan/filter.rs @@ -173,13 +173,12 @@ mod tests { let filter = FilterExpr::new(bound); let conjuncts = filter.conjuncts(); - assert_eq!( - conjuncts - .iter() - .map(|expr| expr.unbind()) - .collect::>(), - vec![root(), not(root()), lit(true)] - ); + let expected = vec![ + root().bind(&dtype)?, + not(root()).bind(&dtype)?, + lit(true).bind(&dtype)?, + ]; + assert_eq!(conjuncts, expected.as_slice()); assert_eq!( conjuncts .iter() From f5ad4f9dac3bdf17a64e19609a22b003106080c4 Mon Sep 17 00:00:00 2001 From: Joe Isaacs Date: Wed, 5 Aug 2026 14:21:20 +0100 Subject: [PATCH 4/4] fix: defer multi-scan filter binding errors Signed-off-by: Joe Isaacs --- vortex-layout/src/scan/multi.rs | 36 ++++++++++++++++++++++++++++----- 1 file changed, 31 insertions(+), 5 deletions(-) diff --git a/vortex-layout/src/scan/multi.rs b/vortex-layout/src/scan/multi.rs index f188a5c34d4..2bc9d5dffe2 100644 --- a/vortex-layout/src/scan/multi.rs +++ b/vortex-layout/src/scan/multi.rs @@ -43,6 +43,7 @@ use vortex_array::stats::StatsSet; use vortex_array::stream::ArrayStreamAdapter; use vortex_array::stream::ArrayStreamExt; use vortex_array::stream::SendableArrayStream; +use vortex_error::SharedVortexResult; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_err; @@ -326,7 +327,7 @@ impl DataSource for MultiLayoutDataSource { #[derive(Clone)] struct BoundScanRequest { projection: BoundExpression, - filter: Option, + filter: SharedVortexResult>, row_range: Option>, selection: Selection, partition_selection: Selection, @@ -352,7 +353,8 @@ impl BoundScanRequest { projection: optimize_and_bind(projection, dtype)?, filter: filter .map(|expr| optimize_and_bind(expr, dtype)) - .transpose()?, + .transpose() + .map_err(Arc::new), row_range, selection, partition_selection, @@ -503,7 +505,7 @@ fn reader_partition( // Check file-level pruning: if the filter can be proven false for the entire row range // using file-level statistics, skip this reader entirely. - if let Some(filter) = &request.filter { + if let Ok(Some(filter)) = &request.filter { let mask_len = usize::try_from(row_range.end - row_range.start).unwrap_or(usize::MAX); let mask = Mask::new_true(mask_len); if let Ok(pruning_future) = reader.pruning_evaluation(&row_range, filter, mask) @@ -559,7 +561,11 @@ impl Partition for MultiLayoutPartition { .limit .map_or(row_count, |limit| row_count.min(limit)); - if self.request.filter.is_some() { + let has_filter = match &self.request.filter { + Ok(filter) => filter.is_some(), + Err(_) => true, + }; + if has_filter { Precision::inexact(row_count) } else { Precision::exact(row_count) @@ -572,10 +578,11 @@ impl Partition for MultiLayoutPartition { fn execute(self: Box) -> VortexResult { let request = self.request; + let filter = request.filter?; let mut builder = ScanBuilder::new(self.session, self.reader) .with_selection(request.selection) .with_projection(request.projection) - .with_some_filter(request.filter) + .with_some_filter(filter) .with_some_limit(request.limit) .with_ordered(request.ordered); @@ -596,6 +603,10 @@ impl Partition for MultiLayoutPartition { mod tests { use rstest::rstest; use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::expr::eq; + use vortex_array::expr::lit; + use vortex_array::expr::root; use super::*; use crate::scan::test::new_session; @@ -630,4 +641,19 @@ mod tests { fn byte_size_precision(#[case] sizes: Vec>, #[case] expected: Precision) { assert_eq!(deferred_source(sizes).byte_size(), expected); } + + #[test] + fn filter_binding_errors_are_deferred() -> VortexResult<()> { + let dtype = DType::Primitive(PType::U8, Nullability::NonNullable); + let request = ScanRequest { + filter: Some(eq(root(), lit(67_i32))), + ..ScanRequest::default() + }; + + let request = BoundScanRequest::try_new(request, &dtype)?; + + assert_eq!(request.projection.dtype(), &dtype); + assert!(request.filter.is_err()); + Ok(()) + } }