[andygrove]

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, routing(), output_schema() and with_work_dir() are new. The executor now stamps the decoded writer with with_work_dir, as its pipelined arm already does, instead of rebuilding it field by field, so it can't drop the routing.
• 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. create_shuffle_writer_with_config now takes an Option<ShuffleRouting> for hash repartitions and the marker alike; the AQE adapter converts its Partitioning with a new TryFrom<Partitioning> for ShuffleRouting. Its output reports UnknownPartitioning, since nothing downstream can assume rows are placed the way DataFusion would place them.
• Wire format. New PartitionIdRepartitionExecNode (tag 21 in BallistaPhysicalPlanNode, with the id column and partition count inline) and PartitionIdRouting messages, so plans containing either can be submitted over gRPC with the physical extension codec. SortShuffleWriterExecNode now holds its routing as a oneof: the existing output_partitioning = 4 or the new partition_id_routing = 10. A oneof member is encoded like a plain field, so plans from older schedulers still decode. An executor that predates the new field sees neither routing and rejects the node, 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 and its statistics describe only 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. The adapter now builds its writer from a ShuffleRouting, so the exchange would only need to carry one, and the pinning and coalesce guard in #2682 look like the right base.

Overlap with #2722

#2722 (plan a range RepartitionExec as a shuffle stage) also generalises how SortShuffleWriterExec routes rows: it keeps a Partitioning (Hash or Range), builds a RowRouter per input partition, and adds partitioning = 11 to SortShuffleWriterExecNode. This PR makes the routing an enum, ShuffleRouting, carried by a oneof. Neither depends on the other, but they conflict in ballista.proto, sort_shuffle/writer.rs, serde/mod.rs, planner.rs and shuffle_writer.rs, so whichever merges second has to reconcile them. Range would most likely become a third ShuffleRouting variant, with partitioning = 11 joining the routing oneof.

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 2342897), and this PR adds tests there. It is its own commit so it can be dropped if a fix lands first; #2721 carries the same fix as its first commit, so whichever of the two merges first makes the other's a no-op.

Tests

• Unit tests for placement (including across spills), the id-validation errors, null counts and statistics, unnamed partitions staying empty, and construction rules; the marker's schema, partitioning and display; codec round trips of the writer's routing, the marker, and a writer node 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.
• ballista/scheduler/tests/partition_id_shuffle.rs: plans a marker query, runs the writer stage through the executor engine's create_query_stage_exec on real rows, resolves the reader, and checks that every partition holds exactly the rows the client placed there and that the column stats cover only the written columns. 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 (with output_schema and TryFrom<Partitioning>), and SortShuffleWriterExec::{try_new_with_routing, routing, with_work_dir}, documented in the tuning guide. Hash shuffles behave and display exactly as before, and the wire encoding of existing plans is unchanged. The one change to an existing public item is in generated code: protobuf::SortShuffleWriterExecNode.output_partitioning is now a member of the routing oneof, which matters only to code that builds or reads that message directly.

[coderfender]

Thanks for implementing this, Andy. A neat way to support Spark-bucketed data without teaching Ballista Spark's hash, and the strict handling of bad ids plus the wire-compat story are great. Two small nits:

• The writer's statistics drop the id column's column_statistics, but total_byte_size still counts it.
• In pipelined mode a plan with the marker fails with "PipelinedShuffleWriterExec only supports Hash partitioning, got: UnknownPartitioning(n)". A message naming PartitionIdRepartitionExec would be clearer. Fine as a follow-up since pipelined mode is experimental.

[andygrove]

Both fixed in 52c427a. The writer's total_byte_size now drops the id column's bytes, and a pipelined job with the marker fails with an error that names PartitionIdRepartitionExec. Pipelined mode is still unsupported for the marker, now documented in the tuning guide.