Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions datafusion/core/tests/expr_api/simplification.rs
Original file line number Diff line number Diff line change
Expand Up @@ -547,9 +547,9 @@ fn test_simplify_with_cycle_count(
};
let simplifier = ExprSimplifier::new(info);
let (simplified_expr, count) = simplifier
.simplify_with_cycle_count(input_expr.clone())
.simplify_with_cycle_count_transformed(input_expr.clone())
.expect("successfully evaluated");

let simplified_expr = simplified_expr.data;
assert_eq!(
simplified_expr, expected_expr,
"Mismatch evaluating {input_expr}\n Expected:{expected_expr}\n Got:{simplified_expr}"
Expand Down
48 changes: 41 additions & 7 deletions datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs
Original file line number Diff line number Diff line change
Expand Up @@ -188,7 +188,7 @@ impl<S: SimplifyInfo> ExprSimplifier<S> {
/// assert_eq!(expr, b_lt_2);
/// ```
pub fn simplify(&self, expr: Expr) -> Result<Expr> {
Ok(self.simplify_with_cycle_count(expr)?.0)
Ok(self.simplify_with_cycle_count_transformed(expr)?.0.data)
Copy link
Member Author

@xudong963 xudong963 Apr 17, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Only changed the simplify_with_cycle_count API, keep the simplify API to avoid two API changes.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe after we merge this one in, it would be worth considering changing the simplify API as well -- but I agree that would be good as a follow on pR

}

/// Like [Self::simplify], simplifies this [`Expr`] as much as possible, evaluating
Expand All @@ -198,7 +198,34 @@ impl<S: SimplifyInfo> ExprSimplifier<S> {
///
/// See [Self::simplify] for details and usage examples.
///
#[deprecated(
since = "48.0.0",
note = "Use `simplify_with_cycle_count_transformed` instead"
)]
#[allow(unused_mut)]
pub fn simplify_with_cycle_count(&self, mut expr: Expr) -> Result<(Expr, u32)> {
let (transformed, cycle_count) =
self.simplify_with_cycle_count_transformed(expr)?;
Ok((transformed.data, cycle_count))
}

/// Like [Self::simplify], simplifies this [`Expr`] as much as possible, evaluating
/// constants and applying algebraic simplifications. Additionally returns a `u32`
/// representing the number of simplification cycles performed, which can be useful for testing
/// optimizations.
///
/// # Returns
///
/// A tuple containing:
/// - The simplified expression wrapped in a `Transformed<Expr>` indicating if changes were made
/// - The number of simplification cycles that were performed
///
/// See [Self::simplify] for details and usage examples.
///
pub fn simplify_with_cycle_count_transformed(
&self,
mut expr: Expr,
) -> Result<(Transformed<Expr>, u32)> {
let mut simplifier = Simplifier::new(&self.info);
let mut const_evaluator = ConstEvaluator::try_new(self.info.execution_props())?;
let mut shorten_in_list_simplifier = ShortenInListSimplifier::new();
Expand All @@ -212,6 +239,7 @@ impl<S: SimplifyInfo> ExprSimplifier<S> {
// simplifications can enable new constant evaluation
// see `Self::with_max_cycles`
let mut num_cycles = 0;
let mut has_transformed = false;
loop {
let Transformed {
data, transformed, ..
Expand All @@ -221,13 +249,18 @@ impl<S: SimplifyInfo> ExprSimplifier<S> {
.transform_data(|expr| expr.rewrite(&mut guarantee_rewriter))?;
expr = data;
num_cycles += 1;
// Track if any transformation occurred
has_transformed = has_transformed || transformed;
if !transformed || num_cycles >= self.max_simplifier_cycles {
break;
}
}
// shorten inlist should be started after other inlist rules are applied
expr = expr.rewrite(&mut shorten_in_list_simplifier).data()?;
Ok((expr, num_cycles))
Ok((
Transformed::new_transformed(expr, has_transformed),
num_cycles,
))
}

/// Apply type coercion to an [`Expr`] so that it can be
Expand Down Expand Up @@ -392,15 +425,15 @@ impl<S: SimplifyInfo> ExprSimplifier<S> {
/// let expr = col("a").is_not_null();
///
/// // When using default maximum cycles, 2 cycles will be performed.
/// let (simplified_expr, count) = simplifier.simplify_with_cycle_count(expr.clone()).unwrap();
/// assert_eq!(simplified_expr, lit(true));
/// let (simplified_expr, count) = simplifier.simplify_with_cycle_count_transformed(expr.clone()).unwrap();
/// assert_eq!(simplified_expr.data, lit(true));
/// // 2 cycles were executed, but only 1 was needed
/// assert_eq!(count, 2);
///
/// // Only 1 simplification pass is necessary here, so we can set the maximum cycles to 1.
/// let (simplified_expr, count) = simplifier.with_max_cycles(1).simplify_with_cycle_count(expr.clone()).unwrap();
/// let (simplified_expr, count) = simplifier.with_max_cycles(1).simplify_with_cycle_count_transformed(expr.clone()).unwrap();
/// // Expression has been rewritten to: (c = a AND b = 1)
/// assert_eq!(simplified_expr, lit(true));
/// assert_eq!(simplified_expr.data, lit(true));
/// // Only 1 cycle was executed
/// assert_eq!(count, 1);
///
Expand Down Expand Up @@ -3320,7 +3353,8 @@ mod tests {
let simplifier = ExprSimplifier::new(
SimplifyContext::new(&execution_props).with_schema(schema),
);
simplifier.simplify_with_cycle_count(expr)
let (expr, count) = simplifier.simplify_with_cycle_count_transformed(expr)?;
Ok((expr.data, count))
}

fn simplify_with_cycle_count(expr: Expr) -> (Expr, u32) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -123,10 +123,11 @@ impl SimplifyExpressions {
let name_preserver = NamePreserver::new(&plan);
let mut rewrite_expr = |expr: Expr| {
let name = name_preserver.save(&expr);
let expr = simplifier.simplify(expr)?;
// TODO it would be nice to have a way to know if the expression was simplified
// or not. For now conservatively return Transformed::yes
Ok(Transformed::yes(name.restore(expr)))
let expr = simplifier.simplify_with_cycle_count_transformed(expr)?.0;
Ok(Transformed::new_transformed(
name.restore(expr.data),
expr.transformed,
))
};

plan.map_expressions(|expr| {
Expand Down