Multi-partition tasks: partition_slice plumbing + SortShuffleWriter refactor by avantgardnerio · Pull Request #2038 · apache/datafusion-ballista

avantgardnerio · 2026-07-14T16:51:39Z

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. The goal is to allow arbitrary re-partitioning models like the upcoming OrderedRangeRepartition and UnorderedRangeRepartion while also sharing them with pure Datafusion (there will be a subsequent PR to remove the Hash branch of ShuffleWriter and just replace it with a Datafusion RepartitionExec(Hash))

This opens us up to things like:

parallel sort
parallel windows
parallel sort merge joins
parallel aggregations
OOMs at cluster threshold, not process threshold

Changes in this PR:

Coordinator+oneshot handoff pattern in both ShuffleWriter and SortShuffleWriter. A task's slice.len() input partitions run concurrently, each producing its own file(s); the internal state fans results out through DataFusion's per-partition execute(N) contract. Sort writer's per-pipeline memory budget is divided across concurrent pipelines to preserve the configured per-task total.
partition_slice on TaskDefinition / TaskId proto so the executor knows the global partition ids each task's restricted plan is covering. Writers use it to name shuffle files with global identity. 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 = partition_slice[local]. ShuffleWriter(Hash) and SortShuffleWriter use their K-space / input-partition space directly.
Scan restriction extended: MemorySourceConfig was silently unrestricted at dispatch and over-reading in multi-partition tasks; restrict_shuffle_reader preserves the Hash/RoundRobinBatch partitioning kind after shrink so InterleaveExec above the reader keeps its "children share hash partitioning" invariant.
Scheduler bookkeeping adapted to append-only task_infos:
- SuccessfulStage::to_running walks task_infos newest-to-oldest so retried partitions aren't double-pushed to pending.
- resubmit_successful_stages / reset_running_stages handle the map_partition_id-is-task_index semantics correctly (renamed for clarity, bounds check 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.
- task_slot renamed to task_index throughout for consistency.

Test plan

cargo build --workspace
cargo test --workspace — all 806 tests pass across 17 binaries
should_execute_submitted_physical_plan_across_shuffle_stages (the correctness case that motivated the sort-writer refactor)
All 51 sort shuffle tests
Manual TPC-H run at scale to validate the sort-shuffle multi-partition path under load

🤖 Generated with Claude Code

andygrove · 2026-07-14T20:28:36Z

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.

andygrove · 2026-07-14T20:33:01Z

TPC-H CI job is failing due to q15 returning incorrect data (returning too many rows).

avantgardnerio · 2026-07-14T20:34:07Z

"max partitions per task"

Yeah, absolutely, strictly speaking I think this is a superset of that, so "max of 1" should be easy - I'll add a config.

milenkovicm · 2026-07-15T06:28:30Z

Will review over weekend, until then i have one unreleased comment: great to see all of you here 😁

andygrove · 2026-07-15T14:18:05Z

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.

avantgardnerio · 2026-07-16T00:25:09Z

TPC-H CI job is failing due to q15 returning incorrect data (returning too many rows).

Fixed!

avantgardnerio · 2026-07-16T19:10:23Z

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):

milenkovicm · 2026-07-17T06:57:30Z

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>

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>

andygrove · 2026-07-17T18:00:32Z

Thanks for this — the slice-per-task direction is compelling. I ran this branch on my k3s TPC-H SF1000 cluster to see how it moves the needle, and wanted to share the numbers plus a root-cause read while the run was going, since a clear pattern showed up early.

I stopped the run after Q7 (single iteration) once the signal was consistent, so this is a partial suite — but it's enough to surface a regression on the join-heavy queries that's worth discussing.

Setup (identical to our published SF1000 baseline, only the image differs)


Branch	`avantgardnerio/arrow-ballista@308f1495` (this PR)
Baseline	`apache/main` (the published benchmarking-guide SF1000 AQE-off numbers, same cluster)
Scale	SF1000, node-local Parquet
Cluster	2 executors, one per node, 8 cores / 56 GiB, `--memory-pool-size=48GB`
Config	`target_partitions=32`, `prefer_hash_join=false`, `enable_dynamic_filter_pushdown=false`, AQE off (`ballista.planner.adaptive.enabled=false`)
Iterations	1

