dongjoon-hyun commented on code in PR #57491:
URL: https://github.com/apache/spark/pull/57491#discussion_r3646794488
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala:
##########
@@ -937,47 +940,76 @@ case class UnionExec(children: Seq[SparkPlan]) extends
SparkPlan with CodegenSup
}
override def outputPartitioning: Partitioning = {
- if (conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING)) {
- val partitionings = prepareOutputPartitioning()
- if (partitionings.forall(comparePartitioning(_, partitionings.head))) {
- val partitioner = partitionings.head
-
- // 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
- val isNarrowed = partitionings.exists {
- case k: KeyedPartitioning => k.isNarrowed
- case _ => false
- }
- KeyedPartitioning(mergedExpressions, mergedKeys, isGrouped,
isNarrowed)
- case e: Expression =>
- e.transform {
- case a: Attribute if attributeMap.contains(a) => attributeMap(a)
- }.asInstanceOf[Partitioning]
- case _ => partitioner
- }
+ if (!conf.getConf(SQLConf.UNION_OUTPUT_PARTITIONING)) {
+ return super.outputPartitioning
+ }
+
+ // Children's partitionings with attributes remapped to the first child's
attributes.
+ val partitionings = prepareOutputPartitioning()
+ // Map from the first child's attributes to this union's own output
attributes.
+ val attributeMap = children.head.output.zip(output).toMap
+ def toUnionOutput(p: Partitioning): Partitioning = p match {
+ case e: Expression =>
+ e.transform {
+ case a: Attribute if attributeMap.contains(a) => attributeMap(a)
+ }.asInstanceOf[Partitioning]
+ case _ => p
+ }
+
+ // Case A: every child is a single `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. This concatenation (numPartitions
= sum) is a
+ // distinct physical strategy from the co-located pass-through below
(numPartitions = N), so
+ // it is kept as a separate case and never folded into a
`PartitioningCollection`.
+ if (partitionings.forall(_.isInstanceOf[KeyedPartitioning])) {
+ val kps = partitionings.map(_.asInstanceOf[KeyedPartitioning])
+ val headKp = kps.head
+ // The `KeyedPartitioning`s must agree on the partition expressions to
merge.
+ val compatible = kps.forall(comparePartitioning(_, headKp))
+ if (compatible) {
+ val mergedKeys = kps.flatMap(_.partitionKeys)
+ val mergedExpressions = headKp.expressions.map(_.transform {
+ case a: Attribute if attributeMap.contains(a) => attributeMap(a)
+ })
+ val isGrouped = mergedKeys.distinct.size == mergedKeys.size
+ val isNarrowed = kps.exists(_.isNarrowed)
+ return KeyedPartitioning(mergedExpressions, mergedKeys, isGrouped,
isNarrowed)
} else {
- super.outputPartitioning
+ return super.outputPartitioning
}
- } else {
- super.outputPartitioning
+ }
+
+ // Case B: treat each child's partitioning as a set of candidate
partitionings (a
+ // `PartitioningCollection` flattens to its members; a single partitioning
is a one-element
+ // set) and pass through the intersection across all children. Only
index-co-locatable
+ // partitionings participate; `KeyedPartitioning` is excluded here because
its concatenation
+ // semantics (Case A) are incompatible with the co-located union RDD.
+ def flattenPartitioning(p: Partitioning): Seq[Partitioning] = p match {
Review Comment:
Please try to use the existing one.
https://github.com/apache/spark/blob/b65ef3e18fbe2fd499e1c48edd212cae79d377ba/sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala#L152-L160
--
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]
