feat: add poll_now_notify to poll_loop and on_work_available callback - #50
Open
Jeadie wants to merge 107 commits into
Open
feat: add poll_now_notify to poll_loop and on_work_available callback#50Jeadie wants to merge 107 commits into
poll_now_notify to poll_loop and on_work_available callback#50Jeadie wants to merge 107 commits into
Conversation
Adds an optional poll_now_notify: Option<Arc<Notify>> parameter to the executor poll_loop so the scheduler can wake an idle executor immediately instead of waiting for the next poll interval. The idle wait becomes a tokio::select! over the poll interval and the notify. Adds an OnWorkAvailableFn callback to SchedulerConfig, invoked from the query stage scheduler when new work becomes available (after JobSubmitted and when new stages become runnable). Ported from #12 for upstreaming.
Jeadie
changed the title
poll_now_notify to poll_loop and on_work_available callback
Author
|
Upstreaming: apache#1893 |
Bumps [uuid](https://github.com/uuid-rs/uuid) from 1.23.3 to 1.23.4. - [Release notes](https://github.com/uuid-rs/uuid/releases) - [Commits](uuid-rs/uuid@v1.23.3...v1.23.4) --- updated-dependencies: - dependency-name: uuid dependency-version: 1.23.4 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
* chore(ci): Deploy WebTUI to nightlies.a.o * Use v8.0.5 of rsync-deployments that is allowed by ASF Infra https://github.com/apache/infrastructure-actions/blob/8d7b86287d0beafc4d68a06fa23e751a04e3976d/actions.yml#L159C3-L160 * burnett01/rsync-deployments works only on Linux * Test whether rsync is pre-installed on macos runners * Use rsync command instead of Github Actions * Load the RSYNC_KEY (a PEM content) into SSH agent * Add RSYNC_HOST to SSH's known hosts Trying to solve error: "Host key verification failed." * Make use of secrets.NIGHTLIES_RSYNC_PATH * Try to figure out the destination PATH * Create the folders for datafusion/ballista/tui/<version> * Debug the copied files * Custom build for nightlies * Deploy to nightlies.a.o only on merge * Merge two CI jobs into one * Extract local env vars for reuse * Rename SSH_OPTIONS to SSH_OPTS * Rename SSH_OPTS to OPTIONS_SSH * Inline the SSH options * Experiment with './' as a public url * Build the WASM32 app once with public_dir="./" This way it could be used at any public url * Deploy at nightlies.a.o only on push to main * Use long name for the rsync options. It is easier to understand the command by humans Also remove some options which are not relevant for us - recursive, links, specials+devices
…pache#1902) upgrade_for_ballista re-applied Ballista's opinionated DataFusion defaults (prefer_hash_join, hash_join_single_partition_threshold, view-type flags) unconditionally. In remote_with_state this runs after the user has set their own values, silently reverting datafusion.* overrides; ballista.* settings were unaffected because the defaults do not touch them. Apply the defaults only when the config has not already been through Ballista setup (no BallistaConfig extension present), so user overrides survive.
Bumps [config](https://github.com/rust-cli/config-rs) from 0.15.24 to 0.15.25. - [Changelog](https://github.com/rust-cli/config-rs/blob/main/CHANGELOG.md) - [Commits](rust-cli/config-rs@v0.15.24...v0.15.25) --- updated-dependencies: - dependency-name: config dependency-version: 0.15.25 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…he#1914) Bumps [taiki-e/install-action](https://github.com/taiki-e/install-action) from 2.81.10 to 2.82.6. - [Release notes](https://github.com/taiki-e/install-action/releases) - [Changelog](https://github.com/taiki-e/install-action/blob/main/CHANGELOG.md) - [Commits](taiki-e/install-action@v2.81.10...9bcaee1) --- updated-dependencies: - dependency-name: taiki-e/install-action dependency-version: 2.82.6 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…apache#1900) * feat(core): broadcast small build sides via non-zero CollectLeft thresholds Set hash_join_single_partition_threshold to 10MB and hash_join_single_partition_threshold_rows to 1M so the AQE join resolver collects small build sides into a broadcast (CollectLeft) hash join instead of repartitioning them. Update the client settings test to assert the new default values. * fix(scheduler): restrict CollectLeft broadcast to safe join types The adaptive join resolver chose CollectLeft purely on size thresholds. A CollectLeft join broadcasts the build (left) side to every probe task, where each task sees only its slice of the probe side, so join types that emit rows on behalf of the build side (Left, Full, LeftSemi, LeftAnti, LeftMark) would have every task emit them independently and produce duplicate or spurious rows. Restrict CollectLeft to join types whose output is driven entirely by the probe side (Inner, Right, RightSemi, RightAnti, RightMark), evaluated on the post-swap join type using the same swap decision as the resolver, and keep unsafe join types repartitioned. Add resolver tests covering left/left-anti (repartitioned) and right (still collected) joins. * fix(scheduler): guard static planner broadcast and add join-type test matrix The static DefaultDistributedPlanner promoted small-side hash joins to CollectLeft on size alone, with no join-type check, so a LEFT/FULL/semi/anti/ mark join whose build side fit under broadcast_join_threshold_bytes was broadcast and produced duplicate or spurious rows — the same hazard already guarded on the adaptive path. Move the broadcast-safety predicate into physical_optimizer::join_selection so both planners share one definition, and apply it in maybe_promote_to_broadcast using the post-swap join type. Add a test matrix covering both planners, both build-side orientations, every join type, prefer_hash_join on/off, and the post-swap decision, plus a unit test that locks down the broadcast-safe join-type set. * test(core): update upgrade defaults test for non-zero CollectLeft threshold apache#1902 added tests asserting the Ballista default hash_join_single_partition_threshold is 0; this PR raises that default to 10 MB (and the row threshold to 1M), so should_apply_defaults_when_upgrading_plain_config now asserts the new values. Give should_preserve_user_overrides_on_upgrade an override distinct from the default so it still proves the user value survives. * fix(scheduler): use dyn ExecutionPlan::downcast_ref in broadcast tests DataFusion 54 removed ExecutionPlan::as_any in favor of the inherent downcast_ref on dyn ExecutionPlan, which auto-derefs through Arc. The new broadcast-safety tests still used as_any().downcast_ref(), which no longer compiles. Drop the as_any() hop to match the rest of the file. * fix: demote DataFusion-promoted unsafe CollectLeft joins to partitioned DataFusion's JoinSelection runs during create_physical_plan and, with the non-zero hash_join_single_partition_threshold default this PR sets, can stamp CollectLeft on a join without restricting by join type. The distributed planner's broadcast-safety guard only covered the Partitioned -> CollectLeft promotion path, so a pre-promoted unsafe join (e.g. a LEFT join with the small side already on the left) reached the broadcast-lowering branch and replicated the outer side across every probe task, producing duplicate/spurious rows. In maybe_promote_to_broadcast, demote any CollectLeft hash join whose join type is not broadcast-safe back to a Partitioned (shuffle) join, hash-partitioning both inputs on the join keys. This is a correctness guard, so it runs regardless of the Ballista broadcast threshold. Add a test that mirrors the production session config so DataFusion's own JoinSelection performs the promotion. * fix: disable DataFusion dynamic filter pushdown in session defaults DataFusion 54 populates dynamic filters at runtime from an upstream operator (a hash join build side, a TopK heap, a partial aggregate) and reads them in a downstream scan within the same plan. Ballista splits a plan into stages at shuffle and broadcast boundaries that execute as independent tasks, so when the producing operator and the consuming scan land in different stages the filter is never populated across the boundary and the consuming scan blocks forever. This deadlocks every multi-join TPC-H query under the adaptive (AQE) planner, where broadcast (CollectLeft) hash joins create the cross-stage producer/ consumer split — the build stages complete but the probe stages hang. The static planner avoids it only because it plans sort-merge joins, which carry no join dynamic filter. Pin datafusion.optimizer.enable_dynamic_filter_pushdown to false in the Ballista session defaults, alongside the other DataFusion options Ballista already disables because they rely on in-process state that cannot cross stage boundaries. Tracked by apache#1375.
…lanner (apache#1904) * feat(scheduler): broadcast small build side of SortMergeJoinExec Add ballista.optimizer.broadcast_sort_merge_join_enabled (default false). When enabled, the static distributed planner converts a SortMergeJoinExec whose smaller side is under broadcast_join_threshold_bytes into a CollectLeft hash join and lowers the build side as a broadcast stage, reusing the existing hash-join broadcast path. Redundant input sorts are dropped during conversion. TPC-H SF10 (AQE off): ~16% faster across the suite (q8 1.6x, q11 1.5x, q2 1.4x), identical row counts. * feat(config): default broadcast_sort_merge_join_enabled to true Enable the SortMergeJoinExec broadcast conversion by default so small build sides are broadcast without opt-in (~16% on TPC-H SF10, AQE off). Disable the flag explicitly in the sort-merge-join client test so it still exercises the plain SortMergeJoinExec execution path. * test: assert SMJ broadcast plans via string snapshots Rewrite the three SortMergeJoinExec broadcast tests to assert the stage plan as a string via the assert_plan! macro instead of walking the plan tree with manual downcasts, making the expected plan explicit and the tests easier to read.
…che#1924) Add a physical_plan field to ExecuteQueryParams.query so clients can submit an ExecutionPlan directly, bypassing logical plan creation on the scheduler. DefaultDistributedPlanner::plan_query_stages already splits physical plans into shuffle stages at RepartitionExec(Hash) boundaries, so this unblocks distributing plans that contain custom/FFI operators with no logical-plan representation. Thread the decoded plan through the existing job-submission event path (SchedulerServer::submit_job -> QueryStageSchedulerEvent::JobQueued -> SchedulerState::submit_job -> TaskManager::submit_job) as a new SubmitPlan enum (Logical | Physical) instead of a bare LogicalPlan. Physical-plan jobs always use the static distributed planner and skip create_physical_plan/EXPLAIN handling entirely; AQE is not available for this path since it requires a logical plan to re-plan against. Add execute_physical_plan() in ballista-core as a client-side helper that mirrors DistributedQueryExec's submission flow (serialize with the physical extension codec, submit, poll status, fetch partitions).
…pache#1903) * feat: failing tests for sibling tasks cancelled on stage failure * feat: cancel running stages and their tasks on job abort * refactor: implement abort_running per ExecutionGraph impl * feat: mark a failed stage's in-flight tasks as cancelled * Revert "refactor: implement abort_running per ExecutionGraph impl" This reverts commit 1626af4. * refactor: use "killed" for aborted in-flight task error
Bumps [ctor](https://github.com/mmastrac/linktime) from 1.0.7 to 1.0.8. - [Release notes](https://github.com/mmastrac/linktime/releases) - [Changelog](https://github.com/mmastrac/linktime/blob/master/CHANGELOG.md) - [Commits](mmastrac/linktime@ctor-1.0.7...ctor-1.0.8) --- updated-dependencies: - dependency-name: ctor dependency-version: 1.0.8 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Bumps [bytesize](https://github.com/bytesize-rs/bytesize) from 2.4.0 to 2.4.2. - [Release notes](https://github.com/bytesize-rs/bytesize/releases) - [Changelog](https://github.com/bytesize-rs/bytesize/blob/master/CHANGELOG.md) - [Commits](bytesize-rs/bytesize@bytesize-v2.4.0...bytesize-v2.4.2) --- updated-dependencies: - dependency-name: bytesize dependency-version: 2.4.2 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…he#2056) Bumps [cryptography](https://github.com/pyca/cryptography) from 46.0.3 to 48.0.1. - [Changelog](https://github.com/pyca/cryptography/blob/main/CHANGELOG.rst) - [Commits](pyca/cryptography@46.0.3...48.0.1) --- updated-dependencies: - dependency-name: cryptography dependency-version: 48.0.1 dependency-type: indirect ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Bumps [pyjwt](https://github.com/jpadilla/pyjwt) from 2.12.0 to 2.13.0. - [Release notes](https://github.com/jpadilla/pyjwt/releases) - [Changelog](https://github.com/jpadilla/pyjwt/blob/master/CHANGELOG.rst) - [Commits](jpadilla/pyjwt@2.12.0...2.13.0) --- updated-dependencies: - dependency-name: pyjwt dependency-version: 2.13.0 dependency-type: indirect ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Bumps [actions/setup-node](https://github.com/actions/setup-node) from 6.4.0 to 7.0.0. - [Release notes](https://github.com/actions/setup-node/releases) - [Commits](actions/setup-node@48b55a0...8207627) --- updated-dependencies: - dependency-name: actions/setup-node dependency-version: 7.0.0 dependency-type: direct:production update-type: version-update:semver-major ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
… their versions are allowed (apache#2058)
The banner was using tui-big-text crate that seems to use too modern UTF-8 symbols which are not yet available for many operating systems and because of this the banner didn't render well.
…tion grow (apache#2061) * fix(shuffle): spill when the memory pool rejects a reservation grow The sort shuffle writer discarded the result of `reservation.try_grow`, treating it as best-effort on the grounds that the absolute buffered-bytes counter would still bound RSS. That counter is per-task and private to the writer, so a rejected grow produced no spill at all: the batch stayed buffered while the pool had not accounted for it. The writer then under-reported its usage to the shared `MemoryPool`, letting other consumers see memory as free that was not, and ignoring the very pressure signal a `with_can_spill(true)` consumer registers to receive. Treat a rejected grow as a spill trigger alongside the existing per-task budget. Spilling flushes the buffered batches to disk and frees the reservation, so no retry is needed — the writer holds nothing afterwards. * test(shuffle): name spill triggers explicitly in sort shuffle writer tests The writer has two independent spill triggers: the shared MemoryPool rejecting a grow, and the private per-task buffered-bytes budget. Naming the pool constant UNCONSTRAINED_POOL described only one knob, so a test passing an unconstrained pool while expecting spills read as a contradiction when the spill in fact came from the per-task budget. Name both knobs by the trigger they arm and state which trigger each test exercises, so each test reads as one trigger armed and the other disarmed.
`client_ttl` defaulted to 0, which disables the client pool entirely. Without a
pool the shuffle-read path opens and discards a new TCP connection per fetch, so
a single shuffle-heavy query exhausts the host's ephemeral ports and fails with
`AddrNotAvailable` — as does every query after it, until the sockets leave
TIME_WAIT.
A single executor never shows this, because shuffle reads are served from local
files and no remote fetch happens. It appears as soon as a second executor
exists, which is every real deployment.
Default `client_ttl` to 30 seconds so the pool is on unless explicitly disabled.
The pooling machinery already exists and the shuffle reader is written assuming
it ("multiplexing over few connections"); it was simply off unless the operator
knew to pass `--client-ttl`.
Measured on the SF1 TPC-DS gate with 1 scheduler + 2 executors, all defaults:
before after
peak TIME_WAIT 16343 12
gate 11 OK/77 FAIL 88 OK/0 FAIL
The ephemeral range on the test host is 16384 ports, so one query consumed all
of it.
* Rename task_slots -> vcores across the workspace
The name `task_slots` implies "how many task instances can run concurrently
on this executor," but it's actually been used two different ways:
- concurrent-task capacity (scheduler-side accounting)
- virtual-core count that a task can drive through one plan-Arc
(what `--concurrent-tasks` actually controlled — the CLI help even
hinted at this)
Rename to `vcores`, following the YARN precedent
(`yarn.nodemanager.resource.cpu-vcores`) and Spark's `spark.executor.cores`.
The `v` prefix is self-documenting: it isn't physical `nproc`, and readers
grepping in a year don't have to hunt for a doc comment to know that.
The new name also leaves room for future execution models that decouple
"how many tasks" from "how many cores per task."
Mechanical rename, no behavior change. Field numbers in the proto are
preserved for wire compat.
Highlights:
- `ExecutorSpecification.task_slots` -> `vcores` (u32), with a doc
comment citing YARN/Spark.
- Proto: `ExecutorResource.task_slots` -> `vcores`,
`AvailableTaskSlots` -> `AvailableVcores`, `ExecutorTaskSlots` ->
`ExecutorVcores`, `PollWorkParams.num_free_slots` -> `num_free_vcores`.
- CLI: `--concurrent-tasks` -> `--vcores` (help text explains it's
scheduler-facing accounting, not physical nproc).
- Cluster bookkeeping: `InMemoryClusterState.task_slots` ->
`available_vcores`. Local variables in bind helpers renamed to match.
- TUI: "Task Slots" column -> "vCores"; `SortColumn::TaskSlots` ->
`Vcores`.
- Executor process: startup log line and doc comment on `vcores`
reworded to drop stale "concurrent tasks" prose.
- Docs, docker-compose, tpch workflow, and CLI README updated so
`--concurrent-tasks` invocations don't fail at boot under the new
flag name. Historical changelogs left as-is.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* Add deprecated aliases for --concurrent-tasks and task_slots API
Keeps the pre-rename operator surface working for one release so in-place
cluster upgrades don't break:
- `--concurrent-tasks` on ballista-executor and ballista-cli: hidden clap
alias that folds into `--vcores` with a stderr deprecation warning.
- `ExecutorSpecification::with_task_slots` / `task_slots()`,
`ExecutorData::{total,available}_task_slots()`, and
`ExecutorDataChange::task_slots()`: `#[deprecated]` shims returning the
renamed fields.
gRPC wire is already binary-compatible (field tags and types unchanged);
no shim needed there.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
* Revert user-personas.md task_slots -> vcores wording
The persona doc is a stable contract about what each user type relies
on; naming refactors shouldn't churn it.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
---------
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…cans (apache#2070) A task executes one partition of its stage, so each file scan beneath it is pinned to the single file group that partition reads — otherwise DataFusion 54's shared work-queue lets a lone `execute(p)` drain the queue and scan the whole table. That pinning applied the stage's partition number directly to every scan, which is only right when the partition passes straight through. A `UnionExec` concatenates its children's partitions: with N-partition children, stage partition `N + k` is the right child's local partition `k`. Applying the stage number there addressed the wrong group, and for partitions past a scan's group count it addressed none at all, leaving that scan unrestricted and free to read its entire table. The result was a silently inflated answer — no error, just several times too much data. Walk the plan from the stage root instead, mirroring `UnionExec::execute`'s own partition arithmetic, so each scan is pinned to the group its local partition reads. Children a partition never reaches are left untouched; this task does not execute them. Minimal case, SF1, executor `-c 8`: select round(sum(p),2) from (select ws_ext_sales_price p from web_sales union all select cs_ext_sales_price p from catalog_sales); -- expected 5492796521.50 -- was 49389026433.70 at --partitions 16 (~9x, varying run to run) Correct now at every partition count tried (4, 5, 6, 8, 16, 20, 32); previously wrong whenever the union stage exceeded one wave of executor task slots. TPC-DS q2 and q5 now match single-process DataFusion and join the correctness gate. Both were skipped as "non-deterministic (ORDER BY ties)"; that diagnosis was wrong — with a total order imposed and LIMIT removed they still mismatched, because they were hitting this bug. q31 and q71 are genuinely ordering-only and stay skipped.
…apache#2062) The adaptive planner's PropagateEmptyExecRule collapses a provably empty sub-plan into an EmptyExec that inherits the partition count of the node it replaces, and the scheduler sizes a stage's task count from that count. datafusion-proto encodes an EmptyExec as its schema alone, so the node decodes on the executor with a single partition. A stage built from a multi-partition EmptyExec therefore launches N tasks against a plan that can only execute partition 0, and every task above it fails with: Internal("Assertion failed: partition < self.partitions: EmptyExec invalid partition 1 (expected less than 1)") Rewrite multi-partition EmptyExec nodes into a round-robin RepartitionExec over a single-partition EmptyExec when a stage plan is built for the wire. The plan is equivalent and a RepartitionExec carries its partitioning through serialization. The rewrite runs after all physical optimizer rules, so no later rule can collapse it, and it is a no-op for the static planner, which never produces a multi-partition EmptyExec. Fixes the assertion for all 15 affected TPC-DS queries at SF1.
…apache#2072) * fix: remove the SortMergeJoinExec broadcast conversion The static distributed planner converted a `SortMergeJoinExec` with a small build side into a broadcast `CollectLeft` hash join, gated by `ballista.optimizer.broadcast_sort_merge_join_enabled` (default `true`). The conversion produced incorrect results, so remove it along with its config key rather than keep a knob whose safe value is always `false`. This affects the default configuration: Ballista sets `datafusion.optimizer.prefer_hash_join = false` (see apache#1648), so joins are planned as sort-merge unless the user opts out, and the conversion then fired on any join with a side under the threshold. The conversion also worked against the reason sort-merge is the default. DataFusion's hash join has no spill support, so a `CollectLeft` build side must fit in memory in every parallel task, while sort-merge spills. Converting one into the other reintroduced the memory behavior that the `prefer_hash_join = false` default exists to avoid. Broadcast promotion of a `HashJoinExec` under `ballista.optimizer.broadcast_join_threshold_bytes` is unchanged. Re-enable TPC-DS q4 and q78 in the correctness gate; they were skipped for this divergence. q38 and q87 stay skipped for unrelated INTERSECT and EXCEPT issues. * test: regenerate TPC-H plan-stability goldens Removing the SortMergeJoinExec broadcast conversion changes the frozen distributed plan for 11 queries. Each change is the same rewrite being undone: a CollectLeft hash join over a broadcast stage becomes a sort-merge join over repartitioned, sorted inputs, dropping the broadcast stage and the column-swap projection. The 20 removed CollectLeft joins (19 Inner, 1 RightAnti) map 1:1 onto 20 added sort-merge joins (19 Inner, 1 LeftAnti); the anti-join flips back because CollectLeft had swapped its inputs.
* docs: add benchmarking page to the contributors guide Document how Ballista is benchmarked at scale and record the current result set in one place. Covers the reference cluster shape and why it is shaped that way, the shared SQLBench-H query set, and self-contained commands for running TPC-H against Ballista, vanilla Spark, and Comet so a result can be reproduced without repo-specific tooling. Explains why Comet is the most informative comparison available -- it executes DataFusion physical plans using the same operators -- along with the caveats that keep it from being an operator-level A/B: Spark's optimizer and AQE produce the plan Comet accelerates, Comet falls back to the JVM for unsupported operators, and the two use different distribution models. Results are tracked as a single current set pinned to the exact commit they came from, with AQE on and AQE off reported side by side. Most entries are TBD pending measurement. * docs: apply prettier formatting to the benchmarking page * docs: clarify that Spark and Comet run an identical copy of the shared query set * docs: record vanilla Spark SF1000 suite results * docs: record Comet 0.17.0 SF1000 suite results * docs: record Ballista AQE-off SF1000 suite results at main@26b29391 * docs: record Ballista AQE-on SF1000 suite results at main@26b29391
The TPC-H, TPC-DS, Rust, Docker, and Web TUI workflows are already gated on path allowlists, but those allowlists use whole-directory globs such as `ballista/**` and `benchmarks/**`. Ten markdown files live inside those directories, so a README-only change fires both benchmark suites and three builds. PR apache#2015 touched only `ballista/client/README.md` and ran the full matrix for nothing. Append `!**/*.md` to each `paths` list. GitHub applies patterns in order and lets the last match win, so the exclusion must come last. Changes that mix code and markdown still trigger every workflow. The prettier check in `dev.yml` is deliberately left ungated since it lints markdown, as is the `docs.yaml` site deploy.
Bumps [tokio](https://github.com/tokio-rs/tokio) from 1.52.3 to 1.52.4. - [Release notes](https://github.com/tokio-rs/tokio/releases) - [Commits](tokio-rs/tokio@tokio-1.52.3...tokio-1.52.4) --- updated-dependencies: - dependency-name: tokio dependency-version: 1.52.4 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Bumps [tokio](https://github.com/tokio-rs/tokio) from 1.52.3 to 1.52.4. - [Release notes](https://github.com/tokio-rs/tokio/releases) - [Commits](tokio-rs/tokio@tokio-1.52.3...tokio-1.52.4) --- updated-dependencies: - dependency-name: tokio dependency-version: 1.52.4 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…e#2076) Bumps [taiki-e/install-action](https://github.com/taiki-e/install-action) from 2.83.2 to 2.83.3. - [Release notes](https://github.com/taiki-e/install-action/releases) - [Changelog](https://github.com/taiki-e/install-action/blob/main/CHANGELOG.md) - [Commits](taiki-e/install-action@43aecc8...ed67fa3) --- updated-dependencies: - dependency-name: taiki-e/install-action dependency-version: 2.83.3 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…ache#2078) * chore(deps): bump github/codeql-action/init from 4.37.0 to 4.37.1 Bumps [github/codeql-action/init](https://github.com/github/codeql-action) from 4.37.0 to 4.37.1. - [Release notes](https://github.com/github/codeql-action/releases) - [Changelog](https://github.com/github/codeql-action/blob/main/CHANGELOG.md) - [Commits](github/codeql-action@99df26d...7188fc3) --- updated-dependencies: - dependency-name: github/codeql-action/init dependency-version: 4.37.1 dependency-type: direct:production update-type: version-update:semver-patch ... Signed-off-by: dependabot[bot] <support@github.com> * Group updates for github/codeql-action/** Use commit prefixes for deps/action updates --------- Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: Martin Tzvetanov Grigorov <mgrigorov@apache.org>
Bumps [tokio](https://github.com/tokio-rs/tokio) from 1.52.4 to 1.53.0. - [Release notes](https://github.com/tokio-rs/tokio/releases) - [Commits](tokio-rs/tokio@tokio-1.52.4...tokio-1.53.0) --- updated-dependencies: - dependency-name: tokio dependency-version: 1.53.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Bumps [tokio](https://github.com/tokio-rs/tokio) from 1.52.4 to 1.53.0. - [Release notes](https://github.com/tokio-rs/tokio/releases) - [Commits](tokio-rs/tokio@tokio-1.52.4...tokio-1.53.0) --- updated-dependencies: - dependency-name: tokio dependency-version: 1.53.0 dependency-type: direct:production update-type: version-update:semver-minor ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…E and the static planner (apache#2086) * feat: make AQE respect broadcast_join_threshold_bytes The ballista.optimizer.broadcast_join_threshold_bytes config is only consumed by the static distributed planner (maybe_promote_to_broadcast). Under adaptive query planning, broadcast (CollectLeft) selection in DynamicJoinSelectionExec::to_actual_join instead used DataFusion's hash_join_single_partition_threshold (1 MiB default), so the Ballista key had no effect and the effective cutoff was silently a different value. Use broadcast_join_threshold_bytes as the byte threshold in the AQE join selection path, keeping DataFusion's row threshold as the absent-stats fallback. A value of 0 disables broadcast promotion, matching the static planner. This gives a single config key consistent behavior under both planners. Closes apache#2085 * feat: add Ballista broadcast_join_threshold_rows and use it in AQE Follow-up within the same change: the AQE join-selection path also used DataFusion's hash_join_single_partition_threshold_rows as the row-count fallback. There was no Ballista equivalent, so the row threshold still escaped the single-config goal. Add ballista.optimizer.broadcast_join_threshold_rows (default 128K, mirroring DataFusion's previous default) plus SessionConfigExt accessors, and use it in DynamicJoinSelectionExec::to_actual_join instead of the DataFusion key. AQE broadcast selection now depends only on Ballista config. Document both broadcast thresholds in the AQE tuning guide. * fix: default broadcast_join_threshold_rows to 1M to match existing behavior SessionConfig::new_with_ballista() already installs a 1,000,000 row threshold for DataFusion's hash_join_single_partition_threshold_rows. Default the new Ballista row key to the same value so consolidating AQE onto the Ballista keys does not silently lower the effective row-count broadcast cutoff. * refactor: derive new_with_ballista join thresholds from Ballista config defaults Instead of hard-coding 10 MiB / 1M for the DataFusion hash_join_single_partition_threshold[_rows] session settings, read them from BallistaConfig::default().broadcast_join_threshold_bytes()/_rows(). The Ballista broadcast-threshold defaults are now the single source of truth for both DataFusion's built-in JoinSelection and Ballista's AQE join selection. Values are unchanged. * feat: demote DataFusion CollectLeft joins over the Ballista threshold In the static planner, maybe_promote_to_broadcast trusted any broadcast-safe HashJoinExec(CollectLeft) that DataFusion's JoinSelection produced. DataFusion decides CollectLeft from its own session threshold, which can exceed a runtime override of ballista.optimizer.broadcast_join_threshold_bytes. Demote such a join back to Partitioned when its build side is not under the current Ballista threshold (or when broadcasts are disabled with threshold 0), so the Ballista key is authoritative in the static path too. Null-aware anti joins are never demoted since they require CollectLeft. * docs: apply prettier formatting to AQE tuning-guide table * test: disable AQE broadcast via Ballista threshold in repartition tests The join-selection repartition tests forced the repartition path by setting DataFusion's hash_join_single_partition_threshold[_rows] to 0. AQE join selection now reads the broadcast cutoff from the Ballista config (broadcast_join_threshold_bytes), so those DataFusion keys no longer gate CollectLeft promotion and the small test tables were promoted to broadcast. Set the Ballista broadcast byte threshold to 0 in the helper, which disables CollectLeft promotion and restores the repartitioned plans the snapshots assert.
* feat: add BALLISTA_PROTOCOL_VERSION + k8s health probes Introduces a compile-time `BALLISTA_PROTOCOL_VERSION` that gates executor↔scheduler communication and adds `/healthz` and `/readyz` endpoints so live-upgrade rollouts have something for Kubernetes to observe. Scheduler rejects executors whose protocol version does not match with `FailedPrecondition`, an `info!` log, and a `ballista_scheduler_rejected_by_protocol_version_total` counter. Missing registration metadata on heartbeats is treated as a mismatch. Executor flips /readyz off on RPC failure and self-terminates after 5 consecutive heartbeat failures, so a mismatched pod crash-loops until its image is bumped rather than silently sitting idle. Health probes are always mounted regardless of the `rest-api` feature: scheduler on the existing axum router, executor on a new `--bind-health-port` (default 50053). Livez is deliberately dumb — scheduler flapping must not restart executors in a cascade. Fixes apache#2071. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * refactor: move executor health probe HTTP server into the binary The library was spinning up its own axum HTTP server on `--bind-health-port` from inside `start_executor_process`, adding a hard axum dep and a second health surface that library embedders had no way to disable short of forking the code. Health probes are a deployment concern, so push them to where the deployment lives: `ballista/executor/src/bin/main.rs`. The library keeps `ExecutorHealth` — a cheap `Arc<AtomicBool>` that the heartbeat loops flip on every RPC outcome — and exposes it on `ExecutorProcessConfig` so callers can observe it however they like. The axum server and its handlers are gated behind `#[cfg(feature = "build-binary")]`, and the `axum` dep moves from unconditional to `optional = true` under `build-binary`. With `--no-default-features`, `ballista-executor` no longer has a direct edge to axum in `cargo tree`. Verified end-to-end against the local native cluster: scheduler `/readyz` reports "ready: 2 executors registered" and both executors' `/healthz` and `/readyz` return 200, so the handle is threaded through bin -> config -> heartbeat loop -> probe correctly. Addresses apache#2088 (comment) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * fix: address CI failures on the protocol-version branch - examples/mtls-cluster.rs: add ballista_protocol_version to ExecutorRegistration and pass ExecutorHealth to poll_loop so the example compiles under `--all-targets` clippy. - docs/.../kubernetes.md: prettier reformatting — _emphasis_ instead of *emphasis*, and pad the rolling-upgrade table columns to match the longest row. - python/Cargo.lock: `cd python && cargo update` to resync the lockfile with the workspace manifests, so `cargo metadata --locked` inside `python/` succeeds again. The TPC-DS SF1 CI failure was infrastructure (`No space left on device` on the runner), not a code issue — re-running the job should clear it. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * feat: expose SchedulerServer::is_ready() for embedders Phillip's fork wraps the Ballista scheduler in a larger process whose health check is a superset of ours — his handler already reports app-level state and would like to AND the scheduler's readiness in as one factor. Add a public `SchedulerServer::is_ready() -> bool` that returns true when at least `min_ready_executors` executors have live heartbeats. The built-in `/readyz` handler now delegates to it (single source of truth), and embedders can call the same method from whatever health surface they expose. This is orthogonal to the `/healthz` and `/readyz` routes mounted on the axum server — those stay unconditional, so operators who don't embed still get the k8s probes for free. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * feat: expose ExecutorHealth::is_ready() for embedders Mirror what SchedulerServer::is_ready() offers on the executor side. The heartbeat loops already flip the internal atomic on every RPC outcome, and the built-in `/readyz` axum handler reads it via this method under `build-binary`. Making the method `pub` and dropping the cfg gate lets an embedder observe the same signal directly and fold it into their own health surface, without linking axum. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…efactor (apache#2038) * Substrate: multi-partition tasks + K-drain ShuffleWriter 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 dfe1f601 (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 * fix: size stage.partitions from the shuffle writer's child output_partitioning The leaf-sum walker undercounted after upstream #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> * fix(task_builder): route each UnionExec parent partition to its child'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 #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> * WIP: merge task_id / task_index into a single per-stage slot id Feedback from #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> * feat(scheduler): add ballista.scheduler.max_partitions_per_task cap Per andygrove's PR #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> * proto: reuse tag 3 for TaskId.global_output_partition_ids Feedback from #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> * fix: preserve per-partition metrics under multi-partition tasks Feedback from #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> * docs: explain multi-partition tasks, Spark comparison, and AQE composition Adds a section to docs/developer/architecture.md addressing PR #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> * docs: prettier — use _emphasis_ not *emphasis* CI's prettier@2.7.1 rewrites *foo* to _foo_. * docs: ASCII threading diagrams on ShuffleWriterExec and SortShuffleWriterExec 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)`. * fix: ShuffleWriter metrics — K buckets per task, keyed by operator-local 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 (fa54364a): the parameter is semantically an operator-local partition index, not a stage-wide task slot. The scheduler's d4f61711 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 8bd2e44d), verified against DataFusion. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * fix: cross-product under-collect metrics against task slice, restoring master-like per-partition UI Second bug surfaced after the ShuffleWriter K-buckets fix (b8f6ca89): 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> * fix(config): default ballista.scheduler.max_partitions_per_task to 1 (pre-branch model) Per @andygrove review on #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> * fix(scheduler): restrict PlaceholderRow/Empty leaves under fan-in splits `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 #2038 CI). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * test(scheduler): thread SessionConfig into aggregation-plan mock 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 bc7a4eda 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> --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…ess (apache#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 apache#2038. Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- Replace the callback's `&str` argument with a `WorkAvailableReason` enum - Fire the stage-resolution callback from the JobUpdated arm after `update_job` has revived the graph, so it only fires for genuinely new runnable work and never before the work is pollable - Add `SchedulerConfig::with_on_work_available` - Document that the callback runs synchronously in the scheduler event loop and must not block - Idle-poll less aggressively (1s fallback) when a poll_now notifier is wired
github-actions
Bot
added
documentation
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
10 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.
What
Adds two related hooks for push-style executor scheduling:
poll_now_notify: Option<Arc<Notify>>parameter on the executorpoll_loop. When provided, the idle wait becomes atokio::select!over the poll interval and the notify, so the scheduler can wake an idle executor immediately instead of waiting for the next poll tick.Nonepreserves the existing timer-only behavior.OnWorkAvailableFncallback onSchedulerConfig(on_work_available), invoked by the query stage scheduler when new work becomes available (afterJobSubmittedand when new stages become runnable).Ported from #12