Results (Q1–Q7, seconds; ratio = baseline / PR, >1 = PR faster)

Query	PR #2038	main (AQE off)	ratio
1	41.3	70.8	1.71×	🟢
2	183.2	155.5	0.85×	🔴
3	256.5	206.2	0.80×	🔴
4	67.2	76.8	1.14×	🟢
5	692.9	542.7	0.78×	🔴
6	15.0	18.3	1.22×	🟢
7	748.7	575.5	0.77×	🔴
Σ 1–7	2004.8	1645.8	0.82×

The direction is consistent: join-free queries improve (Q1, Q4, Q6), join-heavy queries regress ~20–25% (Q2, Q3, Q5, Q7).

Root cause: coarser scheduling granularity → straggler tails on skewed join stages

The DIAG bind_one scheduler traces make the mechanism visible. On a light query, each stage binds cleanly as a few fat slices with no leftover single-partition tasks and a sub-second spread. On a join-heavy query, the same 32-partition stages fragment and dispatch over a long window:

stage kind	tasks	1-partition fragments	first→last bind spread
light query, stage	4	0	<1 s
join query, stage A	6	2	100 s
join query, stage B	9	5	178 s
join query, worst	18	15	—

A slice-task holds all of its vcore slots until its slowest partition finishes. On join stages, per-partition runtime is very uneven (join-key skew + spilling SortExec — one partition spilled 22 GB / 455 M rows in an SMJ with join_time≈413 s). So the binder places a couple of fat slices, then has to dribble the remaining partitions out one at a time as slots free (the frag1 counts), and a whole stage's partitions end up spread across 100–180 s. During each slice's tail, the sibling vcores sit idle and the next partitions can't dispatch. The pre-PR 1-task-per-partition model freed each slot independently, so the 32 partitions load-balanced finely across the 16 vcores. Uniform stages don't hit the tail and actually gain from the lower per-task overhead — hence the wins on Q1/Q4/Q6.

A second-order cost: fragmented stages re-encode a restricted plan per task, so a 15-fragment stage pays ~15× the plan-serialization of a clean 4-task stage.

What this does not appear to be: memory starvation

I initially suspected the shared per-task pool, but memory_pool_policy scales it correctly — an N-vcore slice gets N × per_vcore, so an 8-vcore task gets the full 48 GB FairSpillPool split across its 8 pipelines (≈6 GB each, same as the old per-task budget, and better at the tail as siblings deregister). The heavy spilling looks inherent to a 455 M-row SMJ sort at ~6 GB/partition, not amplified by this PR. So this reads as a scheduling/utilization regression rather than a memory one.

Possible directions (just ideas)

Cap or shrink slice size on stages whose child is a join / sort so partition granularity stays fine where runtimes are skewed.
Let a slice release vcores as individual partitions complete, rather than holding all slots until the slowest finishes.

Happy to re-run the full 22-query suite (and with AQE on) on any follow-up, or share the raw scheduler/executor logs. Really nice work opening up the parallel-sort / parallel-join path — this is about tuning the scheduling granularity, not the idea.

Benchmarking, log analysis, and this write-up were done with LLM assistance (Claude).

Per andygrove's PR apache#2038 review — give Spark-persona users an opt-out knob for the multi-partition-tasks model. `0` (default) leaves this branch's behavior unchanged; setting to `1` restores 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. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Feedback from apache#2038 (andygrove): make the legacy scalar partition_id repeated so the single-partition case keeps working and multi-partition tasks are naturally supported. We kept the more explicit name (`global_output_partition_ids`) but moved it back onto tag 3 with `[packed=false]`, so a single-element list writes exactly one varint at tag 3 — bit-identical to the old `uint32 partition_id = 3` scalar. Old scalar readers parse single- partition messages unchanged; multi-partition messages emit multiple varints on the same tag, which an old reader would collapse to the last value (graceful degradation, not a hard failure). Proto3 requires repeated numeric readers to accept both packed and unpacked encodings, so new readers work with either shape. 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>

