[SPARK-59887][SPARK-59901][SQL][4.3] Fix wrong results when a storage-partitioned join pairs a transform of a join-key expression with one of a column

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala

@@ -123,16 +123,6 @@ case class TransformExpression(
     }
   }
 
-  /**
-   * Re-targets this partition transform expression at `attr`. A partition transform expression
-   * has a single leaf attribute (`KeyedPartitioning.supportsExpressions`), so this replaces that
-   * attribute and keeps the rest of the expression: any field path above the leaf comes from this
-   * expression, not from the re-targeted key. Re-targeting is therefore only faithful when the
-   * source and the target key expressions have the same path shape.
-   */
-  def withReference(attr: Attribute): TransformExpression =
-    transform { case _: AttributeReference => attr }.asInstanceOf[TransformExpression]
-
   // Return a Reducer for a reducible function on another reducible function
   private def reducer(
       thisFunction: ReducibleFunction[_, _],

sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala

@@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst
 
 import org.apache.spark.{SparkFunSuite, SparkUnsupportedOperationException}
 import org.apache.spark.sql.catalyst.dsl.expressions._
-import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, DirectShufflePartitionID, Expression, TransformExpression}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, DirectShufflePartitionID, Expression, GetStructField, TransformExpression}
 import org.apache.spark.sql.catalyst.plans.SQLHelper
 import org.apache.spark.sql.catalyst.plans.physical._
 import org.apache.spark.sql.connector.catalog.functions.{FlipLowBitFunction, Reducer, ReducibleFunction, ScalarFunction}
@@ -685,6 +685,111 @@ class ShuffleSpecSuite extends SparkFunSuite with SQLHelper {
     }
   }
 
