Skip to content

Commit a71a17e

Browse files
zhuqi-lucaszhuqi-lucas
committed
fix: derive a projection's schema again in map_expressions when its expressions change
Rules that rewrite projection expressions via map_expressions, in tree or downstream, kept the old schema. Derive it again there instead of in SimplifyExpressions, and add tests with a downstream-style rule. TypeCoercion skips its own recompute for a projection map_expressions already rewrote, so the analyzer still derives each schema once.
1 parent a365ac3 commit a71a17e

4 files changed

Lines changed: 155 additions & 55 deletions

File tree

‎datafusion/expr/src/logical_plan/tree_node.rs‎

Lines changed: 21 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -544,6 +544,8 @@ impl LogicalPlan {
544544
/// # Notes
545545
/// * Similar to [`TreeNode::map_children`] but for this node's expressions.
546546
/// * Visits only the top level expressions (Does not recurse into each expression)
547+
/// * A `Projection` whose expressions were rewritten gets its schema derived
548+
/// again from the new expressions; other nodes keep their schema.
547549
pub fn map_expressions<F: FnMut(Expr) -> Result<Transformed<Expr>>>(
548550
self,
549551
mut f: F,
@@ -553,13 +555,25 @@ impl LogicalPlan {
553555
expr,
554556
input,
555557
schema,
556-
}) => expr.map_elements(f)?.update_data(|expr| {
557-
LogicalPlan::Projection(Projection {
558-
expr,
559-
input,
560-
schema,
561-
})
562-
}),
558+
}) => {
559+
let expr = expr.map_elements(f)?;
560+
if expr.transformed {
561+
// A projection's output fields are derived from its
562+
// expressions, so rewritten expressions need a freshly
563+
// derived schema.
564+
expr.map_data(|expr| {
565+
Projection::try_new(expr, input).map(LogicalPlan::Projection)
566+
})?
567+
} else {
568+
expr.update_data(|expr| {
569+
LogicalPlan::Projection(Projection {
570+
expr,
571+
input,
572+
schema,
573+
})
574+
})
575+
}
576+
}
563577
LogicalPlan::Values(Values { schema, values }) => values
564578
.map_elements(f)?
565579
.update_data(|values| LogicalPlan::Values(Values { schema, values })),

‎datafusion/optimizer/src/analyzer/type_coercion.rs‎

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -259,19 +259,30 @@ fn analyze_internal(
259259
};
260260

261261
// apply coercion rewrite all expressions in the plan individually
262-
plan.map_expressions(|expr| {
262+
let coerced = plan.map_expressions(|expr| {
263263
// Preserve the names of grouping set members, not of the grouping set itself.
264264
// An aliased grouping set is no longer recognized as a grouping set.
265265
if let Expr::GroupingSet(_) = &expr {
266266
expr.map_children(&mut coerce_expr)
267267
} else {
268268
coerce_expr(expr)
269269
}
270-
})?
271-
// some plans need extra coercion after their expressions are coerced
272-
.map_data(|plan| expr_rewrite.coerce_plan(plan))?
273-
// recompute the schema after the expressions have been rewritten as the types may have changed
274-
.map_data(|plan| plan.recompute_schema())
270+
})?;
271+
// `map_expressions` already derived the schema of a projection whose
272+
// expressions it rewrote, and `coerce_plan` leaves projections alone.
273+
let schema_derived =
274+
coerced.transformed && matches!(coerced.data, LogicalPlan::Projection(_));
275+
coerced
276+
// some plans need extra coercion after their expressions are coerced
277+
.map_data(|plan| expr_rewrite.coerce_plan(plan))?
278+
// recompute the schema after the expressions have been rewritten as the types may have changed
279+
.map_data(|plan| {
280+
if schema_derived {
281+
Ok(plan)
282+
} else {
283+
plan.recompute_schema()
284+
}
285+
})
275286
}
276287

277288
/// Rewrite expressions to apply type coercion.

‎datafusion/optimizer/src/optimize_projections/mod.rs‎

