avantgardnerio opened a new pull request, #2038:
URL: https://github.com/apache/datafusion-ballista/pull/2038
## Summary
Follow-up on the multi-partition-task migration on this branch. Each stage's
tasks now cover a *slice* of partitions instead of a single partition, so that
a 16-vcore executor can drive 16 partitions with one task rather than launching
16 single-partition tasks. This PR completes that migration by:
- Refactoring `SortShuffleWriter` to the coordinator+oneshot handoff pattern
(mirroring `ShuffleWriter` from the prior commit on this branch), so a task's
slice.len() input partitions run concurrently, each producing its own
consolidated `data.arrow` + `.index` file. Per-pipeline memory budget is
divided across concurrent pipelines to preserve the configured per-task budget.
- Threading a global `partition_slice` field through the executor: proto
(`TaskDefinition` / `TaskId`), executor plumbing, and both writers. Writers use
the slice to attach global identity to shuffle files. `ShuffleWriter(None)`
walks its child plan (`walk_child_partition_mapping`) and detects three cases:
`SortPreservingMergeExec` (collapsed, all outputs go to partition 0),
`RepartitionExec::Hash/RoundRobin` (K-space is inherently global), or
pass-through (`global = partition_slice[local]`). `ShuffleWriter(Hash)` and
`SortShuffleWriter` use hash bucket / input partition space directly.
- Fixing scan restriction:
- `MemorySourceConfig` was silently unrestricted at dispatch, causing
over-reads in multi-partition tasks (this was the concrete bug that made
`should_execute_submitted_physical_plan_across_shuffle_stages` lose rows).
- `restrict_shuffle_reader` now preserves the `Hash`/`RoundRobinBatch`
partitioning kind after shrink so `InterleaveExec` above the reader keeps its
"children share hash partitioning" invariant.
- Fixing the five pre-existing scheduler test failures that traced to the
append-only `task_infos` + slice-per-task migration:
- `SuccessfulStage::to_running` walks task_infos newest-to-oldest so
retried partitions aren't double-pushed to `pending`.
- `resubmit_successful_stages` / `reset_running_stages` values are
task_indices (from `map_partition_id`), not partition ids — variables renamed,
bounds check corrected against `task_infos.len()`, TODO for eventual
partition-id semantics.
- Failure error messages name the *partition(s)* that hit
`max_task_failures`, since retry attempts get fresh task_indices.
- `sum_leaf_scan_partitions` treats aligned-input joins (`HashJoinExec`,
`SortMergeJoinExec`) as leaves — walking past them double-counted the two
hash-aligned sides.
- Cluster binding tests count partitions covered rather than task count
(emergent under multi-partition binding).
- Renaming `task_slot` → `task_index` throughout the writer code for
consistency.
- Purging the earlier `substrate` shorthand from all code comments
(conflicts with an unrelated feature name in this repo).
## Test plan
- [x] `cargo build --workspace`
- [x] `cargo test --workspace` — all 806 tests pass across 17 binaries
- [x] Originally failing test:
`should_execute_submitted_physical_plan_across_shuffle_stages` passes
- [x] Sort shuffle tests: all 51 pass
- [ ] Manual TPC-H run at scale (to validate the new sort-shuffle
multi-partition path under load)
🤖 Generated with [Claude Code](https://claude.com/claude-code)
--
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]
