Skip to content
Draft
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
46 changes: 37 additions & 9 deletions datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -146,18 +146,46 @@ impl SimplifyExpressions {
))
};

plan.map_expressions(|expr| {
// Preserve the aliasing of grouping sets.
if let Expr::GroupingSet(_) = &expr {
expr.map_children(&mut rewrite_expr)
} else {
rewrite_expr(expr)
}
})?
.transform_data(rewrite_aggregate_non_aggregate_aggr_expr)
let rewritten = plan
.map_expressions(|expr| {
// Preserve the aliasing of grouping sets.
if let Expr::GroupingSet(_) = &expr {
expr.map_children(&mut rewrite_expr)
} else {
rewrite_expr(expr)
}
})?
.transform_data(rewrite_aggregate_non_aggregate_aggr_expr)?;

if !rewritten.transformed {
return Ok(rewritten);
}

rewritten.map_data(refresh_projection_schema)
}
}

/// Recomputes a `Projection`'s output schema after its expressions were
/// simplified.
///
/// `LogicalPlan::map_expressions` replaces a projection's expressions while
/// keeping its existing schema, so simplification can leave the two out of
/// step: constant folding turns a function call, whose field the planner
/// derived as nullable, into a non-null literal, whose field is not. The
/// schema keeps the pre-folding answer.
///
/// Downstream rules that rebuild the projection with `Projection::try_new`
/// used to paper over that by deriving the schema again, which made the
/// discrepancy invisible but also made it depend on which rules happen to
/// fire. Deriving it here instead keeps a projection's schema in step with
/// its own expressions.
fn refresh_projection_schema(plan: LogicalPlan) -> Result<LogicalPlan> {
let LogicalPlan::Projection(Projection { expr, input, .. }) = plan else {
return Ok(plan);
};
Projection::try_new(expr, input).map(LogicalPlan::Projection)
}

impl SimplifyExpressions {
#[expect(missing_docs)]
pub fn new() -> Self {
Expand Down
Loading