Lines changed: 108 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -953,8 +953,9 @@ fn contains_subquery(expr: &Expr) -> bool {
953953
///
954954
/// Pruning unreferenced sibling columns cannot change the retained fields, so
955955
/// as long as `schema` is current for the projection's expressions and input,
956-
/// the sliced schema is identical to the one [`projection_schema`] would
957-
/// recompute, at O(k) instead of O(exprs * schema_width).
956+
/// the sliced fields are identical to the ones [`projection_schema`] would
957+
/// recompute, at O(k) instead of O(exprs * schema_width). Functional
958+
/// dependencies can be weaker than recomputed ones, never stronger.
958959
///
959960
/// [`projection_schema`]: datafusion_expr::logical_plan::projection_schema
960961
fn project_schema_by_indices(
@@ -1018,13 +1019,14 @@ mod tests {
10181019
use std::vec;
10191020

10201021
use crate::optimize_projections::OptimizeProjections;
1021-
use crate::optimizer::Optimizer;
1022+
use crate::optimizer::{ApplyOrder, Optimizer};
10221023
use crate::test::{
10231024
assert_fields_eq, scan_empty, test_table_scan, test_table_scan_fields,
10241025
test_table_scan_with_name,
10251026
};
1026-
use crate::{OptimizerContext, OptimizerRule};
1027+
use crate::{OptimizerConfig, OptimizerContext, OptimizerRule};
10271028
use arrow::datatypes::{DataType, Field, Schema};
1029+
use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion};
10281030
use datafusion_common::{
10291031
Column, DFSchema, DFSchemaRef, JoinType, Result, ScalarValue, TableReference,
10301032
};
@@ -1035,7 +1037,7 @@ mod tests {
10351037
build_join_schema,
10361038
builder::table_scan_with_filters,
10371039
col,
1038-
expr::{self, Cast},
1040+
expr::{self, Alias, Cast},
10391041
lit,
10401042
logical_plan::{builder::LogicalPlanBuilder, table_scan},
10411043
not, try_cast, when,
@@ -2826,6 +2828,107 @@ mod tests {
28262828
assert_projection_schema_derived(&optimized)
28272829
}
28282830

2831+
/// A downstream rule that rewrites `a AS x` into
2832+
/// `CASE WHEN a > 0 THEN a END AS x` via `map_expressions`, so `x` becomes
2833+
/// nullable. With `transform` false it rewrites nothing.
2834+
#[derive(Debug)]
2835+
struct RewriteToCase {
2836+
transform: bool,
2837+
}
2838+
2839+
impl OptimizerRule for RewriteToCase {
2840+
fn name(&self) -> &str {
2841+
"rewrite_to_case"
2842+
}
2843+
2844+
fn apply_order(&self) -> Option<ApplyOrder> {
2845+
Some(ApplyOrder::BottomUp)
2846+
}
2847+
2848+
fn rewrite(
2849+
&self,
2850+
plan: LogicalPlan,
2851+
_config: &dyn OptimizerConfig,
2852+
) -> Result<Transformed<LogicalPlan>> {
2853+
if !self.transform || !matches!(plan, LogicalPlan::Projection(_)) {
2854+
return Ok(Transformed::no(plan));
2855+
}
2856+
plan.map_expressions(|e| match e {
2857+
Expr::Alias(Alias {
2858+
expr,
2859+
relation,
2860+
name,
2861+
metadata,
2862+
}) if name == "x" && matches!(*expr, Expr::Column(_)) => {
2863+
let new = when(expr.as_ref().clone().gt(lit(0u32)), *expr).end()?;
2864+
Ok(Transformed::yes(Expr::Alias(Alias {
2865+
expr: Box::new(new),
2866+
relation,
2867+
name,
2868+
metadata,
2869+
})))
2870+
}
2871+
e => Ok(Transformed::no(e)),
2872+
})
2873+
}
2874+
}
2875+
2876+
/// `SELECT DISTINCT x FROM (SELECT a AS x, b + 1 AS y FROM test)`
2877+
fn distinct_over_x_and_y() -> Result<LogicalPlan> {
2878+
LogicalPlanBuilder::from(test_table_scan()?)
2879+
.project(vec![col("a").alias("x"), (col("b") + lit(1u32)).alias("y")])?
2880+
.distinct()?
2881+
.project(vec![col("x")])?
2882+
.build()
2883+
}
2884+
2885+
/// A projection whose expressions a downstream rule rewrote carries the
2886+
/// schema the new expressions derive, even after `OptimizeProjections`
2887+
/// prunes it and reuses its schema.
2888+
#[test]
2889+
fn projection_schema_follows_rule_rewrite() -> Result<()> {
2890+
let plan = distinct_over_x_and_y()?;
2891+
assert!(!plan.schema().field(0).is_nullable());
2892+
2893+
let optimizer = Optimizer::with_rules(vec![
2894+
Arc::new(RewriteToCase { transform: true }),
2895+
Arc::new(OptimizeProjections::new()),
2896+
]);
2897+
let optimized = optimizer.optimize(plan, &OptimizerContext::new(), observe)?;
2898+
2899+
let mut projections = 0;
2900+
optimized.apply(|plan| {
2901+
if let LogicalPlan::Projection(_) = plan {
2902+
assert_projection_schema_derived(plan)?;
2903+
projections += 1;
2904+
}
2905+
Ok(TreeNodeRecursion::Continue)
2906+
})?;
2907+
assert!(projections > 0, "expected a projection in {optimized}");
2908+
assert!(optimized.schema().field(0).is_nullable());
2909+
Ok(())
2910+
}
2911+
2912+
/// A rule that rewrites nothing keeps the projection's schema as is.
2913+
#[test]
2914+
fn projection_schema_kept_without_rule_rewrite() -> Result<()> {
2915+
let plan = LogicalPlanBuilder::from(test_table_scan()?)
2916+
.project(vec![col("a").alias("x"), (col("b") + lit(1u32)).alias("y")])?
2917+
.build()?;
2918+
let schema = Arc::clone(plan.schema());
2919+
2920+
let rewritten =
2921+
RewriteToCase { transform: false }.rewrite(plan, &OptimizerContext::new())?;
2922+
assert!(!rewritten.transformed);
2923+
assert!(Arc::ptr_eq(&schema, rewritten.data.schema()));
2924+
2925+
// Expressions visited but left unchanged also keep the schema.
2926+
let rewritten = rewritten.data.map_expressions(|e| Ok(Transformed::no(e)))?;
2927+
assert!(!rewritten.transformed);
2928+
assert!(Arc::ptr_eq(&schema, rewritten.data.schema()));
2929+
Ok(())
2930+
}
2931+
28292932
fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {}
28302933

28312934
fn optimize(plan: LogicalPlan) -> Result<LogicalPlan> {

‎datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs‎

Lines changed: 9 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -118,46 +118,18 @@ impl SimplifyExpressions {
118118
))
119119
};
120120

121-
let rewritten = plan
122-
.map_expressions(|expr| {
123-
// Preserve the aliasing of grouping sets.
124-
if let Expr::GroupingSet(_) = &expr {
125-
expr.map_children(&mut rewrite_expr)
126-
} else {
127-
rewrite_expr(expr)
128-
}
129-
})?
130-
.transform_data(rewrite_aggregate_non_aggregate_aggr_expr)?;
131-
132-
if !rewritten.transformed {
133-
return Ok(rewritten);
134-
}
135-
136-
rewritten.map_data(refresh_projection_schema)
121+
plan.map_expressions(|expr| {
122+
// Preserve the aliasing of grouping sets.
123+
if let Expr::GroupingSet(_) = &expr {
124+
expr.map_children(&mut rewrite_expr)
125+
} else {
126+
rewrite_expr(expr)
127+
}
128+
})?
129+
.transform_data(rewrite_aggregate_non_aggregate_aggr_expr)
137130
}
138131
}
139132

140-
/// Recomputes a `Projection`'s output schema after its expressions were
141-
/// simplified.
142-
///
143-
/// `LogicalPlan::map_expressions` replaces a projection's expressions while
144-
/// keeping its existing schema, so simplification can leave the two out of
145-
/// step: constant folding turns a function call, whose field the planner
146-
/// derived as nullable, into a non-null literal, whose field is not. The
147-
/// schema keeps the pre-folding answer.
148-
///
149-
/// Downstream rules that rebuild the projection with `Projection::try_new`
150-
/// used to paper over that by deriving the schema again, which made the
151-
/// discrepancy invisible but also made it depend on which rules happen to
152-
/// fire. Deriving it here instead keeps a projection's schema in step with
153-
/// its own expressions.
154-
fn refresh_projection_schema(plan: LogicalPlan) -> Result<LogicalPlan> {
155-
let LogicalPlan::Projection(Projection { expr, input, .. }) = plan else {
156-
return Ok(plan);
157-
};
158-
Projection::try_new(expr, input).map(LogicalPlan::Projection)
159-
}
160-
161133
impl SimplifyExpressions {
162134
#[expect(missing_docs)]
163135
pub fn new() -> Self {

0 commit comments

Comments
 (0)