[SPARK-59272][SQL][FOLLOWUP] Keep partition filtering from narrowing a marked layout's key list

Partition filtering shrinks the merged partition key list below a marked
layout's own declared keys, which makes its regrouping non-identity, forfeits
the keyed claim and makes the pairing gate decline the join. Leave a marked
pair unfiltered instead.

Author: peter-toth

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

@@ -541,17 +541,52 @@ case class EnsureRequirements(
     def bothUnprojected(l: KeyedShuffleSpec, r: KeyedShuffleSpec): Boolean =
       l.joinKeyPositions.isEmpty && r.joinKeyPositions.isEmpty
 
+    // Whether the merged key list below may be narrowed to what the join type allows. A marked
+    // layout is left alone: only an identity regrouping keeps its claim (see
+    // `GroupPartitionsExec`), and losing it costs the pair its join at the gate at the end of this
+    // method.
+    //
+    // Filtering is then the only thing that can shrink the merged list *below* the marked side's
+    // own declared keys. The merging arms take the union, and `KeyedShuffleSpec.areKeysCompatible`
+    // pairs a marked layout only with one whose keys are a subset of its declared keys, so the
+    // union is that side's own key set. The count its hash is taken modulo survives the dedup
+    // because a marked layout is always grouped: `canCreatePartitioning` is the only producer of
+    // the marker and it refuses an ungrouped one. A reduce is the other way a merged list comes
+    // out smaller, and it cannot happen here, since the marked arm of `areKeysCompatible` admits
+    // only positions holding the same transform function and `reducersBothWays` finds nothing to
+    // reduce between those.
+    //
+    // What filtering does instead: an intersection with a strictly smaller partner, or the
+    // one-sided arm that keeps the *other* side's keys, drops groups the marked side holds, and
+    // the regrouping stops being the identity. The pair would then trade its whole join for
+    // pruning those groups.
+    //
+    // Sorting is the other way a regrouping stops being the identity, and is not addressed here.
+    // `mergeAndDedupPartitions` sorts, so a marked layout whose declared order is not the sorted
+    // one is relabelled even where the set is unchanged. Handing it its own list verbatim looks
+    // like the same fix and is not, because `KeyedPartitioning.createShuffleSpec` sorts through
+    // `toGrouped` under `v2BucketingAllowKeysSubsetOfPartitionKeys` while it hands a marked layout
+    // back unprojected: the two children would hold one partitioning and report two specs that
+    // `describesSameKeys` calls different, and `ValidateRequirements` rejects the join this method
+    // just allowed. Measured on the generated sweep in `EnsureRequirementsSuite`, cell
+    // `left=id/12 right=id/312/marked Inner`.
+    val partitionFilter = conf.getConf(SQLConf.V2_BUCKETING_PARTITION_FILTER_ENABLED)
+    def filtersKeys(l: KeyedPartitioning, r: KeyedPartitioning): Boolean =
+      partitionFilter && !l.mayContainUnknownPartitionKeys && !r.mayContainUnknownPartitionKeys
+
     // How many key groups the pushdown below would leave this pair. `mergeAndDedupPartitions`
     // keeps one side's keys and drops the other's for the filtered one-sided join types, and there
     // the dropped side's count says nothing, so rank on the side that survives. The arms that
     // really merge have no cheap answer, so they take the larger of the two counts. Keep the join
     // types here in step with `mergeAndDedupPartitions`.
-    val partitionFilter = conf.getConf(SQLConf.V2_BUCKETING_PARTITION_FILTER_ENABLED)
-    def rank(l: KeyedShuffleSpec, r: KeyedShuffleSpec): Int = joinType match {
-      case LeftOuter | LeftAnti | LeftSingle | ExistenceJoin(_) if partitionFilter =>
-        l.numPartitions
-      case RightOuter if partitionFilter => r.numPartitions
-      case _ => l.numPartitions.max(r.numPartitions)
+    def rank(l: KeyedShuffleSpec, r: KeyedShuffleSpec): Int = {
+      val filtered = filtersKeys(l.partitioning, r.partitioning)
+      joinType match {
+        case LeftOuter | LeftAnti | LeftSingle | ExistenceJoin(_) if filtered =>
+          l.numPartitions
+        case RightOuter if filtered => r.numPartitions
+        case _ => l.numPartitions.max(r.numPartitions)
+      }
     }
 
     // Each side may offer several members, and the right one is the one the other side can pair
@@ -560,7 +595,8 @@ case class EnsureRequirements(
     // `ensureDistributionAndOrdering` makes between children when it picks `bestSpecOpt`.
     //
     // Two things keep `rank` from being what the join actually gets, both on the merging arms.
-    // `InnerLike` and `LeftSemi` intersect under `v2BucketingPartitionFilterEnabled`, and an
+    // `InnerLike` and `LeftSemi` intersect under `v2BucketingPartitionFilterEnabled`, unless
+    // `filtersKeys` turned filtering off for the pair, and an
     // intersection is not monotone in member granularity: members cover different clustering keys
     // rather than nested ones, so a finer pair can rank above a coarser one and still meet the
     // other side in fewer groups. And a union does not merely exceed the rank either, because
@@ -655,7 +691,8 @@ case class EnsureRequirements(
 
       // merge values on both sides
       var mergedPartitionKeys =
-        mergeAndDedupPartitions(leftReducedKeys, rightReducedKeys, joinType, reducedKeyOrdering)
+        mergeAndDedupPartitions(leftReducedKeys, rightReducedKeys, joinType, reducedKeyOrdering,
+          filterPartitions = filtersKeys(leftPartitioning, rightPartitioning))
           .map((_, 1))
 
       logInfo(log"After merging, there are " +
@@ -1006,18 +1043,22 @@ case class EnsureRequirements(
   /**
    * Merge, dedup and sort partitions keys for SPJ and optionally enable partition filtering.
    * Both sides must have matching partition expressions.
+   *
    * @param leftPartitionKeys left side partition keys
    * @param rightPartitionKeys right side partition keys
    * @param joinType join type for optional partition filtering
    * @param keyOrdering ordering to sort partition keys
+   * @param filterPartitions whether to narrow the merged list to the keys the join type allows.
+   *                         The caller decides, see `filtersKeys` in `checkKeyGroupCompatible`.
    * @return merged and sorted partition values
    */
   def mergeAndDedupPartitions(
       leftPartitionKeys: Seq[InternalRowComparableWrapper],
       rightPartitionKeys: Seq[InternalRowComparableWrapper],
       joinType: JoinType,
-      keyOrdering: Ordering[InternalRowComparableWrapper]): Seq[InternalRowComparableWrapper] = {
-    val merged = if (SQLConf.get.getConf(SQLConf.V2_BUCKETING_PARTITION_FILTER_ENABLED)) {
+      keyOrdering: Ordering[InternalRowComparableWrapper],
+      filterPartitions: Boolean): Seq[InternalRowComparableWrapper] = {
+    val merged = if (filterPartitions) {
       // Rows with matching join keys land in the same key group. If a group is absent from one
       // side, whether it can produce output depends on which side's unmatched rows the join
       // preserves. Only equi-joins reach this method, since every SMJ/SHJ takes its keys from

sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala

@@ -7628,17 +7628,23 @@ class KeyGroupedPartitioningSuite
           |FROM (SELECT t.id AS id FROM testcat.ns.a a RIGHT OUTER JOIN t ON a.id = t.id) r
           |JOIN testcat.ns.u u ON r.id = u.id
           |""".stripMargin
-      withSQLConf(
-          SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
-          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
-        val df = sql(query)
-        checkAnswer(df, Seq(Row(1, "u1")))
-        // Only the first join's one-side shuffle remains; the second join storage-partitions.
-        assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
-          Seq(true))
-        assert(collectGroupPartitions(df.queryExecution.executedPlan).nonEmpty,
-          s"subset-keyed partner should still storage-partition join, got: " +
-            df.queryExecution.executedPlan)
+      // SPARK-59272: both settings of partition filtering plan the same join. Were the merge free
+      // to narrow a marked layout, the inner join would keep only the [1] intersection, one group
+      // over r's two input partitions, and r would give up its keyed claim on the count clause.
+      Seq(false, true).foreach { partitionFilter =>
+        withSQLConf(
+            SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+            SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+            SQLConf.V2_BUCKETING_PARTITION_FILTER_ENABLED.key -> partitionFilter.toString) {
+          val df = sql(query)
+          checkAnswer(df, Seq(Row(1, "u1")))
+          // Only the first join's one-side shuffle remains; the second join storage-partitions.
+          assertShuffleMayContainUnknownPartitionKeys(df.queryExecution.executedPlan,
+            Seq(true))
+          assert(collectGroupPartitions(df.queryExecution.executedPlan).nonEmpty,
+            s"subset-keyed partner should still storage-partition join, got: " +
+              df.queryExecution.executedPlan)
+        }
       }
     }
   }
@@ -8408,29 +8414,41 @@ class KeyGroupedPartitioningSuite
       withSQLConf(SQLConf.V2_BUCKETING_ENABLED.key -> "false") {
         checkAnswer(sql(query), expected)
       }
-      withSQLConf(
-          SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
-          "spark.sql.autoBroadcastJoinThreshold" -> "-1",
-          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
-        val df = sql(query)
-        checkAnswer(df, expected)
-        val plan = df.queryExecution.executedPlan
-        // A GroupPartitionsExec with no reducers whose output partition i holds exactly input
-        // partition i is an identity grouping; the fix keeps the unknown-keyed claim for it.
-        val identityGpe = collectAllGroupPartitions(plan).find { g =>
-          g.reducers.isEmpty && g.groupedPartitions.zipWithIndex.forall {
-            case ((_, inputIndices), outputIndex) =>
-              inputIndices.lengthCompare(1) == 0 && inputIndices.head == outputIndex
+      // SPARK-59272: partition filtering must not take this join away. On the `true` arm the merge
+      // would otherwise intersect down to [1, 2], two groups over the marked side's three input
+      // partitions, which breaks the count clause of `identityGrouping`, forfeits the claim and
+      // makes the pairing gate decline. Both arms plan the same join.
+      Seq(false, true).foreach { partitionFilter =>
+        withSQLConf(
+            SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+            "spark.sql.autoBroadcastJoinThreshold" -> "-1",
+            SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+            SQLConf.V2_BUCKETING_PARTITION_FILTER_ENABLED.key -> partitionFilter.toString) {
+          val df = sql(query)
+          checkAnswer(df, expected)
+          val plan = df.queryExecution.executedPlan
+          assert(ValidateRequirements.validate(plan),
+            s"the executed plan must satisfy every operator's required distribution:\n$plan")
+          // Locate the marked side's node: no reducers, and output partition i holds exactly input
+          // partition i. That is weaker than `identityGrouping`, which also asks for one output per
+          // input, so the assertions below are what pin the claim and the three keys.
+          val identityGpe = collectAllGroupPartitions(plan).find { g =>
+            g.reducers.isEmpty && g.groupedPartitions.zipWithIndex.forall {
+              case ((_, inputIndices), outputIndex) =>
+                inputIndices.lengthCompare(1) == 0 && inputIndices.head == outputIndex
+            }
+          }
+          assert(identityGpe.isDefined,
+            s"expected a reducer-free identity GroupPartitionsExec, got:\n$plan")
+          identityGpe.get.outputPartitioning match {
+            case k: KeyedPartitioning =>
+              assert(k.mayContainUnknownPartitionKeys,
+                "the identity grouping must keep the unknown-keyed claim")
+              assert(k.numPartitions === 3,
+                s"the marked side keeps all three of its declared keys:\n$plan")
+            case other =>
+              fail(s"expected a KeyedPartitioning output, got $other")
           }
-        }
-        assert(identityGpe.isDefined,
-          s"expected a reducer-free identity GroupPartitionsExec, got:\n$plan")
-        identityGpe.get.outputPartitioning match {
-          case k: KeyedPartitioning =>
-            assert(k.mayContainUnknownPartitionKeys,
-              "the identity grouping must keep the unknown-keyed claim")
-          case other =>
-            fail(s"expected a KeyedPartitioning output, got $other")
         }
       }
     }

sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala

@@ -1178,9 +1178,10 @@ class EnsureRequirementsSuite extends SharedSparkSession {
     val right = partitioning(2, 3, 4)
     val intersected = partitioning(2, 3).partitionKeys
     val union = partitioning(1, 2, 3, 4).partitionKeys
-    def merge(joinType: JoinType): Seq[InternalRow] =
+    def merge(joinType: JoinType, filterPartitions: Boolean): Seq[InternalRow] =
       EnsureRequirements.mergeAndDedupPartitions(
-        left.partitionKeys, right.partitionKeys, joinType, left.keyOrdering).map(_.row)
+        left.partitionKeys, right.partitionKeys, joinType, left.keyOrdering,
+        filterPartitions).map(_.row)
 
     val expected = Seq(
       Inner -> intersected,
@@ -1193,15 +1194,9 @@ class EnsureRequirementsSuite extends SharedSparkSession {
       RightOuter -> right.partitionKeys,
       FullOuter -> union)
 
-    withSQLConf(SQLConf.V2_BUCKETING_PARTITION_FILTER_ENABLED.key -> "true") {
-      expected.foreach { case (joinType, keys) =>
-        assert(merge(joinType) === keys.map(_.row), joinType)
-      }
-    }
-    withSQLConf(SQLConf.V2_BUCKETING_PARTITION_FILTER_ENABLED.key -> "false") {
-      expected.foreach { case (joinType, _) =>
-        assert(merge(joinType) === union.map(_.row), joinType)
-      }
+    expected.foreach { case (joinType, keys) =>
+      assert(merge(joinType, filterPartitions = true) === keys.map(_.row), joinType)
+      assert(merge(joinType, filterPartitions = false) === union.map(_.row), joinType)
     }
   }
 
@@ -2704,6 +2699,51 @@ class EnsureRequirementsSuite extends SharedSparkSession {
     }
   }
 
+  test("SPARK-59272: partition filtering leaves a marked pair's merged keys alone") {
+    // Both arms of the filter would narrow this pair below the marked side's three declared keys:
+    // `Inner` intersects to [1, 2], and `LeftOuter` keeps the unmarked left's [1, 2]. Either way
+    // the marked side regroups two groups over three input partitions, its regrouping stops being
+    // the identity and it forfeits its keyed claim. See `filtersKeys` in
+    // `EnsureRequirements.checkKeyGroupCompatible`. Both settings push all three keys and neither
+    // side is shuffled.
+    //
+    // The two join types put the marker on opposite sides on purpose, because that is the side
+    // `rank` reads for the one-sided arms.
+    val marked = KeyedPartitioning(Seq(exprA), Seq(InternalRow(1), InternalRow(2), InternalRow(3)))
+      .withLayout(_.copy(mayContainUnknownPartitionKeys = true))
+    val plain = KeyedPartitioning(Seq(exprB), Seq(InternalRow(1), InternalRow(2)))
+    val markedRight = KeyedPartitioning(
+      Seq(exprB), Seq(InternalRow(1), InternalRow(2), InternalRow(3)))
+      .withLayout(_.copy(mayContainUnknownPartitionKeys = true))
+    val plainLeft = KeyedPartitioning(Seq(exprA), Seq(InternalRow(1), InternalRow(2)))
+
+    Seq(
+      (Inner: JoinType, marked, plain),
+      (LeftOuter, plainLeft, markedRight)).foreach { case (joinType, left, right) =>
+      Seq(false, true).foreach { partitionFilter =>
+        withSQLConf(
+            SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true",
+            SQLConf.V2_BUCKETING_PARTITION_FILTER_ENABLED.key -> partitionFilter.toString) {
+          val smj = SortMergeJoinExec(Seq(exprA), Seq(exprB), joinType, None,
+            DummySparkPlan(outputPartitioning = left),
+            DummySparkPlan(outputPartitioning = right))
+          val planned = EnsureRequirements.apply(smj)
+
+          assert(planned.collect { case s: ShuffleExchangeExec => s }.isEmpty,
+            s"$joinType: a marked pair keeps its storage-partitioned join either way:\n" +
+              planned.treeString)
+          assert(groupPartitionsNodes(planned).map(_.expectedPartitionKeys.map(_.size)) ===
+            Seq(Some(3), Some(3)),
+            s"$joinType: both sides are pushed the marked side's three keys:\n" +
+              planned.treeString)
+          assert(ValidateRequirements.validate(planned),
+            s"$joinType: the planned join must satisfy both children's required distribution:\n" +
+              planned.treeString)
+        }
+      }
+    }
+  }
+
   private def anyGpeEnabled(plan: SparkPlan): Boolean =
     plan.collectFirst { case gpe: GroupPartitionsExec if gpe.enableSortedMerge => true }.isDefined