jayzhan211 commented on code in PR #24281:
URL: https://github.com/apache/datafusion/pull/24281#discussion_r4412173895
##########
datafusion/expr/src/logical_plan/tree_node.rs:
##########
@@ -553,13 +555,25 @@ impl LogicalPlan {
expr,
input,
schema,
- }) => expr.map_elements(f)?.update_data(|expr| {
- LogicalPlan::Projection(Projection {
- expr,
- input,
- schema,
- })
- }),
+ }) => {
+ let expr = expr.map_elements(f)?;
+ if expr.transformed {
+ // A projection's output fields are derived from its
+ // expressions, so rewritten expressions need a freshly
+ // derived schema.
+ expr.map_data(|expr| {
+ Projection::try_new(expr,
input).map(LogicalPlan::Projection)
Review Comment:
@zhuqi-lucas thanks for moving this into `map_expressions`, the derived
schema now follows rule rewrites as expected. One side effect goes beyond
"errors earlier": a caller that rewrites a projection's expressions to columns
its *next* input will provide, then swaps the input with `with_new_exprs` (the
pattern `ExtractLeafExpressions` uses for other nodes), now gets an error on a
`Projection`. On `main` this call returned `Ok`:
```rs
let plan = LogicalPlanBuilder::from(test_table_scan()?)
.project(vec![col("a").alias("x")])?
.build()?;
// Err(Schema error: No field named __extracted_1.)
plan.map_expressions(|e| match e {
Expr::Alias(Alias { name, .. }) => {
Ok(Transformed::yes(col("__extracted_1").alias(name)))
}
e => Ok(Transformed::no(e)),
})?;
```
Renaming rewrites like `unalias` also change the returned `schema()` now.
Could you add the `api change` label and a short note to
`docs/source/library-user-guide/upgrading/56.0.0.md`? It should say that
`map_expressions` re-derives a `Projection`'s schema when its expressions
change, and that rewrites referencing columns of a not-yet-installed input
should rebuild the projection with `Projection::try_new_with_schema` instead.
##########
datafusion/optimizer/src/analyzer/type_coercion.rs:
##########
@@ -259,19 +259,30 @@ fn analyze_internal(
};
// apply coercion rewrite all expressions in the plan individually
- plan.map_expressions(|expr| {
+ let coerced = plan.map_expressions(|expr| {
// Preserve the names of grouping set members, not of the grouping set
itself.
// An aliased grouping set is no longer recognized as a grouping set.
if let Expr::GroupingSet(_) = &expr {
expr.map_children(&mut coerce_expr)
} else {
coerce_expr(expr)
}
- })?
- // some plans need extra coercion after their expressions are coerced
- .map_data(|plan| expr_rewrite.coerce_plan(plan))?
- // recompute the schema after the expressions have been rewritten as the
types may have changed
- .map_data(|plan| plan.recompute_schema())
+ })?;
+ // `map_expressions` already derived the schema of a projection whose
+ // expressions it rewrote, and `coerce_plan` leaves projections alone.
+ let schema_derived =
Review Comment:
Optional cleanup: this shortcut only works as long as `coerce_plan` never
touches projections. If it ever does, the recompute is skipped without any
error. Unless the extra derivation shows up in a measurement, the original
chain is simpler:
```diff
- let coerced = plan.map_expressions(|expr| {
+ plan.map_expressions(|expr| {
@@
- })?;
- // `map_expressions` already derived the schema of a projection whose
- // expressions it rewrote, and `coerce_plan` leaves projections alone.
- let schema_derived =
- coerced.transformed && matches!(coerced.data,
LogicalPlan::Projection(_));
- coerced
- // some plans need extra coercion after their expressions are
coerced
- .map_data(|plan| expr_rewrite.coerce_plan(plan))?
- // recompute the schema after the expressions have been rewritten
as the types may have changed
- .map_data(|plan| {
- if schema_derived {
- Ok(plan)
- } else {
- plan.recompute_schema()
- }
- })
+ })?
+ // some plans need extra coercion after their expressions are coerced
+ .map_data(|plan| expr_rewrite.coerce_plan(plan))?
+ // recompute the schema after the expressions have been rewritten as
the types may have changed
+ .map_data(|plan| plan.recompute_schema())
```
##########
datafusion/optimizer/src/optimize_projections/mod.rs:
##########
@@ -861,27 +861,131 @@ fn rewrite_projection_given_requirements(
config: &dyn OptimizerConfig,
indices: &RequiredIndices,
) -> Result<Transformed<LogicalPlan>> {
- let Projection { expr, input, .. } = proj;
+ let Projection {
+ expr,
+ input,
+ schema,
+ ..
+ } = proj;
let exprs_used = indices.get_at_indices(&expr);
let required_indices =
RequiredIndices::new().with_exprs(input.schema(), exprs_used.iter());
+ let old_input_schema = Arc::clone(input.schema());
+ let read_indices = required_indices.indices().to_vec();
// rewrite the children projection, and if they are changed rewrite the
// projection down
optimize_projections(Arc::unwrap_or_clone(input), config,
required_indices)?
.transform_data(|input| {
if is_projection_unnecessary(&input, &exprs_used)? {
Ok(Transformed::yes(input))
- } else {
+ } else if exprs_used.iter().any(contains_subquery)
+ || !read_fields_unchanged(
+ &old_input_schema,
+ input.schema(),
+ &read_indices,
+ )
+ {
+ // A subquery is optimized on its own, and rewriting the input
+ // can change the fields this projection reads, so the stored
+ // schema may be stale. Recompute it rather than reuse it.
Projection::try_new(exprs_used, Arc::new(input))
.map(LogicalPlan::Projection)
.map(Transformed::yes)
+ } else {
+ // The retained expressions are a subset of the original ones
and
+ // read unchanged input fields, so their output fields are
+ // unchanged too. Select them from the existing schema instead
of
+ // recomputing each one via `Expr::to_field` (see
+ // `projection_schema`), which costs O(exprs * schema_width).
+ // `merge_consecutive_projections` reuses the schema the same
way.
+ let projected_schema =
+ project_schema_by_indices(&schema, indices.indices())?;
+ Projection::try_new_with_schema(
+ exprs_used,
+ Arc::new(input),
+ projected_schema,
+ )
+ .map(LogicalPlan::Projection)
+ .map(Transformed::yes)
}
})
}
+/// Returns true if the input fields at `read_indices` of `old` are unchanged
in
+/// `new`, the input after it was rewritten for those indices.
+///
+/// The rewritten input either keeps every field of `old` or only the fields at
+/// `read_indices`, in order. Any other shape is treated as changed.
+fn read_fields_unchanged(old: &DFSchema, new: &DFSchema, read_indices:
&[usize]) -> bool {
+ if new.fields().len() == old.fields().len() {
+ read_indices.iter().all(|&i| old.field(i) == new.field(i))
+ } else if new.fields().len() == read_indices.len() {
+ read_indices
+ .iter()
+ .enumerate()
+ .all(|(j, &i)| old.field(i) == new.field(j))
+ } else {
+ false
+ }
+}
+
+/// Returns true if `expr` contains a subquery expression.
+fn contains_subquery(expr: &Expr) -> bool {
+ expr.exists(|e| {
+ Ok(matches!(
+ e,
+ Expr::Exists(_)
+ | Expr::InSubquery(_)
+ | Expr::SetComparison(_)
+ | Expr::ScalarSubquery(_)
+ ))
+ })
+ .unwrap_or(true)
Review Comment:
Optional cleanup: the closure can't fail, so this matches the existing
helpers in `expr.rs`:
```diff
- .unwrap_or(true)
+ .expect("exists closure is infallible")
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
