Adapative Query Runner (AQE) PoC

> **Status:** Draft design, updated September 6, 2026

## 1. Summary

This issue tracks an incremental Adaptive Query Runner (AQE) proof of concept for Pinot's multi-stage query engine (MSE).

The PoC has two complementary tracks:

1. **Executable vertical slice:** make a safe runtime-informed execution decision at an existing pipeline breaker.
2. **Stage-boundary infrastructure:** add the three seams required to adapt an undispatched downstream stage after its producer completes: materialized exchange, staged dispatch, and late-bound routing.

The four PRs deliberately stop before changing downstream worker counts. They prove the execution model and establish the narrow seam where a future deterministic partition-coalescing rule can rewrite `M` logical partitions onto `K` consumer workers.

## 2. Motivation

Pinot MSE currently plans concrete workers and eagerly dispatches streaming-connected stages. Mailbox routing is therefore bound to worker identities before runtime producer statistics are available.

Changing downstream parallelism after a producer finishes is unsafe unless all three conditions hold:

1. Producer output survives independently of the original consumer workers.
2. The broker owns a completion barrier and has not dispatched the consumer yet.
3. Consumer routing is assigned only after the runtime output is known.

Without these seams, a policy object that computes a new worker count is dead code: no production path can safely apply its decision.

## 3. Goals

- Demonstrate one real AQE decision using completed runtime data.
- Materialize HASH exchange output with consumer-independent identities.
- Delay downstream dispatch across a materialized exchange.
- Bind completed materialized outputs to not-yet-dispatched consumer workers.
- Preserve deterministic planning, routing, and failure behavior.
- Keep every phase independently reviewable and disabled or narrowly scoped.
- Preserve the existing execution path when the PoC features are not selected.

## 4. Non-goals

- Join reordering.
- Dynamic stage insertion or removal.
- Rewriting or cancelling a stage that is already running.
- Stage retries or speculative execution.
- Changing worker counts in AQE-PR-2 through AQE-PR-4.
- Production-ready spill quotas, admission control, or long-term storage.
- Supporting multiple materialized inputs for one consumer stage in this PoC.

## 5. Architecture

### 5.1 Executable AQE rule

AQE-PR-1 uses the existing dynamic-broadcast pipeline breaker. The build side is fully materialized before Pinot compiles the probe-side leaf requests, so an empty build result is a safe decision point for a SEMI JOIN.

```text
build-side mailbox exchange
|
v
pipeline breaker completes
|
v
build side empty?
/ \
yes no
| |
skip probe SSE build normal SSE
request creation leaf requests
| |
+-------+------+
|
v
normal LeafOperator lifecycle
|
v
mailbox EOS + existing statistics
```

The optimization is not applied when empty input can still produce output, including global aggregates and grouping sets containing the grand-total `()` set. EXPLAIN also retains normal request construction.

### 5.2 Downstream stage adaptation seam

AQE-PR-2 through AQE-PR-4 create the execution seam for future downstream parallelism changes.

```text
broker planning
|
v
mark one eligible HASH exchange
|
v
build deterministic dispatch-group DAG
|
v
dispatch producer group
|
v
producer workers write logical partition files
|
v
OpChainComplete(handle + rows + bytes) from every worker
|
v
all-worker success + stream-close barrier
|
v
validate exact P x M handle coverage
|
v
bind logical partition p to consumer worker p
|
v
serialize and dispatch the consumer group
|
v
consumer workers read their assigned local/remote handles
```

Where:

- `P` is the number of producer workers.
- `M` is the planned number of logical HASH partitions and consumer workers.
- The current PoC binds partition `p` to worker `p`; it does not yet compute `M -> K`.

## 6. Modules and seams

| Module | Responsibility | Interface exposed to the orchestration path |
|---|---|---|
| `MaterializedExchangePlanner` | Select one deterministic eligible HASH exchange. | Mark the send/receive nodes as materialized without changing assignments. |
| Materialized mailbox store/transport | Atomically publish, fetch, and clean producer-local partition files. | Produce and consume a `MaterializedPartitionHandle`. |
| `StreamingQuerySession` | Accumulate unique worker completion reports, statistics, handles, and stream completion. | Await successful stage coverage and expose completed outputs. |
| `StageDispatchGraph` | Contract streaming-connected stages and preserve materialized edges as dependencies. | Return ready groups and enforce dispatched/completed transitions. |
| `MaterializedPartitionRouter` | Validate exact output coverage and deterministically assign handles. | Return replacement consumer `WorkerMetadata` with bound materialized inputs. |
| `QueryDispatcher` | Own the deadline, lifecycle ordering, binding, serialization, and failure propagation. | Dispatch only ready groups and run stage 0 locally when its group becomes ready. |
| Materialized receive operator | Read assigned partition handles. | Use late-bound handles when present and retain the static fallback. |

The important seam is inside broker dispatch: the consumer plan is already structurally planned but has not been serialized or submitted. Only its materialized input assignment may change.

## 7. Materialized data model

A materialized partition is identified by:

```text
(requestId, producerStageId, producerWorkerId, logicalPartitionId)
```

The identity deliberately excludes consumer stage and worker IDs. This allows producer output to exist before a consumer assignment is chosen.

Each `MaterializedPartitionHandle` carries:

| Field | Purpose |
|---|---|
| `requestId` | Fences output to one query. |
| `producerStageId` / `producerWorkerId` | Identifies the producer op-chain. |
| `logicalPartitionId` | Stable HASH partition identity. |
| `host` / `transferPort` | Locates the serving worker. |
| `opaqueFileId` | Store-owned file reference. |
| `rowCount` / `byteCount` | Runtime statistics available to a future AQE rule. |

