From 7fbb80d59ccd85e0afb9321717aaafbc6a12f077 Mon Sep 17 00:00:00 2001 From: Peter Toth Date: Thu, 3 Sep 2026 12:00:59 +0200 Subject: [PATCH] [SPARK-59187][SQL] Carry the partition key data types on KeyedPartitioning ### What changes were proposed in this pull request? `KeyedPartitioning` gains a `keyDataTypes` field beside `expressions`, and `keyDataTypes` stops sampling the first partition key row. The types are already computed wherever the keys are built, and were dropped there. `KeyedPartitioning.apply` computed them only to build the key wrapper factory. `project` called `projectKeys(positions)._2` and threw away the projected types. `GroupPartitionsExec.grouping` computed the reduced types, `PartitionGrouping` did not carry them, and `outputPartitioning` re-derived them by sampling. All three now pass what they know. Every `copy` keeps the types, which is right at each site: `KeyedShuffleSpec.createPartitioning` and `withNewChildrenInternal` change the expressions and leave the key rows alone, and canonicalisation and `toGrouped` do not change what a key holds. Two copies do change the rows. `project` passes the projected types, and `concat`, the one place that mixes rows from several partitionings, requires the children to agree on them. The field is checked against the thing it describes rather than argued about: the constructor requires one type per expression, and requires the first key row to have been built with them. Every key wrapper carries its own types, so this costs one comparison. `PartitioningCollection`'s invariant deliberately does not gain a clause. Its members share one key list, and a wrapper compares its types before its values, so structurally equal keys force equal types, leaving only empty-key members where the types describe nothing. Two consequences in `EnsureRequirements`. The reduced-types comparison drops SPARK-59176's empty-side guard, which existed only because a keyless side reported the expressions' types. And the two `KeyedPartitioning(clustering, _, _, _)` patterns become type patterns reading `.expressions`, so a future field does not touch them. ### Why are the changes needed? `keyDataTypes` answered by reading the first key row and falling back to `expressionDataTypes` when there was none. Three things followed. A partitioning with no key reported types no key of it would have held. That is SPARK-59176, whose fix keeps such a side out of one comparison rather than giving it an answer. With the field the side answers, and the guard is gone. Nothing checked that the keys after the first agreed with it. There is now nothing to disagree with, because no key is consulted for types at all. The divergence at `KeyedShuffleSpec.createPartitioning` was accidental. `copy(expressions = ...)` silently left the types describing the old expressions. It is now an explicit property of the copy, and the scaladoc says so on `expressionDataTypes`, which is the member whose use needs the warning. ### Does this PR introduce _any_ user-facing change? No. A `stringArgs` override keeps the field out of `explain`, so plan output is unchanged. `KeyedPartitioning` is a `catalyst` class, so this is not public API. It is still a binary- and source-incompatible change to a case class: the constructor, `apply`, `copy`, `unapply` and the companion's `tupled`/`curried` all change shape. MiMa does not flag it, since `MimaExcludes` blanket-excludes `org.apache.spark.sql.catalyst.*`. The two destructuring patterns in the repository are converted to type patterns. ### How was this patch tested? A new `GroupPartitionsExecSuite` test asserts that a reduce's result type reaches the reported partitioning, both with keys and with none, while the expression the node reports keeps its own different type. It uses a both-sides reduce, which is the shape whose reported expression cannot carry the key type: `EnsureRequirements` refuses the same divergence from a one-side reduce. The existing SPARK-59176 tests pin the guard removal: both still pass with the comparison back to its plain form, including the one that asserts a contract-breaking reducer still raises `STORAGE_PARTITION_JOIN_INCOMPATIBLE_REDUCED_TYPES`. The type propagation itself is pinned by twelve existing `KeyGroupedPartitioningSuite` tests, measured by making `GroupPartitionsExec` report the child's types instead of the reduced ones. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 5) --- .../plans/physical/partitioning.scala | 93 ++++++++++++------- .../datasources/v2/GroupPartitionsExec.scala | 6 +- .../exchange/EnsureRequirements.scala | 27 ++---- ...taSourceV2CatalystRuntimeFilterSuite.scala | 7 +- .../KeyGroupedPartitioningSuite.scala | 23 +++-- .../v2/GroupPartitionsExecSuite.scala | 29 +++++- .../exchange/EnsureRequirementsSuite.scala | 6 +- 7 files changed, 112 insertions(+), 79 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 2354cf69205a0..605438d99a030 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 @@ -572,6 +572,14 @@ case class CoalescedNullAwareHashPartitioning( * comparison and grouping. One per partition. Typically in sorted order when * produced by a data source or `GroupPartitionsExec`, but this is not * guaranteed after projection. May contain duplicates when ungrouped. + * @param keyDataTypes The types the `partitionKeys` rows were built with, one per expression. + * Anything reading those rows takes its types from here, and it answers even + * where there is no row to read. A copy that changes the expressions leaves the + * rows alone, so it leaves these alone too. Of the copies that do replace the + * row list, `project` passes new types, `concat` requires the children to agree + * on them, and `toGrouped` and `PartitioningCollection.fromPartitionings` reuse + * rows that already carry these types. See `expressionDataTypes` for the two + * ways the two lists come apart. * @param isGrouped Whether partition keys are unique (no duplicates). Computed on first * creation, then preserved through copy operations to avoid recomputation. * @param isCollapsed Whether a projection or a reduction mapped keys that were distinct in the @@ -582,53 +590,55 @@ case class CoalescedNullAwareHashPartitioning( case class KeyedPartitioning( expressions: Seq[Expression], @transient partitionKeys: Seq[InternalRowComparableWrapper], + keyDataTypes: Seq[DataType], isGrouped: Boolean, isCollapsed: Boolean) extends Expression with Partitioning with Unevaluable { override val numPartitions = partitionKeys.length + // The keys carry their own types, so the field can be checked against the thing it describes + // rather than argued about. One key is enough: `concat` is the only copy that mixes rows from + // several partitionings, and it checks all of them. + require(keyDataTypes.length == expressions.length, + "A KeyedPartitioning must have one key data type per partition expression") + require(partitionKeys.headOption.forall(_.dataTypes == keyDataTypes), + "A KeyedPartitioning's keyDataTypes must be the types its partitionKeys were built with") + override def children: Seq[Expression] = expressions override def nullable: Boolean = false override def dataType: DataType = IntegerType + /** + * Drops the `keyDataTypes`, so that `explain` shows what it showed before the field existed. They + * are the types of the keys printed beside them, which adds nothing a reader of a plan wants. + */ + override protected def stringArgs: Iterator[Any] = + Iterator(expressions, partitionKeys, isGrouped, isCollapsed) + override protected def withNewChildrenInternal( newChildren: IndexedSeq[Expression]): KeyedPartitioning = copy(expressions = newChildren) - /** Need not be what the `partitionKeys` rows hold. See `keyDataTypes`. */ - @transient lazy val expressionDataTypes: Seq[DataType] = expressions.map(_.dataType) - /** - * The types the `partitionKeys` rows were built with. Anything reading those rows should take its - * types from here. It is a driver-side value, since `partitionKeys` is `@transient`. - * - * They differ from the `expressionDataTypes` in two cases. A join that reduced both sides' keys - * onto a key space no transform names leaves a marked expression whose type can be anything, see - * `expressionsDescribeKeys`. A one-side reduce keeps them equal, because the expression the - * partitioning then reports is the target transform and `EnsureRequirements` refuses a reducer - * whose result type disagrees with it. `KeyedShuffleSpec.createPartitioning` is the other case. - * It puts the other child's expressions over these keys with no reducer in sight, so a struct - * field can be named differently on the two sides. With no key at all the expressions are all - * there is, and there is no row to read or to place. + * The types the partition expressions produce. Not what the `partitionKeys` rows hold, whenever + * the expressions have stopped describing the keys, which happens in two ways. * - * The two cases can meet, and then the fallback is not truthful. A marked partitioning can end up - * with no key, for instance when `v2BucketingPartitionFilterEnabled` intersects two sides that - * hold disjoint keys, and this then reports the un-reduced transform's type. What it reports is a - * fact about the key rows, so with no key row there is no fact, and a caller must not hold the - * fallback against a real answer. The reduced-types comparison in `EnsureRequirements` leaves out - * a marked side that has no key for that reason (SPARK-59176). An unmarked one still answers, - * since its expressions describe the keys it would have had, and stays in the comparison. + * A join that reduced both sides' keys onto a key space no transform names leaves a marked + * expression whose type can be anything, see `expressionsDescribeKeys`. A one-side reduce keeps + * the two lists equal, because the expression the partitioning then reports is the target + * transform and `EnsureRequirements` refuses a reducer whose result type disagrees with it. * - * `ShuffleExchangeExec` is the one reader that stays on `expressionDataTypes`. It evaluates the - * expressions to place the other child's rows, and it runs on executors, where this value is not - * available. `expressionsDescribeKeys` is what keeps that site sound. + * `KeyedShuffleSpec.createPartitioning` is the other way. It puts the other child's expressions + * over these keys, so a struct field can be named differently on the two sides. * - * Only the first key's types are read, and nothing enforces that the rest match. SPARK-59187 is - * to carry the types on the partitioning instead of sampling a key row. + * `ShuffleExchangeExec` is the one reader that takes this over `keyDataTypes`, at both of its + * sites. It types rows it has just evaluated the expressions into, so those types are the ones + * that fit, and its two sites have to agree with each other. `expressionsDescribeKeys` is what + * keeps them sound, by refusing a partitioning whose expressions no longer place a row where its + * keys say it belongs. */ - @transient lazy val keyDataTypes: Seq[DataType] = - partitionKeys.headOption.map(_.dataTypes).getOrElse(expressionDataTypes) + @transient lazy val expressionDataTypes: Seq[DataType] = expressions.map(_.dataType) - /** Driver-side, like the `keyDataTypes` it comes from. */ + /** The ordering is compiled, so it is rebuilt after deserialization. */ @transient lazy val keyRowOrdering = KeyedPartitioning.groupedKeyRowOrdering(keyDataTypes) @@ -662,7 +672,7 @@ case class KeyedPartitioning( // from. Two different source keys landing on one projected key is the collapse, and it is // also what makes the projected keys non-unique, so the walk stops at the first one. The // source keys are never hashed, only compared where a projected key repeats. - val projectedKeys = projectKeys(positions)._2 + val (projectedDataTypes, projectedKeys) = projectKeys(positions) val sourceOf = mutable.HashMap.empty[InternalRowComparableWrapper, InternalRowComparableWrapper] var collapses = false @@ -679,6 +689,7 @@ case class KeyedPartitioning( copy( expressions = positions.map(expressions), partitionKeys = projectedKeys, + keyDataTypes = projectedDataTypes, isGrouped = !collapses && sourceOf.size == projectedKeys.length, isCollapsed = isCollapsed || collapses) } @@ -691,8 +702,8 @@ case class KeyedPartitioning( } /** - * Projects this partitioning's expressions by selecting only the specified positions. - * Returns the projected expressions and their data types together with the projected keys. + * Projects this partitioning's partition keys by selecting only the specified positions. + * Returns the types the projected keys were built with, and the projected keys. */ def projectKeys(positions: Seq[Int]): (Seq[DataType], Seq[InternalRowComparableWrapper]) = KeyedPartitioning.projectKeys(partitionKeys, keyDataTypes, positions) @@ -719,7 +730,7 @@ case class KeyedPartitioning( /** * Reduces this partitioning's partition keys by applying the given reducers. - * Returns the reduced keys and their data types. + * Returns the types the reduced keys were built with, and the reduced keys. */ def reduceKeys( reducers: Seq[Option[KeyReducer]]): (Seq[DataType], Seq[InternalRowComparableWrapper]) = @@ -825,7 +836,8 @@ object KeyedPartitioning { val comparablePartitionKeys = partitionKeys.map(comparableKeyWrapperFactory) val isGrouped = comparablePartitionKeys.distinct.size == comparablePartitionKeys.size // Built from scratch, so it is the layout everything else is compared against. - new KeyedPartitioning(expressions, comparablePartitionKeys, isGrouped, isCollapsed = false) + new KeyedPartitioning( + expressions, comparablePartitionKeys, dataTypes, isGrouped, isCollapsed = false) } /** @@ -835,9 +847,16 @@ object KeyedPartitioning { * * Keys repeating across children is not a collapse. Only a child's own collapse carries over, * since such a key still stands for several finer-grained ones in the concatenation. + * + * This is the one place that mixes key rows from several partitionings, so it is the one place + * where the children's `keyDataTypes` have to be checked rather than carried. The caller compares + * the children's expressions, and equal expressions do not by themselves mean equal key types: a + * reduce leaves a partitioning whose keys are typed by the reducer. */ def concat(kps: Seq[KeyedPartitioning]): KeyedPartitioning = { val concatenatedKeys = kps.flatMap(_.partitionKeys) + require(kps.forall(_.keyDataTypes == kps.head.keyDataTypes), + "Concatenated KeyedPartitionings must agree on keyDataTypes") kps.head.copy( partitionKeys = concatenatedKeys, // A child that has duplicates of its own puts them in the concatenation too, which answers @@ -1074,6 +1093,12 @@ case class PartitioningCollection(partitionings: Seq[Partitioning]) * unmarked one. They cannot disagree. The members share one key list, so they describe one * reduce, and a reduce that marks one side's expressions marks the other's in the same step, * while a one-side reduce marks neither. + * + * `keyDataTypes` is not in it either, for a different reason. The members share one key list, and + * a wrapper compares its types before its values, so structurally equal keys force equal types. + * That leaves only members whose key list is empty, where the types describe nothing and cannot + * be read wrong. `KeyedPartitioning`'s own constructor is where the field is checked, against the + * keys it describes. */ private def checkKeyedPartitioningInvariant(): Unit = { firstKeyedPartitioning.foreach { first => 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 eb0142c5a01bd..d1cdfccefa557 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 @@ -96,7 +96,8 @@ case class GroupPartitionsExec( case None => projectedExpressions } KeyedPartitioning( - effectiveExpressions, partitionKeys, grouping.isGrouped, grouping.isCollapsed) + effectiveExpressions, partitionKeys, grouping.keyDataTypes, grouping.isGrouped, + grouping.isCollapsed) }.asInstanceOf[Partitioning] case o => o } @@ -216,7 +217,7 @@ case class GroupPartitionsExec( group.tail.exists(childKeys(_) != first) } } - PartitionGrouping(partitions, isGrouped, isCollapsed) + PartitionGrouping(partitions, reducedDataTypes, isGrouped, isCollapsed) } @transient lazy val groupedPartitions: Seq[(InternalRowComparableWrapper, Seq[Int])] = @@ -403,6 +404,7 @@ case class GroupPartitionsExec( /** What a [[GroupPartitionsExec]] computes once and reports from several members. */ private case class PartitionGrouping( partitions: Seq[(InternalRowComparableWrapper, Seq[Int])], + keyDataTypes: Seq[DataType], isGrouped: Boolean, isCollapsed: Boolean) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala index 2a33280861da0..ee4b99160a802 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala @@ -417,17 +417,17 @@ case class EnsureRequirements( reorder(leftKeys.toIndexedSeq, rightKeys.toIndexedSeq, rightExpressions, rightKeys) .orElse(reorderJoinKeysRecursively( leftKeys, rightKeys, leftPartitioning, None)) - case (Some(KeyedPartitioning(clustering, _, _, _)), _) => + case (Some(kp: KeyedPartitioning), _) => // The single-column invariant in KeyedPartitioning.supportsExpressions guarantees one // attribute per partition expression. - val leafExprs = clustering.flatMap(_.references) + val leafExprs = kp.expressions.flatMap(_.references) reorder(leftKeys.toIndexedSeq, rightKeys.toIndexedSeq, leafExprs, leftKeys) .orElse(reorderJoinKeysRecursively( leftKeys, rightKeys, None, rightPartitioning)) - case (_, Some(KeyedPartitioning(clustering, _, _, _))) => + case (_, Some(kp: KeyedPartitioning)) => // The single-column invariant in KeyedPartitioning.supportsExpressions guarantees one // attribute per partition expression. - val leafExprs = clustering.flatMap(_.references) + val leafExprs = kp.expressions.flatMap(_.references) reorder(leftKeys.toIndexedSeq, rightKeys.toIndexedSeq, leafExprs, rightKeys) .orElse(reorderJoinKeysRecursively( leftKeys, rightKeys, leftPartitioning, None)) @@ -565,22 +565,7 @@ case class EnsureRequirements( val (rightReducedDataTypes, rightReducedKeys) = rightReducers.fold( (rightPartitioning.keyDataTypes, rightPartitioning.partitionKeys) )(rightPartitioning.reduceKeys) - // The reduced types are the types of the key rows the merge below sees. A side with no key - // still answers for them while its expressions describe the keys it would have had, and - // `keyDataTypes` falls back to exactly those types. After a reduce the expressions no - // longer describe them, so the fallback is a type no key of that partitioning would hold, - // and comparing it against a real answer fails a co-partitioned query (SPARK-59176). Only - // such a side is left out. An empty one that is not marked stays in, which is what keeps - // the comparison checking a reducer's result type against the paired transform. - val leftTypesDescribeKeys = - leftReducedKeys.nonEmpty || leftPartitioning.expressionsDescribeKeys - val rightTypesDescribeKeys = - rightReducedKeys.nonEmpty || rightPartitioning.expressionsDescribeKeys - val reducedDataTypes = if (!leftTypesDescribeKeys) { - rightReducedDataTypes - } else if (!rightTypesDescribeKeys || leftReducedDataTypes == rightReducedDataTypes) { - leftReducedDataTypes - } else { + if (leftReducedDataTypes != rightReducedDataTypes) { throw QueryExecutionErrors.storagePartitionJoinIncompatibleReducedTypesError( leftReducers = leftReducers, leftReducedDataTypes = leftReducedDataTypes, @@ -588,7 +573,7 @@ case class EnsureRequirements( rightReducedDataTypes = rightReducedDataTypes) } - val reducedKeyOrdering = KeyedPartitioning.groupedKeyRowOrdering(reducedDataTypes) + val reducedKeyOrdering = KeyedPartitioning.groupedKeyRowOrdering(leftReducedDataTypes) .on((t: InternalRowComparableWrapper) => t.row) // merge values on both sides 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 defa376dac50c..3d579a7d460cd 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 @@ -22,7 +22,6 @@ import org.apache.spark.sql.{AnalysisException, DataFrame, Row} import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Add, AttributeReference, DynamicPruning, DynamicPruningExpression, EqualTo, Expression, GetStructField, GreaterThan, Literal, RLike} import org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning -import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper import org.apache.spark.sql.connector.catalog.{ Column, Identifier, @@ -591,10 +590,8 @@ class DataSourceV2CatalystRuntimeFilterSuite extends SharedSparkSession { val partAttr = AttributeReference("part", IntegerType)() val table = new InMemoryTable("t", Array(Column.create("part", IntegerType)), Array.empty[Transform], java.util.Collections.emptyMap[String, String]) - val partitioning = KeyedPartitioning( - Seq(partAttr), - Seq(InternalRowComparableWrapper(InternalRow(1), Seq(partAttr))), - isGrouped = false, isCollapsed = false) + val partitioning = + KeyedPartitioning(Seq(partAttr), Seq(InternalRow(1))).copy(isGrouped = 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/KeyGroupedPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala index 6a911ce652615..f4ca113785f36 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 @@ -1123,18 +1123,18 @@ class KeyGroupedPartitioningSuite test("SPARK-59176: a leg reduced onto no key at all still joins") { withReducedTsJoinLegs(bothRows, row2020, leg2YearsValues = Some(row2021)) { // The second leg's two sides hold disjoint years, so the partition filter intersects them to - // nothing and the leg reports a reduced partitioning with no key. The reduced types then have - // to come from the first leg. The marked expressions still name the un-reduced `days` and - // `years` transforms, whose types are not the `LongType` the reduced keys hold. + // nothing and the leg reports a reduced partitioning with no key. Its marked expressions name + // the un-reduced `days` and `years` transforms, whose types are not the `LongType` the + // reduced keys hold, so the leg has to report the reduced types rather than derive them. withSQLConf( SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", SQLConf.V2_BUCKETING_PARTITION_FILTER_ENABLED.key -> "true", SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key -> "true", SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") { - // Both orders, since the side that has no key is the one to leave out of the comparison. - // And both join types, since the inner join intersects the two key sets to nothing and so - // has nothing to sort, while the full outer join keeps the other side's keys and sorts them - // by the reported types. + // Both orders, since either side of the join can be the one with no key. And both join + // types, since the inner join intersects the two key sets to nothing and so has nothing to + // sort, while the full outer join keeps the other side's keys and sorts them by the + // reported types. Seq("JOIN" -> Nil, "FULL OUTER JOIN" -> bothTimestamps).foreach { case (joinType, expected) => Seq(false, true).foreach { leg2First => @@ -1149,7 +1149,7 @@ class KeyGroupedPartitioningSuite } } - test("SPARK-59176: an empty side whose expressions describe its keys keeps the reducer check") { + test("SPARK-59176: an empty side does not hide a reducer whose result type disagrees") { withFunction(UnboundDaysFunctionWithToYearsReducerWithDateResult) { createTable(items, itemsColumns, Array(days("arrive_time"))) sql(s"INSERT INTO testcat.ns.$items VALUES " + @@ -1162,10 +1162,9 @@ class KeyGroupedPartitioningSuite } // The inner join intersects two disjoint year key sets, so its leg reports a `years(time)` - // partitioning with no key. Nothing reduced it, so its expressions still describe the keys it - // would have had, and the reduced-types comparison must still run. This `days` function - // breaks the reducer contract, returning `DateType` where the target `years` transform is - // `IntegerType`, and that is what the comparison is there to catch. + // partitioning with no key. Having no key must not cost the reduced-types comparison its + // answer. This `days` function breaks the reducer contract, returning `DateType` where the + // target `years` transform is `IntegerType`, and that is what the comparison catches. withSQLConf( SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", SQLConf.V2_BUCKETING_PARTITION_FILTER_ENABLED.key -> "true", 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 d8ea11c6ea37e..cdd87a14ab985 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 @@ -22,11 +22,11 @@ import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeReference, SortOrder, TransformExpression} import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, KeyedPartitioning, KeyedShuffleSpec, KeyReducer, Partitioning, PartitioningCollection, UnknownPartitioning} import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper -import org.apache.spark.sql.connector.catalog.functions.{BucketFunction, BucketReducer} +import org.apache.spark.sql.connector.catalog.functions.{BucketFunction, BucketReducer, DaysFunctionWithToYearsReducerWithLongResult, DaysToYearsReducerWithLongResult, YearsFunctionWithToYearsReducerWithLongResult} import org.apache.spark.sql.execution.{DummySparkPlan, LeafExecNode, SafeForKWayMerge} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.test.SharedSparkSession -import org.apache.spark.sql.types.IntegerType +import org.apache.spark.sql.types.{DateType, IntegerType, LongType} class GroupPartitionsExecSuite extends SharedSparkSession { @@ -99,6 +99,31 @@ class GroupPartitionsExecSuite extends SharedSparkSession { } } + test("SPARK-59187: a reduce's result types reach the reported partitioning with no key left") { + // A both-sides reduce is the shape whose reported expression cannot carry the key type. The + // reduce lands on `LongType` and the expression it marks stays `DateType`, so the two are told + // apart only by what the partitioning carries. `EnsureRequirements` refuses the same divergence + // from a one-side reduce, so that shape would not be a plan. + val daysExpr = TransformExpression(DaysFunctionWithToYearsReducerWithLongResult, Seq(exprA)) + val yearsExpr = TransformExpression(YearsFunctionWithToYearsReducerWithLongResult, Seq(exprA)) + val markedExpr = daysExpr.reducedTogetherWith(yearsExpr) + val child = DummySparkPlan(outputPartitioning = + KeyedPartitioning(Seq(daysExpr), Seq(row(1), row(2)))) + val reducers = Some(Seq(Some(KeyReducer(DaysToYearsReducerWithLongResult(), markedExpr)))) + + // With keys and without, since the second is the case no key row can answer. + Seq(None, Some(Seq.empty[(InternalRowComparableWrapper, Int)])).foreach { expected => + GroupPartitionsExec(child, expectedPartitionKeys = expected, reducers = reducers) + .outputPartitioning match { + case kp: KeyedPartitioning => + assert(!kp.expressionsDescribeKeys, "test setup: the reported expression is marked") + assert(kp.expressionDataTypes === Seq(DateType), "the marked `days` transform's own type") + assert(kp.keyDataTypes === Seq(LongType), "but the keys hold what the reducer produced") + case other => fail(s"Expected KeyedPartitioning, got $other") + } + } + } + test("SPARK-56241: non-coalescing passes through child ordering unchanged") { // Each partition has a distinct key — no coalescing happens. val partitionKeys = Seq(row(1), row(2), row(3)) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala index 1d877f169605c..fd3ac8b5fb9ea 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala @@ -1159,11 +1159,11 @@ class EnsureRequirementsSuite extends SharedSparkSession { EnsureRequirements.apply(smjExec) match { case ShuffledHashJoinExec(_, _, _, _, _, DummySparkPlan(_, _, left: KeyedPartitioning, _, _), - ShuffleExchangeExec(KeyedPartitioning(attrs, pks, _, _), + ShuffleExchangeExec(shuffled: KeyedPartitioning, DummySparkPlan(_, _, SinglePartition, _, _), _, _, _), _) => assert(left.expressions == a1 :: Nil) - assert(attrs == a1 :: Nil) - assert(partitionKeys == pks.map(_.row)) + assert(shuffled.expressions == a1 :: Nil) + assert(partitionKeys == shuffled.partitionKeys.map(_.row)) case other => fail(other.toString) } }