Multi-partition tasks: partition_slice plumbing + SortShuffleWriter refactor - #2038
Merged
avantgardnerio merged 15 commits intoJul 19, 2026
Merged
Conversation
Member
|
Thanks for the contribution @avantgardnerio. It's going to take me a while to work through this and understand it all. One initial question, though. This PR changes the default execution model from one task per partition to multiple partitions per task. Rather than making this a default change in Ballista, maybe we could consider a config around this? Something like "max partitions per task"? I am concerned that forcing this change on the Spark user persona will cause confusion. If it is something they can opt into for better performance, that is fine. |
Member
|
TPC-H CI job is failing due to q15 returning incorrect data (returning too many rows). |
Contributor
Author
Yeah, absolutely, strictly speaking I think this is a superset of that, so "max of 1" should be easy - I'll add a config. |
andygrove
reviewed
Jul 14, 2026
andygrove
reviewed
Jul 14, 2026
andygrove
reviewed
Jul 15, 2026
Dandandan
reviewed
Jul 15, 2026
Contributor
|
Will review over weekend, until then i have one unreleased comment: great to see all of you here 😁 |
andygrove
reviewed
Jul 15, 2026
andygrove
reviewed
Jul 15, 2026
Member
|
It looks like per-partition metrics break with this PR. Here is what my AI tells me ... Right now OperatorMetricsSet on the wire is a flat list of metrics with no partition field. The per-partition granularity we show in the UI has always been derived from the fact that a task's index and its partition id were the same number. In execution_stage.rs we stamp every metric a task reports with Some(partition), and under this PR that value is now the task index rather than a partition id. The effect is that a task covering a slice of N partitions files all N partitions' metrics under one task-index key, so they get merged into a single bucket. The read side in handlers.rs stays self-consistent so nothing errors, but the per-partition detail is gone before it is ever stored. In practice that means skew inside a slice becomes invisible. If one partition of sixteen carries most of the rows, the stage view shows one large task with no way to see which partition drove it, and that is exactly the view Spark-model users lean on to debug skew. At one partition per task this all looks fine, which is what makes it easy to miss. The regression only shows up once a slice covers more than one partition. I think the clean fix is to make this additive rather than inferred: add an optional partition field to OperatorMetric, have the executor label each metric with its local partition index, and have the scheduler map local to global through partition_slice when it folds task metrics into the stage. That keeps per-partition metrics working in the new model instead of collapsing them to per-task. |
5 tasks
avantgardnerio
force-pushed
the
brent/multi-partition-tasks
branch
from
July 16, 2026 00:20
1b98039 to
7d3c602
Compare
Contributor
Author
Fixed! |
Contributor
Author
|
This branch is now at rough performance parity with main (which is the goal, performance improvements will come as we're able to implement features on top of this work): 
|
Closed
Contributor
|
Would it be possible to have a paragraph or two explaining the approach and compare it with spark, somewhere in documentation? |
Ballista's 1-task-per-partition model shreds DataFusion's within-plan
parallelism primitives — mpsc channels, shared Arcs, Mutex/OnceLock,
per-partition state that assumes a single plan-Arc sees the full sender
set. Any physical operator that coordinates across its input partitions
in-process deadlocks when K independent tasks each hold their own
instance and none of them observes all K partitions.
Motivating follow-up: parallel-window range-repartition (own PR),
which needs an in-process rendezvous across all inputs of a stage.
Substrate: one task per executor slice of the input, task drives all
cores through one plan-Arc.
- get_stage_partitions walks to leaves, sums their partition counts, and
returns ceil(sum / executor_cores) instead of the flat output count.
- New scheduler-side task_builder module rewrites each task's plan tree:
Scan file_groups restricted to the task's slice (empty for others,
partition count preserved) and ShuffleReader partitions filtered to
the assigned indices.
- task_partition_slice(slot) gives task slot N partitions
[N * cores .. (N+1) * cores). Rewriter tolerates over-run indices.
- prepare_task_definition and prepare_multi_task_definition rewrite +
encode fresh per task; the shared encoded_stage_plans cache goes away
since plans differ per task now.
- The old executor-side restrict_scan_to_partition is deleted (it only
worked for the 1-task-1-partition model).
ShuffleWriter changes to complete the picture:
- None branch drains all K child output partitions CONCURRENTLY (one
tokio::spawn per K), writing K files per task with file_id=task_slot.
Sequential drain deadlocks upstream scatter operators that push to
multiple senders — draining one to EOF fills the undrained channel.
- Hash branch iterates all input partitions per task instead of just
input_partition, using task_slot as file_id.
Pre/post-execution plan-dump logging in the executor's
DefaultQueryStageExec — invaluable for debugging distribution shape.
Known limitations left as TODOs in the code:
- Hardcoded TODO_EXECUTOR_CORES=4 in two places (task_manager,
execution_stage) — real cluster-shape plumbing via
ClusterState.registered_executor_metadata() lands in a follow-up.
- task_slots → vcores rename + one-task-per-exec assignment (both tasks
currently land on exec0). Both addressed in follow-up commits in
this PR.
- restrict_scan should collapse file_groups instead of leaving empties.
- Hash branch of ShuffleWriter should be replaced by DF's RepartitionExec
above a None-mode writer.
- Test coverage for task_builder's restrict logic (the executor-side
tests it replaces have been removed).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Rename task_slots → vcores across the workspace
Under the one-task-per-executor substrate (previous commit), the old name
`task_slots` conflates two things that just became orthogonal:
- concurrent-task capacity (now always 1 per executor)
- virtual-core count that a single task drives in parallel through one
plan-Arc (what --concurrent-tasks actually controlled)
Rename to `vcores`, following the YARN precedent
(`yarn.nodemanager.resource.cpu-vcores`) and Spark's `spark.executor.cores`.
The `v` prefix is self-documenting: readers grepping in a year don't have
to hunt for a doc comment to know these aren't nproc.
Mechanical rename, no behavior change. Sets up the follow-up chain (see
TODOs) that plumbs cluster-shape from executor advertisement through
session config into future parallel-window planning.
Highlights:
- `ExecutorSpecification.task_slots` → `vcores` (u32), with a doc
comment citing YARN/Spark.
- Proto: `ExecutorResource.task_slots` → `vcores`,
`AvailableTaskSlots` → `AvailableVcores` (doc: "One instance per
executor"), `ExecutorTaskSlots` → `ExecutorVcores`,
`PollWorkParams.num_free_slots` → `num_free_vcores`. Field numbers
preserved for wire compat.
- CLI: `--concurrent-tasks` → `--vcores` (help text explains it's
scheduler-facing accounting, not physical nproc).
- Cluster bookkeeping: `InMemoryClusterState.task_slots` →
`available_vcores`. Bind helpers rename locals to `budgets`
(Vec<&mut AvailableVcores>) + `current` cursor; loop comments
switched to "advance to a budget with free vcores".
- TUI: "Task Slots" column → "vCores"; `SortColumn::TaskSlots` →
`Vcores`.
- Follow-up TODOs tagged `TODO(c2)`, `TODO(c3)`, `TODO(c4)` so the
remaining cluster-shape plumbing chain is grep-findable.
Known follow-up not in this commit: the executor's memory pool formula
is `memory_pool_size / vcores` (was `memory_pool_size / concurrent_tasks`),
which under-utilizes memory in the one-task-per-exec model — the single
task now internally fair-shares one pool of size pool/N when it could
have the full pool. That's a semantic fix, not a rename; addressed in
a follow-up commit in this PR.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Rename proto task-context `partition_id` → `task_index`
Under the substrate model one task processes a partition slice, not a
single partition, so the `partition_id` field on task-carrier messages
was semantically wrong — it names the task slot within the stage. The
actual partition slice a task processes is baked into the plan bytes
(Scan file_groups / ShuffleReader partitions restricted at dispatch).
Renamed on the wire (field numbers preserved for compat):
- `TaskDefinition.partition_id` → `task_index`
- `TaskId.partition_id` → `task_index` (inside MultiTaskDefinition.task_ids)
- `TaskStatus.partition_id` → `task_index`
- `RunningTaskInfo.partition_id` → `task_index`
Rust-side renames:
- `ballista_core::serde::scheduler::TaskDefinition.partition_id`
→ `task_index` (with a doc comment on the new semantic).
- All proto-field access sites updated.
- Local variables named `partition_id` that carried task-slot semantics
renamed to `task_index` where they trace back to any of the above.
Not touched (deferred to the follow-up TaskKey commit):
- `PartitionId.partition_id` (shared type — genuine shuffle partition
semantic in most contexts; still used as a stand-in task locator in
scheduler-side TaskDescription. Every bridge site where a task_index
value flows into `PartitionId.partition_id` carries a `TODO(c4a.2)`
marker so we can find them when TaskKey lands.)
- Rust `execution_graph::RunningTaskInfo.partition_id` (same story;
also TODO-tagged).
- `execute_query_stage` / `cancel_task` receivers whose `partition_id`
parameter is now task_index-shaped (also TODO-tagged).
Mechanical rename with no behavioral touch.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Introduce TaskKey; retire task-context PartitionId misuse
Under substrate one task processes a partition slice, so where
`PartitionId { job_id, stage_id, partition_id }` was being used as a
task locator (in `TaskDescription`, executor `abort_handles`,
`execute_query_stage` args, `RunningTaskInfo`, etc.), the `partition_id`
field semantically meant "task slot within stage" — a misuse that made
every reader wonder whether they were looking at a partition or a task.
New type in `ballista_core::serde::scheduler`:
pub struct TaskKey {
pub job_id: JobId,
pub stage_id: usize,
pub task_index: usize,
}
Sibling to `PartitionId` — the former identifies an operator output
partition (shuffle-location, `PartitionLocation`, `ShuffleReaderExec`
input), the latter identifies a scheduled task attempt. PartitionId is
untouched in genuine-partition contexts.
Sites migrated (all TODO(c4a.2) markers now gone):
- `TaskDescription.partition: PartitionId` → `key: TaskKey`. Access
sites `task.partition.{job_id, stage_id, partition_id}` become
`task.key.{job_id, stage_id, task_index}`.
- `RunningTaskInfo.partition_id: usize` → `task_index: usize`.
- Executor `abort_handles: DashMap<(usize, PartitionId), _>` →
`DashMap<(usize, TaskKey), _>`.
- `Executor::execute_query_stage(partition: PartitionId, ...)` →
`execute_query_stage(key: TaskKey, ...)`. Internally passes
`key.task_index` to `QueryStageExecutor::execute_query_stage` —
which under substrate calls the plan with the task's index as its
input partition (plan bytes were pre-restricted at dispatch).
- `Executor::cancel_task(partition_id: usize)` → `task_index: usize`.
- `ExecutionEngine::create_query_stage_exec(partition_id)` →
`task_index`.
- `ExecutorMetricsCollector::record_stage(partition)` → `task_index`.
- `as_task_status(partition_id: PartitionId)` → `key: TaskKey`.
Naming leftover: the shared `PartitionId` type in
`ballista_core::serde::scheduler` (line 63) and its counterpart proto
message (ballista.proto:370) stay unchanged — they identify shuffle
partition locations in `PartitionLocation`, `ShuffleReaderExec`, and
the client, where `partition_id` genuinely means partition index.
Pre-existing test failures (18 in `ballista-scheduler`, from the
substrate commit's num_tasks reduction) are unchanged.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
PendingPartitions cursor + append-only task_infos
Under the substrate commit, one task drives all of an executor's vcores
through one plan-Arc — but the number of tasks per stage was still
pre-computed at resolve time via `ceil(leaf_sum / TODO_EXECUTOR_CORES)`
with a hardcoded 4. That tied the scheduling model to a fictional
"executor cores" const and pre-allocated `Vec<Option<TaskInfo>>` slots
whether or not they'd ever be filled.
This commit replaces that with:
- `PendingPartitions`: a bind-time cursor over the stage's real plan
input partitions. `next_slice(exec.vcores)` drains a chunk sized to
whichever executor picks it up. Retries push failed partitions to
the front of the queue so they get re-attempted before any fresh
partition.
- `TaskInfo` gains `partition_slice: Vec<usize>` so executor-loss /
retry paths know which partitions to push back to `pending`, and
per-partition failure bookkeeping can find the right entries in
`task_failure_numbers`.
- `RunningStage.task_infos` becomes append-only `Vec<TaskInfo>` (was
`Vec<Option<TaskInfo>>`). Every bind appends; every retry appends.
task_index is immutable per task, monotonic across the stage's
lifetime.
- `task_failure_numbers` is now indexed by PARTITION id (real plan
input), not by task_index — a task's failure count derives from the
max over its slice.
- `RunningStage.partitions` reflects the real plan input partition
count (frozen at resolve). Number of tasks is emergent from
`task_infos.len()` at each observation point.
- `bind_task_bias` / `bind_task_round_robin` route through a shared
`bind_one` helper that draws the slice from `stage.pending`,
appends `TaskInfo`, builds `TaskDescription`, and consumes the
executor's vcore budget entirely (leftover packing is a follow-up,
deferred).
- `pop_next_task` (static + AQE) now uses `pending.next_slice(1)` —
the size-to-vcores bind path lives only in `cluster/mod.rs`.
- `SuccessfulStage::to_running` rebuilds `pending` from the
partition_slices of Failed(ResultLost) entries; successful entries
stay in place.
- `update_task_info` rejects stale updates either by task_id or by
the task already being in terminal Failed(TaskKilled/ResultLost)
state (i.e., `reset_tasks` already moved partitions back).
Deletes `get_stage_partitions` / `task_partition_slice` /
`TODO_EXECUTOR_CORES`. Replaced by `stage_input_partitions(plan)` —
no ceil-divide, real leaf/output partition count.
Handles heterogeneous vcores, dynamic membership, multi-tenant
contention: one 16-vcore exec covers 16 partitions in a single task;
four 4-vcore execs cover the same 16 partitions in four tasks.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
ShuffleWriter: coordinator+oneshot handoff to respect DF contract
Under substrate the writer's ExecutionPlan::execute(N, ctx) ignored its
partition argument and always did all K writes in one call — a violation
of DataFusion's contract that `execute(N)` returns the stream for output
partition N. In production it was masked because no operator consumes the
writer's stream as partitioned data, but it broke every unit test that
called execute() directly and left the plan lying about its own shape.
Refactor to the range-repartition pattern used upstream in DataFusion's
physical-plan/src/range_repartition.rs:
- `ShuffleWriterExec` gains `state: Arc<Mutex<WriterState>>` holding
K oneshot receivers, initialized lazily on first `execute(N)`.
- First `execute(N)` call:
* takes the lock
* creates K oneshot channels, stashes receivers in state.handoffs
* spawns ONE coordinator task that owns the K senders and drives
the shared write work via `execute_shuffle_write`
* takes its own receiver, returns a stream that awaits it
- Subsequent `execute(N)` calls just take handoffs[N] and return the
awaiting stream. K-drain concurrency is preserved because the
coordinator's None-branch already spawns K `tokio::spawn` writers
internally; the outer K parallel `execute(N)` calls just parallelize
the metadata-batch reads.
- Coordinator dispatches each `ShuffleWritePartition` summary from
`execute_shuffle_write` to the sender matching `summary.partition_id`;
partitions whose writers stayed empty (Hash-mode routing) get a
zero-content sentinel so no execute(N) stalls on an unfilled oneshot.
The launcher filters those out before returning to the scheduler.
Task-launch changes:
- `DefaultQueryStageExec::execute_query_stage(task_index)` now spawns
K parallel `plan.execute(N, ctx)` calls, collects each single-row
metadata batch, and reassembles the `Vec<ShuffleWritePartition>`
the scheduler expects. All K spawn concurrently so the writer's
coordinator sees every oneshot receiver taken before it starts.
- Sort-variant keeps its custom `execute_shuffle_write(task_index)`
entry point (single-call multi-file semantics unchanged pending its
own port to the DF-contract per-partition model).
- `create_query_stage_exec` now uses its `task_index` parameter to
stamp `ShuffleWriterExec::with_task_slot(task_index)` — the writer
needed a home for its file_id namespace now that per-task rewrite
isn't threading it through execute_shuffle_write's argument list.
Wire format changes:
- The metadata batch schema gains a `file_id: UInt64` (nullable)
column between `path` and `stats`, so the reconstructed
`ShuffleWritePartition` values carry file_id through the
single-partition stream boundary. Downstream consumers of the
Vec<ShuffleWritePartition> are unchanged.
Test updates:
- `test`, `test_partitioned`, `test_no_repart_write_failure_propagates`,
`test_hash_repart_write_failure_propagates`: drive all K partitions
via a new `drive_all_partitions` helper (mirrors what
execute_query_stage does in production) instead of calling
execute(0) once and hoping for the best.
- `test_read_local_shuffle`: bump expected batch count from 2 to 4
(writer reads its full input slice — 2 partitions × 2 batches).
Not touched here:
- The internal Hash branch inside `execute_shuffle_write` stays, but
the follow-up commit deletes it and moves hash-repartitioning to
DF's `RepartitionExec::Hash` above a None-mode writer.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Multi-partition tasks: SortShuffleWriter refactor + partition_slice plumbing
Fixes the physical_plan_submission cross-stage test, which lost rows because
SortShuffleWriter was stuck on the old 1-task-1-partition contract and
MemorySourceConfig was silently unrestricted at scheduler-side dispatch.
- SortShuffleWriter mirrors ShuffleWriter's coordinator+oneshot handoff. Each
task drains all its slice's input partitions in parallel pipelines, one
file per input partition, memory budget divided across them. Both writers
now go through drive_shuffle_writer_stage; no more Sort special-case in
execution_engine.
- Unified metadata batch schema (result_schema, summaries_to_batch) shared by
both writers. Coordinator sends a Vec<summary> per handoff slot so a single
execute(N) stream can carry multiple summaries (sort writer's N input files
contributing to hash bucket N).
- Adds partition_slice to TaskDefinition / TaskId proto: the concrete global
partition ids each task covers. Writers attach global identity to shuffle
paths / reported ShuffleWritePartition.partition_id:
- ShuffleWriter(None): walk_child_partition_mapping detects SPM (collapsed
to partition 0), RepartitionExec(Hash/RoundRobin) (K-space, local=global),
or pass-through (global = partition_slice[local]).
- ShuffleWriter(Hash) / SortShuffleWriter: K-space is inherently global.
Sort uses partition_slice[local_input] as file_id.
- restrict_scan handles MemorySourceConfig (was silently skipped → over-read
in multi-partition tasks). Shrink model preserved.
- restrict_shuffle_reader preserves Hash/RoundRobin partitioning kind after
shrink; InterleaveExec above the reader requires matching hash partitioning
and was breaking with UnknownPartitioning downgrade.
- Renames leftover task_slot terminology to task_index throughout the writer
(with_task_slot -> with_task_index).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Fix scheduler tests broken by multi-partition-task model
Fixes 5 pre-existing test failures on this branch that all trace back to
the append-only task_infos + slice-per-task migration.
- `sum_leaf_scan_partitions` now treats aligned-input joins (`HashJoinExec`,
`SortMergeJoinExec`) as leaves. Walking past them double-counted the two
hash-aligned sides (2K leaf partitions for K aligned join-output units).
- `SuccessfulStage::to_running` walks task_infos newest-to-oldest and only
reschedules partitions not yet covered by a later Successful attempt.
The naive iteration pushed the same partition to `pending` multiple times
after retry cycles (Failed original + Failed retry both contributing
their slice).
- `resubmit_successful_stages` / `reset_running_stages` values from
`remove_input_partitions` are `map_partition_id`s, which under this
model equals the source task's `task_index`. Renamed the local variables,
bounds-check against `task_infos.len()` (not `.partitions`, which was
the old input-count and caused "Invalid partition ID N" errors when
task_index legitimately exceeded input count), and fed the values into
`reset_task_info` / `task_infos[…]` directly. Fixed in both
`execution_graph.rs` and `aqe/mod.rs`. TODO for the future switch to
partition-id semantics.
- `execution_graph::update_task_status` renamed the misleading
`partition_id` local to `task_index` throughout (it always held
`task_status.task_index`). Error message now names the *partitions* that
hit `max_task_failures` rather than the retry-attempt task_index, which
isn't user-meaningful under append-only retries.
- Cluster binding tests count partitions covered rather than task count,
which is emergent under multi-partition binding.
- Test `test_max_task_failed_count` asserts `partition_slice` equality
across retries rather than `task_index` (retries get fresh indices).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Purge 'substrate' from comments/docstrings across branch
'substrate' is also a feature/protocol name in this repo, so this branch's
usage of it as shorthand for the multi-partition-task migration is worse
than meaningless to reviewers. Replace with concrete phrasing
("multi-partition-task", "one task processes a slice of partitions", etc.)
in every code comment and docstring introduced on the branch.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
bind_one: keep single-output stages in a single task
For plans that collapse fan-in to one output (e.g. a passthrough
ShuffleWriterExec over SortPreservingMergeExec), splitting the input
partitions across multiple tasks left each task's local SPM merging
only its slice, producing per-task locally-sorted files that
concatenated in nondeterministic order downstream. When the stage's
output collapses to one partition, take the entire pending queue in
one bind regardless of the executor's vcore budget so a single task
sees every input; the per-task plan restriction on the full slice is
a no-op.
Reproduces on CI (few-vcore runners): all six cases of
test_sort_shuffle_group_by_single_column plus
test_shuffle_completes_under_tiny_governor_budget.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Fix lint and clippy
- cargo fmt
- drop redundant `.into_iter()` in two writer coordinators (clippy
useless_conversion)
- silence too_many_arguments on ExecutionEngine::create_query_stage_exec
(multi-arg trait; refactor is out of scope for this PR)
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Drop unused feature-gated imports in task_manager
The disable-stage-plan-cache path in task_manager was removed on this
branch, but its cfg-gated imports were left behind. Clippy with
--all-features (as CI runs it) flags them as unused.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
tpch-sf10: rename --concurrent-tasks to --vcores in workflow
The executor CLI flag was renamed in commit dfe1f60 (task_slots →
vcores) but the tpch-sf10 workflow wasn't updated, so the executor
now rejects the flag with "unexpected argument '--concurrent-tasks'"
and the workflow fails before running any TPC-H query.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
shuffle_writer: panic on RepartitionExec::UnknownPartitioning
The walker previously returned PassThrough(partition_slice) for this
arm, but a RepartitionExec still exchanges rows and freshly numbers its
K outputs regardless of scheme — so slice[local] is a meaningless
mapping. DataFusion's BatchPartitioner rejects UnknownPartitioning
outright, so this arm is unreachable in practice; panicking makes any
future invariant break loud instead of silently mismapping partition
ids.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
chore: gitignore .local/ for dev-only cluster scripts
Lets each developer keep their own iteration scripts (`.local/up.sh`,
`.local/down.sh`, `.local/tpch.sh`, `.local/tail.sh`) in the repo root
without committing them. Mirrors the h2o-window-benchmarks branch layout
adapted for `--vcores` naming.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
sort_shuffle: rename `plan` field to `child`
The writer's `plan` field is redundant with `self.children()[0]` (the
`ExecutionPlan` trait already exposes children), and the name overloads
"plan" — usually the whole tree, here just the writer's sole child. Rename
clarifies that the writer *is* the stage's root plan and this is its
subtree.
Pure rename; no behavior change. Sets up for a follow-up disentangle of
`partition_slice` into `global_input_partition_ids` /
`global_output_partition_ids` where reasoning about "child.output_partitioning
vs writer.output_partitioning" needs to be unambiguous.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
bind_one: detect inner collapses, not just top-level output=1
Previously bind_one only recognized the "single output partition" case
via `plan.output_partitioning().partition_count() == 1` — that catches
`ShuffleWriter(None) + SortPreservingMergeExec` but misses
`SortShuffleWriter(Hash(K))` wrapping an inner collapse
(`AggregateExec(Final, gby=[]) + CoalescePartitionsExec + ...`). The
writer's own output is Hash(K)=K, so the old check saw K != 1 and split
input across tasks; each task's local collapse then only saw its slice,
yielding 16 partial-maxes instead of one global max downstream.
Replace with a walk: from the writer's child down, return true on any
operator whose `output_partitioning().partition_count() == 1`. Walks past
*any* single-child operator so it's forward-compatible with future partition
operators (e.g. `UnorderedRangeRepartitionExec`) that sit between the writer
and the work chain. Explicit stop at `ShuffleReaderExec` — the only
stage-boundary leaf that appears in a resolved stage plan — so the walk
doesn't cross into upstream stages. (General rule: stop at any stage
boundary; if new stage-boundary operators appear, add them here.)
Reproduces on TPC-H SF10 Q15 with `--vcores 4` and
`prefer_hash_join=false`: the max-revenue side (SortShuffleWriter(Hash(16))
over AggregateExec(Final, gby=[]) over CoalescePartitionsExec) currently
returns 16 rows instead of 1 because each of 16 tasks computes its own
partial max.
The detection is correct but the SortShuffleWriter guard
`num_input_partitions == partition_slice.len()` will now fire (slice=16,
child.output=1). Fixed by a follow-up commit that disentangles
`partition_slice` into separate input/output partition id lists.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Rename partition_slice by call-site meaning: input vs output ids
`partition_slice` served two distinct roles depending on where it was
used. Rename each site to reflect its meaning:
Scheduler-side (drives restriction + retry bookkeeping) →
`global_input_partition_ids`:
- `TaskDescription`, `TaskInfo`
- `bind_one`, AQE task binding, execution graph tests
- `restrict_plan_to_partitions` param name
Over the wire + writer + executor plumbing (identifies output partitions
this task contributes to) → `global_output_partition_ids`:
- `TaskDefinition` / `TaskId` proto fields
- Rust `TaskDefinition` mirror, from_proto decoding
- `ShuffleWriterExec` / `SortShuffleWriterExec` fields and
`with_global_output_partition_ids` builder
- `execution_loop`, `executor_server`, `execution_engine`
Bridge in `task_manager::prepare_task_definition`: reads the scheduler-side
input ids and populates the wire's output ids. Comment notes the current
identity — output ids _are_ input ids in passthrough mode; hash/single
modes ignore this field. Follow-up commits will replace the RHS once
`SortShuffleWriter`'s file naming stops depending on input-partition
identity (step: "1 file per task").
Pure rename by site; no behavior change. Q15 still fails until the
`SortShuffleWriter` guard is updated in a subsequent commit.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
WIP: SortShuffleWriter 1-file-per-task refactor
Unit tests pass, TPC-H Q15 passes locally, most other TPC-H queries pass;
Q11 still fails (72 vs 1048 rows) but that failure predates this refactor —
confirmed by testing Q11 on the previous commit. Committing this state so
we can bisect the Q11 regression separately without losing the refactor.
Changes:
- `SortShuffleWriterExec` writes one file per task (`.../{stage}/{task_index}/data.arrow`)
with K sections instead of one file per input pipeline. Pipelines drain
concurrently into per-pipeline (BufferedBatches + SpillManager) state;
`finalize_task_output` walks buckets 0..K and concatenates each pipeline's
spilled + buffered contributions into the merged file. Cross-task uniqueness
is provided by `task_index` alone.
- Guard rewritten: `global_output_partition_ids.len() == num_output_partitions`
(one global output id per hash bucket), matching the K-space intrinsic
to SortShuffle. Default output ids initialize to `[0..K-1]`.
- All metric writes moved to task-level in `finalize_task_output`; drain
accumulates counters (input_rows, spill_events, spill_bytes,
repart/spill nanos) without touching MetricsSet, so a task_index/local
partition-index collision doesn't create duplicate counters.
- `SpillManager::new` now takes a caller-provided `spill_dir: PathBuf`;
callers compose the per-pipeline path (`.../{task}/{local}/spill/`) so
concurrent pipelines within a task don't share writers and tasks across
the stage don't collide.
- `compute_global_output_partition_ids` in `execution_plans/shuffle_writer.rs`
becomes the public helper the scheduler calls before proto-encoding.
Hash writers return `[0..K-1]`; `ShuffleWriter::None` walks the child
via `walk_child_partition_mapping` (Collapsed → [0], HashSpace →
[0..K-1], PassThrough → input ids).
- `task_manager::prepare_task_definition` and `prepare_multi_task_definition`
call the helper instead of relaying input ids as output ids.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
task_builder: scope reader restriction under collapse operators
TPC-H Q11 returned 72 rows instead of 1048 (verified 1048 on merge-base).
`restrict_shuffle_reader` blindly restricted every `ShuffleReaderExec` to
the task's partition slice — including readers whose output feeds a
collapse (`CoalescePartitionsExec`, `SortPreservingMergeExec`, or a
`SinglePartition`-requiring join build side). In Q11 stage 15 that
reduced the subquery-threshold reader from 16 upstream partitions to 1
per task, so each of 16 tasks computed `Final(gby=[])` over 1/16 of the
partial sums, producing a per-task wrong threshold; the HAVING NLJ then
rejected most rows (probe_hit_rate=0%).
Turn `restrict_plan_to_partitions` into a stateful `TreeNodeRewriter`
with a `Scope::Collect` stack frame pushed at Coalesce/SPM (idempotent
w.r.t. an existing Collect on top) and at each `SinglePartition`-
requiring input of `HashJoinExec` / `NestedLoopJoinExec`. Readers/scans
under a `Collect` scope skip restriction so the collapse sees every
upstream partition. Broadcast-reader fast-path preserved.
Pattern ported from dataprime-query-engine's `TaskBuilder`
(scheduler/src/state/execution_graph/v2/builder.rs). We drop DQE's
`ScopedPartitions` variant (Union handling) for now — no TPC-H query
exercises it, and there's a follow-up planned to either simplify the
one-variant enum to a counter or add the Union case properly.
Also: apply `cargo fmt` across the branch — pre-existing formatting
drift picked up alongside the fix.
Verified: TPC-H Q1-Q22 at SF1 all PASS with --verify against DataFusion;
cargo test --workspace --exclude ballista-python: 845 passed, 0 failed;
cargo clippy --workspace --all-targets -- -D warnings: clean.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
task_builder: recursive per-branch scoping, drop shared stack
Follow-up to the previous commit. The `TreeNodeRewriter` design used a
`Vec<Scope>` shared across the whole walk, with a "push N Collects at
each collapse operator, pop-per-leaf" protocol. That works for TPC-H
Q11's plan shape but bakes in fragile assumptions: it silently requires
the SinglePartition-input subtree of a join to contain exactly one leaf,
descended before its siblings — nothing in the DataFusion API guarantees
either, and violations leak Collect tokens sideways into subtrees that
should have been partition-restricted.
Replace with a plain recursive fn `restrict(plan, partitions,
under_collect: bool)`. The call stack is the tree walk's natural stack;
`under_collect` is passed by value, so every recursive call has its own
copy and sibling subtrees can't share state. Per-child scoping via
`child_scopes(&plan, under_collect) -> Vec<bool>`:
- Sticky: once a parent is `under_collect`, every descendant inherits it.
- `CoalescePartitionsExec` / `SortPreservingMergeExec`: all children get
`true`.
- `HashJoinExec` / `NestedLoopJoinExec`: one bool per
`required_input_distribution()` — build gets `true` when the required
distribution is `SinglePartition`, probe inherits `false` and stays
partition-aligned.
- Otherwise: children inherit the parent's `under_collect` unchanged.
This drops the one-variant `Scope` enum in favor of an honest `bool`,
and correctly handles subtrees with multiple leaves under a single
collapse — a `Coalesce → HashJoin(reader_a, reader_b)` build side now
Collect-scopes both readers, not just whichever the walker happened to
visit first.
Adds a `TODO(union)` on `child_scopes` and an `#[ignore]`d test
`restrict_union_splits_partitions_per_child` documenting the assertion
we owe once a query with a partition-aligned `UnionExec` demands the
DQE `ScopedPartitions(start, end)` variant.
Verified: TPC-H Q1-Q22 at SF1 all PASS with --verify; cargo test
--workspace --exclude ballista-python: 845 passed, 0 failed, 7 ignored
(+1 for the new stub); cargo clippy --workspace --all-targets --
-D warnings: clean; cargo fmt --check: clean.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
sort_shuffle: satisfy CI clippy + rustdoc gates
- factor buffered.take() tuple into BufferedTake to appease type_complexity
- drop intra-doc link to private finalize_task_output in public docstring
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
style: cargo fmt
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
perf(scheduler): symmetric vcore accounting in bind/refund
`bind_one` used to zero out the executor's entire vcore budget for every
task, regardless of how many partitions the task actually got. Refund
credited back one vcore per completing task. Under the multi-partition
model that combination monotonically drained executor budgets to 1
vcore-in-flight after the cold start, so warm queries ran effectively
single-threaded per executor.
Fix both sides symmetrically:
- `bind_one` now consumes `min(slice_len, budget.vcores)` vcores. Leftover
budget stays available for another bind on the same executor in the
same scheduling round. Non-collapse paths get real multi-task
concurrency again; collapse paths still monopolize the executor
(their intra-task parallelism is bounded by the collapsing operator).
- The exact consumption is stashed on `TaskInfo.vcores_consumed`.
- Refund site sums `task_vcores(job, stage, task_index)` across the
status batch instead of counting statuses, so bind and refund always
match.
TPC-H SF10 sweep, 2× 4-vcore executors: 60.6 s → 32.1 s total (2.41×
slower than main → 1.28×). Q1/Q6/Q19 fully recover; residual gap is
concentrated on collapse-heavy queries (bug 2, tracked separately).
Also tightens naming in touched code: `slice` → `input_partition_ids`,
closure `|p|` → `|pid|`, opaque `s`/`ss`/`n` in the new refund helper
→ `status`/`job_statuses`/`vcores`.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
perf(scheduler): reserve 1 vcore for collapse tasks
A stage whose plan root has a single output partition is driven by exactly
one tokio task/thread under DataFusion's volcano/pull model, no matter how
many input partitions the collapse operator sits above. `bind_one` was
still charging the executor's full vcore budget for such tasks, blocking
unrelated stages' tasks from running in parallel on the same executor
even though 3 of its 4 vcores were sitting idle.
Charge exactly 1 vcore for collapse-shaped tasks; leave the packing
(entire pending queue into one task) unchanged since that's still
required for correctness — a split collapse yields per-slice partials
downstream can't merge.
Impact on the SF10 single-query sweep is small (32.1 s → 31.2 s),
because in a single query the collapse is the terminal stage and no
other work is waiting on the same executor. The win is on multi-tenant
clusters where a different query's tasks can now be scheduled onto the
executor's remaining vcores during a collapse.
Also updates the `bind_one` doc comment to explain the accounting under
DF's pull model, so future readers don't have to reconstruct it from
the trace.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
feat(executor): size per-task memory pool by vcores_consumed
Multi-partition tasks under this branch's model claim N vcores of the
executor's budget at bind time (see scheduler `bind_one`), but the
executor's memory-pool policy was still handing every task the same
per-vcore share — a 4-vcore task got the same 4 GB pool as a 1-vcore
task. That's `pool = total / total_vcores` when it should be
`pool = total * vcores_consumed / total_vcores`.
Plumb `vcores_consumed` from the scheduler to the executor:
- Add `vcores_consumed` to `TaskId` and `TaskDefinition` proto messages,
and to the corresponding native `TaskDescription` / `TaskDefinition`
structs.
- Populate it in `bind_one` (already computed for `TaskInfo`) and
forward it through `prepare_multi_task_definition` and
`prepare_task_definition`.
- Decode it on the executor side in `get_task_definition` /
`get_task_definition_vec`.
- Extend the `MemoryPoolPolicy` closure signature to take
`vcores_consumed`, and `SessionRuntimeCache::produce_runtime` /
`Executor::produce_runtime_for_session` to pass it through.
- In the production policy, size the `FairSpillPool` as
`per_vcore * vcores_consumed`.
- Always attach the session cache — the previous "override producer
⇒ no cache" branch dropped vcores at the fallback and gave every
task a 1-vcore pool. `pool_policy` composes fine with an override
base producer (e.g. S3-aware), so there's no reason to skip caching.
Verified with a single-query diag on Q3: a stage-5 task with
`vcores_consumed=4` now reports `pool_size=16 GB` (4×), a collapse
task with `vcores_consumed=1` reports `4 GB`, and the executor's
scheduler/executor `has_cache=true`.
Note: this does not fully eliminate spilling on Q3 — the
`SortShuffleWriterExec` writer has its own 256 MB per-task buffer
budget independent of the runtime memory pool, which still triggers
spills at ~200 MB. That's a separate issue tracked outside this
commit.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
perf(sort_shuffle): switch to per-input-partition writes + coordinator
The branch's original `SortShuffleWriterExec` fused all K hash-bucket
outputs of a task into a single sorted file and picked spill vs
in-memory paths based on a per-task budget. Benchmarking on TPC-H SF10
showed the accumulate-and-sort path is ~1.5–2× more thread-time per
partition than main's stream-per-file design (Q3 stage 5: 6.2 s vs
15.9 s of task-time for the same work) — and while dropping the spill
threshold recovered raw parallelism, the write path itself was slower.
Reset `sort_shuffle/writer.rs` and `sort_shuffle/spill.rs` to main's
design (one file per input partition, no K-drain accumulation) and
wrap it in the branch's coordinator handoff so `execute(0..K)` gets a
metadata stream per output partition while the actual write does one
file per input.
Key adaptations for the multi-partition-task model:
- `SortShuffleWriterExec` now carries `task_index` and
`global_output_partition_ids` (with `with_task_index` and
`with_global_output_partition_ids` builders mirroring
`ShuffleWriterExec`), stamped by the executor at
`create_query_stage_exec` time.
- `execute_shuffle_write(input_partition, ctx)` bit-packs
`(task_index, input_partition)` into the `file_id` so the on-disk
path `.../{file_id}/data.arrow` is globally unique — different tasks
in the same stage no longer clobber each other (they otherwise
would, because `compute_global_output_partition_ids` returns the
K-space `0..K` for `SortShuffleWriterExec`, which is intrinsic to
hash shuffle and not per-file-unique).
- New `WriterState` + `run_coordinator` mirror `ShuffleWriterExec`'s
handoff pattern: the first `execute(N)` spawns a coordinator that
drives M concurrent per-input-partition writes (via
`tokio::spawn`), re-buckets the emitted `ShuffleWritePartition`s by
output partition (K), and sends each bucket's summaries through a
matching oneshot. Every `execute(N)` returns a metadata batch with
M rows — one per input file, using the shared `result_schema` /
`summaries_to_batch` helpers from `shuffle_writer.rs`.
TPC-H SF10 sweep (2×4-vcore executors): total 60.6 s → 26.0 s (2.41×
slower than main → 1.04×). No individual query more than 1.13× slower;
several are slightly faster. Zero spilling observed on the sweep.
Also fixes `.local/bench.sh` to tolerate query failures (grep no-match
under `set -o pipefail` was aborting the sweep on the first failing
query).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
� Conflicts:
� ballista-cli/src/main.rs
� ballista/core/src/serde/scheduler/mod.rs
� ballista/executor/src/config.rs
� ballista/executor/src/execution_engine.rs
� ballista/executor/src/executor_process.rs
� ballista/scheduler/src/cluster/mod.rs
…titioning The leaf-sum walker undercounted after upstream apache#2062 (EmptyExec serde-safe rewrite): a plan like AggregateExec → RepartitionExec(RR(N)) → EmptyExec(1) has 1 literal leaf but N tokio-driven output pipelines, so sum-of-leaves returned 1 where the correct partition count is N. Read partition_count directly from the shuffle writer's immediate child's output_partitioning instead. That's the number of independent output- partition polling contexts the writer needs to drive — one per vcore per the "one tokio worker per vcore" invariant. Anything an internal RepartitionExec spawns is per-plan-instance machinery shared across those polls and doesn't become a separate scheduling unit. Also fixes the leaf-sum's HashJoin overcount (build-then-probe is sequential, so peak = max(build, probe) = the operator's own output_partitioning.partition_count, not the sum of both sides). Restores the 7 execution_graph tests that broke on the rebase onto main: test_fetch_failures_in_different_stages, test_normal_fetch_failure, test_long_delayed_failed_task_after_executor_lost, test_many_consecutive_stage_fetch_failures, test_task_update_after_reset_stage, test_reset_resolved_stage_executor_lost, test_reset_completed_stage_executor_lost. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…'s local index Under a `UnionExec`, parent partition `p` maps to exactly one child's local partition — `p` minus the sum of preceding children's counts. The walker previously passed the parent's slice unchanged to every child, so most indices landed out of range and every scan under the union was emptied. Split `partitions` into disjoint per-child sub-slices and recurse; a child that gets `[]` becomes a 0-partition subplan and UnionExec's partition-index math routes `execute(i)` past it correctly. Only applies when `under_collect == false`; under collect, every descendant already reads the full upstream and no split is needed. Ports the fix upstream landed in ballista/executor/src/execution_engine.rs (PR apache#2070) into the scheduler-side plan rewriter, matching the branch's move of scan restriction from executor to scheduler. Adds three tests (mirroring the upstream three) plus a small file-scan test helper. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
avantgardnerio
force-pushed
the
brent/multi-partition-tasks
branch
from
July 17, 2026 16:11
b261c57 to
308f149
Compare
andygrove
reviewed
Jul 17, 2026
andygrove
reviewed
Jul 17, 2026
Feedback from apache#2038 review: task_id (from ExecutionGraph.task_id_gen) and task_index (task_infos slot) always moved in lockstep at bind time, so carrying both was pure clutter. Redefine task_id as the append-order slot in RunningStage.task_infos, drop task_index and task_id_gen everywhere. Behavior-preserving: all 21 TPC-H queries pass and perf is within noise vs the same branch without this commit (37.06s vs 38.00s pristine). Follow-up on the same review comment: stage_attempt_num -> partition_attempt_nums[]. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Feedback from apache#2038 (andygrove): OperatorMetricsSet on the wire was a flat list of MetricValues with no partition tag. Per-partition detail in the UI came from execution_stage stamping every metric with Some(partition), where `partition` used to equal the task's single partition_id. Under this branch that arg is the task_id (append slot), so a task covering N partitions files all N partitions' metrics under one bucket — skew inside a slice becomes invisible. handlers.rs stayed self-consistent (it also keyed by the enumerated task index), so nothing errored; the detail was just gone before storage. Fix, additive rather than inferred: - Proto: add `optional uint32 partition = 16` to OperatorMetric, outside the oneof, carrying the local partition index within the reporting task's plan (0..slice_len). - to_proto: refactor `&MetricValue → OperatorMetric` to `&MetricValue → operator_metric::Metric` and thread the partition at the MetricsSet ↔ OperatorMetricsSet layer, so DataFusion's `Metric::partition()` is preserved on the wire. - from_proto: mirror on the read side — decode the partition and pass it to `Metric::new`. - execution_stage::update_task_metrics: rename the `partition` param to `task_id`, look up `task_infos[task_id].global_input_partition_ids`, and map each incoming metric's local index → global via that slice. Collapse tasks (slice_len=1 with N upstream inputs on the plan side) roll all metrics up to their single output partition. - upsert_metrics_set_for_partition → upsert_metrics_set_for_task: on re-report, drop prior entries for every partition in the task's slice, not just one. - handlers.rs::get_partition_counts: sum across a `&[usize]` slice instead of a single id. Callers pass `info.global_input_partition_ids`. - TaskSummary.partition_id: `u32` → `Vec<u32>`, populated from the task's global partitions. Same JSON key ("keep the ID"), honestly plural — single-partition tasks render as `[N]`, multi-partition as `[a, b, ...]`. Old readers doing `partition_id[0]` still see the first partition. Tests: existing single-partition test refreshed to register task_infos; new `multi_partition_task_preserves_per_partition_detail` verifies a one-task-covers-four-global-partitions case files each partition under its own bucket (skew visible) and that re-reporting the same task replaces snapshots for its whole slice. Verified end-to-end with TPC-H Q1 at SF1 (16 partitions, 2×4-vcore executors → 4 tasks × 4 partitions each in stages 1-2, one collapse task in stage 3). Query completes and verifies against DataFusion; the scheduler API reports distinct per-task output_rows (4.2M / 4.5M / 4.5M / 4.7M in stage 1, summing to the stage aggregate) and the collapse task correctly renders `partition_id=[0]`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…ition Adds a section to docs/developer/architecture.md addressing PR apache#2038 feedback: describes the slice-per-task dispatch model, contrasts it with Spark's one-task-per-partition unit, and shows how the model composes with in-flight DataFusion AQE PoCs (#23026, #23167) to give a single rule library that spans intra-plan, intra-executor slice, and inter-executor stage-boundary levels. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
5 tasks
github-actions
Bot
added
the
documentation
avantgardnerio
commented
Jul 17, 2026
CI's prettier@2.7.1 rewrites *foo* to _foo_.
…iterExec Mirrors DataFusion's RepartitionExec doc style (arrows-up = data-up, child at bottom, executor at top). Both writers use K=3 in the example so the shape difference reads on inspection: - ShuffleWriter (passthrough): K per-output drainers → K files, 1:1 file:partition. - SortShuffleWriter: M per-input writers → M files, K buckets in each (with an index). Both expose the same K-summary contract; on-disk shape is what differs. Notes the coordinator+oneshot idiom is shared (with the Hash branch and with SortShuffle) and could collapse in passthrough once `Hash` is retired in favor of DataFusion's `RepartitionExec(Hash)`.
andygrove
reviewed
Jul 19, 2026
Member
|
I'll benchmark this PR today, both with max of 1 partition per task (current model) and with config set to 0 to allow multiple partitions per task. If there is no regression with setting of 1, I think this is good to go. @phillipleblanc @milenkovicm WDYT? |
…cal input partition Under multi-partition tasks, `ShuffleWriteMetrics::new` was being called with `task_id` as the "partition" arg. This slipped through the task_index → task_id rename (fa54364): the parameter is semantically an operator-local partition index, not a stage-wide task slot. The scheduler's d4f6171 fold added `global_input_partition_ids[local]` as the local→global mapping — with `task_id` in that slot, any stage running more tasks than the slice length reported values ≥ slice_len and the fold errored. execution_graph propagated the Err via `?`, `partition_to_location` never ran, and the downstream stage waited forever. TPC-H Q11 (and any query with enough tasks per stage) hung. Restore the pre-K-drain model: N tasks × 1 bucket each → 1 task × N buckets. Passthrough builds `ShuffleWriteMetrics` inside its drain loop so each spawn owns one bucket keyed by its operator-local input partition (passthrough is 1:1 so local input == local output). Hash pre-allocates a `Vec<ShuffleWriteMetrics>` and a `Vec<BatchPartitioner>` of size num_input_partitions, routes `(local_input_partition, batch)` through the mpsc channel, and the blocking task indexes both by that partition. Total BatchPartitioner count across the stage is unchanged (K per task × 1 task = 1 per task × K tasks pre-K-drain). Metric argument renamed to `input_partition` so it's called what it is. Defense-in-depth: `update_task_metrics` failures in `execution_graph` and `aqe/mod.rs` no longer propagate via `?`. Metrics are observability, not correctness — a folding bug should log and let output-location registration proceed, not hang the whole query. Real bugs still surface via the warn. Verified: TPC-H SF10 Q11 completes in ~8s (matched the 8.6s pre-regression timing at 8bd2e44), verified against DataFusion. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…g master-like per-partition UI Second bug surfaced after the ShuffleWriter K-buckets fix (b8f6ca8): stage 11 of TPC-H Q11 still emitted metrics with local partition indices outside the task's 4-partition slice — up to 15 for a stage feeding a 16-way hash shuffle. Values 4, 6, 8, 11 in the warnings came from a ShuffleReader below a `CoalescePartitionsExec` on the NestedLoopJoin build side. That subtree is `under_collect` per `task_builder::child_scopes` (SinglePartition-required join build ⇒ all descendants keep the full upstream partition count), so its metrics are stamped with 0..upstream_full_count-1, unrelated to the task's slice. `metrics_set_from_task_metrics` errored on those; the previous commit's drop-not-fail wrapper hid the error at the cost of losing the whole task's metrics. Restore the pre-branch behaviour for under-collect: cross-product the metric against every slice member. Aggregate is `metric_value × slice_len` per task — matching master's shape (`N tasks × 1 partition each ⇒ N buckets, each with 1 task's under-collect work`). MetricValue clones share the underlying `Arc<AtomicUsize>`, so N tagged entries pointing at the same counter aggregate correctly (sum → N×; min/max → same value; ratio → same ratio). Restricted-arm metrics still map honestly via `global_input_partition_ids[local]` — cross-product only fires when `local >= slice_len`. Also revert the drop-not-fail wrappers in `execution_graph.rs` and `aqe/mod.rs`: after this fix, `metrics_set_from_task_metrics` has no partition-tag-mismatch error path. `?` on `update_task_metrics` goes back, so real bugs surface loudly again. Documented at length in the TODO on `metrics_set_from_task_metrics`: neither master nor this design is actually *correct*. Master mislabelled every metric with the task's own partition id and over-counted under-collect work by the task count. This design preserves master's per-partition-tag shape (needed by the REST API / TUI which have no way to display arm-scoped metrics) at the cost of hiding the multi-partition-tasks efficiency win (fewer tasks ⇒ fewer redundant under-collect executions) behind the same inflated aggregate. The honest representation is probably one bucket per plan-arm per task with arm identity attached as labels — but that needs API/TUI redesign and is out of scope for this PR. Test: `test_update_task_metrics_under_collect_cross_products_across_slice` covers the new branch. Verified end-to-end: TPC-H SF10 Q11 completes in ~8s with zero warnings; stage 11 now displays its full physical plan with metrics (previously suppressed by dropped-metrics). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…(pre-branch model) Per @andygrove review on apache#2038 discussion_r3610707296: default the cap at 1 so merging this branch is not a breaking change. Behaviour on merge stays "one task per partition" (Spark-style, matching master); multi-partition tasks become opt-in via raising the cap (or `0` for unbounded). Also clears a `cargo doc -D warnings` regression: `(super::sort_shuffle::SortShuffleWriterExec)` link targets on the `ShuffleWriterExec` doc block became redundant once the label resolves through re-exports. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
milenkovicm
reviewed
Jul 19, 2026
milenkovicm
left a comment
milenkovicm
left a comment
Contributor
There was a problem hiding this comment.
few comments regarding config property naming and type. if you're happy with performance please do merge it @andygrove
`restrict_plan_to_partitions` promised in its docstring that a UnionExec child assigned an empty sub-slice becomes a 0-partition subplan, so the executor's `ShuffleWriterExec` iterates `output_partitioning().partition_count()` in agreement with the task's slice. The only rewriters were for `DataSourceExec` and `ShuffleReaderExec`; any other leaf fell through the generic recursion unchanged. `PlaceholderRowExec` (and `EmptyExec`) leaves kept their 1 partition, UnionExec kept reporting the pre-restriction sum, and each task drained every UnionExec output — duplicating rows onto overlapping global partition ids via `PassThrough::resolve`'s `slice.get(local).unwrap_or(local)` fallback. Symptom: `SELECT 1 UNION ALL SELECT 1` returned 4 rows once `max_partitions_per_task=1` split the stage into two tasks. Not exposed by UNION (distinct) — its FinalPartitioned aggregate collapsed the duplicates. Consolidate the two `rewrite_*` helpers into a single `select_output_partitions` that documents the algebraic contract (a projection on the partition-index axis; `Partitioning` kind preserved; `new.execute(j) ≡ plan.execute(indices[j])`), grumbles about DataFusion's missing trait method, and lists which leaf types are covered so future additions have one place to land. For `PlaceholderRowExec` / `EmptyExec`, the algebraic answer is `with_partitions(indices.len())`, but datafusion-proto's codec for these two encodes only the schema — `try_from_placeholder_row_exec` / `try_from_empty_exec` drop the partition count and the decoder always reconstructs with `::new(schema)` (partitions = 1). A restricted 0-partition plan therefore arrives at the executor as 1-partition and the bug reappears. For the empty-slice case we substitute a proto-safe 0-partition `MemorySourceConfig::try_new(&[], schema, None)` — the DataSource codec roundtrips its partition list correctly. Non-empty slices currently only arise as `[0]` on a 1-partition leaf (identity), where returning `None` and letting the leaf pass through unchanged is also proto-safe. Test: `basic::test_union_and_union_all` (previously failing on Linux crates, Linux ballista, and macOS jobs of apache#2038 CI). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The two bind_task_round_robin tests exercise the multi-partition binding path with slice sizes up to 7; that requires max_partitions_per_task unbounded. Since bc7a4ed flipped the default to 1 (master-parity), the mock now takes an explicit SessionConfig so these tests set the knob back to 0. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
andygrove
reviewed
Jul 19, 2026
Member
|
I ran TPC-H @ sf1000 q1-7. q2 was slightly slower other queries were about the same. My benchmarking setup is quite new, so I am not even sure how stable performance numbers are yet. This LGTM. |
andygrove
approved these changes
Jul 19, 2026
andygrove
left a comment
andygrove
left a comment
Member
There was a problem hiding this comment.
LGTM. Thanks @avantgardnerio! I only tested with the default setting for the new config, since that is the use case I am most interested in right now, but look forward to seeing performance improvements with the multi-partition-per-task config set.
Contributor
Author
Thanks @andygrove ! I'm going to go ahead and merge. I'll address the 1 nit in a follow up PR tomorrow. I think we can live with max_partitions_per_task=1 for now. Other settings improve some queries and hurt others, but at least it's up to users to tweak for their workloads. Merging this will unblock me to submit a series of followups to make more small incremental improvements. I really appreciate all the reviews & feedback @milenkovicm @Dandandan & @andygrove ! |
1 task
avantgardnerio
added a commit
that referenced
this pull request
Jul 19, 2026
…ess (#2093) Drop PR-relative phrasing ("pre-multi-partition-tasks execution model", "preserving master's behaviour on merge") and the "Spark user persona" jargon. Describe what the setting does in terms a config-doc reader can act on. Addresses andygrove's review comment on #2038. Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2 tasks
milenkovicm
pushed a commit
that referenced
this pull request
Jul 20, 2026
…2104) Follow-up to #2038 (multi-partition tasks). Four log sites introduced by that PR fire on every scheduling/execution decision and flood scheduler + executor logs at INFO: * `ballista/scheduler/src/cluster/mod.rs`: `stage_has_input_collapse` + `bind_one` — per-bind decision traces (previously prefixed with an ad-hoc `DIAG ` tag, dropped here). * `ballista/executor/src/execution_engine.rs`: `executor plan pre-run` + `executor plan post-run` — per-task DisplayableExecutionPlan dumps. Pre-#2038 this function had no logging; both dumps are noise for anyone not actively debugging. All four move to `debug!`. Turn them back on with `RUST_LOG=ballista_scheduler=debug,ballista_executor=debug`. Closes #2103. Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Migrates Ballista from a 1-task-per-partition model to a slice-per-task model, so that a 16-vcore executor can drive 16 partitions with one task rather than launching 16 single-partition tasks. Goal is to allow arbitrary re-partitioning strategies like the upcoming
OrderedRangeRepartitionandUnorderedRangeRepartitionwhile sharing them with pure DataFusion (a follow-up will remove theHashbranch ofShuffleWriterand replace it with a DataFusionRepartitionExec(Hash)).Opens the door to:
Core changes
ShuffleWriterandSortShuffleWriter. A task'sglobal_output_partition_ids.len()input partitions run concurrently, each producing its own file(s); the internal state fans results out through DataFusion's per-partitionexecute(N)contract. Sort writer's per-pipeline memory budget is divided across concurrent pipelines to preserve the configured per-task total. ASCII threading diagrams live on both writer doc comments.global_output_partition_idsonTaskDefinition/TaskIdproto so the executor knows which global partition ids each task's restricted plan is covering. Encoded asrepeated uint32 ... = 3 [packed=false]so a single-partition message is bit-identical on the wire to the legacyuint32 partition_id = 3scalar.ShuffleWriter(None)walks its child plan (walk_child_partition_mapping) and picks the right mapping:SortPreservingMergeExec-> all outputs go to partition 0,RepartitionExec::Hash/RoundRobin-> K-space is inherently global, otherwise ->global = global_output_partition_ids[local].ShuffleWriter(Hash)andSortShuffleWriteruse their K-space / input-partition space directly.MemorySourceConfigwas silently unrestricted at dispatch and over-reading in multi-partition tasks;restrict_shuffle_readerpreserves theHash/RoundRobinBatchpartitioning kind after shrink soInterleaveExecabove the reader keeps its "children share hash partitioning" invariant. Task builder routes eachUnionExecparent partition to its child's local index.partition_countdirectly from the shuffle writer's immediate child'soutput_partitioning()— that's the number of independent output-partition polling contexts the writer needs to drive. Replaces the earlier leaf-sum walker, which undercounted after upstream fix: keep EmptyExec partition count intact across stage serialization [AQE] #2062 (EmptyExecserde-safe rewrite) and overcounted through aligned-input joins.Scheduler bookkeeping
SuccessfulStage::to_runningwalkstask_infosnewest-to-oldest so retried partitions aren't double-pushed topending.resubmit_successful_stages/reset_running_stageshandle themap_partition_id-is-task_idsemantics correctly (bounds check againsttask_infos.len(), TODO for eventual partition-id semantics).max_task_failures, since retry attempts get fresh task ids.task_idand the previously separatetask_index(which always moved in lockstep at bind time) are merged into a single per-stage append-order slot id.Ops knobs and observability
ballista.scheduler.max_partitions_per_task— Spark-persona escape hatch.0(default) keeps this branch's multi-partition-tasks behavior; setting to1restores the pre-PR one-task-per-partition scheduling for non-collapse stages. Collapse stages ignore the cap since a split collapse would produce partial results downstream can't merge.optional uint32 partition = 16toOperatorMetricon the wire so per-partition detail survives multi-partition tasks.execution_stageused to key metrics by the task's single partition id; under this branch that arg is the task slot, so without this field all N partitions in a slice would file into one bucket. Additive proto change; the UI's per-partition breakdown works again.docs/developer/architecture.mdgains a section on the slice-per-task dispatch model, contrasting it with Spark's per-partition task unit, and showing how the model composes with in-flight DataFusion AQE PoCs (Parallel bounded RANGE-frame window functions without PARTITION BY (draft) datafusion#23026, [PoC/Proposal] AQE-lite: change plan properties at runtime based on stats from pipeline breakers datafusion#23167).Test plan
cargo build --workspacecargo test --workspaceshould_execute_submitted_physical_plan_across_shuffle_stages