fix: Enforce data distribution requirements and fix plans that break them - #582
Open
barbarj wants to merge 1 commit into
Open
fix: Enforce data distribution requirements and fix plans that break them#582barbarj wants to merge 1 commit into
barbarj wants to merge 1 commit into
Conversation
…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.
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. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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. Seevalidate_stages.rsfor a more in-depth explanation and justification.There are 5 invalid plan shapes to fix. The following strategies are used:
CollectLeft(join types: Left/LeftSemi/LeftAnti/LeftMark/Full) → changeCollectLeft->Partitioned(keys exist; no replication needed)CollectLeft,null_aware(join type: LeftAnti) -> cap at 1 taskThe
CollectLeft->Partitionedchange 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 duringprepare_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.