peter-toth commented on code in PR #55885:
URL: https://github.com/apache/spark/pull/55885#discussion_r4370810158
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -92,49 +179,203 @@ 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) || {
+ val (r1, r2) = pairReducers(other)
+ r1.isDefined || r2.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 bare column on
both sides, or the same
+ * expression of one column on both sides (e.g. `b + 1` and `c + 1`, or
`s.a` and `t.a`). These
+ * are the shapes [[functionId]] tells apart, so a pair this admits differs
at most in literal
+ * values and arity, which the connector reducer decides.
+ */
+ private def sameArgumentLayout(other: TransformExpression): Boolean =
+ children.zip(other.children).forall {
+ case (_: Literal, _: Literal) => true
+ case (_: Attribute, _: Attribute) => 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.
+ *
+ * This answers one direction only. To decide a pair, use [[pairReducers]],
which also refuses
+ * the pair when the other direction's reducer threw.
+ */
+ def reducers(other: TransformExpression): Option[Reducer[_, _]] =
+ TransformExpression.reducerOf(reducerOutcome(other))
+
+ /**
+ * Both directions' reducers, `reducers(other)` and `other.reducers(this)`,
or `(None, None)` when
+ * either direction's reducer threw: then the pair is not reducible at all
and shuffles. Reading a
+ * thrown direction as merely "not reducible" would leave the other
direction looking like a
+ * one-side reduce, which is only right for a reducer that satisfies
`r(f1(x)) = f2(x)`; a
+ * two-sided one such as a gcd of bucket counts would then pair keys from
two different spaces.
*/
- def reducers(other: TransformExpression): Option[Reducer[_, _]] = {
- (function, other.function) match {
- case(e1: ReducibleFunction[_, _], e2: ReducibleFunction[_, _]) =>
- reducer(e1, numBucketsOpt, e2, other.numBucketsOpt)
- case _ => None
+ def pairReducers(other: TransformExpression)
+ : (Option[Reducer[_, _]], Option[Reducer[_, _]]) = {
+ import TransformExpression._
+ (reducerOutcome(other), other.reducerOutcome(this)) match {
+ case (Threw(_), _) | (_, Threw(_)) => (None, None)
+ case (o1, o2) => (reducerOf(o1), reducerOf(o2))
}
}
- // 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)
- }
- Option(res)
+ /** What probing the connector for a reducer onto `other` gave, logged when
it failed. */
+ private def reducerOutcome(other: TransformExpression):
TransformExpression.Outcome = {
+ import TransformExpression._
+ val thisFunction: ReducibleFunction[_, _] = function match {
+ case f: ReducibleFunction[_, _] => f
+ case _ => return NotReducible
+ }
+ val otherFunction: ReducibleFunction[_, _] = other.function match {
+ case o: ReducibleFunction[_, _] => o
+ case _ => return NotReducible
+ }
+ if (!sameArgumentLayout(other) ||
+ !literalParamsMatchInputTypes || !other.literalParamsMatchInputTypes ||
+ !noComplexLiteralParams || !other.noComplexLiteralParams ||
+ literalChildren.exists(_.value == null) ||
other.literalChildren.exists(_.value == null)) {
+ return NotReducible
+ }
+
+ 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 =
+ probe(thisFunction.reducer(thisParams, otherFunction, otherParams))
match {
+ // A parameterless pair falls back to the parameterless overload.
+ 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() =>
Review Comment:
**Finding 24.** On `master` this overload was reached only through
`numBucketsOpt`, which only a `BucketTransform` set. So its javadoc's "This
method is for the bucket function" described what Spark did. Here any pair of
the same function with one `IntegerType` literal each gets it, `truncate(data,
2)` against `truncate(data, 4)` included. The same-function check from [finding
5](https://github.com/apache/spark/pull/55885#discussion_r4122106785) keeps
other functions out, but not another transform of the same connector function.
A connector may implement only the deprecated overload, with bucket gcd
logic, on a base class that its `truncate` shares. Its `truncate` then gets a
gcd reducer. With `truncate` bound to such a function, joining `truncate(data,
2)` with `truncate(data, 4)` under `pushPartValues` and
`allowCompatibleTransforms`:
- `master` returns all 6 rows with two shuffles, since such a scan reports
no partitioning there,
- this head fails while planning with `ClassCastException: UTF8String cannot
be cast to Integer`, from the bucket reducer applied to the truncate keys.
Over integer keys the reducer would not throw, so I'd expect it to pair keys
from different key spaces instead. The in-memory catalog truncates strings
only, so that case is not measured. Offering the deprecated overload only to
the `bucket(n, col)` layout fixes the measured case:
```scala
case Unimplemented if isSingleInt(thisParams) &&
isSingleInt(otherParams) &&
function.canonicalName() == other.function.canonicalName() &&
children.headOption.exists(_.isInstanceOf[Literal]) &&
other.children.headOption.exists(_.isInstanceOf[Literal]) =>
```
With it and the change from finding 25, the query returns 6 rows, and the
`SPARK-50593` tests in `KeyGroupedPartitioningSuite`,
`TransformExpressionSuite` and `ShuffleSpecSuite` still pass. The layout is
only a stand-in for `bucket`, though. Iceberg writes its `truncate` width first
(`Spark3Util.truncate`), so a width-first truncate still matches it. Iceberg's
truncate is not a `ReducibleFunction`, so Iceberg is not affected. For the
rest, the deprecated overload's javadoc should say that Spark may now offer it
to such a pair too.
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -7307,6 +7312,815 @@ 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: a sorted bucket with no sort columns resolves as a
bucket") {
+ // The catalog registers `bucket` only. A sort-free `sorted_bucket` is the
same partitioning,
+ // so it must resolve to that function and pair with a `bucket` table
without a shuffle.
+ val sortFree = SortedBucketTransform(
+ LogicalExpressions.literal(4, IntegerType), Seq(FieldReference("id")))
+ createTable("sb_sorted", columns, Array(sortFree))
+ sql("INSERT INTO testcat.ns.sb_sorted VALUES " +
+ "(1, 'a', CAST('2022-01-01' AS timestamp)), (5, 'b', CAST('2022-01-01'
AS timestamp))")
+ createTable("sb_plain", columns, Array(bucket(4, "id")))
+ sql("INSERT INTO testcat.ns.sb_plain VALUES " +
+ "(1, 'c', CAST('2022-01-01' AS timestamp)), (5, 'd', CAST('2022-01-01'
AS timestamp))")
+ val df = sql(
+ s"""
+ |${selectWithMergeJoinHint("sb_sorted", "sb_plain")}
+ |sb_sorted.id FROM testcat.ns.sb_sorted JOIN testcat.ns.sb_plain
+ |ON sb_sorted.id = sb_plain.id
+ |""".stripMargin)
+ assert(collectShuffles(df.queryExecution.executedPlan).isEmpty)
+ checkAnswer(df, Seq(Row(1), Row(5)))
+ }
+
+ 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")))
+ val r = TransformExpression(LegacyBucketFunction, Seq(Literal(4,
DateType), attr("id")))
+ assert(!l.isSameFunction(r))
+ assert(!l.isCompatible(r), "a DateType param must not reach the deprecated
int reducer")
+ assert(l.reducers(r).isEmpty && r.reducers(l).isEmpty)
+ }
+
+ test("SPARK-50593: struct-field arguments reach the reducer only over the
same field") {
+ import org.apache.spark.sql.catalyst.expressions.{Expression,
GetStructField}
+ // `sameArgumentLayout` uses the shapes identity tells apart: `s.a` and
`t.a` are the same
+ // expression of one column, `s.a` and `s.b` are not. BucketFunction
reduces any pair of bucket
+ // counts, so a None here comes from Spark.
+ val structType = StructType(Seq(StructField("a", LongType),
StructField("b", LongType)))
+ val s = AttributeReference("s", structType)()
+ val t = AttributeReference("t", structType)()
+ def bucket(n: Int, arg: Expression): TransformExpression =
+ TransformExpression(BucketFunction, Seq(Literal(n), arg))
+ assert(bucket(8, GetStructField(s, 0)).reducers(bucket(4,
GetStructField(t, 0))).isDefined,
+ "the same field of another struct column reaches the reducer")
+ assert(bucket(8, GetStructField(s, 0)).reducers(bucket(4,
GetStructField(s, 1))).isEmpty,
+ "a different field does not")
+ }
+
+ test("SPARK-50593: mismatched column/literal argument layout is not
reducible") {
+ // Both transforms pass the strict gate (one column-reference non-literal
child), but the column
+ // and literal sit in swapped positions: truncate(id, 2) is (col, lit)
while truncate(4, sid) is
+ // (lit, col). The reducer only sees the literal positions ([2] vs [4]),
so without an
+ // argument-layout check it would wrongly reduce these and co-locate
non-matching rows.
+ // IntegerTruncateFunction has two same-typed (Int) args, which makes this
layout reachable.
+ val l = TransformExpression(IntegerTruncateFunction, Seq(attr("id"),
Literal(2)))
+ val r = TransformExpression(IntegerTruncateFunction, Seq(Literal(4),
attr("store_id")))
+ assert(!l.isSameFunction(r))
+ assert(!l.isCompatible(r), "swapped column/literal layout must not be
reducible")
+ assert(l.reducers(r).isEmpty && r.reducers(l).isEmpty)
+
+ // Control: same layout (col, lit) on both sides remains reducible via
lcm(2, 4).
+ val a = TransformExpression(IntegerTruncateFunction, Seq(attr("id"),
Literal(2)))
+ val b = TransformExpression(IntegerTruncateFunction, Seq(attr("store_id"),
Literal(4)))
+ assert(a.isCompatible(b), "aligned (col, lit) layout must remain
reducible")
+ }
+
+ test("SPARK-50593: a dual-API connector reduces via the generalized
overload") {
+ // DualApiBucketFunction implements both overloads: the deprecated
reducer(int, ...) returns
+ // null, the generalized reducer(Literal[], ...) returns a valid GCD
reducer. Generalized-first
+ // dispatch reduces via the generalized overload directly; the deprecated
overload is not
+ // consulted (its null is irrelevant).
+ val l = TransformExpression(DualApiBucketFunction, Seq(Literal(4),
attr("id")))
+ val r = TransformExpression(DualApiBucketFunction, Seq(Literal(2),
attr("store_id")))
+ assert(l.isCompatible(r), "the generalized overload must produce a
reducer")
+ val red = l.reducers(r)
+ assert(red.isDefined, "generalized reducer must be reached")
+ assert(red.get.asInstanceOf[Reducer[Integer, Integer]].reduce(3) == 1)
+ }
+
+ test("SPARK-50593: the generalized overload's null is authoritative, no
deprecated fallback") {
+ // DualApiGeneralizedNullFunction's generalized reducer returns null (not
reducible) for a
+ // single-int pair its deprecated reducer WOULD reduce (gcd). Under
generalized-first dispatch
+ // the generalized null is authoritative: Spark must not fall back to the
deprecated overload,
+ // so the pair is not reducible. (Deprecated-first would instead
co-partition via gcd(4,2)=2.)
+ val l = TransformExpression(DualApiGeneralizedNullFunction,
Seq(Literal(4), attr("id")))
+ val r = TransformExpression(DualApiGeneralizedNullFunction,
Seq(Literal(2), attr("store_id")))
+ assert(l.reducers(r).isEmpty,
+ "a generalized null must not fall back to the deprecated overload")
+ assert(r.reducers(l).isEmpty, "symmetric")
+ }
+
+ test("SPARK-50593: a complex (non-scalar) literal param is not reducible") {
+ // Reducer parameters must not carry Catalyst-internal containers.
ArrayParamFunction's
+ // generalized reducer returns a reducer unconditionally, so reaching it
at all is the leak; the
+ // guard must refuse the ArrayData-backed literal param first. Different
array values keep
+ // isSameFunction false, forcing the reducer path where the guard applies.
+ val l = TransformExpression(ArrayParamFunction,
+ Seq(Literal.create(Array(1, 2, 3), ArrayType(IntegerType)), attr("id")))
+ val r = TransformExpression(ArrayParamFunction,
+ Seq(Literal.create(Array(4, 5, 6), ArrayType(IntegerType)),
attr("store_id")))
+ assert(!l.isSameFunction(r))
+ assert(!l.isCompatible(r), "a complex literal param must not be reducible")
+ assert(l.reducers(r).isEmpty && r.reducers(l).isEmpty)
+ }
+
+ test("SPARK-50593: a UDT-typed literal param is not reducible (complex-type
guard)") {
+ // noComplexLiteralParams rejects a UDT-typed literal param by its
DataType. UdtParamFunction
+ // declares the UDT as its literal input type, so
literalParamsMatchInputTypes passes (UDT ==
+ // UDT) and only the complex-type guard can reject it. UdtParamFunction
reduces unconditionally,
+ // so reaching its reducer is the leak (here a StructBackedUDT, whose
value is an InternalRow).
+ val udt = new StructBackedUDT
+ val l = TransformExpression(UdtParamFunction,
+ Seq(Literal(udt.serialize(new StructBacked(1)), udt), attr("id")))
+ val r = TransformExpression(UdtParamFunction,
+ Seq(Literal(udt.serialize(new StructBacked(2)), udt), attr("store_id")))
+ assert(!l.isSameFunction(r))
+ assert(!l.isCompatible(r), "a UDT-typed literal param must not be
reducible")
+ assert(l.reducers(r).isEmpty && r.reducers(l).isEmpty)
+ }
+
+ test("SPARK-50593: bundled reducers tolerate a length-mismatched params call
(no AIOOBE)") {
+ // sameArgumentLayout is arity-less, so a 0-vs-1-parameter pair (e.g.
truncate(col) vs
+ // truncate(col, w), whose column slots align) reaches the connector
reducer with
+ // mismatched-length param arrays. The bundled reducers model the
documented contract by
+ // length-checking before indexing -- returning null rather than throwing
ArrayIndexOutOfBounds
+ // (which attempt()'s Try would swallow into a silent missed SPJ + a
misleading warning).
+ val empty =
Array.empty[org.apache.spark.sql.connector.expressions.Literal[_]]
+ val one =
Array[org.apache.spark.sql.connector.expressions.Literal[_]](literal(3))
+ assert(BucketFunction.reducer(empty, BucketFunction, one) == null)
+ assert(TruncateFunction.reducer(empty, TruncateFunction, one) == null)
+ assert(IntegerTruncateFunction.reducer(empty, IntegerTruncateFunction,
one) == null)
+ }
+
+ test("SPARK-50593: a non-UOE reducer exception is logged and treated as not
reducible") {
+ // An UnsupportedOperationException means "overload not implemented"
(silent). Any other
+ // throwable is a bug in an implemented reducer: the dispatch logs it (not
the misleading
+ // "implements no reducer" hint) and treats the pair as not reducible --
it falls back to a
+ // shuffle.
+ val id = attr("id")
+ val l = TransformExpression(ThrowingReducerFunction, Seq(id, Literal(2)))
+ val r = TransformExpression(ThrowingReducerFunction, Seq(id, Literal(4)))
+ val appender = new LogAppender("non-UOE reducer exception")
+ withLogAppender(appender) {
+ assert(l.reducers(r).isEmpty, "a throwing reducer must be treated as not
reducible")
+ }
+ val messages = appender.loggingEvents.map(_.getMessage.getFormattedMessage)
+ assert(messages.exists(_.contains("reducer threw an exception")),
+ "the non-UOE exception must be logged")
+ assert(!messages.exists(_.contains("implements no reducer")),
+ "must not emit the 'implements no reducer' hint for an
implemented-but-throwing reducer")
+ }
+
+ test("SPARK-50593: a parameterless pair is offered the generalized overload
first") {
+ // GeneralizedShimFunction implements only the generalized overload and
forwards a
+ // parameterless pair to the parameterless one. If that default delegated
back, this would
+ // recurse until StackOverflowError, which `Try` does not catch.
+ val a = attr("a")
+ val l = TransformExpression(GeneralizedShimFunction, Seq(a))
+ val r = TransformExpression(DaysFunction, Seq(a))
+ assert(l.reducers(r).isEmpty)
+ // A function implementing only the parameterless overload still reduces
through the fallback.
+ val days = TransformExpression(DaysFunction, Seq(a))
+ val years = TransformExpression(YearsFunction, Seq(a))
+ assert(days.reducers(years).isDefined)
+ }
+
+ test("SPARK-50593: a reducer that throws in one direction makes the pair not
reducible") {
+ // DirectionalBucketFunction reduces bucket(16) onto bucket(12) by gcd 4
from the finer side and
+ // throws when asked from the coarser one. Reading the throw as "this
direction is not
+ // reducible" would leave a one-side reduce: the 16 side reduced by % 4
and reported as
+ // bucket(12), against the 12 side's raw keys, so only keys 0 to 3 would
meet.
+ val id = attr("id")
+ val b16 = TransformExpression(DirectionalBucketFunction, Seq(Literal(16),
id))
+ val b12 = TransformExpression(DirectionalBucketFunction, Seq(Literal(12),
id))
+ assert(b16.reducers(b12).isDefined, "the finer side alone does reduce")
+ assert(b16.pairReducers(b12) == ((None, None)) && b12.pairReducers(b16) ==
((None, None)))
+ assert(!b16.isCompatible(b12) && !b12.isCompatible(b16))
+
+ val data = AttributeReference("id", IntegerType)()
+ def spec(t: TransformExpression): physical.KeyedShuffleSpec =
physical.KeyedShuffleSpec(
+ physical.KeyedPartitioning(Seq(t.copy(children = Seq(t.children.head,
data))), Seq.empty),
+ physical.ClusteredDistribution(Seq(data)))
+ withSQLConf(
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key ->
"false",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ assert(spec(b16).reducersBothWays(spec(b12)) == ((None, None)))
+ }
+
+ withFunction(UnboundDirectionalBucketFunction) {
+ val ids = (0 until 48).map(i => s"($i, 'v$i', CAST('2022-01-01' AS
timestamp))")
+ createTable("dir16", columns, Array(bucket(16, "id")))
+ sql(s"INSERT INTO testcat.ns.dir16 VALUES ${ids.mkString(", ")}")
+ createTable("dir12", columns, Array(bucket(12, "id")))
+ sql(s"INSERT INTO testcat.ns.dir12 VALUES ${ids.mkString(", ")}")
+ withSQLConf(
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key ->
"false",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ val df = sql(
+ s"""
+ |${selectWithMergeJoinHint("dir16", "dir12")}
+ |dir16.id FROM testcat.ns.dir16 JOIN testcat.ns.dir12 ON dir16.id
= dir12.id
+ |""".stripMargin)
+ assert(df.collect().length === 48, "every id must find its match")
+ }
+ }
+ }
+
+ test("SPARK-50593: deprecated overload is not probed once the generalized
one reduces") {
+ // Generalized-first dispatch: the generalized reducer is tried first, and
the deprecated int
+ // overload must not be probed once it reduced. Here the generalized
reducer succeeds and the
+ // deprecated one throws; SPJ must succeed via the generalized path with
NO "reducer threw"
+ // warning (which an eager probe of the deprecated overload would
spuriously log).
+ val l = TransformExpression(GeneralizedOkDeprecatedThrowsFunction,
Seq(Literal(4), attr("id")))
+ val r = TransformExpression(
+ GeneralizedOkDeprecatedThrowsFunction, Seq(Literal(2), attr("store_id")))
+ val appender = new LogAppender("deprecated overload probed eagerly")
+ withLogAppender(appender) {
+ assert(l.reducers(r).isDefined, "the generalized reducer must produce a
reducer")
+ }
+ assert(!appender.loggingEvents.map(_.getMessage.getFormattedMessage)
+ .exists(_.contains("reducer threw an exception")),
+ "the deprecated overload must not be probed (and throw) once the
generalized one reduced")
+ }
+
+ test("SPARK-50593: a throwing generalized overload is surfaced, not masked
by the deprecated " +
+ "one") {
+ // DeprecatedOkGeneralizedThrowsFunction implements both: the generalized
overload throws, the
+ // deprecated one would reduce. Generalized-first dispatch treats the
generalized bug as
+ // authoritative -- it logs the exception and does NOT fall back to the
deprecated overload, so
+ // the pair shuffles. (Deprecated-first would have silently reduced via
the deprecated overload,
+ // masking the new-API bug.)
+ val l = TransformExpression(DeprecatedOkGeneralizedThrowsFunction,
Seq(Literal(4), attr("id")))
+ val r = TransformExpression(
+ DeprecatedOkGeneralizedThrowsFunction, Seq(Literal(2), attr("store_id")))
+ val appender = new LogAppender("generalized bug masked")
+ withLogAppender(appender) {
+ assert(l.reducers(r).isEmpty, "a throwing generalized overload must not
fall back and reduce")
+ }
+ assert(appender.loggingEvents.map(_.getMessage.getFormattedMessage)
+ .exists(_.contains("reducer threw an exception")), "the generalized bug
must be surfaced")
+ }
+
+ test("SPARK-50593: a typed-null integer param is not routed to the
deprecated int reducer") {
+ // LegacyIntReducerFunction implements only the deprecated int reducer and
accepts any int. A
+ // typed-null parameter would reach it as null.asInstanceOf[Int], a
fabricated 0, so Spark
+ // refuses null parameters before asking any overload.
+ val col = AttributeReference("id", IntegerType)()
+ val l = TransformExpression(LegacyIntReducerFunction, Seq(Literal(null,
IntegerType), col))
+ val r = TransformExpression(LegacyIntReducerFunction, Seq(Literal(null,
IntegerType), col))
+ assert(l.reducers(r).isEmpty,
+ "a typed-null int param must not reach the deprecated int fallback")
+ }
+
+ test("SPARK-50593: a typed-null literal param does not reach the generalized
reducer either") {
+ // BucketFunction implements the generalized overload and reads its
parameter with
+ // `value().asInstanceOf[Int]`, which turns null into 0; gcd(0, 4) would
then give a reducer.
+ val col = AttributeReference("id", IntegerType)()
+ val nullBucket = TransformExpression(BucketFunction, Seq(Literal(null,
IntegerType), col))
+ val bucket4 = TransformExpression(BucketFunction, Seq(Literal(4), col))
+ assert(nullBucket.reducers(bucket4).isEmpty &&
bucket4.reducers(nullBucket).isEmpty)
+ assert(!nullBucket.isCompatible(bucket4))
+ }
+
+ test("SPARK-50593: a column whose type differs from the declared input type
is not reducible " +
+ "(identity-vs-transform, exact-typed literal)") {
+ // Models a cross-side join `a = b` on ShortType keys: left is
identity(a), right is
+ // truncate(b, 4). The literal slot matches the declared input type
exactly, but the column does
+ // not -- identityReducer evals the transform (with the identity column
substituted in), feeding
+ // it through SpecificInternalRow(inputTypes), so a ShortType column at an
IntegerType-declared
+ // position would ClassCastException. argsMatchInputTypes checks the
column too, so the pair is
+ // not reducible (shuffle). IntegerTruncateFunction declares (IntegerType,
IntegerType).
+ val a = AttributeReference("a", ShortType)() // left, identity side
+ val b = AttributeReference("b", ShortType)() // right, transform's value
column
+ val identity = physical.KeyedPartitioning(Seq(a), Seq.empty)
+ val truncated = physical.KeyedPartitioning(
+ Seq(TransformExpression(IntegerTruncateFunction, Seq(b, Literal(4)))),
Seq.empty)
+ val idSpec = physical.KeyedShuffleSpec(identity,
physical.ClusteredDistribution(Seq(a)))
+ val trSpec = physical.KeyedShuffleSpec(truncated,
physical.ClusteredDistribution(Seq(b)))
+ assert(idSpec.reducersBothWays(trSpec) == ((None, None)),
Review Comment:
**Finding 26.** `reducersBothWays` starts with `if (!canReduceKeys) return
(None, None)`, and `canReduceKeys` needs `allowCompatibleTransforms`, which is
off by default and not set here. So this assertion, and the ones at L7983,
L8002 and L8067, hold before the check they describe runs. With
`.filter(_.argsMatchInputTypes)` removed from `retargetForIdentity`, all four
still pass, and only the gate test at L8071 fails.
Running the four calls under the confs the gate test sets makes them pin it.
Then they pass on this head, and all four fail without the filter:
```scala
withSQLConf(
SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
SQLConf.V2_BUCKETING_PARTIALLY_CLUSTERED_DISTRIBUTION_ENABLED.key ->
"false",
SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
assert(idSpec.reducersBothWays(trSpec) == ((None, None)), ...)
}
```
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -92,49 +179,203 @@ 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) || {
+ val (r1, r2) = pairReducers(other)
+ r1.isDefined || r2.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 bare column on
both sides, or the same
+ * expression of one column on both sides (e.g. `b + 1` and `c + 1`, or
`s.a` and `t.a`). These
+ * are the shapes [[functionId]] tells apart, so a pair this admits differs
at most in literal
+ * values and arity, which the connector reducer decides.
+ */
+ private def sameArgumentLayout(other: TransformExpression): Boolean =
+ children.zip(other.children).forall {
+ case (_: Literal, _: Literal) => true
+ case (_: Attribute, _: Attribute) => 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.
+ *
+ * This answers one direction only. To decide a pair, use [[pairReducers]],
which also refuses
+ * the pair when the other direction's reducer threw.
+ */
+ def reducers(other: TransformExpression): Option[Reducer[_, _]] =
+ TransformExpression.reducerOf(reducerOutcome(other))
+
+ /**
+ * Both directions' reducers, `reducers(other)` and `other.reducers(this)`,
or `(None, None)` when
+ * either direction's reducer threw: then the pair is not reducible at all
and shuffles. Reading a
+ * thrown direction as merely "not reducible" would leave the other
direction looking like a
+ * one-side reduce, which is only right for a reducer that satisfies
`r(f1(x)) = f2(x)`; a
+ * two-sided one such as a gcd of bucket counts would then pair keys from
two different spaces.
*/
- def reducers(other: TransformExpression): Option[Reducer[_, _]] = {
- (function, other.function) match {
- case(e1: ReducibleFunction[_, _], e2: ReducibleFunction[_, _]) =>
- reducer(e1, numBucketsOpt, e2, other.numBucketsOpt)
- case _ => None
+ def pairReducers(other: TransformExpression)
+ : (Option[Reducer[_, _]], Option[Reducer[_, _]]) = {
+ import TransformExpression._
+ (reducerOutcome(other), other.reducerOutcome(this)) match {
+ case (Threw(_), _) | (_, Threw(_)) => (None, None)
+ case (o1, o2) => (reducerOf(o1), reducerOf(o2))
}
}
- // 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)
- }
- Option(res)
+ /** What probing the connector for a reducer onto `other` gave, logged when
it failed. */
+ private def reducerOutcome(other: TransformExpression):
TransformExpression.Outcome = {
+ import TransformExpression._
+ val thisFunction: ReducibleFunction[_, _] = function match {
+ case f: ReducibleFunction[_, _] => f
+ case _ => return NotReducible
+ }
+ val otherFunction: ReducibleFunction[_, _] = other.function match {
+ case o: ReducibleFunction[_, _] => o
+ case _ => return NotReducible
+ }
+ if (!sameArgumentLayout(other) ||
+ !literalParamsMatchInputTypes || !other.literalParamsMatchInputTypes ||
Review Comment:
**Finding 25.** This chain refuses the pair before any overload is asked,
the deprecated and parameterless ones included. Neither of those gets a literal
it has to interpret by type. The deprecated one gets `Int`s, from `IntegerType`
literals only (`isSingleInt`), and the parameterless one gets none. So the
check protects the generalized overload only. On `master` a bucket pair reduced
whatever `inputTypes()` declared, since nothing read it there.
A connector's `bucket` may declare something else at the count's position.
`bind` hands it `(INT, <column type>)`, and the `BoundFunction` contract lets
it ask for other types. Two such shapes, each implementing only the deprecated
overload, for `bucket(8)` against `bucket(4)` under `pushPartValues` and
`allowCompatibleTransforms`:
- `inputTypes()` with `BIGINT` at the count's position: on `master`,
`areKeysCompatible(allowReduce = true)` is `true` and `reducersBothWays`
returns the connector's reducer. Here it is `false`, with `(None, None)`.
- `inputTypes()` with the column alone, as `FakeBucket` and
`TestBucketFunction` in `ShuffleSpecSuite` did before this PR changed them:
`master` returns 16 rows with no shuffle, and this head 16 rows with two
shuffles.
Iceberg's bucket functions declare `INT` at that position, so they are not
affected.
Asking the check only before the generalized overload restores the reduce.
That means dropping its two terms from the early return above, then:
```scala
val typesMatch = literalParamsMatchInputTypes &&
other.literalParamsMatchInputTypes
val outcome =
(if (typesMatch) probe(thisFunction.reducer(thisParams, otherFunction,
otherParams))
else Unimplemented) match {
// the three existing cases, unchanged
}
```
With it, both shapes reduce as on `master`, and the `SPARK-50593` tests stay
green. It also lets the `DateType` test from finding 21 reach `isSingleInt`'s
type clause, since this check no longer refuses that pair first. If you'd
rather keep the check for every overload, the migration guide's exact-type
entry should say that it also applies to the reduce between two transforms.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -2570,8 +2584,10 @@ case class KeyedShuffleSpec(
// `createPartitioning` replaces a transform's argument with the other
child's cluster key,
// which drops whatever surrounds the key, e.g. the `+ 1` of
`bucket(4, b + 1)` or the `.a`
// of `bucket(4, s.a)`. Such a spec's own child would then be laid out
differently from the
- // children shuffled onto it.
- case t: TransformExpression => argumentsAreColumns(t)
+ // children shuffled onto it. The shuffle also evaluates the rebuilt
transform per row with
+ // no cast, so its arguments must have the declared types, as for the
identity reduce
+ // (`retargetForIdentity`). The key substituted at this position has
this column's type.
+ case t: TransformExpression => argumentsAreColumns(t) &&
t.argsMatchInputTypes
Review Comment:
**Finding 27.** `pickCoPartitionTarget` (`EnsureRequirements.scala:376-423`)
asks `canCreatePartitioning` for two things:
- which layout the other children are shuffled onto,
- which layout a child is read against as it stands.
Only the first one evaluates the transform. An operator that co-partitions
but is not a sort-merge or shuffled hash join has only this path, since
`checkKeyGroupCompatible` declines for it. That is a Pandas or Arrow cogroup,
or `SortMergeAsOfJoinExec`.
With `spark.sql.sources.v2.bucketing.shuffle.enabled`, take two sides that
report the same `bucket(4, id)` and the same keys over `INT` columns, with the
test catalog's `bucket`, which declares `(INT, BIGINT)`. `master` reads them as
they stand, and nothing evaluates the transform. Here neither side can be the
target, so both get a hash shuffle. I measured it in `EnsureRequirementsSuite`,
on the shape of "SPARK-59688: a cogroup does not pair two key spaces as they
stand":
```scala
val keys = Seq(InternalRow(0), InternalRow(1), InternalRow(2),
InternalRow(3))
def bucket(a: Attribute) = TransformExpression(BucketFunction,
Seq(Literal(4), a))
val left = DummySparkPlan(outputPartitioning =
KeyedPartitioning(Seq(bucket(idL)), keys))
val right = DummySparkPlan(outputPartitioning =
KeyedPartitioning(Seq(bucket(idR)), keys))
val cogroup = FlatMapCoGroupsInPandasExec(Seq(idL), Seq(idR), pythonUdf,
AttributeReference("value", IntegerType)() :: Nil, left, right)
withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
// ShuffleExchangeExec count: 0 on master, 2 with this PR
EnsureRequirements.apply(cogroup)
}
```
A `SortMergeJoinExec` over the same two children stays at 0 on both, since
`checkKeyGroupCompatible` pairs it first.
This comes from the clause I asked for in finding 11. #59760 adds
`canEvaluate` to the same gate and has the same shape ([finding 1
there](https://github.com/apache/spark/pull/59760#discussion_r4370153109)), so
whichever lands second needs both checks. The rows are right either way, and
keeping the pairing would mean telling the two questions apart in
`pickCoPartitionTarget`, so it can wait for a follow-up. Could you name this
case in the second migration guide entry, though? It says only that "that join
shuffles instead".
--
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]
