feat: let custom ExecutionPlan nodes restrict themselves to a task's partitions by andygrove · Pull Request #2721 · apache/datafusion-ballista

andygrove · 2026-10-31T14:53:42Z

Which issue does this PR close?

Closes #2679.

Rationale for this change

When the scheduler splits a stage between tasks, restrict_plan_to_partitions rewrites the plan so that each task sees only its own partitions. It could restrict only the leaves select_output_partitions knows and the two operators on the hand-kept as_partition_sliceable list. Any other node that holds data per partition was left as it was, so a task that ran partitions 4 and 5 of it executed its partitions 0 and 1, and the stage returned some partitions twice and others not at all, with no error. A downstream crate could implement PartitionSliceable but had no way to make the scheduler find it, the trait covered only operators with a single child, and a node the crate cannot implement a trait on, such as a DataSourceExec over its own DataSource, could not be restricted at all.

Comet Native Executors (apache/datafusion-comet#7288) need this. They work around it today by giving every fragment a fake MemorySourceConfig leaf with one row per partition.

What changes are included in this PR?

The three things the issue asks for:

PartitionSliceable covers any node. restrict_to_partitions(children, partitions) replaces slice_to_partitions(child, partitions). It receives all of the node's children, already restricted, and is used for leaves too. The children are restricted the way an interior node's always are, so an input declared SinglePartition arrives whole. RangeFilterExec and PrefixMergeExec are ported. The walker also checks that the restricted node has partitions.len() partitions, because a node that gets that wrong would bring the silent duplication back.
A resolver on SchedulerConfig. partition_sliceable_resolver (builder: with_partition_sliceable_resolver) is called with a plan node, its children already restricted, and the task's partitions, and returns the restricted node, or None for a node it does not know. It restricts the node itself rather than returning a PartitionSliceable, so it also covers nodes that cannot implement the trait, such as a DataSourceExec. For a node of your own it typically calls restrict_to_partitions. It is asked after the built-in list. TaskManager passes it to the rewriter. The planner does not use it, so a union branch holding a node Ballista does not know still gets a stage of its own, as before.
A leaf nothing can restrict fails instead of duplicating output. It is accepted only in a task that runs all of its partitions, in order, such as the single partition of a one-partition leaf, so plans made of leaves the scheduler knows are unaffected. Anything else is an error that names the node and what to implement. So that None from select_output_partitions means "unrecognised", EmptyExec and PlaceholderRowExec are now returned as recognised, and the constructor errors of the readers and the memory source propagate instead of turning into None.

Smaller changes: TaskManager::new now delegates to with_launcher, so the config is read in one place. There are entries in the 56.0.0 upgrade guide and a "Custom Physical Plan Nodes" section in the extending-components guide.

Are there any user-facing changes?

Yes, and some are breaking for anyone who embeds the scheduler or defines ExecutionPlan nodes. All are in docs/source/upgrading/56.0.0.md.

PartitionSliceable::slice_to_partitions is now restrict_to_partitions and takes a Vec of children.
task_builder::restrict_plan_to_partitions takes a third argument, Option<&dyn PartitionSliceableResolver>.
SchedulerConfig has a new public field, partition_sliceable_resolver, holding an Arc<dyn PartitionSliceableResolver>. The new trait's restrict_to_partitions takes the node, the restricted children and the partitions, and returns Option<Result<Arc<dyn ExecutionPlan>>>.
A job whose stage is split between tasks now fails when the stage holds a leaf the scheduler cannot restrict. It used to return wrong results. A retry of part of a stage counts as split. A DataSourceExec over a source that is neither a file scan nor an in-memory source used to log a warning and run, and is now refused when its stage is split.
The standalone scheduler builds its own SchedulerConfig, so it cannot be given a resolver.

One limitation: under push-staged scheduling the job's status says only "task could not be prepared or was rejected by executor", and the reason is in the scheduler's log. Pull-staged scheduling puts it in the status. Every task-preparation failure behaves this way today, so it is left alone here.

Testing

New unit tests cover the walker (a leaf found through the resolver, a node without PartitionSliceable restricted by the resolver building a replacement, inputs restricted except one the node reads whole, a node below a collapse, the built-in list answering before the resolver, a union branch owned by another task, a leaf nobody can restrict, a wrong partition count, the EmptyExec and PlaceholderRowExec exemption). Two launch_tasks tests submit a stage that is split between two tasks and holds a custom leaf. Without a resolver the job fails and no task is sent. With one, each task is sent only the payloads of its own partitions, which together are every payload once. Each piece was also broken on purpose, and each break failed the test meant to catch it.

cargo fmt, cargo clippy --all-targets --workspace --all-features -- -D warnings, cargo doc with -D warnings and prettier are clean, and the ballista-scheduler tests pass. The ballista-core tests that bind sockets were not run locally and are left to CI.

apache#2633 added a test to this module that names ExecutionPlan, and apache#2613 then moved FailingPlanRewriteExec out of the module and dropped the imports it had needed, ExecutionPlan among them. Each passed CI against its own base. Together they leave the scheduler's test target uncompilable, so cargo check, Clippy and the test jobs have failed on main since apache#2613 landed.

…partitions When a stage is split between tasks, the scheduler restricts each task's plan to the partitions that task runs. It could only restrict the leaves it knew and the two PartitionSliceable operators on a hand-kept list. Any other node holding data per partition was left as it was, so a task that ran partitions 4 and 5 of it executed its partitions 0 and 1, and the stage returned some partitions twice and others not at all, with no error. - PartitionSliceable::restrict_to_partitions takes any number of children, leaves included. It replaces slice_to_partitions, which took one. - SchedulerConfig::partition_sliceable_resolver finds the PartitionSliceable of nodes Ballista does not list. It is asked after the built-in list. - A leaf that nothing can restrict now fails the task unless the task runs all of its partitions in order, instead of running and duplicating output. The scheduler also rejects a restriction that returns the wrong number of partitions.

…olver test (cherry picked from commit 2efa802)

milenkovicm · 2026-11-01T23:09:39Z

I don't understand motivation at all (which is not blocker on its own) is there any way to explain this pr like I'm 4 😂😀

andygrove · 2026-11-02T00:05:35Z

Sorry. Comet Native needs it. I'll have Claude explain properly tomorrow.

andygrove · 2026-11-02T14:27:10Z

Fair question, sorry the description was dense. Here's the plain version.

The problem

Ballista sometimes splits one stage across several tasks. Say a stage has 4 partitions and 2 tasks. Task A gets partitions 0 and 1, and task B gets 2 and 3. Before sending a task out, the scheduler rewrites the plan so each task only sees its own partitions. It knows how to do this for the nodes it recognises: shuffle readers, Parquet and in-memory scans, and a couple of operators on a hand-kept list.

If a plan contains some other node that holds data per partition, the scheduler leaves it alone. Task B then runs "partitions 0 and 1" of that node instead of 2 and 3. The result is that partitions 0 and 1 come back twice and 2 and 3 never come back. There is no error, just wrong results.

Why Comet needs it

Comet native executors build plans with their own leaf nodes that hold per-partition data. The scheduler doesn't know about them, so it can't slice them. Today Comet works around that by giving every fragment a fake MemorySourceConfig leaf with one row per partition, purely so the scheduler has something it recognises. See apache/datafusion-comet#7288.

What the PR does

PartitionSliceable can now be implemented by any node, including leaves and nodes with several children. The method is renamed from slice_to_partitions to restrict_to_partitions.
SchedulerConfig gets a partition_sliceable_resolver. A downstream crate uses it to tell the scheduler how to find the PartitionSliceable for its own nodes.
A leaf that nobody can restrict now fails the job with a clear message instead of silently duplicating output. The scheduler also checks that a restricted node returns the right number of partitions.

The first commit is unrelated. It's a one-line import fix for main's CI, which has been red since #2613.

Breaking changes

The trait method rename, a third argument on restrict_plan_to_partitions, and a new public field on SchedulerConfig. All are listed in the 56.0.0 upgrade guide.

coderfender · 2026-11-03T05:54:28Z

+    Ok(child)
+}
+
 /// Take `values[global]` for each global partition index, in task-local