Empty logical partitions are published as handles as well. This makes exact coverage validation possible and distinguishes an empty partition from missing output.

## 8. Planning and eligibility

The materialized-exchange planner marks at most one exchange and requires it to be:

- HASH distributed on both send and receive sides.
- Unsorted.
- Not pre-partitioned.
- Single-consumer.
- Backed by a server-side receiver stage.

The planner does not change stage IDs, worker IDs, worker placement, mailbox metadata, or parallelism.

## 9. Dispatch lifecycle

A dispatch group follows this state model:

```text
UNDISPATCHED --dependencies complete--> READY
READY --bind + submit----------> DISPATCHED
DISPATCHED --all workers succeed
+ all streams close------> COMPLETED

worker failure / stream failure / timeout --> query failure; dependents remain undispatched
```

Streaming and pipeline-breaker edges are contracted into the same dispatch group. A materialized edge becomes a dependency between groups.

For each ready group, the broker performs these operations in order:

1. Select the group's remote stages; stage `0` remains broker-local.
2. Read materialized handles produced by completed dependencies.
3. Validate and bind the consumer's materialized inputs.
4. Mark the dispatch group as dispatched.
5. Serialize and submit worker requests.
6. Await every expected `(stageId, workerId)` success report.
7. Await stream closure and then mark the group completed.

One absolute query deadline is shared across all waves; each wave does not receive a fresh timeout.

## 10. Routing and consistency invariants

Before dispatching a materialized consumer, the router enforces:

- Dense producer and consumer worker IDs.
- Exactly `P x M` handles for the producer stage.
- Matching `requestId` and `producerStageId`.
- Valid producer-worker and logical-partition ranges.
- No duplicate handle identity.
- No missing logical partition, including empty partitions.
- No rebinding of consumer metadata that already contains materialized inputs.

The returned worker metadata preserves worker ID, placement/mailbox information, and custom properties. Handle order is deterministic.

The PoC introduces no mutable plan version. The original dispatch plan acts as plan version `0`; stage IDs, worker IDs, and worker counts remain immutable. The only late mutation is the materialized-input list of an undispatched consumer. A future `M -> K` implementation should add explicit plan-version and stage-attempt fencing before it changes assignments.

## 11. Failure, cleanup, and compatibility

- Partition files are atomically published; incomplete output is not visible.
- A failed or missing producer report cannot release the downstream barrier.
- Unknown or duplicate completion reports cannot satisfy expected coverage.
- Query failure fans cancellation out over the open dispatch streams.
- Materialized output is cleaned on consumption, cancellation, failure, and shutdown.
- Local and remote reads use the same logical handle identity.
- `materializedExchange=true` and `stagedDispatch=true` are disabled by default.
- When the streaming/staged path is enabled, all participating servers must support it; the PoC does not silently fall back after partial dispatch.
- Worker metadata without assigned handles continues to use the static materialized-read path for compatibility within the PoC stack.

## 12. Roadmap PRs

- [ ] **AQE-PR-1 — Minimal executable AQE rule:** #19470
- Skip probe-side leaf execution when a completed dynamic-broadcast SEMI JOIN build side is empty.

- [ ] **AQE-PR-2 — Materialized exchange:** #19471
- Add consumer-independent logical partition files, output handles, runtime row/byte statistics, transport, and cleanup.

- [ ] **AQE-PR-3 — Staged dispatch:** #19472
- Add the broker-owned dispatch-group DAG and all-worker completion barriers without changing assignments.

- [ ] **AQE-PR-4 — Late-bound routing:** #19483
- Validate exact producer output coverage and bind handles before consumer serialization and dispatch.

## 13. Validation strategy

Each phase includes focused tests at its seam:

- **AQE-PR-1:** empty-build skip behavior, empty-input-sensitive plans, request construction, operator lifecycle, and statistics propagation.
- **AQE-PR-2:** planner/SerDe, local and remote file reads, backpressure, atomic publication, empty partitions, cleanup, and exact output validation.
- **AQE-PR-3:** deterministic graph construction, cycle rejection, duplicate/unexpected reports, stage barriers, failures, deadlines, and stage-0 ordering.
- **AQE-PR-4:** handle SerDe, exact `P x M` coverage, deterministic assignment, invalid/rebound metadata, assigned-handle reads, and the real staged-dispatch production path.

## 14. Follow-up after this PoC

The next functional AQE rule can use completed `byteCount` statistics to coalesce contiguous logical partitions:

```text
M logical partitions
|
v
completed byte statistics
|
v
deterministic coalescing rule
|
v
K contiguous partition ranges, K <= M
|
v
rewrite only the undispatched consumer assignment
|
v
late-bind ranges and dispatch K workers
```

Production hardening should additionally cover plan-version/stage-attempt fencing, retries, quotas, spill cleanup after broker failure, metrics and tracing, rollout controls, mixed-version testing, and end-to-end performance validation.

Start by reading the four roadmap issues (#19470, #19471, #19472, and #19483) and the named modules: MaterializedExchangePlanner, StreamingQuerySession, StageDispatchGraph, MaterializedPartitionRouter, and QueryDispatcher. Use their focused validation areas as the first tests to run. Done means the selected seam is independently tested, preserves the existing path when disabled, and maintains the stated dispatch and routing invariants.

Written by the indexing model from the issue text.