Skip to content

[SPARK-59187][SQL] Carry the partition key data types on KeyedPartitioning #58501

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
peter-toth wants to merge 1 commit into
base: master
Choose a base branch
from
+112 −79
Draft

[SPARK-59187][SQL] Carry the partition key data types on KeyedPartitioning #58501

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
93 changes: 59 additions & 34 deletions sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

@dongjoon-hyun dongjoon-hyun Sep 4, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: this paragraph, the constructor comment and the concat scaladoc each inventory the copy sites (project, concat, toGrouped, fromPartitionings) and argue the design. The lists will silently rot at the next copy(partitionKeys = ...) (GroupPartitionsExec already builds one directly). I'd keep the contract only, e.g. "The types the partitionKeys rows were built with, one per expression; kept even when there is no key row.", and on concat: "Children must agree on keyDataTypes; the constructor only checks the first key."

* 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
Expand All @@ -582,53 +590,55 @@ case class CoalescedNullAwareHashPartitioning(
case class KeyedPartitioning(
expressions: Seq[Expression],
@transient partitionKeys: Seq[InternalRowComparableWrapper],
keyDataTypes: Seq[DataType],

@dongjoon-hyun dongjoon-hyun Sep 4, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Optional, for consideration: the (types, keys) pair now appears as two constructor args here, two tuple-returning helpers (projectKeys, reduceKeys), two PartitionGrouping fields and the fold seeds in EnsureRequirements, and the pairing is guaranteed only by a head-key require that is vacuous when the key list is empty. A small value object (say TypedKeys(dataTypes, keys)) returned by projectKeys/reduceKeys and held here and in PartitionGrouping would make the pairing structural and is the deeper fix for both issues above. Fine as a follow-up if you prefer to keep this PR small.

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] =

@dongjoon-hyun dongjoon-hyun Sep 4, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: the justification does not hold. InternalRowComparableWrapper has no toString, so the keys print as InternalRowComparableWrapper@<hash> and no type is legible beside them. The hide is also asymmetric: TreeNode.jsonFields / asCode use productIterator, so the field shows up in toJSON but not in explain.

The case this PR exists for (a marked days(...) expression of DateType over LongType keys, possibly with no key at all) is exactly where explain would show a misleading expression type with no way to see the real one, and two partitionings that differ only in keyDataTypes print identically in require/assert messages. No golden file or test asserts this string, so I would drop the override. If it is kept for output stability, a one-line comment saying so would be clearer than arguing the information is worthless.

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.

@dongjoon-hyun dongjoon-hyun Sep 4, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: this reads as if KeyedShuffleSpec.createPartitioning were a case where the expressions have stopped describing the keys, but it only does partitioning.copy(expressions = newExpressions) and sets no marker, so expressionsDescribeKeys stays true there and the last paragraph's "expressionsDescribeKeys is what keeps them sound" only covers the reduce case. Something like: "May differ from keyDataTypes in two cases: (a) a both-sides reduce marks the expressions (expressionsDescribeKeys); (b) KeyedShuffleSpec.createPartitioning re-targets the expressions at the other child's attributes, so struct field names can differ while the expressions still describe the keys."

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

Expand Down Expand Up @@ -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
Expand All @@ -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)
}
Expand All @@ -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)
Expand All @@ -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]) =
Expand Down Expand Up @@ -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)
}