Now that the guide points custom nodes at slice_by_global_partition, its out-of-bounds error reaches users. As a DataFusionError::Internal it renders with "This issue was likely caused by a bug in DataFusion's code. Please help us ... file a bug report" (hit this from a resolver-backed node). Could it, and single_child, use plan_err! like checked_restriction / ensure_leaf_runs_whole?

Good catch. slice_by_global_partition and single_child now return plan_err!, in 29493d9.

coderfender · 2026-11-03T05:54:28Z

+/// 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
+/// it cheap, typically a `downcast_ref` to your own types.
+pub type PartitionSliceableResolver = Arc<


With this PR a split stage whose DataSourceExec wraps a custom DataSource fails, but I don't see a way for an embedder to make it restrictable: PartitionSliceable: ExecutionPlan rules out implementing it on the DataSource, the orphan rule rules out DataSourceExec, and the resolver returns a reference borrowed from the plan, so it can't hand back a wrapper. Would a resolver that does the restriction itself, e.g. Fn(&Arc<dyn ExecutionPlan>, Vec<Arc<dyn ExecutionPlan>>, &[usize]) -> Option<Result<Arc<dyn ExecutionPlan>>>, work? It would cover this case as well as custom nodes.

You're right, the orphan rule and the borrowed return type leave no way to restrict a DataSourceExec over a custom source. The resolver now takes the node, its restricted children and the partitions, and returns the restricted node itself, as you suggested. The walker restricts children once and then asks the built-in list and the resolver. There's a test where the resolver builds a replacement for a node without PartitionSliceable. In 29493d9.

