jayzhan211 commented on code in PR #24281:
URL: https://github.com/apache/datafusion/pull/24281#discussion_r4401974610
##########
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:
@zhuqi-lucas thanks for this, I can reproduce the speedup locally (-44% /
-33% on the #26841 cases). One concern: reusing the schema also drops the
normalization that downstream optimizer rules got for free. A custom rule that
rewrites projection expressions with `map_expressions` keeps the old schema; on
`main` the next `OptimizeProjections` pass re-derived it, but with this PR the
stale schema reaches the physical planner. With a rule registered via
`ctx.add_optimizer_rule` that rewrites `a AS x` into `CASE WHEN a > 0 THEN a
END AS x`:
```sql
CREATE TABLE t(a INT NOT NULL, b INT NOT NULL) AS VALUES (1, 10), (-1, 20);
SELECT DISTINCT x FROM (SELECT a AS x, b + 1 AS y FROM t);
```
`main` returns `1` and `NULL`; this PR fails with `Physical input schema
should be the same as the one converted from logical input schema ... field
nullability at index 0 [x]: (physical) true vs (logical) false`.
Since #24284 already points at `map_expressions` as the root cause, could we
re-derive the schema there when the expressions changed? That covers in-tree
and downstream rules, and makes `refresh_projection_schema` unnecessary. I
tried this on your head: expr and optimizer tests,
`schema_evolution_nested.slt` and the rest of the sqllogictest suite pass (I
skipped the files that need `parquet-testing` data, which I didn't have
locally).
```diff
LogicalPlan::Projection(Projection {
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 {
+ // The output fields are derived from the expressions,
so
+ // rewritten expressions need a freshly derived schema.
+ expr.map_data(|expr| {
+ Projection::try_new(expr,
input).map(LogicalPlan::Projection)
+ })?
+ } else {
+ expr.update_data(|expr| {
+ LogicalPlan::Projection(Projection {
+ expr,
+ input,
+ schema,
+ })
+ })
+ }
+ }
```
Could you add two tests with such a rule: one that returns
`Transformed::yes` and checks the final projection's schema equals
`projection_schema(&p.input, &p.expr)?`, and one that returns `Transformed::no`
and checks the schema `Arc` is kept (`Arc::ptr_eq`). Here is the rule I used:
```rs
#[derive(Debug)]
struct RewriteToCase;
impl OptimizerRule for RewriteToCase {
fn name(&self) -> &str {
"rewrite_to_case"
}
fn apply_order(&self) -> Option<ApplyOrder> {
Some(ApplyOrder::BottomUp)
}
fn rewrite(
&self,
plan: LogicalPlan,
_config: &dyn OptimizerConfig,
) -> Result<Transformed<LogicalPlan>> {
if !matches!(plan, LogicalPlan::Projection(_)) {
return Ok(Transformed::no(plan));
}
plan.map_expressions(|e| match e {
Expr::Alias(Alias {
expr,
relation,
name,
metadata,
}) if name == "x" && matches!(*expr, Expr::Column(_)) => {
let new = when(expr.as_ref().clone().gt(lit(0)),
*expr).end()?;
Ok(Transformed::yes(Expr::Alias(Alias {
expr: Box::new(new),
relation,
name,
metadata,
})))
}
e => Ok(Transformed::no(e)),
})
}
}
```
If changing `map_expressions` feels too broad for this PR, the alternative
is an upgrade-guide entry telling rule authors to rebuild projections with
`Projection::try_new` after rewriting their expressions.
--
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]
