Skip to content

fix: Enforce data distribution requirements and fix plans that break them - #582

Open
barbarj wants to merge 1 commit into
datafusion-contrib:mainfrom
paradedb:barbarj.fix-563
Open

fix: Enforce data distribution requirements and fix plans that break them#582
barbarj wants to merge 1 commit into
datafusion-contrib:mainfrom
paradedb:barbarj.fix-563

Conversation

@barbarj

@barbarj barbarj commented Jul 30, 2026

Copy link
Copy Markdown

Closes #563

The bug in 563 is actually a failure to enforce a plan shape invariant for multi-task stages: for every stage, executing its plan once per task over the per-task input assignment and unioning the outputs must be equivalent to executing it once over all the data.

This PR adds a correctness check after plan finalization (validate_stages.rs) and applies fixes for the problematic shapes (normalize_collect_joins.rs).

This runs validate_distributed_stages() as the final step before returning during physical plan construction. We classify the dataflow across each stage, bottom-up, as partitioned or replicated. Partitioned data flow is safe across multiple tasks. Replicated is not. In cases where the input to a task is replicated data, we (based on our knowledge of node semantics) convert that to partitioned if the join will not emit build-side rows, making the replication (i.e. the broadcast build-side) safe as input and the output meaningfully partitioned. See validate_stages.rs for a more in-depth explanation and justification.

There are 5 invalid plan shapes to fix. The following strategies are used:

  • HashJoin CollectLeft (join types: Left/LeftSemi/LeftAnti/LeftMark/Full) → change
    CollectLeft->Partitioned (keys exist; no replication needed)
  • HashJoinCollectLeft, null_aware (join type: LeftAnti) -> cap at 1 task
  • NLJ (join types: Left/LeftSemi/LeftAnti/LeftMark) -> call swap_inputs(), then broadcast as usual. (makes the join shape safe for broadcast).
  • NLJ (join type: Full) -> cap at 1 task
  • CrossJoin (always safe with broadcast) -> cap at 1 task if broadcasts are disabled.

The CollectLeft->Partitioned change causes the rewriting of a bunch of plan shapes in the plan tests. The problematic stages run over small data, and thus collapse to a single task during
prepare_network_boundaries, avoiding the bug. Due to the shape normalization pass necessarily running before task counts are decided, we can't see that the task consolidation will happen and must modify the plan.

In-repo benchmarks show equivalent or better performance for all queries, save one: TPCH q22. It drops by about 40% (~17ms -> ~24ms on my machine, 8 workers, 2 threads). As far as I can tell, it's losing a dynamic filter due to being cut into more stages.

…them

Closes
[datafusion-contrib#563](datafusion-contrib#563)

The bug in 563 is actually a failure to enforce a plan shape invariant
for multi-task stages: for every stage, executing its plan once per task
over the per-task input assignment and unioning the outputs must be
equivalent to executing it once over all the data.

This PR adds a correctness check after plan finalization
(`validate_stages.rs`) and applies fixes for the problematic shapes
(`normalize_collect_joins.rs`).

This runs `validate_distributed_stages()` as the final step before
returning during physical plan construction. We classify the dataflow
across each stage, bottom-up, as partitioned or replicated. Partitioned
data flow is safe across multiple tasks. Replicated is not. In cases
where the input to a task is replicated data, we (based on our knowledge
of node semantics) convert that to partitioned if the join will not emit
build-side rows, making the replication (i.e. the broadcast build-side)
safe as input and the output meaningfully partitioned. See
`validate_stages.rs` for a more in-depth explanation and justification.

There are 5 invalid plan shapes to fix. The following strategies are
used:
- HashJoin `CollectLeft` (join types:
Left/LeftSemi/LeftAnti/LeftMark/Full) → change
`CollectLeft`->`Partitioned` (keys exist; no replication needed)
- HashJoin`CollectLeft`, `null_aware` (join type: LeftAnti) -> cap at 1
task
- NLJ (join types: Left/LeftSemi/LeftAnti/LeftMark) -> call
swap_inputs(), then broadcast as usual. (makes the join shape safe for
broadcast).
- NLJ (join type: Full) -> cap at 1 task
- CrossJoin (always safe with broadcast) -> cap at 1 task if broadcasts
are disabled.

The `CollectLeft`->`Partitioned` change causes the rewriting of a bunch
of plan shapes in the plan tests. The problematic stages run over small
data, and thus collapse to a single task during
`prepare_network_boundaries`, avoiding the bug. Due to the shape
normalization pass necessarily running before task counts are decided,
we can't see that the task consolidation will happen and must modify the
plan.

In-repo benchmarks show equivalent or better performance for all
queries, save one: TPCH q22. It drops by about 40% (~17ms -> ~24ms on my
machine, 8 workers, 2 threads). As far as I can tell, it's losing a
dynamic filter due to being cut into more stages.
@barbarj

barbarj commented Jul 30, 2026

Copy link
Copy Markdown
Author

My apologies for not surfacing the validation design before cutting this PR. If you think that should be done differently, I'm happy to limit this PR to just the shape fixes and do the validation separately.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Bug: CollectLeft HashJoin, CrossJoin, and NestedLoopJoin all return incomplete result sets when run over mutliple tasks

1 participant