avantgardnerio · 2026-07-17T23:39:29Z


 ![Query Execution Flow](images/query-execution.png)

+## Multi-Partition Tasks


@milenkovicm WDYT? Better, or worse?

github-actions Bot added the python label Jul 14, 2026

avantgardnerio requested review from andygrove and milenkovicm July 14, 2026 16:59

github-actions Bot added the development-process label Jul 14, 2026

andygrove reviewed Jul 14, 2026

View reviewed changes

Comment thread ballista/core/src/execution_plans/sort_shuffle/writer.rs Outdated

andygrove reviewed Jul 14, 2026

View reviewed changes

Comment thread ballista/core/proto/ballista.proto Outdated

andygrove reviewed Jul 15, 2026

View reviewed changes

Comment thread ballista/core/proto/ballista.proto Outdated

Dandandan reviewed Jul 15, 2026

View reviewed changes

Comment thread .github/workflows/tpch.yml

Comment thread ballista/core/proto/ballista.proto Outdated

andygrove reviewed Jul 15, 2026

View reviewed changes

Comment thread ballista/core/proto/ballista.proto

andygrove reviewed Jul 15, 2026

View reviewed changes

Comment thread ballista/core/proto/ballista.proto Outdated

avantgardnerio mentioned this pull request Jul 15, 2026

Rename task_slots -> vcores (executor spec + CLI) #2045

Merged

5 tasks

avantgardnerio force-pushed the brent/multi-partition-tasks branch from 1b98039 to 7d3c602 Compare July 16, 2026 00:20

avantgardnerio mentioned this pull request Jul 16, 2026

Live upgrade strategy #2071

Open

avantgardnerio and others added 3 commits July 17, 2026 01:28

avantgardnerio force-pushed the brent/multi-partition-tasks branch from b261c57 to 308f149 Compare July 17, 2026 16:11

github-actions Bot removed python development-process labels Jul 17, 2026

andygrove reviewed Jul 17, 2026

View reviewed changes

Comment thread ballista/core/proto/ballista.proto Outdated

andygrove reviewed Jul 17, 2026

View reviewed changes

Comment thread ballista/core/proto/ballista.proto

avantgardnerio and others added 2 commits July 17, 2026 12:32

avantgardnerio marked this pull request as ready for review July 17, 2026 21:01

avantgardnerio and others added 2 commits July 17, 2026 15:12

avantgardnerio requested review from Dandandan and andygrove July 17, 2026 21:51

avantgardnerio mentioned this pull request Jul 17, 2026

feat: add BALLISTA_PROTOCOL_VERSION + k8s health probes #2088

Open

5 tasks

github-actions Bot added the documentation Improvements or additions to documentation label Jul 17, 2026

avantgardnerio commented Jul 17, 2026

View reviewed changes


		![Query Execution Flow](images/query-execution.png)

		## Multi-Partition Tasks

Uh oh!

Conversation

avantgardnerio commented Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Uh oh!

andygrove commented Jul 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

andygrove commented Jul 14, 2026

Uh oh!

avantgardnerio commented Jul 14, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

milenkovicm commented Jul 15, 2026

Uh oh!

Uh oh!

Uh oh!

andygrove commented Jul 15, 2026

Uh oh!

avantgardnerio commented Jul 16, 2026

Uh oh!

avantgardnerio commented Jul 16, 2026

Uh oh!

milenkovicm commented Jul 17, 2026

Uh oh!

Uh oh!

Uh oh!

andygrove commented Jul 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Setup (identical to our published SF1000 baseline, only the image differs)

Results (Q1–Q7, seconds; ratio = baseline / PR, >1 = PR faster)

Root cause: coarser scheduling granularity → straggler tails on skewed join stages

What this does not appear to be: memory starvation

Possible directions (just ideas)

Uh oh!

avantgardnerio Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

avantgardnerio commented Jul 14, 2026 •

edited

Loading

andygrove commented Jul 14, 2026 •

edited

Loading

andygrove commented Jul 17, 2026 •

edited

Loading