andygrove opened a new pull request, #2723:
URL: https://github.com/apache/datafusion-ballista/pull/2723
# Which issue does this PR close?
Closes #2681.
# Rationale for this change
A hash shuffle always assigns rows to partitions with DataFusion's hash.
That is consistent within a job, so joins and aggregates over shuffled inputs
are correct, but it can't match data that is already partitioned by another
function. Spark partitions by Murmur3, and bucketed tables and
storage-partitioned joins depend on it. When one side of a join is a
Spark-bucketed table read without a shuffle, the other side has to be shuffled
with Spark's hash for matching rows to meet. Comet Native Executors
(apache/datafusion-comet#7288) turn bucketing off today for this reason.
This implements the first option in the issue: the plan computes each row's
partition into a column and the writer routes by it. A pluggable hasher would
need every executor to have the hasher registered, and Ballista encodes the
expressions in its shuffle nodes with DataFusion's default codec, so a client's
own `PhysicalExpr` couldn't be sent to executors anyway.
# What changes are included in this PR?
- **Writer routing.** `SortShuffleWriterExec` takes a `ShuffleRouting`:
`Hash` (what it does today) or `PartitionId { column, partition_count }`. In
partition-id mode, output partition `i` holds exactly the rows whose id is `i`,
and the id column is dropped before writing, so the shuffle's schema is the
input's without it. Any integer type works. A null, negative or out-of-range id
fails the task rather than being wrapped into some partition. `try_new` and
`shuffle_output_partitioning()` are unchanged; `try_new_with_routing` and
`routing()` are new. Anything that rebuilds a writer carries the routing over,
including the executor's stage construction, which is covered by a test that
fails without that change.
- **Marker node.** `PartitionIdRepartitionExec` goes where a hash
`RepartitionExec` would. It is only a marker: the static distributed planner
turns it into a stage that ends in the new writer, and `execute` returns an
error. Its output reports `UnknownPartitioning`, since nothing downstream can
assume rows are placed the way DataFusion would place them.
- **Wire format.** New `PartitionIdRouting` and
`PartitionIdRepartitionExecNode` messages (the marker is tag 20 in
`BallistaPhysicalPlanNode`), and an optional `partition_id_routing` on
`SortShuffleWriterExecNode`, so plans containing either can be submitted over
gRPC with the physical extension codec. An executor that predates the field
sees a writer node with no `output_partitioning` and rejects it, so the plan
fails to launch instead of being hashed.
- **Adaptive planner.** It rejects the marker with a clear error. Coalescing
or re-planning an exchange's partitions would silently break the placement the
node exists to provide.
- **Per-column stats.** The writer's null counts are computed on the columns
it writes, so they line up with the schema a reader sees.
- **Rejected inputs.** An input whose only column is the id is rejected when
the node is built, since a shuffle with no other columns has nothing to move.
- **Docs.** A "Partition-id shuffles" section in the tuning guide.
## Interaction with #2682
On `main` today, submitted physical plans are always planned statically, so
this works as is. #2682 sends them through the adaptive planner by default,
which would reject the marker. Whichever of the two lands second needs to
handle it. The options are keeping plans that contain
`PartitionIdRepartitionExec` on the static planner in
`TaskManager::submit_plan`, or supporting the marker in the adaptive planner as
an exchange that is never coalesced, for which the pinning and coalesce guard
in #2682 look like the right base.
## A second, unrelated commit
`fix(scheduler): import ExecutionPlan in the execution graph tests` is one
line, and it is a compile error that is on `main` right now:
`losing_own_shuffle_output_restores_the_cache_reader` (#2633) uses `dyn
ExecutionPlan` but relied on an import that #2613 removed along with the helper
that was its only other user. `ballista-scheduler`'s lib tests don't build
without it (reproduced on a clean export of `apache/main` at 2342897f), and
this PR adds tests there. It is its own commit so it can be dropped if a fix
lands first.
# Tests
- Unit tests for placement (including across spills), the id-validation
errors, null counts, unnamed partitions staying empty, and construction rules;
the marker's schema, partitioning and display; codec round trips, including a
node with both routings set and a marker with no routing.
- Planner tests for the marker becoming a partition-id writer, stacking, and
surviving a whole-plan codec round trip; an adaptive-planner test for the
rejection.
- An executor test through `create_query_stage_exec`.
- `ballista/scheduler/tests/partition_id_shuffle.rs`: plans a marker query,
runs the writer stage through the executor engine on real rows, resolves the
reader, and checks that every partition holds exactly the rows the client
placed there. It needs no sockets.
Ran locally: `cargo test` for core, scheduler and executor (including the
TPC-H plan stability snapshots, which are unchanged), `cargo fmt`, `cargo
clippy --all-targets --all-features -- -D warnings` on those crates, rustdoc
with `-D warnings`, a workspace `cargo check --all-targets`, and the pinned
prettier on the docs page.
Not run locally: tests that bind a local socket (the cluster tests in
`ballista/client/tests` and a handful of unit tests), because my environment
blocks those binds. They are left to CI.
# Are there any user-facing changes?
New public API: `PartitionIdRepartitionExec`, `ShuffleRouting`,
`SortShuffleWriterExec::try_new_with_routing` and `routing()`, documented in
the tuning guide. Nothing existing changes, so there are no breaking changes:
hash shuffles behave and display exactly as before, and the proto changes are
additive.
--
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]
