feat: add a broadcast channel for job state event notifications - #52
Merged
Merged
Conversation
Adds a tokio::sync::broadcast Sender to SchedulerServer, subscribable via subscribe_job_updates(). On job state changes (queued/running/completed/ failed/cancelled) the scheduler broadcasts a JobStateEvent carrying the job id and new state. Ported from #15 for upstreaming.
Author
|
Upstreaming here apache#1891 |
…nnel - Add job_state_channel_capacity to SchedulerConfig (default 256) so operators with high job throughput can tune the buffer; removes the hardcoded const from SchedulerServer - Move all job state broadcasts to after their state transitions succeed, so subscribers never see an event for a commit that failed; aligns Queued/Failed/Completed/Cancelled with how Running already behaved - Add From<job_status::Status> for JobState to keep the two enums in sync; document that Cancelled has no protobuf counterpart - Expand module doc with guidance on when to use JobStateEvent vs JobStatusSubscriber, and reference job_state_channel_capacity for tuning - Add integration tests covering the broadcast wiring: lifecycle order (Queued -> Running -> Completed) and cancel path
github-actions
Bot
added
documentation
…2191) Bumps [taiki-e/install-action](https://github.com/taiki-e/install-action) from 2.85.0 to 2.85.2. - [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@7572810...41049aa) --- updated-dependencies: - dependency-name: taiki-e/install-action dependency-version: 2.85.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>
Bumps [aws-config](https://github.com/smithy-lang/smithy-rs) from 1.9.0 to 1.10.1. - [Release notes](https://github.com/smithy-lang/smithy-rs/releases) - [Changelog](https://github.com/smithy-lang/smithy-rs/blob/main/CHANGELOG.md) - [Commits](https://github.com/smithy-lang/smithy-rs/commits) --- updated-dependencies: - dependency-name: aws-config dependency-version: 1.10.1 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 [ctor](https://github.com/mmastrac/linktime) from 1.0.10 to 1.0.11. - [Release notes](https://github.com/mmastrac/linktime/releases) - [Changelog](https://github.com/mmastrac/linktime/blob/master/CHANGELOG.md) - [Commits](mmastrac/linktime@ctor-1.0.10...ctor-1.0.11) --- updated-dependencies: - dependency-name: ctor dependency-version: 1.0.11 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>
…ates (apache#2133) Bumps the github-codeql group with 2 updates in the / directory: [github/codeql-action/init](https://github.com/github/codeql-action) and [github/codeql-action/analyze](https://github.com/github/codeql-action). Updates `github/codeql-action/init` from 4.37.1 to 4.37.3 - [Release notes](https://github.com/github/codeql-action/releases) - [Changelog](https://github.com/github/codeql-action/blob/main/CHANGELOG.md) - [Commits](github/codeql-action@7188fc3...e4fba86) Updates `github/codeql-action/analyze` from 4.37.1 to 4.37.3 - [Release notes](https://github.com/github/codeql-action/releases) - [Changelog](https://github.com/github/codeql-action/blob/main/CHANGELOG.md) - [Commits](github/codeql-action@7188fc3...e4fba86) --- updated-dependencies: - dependency-name: github/codeql-action/analyze dependency-version: 4.37.2 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: github-codeql - dependency-name: github/codeql-action/init dependency-version: 4.37.2 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: github-codeql ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Bumps [astral-sh/setup-uv](https://github.com/astral-sh/setup-uv) from 8.3.2 to 9.0.0. - [Release notes](https://github.com/astral-sh/setup-uv/releases) - [Commits](astral-sh/setup-uv@11f9893...c771a70) --- updated-dependencies: - dependency-name: astral-sh/setup-uv dependency-version: 9.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> Co-authored-by: Martin Grigorov <martin-g@users.noreply.github.com>
…; log merged quantile cuts per stage (apache#2175) * proto: add RuntimeStatsReport wire format on SuccessfulTask Groundwork for shipping RuntimeStatsExec observations (row counts + quantile sketches) back to the scheduler for global-cut selection. Wire format only; behavior unchanged. The executor still emits `runtime_stats: vec![]` at every callsite. Slice B populates it. Proto additions (`ballista.proto`): * `SuccessfulTask.runtime_stats: repeated RuntimeStatsReport`. One entry per RuntimeStatsExec that is still valid at the plan's output (i.e., reachable from the top through distribution- preserving nodes only, per `preserves_distribution` in range_repartition_common). Pre-repartition stats stay local to the executor (used only for the repartitioner's own approximate routing) — the walker's whitelist naturally excludes them. * `RuntimeStatsReport { order_by, partitions }`. `order_by` tags the routing expression so the scheduler groups sketches across tasks that were sampling the same expression. `partitions` carries one entry per observed partition — pre-repartition one per input partition, post-repartition one per output sub- partition; the operator's position determines the interpretation. * `RuntimeStatsPartitionEntry { partition_id, row_count, optional sketch }`. A future `optional MinMaxState min_max` slot is called out as a TODO — the eventual lighter mode for post- repartition bin-pack where full T-Digests are overkill. Note: `QuantileSketchState` and the `sketch_to_proto` / `sketch_from_proto` helpers already landed in main with apache#2094 and are reused here without change. Constructor plumbing: * Five existing `SuccessfulTask { ... }` sites populated with `runtime_stats: vec![]` — the real executor path (executor/src/lib.rs), one scheduler-server integration-test path, and three scheduler test fixtures. Verification: * `cargo build --workspace` — clean. * `cargo test -p ballista-core --lib` — 166 pass (existing wire tests over `sketch_to_proto` / `sketch_from_proto` cover the payload; prost codegen refuses field-number clashes at build time so no new roundtrip test is warranted for a purely structural addition). * `cargo test -p ballista-scheduler --lib` — 269 pass. * `cargo test -p ballista-executor --lib` — 53 pass. * `cargo clippy --all-targets` — clean. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> executor: ship RuntimeStatsReports to scheduler; log on arrival End-to-end path for stage-1 quantile sketches — extract on the executor at task completion, transport in `SuccessfulTask.runtime_stats` (landed in the previous commit), and log a per-report summary on the scheduler as proof-of-life. Per-stage accumulation + merged-quantile-cut logging comes next. # Executor extraction (runtime_stats.rs) `collect_reports(plan)` walks the plan through the `preserves_distribution` whitelist and collects one `RuntimeStatsReport` per reachable `RuntimeStatsExec`. Stats-taps sitting below any distribution-changing operator (e.g. `UnorderedRangeRepartitionExec`) are excluded automatically because the walker stops at that boundary; their sketches describe data the repartitioner then routed away and are no longer meaningful at the plan's output. Each emitted report carries the operator's `order_by` (serialised as a `PhysicalSortExprNode` — the wire tag the scheduler will group on) and one `RuntimeStatsPartitionEntry` per partition slot with `{partition_id, row_count, optional sketch}`. Empty sketches are elided (`sketch: None`) so bin-pack space is proportional to actual samples, not slot count. `sketch_to_proto` / `sketch_from_proto`, `partition_count()`, and the widened `preserves_distribution` whitelist already exist in main — introduced with the RuntimeStatsExec / URRE PRs. This commit reuses them. # Trait plumbing (execution_engine.rs) `QueryStageExecutor::collect_runtime_stats_reports()` — default returns empty. `DefaultQueryStageExec` overrides to call `collect_reports` against whichever `ShuffleWriterVariant` it holds (`Passthrough` or `Sort`). Serialisation errors are logged and the report dropped rather than failing the task; the task's data was already produced correctly, telemetry loss shouldn't tank the query. # Transport (executor lib + call sites) `as_task_status` gains a `runtime_stats: Vec<RuntimeStatsReport>` parameter (marked `#[allow(clippy::too_many_arguments)]`) and populates `SuccessfulTask.runtime_stats` with it. Both executor call-sites (`execution_loop.rs`, `executor_server.rs`) call `collect_runtime_stats_reports()` on the query-stage executor and thread the result through. Failure paths pass `vec![]`. # Scheduler receive (task_manager.rs) `TaskManager::update_task_statuses` calls a new `log_runtime_stats_arrival` helper on each status before batching. Logged at `debug!` — RUST_LOG can promote it locally during verification. Fields: executor / job / stage / task + per-report summary (order_by_len, partition count, non-empty partitions, total rows, sketches). # Tests Four new walker tests in `execution_plans::runtime_stats::collect_tests`: * `collect_reports_finds_stats_and_ships_sketch` — stats at plan root, populated sketch survives `sketch_from_proto`, row_count matches drained data. * `collect_reports_row_count_only_emits_report_without_sketch` — row-count-only mode emits a report but no sketch payload. * `collect_reports_descends_through_whitelisted_op` — a `BufferExec` intermediary does not block the walker. * `collect_reports_stops_at_sort_that_collapses_partitions` — N→1 SortExec is explicitly excluded from the whitelist; the walker respects that. # Verification * `cargo test -p ballista-core --lib` — 170 pass (was 166). * `cargo test -p ballista-scheduler --lib` — 269 pass, 2 ignored. * `cargo test -p ballista-executor --lib` — 53 pass. * `cargo clippy --all-targets` — clean. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> scheduler: accumulate runtime-stats reports per stage; log merged cuts Barrier 1.5 prep. On the scheduler, each stage-attempt now accumulates the `RuntimeStatsReport`s that arrive with successful task statuses. When the stage transitions to Successful, the reports are grouped by `order_by` wire tag, sketches merged via `TDigest::merge_digests`, and one debug! line emitted per group with the `K-1` quantile cuts a globally-informed router would have used for `K` output partitions. `K` is read from the reports themselves — every report carries one partition entry per repartitioner-output slot, so `K == report.partitions.len()`. Reports with mismatched `K` within a group are surfaced as a `warn!` and the group is skipped (planner invariant break). # ballista_core::execution_plans::runtime_stats Two new symbols, exported from `execution_plans::mod`: * `merge_reports(reports) -> Vec<MergedRuntimeStats>`. Groups by `order_by` (prost-encoded bytes as HashMap key), merges T-Digests within each group, computes `K-1` cuts at quantiles `i/K`, returns `MergedRuntimeStats { order_by_len, k, task_count, total_rows, cuts, min, max }`. Empty input → empty output. Sketches with `count() == 0` are dropped from the merge input rather than corrupting the result. * `log_merged_runtime_stats(job_id, stage_id, reports)`. Calls `merge_reports` and formats each group at `debug!` — `RUST_LOG` can promote when verifying against a live cluster. Splits into two log lines depending on whether any sketches were present (drops `min` / `max` in row-count-only mode). `MergedRuntimeStats` is `pub` so slice-D consumers (an AQE rule that plans stage 2 from the cuts) can call `merge_reports` directly without paying the log-formatting cost. # ballista_scheduler::state::execution_stage `RunningStage` gains a `runtime_stats_reports: Vec<RuntimeStatsReport>` field, initialized empty in `RunningStage::new` and in the retry path `FailedStage::to_running` (a fresh attempt discards the previous attempt's stats — merged-cut logging fires per-attempt). New method `append_runtime_stats_reports(reports)` for the graph to push into on each successful task update. # ballista_scheduler::state::execution_graph In the `Successful` arm of `update_task_status`, destructure `SuccessfulTask` to grab both `partitions` and `runtime_stats` and push the reports onto the running stage. When `is_final_successful` flips true (stage transitions to Successful), call `log_merged_runtime_stats` alongside the existing `print_stage_metrics` hook. Reports are dropped on the floor after logging — slice D propagates them onto `SuccessfulStage` when the AQE consumer needs them at plan-time. # Tests Five new tests in `execution_plans::runtime_stats::merge_tests`: * `merge_reports_combines_disjoint_ranges` — two reports over disjoint value ranges; merged sketch spans the union, K=2 midpoint cut falls between the two ranges, total_rows sums. * `merge_reports_k_of_4_produces_three_quartile_cuts` — uniform [0, 100) sample over K=4; the three cuts land in the expected quartile bands. * `merge_reports_row_count_only_emits_empty_cuts` — reports with only row_counts (no sketches) sum correctly with empty `cuts` and `None` `min` / `max`. * `merge_reports_skips_group_with_mismatched_partition_counts` — mismatched `K` within a group drops the whole group. * `merge_reports_empty_input_is_empty_output` — degenerate case. # Verification * `cargo test -p ballista-core --lib` — 175 pass (was 170). * `cargo test -p ballista-scheduler --lib` — 269 pass, 2 ignored. * `cargo test -p ballista-executor --lib` — 53 pass. * `cargo clippy --all-targets` — clean. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> review: merge_reports returns Result; drop silent per-group fallbacks Address review feedback on the slice-C merge helper: * merge_reports is fallible now. Two failure modes it used to swallow silently — a group with mismatched partition counts, and a corrupted sketch that won't decode — both surface as errors. The scheduler-side caller (log_merged_runtime_stats) catches at warn! and returns; a slice-D consumer that actually plans off the merged cuts can propagate or fail the stage as it sees fit. If the function's job is to merge, it either merges or errors — no "returns empty because the third group was funky" middle ground. * The group-loop body moved into its own `merge_group(&[&Report]) -> Result<MergedRuntimeStats>` helper. The outer loop is one `?` per group; each group's own failure modes live in one place. * `MergedRuntimeStats.k` → `.partition_count`. Struct is public; slice D will consume it. `k` is a math-notation letter, not a field name. * Local single-letter names retired: `r1`/`r2` → `low_range` / `high_range`, `k`/`vs` inside `sketching_report` → `slot_id` / `slot_values`, `make` closure → `make_report`, closure params `|r|` → `|report|`. Per the max-specificity naming preference in memory. * No array accessors in the merge helper or its tests. The empty- group branch of `merge_group` uses `let [first, rest @ ..] = group` and returns an internal error rather than indexing at `[0]`. Tests unpack `merge_reports(...)?` results with `match cuts.as_slice()` and `only_group(...)` helper that panics with a real message rather than `groups[0]`. New test: `merge_reports_propagates_sketch_decode_errors` — builds a wire-shape-corrupt `QuantileSketchState` (3 scalars instead of 6) and asserts the decode error propagates through `merge_reports`. This was the case slice C's original `warn!` + `continue` was hiding. Existing test renamed `merge_reports_skips_group_with_mismatched_...` → `merge_reports_errors_on_mismatched_partition_counts` to reflect the new contract. Verification: * cargo test -p ballista-core --lib — 176 pass (was 175; +1 test). * cargo clippy --all-targets — clean. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> fix(docs): drop rustdoc links to items outside the doc'd crate CI's `cargo doc --document-private-items --no-deps --workspace` runs with `-D warnings`, which promotes two rustdoc lint classes to errors: * `rustdoc::private-intra-doc-links` — a public item's doc can't link to a private one. `collect_reports` (and its re-export `collect_runtime_stats_reports`) linked to `super::range_repartition_common::preserves_distribution`, but `range_repartition_common` is `mod`, not `pub mod`. * `rustdoc::broken-intra-doc-links` — a link target must resolve in the doc'd crate. `RunningStage::append_runtime_stats_reports` linked to `log_merged_runtime_stats`, which now lives in ballista-core (moved there in the slice-C refactor because ballista-scheduler doesn't depend on `datafusion_functions_aggregate_common`). Both docstrings switched from intra-doc `[...]` links to plain backtick prose that names the target function fully. The prose still points a reader to the right place; the compiler stops complaining. Verified locally with the exact CI invocation: `RUSTDOCFLAGS='-D warnings' cargo doc --document-private-items \ --no-deps --workspace` — clean. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> review: reshape as_task_status to take TaskCompletionExtras Fold operator_metrics + runtime_stats into a TaskCompletionExtras struct marked #[non_exhaustive] + Default, so future additions are non-breaking for external callers via ..Default::default(). Drops the too_many_arguments allow on as_task_status. Documents the break in docs/source/upgrading/55.0.0.md since the parameter list still changes shape (execution_times now precedes extras, and operator_metrics moves from a positional param into the struct). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * fix: AQE path forwards SuccessfulTask.runtime_stats to the scheduler `AdaptiveExecutionGraph::update_task_status` in `aqe/mod.rs:774` unpacked `SuccessfulTask.partitions` from an incoming task-success message but silently dropped the `runtime_stats` field. It also skipped `log_merged_runtime_stats` at final-success. The RSE-follow-up wire path was only plumbed on `StaticExecutionGraph`; any tap running under AQE mode (`ballista.planner.adaptive.enabled=true`) had its sketches discarded at the scheduler. Symptom on the adaptive-range-shuffle path: stage-1 tasks correctly harvested one `RuntimeStatsReport` each (visible in executor logs as `Task N finished with … 1 runtime-stats report(s)`), the wire message carried them, but nothing on the scheduler side ever consumed them — `log_merged_runtime_stats` never fired, so the merged cut points that downstream stages need weren't visible. The fix mirrors the static path (`execution_graph.rs:956`): unpack the `runtime_stats` field, hand it to `running_stage.append_runtime_stats_reports`, and call `log_merged_runtime_stats` once the stage is final-successful. Verified on Q20 SF10 — merged output now appears: ``` merged runtime stats: job=… stage=1 order_by_len=1 partition_count=16 task_count=4 total_rows=9088057 cuts=[125078.52, 249995.47, …, 1874842.86] min=1 max=2000000 ``` Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * key runtime-stats reports by producer task_id Wrap each incoming `RuntimeStatsReport` in a `TaskRuntimeStats { producer_task_id, report }` when the scheduler accumulates it on `RunningStage.runtime_stats_reports`. Threads the producer `task_id` through both accumulation paths (static- and adaptive-graph `update_task_status`). Motivation: reports are appended on every successful-task update, but a task's "success" can be invalidated later (executor loss → task marked `Failed(ResultLost)`, partition rescheduled). Without a producer key on the accumulated entries, the merged view at stage-final-success would double-count the reset producer's contribution — the old ghost plus the retry's fresh report. Keying by producer `task_id` gives the follow-up purge helper something to filter on. No behavior change here (purge lands in the next commit). `log_merged_runtime_stats` extracts the raw reports internally for the existing merge, so the logged output is unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * purge stale runtime-stats reports when a task is reset for retry `RunningStage.runtime_stats_reports` accumulates per-task sketches on every successful-task update, but a task's "success" is not permanent — an executor loss flips previously-successful tasks to `Failed(ResultLost)` and reschedules their partitions on another executor. Without a purge, the ghost reports stayed in the vec; the retry appended its fresh reports under a new `task_id`; the merged view at stage-final-success double-counted the bounced slices. Addresses reviewer feedback on this PR (Phillip LeBlanc, apache#2175): "Reports are not associated with task attempts. When successful work is invalidated and rerun (i.e. lost executor, and the scheduler requeues the task), the old report cannot be removed." They are now — each accumulated entry carries the producer `task_id`, so both reset paths on `RunningStage` filter out stale contributions: - `reset_task_info(task_id)`: retryable single-task reset (task killed / result lost). Drops entries where `producer_task_id == task_id`. - `RunningStage::reset_tasks(executor)`: whole-executor loss on a running stage. Collects the set of reset task_ids while flipping their statuses, then filters the reports vec against that set in one pass. `SuccessfulStage::reset_tasks` doesn't need matching treatment — `SuccessfulStage` doesn't carry the reports (they were consumed by `log_merged_runtime_stats` at stage-final-success), and `to_running` starts the next attempt with a fresh empty vec. Two unit tests cover both paths: single-task reset purges only the matching entry; executor-loss reset purges every entry from the lost executor while leaving surviving executors' entries alone. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…ts (apache#2186) * fix: stop every task running every union branch at low partition counts Per-task restriction gives each branch of a `UnionExec` the partitions that task owns and an empty slice to the rest, relying on an unowned branch becoming 0-partition so the union's index arithmetic routes past it. That fails when a branch's partition count does not come from restrictable leaves. A `CoalescePartitionsExec` or a non-preserving `SortExec` reports one partition whatever its leaves do, and a broadcast `ShuffleReaderExec` is deliberately never pruned. The branch keeps reporting a partition and keeps producing all of its data, so the executor -- which runs every partition of the plan it is handed -- runs every branch for every task. The result is a silently inflated answer, or a task executing a branch whose scan was emptied. At SF1 with `target_partitions=1`, TPC-DS q5 came back 4x and q80 2x, and q49 failed with `FileStreamBuilder invalid partition index`. It does not show at higher partition counts because those union stages are capped by a hash repartition, where the writer's output is an intrinsic K-space and inner partition counts do not matter. Cut a stage boundary beneath a union branch that cannot be restricted away. The branch becomes a `ShuffleReaderExec`, which restricts to nothing cleanly, and the collapse moves into its own stage where it still sees its whole input. Branches that already restrict away stay inline, so a union of plain scans is unaffected. Whether a branch can be restricted away is decided by asking the rewriter rather than by matching operator types, so it stays correct as the rewriter evolves. Branches holding a broadcast input or a `CollectLeft` join are treated as unrestrictable regardless: those build and probe sides can still be swapped after stage planning, and the two orders restrict differently, so a branch that looks restrictable at planning time may not be by the time the task is built. Also take restriction scope from `required_input_distribution()` for every operator instead of a list of join types. `CrossJoinExec` collects its left input but was not on that list, so each task of a multi-partition stage rebuilt the cross join's left side from its own slice. That surfaced as an inflated q77 (`from cs, cr`) once its branch got a stage of its own. Verified at SF1 against single-process DataFusion: the full suite passes at `--partitions 16`, and at `--partitions 1` only q72 remains, exhausting its per-task memory share rather than returning a wrong answer. The adaptive planner builds stages through its own path and is not covered by this change; it still fails these queries at one partition. * refactor(scheduler): tidy the union stage-boundary fix after review No behaviour change; verified by re-running TPC-DS SF1 at both partition counts (full suite passes at 16, only the pre-existing q72 memory failure at 1). - Reuse the generic child-recursion loop instead of repeating it in the UnionExec branch. The union case now rewrites the collected children in place, so the recursion contract lives in one spot. - Rename the branch predicate to `can_stay_inline`. It answered "can this be restricted away?" in its name while also folding in the broadcast guard, so the name promised less than the function did. `restricts_to_nothing` is now just the rewriter query it claims to be. - Share `scan_with_file_groups` from `test_utils` rather than keeping a near-identical copy in each of the two test modules. - Make the broadcast test exercise the guard. It asserted on a bare broadcast node, which the rewriter already reports as unrestrictable, so it passed with the guard deleted. It now checks detection through a nested branch, which is where the guard actually earns its keep. - Record why the `CoalescePartitionsExec` / `SortPreservingMergeExec` case cannot be folded into the rule below it: both declare `UnspecifiedDistribution`, so the general path would miss them.
Standalone operator carrying `Vec<Arc<dyn PhysicalExpr>>` — one boolean predicate per input partition — applied on `execute(k)`. Preserves vcore packing over the one-task-per-partition alternative, which would explode task count to K and defeat the range-repartition family. Not yet wired into any planner/rule; that lands in a follow-up slice alongside cuts-on-ExchangeExec and injection at the URRE-consuming stage. PerPartitionFilterExec: bounds-check partition and eager-release input on EOS Two small defensive tweaks to match DataFusion's FilterExec pattern: - `execute(partition)` now returns `internal_err!` if `partition` is out of range instead of panicking on the Vec index. The invariant is `partition < output_partitioning().partition_count()`, which construction guarantees against the predicate count — but ExecutionPlan callers indexing wrongly should get an error, not a panic. Uses `let-else` on `predicates.get(...).cloned()` so the guard and the clone happen in one step. - On end-of-stream in `poll_next`, replace `self.input` with an `EmptyRecordBatchStream` before returning `Poll::Ready(None)`. Mirrors DF's FilterExec (see `filter.rs:1027`) so the input's operator chain doesn't stay alive on the heap while the outer stream is still held. Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Null-aware LeftAnti joins cannot be swapped and DataFusion coordinates their visited-row and NULL state only within one process. Preserve CollectLeft, collect the build side, and coalesce the probe side so static and AQE plans run the join in exactly one task. Reject disabled or known oversized single-task builds through the Ballista broadcast threshold. Add optimizer, stage-lowering, threshold, and standalone distributed NOT IN regression tests. Closes apache#2187 Signed-off-by: Phillip LeBlanc <879445+phillipleblanc@users.noreply.github.com>
…embedded projections, and dropped exchange partitioning (apache#2200) * test(scheduler): show PropagateEmptyExecRule collapses GROUPING SETS(()) DataFusion's logical PropagateEmptyRelation guards on both `!agg.group_expr.is_empty()` AND `!has_empty_grouping_set(&agg.group_expr)`. The physical port here only carries the first guard, so a GROUPING SETS / ROLLUP / CUBE containing the empty subset `()` still collapses to EmptyExec over an empty input — losing the all-NULL row(s) the empty subset would emit. Adds a failing test that constructs `GROUPING SETS ((a), ())` over an EmptyExec and expects the aggregate to be preserved. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * fix(scheduler): harden PropagateEmptyExecRule against grouping sets, embedded projections, and dropped exchange partitioning Preserve aggregates whose grouping contains the empty grouping set () from GROUPING SETS/ROLLUP/CUBE, mirroring has_empty_grouping_set in DataFusion's logical PropagateEmptyRelation rule. On DataFusion 54 the empty subset emits one grand-total row even over zero input rows, so collapsing the aggregate to EmptyExec drops that row. Skip the join rewrites that keep one input alive (null-padded projection for Left/Right/Full, raw-child replacement for anti joins) when the hash join carries an embedded projection: the join schema then no longer maps positionally onto the input schemas, so null-padding reads the wrong columns and child replacement changes the schema. Arms that produce an EmptyExec use the join's own projected schema and remain unguarded. Replace a FilterExec over an empty input with an EmptyExec using the filter's schema instead of returning the input, since a filter with an embedded projection has a different schema than its input. Replace an ExchangeExec over an empty input with an EmptyExec that keeps the exchange's output partition count instead of returning the input, matching the stats-based exchange arm below it. --------- Co-authored-by: Brent Gardner <bgardner@squarelabs.net> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…ans (apache#2199) * feat(core): rewrite NOT IN subqueries into distributable anti join plans Add a NotInSubqueryRewrite logical optimizer rule that rewrites uncorrelated NOT IN (subquery) filter conjuncts into a plain LeftAnti join plus a one-row count aggregate before DataFusion's subquery decorrelation can plan a null-aware hash join. The null-aware join coordinates probe-side state through in-process atomics and therefore forces single-task execution in Ballista; the rewritten plan uses only ordinary joins and aggregates and distributes normally. The rule runs first in the logical optimizer chain and is registered in every Ballista session builder (client state creation and upgrade, and the scheduler's default session builder). It is enabled by default and controlled by ballista.optimizer.not_in_subquery_rewrite. The rewrite only fires where null-aware semantics would be needed: uncorrelated single-column subqueries with a nullable key on either side. Correlated, multi-column, volatile, and provably non-nullable predicates are left to DataFusion's own decorrelation. This also fixes wrong NOT IN results under Ballista's default sort-merge-join configuration, where the null-aware flag was previously dropped before scheduler lowering. * fix(core): use explicit path for intra-doc link in optimizer module docs The outer doc comment on 'pub mod optimizer' in lib.rs is merged with the module's inner docs, and rustdoc resolves the merged links in the parent scope where the bare item name is not visible.
…e paths on a real cluster (apache#2026)
…2203) Bumps [taiki-e/install-action](https://github.com/taiki-e/install-action) from 2.85.2 to 2.85.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@41049aa...18b1216) --- updated-dependencies: - dependency-name: taiki-e/install-action dependency-version: 2.85.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>
Bumps [rustls](https://github.com/rustls/rustls) from 0.23.42 to 0.23.43. - [Release notes](https://github.com/rustls/rustls/releases) - [Changelog](https://github.com/rustls/rustls/blob/main/CHANGELOG.md) - [Commits](rustls/rustls@v/0.23.42...v/0.23.43) --- updated-dependencies: - dependency-name: rustls dependency-version: 0.23.43 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 [http](https://github.com/hyperium/http) from 1.4.2 to 1.5.0. - [Release notes](https://github.com/hyperium/http/releases) - [Changelog](https://github.com/hyperium/http/blob/master/CHANGELOG.md) - [Commits](hyperium/http@v1.4.2...v1.5.0) --- updated-dependencies: - dependency-name: http dependency-version: 1.5.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 [ctor](https://github.com/mmastrac/linktime) from 1.0.11 to 1.0.12. - [Release notes](https://github.com/mmastrac/linktime/releases) - [Changelog](https://github.com/mmastrac/linktime/blob/master/CHANGELOG.md) - [Commits](mmastrac/linktime@ctor-1.0.11...ctor-1.0.12) --- updated-dependencies: - dependency-name: ctor dependency-version: 1.0.12 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>
* fix(cli): sort jobs by stage completion ratio Signed-off-by: QuakeWang <wangfuzheng0814@foxmail.com> * fix(cli): simplify stage completion sorting Signed-off-by: QuakeWang <wangfuzheng0814@foxmail.com> --------- Signed-off-by: QuakeWang <wangfuzheng0814@foxmail.com>
…2205) Bumps [taiki-e/install-action](https://github.com/taiki-e/install-action) from 2.85.3 to 2.85.4. - [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@18b1216...065d6a0) --- updated-dependencies: - dependency-name: taiki-e/install-action dependency-version: 2.85.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 [nix](https://github.com/nix-rust/nix) from 0.29.0 to 0.31.3. - [Changelog](https://github.com/nix-rust/nix/blob/master/CHANGELOG.md) - [Commits](nix-rust/nix@v0.29.0...v0.31.3) --- updated-dependencies: - dependency-name: nix dependency-version: 0.31.3 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 [reqwest](https://github.com/seanmonstar/reqwest) from 0.12.28 to 0.13.4. - [Release notes](https://github.com/seanmonstar/reqwest/releases) - [Changelog](https://github.com/seanmonstar/reqwest/blob/master/CHANGELOG.md) - [Commits](seanmonstar/reqwest@v0.12.28...v0.13.4) --- updated-dependencies: - dependency-name: reqwest dependency-version: 0.13.4 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>
Build replacement stages before updating the execution graph so failed plan rewrites leave the original stage intact. Add regression coverage for resolve and rollback transitions. Signed-off-by: QuakeWang <wangfuzheng0814@foxmail.com>
…d) (apache#2196) * feat(scheduler): range-repartition rule primitives (batteries-included) Scheduler-side AQE machinery so any optimizer rule that inserts an `UnorderedRangeRepartitionExec` or `OrderedRangeRepartitionExec` gets end-to-end correct routing for free. A range-partitioning rule now needs to do only one thing: splice a `RuntimeStatsExec` + URRE/ORRE above whatever it wants range-repartitioned. Everything downstream is handled: - Executor already ships per-sub-part quantile sketches to the scheduler (apache#2094 / apache#2175). - Scheduler recognizes the range repartition at stage-completion time, merges sketches into `K − 1` global cuts, and rewrites the downstream location map so partition `k` pulls from every producer file whose sketched `[min, max]` overlaps `k`'s assigned cut range. - Cuts + routing expression are parked on the boundary `ExchangeExec`; at task-specialization time the adapter wraps the `ShuffleReader` in a `PerPartitionFilterExec` (apache#2195) that trims straddling sub-parts to their assigned slice. Primitives a rule writer now has: - `plan_contains_range_repartition` / `find_range_repartition_routing_expr` (runtime_stats.rs) — detect URRE or ORRE anywhere in a plan subtree. - `compute_overlapping_locations` / `overlap_remap_partitions` (runtime_stats.rs) — merged sketches + cuts → per-downstream-partition location map with correct producer-file assignment. - `ExchangeExec.range_repartition_routing` slot + `resolve_*` / `.range_repartition_routing()` accessors — scheduler → adapter handoff. Mirrors the existing `coalesce: Arc<Mutex<Option<Arc<CoalescePlan>>>>` slot pattern (same lifecycle, same `with_new_children` carry-through, same idempotent-overwrite semantics — see `b839f036` /`232d7611` for the coalesce refactors that established the pattern). - DER classifier arm (distributed_exchange.rs) — recognizes `RuntimeStatsExec → URRE/ORRE` as a stage boundary so the sketches ship to the scheduler at stage-N completion. All of this fires only when a rule has actually inserted an RRE. On plans without one, `plan_contains_range_repartition` returns `false` and the machinery no-ops — zero overhead for the 99%+ non-RRE case. Test coverage (15 new unit tests, all pure-function): - `plan_walker_tests` (6): URRE and ORRE at root and nested, bare source returns `None`/`false`. - `overlap_remap_tests` (5): disjoint, straddling, missing-file_id error, default/report mismatch error, empty-sketch passthrough. - `range_repartition_routing_tests` (4): unresolved returns `None`, resolve→get roundtrip, second resolve overwrites, `with_new_children` preserves the slot. Full end-to-end wiring is exercised in the follow-up consumer rule PRs (AdaptiveRangeShuffleRule for join-side rewrites, ParallelWindowDetectRule for parallel windows) where a live URRE-inserting fixture is naturally available; deferring the ~300 LOC of AQE-graph scaffolding here in favor of pure-function coverage. Existing tests: 222 in ballista-core, 281 in ballista-scheduler — all pass. Clippy clean, fmt clean. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> docs: drop intra-doc link to private range_repartition_common module Rustdoc CI (`RUSTDOCFLAGS=-D warnings`) rejects public docs linking to private items. `range_repartition_common` is `mod` (not `pub mod`) and `split_batch_by_range` is `pub(super)`, so neither is reachable from the public rustdoc surface for `range_partition_predicates`. Convert the intra-doc link to plain text; the code cross-reference still reads for a curious reviewer grepping the tree. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> test(scheduler): cover the fourth-arm chain-top insertion and its idempotency The DER classifier's fourth arm inserts a passthrough `ExchangeExec(None)` above a `RuntimeStatsExec → URRE/ORRE` chain top, and its docstring claims a second pass over the resulting plan is a no-op — the freshly-inserted Exchange is caught by the `.downcast_ref::<ExchangeExec>().is_none()` outer guard, and from any upstream node's perspective an `ExchangeExec` is not a chain top (only `RuntimeStatsExec` is). Both halves of that claim are load-bearing (a bug in either would double-wrap the chain top on every AQE replan) but were only covered by the follow-up consumer-rule PRs. Two direct tests here: - `range_repartition_chain_top_gets_exchange_inserted` — verifies that arm 4 fires on `RoundRobin → RSE → URRE → leaf` and inserts the Exchange between the parent and the chain top, leaving the chain itself intact underneath. - `range_repartition_chain_top_insertion_is_idempotent` — runs the rule twice against the same starting plan and asserts the second pass adds no additional Exchange and produces an identical plan text. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> refactor: name single-letter loop vars in the range-repartition helpers `i` and `k` say nothing about what they index. In these helpers both are partition-space indices, so name them for what they are: - `k` → `partition_count` (in `range_partition_predicates`, `compute_overlapping_locations`) - `|i|` closure param → `|partition_idx|`; the inner `|j|` cut lookup → `|cut_idx|` - `partition_k` → `partition_idx` (matches the closure param above) - `out` → `overlaps_by_partition` (matches what the return doc calls it) - `smin` / `smax` → `sketch_min` / `sketch_max` - `spl` → `assignment` (it's one element of `assignments`) - `inner` → `remapped_bucket` - Plan walkers' `|c|` / `for c in ...` → `|child|` / `for child in ...` `lit`, `ge`, `lt`, `lo`, `hi` left as-is — those are established domain abbreviations, not opaque one-letter loop vars. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> refactor: trim chain-top comment, drop arm-numbering, tighten helpers - Cut the 17-line "Fourth classifier arm" block down to the one thing the code doesn't already say: why we insert a boundary here at all (executor's report walker collects sketches at stage boundaries). - Drop "arm 4 / arm 3" numbering in the test-section header, test docstrings, and assertion messages — positional numbering rots the moment someone reorders the classifier's else-if chain or adds a case. - Rename the fourth arm's `|c|` closure params to `|child|`. - Fully railroad `is_range_repartition_chain_top` — every branch has an explicit early return, no `||` chain at the tail. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> refactor: trim adapter comment, fail hard on missing range-repartition metadata Comment: - Cut the 13-line block above the PerPartitionFilterExec injection down to the one non-obvious why: straddling sub-parts + FinalPartitioned splitting partial sums. Log levels: - Downgrade the "injecting PerPartitionFilterExec" adapter trace, the "no non-empty sketches" fallback, and the happy-path overlap-remap summary from info! to debug!. All are per-range-repartition-stage events, useful for debug but not for baseline INFO output. Two of Brent's TODOs on the AQE hook flagged silent-degradation paths: - Empty `runtime_stats_reports` for a stage the planner detected as a range repartition: passthrough routing would silently misroute straddling rows. Now returns `Err` — no sketches means no cuts and we can't route. - `plan_contains_range_repartition` returns true but the routing-expr walker returns None: previously logged and continued with `routing = None`, which skipped the downstream filter injection and produced wrong aggregates. Now returns `Err` — this is a plan-shape invariant violation, not a soft-fallback case. Both errors propagate through `update_stage_progress` and fail the job rather than produce incorrect results. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> refactor: import proto types and tighten range-repartition remap fn Bring RuntimeStatsReport, RuntimeStatsPartitionEntry, and QuantileSketchState into runtime_stats.rs at file scope so signatures for merge_reports and kin drop the crate::serde::protobuf:: prefix. Bring the range-repartition helpers into aqe/mod.rs the same way; rename single-letter closure vars (t/m/c/e/v) and trim the error and debug messages to one line each. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> refactor: railroad plan-contains walker, return Result from routing-expr finder, rename overlap-remap arg to original_partitions Split plan_contains_range_repartition into two early-return ifs plus the trailing recursive .any(). Switch find_range_repartition_routing_expr from Option to Result: the outer public fn returns Result<Arc<dyn PhysicalExpr>>; recursion moves into a private find_routing_expr returning Result<Option<_>>, so "not found in this subtree" (Ok(None)) and "found but order_by is empty" (Err) stay distinguishable. Match-on-slice replaces the silent order_by()[0] panics. Rename overlap_remap_partitions' `default` argument (and the scheduler's local) to `original_partitions` — the semantics are "before remap", not "fallback default". Import PhysicalExpr and PartitionLocation into runtime_stats.rs to keep the affected signatures on one line. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> refactor refactor: split range-repartition routing recovery from remap, rename to cut_partitions Break maybe_range_repartition_overlap_remap into two orthogonal pieces: range_repartition_routing (recover cuts + routing expression from merged sketches; Ok(None) for legit passthrough cases, Err for invariant breaks) and the caller composing that with cut_partitions (previously overlap_remap_partitions). The old fn conflated plan-shape detection, sketch merging, and partition remap; now the caller uses plan_contains_range_repartition as a guard and slots in the routing where it belongs. Distinguish two "empty cuts" cases that used to collapse into one silent debug! + passthrough: total_rows == 0 (no data, passthrough is safe) and partition_count < 2 (K=1, single-partition range-repartition is degenerate but valid) both return Ok(None); any other empty-cuts case is a real sketch bug and now errors out. Also match on merged.as_slice() rather than merged.first() to catch the 2+ groups shape bug that first() silently ignored. Rename overlap_remap_partitions -> cut_partitions to match the "cut the K-space by these boundaries" mental model. Terse one-liner error messages throughout. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> refactor: collapse two range-repartition plan walkers into one plan_contains_range_repartition and find_range_repartition_routing_expr walked the same tree looking for URRE/ORRE, called back-to-back at the same call site. Merge into range_repartition_routing_expr returning Result<Option<Arc<dyn PhysicalExpr>>>: Ok(None) means no operator in the plan (single-mode absence), Err means one was found but its order_by was empty (invariant break). Caller uses .is_some() for the guard in the let-chain and passes the recovered expr into range_repartition_routing so the routing recovery no longer walks the plan itself. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> refactor: use .is::<T>() over .downcast_ref::<T>().is_some() Sweep the sites this PR introduced (across shuffle_writer.rs and distributed_exchange.rs) that reached for the boolean via the longer form. DataFusion provides is::<T: ExecutionPlan>() directly on dyn ExecutionPlan alongside downcast_ref, so the terse form is the right fit when we only need the check. Where the negation reads better as !.is::<T>() it replaces .downcast_ref().is_none(). Collapses several 6-line if-blocks (imports for UnorderedRangeRepartitionExec / OrderedRangeRepartitionExec at the top of shuffle_writer.rs let the ||-chain fit on one line) and lets the symmetric arms of is_range_repartition_chain_top stay as two one-liners each. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> comment cleanup less verbose refactor: single-pass cut_partitions, err on data loss instead of silent skip Drive cut_partitions from original_partitions directly; look up sketches by (task_id, sub_part_id) from a small map built once out of the reports. Deletes compute_overlapping_locations and SubPartLocation — the intermediate identity vector isn't needed once the walk is driven from the physical records. Files without a usable sketch (missing entry, or count == 0) are safe to skip only when partition_stats.num_rows == Some(0). Some(n > 0) or None now surface as an error rather than silently dropping data. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> perf(cut_partitions): binary-search cuts instead of scanning K buckets per file The inner K-loop over `remapped` made the walk O(F·K) — quadratic when the file count and output partition count are of similar scale (typical for a range shuffle). Since `global_cuts` is monotone and each file's sketch is a single [min, max] interval, the set of overlapping buckets is contiguous. Two `partition_point` calls locate `[b_lo, b_hi]` and we push into that slice, taking each file to O(log K + overlap). Adds a debug_assert on cuts monotonicity and a multi-cut test that exercises the range-walk boundaries (existing tests only had a single cut and couldn't catch an off-by-one). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> test(range-repartition): red regression tests for root-level chain top Two failing tests documenting the routing-park gap Andy flagged in PR apache#2196 review: when an RSE → URRE splice sits at the plan root, DER never wraps an ExchangeExec above it (transform_up only inspects children) and set_repartition_routing has nothing to park cuts on, so cut_partitions' straddler duplication reaches output_locations unfiltered. Both new tests assert the corrected behavior — exchange inserted, cuts parked — so they go red on current main and will go green when the fix lands. The FilterExec-parent sibling stays green as the positive control. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> fix(planner): set_repartition_routing errors when no ExchangeExec to park on The routing park silently returned Ok(()) when the stage's cached boundary wasn't an ExchangeExec, leaving cut_partitions' straddler duplication in place with no downstream filter to trim it — the silent outlier in a code path that otherwise fails loud on missing file_id, absent sketch, or bad merge_reports counts. Fail loud with a diagnostic error naming the actual boundary type instead. The one reachable trigger for this today (a range-repartition splice at the plan root) still leaves the URRE-at-root regression tests red; the DER-side fix that closes that shape lands next. Adds a synthetic bare-leaf regression test that pins the contract independent of DER: whenever the stage cache holds something other than an ExchangeExec, calling set_repartition_routing must error. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> fix(range-repartition): wrap root-level chain at DER, split plan-shape whitelist `DistributedExchangeRule::optimize` walks `transform_up`, which never visits the plan root as a child. A range-repartitioned splice at the root therefore got no boundary exchange wrapped above it, so `set_repartition_routing` had nothing to park cuts on. Wrap the root as a post-walk step, before the outer `AdaptiveDatafusionExec` goes on. Also splits the plan-shape whitelist out into a shared `plan_algebra` module: `preserves_distribution` stays where it was semantically (used by the sketch walker to descend past passthrough operators), and a looser `preserves_partitioning` sibling covers cases where only partition boundaries need to survive (rows and values within a partition are fair game). `is_range_repartitioned` drops its hardcoded `RuntimeStatsExec` check in favor of `preserves_partitioning` — the sketch tap is one of several partitioning-preserving nodes that can sit above a URRE/ORRE, and DF doesn't expose these as algebraic properties on `ExecutionPlan`. Flips the two red URRE-at-plan-root regression tests green; existing parent-based, idempotency, and control tests keep passing. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> test(range-repartition): drop RoundRobin-parent test, rewrite idempotency at root Andy's PR apache#2196 review flagged that the two DER-only tests exercised a shape the real optimizer pipeline strips: a `RoundRobinBatch` `RepartitionExec` parent that survives `rule.optimize()` in isolation but gets removed by other passes before DER runs in the pipeline. So the tests passed on a plan shape production never produces, and a rule author could write something that looks protected and lose the protection at runtime. - Delete `range_repartition_chain_top_gets_exchange_inserted`. The parent case is already covered end-to-end by `routing_parks_when_range_repartition_has_a_parent`, which uses a `FilterExec` parent (survives the pipeline) and drives through `AdaptivePlanner` instead of calling the rule alone. - Rewrite the idempotency test to use the root-level shape directly (`stats_over_urre_over_leaf()` with no artificial parent) — this is a shape production actually produces, and running `rule.optimize()` twice on it still validates DER's short-circuit for `AdaptiveDatafusionExec` inputs. Also drops the last of the "chain top" jargon from test names and docstrings. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> refactor(distributed-exchange): symmetric range-repart arm + railroad optimize Two orthogonal cleanups on the DER rule: - Collapse the range-repartition arm to the same shape as Coalesce/SPM: identify pattern, grab the (single) child, wrap in `ExchangeExec`, return via `with_new_children`. The old `.any()` + `.map()` over children was defending against a multi-child parent case that can't occur — URRE only ever sits below a stage boundary, so its parent is unary in practice. - Flatten `optimize()`: early-return for the idempotent `AdaptiveDatafusionExec` input case so the root-wrap logic reads as a linear sequence rather than a nested `if/else`. Adds a TODO on the range-repart arm noting the intended endgame — kill it entirely when `ExchangeExec` carries a range-cuts partitioning variant the way it carries `Partitioning::Hash`, at which point URRE can replace itself the way the Hash arm above already does. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> fix(distributed-exchange): move TODO to the correct arm The previous commit put the TODO on the `transform()` range-repart arm, but the arm the user wanted killed is the root-wrap block in `optimize()` — that's the one that goes away when we wrap in `AdaptiveDatafusionExec` up front and let `transform_up` visit the plan root as a child. Moving the TODO there with the specific unblock path (mechanical plan_id snapshot updates). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> fix(distributed-exchange): TODO points at Partitioning::Range endgame Correct my earlier reasoning: the wrap-first-in-AdaptiveDatafusionExec approach isn't the right endgame — it shifts plan_id numbering because the wrapper gets allocated before the inner ops instead of after. The actual endgame is the same as the `transform()` range-repart arm: `ExchangeExec` learning a range-cuts partitioning variant like it already carries `Partitioning::Hash`. Once URRE can replace itself with an ExchangeExec the way the Hash arm does, `transform_up`'s output already has an ExchangeExec at the range-repart position, the plain `AdaptiveDatafusionExec` wrap handles the root case with no ceremony, plan_id numbering stays unchanged, and both this root-wrap branch and the range-repart arm in `transform()` above collapse together. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> test(distributed-exchange): pin single-child scope with Union+SMJ probes Reviewer flagged that after the df951f3 simplification, DER's range-repart arm narrows to single-child parents (`let [child] = children.as_slice()`). A `UnionExec` or `SortMergeJoinExec` with a range-repartitioned chain on each side is silently left untouched, and the eventual runtime error ("no runtime-stats reports") points at stats collection rather than at the actual problem: the splice position isn't supported yet. Supporting multi-child parents isn't a code-only change — both sides need coordinated cuts (one cut set derived from all sides' sketches, parked on all boundary exchanges), which is cross-stage machinery we haven't built. The correct near-term move is to make the scope explicit, not to paper over it. - Doc-comment on `is_range_repartitioned` states single-child-only is deliberate and names the missing primitive. - `range_repartition_under_union_is_untouched` pins the Union case. - `range_repartition_under_sort_merge_join_is_untouched` pins the SMJ case with a TODO calling out cross-stage cut coordination as the motivating consumer for the eventual fix. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> fix(distributed-exchange): reject multi-child parent at plan time The previous commit pinned the single-child scope with tests but let the offending shape fall through to a runtime error at routing recovery ("no runtime-stats reports") — which points at stats collection rather than at the actual problem, exactly what Andy flagged in review. Turn DER's fall-through arm into a `match`: single-child parent with a range-repartitioned child still splices `ExchangeExec` as before; multi-child parent with any range-repartitioned child now returns a `plan_err!` naming the operator and pointing at cross-stage cut coordination as the missing primitive. Same functional outcome (the job fails), but the message now surfaces where the splice decision is made and tells a rule author precisely what shape isn't supported. The two pinning tests flip from asserting-no-op to asserting the error message + operator name; the doc-comment on `is_range_repartitioned` updates from "fails loudly at routing recovery" to "rejected at plan time." Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> style(distributed-exchange): match by slice shape, drop redundant guards The previous transform() fall-through arm mixed slice patterns with guards (`many if many.len() > 1 && ...` + a wildcard `_ => {}`), which was hard to read and had a redundant length check (the single-child case was already covered by an earlier arm). Recast as an exhaustive match on `[]` / `[child]` / `many`, with the range-repartition predicate as a plain `if` inside each arm body. No behavior change, 22 DER tests still green. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> PR feedback fix(distributed-exchange): reject ProjectionExec above range-repartition `is_range_repartitioned` gated on `preserves_partitioning`, which accepts `ProjectionExec` — but a projection between URRE/ORRE and the boundary reindexes, drops, or shadows the routing expression's referenced columns. The read-side `PerPartitionFilterExec` then evaluates against the wrong column and silently misroutes. Rename to `can_be_range_repartitioned`, return `Result<bool>`: `Ok(true)` splices as before, `Ok(false)` leaves alone, `Err(_)` when a `ProjectionExec` sits above a URRE/ORRE. `preserves_partitioning` stays as the algebraic property; the projection check is layered on top. TODO left in place for the arbitrary-routing-expression follow-up when we can walk the expression tree and prove survival. Ref: apache#2196 (comment) Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> fix(coalesce-rule): bail when leaf exchange has range-repartition routing Coalesce K and range-repartition K' (= cuts.len() + 1) can both land on the same `ExchangeExec` today: `CoalescePartitionsRule` collects every leaf and bails only on `broadcast`, so nothing stops the two slots from racing. When both are set, the adapter builds a reader with `cp.groups.len()` partitions but hands `cuts.len() + 1` predicates to `PerPartitionFilterExec::try_new` — hard error at plan time on the count mismatch. Mirror the broadcast bail: skip the entire alignment group when any leaf has `range_repartition_routing().is_some()`. Contiguous-group coalescing is compatible with range partitioning in principle (merge adjacent cut buckets alongside groups), tracked as apache#2220; explicit bail until that lands. Integration-level regression check parked next to `should_skip_coalesce_when_rule_disabled`: same happy-path inputs, routing slot parked before finalization, snapshot verifies no `coalesce=` on plan_id=0. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> * style(distributed-exchange): sort physical_plan imports alphabetically `cargo fmt --check` failed after the rebase pulled in `physical_plan::joins::{HashJoinExec, PartitionMode}` — rustfmt puts it before `projection`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Bumps the github-codeql group with 2 updates: [github/codeql-action/init](https://github.com/github/codeql-action) and [github/codeql-action/analyze](https://github.com/github/codeql-action). Updates `github/codeql-action/init` from 4.37.3 to 4.37.4 - [Release notes](https://github.com/github/codeql-action/releases) - [Changelog](https://github.com/github/codeql-action/blob/main/CHANGELOG.md) - [Commits](github/codeql-action@e4fba86...f205ea1) Updates `github/codeql-action/analyze` from 4.37.3 to 4.37.4 - [Release notes](https://github.com/github/codeql-action/releases) - [Changelog](https://github.com/github/codeql-action/blob/main/CHANGELOG.md) - [Commits](github/codeql-action@e4fba86...f205ea1) --- updated-dependencies: - dependency-name: github/codeql-action/init dependency-version: 4.37.4 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: github-codeql - dependency-name: github/codeql-action/analyze dependency-version: 4.37.4 dependency-type: direct:production update-type: version-update:semver-patch dependency-group: github-codeql ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
…eduler (apache#2225) * fix(executor): bind the gRPC listener before registering with the scheduler The executor spawned its tonic server and registered with the scheduler immediately afterwards. Server::serve binds inside the future it returns, so the socket did not exist yet, while the scheduler's registration handler dials that port back via a single connect with no retry. Losing that race gave ECONNREFUSED, failed registration, and killed the executor at startup. Bind the listener on the current task before spawning, and serve from it with serve_with_incoming_shutdown, so registration cannot run before the port accepts connections. tonic ignores the builder's tcp_nodelay and tcp_keepalive when serving from a pre-bound listener, so the new create_grpc_server_incoming helper applies them itself. * refactor: tighten create_grpc_server_incoming docs and tests * test: restore port-in-use test for create_grpc_server_incoming
Bumps [bytesize](https://github.com/bytesize-rs/bytesize) from 2.4.2 to 2.7.0. - [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.2...bytesize-v2.7.0) --- updated-dependencies: - dependency-name: bytesize dependency-version: 2.7.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>
…2228) Bumps [taiki-e/install-action](https://github.com/taiki-e/install-action) from 2.85.4 to 2.85.7. - [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@065d6a0...67729d5) --- updated-dependencies: - dependency-name: taiki-e/install-action dependency-version: 2.85.7 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#2221) Bumps [cryptography](https://github.com/pyca/cryptography) from 48.0.1 to 50.0.0. - [Changelog](https://github.com/pyca/cryptography/blob/main/CHANGELOG.rst) - [Commits](pyca/cryptography@48.0.1...50.0.0) --- updated-dependencies: - dependency-name: cryptography dependency-version: 50.0.0 dependency-type: indirect ... Signed-off-by: dependabot[bot] <support@github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
lukekim
approved these changes
Aug 12, 2026
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
8 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a
tokio::sync::broadcastsender toSchedulerServer, subscribable viasubscribe_job_updates(). On job state changes (queued / running / completed / failed / cancelled) the scheduler broadcasts aJobStateEventcarrying the job id and new state. Lets components holding a handle to the scheduler observe job lifecycle transitions without polling.Ported from #15.
Relationship to the existing JobStatusSubscriber
Similar to the existing
JobStatusSubscriber(mpsc) mechanism wired throughsubmit_job(..., subscriber). This broadcast channel is a separate, fan-out style notification path; worth deciding the best DX between these.JobStatusSubscriber(existing)JobStateEventbroadcast (this PR)mpsc- single consumerbroadcast- fan-outExecuteQuery/GetJobStatusResult.subscribe_job_updates()JobStatusJobStateEvent { job_id: String, state }Both fire from the same lifecycle points in
query_stage_scheduler.rs.