NoahKusaba commented on code in PR #2721:
URL:
https://github.com/apache/datafusion-ballista/pull/2721#discussion_r4521080854
##########
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:
Only leaves get this check. An unknown node with children that holds
per-partition state, for example a third-party operator shaped like
`RangeFilterExec`, still goes through `replace_children_if_necessary` without
any check. Its children get restricted but its own state doesn't, which is the
original bug one level up.
I don't think the scheduler can tell that node apart from a stateless
`FilterExec`, so this may not be fixable here. Could the docs and upgrade guide
say explicitly that only leaves are guarded? "Fails instead of duplicating
output" reads broader than what this does.
##########
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:
A question about where this is registered. A custom node can only reach a
task if the physical extension codec can serialize it, so whoever registers the
codec also knows how to restrict the node. Keeping the two in separate places
means they can disagree, and the standalone scheduler already shows that: it
builds its own `SchedulerConfig`, so it can get the codec but not the resolver.
For integrations like Iceberg (#2638), a `UNION ALL` over a custom scan on
standalone would then fail with no way to fix it.
Could the standalone builder accept a resolver too, or could the restriction
be registered alongside the codec? A follow-up issue would be fine.
##########
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:
Minor: `Arc<dyn Fn(&Arc<dyn ExecutionPlan>, Vec<…>, &[usize]) ->
Option<Result<…>>>` is hard to read in rustdoc, and it doesn't match the
trait-based extension points next to it on `SchedulerConfig` (e.g.
`adaptive_rules: Arc<dyn AdaptiveRules>`). A small trait would also let methods
be added later without another breaking change.
Since the doc comment says an upstream `select_partitions` on
`ExecutionPlan` would replace this, it might be worth filing that in DataFusion
and linking it here, so this hook has a planned replacement.
##########
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:
Typo in the user-facing error: "A node of your own be restricted by" is
missing "can".
##########
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:
Agreed that passing `None` here is safe, since an unknown node gets its own
stage. It does mean a union branch with a resolvable custom node (Comet,
Iceberg) loses the inline parallelism. A tracking issue for passing the
resolver to the planner would be nice.
##########
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:
This is the right trade (an error instead of wrong results), but these
failures depend on cluster load. Whether a stage gets split depends on free
vcores, and partial retries count as split, so a query with a custom
`DataSource` or something like `StreamingTableExec` can pass in dev and fail in
production. Worth saying that explicitly here.
Related: as the PR description notes, push-staged scheduling hides this
message behind "task could not be prepared or was rejected by executor". Now
that this is an expected user-facing error, a follow-up issue to surface the
reason would help.
##########
ballista/scheduler/src/state/task_builder.rs:
##########
@@ -795,4 +864,204 @@ mod tests {
assert_eq!(reader.partition[0][0].partition_id.partition_id, 1);
assert_eq!(reader.partition[1][0].partition_id.partition_id, 3);
}
+
+ // --- Nodes the scheduler does not know: `PayloadExec` stands in for one
---
+
+ use crate::test_utils::payload::{PayloadExec, payload_resolver,
payloads_of};
+ use
ballista_core::execution_plans::plan_algebra::PartitionSliceableResolver;
+ use std::sync::Mutex;
+
+ fn partition_count(plan: &Arc<dyn ExecutionPlan>) -> usize {
+ plan.properties().output_partitioning().partition_count()
+ }
+
+ /// The file of each group of a scan (`group_file_counts` only counts
them).
+ fn group_files(plan: &Arc<dyn ExecutionPlan>) -> Vec<String> {
+ file_groups(plan)
+ .iter()
+ .map(|group| group.files()[0].object_meta.location.to_string())
+ .collect()
+ }
+
+ #[test]
+ fn resolver_restricts_a_node_the_scheduler_does_not_know() {
+ let resolver = payload_resolver();
+ let leaf = PayloadExec::leaf(&["p0", "p1", "p2", "p3"]);
+
+ let restricted =
+ restrict_plan_to_partitions(leaf, &[1, 3],
Some(&resolver)).unwrap();
+
+ assert_eq!(payloads_of(&restricted), ["p1", "p3"]);
+ assert_eq!(partition_count(&restricted), 2);
+ }
+
+ /// A node that cannot implement `PartitionSliceable`, such as a
+ /// `DataSourceExec` over a source of your own, is restricted by the
+ /// resolver building a replacement.
+ #[test]
+ fn resolver_restricts_a_node_without_partition_sliceable() {
+ let resolver: PartitionSliceableResolver =
+ Arc::new(|plan, children, partitions| {
+ assert!(children.is_empty());
+ plan.is::<PayloadExec>().then(|| {
+ Ok(Arc::new(
+
EmptyExec::new(plan.schema()).with_partitions(partitions.len()),
+ ) as Arc<dyn ExecutionPlan>)
+ })
+ });
+ let leaf = PayloadExec::leaf(&["p0", "p1", "p2", "p3"]);
+
+ let restricted =
+ restrict_plan_to_partitions(leaf, &[1, 3],
Some(&resolver)).unwrap();
+
+ assert!(restricted.is::<EmptyExec>());
+ assert_eq!(partition_count(&restricted), 2);
+ }
+
+ #[test]
+ fn a_leaf_nobody_can_restrict_is_refused_unless_the_task_runs_all_of_it() {
+ let leaf = PayloadExec::leaf(&["p0", "p1"]);
+
+ let kept = restrict_plan_to_partitions(leaf.clone(), &[0, 1],
None).unwrap();
+ assert!(
+ Arc::ptr_eq(&kept, &leaf),
+ "an unrestricted leaf is left alone"
+ );
+
+ // A one-partition leaf is right for the task that owns it.
+ let single = PayloadExec::leaf(&["p0"]);
+ assert!(restrict_plan_to_partitions(single, &[0], None).is_ok());
+
+ for partitions in [&[0][..], &[1], &[1, 0], &[]] {
+ assert!(
+ restrict_plan_to_partitions(leaf.clone(), partitions,
None).is_err(),
+ "{partitions:?} is not all of the leaf, in order"
+ );
+ }
+
+ let error = restrict_plan_to_partitions(leaf, &[1], None)
+ .unwrap_err()
+ .to_string();
+ assert!(error.contains("PayloadExec"), "{error}");
Review Comment:
Nit: comparing the full error message instead of using `contains` would
catch wording regressions like the typo in `ensure_leaf_runs_whole`.
--
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]
