adriangb opened a new pull request, #26428:
URL: https://github.com/apache/datafusion/pull/26428
## Which issue does this PR close?
- Closes apache/datafusion#26424.
Stacked on apache/datafusion#26271. Only the last commit is new.
## Rationale for this change
A reference to a `MATERIALIZED` CTE reports no ordering and unknown
partitioning. The continuation then sorts or repartitions rows that the body
already sorted or partitioned:
```sql
WITH c AS MATERIALIZED (SELECT g, count(*) AS n FROM p GROUP BY g)
SELECT g, sum(n) FROM c GROUP BY g;
```
Before:
```
AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[sum(c.n)]
RepartitionExec: partitioning=Hash([g@0], 12), input_partitions=12
AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[sum(c.n)]
MaterializedCteScanExec: name=c
```
After:
```
AggregateExec: mode=SinglePartitioned, gby=[g@0 as g], aggr=[sum(c.n)]
MaterializedCteScanExec: name=c, partitioning=Hash([g@0], 12)
```
On 10 million rows with 1 million groups (12 partitions, `datafusion-cli`,
release build, median of 30 runs), the query takes about 35 ms instead of about
41 ms.
## What changes are included in this PR?
`DefaultPhysicalPlanner::optimize_physical_plan` now does these steps for a
plan with materialized CTEs:
1. It optimizes each body on its own, as the child of a
`MaterializedCteExec` with an empty continuation, so the body has the same
parent as in the full plan. A body that reads another CTE goes after that CTE.
2. It gives each scan the partition count, the output ordering and the hash
partitioning of its optimized body.
3. It optimizes the full plan one time.
4. It keeps the result only if each body in it still has the properties that
its scans report. If not, it optimizes and uses the plan with statistics only.
There is no retry loop. The cost is one full optimizer pass plus one pass
for each body, and one more full pass only when step 4 fails.
A scan with the properties of the body replays one body partition per
partition, and returns an internal error if the partition counts are not the
same at execution. A scan that comes from a separate planning run, such as a
view, is bound again, so that it loses properties from that run.
Planning time per query, physical and logical (`EXPLAIN`, 2000 runs, median
of 9):
| Query | Statistics only (#26271) | Repeated full passes (first version of
#26271) | This PR |
| --- | --- | --- | --- |
| One CTE | 0.36 ms | 0.88 ms | 0.40 ms |
| Chain of 6 CTEs | 1.43 ms | 1.93 ms | 1.60 ms |
Limits:
- An `ORDER BY` at the top of a body without `LIMIT` is removed by
`EnforceSorting`, as before, so the scan has no ordering in that case.
- `EXPLAIN VERBOSE` shows the rules of the full optimizer run that made the
final plan, not the runs on the bodies.
## Are these changes tested?
Yes. `cte_materialized.slt`: a sort that is removed, a sort that is still
needed, a body `ORDER BY` without `LIMIT`, a repartition and a partial
aggregate that are removed, a partitioned join of two references, a chain of
two CTEs, a CTE defined inside a CTE body, a scalar subquery and a view. Unit
tests for the copied properties, the replay order and the partition-count
error. Planner tests check that the optimizer runs one time on the body and one
time on the full plan, and that an optimizer rule that changes the body after
the scans got its properties makes the planner use the plan with statistics
only. A view test checks a view that is not inlined.
## Are there any user-facing changes?
Plans of queries with `MATERIALIZED` CTEs can change.
`MaterializedCteScanExec` shows `partitioning=` and `output_ordering=` in
EXPLAIN when it has them. New public functions:
`MaterializedCteScanExec::with_body_properties`, `body_properties_hold`,
`mirrors_body`, and `MaterializedCteExec::register`.
--
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]