stuhood opened a new pull request, #24600: URL: https://github.com/apache/datafusion/pull/24600
## Which issue does this PR close? - Closes #24599. ## Rationale for this change When multi-child operators (such as `HashJoinExec` or `SortMergeJoinExec`) declare co-partitioning requirements via `InputDistributionRequirements`, all children must share a compatible partitioning scheme on their respective join keys. Prior to this change, joining an unpartitioned stream with an input that was already partitioned (such as `Partitioning::Range` or `Partitioning::Hash` from a pre-sorted/partitioned data source scan): 1. During per-child requirement enforcement in [`ensure_distribution`](https://github.com/apache/datafusion/blob/d5552342012888b7d1a3ab88d92e3d292fc0cde0/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs#L1285-L1552), the unpartitioned input failed its `Distribution::KeyPartitioned` requirement and was wrapped in `RepartitionExec: Hash(keys, target_partitions)`. 2. When [`enforce_distribution_relationships`](https://github.com/apache/datafusion/blob/d5552342012888b7d1a3ab88d92e3d292fc0cde0/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs#L1078-L1271) inspected both children, it saw a layout mismatch between `Hash` and `Range`. The code would only create `Hash` partitionings, and so it discarded the existing native `Range` partitioning on the second child and wrapped both children in `RepartitionExec: Hash`, forcing a double shuffle. ## What changes are included in this PR? This PR adjusts `enforce_distribution_relationships` to recognize when one of the children in a co-partitioned group already satisfies its distribution requirement, selecting it as a reference partitioning and adapting unsatisfied peer inputs to match: * Candidate Identification & Tie-Breaking: - Inspect co-partitioned children to find those that satisfy [`input_distributions.child_satisfaction()`](https://github.com/apache/datafusion/blob/d5552342012888b7d1a3ab88d92e3d292fc0cde0/datafusion/physical-plan/src/distribution_requirements.rs#L164-L181). - Prefer native partitioned inputs (e.g. data source scans) over newly injected `RepartitionExec` exchange nodes (since their statistics and partition balance are more likely to be accurate). - For tie-breaking among multiple satisfied candidates, rank candidates using `PlanSize` (`total_byte_size`, then `num_rows` computed via `StatisticsContext`) and select the strictly largest input. * Partitioning Adaptation: - `RangePartitioning`: If the satisfied reference child has `Partitioning::Range(ref_range)`, adapt `ref_range`'s `split_points` and sort options to the unsatisfied child's join expressions after verifying that expression data types match split point scalar values via `split_points_match_expr_types`. - `HashPartitioning`: If the reference child has `Partitioning::Hash(_, count)`, adapt to `Partitioning::Hash(exprs, count)`. - Validate that the reference candidate can be adapted across all unsatisfied children in the co-partitioned group; if any child is incompatible, safely fall back to standard two-sided hash repartitioning. * In-Place Repartition Rewiring: - If the unsatisfied child was already wrapped in an eager `RepartitionExec` during Phase 1, replace its input directly rather than nesting redundant `RepartitionExec` nodes. ## Are these changes tested? Yes, integration tests in `datafusion/core/tests/physical_optimizer/enforce_distribution.rs`: - `range_hash_join_repartitions_unsatisfied_side_to_match_range`: verifies unsatisfied range input adapts to peer reference range partitioning. - `range_hash_join_repartitions_unpartitioned_side_to_match_range`: verifies unpartitioned scan input adapts directly to peer range partitioning without modifying the partitioned side. - `range_hash_join_rehashes_incompatible_data_type`: verifies type mismatch (`Int32` vs `Int64`) safely falls back to two-sided hash repartitioning. ## Are there any user-facing changes? No API changes. -- 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]
