vortex_array/expr/transform/
optimizer.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_dtype::DType;
5use vortex_error::VortexResult;
6
7use crate::expr::Expression;
8use crate::expr::session::ExprSession;
9use crate::expr::session::RewriteRuleRegistry;
10use crate::expr::transform::simplify::simplify;
11use crate::expr::transform::simplify_typed::simplify_typed;
12
13/// A unified optimizer for expressions that can work with or without type information.
14#[derive(Debug, Clone)]
15pub struct ExprOptimizer {
16    rule_registry: RewriteRuleRegistry,
17}
18
19impl ExprOptimizer {
20    /// Creates a new optimizer with the rules in `ExprSession`.
21    pub fn new(session: &ExprSession) -> Self {
22        Self {
23            rule_registry: session.rewrite_rules().clone(),
24        }
25    }
26
27    /// Optimize the given expression without a dtype.
28    pub fn optimize(&self, expr: Expression) -> VortexResult<Expression> {
29        simplify(expr, &self.rule_registry)
30    }
31
32    /// Optimize the expression, with a known dtype. This will also apply rules in `optimize`.
33    pub fn optimize_typed(&self, expr: Expression, dtype: &DType) -> VortexResult<Expression> {
34        simplify_typed(expr, dtype, &self.rule_registry)
35    }
36}