feat: allow fractional desired task counts by shinzoxD · Pull Request #617 · datafusion-contrib/datafusion-distributed

shinzoxD · 2026-08-09T15:51:30Z

Summary

allow DesiredTaskCountHandler responses to use fractional f64 task-count hints
preserve fractions through hint merging and union aggregation, then round up when resolving an integer task count
migrate built-in handlers, examples, public docs, and the 3.0 upgrade guide

Closes #562

Validation

cargo fmt --all -- --check
cargo check --lib
cargo clippy --lib -- -D warnings
cargo test --features integration --lib -- --skip protocol::grpc::channel_resolver::tests::fails_establishing_connection (289 passed, 0 failed, 1 ignored, 1 filtered out)
cargo check --example custom_worker_url_routing --features integration

The library suite and example check passed locally with the benchmark-only dev-dependency edge temporarily excluded to avoid its unrelated vendored-OpenSSL Windows build. The snapshot path filter was normalized only for the Windows test run, and the Windows-only connection-failure assertion was skipped because its socket error text differs by platform. Cargo.toml, Cargo.lock, and the snapshot helper were restored with no diff before commit.

Assisted-by: OpenAI Codex

Copilot

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

gabotechs

This looks mostly great, thanks @shinzoxD! just left some minor comments

gabotechs · 2026-08-09T16:18:02Z

+        ChildrenIsolatorUnionExec: task_count=Desired(1)
+          FilterExec: task_count=Maximum(1)
+            RepartitionExec: task_count=Maximum(1)
+              DistributedLeafExec: task_count=Maximum(1)
+          ProjectionExec: task_count=Maximum(1)
+            FilterExec: task_count=Maximum(1)
+              RepartitionExec: task_count=Maximum(1)
+                DistributedLeafExec: task_count=Maximum(1)


🤔 I'm not sure why there's a Maximum(1) getting stamped in the two separate branches.

How it is that we are not seeing a Desired(0.5) instead?

The Maximum(1) annotations are the final per-child allocations: after the union resolves to one task, propagate_task_count_until_network_boundaries stamps each allocated child count as a hard cap. The union's Desired(1) is the key assertion that the two 0.5 hints were summed before rounding; early rounding would produce Desired(2). I added that explanation beside the snapshot.

Sounds good. Can we add another test that double checks what happens when the two leaves of the union report 0.0 task count?

Added in 01c6441 as test_union_all_zero_task_count_leaves. Two Desired(0.0) leaves still sum before rounding, then ChildrenIsolatorUnionExec rejects the resulting zero-task stage as an internal planning error rather than silently collapsing to one task.

gabotechs · 2026-08-09T16:44:23Z

