Skip to content

Commit 30f904a

Browse files
gabotechsgabotechs
committed
Add aqe docs
1 parent 9867814 commit 30f904a

3 files changed

Lines changed: 297 additions & 0 deletions

File tree

‎docs/source/advanced/08-adaptive-query-execution.md‎

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
# Using Adaptive Query Execution
2+
3+
Adaptive Query Execution (AQE) lets Distributed DataFusion choose the number of
4+
tasks for each stage while the query is running. Instead of assigning every
5+
stage a task count statically at planning time, the coordinator samples each
6+
producer at its stage boundary and uses the observed data to size the stage
7+
above it.
8+
9+
AQE currently adapts **distributed task counts**. It does not rerun DataFusion's
10+
physical optimizer or replace joins, aggregates, or other operators during
11+
execution.
12+
See [How Adaptive Query Execution Works](../learn/03-how-adaptive-query-execution-works.md)
13+
for the execution flow and sampling model.
14+
15+
## Enable AQE
16+
17+
Enable AQE on the coordinating session with
18+
`with_distributed_dynamic_task_count(true)`:
19+
20+
```rust
21+
let state = SessionStateBuilder::new()
22+
.with_default_features()
23+
.with_distributed_worker_resolver(worker_resolver)
24+
.with_distributed_planner()
25+
+ .with_distributed_dynamic_task_count(true)?
26+
.build();
27+
```
28+
29+
## Configuration
30+
31+
The following settings affect AQE decisions:
32+
33+
| Setting | Default | Effect |
34+
|-------------------------------------------|--------:|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
35+
| `distributed.dynamic_task_count` | `false` | Enables runtime task-count selection. |
36+
| `distributed.dynamic_bytes_per_partition` | 16 MiB | Sets the compute-cost budget per output partition. Lower values generally produce more tasks; higher values produce fewer. Configure it with `with_distributed_dynamic_bytes_per_partition`. |
37+
38+
## Custom data sources
39+
40+
AQE depends on both planning-time statistics and execution-time metrics. A
41+
custom leaf should provide all of the following:
42+
43+
- A registered `TaskEstimator` that implements `scale_up_leaf_node` for the task
44+
count selected by AQE. Without an estimator, the planner treats an unknown
45+
leaf as limited to one task.
46+
- A useful implementation of `ExecutionPlan::partition_statistics` in the custom
47+
data source. The more accurate the statistics are, the better the decisions
48+
AQE can make.
49+
- Standard DataFusion execution metrics for the custom data source, including an
50+
`output_rows` metric. AQE uses the output-row count of leaves on the stage's
51+
driver path to measure sampling progress. If that metric is absent or not
52+
updated as batches flow, the coordinator cannot reliably extrapolate the final
53+
stage output.
54+
55+
Most of these requirements benefit the custom data source even when Distributed
56+
DataFusion is not involved, so taking the time to provide good implementations
57+
is always worthwhile.
58+
59+
## Considerations
60+
61+
There are a few things to take into account when using AQE:
62+
63+
### Visualizing physical plans
64+
65+
Plan visualization works out of the box with queries using AQE, but visualizing
66+
an unexecuted plan is not very useful. Task-count decisions have not yet been
67+
made, and network boundaries have not yet been injected, so you will effectively
68+
be visualizing the single-node plan before distribution.
69+
70+
This is not a bug; it is how AQE works: the final physical plan is decided
71+
dynamically at runtime.
72+
73+
When using AQE, prefer reading the final plan after the query has executed and
74+
all metrics have been collected from remote workers. See
75+
[Collecting metrics from workers](../user-guide/05-metrics.md).
76+
77+
### Distributed UNION operations
78+
79+
AQE relies on runtime sampling to estimate the size of the different stages
80+
involved in a query. When a `UNION` has multiple children and they all pull data
81+
from remote stages, the stage cannot eagerly yield data from faster children.
82+
It must wait for every child to be sampled before it can start.
83+
84+
This typically happens in systems that abuse `UNION` operations to model range
85+
partitioning. There are many reasons to move away from this pattern, including
86+
schema mismatches between children that silently introduce data loss, discarded
87+
partitioning information, too many Tokio tasks caused by large child counts,
88+
huge plans that are expensive to serialize, and unreadable `EXPLAIN` output. If
89+
your system relies on this pattern, note that one consequence is an increased
90+
time to first batch, even if the total query duration is unaffected.

