diff --git a/vortex-array/Cargo.toml b/vortex-array/Cargo.toml index 16f2db98554..1e827e61302 100644 --- a/vortex-array/Cargo.toml +++ b/vortex-array/Cargo.toml @@ -186,6 +186,11 @@ name = "expr_optimize" path = "benches/expr/optimize_bench.rs" harness = false +[[bench]] +name = "expr_optimizer_comparison" +path = "benches/expr/optimizer_comparison.rs" +harness = false + [[bench]] name = "expr_optimize_predicate" path = "benches/expr/optimize_predicate.rs" diff --git a/vortex-array/benches/expr/optimizer_comparison.rs b/vortex-array/benches/expr/optimizer_comparison.rs new file mode 100644 index 00000000000..78a97db655c --- /dev/null +++ b/vortex-array/benches/expr/optimizer_comparison.rs @@ -0,0 +1,309 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Compares the legacy `Expression` optimizer with the rule-driven `BoundExpression` optimizer. +//! +//! Expressions are constructed and bound outside the timed region. `builtins` varies both tree +//! size and the number of nodes that the built-in rules can rewrite. `rule_dispatch` isolates the +//! bound optimizer and varies the number of rules checked before a successful rewrite. + +#![expect(clippy::unwrap_used)] + +use std::fmt::Display; +use std::fmt::Formatter; +use std::fmt::Result as FmtResult; + +use divan::Bencher; +use divan::black_box; +use divan::counter::ItemsCount; +use mimalloc::MiMalloc; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::dtype::StructFields; +use vortex_array::expr::BoundExpression; +use vortex_array::expr::BoundExpressionOptimizer; +use vortex_array::expr::Expression; +use vortex_array::expr::ExpressionId; +use vortex_array::expr::OptimizerRule; +use vortex_array::expr::OptimizerRuleRegistry; +use vortex_array::expr::and; +use vortex_array::expr::col; +use vortex_array::expr::eq; +use vortex_array::expr::lit; +use vortex_array::expr::or_collect; +use vortex_array::scalar_fn::ScalarFnVTable; +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_error::VortexResult; + +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +fn main() { + divan::main(); +} + +fn struct_scope() -> DType { + DType::Struct( + StructFields::new( + ["x"].into(), + vec![DType::Primitive(PType::I32, Nullability::NonNullable)], + ), + Nullability::NonNullable, + ) +} + +#[derive(Clone, Copy, Debug)] +struct RewriteCase { + terms: usize, + rewrite_sites: usize, +} + +impl RewriteCase { + fn node_count(self) -> usize { + // Each term is `eq(get_item("x", root()), literal)`, the terms are joined by OR nodes, + // and each rewrite site adds `and(term, true)`. + 5 * self.terms - 1 + 2 * self.rewrite_sites + } +} + +impl Display for RewriteCase { + fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { + write!( + f, + "nodes={}, rewrites={}", + self.node_count(), + self.rewrite_sites + ) + } +} + +const REWRITE_CASES: &[RewriteCase] = &[ + RewriteCase { + terms: 1, + rewrite_sites: 0, + }, + RewriteCase { + terms: 1, + rewrite_sites: 1, + }, + RewriteCase { + terms: 16, + rewrite_sites: 0, + }, + RewriteCase { + terms: 16, + rewrite_sites: 4, + }, + RewriteCase { + terms: 16, + rewrite_sites: 16, + }, + RewriteCase { + terms: 128, + rewrite_sites: 0, + }, + RewriteCase { + terms: 128, + rewrite_sites: 32, + }, + RewriteCase { + terms: 128, + rewrite_sites: 128, + }, + RewriteCase { + terms: 512, + rewrite_sites: 0, + }, + RewriteCase { + terms: 512, + rewrite_sites: 128, + }, + RewriteCase { + terms: 512, + rewrite_sites: 512, + }, +]; + +const NO_REWRITE_CASES: &[RewriteCase] = &[ + REWRITE_CASES[0], + REWRITE_CASES[2], + REWRITE_CASES[5], + REWRITE_CASES[8], +]; + +fn build_expression(case: RewriteCase) -> Expression { + or_collect((0..case.terms).map(|idx| { + let term = eq(col("x"), lit(i32::try_from(idx).unwrap())); + if idx < case.rewrite_sites { + and(term, lit(true)) + } else { + term + } + })) + .unwrap() +} + +mod builtins { + use super::*; + + #[divan::bench(args = REWRITE_CASES)] + fn expression(bencher: Bencher, case: &RewriteCase) { + let scope = struct_scope(); + let expr = build_expression(*case); + + bencher + .counter(ItemsCount::new(case.node_count())) + .bench(|| black_box(expr.optimize_recursive(&scope).unwrap())); + } + + #[divan::bench(args = REWRITE_CASES)] + fn bound_expression(bencher: Bencher, case: &RewriteCase) { + let scope = struct_scope(); + let unbound = build_expression(*case); + let expr = unbound.bind(&scope).unwrap(); + let optimizer = BoundExpressionOptimizer::default(); + + let expected = unbound + .optimize_recursive(&scope) + .unwrap() + .bind(&scope) + .unwrap(); + assert_eq!(optimizer.optimize(&expr).unwrap(), expected); + + bencher + .counter(ItemsCount::new(case.node_count())) + .bench(|| black_box(optimizer.optimize(&expr).unwrap())); + } + + #[divan::bench(args = NO_REWRITE_CASES)] + fn bound_expression_empty(bencher: Bencher, case: &RewriteCase) { + let scope = struct_scope(); + let expr = build_expression(*case).bind(&scope).unwrap(); + let optimizer = BoundExpressionOptimizer::new(OptimizerRuleRegistry::empty()); + + bencher + .counter(ItemsCount::new(case.node_count())) + .bench(|| black_box(optimizer.optimize(&expr).unwrap())); + } +} + +/// A binary rule that either declines every node or simplifies `value AND true`. +#[derive(Debug)] +struct AndTrueRule { + enabled: bool, +} + +impl OptimizerRule for AndTrueRule { + fn expression_id(&self) -> ExpressionId { + Binary.id() + } + + fn rewrite(&self, expr: &BoundExpression) -> VortexResult> { + if !self.enabled || expr.as_opt::() != Some(&Operator::And) { + return Ok(None); + } + let rhs_is_true = expr + .child(1) + .as_opt::() + .and_then(|scalar| scalar.as_bool_opt()) + .is_some_and(|value| value.value() == Some(true)); + Ok(rhs_is_true.then(|| expr.child(0).clone())) + } +} + +#[derive(Clone, Copy, Debug)] +struct RuleDispatchCase { + terms: usize, + candidate_rules: usize, +} + +impl RuleDispatchCase { + fn rewrite_case(self) -> RewriteCase { + RewriteCase { + terms: self.terms, + rewrite_sites: self.terms, + } + } +} + +impl Display for RuleDispatchCase { + fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult { + write!( + f, + "nodes={}, candidate_rules={}", + self.rewrite_case().node_count(), + self.candidate_rules + ) + } +} + +const RULE_DISPATCH_CASES: &[RuleDispatchCase] = &[ + RuleDispatchCase { + terms: 16, + candidate_rules: 1, + }, + RuleDispatchCase { + terms: 16, + candidate_rules: 4, + }, + RuleDispatchCase { + terms: 16, + candidate_rules: 16, + }, + RuleDispatchCase { + terms: 16, + candidate_rules: 64, + }, + RuleDispatchCase { + terms: 128, + candidate_rules: 1, + }, + RuleDispatchCase { + terms: 128, + candidate_rules: 4, + }, + RuleDispatchCase { + terms: 128, + candidate_rules: 16, + }, + RuleDispatchCase { + terms: 128, + candidate_rules: 64, + }, + RuleDispatchCase { + terms: 512, + candidate_rules: 1, + }, + RuleDispatchCase { + terms: 512, + candidate_rules: 4, + }, + RuleDispatchCase { + terms: 512, + candidate_rules: 16, + }, + RuleDispatchCase { + terms: 512, + candidate_rules: 64, + }, +]; + +#[divan::bench(args = RULE_DISPATCH_CASES)] +fn rule_dispatch(bencher: Bencher, case: &RuleDispatchCase) { + let rewrite_case = case.rewrite_case(); + let scope = struct_scope(); + let expr = build_expression(rewrite_case).bind(&scope).unwrap(); + let mut registry = OptimizerRuleRegistry::empty(); + for _ in 1..case.candidate_rules { + registry.register(AndTrueRule { enabled: false }); + } + registry.register(AndTrueRule { enabled: true }); + let optimizer = BoundExpressionOptimizer::new(registry); + + bencher + .counter(ItemsCount::new(rewrite_case.node_count())) + .bench(|| black_box(optimizer.optimize(&expr).unwrap())); +} diff --git a/vortex-array/src/expr/bound_expression.rs b/vortex-array/src/expr/bound_expression.rs index 4bd276e191d..fb02dacd625 100644 --- a/vortex-array/src/expr/bound_expression.rs +++ b/vortex-array/src/expr/bound_expression.rs @@ -13,9 +13,11 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; use vortex_session::VortexSession; +use vortex_session::registry::CachedId; use crate::dtype::DType; use crate::expr::Expression; +use crate::expr::ExpressionId; use crate::expr::display::DisplayTreeExpr; use crate::expr::scope::Scope; use crate::expr::traversal::TraversalOrder; @@ -106,6 +108,17 @@ impl Hash for ExactBoundExpr { } impl BoundExpression { + /// Return the globally unique ID of this expression node implementation. + pub fn id(&self) -> ExpressionId { + match self { + Self::Scalar { scalar_fn, .. } => scalar_fn.id(), + Self::Root { .. } => { + static ID: CachedId = CachedId::new("vortex.expr.root"); + *ID + } + } + } + /// Create a bound root expression with the given dtype. pub fn new_root(dtype: DType) -> Self { Self::Root { dtype } diff --git a/vortex-array/src/expr/mod.rs b/vortex-array/src/expr/mod.rs index 59fd21c46ee..f20a4012f1b 100644 --- a/vortex-array/src/expr/mod.rs +++ b/vortex-array/src/expr/mod.rs @@ -44,6 +44,7 @@ use std::hash::Hasher; use std::sync::Arc; use vortex_error::VortexExpect; +use vortex_session::registry::Id; use vortex_utils::aliases::hash_set::HashSet; use crate::dtype::FieldName; @@ -63,6 +64,7 @@ mod exprs; pub(crate) mod field; pub mod forms; mod optimize; +pub mod optimizer; pub mod proto; pub mod scope; pub mod stats; @@ -120,8 +122,18 @@ pub use exprs::select_exclude; pub use exprs::union_child_validities; pub use exprs::variant_get; pub use exprs::zip_expr; +pub use optimizer::BoundExpressionOptimizer; +pub use optimizer::OptimizerRule; +pub use optimizer::OptimizerRuleRef; +pub use optimizer::OptimizerRuleRegistry; pub use scope::*; +/// A globally unique identifier for an expression node implementation. +/// +/// Scalar-function nodes reuse their scalar-function ID. Other expression node implementations, +/// such as higher-order functions and lambda nodes, use IDs from the same global namespace. +pub type ExpressionId = Id; + pub trait VortexExprExt { /// Accumulate all field references from this expression and its children in a set fn field_references(&self) -> HashSet; diff --git a/vortex-array/src/expr/optimizer/mod.rs b/vortex-array/src/expr/optimizer/mod.rs new file mode 100644 index 00000000000..ddbae4314f7 --- /dev/null +++ b/vortex-array/src/expr/optimizer/mod.rs @@ -0,0 +1,458 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Rule-driven optimization for [`BoundExpression`] trees. +//! +//! This optimizer is independent from the legacy [`Expression`](super::Expression) optimizer. +//! Rules operate only on already-bound expressions, which makes every node's dtype available in +//! constant time and lets the driver verify that rewrites preserve the tree's type proof. +//! +//! # How optimization works +//! +//! [`OptimizerRuleRegistry`] holds the reusable rewrite rules. A +//! [`BoundExpressionOptimizer`] takes ownership of a configured registry and applies its rules +//! with a configurable rewrite limit. Each call to [`BoundExpressionOptimizer::try_optimize`] +//! creates an `OptimizationRun` containing a reference to the registry and the mutable state for +//! that traversal. +//! +//! A run recursively walks the tree with copy-on-write rebuilding: +//! +//! 1. Rules registered for a node's expression ID run in registration order before its children. +//! The first matching rule wins, and root rewrites repeat to convergence. +//! 2. Children are optimized in evaluation order. The node is rebuilt only if a child changed. +//! 3. A rebuilt node is rewritten again. Any replacement is walked as a new subtree so rules may +//! safely introduce expressions that need further optimization. +//! +//! Every replacement must preserve the node's dtype and differ from the expression it replaces. +//! A per-run rewrite limit terminates rule cycles, and a depth limit prevents stack overflow. + +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; + +use crate::expr::BoundExpression; + +mod rules; + +pub use rules::OptimizerRule; +pub use rules::OptimizerRuleRef; +pub use rules::OptimizerRuleRegistry; + +const DEFAULT_MAX_REWRITES: usize = 10_000; +const MAX_DEPTH: usize = 256; + +/// A deterministic optimizer for [`BoundExpression`] trees. +/// +/// The optimizer uses the rules in its [`OptimizerRuleRegistry`]. Rules are +/// grouped by root expression ID and run in registration order. Nodes are rewritten before their +/// children and again after changed children are installed. A global rewrite budget terminates +/// cyclic rule sets. +#[derive(Debug)] +pub struct BoundExpressionOptimizer { + registry: OptimizerRuleRegistry, + max_rewrites: usize, +} + +impl Default for BoundExpressionOptimizer { + fn default() -> Self { + Self::new(OptimizerRuleRegistry::default()) + } +} + +impl BoundExpressionOptimizer { + /// Create an optimizer from a configured rule registry. + pub fn new(registry: OptimizerRuleRegistry) -> Self { + Self { + registry, + max_rewrites: DEFAULT_MAX_REWRITES, + } + } + + /// Set the maximum number of successful rewrites allowed in one optimization. + pub fn with_max_rewrites(mut self, max_rewrites: usize) -> Self { + self.max_rewrites = max_rewrites; + self + } + + /// Optimize an entire bound expression tree, cloning the input when it remains unchanged. + pub fn optimize(&self, expr: &BoundExpression) -> VortexResult { + Ok(self.try_optimize(expr)?.unwrap_or_else(|| expr.clone())) + } + + /// Optimize an entire bound expression tree, returning `None` when no subtree changed. + pub fn try_optimize(&self, expr: &BoundExpression) -> VortexResult> { + OptimizationRun::new(&self.registry, self.max_rewrites).run(expr) + } +} + +/// Mutable state for one optimizer invocation. +struct OptimizationRun<'rules> { + registry: &'rules OptimizerRuleRegistry, + max_rewrites: usize, + rewrite_count: usize, +} + +impl<'rules> OptimizationRun<'rules> { + /// Create a run using the given rule registry and rewrite limit. + fn new(registry: &'rules OptimizerRuleRegistry, max_rewrites: usize) -> Self { + Self { + registry, + max_rewrites, + rewrite_count: 0, + } + } + + /// Optimize `expr`, returning `None` when the complete tree is unchanged. + fn run(mut self, expr: &BoundExpression) -> VortexResult> { + self.optimize_subtree(expr, 0) + } + + /// Apply the first matching rule to `expression`. + fn try_apply_rule( + &mut self, + expression: &BoundExpression, + ) -> VortexResult> { + let Some(rules) = self.registry.get(expression.id()) else { + return Ok(None); + }; + for rule in rules.iter() { + let Some(replacement) = rule.rewrite(expression)? else { + continue; + }; + let rule_name = rule.name(); + + vortex_ensure!( + replacement.dtype() == expression.dtype(), + "bound-expression rewrite rule {rule_name} changed dtype from {} to {}", + expression.dtype(), + replacement.dtype() + ); + vortex_ensure!( + replacement != *expression, + "bound-expression rewrite rule {rule_name} returned an unchanged expression" + ); + if self.rewrite_count >= self.max_rewrites { + vortex_bail!( + "Exceeded bound-expression rewrite limit of {} while applying {rule_name} \ + (possible rewrite cycle)", + self.max_rewrites + ); + } + self.rewrite_count += 1; + return Ok(Some(replacement)); + } + Ok(None) + } + + /// Optimize a subtree to convergence using copy-on-write rebuilding. + fn optimize_subtree( + &mut self, + original: &BoundExpression, + depth: usize, + ) -> VortexResult> { + if depth >= MAX_DEPTH { + vortex_bail!( + "Exceeded bound-expression optimization depth limit of \ + {MAX_DEPTH}" + ); + } + + let mut current = None; + loop { + loop { + let expression = current.as_ref().unwrap_or(original); + let Some(replacement) = self.try_apply_rule(expression)? else { + break; + }; + current = Some(replacement); + } + + let expression = current.as_ref().unwrap_or(original); + let Some(children) = self.optimize_children(expression, depth)? else { + return Ok(current); + }; + + let original_dtype = expression.dtype().clone(); + let expression = current.take().unwrap_or_else(|| original.clone()); + let rebuilt = expression.with_children(children)?; + vortex_ensure!( + rebuilt.dtype() == &original_dtype, + "optimizing children changed a node dtype from {original_dtype} to {}", + rebuilt.dtype() + ); + + let Some(replacement) = self.try_apply_rule(&rebuilt)? else { + return Ok(Some(rebuilt)); + }; + current = Some(replacement); + } + } + + /// Optimize a node's children, allocating a replacement vector only after one changes. + fn optimize_children( + &mut self, + expression: &BoundExpression, + depth: usize, + ) -> VortexResult>> { + let children = expression.children(); + let mut optimized_children = None; + + for (index, child) in children.iter().enumerate() { + match self.optimize_subtree(child, depth + 1)? { + Some(optimized_child) => { + optimized_children + .get_or_insert_with(|| { + let mut optimized = Vec::with_capacity(children.len()); + optimized.extend_from_slice(&children[..index]); + optimized + }) + .push(optimized_child); + } + None => { + if let Some(optimized) = &mut optimized_children { + optimized.push(child.clone()); + } + } + } + } + + Ok(optimized_children) + } +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use super::BoundExpressionOptimizer; + use super::OptimizerRule; + use super::OptimizerRuleRegistry; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::expr::BoundExpression; + use crate::expr::ExpressionId; + use crate::expr::bound; + use crate::scalar::Scalar; + use crate::scalar_fn::ScalarFnVTable; + use crate::scalar_fn::fns::binary::Binary; + use crate::scalar_fn::fns::is_null::IsNull; + use crate::scalar_fn::fns::literal::Literal; + + #[derive(Debug)] + struct IsNullTo(bool); + + impl OptimizerRule for IsNullTo { + fn expression_id(&self) -> ExpressionId { + IsNull.id() + } + + fn rewrite(&self, _expr: &BoundExpression) -> VortexResult> { + Ok(Some(bound::lit(self.0))) + } + } + + #[test] + fn rules_run_in_registration_order() -> VortexResult<()> { + let input = bound::is_null(bound::lit(1i32)); + let mut registry = OptimizerRuleRegistry::empty(); + registry.register(IsNullTo(true)); + registry.register(IsNullTo(false)); + let optimizer = BoundExpressionOptimizer::new(registry); + + assert_eq!(optimizer.optimize(&input)?, bound::lit(true)); + Ok(()) + } + + #[derive(Debug)] + struct IsNullToReducibleTree; + + impl OptimizerRule for IsNullToReducibleTree { + fn expression_id(&self) -> ExpressionId { + IsNull.id() + } + + fn rewrite(&self, _expr: &BoundExpression) -> VortexResult> { + Ok(Some(bound::and(bound::lit(false), bound::lit(true)))) + } + } + + #[derive(Debug)] + struct AndFalse; + + impl OptimizerRule for AndFalse { + fn expression_id(&self) -> ExpressionId { + Binary.id() + } + + fn rewrite(&self, expr: &BoundExpression) -> VortexResult> { + Ok( + (expr.child(0).as_opt::() == Some(&Scalar::from(false))) + .then(|| bound::lit(false)), + ) + } + } + + #[test] + fn optimizes_subtrees_introduced_by_rules() -> VortexResult<()> { + let input = bound::is_null(bound::lit(1i32)); + let mut registry = OptimizerRuleRegistry::empty(); + registry.register(IsNullToReducibleTree); + registry.register(AndFalse); + let optimizer = BoundExpressionOptimizer::new(registry); + + assert_eq!(optimizer.optimize(&input)?, bound::lit(false)); + Ok(()) + } + + #[test] + fn retries_root_rules_after_optimizing_children() -> VortexResult<()> { + let input = bound::and(bound::is_null(bound::lit(1i32)), bound::lit(true)); + let mut registry = OptimizerRuleRegistry::empty(); + registry.register(IsNullTo(false)); + registry.register(AndFalse); + let optimizer = BoundExpressionOptimizer::new(registry); + + assert_eq!(optimizer.optimize(&input)?, bound::lit(false)); + Ok(()) + } + + #[derive(Debug)] + struct WrongDType; + + impl OptimizerRule for WrongDType { + fn expression_id(&self) -> ExpressionId { + IsNull.id() + } + + fn rewrite(&self, _expr: &BoundExpression) -> VortexResult> { + Ok(Some(bound::lit(1i32))) + } + } + + #[test] + fn rejects_dtype_changes() { + let input = bound::is_null(bound::lit(1i32)); + let mut registry = OptimizerRuleRegistry::empty(); + registry.register(WrongDType); + let optimizer = BoundExpressionOptimizer::new(registry); + + assert!(optimizer.optimize(&input).is_err()); + } + + #[derive(Debug)] + struct Unchanged; + + impl OptimizerRule for Unchanged { + fn expression_id(&self) -> ExpressionId { + IsNull.id() + } + + fn rewrite(&self, expr: &BoundExpression) -> VortexResult> { + Ok(Some(expr.clone())) + } + } + + #[test] + fn rejects_unchanged_replacements() { + let input = bound::is_null(bound::lit(1i32)); + let mut registry = OptimizerRuleRegistry::empty(); + registry.register(Unchanged); + let optimizer = BoundExpressionOptimizer::new(registry); + + assert!(optimizer.optimize(&input).is_err()); + } + + #[derive(Debug)] + struct ToggleBoolean; + + impl OptimizerRule for ToggleBoolean { + fn expression_id(&self) -> ExpressionId { + Literal.id() + } + + fn rewrite(&self, expr: &BoundExpression) -> VortexResult> { + let Some(value) = expr + .as_opt::() + .and_then(|scalar| scalar.as_bool_opt()) + else { + return Ok(None); + }; + Ok(value.value().map(|value| bound::lit(!value))) + } + } + + #[test] + fn rewrite_budget_terminates_cycles() { + let mut registry = OptimizerRuleRegistry::empty(); + registry.register(ToggleBoolean); + let optimizer = BoundExpressionOptimizer::new(registry).with_max_rewrites(4); + + assert!(optimizer.optimize(&bound::lit(true)).is_err()); + } + + #[derive(Debug)] + struct RootToOne(ExpressionId); + + impl OptimizerRule for RootToOne { + fn expression_id(&self) -> ExpressionId { + self.0 + } + + fn rewrite(&self, _expr: &BoundExpression) -> VortexResult> { + Ok(Some(bound::lit(1i32))) + } + } + + #[test] + fn rules_can_target_non_scalar_expression_nodes() -> VortexResult<()> { + let input = bound::root(DType::Primitive(PType::I32, Nullability::NonNullable)); + let mut registry = OptimizerRuleRegistry::empty(); + registry.register(RootToOne(input.id())); + let optimizer = BoundExpressionOptimizer::new(registry); + + assert_eq!(optimizer.optimize(&input)?, bound::lit(1i32)); + Ok(()) + } + + #[test] + fn depth_limit_prevents_stack_overflow() { + let mut expr = bound::root(DType::Bool(Nullability::NonNullable)); + for _ in 0..1_000 { + expr = bound::not(expr); + } + + assert!( + BoundExpressionOptimizer::new(OptimizerRuleRegistry::empty()) + .try_optimize(&expr) + .is_err() + ); + } + + #[test] + fn root_rewrite_can_discard_a_deep_subtree() -> VortexResult<()> { + let mut discarded = bound::root(DType::Bool(Nullability::NonNullable)); + for _ in 0..1_000 { + discarded = bound::not(discarded); + } + let input = bound::and(bound::lit(false), discarded); + let mut registry = OptimizerRuleRegistry::empty(); + registry.register(AndFalse); + let optimizer = BoundExpressionOptimizer::new(registry); + + assert_eq!(optimizer.optimize(&input)?, bound::lit(false)); + Ok(()) + } + + #[test] + fn default_optimizer_folds_literal_cast() -> VortexResult<()> { + let target = DType::Primitive(PType::I64, Nullability::NonNullable); + let expr = bound::cast(bound::lit(1i32), target); + + assert_eq!( + BoundExpressionOptimizer::default().optimize(&expr)?, + bound::lit(1i64) + ); + Ok(()) + } +} diff --git a/vortex-array/src/expr/optimizer/rules/binary.rs b/vortex-array/src/expr/optimizer/rules/binary.rs new file mode 100644 index 00000000000..68deb143144 --- /dev/null +++ b/vortex-array/src/expr/optimizer/rules/binary.rs @@ -0,0 +1,371 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::OptimizerRule; +use super::preserve_dtype; +use crate::expr::BoundExpression; +use crate::expr::ExpressionId; +use crate::expr::bound; +use crate::scalar::Scalar; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::ScalarFnVTableExt; +use crate::scalar_fn::fns::between::Between; +use crate::scalar_fn::fns::between::BetweenOptions; +use crate::scalar_fn::fns::between::StrictComparison; +use crate::scalar_fn::fns::binary::Binary; +use crate::scalar_fn::fns::get_item::GetItem; +use crate::scalar_fn::fns::literal::Literal; +use crate::scalar_fn::fns::operators::Operator; + +/// Simplifies `AND` and `OR` with literal operands using Kleene boolean semantics. +/// +/// # Example +/// +/// ```text +/// original: and(value, lit(true)) +/// rewritten: value +/// ``` +#[derive(Debug)] +pub(crate) struct BinaryBoolean; + +impl OptimizerRule for BinaryBoolean { + fn expression_id(&self) -> ExpressionId { + Binary.id() + } + + fn rewrite(&self, expr: &BoundExpression) -> VortexResult> { + let operator = expr.as_::(); + let lhs = expr.child(0); + let rhs = expr.child(1); + let bool_literal = |expr: &BoundExpression| { + expr.as_opt::()? + .as_bool_opt() + .map(|value| value.value()) + }; + + let replacement = match operator { + Operator::And => match (bool_literal(lhs), bool_literal(rhs)) { + (Some(Some(false)), _) | (_, Some(Some(false))) => { + Some(bound::lit(Scalar::bool(false, expr.dtype().nullability()))) + } + (Some(Some(true)), _) => Some(preserve_dtype(rhs.clone(), expr.dtype())?), + (_, Some(Some(true))) => Some(preserve_dtype(lhs.clone(), expr.dtype())?), + (Some(None), Some(None)) => Some(lhs.clone()), + _ => None, + }, + Operator::Or => match (bool_literal(lhs), bool_literal(rhs)) { + (Some(Some(true)), _) | (_, Some(Some(true))) => { + Some(bound::lit(Scalar::bool(true, expr.dtype().nullability()))) + } + (Some(Some(false)), _) => Some(preserve_dtype(rhs.clone(), expr.dtype())?), + (_, Some(Some(false))) => Some(preserve_dtype(lhs.clone(), expr.dtype())?), + (Some(None), Some(None)) => Some(lhs.clone()), + _ => None, + }, + _ => None, + }; + Ok(replacement) + } +} + +/// Replaces a comparison against a null literal with a null boolean literal. +/// +/// # Example +/// +/// ```text +/// original: eq(value, lit(Scalar::null(nullable_i32))) +/// rewritten: lit(Scalar::null(nullable_bool)) +/// ``` +#[derive(Debug)] +pub(crate) struct BinaryNullComparison; + +impl OptimizerRule for BinaryNullComparison { + fn expression_id(&self) -> ExpressionId { + Binary.id() + } + + fn rewrite(&self, expr: &BoundExpression) -> VortexResult> { + if !expr.as_::().is_comparison() { + return Ok(None); + } + let is_null_literal = + |child: &BoundExpression| child.as_opt::().is_some_and(Scalar::is_null); + if !is_null_literal(expr.child(0)) && !is_null_literal(expr.child(1)) { + return Ok(None); + } + + Ok(Some(bound::lit(Scalar::null(expr.dtype().clone())))) + } +} + +/// Combines compatible lower- and upper-bound conjuncts into `between` expressions. +/// +/// # Example +/// +/// ```text +/// original: and(gt_eq(x, lit(1)), lt(x, lit(10))) +/// rewritten: between(x, lit(1), lit(10), BetweenOptions { lower_strict: NonStrict, upper_strict: Strict }) +/// ``` +#[derive(Debug)] +pub(crate) struct FindBetween; + +impl OptimizerRule for FindBetween { + fn expression_id(&self) -> ExpressionId { + Binary.id() + } + + fn rewrite(&self, expr: &BoundExpression) -> VortexResult> { + if expr.as_opt::() != Some(&Operator::And) { + return Ok(None); + } + + let mut conjuncts = collect_conjuncts(expr); + let mut rewritten = Vec::with_capacity(conjuncts.len()); + let mut changed = false; + for idx in 0..conjuncts.len() { + let Some(conjunct) = conjuncts.get(idx).cloned() else { + continue; + }; + let mut matched = false; + for other_idx in (idx + 1)..conjuncts.len() { + let Some(other) = conjuncts.get(other_idx) else { + continue; + }; + if let Some(between) = match_between(&conjunct, other)? { + rewritten.push(between); + conjuncts.remove(other_idx); + matched = true; + changed = true; + break; + } + } + if !matched { + rewritten.push(conjunct); + } + } + + if !changed { + return Ok(None); + } + Ok(bound::and_collect(rewritten)) + } +} + +fn collect_conjuncts(expr: &BoundExpression) -> Vec { + let mut stack = vec![expr]; + let mut conjuncts = Vec::new(); + while let Some(expr) = stack.pop() { + if expr.as_opt::() == Some(&Operator::And) { + stack.push(expr.child(1)); + stack.push(expr.child(0)); + } else { + conjuncts.push(expr.clone()); + } + } + conjuncts +} + +fn match_between( + lhs: &BoundExpression, + rhs: &BoundExpression, +) -> VortexResult> { + let (Some(lhs_op), Some(rhs_op)) = (lhs.as_opt::(), rhs.as_opt::()) else { + return Ok(None); + }; + if lhs.child(0) == lhs.child(1) || rhs.child(0) == rhs.child(1) { + return Ok(None); + } + + let lhs = normalize_get_item_comparison(lhs, *lhs_op)?; + let rhs = normalize_get_item_comparison(rhs, *rhs_op)?; + let (Some(lhs), Some(rhs)) = (lhs, rhs) else { + return Ok(None); + }; + if lhs.child(0) != rhs.child(0) { + return Ok(None); + } + + let (lower, upper) = match (lhs.as_::(), rhs.as_::()) { + (Operator::Lt | Operator::Lte, Operator::Gt | Operator::Gte) => (rhs, lhs), + (Operator::Gt | Operator::Gte, Operator::Lt | Operator::Lte) => (lhs, rhs), + _ => return Ok(None), + }; + let lower_lit = lower.child(1).as_opt::(); + let upper_lit = upper.child(1).as_opt::(); + if lower_lit.is_none_or(Scalar::is_null) || upper_lit.is_none_or(Scalar::is_null) { + return Ok(None); + } + + // Binary comparisons permit an extension value against its raw storage dtype, but Between + // requires all three logical dtypes to match. + let value_dtype = lower.child(0).dtype(); + if !value_dtype.eq_ignore_nullability(lower.child(1).dtype()) + || !value_dtype.eq_ignore_nullability(upper.child(1).dtype()) + { + return Ok(None); + } + + let lower_strict = comparison_strictness(*lower.as_::())?; + let upper_strict = comparison_strictness(*upper.as_::())?; + Ok(Some(Between.try_new_bound_expr( + BetweenOptions { + lower_strict, + upper_strict, + }, + [ + lower.child(0).clone(), + lower.child(1).clone(), + upper.child(1).clone(), + ], + )?)) +} + +fn normalize_get_item_comparison( + expr: &BoundExpression, + operator: Operator, +) -> VortexResult> { + match (expr.child(0).is::(), expr.child(1).is::()) { + (true, false) => Ok(Some(expr.clone())), + (false, true) => { + let Some(swapped) = operator.swap() else { + return Ok(None); + }; + Ok(Some(Binary.try_new_bound_expr( + swapped, + [expr.child(1).clone(), expr.child(0).clone()], + )?)) + } + _ => Ok(None), + } +} + +fn comparison_strictness(operator: Operator) -> VortexResult { + match operator { + Operator::Lt | Operator::Gt => Ok(StrictComparison::Strict), + Operator::Lte | Operator::Gte => Ok(StrictComparison::NonStrict), + _ => Err(vortex_err!( + "expected an inequality operator, got {operator}" + )), + } +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::expr::BoundExpression; + use crate::expr::bound; + use crate::expr::optimizer::BoundExpressionOptimizer; + use crate::extension::datetime::TimeUnit; + use crate::extension::datetime::Timestamp; + use crate::scalar::Scalar; + use crate::scalar_fn::fns::between::Between; + use crate::scalar_fn::fns::between::BetweenOptions; + use crate::scalar_fn::fns::between::StrictComparison; + + fn optimize(expr: &BoundExpression) -> VortexResult { + BoundExpressionOptimizer::default().optimize(expr) + } + + #[test] + fn boolean_annihilator_preserves_nullable_dtype() -> VortexResult<()> { + let nullable_bool = Scalar::bool(true, Nullability::Nullable); + let expr = bound::and(bound::lit(false), bound::lit(nullable_bool)); + + assert_eq!( + optimize(&expr)?, + bound::lit(Scalar::bool(false, Nullability::Nullable)) + ); + Ok(()) + } + + #[test] + fn null_comparison_folds_to_nullable_null() -> VortexResult<()> { + let nullable_i32 = DType::Primitive(PType::I32, Nullability::Nullable); + let expr = bound::eq( + bound::root(nullable_i32.clone()), + bound::lit(Scalar::null(nullable_i32)), + ); + + assert_eq!( + optimize(&expr)?, + bound::lit(Scalar::null(DType::Bool(Nullability::Nullable))) + ); + Ok(()) + } + + fn comparison_scope() -> DType { + DType::struct_( + [("x", DType::Primitive(PType::I32, Nullability::NonNullable))], + Nullability::NonNullable, + ) + } + + #[test] + fn comparison_pair_lowers_to_between() -> VortexResult<()> { + let scope = comparison_scope(); + let x = bound::col("x", scope); + let expr = bound::and( + bound::gt_eq(x.clone(), bound::lit(2i32)), + bound::lt(x.clone(), bound::lit(5i32)), + ); + + assert_eq!( + optimize(&expr)?, + bound::between( + x, + bound::lit(2i32), + bound::lit(5i32), + BetweenOptions { + lower_strict: StrictComparison::NonStrict, + upper_strict: StrictComparison::Strict, + } + ) + ); + Ok(()) + } + + #[test] + fn null_bound_does_not_lower_to_between() -> VortexResult<()> { + let scope = comparison_scope(); + let x = bound::col("x", scope); + let null = bound::lit(Scalar::null(DType::Primitive( + PType::I32, + Nullability::Nullable, + ))); + let expr = bound::and( + bound::gt_eq(x.clone(), null), + bound::lt(x, bound::lit(5i32)), + ); + + assert!(!optimize(&expr)?.contains::()?); + Ok(()) + } + + #[test] + fn extension_storage_bounds_do_not_lower_to_between() -> VortexResult<()> { + let scope = DType::struct_( + [( + "x", + DType::Extension( + Timestamp::new(TimeUnit::Milliseconds, Nullability::NonNullable).erased(), + ), + )], + Nullability::NonNullable, + ); + let x = bound::col("x", scope); + let expr = bound::and( + bound::gt_eq(x.clone(), bound::lit(2i64)), + bound::lt(x, bound::lit(5i64)), + ); + + assert_eq!(optimize(&expr)?, expr); + Ok(()) + } +} diff --git a/vortex-array/src/expr/optimizer/rules/cast.rs b/vortex-array/src/expr/optimizer/rules/cast.rs new file mode 100644 index 00000000000..12ed66f610b --- /dev/null +++ b/vortex-array/src/expr/optimizer/rules/cast.rs @@ -0,0 +1,41 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use super::OptimizerRule; +use crate::expr::BoundExpression; +use crate::expr::ExpressionId; +use crate::expr::bound; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::cast::Cast; +use crate::scalar_fn::fns::literal::Literal; + +/// Removes identity casts and evaluates casts of literal values during optimization. +/// +/// # Example +/// +/// ```text +/// original: cast(lit(1_i32), i64) +/// rewritten: lit(1_i64) +/// ``` +#[derive(Debug)] +pub(crate) struct CastLiteralOrIdentity; + +impl OptimizerRule for CastLiteralOrIdentity { + fn expression_id(&self) -> ExpressionId { + Cast.id() + } + + fn rewrite(&self, expr: &BoundExpression) -> VortexResult> { + let target = expr.as_::(); + let child = expr.child(0); + if child.dtype() == target { + return Ok(Some(child.clone())); + } + let Some(scalar) = child.as_opt::() else { + return Ok(None); + }; + Ok(scalar.cast(target).ok().map(bound::lit)) + } +} diff --git a/vortex-array/src/expr/optimizer/rules/conditional.rs b/vortex-array/src/expr/optimizer/rules/conditional.rs new file mode 100644 index 00000000000..aebf5edf9f2 --- /dev/null +++ b/vortex-array/src/expr/optimizer/rules/conditional.rs @@ -0,0 +1,105 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use super::OptimizerRule; +use super::preserve_dtype; +use crate::expr::BoundExpression; +use crate::expr::ExpressionId; +use crate::expr::bound; +use crate::scalar::Scalar; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::literal::Literal; +use crate::scalar_fn::fns::mask::Mask; +use crate::scalar_fn::fns::zip::Zip; + +/// Evaluates a mask whose mask argument is a non-null boolean literal. +/// +/// # Example +/// +/// ```text +/// original: mask(nullable_value, lit(true)) +/// rewritten: nullable_value +/// ``` +#[derive(Debug)] +pub(crate) struct ConstantMask; + +impl OptimizerRule for ConstantMask { + fn expression_id(&self) -> ExpressionId { + Mask.id() + } + + fn rewrite(&self, expr: &BoundExpression) -> VortexResult> { + let Some(mask) = expr + .child(1) + .as_opt::() + .and_then(|scalar| scalar.as_bool_opt()) + .and_then(|value| value.value()) + else { + return Ok(None); + }; + + if mask { + return Ok(Some(preserve_dtype(expr.child(0).clone(), expr.dtype())?)); + } + Ok(Some(bound::lit(Scalar::null(expr.dtype().clone())))) + } +} + +/// Selects the reachable branch of a zip with a non-null literal mask. +/// +/// # Example +/// +/// ```text +/// original: zip_expr(lit(true), if_true, if_false) +/// rewritten: if_true +/// ``` +#[derive(Debug)] +pub(crate) struct ConstantZip; + +impl OptimizerRule for ConstantZip { + fn expression_id(&self) -> ExpressionId { + Zip.id() + } + + fn rewrite(&self, expr: &BoundExpression) -> VortexResult> { + let Some(mask) = expr + .child(2) + .as_opt::() + .and_then(|scalar| scalar.as_bool_opt()) + .and_then(|value| value.value()) + else { + return Ok(None); + }; + let child = if mask { expr.child(0) } else { expr.child(1) }; + Ok(Some(preserve_dtype(child.clone(), expr.dtype())?)) + } +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use crate::dtype::Nullability; + use crate::expr::BoundExpression; + use crate::expr::bound; + use crate::expr::optimizer::BoundExpressionOptimizer; + use crate::scalar::Scalar; + + fn optimize(expr: &BoundExpression) -> VortexResult { + BoundExpressionOptimizer::default().optimize(expr) + } + + #[test] + fn constant_mask_and_zip_preserve_nullable_output() -> VortexResult<()> { + let nullable_two = bound::lit(Scalar::primitive(2i32, Nullability::Nullable)); + let masked = bound::mask(bound::lit(1i32), bound::lit(true)); + let zipped = bound::zip_expr(bound::lit(true), bound::lit(1i32), nullable_two); + let expected = bound::lit(Scalar::primitive(1i32, Nullability::Nullable)); + + assert_eq!(optimize(&masked)?, expected); + assert_eq!(optimize(&zipped)?, expected); + Ok(()) + } +} diff --git a/vortex-array/src/expr/optimizer/rules/mod.rs b/vortex-array/src/expr/optimizer/rules/mod.rs new file mode 100644 index 00000000000..7c8deaba29a --- /dev/null +++ b/vortex-array/src/expr/optimizer/rules/mod.rs @@ -0,0 +1,109 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use std::fmt::Debug; +use std::sync::Arc; + +use vortex_error::VortexResult; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::dtype::DType; +use crate::expr::BoundExpression; +use crate::expr::ExpressionId; +use crate::scalar_fn::ScalarFnVTableExt; +use crate::scalar_fn::fns::cast::Cast; + +mod binary; +mod cast; +mod conditional; +mod nulls; +mod structural; + +pub(crate) use binary::BinaryBoolean; +pub(crate) use binary::BinaryNullComparison; +pub(crate) use binary::FindBetween; +pub(crate) use cast::CastLiteralOrIdentity; +pub(crate) use conditional::ConstantMask; +pub(crate) use conditional::ConstantZip; +pub(crate) use nulls::CaseWhenToFillNull; +pub(crate) use nulls::RemoveRedundantFillNull; +pub(crate) use structural::GetItemFromPack; +pub(crate) use structural::MergeToPack; +pub(crate) use structural::SelectFromPack; + +/// Shared reference to a rewrite rule. +pub type OptimizerRuleRef = Arc; + +/// An equivalence rewrite for bound expressions with a particular root node implementation. +/// +/// The optimizer invokes a rule only when the expression's root ID equals +/// [`Self::expression_id`]. Returning `None` means the rule does not match. A replacement must be +/// semantically equivalent to the input and have exactly the same dtype, including nullability. +/// The optimizer verifies the dtype and rejects unchanged replacements. +pub trait OptimizerRule: Debug + Send + Sync + 'static { + /// Returns a diagnostic name for this rule. + fn name(&self) -> &'static str { + std::any::type_name::() + } + + /// Returns the expression node ID handled by this rule. + fn expression_id(&self) -> ExpressionId; + + /// Try to rewrite `expr` to a semantically equivalent bound expression. + fn rewrite(&self, expr: &BoundExpression) -> VortexResult>; +} + +/// Rewrite rules grouped by the expression node ID they handle. +#[derive(Debug)] +pub struct OptimizerRuleRegistry { + rules: HashMap>, +} + +impl OptimizerRuleRegistry { + /// Create an empty rule registry. + pub fn empty() -> Self { + Self { + rules: HashMap::default(), + } + } + + /// Register a rule after existing rules for the same expression node ID. + pub fn register(&mut self, rule: R) { + self.rules + .entry(rule.expression_id()) + .or_default() + .push(Arc::new(rule)); + } + + /// Return rules for `expression_id` in registration order. + pub(super) fn get(&self, expression_id: ExpressionId) -> Option<&[OptimizerRuleRef]> { + self.rules.get(&expression_id).map(Vec::as_slice) + } +} + +impl Default for OptimizerRuleRegistry { + fn default() -> Self { + let mut registry = Self::empty(); + + registry.register(BinaryBoolean); + registry.register(BinaryNullComparison); + registry.register(FindBetween); + registry.register(CastLiteralOrIdentity); + registry.register(GetItemFromPack); + registry.register(MergeToPack); + registry.register(SelectFromPack); + registry.register(RemoveRedundantFillNull); + registry.register(CaseWhenToFillNull); + registry.register(ConstantMask); + registry.register(ConstantZip); + + registry + } +} + +fn preserve_dtype(replacement: BoundExpression, dtype: &DType) -> VortexResult { + if replacement.dtype() == dtype { + return Ok(replacement); + } + Cast.try_new_bound_expr(dtype.clone(), [replacement]) +} diff --git a/vortex-array/src/expr/optimizer/rules/nulls.rs b/vortex-array/src/expr/optimizer/rules/nulls.rs new file mode 100644 index 00000000000..6f81c73c169 --- /dev/null +++ b/vortex-array/src/expr/optimizer/rules/nulls.rs @@ -0,0 +1,146 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use super::OptimizerRule; +use super::preserve_dtype; +use crate::expr::BoundExpression; +use crate::expr::ExpressionId; +use crate::expr::bound; +use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::ScalarFnVTableExt; +use crate::scalar_fn::fns::case_when::CaseWhen; +use crate::scalar_fn::fns::fill_null::FillNull; +use crate::scalar_fn::fns::is_not_null::IsNotNull; +use crate::scalar_fn::fns::is_null::IsNull; +use crate::scalar_fn::fns::literal::Literal; + +/// Removes `fill_null` when its input is already non-nullable. +/// +/// # Example +/// +/// ```text +/// original: fill_null(non_nullable_value, lit(0)) +/// rewritten: non_nullable_value +/// ``` +#[derive(Debug)] +pub(crate) struct RemoveRedundantFillNull; + +impl OptimizerRule for RemoveRedundantFillNull { + fn expression_id(&self) -> ExpressionId { + FillNull.id() + } + + fn rewrite(&self, expr: &BoundExpression) -> VortexResult> { + if expr.child(0).dtype().is_nullable() { + return Ok(None); + } + Ok(Some(preserve_dtype(expr.child(0).clone(), expr.dtype())?)) + } +} + +/// Lowers a single-branch null-checking `case_when` into `fill_null` or its input. +/// +/// # Example +/// +/// ```text +/// original: case_when(is_null(value), fill, value) +/// rewritten: fill_null(value, fill) +/// ``` +#[derive(Debug)] +pub(crate) struct CaseWhenToFillNull; + +impl OptimizerRule for CaseWhenToFillNull { + fn expression_id(&self) -> ExpressionId { + CaseWhen.id() + } + + fn rewrite(&self, expr: &BoundExpression) -> VortexResult> { + let options = expr.as_::(); + if options.num_when_then_pairs != 1 || !options.has_else { + return Ok(None); + } + + let when = expr.child(0); + let then = expr.child(1); + let els = expr.child(2); + let (value, fill) = if when.is::() && when.child(0) == els { + (els, then) + } else if when.is::() && when.child(0) == then { + (then, els) + } else { + return Ok(None); + }; + let Some(fill_scalar) = fill.as_opt::() else { + return Ok(None); + }; + + if fill_scalar.is_null() { + return Ok(Some(preserve_dtype(value.clone(), expr.dtype())?)); + } + let fill = if fill.dtype() == expr.dtype() { + fill.clone() + } else { + bound::lit(fill_scalar.cast(expr.dtype())?) + }; + Ok(Some( + FillNull.try_new_bound_expr(EmptyOptions, [value.clone(), fill])?, + )) + } +} + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; + use crate::expr::BoundExpression; + use crate::expr::bound; + use crate::expr::optimizer::BoundExpressionOptimizer; + use crate::scalar::Scalar; + use crate::scalar_fn::fns::fill_null::FillNull; + + fn optimize(expr: &BoundExpression) -> VortexResult { + BoundExpressionOptimizer::default().optimize(expr) + } + + #[test] + fn nonnullable_fill_null_input_is_removed() -> VortexResult<()> { + let expr = bound::fill_null( + bound::lit(1i32), + bound::lit(Scalar::primitive(0i32, Nullability::Nullable)), + ); + + assert_eq!( + optimize(&expr)?, + bound::lit(Scalar::primitive(1i32, Nullability::Nullable)) + ); + Ok(()) + } + + #[test] + fn coalesce_shaped_case_lowers_to_fill_null() -> VortexResult<()> { + let value_dtype = DType::Primitive(PType::I64, Nullability::Nullable); + let value = bound::root(value_dtype); + let expr = bound::case_when( + bound::is_null(value.clone()), + bound::lit(0i64), + value.clone(), + ); + let optimized = optimize(&expr)?; + + assert!(optimized.is::()); + assert_eq!( + optimized, + bound::fill_null( + value, + bound::lit(Scalar::primitive(0i64, Nullability::Nullable)) + ) + ); + Ok(()) + } +} diff --git a/vortex-array/src/expr/optimizer/rules/structural.rs b/vortex-array/src/expr/optimizer/rules/structural.rs new file mode 100644 index 00000000000..65e47420ac2 --- /dev/null +++ b/vortex-array/src/expr/optimizer/rules/structural.rs @@ -0,0 +1,275 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use itertools::Itertools; +use vortex_error::VortexExpect; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; +use vortex_utils::aliases::hash_set::HashSet; + +use super::OptimizerRule; +use crate::dtype::FieldNames; +use crate::expr::BoundExpression; +use crate::expr::ExpressionId; +use crate::expr::bound; +use crate::scalar_fn::EmptyOptions; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::ScalarFnVTableExt; +use crate::scalar_fn::fns::get_item::GetItem; +use crate::scalar_fn::fns::mask::Mask; +use crate::scalar_fn::fns::merge::DuplicateHandling; +use crate::scalar_fn::fns::merge::Merge; +use crate::scalar_fn::fns::pack::Pack; +use crate::scalar_fn::fns::pack::PackOptions; +use crate::scalar_fn::fns::select::Select; + +/// Replaces a field access on a pack with the corresponding packed expression. +/// +/// # Example +/// +/// ```text +/// original: get_item("b", pack([("a", a), ("b", b)], NonNullable)) +/// rewritten: b +/// ``` +#[derive(Debug)] +pub(crate) struct GetItemFromPack; + +impl OptimizerRule for GetItemFromPack { + fn expression_id(&self) -> ExpressionId { + GetItem.id() + } + + fn rewrite(&self, expr: &BoundExpression) -> VortexResult> { + let field_name = expr.as_::(); + let child = expr.child(0); + let Some(pack) = child.as_opt::() else { + return Ok(None); + }; + let Some(idx) = pack.names.find(field_name) else { + return Err(vortex_err!( + "Cannot find field {field_name} in pack fields {:?}", + pack.names + )); + }; + + let mut field = child.child(idx).clone(); + if pack.nullability.is_nullable() { + field = Mask.try_new_bound_expr(EmptyOptions, [field, bound::lit(true)])?; + } + Ok(Some(field)) + } +} + +/// Lowers a merge of struct expressions into a pack while honoring its duplicate-field policy. +/// +/// # Example +/// +/// When `left` contains `a` and `right` contains `b`: +/// +/// ```text +/// original: merge([left, right]) +/// rewritten: pack([("a", get_item("a", left)), ("b", get_item("b", right))], NonNullable) +/// ``` +#[derive(Debug)] +pub(crate) struct MergeToPack; + +impl OptimizerRule for MergeToPack { + fn expression_id(&self) -> ExpressionId { + Merge.id() + } + + fn rewrite(&self, expr: &BoundExpression) -> VortexResult> { + let options = expr.as_::(); + let mut names = Vec::with_capacity(expr.children().len() * 2); + let mut sources = Vec::with_capacity(expr.children().len() * 2); + let mut duplicate_names = HashSet::new(); + + for child in expr.children() { + let fields = child.dtype().as_struct_fields_opt().ok_or_else(|| { + vortex_err!( + "Merge child must return a struct dtype, got {}", + child.dtype() + ) + })?; + for name in fields.names().iter() { + if let Some(idx) = names.iter().position(|existing| existing == name) { + duplicate_names.insert(name.clone()); + sources[idx] = child.clone(); + } else { + names.push(name.clone()); + sources.push(child.clone()); + } + } + } + + if options == &DuplicateHandling::Error && !duplicate_names.is_empty() { + vortex_bail!( + "merge: duplicate fields in children: {}", + duplicate_names.into_iter().format(", ") + ) + } + + let children = names + .iter() + .zip(sources) + .map(|(name, source)| GetItem.try_new_bound_expr(name.clone(), [source])) + .collect::>>()?; + Ok(Some(Pack.try_new_bound_expr( + PackOptions { + names: FieldNames::from(names), + nullability: expr.dtype().nullability(), + }, + children, + )?)) + } +} + +/// Lowers a selection from a pack into a smaller pack when struct validity is preserved. +/// +/// # Example +/// +/// ```text +/// original: select(["b"], pack([("a", a), ("b", b)], NonNullable)) +/// rewritten: pack([("b", b)], NonNullable) +/// ``` +#[derive(Debug)] +pub(crate) struct SelectFromPack; + +impl OptimizerRule for SelectFromPack { + fn expression_id(&self) -> ExpressionId { + Select.id() + } + + fn rewrite(&self, expr: &BoundExpression) -> VortexResult> { + let selection = expr.as_::