/**
Expand All @@ -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),

@dongjoon-hyun dongjoon-hyun Sep 4, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This require is reachable from UnionExec.comparePartitioning, the only caller, which compares children by semanticEquals on expressions alone and otherwise falls back to super.outputPartitioning. Semantically equal expressions do not imply equal keyDataTypes; this PR itself documents KeyedShuffleSpec.createPartitioning keeping the keyed side's struct field names. So a query that used to plan now fails with IllegalArgumentException.

Concrete case with shuffle-one-side on: purchases p LEFT JOIN items i ON p.item_id = i.id, items identity-partitioned on id: struct<a:int>, purchases unkeyed. The purchases side is shuffled via createPartitioning, and being LeftOuter its KeyedPartitioning(item_id, <items keys>, [struct<a:int>]) becomes the join output. SELECT item_id ... UNION ALL SELECT c FROM t3 with t3 identity-partitioned on c: struct<b:int>. BinaryComparison.sameType ignores struct field names so no Cast is inserted (the existing test SPARK-59054: shuffle one side: struct partition keys with different field names plans exactly this join). UnionExec remaps both expressions to the union output attribute, semanticEquals holds, and concat receives [struct<a:int>] vs [struct<b:int>].

Before this PR the mixed-type concat ran (with a latent isGrouped miscount, since wrappers of different types never compare equal). I think the check belongs in UnionExec.comparePartitioning next to the expression comparison, so that a mismatch takes the existing fallback instead of throwing here.

"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
Expand Down Expand Up @@ -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

@dongjoon-hyun dongjoon-hyun Sep 4, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think this premise holds. An empty-key member's keyDataTypes is read: by GroupPartitionsExec.grouping (via collectFirst), by the reduced-types comparison in EnsureRequirements, and by PushDownUtils. Those readers pick a member by different rules (collectFirst vs. the first member that satisfies), so if members disagree the answer depends on which one is consulted. That is what produces the GroupPartitionsExec failure I described above.

A one-line require(rep.keyDataTypes == first.keyDataTypes, ...) next to the existing isCollapsed check (same O(members) cost) would make this structural, or fromPartitionings could normalize the field the way it interns partitionKeys.

* 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 =>
Expand Down
6 changes: 4 additions & 2 deletions ...re/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,7 +96,8 @@ case class GroupPartitionsExec(
case None => projectedExpressions
}
KeyedPartitioning(
effectiveExpressions, partitionKeys, grouping.isGrouped, grouping.isCollapsed)
effectiveExpressions, partitionKeys, grouping.keyDataTypes, grouping.isGrouped,

@dongjoon-hyun dongjoon-hyun Sep 4, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks like a regression. grouping takes keyDataTypes from the member found by collectFirst (line ~177) and this line stamps it on every member, while EnsureRequirements.createKeyedShuffleSpec builds expectedPartitionKeys from the first member that satisfies the distribution. Members of a PartitioningCollection with empty key lists can carry different keyDataTypes (nothing checks or normalizes them, see my comment on checkKeyedPartitioningInvariant), so the two can disagree and the new constructor require fires.

Reachable shape, with pushPartValues, partitionFilter and allowCompatibleTransforms on:

  1. leg1 = t1 JOIN t2, both identity-partitioned on a: string with disjoint keys -> partition filter intersects to nothing, two members KP(a, [], [String]).
  2. leg2 = t3 JOIN t4, both bucket(4, b: string), disjoint -> two members KP(bucket(4,b), [], [Int]).
  3. leg1 JOIN leg2 ON a = b: KeyedShuffleSpec.isCompatibleWith is true (Nil == Nil, numPartitions 0 == 0, attribute vs transform is compatible via canReduceKeys), so the push-down branch and its type check are skipped and fromPartitionings produces a collection mixing [String] and [Int].
  4. ... FULL OUTER JOIN t5(bucket(4, c)) ON b = c: EnsureRequirements picks the bucket(4,b) member, merged keys are Int rows. grouping picks member a, so reducedDataTypes = [String], and this line builds KeyedPartitioning(bucket(4,b), <Int keys>, [String], ...) -> IllegalArgumentException from the require at partitioning.scala:603.

Before this PR the types were read from the key rows, so the query ran. The inner-join variant passes the require (merged keys are empty) but carries [String] on the bucket member, and a later push-down join then throws STORAGE_PARTITION_JOIN_INCOMPATIBLE_REDUCED_TYPES.

grouping.isCollapsed)
}.asInstanceOf[Partitioning]
case o => o
}
Expand Down Expand Up @@ -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])] =
Expand Down Expand Up @@ -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)

Expand Down
27 changes: 6 additions & 21 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down Expand Up @@ -565,30 +565,15 @@ 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,
rightReducers = rightReducers,
rightReducedDataTypes = rightReducedDataTypes)
}

val reducedKeyOrdering = KeyedPartitioning.groupedKeyRowOrdering(reducedDataTypes)
val reducedKeyOrdering = KeyedPartitioning.groupedKeyRowOrdering(leftReducedDataTypes)
.on((t: InternalRowComparableWrapper) => t.row)

// merge values on both sides
Expand Down
7 changes: 2 additions & 5 deletions ...rc/test/scala/org/apache/spark/sql/connector/DataSourceV2CatalystRuntimeFilterSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand Down
Loading