From 83e8bba813d92b05172ed8a2246b137dbe483447 Mon Sep 17 00:00:00 2001 From: Matt Katz Date: Tue, 25 Aug 2026 15:21:32 -0400 Subject: [PATCH] optimizer first pass Signed-off-by: Matt Katz --- vortex-array/src/expr/bound_expression.rs | 13 + vortex-array/src/expr/mod.rs | 11 + vortex-array/src/expr/optimizer/mod.rs | 545 ++++++++++++++++++ .../src/expr/optimizer/rules/binary.rs | 346 +++++++++++ vortex-array/src/expr/optimizer/rules/cast.rs | 46 ++ .../src/expr/optimizer/rules/conditional.rs | 111 ++++ vortex-array/src/expr/optimizer/rules/mod.rs | 31 + .../src/expr/optimizer/rules/nulls.rs | 152 +++++ .../src/expr/optimizer/rules/structural.rs | 261 +++++++++ 9 files changed, 1516 insertions(+) create mode 100644 vortex-array/src/expr/optimizer/mod.rs create mode 100644 vortex-array/src/expr/optimizer/rules/binary.rs create mode 100644 vortex-array/src/expr/optimizer/rules/cast.rs create mode 100644 vortex-array/src/expr/optimizer/rules/conditional.rs create mode 100644 vortex-array/src/expr/optimizer/rules/mod.rs create mode 100644 vortex-array/src/expr/optimizer/rules/nulls.rs create mode 100644 vortex-array/src/expr/optimizer/rules/structural.rs 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..e1269411ac8 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,17 @@ 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::BoundExpressionRewriteRule; +pub use optimizer::BoundExpressionRewriteRuleRef; 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..a923d1f4d57 --- /dev/null +++ b/vortex-array/src/expr/optimizer/mod.rs @@ -0,0 +1,545 @@ +// 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 +//! +//! [`BoundExpressionOptimizer`] holds the reusable rule registry and rewrite limit. Each call to +//! [`BoundExpressionOptimizer::try_optimize`] creates an `OptimizationRun` containing the mutable +//! state for that traversal. +//! +//! A run walks the tree bottom-up without recursion, using one stack of pending tasks and another +//! stack of completed node results: +//! +//! 1. A `Visit` task schedules the node to be finished, then schedules its children in evaluation +//! order. +//! 2. A `Finish` task collects the optimized children and rebuilds the node only if a child +//! changed. +//! 3. Rules registered for the node's expression ID run in registration order. The first matching +//! rule wins. +//! 4. A replacement is visited as a new subtree so that rewrites introduced by other rewrites are +//! also optimized. `FinishRewrite` then records that the original subtree changed. +//! +//! Every replacement must preserve the node's dtype and differ from the expression it replaces. +//! A per-run rewrite limit terminates rule cycles. + +use std::fmt::Debug; +use std::sync::Arc; + +use smallvec::SmallVec; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_utils::aliases::hash_map::HashMap; + +use crate::expr::BoundExpression; +use crate::expr::ExpressionId; + +mod rules; + +const DEFAULT_MAX_REWRITES: usize = 10_000; + +/// Shared reference to a bound-expression rewrite rule. +pub type BoundExpressionRewriteRuleRef = 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 BoundExpressionRewriteRule: 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>; +} + +/// A deterministic, bottom-up optimizer for [`BoundExpression`] trees. +/// +/// Rules are grouped by root expression ID and run in registration order. After the first +/// matching rule rewrites a node, the replacement subtree is optimized from the bottom up before +/// its parent is revisited. A global rewrite budget terminates cyclic rule sets. +#[derive(Clone, Debug)] +pub struct BoundExpressionOptimizer { + rules: HashMap>, + max_rewrites: usize, +} + +impl Default for BoundExpressionOptimizer { + fn default() -> Self { + let mut optimizer = Self::empty(); + rules::register(&mut optimizer); + optimizer + } +} + +impl BoundExpressionOptimizer { + /// Create an optimizer containing the built-in Vortex expression rewrites. + pub fn new() -> Self { + Self::default() + } + + /// Create an optimizer with no registered rules. + pub fn empty() -> Self { + Self { + rules: HashMap::default(), + 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 + } + + /// Register a rewrite after existing rules for the same expression node ID. + pub fn register_rule(&mut self, rule: R) { + self.register_rule_ref(Arc::new(rule)); + } + + /// Register a shared rewrite after existing rules for the same expression node ID. + pub fn register_rule_ref(&mut self, rule: BoundExpressionRewriteRuleRef) { + self.rules + .entry(rule.expression_id()) + .or_default() + .push(rule); + } + + /// 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, expr).run() + } + + /// Run rules for the expression's root ID and return the first successful rewrite. + fn rewrite_root( + &self, + expr: &BoundExpression, + ) -> VortexResult> { + let Some(rules) = self.rules.get(&expr.id()) else { + return Ok(None); + }; + + for rule in rules { + if let Some(replacement) = rule.rewrite(expr)? { + return Ok(Some((rule.name(), replacement))); + } + } + Ok(None) + } +} + +/// Mutable state for one optimizer invocation. +struct OptimizationRun<'a> { + optimizer: &'a BoundExpressionOptimizer, + tasks: TaskStack, + results: ResultStack, + rewrite_count: usize, +} + +impl<'a> OptimizationRun<'a> { + /// Create a run with the root expression as its first pending visit. + fn new(optimizer: &'a BoundExpressionOptimizer, expr: &BoundExpression) -> Self { + let mut tasks = TaskStack::new(); + tasks.push(Task::Visit(expr.clone())); + Self { + optimizer, + tasks, + results: ResultStack::new(), + rewrite_count: 0, + } + } + + /// Execute pending tasks until the root produces its single optimization result. + fn run(mut self) -> VortexResult> { + while let Some(task) = self.tasks.pop() { + match task { + Task::Visit(expr) => self.visit(expr)?, + Task::Finish { expr, child_count } => self.finish(expr, child_count)?, + Task::FinishRewrite => self.finish_rewrite()?, + } + } + + vortex_ensure!( + self.results.len() == 1, + "bound-expression optimizer produced {} root results", + self.results.len() + ); + let result = self + .results + .pop() + .ok_or_else(|| vortex_err!("bound-expression optimizer produced no root result"))?; + Ok(result.changed.then_some(result.expr)) + } + + /// Schedule an expression's children followed by the task that finishes the expression. + /// + /// A leaf has no child work to schedule, so it is finished immediately. + fn visit(&mut self, expr: BoundExpression) -> VortexResult<()> { + let child_count = expr.children().len(); + if child_count == 0 { + return self.finish_node(expr, OptimizedChildren::Unchanged); + } + + self.tasks.reserve(child_count + 1); + self.results.reserve(child_count); + self.tasks.push(Task::Finish { + expr: expr.clone(), + child_count, + }); + self.tasks + .extend(expr.children().iter().rev().cloned().map(Task::Visit)); + Ok(()) + } + + /// Collect an expression's child results and finish the expression itself. + /// + /// When every child is unchanged, their results are discarded and the original expression can + /// be reused. Otherwise, all optimized children are collected to rebuild the expression. + fn finish(&mut self, expr: BoundExpression, child_count: usize) -> VortexResult<()> { + vortex_ensure!( + self.results.len() >= child_count, + "bound-expression optimizer lost a child result" + ); + let first_child = self.results.len() - child_count; + let children = if self.results[first_child..] + .iter() + .any(|result| result.changed) + { + let mut children = Vec::with_capacity(child_count); + children.extend(self.results.drain(first_child..).map(|result| result.expr)); + OptimizedChildren::Changed(children) + } else { + self.results.truncate(first_child); + OptimizedChildren::Unchanged + }; + self.finish_node(expr, children) + } + + /// Mark an optimized replacement subtree as changed from the expression it replaced. + /// + /// Optimizing the replacement may itself report no changes, but applying the rule that produced + /// it must remain visible to the original expression's parent. + fn finish_rewrite(&mut self) -> VortexResult<()> { + let Some(mut result) = self.results.pop() else { + vortex_bail!("bound-expression optimizer lost a rewrite result") + }; + result.changed = true; + self.results.push(result); + Ok(()) + } + + /// Rebuild a node when necessary, then apply the first matching root rewrite. + /// + /// A node with no matching rule is pushed onto the result stack. A successful replacement is + /// validated and scheduled for its own bottom-up optimization. + fn finish_node( + &mut self, + original: BoundExpression, + children: OptimizedChildren, + ) -> VortexResult<()> { + let (current, children_changed) = match children { + OptimizedChildren::Unchanged => (original, false), + OptimizedChildren::Changed(children) => { + let original_dtype = original.dtype().clone(); + let rebuilt = original.with_children(children)?; + vortex_ensure!( + rebuilt.dtype() == &original_dtype, + "optimizing children changed a node dtype from {original_dtype} to {}", + rebuilt.dtype() + ); + (rebuilt, true) + } + }; + + let Some((rule_name, replacement)) = self.optimizer.rewrite_root(¤t)? else { + self.results.push(OptimizationResult { + expr: current, + changed: children_changed, + }); + return Ok(()); + }; + + vortex_ensure!( + replacement.dtype() == current.dtype(), + "bound-expression rewrite rule {rule_name} changed dtype from {} to {}", + current.dtype(), + replacement.dtype() + ); + vortex_ensure!( + replacement != current, + "bound-expression rewrite rule {rule_name} returned an unchanged expression" + ); + if self.rewrite_count >= self.optimizer.max_rewrites { + vortex_bail!( + "Exceeded bound-expression rewrite limit of {} while applying {rule_name} \ + (possible rewrite cycle)", + self.optimizer.max_rewrites + ); + } + self.rewrite_count += 1; + + self.tasks.reserve(2); + self.tasks.push(Task::FinishRewrite); + self.tasks.push(Task::Visit(replacement)); + Ok(()) + } +} + +enum Task { + /// Visit an expression by scheduling its children followed by [`Task::Finish`]. + /// + /// Childless expressions are finished immediately without scheduling another task. + Visit(BoundExpression), + /// Finish an original expression after all of its children have been optimized. + /// + /// The `child_count` most recent results belong to this expression. They are used to rebuild + /// the node when any child changed, after which the optimizer tries to rewrite the node itself. + Finish { + expr: BoundExpression, + child_count: usize, + }, + /// Finish a successful root rewrite after its replacement subtree has been optimized. + /// + /// Unlike [`Task::Finish`], this task has no children to collect or node to rebuild: the + /// replacement's result is already on the result stack. It marks that result as changed so the + /// rewrite remains visible even when optimizing the replacement made no further changes. + FinishRewrite, +} + +/// Optimized children, distinguishing reusable originals from a changed child list. +enum OptimizedChildren { + Unchanged, + Changed(Vec), +} + +/// A completed subtree and whether any rewrite changed it. +struct OptimizationResult { + expr: BoundExpression, + changed: bool, +} + +type TaskStack = SmallVec<[Task; 16]>; +type ResultStack = SmallVec<[OptimizationResult; 16]>; + +#[cfg(test)] +mod tests { + use vortex_error::VortexResult; + + use super::BoundExpressionOptimizer; + use super::BoundExpressionRewriteRule; + 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 BoundExpressionRewriteRule 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 optimizer = BoundExpressionOptimizer::empty(); + optimizer.register_rule(IsNullTo(true)); + optimizer.register_rule(IsNullTo(false)); + + assert_eq!(optimizer.optimize(&input)?, bound::lit(true)); + Ok(()) + } + + #[derive(Debug)] + struct IsNullToReducibleTree; + + impl BoundExpressionRewriteRule 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 BoundExpressionRewriteRule 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 optimizer = BoundExpressionOptimizer::empty(); + optimizer.register_rule(IsNullToReducibleTree); + optimizer.register_rule(AndFalse); + + assert_eq!(optimizer.optimize(&input)?, bound::lit(false)); + Ok(()) + } + + #[derive(Debug)] + struct WrongDType; + + impl BoundExpressionRewriteRule 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 optimizer = BoundExpressionOptimizer::empty(); + optimizer.register_rule(WrongDType); + + assert!(optimizer.optimize(&input).is_err()); + } + + #[derive(Debug)] + struct Unchanged; + + impl BoundExpressionRewriteRule 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 optimizer = BoundExpressionOptimizer::empty(); + optimizer.register_rule(Unchanged); + + assert!(optimizer.optimize(&input).is_err()); + } + + #[derive(Debug)] + struct ToggleBoolean; + + impl BoundExpressionRewriteRule 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 optimizer = BoundExpressionOptimizer::empty().with_max_rewrites(4); + optimizer.register_rule(ToggleBoolean); + + assert!(optimizer.optimize(&bound::lit(true)).is_err()); + } + + #[derive(Debug)] + struct RootToOne(ExpressionId); + + impl BoundExpressionRewriteRule 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 optimizer = BoundExpressionOptimizer::empty(); + optimizer.register_rule(RootToOne(input.id())); + + assert_eq!(optimizer.optimize(&input)?, bound::lit(1i32)); + Ok(()) + } + + #[test] + fn unchanged_deep_tree_does_not_recurse() -> VortexResult<()> { + let mut expr = bound::root(DType::Bool(Nullability::NonNullable)); + for _ in 0..20_000 { + expr = bound::not(expr); + } + + assert!( + BoundExpressionOptimizer::empty() + .try_optimize(&expr)? + .is_none() + ); + 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..1a026551288 --- /dev/null +++ b/vortex-array/src/expr/optimizer/rules/binary.rs @@ -0,0 +1,346 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::preserve_dtype; +use crate::expr::BoundExpression; +use crate::expr::ExpressionId; +use crate::expr::bound; +use crate::expr::optimizer::BoundExpressionOptimizer; +use crate::expr::optimizer::BoundExpressionRewriteRule; +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; + +pub(super) fn register(optimizer: &mut BoundExpressionOptimizer) { + optimizer.register_rule(BinaryBoolean); + optimizer.register_rule(BinaryNullComparison); + optimizer.register_rule(FindBetween); +} + +/// Simplifies `AND` and `OR` with literal operands using Kleene boolean semantics. +/// +/// # Example +/// +/// ```text +/// original: and(value, lit(true)) +/// rewritten: value +/// ``` +#[derive(Debug)] +struct BinaryBoolean; + +impl BoundExpressionRewriteRule 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)] +struct BinaryNullComparison; + +impl BoundExpressionRewriteRule 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)] +struct FindBetween; + +impl BoundExpressionRewriteRule 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); + } + + 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::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(()) + } +} 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..0c7992731ce --- /dev/null +++ b/vortex-array/src/expr/optimizer/rules/cast.rs @@ -0,0 +1,46 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use crate::expr::BoundExpression; +use crate::expr::ExpressionId; +use crate::expr::bound; +use crate::expr::optimizer::BoundExpressionOptimizer; +use crate::expr::optimizer::BoundExpressionRewriteRule; +use crate::scalar_fn::ScalarFnVTable; +use crate::scalar_fn::fns::cast::Cast; +use crate::scalar_fn::fns::literal::Literal; + +pub(super) fn register(optimizer: &mut BoundExpressionOptimizer) { + optimizer.register_rule(CastLiteralOrIdentity); +} + +/// 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)] +struct CastLiteralOrIdentity; + +impl BoundExpressionRewriteRule 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..eb485fcdde6 --- /dev/null +++ b/vortex-array/src/expr/optimizer/rules/conditional.rs @@ -0,0 +1,111 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use super::preserve_dtype; +use crate::expr::BoundExpression; +use crate::expr::ExpressionId; +use crate::expr::bound; +use crate::expr::optimizer::BoundExpressionOptimizer; +use crate::expr::optimizer::BoundExpressionRewriteRule; +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; + +pub(super) fn register(optimizer: &mut BoundExpressionOptimizer) { + optimizer.register_rule(ConstantMask); + optimizer.register_rule(ConstantZip); +} + +/// 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)] +struct ConstantMask; + +impl BoundExpressionRewriteRule 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)] +struct ConstantZip; + +impl BoundExpressionRewriteRule 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..5d820dad557 --- /dev/null +++ b/vortex-array/src/expr/optimizer/rules/mod.rs @@ -0,0 +1,31 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use super::BoundExpressionOptimizer; +use crate::dtype::DType; +use crate::expr::BoundExpression; +use crate::scalar_fn::ScalarFnVTableExt; +use crate::scalar_fn::fns::cast::Cast; + +mod binary; +mod cast; +mod conditional; +mod nulls; +mod structural; + +pub(super) fn register(optimizer: &mut BoundExpressionOptimizer) { + binary::register(optimizer); + cast::register(optimizer); + structural::register(optimizer); + nulls::register(optimizer); + conditional::register(optimizer); +} + +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..39c6d9d91fd --- /dev/null +++ b/vortex-array/src/expr/optimizer/rules/nulls.rs @@ -0,0 +1,152 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use vortex_error::VortexResult; + +use super::preserve_dtype; +use crate::expr::BoundExpression; +use crate::expr::ExpressionId; +use crate::expr::bound; +use crate::expr::optimizer::BoundExpressionOptimizer; +use crate::expr::optimizer::BoundExpressionRewriteRule; +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; + +pub(super) fn register(optimizer: &mut BoundExpressionOptimizer) { + optimizer.register_rule(RemoveRedundantFillNull); + optimizer.register_rule(CaseWhenToFillNull); +} + +/// 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)] +struct RemoveRedundantFillNull; + +impl BoundExpressionRewriteRule 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)] +struct CaseWhenToFillNull; + +impl BoundExpressionRewriteRule 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..237d23476c5 --- /dev/null +++ b/vortex-array/src/expr/optimizer/rules/structural.rs @@ -0,0 +1,261 @@ +// 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 crate::dtype::FieldNames; +use crate::expr::BoundExpression; +use crate::expr::ExpressionId; +use crate::expr::bound; +use crate::expr::optimizer::BoundExpressionOptimizer; +use crate::expr::optimizer::BoundExpressionRewriteRule; +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; + +pub(super) fn register(optimizer: &mut BoundExpressionOptimizer) { + optimizer.register_rule(GetItemFromPack); + optimizer.register_rule(MergeToPack); + optimizer.register_rule(SelectFromPack); +} + +/// 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)] +struct GetItemFromPack; + +impl BoundExpressionRewriteRule 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)] +struct MergeToPack; + +impl BoundExpressionRewriteRule 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)] +struct SelectFromPack; + +impl BoundExpressionRewriteRule for SelectFromPack { + fn expression_id(&self) -> ExpressionId { + Select.id() + } + + fn rewrite(&self, expr: &BoundExpression) -> VortexResult> { + let selection = expr.as_::