@@ -30,7 +30,9 @@ pub(crate) fn file_scan_config_desired_task_count(
        .div_ceil(d_cfg.file_scan_config_bytes_per_partition)
        .div_ceil(cfg.target_partitions());

-    Some(Ok(DesiredTaskCountEventResponse::desired(task_count)))
+    Some(Ok(DesiredTaskCountEventResponse::desired(
+        task_count as f64,
+    )))


Here, we should no longer be doing a .div_ceil(), we should be doing normal f64 divisions so that we get true fractional numbers.

Done in 01c6441. The default file-scan handler now does total_bytes as f64 / bytes_per_partition / target_partitions (and the work-unit scan helper matches). as_usize() still ceils only when the planner resolves an integer task count.

shinzoxD · 2026-08-09T17:03:05Z

CI exposed three stale inline snapshot expectations that should now remain fractional. I updated only those expectations in 849f66d; the broad local library suite passes with 289 tests passed, 0 failed, 1 ignored, and the known Windows-only socket-text assertion filtered out. The replacement workflow is waiting for external-contributor approval: https://github.com/datafusion-contrib/datafusion-distributed/actions/runs/31325189612

gabotechs · 2026-08-09T17:02:44Z

        assert_snapshot!(annotated, @"
-        HashJoinExec: task_count=Desired(2)
-          NetworkShuffleExec: task_count=Desired(2)
+        HashJoinExec: task_count=Desired(1.3333333333333333)


For the Debug implementation, it might be enough to just show 2 decimal numbers.

TaskCountAnnotation::Debug now prints whole numbers as Desired(4) and fractions with two decimals (Desired(1.33)). Snapshot expectations in this file were updated to match.

gabotechs · 2026-08-14T07:52:58Z

@shinzoxD let me know if you still want to work on this, otherwise we can have other contributors help here

shinzoxD · 2026-08-15T18:46:25Z

Yes, still working on this. The remaining review items are:

a union-of-zero-task-count leaves test
dropping .div_ceil() in file_scan_config so the default handler returns a true fractional f64
two-decimal Debug for the task-count annotation

I'll push those shortly.

gabotechs · 2026-08-16T14:54:16Z

+        // Two 0.0 leaf hints sum to Desired(0) before rounding. The isolator
+        // then rejects a zero-task union as an internal planning error rather
+        // than silently collapsing to a single task.
+        let err = annotate_test_plan_result(test_plan_builder, query)
+            .await
+            .expect_err("zero-task union should fail planning");


This should not error out. A Desired(0) is a valid outcome (e.g. all datasets in the leaves of the union are empty). Sounds like something we need to fix.

Addressed in 84bacd0: Desired(0) is now a valid empty-stage outcome. Two zero-weight union leaves plan as Desired(0) / Maximum(0) instead of an internal error.

gabotechs · 2026-08-16T14:55:44Z

    async fn annotate_test_plan(test_plan_builder: TestPlanBuilder, query: &str) -> String {
+        annotate_test_plan_result(test_plan_builder, query)
+            .await
+            .expect("failed to annotate plan")
+    }
+
+    async fn annotate_test_plan_result(
+        test_plan_builder: TestPlanBuilder,
+        query: &str,
+    ) -> datafusion::error::Result<String> {


We never want to assert errors here, I'd rollback this change.

Addressed in 84bacd0: rolled back annotate_test_plan_result. The helper is success-only again; the zero-task case now snapshots the valid Desired(0) plan instead of asserting an error.

gabotechs · 2026-08-17T09:37:23Z

                    // about the task count assigned to the network boundary in the consumer stage,
                    // and we don't want it to affect other task count decisions.
-                    (Some(Arc::new(stats)), Desired(1))
+                    (Some(Arc::new(stats)), Desired(1.0))


🤔 I think we should be returning 0.0 here. The dynamic planner might now decide to that maybe 0.2 tasks are enough, and this might get added between children of a UNION, so we would not want the 0.2 to be masked by the greater 1.0 in this case, we want the lower fractional values to be summed across children in a UNION.

Addressed in 84bacd0: the dynamic planner now returns Desired(0.0) for the consumer-side network-boundary hint so small fractional values are not masked by 1.0 and can be summed across UNION children.

gabotechs · 2026-08-17T09:38:20Z

-    let task_count = total_bytes
-        .div_ceil(d_cfg.file_scan_config_bytes_per_partition)
-        .div_ceil(cfg.target_partitions());
+    let bytes_per_partition = d_cfg.file_scan_config_bytes_per_partition.max(1) as f64;
+    let target_partitions = cfg.target_partitions().max(1) as f64;
+    let task_count = total_bytes as f64 / bytes_per_partition / target_partitions;


shinzoxD · 2026-08-19T18:39:28Z

Rebased onto main (now at ee561b5) to resolve the conflict.

The only conflict was in src/distributed_planner/inject_network_boundaries.rs. Resolution keeps main's newer structure from #655 (unified leaf handling, NetworkCoalesceExec above leaves) and this PR's fractional task-count behavior:

Union children are summed as f64 before rounding
Absent non-leaf hints default to Desired(0.0) so they do not inflate fractional children via max-merge
Desired(0) remains a valid empty-stage outcome
Dynamic planner still returns Desired(0.0) so a 0.0 hint does not mask smaller fractional values or inflate UNION sums

New HEAD: fe459198204b1bed843f36e882651af077ed8b3c

shinzoxD · 2026-08-24T16:22:48Z

Remaining review items from that thread are in: union zero-task leaves, fractional file-scan, two-decimal Debug, Desired(0) valid, and dynamic planner 0.0. PR is mergeable on current main.

shinzoxD · 2026-08-24T17:51:32Z

Rebased onto latest main (includes DF 55 / 4.0.0 and #661). Remaining review items are still in; not merging.

Fractional behavior is unchanged: union children sum as f64 before rounding, absent non-leaf hints are Desired(0.0), file-scan uses true f64 division, Desired(0) remains a valid empty-stage outcome. The Desired(f64) migration note moved to docs/upgrade/5.0.0.md now that 4.0.0 is released.

shinzoxD · 2026-08-24T18:14:32Z

Rebase onto main pulled in build_side_fetch_is_preserved_by_broadcast (#618). Fractional file-scan hints were giving that broadcast producer 4 tasks, so LIMIT 50 applied per task (3750 vs 2500). The producer stage is now capped at 1 when the subtree has a fetch/limit.

Keep Desired Debug to two decimal places for non-integers, compute default file-scan task counts with true f64 division, and add a union regression for 0.0 leaf hints. Two zero hints still sum before rounding and are then rejected as a zero-task ChildrenIsolatorUnionExec.

Empty union children can sum to Desired(0). That is a valid empty stage, not an internal planning error. Return Desired(0.0) from the dynamic planner so UNION children can sum small fractional hints instead of being masked by 1.0.

v4.0.0 is now released, so the Desired(f64) public break belongs in the next major upgrade guide rather than the historical 3.0.0 notes.

shinzoxD · 2026-08-24T18:30:52Z

Rebased onto main after #670. Dropped the 1-task producer cap — that was papering over the missing topology. #670 is the right fix (fetch-bearing coalesce stays below BroadcastExec).

Copilot AI lite review requested due to automatic review settings August 9, 2026 15:51

Copilot AI reviewed Aug 9, 2026

gabotechs reviewed Aug 9, 2026

View reviewed changes

gabotechs reviewed Aug 16, 2026

View reviewed changes

gabotechs reviewed Aug 17, 2026

View reviewed changes

shinzoxD force-pushed the feat-fractional-task-counts-562 branch from 84bacd0 to fe45919 Compare August 19, 2026 18:38

shinzoxD force-pushed the feat-fractional-task-counts-562 branch from fe45919 to 4446393 Compare August 24, 2026 17:51

shinzoxD added 6 commits August 24, 2026 23:59

feat: allow fractional desired task counts

253f356

refactor: address fractional task count review

074debf

test: preserve fractional task count snapshots

f0d4ce7

fix: treat Desired(0) as a valid task count

b949ec4

Empty union children can sum to Desired(0). That is a valid empty stage, not an internal planning error. Return Desired(0.0) from the dynamic planner so UNION children can sum small fractional hints instead of being masked by 1.0.

docs: move fractional task-count migration to 5.0.0

f0f0642

v4.0.0 is now released, so the Desired(f64) public break belongs in the next major upgrade guide rather than the historical 3.0.0 notes.

shinzoxD force-pushed the feat-fractional-task-counts-562 branch from 037b3fd to f0f0642 Compare August 24, 2026 18:30

gabotechs closed this Aug 24, 2026

Uh oh!

Conversation

shinzoxD commented Aug 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Validation

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

gabotechs left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

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!

shinzoxD commented Aug 9, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gabotechs commented Aug 14, 2026

Uh oh!

shinzoxD commented Aug 15, 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!

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!

shinzoxD commented Aug 19, 2026

Uh oh!

shinzoxD commented Aug 24, 2026

Uh oh!

shinzoxD commented Aug 24, 2026

Uh oh!

shinzoxD commented Aug 24, 2026

Uh oh!

shinzoxD commented Aug 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

shinzoxD commented Aug 9, 2026 •

edited

Loading