ulysses-you commented on code in PR #58262:
URL: https://github.com/apache/spark/pull/58262#discussion_r3878034637
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -867,42 +879,259 @@ case class EnsureRequirements(
}
/**
- * Splits a partitioning into three categories:
- * 1. Non-KeyedPartitioning (HashPartitioning, RangePartitioning, etc.)
- * 2. Grouped KeyedPartitioning (isGrouped = true)
- * 3. Non-grouped KeyedPartitioning (isGrouped = false)
+ * The positions of `kp`'s partition expressions that are operation keys of
`distribution`, and so
+ * have to survive a projection. All of them when nothing needs projecting.
+ *
+ * Under `v2BucketingAllowKeysSubsetOfPartitionKeys` a [[KeyedPartitioning]]
may be grouped on
+ * more keys than the operation requires, in which case partitions sharing
an operation key are
+ * still separate. A partition expression is an operation key in two ways:
one of its *references*
+ * is a cluster key - the form `groupedSatisfies` and
`KeyedShuffleSpec.keyPositions` both use,
+ * where a `bucket(4, a)` transform covers the cluster key `a` - or the
expression *itself* is a
+ * cluster key. The second is never decisive in practice, because
`IdentityTransform` resolves to
+ * the attribute itself and then the reference-level test matches the same
position anyway; it is
+ * kept so that a partition expression which is a cluster key can never be
projected away.
+ *
+ * Returns every position for a co-partitioned operator: there the
multi-child block owns the
+ * projection, and doing it here as well would leave that block deriving
positions from an already
+ * projected partitioning and applying them to the unprojected partition
expressions.
+ *
+ * An empty result means no partition expression covers an operation key, so
there is nothing to
+ * project onto. That is the answer for a member that cannot satisfy
`distribution`, which is why
+ * the caller only asks for members that can. It also happens for a member
that can: one whose
+ * expressions have no references at all makes every `groupedSatisfies`
branch vacuously true, and
+ * nothing at `KeyedPartitioning` construction rejects that. The caller
skips such a member rather
+ * than project it to no position, which would collapse every partition into
one.
+ *
+ * Keeping a position is only sound because `groupedSatisfies`' subset
branch also requires
+ * `expressions.forall(_.references.size == 1)`: a kept expression is then a
function of a single
+ * cluster key, so coalescing on the projected keys cannot put rows that
share an operation key on
+ * different partitions.
+ */
+ private def clusterKeyPositions(
+ kp: KeyedPartitioning,
+ distribution: Distribution,
+ isCoPartitioned: Boolean): BitSet = distribution match {
+ case c: ClusteredDistribution if !isCoPartitioned =>
+ val positions = kp.expressions.indices.filter { i =>
+ val e = kp.expressions(i)
+ c.clustering.exists(_.semanticEquals(e)) ||
+ e.references.exists(ref =>
c.clustering.exists(_.semanticEquals(ref)))
+ }.to(BitSet)
+ positions
Review Comment:
Redundant local val
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -544,6 +550,21 @@ case class KeyedPartitioning(
@transient lazy val expressionDataTypes: Seq[DataType] =
expressions.map(_.dataType)
+ /**
+ * The data types of the `partitionKeys` rows: the types each key was built
with, and the ones it
+ * is hashed and compared under (`InternalRowComparableWrapper.dataTypes`).
+ *
+ * These are the `expressionDataTypes` unless a reducer has rewritten the
keys. With
+ * `v2BucketingAllowCompatibleTransforms` a storage-partitioned join reduces
one or both sides'
+ * keys onto a common key space, while the partitioning keeps reporting the
expressions it was
+ * built from: joining an `identity(ts)`-partitioned table to a
`years(ts)`-partitioned one leaves
+ * `IntegerType` year values under a `TimestampType`-declared expression.
Reading a key row is
+ * only sound at the types the row was written with, which is what these are
for; use
+ * `expressionDataTypes` only where the question is about the expressions
themselves.
+ */
+ @transient lazy val keyDataTypes: Seq[DataType] =
+ partitionKeys.headOption.map(_.dataTypes).getOrElse(expressionDataTypes)
Review Comment:
the key-type axis is fixed in `projectKeys` but not in the ordering
consumers — `toGrouped`, `keyRowOrdering`, and the OrderedDistribution arm
still read at `expressionDataTypes`
- **Where:** `partitioning.scala:547-556`
(`keyRowOrdering`/`keyOrdering`/`toGrouped`), reached
via `KeyedPartitioning.createShuffleSpec` subset branch at
`partitioning.scala:620-629`; also
`EnsureRequirements.scala:93-94`.
- **Defect:** `toGrouped` sorts with
`groupedKeyRowOrdering(expressionDataTypes)` — the declared
expression types — while this PR establishes that keys must be read at
`keyDataTypes`. For a
reducer-rewritten partitioning (Integer year values under a
`TimestampType`-declared `ts`
expression), sorting compares `GenericInternalRow.getLong` against a
boxed `Integer`
(`rows.scala`: `getLong(ordinal) = getAs(ordinal)`) ->
`ClassCastException: java.lang.Integer cannot be cast to java.lang.Long`
— the exact error this
PR fixes elsewhere.
- **Concrete failure:** `v2BucketingAllowCompatibleTransforms` +
`v2BucketingAllowKeysSubsetOfPartitionKeys` on;
`identity(ts)`/`years(ts)` SPJ join (the PR's
own test shape) -> any `PartitioningPreservingUnaryExecNode` above it
(final `HashAggregate`)
preserving the stale-typed KP -> a join above *that* calls
`createShuffleSpec` -> subset branch
runs `projectKeys` (now fixed) then `.toGrouped` -> CCE at planning with
>=2 distinct projected
keys. Same mismatch in the `OrderedDistribution` arm
(`EnsureRequirements.scala:93`):
`RowOrdering.create(o.ordering, attrs)` compares reduced keys with an
ordering bound to the
stale expression types, under `v2BucketingAllowSorting`.
- **Peer code / invariant:** the PR's own `keyDataTypes` scaladoc
(`partitioning.scala:550-564`)
states "the types each key was built with, and the ones it is hashed and
compared under" — the
hash/compare path (`InternalRowComparableWrapper.equals/hashCode`) uses
the wrapper's own types
and is safe; `toGrouped`'s sort does not go through the wrapper's
ordering, so it violates the
stated invariant.
- **Severity honesty:** *not a regression* — on master these plans CCE'd
even earlier (in
`projectKeys`). The PR widens the set of plannable queries past the
first crash and lands on
the next instance of the same bug. Recommend either deriving
`keyRowOrdering` from
`keyDataTypes` (one line, consistent with the new scaladoc), or
explicitly listing the
remaining consumers in the follow-up JIRA alongside the
`ValidateRequirements` gap.
- **Verdict: CONFIRMED** (line-by-line trace; cast mechanics verified in
source). Runnable repro,
works on master *without* the PR, and would still fail after it —
appendable to
`EnsureRequirementsSuite`:
```scala
test("SPARK-58968: createShuffleSpec must sort keys at the types they were
built with") {
val exprTs = AttributeReference("ts", TimestampType)()
val exprId = AttributeReference("id", IntegerType)()
val factory = InternalRowComparableWrapper
.getInternalRowComparableWrapperFactory(Seq(IntegerType, IntegerType))
// Reduced keys (year, bucket) under stale expressions (ts:
TimestampType), as
// GroupPartitionsExec.outputPartitioning reports after applying
reducers.
val keys = Seq(InternalRow(2020, 0), InternalRow(2021, 1)).map(factory)
val kp = new KeyedPartitioning(Seq(exprTs, exprId), keys, isGrouped =
true)
withSQLConf(SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key
-> "true") {
val e =
intercept[ClassCastException](kp.createShuffleSpec(ClusteredDistribution(Seq(exprTs))))
assert(e.getMessage.contains("java.lang.Integer cannot be cast to
java.lang.Long"))
}
}
```
Run:
`build/sbt 'sql/testOnly
org.apache.spark.sql.execution.exchange.EnsureRequirementsSuite -- -z
"createShuffleSpec must sort keys"'`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
