zhuqi-lucas commented on code in PR #24281:
URL: https://github.com/apache/datafusion/pull/24281#discussion_r4412972996
##########
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:
Thanks, added a 56.0.0 upgrade note covering both cases, with the migration
path, and the `api change` label. In-tree, `ExtractLeafExpressions` skips
projections, so nothing here hits it.
##########
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:
Agreed, reverted to the original chain.
##########
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:
Done.
--
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]
