peter-toth commented on code in PR #56961:
URL: https://github.com/apache/spark/pull/56961#discussion_r3521037274
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -4183,4 +4190,191 @@ class KeyGroupedPartitioningSuite extends
DistributionAndOrderingSuiteBase with
}
}
}
+
+ test("SPARK-57881: storage-partitioned join leverages union output
KeyedPartitioning to " +
Review Comment:
All four tests partition by `identity("id")`. A variant over a bucket
transform (`Array(bucket(4, "id"))`) would cover the PR's motivating
bucketed-source case end-to-end.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:
##########
@@ -938,6 +945,24 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
// Take the output attributes of this union and map the partitioner to
them.
val attributeMap = children.head.output.zip(output).toMap
partitioner match {
+ case headKp: KeyedPartitioning =>
+ // A `UnionExec` concatenates its children's partitions in order
(one child's
+ // partitions after another's), so the merged `KeyedPartitioning`
carries the
+ // concatenation of the children's partition keys, one key per
physical output
+ // partition. Children usually hold different key sets, so the
merged keys often
+ // contain duplicates and `isGrouped` is false; a downstream
`GroupPartitionsExec`
+ // regroups partitions that share a key. The children's
expressions have already
+ // been remapped to the first child's attributes by
`prepareOutputPartitioning`;
+ // here they are remapped to the union's output attributes.
+ val mergedKeys = partitionings.flatMap {
+ case k: KeyedPartitioning => k.partitionKeys
+ case _ => return super.outputPartitioning
+ }
+ val mergedExpressions = headKp.expressions.map(_.transform {
+ case a: Attribute if attributeMap.contains(a) => attributeMap(a)
+ })
+ val isGrouped = mergedKeys.distinct.size == mergedKeys.size
+ KeyedPartitioning(mergedExpressions, mergedKeys, isGrouped)
Review Comment:
The 3-arg constructor leaves `isNarrowed = false`, so the merge drops the
flag whenever a child `KeyedPartitioning` is narrowed (e.g. a union leg that is
a `ProjectExec`/aggregate narrowing a composite partition key -- `ProjectExec`
is a `PartitioningPreservingUnaryExecNode`, `basicPhysicalOperators.scala:49`).
`comparePartitioning` only matches on expressions, so a narrowed child and a
non-narrowed child with the same expressions are still considered compatible
and merged here.
`AliasAwareOutputExpression.scala:135-139` calls this flag out as needing to
be "sticky ... a subsequent [node] that passes all positions through would
otherwise recompute isNarrowed=false, silently dropping the protection" -- this
merge is exactly that dropping. The consequence is in
`KeyedPartitioning.groupedSatisfies` (`partitioning.scala:588-593`): a
narrowed, non-grouped KP is meant to satisfy `ClusteredDistribution` only under
`v2BucketingAllowKeysSubsetOfPartitionKeys`, precisely because
`GroupPartitionsExec` will then merge partitions that held distinct keys in the
finer partitioning (skew risk). Reporting `isNarrowed=false` lets that SPJ
proceed without the opt-in.
Propagate the flag from the children:
```suggestion
val isNarrowed = partitionings.exists {
case k: KeyedPartitioning => k.isNarrowed
case _ => false
}
KeyedPartitioning(mergedExpressions, mergedKeys, isGrouped,
isNarrowed)
```
--
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]
