[SPARK-53406][SQL] Avoid unnecessary shuffle join in direct passthrough shuffle id #52443
Conversation
zhengruifeng
changed the title
| numPartitions: Int) extends Expression with Partitioning with Unevaluable { | ||
|
|
||
| // TODO(SPARK-53401): Support Shuffle Spec in Direct Partition ID Pass Through | ||
| // We don't support creating partitioning for ShufflePartitionIdPassThrough. |
There was a problem hiding this comment.
This comment should be put in ShufflePartitionIdPassThroughSpec#canCreatePartitioning
| ) | ||
|
|
||
| // Mismatched key positions should be incompatible | ||
| val dist1 = ClusteredDistribution(Seq($"a", $"b")) |
There was a problem hiding this comment.
This is the same as dist?
There was a problem hiding this comment.
I was going to wait for address of @cloud-fan's comments :) but i agree this test is a bit hard to read. The variable declaration is not so consistent (ie, on very top or above the check). Also as wenchen point out, some variable like ClusteredDistribution(Seq(
How about, we make some better names too, like ab, cd, a, b? (or a bit longer if necessary)
checkCompatible(
ShufflePartitionIdPassThrough(b, 10).createDist(ab),
ShuffelPartitionIdPassThrough(c, 10).createDist(cd)
expected = false
)
There was a problem hiding this comment.
thanks for calling it out! I followed this and updated all tests in this pr
| newChildren.isDefined | ||
| } | ||
|
|
||
| val isShufflePassThroughCompatible = !isKeyGroupCompatible && |
There was a problem hiding this comment.
shall we combine this with the key group compatibility check?
// Check if the following conditions are satisfied:
// 1. There are exactly two children (e.g., join). Note that Spark doesn't support
// multi-way join at the moment, so this check should be sufficient.
// 2. All children are of the same partitioning, and they are compatible with each other
// If both are true, skip shuffle.
val areChildrenCompatible = parent.isDefined &&
children.length == 2 && childrenIndexes.length == 2 && {
key group check and id pass through check
}
| TransformExpression(BucketFunction, expr, Some(numBuckets)) | ||
| } | ||
|
|
||
| test("ShufflePartitionIdPassThrough - avoid necessary shuffle when they are compatible") { |
There was a problem hiding this comment.
| test("ShufflePartitionIdPassThrough - avoid necessary shuffle when they are compatible") { | |
| test("ShufflePartitionIdPassThrough - avoid unnecessary shuffle when children are compatible") { |
| val plan2 = DummySparkPlan( | ||
| outputPartitioning = ShufflePartitionIdPassThrough(DirectShufflePartitionID(exprC), 5)) | ||
| // Join on different keys than partitioning keys | ||
| val smjExec = SortMergeJoinExec(exprB :: Nil, exprD :: Nil, Inner, None, plan1, plan2) |
There was a problem hiding this comment.
| val smjExec = SortMergeJoinExec(exprB :: Nil, exprD :: Nil, Inner, None, plan1, plan2) | |
| val smjExec = SortMergeJoinExec(exprA :: exprB :: Nil, exprD :: exprC :: Nil, Inner, None, plan1, plan2) |
| case SortMergeJoinExec(_, _, _, _, | ||
| SortExec(_, _, _: DummySparkPlan, _), | ||
| SortExec(_, _, ShuffleExchangeExec(_: HashPartitioning, _, _, _), _), _) => | ||
| // Right side shuffled, left side kept as-is |
There was a problem hiding this comment.
the test result is random and these two cases can both happen?
There was a problem hiding this comment.
Agree, i dont get why the return value match both
| } | ||
| } | ||
|
|
||
| test("ShufflePartitionIdPassThrough - cross position matching behavior") { |
There was a problem hiding this comment.
This looks the same as ShufflePartitionIdPassThrough incompatibility - key position mismatch?
| } | ||
| } | ||
|
|
||
| test("ShufflePartitionIdPassThrough - compatible when partition key matches at any position") { |
There was a problem hiding this comment.
can we merge this test case into ShufflePartitionIdPassThrough - compatible with multiple clustering keys?
| // as the partitioning expression, we check compatibility as follows: | ||
| // 1. Same number of clustering expressions | ||
| // 2. Same number of partitions | ||
| // 3. each pair of partitioning expression from both sides has overlapping positions in their |
There was a problem hiding this comment.
sorry, im a bit new to this and curious, do we need to actually check the partition expression of repartitionByExpression() is compatible or not (like for KeyGroupedPartitioning) before we can skip shuffle, or I miss something ?
There was a problem hiding this comment.
We don't need to, as it's just an expression without any data so it's more similar to the hash partitioning, while key grouped partitioning has actual partition data.
There was a problem hiding this comment.
ah got it, i read the comment on DirectShufflePartitionID that explain the limitation on the child expression, didnt see that earlier.
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
Outdated
Show resolved
Hide resolved
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala
Outdated
Show resolved
Hide resolved
What changes were proposed in this pull request?
This PR implements compatibility checking for ShufflePartitionIdPassThrough partitioning to avoid unnecessary shuffle operations when both sides of a join use compatible direct partition ID pass-through.
Why are the changes needed?
Improve performance
Does this PR introduce any user-facing change?
No
How was this patch tested?
New unit tests
Was this patch authored or co-authored using generative AI tooling?
Yes