diff --git a/benchmarks/compress-bench/src/vortex.rs b/benchmarks/compress-bench/src/vortex.rs index fae0cc44189..0726886bc4c 100644 --- a/benchmarks/compress-bench/src/vortex.rs +++ b/benchmarks/compress-bench/src/vortex.rs @@ -64,14 +64,17 @@ 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 = select(names, root()) + .optimize_recursive(&source_dtype)? + .bind(&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..e0dddfb4b74 100644 --- a/docs/developer-guide/internals/session.md +++ b/docs/developer-guide/internals/session.md @@ -78,8 +78,11 @@ session.write_options() .await?; // Scanning a layout +let filter = expr + .optimize_recursive(layout_reader.dtype())? + .bind(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..38a9016e4b8 100644 --- a/fuzz/fuzz_targets/file_io.rs +++ b/fuzz/fuzz_targets/file_io.rs @@ -78,14 +78,28 @@ 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 = projection_expr + .unwrap_or_else(root) + .optimize_recursive(file.dtype()) + .and_then(|expr| expr.bind(file.dtype())) + .vortex_expect("projection should bind in fuzz test"); + let filter = filter_expr + .map(|filter| { + filter + .optimize_recursive(file.dtype()) + .and_then(|expr| expr.bind(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-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index 26173171d91..bd0e02d8ebd 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -12,13 +12,18 @@ 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::expr::traversal::TraversalOrder; +use crate::expr::traversal::pre_order_visit_down; 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 +176,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,50 +189,57 @@ 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 whether this expression tree contains a node using the given scalar-function vtable. + pub fn contains(&self) -> VortexResult { + let mut contains = false; + pre_order_visit_down(self, |node| { + if node.is::() { + contains = true; + return Ok(TraversalOrder::Stop); + } + Ok(TraversalOrder::Continue) + })?; + Ok(contains) + } + + /// 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) } } @@ -293,6 +310,7 @@ mod tests { use crate::expr::lit; use crate::expr::root; use crate::expr::test_harness::struct_dtype; + use crate::scalar_fn::fns::literal::Literal; fn scope() -> Scope { Scope::new(struct_dtype()) @@ -303,7 +321,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(()) } @@ -335,6 +353,14 @@ mod tests { Ok(()) } + #[test] + fn contains_scalar_function() -> VortexResult<()> { + let bound = eq(col("a"), lit(1_i32)).bind_scope(&scope())?; + assert!(bound.contains::()?); + assert!(!root().bind_scope(&scope())?.contains::()?); + Ok(()) + } + #[test] fn bound_display_matches_unbound() -> VortexResult<()> { for expr in [root(), col("a"), eq(col("a"), lit(1_i32)), lit(true)] { @@ -379,11 +405,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/expression.rs b/vortex-array/src/expr/expression.rs index e53474ad040..27b47ad3d6e 100644 --- a/vortex-array/src/expr/expression.rs +++ b/vortex-array/src/expr/expression.rs @@ -12,7 +12,6 @@ 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; @@ -21,7 +20,6 @@ use crate::expr::traversal::pre_order_visit_down; use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::fns::root::Root; -use crate::stats::rewrite::StatsRewriteCtx; /// A node in a Vortex expression tree. /// @@ -116,40 +114,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/exprs.rs b/vortex-array/src/expr/exprs.rs index 9843277ede7..71225c4bf43 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -16,6 +16,7 @@ use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::FieldNames; use crate::dtype::Nullability; +use crate::expr::BoundExpression; use crate::expr::Expression; use crate::scalar::Scalar; use crate::scalar::ScalarValue; @@ -71,6 +72,11 @@ pub fn root() -> Expression { ROOT.clone() } +/// Creates a bound expression that references a root scope with the given dtype. +pub fn bound_root(dtype: DType) -> BoundExpression { + BoundExpression::new_root(dtype) +} + /// Return whether the expression is a root expression. pub fn is_root(expr: &Expression) -> bool { // root doesn't have any children, and scalar_fns have distinct ids @@ -101,6 +107,13 @@ pub fn lit(value: impl Into) -> Expression { Literal.new_expr(value.into(), []) } +/// Creates a bound literal expression. +pub fn bound_lit(value: impl Into) -> BoundExpression { + Literal + .try_new_bound_expr(value.into(), []) + .vortex_expect("literal expressions are always well-typed") +} + // ---- GetItem / Col ---- /// Creates an expression that accesses a field from the root array. @@ -115,6 +128,11 @@ pub fn col(field: impl Into) -> Expression { GetItem.new_expr(field.into(), vec![root()]) } +/// Creates a bound expression that accesses a field from a root scope with the given dtype. +pub fn bound_col(field: impl Into, scope: DType) -> BoundExpression { + bound_get_item(field, bound_root(scope)) +} + /// Creates an expression that extracts a named field from a struct expression. /// /// Accesses the specified field from the result of the child expression. @@ -127,6 +145,13 @@ pub fn get_item(field: impl Into, child: Expression) -> Expression { GetItem.new_expr(field.into(), vec![child]) } +/// Creates a bound expression that extracts a named field from a struct expression. +pub fn bound_get_item(field: impl Into, child: BoundExpression) -> BoundExpression { + GetItem + .try_new_bound_expr(field.into(), [child]) + .vortex_expect("get-item expressions must reference a field in the child dtype") +} + // ---- VariantGet ---- /// Creates an expression that extracts a path from a Variant expression. @@ -141,6 +166,17 @@ pub fn variant_get( VariantGet.new_expr(VariantGetOptions::new(path.into(), dtype), vec![child]) } +/// Creates a bound expression that extracts a path from a Variant expression. +pub fn bound_variant_get( + child: BoundExpression, + path: impl Into, + dtype: Option, +) -> BoundExpression { + VariantGet + .try_new_bound_expr(VariantGetOptions::new(path.into(), dtype), [child]) + .vortex_expect("variant-get expressions require a Variant child") +} + // ---- CaseWhen ---- /// Creates a CASE WHEN expression with one WHEN/THEN pair and an ELSE value. @@ -156,6 +192,21 @@ pub fn case_when( CaseWhen.new_expr(options, [condition, then_value, else_value]) } +/// Creates a bound CASE WHEN expression with one WHEN/THEN pair and an ELSE value. +pub fn bound_case_when( + condition: BoundExpression, + then_value: BoundExpression, + else_value: BoundExpression, +) -> BoundExpression { + let options = CaseWhenOptions { + num_when_then_pairs: 1, + has_else: true, + }; + CaseWhen + .try_new_bound_expr(options, [condition, then_value, else_value]) + .vortex_expect("case expressions must have boolean conditions and matching branch dtypes") +} + /// Creates a CASE WHEN expression with one WHEN/THEN pair and no ELSE value. pub fn case_when_no_else(condition: Expression, then_value: Expression) -> Expression { let options = CaseWhenOptions { @@ -165,6 +216,20 @@ pub fn case_when_no_else(condition: Expression, then_value: Expression) -> Expre CaseWhen.new_expr(options, [condition, then_value]) } +/// Creates a bound CASE WHEN expression with one WHEN/THEN pair and no ELSE value. +pub fn bound_case_when_no_else( + condition: BoundExpression, + then_value: BoundExpression, +) -> BoundExpression { + let options = CaseWhenOptions { + num_when_then_pairs: 1, + has_else: false, + }; + CaseWhen + .try_new_bound_expr(options, [condition, then_value]) + .vortex_expect("case expressions must have boolean conditions") +} + /// Creates an n-ary CASE WHEN expression from WHEN/THEN pairs and an optional ELSE value. pub fn nested_case_when( when_then_pairs: Vec<(Expression, Expression)>, @@ -195,8 +260,58 @@ pub fn nested_case_when( CaseWhen.new_expr(options, children) } +/// Creates a bound n-ary CASE WHEN expression from WHEN/THEN pairs and an optional ELSE value. +pub fn bound_nested_case_when( + when_then_pairs: Vec<(BoundExpression, BoundExpression)>, + else_value: Option, +) -> BoundExpression { + assert!( + !when_then_pairs.is_empty(), + "nested_case_when requires at least one when/then pair" + ); + + let Ok(num_when_then_pairs) = u32::try_from(when_then_pairs.len()) else { + vortex_panic!("nested_case_when has too many when/then pairs"); + }; + let has_else = else_value.is_some(); + let mut children = Vec::with_capacity(when_then_pairs.len() * 2 + usize::from(has_else)); + for (condition, then_value) in when_then_pairs { + children.push(condition); + children.push(then_value); + } + if let Some(else_expr) = else_value { + children.push(else_expr); + } + + let options = CaseWhenOptions { + num_when_then_pairs, + has_else, + }; + CaseWhen + .try_new_bound_expr(options, children) + .vortex_expect("case expressions must have boolean conditions and matching branch dtypes") +} + // ---- Binary operators ---- +/// Creates a binary expression with the given operator. +pub fn binary(operator: Operator, lhs: Expression, rhs: Expression) -> Expression { + Binary + .try_new_expr(operator, [lhs, rhs]) + .vortex_expect("Failed to create binary expression") +} + +/// Creates a bound binary expression with the given operator. +pub fn bound_binary( + operator: Operator, + lhs: BoundExpression, + rhs: BoundExpression, +) -> BoundExpression { + Binary + .try_new_bound_expr(operator, [lhs, rhs]) + .vortex_expect("binary expressions must have compatible operand dtypes") +} + /// Create a new [`Binary`] using the [`Eq`](Operator::Eq) operator. /// /// ## Example usage @@ -224,6 +339,11 @@ pub fn eq(lhs: Expression, rhs: Expression) -> Expression { .vortex_expect("Failed to create Eq binary expression") } +/// Creates a bound equality expression. +pub fn bound_eq(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { + bound_binary(Operator::Eq, lhs, rhs) +} + /// Create a new [`Binary`] using the [`NotEq`](Operator::NotEq) operator. /// /// ## Example usage @@ -251,6 +371,11 @@ pub fn not_eq(lhs: Expression, rhs: Expression) -> Expression { .vortex_expect("Failed to create NotEq binary expression") } +/// Creates a bound inequality expression. +pub fn bound_not_eq(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { + bound_binary(Operator::NotEq, lhs, rhs) +} + /// Create a new [`Binary`] using the [`Gte`](Operator::Gte) operator. /// /// ## Example usage @@ -278,6 +403,11 @@ pub fn gt_eq(lhs: Expression, rhs: Expression) -> Expression { .vortex_expect("Failed to create Gte binary expression") } +/// Creates a bound greater-than-or-equal expression. +pub fn bound_gt_eq(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { + bound_binary(Operator::Gte, lhs, rhs) +} + /// Create a new [`Binary`] using the [`Gt`](Operator::Gt) operator. /// /// ## Example usage @@ -305,6 +435,11 @@ pub fn gt(lhs: Expression, rhs: Expression) -> Expression { .vortex_expect("Failed to create Gt binary expression") } +/// Creates a bound greater-than expression. +pub fn bound_gt(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { + bound_binary(Operator::Gt, lhs, rhs) +} + /// Create a new [`Binary`] using the [`Lte`](Operator::Lte) operator. /// /// ## Example usage @@ -332,6 +467,11 @@ pub fn lt_eq(lhs: Expression, rhs: Expression) -> Expression { .vortex_expect("Failed to create Lte binary expression") } +/// Creates a bound less-than-or-equal expression. +pub fn bound_lt_eq(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { + bound_binary(Operator::Lte, lhs, rhs) +} + /// Create a new [`Binary`] using the [`Lt`](Operator::Lt) operator. /// /// ## Example usage @@ -359,6 +499,11 @@ pub fn lt(lhs: Expression, rhs: Expression) -> Expression { .vortex_expect("Failed to create Lt binary expression") } +/// Creates a bound less-than expression. +pub fn bound_lt(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { + bound_binary(Operator::Lt, lhs, rhs) +} + /// Create a new [`Binary`] using the [`Or`](Operator::Or) operator. /// /// ## Example usage @@ -384,6 +529,11 @@ pub fn or(lhs: Expression, rhs: Expression) -> Expression { .vortex_expect("Failed to create Or binary expression") } +/// Creates a bound boolean OR expression. +pub fn bound_or(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { + bound_binary(Operator::Or, lhs, rhs) +} + /// Collects a list of `or`ed values into a single expression using a balanced tree. /// /// This creates a balanced binary tree to avoid deep nesting that could cause @@ -397,6 +547,14 @@ where iter.into_iter().reduce_balanced(or) } +/// Collects bound expressions into a balanced tree of boolean OR expressions. +pub fn bound_or_collect(iter: I) -> Option +where + I: IntoIterator, +{ + iter.into_iter().reduce_balanced(bound_or) +} + /// Create a new [`Binary`] using the [`And`](Operator::And) operator. /// /// ## Example usage @@ -422,6 +580,11 @@ pub fn and(lhs: Expression, rhs: Expression) -> Expression { .vortex_expect("Failed to create And binary expression") } +/// Creates a bound boolean AND expression. +pub fn bound_and(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { + bound_binary(Operator::And, lhs, rhs) +} + /// Collects a list of `and`ed values into a single expression using a balanced tree. /// /// This creates a balanced binary tree to avoid deep nesting that could cause @@ -435,6 +598,14 @@ where iter.into_iter().reduce_balanced(and) } +/// Collects bound expressions into a balanced tree of boolean AND expressions. +pub fn bound_and_collect(iter: I) -> Option +where + I: IntoIterator, +{ + iter.into_iter().reduce_balanced(bound_and) +} + /// The conjunction of an expression's child validities — i.e. the validity of a scalar function /// whose result is null exactly when any operand is null. /// @@ -475,6 +646,11 @@ pub fn checked_add(lhs: Expression, rhs: Expression) -> Expression { .vortex_expect("Failed to create Add binary expression") } +/// Creates a bound checked-add expression. +pub fn bound_checked_add(lhs: BoundExpression, rhs: BoundExpression) -> BoundExpression { + bound_binary(Operator::Add, lhs, rhs) +} + // ---- Not ---- /// Creates an expression that logically inverts boolean values. @@ -489,6 +665,12 @@ pub fn not(operand: Expression) -> Expression { Not.new_expr(EmptyOptions, vec![operand]) } +/// Creates a bound expression that logically inverts boolean values. +pub fn bound_not(operand: BoundExpression) -> BoundExpression { + Not.try_new_bound_expr(EmptyOptions, [operand]) + .vortex_expect("not expressions require a boolean operand") +} + // ---- Between ---- /// Creates an expression that checks if values are between two bounds. @@ -517,6 +699,18 @@ pub fn between( .vortex_expect("Failed to create Between expression") } +/// Creates a bound expression that checks if values are between two bounds. +pub fn bound_between( + arr: BoundExpression, + lower: BoundExpression, + upper: BoundExpression, + options: BetweenOptions, +) -> BoundExpression { + Between + .try_new_bound_expr(options, [arr, lower, upper]) + .vortex_expect("between expressions require compatible operand dtypes") +} + // ---- Select ---- /// Creates an expression that selects (includes) specific fields from an array. @@ -532,6 +726,13 @@ pub fn select(field_names: impl Into, child: Expression) -> Expressi .vortex_expect("Failed to create Select expression") } +/// Creates a bound expression that selects specific fields from a struct expression. +pub fn bound_select(field_names: impl Into, child: BoundExpression) -> BoundExpression { + Select + .try_new_bound_expr(FieldSelection::Include(field_names.into()), [child]) + .vortex_expect("select expressions require fields from a struct child") +} + /// Creates an expression that excludes specific fields from an array. /// /// Projects all fields except the specified ones from the input struct expression. @@ -546,6 +747,16 @@ pub fn select_exclude(fields: impl Into, child: Expression) -> Expre .vortex_expect("Failed to create Select expression") } +/// Creates a bound expression that excludes specific fields from a struct expression. +pub fn bound_select_exclude( + fields: impl Into, + child: BoundExpression, +) -> BoundExpression { + Select + .try_new_bound_expr(FieldSelection::Exclude(fields.into()), [child]) + .vortex_expect("select expressions require fields from a struct child") +} + // ---- Pack ---- /// Creates an expression that packs values into a struct with named fields. @@ -572,6 +783,25 @@ pub fn pack( ) } +/// Creates a bound expression that packs values into a struct with named fields. +pub fn bound_pack( + elements: impl IntoIterator, BoundExpression)>, + nullability: Nullability, +) -> BoundExpression { + let (names, values): (Vec<_>, Vec<_>) = elements + .into_iter() + .map(|(name, value)| (name.into(), value)) + .unzip(); + Pack.try_new_bound_expr( + PackOptions { + names: names.into(), + nullability, + }, + values, + ) + .vortex_expect("pack expressions must have one name per child") +} + // ---- Cast ---- /// Creates an expression that casts values to a target data type. @@ -588,6 +818,12 @@ pub fn cast(child: Expression, target: DType) -> Expression { .vortex_expect("Failed to create Cast expression") } +/// Creates a bound expression that casts values to a target dtype. +pub fn bound_cast(child: BoundExpression, target: DType) -> BoundExpression { + Cast.try_new_bound_expr(target, [child]) + .vortex_expect("cast expressions require a supported source and target dtype") +} + // ---- FillNull ---- /// Creates an expression that replaces null values with a fill value. @@ -600,6 +836,13 @@ pub fn fill_null(child: Expression, fill_value: Expression) -> Expression { FillNull.new_expr(EmptyOptions, [child, fill_value]) } +/// Creates a bound expression that replaces null values with a fill value. +pub fn bound_fill_null(child: BoundExpression, fill_value: BoundExpression) -> BoundExpression { + FillNull + .try_new_bound_expr(EmptyOptions, [child, fill_value]) + .vortex_expect("fill-null expressions require compatible child and fill dtypes") +} + // ---- IsNull ---- /// Creates an expression that checks for null values. @@ -614,6 +857,13 @@ pub fn is_null(child: Expression) -> Expression { IsNull.new_expr(EmptyOptions, vec![child]) } +/// Creates a bound expression that checks for null values. +pub fn bound_is_null(child: BoundExpression) -> BoundExpression { + IsNull + .try_new_bound_expr(EmptyOptions, [child]) + .vortex_expect("is-null expressions are always well-typed") +} + // ---- IsNotNull ---- /// Creates an expression that checks for non-null values. @@ -628,6 +878,13 @@ pub fn is_not_null(child: Expression) -> Expression { IsNotNull.new_expr(EmptyOptions, vec![child]) } +/// Creates a bound expression that checks for non-null values. +pub fn bound_is_not_null(child: BoundExpression) -> BoundExpression { + IsNotNull + .try_new_bound_expr(EmptyOptions, [child]) + .vortex_expect("is-not-null expressions are always well-typed") +} + // ---- Like ---- /// Creates a SQL LIKE expression. @@ -641,6 +898,11 @@ pub fn like(child: Expression, pattern: Expression) -> Expression { ) } +/// Creates a bound SQL LIKE expression. +pub fn bound_like(child: BoundExpression, pattern: BoundExpression) -> BoundExpression { + bound_like_with_options(child, pattern, false, false) +} + /// Creates a case-insensitive SQL ILIKE expression. pub fn ilike(child: Expression, pattern: Expression) -> Expression { Like.new_expr( @@ -652,6 +914,11 @@ pub fn ilike(child: Expression, pattern: Expression) -> Expression { ) } +/// Creates a bound case-insensitive SQL ILIKE expression. +pub fn bound_ilike(child: BoundExpression, pattern: BoundExpression) -> BoundExpression { + bound_like_with_options(child, pattern, false, true) +} + /// Creates a negated SQL NOT LIKE expression. pub fn not_like(child: Expression, pattern: Expression) -> Expression { Like.new_expr( @@ -663,6 +930,11 @@ pub fn not_like(child: Expression, pattern: Expression) -> Expression { ) } +/// Creates a bound negated SQL NOT LIKE expression. +pub fn bound_not_like(child: BoundExpression, pattern: BoundExpression) -> BoundExpression { + bound_like_with_options(child, pattern, true, false) +} + /// Creates a negated case-insensitive SQL NOT ILIKE expression. pub fn not_ilike(child: Expression, pattern: Expression) -> Expression { Like.new_expr( @@ -674,6 +946,27 @@ pub fn not_ilike(child: Expression, pattern: Expression) -> Expression { ) } +/// Creates a bound negated case-insensitive SQL NOT ILIKE expression. +pub fn bound_not_ilike(child: BoundExpression, pattern: BoundExpression) -> BoundExpression { + bound_like_with_options(child, pattern, true, true) +} + +fn bound_like_with_options( + child: BoundExpression, + pattern: BoundExpression, + negated: bool, + case_insensitive: bool, +) -> BoundExpression { + Like.try_new_bound_expr( + LikeOptions { + negated, + case_insensitive, + }, + [child, pattern], + ) + .vortex_expect("like expressions require UTF-8 or binary operands") +} + // ---- Mask ---- /// Creates a mask expression that applies the given boolean mask to the input array. @@ -681,6 +974,12 @@ pub fn mask(array: Expression, mask: Expression) -> Expression { Mask.new_expr(EmptyOptions, [array, mask]) } +/// Creates a bound mask expression. +pub fn bound_mask(array: BoundExpression, mask: BoundExpression) -> BoundExpression { + Mask.try_new_bound_expr(EmptyOptions, [array, mask]) + .vortex_expect("mask expressions require a boolean mask") +} + // ---- Merge ---- /// Creates an expression that merges struct expressions into a single struct. @@ -699,6 +998,11 @@ pub fn merge(elements: impl IntoIterator>) -> Expre Merge.new_expr(DuplicateHandling::default(), values) } +/// Creates a bound expression that merges struct expressions. +pub fn bound_merge(elements: impl IntoIterator) -> BoundExpression { + bound_merge_opts(elements, DuplicateHandling::default()) +} + /// Creates a merge expression with explicit duplicate handling. pub fn merge_opts( elements: impl IntoIterator>, @@ -709,6 +1013,16 @@ pub fn merge_opts( Merge.new_expr(duplicate_handling, values) } +/// Creates a bound merge expression with explicit duplicate handling. +pub fn bound_merge_opts( + elements: impl IntoIterator, + duplicate_handling: DuplicateHandling, +) -> BoundExpression { + Merge + .try_new_bound_expr(duplicate_handling, elements) + .vortex_expect("merge expressions require non-nullable struct children") +} + // ---- Zip ---- /// Creates a zip expression that conditionally selects between two arrays. @@ -721,8 +1035,33 @@ pub fn zip_expr(mask: Expression, if_true: Expression, if_false: Expression) -> Zip.new_expr(EmptyOptions, [if_true, if_false, mask]) } +/// Creates a bound zip expression that conditionally selects between two arrays. +pub fn bound_zip_expr( + mask: BoundExpression, + if_true: BoundExpression, + if_false: BoundExpression, +) -> BoundExpression { + Zip.try_new_bound_expr(EmptyOptions, [if_true, if_false, mask]) + .vortex_expect("zip expressions require a boolean mask and compatible value dtypes") +} + // ---- Dynamic ---- +/// Creates a dynamic comparison expression from its complete options. +pub fn dynamic_with_options(options: DynamicComparisonExpr, lhs: Expression) -> Expression { + DynamicComparison.new_expr(options, [lhs]) +} + +/// Creates a bound dynamic comparison expression from its complete options. +pub fn bound_dynamic_with_options( + options: DynamicComparisonExpr, + lhs: BoundExpression, +) -> BoundExpression { + DynamicComparison + .try_new_bound_expr(options, [lhs]) + .vortex_expect("dynamic comparisons require a compatible left-hand dtype") +} + /// Creates a dynamic comparison expression. pub fn dynamic( operator: CompareOperator, @@ -731,7 +1070,7 @@ pub fn dynamic( default: bool, lhs: Expression, ) -> Expression { - DynamicComparison.new_expr( + dynamic_with_options( DynamicComparisonExpr { operator, rhs: Arc::new(Rhs { @@ -740,7 +1079,28 @@ pub fn dynamic( }), default, }, - [lhs], + lhs, + ) +} + +/// Creates a bound dynamic comparison expression. +pub fn bound_dynamic( + operator: CompareOperator, + rhs_value: impl Fn() -> Option + Send + Sync + 'static, + rhs_dtype: DType, + default: bool, + lhs: BoundExpression, +) -> BoundExpression { + bound_dynamic_with_options( + DynamicComparisonExpr { + operator, + rhs: Arc::new(Rhs { + value: Arc::new(rhs_value), + dtype: rhs_dtype, + }), + default, + }, + lhs, ) } @@ -758,6 +1118,13 @@ pub fn list_contains(list: Expression, value: Expression) -> Expression { ListContains.new_expr(EmptyOptions, [list, value]) } +/// Creates a bound expression that checks if a value is contained in a list. +pub fn bound_list_contains(list: BoundExpression, value: BoundExpression) -> BoundExpression { + ListContains + .try_new_bound_expr(EmptyOptions, [list, value]) + .vortex_expect("list-contains expressions require a compatible list and value dtype") +} + // ---- ByteLength ---- /// Creates an expression that computes the byte length of each element. @@ -771,6 +1138,13 @@ pub fn byte_length(input: Expression) -> Expression { ByteLength.new_expr(EmptyOptions, [input]) } +/// Creates a bound expression that computes each element's byte length. +pub fn bound_byte_length(input: BoundExpression) -> BoundExpression { + ByteLength + .try_new_bound_expr(EmptyOptions, [input]) + .vortex_expect("byte-length expressions require a variable-length binary child") +} + // ---- ExtStorage ---- /// Creates an expression that extracts the storage values from an extension array. @@ -783,6 +1157,13 @@ pub fn ext_storage(input: Expression) -> Expression { ExtStorage.new_expr(EmptyOptions, [input]) } +/// Creates a bound expression that extracts an extension array's storage values. +pub fn bound_ext_storage(input: BoundExpression) -> BoundExpression { + ExtStorage + .try_new_bound_expr(EmptyOptions, [input]) + .vortex_expect("extension-storage expressions require an extension child") +} + // ---- ListLength ---- /// Creates an expression that computes the number of elements in each list @@ -797,6 +1178,13 @@ pub fn list_length(input: Expression) -> Expression { ListLength.new_expr(EmptyOptions, [input]) } +/// Creates a bound expression that computes the number of elements in each list. +pub fn bound_list_length(input: BoundExpression) -> BoundExpression { + ListLength + .try_new_bound_expr(EmptyOptions, [input]) + .vortex_expect("list-length expressions require a list child") +} + // ---- ListSum ---- /// Creates an expression that sums the elements of each list for `List` and @@ -815,8 +1203,78 @@ pub fn list_sum(input: Expression) -> Expression { ListSum.new_expr(NumericalAggregateOpts::default(), [input]) } +/// Creates a bound expression that sums the elements of each list. +pub fn bound_list_sum(input: BoundExpression) -> BoundExpression { + ListSum + .try_new_bound_expr(NumericalAggregateOpts::default(), [input]) + .vortex_expect("list-sum expressions require a numeric list child") +} + /// Creates a [`list_sum`] expression with explicit [`NumericalAggregateOpts`], controlling /// whether NaN float elements are skipped (the default) or poison the list's sum to NaN. pub fn list_sum_opts(input: Expression, options: NumericalAggregateOpts) -> Expression { ListSum.new_expr(options, [input]) } + +/// Creates a bound list-sum expression with explicit aggregate options. +pub fn bound_list_sum_opts( + input: BoundExpression, + options: NumericalAggregateOpts, +) -> BoundExpression { + ListSum + .try_new_bound_expr(options, [input]) + .vortex_expect("list-sum expressions require a numeric list child") +} + +/// Constructors for expressions whose children have already been bound and type-checked. +/// +/// These mirror the constructors in [`crate::expr`] and panic when the supplied children do not +/// form a well-typed expression. Use [`BoundExpression::try_new`] when construction must be +/// fallible. +pub mod bound { + pub use super::bound_and as and; + pub use super::bound_and_collect as and_collect; + pub use super::bound_between as between; + pub use super::bound_binary as binary; + pub use super::bound_byte_length as byte_length; + pub use super::bound_case_when as case_when; + pub use super::bound_case_when_no_else as case_when_no_else; + pub use super::bound_cast as cast; + pub use super::bound_checked_add as checked_add; + pub use super::bound_col as col; + pub use super::bound_dynamic as dynamic; + pub use super::bound_dynamic_with_options as dynamic_with_options; + pub use super::bound_eq as eq; + pub use super::bound_ext_storage as ext_storage; + pub use super::bound_fill_null as fill_null; + pub use super::bound_get_item as get_item; + pub use super::bound_gt as gt; + pub use super::bound_gt_eq as gt_eq; + pub use super::bound_ilike as ilike; + pub use super::bound_is_not_null as is_not_null; + pub use super::bound_is_null as is_null; + pub use super::bound_like as like; + pub use super::bound_list_contains as list_contains; + pub use super::bound_list_length as list_length; + pub use super::bound_list_sum as list_sum; + pub use super::bound_list_sum_opts as list_sum_opts; + pub use super::bound_lit as lit; + pub use super::bound_lt as lt; + pub use super::bound_lt_eq as lt_eq; + pub use super::bound_mask as mask; + pub use super::bound_merge as merge; + pub use super::bound_merge_opts as merge_opts; + pub use super::bound_nested_case_when as nested_case_when; + pub use super::bound_not as not; + pub use super::bound_not_eq as not_eq; + pub use super::bound_not_ilike as not_ilike; + pub use super::bound_not_like as not_like; + pub use super::bound_or as or; + pub use super::bound_or_collect as or_collect; + pub use super::bound_pack as pack; + pub use super::bound_root as root; + pub use super::bound_select as select; + pub use super::bound_select_exclude as select_exclude; + pub use super::bound_variant_get as variant_get; + pub use super::bound_zip_expr as zip_expr; +} diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index 2ae61509ceb..0ddd003e65d 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -72,7 +72,54 @@ pub mod traversal; pub use analysis::*; pub use bound_expression::*; pub use expression::*; -pub use exprs::*; +pub use exprs::and; +pub use exprs::and_collect; +pub use exprs::between; +pub use exprs::binary; +pub use exprs::bound; +pub use exprs::byte_length; +pub use exprs::case_when; +pub use exprs::case_when_no_else; +pub use exprs::cast; +pub use exprs::checked_add; +pub use exprs::col; +pub use exprs::dynamic; +pub use exprs::dynamic_with_options; +pub use exprs::eq; +pub use exprs::ext_storage; +pub use exprs::fill_null; +pub use exprs::get_item; +pub use exprs::gt; +pub use exprs::gt_eq; +pub use exprs::ilike; +pub use exprs::is_not_null; +pub use exprs::is_null; +pub use exprs::is_root; +pub use exprs::like; +pub use exprs::list_contains; +pub use exprs::list_length; +pub use exprs::list_sum; +pub use exprs::list_sum_opts; +pub use exprs::lit; +pub use exprs::lt; +pub use exprs::lt_eq; +pub use exprs::mask; +pub use exprs::merge; +pub use exprs::merge_opts; +pub use exprs::nested_case_when; +pub use exprs::not; +pub use exprs::not_eq; +pub use exprs::not_ilike; +pub use exprs::not_like; +pub use exprs::or; +pub use exprs::or_collect; +pub use exprs::pack; +pub use exprs::root; +pub use exprs::select; +pub use exprs::select_exclude; +pub use exprs::union_child_validities; +pub use exprs::variant_get; +pub use exprs::zip_expr; pub use scope::*; pub trait VortexExprExt { @@ -167,6 +214,8 @@ mod tests { use crate::dtype::PType; use crate::dtype::StructFields; use crate::expr::and; + use crate::expr::bound; + use crate::expr::case_when; use crate::expr::col; use crate::expr::get_item; use crate::expr::gt; @@ -215,6 +264,46 @@ mod tests { assert_ne!(a, rebuilt); } + #[test] + fn bound_constructors_preserve_order_and_types() -> vortex_error::VortexResult<()> { + let value_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let scope = DType::Struct( + StructFields::from_iter([("value", value_dtype.clone())]), + Nullability::NonNullable, + ); + + let root = bound::root(scope.clone()); + let value = bound::get_item("value", root); + let literal = bound::lit(5i32); + let condition = bound::gt(value.clone(), literal.clone()); + assert_eq!(condition.dtype(), &DType::Bool(Nullability::NonNullable)); + assert_eq!(condition.children(), &[value.clone(), literal.clone()]); + + let case = bound::case_when(condition.clone(), value.clone(), literal.clone()); + assert_eq!(case.dtype(), &value_dtype); + assert_eq!(case.children(), &[condition.clone(), value, literal]); + + let packed = bound::pack( + [("condition", condition.clone()), ("value", case.clone())], + Nullability::NonNullable, + ); + assert_eq!(packed.children(), &[condition, case.clone()]); + assert_eq!( + packed.dtype(), + &DType::Struct( + StructFields::from_iter([ + ("condition", DType::Bool(Nullability::NonNullable)), + ("value", value_dtype), + ]), + Nullability::NonNullable, + ) + ); + + let unbound = case_when(gt(col("value"), lit(5i32)), col("value"), lit(5i32)); + assert_eq!(unbound.bind(&scope)?, case); + Ok(()) + } + #[test] fn expr_display() { assert_eq!(col("a").to_string(), "$.a"); diff --git a/vortex-array/src/expr/transform/bound_partition.rs b/vortex-array/src/expr/transform/bound_partition.rs index 95fe3f4a7ea..3272be862bc 100644 --- a/vortex-array/src/expr/transform/bound_partition.rs +++ b/vortex-array/src/expr/transform/bound_partition.rs @@ -22,14 +22,12 @@ use crate::expr::analysis::Annotation; use crate::expr::analysis::AnnotationFn; use crate::expr::analysis::BoundAnnotations; use crate::expr::analysis::descendent_bound_annotations; +use crate::expr::bound::get_item; +use crate::expr::bound::pack; use crate::expr::traversal::NodeExt; use crate::expr::traversal::NodeRewriter; use crate::expr::traversal::Transformed; use crate::expr::traversal::TraversalOrder; -use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::fns::get_item::GetItem; -use crate::scalar_fn::fns::pack::Pack; -use crate::scalar_fn::fns::pack::PackOptions; /// Partition an expression into sub-expressions that are uniquely associated with an annotation. /// A root expression is also returned that can be used to recombine the results of the partitions @@ -75,12 +73,12 @@ where for (annotation, exprs) in collector.sub_expressions { // We pack all sub-expressions for the same annotation into a single expression. - let names = exprs + let names: FieldNames = exprs .iter() .enumerate() .map(|(idx, _)| PartitionCollector::field_name(&annotation, idx)) .collect(); - let expr = bound_pack(names, exprs)?; + let expr = pack(names.into_iter().zip(exprs), Nullability::NonNullable); partitions.push(expr); partition_annotations.push(annotation); @@ -259,11 +257,11 @@ where let field_name = PartitionCollector::field_name(annotation, *offset); *offset += 1; - let partition = bound_get_item( + let partition = get_item( FieldName::from(annotation.clone()), BoundExpression::new_root(self.root_dtype.clone()), - )?; - let value = bound_get_item(field_name, partition)?; + ); + let value = get_item(field_name, partition); Ok(Transformed { value, @@ -273,20 +271,6 @@ where } } -fn bound_get_item(field_name: FieldName, child: BoundExpression) -> VortexResult { - BoundExpression::try_new(GetItem.bind(field_name), [child]) -} - -fn bound_pack(names: FieldNames, children: Vec) -> VortexResult { - BoundExpression::try_new( - Pack.bind(PackOptions { - names, - nullability: Nullability::NonNullable, - }), - children, - ) -} - fn partition_root_dtype(names: &FieldNames, partitions: &[BoundExpression]) -> DType { DType::Struct( StructFields::new( @@ -372,7 +356,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 +370,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 +394,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 +431,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 +445,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/fns/dynamic.rs b/vortex-array/src/scalar_fn/fns/dynamic.rs index 00e59796e05..be910629225 100644 --- a/vortex-array/src/scalar_fn/fns/dynamic.rs +++ b/vortex-array/src/scalar_fn/fns/dynamic.rs @@ -19,7 +19,7 @@ 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; @@ -205,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/is_not_null.rs b/vortex-array/src/scalar_fn/fns/is_not_null.rs index dcf00f62c08..881bc32c306 100644 --- a/vortex-array/src/scalar_fn/fns/is_not_null.rs +++ b/vortex-array/src/scalar_fn/fns/is_not_null.rs @@ -253,13 +253,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/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/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index 55c897ce1a1..5d3561ff039 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -19,6 +19,7 @@ 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::display::ExprDisplay; use crate::scalar_fn::ScalarFnId; @@ -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..e07a588de94 100644 --- a/vortex-array/src/stats/bind.rs +++ b/vortex-array/src/stats/bind.rs @@ -17,8 +17,8 @@ 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::bound::lit; use crate::expr::traversal::NodeExt; use crate::expr::traversal::Transformed; use crate::scalar::Scalar; @@ -31,26 +31,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 +57,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 { + Ok(lit(Scalar::null(dtype.as_nullable()))) } #[cfg(test)] @@ -111,6 +102,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 +111,7 @@ mod tests { struct TestBinder { input_scope: DType, + stats_scope: DType, bind_nan_count: bool, } @@ -132,28 +125,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 +162,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 +172,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 +186,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 +209,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/expr.rs b/vortex-array/src/stats/expr.rs index 2038df0b613..1e0ceef02d3 100644 --- a/vortex-array/src/stats/expr.rs +++ b/vortex-array/src/stats/expr.rs @@ -3,6 +3,8 @@ //! Expression constructors for statistics backed by aggregate functions. +use vortex_error::VortexExpect; + use crate::aggregate_fn::AggregateFnRef; use crate::aggregate_fn::AggregateFnVTableExt; use crate::aggregate_fn::EmptyOptions; @@ -15,6 +17,7 @@ use crate::aggregate_fn::fns::min_max::MinMax; use crate::aggregate_fn::fns::nan_count::NanCount; use crate::aggregate_fn::fns::null_count::NullCount; use crate::aggregate_fn::fns::sum::Sum; +use crate::expr::BoundExpression; use crate::expr::Expression; use crate::scalar_fn::ScalarFnVTableExt; pub use crate::scalar_fn::fns::stat::StatFn; @@ -28,48 +31,140 @@ pub fn stat(expr: Expression, aggregate_fn: AggregateFnRef) -> Expression { StatFn.new_expr(StatOptions::new(aggregate_fn), [expr]) } +fn bound_stat(expr: BoundExpression, aggregate_fn: AggregateFnRef) -> BoundExpression { + StatFn + .try_new_bound_expr(StatOptions::new(aggregate_fn), [expr]) + .vortex_expect("stat expressions must use an aggregate supported by the child dtype") +} + /// Creates `stat(expr, min_max)`, returning a nullable `{ min, max }` struct statistic. pub fn min_max(expr: Expression) -> Expression { // Statistics follow NaN-skipping semantics; request it explicitly rather than via the default. stat(expr, MinMax.bind(NumericalAggregateOpts::skip_nans())) } +fn bound_min_max(expr: BoundExpression) -> BoundExpression { + bound_stat(expr, MinMax.bind(NumericalAggregateOpts::skip_nans())) +} + /// Creates `stat(expr, sum)`, returning a nullable sum statistic. pub fn sum(expr: Expression) -> Expression { // Statistics follow NaN-skipping semantics; request it explicitly rather than via the default. stat(expr, Sum.bind(NumericalAggregateOpts::skip_nans())) } +fn bound_sum(expr: BoundExpression) -> BoundExpression { + bound_stat(expr, Sum.bind(NumericalAggregateOpts::skip_nans())) +} + /// Creates `stat(expr, null_count)`, returning a nullable null-count statistic. pub fn null_count(expr: Expression) -> Expression { stat(expr, NullCount.bind(EmptyOptions)) } +fn bound_null_count(expr: BoundExpression) -> BoundExpression { + bound_stat(expr, NullCount.bind(EmptyOptions)) +} + /// Creates `stat(expr, all_null)`, returning a nullable all-null statistic. pub fn all_null(expr: Expression) -> Expression { stat(expr, AllNull.bind(EmptyOptions)) } +fn bound_all_null(expr: BoundExpression) -> BoundExpression { + bound_stat(expr, AllNull.bind(EmptyOptions)) +} + /// Creates `stat(expr, all_nan)`, returning a nullable all-NaN statistic. pub fn all_nan(expr: Expression) -> Expression { stat(expr, AllNan.bind(EmptyOptions)) } +fn bound_all_nan(expr: BoundExpression) -> BoundExpression { + bound_stat(expr, AllNan.bind(EmptyOptions)) +} + /// Creates `stat(expr, all_non_null)`, returning a nullable all-non-null statistic. pub fn all_non_null(expr: Expression) -> Expression { stat(expr, AllNonNull.bind(EmptyOptions)) } +fn bound_all_non_null(expr: BoundExpression) -> BoundExpression { + bound_stat(expr, AllNonNull.bind(EmptyOptions)) +} + /// Creates `stat(expr, all_non_nan)`, returning a nullable all-non-NaN statistic. pub fn all_non_nan(expr: Expression) -> Expression { stat(expr, AllNonNan.bind(EmptyOptions)) } +fn bound_all_non_nan(expr: BoundExpression) -> BoundExpression { + bound_stat(expr, AllNonNan.bind(EmptyOptions)) +} + /// Creates `stat(expr, nan_count)`, returning a nullable NaN-count statistic. pub fn nan_count(expr: Expression) -> Expression { stat(expr, NanCount.bind(EmptyOptions)) } +fn bound_nan_count(expr: BoundExpression) -> BoundExpression { + bound_stat(expr, NanCount.bind(EmptyOptions)) +} + +/// Constructors for statistic expressions whose input has already been bound. +/// +/// These mirror the constructors in [`crate::stats`] and panic when the aggregate does not support +/// the input dtype. +pub mod bound { + use crate::aggregate_fn::AggregateFnRef; + use crate::expr::BoundExpression; + + /// Creates a bound expression that reads a stored aggregate statistic. + pub fn stat(expr: BoundExpression, aggregate_fn: AggregateFnRef) -> BoundExpression { + super::bound_stat(expr, aggregate_fn) + } + + /// Creates a bound nullable `{ min, max }` statistic expression. + pub fn min_max(expr: BoundExpression) -> BoundExpression { + super::bound_min_max(expr) + } + + /// Creates a bound nullable sum statistic expression. + pub fn sum(expr: BoundExpression) -> BoundExpression { + super::bound_sum(expr) + } + + /// Creates a bound nullable null-count statistic expression. + pub fn null_count(expr: BoundExpression) -> BoundExpression { + super::bound_null_count(expr) + } + + /// Creates a bound nullable all-null statistic expression. + pub fn all_null(expr: BoundExpression) -> BoundExpression { + super::bound_all_null(expr) + } + + /// Creates a bound nullable all-NaN statistic expression. + pub fn all_nan(expr: BoundExpression) -> BoundExpression { + super::bound_all_nan(expr) + } + + /// Creates a bound nullable all-non-null statistic expression. + pub fn all_non_null(expr: BoundExpression) -> BoundExpression { + super::bound_all_non_null(expr) + } + + /// Creates a bound nullable all-non-NaN statistic expression. + pub fn all_non_nan(expr: BoundExpression) -> BoundExpression { + super::bound_all_non_nan(expr) + } + + /// Creates a bound nullable NaN-count statistic expression. + pub fn nan_count(expr: BoundExpression) -> BoundExpression { + super::bound_nan_count(expr) + } +} + #[cfg(test)] mod tests { use std::sync::LazyLock; @@ -83,6 +178,7 @@ mod tests { use super::all_non_nan; use super::all_non_null; use super::all_null; + use super::bound as bound_stats; use super::null_count; use super::stat; use super::sum; @@ -99,6 +195,7 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::expr::bound as bound_expr; use crate::expr::root; use crate::expr::stats::Precision; use crate::expr::stats::Stat; @@ -108,6 +205,21 @@ mod tests { static SESSION: LazyLock = LazyLock::new(array_session); + #[test] + fn bound_stats_constructor_preserves_child_and_dtype() -> VortexResult<()> { + let input_dtype = DType::Primitive(PType::I32, Nullability::NonNullable); + let root = bound_expr::root(input_dtype.clone()); + let bound = bound_stats::sum(root.clone()); + + assert_eq!(bound.children(), &[root]); + assert_eq!( + bound.dtype(), + &DType::Primitive(PType::I64, Nullability::Nullable) + ); + assert_eq!(bound, sum(crate::expr::root()).bind(&input_dtype)?); + Ok(()) + } + #[test] fn stat_expr_reads_cached_sum() -> VortexResult<()> { let array = buffer![1i32, 2, 3].into_array(); diff --git a/vortex-array/src/stats/mod.rs b/vortex-array/src/stats/mod.rs index 5f5684dbde2..f7610fc71e4 100644 --- a/vortex-array/src/stats/mod.rs +++ b/vortex-array/src/stats/mod.rs @@ -11,6 +11,7 @@ pub use expr::all_nan; pub use expr::all_non_nan; pub use expr::all_non_null; pub use expr::all_null; +pub use expr::bound; pub use expr::min_max; pub use expr::nan_count; pub use expr::null_count; 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..571c8b5ff84 100644 --- a/vortex-array/src/stats/rewrite/builtins.rs +++ b/vortex-array/src/stats/rewrite/builtins.rs @@ -3,6 +3,7 @@ use std::sync::Arc; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::aggregate_fn::AggregateFnRef; @@ -12,18 +13,20 @@ 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; +use crate::expr::bound::and; +use crate::expr::bound::and_collect; +use crate::expr::bound::binary; +use crate::expr::bound::cast; +use crate::expr::bound::dynamic_with_options; +use crate::expr::bound::eq; +use crate::expr::bound::gt; +use crate::expr::bound::gt_eq; +use crate::expr::bound::lit; +use crate::expr::bound::lt; +use crate::expr::bound::lt_eq; +use crate::expr::bound::or; +use crate::expr::bound::or_collect; use crate::expr::stats::Stat; use crate::scalar::StringLike; use crate::scalar_fn::EmptyOptions; @@ -44,8 +47,7 @@ use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::fns::operators::Operator; use crate::scalar_fn::internal::row_count::RowCount; -use crate::stats::expr::StatFn; -use crate::stats::expr::StatOptions; +use crate::stats::bound::stat; use crate::stats::rewrite::StatsRewriteCtx; use crate::stats::rewrite::StatsRewriteRule; use crate::stats::session::StatsSession; @@ -68,6 +70,12 @@ pub(crate) fn register_builtins(session: &StatsSession) { session.register_rewrite(DynamicComparisonAllNonNanStatsRewrite); } +fn row_count() -> BoundExpression { + RowCount + .try_new_bound_expr(EmptyOptions, []) + .vortex_expect("row-count expressions are always well-typed") +} + #[derive(Debug)] struct BinaryNanCountStatsRewrite; @@ -78,9 +86,9 @@ impl StatsRewriteRule for BinaryNanCountStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { binary_falsify::(expr, ctx) } } @@ -95,17 +103,17 @@ impl StatsRewriteRule for BinaryAllNonNanStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { binary_falsify::(expr, ctx) } } fn binary_falsify( - expr: &Expression, + expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>, -) -> VortexResult> { +) -> VortexResult> { let operator = expr.as_::(); let lhs = expr.child(0); let rhs = expr.child(1); @@ -178,16 +186,16 @@ impl StatsRewriteRule for BetweenStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpression, 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 +210,18 @@ impl StatsRewriteRule for IsNullNullCountStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpression, 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: &BoundExpression, 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 +235,9 @@ impl StatsRewriteRule for IsNullAllNonNullStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpression, _ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(Some(all_non_null(expr.child(0)))) } } @@ -245,9 +252,9 @@ impl StatsRewriteRule for IsNullAllNullStatsRewrite { fn satisfy( &self, - expr: &Expression, + expr: &BoundExpression, _ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(Some(all_null(expr.child(0)))) } } @@ -262,18 +269,17 @@ impl StatsRewriteRule for IsNotNullNullCountStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpression, 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: &BoundExpression, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(null_count(expr.child(0), ctx).map(|null_count| eq(null_count, lit(0u64)))) } } @@ -288,9 +294,9 @@ impl StatsRewriteRule for IsNotNullAllNullStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpression, _ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(Some(all_null(expr.child(0)))) } } @@ -305,9 +311,9 @@ impl StatsRewriteRule for IsNotNullAllNonNullStatsRewrite { fn satisfy( &self, - expr: &Expression, + expr: &BoundExpression, _ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { Ok(Some(all_non_null(expr.child(0)))) } } @@ -322,9 +328,9 @@ impl StatsRewriteRule for LikeStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { let like_options = expr.as_::(); if like_options.negated || like_options.case_insensitive { return Ok(None); @@ -377,9 +383,9 @@ impl StatsRewriteRule for ListContainsNanCountStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { list_contains_falsify::(expr, ctx) } } @@ -394,17 +400,17 @@ impl StatsRewriteRule for ListContainsAllNonNanStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { list_contains_falsify::(expr, ctx) } } fn list_contains_falsify( - expr: &Expression, + expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>, -) -> VortexResult> { +) -> VortexResult> { let list = expr.child(0); let needle = expr.child(1); @@ -451,9 +457,9 @@ impl StatsRewriteRule for DynamicComparisonNanCountStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { dynamic_comparison_falsify::(expr, ctx) } } @@ -468,17 +474,17 @@ impl StatsRewriteRule for DynamicComparisonAllNonNanStatsRewrite { fn falsify( &self, - expr: &Expression, + expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>, - ) -> VortexResult> { + ) -> VortexResult> { dynamic_comparison_falsify::(expr, ctx) } } fn dynamic_comparison_falsify( - expr: &Expression, + expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>, -) -> VortexResult> { +) -> VortexResult> { let dynamic = expr.as_::(); let lhs = expr.child(0); @@ -492,47 +498,47 @@ fn dynamic_comparison_falsify( return Ok(None); }; - let value_predicate = DynamicComparison.new_expr( + let value_predicate = dynamic_with_options( DynamicComparisonExpr { operator, rhs: Arc::clone(&dynamic.rhs), default: !dynamic.default, }, - [lhs_stat], + lhs_stat, ); with_non_nan_guards::

(ctx, [lhs], value_predicate) } -fn min(expr: &Expression, ctx: &StatsRewriteCtx<'_>) -> Option { +fn min(expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>) -> Option { stat_expr(expr, Stat::Min, ctx) } -fn max(expr: &Expression, ctx: &StatsRewriteCtx<'_>) -> Option { +fn max(expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>) -> Option { stat_expr(expr, Stat::Max, ctx) } -fn null_count(expr: &Expression, ctx: &StatsRewriteCtx<'_>) -> Option { +fn null_count(expr: &BoundExpression, ctx: &StatsRewriteCtx<'_>) -> Option { stat_expr(expr, Stat::NullCount, ctx) } -fn all_null(expr: &Expression) -> Expression { +fn all_null(expr: &BoundExpression) -> BoundExpression { stat_fn(expr.clone(), AllNull.bind(AggregateEmptyOptions)) } -fn all_non_null(expr: &Expression) -> Expression { +fn all_non_null(expr: &BoundExpression) -> BoundExpression { stat_fn(expr.clone(), AllNonNull.bind(AggregateEmptyOptions)) } enum NanCheck { NotNeeded, - Check(Expression), + Check(BoundExpression), Unavailable, } trait NonNanProof { const EMIT_UNGUARDED_REWRITES: bool; - fn check(ctx: &StatsRewriteCtx<'_>, expr: &Expression) -> VortexResult; + fn check(ctx: &StatsRewriteCtx<'_>, expr: &BoundExpression) -> VortexResult; } struct NanCountProof; @@ -540,7 +546,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: &BoundExpression) -> 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 +561,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: &BoundExpression) -> VortexResult { non_nan_check(ctx, expr, |expr| { NanCheck::Check(stat_fn(expr.clone(), AllNonNan.bind(AggregateEmptyOptions))) }) @@ -567,8 +573,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: &BoundExpression, + proof: impl FnOnce(&BoundExpression) -> NanCheck, ) -> VortexResult { if let Some(scalar) = expr.as_opt::() { if !scalar.dtype().is_float() { @@ -600,7 +606,11 @@ fn has_nans(dtype: &DType) -> bool { dtype.is_float() } -fn stat_expr(expr: &Expression, stat: Stat, ctx: &StatsRewriteCtx<'_>) -> Option { +fn stat_expr( + expr: &BoundExpression, + stat: Stat, + ctx: &StatsRewriteCtx<'_>, +) -> Option { if let Some(literal) = literal_stat(expr, stat) { return Some(literal); } @@ -629,9 +639,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: BoundExpression, +) -> VortexResult> { let mut nan_checks = Vec::new(); for expr in exprs { match P::check(ctx, expr)? { @@ -652,7 +662,7 @@ fn with_non_nan_guards<'a, P: NonNanProof>( }) } -fn literal_stat(expr: &Expression, stat: Stat) -> Option { +fn literal_stat(expr: &BoundExpression, stat: Stat) -> Option { let scalar = expr.as_opt::()?; match stat { Stat::Min | Stat::Max => Some(lit(scalar.clone())), @@ -674,11 +684,11 @@ fn literal_stat(expr: &Expression, stat: Stat) -> Option { } fn cast_stat( - expr: &Expression, + expr: &BoundExpression, 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 +696,8 @@ fn cast_stat( } } -fn stat_fn(expr: Expression, aggregate_fn: AggregateFnRef) -> Expression { - StatFn.new_expr(StatOptions::new(aggregate_fn), [expr]) +fn stat_fn(expr: BoundExpression, aggregate_fn: AggregateFnRef) -> BoundExpression { + stat(expr, aggregate_fn) } #[cfg(test)] @@ -698,10 +708,6 @@ mod tests { use vortex_error::VortexResult; use vortex_session::VortexSession; - 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,6 +716,7 @@ mod tests { use crate::dtype::Nullability; use crate::dtype::PType; use crate::dtype::StructFields; + use crate::expr::BoundExpression; use crate::expr::Expression; use crate::expr::and; use crate::expr::between; @@ -737,6 +744,8 @@ mod tests { use crate::scalar_fn::fns::dynamic::DynamicComparisonExpr; use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::internal::row_count::RowCount; + use crate::stats::expr::StatFn; + use crate::stats::expr::StatOptions; static SESSION: LazyLock = LazyLock::new(crate::array_session); @@ -770,12 +779,30 @@ mod tests { ) } - fn falsify(expr: &Expression) -> VortexResult> { - expr.falsify(&test_scope(), &SESSION) + fn falsify(expr: &Expression) -> VortexResult> { + expr.bind(&test_scope())?.falsify(&SESSION) } - fn satisfy(expr: &Expression) -> VortexResult> { - expr.satisfy(&test_scope(), &SESSION) + fn satisfy(expr: &Expression) -> 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: &Expression) -> Expression { + crate::stats::all_null(expr.clone()) + } + + fn all_non_null(expr: &Expression) -> Expression { + 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 { @@ -794,13 +821,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 +836,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 +849,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 +871,7 @@ mod tests { }, ); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(or( gt(lit(10), stat(col("a"), Stat::Max)), @@ -856,7 +883,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 +891,7 @@ mod tests { )) ); - assert_eq!( + assert_rewrite_eq!( falsify(&is_not_null(col("a")))?, Some(or( eq( @@ -879,7 +906,7 @@ mod tests { #[test] fn rewrites_null_satisfiers() -> VortexResult<()> { - assert_eq!( + assert_rewrite_eq!( satisfy(&is_null(col("a")))?, Some(or( eq( @@ -890,7 +917,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 +936,7 @@ mod tests { ); let expr = list_contains(lit(list), col("a")); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(and( and( @@ -934,7 +961,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 +970,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 +979,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 +988,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 +997,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 +1006,7 @@ mod tests { ); let expr = like(col("s"), lit("%suffix")); - assert_eq!(falsify(&expr)?, None); + assert_rewrite_eq!(falsify(&expr)?, None); Ok(()) } @@ -994,7 +1021,7 @@ mod tests { ); let dynamic = expr.as_::(); - assert_eq!( + assert_rewrite_eq!( falsify(&expr)?, Some(DynamicComparison.new_expr( DynamicComparisonExpr { @@ -1013,7 +1040,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 +1058,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 +1068,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-bench/src/datasets/tpch_l_comment.rs b/vortex-bench/src/datasets/tpch_l_comment.rs index 55d8497c2b3..c57bc91a65d 100644 --- a/vortex-bench/src/datasets/tpch_l_comment.rs +++ b/vortex-bench/src/datasets/tpch_l_comment.rs @@ -66,9 +66,12 @@ impl Dataset for TPCHLCommentChunked { let path = data_dir.join("lineitem.vortex"); let file = SESSION.open_options().open_path(path).await?; + let projection = pack(vec![("l_comment", col("l_comment"))], NonNullable) + .optimize_recursive(file.dtype())? + .bind(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..f70de847966 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -302,9 +302,13 @@ 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 = scan_projection + .optimize_recursive(vxf.dtype()) + .and_then(|projection| projection.bind(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 +423,10 @@ impl FileOpener for VortexOpener { make_vortex_predicate(expr_convertor.as_ref(), &pushed).transpose() }) .transpose()?; + let filter = filter + .map(|filter| filter.optimize_recursive(vxf.dtype())?.bind(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/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..df327638d00 100644 --- a/vortex-file/src/pruning.rs +++ b/vortex-file/src/pruning.rs @@ -10,9 +10,8 @@ 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::bound::lit; use vortex_array::expr::stats::Stat; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::fns::cast::Cast; @@ -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; @@ -111,8 +102,8 @@ impl FileStatsBinder<'_> { } } -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/tests.rs b/vortex-file/src/tests.rs index b11c84e55aa..1c66c8b2f7c 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -38,6 +38,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; @@ -108,6 +110,12 @@ static SESSION: LazyLock = LazyLock::new(|| { session }); +fn bind_scan_expr(file: &VortexFile, expr: Expression) -> BoundExpression { + expr.optimize_recursive(file.dtype()) + .and_then(|expr| expr.bind(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 @@ -315,7 +323,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() @@ -341,7 +349,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() @@ -512,18 +520,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() @@ -564,13 +573,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() @@ -624,17 +634,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() @@ -692,15 +703,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() @@ -906,7 +918,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() @@ -924,7 +939,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() @@ -952,7 +970,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() @@ -1014,7 +1035,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() @@ -1104,9 +1128,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() @@ -1179,7 +1206,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() @@ -1293,7 +1323,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 @@ -1356,15 +1386,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 = Pack + .new_expr( PackOptions { names: Default::default(), nullability: Nullability::Nullable, }, [], - ))) + ) + .optimize_recursive(array.dtype())? + .bind(array.dtype())?; + + let result = round_trip(&array.clone().into_array(), |scan| { + Ok(scan.with_projection(projection)) }) .await?; @@ -2188,13 +2222,11 @@ 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 = filter_expr + .optimize_recursive(file.dtype())? + .bind(file.dtype())?; + let mut stream = file.scan()?.with_filter(filter).into_array_stream()?; let result = stream.try_next().await; assert!(result.is_err()); @@ -2241,13 +2273,11 @@ 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 = filter_expr + .optimize_recursive(file.dtype())? + .bind(file.dtype())?; + let stream = file.scan()?.with_filter(filter).into_array_stream()?; let results = stream.try_collect::>().await; assert!( @@ -2340,7 +2370,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?; @@ -2387,7 +2417,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?; @@ -2750,9 +2780,11 @@ 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 = filter + .optimize_recursive(file.dtype())? + .bind(file.dtype())?; + let result = file .scan()? .with_filter(filter) .into_array_stream()? 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..de4f8e267fc 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; 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: &BoundExpression, 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: &BoundExpression, 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..5d238cb955f 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; 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: &BoundExpression, 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..1e718fa2016 100644 --- a/vortex-geo/src/prune/mod.rs +++ b/vortex-geo/src/prune/mod.rs @@ -22,23 +22,22 @@ 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::expr::BoundExpression; +use vortex_array::expr::bound::binary; +use vortex_array::expr::bound::case_when; +use vortex_array::expr::bound::checked_add; +use vortex_array::expr::bound::ext_storage; +use vortex_array::expr::bound::get_item; +use vortex_array::expr::bound::gt; +use vortex_array::expr::bound::gt_eq; +use vortex_array::expr::bound::lit; +use vortex_array::expr::bound::lt; +use vortex_array::expr::bound::lt_eq; use vortex_array::scalar::Scalar; -use vortex_array::scalar_fn::ScalarFnVTableExt; -use vortex_array::scalar_fn::fns::binary::Binary; use vortex_array::scalar_fn::fns::literal::Literal; use vortex_array::scalar_fn::fns::operators::Operator; +use vortex_array::stats::bound::stat; use vortex_array::stats::rewrite::StatsRewriteCtx; -use vortex_array::stats::stat; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::aggregate_fn::GeometryAabb; @@ -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 BoundExpression, 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,7 +95,7 @@ fn query_aabb(constant: &Scalar, ctx: &StatsRewriteCtx<'_>) -> VortexResult Expression { +fn aabb_stat(geom: &BoundExpression) -> BoundExpression { // `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))) @@ -106,17 +105,17 @@ fn aabb_stat(geom: &Expression) -> Expression { /// 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: &BoundExpression, query: GeoRect) -> BoundExpression { 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: BoundExpression, hi: BoundExpression| { maximum( lit(0.0), maximum( - binop(Operator::Sub, lit(q_lo), hi), - binop(Operator::Sub, lo, lit(q_hi)), + binary(Operator::Sub, lit(q_lo), hi), + binary(Operator::Sub, lo, lit(q_hi)), ), ) }; @@ -128,13 +127,13 @@ 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: &BoundExpression, query: GeoRect) -> BoundExpression { 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| { - binop( + let span = |q_lo: f64, q_hi: f64, lo: BoundExpression, hi: BoundExpression| { + binary( Operator::Sub, maximum(lit(q_hi), hi), minimum(lit(q_lo), lo), @@ -145,24 +144,17 @@ fn max_dist_sq(aabb: &Expression, query: GeoRect) -> Expression { checked_add(square(dx), square(dy)) } -/// `a b`. -fn binop(op: Operator, a: Expression, b: Expression) -> Expression { - Binary - .try_new_expr(op, [a, b]) - .vortex_expect("binary expression") -} - /// `e * e`. -fn square(e: Expression) -> Expression { - binop(Operator::Mul, e.clone(), e) +fn square(e: BoundExpression) -> BoundExpression { + binary(Operator::Mul, e.clone(), e) } /// `max(a, b)`. -fn maximum(a: Expression, b: Expression) -> Expression { +fn maximum(a: BoundExpression, b: BoundExpression) -> BoundExpression { case_when(gt(a.clone(), b.clone()), a, b) } /// `min(a, b)`. -fn minimum(a: Expression, b: Expression) -> Expression { +fn minimum(a: BoundExpression, b: BoundExpression) -> BoundExpression { 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..c3033063d9c 100644 --- a/vortex-layout/src/layouts/dict/reader.rs +++ b/vortex-layout/src/layouts/dict/reader.rs @@ -21,14 +21,12 @@ use vortex_array::dtype::FieldMask; use vortex_array::dtype::Nullability; use vortex_array::expr::BoundExpression; use vortex_array::expr::ExactBoundExpr; +use vortex_array::expr::bound::pack as bound_pack; use vortex_array::expr::direct_bound_annotations; use vortex_array::expr::label_bound_tree; use vortex_array::expr::root; use vortex_array::expr::transform::partition_bound_annotations; use vortex_array::optimizer::ArrayOptimizer; -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::scalar_fn::is_negative_cost; use vortex_error::VortexError; use vortex_error::VortexExpect; @@ -55,7 +53,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 +151,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| { @@ -301,13 +301,7 @@ impl LayoutReader for DictReader { let values_eval = if let Some(inner) = expr_inner { // "outer" takes a struct field with PUSHDOWN_ANNOTATION name, so // pack inner with this name as well - let inner = BoundExpression::try_new( - Pack.bind(PackOptions { - names: [PUSHDOWN_ANNOTATION].into(), - nullability: Nullability::NonNullable, - }), - [inner], - )?; + let inner = bound_pack([(PUSHDOWN_ANNOTATION, inner)], Nullability::NonNullable); // We can't use values_eval as it uses values_array_uncanonical // which in turn gets populated from self.values. If @@ -376,7 +370,10 @@ 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::bound::pack as bound_pack; use vortex_array::expr::byte_length; use vortex_array::expr::cast; use vortex_array::expr::eq; @@ -746,7 +743,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 +755,36 @@ mod tests { ) .into_array(); - let pushed = array.clone().apply(&pack( - [(PUSHDOWN_ANNOTATION, inner)], + let pushed_expr = bound_pack( + [(FieldName::from(PUSHDOWN_ANNOTATION), inner)], Nullability::NonNullable, - ))?; - let actual = pushed.apply(&outer)?; + ); + 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 +793,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 +821,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/list/expr.rs b/vortex-layout/src/layouts/list/expr.rs index f446ad11ab6..8c10e400ad8 100644 --- a/vortex-layout/src/layouts/list/expr.rs +++ b/vortex-layout/src/layouts/list/expr.rs @@ -4,12 +4,10 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::expr::BoundExpression; -use vortex_array::scalar_fn::EmptyOptions; -use vortex_array::scalar_fn::ScalarFnVTableExt; +use vortex_array::expr::bound::not; use vortex_array::scalar_fn::fns::is_not_null::IsNotNull; use vortex_array::scalar_fn::fns::is_null::IsNull; use vortex_array::scalar_fn::fns::list_length::ListLength; -use vortex_array::scalar_fn::fns::not::Not; use vortex_error::VortexResult; /// The minimal set of list children an expression needs for evaluation. @@ -84,10 +82,7 @@ fn rewrite_validity_expr_with_root( && expr.children().len() == 1 && expr.children()[0].is_root() { - return BoundExpression::try_new( - Not.bind(EmptyOptions), - [BoundExpression::new_root(root_dtype.clone())], - ); + return Ok(not(BoundExpression::new_root(root_dtype.clone()))); } if expr.is_root() { return Ok(BoundExpression::new_root(root_dtype.clone())); diff --git a/vortex-layout/src/layouts/struct_/reader.rs b/vortex-layout/src/layouts/struct_/reader.rs index 1c0a9006f48..db9bf2e787d 100644 --- a/vortex-layout/src/layouts/struct_/reader.rs +++ b/vortex-layout/src/layouts/struct_/reader.rs @@ -17,11 +17,12 @@ use vortex_array::builtins::ArrayBuiltins; use vortex_array::dtype::DType; use vortex_array::dtype::FieldMask; use vortex_array::dtype::FieldName; -use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::StructFields; use vortex_array::expr::BoundExpression; use vortex_array::expr::ExactBoundExpr; +use vortex_array::expr::bound::get_item; +use vortex_array::expr::bound::pack; use vortex_array::expr::make_bound_free_field_annotator; use vortex_array::expr::root; use vortex_array::expr::transform::BoundPartitionedExpr; @@ -29,11 +30,9 @@ use vortex_array::expr::transform::partition_bound; use vortex_array::expr::traversal::NodeExt; use vortex_array::expr::traversal::Transformed; use vortex_array::expr::traversal::TraversalOrder; -use vortex_array::scalar_fn::ScalarFnVTableExt; use vortex_array::scalar_fn::fns::get_item::GetItem; use vortex_array::scalar_fn::fns::merge::Merge; use vortex_array::scalar_fn::fns::pack::Pack; -use vortex_array::scalar_fn::fns::pack::PackOptions; use vortex_array::scalar_fn::fns::select::Select; use vortex_error::VortexExpect; use vortex_error::VortexResult; @@ -243,9 +242,12 @@ fn expanded_struct_root( let children = fields .names() .iter() - .map(|name| BoundExpression::try_new(GetItem.bind(name.clone()), [root.clone()])) - .try_collect()?; - bound_pack(fields.names().clone(), children) + .map(|name| get_item(name.clone(), root.clone())) + .collect::>(); + Ok(pack( + fields.names().iter().cloned().zip(children), + Nullability::NonNullable, + )) } fn expand_struct_root( @@ -287,7 +289,7 @@ fn expand_struct_root( if let Some(selection) = scalar_fn.as_opt::