You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by pt...@apache.org on 2026/10/21 09:05:35 UTC
(spark) branch branch-4.x updated: [SPARK-59234][SQL] Don't produce key reducers for a pair admitted without reduction
This is an automated email from the ASF dual-hosted git repository.
peter-toth pushed a commit to branch branch-4.x
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.x by this push:
new fabd45cb92eb [SPARK-59234][SQL] Don't produce key reducers for a pair admitted without reduction
fabd45cb92eb is described below
commit fabd45cb92ebd5f2f91d6c15344dd7e3bbf6f7b2
Author: naveenp2708 <na...@gmail.com>
AuthorDate: Wed Oct 21 11:05:09 2026 +0200
[SPARK-59234][SQL] Don't produce key reducers for a pair admitted without reduction
### What changes were proposed in this pull request?
Under partially clustered distribution, a join whose partition transform reduces its own keys could miscount the spread side’s expected partitions: reducersBothWays produced a reducer even though the pair was admitted without reduction, so a key-merging reducer made GroupPartitionsExec group more splits under a merged key than the expected count declared, and the aligner over-emitted.
This fixes it at the source: KeyedShuffleSpec.reducersBothWays now returns no reducers when canReduceKeys is false. That is the same condition areKeysCompatible consults — when it is false (as under partial clustering) only same-function pairs are admitted, and those already share a key space, so neither side needs a reducer. The join then stays on the raw keys, with no merging to miscount and no extra replicated-partition reads.
A defensive assert(splits.size <= numSplits) in GroupPartitionsExec’s distribute branch is kept as defense in depth to localise any future miscount at the node.
### Why are the changes needed?
With partially clustered distribution on, a same-function pair is admitted through isSameFunction without consulting canReduceKeys, but reducersBothWays ran regardless and asked the connector, so a function that reduces itself returned a merging reducer. On master the committed describesSameKeys gate in EnsureRequirements then fails and the join silently falls back to a shuffle (correct results, no SPJ). Not producing the reducer keeps the SPJ and avoids the miscount entirely.
### Does this PR introduce any user-facing change?
No. Output is unchanged. A same-function pair now stays on its raw keys whenever canReduceKeys is false — not only under partial clustering, but also on the default config (V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS off) — instead of reducing onto a merged key space. Partition counts can differ in those cases; results do not. A partially-clustered join that silently shuffled on master now keeps the storage-partitioned join.
### How was this patch tested?
An end-to-end case in KeyGroupedPartitioningSuite: a self-reducing test transform (self_reduce) over two four-row tables with partially clustered distribution. It asserts the same-function pair is admitted without reducing (reducers.isEmpty), the partially-clustered layout is produced, rows are correct, and no shuffle is added. Reverting the fix makes it fail planning with “more splits (2) than its expected count (1)”.
GroupPartitionsExecSuite: the distribute guard fails fast on an over-count, and the REPLICATED_FOR_JOIN branch keeps both splits in its one slot without over-emitting.
GroupPartitionsExecSuite, KeyGroupedPartitioningSuite, and EnsureRequirementsSuite pass (339 tests).
### Was this patch authored or co-authored using generative AI tooling?
No
Closes #59317 from naveenp2708/SPARK-59234-distribute-guard.
Authored-by: naveenp2708 <na...@gmail.com>
Signed-off-by: Peter Toth <pe...@gmail.com>
(cherry picked from commit f32a1a97373a206e07f4377b40d89697217c3fec)
Signed-off-by: Peter Toth <pe...@gmail.com>
---
.../sql/catalyst/plans/physical/partitioning.scala | 6 +++
.../sql/connector/catalog/InMemoryBaseTable.scala | 9 +++-
.../catalog/functions/SelfReduceFunction.scala | 51 ++++++++++++++++++++++
.../datasources/v2/GroupPartitionsExec.scala | 15 ++++++-
.../connector/KeyGroupedPartitioningSuite.scala | 39 +++++++++++++++++
.../datasources/v2/GroupPartitionsExecSuite.scala | 32 ++++++++++++++
6 files changed, 149 insertions(+), 3 deletions(-)
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 4c3c6857d38c..ad6d67a19d9d 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
@@ -2409,6 +2409,12 @@ case class KeyedShuffleSpec(
*/
def reducersBothWays(other: KeyedShuffleSpec)
: (Option[Seq[Option[KeyReducer]]], Option[Seq[Option[KeyReducer]]]) = {
+ // `canReduceKeys` is the same condition `areKeysCompatible` consults: when it is false, for
+ // example under partial clustering, only same-function pairs are admitted, and those already
+ // share a key space, so neither side needs a reducer. Deriving one anyway lets a function that
+ // reduces itself merge keys that were never meant to merge, so the spread side's expected
+ // counts no longer match the groups it builds (SPARK-59234).
+ if (!canReduceKeys) return (None, None)
val results: Seq[(Option[KeyReducer], Option[KeyReducer])] =
partitioning.expressions.zip(other.partitioning.expressions).map {
// Keys an earlier join already reduced live in a key space that neither expression
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala
index 091a1a441f3b..adf1e655ab10 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/InMemoryBaseTable.scala
@@ -32,7 +32,7 @@ import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{AttributeReference, BoundReference, Cast, EvalMode, Expression => CatalystExpression, GenericInternalRow, GetStructField, JoinedRow, Literal, MetadataStructFieldWithLogicalName, Predicate => CatalystPredicate}
import org.apache.spark.sql.catalyst.util.{ArrayBasedMapData, ArrayData, CaseInsensitiveMap, CharVarcharUtils, DateTimeUtils, GenericArrayData, MapData, ResolveDefaultColumns}
import org.apache.spark.sql.connector.catalog.constraints.Constraint
-import org.apache.spark.sql.connector.catalog.functions.FlipLowBitFunction
+import org.apache.spark.sql.connector.catalog.functions.{FlipLowBitFunction, SelfReduceFunction}
import org.apache.spark.sql.connector.distributions.{Distribution, Distributions}
import org.apache.spark.sql.connector.expressions._
import org.apache.spark.sql.connector.expressions.{Literal => V2Literal}
@@ -217,6 +217,7 @@ abstract class InMemoryBaseTable(
case NamedTransform("truncate", Seq(_: NamedReference, _: V2Literal[_])) =>
case NamedTransform("signed_zeros", Seq(_: NamedReference)) =>
case NamedTransform("flip_low_bit", Seq(_: NamedReference)) =>
+ case NamedTransform("self_reduce", Seq(_: NamedReference)) =>
case t if !allowUnsupportedTransforms =>
throw new IllegalArgumentException(s"Transform $t is not a supported transform")
}
@@ -353,6 +354,12 @@ abstract class InMemoryBaseTable(
case (v, t) =>
throw new IllegalArgumentException(s"Match: unsupported argument(s) type - ($v, $t)")
}
+ case NamedTransform("self_reduce", Seq(ref: NamedReference)) =>
+ extractor(ref.fieldNames, cleanedSchema, row) match {
+ case (value: Long, LongType) => SelfReduceFunction.produceResult(InternalRow(value))
+ case (v, t) =>
+ throw new IllegalArgumentException(s"Match: unsupported argument(s) type - ($v, $t)")
+ }
case ClusterByTransform(columnNames) =>
columnNames.map { colName =>
extractor(colName.fieldNames, cleanedSchema, row)._1
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/functions/SelfReduceFunction.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/functions/SelfReduceFunction.scala
new file mode 100644
index 000000000000..de10d6f95ba9
--- /dev/null
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/connector/catalog/functions/SelfReduceFunction.scala
@@ -0,0 +1,51 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.connector.catalog.functions
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.types.{DataType, IntegerType, LongType}
+
+/**
+ * A partition transform on a long key that reduces itself: `reducer` returns a merging reducer
+ * (x -> x % 2) when the other side is this same function. The both-sides reduce contract
+ * r1(f(x)) = r2(f(x)) is vacuous when the two sides share the function, so any merging reducer is
+ * admitted. No in-tree reducible function reduces itself. This one lets a test check that a
+ * same-function pair is not reduced under partial clustering (SPARK-59234).
+ *
+ * Like [[FlipLowBitFunction]], it lives in catalyst's test sources so `InMemoryBaseTable` can
+ * compute its partition keys by calling `produceResult` here instead of repeating the arithmetic.
+ */
+object SelfReduceFunction extends SimpleFunction
+ with ScalarFunction[Int] with ReducibleFunction[Int, Int] {
+ override def inputTypes(): Array[DataType] = Array(LongType)
+ override def resultType(): DataType = IntegerType
+ override def name(): String = "self_reduce"
+ override def canonicalName(): String = name()
+ override def description(): String = name()
+ override def toString: String = name()
+ override def produceResult(input: InternalRow): Int = input.getLong(0).toInt
+
+ override def reducer(otherFunc: ReducibleFunction[_, _]): Reducer[Int, Int] =
+ if (otherFunc == SelfReduceFunction) SelfReducer else null
+}
+
+object SelfReducer extends Reducer[Int, Int] {
+ override def reduce(value: Int): Int = value % 2
+ override def resultType(): DataType = IntegerType
+ override def displayName(): String = "self_reduce_reducer"
+}
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 d99a40fa68c4..bed5ae47dd0a 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
@@ -722,12 +722,23 @@ private[sql] object GroupPartitionsExec {
val splits = keyMap.getOrElse(key, Seq.empty)
numMatchedPartitions += splits.size
if (distributePartitions) {
+ // `padTo` only pads, never truncates, so a key holding more splits than its expected
+ // count would be emitted more times than the count says, while `isGrouped` can still
+ // read as grouped. No producer does this today: under partial clustering the join split
+ // carries no reducers (`reducersBothWays` returns none), so it never merges keys and its
+ // counts match its splits, and the ordering split counts the same keys it groups. The
+ // assert surfaces any future miscount here instead of as a shuffle fallback at the
+ // `committed` gate in `EnsureRequirements` or as a bad grouped layout.
+ assert(splits.size <= numSplits,
+ s"a distributed partition key holds more splits (${splits.size}) than its expected " +
+ s"count ($numSplits)")
// Distribute splits across expected partitions, padding with empty sequences
val paddedSplits = splits.map(Seq(_)).padTo(numSplits, Seq.empty)
paddedSplits.map((key, _))
} else {
- // The slots beyond the first re-read the whole group. A key this side does not hold
- // has no splits to re-read, so the product is 0 and no explicit guard is needed.
+ // `Seq.fill(numSplits)` emits exactly `numSplits` entries whatever the split count, so this
+ // branch cannot over-emit and needs no size guard. The slots beyond the first re-read the
+ // whole group; a key this side does not hold has no splits to re-read, so the product is 0.
numReplicatedPartitionReads += (numSplits - 1) * splits.size
// Replicate all splits to each expected partition
Seq.fill(numSplits)((key, splits))
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 4c08446a230d..51504692a265 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
@@ -1328,6 +1328,45 @@ class KeyGroupedPartitioningSuite
}
}
+ test("SPARK-59234: a self-reducing function is not reduced under partial clustering") {
+ // self_reduce(id) is a transform that reduces itself. The join is a same-function pair, which
+ // `areKeysCompatible` admits, and `V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS` is on so partial
+ // clustering is the only reason `canReduceKeys` is false. `reducersBothWays` then returns none
+ // (SPARK-59234), so the join stays on the raw keys with no merging to miscount. t1 holds key 0
+ // twice, so partial clustering really spreads it across two slots (`SPLIT_FOR_JOIN`) while t2
+ // repeats it (`REPLICATED_FOR_JOIN`). Assert that, plus correct rows, no shuffle, and that
+ // neither side reduced. If a change brought the reducer back, the merged keys would overflow
+ // the spread side's groups and trip the fail-fast guard.
+ withFunction(SelfReduceFunction) {
+ val cols = Array(Column.create("id", LongType), Column.create("data", StringType))
+ createTable("t1", cols, Array(Expressions.apply("self_reduce", Expressions.column("id"))))
+ sql("INSERT INTO testcat.ns.t1 VALUES (0, 'a'), (0, 'e'), (1, 'b'), (2, 'c'), (3, 'd')")
+ createTable("t2", cols, Array(Expressions.apply("self_reduce", Expressions.column("id"))))
+ sql("INSERT INTO testcat.ns.t2 VALUES (0, 'w'), (1, 'x'), (2, 'y'), (3, 'z')")
+ val df = sql(
+ "SELECT t1.id FROM testcat.ns.t1 JOIN testcat.ns.t2 ON t1.id = t2.id")
+ withSQLConf(
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true",
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key -> "true") {
+ checkAnswer(df, Seq(Row(0L), Row(0L), Row(1L), Row(2L), Row(3L)))
+ val plan = stripAQEPlan(df.queryExecution.executedPlan)
+ assert(collectShuffles(plan).isEmpty,
+ "storage-partitioned join should be kept, not shuffled")
+ val gpes = collectGroupPartitions(plan)
+ assert(gpes.flatMap(_.ungroupingOrigin).toSet == Set(SPLIT_FOR_JOIN, REPLICATED_FOR_JOIN),
+ s"one side spreads and the other repeats:\n$plan")
+ assert(gpes.forall(_.reducers.isEmpty),
+ s"a same-function pair is admitted without reducing under partial clustering:\n$plan")
+ assert(gpes.map(_.outputPartitioning.numPartitions) == Seq(5, 5),
+ s"key 0 gets two slots on both sides:\n$plan")
+ val joins = collect(plan) { case smj: SortMergeJoinExec => smj }
+ assert(joins.size == 1 && ValidateRequirements.validate(joins.head),
+ s"the pair holds up:\n$plan")
+ }
+ }
+ }
+
test("SPARK-59045: reduced expression is retargeted per KeyedPartitioning") {
// A chained SPJ's output partitioning reports one `KeyedPartitioning` per join side, but the
// reduced expression is derived from the one member `checkKeyGroupCompatible` paired this side
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala
index a39effa38cd9..18fe4e93c796 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExecSuite.scala
@@ -74,6 +74,38 @@ class GroupPartitionsExecSuite extends SharedSparkSession {
}
}
+ test("SPARK-59234: a distributed key with more splits than its expected count fails fast") {
+ // Keys [(1,1), (1,2)] projected onto position 0 both become key 1, so the child holds two
+ // splits for it. Declaring an expected count of 1 for key 1 while the node is ungrouped for a
+ // join split is the inconsistency this guard catches: `padTo` would emit both splits while the
+ // count says one, leaving duplicate keys under a grouped report.
+ val childKp = KeyedPartitioning(Seq(exprA, exprB), Seq(row(1, 1), row(1, 2)))
+ def keyOf(a: Int): InternalRowComparableWrapper =
+ InternalRowComparableWrapper(row(a), Seq(exprA))
+ val e = intercept[AssertionError] {
+ GroupPartitionsExec(DummySparkPlan(outputPartitioning = childKp),
+ joinKeyPositions = Some(Seq(0)),
+ expectedPartitionKeys = Some(Seq(keyOf(1) -> 1)),
+ ungroupingOrigin = Some(SPLIT_FOR_JOIN))
+ }
+ assert(e.getMessage.contains("more splits"))
+ }
+
+ test("SPARK-59234: the replicate branch emits numSplits entries and does not overflow") {
+ // The same over-count input as the distribute case, but ungrouped for a replicate rather than a
+ // split. REPLICATED_FOR_JOIN takes the else branch, which emits exactly numSplits entries per
+ // key whatever the split count, so it cannot overflow and builds without the guard firing.
+ val childKp = KeyedPartitioning(Seq(exprA, exprB), Seq(row(1, 1), row(1, 2)))
+ def keyOf(a: Int): InternalRowComparableWrapper =
+ InternalRowComparableWrapper(row(a), Seq(exprA))
+ val gpe = GroupPartitionsExec(DummySparkPlan(outputPartitioning = childKp),
+ joinKeyPositions = Some(Seq(0)),
+ expectedPartitionKeys = Some(Seq(keyOf(1) -> 1)),
+ ungroupingOrigin = Some(REPLICATED_FOR_JOIN))
+ assert(gpe.groupedPartitions.map(_._2) === Seq(Seq(0, 1)),
+ "the one expected partition replicates both splits")
+ }
+
test("SPARK-59121: a node with no reducers keeps the child's reduced key marker") {
// This node re-reports the child's expressions, projected to `joinKeyPositions`. A reduce that
// happened below it has to survive that, or a further join above reads the reported transform
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org