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) } }