2196: feat(scheduler): range-repartition rule primitives (batteries-included) - #88
2196: feat(scheduler): range-repartition rule primitives (batteries-included)#88martin-augment wants to merge 3 commits into
Conversation
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>
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>
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
WalkthroughThe change adds half-open range predicate generation and runtime-statistics helpers for detecting range repartition plans, recovering routing expressions, computing sketch overlaps, and remapping partition locations. AQE exchanges and planners now carry recovered cuts and expressions through stage completion and reader transformation. Shuffle writers recognize range repartition outputs as K-space, distributed exchange handling adds required boundaries, and task restriction slices predicates with selected partitions. Review guidance files also receive new review rules. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Review: #88 — range-repartition rule primitives
|
🤖 Augment PR SummarySummary: This PR adds “batteries-included” primitives for adaptive range-repartition routing in Ballista’s scheduler/execution-plan layer. Changes:
Technical Notes: This wires up overlap-based routing metadata propagation (cuts + routing expr) and downstream filtering to preserve correctness when producer-local range bins straddle globally-assigned cut ranges. 🤖 Was this summary useful? React with 👍 or 👎 |
![augmentcode[bot]](https://avatars.githubusercontent.com/in/1027498?s=60&v=4){kind=small}
| // downstream filters need to reference. Missing it here is a | ||
| // plan-shape bug (range repartition detected by | ||
| // `plan_contains_range_repartition` but the walker can't find it) — | ||
| // surface as the ordinary passthrough case so the query still |
There was a problem hiding this comment.
maybe_range_repartition_overlap_remap: the comment says missing a routing expr should "surface as the ordinary passthrough case", but the function still returns remapped while downstream filters won’t be injected. That edge case seems likely to duplicate rows across downstream partitions (overlap mapping without PerPartitionFilterExec) rather than behaving like passthrough.
Severity: medium
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
| let new_child = restrict(child, partitions, false)?; | ||
| let new_predicates: Vec<_> = partitions | ||
| .iter() | ||
| .map(|&p| ppf.predicates()[p].clone()) |
There was a problem hiding this comment.
restrict()’s PerPartitionFilterExec handling indexes ppf.predicates()[p] directly; if partitions ever contains an out-of-range index, this will panic the scheduler instead of returning an error. Even if callers “should” pass valid partition IDs, it may be worth guarding to keep task specialization robust to upstream bookkeeping bugs.
Severity: medium
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (4)
ballista/scheduler/src/state/aqe/execution_plan/exchange.rs (1)
348-350: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
TreeRenderoutput omits the routing info the default formatter now prints. Cheap to mirror so EXPLAIN tree mode doesn't hide that a stage is range-routed.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@ballista/scheduler/src/state/aqe/execution_plan/exchange.rs` around lines 348 - 350, Update the TreeRender formatting logic in the execution plan display to include range repartition routing information, matching the default formatter’s output. Reuse self.range_repartition_routing and report the routing cut count when present, alongside the existing stage details.ballista/core/src/execution_plans/shuffle_writer.rs (1)
76-84: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueEnum doc lists only
UnorderedRangeRepartitionExec; the arm at Line 144 also matchesOrderedRangeRepartitionExec.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@ballista/core/src/execution_plans/shuffle_writer.rs` around lines 76 - 84, Update the KSpace enum documentation to also identify OrderedRangeRepartitionExec as a qualifying range-routed operator, matching the behavior handled by the corresponding match arm while preserving the existing description of partition indexing.ballista/scheduler/src/state/aqe/mod.rs (1)
322-334: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueFull deep-clone of every runtime-stats report just to satisfy
merge_reports' slice signature. For wide stages this copies every serialized T-Digest. Changingmerge_reportsto take&[&RuntimeStatsReport](or an iterator) would drop the copy.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@ballista/scheduler/src/state/aqe/mod.rs` around lines 322 - 334, Update merge_runtime_stats_reports and its callers to accept borrowed RuntimeStatsReport references (such as &[&RuntimeStatsReport] or an iterator) instead of owned reports. In the range-repartition remap flow around runtime_stats_reports, remove the cloned raw Vec and pass borrowed report references directly while preserving the existing merge error handling.ballista/core/src/execution_plans/per_partition_filter.rs (1)
297-337: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueNaN routing values fall through every predicate.
All comparisons against NaN evaluate false, so a NaN row matches no partition and is silently dropped downstream (the write side would have routed it somewhere via
partition_point). Worth documenting as an explicit non-goal or handling in the last bucket if NaN can reach the routing expression.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@ballista/core/src/execution_plans/per_partition_filter.rs` around lines 297 - 337, Update range_partition_predicates so NaN routing values are assigned to the final bucket, matching the write-side partition_point behavior. Ensure the last predicate explicitly accepts NaN in addition to its existing lower-bound condition, while preserving all other bucket predicates and documenting the handling if needed.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@ballista/core/src/execution_plans/runtime_stats.rs`:
- Around line 985-1015: Update the report-partition routing loop around
entry.sketch and sketch.count() to handle entries with row_count > 0 but no
usable sketch: do not silently skip them; either return an error or route the
sub-part to every output partition so downstream filtering can trim it. Preserve
skipping genuinely empty entries, and ensure zero-count or absent sketches with
positive row counts follow the selected fallback behavior.
In `@ballista/scheduler/src/state/aqe/adapter.rs`:
- Around line 124-136: Guard range-repartition routing in the reader
transformation before calling PerPartitionFilterExec::try_new: apply the
predicates only when the reader’s partition count matches predicates.len(). For
coalesced or broadcast reader variants, skip filtering with a warning (or return
an explicit error naming the incompatible combination) instead of allowing the
constructor’s generic count-mismatch failure. Keep the existing filtered path
for compatible plain readers.
In `@ballista/scheduler/src/state/aqe/mod.rs`:
- Around line 338-346: Validate the recovered cuts in the range-repartition flow
before returning them from the merged sketch selection: ensure the sequence is
monotonically non-decreasing, including handling invalid values such as NaN
according to the existing ordering semantics. If validation fails, log the same
passthrough fallback and return Ok((default, None)); otherwise preserve the
current cloned-cuts path.
In `@ballista/scheduler/src/state/aqe/planner.rs`:
- Around line 212-234: Move the stale “Resolves a stage by its ID and updates
its partitions” documentation from set_range_repartition_routing to
finalise_stage_internal, where its description and partitions argument apply. In
set_range_repartition_routing, add a debug log when the stage is missing or is
not an ExchangeExec, while preserving the existing successful routing resolution
and no-op return behavior.
In `@ballista/scheduler/src/state/task_builder.rs`:
- Around line 95-106: Update the PerPartitionFilterExec handling in restrict to
perform a checked lookup on ppf.predicates() for each partition index instead of
indexing directly. Convert a missing predicate into the existing internal plan
error via internal_err!, while preserving cloning and collection of valid
predicates.
---
Nitpick comments:
In `@ballista/core/src/execution_plans/per_partition_filter.rs`:
- Around line 297-337: Update range_partition_predicates so NaN routing values
are assigned to the final bucket, matching the write-side partition_point
behavior. Ensure the last predicate explicitly accepts NaN in addition to its
existing lower-bound condition, while preserving all other bucket predicates and
documenting the handling if needed.
In `@ballista/core/src/execution_plans/shuffle_writer.rs`:
- Around line 76-84: Update the KSpace enum documentation to also identify
OrderedRangeRepartitionExec as a qualifying range-routed operator, matching the
behavior handled by the corresponding match arm while preserving the existing
description of partition indexing.
In `@ballista/scheduler/src/state/aqe/execution_plan/exchange.rs`:
- Around line 348-350: Update the TreeRender formatting logic in the execution
plan display to include range repartition routing information, matching the
default formatter’s output. Reuse self.range_repartition_routing and report the
routing cut count when present, alongside the existing stage details.
In `@ballista/scheduler/src/state/aqe/mod.rs`:
- Around line 322-334: Update merge_runtime_stats_reports and its callers to
accept borrowed RuntimeStatsReport references (such as &[&RuntimeStatsReport] or
an iterator) instead of owned reports. In the range-repartition remap flow
around runtime_stats_reports, remove the cloned raw Vec and pass borrowed report
references directly while preserving the existing merge error handling.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 08b72a0a-b736-4a23-9e9d-ffc85f9fca18
📒 Files selected for processing (14)
.cursor/rules.md.gemini/rules.mdAGENTS.mdCLAUDE.mdballista/core/src/execution_plans/mod.rsballista/core/src/execution_plans/per_partition_filter.rsballista/core/src/execution_plans/runtime_stats.rsballista/core/src/execution_plans/shuffle_writer.rsballista/scheduler/src/state/aqe/adapter.rsballista/scheduler/src/state/aqe/execution_plan/exchange.rsballista/scheduler/src/state/aqe/mod.rsballista/scheduler/src/state/aqe/optimizer_rule/distributed_exchange.rsballista/scheduler/src/state/aqe/planner.rsballista/scheduler/src/state/task_builder.rs
| for stats in reports { | ||
| for entry in &stats.report.partitions { | ||
| let Some(sketch_proto) = entry.sketch.as_ref() else { | ||
| continue; | ||
| }; | ||
| let sketch = sketch_from_proto(sketch_proto)?; | ||
| if sketch.count() == 0.0 { | ||
| continue; | ||
| } | ||
| let smin = sketch.min(); | ||
| let smax = sketch.max(); | ||
| for (partition_k, bucket) in out.iter_mut().enumerate() { | ||
| let lower = if partition_k == 0 { | ||
| f64::NEG_INFINITY | ||
| } else { | ||
| global_cuts[partition_k - 1] | ||
| }; | ||
| let upper = if partition_k == k - 1 { | ||
| f64::INFINITY | ||
| } else { | ||
| global_cuts[partition_k] | ||
| }; | ||
| if smax >= lower && smin < upper { | ||
| bucket.push(SubPartLocation { | ||
| producer_task_id: stats.producer_task_id, | ||
| sub_part_id: entry.partition_id, | ||
| }); | ||
| } | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Sub-parts with rows but no sketch are silently dropped from every downstream partition.
entry.sketch == None is treated as "nothing to route", but the entry's row_count isn't consulted. If a producer ever reports a non-zero row_count with an absent/zero-count sketch (mixed sketch/row-count-only reports, or an encode failure upstream), those rows disappear from the query result with no error and no log. Consider erroring — or at minimum routing such a sub-part to every partition (the filter would then trim it) — when row_count > 0 and no usable sketch is present.
🛡️ Suggested guard
let Some(sketch_proto) = entry.sketch.as_ref() else {
+ if entry.row_count > 0 {
+ return internal_err!(
+ "range-repartition overlap: producer {} sub-part {} reported \
+ {} rows with no sketch — cannot route without a value range",
+ stats.producer_task_id,
+ entry.partition_id,
+ entry.row_count
+ );
+ }
continue;
};
let sketch = sketch_from_proto(sketch_proto)?;
- if sketch.count() == 0.0 {
+ if sketch.count() == 0.0 {
continue;
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| for stats in reports { | |
| for entry in &stats.report.partitions { | |
| let Some(sketch_proto) = entry.sketch.as_ref() else { | |
| continue; | |
| }; | |
| let sketch = sketch_from_proto(sketch_proto)?; | |
| if sketch.count() == 0.0 { | |
| continue; | |
| } | |
| let smin = sketch.min(); | |
| let smax = sketch.max(); | |
| for (partition_k, bucket) in out.iter_mut().enumerate() { | |
| let lower = if partition_k == 0 { | |
| f64::NEG_INFINITY | |
| } else { | |
| global_cuts[partition_k - 1] | |
| }; | |
| let upper = if partition_k == k - 1 { | |
| f64::INFINITY | |
| } else { | |
| global_cuts[partition_k] | |
| }; | |
| if smax >= lower && smin < upper { | |
| bucket.push(SubPartLocation { | |
| producer_task_id: stats.producer_task_id, | |
| sub_part_id: entry.partition_id, | |
| }); | |
| } | |
| } | |
| } | |
| } | |
| for stats in reports { | |
| for entry in &stats.report.partitions { | |
| let Some(sketch_proto) = entry.sketch.as_ref() else { | |
| if entry.row_count > 0 { | |
| return internal_err!( | |
| "range-repartition overlap: producer {} sub-part {} reported \ | |
| {} rows with no sketch — cannot route without a value range", | |
| stats.producer_task_id, | |
| entry.partition_id, | |
| entry.row_count | |
| ); | |
| } | |
| continue; | |
| }; | |
| let sketch = sketch_from_proto(sketch_proto)?; | |
| if sketch.count() == 0.0 { | |
| continue; | |
| } | |
| let smin = sketch.min(); | |
| let smax = sketch.max(); | |
| for (partition_k, bucket) in out.iter_mut().enumerate() { | |
| let lower = if partition_k == 0 { | |
| f64::NEG_INFINITY | |
| } else { | |
| global_cuts[partition_k - 1] | |
| }; | |
| let upper = if partition_k == k - 1 { | |
| f64::INFINITY | |
| } else { | |
| global_cuts[partition_k] | |
| }; | |
| if smax >= lower && smin < upper { | |
| bucket.push(SubPartLocation { | |
| producer_task_id: stats.producer_task_id, | |
| sub_part_id: entry.partition_id, | |
| }); | |
| } | |
| } | |
| } | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@ballista/core/src/execution_plans/runtime_stats.rs` around lines 985 - 1015,
Update the report-partition routing loop around entry.sketch and sketch.count()
to handle entries with row_count > 0 but no usable sketch: do not silently skip
them; either return an error or route the sub-part to every output partition so
downstream filtering can trim it. Preserve skipping genuinely empty entries, and
ensure zero-count or absent sketches with positive row counts follow the
selected fallback behavior.
| if let Some(routing) = exchange.range_repartition_routing() { | ||
| let predicates = | ||
| range_partition_predicates(routing.routing_expr, &routing.cuts); | ||
| info!( | ||
| "range-repartition: injecting PerPartitionFilterExec above \ | ||
| ShuffleReader for stage {} — {} predicates over {} cuts", | ||
| stage_id, | ||
| predicates.len(), | ||
| routing.cuts.len(), | ||
| ); | ||
| let filtered = PerPartitionFilterExec::try_new(reader, predicates)?; | ||
| return Ok(Transformed::yes(Arc::new(filtered))); | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Routing is applied regardless of which reader variant was built — coalesced/broadcast readers will hard-fail here.
PerPartitionFilterExec::try_new requires predicates.len() == reader.partition_count(). That holds only for the plain (None, false) reader, where the remapped location vector is cuts.len() + 1 long. With a CoalescePlan the reader exposes cp.groups.len() partitions, and a broadcast reader exposes 1 — either combination turns into an internal error that fails the job rather than degrading to unfiltered reads. Worth an explicit guard (skip + warn, or error with a message naming the incompatible combination) instead of relying on the constructor's generic count mismatch.
🛡️ Suggested guard
- if let Some(routing) = exchange.range_repartition_routing() {
+ if let Some(routing) = exchange.range_repartition_routing() {
+ if exchange.coalesce().is_some() || exchange.broadcast {
+ return internal_err!(
+ "range-repartition routing on stage {stage_id} is not \
+ supported with coalesced/broadcast shuffle readers"
+ );
+ }
let predicates =
range_partition_predicates(routing.routing_expr, &routing.cuts);📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if let Some(routing) = exchange.range_repartition_routing() { | |
| let predicates = | |
| range_partition_predicates(routing.routing_expr, &routing.cuts); | |
| info!( | |
| "range-repartition: injecting PerPartitionFilterExec above \ | |
| ShuffleReader for stage {} — {} predicates over {} cuts", | |
| stage_id, | |
| predicates.len(), | |
| routing.cuts.len(), | |
| ); | |
| let filtered = PerPartitionFilterExec::try_new(reader, predicates)?; | |
| return Ok(Transformed::yes(Arc::new(filtered))); | |
| } | |
| if let Some(routing) = exchange.range_repartition_routing() { | |
| if exchange.coalesce().is_some() || exchange.broadcast { | |
| return internal_err!( | |
| "range-repartition routing on stage {stage_id} is not \ | |
| supported with coalesced/broadcast shuffle readers" | |
| ); | |
| } | |
| let predicates = | |
| range_partition_predicates(routing.routing_expr, &routing.cuts); | |
| info!( | |
| "range-repartition: injecting PerPartitionFilterExec above \ | |
| ShuffleReader for stage {} — {} predicates over {} cuts", | |
| stage_id, | |
| predicates.len(), | |
| routing.cuts.len(), | |
| ); | |
| let filtered = PerPartitionFilterExec::try_new(reader, predicates)?; | |
| return Ok(Transformed::yes(Arc::new(filtered))); | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@ballista/scheduler/src/state/aqe/adapter.rs` around lines 124 - 136, Guard
range-repartition routing in the reader transformation before calling
PerPartitionFilterExec::try_new: apply the predicates only when the reader’s
partition count matches predicates.len(). For coalesced or broadcast reader
variants, skip filtering with a warning (or return an explicit error naming the
incompatible combination) instead of allowing the constructor’s generic
count-mismatch failure. Keep the existing filtered path for compatible plain
readers.
| let Some(cuts) = merged.iter().map(|m| &m.cuts).find(|c| !c.is_empty()) else { | ||
| info!( | ||
| "range-repartition stage {} produced no non-empty sketches; \ | ||
| falling back to passthrough partition mapping", | ||
| stage_id, | ||
| ); | ||
| return Ok((default, None)); | ||
| }; | ||
| let cuts = cuts.clone(); |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟡 Minor | ⚡ Quick win
No monotonicity check on the recovered cuts.
RangeRepartitionRouting documents cuts as monotone, and range_partition_predicates depends on it for exhaustive, disjoint coverage. If estimate_quantile ever returns a non-increasing pair (degenerate digest, NaN), some values match no predicate and rows are silently dropped. A cheap is_sorted-style check here that falls back to passthrough would keep the failure mode benign.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@ballista/scheduler/src/state/aqe/mod.rs` around lines 338 - 346, Validate the
recovered cuts in the range-repartition flow before returning them from the
merged sketch selection: ensure the sequence is monotonically non-decreasing,
including handling invalid values such as NaN according to the existing ordering
semantics. If validation fails, log the same passthrough fallback and return
Ok((default, None)); otherwise preserve the current cloned-cuts path.
| /// Attaches range-repartition-recovered range boundaries to the | ||
| /// boundary `ExchangeExec` for `stage_id`. Called right after the | ||
| /// completed range-repartition stage's partition mapping is resolved; | ||
| /// the routing carries the cuts + routing expression that downstream | ||
| /// task specialization needs to build per-partition range filters. | ||
| /// | ||
| /// No-op if the stage's boundary root is an `AdaptiveDatafusionExec` | ||
| /// (query root) — a range repartition writes at a shuffle boundary, so | ||
| /// the parking slot only exists on the `ExchangeExec` variant. | ||
| pub(super) fn set_range_repartition_routing( | ||
| &mut self, | ||
| stage_id: usize, | ||
| routing: RangeRepartitionRouting, | ||
| ) -> common::Result<()> { | ||
| if let Some(exchange) = self | ||
| .runnable_stage_cache | ||
| .get(&stage_id) | ||
| .and_then(|stage| stage.downcast_ref::<ExchangeExec>()) | ||
| { | ||
| exchange.resolve_range_repartition_routing(routing); | ||
| } | ||
| Ok(()) | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Stale doc block above this method.
The pre-existing comment at lines 204-211 ("Resolves a stage by its ID and updates its partitions", with a * partitions argument) is now attached to set_range_repartition_routing, which takes no partitions. It belongs with finalise_stage_internal below. Also consider a debug! on the miss path — a silent no-op here is exactly the case where downstream filters never appear.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@ballista/scheduler/src/state/aqe/planner.rs` around lines 212 - 234, Move the
stale “Resolves a stage by its ID and updates its partitions” documentation from
set_range_repartition_routing to finalise_stage_internal, where its description
and partitions argument apply. In set_range_repartition_routing, add a debug log
when the stage is missing or is not an ExchangeExec, while preserving the
existing successful routing resolution and no-op return behavior.
| if !under_collect && let Some(ppf) = plan.downcast_ref::<PerPartitionFilterExec>() { | ||
| let child = plan.children()[0].clone(); | ||
| let new_child = restrict(child, partitions, false)?; | ||
| let new_predicates: Vec<_> = partitions | ||
| .iter() | ||
| .map(|&p| ppf.predicates()[p].clone()) | ||
| .collect(); | ||
| return Ok(Arc::new(PerPartitionFilterExec::try_new( | ||
| new_child, | ||
| new_predicates, | ||
| )?)); | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Unchecked index into predicates can panic the scheduler's task builder.
ppf.predicates()[p] assumes every requested global index is in range. If a PerPartitionFilterExec ever ends up at a depth where partitions was rebased for a different width (e.g. beneath a union split, or a mismatch between the reader's K and the recovered cut count), this panics inside task construction instead of surfacing as a plan error. .get(p) + internal_err! keeps the failure recoverable.
🛡️ Checked lookup
- let new_predicates: Vec<_> = partitions
- .iter()
- .map(|&p| ppf.predicates()[p].clone())
- .collect();
+ let new_predicates: Vec<_> = partitions
+ .iter()
+ .map(|&p| {
+ ppf.predicates().get(p).cloned().ok_or_else(|| {
+ DataFusionError::Internal(format!(
+ "PerPartitionFilterExec: partition {p} out of range \
+ ({} predicates)",
+ ppf.predicates().len()
+ ))
+ })
+ })
+ .collect::<Result<Vec<_>>>()?;📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if !under_collect && let Some(ppf) = plan.downcast_ref::<PerPartitionFilterExec>() { | |
| let child = plan.children()[0].clone(); | |
| let new_child = restrict(child, partitions, false)?; | |
| let new_predicates: Vec<_> = partitions | |
| .iter() | |
| .map(|&p| ppf.predicates()[p].clone()) | |
| .collect(); | |
| return Ok(Arc::new(PerPartitionFilterExec::try_new( | |
| new_child, | |
| new_predicates, | |
| )?)); | |
| } | |
| if !under_collect && let Some(ppf) = plan.downcast_ref::<PerPartitionFilterExec>() { | |
| let child = plan.children()[0].clone(); | |
| let new_child = restrict(child, partitions, false)?; | |
| let new_predicates: Vec<_> = partitions | |
| .iter() | |
| .map(|&p| { | |
| ppf.predicates().get(p).cloned().ok_or_else(|| { | |
| DataFusionError::Internal(format!( | |
| "PerPartitionFilterExec: partition {p} out of range \ | |
| ({} predicates)", | |
| ppf.predicates().len() | |
| )) | |
| }) | |
| }) | |
| .collect::<Result<Vec<_>>>()?; | |
| return Ok(Arc::new(PerPartitionFilterExec::try_new( | |
| new_child, | |
| new_predicates, | |
| )?)); | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@ballista/scheduler/src/state/task_builder.rs` around lines 95 - 106, Update
the PerPartitionFilterExec handling in restrict to perform a checked lookup on
ppf.predicates() for each partition index instead of indexing directly. Convert
a missing predicate into the existing internal plan error via internal_err!,
while preserving cloning and collection of valid predicates.
3 participants
2196: To review by AI