zhuqi-lucas commented on code in PR #24281:
URL: https://github.com/apache/datafusion/pull/24281#discussion_r4402755080
##########
datafusion/optimizer/src/optimize_projections/mod.rs:
##########
@@ -861,27 +861,130 @@ 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())?;
Review Comment:
Thanks @jayzhan211, great catch. Done: `map_expressions` now derives the
projection schema again when its expressions were rewritten, and the
`SimplifyExpressions` helper is removed. CSE builds its projection itself, so
its fix stays. Added both tests with your rule
(`projection_schema_follows_rule_rewrite`,
`projection_schema_kept_without_rule_rewrite`); the first fails without the
change.
`TypeCoercion` already recomputes every schema after coercing, so it now
skips that for a projection `map_expressions` just rewrote, to avoid deriving
it twice. One side effect: a rewrite whose schema can't be derived now errors
inside `map_expressions` rather than at the next `OptimizeProjections`.
##########
datafusion/optimizer/src/optimize_projections/mod.rs:
##########
@@ -861,27 +861,130 @@ 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)
+}
+
+/// Builds the output schema of a projection that keeps only the fields at
+/// `indices` (a sorted, deduplicated subset produced by [`RequiredIndices`])
of
+/// `schema`, reusing the already-computed fields instead of recomputing each
+/// field's type/nullability from the expressions.
+///
+/// Pruning unreferenced sibling columns cannot change the retained fields, so
+/// as long as `schema` is current for the projection's expressions and input,
Review Comment:
Applied, thanks.
--
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]
