avantgardnerio commented on code in PR #2211:
URL:
https://github.com/apache/datafusion-ballista/pull/2211#discussion_r3786019136
##########
ballista/scheduler/src/state/task_builder.rs:
##########
@@ -87,41 +86,20 @@ fn restrict(
return Ok(rewritten);
}
- // RangeFilterExec: raw_bounds is indexed by input partition; restriction
- // slices bounds parallel to the input's partition subset. Halos + routing
- // are carried over verbatim; RFE re-widens on the fresh operator.
- if !under_collect && let Some(rf) = plan.downcast_ref::<RangeFilterExec>()
{
+ // Operators carrying data indexed by global input partition slice it
+ // parallel to the input restriction. Each one implements the slicing
+ // beside its own fields; this walker only supplies the restricted child.
+ if !under_collect && let Some(sliceable) = as_partition_sliceable(&plan) {
let children = plan.children();
let [child] = children.as_slice() else {
return internal_err!(
- "RangeFilterExec must have exactly 1 child, got {}",
+ "{} is PartitionSliceable but has {} children, expected 1",
+ plan.name(),
children.len()
);
};
let new_child = restrict((*child).clone(), partitions, false)?;
- let raw_bounds = rf.raw_bounds().ok_or_else(|| {
- datafusion::common::DataFusionError::Internal(
- "RangeFilterExec: task-restriction before
resolve_bounds()".into(),
- )
- })?;
- let sliced_bounds: Vec<_> = partitions
- .iter()
- .map(|&global| {
- raw_bounds.get(global).cloned().ok_or_else(|| {
- datafusion::common::DataFusionError::Internal(format!(
- "RangeFilterExec: partition index {global} out of
bounds ({} raw bounds)",
- raw_bounds.len()
- ))
- })
- })
- .collect::<datafusion::common::Result<_>>()?;
- return Ok(Arc::new(RangeFilterExec::try_new_resolved(
- new_child,
- rf.routing_expr().clone(),
- rf.halo_lo().clone(),
- rf.halo_hi().clone(),
- sliced_bounds,
- )?));
+ return sliceable.slice_to_partitions(new_child, partitions);
}
Review Comment:
> any rule that repartitions the input to the same partition count would
silently attach each partition's offsets to the wrong rows
This was real and it fired on the first end to end run, though not by the
route you predicted. Task specialization already restricts a stage plan to one
task's partition slice, so `PrefixMergeExec` was rebuilt over a narrower input
while still holding state keyed by global partition index. It failed loudly
only because `try_new` validates the length: `per_partition_state.len() 4 does
not match input partition count 1`. Without that check it is exactly the silent
misattachment you describe, since a task's `execute(k)` numbers its own slice
from zero.
Fixed by slicing state and `Scalar.offset` parallel to the input
restriction, following the precedent `RangeFilterExec` already set for its
bounds. It is behind a `PartitionSliceable` trait so each operator implements
the slicing next to the fields being sliced, rather than the task builder
knowing their internals: add a partition indexed field and the code that has to
slice it is in the same file.
The e2e runs with `max_partitions_per_task = 2`, so tasks carry a genuine
multi partition slice rather than the degenerate one partition case.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
