diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala index 688968bda16bf..3567806a04f36 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala @@ -491,14 +491,40 @@ case class CoalescedNullAwareHashPartitioning( * `ClusteredDistribution`, this is the only route by which a grouped KP satisfies one, and no * grouping is involved: the keys are already unique. * - * That second caller is why the narrowing guard in `groupedSatisfies()` is a conjunction with - * `!isGrouped`. A narrowed KP whose projected keys stayed distinct is grouped, and dropping the - * `!isGrouped` term would stop it from satisfying a `ClusteredDistribution` and cost it a shuffle, - * even though grouping it would merge nothing. + * That second caller is why the collapse guard in `groupedSatisfies()` is a conjunction with + * `!isGrouped`. A collapsed KP can end up grouped -- by `GroupPartitionsExec`, or by reducing its + * keys onto a coarser transform -- and dropping the `!isGrouped` term would stop such a KP from + * satisfying a `ClusteredDistribution` and cost it a shuffle, even though grouping it would merge + * nothing. * * For `OrderedDistribution`, `GroupPartitionsExec` must also sort the partition keys to meet the * ordering requirement. * + * == Key Collapse == + * Two things happen to partition keys, and only the first is a loss of granularity: + * + * - '''Key collapse''': a projection or a reduction maps two keys that were distinct onto the same + * new key. `[(1, 'a'), (1, 'b'), (2, 'c')]` projected onto the first position gives `[1, 1, 2]`: + * three distinct keys became two. `isCollapsed` records this. + * - '''Grouping''': `GroupPartitionsExec` physically combines the partitions that share a key. + * `[1, 1, 2]` becomes `[1, 2]`. `isGrouped` says the keys are unique, however they got that way: + * a source with natively unique keys reports it too. + * + * Grouping after a collapse is what produces a partition holding more data than any the source + * declared -- the two `1` partitions above came from different `(1, 'a')` and `(1, 'b')` keys -- so + * it needs `allowKeysSubsetOfPartitionKeys`. Grouping without a collapse only merges partitions + * that already shared a key (a source reporting several splits per key, or a union of children that + * overlap), and needs no opt-in. `OrderedDistribution` is not gated at all: `GroupPartitionsExec` + * pads that path out to the expected split counts rather than coalescing, so nothing is merged. + * + * A collapsed partitioning is still kept rather than dropped to `UnknownPartitioning`, because + * whether it is usable depends on the required distribution, which only `EnsureRequirements` knows. + * The same partitioning is at once: fine for `UnspecifiedDistribution`; fine for + * `ClusteredDistribution` once its keys are unique; usable as the reference side of a one-side + * shuffle, which never consults the flag; and refused only for `ClusteredDistribution` while + * duplicate keys remain and the config is off. Dropping it would discard all of those and make the + * plan shape depend on a config. + * * == Example == * Consider a data source with partition transform `[years(ts_col)]` and 4 input splits: * @@ -528,22 +554,25 @@ case class CoalescedNullAwareHashPartitioning( * guaranteed after projection. May contain duplicates when ungrouped. * @param isGrouped Whether partition keys are unique (no duplicates). Computed on first * creation, then preserved through copy operations to avoid recomputation. - * @param isNarrowed Whether this partitioning was derived from a finer-grained one by dropping key - * positions (e.g. via `PartitioningPreservingUnaryExecNode`). When true and the - * keys are no longer unique, `GroupPartitionsExec` may merge partitions that held - * distinct keys in the original partitioning, carrying the same skew risk as - * `allowKeysSubsetOfPartitionKeys`. "May", because the condition is a proxy: the - * duplicate keys can also come from a source that reports several splits per - * partition key, in which case grouping merges only same-key partitions. Such a - * partitioning can only satisfy `ClusteredDistribution` by being grouped, and - * `groupedSatisfies` refuses that unless the config is enabled, regardless of - * `requireAllClusterKeysForDistribution`. + * @param isCollapsed Whether a projection or a reduction mapped keys that were distinct in the + * partitioning this one was derived from onto the same key, so one partition + * here can stand for several of the original ones -- see "Key Collapse" above. + * Dropping key positions does not set it on its own; the projected keys have to + * actually lose distinctness. Sticky, because neither grouping nor a further + * projection can make a partitioning finer again. One case sets it without a + * collapse of its own: the side shuffled onto a collapsed partitioning's keys + * inherits it, because the two are then co-located on that key set -- see + * `KeyedShuffleSpec.createPartitioning`. + * Together with `!isGrouped` it decides whether `groupedSatisfies` may coalesce + * the duplicate keys without `allowKeysSubsetOfPartitionKeys`: `isCollapsed` + * says the collapse happened, `!isGrouped` says there is still something left + * to merge, and only both together mean there is an outstanding risk to gate. */ case class KeyedPartitioning( expressions: Seq[Expression], @transient partitionKeys: Seq[InternalRowComparableWrapper], isGrouped: Boolean, - isNarrowed: Boolean = false) extends Expression with Partitioning with Unevaluable { + isCollapsed: Boolean) extends Expression with Partitioning with Unevaluable { override val numPartitions = partitionKeys.length override def children: Seq[Expression] = expressions @@ -561,10 +590,30 @@ case class KeyedPartitioning( @transient lazy val keyOrdering = keyRowOrdering.on((t: InternalRowComparableWrapper) => t.row) + /** + * Number of distinct partition keys. Free when this partitioning is grouped, and computed on + * demand otherwise. Used to tell an actual key collapse from a projection that merely dropped key + * positions. + */ + @transient lazy val distinctKeyCount: Int = + if (isGrouped) numPartitions else partitionKeys.distinct.length + + /** + * Whether a projection of this partitioning that leaves `projectedDistinctKeyCount` distinct keys + * mapped keys that were distinct here onto the same key. Callers carry the inherited flag + * themselves, since a projection can also read it from partitionings other than this one; + * this is the count comparison they share. `GroupPartitionsExec` does not use it: it answers the + * same question exactly, from the key groups it keeps. + */ + def collapsesOnProjection(projectedDistinctKeyCount: Int): Boolean = + projectedDistinctKeyCount < distinctKeyCount + def toGrouped: KeyedPartitioning = { val groupedPartitionKeys = partitionKeys.distinct.sorted(keyOrdering) - new KeyedPartitioning(expressions, groupedPartitionKeys, isGrouped = true) + // Grouping removes the duplicate keys; it does not make this partitioning finer than the layout + // it came from, so `isCollapsed` travels with it. + new KeyedPartitioning(expressions, groupedPartitionKeys, isGrouped = true, isCollapsed) } /** @@ -595,18 +644,12 @@ case class KeyedPartitioning( val allowKeysSubsetOfPartitionKeys = SQLConf.get.v2BucketingAllowKeysSubsetOfPartitionKeys - if (isNarrowed && !isGrouped && !allowKeysSubsetOfPartitionKeys) { - // A narrowed, non-grouped partitioning carries the same skew risk as using a subset of - // partition keys for a join: GroupPartitionsExec may merge partitions that held distinct - // keys in the original finer-grained partitioning. Require the same config to opt in. - // + if (isCollapsed && !isGrouped && !allowKeysSubsetOfPartitionKeys) { + // Coarser than the layout it came from, and duplicate keys are still there, so grouping + // would merge partitions the finer-grained partitioning held apart. Both terms are + // load-bearing and `OrderedDistribution` is deliberately not gated -- see the class doc. // Checked before the `requireAllClusterKeys` branch, because the risk is independent of // which key sets count as matching (SPARK-58974). - // - // The `!isGrouped` term is required, not redundant: a narrowing projection whose - // projected keys stayed distinct is grouped, and `satisfies0` routes such a partitioning - // through here (`nonGroupedSatisfies` is false for a `ClusteredDistribution`). Grouping - // it would merge nothing, so refusing it would only cost a shuffle. false } else if (requireAllClusterKeys) { // Checks whether this partitioning is partitioned on exactly same clustering keys of @@ -643,12 +686,23 @@ case class KeyedPartitioning( val joinKeyPositions = result.keyPositions.map(_.nonEmpty).zipWithIndex.filter(_._1).map(_._2) val projectedExpressions = joinKeyPositions.map(expressions) val projectedKeys = projectKeys(joinKeyPositions)._2 - // Sort the distinct projected keys the same way `GroupPartitionsExec` does (both sort with - // `KeyedPartitioning.groupedKeyRowOrdering`). Otherwise, when only the keyed side is grouped - // and the other side is re-shuffled using this spec, the two `KeyedPartitioning`s carry the - // same keys in a different order and `PartitioningCollection.fromPartitionings` rejects them. - val projectedPartitioning = - new KeyedPartitioning(projectedExpressions, projectedKeys, isGrouped = false).toGrouped + // `toGrouped` dedups and sorts the keys the same way `GroupPartitionsExec` does (both sort + // with `KeyedPartitioning.groupedKeyRowOrdering`). Otherwise, when only the keyed side is + // grouped and the other side is re-shuffled using this spec, the two `KeyedPartitioning`s + // carry the same keys in a different order and `PartitioningCollection.fromPartitionings` + // rejects them. Its `distinct` is also the only one needed here: the partition count it + // leaves is the projected distinct key count the collapse test asks for. + val grouped = new KeyedPartitioning( + projectedExpressions, projectedKeys, isGrouped = false, isCollapsed = false).toGrouped + // Projecting onto the operation keys can collapse keys in its own right. Dropping no position + // cannot, so the counts are only compared when one was dropped. The gate in + // `groupedSatisfies` is bypassed while this config is on, so the flag decides nothing here + // today, but it travels with the partitioning, and leaving a producer to launder it is how + // the protection went missing. + val projectedCollapsed = isCollapsed || + (joinKeyPositions.length < expressions.length && + collapsesOnProjection(grouped.numPartitions)) + val projectedPartitioning = grouped.copy(isCollapsed = projectedCollapsed) result.copy(partitioning = projectedPartitioning, joinKeyPositions = Some(joinKeyPositions)) } else { result @@ -669,7 +723,9 @@ object KeyedPartitioning { InternalRowComparableWrapper.getInternalRowComparableWrapperFactory(dataTypes) val comparablePartitionKeys = partitionKeys.map(comparableKeyWrapperFactory) val isGrouped = comparablePartitionKeys.distinct.size == comparablePartitionKeys.size - new KeyedPartitioning(expressions, comparablePartitionKeys, isGrouped) + // A partitioning built from scratch, e.g. as a data source reports it, is the layout everything + // else is compared against, so nothing has collapsed yet. + new KeyedPartitioning(expressions, comparablePartitionKeys, isGrouped, isCollapsed = false) } def supportsExpressions(expressions: Seq[Expression]): Boolean = { @@ -833,10 +889,14 @@ case class RangePartitioning(ordering: Seq[SortOrder], numPartitions: Int) * Outer Join operators. * * [[KeyedPartitioning]]s within a `PartitioningCollection` describe the same physical partitioning - * and so must share the same `partitionKeys` reference, differing only in their `expressions` (with - * matching arity). Use [[PartitioningCollection.fromPartitionings]] to build a collection from - * independently-computed partitionings (e.g. join `outputPartitioning`); it interns `partitionKeys` - * references (including across nested collections) so the invariant holds. + * and so must share the same `partitionKeys` reference and `isCollapsed` flag, differing only in + * their `expressions` (with matching arity). Use + * [[PartitioningCollection.fromPartitionings]] to build a collection from independently-computed + * partitionings (e.g. join `outputPartitioning`); it interns `partitionKeys` references and + * normalizes `isCollapsed` by OR (including across nested collections) so the invariant holds. + * Uniformity matters because consumers read the flag off one member: `satisfies0` and + * `EnsureRequirements` accept when any single member satisfies the distribution, so a member that + * under-reported the collapse would let the gate through. */ case class PartitioningCollection(partitionings: Seq[Partitioning]) extends Expression with Partitioning with Unevaluable { @@ -851,8 +911,10 @@ case class PartitioningCollection(partitionings: Seq[Partitioning]) * First [[KeyedPartitioning]] reachable from this collection through direct members or nested * collections, if any. Since every collection validates the invariant on construction, this * single representative stands for all [[KeyedPartitioning]]s in the subtree: they all share - * its `partitionKeys` reference and expression arity. The invariant check forces this lazy val - * during construction, so it is only recomputed after deserialization. + * its `partitionKeys` reference, its expression arity and its `isCollapsed` flag. The invariant + * check forces this lazy val during construction, so it is only recomputed after deserialization. + * Consumers that need the flag or the distinct key count should read one representative rather + * than every member, since the answer is the same and the count costs a pass over the keys. */ @transient private[physical] lazy val firstKeyedPartitioning: Option[KeyedPartitioning] = partitionings.view.map { @@ -883,6 +945,8 @@ case class PartitioningCollection(partitionings: Seq[Partitioning]) require(representative.partitionKeys eq first.partitionKeys, "All KeyedPartitionings in a PartitioningCollection must share the same " + "partitionKeys reference") + require(representative.isCollapsed == first.isCollapsed, + "All KeyedPartitionings in a PartitioningCollection must agree on isCollapsed") } } } @@ -929,9 +993,25 @@ object PartitioningCollection { * Note: this can't be implemented with `TreeNode.transform`. */ def fromPartitionings(partitionings: Seq[Partitioning]): PartitioningCollection = { + def representativeOf(p: Partitioning): Option[KeyedPartitioning] = p match { + case k: KeyedPartitioning => Some(k) + case pc: PartitioningCollection => pc.firstKeyedPartitioning + case _ => None + } + + // `isCollapsed` describes the shared physical layout, not one member's naming of it: if any + // member is coarser than the layout it was derived from, an output partition really does cover + // several of the finer ones, whichever member's expressions name it. So normalize the flag by + // OR, the same way `partitionKeys` references are interned below. One representative per member + // is enough, because every collection agrees on the flag internally by this same construction, + // which keeps this pass O(members) for an already-uniform collection; a subtree that disagrees + // is rebuilt. + val anyCollapsed = partitionings.exists(representativeOf(_).exists(_.isCollapsed)) + var canonicalKeys: Seq[InternalRowComparableWrapper] = null def intern(p: Partitioning): Partitioning = p match { - case k: KeyedPartitioning => + case keyed: KeyedPartitioning => + val k = if (anyCollapsed && !keyed.isCollapsed) keyed.copy(isCollapsed = true) else keyed if (canonicalKeys == null) { canonicalKeys = k.partitionKeys k @@ -948,17 +1028,23 @@ object PartitioningCollection { // collection as-is keeps repeated outputPartitioning computations over deeply nested // collections (e.g. chains of same-key joins) O(1) per level. case None => pc - case Some(representative) if canonicalKeys == null => - canonicalKeys = representative.partitionKeys - pc - // The collection's own invariant guarantees all its KeyedPartitionings share the - // representative's `partitionKeys` reference, so reference-equality of the - // representative's keys means the whole subtree is already interned. - case Some(representative) if representative.partitionKeys eq canonicalKeys => pc case Some(representative) => - require(representative.partitionKeys == canonicalKeys, - "All KeyedPartitionings in a PartitioningCollection must have equal partitionKeys") - new PartitioningCollection(pc.partitionings.map(intern)) + // Descend only when this subtree disagrees with the collection, on the keys or on the + // flag; otherwise it is already normalized and can be returned as-is. + val fixCollapsed = anyCollapsed && !representative.isCollapsed + if (canonicalKeys == null) { + canonicalKeys = representative.partitionKeys + if (fixCollapsed) new PartitioningCollection(pc.partitionings.map(intern)) else pc + } else if ((representative.partitionKeys eq canonicalKeys) && !fixCollapsed) { + // The collection's own invariant guarantees all its KeyedPartitionings share the + // representative's `partitionKeys` reference and its flag, so reference-equality of + // the representative's keys means the whole subtree is already interned. + pc + } else { + require(representative.partitionKeys == canonicalKeys, + "All KeyedPartitionings in a PartitioningCollection must have equal partitionKeys") + new PartitioningCollection(pc.partitionings.map(intern)) + } } case other => other } @@ -1445,7 +1531,13 @@ case class KeyedShuffleSpec( te.copy(children = te.children.map(_ => clustering(positionSet.head))) case (_, positionSet) => clustering(positionSet.head) } - KeyedPartitioning(newExpressions, partitioning.partitionKeys, partitioning.isGrouped) + // The shuffled side is laid out on this side's partition keys, so it inherits the flag. + // Strictly nothing collapsed on this side -- its partitions are what a hash partitioning would + // give -- so this is deliberate conservatism: the two sides are co-located on one key set, and + // a later grouping of that key set carries the collapsed side's risk. It can only add shuffles, + // never remove one. + KeyedPartitioning(newExpressions, partitioning.partitionKeys, partitioning.isGrouped, + partitioning.isCollapsed) } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala index 9398d14eb94e8..d15fc0fd1af65 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala @@ -2515,9 +2515,11 @@ object SQLConf { "tables. At planning time, Spark will group the partitions by only those keys that are " + "in the operation's keys. That is currently enabled only if " + s"${REQUIRE_ALL_CLUSTER_KEYS_FOR_DISTRIBUTION.key} is false. This config also gates " + - "grouping a partitioning that was narrowed to a subset of its keys and whose keys are no " + - "longer distinct, which carries the same risk of skew; that applies regardless of " + - s"${REQUIRE_ALL_CLUSTER_KEYS_FOR_DISTRIBUTION.key}." + "grouping a partitioning whose keys collapsed, that is, where a projection or a " + + "reduction mapped keys that were distinct in the source onto the same key, so that " + + "grouping them would produce a partition larger than any the source declared. That " + + s"applies regardless of ${REQUIRE_ALL_CLUSTER_KEYS_FOR_DISTRIBUTION.key}. It does not " + + "apply to duplicate keys the source itself reported, which are grouped without this config." ) .version("4.0.0") .withBindingPolicy(ConfigBindingPolicy.SESSION) diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/DistributionSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/DistributionSuite.scala index 2e9e1270b3fc5..0a8a364c2d2a6 100644 --- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/DistributionSuite.scala +++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/DistributionSuite.scala @@ -466,4 +466,22 @@ class DistributionSuite extends SparkFunSuite { } assert(arityMismatch.getMessage.contains("matching expression arity")) } + + test("SPARK-59057: toGrouped and KeyedShuffleSpec.createPartitioning keep isCollapsed sticky") { + val x = AttributeReference("x", IntegerType)() + val y = AttributeReference("y", IntegerType)() + + val collapsedKP = KeyedPartitioning(Seq(x), Seq(InternalRow(1), InternalRow(1), InternalRow(2))) + .copy(isCollapsed = true) + // The protection this flag provides: a coarsened, non-grouped KP must not satisfy + // ClusteredDistribution via grouping without allowKeysSubsetOfPartitionKeys. + assert(!collapsedKP.groupedSatisfies(ClusteredDistribution(Seq(x)))) + + assert(collapsedKP.toGrouped.isCollapsed, "toGrouped must keep isCollapsed sticky") + + val spec = KeyedShuffleSpec(collapsedKP, ClusteredDistribution(Seq(x))) + val created = spec.createPartitioning(Seq(y)).asInstanceOf[KeyedPartitioning] + assert(created.isCollapsed, "createPartitioning must keep isCollapsed sticky: the shuffled " + + "side is laid out on the collapsed side's partition keys and inherits its skew risk") + } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala index e5d293f6e2193..99c5ced678d14 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/AliasAwareOutputExpression.scala @@ -137,12 +137,24 @@ trait PartitioningPreservingUnaryExecNode extends UnaryExecNode if (projectablePositions.length == numPositions) keySource.partitionKeys else keySource.projectKeys(projectablePositions)._2 - val isGrouped = sharedKeys.distinct.size == sharedKeys.size - // A KP is narrowed if this node drops positions, or if the input KPs were already narrowed - // (i.e. came from a finer-grained partitioning). The flag must be sticky: a subsequent - // PartitioningPreservingUnaryExecNode that passes all positions through would otherwise - // recompute isNarrowed=false, silently dropping the protection. - val isNarrowed = projectablePositions.length < numPositions || keySource.isNarrowed + val distinctSharedKeys = sharedKeys.distinct + val isGrouped = distinctSharedKeys.size == sharedKeys.size + // This projection collapses keys when it maps keys that were distinct in the input + // onto the same projected key -- dropping positions is not enough on its own, since the + // projected keys can stay just as distinct as the originals. The flag is sticky: a subsequent + // PartitioningPreservingUnaryExecNode that passes all positions through must not recompute it + // as false and drop the protection, and no projection can make a partitioning finer again. + // + // Both cheap terms come first: an inherited flag or a projection that drops no position + // settles the question without counting distinct keys. A pass-through projection cannot + // collapse anything, since it keeps the input's keys as they are. + // + // The inherited flag is read from all inputs rather than from the key source alone. A + // `PartitioningCollection` normalizes it across its members, so the two agree today; reading + // all of them keeps this producer correct without depending on that. + val isCollapsed = kps.exists(_.isCollapsed) || + (projectablePositions.length < numPositions && + keySource.collapsesOnProjection(distinctSharedKeys.size)) // Cross-product the per-position alternatives to produce all concrete KPs. // Note: generateCartesianProduct expects thunks () => Seq[T], but wrapping LazyLists in thunks @@ -152,7 +164,7 @@ trait PartitioningPreservingUnaryExecNode extends UnaryExecNode MultiTransform.generateCartesianProduct( projectablePositions.map(i => () => alternativesPerPosition(i))) .map(projectedExprs => - new KeyedPartitioning(projectedExprs, sharedKeys, isGrouped, isNarrowed)) + new KeyedPartitioning(projectedExprs, sharedKeys, isGrouped, isCollapsed)) } } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala index bea86501e6f3a..9e161d6f16173 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala @@ -989,8 +989,12 @@ case class UnionExec(children: Seq[SparkPlan]) extends SparkPlan with CodegenSup val mergedKeys = kps.flatMap(_.partitionKeys) val mergedExpressions = headKp.expressions val isGrouped = mergedKeys.distinct.size == mergedKeys.size - val isNarrowed = kps.exists(_.isNarrowed) - return KeyedPartitioning(mergedExpressions, mergedKeys, isGrouped, isNarrowed) + // A collapse is a property of each child's own keys: a child key that stands for several + // finer-grained keys keeps standing for them in the concatenation. Keys repeating *across* + // children is not a collapse -- grouping merges partitions that share a key, which is what + // it does for any partitioning that never went through one, and needs no opt-in. + val isCollapsed = kps.exists(_.isCollapsed) + return KeyedPartitioning(mergedExpressions, mergedKeys, isGrouped, isCollapsed) } else { return super.outputPartitioning } diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala index 6f671bfcb80be..3b37375328154 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala @@ -68,13 +68,20 @@ case class GroupPartitionsExec( child.outputPartitioning match { case p: Partitioning with Expression => // There can be multiple `KeyedPartitioning`s in an output partitioning of a join, but they - // can only differ in `expressions`; their `partitionKeys` reference is shared (enforced by - // `PartitioningCollection`), so `groupedPartitions` is computed only once. + // can only differ in `expressions`; their `partitionKeys` reference and `isCollapsed` flag + // are shared (enforced by `PartitioningCollection`), so both `groupedPartitions` and the + // new flag are computed once, outside the transform. val partitionKeys = groupedPartitions.map(_._1) + // `isCollapsed` is sticky: grouping removes the duplicate keys, but it does not make this + // partitioning any finer than the layout it came from. Projecting onto the join key + // positions, or reducing the keys onto a coarser transform, can collapse keys in its own + // right, which `collapsesKeys` answers exactly -- it is computed with the grouping itself, + // so it costs nothing per call and needs no comparison of key counts. + val isCollapsed = childIsCollapsed || collapsesKeys p.transform { case k: KeyedPartitioning => val projectedExpressions = joinKeyPositions.fold(k.expressions)(_.map(k.expressions)) - KeyedPartitioning(projectedExpressions, partitionKeys, isGrouped = isGrouped) + KeyedPartitioning(projectedExpressions, partitionKeys, isGrouped, isCollapsed) }.asInstanceOf[Partitioning] case o => o } @@ -112,6 +119,16 @@ case class GroupPartitionsExec( keyMap.toSeq.sorted(keyOrdering.on((t: (InternalRowComparableWrapper, _)) => t._1.row)) } + // There must be a `KeyedPartitioning` in the child's output partitioning, as a + // `GroupPartitionsExec` node is added to a plan only in that case. A collection's members share + // their `partitionKeys` reference and their `isCollapsed` flag, so any one of them will do. + @transient private lazy val childKeyedPartitioning: KeyedPartitioning = + child.outputPartitioning + .asInstanceOf[Partitioning with Expression] + .collectFirst { case k: KeyedPartitioning => k } + .getOrElse( + throw new SparkException("GroupPartitionsExec requires a child with KeyedPartitioning")) + /** * Computes the grouped partitions by: * 1. Projecting partition keys if joinKeyPositions is specified @@ -119,19 +136,15 @@ case class GroupPartitionsExec( * 3. Grouping input partition indices by their (possibly projected/reduced) keys * 4. Sorting or distributing based on whether partial clustering is enabled * - * Returns a tuple of (partitions, isGrouped) where: + * Returns a tuple of (partitions, isGrouped, collapsesKeys) where: * - partitions: sequence of (partitionKey, inputPartitionIndices) pairs representing * how input partitions should be grouped together * - isGrouped: whether the output partitioning is grouped (no duplicates in partition keys) + * - collapsesKeys: whether any group this node outputs covers more than one of the child's own + * partition keys, i.e. whether the projection or the reduction actually merged keys */ @transient private lazy val groupedPartitionsTuple = { - // There must be a `KeyedPartitioning` in child's output partitioning as a - // `GroupPartitionsExec` node is added to a plan only in that case. - val keyedPartitioning = child.outputPartitioning - .asInstanceOf[Partitioning with Expression] - .collectFirst { case k: KeyedPartitioning => k } - .getOrElse( - throw new SparkException("GroupPartitionsExec requires a child with KeyedPartitioning")) + val keyedPartitioning = childKeyedPartitioning // Project partition keys if join key positions are specified val (projectedDataTypes, projectedKeys) = @@ -145,10 +158,26 @@ case class GroupPartitionsExec( val keyToPartitionIndices = reducedKeys.zipWithIndex.groupMap(_._1)(_._2) + // Whether this node collapses keys: does any key it keeps stand for more than one of the + // child's own partition keys? Counting the child's *keys* rather than its partitions is what + // tells a collapse from a source reporting several splits per key, and asking it of the keys + // this node keeps is what tells it from `alignToExpectedKeys` dropping keys, which merges + // nothing. Ask the key groups, not the partitions finally emitted: `distributePartitions` + // spreads a group's splits over one partition each, which would hide the merge, and + // replication would ask about the same group repeatedly. + val childKeys = keyedPartitioning.partitionKeys.toIndexedSeq + def coversSeveralChildKeys(indices: Seq[Int]): Boolean = + indices.map(childKeys).distinct.size > 1 + if (expectedPartitionKeys.isDefined) { - alignToExpectedKeys(keyToPartitionIndices) + val (alignedPartitions, grouped) = alignToExpectedKeys(keyToPartitionIndices) + val keptGroups = expectedPartitionKeys.get.map { case (key, _) => + keyToPartitionIndices.getOrElse(key, Seq.empty) + } + (alignedPartitions, grouped, keptGroups.exists(coversSeveralChildKeys)) } else { - (groupAndSortByKeys(keyToPartitionIndices, reducedDataTypes), true) + val sorted = groupAndSortByKeys(keyToPartitionIndices, reducedDataTypes) + (sorted, true, sorted.map(_._2).exists(coversSeveralChildKeys)) } } @@ -157,6 +186,12 @@ case class GroupPartitionsExec( @transient lazy val isGrouped: Boolean = groupedPartitionsTuple._2 + /** Whether the grouping this node performs merges keys the child held apart. */ + @transient private lazy val collapsesKeys: Boolean = groupedPartitionsTuple._3 + + /** The child's own flag, from the same `KeyedPartitioning` the grouping was computed from. */ + @transient private lazy val childIsCollapsed: Boolean = childKeyedPartitioning.isCollapsed + @transient private lazy val hasCoalescing: Boolean = groupedPartitions.exists(_._2.size > 1) // Whether the child subtree is safe to use with SortedMergeCoalescedRDD (k-way merge). diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala index c0a1faa032635..506f781ed471a 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala @@ -499,7 +499,7 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession { val partitioning = KeyedPartitioning( Seq(partAttr), Seq(InternalRowComparableWrapper(InternalRow(1), Seq(partAttr))), - isGrouped = false) + isGrouped = false, isCollapsed = false) def replanAfterFiltering(afterFilter: Seq[InputPartition]): Unit = { val scan = new PartitioningBreakingScan(Seq(KeyedInputPartition(1)), afterFilter) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/DistributionAndOrderingSuiteBase.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/DistributionAndOrderingSuiteBase.scala index 98ed8bf0b3d0d..f2a51cc5ce4d6 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/DistributionAndOrderingSuiteBase.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/DistributionAndOrderingSuiteBase.scala @@ -53,8 +53,9 @@ abstract class DistributionAndOrderingSuiteBase plan: QueryPlan[T]): Partitioning = partitioning match { case HashPartitioning(exprs, numPartitions) => HashPartitioning(exprs.map(resolveAttrs(_, plan)), numPartitions) - case KeyedPartitioning(expressions, partitionKeys, isGrouped, _) => - KeyedPartitioning(expressions.map(resolveAttrs(_, plan)), partitionKeys, isGrouped) + case KeyedPartitioning(expressions, partitionKeys, isGrouped, isCollapsed) => + KeyedPartitioning( + expressions.map(resolveAttrs(_, plan)), partitionKeys, isGrouped, isCollapsed) case PartitioningCollection(partitionings) => PartitioningCollection(partitionings.map(resolvePartitioning(_, plan))) case RangePartitioning(ordering, numPartitions) => diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala index 948de99959544..046f10bad8008 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala @@ -4243,7 +4243,7 @@ class KeyGroupedPartitioningSuite test("SPARK-46367: narrowing projection requires allowKeysSubsetOfPartitionKeys") { // items is partitioned by (id, name). The subquery projects away 'name', narrowing - // KeyedPartitioning([id, name]) -> KeyedPartitioning([id]) with isNarrowed=true. + // KeyedPartitioning([id, name]) -> KeyedPartitioning([id]) with isCollapsed=true. // Because id=1 maps to two original partitions ("aa" and "bb"), isGrouped=false. // GroupPartitionsExec would merge them, carrying the same skew risk as subset partition // keys -- so SPJ requires allowKeysSubsetOfPartitionKeys to be enabled. @@ -4290,9 +4290,9 @@ class KeyGroupedPartitioningSuite test("SPARK-46367: narrowing projection with distinct projected keys does not require " + "allowKeysSubsetOfPartitionKeys") { // items is partitioned by (id, name) but each id value is unique, so projecting away 'name' - // produces KeyedPartitioning([id]) with isNarrowed=true but isGrouped=true. - // Because no two original partitions share the same projected key, GroupPartitionsExec does not - // merge any partitions -- no skew risk -- so SPJ works without config. + // produces KeyedPartitioning([id]) with isGrouped=true and isCollapsed=false: no two original + // partitions share the same projected key, so the projection lost no distinct key and + // GroupPartitionsExec would merge nothing. No skew risk, so SPJ works without config. val items_partitions = Array(identity("id"), identity("name")) createTable(items, itemsColumns, items_partitions) sql(s"INSERT INTO testcat.ns.$items VALUES " + @@ -4331,7 +4331,7 @@ class KeyGroupedPartitioningSuite "with allowKeysSubsetOfPartitionKeys") { // Table partitioned by (id, name): id=1 maps to two distinct partition keys (1,'aa') and // (1,'bb'). The partial HashAggregate (a PartitioningPreservingUnaryExecNode) projects away - // 'name', narrowing the KP from [id,name] to KP([id], isNarrowed=true, isGrouped=false). + // 'name', narrowing the KP from [id,name] to KP([id], isCollapsed=true, isGrouped=false). // By default a shuffle is required; with allowKeysSubsetOfPartitionKeys enabled, // EnsureRequirements inserts GroupPartitionsExec to coalesce both id=1 partitions so the final // aggregate sees all id=1 partial results in one task -- correct and shuffle-free. @@ -4367,7 +4367,7 @@ class KeyGroupedPartitioningSuite "with allowKeysSubsetOfPartitionKeys") { // Same narrowing mechanism as the aggregate test: the partial HashAggregate (a // PartitioningPreservingUnaryExecNode) for the inner GROUP BY id, price projects away 'name', - // narrowing KP([id,name]) to KP([id], isNarrowed=true, isGrouped=false). With + // narrowing KP([id,name]) to KP([id], isCollapsed=true, isGrouped=false). With // allowKeysSubsetOfPartitionKeys enabled, EnsureRequirements inserts GroupPartitionsExec to // coalesce both id=1 partitions for the final aggregate. The window PARTITION BY id then sees // KP([id], isGrouped=true) from the aggregate output and needs no further exchange. @@ -4915,11 +4915,11 @@ class KeyGroupedPartitioningSuite val df = sql(query) val plan = df.queryExecution.executedPlan - val narrowed = collect(plan) { case p: ProjectExec => p } + val collapsed = collect(plan) { case p: ProjectExec => p } .map(_.outputPartitioning) - .collect { case kp: physical.KeyedPartitioning if kp.isNarrowed => kp } - assert(narrowed.exists(!_.isGrouped), - "this test needs a narrowed, ungrouped partitioning to be meaningful") + .collect { case kp: physical.KeyedPartitioning if kp.isCollapsed => kp } + assert(collapsed.exists(!_.isGrouped), + "this test needs a collapsed, ungrouped partitioning to be meaningful") // Count over the whole plan, so a shuffle or a grouping anywhere in it is caught, not // only the ones inside the join subtree. @@ -4946,6 +4946,261 @@ class KeyGroupedPartitioningSuite } } } + + test("several splits per partition key are grouped without allowKeysSubsetOfPartitionKeys") { + // The scan reports one partition key per split, so a table with two splits for the same + // (id, dept) value already has duplicate keys before any projection. Dropping dept maps + // (1, 'x'), (1, 'x') onto 1 and (2, 'y') onto 2: two distinct keys before, two after, so the + // projection coarsened nothing. Grouping merges only the two splits that already shared a key, + // which is what happens for any partitioning that never went through a projection, so it must + // not need the opt-in. + val cols = Array( + Column.create("id", LongType), + Column.create("dept", StringType), + Column.create("data", StringType)) + val t2cols = Array(Column.create("id", LongType), Column.create("data", StringType)) + withTable("t1", "t2") { + createTable("t1", cols, Array(identity("id"), identity("dept"))) + sql("INSERT INTO testcat.ns.t1 VALUES (1, 'x', 'a1'), (1, 'x', 'a2'), (2, 'y', 'a3')") + createTable("t2", t2cols, Array(identity("id"))) + sql("INSERT INTO testcat.ns.t2 VALUES (1, 'b1'), (2, 'b2')") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "false") { + // `dept RLIKE ...` has no V2 translation, so the scan must output dept and the Project + // above the Filter is what drops it. + val df = sql( + """SELECT /*+ MERGE(u, t2) */ u.id, t2.data + |FROM (SELECT id FROM testcat.ns.t1 WHERE dept RLIKE 'x|y') u + |JOIN testcat.ns.t2 ON u.id = t2.id + |""".stripMargin) + val plan = df.queryExecution.executedPlan + + val projected = collect(plan) { case p: ProjectExec => p } + .map(_.outputPartitioning) + .collect { case kp: physical.KeyedPartitioning => kp } + assert(projected.exists(kp => !kp.isGrouped && !kp.isCollapsed), + "this test needs an ungrouped partitioning that the projection did not coarsen") + + assert(collectAllGroupPartitions(plan).nonEmpty, + "the duplicate splits must be grouped, no opt-in needed") + assert(collectAllShuffles(plan).isEmpty, "grouping must replace the shuffles") + checkAnswer(df, Seq(Row(1, "b1"), Row(1, "b1"), Row(2, "b2"))) + } + } + } + + test("filtering partition keys out is not a key collapse") { + // With partition filtering on, an inner join plans both sides on the intersection of their + // partition keys, so a side that had more keys ends up with fewer than it reported. That is + // pruning, not merging: every remaining partition still holds exactly one of its own keys. + // Counting it as a collapse would make the sticky flag refuse grouping further up the plan. + val cols = Array(Column.create("id", LongType), Column.create("data", StringType)) + withTable("t1", "t2", "t3") { + createTable("t1", cols, Array(identity("id"))) + sql("INSERT INTO testcat.ns.t1 VALUES (1, 'a1'), (2, 'a2'), (3, 'a3')") + createTable("t2", cols, Array(identity("id"))) + sql("INSERT INTO testcat.ns.t2 VALUES (1, 'b1'), (2, 'b2')") + createTable("t3", cols, Array(identity("id"))) + sql("INSERT INTO testcat.ns.t3 VALUES (1, 'c1'), (2, 'c2')") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.V2_BUCKETING_PARTITION_FILTER_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "false") { + val df = sql( + """SELECT id, count(*) AS cnt FROM ( + | SELECT /*+ MERGE(t1, t2) */ t1.id FROM testcat.ns.t1 JOIN testcat.ns.t2 + | ON t1.id = t2.id + | UNION ALL + | SELECT id FROM testcat.ns.t3 + |) GROUP BY id + |""".stripMargin) + val plan = df.queryExecution.executedPlan + + val groupPartitions = collectAllGroupPartitions(plan) + val keyed = groupPartitions.map(_.outputPartitioning) + .flatMap(physical.PartitioningCollection.flatten) + .collect { case kp: physical.KeyedPartitioning => kp } + assert(keyed.nonEmpty, "the join must be planned as a storage-partitioned join") + assert(keyed.forall(!_.isCollapsed), + "key 3 was filtered out, not merged into another partition") + assert(collectAllShuffles(plan).isEmpty, + "nothing was coarsened, so the aggregate above the union must not need a shuffle") + checkAnswer(df, Seq(Row(1, 2), Row(2, 2))) + } + } + } + + test("SPARK-59057: a shuffle built from a template that collapsed nothing needs no opt-in") { + // The chain apache#58316 built for the shuffle-template path, with the expectation the collapse + // semantics call for. items is partitioned by (id, name) with *unique* ids, so projecting + // `name` away drops a key position but loses no distinct key: nothing collapsed. purchases + // reports no + // partitioning, so with v2BucketingShuffleEnabled its side is shuffled using the projected + // partitioning as the template; a RIGHT OUTER join exposes only that shuffled side, and the + // union with t3 (overlapping keys) makes the merged partitioning ungrouped. The duplicate keys + // there come from the two children holding the same ids, not from the projection, so the final + // aggregate may group them with the opt-in off. + createTable(items, itemsColumns, Array(identity("id"), identity("name"))) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + + s"(2, 'bb', 10.0, cast('2020-01-01' as timestamp)), " + + s"(3, 'cc', 15.5, cast('2020-02-01' as timestamp))") + + createTable(purchases, purchasesColumns, Array.empty) + sql(s"INSERT INTO testcat.ns.$purchases VALUES " + + s"(1, 42.0, cast('2020-01-01' as timestamp)), " + + s"(2, 11.0, cast('2020-01-01' as timestamp)), " + + s"(3, 19.5, cast('2020-02-01' as timestamp))") + + createTable("t3", Array(Column.create("id", LongType)), Array(identity("id"))) + sql("INSERT INTO testcat.ns.t3 VALUES (1), (2)") + + val query = + s""" + |SELECT id, COUNT(*) AS cnt FROM ( + | ${selectWithMergeJoinHint("sub", "p")} + | p.item_id AS id + | FROM (SELECT id FROM testcat.ns.$items WHERE name >= 'aa') sub + | RIGHT OUTER JOIN testcat.ns.$purchases p + | ON sub.id = p.item_id + | UNION ALL + | SELECT id FROM testcat.ns.t3 + |) GROUP BY id + |""".stripMargin + + Seq(false, true).foreach { allowSubset => + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.UNION_OUTPUT_PARTITIONING.key -> "true", + SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> allowSubset.toString) { + val plan = sql(query).queryExecution.executedPlan + + val union = collect(plan) { case u: UnionExec => u }.head + val kp = union.outputPartitioning.asInstanceOf[physical.KeyedPartitioning] + assert(!kp.isGrouped, "keys 1 and 2 repeat across the union children") + assert(!kp.isCollapsed, + "the ids were unique, so dropping `name` collapsed nothing to carry down this chain") + assert(collectAllGroupPartitions(plan).nonEmpty, + s"allowSubset=$allowSubset: grouping merges only partitions that already shared a key") + assert(collectAllShuffles(plan).size == 1, + s"allowSubset=$allowSubset: only the purchases side shuffles") + checkAnswer(sql(query), Seq(Row(1L, 2L), Row(2L, 2L), Row(3L, 1L))) + } + } + } + + test("SPARK-59057: a collapse is reported when the splits are distributed, not replicated") { + // Under partially clustered distribution `GroupPartitionsExec` spreads a key's splits over one + // partition each instead of replicating them, so the partitions it emits never hold more than + // one split. The collapse has to be read off the key groups it keeps, or this whole mode + // silently reports nothing collapsed: here dropping `name` maps the distinct keys (1, 'aa') and + // (1, 'bb') onto key 1, which is exactly the state the gate exists for. + createTable(items, itemsColumns, Array(identity("id"), identity("name"))) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + s"(1, 'aa', 40.0, cast('2020-01-01' as timestamp)), " + + s"(1, 'bb', 41.0, cast('2020-01-01' as timestamp)), " + + s"(2, 'cc', 10.0, cast('2020-01-01' as timestamp))") + + createTable(purchases, purchasesColumns, Array(identity("item_id"))) + sql(s"INSERT INTO testcat.ns.$purchases VALUES " + + s"(1, 42.0, cast('2020-01-01' as timestamp)), " + + s"(1, 44.0, cast('2020-01-02' as timestamp)), " + + s"(2, 11.0, cast('2020-01-01' as timestamp))") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true") { + val df = sql(s"${selectWithMergeJoinHint("i", "p")} i.id, i.name, p.price " + + s"FROM testcat.ns.$items i JOIN testcat.ns.$purchases p ON p.item_id = i.id") + val plan = df.queryExecution.executedPlan + + val distributing = collectAllGroupPartitions(plan).filter(_.distributePartitions) + assert(distributing.nonEmpty, "this test needs a distributing GroupPartitionsExec") + val keyed = distributing.map(_.outputPartitioning) + .flatMap(physical.PartitioningCollection.flatten) + .collect { case kp: physical.KeyedPartitioning => kp } + assert(keyed.nonEmpty && keyed.forall(_.isCollapsed), + "dropping `name` merged two distinct keys, however the splits are laid out afterwards") + checkAnswer(df, Seq(Row(1, "aa", 42.0), Row(1, "aa", 44.0), + Row(1, "bb", 42.0), Row(1, "bb", 44.0), Row(2, "cc", 11.0))) + } + } + + test("SPARK-59057: a collapse confined to keys the other side filters out is not one") { + // items is partitioned by `identity(id)` with ids 0, 4, 5 and purchases by `bucket(4, item_id)` + // holding only bucket 1, so items' keys are reduced onto buckets [0, 0, 1] -- three distinct + // ids onto two buckets. Bucket 0 is then dropped by partition filtering, since purchases has + // no rows for it, so the only key the grouping outputs is bucket 1, covering the single id 5. + // Counting distinct keys before the filtering reports a collapse (2 < 3) for a partitioning + // where nothing was merged, and the flag is sticky, so it would keep saying so downstream. + createTable(items, itemsColumns, Array(identity("id"))) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + s"(0, 'aa', 39.0, cast('2020-01-01' as timestamp)), " + + s"(4, 'bb', 40.0, cast('2020-01-01' as timestamp)), " + + s"(5, 'cc', 41.0, cast('2020-01-01' as timestamp))") + + createTable(purchases, purchasesColumns, Array(bucket(4, "item_id"))) + sql(s"INSERT INTO testcat.ns.$purchases VALUES " + + s"(5, 42.0, cast('2020-01-01' as timestamp))") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_PARTITION_FILTER_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "false") { + val df = sql(s"${selectWithMergeJoinHint("i", "p")} i.id, p.price " + + s"FROM testcat.ns.$items i JOIN testcat.ns.$purchases p ON p.item_id = i.id") + val plan = df.queryExecution.executedPlan + + val keyed = collectAllGroupPartitions(plan).map(_.outputPartitioning) + .flatMap(physical.PartitioningCollection.flatten) + .collect { case kp: physical.KeyedPartitioning => kp } + assert(keyed.nonEmpty, "the reduced join must be planned as a storage-partitioned join") + assert(keyed.forall(!_.isCollapsed), + "the only surviving key covers one source id, so nothing was merged") + checkAnswer(df, Seq(Row(5, 42.0))) + } + } + + test("SPARK-59057: reducing keys onto a coarser transform collapses keys") { + // `identity(item_id)` reduced onto `bucket(4, id)` maps several ids onto one bucket, so keys + // that were distinct in the source land on the same key: one output partition covers what were + // separate id partitions. That is a collapse, and unlike a projection it drops no key position, + // so the old provenance flag missed it. + createTable(items, itemsColumns, Array(bucket(4, "id"))) + sql(s"INSERT INTO testcat.ns.$items VALUES " + + s"(0, 'aa', 39.0, cast('2020-01-01' as timestamp)), " + + s"(4, 'bb', 40.0, cast('2020-01-01' as timestamp))") + + createTable(purchases, purchasesColumns, Array(identity("item_id"))) + sql(s"INSERT INTO testcat.ns.$purchases VALUES " + + s"(0, 42.0, cast('2020-01-01' as timestamp)), " + + s"(4, 44.0, cast('2020-01-15' as timestamp))") + + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") { + val df = sql(s"${selectWithMergeJoinHint("i", "p")} i.id, p.price " + + s"FROM testcat.ns.$items i JOIN testcat.ns.$purchases p ON p.item_id = i.id") + val plan = df.queryExecution.executedPlan + + val keyed = collectAllGroupPartitions(plan).map(_.outputPartitioning) + .flatMap(physical.PartitioningCollection.flatten) + .collect { case kp: physical.KeyedPartitioning => kp } + assert(keyed.nonEmpty, "the reduced join must be planned as a storage-partitioned join") + assert(keyed.exists(_.isCollapsed), + "the side whose keys were reduced onto a coarser transform is coarsened") + checkAnswer(df, Seq(Row(0, 42.0), Row(4, 44.0))) + } + } } /** diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala index 408f116540bba..8411e22f56cc8 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala @@ -359,7 +359,7 @@ class ProjectedOrderingAndPartitioningSuite case kp: KeyedPartitioning => assert(kp.expressions.map(_.asInstanceOf[Attribute].name) === Seq("y", "z"), "expressions must follow original KP position order [y, z], not output order [z, y]") - assert(kp.isNarrowed, "dropping x must mark the KP as narrowed") + assert(kp.isCollapsed, "dropping x maps (1,1,1) and (2,1,1) onto the same key (1,1)") assert(!kp.isGrouped, "projected keys have duplicate (1,1) entries") case other => fail(s"Expected KeyedPartitioning, got $other") @@ -393,7 +393,7 @@ class ProjectedOrderingAndPartitioningSuite "expressions must follow original KP position order [y, z/z_alias], not output order") assert(kps.tail.forall(_.partitionKeys eq kps.head.partitionKeys), "all narrowed KPs must share the same partitionKeys object") - assert(kps.forall(_.isNarrowed), "all KPs must be marked as narrowed") + assert(kps.forall(_.isCollapsed), "all KPs must be marked as collapsed") assert(kps.forall(!_.isGrouped), "projected keys have duplicate (1,1) entries") case other => fail(s"Expected PartitioningCollection, got $other") @@ -455,7 +455,7 @@ class ProjectedOrderingAndPartitioningSuite project.outputPartitioning match { case kp: KeyedPartitioning => assert(!kp.isGrouped, "narrowed keys must have duplicates (1 appears twice)") - assert(kp.isNarrowed, "projection must mark the KP as narrowed") + assert(kp.isCollapsed, "dropping y maps (1,1) and (1,2) onto the same key 1") assert(!kp.groupedSatisfies(ClusteredDistribution(Seq(x))), "narrowed ungrouped KP must not satisfy via groupedSatisfies without config") case other => fail(s"Expected KeyedPartitioning, got $other") @@ -473,7 +473,9 @@ class ProjectedOrderingAndPartitioningSuite // Scenario 2: projected keys are distinct (x-values: 1, 2, 3) -> isGrouped=true. // Each projected key maps to exactly one original partition so GroupPartitionsExec does not - // merge any partitions. No skew risk: must satisfy ClusteredDistribution regardless of config. + // merge any partitions. Nothing collapsed either -- the projection dropped a position but kept + // every key distinct -- so the partitioning is not coarser than the layout it came from. No + // skew risk: must satisfy ClusteredDistribution regardless of config. val keys2dDistinct = Seq(InternalRow(1, 1), InternalRow(2, 2), InternalRow(3, 3)) val projectDistinct = ProjectExec(Seq(x), DummyLeafExecWithPartitioning(output = Seq(x, y), @@ -483,9 +485,63 @@ class ProjectedOrderingAndPartitioningSuite projectDistinct.outputPartitioning match { case kp: KeyedPartitioning => assert(kp.isGrouped, "distinct projected keys must be grouped") - assert(kp.isNarrowed, "projection must mark the KP as narrowed") + assert(!kp.isCollapsed, + "dropping y lost no distinct key, so this partitioning is not coarser than its source") assert(kp.satisfies(ClusteredDistribution(Seq(x))), - "grouped narrowed KP must satisfy ClusteredDistribution without config (no merging)") + "a grouped KP must satisfy ClusteredDistribution without config (no merging)") + case other => fail(s"Expected KeyedPartitioning, got $other") + } + } + } + + test("PartitioningCollection normalizes isCollapsed across its members") { + val x = AttributeReference("x", IntegerType)() + val y = AttributeReference("y", IntegerType)() + val z = AttributeReference("z", IntegerType)() + val keys = Seq(InternalRow(1), InternalRow(2)) + + // The flag describes the shared physical layout, not one member's naming of it, so a collection + // built from a coarsened and a plain member must report the coarsening on both. Otherwise a + // consumer that reads one member -- `satisfies0` and `EnsureRequirements` accept when any + // single member satisfies the distribution -- could get through the gate on the plain one. + val collection = PartitioningCollection.fromPartitionings(Seq( + KeyedPartitioning(Seq(x), keys), + KeyedPartitioning(Seq(y), keys).copy(isCollapsed = true))) + assert(PartitioningCollection.flatten(collection) + .collect { case kp: KeyedPartitioning => kp }.forall(_.isCollapsed), + "a coarsened member must mark the whole collection") + + // Nested collections are normalized too: a collapsed sibling reaches into them. + val nested = PartitioningCollection.fromPartitionings(Seq( + PartitioningCollection.fromPartitionings(Seq( + KeyedPartitioning(Seq(x), keys), KeyedPartitioning(Seq(y), keys))), + KeyedPartitioning(Seq(z), keys).copy(isCollapsed = true))) + assert(PartitioningCollection.flatten(nested) + .collect { case kp: KeyedPartitioning => kp }.forall(_.isCollapsed), + "a coarsened sibling must mark the members of a nested collection") + } + + test("a projection that keeps every distinct key collapses nothing") { + val x = AttributeReference("x", IntegerType)() + val y = AttributeReference("y", IntegerType)() + + // The source reports two splits for the same partition value, so its keys already contain a + // duplicate before any projection. Dropping y maps (1,1),(1,1) onto 1 and (2,2) onto 2: two + // distinct keys before, two after. Grouping would merge only the two splits that already + // shared a key, which is what GroupPartitionsExec does for any partitioning and needs no + // opt-in, so `groupedSatisfies` must accept it with the config off. + val keys = Seq(InternalRow(1, 1), InternalRow(1, 1), InternalRow(2, 2)) + val project = ProjectExec(Seq(x), + DummyLeafExecWithPartitioning(output = Seq(x, y), + partitioning = KeyedPartitioning(Seq(x, y), keys))) + + withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "false") { + project.outputPartitioning match { + case kp: KeyedPartitioning => + assert(!kp.isGrouped, "the duplicate key comes from the source's two splits") + assert(!kp.isCollapsed, "the projection lost no distinct key") + assert(kp.groupedSatisfies(ClusteredDistribution(Seq(x))), + "grouping merges only splits that already shared a key, so no opt-in is needed") case other => fail(s"Expected KeyedPartitioning, got $other") } } @@ -503,7 +559,7 @@ class ProjectedOrderingAndPartitioningSuite DummyLeafExecWithPartitioning(output = Seq(x, y), partitioning = KeyedPartitioning(Seq(x, y), keys))) val kp = project.outputPartitioning.asInstanceOf[KeyedPartitioning] - assert(kp.isNarrowed && !kp.isGrouped) + assert(kp.isCollapsed && !kp.isGrouped) Seq(true, false).foreach { requireAll => // Both values on purpose: the whole claim of the fix is that the guard answers the same @@ -523,7 +579,7 @@ class ProjectedOrderingAndPartitioningSuite } } - test("SPARK-46367: isNarrowed is sticky across chained PartitioningPreservingUnaryExecNodes") { + test("SPARK-46367: isCollapsed is sticky across chained PartitioningPreservingUnaryExecNodes") { val x = AttributeReference("x", IntegerType)() val y = AttributeReference("y", IntegerType)() @@ -537,8 +593,8 @@ class ProjectedOrderingAndPartitioningSuite outerProject.outputPartitioning match { case kp: KeyedPartitioning => assert(!kp.isGrouped, "duplicate keys must survive the second hop") - assert(kp.isNarrowed, - "isNarrowed must be sticky: a second hop that keeps all positions must not reset it") + assert(kp.isCollapsed, + "isCollapsed must be sticky: a second hop that keeps all positions must not reset it") assert(!kp.groupedSatisfies(ClusteredDistribution(Seq(x))), "narrowed ungrouped KP must still not satisfy ClusteredDistribution without config " + "after a second PartitioningPreservingUnaryExecNode hop") @@ -572,7 +628,7 @@ class ProjectedOrderingAndPartitioningSuite } assert(kp.partitionKeys eq child.partitioning.asInstanceOf[KeyedPartitioning].partitionKeys, "partition keys must be unchanged") - assert(!kp.isNarrowed, "same number of positions: not narrowed") + assert(!kp.isCollapsed, "no position dropped: nothing collapsed") case other => fail(s"Expected KeyedPartitioning, got $other") } } @@ -580,7 +636,7 @@ class ProjectedOrderingAndPartitioningSuite test("SPARK-46367: narrowing projection drops transform when its column is absent") { // KP([bucket(32, id), years(ts)], keys2d) through Project(id) -- ts is dropped. // bucket(32, id) is projectable (id in output); years(ts) is not (ts absent). - // Result: KP([bucket(32, id)], keys1d, isNarrowed=true, isGrouped=false). + // Result: KP([bucket(32, id)], keys1d, isCollapsed=true, isGrouped=false). val id = AttributeReference("id", IntegerType)() val ts = AttributeReference("ts", IntegerType)() val bucketExpr = TransformExpression(BucketFunction, Seq(id), Some(32)) @@ -601,7 +657,7 @@ class ProjectedOrderingAndPartitioningSuite assert(te.children.head.asInstanceOf[Attribute].name === "id") case other => fail(s"Expected TransformExpression, got $other") } - assert(kp.isNarrowed, "dropping years(ts) position must mark the KP as narrowed") + assert(kp.isCollapsed, "dropping years(ts) maps (0,2020) and (0,2021) onto bucket 0") assert(!kp.isGrouped, "projected bucket keys (0,1,0) have duplicates") case other => fail(s"Expected KeyedPartitioning, got $other") } @@ -641,7 +697,7 @@ class ProjectedOrderingAndPartitioningSuite } assert(kp.partitionKeys eq child.partitioning.asInstanceOf[KeyedPartitioning].partitionKeys, "partition keys must be unchanged") - assert(!kp.isNarrowed, "both positions projected: not narrowed") + assert(!kp.isCollapsed, "both positions projected: nothing collapsed") case other => fail(s"Expected KeyedPartitioning, got $other") } }