diff --git a/benchmarks/compress-bench/src/vortex.rs b/benchmarks/compress-bench/src/vortex.rs index 20b4b9f1402..77f579c7333 100644 --- a/benchmarks/compress-bench/src/vortex.rs +++ b/benchmarks/compress-bench/src/vortex.rs @@ -72,8 +72,8 @@ impl Compressor for VortexCompressor { // Columns are named "0".."num_columns-1"; project the given subset. let names: FieldNames = cols.iter().map(|i| i.to_string()).collect(); let projection = select(names, root()) - .optimize_recursive(&source_dtype)? - .bind(&source_dtype)?; + .bind(&source_dtype)? + .optimize_recursive()?; scan = scan.with_projection(projection); } let schema = Arc::new(SESSION.arrow().to_arrow_schema(&scan.dtype()?)?); diff --git a/fuzz/fuzz_targets/file_io.rs b/fuzz/fuzz_targets/file_io.rs index 38a9016e4b8..2f008707b64 100644 --- a/fuzz/fuzz_targets/file_io.rs +++ b/fuzz/fuzz_targets/file_io.rs @@ -84,14 +84,14 @@ fuzz_target!(|fuzz: FuzzFileAction| -> Corpus { .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())) + .bind(file.dtype()) + .and_then(|expr| expr.optimize_recursive()) .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())) + .bind(file.dtype()) + .and_then(|expr| expr.optimize_recursive()) }) .transpose() .vortex_expect("filter should bind in fuzz test"); diff --git a/vortex-array/benches/expr/case_when_bench.rs b/vortex-array/benches/expr/case_when_bench.rs index fa5af4bc0ef..819111d2861 100644 --- a/vortex-array/benches/expr/case_when_bench.rs +++ b/vortex-array/benches/expr/case_when_bench.rs @@ -72,13 +72,14 @@ fn case_when_simple(bencher: Bencher, size: usize) { lit(100i32), lit(0i32), ); + let expr = expr.bind(array.dtype()).unwrap(); bencher .with_inputs(|| (&expr, &array, SESSION.create_execution_ctx())) .bench_refs(|(expr, array, ctx)| { array .clone() - .apply(expr) + .apply_bound(expr) .unwrap() .execute::(ctx) .unwrap() @@ -99,13 +100,14 @@ fn case_when_nary_3_conditions(bencher: Bencher, size: usize) { ], Some(lit(0i32)), ); + let expr = expr.bind(array.dtype()).unwrap(); bencher .with_inputs(|| (&expr, &array, SESSION.create_execution_ctx())) .bench_refs(|(expr, array, ctx)| { array .clone() - .apply(expr) + .apply_bound(expr) .unwrap() .execute::(ctx) .unwrap() @@ -127,13 +129,14 @@ fn case_when_nary_10_conditions(bencher: Bencher, size: usize) { }) .collect(); let expr = nested_case_when(pairs, Some(lit(0i32))); + let expr = expr.bind(array.dtype()).unwrap(); bencher .with_inputs(|| (&expr, &array, SESSION.create_execution_ctx())) .bench_refs(|(expr, array, ctx)| { array .clone() - .apply(expr) + .apply_bound(expr) .unwrap() .execute::(ctx) .unwrap() @@ -150,13 +153,14 @@ fn case_when_nary_equality_lookup(bencher: Bencher, size: usize) { .map(|i| (eq(get_item("value", root()), lit(i)), lit(i * 10))) .collect(); let expr = nested_case_when(pairs, Some(lit(-1i32))); + let expr = expr.bind(array.dtype()).unwrap(); bencher .with_inputs(|| (&expr, &array, SESSION.create_execution_ctx())) .bench_refs(|(expr, array, ctx)| { array .clone() - .apply(expr) + .apply_bound(expr) .unwrap() .execute::(ctx) .unwrap() @@ -170,13 +174,14 @@ fn case_when_without_else(bencher: Bencher, size: usize) { // CASE WHEN value > 500 THEN 100 END let expr = case_when_no_else(gt(get_item("value", root()), lit(500i32)), lit(100i32)); + let expr = expr.bind(array.dtype()).unwrap(); bencher .with_inputs(|| (&expr, &array, SESSION.create_execution_ctx())) .bench_refs(|(expr, array, ctx)| { array .clone() - .apply(expr) + .apply_bound(expr) .unwrap() .execute::(ctx) .unwrap() @@ -194,13 +199,14 @@ fn case_when_all_true(bencher: Bencher, size: usize) { lit(100i32), lit(0i32), ); + let expr = expr.bind(array.dtype()).unwrap(); bencher .with_inputs(|| (&expr, &array, SESSION.create_execution_ctx())) .bench_refs(|(expr, array, ctx)| { array .clone() - .apply(expr) + .apply_bound(expr) .unwrap() .execute::(ctx) .unwrap() @@ -227,13 +233,14 @@ fn case_when_nary_early_dominant(bencher: Bencher, size: usize) { ], Some(lit(4i32)), ); + let expr = expr.bind(array.dtype()).unwrap(); bencher .with_inputs(|| (&expr, &array, SESSION.create_execution_ctx())) .bench_refs(|(expr, array, ctx)| { array .clone() - .apply(expr) + .apply_bound(expr) .unwrap() .execute::(ctx) .unwrap() @@ -251,13 +258,14 @@ fn case_when_all_false(bencher: Bencher, size: usize) { lit(100i32), lit(0i32), ); + let expr = expr.bind(array.dtype()).unwrap(); bencher .with_inputs(|| (&expr, &array, SESSION.create_execution_ctx())) .bench_refs(|(expr, array, ctx)| { array .clone() - .apply(expr) + .apply_bound(expr) .unwrap() .execute::(ctx) .unwrap() @@ -278,13 +286,14 @@ fn case_when_fragmented(bencher: Bencher, size: usize) { ], Some(lit(2i32)), ); + let expr = expr.bind(array.dtype()).unwrap(); bencher .with_inputs(|| (&expr, &array, SESSION.create_execution_ctx())) .bench_refs(|(expr, array, ctx)| { array .clone() - .apply(expr) + .apply_bound(expr) .unwrap() .execute::(ctx) .unwrap() diff --git a/vortex-array/benches/expr/optimize_bench.rs b/vortex-array/benches/expr/optimize_bench.rs index d86e115c643..e94c8f35c58 100644 --- a/vortex-array/benches/expr/optimize_bench.rs +++ b/vortex-array/benches/expr/optimize_bench.rs @@ -45,5 +45,6 @@ fn build_or_chain(n: usize) -> Expression { fn optimize_or_chain(bencher: Bencher, n: usize) { let expr = build_or_chain(n); let scope = struct_scope(); - bencher.bench(|| expr.optimize_recursive(&scope).unwrap()); + let expr = expr.bind(&scope).unwrap(); + bencher.bench(|| expr.optimize_recursive().unwrap()); } diff --git a/vortex-array/benches/expr/optimize_predicate.rs b/vortex-array/benches/expr/optimize_predicate.rs index 731c7c2b06f..24ec28f2958 100644 --- a/vortex-array/benches/expr/optimize_predicate.rs +++ b/vortex-array/benches/expr/optimize_predicate.rs @@ -1,9 +1,9 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Benchmarks `Expression::optimize_recursive` on lookup-style pushdown predicates: -//! an id membership test (either `list_contains` or a balanced OR of equalities) conjoined -//! with timestamp range bounds and kind filters. +//! Benchmarks bound optimization on lookup-style pushdown predicates: an id membership test +//! (either `list_contains` or a balanced OR of equalities) conjoined with timestamp range bounds +//! and kind filters. #![expect(clippy::unwrap_used)] @@ -207,7 +207,7 @@ fn lookup_predicate(predicate_case: PredicateCase) -> Expression { #[divan::bench(args = PREDICATE_CASES)] fn optimize_lookup_predicate(bencher: Bencher, predicate_case: &PredicateCase) { let scope = scope(); - let predicate = lookup_predicate(*predicate_case); + let predicate = lookup_predicate(*predicate_case).bind(&scope).unwrap(); - bencher.bench(|| black_box(predicate.optimize_recursive(&scope))); + bencher.bench(|| black_box(predicate.optimize_recursive())); } diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index 4bd276e191d..bfb990ebd05 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -9,6 +9,7 @@ use std::hash::Hasher; use std::sync::Arc; use itertools::Itertools; +use smallvec::SmallVec; use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; @@ -127,10 +128,8 @@ impl BoundExpression { children.len() ); - let arg_dtypes = children - .iter() - .map(|child| child.dtype().clone()) - .collect_vec(); + let arg_dtypes: SmallVec<[DType; 4]> = + children.iter().map(|child| child.dtype().clone()).collect(); let dtype = scalar_fn.return_dtype(&arg_dtypes)?; Ok(Self::Scalar { diff --git a/vortex-array/src/expr/optimize.rs b/vortex-array/src/expr/optimize.rs index 3e625324a18..1c1fef3fbed 100644 --- a/vortex-array/src/expr/optimize.rs +++ b/vortex-array/src/expr/optimize.rs @@ -1,65 +1,36 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use std::cell::RefCell; - -use itertools::Itertools; use vortex_error::VortexResult; -use vortex_error::vortex_err; -use vortex_utils::aliases::hash_map::HashMap; -use crate::dtype::DType; -use crate::expr::Expression; -use crate::expr::transform::match_between::find_between; -use crate::scalar_fn::ExpressionReduceNode; -use crate::scalar_fn::SimplifyCtx; +use crate::expr::BoundExpression; +use crate::expr::transform::match_between::find_between_bound; -impl Expression { - /// Optimize the root expression node only, iterating to convergence. +impl BoundExpression { + /// Optimize the root node with simplification and abstract reduction. /// - /// This applies optimization rules repeatedly until no more changes occur: - /// 1. `simplify_untyped` - type-independent simplifications - /// 2. `simplify` - type-aware simplifications - /// 3. `reduce` - abstract reduction rules via `ReduceNode` - pub fn optimize(&self, scope: &DType) -> VortexResult { - let cache = SimplifyCache::new(scope); - Ok(self.try_optimize(&cache)?.unwrap_or_else(|| self.clone())) + /// Expressions must be bound before they can be optimized, even for rules that do not inspect + /// dtypes. This keeps the optimizer's input and output in the same typed representation. + pub fn optimize(&self) -> VortexResult { + Ok(self.try_optimize()?.unwrap_or_else(|| self.clone())) } - /// Apply this node's own untyped simplification rule, if it has one. - /// - /// Non-scalar nodes carry no rules, so they never simplify. - fn simplify_untyped_node(&self) -> VortexResult> { + fn simplify_node(&self) -> VortexResult> { match self { - Expression::Scalar { scalar_fn, .. } => scalar_fn.simplify_untyped(self), - Expression::Root => Ok(None), + BoundExpression::Scalar { scalar_fn, .. } => scalar_fn.simplify(self), + BoundExpression::Root { .. } => Ok(None), } } - /// Apply this node's own type-aware simplification rule, if it has one. - fn simplify_node(&self, ctx: &dyn SimplifyCtx) -> VortexResult> { + fn reduce_node(&self) -> VortexResult> { match self { - Expression::Scalar { scalar_fn, .. } => scalar_fn.simplify(self, ctx), - Expression::Root => Ok(None), + BoundExpression::Scalar { scalar_fn, .. } => scalar_fn.reduce_bound_expression(self), + BoundExpression::Root { .. } => Ok(None), } } - /// Apply this node's own abstract reduction rule, if it has one. - fn reduce_node<'a>( - &self, - node: &ExpressionReduceNode<'a>, - ) -> VortexResult>> { - match self { - Expression::Scalar { scalar_fn, .. } => scalar_fn.reduce_expression(node), - Expression::Root => Ok(None), - } - } - - /// Try to optimize the root expression node only, returning None if no optimizations applied. - fn try_optimize(&self, cache: &SimplifyCache<'_>) -> VortexResult> { - // Copy-on-write: `current` stays None until a rule fires, so unchanged nodes (the common - // case) are never cloned. - let mut current: Option = None; + fn try_optimize(&self) -> VortexResult> { + let mut current: Option = None; let mut loop_counter = 0; loop { @@ -73,29 +44,14 @@ impl Expression { let expr = current.as_ref().unwrap_or(self); let mut changed = false; - // Try simplify_untyped - if let Some(simplified) = expr.simplify_untyped_node()? { + if let Some(simplified) = expr.simplify_node()? { current = Some(simplified); changed = true; } - // Try simplify (typed) let expr = current.as_ref().unwrap_or(self); - if let Some(simplified) = expr.simplify_node(cache)? { - current = Some(simplified); - changed = true; - } - - // Try reduce via ReduceNode. The node borrows the expression and scope, so - // constructing it is free; the block scopes the borrows so `current` can be updated. - let reduced = { - let expr = current.as_ref().unwrap_or(self); - let reduce_node = ExpressionReduceNode::new(expr, cache.scope); - expr.reduce_node(&reduce_node)? - .map(ExpressionReduceNode::into_expression) - }; - if let Some(reduced_expr) = reduced { - current = Some(reduced_expr); + if let Some(reduced) = expr.reduce_node()? { + current = Some(reduced); changed = true; } @@ -107,41 +63,33 @@ impl Expression { Ok(current) } - /// Optimize the entire expression tree recursively. - /// - /// Optimizes children first (bottom-up), then optimizes the root. - pub fn optimize_recursive(&self, scope: &DType) -> VortexResult { + /// Optimize the entire bound expression tree recursively. + pub fn optimize_recursive(&self) -> VortexResult { Ok(self .clone() - .try_optimize_recursive(scope)? + .try_optimize_recursive()? .unwrap_or_else(|| self.clone())) } - /// Try to optimize the entire expression tree recursively. - pub fn try_optimize_recursive(&self, scope: &DType) -> VortexResult> { - let cache = SimplifyCache::new(scope); - let result = self.try_optimize_recursive_inner(&cache)?; + pub fn try_optimize_recursive(&self) -> VortexResult> { + let result = self.try_optimize_recursive_inner()?; // Apply the between optimization once at the top level only. // TODO(ngates): remove the "between" optimization, or rewrite it to not always convert - // to CNF? - Ok(Some(find_between(result.unwrap_or_else(|| self.clone())))) + // to CNF? + Ok(Some(find_between_bound( + result.unwrap_or_else(|| self.clone()), + ))) } - fn try_optimize_recursive_inner( - &self, - cache: &SimplifyCache<'_>, - ) -> VortexResult> { - // First optimize the root - let mut current = self.try_optimize(cache)?; + fn try_optimize_recursive_inner(&self) -> VortexResult> { + let mut current = self.try_optimize()?; - // Then recursively optimize children. The new children vector is only allocated once a - // child actually changes, so fully-optimized subtrees cost no allocations. let expr = current.as_ref().unwrap_or(self); let children = expr.children(); - let mut new_children: Option> = None; + let mut new_children: Option> = None; for (idx, child) in children.iter().enumerate() { - if let Some(optimized) = child.try_optimize_recursive_inner(cache)? { + if let Some(optimized) = child.try_optimize_recursive_inner()? { new_children .get_or_insert_with(|| children[..idx].to_vec()) .push(optimized); @@ -152,58 +100,13 @@ impl Expression { if let Some(new_children) = new_children { let updated = expr.clone().with_children(new_children)?; - - // After updating children, try to optimize root again - current = Some(updated.try_optimize(cache)?.unwrap_or(updated)); + current = Some(updated.try_optimize()?.unwrap_or(updated)); } Ok(current) } } -struct SimplifyCache<'a> { - scope: &'a DType, - dtype_cache: RefCell>, -} - -impl<'a> SimplifyCache<'a> { - fn new(scope: &'a DType) -> Self { - Self { - scope, - dtype_cache: RefCell::new(HashMap::new()), - } - } -} - -impl SimplifyCtx for SimplifyCache<'_> { - fn return_dtype(&self, expr: &Expression) -> VortexResult { - // If the expression is "root", return the scope dtype - if expr.is_root() { - return Ok(self.scope.clone()); - } - - if let Some(dtype) = self.dtype_cache.borrow().get(expr) { - return Ok(dtype.clone()); - } - - // Otherwise, compute dtype from children - let input_dtypes: Vec<_> = expr - .children() - .iter() - .map(|c| self.return_dtype(c)) - .try_collect()?; - let dtype = expr - .as_scalar() - .ok_or_else(|| vortex_err!("cannot type a non-scalar expression: {expr}"))? - .return_dtype(&input_dtypes)?; - self.dtype_cache - .borrow_mut() - .insert(expr.clone(), dtype.clone()); - - Ok(dtype) - } -} - #[cfg(test)] mod tests { use vortex_error::VortexResult; @@ -236,7 +139,7 @@ mod tests { ), Nullability::NonNullable, ); - let optimized = expr.optimize_recursive(&scope)?; + let optimized = expr.bind(&scope)?.optimize_recursive()?; let s = optimized.to_string(); assert!(s.contains("$.x"), "expected $.x in {s}"); @@ -261,7 +164,7 @@ mod tests { ), Nullability::NonNullable, ); - let optimized = expr.optimize_recursive(&scope)?; + let optimized = expr.bind(&scope)?.optimize_recursive()?; // Prune rules pattern-match a bare Literal on the comparison RHS; a cast wrapper // silently disables pruning. diff --git a/vortex-array/src/expr/transform/match_between.rs b/vortex-array/src/expr/transform/match_between.rs index 7cb427596f4..4ac38ecd290 100644 --- a/vortex-array/src/expr/transform/match_between.rs +++ b/vortex-array/src/expr/transform/match_between.rs @@ -1,12 +1,11 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -use crate::expr::Expression; -use crate::expr::and_collect; -use crate::expr::forms::conjuncts; -use crate::expr::lit; -use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::fns::between::Between; +use crate::expr::BoundExpression; +use crate::expr::bound::and_collect as bound_and_collect; +use crate::expr::bound::between as bound_between; +use crate::expr::bound::binary as bound_binary; +use crate::expr::bound::lit as bound_lit; use crate::scalar_fn::fns::between::BetweenOptions; use crate::scalar_fn::fns::between::StrictComparison; use crate::scalar_fn::fns::binary::Binary; @@ -14,26 +13,21 @@ use crate::scalar_fn::fns::get_item::GetItem; use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::operators::Operator; -/// This pass looks for expression of the form -/// `x >= a && x < b` and converts them into x between a and b` -pub fn find_between(expr: Expression) -> Expression { - // We search all pairs of cnfs to find any pair of expressions can be converted into a between - // expression. - let mut conjuncts = conjuncts(&expr); +/// Look for `x >= a AND x < b` and replace it with a bound `between` expression. +pub(crate) fn find_between_bound(expr: BoundExpression) -> BoundExpression { + let mut conjuncts = bound_conjuncts(&expr); let mut rest = vec![]; for idx in 0..conjuncts.len() { - let Some(c) = conjuncts.get(idx).cloned() else { + let Some(conjunct) = conjuncts.get(idx).cloned() else { continue; }; let mut matched = false; for idx2 in (idx + 1)..conjuncts.len() { - // Since values are removed in iterations there might not be a value at idx2, - // but all values will have been considered. - let Some(c2) = conjuncts.get(idx2) else { + let Some(other) = conjuncts.get(idx2) else { continue; }; - if let Some(expr) = maybe_match(&c, c2) { + if let Some(expr) = maybe_match_bound(&conjunct, other) { rest.push(expr); conjuncts.remove(idx2); matched = true; @@ -41,33 +35,48 @@ pub fn find_between(expr: Expression) -> Expression { } } if !matched { - rest.push(c.clone()) + rest.push(conjunct); } } - and_collect(rest).unwrap_or_else(|| lit(true)) + bound_and_collect(rest).unwrap_or_else(|| bound_lit(true)) } -fn maybe_match(lhs: &Expression, rhs: &Expression) -> Option { +fn bound_conjuncts(expr: &BoundExpression) -> Vec { + let mut conjuncts = vec![]; + bound_conjuncts_impl(expr, &mut conjuncts); + conjuncts +} + +fn bound_conjuncts_impl(expr: &BoundExpression, conjuncts: &mut Vec) { + if expr + .as_opt::() + .is_some_and(|operator| *operator == Operator::And) + { + bound_conjuncts_impl(expr.child(0), conjuncts); + bound_conjuncts_impl(expr.child(1), conjuncts); + } else { + conjuncts.push(expr.clone()); + } +} + +fn maybe_match_bound(lhs: &BoundExpression, rhs: &BoundExpression) -> Option { let (Some(lhs_op), Some(rhs_op)) = (lhs.as_opt::(), rhs.as_opt::()) else { return None; }; - // Extract the grandchildren let lhs_lhs = lhs.child(0); let lhs_rhs = lhs.child(1); let rhs_lhs = rhs.child(0); let rhs_rhs = rhs.child(1); - // Cannot compare to self if lhs_lhs.eq(lhs_rhs) || rhs_lhs.eq(rhs_rhs) { return None; } - // First, get both halves to have GetItem on the left let lhs = match (lhs_lhs.is::(), lhs_rhs.is::()) { (true, false) => lhs.clone(), - (false, true) => Binary.new_expr(lhs_op.swap()?, [lhs_rhs.clone(), lhs_lhs.clone()]), + (false, true) => bound_binary(lhs_op.swap()?, lhs_rhs.clone(), lhs_lhs.clone()), _ => return None, }; let lhs_op = lhs.as_::(); @@ -75,20 +84,17 @@ fn maybe_match(lhs: &Expression, rhs: &Expression) -> Option { let rhs = match (rhs_lhs.is::(), rhs_rhs.is::()) { (true, false) => rhs.clone(), - (false, true) => Binary.new_expr(rhs_op.swap()?, [rhs_rhs.clone(), rhs_lhs.clone()]), + (false, true) => bound_binary(rhs_op.swap()?, rhs_rhs.clone(), rhs_lhs.clone()), _ => return None, }; let rhs_op = rhs.as_::(); let rhs_lhs = rhs.child(0); - // Both conjuncts must reference the same GetItem column if !lhs_lhs.eq(rhs_lhs) { return None; } let target = lhs_lhs.clone(); - - // Find the lower bound let (lower, upper) = match (lhs_op, rhs_op) { (Operator::Lt | Operator::Lte, Operator::Gt | Operator::Gte) => (rhs, lhs), (Operator::Gt | Operator::Gte, Operator::Lt | Operator::Lte) => (lhs, rhs), @@ -99,19 +105,20 @@ fn maybe_match(lhs: &Expression, rhs: &Expression) -> Option { let upper_op = upper.as_::(); let upper_rhs = upper.child(1); - // Ensure bounds are literals - let _ = lower_rhs.as_opt::()?; - let _ = upper_rhs.as_opt::()?; + lower_rhs.as_opt::()?; + upper_rhs.as_opt::()?; let lower_strict = is_strict_comparison(*lower_op)?; let upper_strict = is_strict_comparison(*upper_op)?; - Some(Between.new_expr( + Some(bound_between( + target, + lower_rhs.clone(), + upper_rhs.clone(), BetweenOptions { lower_strict, upper_strict, }, - [target, lower_rhs.clone(), upper_rhs.clone()], )) } @@ -128,7 +135,6 @@ mod tests { use vortex_buffer::buffer; use vortex_error::VortexResult; - use super::find_between; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; @@ -137,6 +143,8 @@ mod tests { use crate::dtype::DType; use crate::dtype::Nullability; use crate::dtype::PType; + use crate::expr::BoundExpression; + use crate::expr::Expression; use crate::expr::and; use crate::expr::between; use crate::expr::col; @@ -149,6 +157,22 @@ mod tests { use crate::scalar_fn::fns::between::BetweenOptions; use crate::scalar_fn::fns::between::StrictComparison; + fn scope(fields: &[&str]) -> DType { + DType::struct_( + fields.iter().map(|name| { + ( + *name, + DType::Primitive(PType::I32, Nullability::NonNullable), + ) + }), + Nullability::NonNullable, + ) + } + + fn optimize(expr: Expression, scope: &DType) -> VortexResult { + expr.bind(scope)?.optimize_recursive() + } + /// A null literal bound must not change the values of the rewritten expression. Kleene `AND` /// keeps a row false when the surviving comparison is false, so the rewrite cannot null it. #[test] @@ -169,8 +193,9 @@ mod tests { .execute::(ctx)? .opt_bool_vec(ctx); + let optimized = optimize(expr, data.dtype())?; let after = data - .apply(&find_between(expr))? + .apply_bound(&optimized)? .execute::(ctx)? .opt_bool_vec(ctx); @@ -182,14 +207,15 @@ mod tests { } #[test] - fn test_bad_match() { + fn test_bad_match() -> VortexResult<()> { // An impossible expression let expr = and(lt_eq(lit(100), col("x")), gt(lit(-100), col("x"))); - let find = find_between(expr); + let scope = scope(&["x"]); + let find = optimize(expr, &scope)?; assert_eq!( - &find, - &between( + find, + between( col("x"), lit(100), lit(-100), @@ -198,17 +224,20 @@ mod tests { upper_strict: StrictComparison::Strict, } ) + .bind(&scope)? ); + Ok(()) } #[test] - fn test_match_between() { + fn test_match_between() -> VortexResult<()> { let expr = and(lt(lit(2), col("x")), gt_eq(lit(5), col("x"))); - let find = find_between(expr); + let scope = scope(&["x"]); + let find = optimize(expr, &scope)?; // 2 < x <= 5 assert_eq!( - &between( + between( col("x"), lit(2), lit(5), @@ -216,19 +245,22 @@ mod tests { lower_strict: StrictComparison::Strict, upper_strict: StrictComparison::NonStrict, } - ), - &find + ) + .bind(&scope)?, + find ); + Ok(()) } #[test] - fn test_match_2_between() { + fn test_match_2_between() -> VortexResult<()> { let expr = and(gt_eq(col("x"), lit(2)), lt(col("x"), lit(5))); - let find = find_between(expr); + let scope = scope(&["x"]); + let find = optimize(expr, &scope)?; // 2 <= x < 5 assert_eq!( - &between( + between( col("x"), lit(2), lit(5), @@ -236,19 +268,22 @@ mod tests { lower_strict: StrictComparison::NonStrict, upper_strict: StrictComparison::Strict, } - ), - &find + ) + .bind(&scope)?, + find ); + Ok(()) } #[test] - fn test_match_3_between() { + fn test_match_3_between() -> VortexResult<()> { let expr = and(gt_eq(col("x"), lit(2)), gt_eq(lit(5), col("x"))); - let find = find_between(expr); + let scope = scope(&["x"]); + let find = optimize(expr, &scope)?; // 2 <= x < 5 assert_eq!( - &between( + between( col("x"), lit(2), lit(5), @@ -256,19 +291,22 @@ mod tests { lower_strict: StrictComparison::NonStrict, upper_strict: StrictComparison::NonStrict, } - ), - &find + ) + .bind(&scope)?, + find ); + Ok(()) } #[test] - fn test_match_4_between() { + fn test_match_4_between() -> VortexResult<()> { let expr = and(gt_eq(lit(5), col("x")), lt(lit(2), col("x"))); - let find = find_between(expr); + let scope = scope(&["x"]); + let find = optimize(expr, &scope)?; // 2 < x <= 5 assert_eq!( - &between( + between( col("x"), lit(2), lit(5), @@ -276,22 +314,25 @@ mod tests { lower_strict: StrictComparison::Strict, upper_strict: StrictComparison::NonStrict, } - ), - &find + ) + .bind(&scope)?, + find ); + Ok(()) } #[test] - fn test_match_5_between() { + fn test_match_5_between() -> VortexResult<()> { let expr = and( and(gt_eq(col("y"), lit(10)), gt_eq(lit(5), col("x"))), lt(lit(2), col("x")), ); - let find = find_between(expr); + let scope = scope(&["x", "y"]); + let find = optimize(expr, &scope)?; // $.y >= 10 /\ 2 < $.x <= 5 assert_eq!( - &and( + and( gt_eq(col("y"), lit(10)), between( col("x"), @@ -302,22 +343,25 @@ mod tests { upper_strict: StrictComparison::NonStrict, } ) - ), - &find + ) + .bind(&scope)?, + find ); + Ok(()) } #[test] - fn test_match_6_between() { + fn test_match_6_between() -> VortexResult<()> { let expr = and( and(gt_eq(lit(5), col("x")), gt_eq(col("y"), lit(10))), lt(lit(2), col("x")), ); - let find = find_between(expr); + let scope = scope(&["x", "y"]); + let find = optimize(expr, &scope)?; // $.y >= 10 /\ 2 < $.x <= 5 assert_eq!( - &and( + and( between( col("x"), lit(2), @@ -328,8 +372,10 @@ mod tests { } ), gt_eq(col("y"), lit(10)), - ), - &find + ) + .bind(&scope)?, + find ); + Ok(()) } } diff --git a/vortex-array/src/expr/transform/partition.rs b/vortex-array/src/expr/transform/partition.rs index 92dca145d47..7c70cf08618 100644 --- a/vortex-array/src/expr/transform/partition.rs +++ b/vortex-array/src/expr/transform/partition.rs @@ -14,7 +14,6 @@ use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::FieldNames; use crate::dtype::Nullability; -use crate::dtype::StructFields; use crate::expr::Expression; use crate::expr::analysis::Annotation; use crate::expr::analysis::AnnotationFn; @@ -83,7 +82,6 @@ where Nullability::NonNullable, ); - let expr = expr.optimize_recursive(scope)?; let expr_dtype = expr.return_dtype(scope)?; partitions.push(expr); @@ -95,13 +93,8 @@ where .iter() .map(|id| FieldName::from(id.clone())) .collect::(); - let root_scope = DType::Struct( - StructFields::new(partition_names.clone(), partition_dtypes.clone()), - Nullability::NonNullable, - ); - Ok(PartitionedExpr { - root: root.optimize_recursive(&root_scope)?, + root, partitions: partitions.into_boxed_slice(), partition_names, partition_dtypes: partition_dtypes.into_boxed_slice(), @@ -296,7 +289,7 @@ mod tests { let split_a = partitioned.find_partition(&"a".into()).unwrap(); assert_eq!( - &split_a.optimize_recursive(&dtype).unwrap(), + split_a, &pack( [ ("a_0", get_item("x", get_item("a", root()))), @@ -336,14 +329,7 @@ mod tests { let expr = merge([col("a"), pack([("b", col("b"))], NonNullable)]); let partitioned = partition(expr, &dtype, make_free_field_annotator(fields)).unwrap(); - let expected = pack( - [ - ("x", get_item("x", get_item("a_0", col("a")))), - ("y", get_item("y", get_item("a_0", col("a")))), - ("b", get_item("b", get_item("b_0", col("b")))), - ], - NonNullable, - ); + let expected = merge([get_item("a_0", col("a")), get_item("b_0", col("b"))]); assert_eq!( &partitioned.root, &expected, "{} {}", diff --git a/vortex-array/src/expression.rs b/vortex-array/src/expression.rs index d0590f2bf58..e0648d9e3e0 100644 --- a/vortex-array/src/expression.rs +++ b/vortex-array/src/expression.rs @@ -2,7 +2,6 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use itertools::Itertools; -use vortex_error::VortexExpect; use vortex_error::VortexResult; use crate::ArrayRef; @@ -43,31 +42,7 @@ impl ArrayRef { /// Apply the expression to this array, producing a new array in constant time. pub fn apply(self, expr: &Expression) -> VortexResult { - // If the expression is a root, return self. - if expr.is_root() { - return Ok(self); - } - - // Manually convert literals to ConstantArray. - if let Some(scalar) = expr.as_opt::() { - return Ok(ConstantArray::new(scalar.clone(), self.len()).into_array()); - } - - // Otherwise, collect the child arrays. - let children: Vec<_> = expr - .children() - .iter() - .map(|e| self.clone().apply(e)) - .try_collect()?; - - // And wrap the scalar function up in an array. - let scalar_fn = expr - .as_scalar() - .vortex_expect("root and literal were handled above, so this is a scalar node"); - let array = - ScalarFnArray::try_new_with_len(scalar_fn.clone(), children, self.len())?.into_array(); - - // Optimize the resulting array's root. - array.optimize() + let bound = expr.bind(self.dtype())?; + self.apply_bound(&bound) } } diff --git a/vortex-array/src/scalar_fn/erased.rs b/vortex-array/src/scalar_fn/erased.rs index 8b25398c324..35e9112c30d 100644 --- a/vortex-array/src/scalar_fn/erased.rs +++ b/vortex-array/src/scalar_fn/erased.rs @@ -19,16 +19,15 @@ use vortex_utils::debug_with::DebugWith; 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::ArrayReduceNode; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ExpressionReduceNode; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::SimplifyCtx; use crate::scalar_fn::fns::is_not_null::IsNotNull; use crate::scalar_fn::options::ScalarFnOptions; use crate::scalar_fn::signature::ScalarFnSignature; @@ -142,12 +141,12 @@ impl ScalarFnRef { self.0.execute(args, ctx) } - /// Perform abstract reduction on this scalar function node in an expression tree. - pub fn reduce_expression<'a>( + /// Perform abstract reduction on this scalar function node in a bound expression tree. + pub fn reduce_bound_expression( &self, - node: &ExpressionReduceNode<'a>, - ) -> VortexResult>> { - self.0.reduce_expression(node) + node: &BoundExpression, + ) -> VortexResult> { + self.0.reduce_bound_expression(node) } /// Perform abstract reduction on this scalar function node in an array tree. @@ -171,18 +170,9 @@ impl ScalarFnRef { self.0.fmt_sql(expr, f) } - /// Simplify the expression using type information. - pub(crate) fn simplify( - &self, - expr: &Expression, - ctx: &dyn SimplifyCtx, - ) -> VortexResult> { - self.0.simplify(expr, ctx) - } - - /// Simplify the expression without type information. - pub(crate) fn simplify_untyped(&self, expr: &Expression) -> VortexResult> { - self.0.simplify_untyped(expr) + /// Simplify a bound expression using its type information. + pub(crate) fn simplify(&self, expr: &BoundExpression) -> VortexResult> { + self.0.simplify(expr) } } diff --git a/vortex-array/src/scalar_fn/fns/binary/mod.rs b/vortex-array/src/scalar_fn/fns/binary/mod.rs index ed31cdf8d46..6ae2bc087ce 100644 --- a/vortex-array/src/scalar_fn/fns/binary/mod.rs +++ b/vortex-array/src/scalar_fn/fns/binary/mod.rs @@ -20,17 +20,17 @@ use crate::ExecutionCtx; use crate::arrays::ScalarFnArray; use crate::dtype::DType; use crate::dtype::Nullability; +use crate::expr::BoundExpression; use crate::expr::and; +use crate::expr::bound::lit as bound_lit; use crate::expr::display::ExprDisplay; use crate::expr::expression::Expression; -use crate::expr::lit; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::SimplifyCtx; use crate::scalar_fn::fns::literal::Literal; use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::fns::operators::Operator; @@ -180,15 +180,15 @@ impl ScalarFnVTable for Binary { } } - fn simplify_untyped( + fn simplify( &self, operator: &Operator, - expr: &Expression, - ) -> VortexResult> { + expr: &BoundExpression, + ) -> VortexResult> { let lhs = expr.child(0); let rhs = expr.child(1); - let bool_literal = |expr: &Expression| { + let bool_literal = |expr: &BoundExpression| { expr.as_opt::()? .as_bool_opt() .map(|value| value.value()) @@ -209,41 +209,37 @@ impl ScalarFnVTable for Binary { // Other null cases either fall out of the identity/annihilator rules // above (`null AND true`, `null OR false`) or cannot be simplified under // Kleene semantics (`null AND x`, `null OR x` for non-literal `x`). - Ok(match operator { + let simplified = match operator { Operator::And => match (bool_literal(lhs), bool_literal(rhs)) { - (Some(Some(false)), _) | (_, Some(Some(false))) => Some(lit(false)), + (Some(Some(false)), _) | (_, Some(Some(false))) => Some(bound_lit(false)), (Some(Some(true)), _) => Some(rhs.clone()), (_, Some(Some(true))) => Some(lhs.clone()), (Some(None), Some(None)) => Some(lhs.clone()), _ => None, }, Operator::Or => match (bool_literal(lhs), bool_literal(rhs)) { - (Some(Some(true)), _) | (_, Some(Some(true))) => Some(lit(true)), + (Some(Some(true)), _) | (_, Some(Some(true))) => Some(bound_lit(true)), (Some(Some(false)), _) => Some(rhs.clone()), (_, Some(Some(false))) => Some(lhs.clone()), (Some(None), Some(None)) => Some(lhs.clone()), _ => None, }, _ => None, - }) - } + }; + + if simplified.is_some() { + return Ok(simplified); + } - fn simplify( - &self, - operator: &Operator, - expr: &Expression, - ctx: &dyn SimplifyCtx, - ) -> VortexResult> { let is_literal_null = - |expr: &Expression| expr.as_opt::().is_some_and(Scalar::is_null); + |expr: &BoundExpression| expr.as_opt::().is_some_and(Scalar::is_null); if operator.is_comparison() && (is_literal_null(expr.child(0)) || is_literal_null(expr.child(1))) { - // Validate the comparison before reducing it. This preserves type - // errors for expressions like `int_col = null_utf8`. - ctx.return_dtype(expr)?; - return Ok(Some(lit(Scalar::null(DType::Bool(Nullability::Nullable))))); + return Ok(Some(bound_lit(Scalar::null(DType::Bool( + Nullability::Nullable, + ))))); } Ok(None) @@ -440,21 +436,22 @@ mod tests { ); assert_eq!( - expr.optimize_recursive(&dtype)?, - lit(Scalar::null(DType::Bool(Nullability::Nullable))) + expr.bind(&dtype)?.optimize_recursive()?, + lit(Scalar::null(DType::Bool(Nullability::Nullable))).bind(&dtype)? ); Ok(()) } #[test] - fn comparison_with_incompatible_null_still_type_checks() { + fn comparison_with_incompatible_null_still_type_checks() -> VortexResult<()> { let dtype = test_harness::struct_dtype(); let expr = eq( col("col1"), lit(Scalar::null(DType::Utf8(Nullability::Nullable))), ); - assert!(expr.optimize_recursive(&dtype).is_err()); + assert!(expr.bind(&dtype).is_err()); + Ok(()) } #[test] diff --git a/vortex-array/src/scalar_fn/fns/case_when.rs b/vortex-array/src/scalar_fn/fns/case_when.rs index d5ea8702c53..88d9b0b3b6b 100644 --- a/vortex-array/src/scalar_fn/fns/case_when.rs +++ b/vortex-array/src/scalar_fn/fns/case_when.rs @@ -34,7 +34,8 @@ use crate::builders::ArrayBuilder; use crate::builders::builder_with_capacity; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; -use crate::expr::Expression; +use crate::expr::BoundExpression; +use crate::expr::bound::fill_null as bound_fill_null; use crate::expr::display::ExprDisplay; use crate::scalar::Scalar; use crate::scalar_fn::Arity; @@ -42,7 +43,6 @@ use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::SimplifyCtx; use crate::scalar_fn::fns::is_not_null::IsNotNull; use crate::scalar_fn::fns::is_null::IsNull; use crate::scalar_fn::fns::literal::Literal; @@ -259,9 +259,8 @@ impl ScalarFnVTable for CaseWhen { fn simplify( &self, options: &Self::Options, - expr: &Expression, - _ctx: &dyn SimplifyCtx, - ) -> VortexResult> { + expr: &BoundExpression, + ) -> VortexResult> { // Rewrite the COALESCE-shaped CASE WHEN into `fill_null`, which references `x` // once and lowers to a single fill kernel instead of a `zip`/merge that resolves // `x` twice (once for the `is_null` predicate, once for the value branch). @@ -298,7 +297,7 @@ impl ScalarFnVTable for CaseWhen { return Ok(Some(x.clone())); } - Ok(Some(crate::expr::fill_null(x.clone(), fill.clone()))) + Ok(Some(bound_fill_null(x.clone(), fill.clone()))) } fn is_strict(&self, _options: &Self::Options) -> bool { @@ -461,6 +460,8 @@ 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::case_when; use crate::expr::case_when_no_else; use crate::expr::col; @@ -489,6 +490,17 @@ mod tests { .into_array() } + fn evaluate_bound_expr(expr: &BoundExpression, array: &ArrayRef) -> ArrayRef { + let mut ctx = SESSION.create_execution_ctx(); + array + .clone() + .apply_bound(expr) + .unwrap() + .execute::(&mut ctx) + .unwrap() + .into_array() + } + // ==================== Serialization Tests ==================== #[test] @@ -1304,7 +1316,9 @@ mod tests { fn test_simplify_coalesce_is_null_rewrites_to_fill_null() -> VortexResult<()> { // CASE WHEN is_null(x) THEN 0 ELSE x END ==> fill_null(x, 0) let expr = case_when(is_null(col("x")), lit(0i64), col("x")); - let optimized = expr.optimize_recursive(&nullable_i64_scope(&["x"]))?; + let optimized = expr + .bind(&nullable_i64_scope(&["x"]))? + .optimize_recursive()?; assert!( optimized.to_string().starts_with("vortex.fill_null"), "expected fill_null, got {optimized}" @@ -1316,7 +1330,9 @@ mod tests { fn test_simplify_coalesce_is_not_null_rewrites_to_fill_null() -> VortexResult<()> { // CASE WHEN is_not_null(x) THEN x ELSE 0 END ==> fill_null(x, 0) let expr = case_when(is_not_null(col("x")), col("x"), lit(0i64)); - let optimized = expr.optimize_recursive(&nullable_i64_scope(&["x"]))?; + let optimized = expr + .bind(&nullable_i64_scope(&["x"]))? + .optimize_recursive()?; assert!( optimized.to_string().starts_with("vortex.fill_null"), "expected fill_null, got {optimized}" @@ -1328,7 +1344,9 @@ mod tests { fn test_simplify_does_not_fire_when_operands_differ() -> VortexResult<()> { // The is_null operand (x) and the ELSE (y) are different columns: not a COALESCE. let expr = case_when(is_null(col("x")), lit(0i64), col("y")); - let optimized = expr.optimize_recursive(&nullable_i64_scope(&["x", "y"]))?; + let optimized = expr + .bind(&nullable_i64_scope(&["x", "y"]))? + .optimize_recursive()?; let s = optimized.to_string(); assert!(s.contains("CASE"), "expected CASE WHEN to remain, got {s}"); assert!(!s.contains("fill_null"), "must not rewrite, got {s}"); @@ -1340,7 +1358,9 @@ mod tests { // COALESCE(x, c) with a *column* fill: fill_null cannot consume a non-constant // fill value, so the rewrite must not fire. let expr = case_when(is_null(col("x")), col("c"), col("x")); - let optimized = expr.optimize_recursive(&nullable_i64_scope(&["x", "c"]))?; + let optimized = expr + .bind(&nullable_i64_scope(&["x", "c"]))? + .optimize_recursive()?; let s = optimized.to_string(); assert!(s.contains("CASE"), "expected CASE WHEN to remain, got {s}"); assert!(!s.contains("fill_null"), "must not rewrite, got {s}"); @@ -1363,7 +1383,9 @@ mod tests { case_when(is_null(col("x")), null_fill(), col("x")), case_when(is_not_null(col("x")), col("x"), null_fill()), ] { - let optimized = expr.optimize_recursive(&nullable_i64_scope(&["x"]))?; + let optimized = expr + .bind(&nullable_i64_scope(&["x"]))? + .optimize_recursive()?; assert_eq!( optimized.to_string(), "$.x", @@ -1385,7 +1407,7 @@ mod tests { ))); let original = case_when(is_null(root()), null_fill, root()); - let optimized = original.optimize_recursive(&scope)?; + let optimized = original.bind(&scope)?.optimize_recursive()?; assert_eq!( optimized.to_string(), "$", @@ -1394,14 +1416,16 @@ mod tests { let expected = PrimitiveArray::from_option_iter([Some(1i64), None, Some(3)]).into_array(); assert_arrays_eq!(evaluate_expr(&original, &array), expected, &mut ctx); - assert_arrays_eq!(evaluate_expr(&optimized, &array), expected, &mut ctx); + assert_arrays_eq!(evaluate_bound_expr(&optimized, &array), expected, &mut ctx); Ok(()) } #[test] fn test_simplify_does_not_fire_without_else() -> VortexResult<()> { let expr = case_when_no_else(is_null(col("x")), lit(0i64)); - let optimized = expr.optimize_recursive(&nullable_i64_scope(&["x"]))?; + let optimized = expr + .bind(&nullable_i64_scope(&["x"]))? + .optimize_recursive()?; assert!( !optimized.to_string().contains("fill_null"), "must not rewrite a no-ELSE case_when, got {optimized}" @@ -1418,7 +1442,9 @@ mod tests { ], Some(col("x")), ); - let optimized = expr.optimize_recursive(&nullable_i64_scope(&["x"]))?; + let optimized = expr + .bind(&nullable_i64_scope(&["x"]))? + .optimize_recursive()?; assert!( !optimized.to_string().contains("fill_null"), "must not rewrite a multi-pair case_when, got {optimized}" @@ -1434,7 +1460,7 @@ mod tests { let scope = DType::Primitive(PType::I64, Nullability::Nullable); let original = case_when(is_null(root()), lit(0i64), root()); - let optimized = original.optimize_recursive(&scope)?; + let optimized = original.bind(&scope)?.optimize_recursive()?; assert!( optimized.to_string().starts_with("vortex.fill_null"), "expected fill_null, got {optimized}" @@ -1448,7 +1474,7 @@ mod tests { &mut ctx ); assert_arrays_eq!( - evaluate_expr(&optimized, &array), + evaluate_bound_expr(&optimized, &array), buffer![1i64, 0, 3].into_array(), &mut ctx ); diff --git a/vortex-array/src/scalar_fn/fns/cast/mod.rs b/vortex-array/src/scalar_fn/fns/cast/mod.rs index 16802d22d32..f960e87ddb4 100644 --- a/vortex-array/src/scalar_fn/fns/cast/mod.rs +++ b/vortex-array/src/scalar_fn/fns/cast/mod.rs @@ -36,6 +36,8 @@ use crate::arrays::VarBinView; use crate::arrays::struct_::compute::cast::struct_cast; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; +use crate::expr::BoundExpression; +use crate::expr::bound::lit as bound_lit; use crate::expr::display::ExprDisplay; use crate::expr::expression::Expression; use crate::expr::lit; @@ -156,24 +158,24 @@ impl ScalarFnVTable for Cast { fn reduce(&self, target_dtype: &DType, node: &T) -> VortexResult> { // Collapse node if child is already the target type - let child = node.child(0); + let child = node.reduce_child(0); if &child.node_dtype()? == target_dtype { return Ok(Some(child)); } Ok(None) } - fn simplify_untyped( + fn simplify( &self, target_dtype: &DType, - expr: &Expression, - ) -> VortexResult> { + expr: &BoundExpression, + ) -> VortexResult> { let Some(scalar) = expr.child(0).as_opt::() else { return Ok(None); }; // A failing cast (e.g. null to a non-nullable dtype) is left in place so the error // surfaces at execution time rather than during optimization. - Ok(scalar.cast(target_dtype).ok().map(lit)) + Ok(scalar.cast(target_dtype).ok().map(bound_lit)) } fn validity(&self, dtype: &DType, expression: &Expression) -> VortexResult> { @@ -300,7 +302,7 @@ mod tests { lit(3i32), DType::Primitive(PType::F64, Nullability::NonNullable), ); - let optimized = expr.optimize(&test_harness::struct_dtype())?; + let optimized = expr.bind(&DType::Null)?.optimize()?; let scalar = optimized .as_opt::() @@ -320,7 +322,7 @@ mod tests { lit(decimal), DType::Primitive(PType::F64, Nullability::NonNullable), ); - let optimized = expr.optimize(&test_harness::struct_dtype())?; + let optimized = expr.bind(&DType::Null)?.optimize()?; let scalar = optimized .as_opt::() @@ -342,7 +344,7 @@ mod tests { ))), target.clone(), ); - let optimized = expr.optimize(&test_harness::struct_dtype())?; + let optimized = expr.bind(&DType::Null)?.optimize()?; assert!(optimized.as_opt::().is_none()); assert_eq!(optimized.as_opt::(), Some(&target)); diff --git a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs index 8536b4bf7da..00a9b38ab2e 100644 --- a/vortex-array/src/scalar_fn/fns/fill_null/mod.rs +++ b/vortex-array/src/scalar_fn/fns/fill_null/mod.rs @@ -22,6 +22,7 @@ use crate::arrays::Primitive; use crate::arrays::ScalarFnArray; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; +use crate::expr::BoundExpression; use crate::expr::Expression; use crate::scalar::Scalar; use crate::scalar_fn::Arity; @@ -123,10 +124,9 @@ impl ScalarFnVTable for FillNull { fn simplify( &self, _options: &Self::Options, - expr: &Expression, - ctx: &dyn crate::scalar_fn::SimplifyCtx, - ) -> VortexResult> { - let input_dtype = ctx.return_dtype(expr.child(0))?; + expr: &BoundExpression, + ) -> VortexResult> { + let input_dtype = expr.child(0).dtype(); if !input_dtype.is_nullable() { return Ok(Some(expr.child(0).clone())); diff --git a/vortex-array/src/scalar_fn/fns/get_item.rs b/vortex-array/src/scalar_fn/fns/get_item.rs index 986dd620460..92ab18cfe84 100644 --- a/vortex-array/src/scalar_fn/fns/get_item.rs +++ b/vortex-array/src/scalar_fn/fns/get_item.rs @@ -17,13 +17,10 @@ use crate::arrays::ScalarFnArray; use crate::arrays::StructArray; use crate::arrays::struct_::StructArrayExt; use crate::builtins::ArrayBuiltins; -use crate::builtins::ExprBuiltins; use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::Nullability; -use crate::expr::Expression; use crate::expr::display::ExprDisplay; -use crate::expr::lit; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::EmptyOptions; @@ -136,12 +133,12 @@ impl ScalarFnVTable for GetItem { } fn reduce(&self, field_name: &FieldName, node: &T) -> VortexResult> { - let child = node.child(0); + let child = node.reduce_child(0); if let Some(child_fn) = child.scalar_fn() && let Some(pack) = child_fn.as_opt::() && let Some(idx) = pack.names.find(field_name) { - let mut field = child.child(idx); + let mut field = child.reduce_child(idx); // Possibly mask the field if the pack is nullable if pack.nullability.is_nullable() { @@ -157,44 +154,6 @@ impl ScalarFnVTable for GetItem { Ok(None) } - fn simplify_untyped( - &self, - field_name: &FieldName, - expr: &Expression, - ) -> VortexResult> { - let child = expr.child(0); - - // If the child is a Pack expression, we can directly return the corresponding child. - if let Some(pack) = child.as_opt::() { - let idx = pack - .names - .iter() - .position(|name| name == field_name) - .ok_or_else(|| { - vortex_err!( - "Cannot find field {} in pack fields {:?}", - field_name, - pack.names - ) - })?; - - let mut field = child.child(idx).clone(); - - // It's useful to simplify this node without type info, but we need to make sure - // the nullability is correct. We cannot cast since we don't have the dtype info here, - // so instead we insert a Mask expression that we know converts a child's dtype to - // nullable. - if pack.nullability.is_nullable() { - // Mask with an all-true array to ensure the field DType is nullable. - field = field.mask(lit(true))?; - } - - return Ok(Some(field)); - } - - Ok(None) - } - fn is_strict(&self, _field_name: &FieldName) -> bool { true } @@ -216,7 +175,6 @@ mod tests { use crate::dtype::Nullability; use crate::dtype::Nullability::NonNullable; use crate::dtype::PType; - use crate::dtype::StructFields; use crate::expr::checked_add; use crate::expr::get_item; use crate::expr::lit; @@ -292,34 +250,32 @@ mod tests { } #[test] - fn test_pack_get_item_rule() { + fn test_pack_get_item_rule() -> VortexResult<()> { // Create: pack(a: lit(1), b: lit(2)).get_item("b") let pack_expr = pack([("a", lit(1)), ("b", lit(2))], NonNullable); let get_item_expr = get_item("b", pack_expr); - let result = get_item_expr - .optimize_recursive(&DType::Struct(StructFields::empty(), NonNullable)) - .unwrap(); + let result = get_item_expr.bind(&DType::Null)?.optimize_recursive()?; - assert_eq!(result, lit(2)); + assert_eq!(result, lit(2).bind(&DType::Null)?); + Ok(()) } #[test] - fn test_multi_level_pack_get_item_simplify() { + fn test_multi_level_pack_get_item_simplify() -> VortexResult<()> { let inner_pack = pack([("a", lit(1)), ("b", lit(2))], NonNullable); let get_a = get_item("a", inner_pack); let outer_pack = pack([("x", get_a), ("y", lit(3)), ("z", lit(4))], NonNullable); let get_z = get_item("z", outer_pack); - let dtype = DType::Primitive(PType::I32, NonNullable); - - let result = get_z.optimize_recursive(&dtype).unwrap(); - assert_eq!(result, lit(4)); + let result = get_z.bind(&DType::Null)?.optimize_recursive()?; + assert_eq!(result, lit(4).bind(&DType::Null)?); + Ok(()) } #[test] - fn test_deeply_nested_pack_get_item() { + fn test_deeply_nested_pack_get_item() -> VortexResult<()> { let innermost = pack([("a", lit(42))], NonNullable); let get_a = get_item("a", innermost); @@ -332,14 +288,13 @@ mod tests { let outermost = pack([("final", get_c)], NonNullable); let get_final = get_item("final", outermost); - let dtype = DType::Primitive(PType::I32, NonNullable); - - let result = get_final.optimize_recursive(&dtype).unwrap(); - assert_eq!(result, lit(42)); + let result = get_final.bind(&DType::Null)?.optimize_recursive()?; + assert_eq!(result, lit(42).bind(&DType::Null)?); + Ok(()) } #[test] - fn test_partial_pack_get_item_simplify() { + fn test_partial_pack_get_item_simplify() -> VortexResult<()> { let inner_pack = pack([("x", lit(1)), ("y", lit(2))], NonNullable); let get_x = get_item("x", inner_pack); let add_expr = checked_add(get_x, lit(10)); @@ -347,11 +302,10 @@ mod tests { let outer_pack = pack([("result", add_expr)], NonNullable); let get_result = get_item("result", outer_pack); - let dtype = DType::Primitive(PType::I32, NonNullable); - - let result = get_result.optimize_recursive(&dtype).unwrap(); + let result = get_result.bind(&DType::Null)?.optimize_recursive()?; let expected = checked_add(lit(1), lit(10)); - assert_eq!(&result, &expected); + assert_eq!(result, expected.bind(&DType::Null)?); + Ok(()) } #[test] diff --git a/vortex-array/src/scalar_fn/fns/mask/mod.rs b/vortex-array/src/scalar_fn/fns/mask/mod.rs index d2353b17091..7a0ad496daa 100644 --- a/vortex-array/src/scalar_fn/fns/mask/mod.rs +++ b/vortex-array/src/scalar_fn/fns/mask/mod.rs @@ -22,9 +22,10 @@ use crate::builtins::ArrayBuiltins; use crate::child_to_validity; use crate::dtype::DType; use crate::dtype::Nullability; +use crate::expr::BoundExpression; use crate::expr::Expression; use crate::expr::and; -use crate::expr::lit; +use crate::expr::bound::lit as bound_lit; use crate::scalar::Scalar; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -33,7 +34,6 @@ use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::SimplifyCtx; use crate::scalar_fn::fns::literal::Literal; /// An expression that masks an input based on a boolean mask. @@ -115,9 +115,8 @@ impl ScalarFnVTable for Mask { fn simplify( &self, _options: &Self::Options, - expr: &Expression, - ctx: &dyn SimplifyCtx, - ) -> VortexResult> { + expr: &BoundExpression, + ) -> VortexResult> { let Some(mask_lit) = expr.child(1).as_opt::() else { return Ok(None); }; @@ -132,8 +131,9 @@ impl ScalarFnVTable for Mask { Ok(Some(expr.child(0).clone())) } else { // Mask is all false, so the output is all nulls. - let input_dtype = ctx.return_dtype(expr.child(0))?; - Ok(Some(lit(Scalar::null(input_dtype.as_nullable())))) + Ok(Some(bound_lit(Scalar::null( + expr.child(0).dtype().as_nullable(), + )))) } } @@ -194,7 +194,7 @@ fn execute_canonical( #[cfg(test)] mod test { - use vortex_error::VortexExpect; + use vortex_error::VortexResult; use crate::dtype::DType; use crate::dtype::Nullability::Nullable; @@ -204,22 +204,19 @@ mod test { use crate::scalar::Scalar; #[test] - fn test_simplify() { + fn test_simplify() -> VortexResult<()> { let input_expr = lit(42u32); let true_mask_expr = lit(true); let false_mask_expr = lit(false); let mask_true_expr = mask(input_expr.clone(), true_mask_expr); - let simplified_true = mask_true_expr - .optimize(&DType::Null) - .vortex_expect("Simplification"); - assert_eq!(&simplified_true, &input_expr); + let simplified_true = mask_true_expr.bind(&DType::Null)?.optimize()?; + assert_eq!(simplified_true, input_expr.bind(&DType::Null)?); let mask_false_expr = mask(input_expr, false_mask_expr); - let simplified_false = mask_false_expr - .optimize(&DType::Null) - .vortex_expect("Simplification"); + let simplified_false = mask_false_expr.bind(&DType::Null)?.optimize()?; let expected_null_expr = lit(Scalar::null(DType::Primitive(PType::U32, Nullable))); - assert_eq!(&simplified_false, &expected_null_expr); + assert_eq!(simplified_false, expected_null_expr.bind(&DType::Null)?); + Ok(()) } } diff --git a/vortex-array/src/scalar_fn/fns/merge.rs b/vortex-array/src/scalar_fn/fns/merge.rs index 178075e936c..eb948f25a91 100644 --- a/vortex-array/src/scalar_fn/fns/merge.rs +++ b/vortex-array/src/scalar_fn/fns/merge.rs @@ -178,7 +178,7 @@ impl ScalarFnVTable for Merge { let mut children = Vec::with_capacity(node.child_count() * 2); let mut duplicate_names = HashSet::<_>::new(); - for child in (0..node.child_count()).map(|i| node.child(i)) { + for child in (0..node.child_count()).map(|i| node.reduce_child(i)) { let child_dtype = child.node_dtype()?; if !child_dtype.is_struct() { vortex_bail!( @@ -574,12 +574,12 @@ mod tests { DuplicateHandling::RightMost, ); - let result = e.optimize(&dtype).unwrap(); + let result = e.bind(&dtype).unwrap().optimize().unwrap(); assert!(result.is::()); assert_eq!( - result.return_dtype(&dtype).unwrap(), - DType::struct_([("a", I32), ("b", U32), ("c", U64)], NonNullable) + result.dtype(), + &DType::struct_([("a", I32), ("b", U32), ("c", U64)], NonNullable) ); } } diff --git a/vortex-array/src/scalar_fn/fns/select.rs b/vortex-array/src/scalar_fn/fns/select.rs index b0ac654a2f1..0b602f923a4 100644 --- a/vortex-array/src/scalar_fn/fns/select.rs +++ b/vortex-array/src/scalar_fn/fns/select.rs @@ -24,17 +24,16 @@ use crate::arrays::struct_::StructArrayExt; use crate::dtype::DType; use crate::dtype::FieldName; use crate::dtype::FieldNames; +use crate::expr::BoundExpression; +use crate::expr::bound::get_item as bound_get_item; +use crate::expr::bound::pack as bound_pack; use crate::expr::display::ExprDisplay; -use crate::expr::expression::Expression; use crate::expr::field::DisplayFieldNames; -use crate::expr::get_item; -use crate::expr::pack; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::SimplifyCtx; use crate::scalar_fn::fns::pack::Pack; #[derive(Debug, Clone, PartialEq, Eq, Hash)] @@ -170,11 +169,10 @@ impl ScalarFnVTable for Select { fn simplify( &self, selection: &FieldSelection, - expr: &Expression, - ctx: &dyn SimplifyCtx, - ) -> VortexResult> { + expr: &BoundExpression, + ) -> VortexResult> { let child_struct = expr.child(0); - let struct_dtype = ctx.return_dtype(child_struct)?; + let struct_dtype = child_struct.dtype(); let struct_nullability = struct_dtype.nullability(); let struct_fields = struct_dtype.as_struct_fields_opt().ok_or_else(|| { @@ -201,8 +199,10 @@ impl ScalarFnVTable for Select { // special-casing for pack, but not for select. We will fix this up when we revisit the // layout APIs. if included_fields.is_empty() { - let empty: Vec<(FieldName, Expression)> = vec![]; - return Ok(Some(pack(empty, struct_nullability))); + return Ok(Some(bound_pack( + Vec::<(FieldName, BoundExpression)>::new(), + struct_nullability, + ))); } // We cannot always convert a `select` into a `pack(get_item(f1), get_item(f2), ...)`. @@ -220,10 +220,10 @@ impl ScalarFnVTable for Select { struct_nullability.is_nullable() && !all_included_fields_are_nullable; if child_is_pack && !would_intersect_validity { - let pack_expr = pack( + let pack_expr = bound_pack( included_fields .into_iter() - .map(|name| (name.clone(), get_item(name, child_struct.clone()))), + .map(|name| (name.clone(), bound_get_item(name, child_struct.clone()))), struct_nullability, ); @@ -436,9 +436,9 @@ mod tests { ); let e = select(["a", "b"], root()); - let result = e.optimize_recursive(&dtype).unwrap(); + let result = e.bind(&dtype).unwrap().optimize_recursive().unwrap(); - assert!(result.return_dtype(&dtype).unwrap().is_nullable()); + assert!(result.dtype().is_nullable()); } #[test] @@ -454,10 +454,10 @@ mod tests { ); let e = select_exclude(["c"], root()); - let result = e.optimize_recursive(&dtype).unwrap(); + let result = e.bind(&dtype).unwrap().optimize_recursive().unwrap(); // Should exclude "c" and include "a" and "b" - let result_dtype = result.return_dtype(&dtype).unwrap(); + let result_dtype = result.dtype(); assert!(result_dtype.is_nullable()); let fields = result_dtype.as_struct_fields_opt().unwrap(); assert_eq!(fields.names().as_ref(), &["a", "b"]); diff --git a/vortex-array/src/scalar_fn/fns/zip/mod.rs b/vortex-array/src/scalar_fn/fns/zip/mod.rs index 440a7e99b5a..03bcd5bdbbe 100644 --- a/vortex-array/src/scalar_fn/fns/zip/mod.rs +++ b/vortex-array/src/scalar_fn/fns/zip/mod.rs @@ -27,7 +27,7 @@ use crate::builders::builder_with_capacity; use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::StructFields; -use crate::expr::Expression; +use crate::expr::BoundExpression; use crate::expr::display::ExprDisplay; use crate::scalar_fn::Arity; use crate::scalar_fn::ChildName; @@ -36,7 +36,6 @@ use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; -use crate::scalar_fn::SimplifyCtx; use crate::scalar_fn::fns::literal::Literal; use crate::validity::Validity; @@ -159,9 +158,8 @@ impl ScalarFnVTable for Zip { fn simplify( &self, _options: &Self::Options, - expr: &Expression, - _ctx: &dyn SimplifyCtx, - ) -> VortexResult> { + expr: &BoundExpression, + ) -> VortexResult> { let Some(mask_lit) = expr.child(2).as_opt::() else { return Ok(None); }; diff --git a/vortex-array/src/scalar_fn/typed.rs b/vortex-array/src/scalar_fn/typed.rs index 84670b537d7..d5f6ed39eff 100644 --- a/vortex-array/src/scalar_fn/typed.rs +++ b/vortex-array/src/scalar_fn/typed.rs @@ -23,17 +23,16 @@ use vortex_error::VortexResult; 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::Arity; use crate::scalar_fn::ArrayReduceNode; use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; -use crate::scalar_fn::ExpressionReduceNode; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnRef; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::SimplifyCtx; /// A typed scalar function instance, parameterized by a concrete [`ScalarFnVTable`]. /// @@ -81,10 +80,10 @@ pub(super) trait DynScalarFn: 'static + Send + Sync + super::sealed::Sealed { // Bound methods — options accessed from self fn execute(&self, args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx) -> VortexResult; fn return_dtype(&self, arg_types: &[DType]) -> VortexResult; - fn reduce_expression<'a>( + fn reduce_bound_expression( &self, - node: &ExpressionReduceNode<'a>, - ) -> VortexResult>>; + node: &BoundExpression, + ) -> VortexResult>; fn reduce_array<'a>( &self, node: &ArrayReduceNode<'a>, @@ -96,12 +95,7 @@ pub(super) trait DynScalarFn: 'static + Send + Sync + super::sealed::Sealed { // Expression methods — take expressions for tree traversal fn fmt_sql(&self, expression: &dyn ExprDisplay, f: &mut Formatter<'_>) -> fmt::Result; - fn simplify( - &self, - expression: &Expression, - ctx: &dyn SimplifyCtx, - ) -> VortexResult>; - fn simplify_untyped(&self, expression: &Expression) -> VortexResult>; + fn simplify(&self, expression: &BoundExpression) -> VortexResult>; fn validity(&self, expression: &Expression) -> VortexResult>; // Options operations — self-contained @@ -166,10 +160,10 @@ impl DynScalarFn for TypedScalarFnInstance { V::return_dtype(&self.vtable, &self.options, arg_dtypes) } - fn reduce_expression<'a>( + fn reduce_bound_expression( &self, - node: &ExpressionReduceNode<'a>, - ) -> VortexResult>> { + node: &BoundExpression, + ) -> VortexResult> { V::reduce(&self.vtable, &self.options, node) } @@ -200,16 +194,8 @@ impl DynScalarFn for TypedScalarFnInstance { V::fmt_sql(&self.vtable, &self.options, expression, f) } - fn simplify( - &self, - expression: &Expression, - ctx: &dyn SimplifyCtx, - ) -> VortexResult> { - V::simplify(&self.vtable, &self.options, expression, ctx) - } - - fn simplify_untyped(&self, expression: &Expression) -> VortexResult> { - V::simplify_untyped(&self.vtable, &self.options, expression) + fn simplify(&self, expression: &BoundExpression) -> VortexResult> { + V::simplify(&self.vtable, &self.options, expression) } fn validity(&self, expression: &Expression) -> VortexResult> { diff --git a/vortex-array/src/scalar_fn/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index da67f6231d7..34ce4228a68 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -138,25 +138,12 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { Ok(None) } - /// Simplify the expression if possible. + /// Simplify a bound expression if possible. fn simplify( &self, options: &Self::Options, - expr: &Expression, - ctx: &dyn SimplifyCtx, - ) -> VortexResult> { - _ = options; - _ = expr; - _ = ctx; - Ok(None) - } - - /// Simplify the expression if possible, without type information. - fn simplify_untyped( - &self, - options: &Self::Options, - expr: &Expression, - ) -> VortexResult> { + expr: &BoundExpression, + ) -> VortexResult> { _ = options; _ = expr; Ok(None) @@ -225,10 +212,8 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { /// A node used for implementing abstract reduction rules over a tree of scalar functions. /// /// Reduction rules are generic over the node type, so a rule is written once and monomorphized -/// per reducible tree kind: [`ExpressionReduceNode`] for expression trees and -/// [`ArrayReduceNode`] for array trees. Nodes borrow from the tree being reduced, making -/// traversal allocation-free, while nodes produced by [`ReduceNode::new_node`] own their -/// freshly-built subtrees. +/// per reducible tree kind: [`BoundExpression`] for bound expression trees and +/// [`ArrayReduceNode`] for array trees. pub trait ReduceNode: Clone { /// Return the data type of this node. fn node_dtype(&self) -> VortexResult; @@ -236,8 +221,8 @@ pub trait ReduceNode: Clone { /// Return this node's scalar function if it is indeed a scalar fn. fn scalar_fn(&self) -> Option<&ScalarFnRef>; - /// Descend to the child of this node. - fn child(&self, idx: usize) -> Self; + /// Clone a child of this node for use in a reduction. + fn reduce_child(&self, idx: usize) -> Self; /// Returns the number of children of this node. fn child_count(&self) -> usize; @@ -247,69 +232,25 @@ pub trait ReduceNode: Clone { fn new_node(&self, scalar_fn: ScalarFnRef, children: &[Self]) -> VortexResult; } -/// A [`ReduceNode`] over an expression tree, typed within a scope. -#[derive(Clone)] -pub struct ExpressionReduceNode<'a> { - expression: Cow<'a, Expression>, - scope: &'a DType, -} - -impl<'a> ExpressionReduceNode<'a> { - /// Creates a node borrowing the given expression and scope. - pub fn new(expression: &'a Expression, scope: &'a DType) -> Self { - Self { - expression: Cow::Borrowed(expression), - scope, - } - } - - /// Returns the expression backing this node. - pub fn expression(&self) -> &Expression { - &self.expression - } - - /// Consumes this node and returns the backing expression. - pub fn into_expression(self) -> Expression { - self.expression.into_owned() - } -} - -impl ReduceNode for ExpressionReduceNode<'_> { +impl ReduceNode for BoundExpression { fn node_dtype(&self) -> VortexResult { - self.expression.return_dtype(self.scope) + Ok(self.dtype().clone()) } fn scalar_fn(&self) -> Option<&ScalarFnRef> { - self.expression.as_scalar() + self.as_scalar() } - fn child(&self, idx: usize) -> Self { - let expression = match &self.expression { - Cow::Borrowed(expression) => Cow::Borrowed(expression.child(idx)), - Cow::Owned(expression) => Cow::Owned(expression.child(idx).clone()), - }; - Self { - expression, - scope: self.scope, - } + fn reduce_child(&self, idx: usize) -> Self { + self.child(idx).clone() } fn child_count(&self) -> usize { - self.expression.children().len() + self.children().len() } fn new_node(&self, scalar_fn: ScalarFnRef, children: &[Self]) -> VortexResult { - let expression = Expression::try_new( - scalar_fn, - children - .iter() - .map(|c| c.expression.as_ref().clone()) - .collect::>(), - )?; - Ok(Self { - expression: Cow::Owned(expression), - scope: self.scope, - }) + BoundExpression::try_new(scalar_fn, children.iter().cloned()) } } @@ -349,7 +290,7 @@ impl ReduceNode for ArrayReduceNode<'_> { .map(|a| a.data().scalar_fn()) } - fn child(&self, idx: usize) -> Self { + fn reduce_child(&self, idx: usize) -> Self { let array = match &self.array { Cow::Borrowed(array) => Cow::Borrowed( array @@ -422,14 +363,6 @@ impl Arity { } } -/// Context for simplification. -/// -/// Used to lazily compute input data types where simplification requires them. -pub trait SimplifyCtx { - /// Get the data type of the given expression. - fn return_dtype(&self, expr: &Expression) -> VortexResult; -} - /// Arguments for expression execution. pub trait ExecutionArgs { /// Returns the input array at the given index. diff --git a/vortex-bench/src/datasets/tpch_l_comment.rs b/vortex-bench/src/datasets/tpch_l_comment.rs index c57bc91a65d..085ca75c2a0 100644 --- a/vortex-bench/src/datasets/tpch_l_comment.rs +++ b/vortex-bench/src/datasets/tpch_l_comment.rs @@ -67,8 +67,8 @@ 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())?; + .bind(file.dtype())? + .optimize_recursive()?; let chunks: Vec<_> = file .scan()? .with_projection(projection) diff --git a/vortex-datafusion/src/persistent/opener.rs b/vortex-datafusion/src/persistent/opener.rs index 89abfe68d18..754caeff4db 100644 --- a/vortex-datafusion/src/persistent/opener.rs +++ b/vortex-datafusion/src/persistent/opener.rs @@ -300,8 +300,8 @@ 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_projection = scan_projection - .optimize_recursive(vxf.dtype()) - .and_then(|projection| projection.bind(vxf.dtype())) + .bind(vxf.dtype()) + .and_then(|projection| projection.optimize_recursive()) .map_err(|_e| { exec_datafusion_err!("Couldn't get the dtype for the underlying Vortex scan") })?; @@ -393,7 +393,7 @@ impl FileOpener for VortexOpener { }) .transpose()?; let filter = filter - .map(|filter| filter.optimize_recursive(vxf.dtype())?.bind(vxf.dtype())) + .map(|filter| filter.bind(vxf.dtype())?.optimize_recursive()) .transpose() .map_err(|e| exec_datafusion_err!("Couldn't bind Vortex scan filter: {e}"))?; diff --git a/vortex-duckdb/src/projection.rs b/vortex-duckdb/src/projection.rs index ef91aaae52c..d6778acd8af 100644 --- a/vortex-duckdb/src/projection.rs +++ b/vortex-duckdb/src/projection.rs @@ -238,7 +238,7 @@ impl Filter { }; let filter = and_collect(table_filter_exprs) - .map(|expr| expr.optimize_recursive(dtype)?.bind(dtype)) + .map(|expr| expr.bind(dtype)?.optimize_recursive()) .transpose()?; let out = Self { diff --git a/vortex-duckdb/src/table_function.rs b/vortex-duckdb/src/table_function.rs index d7d45ef8798..9e19b9583aa 100644 --- a/vortex-duckdb/src/table_function.rs +++ b/vortex-duckdb/src/table_function.rs @@ -371,7 +371,7 @@ pub fn init_local(bind_data: &BindState, global: &GlobalState) -> LocalState { } pub(crate) fn optimize_and_bind(expr: Expression, dtype: &DType) -> VortexResult { - expr.optimize_recursive(dtype)?.bind(dtype) + expr.bind(dtype)?.optimize_recursive() } pub(crate) fn convert_result(array: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { diff --git a/vortex-file/src/tests.rs b/vortex-file/src/tests.rs index f5c177c9cdf..d3f1d4dda82 100644 --- a/vortex-file/src/tests.rs +++ b/vortex-file/src/tests.rs @@ -116,8 +116,8 @@ fn strict_sorted(indices: Buffer) -> StrictSortedBuffer { } fn bind_scan_expr(file: &VortexFile, expr: Expression) -> BoundExpression { - expr.optimize_recursive(file.dtype()) - .and_then(|expr| expr.bind(file.dtype())) + expr.bind(file.dtype()) + .and_then(|expr| expr.optimize_recursive()) .vortex_expect("scan expression should bind") } @@ -1399,8 +1399,8 @@ async fn scan_empty_fields() -> VortexResult<()> { }, [], ) - .optimize_recursive(array.dtype())? - .bind(array.dtype())?; + .bind(array.dtype())? + .optimize_recursive()?; let result = round_trip(&array.clone().into_array(), |scan| { Ok(scan.with_projection(projection)) @@ -2228,9 +2228,7 @@ async fn timestamp_unit_mismatch() -> Result<(), Box> { ); let file = SESSION.open_options().open_buffer(buf)?; - let filter = filter_expr - .optimize_recursive(file.dtype())? - .bind(file.dtype())?; + let filter = filter_expr.bind(file.dtype())?.optimize_recursive()?; let mut stream = file.scan()?.with_filter(filter).into_array_stream()?; let result = stream.try_next().await; @@ -2279,9 +2277,7 @@ async fn timestamp_unit_mismatch_errors_with_constant_children() ); let file = SESSION.open_options().open_buffer(buf)?; - let filter = filter_expr - .optimize_recursive(file.dtype())? - .bind(file.dtype())?; + let filter = filter_expr.bind(file.dtype())?.optimize_recursive()?; let stream = file.scan()?.with_filter(filter).into_array_stream()?; let results = stream.try_collect::>().await; @@ -2786,9 +2782,7 @@ async fn repro_8166_binary_gt_all_ff_max() -> VortexResult<()> { ); let file = SESSION.open_options().open_buffer(buf)?; - let filter = filter - .optimize_recursive(file.dtype())? - .bind(file.dtype())?; + let filter = filter.bind(file.dtype())?.optimize_recursive()?; let result = file .scan()? .with_filter(filter) diff --git a/vortex-layout/src/plan/tests.rs b/vortex-layout/src/plan/tests.rs index 2b87f7ae3c5..783eb455327 100644 --- a/vortex-layout/src/plan/tests.rs +++ b/vortex-layout/src/plan/tests.rs @@ -58,16 +58,12 @@ fn make_plan(layout: LayoutRef) -> VortexResult { } fn make_eval(expression: Expression, child: PlanRef) -> VortexResult { - let expression = expression - .optimize_recursive(child.dtype())? - .bind(child.dtype())?; + let expression = expression.bind(child.dtype())?.optimize_recursive()?; EvalPlan::try_new(expression, child) } fn make_row_idx_plan(expression: Expression, child: PlanRef) -> VortexResult { - let expression = expression - .optimize_recursive(child.dtype())? - .bind(child.dtype())?; + let expression = expression.bind(child.dtype())?.optimize_recursive()?; plan_row_idx_expression(expression, child) } diff --git a/vortex-layout/src/scan/layout.rs b/vortex-layout/src/scan/layout.rs index 94ce7acf277..d5af9a4a037 100644 --- a/vortex-layout/src/scan/layout.rs +++ b/vortex-layout/src/scan/layout.rs @@ -117,14 +117,11 @@ impl DataSource for LayoutReaderDataSource { let projection = scan_request .projection - .optimize_recursive(self.reader.dtype())? - .bind(self.reader.dtype())?; + .bind(self.reader.dtype())? + .optimize_recursive()?; let filter = scan_request .filter - .map(|expr| { - expr.optimize_recursive(self.reader.dtype())? - .bind(self.reader.dtype()) - }) + .map(|expr| expr.bind(self.reader.dtype())?.optimize_recursive()) .transpose()?; let dtype = projection.dtype().clone(); diff --git a/vortex-layout/src/scan/multi.rs b/vortex-layout/src/scan/multi.rs index f0b7a506e08..77b5bd4c894 100644 --- a/vortex-layout/src/scan/multi.rs +++ b/vortex-layout/src/scan/multi.rs @@ -348,9 +348,9 @@ impl BoundScanRequest { } = request; Ok(Self { - projection: projection.optimize_recursive(dtype)?.bind(dtype)?, + projection: projection.bind(dtype)?.optimize_recursive()?, filter: filter - .map(|expr| expr.optimize_recursive(dtype)?.bind(dtype)) + .map(|expr| expr.bind(dtype)?.optimize_recursive()) .transpose() .map_err(Arc::new), row_range, diff --git a/vortex-python/src/dataset.rs b/vortex-python/src/dataset.rs index b9fde0ac806..d8adfd14832 100644 --- a/vortex-python/src/dataset.rs +++ b/vortex-python/src/dataset.rs @@ -59,15 +59,9 @@ pub fn read_array_from_reader( row_range: Option<(u64, u64)>, ctx: &mut ExecutionCtx, ) -> VortexResult { - let projection = projection - .optimize_recursive(vortex_file.dtype())? - .bind(vortex_file.dtype())?; + let projection = projection.bind(vortex_file.dtype())?.optimize_recursive()?; let filter = filter - .map(|filter| { - filter - .optimize_recursive(vortex_file.dtype())? - .bind(vortex_file.dtype()) - }) + .map(|filter| filter.bind(vortex_file.dtype())?.optimize_recursive()) .transpose()?; let mut scan = vortex_file.scan()?.with_projection(projection); @@ -197,11 +191,9 @@ impl PyVortexDataset { let filter = filter_from_python(row_filter); let reader = self_.py().detach(move || { - let projection = projection - .optimize_recursive(vxf.dtype())? - .bind(vxf.dtype())?; + let projection = projection.bind(vxf.dtype())?.optimize_recursive()?; let filter = filter - .map(|filter| filter.optimize_recursive(vxf.dtype())?.bind(vxf.dtype())) + .map(|filter| filter.bind(vxf.dtype())?.optimize_recursive()) .transpose()?; let mut scan = vxf .scan()? @@ -244,10 +236,10 @@ impl PyVortexDataset { let filter = filter_from_python(row_filter); let n_rows: usize = self_.py().detach(move || { let projection = select(FieldNames::empty(), root()) - .optimize_recursive(vxf.dtype())? - .bind(vxf.dtype())?; + .bind(vxf.dtype())? + .optimize_recursive()?; let filter = filter - .map(|filter| filter.optimize_recursive(vxf.dtype())?.bind(vxf.dtype())) + .map(|filter| filter.bind(vxf.dtype())?.optimize_recursive()) .transpose()?; let mut scan = vxf .scan()? diff --git a/vortex-python/src/file.rs b/vortex-python/src/file.rs index 8786546dce4..d2a797eb32e 100644 --- a/vortex-python/src/file.rs +++ b/vortex-python/src/file.rs @@ -229,17 +229,13 @@ impl PyVortexFile { let runtime = current_runtime(); let reader = slf.py().detach(|| { let filter = expr - .map(|e| { - e.into_inner() - .optimize_recursive(vxf.dtype())? - .bind(vxf.dtype()) - }) + .map(|e| e.into_inner().bind(vxf.dtype())?.optimize_recursive()) .transpose()?; let projection = projection .map(|p| p.0) .unwrap_or_else(root) - .optimize_recursive(vxf.dtype())? - .bind(vxf.dtype())?; + .bind(vxf.dtype())? + .optimize_recursive()?; let mut builder = vxf .scan()? .with_some_filter(filter) @@ -290,10 +286,10 @@ fn scan_builder( ) -> VortexResult> { let projection = projection .unwrap_or_else(root) - .optimize_recursive(vxf.dtype())? - .bind(vxf.dtype())?; + .bind(vxf.dtype())? + .optimize_recursive()?; let expr = expr - .map(|expr| expr.optimize_recursive(vxf.dtype())?.bind(vxf.dtype())) + .map(|expr| expr.bind(vxf.dtype())?.optimize_recursive()) .transpose()?; let mut builder = vxf .scan()? diff --git a/vortex/src/lib.rs b/vortex/src/lib.rs index 1d1ab1252ac..03f2a7b31cc 100644 --- a/vortex/src/lib.rs +++ b/vortex/src/lib.rs @@ -86,8 +86,8 @@ //! .open_options() //! .open_buffer(bytes)?; //! let filter = gt(root(), lit(2u64)) -//! .optimize_recursive(file.dtype())? -//! .bind(file.dtype())?; +//! .bind(file.dtype())? +//! .optimize_recursive()?; //! let filtered = file //! .scan()? //! .with_filter(filter) @@ -449,8 +449,8 @@ mod test { // [read] let file = session.open_options().open_path(path.clone()).await?; let filter = gt(root(), lit(2u64)) - .optimize_recursive(file.dtype())? - .bind(file.dtype())?; + .bind(file.dtype())? + .optimize_recursive()?; let array = file .scan()? .with_filter(filter) @@ -549,8 +549,8 @@ mod test { // Read the file back, but project down to just the "value" column. let file = session.open_options().open_path(path.clone()).await?; let projection = select(["value"], root()) - .optimize_recursive(file.dtype())? - .bind(file.dtype())?; + .bind(file.dtype())? + .optimize_recursive()?; let projected = file .scan()? .with_projection(projection)