peter-toth commented on code in PR #55885:
URL: https://github.com/apache/spark/pull/55885#discussion_r4340180541
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -2389,6 +2389,17 @@ case class KeyedShuffleSpec(
conf.v2BucketingAllowCompatibleTransforms
}
+ /**
+ * `t` retargeted at the identity side's key `col`, or None if its argument
types do not match
+ * the bound function's declared types, since evaluating it would then fail.
The compatibility
+ * gate and `reducersBothWays` must agree on this -- a gate admitting a pair
the reducer refuses
+ * would pair raw keys and drop matches -- so both call it.
+ */
+ private def retargetForIdentity(
+ t: TransformExpression,
+ col: AttributeReference): Option[TransformExpression] =
+ Some(rebuiltOver(t, col)).filter(_.argsMatchInputTypes)
Review Comment:
**Finding 11.** This gate is what
[r3470467090](https://github.com/apache/spark/pull/55885#discussion_r3470467090)
asked for: the identity reduce no longer evaluates a transform whose arguments
don't have the declared types. The one-side shuffle evaluates the same kind of
rebuilt transform, but `canCreatePartitioning` (`partitioning.scala:2516`) only
checks `argumentsAreColumns(t)` on the template. `createPartitioning` then
rebuilds it over the other child's key, and `ShuffleExchangeExec` evaluates it
per row.
Measured on this head with `spark.sql.sources.v2.bucketing.shuffle.enabled`.
The scan reports `truncate(data, 2S)` against the fixtures' `truncate`, which
declares `(StringType, IntegerType)`. That is the legal implicit cast the test
at `KeyGroupedPartitioningSuite.scala:7925` describes:
```
width=IntegerType canCreatePartitioning=true argsMatchInputTypes=true
eval -> "ab"
width=ShortType canCreatePartitioning=true argsMatchInputTypes=false
eval -> ClassCastException: java.lang.Short cannot be cast to java.lang.Integer
```
On master, `supportsExpressions` refused a second child, so this scan
reported no partitioning and the join shuffled both sides. An arity-flexible
transform such as `zero_or_one(col, 2)` takes the same path and fails with an
index error instead.
The same predicate on the template closes it, at `partitioning.scala:2516`:
```scala
case t: TransformExpression => argumentsAreColumns(t) &&
t.argsMatchInputTypes
```
The template's own argument types are enough. The key `createPartitioning`
substitutes is the other side's join key at this position, so it has this
column's type.
One fixture has to follow. `ShuffleSpecSuite`'s `TestBucketFunction` still
declares `Array(IntegerType)` but is now built with two children, so it should
declare `Array(IntegerType, IntegerType)`. Otherwise the `SPARK-59121` and
`SPARK-59971` `canCreatePartitioning` tests fail.
With both changes, the Short width gives `canCreatePartitioning=false` and
the Int width stays `true`. All 466 tests stay green, together with the changes
from findings 13 and 14. The suites are `TransformExpressionSuite`,
`ShuffleSpecSuite`, `KeyGroupedPartitioningSuite`, `GroupPartitionsExecSuite`,
`EnsureRequirementsSuite` and `ProjectedOrderingAndPartitioningSuite`.
A sibling of the two `mismatchedIdVsTransformSpecs()` tests pins it. It
fails on this head and passes with the clause:
```scala
val (_, trSpec) = mismatchedIdVsTransformSpecs()
withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
assert(!trSpec.canCreatePartitioning,
"a template whose arguments don't match the declared types must not
be shuffled onto")
}
```
---
_Generated by [Claude Code](https://claude.ai/code)_
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -92,49 +179,180 @@ case class TransformExpression(
* @param other the transform expression to compare to
* @return true if compatible, false if not
*/
- def isCompatible(other: TransformExpression): Boolean = {
- if (isSameFunction(other)) {
- true
- } else {
- (function, other.function) match {
- case (f: ReducibleFunction[_, _], o: ReducibleFunction[_, _]) =>
- val thisReducer = reducer(f, numBucketsOpt, o, other.numBucketsOpt)
- val otherReducer = reducer(o, other.numBucketsOpt, f, numBucketsOpt)
- thisReducer.isDefined || otherReducer.isDefined
- case _ => false
- }
+ def isCompatible(other: TransformExpression): Boolean =
+ isSameFunction(other) || reducers(other).isDefined ||
other.reducers(this).isDefined
+
+ /**
+ * Rewrites this transform's column arguments with `rewriteColumn`, leaving
its literal parameters
+ * untouched.
+ *
+ * Since literal parameters (the bucket count, the truncate width) live in
`children` rather than
+ * in a field of their own, any path that rewrites a transform's children
has to skip them. A
+ * parameter is not an aliasable value: substituting it changes what the
transform computes, drops
+ * it out of `functionId`, adds a second entry to `references` (tripping
+ * `KeyedShuffleSpec.keyPositions`' single-reference assert), and makes
+ * `KeyedPartitioning.supportsExpressions` reject the partitioning outright
-- so SPJ is silently
+ * lost. Use it for any retargeting of a partition transform.
+ *
+ * Callers: `KeyedShuffleSpec.rebuiltOver`, which rebuilds a transform over
the other side's key,
+ * and `GroupPartitionsExec`, which rebuilds a reduced expression over each
member's argument.
+ */
+ def rewriteColumnSlots(rewriteColumn: Expression => Expression):
TransformExpression =
+ copy(children = children.map {
+ case l: Literal => l
+ case c => rewriteColumn(c)
+ })
+
+ /** The non-literal arguments. */
+ def columnSlots: Seq[Expression] =
children.filterNot(_.isInstanceOf[Literal])
+
+ /**
+ * Extract all literal parameters of this transform as V2 [[V2Literal]]s,
preserving each value's
+ * internal representation and its `DataType`. Only consulted once a reducer
path has confirmed
+ * the literal params already match the declared input types (see
+ * [[literalParamsMatchInputTypes]]), so no type coercion happens here.
Memoized.
+ *
+ * Examples:
+ * bucket(4, col) => [Literal(4, IntegerType)]
+ * truncate(col, 3) => [Literal(3, IntegerType)]
+ * days(col) => [] (no literals)
+ */
+ private lazy val extractParameters: Array[V2Literal[_]] =
+ literalChildren.map(l => LiteralValue(l.value, l.dataType):
V2Literal[_]).toArray
+
+ /**
+ * Whether the `select`ed children match the bound function's declared input
type at their
+ * positions. A child beyond the declared arity has no declared type to
compare against (e.g. an
+ * arity-flexible function), so it is left to the connector reducer / other
guards. The DataType
+ * match is exact by design: any mismatch (including cosmetic ones like
Array `containsNull` or
+ * Decimal precision/scale) fails safe to a shuffle. See the two predicates
below for the callers.
+ *
+ * Reading `dataType` can throw rather than answer, e.g. for a
`GetStructField` over a non-struct
+ * column. Both callers are gates that must answer "not reducible" rather
than fail the query, so
+ * a child whose type cannot be read is treated as not matching.
+ */
+ private def inputTypesMatch(select: Expression => Boolean): Boolean = {
+ val declaredTypes = function.inputTypes()
+ children.zipWithIndex.forall {
+ case (c, i) =>
+ !select(c) || i >= declaredTypes.length ||
+ Try(c.dataType).toOption.contains(declaredTypes(i))
}
}
/**
- * Return a [[Reducer]] for this transform expression on another
- * on the transform expression.
- * <p>
- * A [[Reducer]] exists for a transform expression function if it is
- * 'reducible' on the other expression function.
- * <p>
- * @return reducer function or None if not reducible on the other transform
expression
+ * Whether every literal parameter matches its declared input type. Used by
the transform-vs-
+ * transform reducer path, which hands literal *values* to the connector
reducer without Analyzer
+ * type coercion. A literal whose type differs from the declared input type
(a legal implicit cast
+ * under [[BoundFunction]]) is not reducible: the join falls back to a
shuffle rather than handing
+ * the connector a value the partitions were not built on, which it would
then mis-cast. Column
+ * slots are not checked (they are not passed to the connector); the eval
path uses
+ * [[argsMatchInputTypes]] instead.
+ */
+ lazy val literalParamsMatchInputTypes: Boolean =
inputTypesMatch(_.isInstanceOf[Literal])
+
+ /**
+ * Whether every argument has its declared type, at exactly the declared
arity. Required before
+ * evaluating the transform directly: no cast is applied, and a mismatch
would fail at evaluation.
+ */
+ lazy val argsMatchInputTypes: Boolean =
+ children.length == function.inputTypes().length && inputTypesMatch(_ =>
true)
+
+ /**
+ * Whether each position holds a literal on both sides, a column reference
on both sides, or the
+ * same expression of a column on both sides (e.g. `b + 1` and `c + 1`).
Literal values and arity
+ * may differ; the connector reducer decides those.
+ */
+ private def sameArgumentLayout(other: TransformExpression): Boolean =
+ children.zip(other.children).forall {
+ case (_: Literal, _: Literal) => true
+ case (c1, c2) if TransformExpression.isColumnRef(c1) &&
TransformExpression.isColumnRef(c2) =>
+ true
+ case (c1, c2) => TransformExpression.sameShape(c1, c2)
+ }
+
+ /**
+ * Whether no literal parameter has an array, map, struct or UDT type; those
are not passed to a
+ * connector reducer. Other scalars, such as intervals, are.
+ */
+ private def noComplexLiteralParams: Boolean =
+ literalChildren.forall(_.dataType match {
+ case _: ArrayType | _: MapType | _: StructType | _: UserDefinedType[_]
=> false
+ case _ => true
+ })
+
+ /**
+ * Return a [[Reducer]] that maps this transform's partition keys onto
`other`'s, or None if there
+ * is none: when either function is not a [[ReducibleFunction]], when the
argument layouts or
+ * literal parameters rule a reducer out, or when the connector reports it
is not reducible.
+ * Handles both parameterized (bucket, truncate) and non-parameterized
(days, hours) functions.
*/
def reducers(other: TransformExpression): Option[Reducer[_, _]] = {
- (function, other.function) match {
- case(e1: ReducibleFunction[_, _], e2: ReducibleFunction[_, _]) =>
- reducer(e1, numBucketsOpt, e2, other.numBucketsOpt)
- case _ => None
+ import TransformExpression._
+ val thisFunction: ReducibleFunction[_, _] = function match {
+ case f: ReducibleFunction[_, _] => f
+ case _ => return None
+ }
+ val otherFunction: ReducibleFunction[_, _] = other.function match {
+ case o: ReducibleFunction[_, _] => o
+ case _ => return None
+ }
+ if (!sameArgumentLayout(other) ||
+ !literalParamsMatchInputTypes || !other.literalParamsMatchInputTypes ||
+ !noComplexLiteralParams || !other.noComplexLiteralParams ||
+ literalChildren.exists(_.value == null) ||
other.literalChildren.exists(_.value == null)) {
+ return None
}
- }
- // Return a Reducer for a reducible function on another reducible function
- private def reducer(
- thisFunction: ReducibleFunction[_, _],
- thisNumBucketsOpt: Option[Int],
- otherFunction: ReducibleFunction[_, _],
- otherNumBucketsOpt: Option[Int]): Option[Reducer[_, _]] = {
- val res = (thisNumBucketsOpt, otherNumBucketsOpt) match {
- case (Some(numBuckets), Some(otherNumBuckets)) =>
- thisFunction.reducer(numBuckets, otherFunction, otherNumBuckets)
- case _ => thisFunction.reducer(otherFunction)
+ val thisParams = extractParameters
+ val otherParams = other.extractParameters
+ val thisName = function.canonicalName()
+
+ // A single IntegerType parameter on each side is the shape the deprecated
+ // reducer(int, ..., int) overload accepts. Null parameters were refused
above.
+ def isSingleInt(p: Array[V2Literal[_]]): Boolean = {
+ p.length == 1 && p(0).dataType == IntegerType
+ }
+
+ // Probe one reducer overload into an Outcome. Pure -- logging is decided
once, below.
+ def probe(call: => Reducer[_, _]): Outcome = Try(Option(call)) match {
+ case Success(Some(r)) => Reducible(r)
+ case Success(None) => NotReducible
+ case Failure(_: UnsupportedOperationException) => Unimplemented
+ case Failure(e) => Threw(e)
+ }
+
+ // Prefer the generalized overload. Use the deprecated one only if the
generalized one is not
+ // implemented; any other outcome is final.
+ val outcome =
+ if (thisParams.isEmpty && otherParams.isEmpty) {
+ probe(thisFunction.reducer(otherFunction))
+ } else {
+ probe(thisFunction.reducer(thisParams, otherFunction, otherParams))
match {
+ // The deprecated overload is documented for bucket against bucket,
so only offer it a
+ // pair of the same function.
+ case Unimplemented if isSingleInt(thisParams) &&
isSingleInt(otherParams) &&
+ function.canonicalName() == other.function.canonicalName() =>
+ probe(thisFunction.reducer(
+ thisParams(0).value().asInstanceOf[Int], otherFunction,
+ otherParams(0).value().asInstanceOf[Int]))
+ case other => other
+ }
+ }
+
+ outcome match {
+ case Reducible(r) => Some(r)
+ case NotReducible => None
+ case Threw(e) =>
Review Comment:
**Finding 12.** `isCompatible` and `reducersBothWays` each ask the two
directions separately, and a `Threw` here only empties its own direction. So
when one direction throws and the other returns a reducer, the pair is read as
a one-side reduce. `reducersBothWays` keeps the working side's reducer and
reports the throwing side's unreduced transform for its keys. The javadoc
promises the opposite at `ReducibleFunction.java:90`: "any other exception is
logged and the pair is treated as not reducible (the join falls back to a
shuffle)".
A one-side reading is right only when the surviving reducer satisfies
`r(f1(x)) = f2(x)`. Against a two-sided one, such as the gcd of two bucket
counts or the int truncate's lcm, it pairs keys from two different spaces.
Measured on this head with a copy of `BucketFunction` whose generalized
reducer guards its direction with `require(thisNumBuckets >= otherNumBuckets)`.
That is a natural guard for an implementor who expects to be asked only from
the coarser side. The query joins `bucket(16, id)` and `bucket(12, id)` over
the same 48 ids, under `allowCompatibleTransforms`:
```
BucketFunction: rows=48 expected=48 shuffles=0
require-guarded copy: rows=16 expected=48 shuffles=0
```
The 16 side reduces by `% 4` and is reported as `bucket(12, id)`. The 12
side keeps its raw keys `0..11`, so only `0..3` meet. On master a throwing
reducer failed the query. Here the only trace is the WARN.
To match the javadoc, a throw has to veto the pair, not just its own
direction. One entry point can answer for both directions and return `None`
when either threw. `isCompatible` and the `(TransformExpression,
TransformExpression)` arm of `reducersBothWays` would call it instead of
calling `reducers` twice:
```scala
/** Both directions' reducers, or None when either direction threw: then
the pair shuffles. */
private[sql] def reducersWith(other: TransformExpression)
: Option[(Option[Reducer[_, _]], Option[Reducer[_, _]])]
```
`reducers` keeps the `Outcome` internally, so `Threw` stays distinguishable
from `NotReducible`. To pin it, use a fixture like the one above and assert
that `reducersBothWays` returns `(None, None)` and the row count is complete.
---
_Generated by [Claude Code](https://claude.ai/code)_
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -675,11 +675,13 @@ private[sql] object GroupPartitionsExec {
// argument takes the place of `reduced`'s. A side shuffled
onto this layout
// reports `bucket(8, b + 1)` next to `bucket(8, id)`, or `b
+ 1` next to an
// identity `id`.
+ // Only the column argument is replaced; `reduced` keeps its
own literal
+ // parameters, e.g. the reduced bucket count.
val argument = expr match {
- case t: TransformExpression => t.children
- case e => Seq(e)
+ case t: TransformExpression => t.columnSlots.head
+ case e => e
}
- reduced.copy(children = argument)
+ reduced.rewriteColumnSlots(_ => argument)
Review Comment:
**Finding 14.** On point 2 of [your
note](https://github.com/apache/spark/pull/55885#issuecomment-6305985451):
identity and the gate agree on every pair the planner compares, with one
exception. The exception is the reduced key space this rebuild produces.
`reduced` carries `reducedWith`, the partner's identity at the time of the
reduce, which paired two columns. Rebuilding `reduced` over a member's
expression argument gives its own identity an `Expr` shape, but the partner
stays `ColumnRef`. The other side does the same in mirror image, so the two
members of one reduce stop describing the same key space.
Measured on this head. Each side is a `GroupPartitionsExec` over a
two-member collection, reduced together by `% 4`. The left side holds
`bucket(8, id)` and `bucket(8, b + 1)`, the right side `bucket(12, y)` and
`bucket(12, c + 1)`:
```
column members hasSameReducedKeys = true
expression members hasSameReducedKeys = false (true on master)
```
The comment above names that shape, a side shuffled onto this layout over `b
+ 1`, as the reason for the rebuild. A join above that pairs the two expression
members passes the `argumentShape` gate, since both are `x + 1`. It now
shuffles where master paired them.
Rebuilding the partner the same way keeps the pair together. The member that
pairs with this one on the other side has the same argument shape, since
`argumentShape` admits nothing else:
```scala
reduced.rewriteColumnSlots(_ => argument).copy(reducedWith
=
reduced.reducedWith.map(TransformExpression.withColumnShape(_, argument)))
```
```scala
/** `id` with each column argument's shape replaced by `argument`'s. */
private[sql] def withColumnShape(id: TransformFunctionId, argument:
Expression)
: TransformFunctionId = {
import TransformFunctionId.ArgumentShape
val shape = argument match {
case _: Attribute => Some(ArgumentShape.ColumnRef)
case e => expressionShape(e).map(ArgumentShape.Expr(_))
}
shape.fold(id) { s =>
id.copy(argumentShapes = id.argumentShapes.map {
case ArgumentShape.ColumnRef | ArgumentShape.Expr(_) => s
case other => other
})
}
}
```
With it the expression members match, the column-vs-expression cross pair
still does not, and the six suites stay green.
Erasing the column shapes from `reducedKeySpace` looks simpler but does not
work. It breaks "isSameFunction and hasSameReducedKeys agree on every pair" on
the struct-field fixture. The probe above, turned into a
`GroupPartitionsExecSuite` case, pins this.
---
_Generated by [Claude Code](https://claude.ai/code)_
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -92,49 +179,180 @@ case class TransformExpression(
* @param other the transform expression to compare to
* @return true if compatible, false if not
*/
- def isCompatible(other: TransformExpression): Boolean = {
- if (isSameFunction(other)) {
- true
- } else {
- (function, other.function) match {
- case (f: ReducibleFunction[_, _], o: ReducibleFunction[_, _]) =>
- val thisReducer = reducer(f, numBucketsOpt, o, other.numBucketsOpt)
- val otherReducer = reducer(o, other.numBucketsOpt, f, numBucketsOpt)
- thisReducer.isDefined || otherReducer.isDefined
- case _ => false
- }
+ def isCompatible(other: TransformExpression): Boolean =
+ isSameFunction(other) || reducers(other).isDefined ||
other.reducers(this).isDefined
+
+ /**
+ * Rewrites this transform's column arguments with `rewriteColumn`, leaving
its literal parameters
+ * untouched.
+ *
+ * Since literal parameters (the bucket count, the truncate width) live in
`children` rather than
+ * in a field of their own, any path that rewrites a transform's children
has to skip them. A
+ * parameter is not an aliasable value: substituting it changes what the
transform computes, drops
+ * it out of `functionId`, adds a second entry to `references` (tripping
+ * `KeyedShuffleSpec.keyPositions`' single-reference assert), and makes
+ * `KeyedPartitioning.supportsExpressions` reject the partitioning outright
-- so SPJ is silently
+ * lost. Use it for any retargeting of a partition transform.
+ *
+ * Callers: `KeyedShuffleSpec.rebuiltOver`, which rebuilds a transform over
the other side's key,
+ * and `GroupPartitionsExec`, which rebuilds a reduced expression over each
member's argument.
+ */
+ def rewriteColumnSlots(rewriteColumn: Expression => Expression):
TransformExpression =
+ copy(children = children.map {
+ case l: Literal => l
+ case c => rewriteColumn(c)
+ })
+
+ /** The non-literal arguments. */
+ def columnSlots: Seq[Expression] =
children.filterNot(_.isInstanceOf[Literal])
+
+ /**
+ * Extract all literal parameters of this transform as V2 [[V2Literal]]s,
preserving each value's
+ * internal representation and its `DataType`. Only consulted once a reducer
path has confirmed
+ * the literal params already match the declared input types (see
+ * [[literalParamsMatchInputTypes]]), so no type coercion happens here.
Memoized.
+ *
+ * Examples:
+ * bucket(4, col) => [Literal(4, IntegerType)]
+ * truncate(col, 3) => [Literal(3, IntegerType)]
+ * days(col) => [] (no literals)
+ */
+ private lazy val extractParameters: Array[V2Literal[_]] =
+ literalChildren.map(l => LiteralValue(l.value, l.dataType):
V2Literal[_]).toArray
+
+ /**
+ * Whether the `select`ed children match the bound function's declared input
type at their
+ * positions. A child beyond the declared arity has no declared type to
compare against (e.g. an
+ * arity-flexible function), so it is left to the connector reducer / other
guards. The DataType
+ * match is exact by design: any mismatch (including cosmetic ones like
Array `containsNull` or
+ * Decimal precision/scale) fails safe to a shuffle. See the two predicates
below for the callers.
+ *
+ * Reading `dataType` can throw rather than answer, e.g. for a
`GetStructField` over a non-struct
+ * column. Both callers are gates that must answer "not reducible" rather
than fail the query, so
+ * a child whose type cannot be read is treated as not matching.
+ */
+ private def inputTypesMatch(select: Expression => Boolean): Boolean = {
+ val declaredTypes = function.inputTypes()
+ children.zipWithIndex.forall {
+ case (c, i) =>
+ !select(c) || i >= declaredTypes.length ||
+ Try(c.dataType).toOption.contains(declaredTypes(i))
}
}
/**
- * Return a [[Reducer]] for this transform expression on another
- * on the transform expression.
- * <p>
- * A [[Reducer]] exists for a transform expression function if it is
- * 'reducible' on the other expression function.
- * <p>
- * @return reducer function or None if not reducible on the other transform
expression
+ * Whether every literal parameter matches its declared input type. Used by
the transform-vs-
+ * transform reducer path, which hands literal *values* to the connector
reducer without Analyzer
+ * type coercion. A literal whose type differs from the declared input type
(a legal implicit cast
+ * under [[BoundFunction]]) is not reducible: the join falls back to a
shuffle rather than handing
+ * the connector a value the partitions were not built on, which it would
then mis-cast. Column
+ * slots are not checked (they are not passed to the connector); the eval
path uses
+ * [[argsMatchInputTypes]] instead.
+ */
+ lazy val literalParamsMatchInputTypes: Boolean =
inputTypesMatch(_.isInstanceOf[Literal])
+
+ /**
+ * Whether every argument has its declared type, at exactly the declared
arity. Required before
+ * evaluating the transform directly: no cast is applied, and a mismatch
would fail at evaluation.
+ */
+ lazy val argsMatchInputTypes: Boolean =
+ children.length == function.inputTypes().length && inputTypesMatch(_ =>
true)
+
+ /**
+ * Whether each position holds a literal on both sides, a column reference
on both sides, or the
+ * same expression of a column on both sides (e.g. `b + 1` and `c + 1`).
Literal values and arity
+ * may differ; the connector reducer decides those.
+ */
+ private def sameArgumentLayout(other: TransformExpression): Boolean =
+ children.zip(other.children).forall {
+ case (_: Literal, _: Literal) => true
+ case (c1, c2) if TransformExpression.isColumnRef(c1) &&
TransformExpression.isColumnRef(c2) =>
+ true
+ case (c1, c2) => TransformExpression.sameShape(c1, c2)
+ }
+
+ /**
+ * Whether no literal parameter has an array, map, struct or UDT type; those
are not passed to a
+ * connector reducer. Other scalars, such as intervals, are.
+ */
+ private def noComplexLiteralParams: Boolean =
+ literalChildren.forall(_.dataType match {
+ case _: ArrayType | _: MapType | _: StructType | _: UserDefinedType[_]
=> false
+ case _ => true
+ })
+
+ /**
+ * Return a [[Reducer]] that maps this transform's partition keys onto
`other`'s, or None if there
+ * is none: when either function is not a [[ReducibleFunction]], when the
argument layouts or
+ * literal parameters rule a reducer out, or when the connector reports it
is not reducible.
+ * Handles both parameterized (bucket, truncate) and non-parameterized
(days, hours) functions.
*/
def reducers(other: TransformExpression): Option[Reducer[_, _]] = {
- (function, other.function) match {
- case(e1: ReducibleFunction[_, _], e2: ReducibleFunction[_, _]) =>
- reducer(e1, numBucketsOpt, e2, other.numBucketsOpt)
- case _ => None
+ import TransformExpression._
+ val thisFunction: ReducibleFunction[_, _] = function match {
+ case f: ReducibleFunction[_, _] => f
+ case _ => return None
+ }
+ val otherFunction: ReducibleFunction[_, _] = other.function match {
+ case o: ReducibleFunction[_, _] => o
+ case _ => return None
+ }
+ if (!sameArgumentLayout(other) ||
+ !literalParamsMatchInputTypes || !other.literalParamsMatchInputTypes ||
+ !noComplexLiteralParams || !other.noComplexLiteralParams ||
+ literalChildren.exists(_.value == null) ||
other.literalChildren.exists(_.value == null)) {
+ return None
}
- }
- // Return a Reducer for a reducible function on another reducible function
- private def reducer(
- thisFunction: ReducibleFunction[_, _],
- thisNumBucketsOpt: Option[Int],
- otherFunction: ReducibleFunction[_, _],
- otherNumBucketsOpt: Option[Int]): Option[Reducer[_, _]] = {
- val res = (thisNumBucketsOpt, otherNumBucketsOpt) match {
- case (Some(numBuckets), Some(otherNumBuckets)) =>
- thisFunction.reducer(numBuckets, otherFunction, otherNumBuckets)
- case _ => thisFunction.reducer(otherFunction)
+ val thisParams = extractParameters
+ val otherParams = other.extractParameters
+ val thisName = function.canonicalName()
+
+ // A single IntegerType parameter on each side is the shape the deprecated
+ // reducer(int, ..., int) overload accepts. Null parameters were refused
above.
+ def isSingleInt(p: Array[V2Literal[_]]): Boolean = {
+ p.length == 1 && p(0).dataType == IntegerType
+ }
+
+ // Probe one reducer overload into an Outcome. Pure -- logging is decided
once, below.
+ def probe(call: => Reducer[_, _]): Outcome = Try(Option(call)) match {
+ case Success(Some(r)) => Reducible(r)
+ case Success(None) => NotReducible
+ case Failure(_: UnsupportedOperationException) => Unimplemented
+ case Failure(e) => Threw(e)
+ }
+
+ // Prefer the generalized overload. Use the deprecated one only if the
generalized one is not
+ // implemented; any other outcome is final.
+ val outcome =
+ if (thisParams.isEmpty && otherParams.isEmpty) {
+ probe(thisFunction.reducer(otherFunction))
+ } else {
+ probe(thisFunction.reducer(thisParams, otherFunction, otherParams))
match {
+ // The deprecated overload is documented for bucket against bucket,
so only offer it a
+ // pair of the same function.
+ case Unimplemented if isSingleInt(thisParams) &&
isSingleInt(otherParams) &&
+ function.canonicalName() == other.function.canonicalName() =>
+ probe(thisFunction.reducer(
+ thisParams(0).value().asInstanceOf[Int], otherFunction,
+ otherParams(0).value().asInstanceOf[Int]))
+ case other => other
+ }
+ }
+
+ outcome match {
+ case Reducible(r) => Some(r)
+ case NotReducible => None
+ case Threw(e) =>
+ logWarning(log"V2 function ${MDC(FUNCTION_NAME, thisName)} reducer
threw an exception; " +
+ log"treating as not reducible.", e)
+ None
+ case Unimplemented =>
Review Comment:
**Finding 15.** `reducers` runs from `isCompatible` in both directions, and
again from `reducersBothWays`. `EnsureRequirements` asks that for each
candidate spec pair, and AQE re-plans each stage. So a probe that ends here, or
in `Threw` with its stack trace, logs a WARN several times per join per
planning pass. Before this PR, a pair that could not be reduced logged nothing.
The text can also be wrong. Take a legacy bucket that implements only the
deprecated overload, paired with a different function of the same `[literal,
column]` layout. The pair ends here, because the fallback needs the same
canonical name. The warning then says the function "implements no reducer" and
asks for an override of the generalized one. But the function does implement
one, and that pair would not reduce either way.
Logging once per canonical name (a set of names already warned about), or at
debug, keeps the hint without the volume. "Has no reducer for this pair" says
what was found.
---
_Generated by [Claude Code](https://claude.ai/code)_
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/GroupPartitionsExec.scala:
##########
@@ -675,11 +675,13 @@ private[sql] object GroupPartitionsExec {
// argument takes the place of `reduced`'s. A side shuffled
onto this layout
// reports `bucket(8, b + 1)` next to `bucket(8, id)`, or `b
+ 1` next to an
// identity `id`.
+ // Only the column argument is replaced; `reduced` keeps its
own literal
+ // parameters, e.g. the reduced bucket count.
val argument = expr match {
- case t: TransformExpression => t.children
- case e => Seq(e)
+ case t: TransformExpression => t.columnSlots.head
Review Comment:
**Finding 17.** `t.columnSlots.head` throws `NoSuchElementException` for a
transform with no column argument. With two column arguments,
`rewriteColumnSlots(_ => argument)` (here and in `rebuiltOver`) would put this
one argument into both slots. `supportsExpressions` admits exactly one column
slot, so nothing reaches either case today. A match that says so keeps the
invariant next to the code that relies on it:
```scala
val argument = expr match {
case t: TransformExpression => t.columnSlots match {
case Seq(slot) => slot
case slots => throw SparkException.internalError(
s"Expected one column argument in $t, got
${slots.length}")
}
case e => e
}
```
---
_Generated by [Claude Code](https://claude.ai/code)_
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -92,49 +179,180 @@ case class TransformExpression(
* @param other the transform expression to compare to
* @return true if compatible, false if not
*/
- def isCompatible(other: TransformExpression): Boolean = {
- if (isSameFunction(other)) {
- true
- } else {
- (function, other.function) match {
- case (f: ReducibleFunction[_, _], o: ReducibleFunction[_, _]) =>
- val thisReducer = reducer(f, numBucketsOpt, o, other.numBucketsOpt)
- val otherReducer = reducer(o, other.numBucketsOpt, f, numBucketsOpt)
- thisReducer.isDefined || otherReducer.isDefined
- case _ => false
- }
+ def isCompatible(other: TransformExpression): Boolean =
+ isSameFunction(other) || reducers(other).isDefined ||
other.reducers(this).isDefined
+
+ /**
+ * Rewrites this transform's column arguments with `rewriteColumn`, leaving
its literal parameters
+ * untouched.
+ *
+ * Since literal parameters (the bucket count, the truncate width) live in
`children` rather than
+ * in a field of their own, any path that rewrites a transform's children
has to skip them. A
+ * parameter is not an aliasable value: substituting it changes what the
transform computes, drops
+ * it out of `functionId`, adds a second entry to `references` (tripping
+ * `KeyedShuffleSpec.keyPositions`' single-reference assert), and makes
+ * `KeyedPartitioning.supportsExpressions` reject the partitioning outright
-- so SPJ is silently
+ * lost. Use it for any retargeting of a partition transform.
+ *
+ * Callers: `KeyedShuffleSpec.rebuiltOver`, which rebuilds a transform over
the other side's key,
+ * and `GroupPartitionsExec`, which rebuilds a reduced expression over each
member's argument.
+ */
+ def rewriteColumnSlots(rewriteColumn: Expression => Expression):
TransformExpression =
+ copy(children = children.map {
+ case l: Literal => l
+ case c => rewriteColumn(c)
+ })
+
+ /** The non-literal arguments. */
+ def columnSlots: Seq[Expression] =
children.filterNot(_.isInstanceOf[Literal])
+
+ /**
+ * Extract all literal parameters of this transform as V2 [[V2Literal]]s,
preserving each value's
+ * internal representation and its `DataType`. Only consulted once a reducer
path has confirmed
+ * the literal params already match the declared input types (see
+ * [[literalParamsMatchInputTypes]]), so no type coercion happens here.
Memoized.
+ *
+ * Examples:
+ * bucket(4, col) => [Literal(4, IntegerType)]
+ * truncate(col, 3) => [Literal(3, IntegerType)]
+ * days(col) => [] (no literals)
+ */
+ private lazy val extractParameters: Array[V2Literal[_]] =
+ literalChildren.map(l => LiteralValue(l.value, l.dataType):
V2Literal[_]).toArray
+
+ /**
+ * Whether the `select`ed children match the bound function's declared input
type at their
+ * positions. A child beyond the declared arity has no declared type to
compare against (e.g. an
+ * arity-flexible function), so it is left to the connector reducer / other
guards. The DataType
+ * match is exact by design: any mismatch (including cosmetic ones like
Array `containsNull` or
+ * Decimal precision/scale) fails safe to a shuffle. See the two predicates
below for the callers.
+ *
+ * Reading `dataType` can throw rather than answer, e.g. for a
`GetStructField` over a non-struct
+ * column. Both callers are gates that must answer "not reducible" rather
than fail the query, so
+ * a child whose type cannot be read is treated as not matching.
+ */
+ private def inputTypesMatch(select: Expression => Boolean): Boolean = {
+ val declaredTypes = function.inputTypes()
+ children.zipWithIndex.forall {
+ case (c, i) =>
+ !select(c) || i >= declaredTypes.length ||
+ Try(c.dataType).toOption.contains(declaredTypes(i))
}
}
/**
- * Return a [[Reducer]] for this transform expression on another
- * on the transform expression.
- * <p>
- * A [[Reducer]] exists for a transform expression function if it is
- * 'reducible' on the other expression function.
- * <p>
- * @return reducer function or None if not reducible on the other transform
expression
+ * Whether every literal parameter matches its declared input type. Used by
the transform-vs-
+ * transform reducer path, which hands literal *values* to the connector
reducer without Analyzer
+ * type coercion. A literal whose type differs from the declared input type
(a legal implicit cast
+ * under [[BoundFunction]]) is not reducible: the join falls back to a
shuffle rather than handing
+ * the connector a value the partitions were not built on, which it would
then mis-cast. Column
+ * slots are not checked (they are not passed to the connector); the eval
path uses
+ * [[argsMatchInputTypes]] instead.
+ */
+ lazy val literalParamsMatchInputTypes: Boolean =
inputTypesMatch(_.isInstanceOf[Literal])
+
+ /**
+ * Whether every argument has its declared type, at exactly the declared
arity. Required before
+ * evaluating the transform directly: no cast is applied, and a mismatch
would fail at evaluation.
+ */
+ lazy val argsMatchInputTypes: Boolean =
+ children.length == function.inputTypes().length && inputTypesMatch(_ =>
true)
+
+ /**
+ * Whether each position holds a literal on both sides, a column reference
on both sides, or the
+ * same expression of a column on both sides (e.g. `b + 1` and `c + 1`).
Literal values and arity
+ * may differ; the connector reducer decides those.
+ */
+ private def sameArgumentLayout(other: TransformExpression): Boolean =
+ children.zip(other.children).forall {
+ case (_: Literal, _: Literal) => true
+ case (c1, c2) if TransformExpression.isColumnRef(c1) &&
TransformExpression.isColumnRef(c2) =>
+ true
+ case (c1, c2) => TransformExpression.sameShape(c1, c2)
+ }
+
+ /**
+ * Whether no literal parameter has an array, map, struct or UDT type; those
are not passed to a
+ * connector reducer. Other scalars, such as intervals, are.
+ */
+ private def noComplexLiteralParams: Boolean =
+ literalChildren.forall(_.dataType match {
+ case _: ArrayType | _: MapType | _: StructType | _: UserDefinedType[_]
=> false
+ case _ => true
+ })
+
+ /**
+ * Return a [[Reducer]] that maps this transform's partition keys onto
`other`'s, or None if there
+ * is none: when either function is not a [[ReducibleFunction]], when the
argument layouts or
+ * literal parameters rule a reducer out, or when the connector reports it
is not reducible.
+ * Handles both parameterized (bucket, truncate) and non-parameterized
(days, hours) functions.
*/
def reducers(other: TransformExpression): Option[Reducer[_, _]] = {
- (function, other.function) match {
- case(e1: ReducibleFunction[_, _], e2: ReducibleFunction[_, _]) =>
- reducer(e1, numBucketsOpt, e2, other.numBucketsOpt)
- case _ => None
+ import TransformExpression._
+ val thisFunction: ReducibleFunction[_, _] = function match {
+ case f: ReducibleFunction[_, _] => f
+ case _ => return None
+ }
+ val otherFunction: ReducibleFunction[_, _] = other.function match {
+ case o: ReducibleFunction[_, _] => o
+ case _ => return None
+ }
+ if (!sameArgumentLayout(other) ||
+ !literalParamsMatchInputTypes || !other.literalParamsMatchInputTypes ||
+ !noComplexLiteralParams || !other.noComplexLiteralParams ||
+ literalChildren.exists(_.value == null) ||
other.literalChildren.exists(_.value == null)) {
+ return None
}
- }
- // Return a Reducer for a reducible function on another reducible function
- private def reducer(
- thisFunction: ReducibleFunction[_, _],
- thisNumBucketsOpt: Option[Int],
- otherFunction: ReducibleFunction[_, _],
- otherNumBucketsOpt: Option[Int]): Option[Reducer[_, _]] = {
- val res = (thisNumBucketsOpt, otherNumBucketsOpt) match {
- case (Some(numBuckets), Some(otherNumBuckets)) =>
- thisFunction.reducer(numBuckets, otherFunction, otherNumBuckets)
- case _ => thisFunction.reducer(otherFunction)
+ val thisParams = extractParameters
+ val otherParams = other.extractParameters
+ val thisName = function.canonicalName()
+
+ // A single IntegerType parameter on each side is the shape the deprecated
+ // reducer(int, ..., int) overload accepts. Null parameters were refused
above.
+ def isSingleInt(p: Array[V2Literal[_]]): Boolean = {
+ p.length == 1 && p(0).dataType == IntegerType
+ }
+
+ // Probe one reducer overload into an Outcome. Pure -- logging is decided
once, below.
+ def probe(call: => Reducer[_, _]): Outcome = Try(Option(call)) match {
+ case Success(Some(r)) => Reducible(r)
+ case Success(None) => NotReducible
+ case Failure(_: UnsupportedOperationException) => Unimplemented
+ case Failure(e) => Threw(e)
+ }
+
+ // Prefer the generalized overload. Use the deprecated one only if the
generalized one is not
+ // implemented; any other outcome is final.
+ val outcome =
+ if (thisParams.isEmpty && otherParams.isEmpty) {
Review Comment:
**Finding 13.** The javadoc on the generalized overload says "Spark tries
this generalized overload first" (`ReducibleFunction.java:85`). For a pair with
no parameters on either side, this branch calls `reducer(ReducibleFunction)`
instead. That works only because that method's default now delegates back, and
the delegation is what finding 9's paragraph has to warn about. A connector
that overrides both methods gets the old method's answer, which is the reverse
of the documented order.
Probing the generalized overload first for every shape fixes the order. The
one-argument overload then becomes the fallback on `Unimplemented` when both
parameter lists are empty, and its default can throw again, as on master:
```scala
val outcome =
probe(thisFunction.reducer(thisParams, otherFunction, otherParams))
match {
case Unimplemented if thisParams.isEmpty && otherParams.isEmpty =>
probe(thisFunction.reducer(otherFunction))
// The deprecated overload is documented for bucket against bucket,
so only offer it a
// pair of the same function.
case Unimplemented if isSingleInt(thisParams) &&
isSingleInt(otherParams) &&
function.canonicalName() == other.function.canonicalName() =>
probe(thisFunction.reducer(
thisParams(0).value().asInstanceOf[Int], otherFunction,
otherParams(0).value().asInstanceOf[Int]))
case outcome => outcome
}
```
In `ReducibleFunction.java`, `reducer(ReducibleFunction)` goes back to
`throw new UnsupportedOperationException();`, and the delegation paragraph goes.
Measured with finding 9's shim, a connector that implements only the
generalized overload and returns `reducer(other)` for empty parameters. On this
head, `days(a).reducers(hours(a))` dies with `StackOverflowError`. With the
change it returns `None`, so the pair shuffles. The fixtures' `DaysFunction`,
which implements only `reducer(ReducibleFunction)`, reaches it through the
`Unimplemented` arm. The six suites in finding 11 stay green. This also takes
care of finding 20.
---
_Generated by [Claude Code](https://claude.ai/code)_
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -92,49 +179,180 @@ case class TransformExpression(
* @param other the transform expression to compare to
* @return true if compatible, false if not
*/
- def isCompatible(other: TransformExpression): Boolean = {
- if (isSameFunction(other)) {
- true
- } else {
- (function, other.function) match {
- case (f: ReducibleFunction[_, _], o: ReducibleFunction[_, _]) =>
- val thisReducer = reducer(f, numBucketsOpt, o, other.numBucketsOpt)
- val otherReducer = reducer(o, other.numBucketsOpt, f, numBucketsOpt)
- thisReducer.isDefined || otherReducer.isDefined
- case _ => false
- }
+ def isCompatible(other: TransformExpression): Boolean =
+ isSameFunction(other) || reducers(other).isDefined ||
other.reducers(this).isDefined
+
+ /**
+ * Rewrites this transform's column arguments with `rewriteColumn`, leaving
its literal parameters
+ * untouched.
+ *
+ * Since literal parameters (the bucket count, the truncate width) live in
`children` rather than
+ * in a field of their own, any path that rewrites a transform's children
has to skip them. A
+ * parameter is not an aliasable value: substituting it changes what the
transform computes, drops
+ * it out of `functionId`, adds a second entry to `references` (tripping
+ * `KeyedShuffleSpec.keyPositions`' single-reference assert), and makes
+ * `KeyedPartitioning.supportsExpressions` reject the partitioning outright
-- so SPJ is silently
+ * lost. Use it for any retargeting of a partition transform.
+ *
+ * Callers: `KeyedShuffleSpec.rebuiltOver`, which rebuilds a transform over
the other side's key,
+ * and `GroupPartitionsExec`, which rebuilds a reduced expression over each
member's argument.
+ */
+ def rewriteColumnSlots(rewriteColumn: Expression => Expression):
TransformExpression =
+ copy(children = children.map {
+ case l: Literal => l
+ case c => rewriteColumn(c)
+ })
+
+ /** The non-literal arguments. */
+ def columnSlots: Seq[Expression] =
children.filterNot(_.isInstanceOf[Literal])
+
+ /**
+ * Extract all literal parameters of this transform as V2 [[V2Literal]]s,
preserving each value's
+ * internal representation and its `DataType`. Only consulted once a reducer
path has confirmed
+ * the literal params already match the declared input types (see
+ * [[literalParamsMatchInputTypes]]), so no type coercion happens here.
Memoized.
+ *
+ * Examples:
+ * bucket(4, col) => [Literal(4, IntegerType)]
+ * truncate(col, 3) => [Literal(3, IntegerType)]
+ * days(col) => [] (no literals)
+ */
+ private lazy val extractParameters: Array[V2Literal[_]] =
+ literalChildren.map(l => LiteralValue(l.value, l.dataType):
V2Literal[_]).toArray
+
+ /**
+ * Whether the `select`ed children match the bound function's declared input
type at their
+ * positions. A child beyond the declared arity has no declared type to
compare against (e.g. an
+ * arity-flexible function), so it is left to the connector reducer / other
guards. The DataType
+ * match is exact by design: any mismatch (including cosmetic ones like
Array `containsNull` or
+ * Decimal precision/scale) fails safe to a shuffle. See the two predicates
below for the callers.
+ *
+ * Reading `dataType` can throw rather than answer, e.g. for a
`GetStructField` over a non-struct
+ * column. Both callers are gates that must answer "not reducible" rather
than fail the query, so
+ * a child whose type cannot be read is treated as not matching.
+ */
+ private def inputTypesMatch(select: Expression => Boolean): Boolean = {
+ val declaredTypes = function.inputTypes()
+ children.zipWithIndex.forall {
+ case (c, i) =>
+ !select(c) || i >= declaredTypes.length ||
+ Try(c.dataType).toOption.contains(declaredTypes(i))
}
}
/**
- * Return a [[Reducer]] for this transform expression on another
- * on the transform expression.
- * <p>
- * A [[Reducer]] exists for a transform expression function if it is
- * 'reducible' on the other expression function.
- * <p>
- * @return reducer function or None if not reducible on the other transform
expression
+ * Whether every literal parameter matches its declared input type. Used by
the transform-vs-
+ * transform reducer path, which hands literal *values* to the connector
reducer without Analyzer
+ * type coercion. A literal whose type differs from the declared input type
(a legal implicit cast
+ * under [[BoundFunction]]) is not reducible: the join falls back to a
shuffle rather than handing
+ * the connector a value the partitions were not built on, which it would
then mis-cast. Column
+ * slots are not checked (they are not passed to the connector); the eval
path uses
+ * [[argsMatchInputTypes]] instead.
+ */
+ lazy val literalParamsMatchInputTypes: Boolean =
inputTypesMatch(_.isInstanceOf[Literal])
+
+ /**
+ * Whether every argument has its declared type, at exactly the declared
arity. Required before
+ * evaluating the transform directly: no cast is applied, and a mismatch
would fail at evaluation.
+ */
+ lazy val argsMatchInputTypes: Boolean =
+ children.length == function.inputTypes().length && inputTypesMatch(_ =>
true)
+
+ /**
+ * Whether each position holds a literal on both sides, a column reference
on both sides, or the
+ * same expression of a column on both sides (e.g. `b + 1` and `c + 1`).
Literal values and arity
+ * may differ; the connector reducer decides those.
+ */
+ private def sameArgumentLayout(other: TransformExpression): Boolean =
Review Comment:
**Finding 19.** This follows up
[r3441967481](https://github.com/apache/spark/pull/55885#discussion_r3441967481),
which asked for one rule here and in `isSameFunction`. The two now disagree on
struct fields. This method counts a `GetStructField` chain as a column
(`isColumnRef`), but `functionId` gives it an `Expr` shape. So
`TransformExpression.isCompatible` says true for `bucket(8, s.a)` against
`bucket(4, x)`, while `isSameFunction` says false for `bucket(4, s.a)` against
`bucket(4, x)`. Nothing pairs them today, because
`KeyedShuffleSpec.argumentShape` refuses the pair first.
The comment at `KeyGroupedPartitioningSuite.scala:7513` states this method's
rule, "A GetStructField chain IS a column reference". The `isSameFunction`
assertion below it compares `sf` with itself, so it cannot tell the two rules
apart.
Deriving the layout check from `functionId`'s shapes, with literal values
ignored, would give one rule. `argumentShape` (`partitioning.scala:2377`) could
then call `expressionShape`. Today it repeats it without the `resolved` check.
---
_Generated by [Claude Code](https://claude.ai/code)_
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/functions/ReducibleFunction.java:
##########
@@ -60,6 +61,56 @@
@Evolving
public interface ReducibleFunction<I, O> {
+ /**
+ * Generic reducer for parameterized functions (bucket, truncate, etc.).
+ *
+ * If this function is 'reducible' on another function, return the {@link
Reducer}.
+ * <p>
+ * Each parameter is a {@link Literal} carrying both its value and data
type, such as a bucket
+ * count, a truncate width, or a {@code CalendarInterval}. {@link
Literal#value()} is Spark's
+ * internal representation (e.g. {@code UTF8String} for strings, {@code
Decimal} for decimals);
+ * use {@link Literal#dataType()} to interpret it rather than assuming a JVM
type. Spark does not
+ * call this method for a pair where either side has a null literal, a
literal of array, map,
+ * struct or UDT type, or one of a type other than the function declares;
the join shuffles
Review Comment:
**Finding 23.** "or one of a type other than the function declares" holds
only up to `inputTypes().length`. `inputTypesMatch` skips the children past
that on purpose (`TransformExpression.scala:238`), so an arity-flexible
function such as the fixtures' `ZeroOrOneParamFunction` receives those literals
unchecked. A connector that trusts this sentence and casts `value()` gets a
`ClassCastException`, and Spark logs it as "reducer threw". Wording such as "or
a literal whose type differs from the one the function declares at its
position; literals past the declared arity are passed as they are" would
describe what happens.
---
_Generated by [Claude Code](https://claude.ai/code)_
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -92,49 +179,180 @@ case class TransformExpression(
* @param other the transform expression to compare to
* @return true if compatible, false if not
*/
- def isCompatible(other: TransformExpression): Boolean = {
- if (isSameFunction(other)) {
- true
- } else {
- (function, other.function) match {
- case (f: ReducibleFunction[_, _], o: ReducibleFunction[_, _]) =>
- val thisReducer = reducer(f, numBucketsOpt, o, other.numBucketsOpt)
- val otherReducer = reducer(o, other.numBucketsOpt, f, numBucketsOpt)
- thisReducer.isDefined || otherReducer.isDefined
- case _ => false
- }
+ def isCompatible(other: TransformExpression): Boolean =
+ isSameFunction(other) || reducers(other).isDefined ||
other.reducers(this).isDefined
+
+ /**
+ * Rewrites this transform's column arguments with `rewriteColumn`, leaving
its literal parameters
+ * untouched.
+ *
+ * Since literal parameters (the bucket count, the truncate width) live in
`children` rather than
+ * in a field of their own, any path that rewrites a transform's children
has to skip them. A
+ * parameter is not an aliasable value: substituting it changes what the
transform computes, drops
+ * it out of `functionId`, adds a second entry to `references` (tripping
+ * `KeyedShuffleSpec.keyPositions`' single-reference assert), and makes
+ * `KeyedPartitioning.supportsExpressions` reject the partitioning outright
-- so SPJ is silently
+ * lost. Use it for any retargeting of a partition transform.
+ *
+ * Callers: `KeyedShuffleSpec.rebuiltOver`, which rebuilds a transform over
the other side's key,
+ * and `GroupPartitionsExec`, which rebuilds a reduced expression over each
member's argument.
+ */
+ def rewriteColumnSlots(rewriteColumn: Expression => Expression):
TransformExpression =
+ copy(children = children.map {
+ case l: Literal => l
+ case c => rewriteColumn(c)
+ })
+
+ /** The non-literal arguments. */
+ def columnSlots: Seq[Expression] =
children.filterNot(_.isInstanceOf[Literal])
+
+ /**
+ * Extract all literal parameters of this transform as V2 [[V2Literal]]s,
preserving each value's
+ * internal representation and its `DataType`. Only consulted once a reducer
path has confirmed
+ * the literal params already match the declared input types (see
+ * [[literalParamsMatchInputTypes]]), so no type coercion happens here.
Memoized.
+ *
+ * Examples:
+ * bucket(4, col) => [Literal(4, IntegerType)]
+ * truncate(col, 3) => [Literal(3, IntegerType)]
+ * days(col) => [] (no literals)
+ */
+ private lazy val extractParameters: Array[V2Literal[_]] =
+ literalChildren.map(l => LiteralValue(l.value, l.dataType):
V2Literal[_]).toArray
+
+ /**
+ * Whether the `select`ed children match the bound function's declared input
type at their
+ * positions. A child beyond the declared arity has no declared type to
compare against (e.g. an
+ * arity-flexible function), so it is left to the connector reducer / other
guards. The DataType
+ * match is exact by design: any mismatch (including cosmetic ones like
Array `containsNull` or
+ * Decimal precision/scale) fails safe to a shuffle. See the two predicates
below for the callers.
+ *
+ * Reading `dataType` can throw rather than answer, e.g. for a
`GetStructField` over a non-struct
+ * column. Both callers are gates that must answer "not reducible" rather
than fail the query, so
+ * a child whose type cannot be read is treated as not matching.
+ */
+ private def inputTypesMatch(select: Expression => Boolean): Boolean = {
+ val declaredTypes = function.inputTypes()
+ children.zipWithIndex.forall {
+ case (c, i) =>
+ !select(c) || i >= declaredTypes.length ||
+ Try(c.dataType).toOption.contains(declaredTypes(i))
}
}
/**
- * Return a [[Reducer]] for this transform expression on another
- * on the transform expression.
- * <p>
- * A [[Reducer]] exists for a transform expression function if it is
- * 'reducible' on the other expression function.
- * <p>
- * @return reducer function or None if not reducible on the other transform
expression
+ * Whether every literal parameter matches its declared input type. Used by
the transform-vs-
+ * transform reducer path, which hands literal *values* to the connector
reducer without Analyzer
+ * type coercion. A literal whose type differs from the declared input type
(a legal implicit cast
+ * under [[BoundFunction]]) is not reducible: the join falls back to a
shuffle rather than handing
+ * the connector a value the partitions were not built on, which it would
then mis-cast. Column
+ * slots are not checked (they are not passed to the connector); the eval
path uses
+ * [[argsMatchInputTypes]] instead.
+ */
+ lazy val literalParamsMatchInputTypes: Boolean =
inputTypesMatch(_.isInstanceOf[Literal])
+
+ /**
+ * Whether every argument has its declared type, at exactly the declared
arity. Required before
+ * evaluating the transform directly: no cast is applied, and a mismatch
would fail at evaluation.
+ */
+ lazy val argsMatchInputTypes: Boolean =
+ children.length == function.inputTypes().length && inputTypesMatch(_ =>
true)
+
+ /**
+ * Whether each position holds a literal on both sides, a column reference
on both sides, or the
+ * same expression of a column on both sides (e.g. `b + 1` and `c + 1`).
Literal values and arity
+ * may differ; the connector reducer decides those.
+ */
+ private def sameArgumentLayout(other: TransformExpression): Boolean =
+ children.zip(other.children).forall {
+ case (_: Literal, _: Literal) => true
+ case (c1, c2) if TransformExpression.isColumnRef(c1) &&
TransformExpression.isColumnRef(c2) =>
+ true
+ case (c1, c2) => TransformExpression.sameShape(c1, c2)
+ }
+
+ /**
+ * Whether no literal parameter has an array, map, struct or UDT type; those
are not passed to a
+ * connector reducer. Other scalars, such as intervals, are.
+ */
+ private def noComplexLiteralParams: Boolean =
+ literalChildren.forall(_.dataType match {
+ case _: ArrayType | _: MapType | _: StructType | _: UserDefinedType[_]
=> false
+ case _ => true
+ })
+
+ /**
+ * Return a [[Reducer]] that maps this transform's partition keys onto
`other`'s, or None if there
+ * is none: when either function is not a [[ReducibleFunction]], when the
argument layouts or
+ * literal parameters rule a reducer out, or when the connector reports it
is not reducible.
+ * Handles both parameterized (bucket, truncate) and non-parameterized
(days, hours) functions.
*/
def reducers(other: TransformExpression): Option[Reducer[_, _]] = {
- (function, other.function) match {
- case(e1: ReducibleFunction[_, _], e2: ReducibleFunction[_, _]) =>
- reducer(e1, numBucketsOpt, e2, other.numBucketsOpt)
- case _ => None
+ import TransformExpression._
+ val thisFunction: ReducibleFunction[_, _] = function match {
+ case f: ReducibleFunction[_, _] => f
+ case _ => return None
+ }
+ val otherFunction: ReducibleFunction[_, _] = other.function match {
+ case o: ReducibleFunction[_, _] => o
+ case _ => return None
+ }
+ if (!sameArgumentLayout(other) ||
+ !literalParamsMatchInputTypes || !other.literalParamsMatchInputTypes ||
+ !noComplexLiteralParams || !other.noComplexLiteralParams ||
+ literalChildren.exists(_.value == null) ||
other.literalChildren.exists(_.value == null)) {
+ return None
}
- }
- // Return a Reducer for a reducible function on another reducible function
- private def reducer(
- thisFunction: ReducibleFunction[_, _],
- thisNumBucketsOpt: Option[Int],
- otherFunction: ReducibleFunction[_, _],
- otherNumBucketsOpt: Option[Int]): Option[Reducer[_, _]] = {
- val res = (thisNumBucketsOpt, otherNumBucketsOpt) match {
- case (Some(numBuckets), Some(otherNumBuckets)) =>
- thisFunction.reducer(numBuckets, otherFunction, otherNumBuckets)
- case _ => thisFunction.reducer(otherFunction)
+ val thisParams = extractParameters
+ val otherParams = other.extractParameters
+ val thisName = function.canonicalName()
+
+ // A single IntegerType parameter on each side is the shape the deprecated
+ // reducer(int, ..., int) overload accepts. Null parameters were refused
above.
+ def isSingleInt(p: Array[V2Literal[_]]): Boolean = {
+ p.length == 1 && p(0).dataType == IntegerType
+ }
+
+ // Probe one reducer overload into an Outcome. Pure -- logging is decided
once, below.
+ def probe(call: => Reducer[_, _]): Outcome = Try(Option(call)) match {
+ case Success(Some(r)) => Reducible(r)
+ case Success(None) => NotReducible
+ case Failure(_: UnsupportedOperationException) => Unimplemented
+ case Failure(e) => Threw(e)
+ }
+
+ // Prefer the generalized overload. Use the deprecated one only if the
generalized one is not
+ // implemented; any other outcome is final.
+ val outcome =
+ if (thisParams.isEmpty && otherParams.isEmpty) {
+ probe(thisFunction.reducer(otherFunction))
+ } else {
+ probe(thisFunction.reducer(thisParams, otherFunction, otherParams))
match {
+ // The deprecated overload is documented for bucket against bucket,
so only offer it a
+ // pair of the same function.
+ case Unimplemented if isSingleInt(thisParams) &&
isSingleInt(otherParams) &&
+ function.canonicalName() == other.function.canonicalName() =>
+ probe(thisFunction.reducer(
+ thisParams(0).value().asInstanceOf[Int], otherFunction,
+ otherParams(0).value().asInstanceOf[Int]))
+ case other => other
Review Comment:
**Finding 20.** `case other => other` rebinds `other`, so this arm and the
guard two lines up (`other.function.canonicalName()`) mean different things by
the same name.
```suggestion
case outcome => outcome
```
The guard can also use `thisName` instead of calling `canonicalName()` again.
---
_Generated by [Claude Code](https://claude.ai/code)_
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -92,49 +179,180 @@ case class TransformExpression(
* @param other the transform expression to compare to
* @return true if compatible, false if not
*/
- def isCompatible(other: TransformExpression): Boolean = {
- if (isSameFunction(other)) {
- true
- } else {
- (function, other.function) match {
- case (f: ReducibleFunction[_, _], o: ReducibleFunction[_, _]) =>
- val thisReducer = reducer(f, numBucketsOpt, o, other.numBucketsOpt)
- val otherReducer = reducer(o, other.numBucketsOpt, f, numBucketsOpt)
- thisReducer.isDefined || otherReducer.isDefined
- case _ => false
- }
+ def isCompatible(other: TransformExpression): Boolean =
+ isSameFunction(other) || reducers(other).isDefined ||
other.reducers(this).isDefined
+
+ /**
+ * Rewrites this transform's column arguments with `rewriteColumn`, leaving
its literal parameters
+ * untouched.
+ *
+ * Since literal parameters (the bucket count, the truncate width) live in
`children` rather than
+ * in a field of their own, any path that rewrites a transform's children
has to skip them. A
+ * parameter is not an aliasable value: substituting it changes what the
transform computes, drops
+ * it out of `functionId`, adds a second entry to `references` (tripping
+ * `KeyedShuffleSpec.keyPositions`' single-reference assert), and makes
Review Comment:
**Finding 22.** A few comments name code that isn't there:
- `keyPositions` has no single-reference assert, so this line's claim is
wrong. A substituted parameter gives the transform two references, so it maps
to no position (`partitioning.scala:2174`). That is the silent loss this
sentence ends with. `supportsExpressions` runs on the scan only, so it never
sees a projected `truncate(d, w)`.
- L162 and `TransformExpressionSuite.scala:137` say `sameArgumentLayout`
runs "only inside `reducer`". The method is `reducers`.
- `KeyGroupedPartitioningSuite.scala:7734` names `attempt()`, and L7821,
L7836 (a test name), L7839, L7842 and L7857 name `identityReducer`. Those are
now `probe` and `retargetForIdentity`.
- `ShuffleSpecSuite.scala:864` still says `isSameFunction` "compares names
and bucket counts".
---
_Generated by [Claude Code](https://claude.ai/code)_
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -7268,6 +7273,710 @@ class KeyGroupedPartitioningSuite
}
}
+ test("SPARK-50593: cross-function truncate vs bucket should NOT trigger
SPJ") {
+ val partitions1 = Array(
+ Expressions.apply("truncate", Expressions.column("data"),
Expressions.literal(3))
+ )
+ val partitions2 = Array(
+ Expressions.bucket(4, "data")
+ )
+
+ createTable("trunc_cross1", columns, partitions1)
+ sql("INSERT INTO testcat.ns.trunc_cross1 VALUES " +
+ "(0, 'aaa', CAST('2022-01-01' AS timestamp)), " +
+ "(1, 'bbb', CAST('2021-01-01' AS timestamp))")
+
+ createTable("trunc_cross2", columns2, partitions2)
+ sql("INSERT INTO testcat.ns.trunc_cross2 VALUES " +
+ "(1, 5, 'aaa'), " +
+ "(5, 10, 'bbb')")
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+
+ val df = sql(
+ s"""
+ |${selectWithMergeJoinHint("trunc_cross1", "trunc_cross2")}
+ |trunc_cross1.id, trunc_cross2.store_id
+ |FROM testcat.ns.trunc_cross1 JOIN testcat.ns.trunc_cross2
+ |ON trunc_cross1.data = trunc_cross2.data
+ |ORDER BY trunc_cross1.id
+ |""".stripMargin)
+
+ // Different functions (truncate vs bucket) are not mutually reducible,
so a shuffle
+ // must still be planned.
+ val shuffles = collectShuffles(df.queryExecution.executedPlan)
+ assert(shuffles.nonEmpty,
+ "truncate vs bucket are not compatible - a shuffle should be present,
" +
+ "but none was planned")
+ checkAnswer(df, Seq(Row(0, 1), Row(1, 5)))
+ }
+ }
+
+ test("SPARK-50593: truncate(3) vs truncate(5) triggers SPJ via width
reducer") {
+ // Exercises the Literal[]-based reducer path end-to-end: truncate widths
3 and 5
+ // are mutually reducible (reduce the larger to the smaller), so SPJ must
avoid the shuffle.
+ val table1 = "trunc_three"
+ val table2 = "trunc_five"
+
+ val partitions1 = Array(
+ Expressions.apply("truncate", Expressions.column("data"),
Expressions.literal(3)))
+ val partitions2 = Array(
+ Expressions.apply("truncate", Expressions.column("data"),
Expressions.literal(5)))
+
+ createTable(table1, columns, partitions1)
+ sql(s"INSERT INTO testcat.ns.$table1 VALUES " +
+ "(0, 'apple', CAST('2022-01-01' AS timestamp)), " +
+ "(1, 'grape', CAST('2021-01-01' AS timestamp)), " +
+ "(2, 'orange', CAST('2020-01-01' AS timestamp))")
+
+ createTable(table2, columns, partitions2)
+ sql(s"INSERT INTO testcat.ns.$table2 VALUES " +
+ "(10, 'apple', CAST('2022-01-01' AS timestamp)), " +
+ "(20, 'grape', CAST('2021-01-01' AS timestamp)), " +
+ "(30, 'orange', CAST('2020-01-01' AS timestamp))")
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+
+ val df = sql(
+ s"""
+ |${selectWithMergeJoinHint(table1, table2)}
+ |$table1.id AS left_id, $table2.id AS right_id
+ |FROM testcat.ns.$table1 JOIN testcat.ns.$table2
+ |ON $table1.data = $table2.data
+ |ORDER BY $table1.id
+ |""".stripMargin)
+
+ val shuffles = collectShuffles(df.queryExecution.executedPlan)
+ assert(shuffles.isEmpty,
+ "truncate(3) vs truncate(5) should avoid shuffle via the width
reducer, " +
+ "but a shuffle was planned")
+ checkAnswer(df, Seq(Row(0, 10), Row(1, 20), Row(2, 30)))
+ }
+ }
+
+ test("SPARK-50593: identical truncate widths trigger SPJ under the default
config") {
+ val table1 = "trunc_same_a"
+ val table2 = "trunc_same_b"
+ val partitions = Array(
+ Expressions.apply("truncate", Expressions.column("data"),
Expressions.literal(3)))
+
+ createTable(table1, columns, partitions)
+ sql(s"INSERT INTO testcat.ns.$table1 VALUES " +
+ "(0, 'apple', CAST('2022-01-01' AS timestamp)), " +
+ "(1, 'grape', CAST('2021-01-01' AS timestamp)), " +
+ "(2, 'orange', CAST('2020-01-01' AS timestamp))")
+
+ createTable(table2, columns, partitions)
+ sql(s"INSERT INTO testcat.ns.$table2 VALUES " +
+ "(10, 'apple', CAST('2022-01-01' AS timestamp)), " +
+ "(20, 'grape', CAST('2021-01-01' AS timestamp)), " +
+ "(30, 'orange', CAST('2020-01-01' AS timestamp))")
+
+ // No withSQLConf: identical widths are the same transform, so this has to
work on the
+ // default config, where allowCompatibleTransforms is off.
+ val df = sql(
+ s"""
+ |${selectWithMergeJoinHint(table1, table2)}
+ |$table1.id AS left_id, $table2.id AS right_id
+ |FROM testcat.ns.$table1 JOIN testcat.ns.$table2
+ |ON $table1.data = $table2.data
+ |ORDER BY $table1.id
+ |""".stripMargin)
+
+ assert(collectShuffles(df.queryExecution.executedPlan).isEmpty,
+ "identical truncate widths are the same function, so no shuffle is
needed")
+ checkAnswer(df, Seq(Row(0, 10), Row(1, 20), Row(2, 30)))
+ }
+
+ test("SPARK-50593: existing bucket SPJ still works with Literal[] API") {
+ // Exercises the new Literal[]-based reducer path end-to-end: bucket(4) and
+ // bucket(2) differ, so SPJ can only avoid the shuffle if BucketFunction's
reducer
+ // (now implemented via Literal[] params) correctly returns a GCD-based
Reducer.
+ // BucketFunction overrides only the new API, which the dispatch tries
first, so no fallback to
+ // the deprecated overload is involved.
+ val table1 = "bucket_compat1"
+ val table2 = "bucket_compat2"
+
+ val partitions1 = Array(Expressions.bucket(4, "id"))
+ val partitions2 = Array(Expressions.bucket(2, "store_id"))
+
+ createTable(table1, columns, partitions1)
+ sql(s"INSERT INTO testcat.ns.$table1 VALUES " +
+ "(0, 'aaa', CAST('2022-01-01' AS timestamp)), " +
+ "(1, 'bbb', CAST('2021-01-01' AS timestamp)), " +
+ "(2, 'ccc', CAST('2020-01-01' AS timestamp)), " +
+ "(3, 'ddd', CAST('2019-01-01' AS timestamp))")
+
+ createTable(table2, columns2, partitions2)
+ sql(s"INSERT INTO testcat.ns.$table2 VALUES " +
+ "(0, 5, 'aaa'), " +
+ "(1, 10, 'bbb'), " +
+ "(2, 15, 'ccc'), " +
+ "(3, 20, 'ddd')")
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ val df = sql(
+ s"""
+ |${selectWithMergeJoinHint(table1, table2)}
+ |$table1.id, $table2.store_id
+ |FROM testcat.ns.$table1 JOIN testcat.ns.$table2
+ |ON $table1.id = $table2.store_id
+ |ORDER BY $table1.id
+ |""".stripMargin)
+
+ val shuffles = collectShuffles(df.queryExecution.executedPlan)
+ assert(shuffles.isEmpty,
+ "bucket(4) vs bucket(2) should avoid shuffle via the GCD reducer, " +
+ "but a shuffle was planned")
+ checkAnswer(df, Seq(Row(0, 0), Row(1, 1), Row(2, 2), Row(3, 3)))
+ }
+ }
+
+ test("SPARK-50593: bucket(4) vs bucket(3) - no common divisor, must
shuffle") {
+ // GCD(4, 3) = 1 -- BucketFunction.reducer returns null. Spark must NOT
enable SPJ.
+ // Regression guard: a buggy null-handling in TransformExpression.reducer
(e.g.,
+ // Try(...).toOption instead of Try(Option(...))) would treat null as
Some(null),
+ // enable SPJ, and produce wrong join results for incompatible bucket
layouts.
+ val table1 = "bucket_gcd1_a"
+ val table2 = "bucket_gcd1_b"
+
+ val partitions1 = Array(Expressions.bucket(4, "id"))
+ val partitions2 = Array(Expressions.bucket(3, "store_id"))
+
+ createTable(table1, columns, partitions1)
+ sql(s"INSERT INTO testcat.ns.$table1 VALUES " +
+ "(0, 'aaa', CAST('2022-01-01' AS timestamp)), " +
+ "(1, 'bbb', CAST('2021-01-01' AS timestamp)), " +
+ "(2, 'ccc', CAST('2020-01-01' AS timestamp))")
+
+ createTable(table2, columns2, partitions2)
+ sql(s"INSERT INTO testcat.ns.$table2 VALUES " +
+ "(0, 5, 'aaa'), " +
+ "(1, 10, 'bbb'), " +
+ "(2, 15, 'ccc')")
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ val df = sql(
+ s"""
+ |${selectWithMergeJoinHint(table1, table2)}
+ |$table1.id, $table2.store_id
+ |FROM testcat.ns.$table1 JOIN testcat.ns.$table2
+ |ON $table1.id = $table2.store_id
+ |ORDER BY $table1.id
+ |""".stripMargin)
+
+ val shuffles = collectShuffles(df.queryExecution.executedPlan)
+ assert(shuffles.nonEmpty,
+ "bucket(4) vs bucket(3) have no common divisor (GCD=1), so the reducer
" +
+ "returns null. SPJ must NOT be enabled; a shuffle is required.")
+ checkAnswer(df, Seq(Row(0, 0), Row(1, 1), Row(2, 2)))
+ }
+ }
+
+ test("SPARK-50593: isSameFunction is a structural per-position walk, not a
literal bag") {
+ import org.apache.spark.sql.catalyst.expressions.{Add, Expression,
GetStructField}
+ val a = attr("a")
+ val b = attr("b")
+ def bucket(n: Int, e: Expression): TransformExpression =
+ TransformExpression(BucketFunction, Seq(Literal(n), e))
+ def years(e: Expression): TransformExpression =
TransformExpression(YearsFunction, Seq(e))
+ def days(e: Expression): TransformExpression =
TransformExpression(DaysFunction, Seq(e))
+
+ // Column identity is ignored -- it is reconciled separately, via
keyPositions.
+ assert(bucket(4, a).isSameFunction(bucket(4, b)))
+ assert(!bucket(4, a).isSameFunction(bucket(2, b)))
+
+ // Nested transforms are compared recursively, although both literal lists
here are [4].
+ // supportsExpressions rejects nested shapes for SPJ, but isCompatible and
ValidateRequirements
+ // consult identity too.
+ assert(!bucket(4, years(a)).isSameFunction(bucket(4, days(a))),
+ "a different nested transform is a different function")
+ assert(bucket(4, years(a)).isSameFunction(bucket(4, years(b))),
+ "the same nested transform over a different column is still the same
function")
+ assert(!bucket(4, years(a)).isSameFunction(bucket(2, years(a))), "the
literal still matters")
+
+ // An expression argument is compared by its shape, the same rule the
join's pairing uses.
+ val ra = AttributeReference("a", IntegerType)()
+ val rb = AttributeReference("b", IntegerType)()
+ val add = bucket(4, Add(ra, Literal(1)))
+ assert(add.isSameFunction(bucket(4, Add(rb, Literal(1)))), "the same shape
over another column")
+ assert(!add.isSameFunction(bucket(4, ra)), "an expression is not a bare
column")
+
+ // A GetStructField chain IS a column reference, so it is ignored like any
other column slot.
+ val s = AttributeReference("s", StructType(Seq(StructField("f",
IntegerType))))()
+ val sf = GetStructField(s, 0)
+ assert(bucket(4, sf).isSameFunction(bucket(4, sf)))
+ }
+
+ test("SPARK-50593: generalized identity is strictly more discriminating than
" +
+ "bucket-only identity") {
+ // Identity must separate truncate widths, and still separate everything a
bucket count alone
+ // separated.
+ val a = attr("a")
+ import org.apache.spark.sql.catalyst.expressions.Expression
+ def truncate(e: Expression, width: Int): TransformExpression =
+ TransformExpression(TruncateFunction, Seq(e, Literal(width)))
+ def bucket(n: Int, e: Expression): TransformExpression =
+ TransformExpression(BucketFunction, Seq(Literal(n), e))
+
+ // Truncate widths are distinguished.
+ assert(!truncate(a, 3).isSameFunction(truncate(a, 5)))
+ assert(truncate(a, 3).isSameFunction(truncate(a, 3)))
+
+ // Bucket counts and function names are still distinguished, and the
column is still ignored.
+ assert(bucket(4, a).isSameFunction(bucket(4, attr("b"))), "bucket count
equal, column ignored")
+ assert(!bucket(4, a).isSameFunction(bucket(8, a)), "bucket count differs")
+ assert(!TransformExpression(DaysFunction, Seq(a)).isSameFunction(
+ TransformExpression(YearsFunction, Seq(a))), "different function name")
+ }
+
+ test("SPARK-50593: supportsExpressions admits flat parameterized transforms,
" +
+ "rejects nested and non-reference slots") {
+ import org.apache.spark.sql.catalyst.expressions.{Add, Expression}
+ val a = AttributeReference("a", IntegerType)()
+ val b = AttributeReference("b", IntegerType)()
+ def bucket(n: Int, e: Expression): TransformExpression =
+ TransformExpression(BucketFunction, Seq(Literal(n), e))
+
+ // Flat parameterized transform over a bare column -> admitted (one
non-literal child = column).
+ assert(physical.KeyedPartitioning.supportsExpressions(Seq(bucket(4, a))))
+ // Bare identity column -> admitted.
+ assert(physical.KeyedPartitioning.supportsExpressions(Seq(a)))
+
+ // Nested transform -> rejected: the non-literal child is a transform, not
a column reference.
+ // SPJ reasons about a transform via its function and literal params
alone, which is unsound
+ // when the remaining argument is itself a transform.
+ val nested = bucket(4, TransformExpression(YearsFunction, Seq(a)))
+ assert(!physical.KeyedPartitioning.supportsExpressions(Seq(nested)))
+
+ // Value-changing slot (a + 1) -> rejected: not a plain column reference.
+ assert(!physical.KeyedPartitioning.supportsExpressions(Seq(bucket(4,
Add(a, Literal(1))))))
+
+ // Two non-literal column references -> rejected: a partition expression
must map to exactly one
+ // clustering column (the positional keyPositions model needs a single
column per transform).
+ assert(!physical.KeyedPartitioning.supportsExpressions(
+ Seq(TransformExpression(BucketFunction, Seq(Literal(4), a, b)))))
+ }
+
+ test("SPARK-50593: integer truncate is reducible via lcm (generalized
reducer, non-bucket)") {
+ // A second reducible transform exercising the generalized Literal[]
reducer API with reducer
+ // math distinct from bucket (GCD) and string truncate (prefix-min):
integer truncate snaps to
+ // a coarser grid, so truncate(v, W1) and truncate(v, W2) reduce onto
multiples of lcm(W1, W2).
+ import org.apache.spark.sql.catalyst.expressions.Expression
+ val id = attr("id")
+ def itrunc(e: Expression, w: Int): TransformExpression =
+ TransformExpression(IntegerTruncateFunction, Seq(e, Literal(w)))
+
+ // Same width -> same function (no reduction needed).
+ assert(itrunc(id, 4).isSameFunction(itrunc(id, 4)))
+
+ // W2 is a multiple of W1: the finer side (W1=2) reduces onto the coarser
grid (W2=4).
+ assert(itrunc(id, 2).isCompatible(itrunc(id, 4)))
+ val r = itrunc(id, 2).reducers(itrunc(id, 4))
+ assert(r.isDefined, "truncate(2) must reduce onto truncate(4)")
+ val red = r.get.asInstanceOf[Reducer[Integer, Integer]]
+ // truncate(.,2) values snapped to multiples of 4: 6 -> 4, 2 -> 0, 8 -> 8
+ assert(red.reduce(6) == 4 && red.reduce(2) == 0 && red.reduce(8) == 8)
+ // The coarser side (4) is already the common grid -> no reducer.
+ assert(itrunc(id, 4).reducers(itrunc(id, 2)).isEmpty)
+
+ // Neither divides the other: both sides reduce to the lcm grid.
+ assert(itrunc(id, 6).isCompatible(itrunc(id, 4))) // lcm(6, 4) =
12
+ assert(itrunc(id, 6).reducers(itrunc(id, 4)).isDefined)
+ assert(itrunc(id, 4).reducers(itrunc(id, 6)).isDefined)
+ assert(itrunc(id, 3).isCompatible(itrunc(id, 5))) // coprime ->
lcm(3, 5) = 15
+ }
+
+ test("SPARK-50593: deprecated int reducer API still works (legacy connector
backward compat)") {
+ // The dispatch falls back to the deprecated reducer(int, func, int) when
the generalized
+ // overload is not implemented, so a ReducibleFunction that overrides ONLY
the deprecated
+ // method still reduces.
+ // This mirrors how Iceberg 1.10.0 (and earlier) ship -- they predate the
Literal[] API.
+ val bucketExpr4 = TransformExpression(LegacyBucketFunction,
Seq(Literal(4), attr("id")))
+ val bucketExpr2 = TransformExpression(LegacyBucketFunction,
Seq(Literal(2), attr("id")))
+
+ val reducer = bucketExpr4.reducers(bucketExpr2)
+ assert(reducer.isDefined, "Expected a reducer for legacy_bucket(4) on
legacy_bucket(2)")
+
+ // Verify the returned Reducer actually reduces bucket 4 -> bucket 2 (GCD
= 2).
+ // bucket(4, x) produces values in [0, 4); reducing by GCD=2 gives v % 2.
+ val r = reducer.get.asInstanceOf[Reducer[Integer, Integer]]
+ assert(r.reduce(3) == 1, s"Expected reduce(3) == 1, got ${r.reduce(3)}")
+ assert(r.reduce(2) == 0, s"Expected reduce(2) == 0, got ${r.reduce(2)}")
+ }
+
+ test("SPARK-50593: the deprecated int reducer is only offered a pair of the
same function") {
+ // The deprecated overload is documented for bucket against bucket, so a
connector may not
+ // check the other function. CarelessLegacyBucketFunction does not, and
would return a gcd
+ // reducer for any pair of ints; offered a different function with the
same argument layout, it
+ // would pair partitions that do not hold the same rows.
+ val careless4 = TransformExpression(CarelessLegacyBucketFunction,
Seq(Literal(4), attr("id")))
+ val careless2 = TransformExpression(CarelessLegacyBucketFunction,
Seq(Literal(2), attr("id")))
+ val other2 = TransformExpression(LiteralFirstIntFunction, Seq(Literal(2),
attr("id")))
+
+ assert(careless4.reducers(other2).isEmpty &&
other2.reducers(careless4).isEmpty,
+ "a different function must not be offered the deprecated overload")
+ assert(!careless4.isCompatible(other2))
+ assert(careless4.reducers(careless2).isDefined, "the same function still
reduces")
+
+ // Same function means the same canonicalName; name() is only a display
name.
+ val lookalike2 = TransformExpression(CarelessLookalikeFunction,
Seq(Literal(2), attr("id")))
+ assert(careless4.reducers(lookalike2).isEmpty,
+ "a different function sharing only name() must not be offered the
deprecated overload")
+ }
+
+ test("SPARK-50593: truncate over a string and over an int are different
transforms") {
+ // The two partition data differently, so they must not share a canonical
name.
+ val strTrunc = TransformExpression(TruncateFunction, Seq(attr("data"),
Literal(3)))
+ val intTrunc = TransformExpression(IntegerTruncateFunction,
Seq(attr("id"), Literal(3)))
+ assert(!strTrunc.isSameFunction(intTrunc) &&
!intTrunc.isSameFunction(strTrunc))
+ }
+
+ test("SPARK-50593: a non-IntegerType param (DateType) does not reach the
deprecated " +
+ "int reducer") {
+ // DateType is stored as a boxed Integer (epoch days) internally, so the
reducer dispatch must
+ // key off the DataType, not the runtime class -- otherwise a DateType
param is mistaken for the
+ // bucket-style int param and routed to the deprecated reducer(int, ...).
LegacyBucketFunction
+ // overrides ONLY that deprecated method, so with a DateType param it must
be unreachable,
+ // leaving the pair not reducible (rather than producing a bogus GCD
reducer over epoch-days).
+ val l = TransformExpression(LegacyBucketFunction, Seq(Literal(8,
DateType), attr("id")))
Review Comment:
**Finding 21.** `LegacyBucketFunction` declares `IntegerType` at position 0,
so `literalParamsMatchInputTypes` refuses this `DateType` literal before
`isSingleInt` is reached. The test therefore still passes with `&&
p(0).dataType == IntegerType` deleted, which I checked by mutation.
A legacy-only fixture that declares `DateType` at that position would make
the clause load-bearing. The gate then admits the literal, and only the type
clause keeps it out of the deprecated overload.
---
_Generated by [Claude Code](https://claude.ai/code)_
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -92,49 +179,180 @@ case class TransformExpression(
* @param other the transform expression to compare to
* @return true if compatible, false if not
*/
- def isCompatible(other: TransformExpression): Boolean = {
- if (isSameFunction(other)) {
- true
- } else {
- (function, other.function) match {
- case (f: ReducibleFunction[_, _], o: ReducibleFunction[_, _]) =>
- val thisReducer = reducer(f, numBucketsOpt, o, other.numBucketsOpt)
- val otherReducer = reducer(o, other.numBucketsOpt, f, numBucketsOpt)
- thisReducer.isDefined || otherReducer.isDefined
- case _ => false
- }
+ def isCompatible(other: TransformExpression): Boolean =
+ isSameFunction(other) || reducers(other).isDefined ||
other.reducers(this).isDefined
+
+ /**
+ * Rewrites this transform's column arguments with `rewriteColumn`, leaving
its literal parameters
+ * untouched.
+ *
+ * Since literal parameters (the bucket count, the truncate width) live in
`children` rather than
+ * in a field of their own, any path that rewrites a transform's children
has to skip them. A
+ * parameter is not an aliasable value: substituting it changes what the
transform computes, drops
+ * it out of `functionId`, adds a second entry to `references` (tripping
+ * `KeyedShuffleSpec.keyPositions`' single-reference assert), and makes
+ * `KeyedPartitioning.supportsExpressions` reject the partitioning outright
-- so SPJ is silently
+ * lost. Use it for any retargeting of a partition transform.
+ *
+ * Callers: `KeyedShuffleSpec.rebuiltOver`, which rebuilds a transform over
the other side's key,
+ * and `GroupPartitionsExec`, which rebuilds a reduced expression over each
member's argument.
+ */
+ def rewriteColumnSlots(rewriteColumn: Expression => Expression):
TransformExpression =
+ copy(children = children.map {
+ case l: Literal => l
+ case c => rewriteColumn(c)
+ })
+
+ /** The non-literal arguments. */
+ def columnSlots: Seq[Expression] =
children.filterNot(_.isInstanceOf[Literal])
+
+ /**
+ * Extract all literal parameters of this transform as V2 [[V2Literal]]s,
preserving each value's
+ * internal representation and its `DataType`. Only consulted once a reducer
path has confirmed
+ * the literal params already match the declared input types (see
+ * [[literalParamsMatchInputTypes]]), so no type coercion happens here.
Memoized.
+ *
+ * Examples:
+ * bucket(4, col) => [Literal(4, IntegerType)]
+ * truncate(col, 3) => [Literal(3, IntegerType)]
+ * days(col) => [] (no literals)
+ */
+ private lazy val extractParameters: Array[V2Literal[_]] =
Review Comment:
**Finding 18.** This array is memoized and handed to the connector by
reference, in both directions and on every probe. A connector that writes into
it changes what every later call on this expression sees, for the rest of
planning, and the new javadoc can only ask it not to. Passing
`extractParameters.clone()` at the two call sites costs one shallow copy of a
one-element array per probe, and makes the contract unnecessary.
---
_Generated by [Claude Code](https://claude.ai/code)_
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -1223,20 +1222,16 @@ object KeyedPartitioning {
}
def supportsExpressions(expressions: Seq[Expression]): Boolean = {
- def isSupportedTransform(transform: TransformExpression): Boolean = {
- transform.children.size == 1 && isReference(transform.children.head)
- }
-
- @tailrec
- def isReference(e: Expression): Boolean = e match {
- case _: Attribute => true
- case g: GetStructField => isReference(g.child)
- case _ => false
- }
+ // Exactly one column argument, and it must be a plain column reference.
+ def isSupportedTransform(transform: TransformExpression): Boolean =
+ transform.columnSlots match {
Review Comment:
**Finding 16.** This is where the change reaches past joins. Any transform
with literal parameters and one column now passes. So a scan partitioned by
`truncate(data, 2)`, or by any connector function with a literal argument, now
reports `KeyedPartitioning` where it reported none. Aggregates and windows over
such a scan then skip their shuffle and run on the scan's partitions. Measured
over a `truncate(data, 2)` table:
```
master
this head
SELECT data, count(*) ... GROUP BY data 1 shuffle, 5 partitions
0 shuffles, 4 (one per key)
row_number() OVER (PARTITION BY data ORDER BY id) 1 shuffle
0 shuffles
```
The user-facing section mentions joins only. A line there would cover it. So
would a migration-guide entry, like the other 4.4 SPJ ones, if you count the
change in partition count as user-visible.
One smaller point for the same section. "No action required for existing
connectors" holds only when a bound function's `inputTypes()` lists every
argument at its exact type. Both the reduce and the identity retarget now
refuse anything else, as decided in
[r3470467090](https://github.com/apache/spark/pull/55885#discussion_r3470467090).
So a legacy bucket whose `inputTypes()` leaves out the count now loses its
reduce silently, which is why `ShuffleSpecSuite`'s `FakeBucket` had to change.
One sentence would cover it.
---
_Generated by [Claude Code](https://claude.ai/code)_
--
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]
