Skip to content

Commit 1a5b658

Browse files
gabotechsgabotechs
authored
Add aqe docs (#561)
Adds docs about AQE
1 parent 9867814 commit 1a5b658

3 files changed

Lines changed: 361 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: 269 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,269 @@
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 sent to its workers. In this example, Stage 2 feeds a
108+
`NetworkCoalesceExec`, whose consumer is the single-task head stage. The head
109+
task count is already fixed, so Stage 2 does not need to be sampled. It begins
110+
normal execution when the head requests its output. In plans with more
111+
intermediate stages, the sampling process repeats until the coordinator reaches
112+
such a final coalescing boundary.
113+
114+
The resulting distributed plan might look like this, with each task count
115+
decided using the best statistics available at that point:
116+
117+
```text
118+
┌───── Head stage ── tasks=1 ───────┐
119+
│ SortPreservingMergeExec │
120+
│ [Stage 2] NetworkCoalesceExec │
121+
└───────────────────────────────────┘
122+
┌───── Stage 2 ── tasks=2 ──────────┐
123+
│ SortExec │
124+
│ AggregateExec: Final │
125+
│ [Stage 1] NetworkShuffleExec │
126+
└───────────────────────────────────┘
127+
┌───── Stage 1 ── tasks=4 ──────────┐
128+
│ RepartitionExec │
129+
│ SamplerExec │
130+
│ AggregateExec: Partial │
131+
│ DataSourceExec │
132+
└───────────────────────────────────┘
133+
```
134+
135+
The task counts in this example are illustrative. Depending on the observed
136+
data and the operators in each stage, AQE may make an intermediate stage wider
137+
or narrower than the stage below it.
138+
139+
## Implications of progressive planning
140+
141+
Interleaving planning, sampling, and execution has several consequences:
142+
143+
1. Plan fragments are sent to workers progressively. A producer stage is sent
144+
and started first, but fragments that consume its output are not sent until
145+
sampling has produced the runtime statistics needed to size them.
146+
2. Producer tasks may start before they have consumers. The sampler buffers the
147+
batches produced during this period and releases them when the downstream
148+
fragment starts consuming the stage.
149+
3. The final distributed plan does not exist at the beginning of the query.
150+
Network boundaries, task counts, and worker assignments for later stages are
151+
decided as execution moves up the plan.
152+
4. Independent branches can be sampled concurrently, but a stage that consumes
153+
several branches must wait until the required runtime statistics are
154+
available from all of them.
155+
156+
## Deep dive into sampling
157+
158+
Sampling needs to answer two separate questions:
159+
160+
1. How much output has the stage produced so far?
161+
2. How far has the stage progressed through the input that drives that output?
162+
163+
`SamplerExec` answers the first question from the batches that reach it. It
164+
buffers those batches and measures their row count, total and per-column byte
165+
sizes, distinct-value percentages, and null percentages.
166+
167+
The leaves on the stage's **driver path** answer the second question. Their
168+
`output_rows` metrics report how many input rows have been pulled so far, while
169+
`ExecutionPlan::partition_statistics` provides an estimate of how many rows
170+
they will produce in total:
171+
172+
```text
173+
┌───── Stage 1 ─────────────────────────────────────────────┐
174+
│ RepartitionExec │
175+
│ SamplerExec ◀── output produced so far │
176+
│ AggregateExec: Partial │
177+
│ DataSourceExec ◀── driver rows consumed │
178+
└───────────────────────────────────────────────────────────┘
179+
▲
180+
└── estimated total rows
181+
from partition_statistics
182+
```
183+
184+
These measurements deliberately count rows at different points in the plan.
185+
An aggregate or filter may consume many input rows while yielding only a small
186+
number of output rows. AQE therefore does not compare the sampler's output row
187+
count directly with the leaf estimate. Instead, it uses the leaf measurements
188+
to estimate progress, then applies that progress to the output observed by the
189+
sampler.
190+
191+
### The driver path
192+
193+
The driver path contains the operators whose continued input makes the stage
194+
produce more output. For most operators, it follows every child. For
195+
pipeline-breaking joins such as `HashJoinExec`, `NestedLoopJoinExec`, and
196+
`CrossJoinExec`, it follows only the right-hand probe side:
197+
198+
```text
199+
SamplerExec
200+
HashJoinExec
201+
left: build side ── ignored for progress
202+
right: probe side ◀─ driver path
203+
DataSourceExec
204+
```
205+
206+
The build side is excluded because it must be materialized before the join can
207+
yield its first output batch. Consuming build-side rows is setup work, not an
208+
indication of how far the join has progressed through the stream that drives
209+
its output. When a driver path reaches several leaves, their estimated and
210+
consumed row counts are added together.
211+
212+
### Reporting `LoadInfo`
213+
214+
Each partition inside a `SamplerExec` produces one `LoadInfo` report containing
215+
the measurements needed by the coordinator:
216+
217+
```text
218+
LoadInfo
219+
├── partition
220+
├── rows_ready
221+
├── per_column_bytes_ready
222+
├── per_column_ndv_percentage
223+
├── per_column_null_percentage
224+
├── rows_pulled_from_leaf
225+
└── reached_eos
226+
```
227+
228+
Every task sends its partition reports through its worker-to-coordinator
229+
channel. The coordinator merges reports from all tasks as they arrive:
230+
231+
```text
232+
Worker A / task 0
233+
sampler partition 0 ── LoadInfo ──┐
234+
sampler partition 1 ── LoadInfo ──┤
235+
├──▶ coordinator
236+
Worker B / task 1 │ merge reports
237+
sampler partition 0 ── LoadInfo ──┤ stop at threshold
238+
sampler partition 1 ── LoadInfo ──┘
239+
```
240+
241+
The total number of sampler partitions is the number of partitions per task
242+
multiplied by the number of tasks in the stage. The coordinator collects
243+
reports until enough partitions have produced non-empty output, or until every
244+
partition has reported:
245+
246+
```text
247+
total sampler partitions = partitions per task × stage tasks
248+
249+
enough reports = sampling threshold reached
250+
or all sampler partitions have reported
251+
```
252+
253+
Because sampling may stop before every partition reports, the observed rows,
254+
bytes, and consumed driver rows are first scaled from the number of reporting
255+
partitions to the total partition count.
256+
257+
The coordinator then estimates how much of the driver input has been consumed
258+
and uses that fraction to extrapolate the sampler output:
259+
260+
```text
261+
estimated completion = consumed driver rows / estimated driver rows
262+
estimated stage output = normalized sampled output / estimated completion
263+
```
264+
265+
The resulting row counts, byte sizes, distinct counts, and null counts become
266+
the runtime `Statistics` exposed by the network boundary to the stage above. If
267+
every sampled partition has reached end-of-stream, the sample is treated as
268+
100% complete. If the driver leaves do not provide a row-count estimate, AQE
269+
falls back to a default completion estimate.

0 commit comments

Comments
 (0)