andygrove commented on code in PR #2721:
URL:
https://github.com/apache/datafusion-ballista/pull/2721#discussion_r4529891723
##########
ballista/scheduler/src/state/task_builder.rs:
##########
@@ -151,21 +141,97 @@ fn restrict(
.into_iter()
.cloned()
.zip(per_child)
- .map(|(child, sub)| restrict(child, &sub, false))
+ .map(|(child, sub)| restrict(child, &sub, false, resolver))
.collect::<Result<Vec<_>>>()?;
return replace_children_if_necessary(plan, new_children);
}
- // Interior: compute per-child scope, recurse, rebuild the node.
- let per_child = child_scopes(&plan, under_collect);
- let new_children: Vec<Arc<dyn ExecutionPlan>> = plan
- .children()
+ let children = restrict_children(&plan, partitions, under_collect,
resolver)?;
+ if under_collect {
+ return replace_children_if_necessary(plan, children);
+ }
+
+ // Operators carrying data indexed by global partition slice it parallel to
+ // the restriction. Each one implements the slicing beside its own fields;
+ // this walker only supplies the restricted children. A node Ballista does
+ // not list may be restricted by the resolver.
+ let restricted = match as_partition_sliceable(&plan) {
+ Some(sliceable) => {
+ Some(sliceable.restrict_to_partitions(children.clone(),
partitions))
+ }
+ None => resolver.and_then(|resolve| resolve(&plan, children.clone(),
partitions)),
+ };
+ if let Some(restricted) = restricted {
+ return checked_restriction(plan.name(), restricted?, partitions);
+ }
+
+ // A leaf nothing above could restrict is only right if the task runs all
+ // of it.
+ if children.is_empty() {
+ ensure_leaf_runs_whole(&plan, partitions)?;
+ return Ok(plan);
+ }
+ replace_children_if_necessary(plan, children)
+}
+
+/// Restrict each child of `plan`, in the scope `child_scopes` gives it.
+fn restrict_children(
+ plan: &Arc<dyn ExecutionPlan>,
+ partitions: &[usize],
+ under_collect: bool,
+ resolver: Option<&PartitionSliceableResolver>,
+) -> Result<Vec<Arc<dyn ExecutionPlan>>> {
+ let per_child = child_scopes(plan, under_collect);
+ plan.children()
.into_iter()
.cloned()
.zip(per_child)
- .map(|(child, collect)| restrict(child, partitions, collect))
- .collect::<Result<Vec<_>>>()?;
- replace_children_if_necessary(plan, new_children)
+ .map(|(child, collect)| restrict(child, partitions, collect, resolver))
+ .collect()
+}
+
+/// Hold a `PartitionSliceable` to its contract: one partition for each
+/// partition the task runs. Another count would run the wrong partitions.
+fn checked_restriction(
+ name: &str,
+ restricted: Arc<dyn ExecutionPlan>,
+ partitions: &[usize],
+) -> Result<Arc<dyn ExecutionPlan>> {
+ let count = restricted
+ .properties()
+ .output_partitioning()
+ .partition_count();
+ if count != partitions.len() {
+ return plan_err!(
+ "{name}::restrict_to_partitions returned {count} partitions for \
+ partitions {partitions:?}; it has to return exactly one for each"
+ );
+ }
+ Ok(restricted)
+}
+
+/// Accept a leaf that nothing could restrict only if the task runs all of it,
+/// in order. Left as it is, it numbers its partitions from zero, so any other
+/// slice would run the wrong ones. A one-partition leaf therefore passes in
the
+/// task that owns it, and fails in a `UnionExec` branch that belongs to
another.
+fn ensure_leaf_runs_whole(
+ plan: &Arc<dyn ExecutionPlan>,
+ partitions: &[usize],
+) -> Result<()> {
+ let count = plan.properties().output_partitioning().partition_count();
+ if partitions.iter().copied().eq(0..count) {
+ return Ok(());
+ }
+ let name = plan.name();
+ let described = displayable(plan.as_ref()).one_line().to_string();
+ let described = described.trim_end();
+ plan_err!(
+ "cannot restrict `{described}` to partitions {partitions:?} of its
{count}: \
+ the scheduler does not know how to restrict a `{name}`, and left
alone it \
+ would number its partitions from zero in every task. A node of your
own \
Review Comment:
Fixed in aef82ff16.
##########
ballista/scheduler/src/state/task_builder.rs:
##########
@@ -151,21 +141,97 @@ fn restrict(
.into_iter()
.cloned()
.zip(per_child)
- .map(|(child, sub)| restrict(child, &sub, false))
+ .map(|(child, sub)| restrict(child, &sub, false, resolver))
.collect::<Result<Vec<_>>>()?;
return replace_children_if_necessary(plan, new_children);
}
- // Interior: compute per-child scope, recurse, rebuild the node.
- let per_child = child_scopes(&plan, under_collect);
- let new_children: Vec<Arc<dyn ExecutionPlan>> = plan
- .children()
+ let children = restrict_children(&plan, partitions, under_collect,
resolver)?;
+ if under_collect {
+ return replace_children_if_necessary(plan, children);
+ }
+
+ // Operators carrying data indexed by global partition slice it parallel to
+ // the restriction. Each one implements the slicing beside its own fields;
+ // this walker only supplies the restricted children. A node Ballista does
+ // not list may be restricted by the resolver.
+ let restricted = match as_partition_sliceable(&plan) {
+ Some(sliceable) => {
+ Some(sliceable.restrict_to_partitions(children.clone(),
partitions))
+ }
+ None => resolver.and_then(|resolve| resolve(&plan, children.clone(),
partitions)),
+ };
+ if let Some(restricted) = restricted {
+ return checked_restriction(plan.name(), restricted?, partitions);
+ }
+
+ // A leaf nothing above could restrict is only right if the task runs all
+ // of it.
+ if children.is_empty() {
Review Comment:
Agreed, the scheduler can't tell such a node from a stateless `FilterExec`,
so only leaves are guarded. The upgrade guide and the extending-components
guide now say so explicitly: an interior node of your own that holds data per
partition still returns wrong results in a split stage unless a resolver
restricts it. In aef82ff16. The longer-term fix is a method on `ExecutionPlan`
itself; see #2842.
##########
ballista/scheduler/src/config.rs:
##########
@@ -440,6 +441,14 @@ pub struct SchedulerConfig {
pub override_physical_codec: Option<Arc<dyn PhysicalExtensionCodec>>,
/// Override function for customizing gRPC client endpoints before they
are used
pub override_create_grpc_client_endpoint: Option<EndpointOverrideFn>,
+ /// Restricts `ExecutionPlan` nodes of your own that hold data per
+ /// partition to a task's partitions, using
+ ///
[`PartitionSliceable`](ballista_core::execution_plans::plan_algebra::PartitionSliceable)
+ /// or by building the restricted node itself. Without it, a
+ /// stage split between tasks that holds such a leaf fails the job. Unlike
+ /// the `override_*` fields it adds to Ballista's own list rather than
+ /// replacing it. See [`PartitionSliceableResolver`].
+ pub partition_sliceable_resolver: Option<PartitionSliceableResolver>,
Review Comment:
Good point about keeping it next to the codec. Filed #2839 to make the
resolver reachable from standalone, preferably by registering it on the
`SessionConfig` alongside the physical codec.
##########
ballista/core/src/execution_plans/plan_algebra.rs:
##########
@@ -125,6 +140,55 @@ pub fn as_partition_sliceable(
None
}
+/// Restricts a node that [`as_partition_sliceable`] does not list to a task's
+/// partitions, so the scheduler can restrict `ExecutionPlan` nodes of your
own,
+/// including ones you cannot implement [`PartitionSliceable`] on, such as a
+/// `DataSourceExec` over a `DataSource` of your own. Set it with
+/// `SchedulerConfig::with_partition_sliceable_resolver`.
+///
+/// It receives the node, its children already restricted as
+/// [`PartitionSliceable::restrict_to_partitions`] receives them, and the
global
+/// partition indices the task runs. It returns `None` for a node it does not
+/// know, and otherwise the node rebuilt over `children` with exactly
+/// `partitions.len()` output partitions; any other count is an error. For a
node
+/// of your own, that is typically a `downcast_ref` to your type followed by a
+/// call to its [`PartitionSliceable::restrict_to_partitions`].
+///
+/// It is asked after the built-in list, and never about a node that is read
+/// whole: one below a collapse, or declared `SinglePartition` by its parent.
It
+/// is called for most nodes of every task's plan as tasks are launched, so
keep
+/// the `None` path cheap.
+pub type PartitionSliceableResolver = Arc<
Review Comment:
Done: `PartitionSliceableResolver` is now a trait (`Debug + Send + Sync`,
one `restrict_to_partitions` method), like `AdaptiveRules`, and
`SchedulerConfig::with_partition_sliceable_resolver` takes an `Arc<dyn
PartitionSliceableResolver>`. In aef82ff16. For the upstream method, I'd like
to settle the design here before proposing it to DataFusion: #2842.
##########
docs/source/upgrading/56.0.0.md:
##########
@@ -178,6 +178,28 @@ early, so a straggler partition leaves them idle. Set
`ballista.scheduler.max_partitions_per_task` to a positive value to dispatch
smaller tasks, which release their vcores sooner.
+#### Stages split between tasks fail on a leaf the scheduler cannot restrict
+
+When a stage is split between tasks, each task gets a plan restricted to its
+own partitions. A leaf the scheduler did not know how to restrict used to be
+left as it was, so a task that ran partitions 4 and 5 of it ran its partitions 0
+and 1, and the stage returned some partitions twice and others not at all, with
+no error.
+
+The scheduler now fails the job instead. Such a leaf is accepted only in a task
Review Comment:
Added a paragraph to the upgrade guide saying that splitting depends on free
task slots and on retries, so a query can pass on an idle cluster and fail on a
busy one, and suggesting `ballista.scheduler.max_partitions_per_task` set low
for testing. In aef82ff16. Filed #2840 for surfacing the reason under
push-staged scheduling.
##########
ballista/scheduler/src/planner.rs:
##########
@@ -669,9 +669,11 @@ fn can_stay_inline(plan: &Arc<dyn ExecutionPlan>) -> bool {
/// Whether per-task restriction can reduce `plan` to zero partitions. An empty
/// slice means "this task polls nothing here", so a branch that still reports
a
-/// partition afterwards is one that cannot be restricted away.
+/// partition afterwards, or that the rewriter refuses, cannot be restricted
+/// away. No resolver is passed, so a node Ballista does not know gets a stage
+/// of its own.
fn restricts_to_nothing(plan: &Arc<dyn ExecutionPlan>) -> bool {
- restrict_plan_to_partitions(plan.clone(), &[])
+ restrict_plan_to_partitions(plan.clone(), &[], None)
Review Comment:
Filed #2841.
--
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]