‎docs/source/index.rst‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@ side of the screen for the answer.
8888
advanced/05-custom-distributed-plans
8989
advanced/06-worker-routing
9090
advanced/07-worker-versioning
91+
advanced/08-adaptive-query-execution
9192

9293
.. toctree::
9394
:maxdepth: 1
@@ -96,6 +97,7 @@ side of the screen for the answer.
9697

9798
learn/01-concepts
9899
learn/02-how-a-distributed-plan-is-built
100+
learn/03-how-adaptive-query-execution-works
99101

100102
.. toctree::
101103
:maxdepth: 1

‎docs/source/learn/03-how-adaptive-query-execution-works.md‎

Lines changed: 205 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,205 @@
1+
# How Adaptive Query Execution Works
2+
3+
The normal distributed planner decides all stage boundaries and task counts
4+
before execution begins. Adaptive Query Execution (AQE) delays those decisions:
5+
it starts at the leaves, executes enough of each producer stage to measure its
6+
output, and uses those measurements to size the next stage.
7+
8+
AQE sizes each stage using a cost model that combines the amount of data flowing
9+
through each node with the estimated compute cost of that node. A stage that
10+
moves a large amount of data through fully streaming operators with little
11+
per-row computation may need only a few tasks. In contrast, a stage containing
12+
compute-intensive operators may be assigned more tasks even when relatively
13+
little data flows through it.
14+
15+
AQE estimates the amount of data flowing through a stage in two ways:
16+
17+
1. For leaf stages, it uses the statistics returned by
18+
`ExecutionPlan::partition_statistics`, relying on DataFusion's upstream
19+
statistics machinery.
20+
2. For intermediate stages, it infers statistics by sampling data at runtime as
21+
the stages below them execute.
22+
23+
## Building the plan from the leaves up
24+
25+
Unlike the static distributed planner, AQE does not build every stage before
26+
execution starts. It begins with the original DataFusion physical plan, where no
27+
network boundaries or distributed task counts have been assigned yet. A
28+
simplified plan might look like this:
29+
30+
```text
31+
SortPreservingMergeExec
32+
SortExec
33+
AggregateExec: Final
34+
RepartitionExec
35+
AggregateExec: Partial
36+
DataSourceExec
37+
```
38+
39+
The coordinator walks this plan from the leaves upward. When it reaches the
40+
first point where a network boundary is required, it has enough information to
41+
close the plan below that boundary as the first stage. Because this stage
42+
contains a data-source leaf, its cost is calculated from
43+
`ExecutionPlan::partition_statistics`. The cost model and the registered
44+
`TaskEstimator` are then used to choose its task count:
45+
46+
```text
47+
SortPreservingMergeExec
48+
SortExec
49+
AggregateExec: Final
50+
▲
51+
│ future network boundary
52+
│
53+
┌───── Stage 1 ── tasks=4 ──────────┐
54+
│ RepartitionExec │
55+
│ AggregateExec: Partial │
56+
│ DataSourceExec │
57+
└───────────────────────────────────┘
58+
▲
59+
└── sized from leaf statistics
60+
```
61+
62+
The coordinator inserts a `SamplerExec` directly below the producer-head
63+
`RepartitionExec` and sends Stage 1 to its workers. The workers begin executing
64+
it before the stage above has been sized. The sampler holds on to the batches
65+
that reach it and reports information about them back to the coordinator:
66+
67+
```text
68+
coordinator
69+
▲
70+
│ sampled rows, bytes,
71+
│ distinct values, and nulls
72+
│
73+
┌───── Stage 1 ── tasks=4 ──────────┐
74+
│ RepartitionExec │
75+
│ SamplerExec │◀── injected by AQE
76+
│ AggregateExec: Partial │
77+
│ DataSourceExec │
78+
└───────────────────────────────────┘
79+
```
80+
81+
The sampled batches are not discarded. They remain buffered until the consumer
82+
stage starts and are then returned as part of the normal execution stream.
83+
84+
The coordinator turns the sample into DataFusion `Statistics` and attaches them
85+
to the network boundary that represents Stage 1. From the point of view of the
86+
operators above it, that boundary now behaves like a leaf with runtime-derived
87+
statistics. AQE can therefore calculate the cost of Stage 2 and choose its task
88+
count using the runtime behavior observed from Stage 1:
89+
90+
```text
91+
┌───── Stage 2 ── tasks=2 ──────────┐
92+
│ SortExec │
93+
│ AggregateExec: Final │
94+
│ [Stage 1] NetworkShuffleExec │◀── runtime statistics
95+
└───────────────────────────────────┘
96+
▲
97+
│ reads from
98+
│
99+
┌───── Stage 1 ── tasks=4 ──────────┐
100+
│ RepartitionExec │
101+
│ SamplerExec │
102+
│ AggregateExec: Partial │
103+
│ DataSourceExec │
104+
└───────────────────────────────────┘
105+
```
106+
107+
Stage 2 is then started and sampled in the same way. This process repeats until
108+
the coordinator reaches the head of the plan. The resulting distributed plan
109+
might look like this, with each task count decided using the best statistics
110+
available at that point:
111+
112+
```text
113+
┌───── Head stage ── tasks=1 ───────┐
114+
│ SortPreservingMergeExec │
115+
│ [Stage 2] NetworkCoalesceExec │
116+
└───────────────────────────────────┘
117+
┌───── Stage 2 ── tasks=2 ──────────┐
118+
│ SortExec │
119+
│ AggregateExec: Final │
120+
│ [Stage 1] NetworkShuffleExec │
121+
└───────────────────────────────────┘
122+
┌───── Stage 1 ── tasks=4 ──────────┐
123+
│ RepartitionExec │
124+
│ SamplerExec │
125+
│ AggregateExec: Partial │
126+
│ DataSourceExec │
127+
└───────────────────────────────────┘
128+
```
129+
130+
The task counts in this example are illustrative. Depending on the observed
131+
data and the operators in each stage, AQE may make an intermediate stage wider
132+
or narrower than the stage below it.
133+
134+
## Implications of progressive planning
135+
136+
Interleaving planning, sampling, and execution has several consequences:
137+
138+
1. Plan fragments are sent to workers progressively. A producer stage is sent
139+
and started first, but fragments that consume its output are not sent until
140+
sampling has produced the runtime statistics needed to size them.
141+
2. Producer tasks may start before they have consumers. The sampler buffers the
142+
batches produced during this period and releases them when the downstream
143+
fragment starts consuming the stage.
144+
3. The final distributed plan does not exist at the beginning of the query.
145+
Network boundaries, task counts, and worker assignments for later stages are
146+
decided as execution moves up the plan.
147+
4. Independent branches can be sampled concurrently, but a stage that consumes
148+
several branches must wait until the required runtime statistics are
149+
available from all of them.
150+
151+
## Deep dive into sampling
152+
153+
Sampling needs to answer two separate questions:
154+
155+
1. How much output has the stage produced so far?
156+
2. How far has the stage progressed through the input that drives that output?
157+
158+
`SamplerExec` answers the first question from the batches that reach it. It
159+
buffers those batches and measures their row count, total and per-column byte
160+
sizes, distinct-value percentages, and null percentages.
161+
162+
The leaves on the stage's **driver path** answer the second question. Their
163+
`output_rows` metrics report how many input rows have been pulled so far, while
164+
`ExecutionPlan::partition_statistics` provides an estimate of how many rows
165+
they will produce in total:
166+
167+
```text
168+
┌───── Stage 1 ─────────────────────────────────────────────┐
169+
│ RepartitionExec │
170+
│ SamplerExec ◀── output produced so far │
171+
│ AggregateExec: Partial │
172+
│ DataSourceExec ◀── driver rows consumed │
173+
└───────────────────────────────────────────────────────────┘
174+
▲
175+
└── estimated total rows
176+
from partition_statistics
177+
```
178+
179+
These measurements deliberately count rows at different points in the plan.
180+
An aggregate or filter may consume many input rows while yielding only a small
181+
number of output rows. AQE therefore does not compare the sampler's output row
182+
count directly with the leaf estimate. Instead, it uses the leaf measurements
183+
to estimate progress, then applies that progress to the output observed by the
184+
sampler.
185+
186+
### The driver path
187+
188+
The driver path contains the operators whose continued input makes the stage
189+
produce more output. For most operators, it follows every child. For
190+
pipeline-breaking joins such as `HashJoinExec`, `NestedLoopJoinExec`, and
191+
`CrossJoinExec`, it follows only the right-hand probe side:
192+
193+
```text
194+
SamplerExec
195+
HashJoinExec
196+
left: build side ── ignored for progress
197+
right: probe side ◀─ driver path
198+
DataSourceExec
199+
```
200+
201+
The build side is excluded because it must be materialized before the join can
202+
yield its first output batch. Consuming build-side rows is setup work, not an
203+
indication of how far the join has progressed through the stream that drives
204+
its output. When a driver path reaches several leaves, their estimated and
205+
consumed row counts are added together.

0 commit comments

Comments
 (0)