andygrove opened a new pull request, #1954:
URL: https://github.com/apache/datafusion-ballista/pull/1954
# Which issue does this PR close?
Closes #1953.
> Stacked on #1951 (`shuffle-inflight-governor`). Until #1951 merges, the
diff here also shows that PR's commits; it will narrow to just the
exchange-reuse commits once #1951 lands.
# Rationale for this change
The distributed planner (`ballista/scheduler/src/planner.rs`,
`plan_query_stages`) inserts shuffle boundaries by walking the physical plan,
but has **no mechanism to detect and reuse identical repeated subplans** —
`grep -r "Reuse" ballista/` finds nothing. When a query contains the same
subtree more than once (common in TPC-H — Q2, Q11, Q15 reference a shared
aggregate/scan subquery), Ballista plans and **executes each occurrence
independently**: redundant scans, shuffles, and compute for work that produces
identical results.
Spark solves this with `ReuseExchange`/`ReuseSubquery`: materialize a shared
exchange once and fan its output out to every consumer. This PR ports the
exchange-reuse half of that model onto Ballista's stage DAG.
The key observation is that **Ballista already has the `ReusedExchangeExec`
indirection built in**. A Ballista exchange is a `ShuffleWriterExec` *stage*; a
consumer references it *by `stage_id`* via `UnresolvedShuffleExec` →
`ShuffleReaderExec`, and `ExecutionStageBuilder` already fans one stage out to
many consumers through `output_links: Vec<usize>`. Reuse never happened only
because the planner minted a fresh `stage_id` for every writer, even for
structurally identical subtrees. So no new operator and no output-attribute
remapping are needed (Ballista physical plans are positional — no exprIds):
reuse reduces to making identical writer subtrees share one `stage_id`.
# What changes are included in this PR?
1. **A post-planning reuse pass**
(`ballista/scheduler/src/physical_optimizer/reuse_exchange.rs`, the analog of
Spark's standalone `ReuseExchangeAndSubquery` rule). `reuse_shuffle_stages`
deduplicates `ShuffleWriter` stages whose **stage-id-normalized protobuf
serialization** is byte-identical, and rewires every dropped stage's consumers'
`UnresolvedShuffleExec` to the surviving `stage_id`.
- **Identity test = exact protobuf bytes** of a stage-id-normalized
writer, produced with the configured `PhysicalExtensionCodec` (the same codec
that already ships stages to executors). Byte-equality is semantic equality
here because plans are positional. It captures full `DataSourceExec` file
groups, expressions, partitioning, and nested `stage_id`s — a `Display`-string
key would elide scan file groups and risk a false merge. If a subtree can't be
canonicalized (e.g. a custom op the codec can't encode alone), it is treated as
unique and **never merged** — a missed optimization, never a wrong result.
- Stages are processed in **ascending `stage_id`** order. Because the
planner assigns ids bottom-up, a dependency is finalized before any dependent
is keyed, giving a single-pass fixed point (inner duplicates collapse, then the
outer subtrees that now carry the collapsed inner id collapse in turn). The
**root stage is never dropped**.
2. **Wiring.** `StaticExecutionGraph::new` is split into `new` (unchanged
behavior, delegates with reuse off) and `new_with_reuse`, which runs the pass
between `plan_query_stages` and `ExecutionStageBuilder::build`. The scheduler
builds the canonicalizer closure from its configured codec; any serialization
failure degrades to "not reusable" and never aborts planning.
3. **Config** `ballista.optimizer.reuse_exchange_enabled` (default `true`,
parity with Spark's `spark.sql.exchange.reuse`).
# Are there any user-facing changes?
One new configuration key, `ballista.optimizer.reuse_exchange_enabled`
(default `true`), so existing deployments get exchange reuse automatically and
can turn it off if needed. No public API breakage.
## How are these changes tested?
- Unit tests on `reuse_shuffle_stages` with synthetic stage DAGs: an
identical pair collapses; differing input or partitioning is not merged; a
nested duplicate collapses in a single pass (with a stage-id-sensitive stub
proven to fail if the pre-key ref-rewrite is skipped); the root is never
dropped; an un-canonicalizable stage is never merged; and after
`ExecutionStageBuilder::build`, a shared stage fans out to **distinct**
consumer stages (`output_links` contains each consumer).
- A TPC-H detection test plans the real query SQL from
`benchmarks/queries/qN.sql` against the test context and runs the pass with the
real protobuf canonicalizer, asserting the exact stage-count delta. Exchange
reuse fires on **Q2 (19→18 stages), Q11 (12→10), and Q15 (6→5)**; it is
correctly a no-op on Q14/Q17/Q20/Q21, whose plans contain no
structurally-identical exchanges. A guard asserts at least one query exercises
reuse, so a silent no-op regression fails the test.
--
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]