+  test("SPARK-59887: a spec over a transform of an expression pairs only with the same shape") {
+    // A join pairs its two sides up by the column each partition expression references. Over
+    // `a = b` it would pair `bucket(4, a + 1)` with `bucket(4, b)`, and over `s = t` it would pair
+    // `bucket(4, s.x)` with `bucket(4, t.y)`. A row with `a = b` sits in a different bucket on each
+    // side. `bucket(4, a + 1)` and `bucket(4, b + 1)` are the same function of `a` and of `b`, so
+    // they pair up.
+    val fn = new FakeBucket
+    val a = $"a".long
+    val b = $"b".long
+    val structType = new StructType().add("x", LongType).add("y", LongType)
+    val s = $"s".struct(structType)
+    val t = $"t".struct(structType)
+    // A spec clustered on the column its one partition expression references.
+    def keyed(expression: Expression): KeyedShuffleSpec =
+      KeyedShuffleSpec(
+        KeyedPartitioning(Seq(expression), Seq(InternalRow(0L), InternalRow(1L))),
+        ClusteredDistribution(expression.references.toSeq))
+    def bucket(numBuckets: Int, argument: Expression): TransformExpression =
+      TransformExpression(fn, Seq(argument), Some(numBuckets))
+    def spec(numBuckets: Int, argument: Expression): KeyedShuffleSpec =
+      keyed(bucket(numBuckets, argument))
+    // Whether the two sides of one reduce, `bucket(8, left)` with `bucket(4, right)`, pair up.
+    def pairedByReduce(left: Expression, right: Expression): Boolean =
+      keyed(bucket(8, left).reducedTogetherWith(bucket(4, right)))
+        .isCompatibleWith(keyed(bucket(4, right).reducedTogetherWith(bucket(8, left))))
+    val column = spec(4, b)
+
+    withSQLConf(
+        SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+        SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+        SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "false",
+        SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+      assert(spec(4, a).isCompatibleWith(column), "two columns pair up")
+      assert(spec(4, a).canCreatePartitioning)
+      assert(spec(8, a).areKeysCompatible(column, allowReduce = true),
+        "the fixture reduces two columns")
+      assert(pairedByReduce(a, b), "two columns reduced together share one key space")
+      assert(keyed(b).areKeysCompatible(spec(4, a), allowReduce = true),
+        "an identity side reduces onto a column's transform")
+
+      Seq(
+        ("an expression", a + 1L, b + 1L, b + 2L),
+        ("a struct field", GetStructField(s, 0), GetStructField(t, 0), GetStructField(t, 1))
+      ).foreach { case (shape, argument, sameShape, otherShape) =>
+        val over = spec(4, argument)
+        assert(over.isCompatibleWith(spec(4, sameShape)), s"over $shape, the same shape")
+        assert(over.isCompatibleWith(over), s"over $shape, with itself")
+        assert(!over.isCompatibleWith(spec(4, otherShape)), s"over $shape, another shape")
+        assert(!over.isCompatibleWith(column), s"over $shape, against a column")
+        assert(!column.isCompatibleWith(over), s"over $shape, the other way round")
+        assert(!over.canCreatePartitioning, s"over $shape, no layout to shuffle onto")
+
+        // Over the same shape, such a pair reduces as two columns do, whichever side reduces. It
+        // does not reduce onto another shape or a column, and an identity side does not reduce
+        // onto it.
+        assert(spec(8, argument).areKeysCompatible(spec(4, sameShape), allowReduce = true),
+          s"over $shape, reducing the same shape")
+        assert(spec(4, sameShape).areKeysCompatible(spec(8, argument), allowReduce = true),
+          s"over $shape, reducing the same shape the other way round")
+        assert(!spec(8, argument).areKeysCompatible(spec(4, otherShape), allowReduce = true),
+          s"over $shape, reducing another shape")
+        assert(!spec(8, argument).areKeysCompatible(column, allowReduce = true), s"over $shape")
+        assert(!column.areKeysCompatible(spec(8, argument), allowReduce = true), s"over $shape")
+        assert(!over.areKeysCompatible(spec(8, b), allowReduce = true), s"over $shape")
+        assert(!keyed(b).areKeysCompatible(over, allowReduce = true),
+          s"over $shape, reducing an identity side")
+        // Keys an earlier join reduced are in the space of that reduce. Two transforms reduced
+        // through the same pairing pair over the same shape, as two columns do. `bucket(12)`
+        // reduced with `bucket(8)` holds buckets of 4, with `bucket(18)` buckets of 6, and
+        // unreduced buckets of 12.
+        assert(pairedByReduce(argument, sameShape), s"over $shape, reduced together")
+        assert(!pairedByReduce(argument, otherShape),
+          s"over $shape, another shape reduced together")
+        val reducedWith8 = keyed(bucket(12, argument).reducedTogetherWith(bucket(8, a)))
+        assert(!reducedWith8.isCompatibleWith(
+          keyed(bucket(12, sameShape).reducedTogetherWith(bucket(18, b)))),
+          s"over $shape, reduced through different pairings")
+        assert(!reducedWith8.isCompatibleWith(spec(12, sameShape)),
+          s"over $shape, reduced against unreduced")
+        assert(!spec(12, sameShape).isCompatibleWith(reducedWith8),
+          s"over $shape, unreduced against reduced")
+
+        // One such member does not keep a sibling from serving as the layout.
+        assert(ShuffleSpecCollection(Seq(over, spec(4, a))).canCreatePartitioning)
+      }
+    }
+  }
+
+  test("SPARK-59901: an expression over two columns maps to no cluster key") {
+    // A one-side shuffle builds its partition expression over the other side's join key, which can
+    // reference two columns. No single cluster key stands for it.
+    val a = $"a".long
+    val b = $"b".long
+    Seq(a + b, TransformExpression(new FakeBucket, Seq(a + b), Some(4))).foreach { e =>
+      val twoColumns = KeyedShuffleSpec(
+        KeyedPartitioning(Seq(e), Seq(InternalRow(0L), InternalRow(1L))),
+        ClusteredDistribution(Seq(a, b)))
+      assert(twoColumns.keyPositions.forall(_.isEmpty), s"$e")
+      assert(!twoColumns.isCompatibleWith(twoColumns), s"$e")
+      withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+        assert(!twoColumns.canCreatePartitioning, s"$e")
+      }
+    }
+  }
+
   test("createShuffleSpec: a marked narrowing projection yields an unusable spec") {
     val a = $"a".int
     val b = $"b".int

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala

@@ -76,8 +76,8 @@ case class GroupPartitionsExec(
         // There can be multiple `KeyedPartitioning`s in an output partitioning of a join, but they
         // can only differ in `expressions`. Their `partitionKeys` reference and `isCollapsed` flag
         // are shared (enforced by `PartitioningCollection`), so the grouping is computed once.
-        // When reducers are applied, the stored reduced expressions are re-targeted at each
-        // `KeyedPartitioning`'s own key attribute and reported instead of the original ones. Their
+        // When reducers are applied, the stored reduced expressions are rebuilt over each
+        // `KeyedPartitioning`'s own argument and reported instead of the original ones. Their
         // data types match the reduced partition keys for the identity-vs-transform and
         // single-side-transform reducers; for the both-sides-reduce shape no single transform
         // describes the keys (see `KeyedShuffleSpec.reducersBothWays`).
@@ -107,9 +107,16 @@ case class GroupPartitionsExec(
                 projectedExpressions.zip(exprs).map {
                   case (expr, Some(KeyReducer(_, reduced))) =>
                     // `reduced` was stored from the single spec that `createKeyedShuffleSpec`
-                    // picked (`collectFirst`); re-target it at this `KeyedPartitioning`'s own key
-                    // attribute so that every `KeyedPartitioning` in a collection keeps its own.
-                    reduced.withReference(expr.references.head)
+                    // picked (`collectFirst`). The keys are reduced once, from the shared key rows,
+                    // so `reduced` describes them whichever member this is. Only its function says
+                    // that, so the member's own argument takes the place of `reduced`'s. A side
+                    // shuffled onto this layout reports `bucket(8, b + 1)` next to `bucket(8, id)`,
+                    // or `b + 1` next to an identity `id`.
+                    val argument = expr match {
+                      case t: TransformExpression => t.children
+                      case e => Seq(e)
+                    }
+                    reduced.copy(children = argument)
                   case (expr, None) => expr
                 }
               case None => projectedExpressions

sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala

@@ -206,7 +206,11 @@ case class EnsureRequirements(
       val shouldConsiderMinParallelism = specs.forall(p =>
         !p._2.canCreatePartitioning || children(p._1).isInstanceOf[ShuffleExchangeLike]
       )
-      // Choose all the specs that can be used to shuffle other children
+      // Choose all the specs that can be used to shuffle other children. For a collection only the
+      // members that can serve as the layout count, here and below. Any other member would build a
+      // partitioning its own child is not laid out on.
+      def servingMembers(spec: ShuffleSpec): Seq[ShuffleSpec] =
+        flattenSpec(spec).filter(_.canCreatePartitioning)
       val candidateSpecs = specs.filter { case (index, spec) =>
         spec.canCreatePartitioning &&
           (!shouldConsiderMinParallelism ||
@@ -228,8 +232,9 @@ case class EnsureRequirements(
         } else {
           candidateSpecs
         }
-        // Pick the spec with the best parallelism
-        Some(finalCandidateSpecs.values.maxBy(_.numPartitions))
+        // Pick the spec with the best parallelism. A collection counts with its first member that
+        // can serve as the layout, which is its first member unless that one cannot.
+        Some(finalCandidateSpecs.values.maxBy(servingMembers(_).head.numPartitions))
       }
 
       // Check if the following conditions are satisfied:
@@ -263,14 +268,14 @@ case class EnsureRequirements(
       // child does not even have. Pick that member once, preferring the finest when several
       // qualify. Only the branch that shuffles a child reads these, hence `lazy`.
       lazy val matchedIndexes = bestSpecOpt.toSeq.flatMap { best =>
-        childrenIndexes.filter(i => best.isCompatibleWith(specs(i)))
+        childrenIndexes.filter(i => servingMembers(best).exists(_.isCompatibleWith(specs(i))))
       }
       lazy val bestMemberOpt = bestSpecOpt.flatMap { best =>
         val matchedMembers = matchedIndexes.map(i => flattenSpec(specs(i)))
         // No member serving every matched child means there is no layout to align them on, so they
         // all take the ordinary shuffle. That needs three or more clustered children, since with
         // two the member that reported the match serves both, and no operator has three today.
-        flattenSpec(best)
+        servingMembers(best)
           .filter(m => matchedMembers.forall(_.exists(m.isCompatibleWith)))
           .maxByOption(_.numPartitions)
       }
@@ -929,15 +934,17 @@ case class EnsureRequirements(
       partitioning: Partitioning,
       distribution: ClusteredDistribution): Option[KeyedShuffleSpec] = {
     def tryCreate(partitioning: KeyedPartitioning): Option[KeyedShuffleSpec] = {
-      // The single-column invariant in KeyedPartitioning.supportsExpressions guarantees one
-      // attribute per partition expression.
       val attributes = partitioning.expressions.flatMap(_.references)
       val clustering = distribution.clustering
 
       val satisfies = if (SQLConf.get.getConf(SQLConf.REQUIRE_ALL_CLUSTER_KEYS_FOR_CO_PARTITION)) {
-        attributes.length == clustering.length && attributes.zip(clustering).forall {
-          case (l, r) => l.semanticEquals(r)
-        }
+        // Only an expression over a single column covers that column. One over several, e.g.
+        // `b + c`, maps to no position (`KeyedShuffleSpec.keyPositions`). The spec turns it away
+        // or projects it away, so a layout with one does not cover the cluster keys.
+        partitioning.expressions.forall(_.references.size == 1) &&
+          attributes.length == clustering.length && attributes.zip(clustering).forall {
+            case (l, r) => l.semanticEquals(r)
+          }
       } else {
         partitioning.satisfies(distribution)
       }