…-to-task-partitions Resolve conflicts in the scheduler config, task manager, task builder and state tests. `stages_read_whole` no longer special-cases partition-sliceable nodes, since their children now get the same scopes as any other node. Also address review feedback on apache#2721: - The partition sliceable resolver now restricts the node itself, given the node, its restricted children and the task's partitions, so it can restrict nodes that cannot implement `PartitionSliceable`, such as a `DataSourceExec` over a custom `DataSource`. - `slice_by_global_partition` and `single_child` return plan errors rather than internal errors, so users are not told to file a bug.

…-to-task-partitions # Conflicts: # ballista/scheduler/src/state/mod.rs

NoahKusaba · 2026-11-10T04:19:57Z

+    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 \


Nit: typo in the user-facing error: "A node of your own be restricted by" is missing "can".

Fixed in aef82ff.

NoahKusaba · 2026-11-10T04:19:57Z

+
+    // A leaf nothing above could restrict is only right if the task runs all
+    // of it.
+    if children.is_empty() {


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.

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 aef82ff. The longer-term fix is a method on ExecutionPlan itself; see #2842.

NoahKusaba · 2026-11-10T04:19:57Z

+    /// 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>,


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.

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.

NoahKusaba · 2026-11-10T04:19:57Z

+/// 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<


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.

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 aef82ff. For the upstream method, I'd like to settle the design here before proposing it to DataFusion: #2842.

NoahKusaba · 2026-11-10T04:19:57Z

+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


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.

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 aef82ff. Filed #2840 for surfacing the reason under push-staged scheduling.

Would a query fail? I thought it would just get its tasks restricted down to available slots, including to 1, and then succeed, but slower?

It does fail, but only for plans that contain a leaf the scheduler can't restrict. Your picture is right for everything else. When a stage gets fewer slots, its tasks get fewer partitions each. For a shuffle reader, a Parquet/memory scan, EmptyExec or PlaceholderRowExec, the task's plan is narrowed to those partitions and the query just runs slower, as you describe.

The problem is a leaf the scheduler doesn't know how to narrow, such as a DataSourceExec over a custom DataSource or a third-party scan with no resolver. Cutting the task down to the free slots is exactly when that leaf has to be narrowed. Before this PR it was left as it was, so a task given partitions 2 and 3 actually ran its partitions 0 and 1. The query "succeeded" but returned some partitions twice and others not at all. Now ensure_leaf_runs_whole accepts such a leaf only when the task runs all of its partitions (0..count). Otherwise the job fails with an error naming the node.

So on an idle cluster the stage goes out as one task holding every partition and passes. On a busy cluster, or when part of the stage is retried, it gets split and fails, where before it silently gave wrong results. That's the load dependence the paragraph is warning about. The fix for an affected plan is a resolver (or PartitionSliceable) for the node, and then it behaves like the built-in leaves: restricted, slower, correct.

NoahKusaba · 2026-11-10T04:19:57Z

+/// 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)


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.

Filed #2841.

NoahKusaba · 2026-11-10T04:19:57Z

+        let error = restrict_plan_to_partitions(leaf, &[1], None)
+            .unwrap_err()
+            .to_string();
+        assert!(error.contains("PayloadExec"), "{error}");


Nit: comparing the full error message instead of using contains would catch wording regressions like the typo in ensure_leaf_runs_whole.

The test now compares the whole message. In aef82ff.

NoahKusaba · 2026-11-10T04:33:05Z

I generally agree with the design, and I'm approving to unblock Iceberg.

I couldn't find anything wrong, but I had my AI go through it and it found some nits / doc cleanups + we should open that datafusion Github issue.

avantgardnerio

I've been wanting to go this direction anyway, thank you for doing this!

CC @phillipleblanc as there's downstream breaking changes.

…-to-task-partitions Keep both 56.0.0 upgrade notes: the PartitionSliceable resolver and the ExecutorEndpoint change to partition locations.

… guarded The resolver is now a trait, like the other extension points on SchedulerConfig, so methods can be added later without another breaking change. Fix the missing "can" in the unrestrictable-leaf error and compare the whole message in its test. The upgrade guide and the extending-components guide now say that only leaves are checked, and that whether a stage is split depends on cluster load and retries. Co-Authored-By: Claude Opus 5.5 <[email protected]>

andygrove · 2026-11-10T21:22:31Z

Thanks for the reviews @NoahKusaba @avantgardnerio @coderfender @milenkovicm

andygrove added the api-change label Oct 31, 2026

github-actions Bot added the documentation Improvements or additions to documentation label Oct 31, 2026

This was referenced Oct 31, 2026

feat: let hash shuffles place rows by a client-computed partition id #2723

Open

fix: plan a range RepartitionExec as a shuffle stage instead of dropping it #2722

Open

andygrove force-pushed the restrict-custom-nodes-to-task-partitions branch from 089383e to 03c13f1 Compare October 31, 2026 15:48

This was referenced Oct 31, 2026

feat: experimental MVP of Comet Native Executors apache/datafusion-comet#7289

Draft

Let custom operators declare that they tolerate a coalesced partition count #2728

Open

test: look through CooperativeExec for the restricted leaf in the res…

ad87057

…olver test (cherry picked from commit 2efa802)

andygrove requested review from avantgardnerio and milenkovicm November 1, 2026 14:55

andygrove mentioned this pull request Nov 1, 2026

chore: pin DataFusion to a commit on branch-56 ahead of DataFusion 56 #2718

Open

coderfender reviewed Nov 3, 2026

View reviewed changes

This was referenced Nov 3, 2026

chore: comet-native integration branch (testing with Comet native executors) #2755

Draft

feat(iceberg): scan Iceberg tables in Ballista #2638

Open

Merge remote-tracking branch 'apache/main' into restrict-custom-nodes…

7b4f699

…-to-task-partitions # Conflicts: # ballista/scheduler/src/state/mod.rs

This comment was marked as duplicate.

Sign in to view

NoahKusaba approved these changes Nov 10, 2026

View reviewed changes

avantgardnerio approved these changes Nov 10, 2026

View reviewed changes

sim and others added 2 commits November 10, 2026 08:25

Merge remote-tracking branch 'apache/main' into restrict-custom-nodes…

c21e1d1

…-to-task-partitions Keep both 56.0.0 upgrade notes: the PartitionSliceable resolver and the ExecutorEndpoint change to partition locations.

This was referenced Nov 10, 2026

Let the standalone scheduler take a PartitionSliceableResolver #2839

Open

Surface task-preparation errors in the job status under push-staged scheduling #2840

Open

Pass the PartitionSliceableResolver to the planner's union-inlining check #2841

Open

andygrove mentioned this pull request Nov 10, 2026

Discuss proposing a partition-restriction method on DataFusion's ExecutionPlan #2842

Open

andygrove added this pull request to the merge queue Nov 10, 2026

Merged via the queue into apache:main with commit 6ecc5c3 Nov 10, 2026
24 checks passed

andygrove deleted the restrict-custom-nodes-to-task-partitions branch November 10, 2026 21:46

NoahKusaba added a commit to NoahKusaba/datafusion-ballista that referenced this pull request Nov 11, 2026

Pin ballista-iceberg to a Ballista commit with apache#2721

8c19489

Conversation

andygrove commented Oct 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are there any user-facing changes?

Testing

Uh oh!

milenkovicm commented Nov 1, 2026

Uh oh!

andygrove commented Nov 2, 2026

Uh oh!

andygrove commented Nov 2, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This comment was marked as duplicate.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NoahKusaba commented Nov 10, 2026

Uh oh!

avantgardnerio left a comment

Choose a reason for hiding this comment

Uh oh!

andygrove commented Nov 10, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

andygrove commented Oct 31, 2026 •

edited

Loading