dongjoon-hyun commented on code in PR #59760:
URL: https://github.com/apache/spark/pull/59760#discussion_r4433579770
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -202,4 +333,60 @@ object TransformExpression {
case t: TransformExpression => t.reducedWith.isDefined
case _ => false
}
+
+ /**
+ * `argument` at `declaredType`, or `None` when the function could not take
it. See
+ * [[TransformExpression.evaluableFunction]].
+ *
+ * An argument goes as it is when it is stored the way the declared type is
(`storedAlike`),
+ * which is how a function read it before this check. A narrower integral
argument is held as
+ * another JVM type, so it is cast. Anything else is refused, a `DATE`
against a `TIMESTAMP` in
+ * particular, whose cast takes the session time zone.
+ */
+ private def atDeclaredType(argument: Expression, declaredType: DataType):
Option[Expression] =
+ (argument.dataType, declaredType) match {
+ case (argumentType, _) if storedAlike(argumentType, declaredType, nested
= false) =>
+ Some(argument)
+ case (argumentType: IntegralType, declared: IntegralType)
Review Comment:
**[2] A narrower integral argument against a declared `FLOAT` or `DOUBLE` is
now refused, although master evaluated it correctly through a magic `invoke`.**
`atDeclaredType` casts only integral to integral, so `bucket(4, id)` over an
`INT` column, bound to a function that declares `(INT, DOUBLE)` with
`invoke(int, double)`, falls to `case _` at L353. On master,
`getMatchingAccessibleMethod` matched `invoke(int, double)` for the `int`
argument and reflection widened the value, so with `shuffle.enabled` only the
plain side was shuffled. Now both sides are, and the identity reduce is refused
the same way at partitioning.scala:2405. The rows are the same and the
connector shape is rare, but the description says "Plans change only where
master failed, could misread the value, or relied on a lookup that crashes",
which does not cover this case. Could `atDeclaredType` widen an integral to
`FLOAT` or `DOUBLE` with a `Cast` as well, which needs no time zone either, or
could the description list this case?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -2570,17 +2602,27 @@ 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.
+ // children shuffled onto it. Whether Spark can evaluate it is
`canEvaluatePartitioning`.
case t: TransformExpression => argumentsAreColumns(t)
case _ => false
} &&
// Shuffling another child onto these keys evaluates the partition
expressions per row to
// decide where each row goes, and reduced keys are not what those
expressions compute.
partitioning.expressionsDescribeKeys
+ /** Whether Spark can evaluate each transform of this layout, see
`canEvaluate`. */
+ override lazy val canEvaluatePartitioning: Boolean =
partitioning.expressions.forall {
Review Comment:
**[1] `canEvaluatePartitioning` asks the transform over this layout's own
column, but the one-side shuffle evaluates the one rebuilt over the shuffled
child's key (L2630).**
The `canEvaluate` doc (L2435-2439) justifies this by saying a join pairs two
keys only when their types match. That holds for an equi-join, but not for a
Python cogroup: `flatMapCoGroupsInPandas` checks only the number of keys, and
`checkKeyGroupCompatible` declines a cogroup, so it reaches
`pickCoPartitionTarget` with keys of different types. For example, with
`spark.sql.sources.v2.bucketing.shuffle.enabled=true`:
- `t1` is bucketed by `bucket(4, amount)` over `DECIMAL(10,2)`, and the
function declares `(INT, DECIMAL(10,2))` with a static `invoke(int, Decimal)`,
as Iceberg's does.
- `t1.groupby('amount').cogroup(df2.groupby('amount')).applyInPandas(...)`
runs with `df2.amount` as `DECIMAL(12,2)`.
The keyed side serves as the layout, and `df2` is shuffled onto `bucket(4,
amount: DECIMAL(12,2))`. That argument fails `f.precision <= t.precision` in
`storedAlike`, so `eval` throws `INTERNAL_ERROR` in every map task. On master
the same `StaticInvoke` got the `Decimal`, Iceberg hashes only the unscaled
value, and the cogroup returned the right rows. An `INT` bucket declaring
`(INT, INT)` against a `BIGINT` key failed on master too, but it still fails on
an executor, not while planning as L2622-2623 promise.
Also, the doc cites `DataTypeUtils.sameType`, but `EqualTo` compares with
`BinaryComparison.sameType`, which is
`DataType.equalsStructurally(ignoreNullability = true)` and ignores field names
altogether.
Could `pickCoPartitionTarget` ask `evaluableFunction` of the transforms
rebuilt over each child it shuffles, as `createPartitioning` builds them,
rather than of the reported ones?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala:
##########
@@ -1792,6 +1792,251 @@ class KeyGroupedPartitioningSuite
}
}
+ test("SPARK-60522: a transform Spark cannot evaluate is not reduced onto or
shuffled onto") {
+ // With `UnboundBucketFunction`, the first join reduces `ident` onto
`bucket(4, id)`, and the
+ // second shuffles `plain` onto it, both by evaluating `bucket`. A
connector may bind `bucket`
+ // to a function that is not a `ScalarFunction`, or to one Spark finds
nothing to call on, and
+ // Spark cannot evaluate either. Then both joins shuffle both sides.
+ Seq(
+ UnboundBucketFunction -> (0, 1),
+ UnboundNonScalarBucketFunction -> (2, 2),
+ UnboundInheritedBucketFunction -> (2, 2)
+ ).foreach {
+ case (fn, (reduceShuffles, oneSideShuffles)) =>
+ withFunction(fn) {
+ createBucketedIdTable("bucket4", 4, numIds = 8)
+ createIdentityIdTable()
+ createPlainTable()
+
+ withSQLConf(
+ SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true",
+ SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") {
+ val df = sql(
+ "SELECT k.id FROM testcat.ns.bucket4 k JOIN testcat.ns.ident i
ON k.id = i.id")
+ checkAnswer(df, (0 until 8).map(i => Row(i.toLong)))
+ assert(collectAllShuffles(df.queryExecution.executedPlan).size ==
reduceShuffles,
+ fn.name())
Review Comment:
**[8] Nit: the clue `fn.name()` is "bucket" for all three rows.**
`UnboundBucketFunction`, `UnboundNonScalarBucketFunction` and
`UnboundInheritedBucketFunction` are all named "bucket", and the two rows this
test exists for both expect `(2, 2)`, so a regression in only the
inherited-`produceResult` case fails with the same message as the non-scalar
one. The clues at L1917 and L1924 have the same problem for their two
`IntegerType` rows, and a `checkAnswer` failure carries no clue at all.
`createIdTables` (L1886) also repeats `createBucketedIdTable`,
`createIdentityIdTable` and `createPlainTable` except for the column type.
Could these use `fn.getClass.getSimpleName`, as L1954 and L1960 already do, and
could the five copied blocks (L1811-1827, L1910-1925, L1948-1961, L1984-1997,
L2018-2033) share one helper that takes the clue?
##########
sql/core/src/test/scala/org/apache/spark/sql/connector/catalog/functions/transformFunctions.scala:
##########
@@ -235,6 +235,128 @@ case class BucketReducer(divisor: Int) extends
Reducer[Int, Int] {
override def displayName(): String = toString
}
+object UnboundNonScalarBucketFunction extends UnboundFunction {
+ override def bind(inputType: StructType): BoundFunction =
NonScalarBucketFunction
+ override def description(): String = name()
+ override def name(): String = "bucket"
+}
+
+// A `bucket` that is neither a ScalarFunction nor an AggregateFunction, so
Spark cannot evaluate
+// it. The in-memory table computes the bucket values itself, see
BucketTransform in
+// InMemoryBaseTable.scala.
+object NonScalarBucketFunction extends BoundFunction {
+ override def inputTypes(): Array[DataType] = Array(IntegerType, LongType)
+ override def resultType(): DataType = IntegerType
+ override def name(): String = "bucket"
+ override def canonicalName(): String = name()
+ override def toString: String = name()
+}
+
+object UnboundInheritedBucketFunction extends UnboundFunction {
+ override def bind(inputType: StructType): BoundFunction =
InheritedBucketFunction
+ override def description(): String = name()
+ override def name(): String = "bucket"
+}
+
+abstract class BucketFunctionBase extends ScalarFunction[Int] {
+ override def produceResult(input: InternalRow): Int = {
+ Math.floorMod(input.getLong(1), input.getInt(0))
+ }
+}
+
+// A `bucket` ScalarFunction that inherits `produceResult`, which Spark does
not find when it
+// resolves the function, so Spark cannot evaluate it either.
+object InheritedBucketFunction extends BucketFunctionBase {
+ override def inputTypes(): Array[DataType] = Array(IntegerType, LongType)
+ override def resultType(): DataType = IntegerType
+ override def name(): String = "bucket"
+ override def canonicalName(): String = name()
+ override def toString: String = name()
+}
+
+object UnboundIntBucketFunction extends UnboundFunction {
+ override def bind(inputType: StructType): BoundFunction = IntBucketFunction
+ override def description(): String = name()
+ override def name(): String = "bucket"
+}
+
+// A `bucket` that declares `INT` for the value whatever it is bound to, with
a magic `invoke`, as
+// Apache Iceberg's does for a `TINYINT`, `SMALLINT` or `INT` column.
Consistent with
+// BucketTransform in InMemoryBaseTable.scala.
+object IntBucketFunction extends ScalarFunction[Int] {
+ override def inputTypes(): Array[DataType] = Array(IntegerType, IntegerType)
+ override def resultType(): DataType = IntegerType
+ override def name(): String = "bucket"
+ override def canonicalName(): String = name()
+ override def toString: String = name()
+ def invoke(numBuckets: Int, value: Int): Int = Math.floorMod(value,
numBuckets)
+}
+
+object UnboundStringBucketFunction extends UnboundFunction {
+ override def bind(inputType: StructType): BoundFunction =
StringBucketFunction
+ override def description(): String = name()
+ override def name(): String = "bucket"
+}
+
+// A `bucket` that declares a plain `STRING` for the value whatever it is
bound to. Consistent with
+// BucketTransform in InMemoryBaseTable.scala, which hashes any string type.
+object StringBucketFunction extends ScalarFunction[Int] {
+ override def inputTypes(): Array[DataType] = Array(IntegerType, StringType)
+ override def resultType(): DataType = IntegerType
+ override def name(): String = "bucket"
+ override def canonicalName(): String = name()
+ override def toString: String = name()
+ override def produceResult(input: InternalRow): Int =
+ Math.floorMod(input.getUTF8String(1).hashCode.toLong,
input.getInt(0).toLong).toInt
+}
+
+class UnboundPhysicalDaysFunction(bound: BoundFunction) extends
UnboundFunction {
+ override def bind(inputType: StructType): BoundFunction = bound
+ override def description(): String = name()
+ override def name(): String = "days"
+}
+
+// `days` that declares `INT` for a `DATE` column, which is held as the days
since the epoch.
+// Consistent with DaysTransform in InMemoryBaseTable.scala.
+object IntDaysFunction extends ScalarFunction[Int] {
+ override def inputTypes(): Array[DataType] = Array(IntegerType)
+ override def resultType(): DataType = DateType
+ override def name(): String = "days"
+ override def canonicalName(): String = name()
+ override def toString: String = name()
+ def invoke(days: Int): Int = days
+}
+
+// `days` that declares `BIGINT` for a `TIMESTAMP` column, which is held as
the micros since the
+// epoch. Consistent with DaysTransform in InMemoryBaseTable.scala.
+object LongDaysFunction extends ScalarFunction[Int] {
+ override def inputTypes(): Array[DataType] = Array(LongType)
+ override def resultType(): DataType = DateType
+ override def name(): String = "days"
+ override def canonicalName(): String = name()
+ override def toString: String = name()
+ def invoke(micros: Long): Int = Math.floorDiv(micros, 86400000000L).toInt
Review Comment:
**[4] `LongDaysFunction` floors, but the `DaysTransform` it says it is
consistent with (L330-331) truncates toward zero.**
The table's `TIMESTAMP` arm (InMemoryBaseTable.scala:301-302) computes
`ChronoUnit.DAYS.between(Instant.EPOCH, microsToInstant(micros))`, and
`Math.floorDiv` is one day lower for every pre-1970 timestamp off a UTC day
boundary: `-43200000000` micros (1969-12-31T12:00Z) is day 0 in the table and
day -1 here. The SPARK-60027 test passes only because it inserts 2021 values.
With `TIMESTAMP'1969-12-0N 10:00:00Z'` rows, the table reports days -30 to -28
while Spark computes -31 to -29, so the identity reduce and the one-side
shuffle pair rows of different days, and `checkAnswer` fails because of the
fixture, not Spark. Could `invoke` use the table's expression, or could the
comment say that they agree only from the epoch on?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala:
##########
@@ -1599,6 +1599,80 @@ class EnsureRequirementsSuite extends SharedSparkSession
{
}
}
+ test("SPARK-60522: a cogroup over a transform Spark cannot evaluate still
pairs as it stands") {
+ // `pickCoPartitionTarget` decides both which layout the other children
are shuffled onto and
+ // which a child is read against as it stands. Only the first evaluates
the transform, so a
+ // transform Spark cannot evaluate is refused only for that. A cogroup has
only this path,
+ // since `checkKeyGroupCompatible` declines for anything that is not a
`ShuffledJoin`. So two
+ // sides on the same such transform and keys, which nothing has to
evaluate, keep their
+ // layout, while a plain side is not shuffled onto it.
+ val nonScalar = new BoundFunction {
Review Comment:
**[12] Nit: both new SPARK-60522 tests define the same anonymous non-scalar
`BoundFunction` (L1609-1614 and L1651-1656).**
Its members match `NonScalarBucketFunction`, which this PR adds at
transformFunctions.scala:247 and which the `functions._` import (L32) already
brings in. Using it would drop both blocks and the `DataType` import (L41),
which exists only for them; a local `val nonScalar = NonScalarBucketFunction`
keeps L1660 under 100 columns. ShuffleSpecSuite.scala repeats the pattern: the
anonymous functions at 722 and 863, `nestedBucket` (885) and `structBucket`
(928), and the blocks at 750-760 and 820-830. Could `NonScalarBucketFunction`
move next to `FlipLowBitFunction` in
sql/catalyst/src/test/.../connector/catalog/functions/, as `FlipLowBitFunction`
did in #58943, so that both suites can use it?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -167,30 +175,153 @@ case class TransformExpression(
copy(children = newChildren)
/**
- * The scalar function call this transform stands for, with the bucket count
prepended as a
- * literal argument. `None` when the bound function is not a
[[ScalarFunction]], and also when a
- * join reduced this expression's keys, since then the call no longer
computes them. Evaluating it
- * to place a row would send the row to a partition it does not belong to.
That second arm is a
- * local gate. No caller reaches it today, because every consumer of a
reduced partitioning
- * refuses it first, and the write path never sees one.
+ * The bound function, when Spark could call it for this transform: a
[[ScalarFunction]], over
+ * keys no join reduced. See [[resolvedFunction]].
*/
- lazy val resolvedFunction: Option[Expression] = function match {
- case scalarFunc: ScalarFunction[_] if reducedWith.isEmpty =>
- val arguments = numBucketsOpt.fold(children)(n => Literal(n) +: children)
- Some(V2ExpressionUtils.resolveScalarFunction(scalarFunc, arguments))
+ private def scalarFunction: Option[ScalarFunction[_]] = function match {
+ case scalarFunc: ScalarFunction[_] if reducedWith.isEmpty =>
Some(scalarFunc)
case _ => None
}
- override def eval(input: InternalRow): Any = resolvedFunction match {
+ /** The arguments of the call: the bucket count, if any, as a literal, then
the children. */
+ private def callArguments: Seq[Expression] =
+ numBucketsOpt.fold(children)(n => Literal(n) +: children)
+
+ /**
+ * The scalar function call this transform stands for, over its arguments as
they are. `None`
+ * when the bound function is not a [[ScalarFunction]], and also when a join
reduced this
+ * expression's keys, since then the call no longer computes them.
Evaluating it to place a row
+ * would send the row to a partition it does not belong to. That second arm
is a local gate. No
+ * caller reaches it today, because every consumer of a reduced partitioning
refuses it first,
+ * and the write path never sees one. The write path
(`DistributionAndOrderingUtils`) casts the
+ * arguments to the function's input types itself. `eval` uses
[[evaluableFunction]] instead.
+ */
+ lazy val resolvedFunction: Option[Expression] =
+ scalarFunction.map(V2ExpressionUtils.resolveScalarFunction(_,
callArguments))
+
+ /**
+ * The call `eval` runs: [[resolvedFunction]], with each argument at the
type the function
+ * declares for it in `inputTypes()`. Nothing casts the arguments after
analysis, so the call
+ * gets the values as the argument holds them. Most arguments go as they
are, when they are held
+ * the way the declared type is, which is what a function read before this
check. One kind is
+ * cast: an integral argument narrower than declared is widened with a
[[Cast]], which keeps its
+ * value and needs no time zone. A function bound to a `SMALLINT` column but
declaring `INT`
+ * stays evaluable that way. Any other mismatch, or a different number of
arguments, gives
+ * `None`, as does everything `resolvedFunction` is `None` for. See
`atDeclaredType`. A mismatch
+ * is logged at info level once per function class and argument types, since
it makes a
+ * storage-partitioned join shuffle where it would not otherwise.
+ *
+ * Resolving the call can throw, e.g. when Spark finds no method to call on
a `ScalarFunction`
+ * that inherits `produceResult` from a base class, or when connector code
it calls throws. That
+ * gives `None` too, and `eval` rethrows it. So does a `LinkageError`, e.g.
a connector class
+ * whose dependency is missing, which the planner would otherwise hit for a
layout it does not
+ * pick. The planner asks this of such layouts on every query, so a failure
is logged as a
+ * warning once per function class and kind of failure, and with its stack
trace at debug level.
+ * Not serialized, so an executor resolves it again rather than receiving
the call (or the
Review Comment:
**[9] Nit: "Not serialized, so an executor resolves it again ..." (L220-221)
documents the `@transient` on `evaluableFunctionAttempt` (L229), but it sits on
the `def`.**
The field's own doc (L225-228) says nothing about serialization. The def's
doc also restates `atDeclaredType`'s rules (L341-344) and the log-once keys
documented on the two sets (L270-281), so a change such as another widening
needs edits in two docs. Could L220-221 move to the field's doc, and the def's
doc keep to its contract (the call `eval` runs, when it is `None`, the rethrow)
and why it logs?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -167,30 +175,153 @@ case class TransformExpression(
copy(children = newChildren)
/**
- * The scalar function call this transform stands for, with the bucket count
prepended as a
- * literal argument. `None` when the bound function is not a
[[ScalarFunction]], and also when a
- * join reduced this expression's keys, since then the call no longer
computes them. Evaluating it
- * to place a row would send the row to a partition it does not belong to.
That second arm is a
- * local gate. No caller reaches it today, because every consumer of a
reduced partitioning
- * refuses it first, and the write path never sees one.
+ * The bound function, when Spark could call it for this transform: a
[[ScalarFunction]], over
+ * keys no join reduced. See [[resolvedFunction]].
*/
- lazy val resolvedFunction: Option[Expression] = function match {
- case scalarFunc: ScalarFunction[_] if reducedWith.isEmpty =>
- val arguments = numBucketsOpt.fold(children)(n => Literal(n) +: children)
- Some(V2ExpressionUtils.resolveScalarFunction(scalarFunc, arguments))
+ private def scalarFunction: Option[ScalarFunction[_]] = function match {
+ case scalarFunc: ScalarFunction[_] if reducedWith.isEmpty =>
Some(scalarFunc)
case _ => None
}
- override def eval(input: InternalRow): Any = resolvedFunction match {
+ /** The arguments of the call: the bucket count, if any, as a literal, then
the children. */
+ private def callArguments: Seq[Expression] =
+ numBucketsOpt.fold(children)(n => Literal(n) +: children)
+
+ /**
+ * The scalar function call this transform stands for, over its arguments as
they are. `None`
+ * when the bound function is not a [[ScalarFunction]], and also when a join
reduced this
+ * expression's keys, since then the call no longer computes them.
Evaluating it to place a row
+ * would send the row to a partition it does not belong to. That second arm
is a local gate. No
+ * caller reaches it today, because every consumer of a reduced partitioning
refuses it first,
+ * and the write path never sees one. The write path
(`DistributionAndOrderingUtils`) casts the
+ * arguments to the function's input types itself. `eval` uses
[[evaluableFunction]] instead.
+ */
+ lazy val resolvedFunction: Option[Expression] =
+ scalarFunction.map(V2ExpressionUtils.resolveScalarFunction(_,
callArguments))
+
+ /**
+ * The call `eval` runs: [[resolvedFunction]], with each argument at the
type the function
+ * declares for it in `inputTypes()`. Nothing casts the arguments after
analysis, so the call
+ * gets the values as the argument holds them. Most arguments go as they
are, when they are held
+ * the way the declared type is, which is what a function read before this
check. One kind is
+ * cast: an integral argument narrower than declared is widened with a
[[Cast]], which keeps its
+ * value and needs no time zone. A function bound to a `SMALLINT` column but
declaring `INT`
+ * stays evaluable that way. Any other mismatch, or a different number of
arguments, gives
+ * `None`, as does everything `resolvedFunction` is `None` for. See
`atDeclaredType`. A mismatch
+ * is logged at info level once per function class and argument types, since
it makes a
+ * storage-partitioned join shuffle where it would not otherwise.
+ *
+ * Resolving the call can throw, e.g. when Spark finds no method to call on
a `ScalarFunction`
+ * that inherits `produceResult` from a base class, or when connector code
it calls throws. That
+ * gives `None` too, and `eval` rethrows it. So does a `LinkageError`, e.g.
a connector class
+ * whose dependency is missing, which the planner would otherwise hit for a
layout it does not
+ * pick. The planner asks this of such layouts on every query, so a failure
is logged as a
+ * warning once per function class and kind of failure, and with its stack
trace at debug level.
+ * Not serialized, so an executor resolves it again rather than receiving
the call (or the
+ * failure) the driver resolved while planning.
+ */
+ def evaluableFunction: Option[Expression] =
evaluableFunctionAttempt.toOption.flatten
Review Comment:
**[3] Minor: unwrapping the cached `Try` costs allocations that master did
not pay.**
`evaluableFunctionAttempt.toOption.flatten` allocates a `Some` on each call
(`Success.toOption` is `Some(value)`), and with `shuffle.enabled`, `eval`
(L258) runs once per row of the shuffled child (ShuffleExchangeExec.scala:487),
where master matched the cached `resolvedFunction`. Also,
`attempt.failed.foreach(...)` at L234 builds and drops `Failure(new
UnsupportedOperationException("Success.failed"))`, with a filled-in stack
trace, on every successful resolution: once per instance, including each
per-task bound copy and each `IdentityReducer.bound`, and again at L261 on the
`Success(None)` path. Could you log in the `catch` arms, and match the `Try`
directly in `eval` (`case Success(Some(fn)) => fn.eval(input)`, `case
Failure(e) => throw e`, `case _ => throw
QueryExecutionErrors.cannotEvaluateExpressionError(this)`)?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/expressions/TransformExpressionSuite.scala:
##########
@@ -112,4 +118,277 @@ class TransformExpressionSuite extends SparkFunSuite {
assert(!b12.hasSameReducedKeys(left))
assert(!b12.hasSameReducedKeys(b8), "nor do two unreduced ones")
}
+
+ test("SPARK-60027: eval casts only a narrower integral argument to the
function's input type") {
+ // Reads the value at the declared `BIGINT`, so it fails on an `INT` that
is not cast.
+ val longBucket = new ScalarFunction[Int] {
+ override def inputTypes(): Array[DataType] = Array(IntegerType, LongType)
+ override def resultType(): DataType = IntegerType
+ override def name(): String = "bucket"
+ override def canonicalName(): String = "test.bucket"
+ override def produceResult(input: InternalRow): Int =
+ Math.floorMod(input.getLong(1), input.getInt(0).toLong).toInt
+ }
+ Seq(Literal(7), Literal(7L), Literal(7.toShort)).foreach { value =>
+ val transform = bucket(longBucket, value)
+ assert(transform.evaluableFunction.isDefined, value.dataType)
+ assert(transform.eval(InternalRow.empty) == 3, value.dataType)
+ }
+
+ // Wider than declared, not integral, or the wrong number of arguments:
refused.
+ val intBucket = new ScalarFunction[Int] {
+ override def inputTypes(): Array[DataType] = Array(IntegerType,
IntegerType)
+ override def resultType(): DataType = IntegerType
+ override def name(): String = "bucket"
+ override def canonicalName(): String = "test.bucket"
+ override def produceResult(input: InternalRow): Int = input.getInt(1)
+ }
+ Seq(
+ bucket(intBucket, Literal(7L)),
+ bucket(longBucket, Literal("7")),
+ TransformExpression(longBucket, Seq(Literal(7L))),
+ TransformExpression(longBucket, Seq(Literal(7L), Literal(7L)), Some(4))
+ ).foreach { transform =>
+ assert(transform.evaluableFunction.isEmpty, transform)
+ intercept[SparkException](transform.eval(InternalRow.empty))
+ }
+
+ // A decimal that a declared decimal of its scale can hold, and a
`VARCHAR` against a plain
+ // string, go as they are, with the value a function read before this
check.
+ val decimalBucket = new ScalarFunction[Int] {
+ override def inputTypes(): Array[DataType] = Array(IntegerType,
DecimalType(38, 2))
+ override def resultType(): DataType = IntegerType
+ override def name(): String = "bucket"
+ override def canonicalName(): String = "test.bucket"
+ override def produceResult(input: InternalRow): Int =
+ input.getDecimal(1, 38, 2).toUnscaledLong.toInt
+ }
+ val decimal = bucket(decimalBucket, Literal(Decimal(BigDecimal("12.34"),
10, 2)))
+ assert(!decimal.evaluableFunction.get.exists(_.isInstanceOf[Cast]))
+ assert(decimal.eval(InternalRow.empty) == 1234, "the unscaled digits at
the declared scale")
+ // At another scale its unscaled digits would be another number, so it is
refused.
+ assert(bucket(decimalBucket, Literal(Decimal(BigDecimal("12.3"), 10, 1)))
+ .evaluableFunction.isEmpty)
+ val stringBucket = new ScalarFunction[Int] {
+ override def inputTypes(): Array[DataType] = Array(IntegerType,
StringType)
+ override def resultType(): DataType = IntegerType
+ override def name(): String = "bucket"
+ override def canonicalName(): String = "test.bucket"
+ override def produceResult(input: InternalRow): Int =
input.getUTF8String(1).numChars()
+ }
+ val varchar = bucket(stringBucket,
+ Literal.create(UTF8String.fromString("abc"), VarcharType(10)))
+ assert(!varchar.evaluableFunction.get.exists(_.isInstanceOf[Cast]))
+ assert(varchar.eval(InternalRow.empty) == 3)
+
+ // The write path coerces the arguments itself, so it keeps the uncast
call.
+ val stringArgument = bucket(longBucket, Literal("7"))
+ assert(stringArgument.resolvedFunction.isDefined)
+ assert(!stringArgument.resolvedFunction.get.collectFirst { case c: Cast =>
c }.isDefined)
+ }
+
+ test("SPARK-60027: an argument is taken as it is when it is stored as the
declared type is") {
+ // A function declaring `declared` reads a value of `argumentType` as it
is, field by field and
+ // by position, when the two are stored alike, as it did before the input
type check.
+ class IntUDT extends UserDefinedType[Integer] {
+ override def sqlType: DataType = IntegerType
+ override def serialize(obj: Integer): Any = obj
+ override def deserialize(datum: Any): Integer =
datum.asInstanceOf[Integer]
+ override def userClass: Class[Integer] = classOf[Integer]
+ }
+ def function(declared: DataType)(read: InternalRow => Int):
ScalarFunction[Int] =
+ new ScalarFunction[Int] {
+ override def inputTypes(): Array[DataType] = Array(IntegerType,
declared)
+ override def resultType(): DataType = IntegerType
+ override def name(): String = "bucket"
+ override def canonicalName(): String = "test.bucket"
+ override def produceResult(input: InternalRow): Int =
+ Math.floorMod(read(input), input.getInt(0))
+ }
+ val structField = (row: InternalRow) => row.getStruct(1, 1).getInt(0)
+ // (what, the argument's type, the declared type, the value, how the
function reads it,
+ // `spark.sql.caseSensitive`, whether it is taken)
+ Seq(
+ ("a field name in another case", new StructType().add("a", IntegerType),
+ new StructType().add("A", IntegerType), InternalRow(5), structField,
true, true),
+ ("a renamed field", new StructType().add("x", IntegerType),
+ new StructType().add("y", IntegerType), InternalRow(5), structField,
false, true),
+ ("a DATE field against an INT one", new StructType().add("d", DateType),
+ new StructType().add("d", IntegerType), InternalRow(5), structField,
false, true),
+ ("a nested decimal with more precision on the long side",
Review Comment:
**[6] Nothing evaluates a nested decimal through an unsafe row, which is
where the storage rule matters.**
The rows are checked at the decision level, and the accepted row's `eval`
runs on a `GenericInternalRow`, whose `getDecimal` ignores the precision it is
asked for. So if the array arm at TransformExpression.scala:380 passed `nested`
instead of `nested = true`, an `array<decimal(10,2)>` argument against a
declared `array<decimal(38,2)>` would be taken as is, and an executor's
`getArray(1).getDecimal(0, 38, 2)` would read the element's 8-byte slot as an
offset and size, while every test in this PR still passes. Dropping
`f.precision <= t.precision` would go unnoticed too, since no row has the same
scale with more precision on the argument's side. Also, `storedAlike` never
reads `SQLConf`, so the `caseSensitive` column changes nothing. Could the rows
be evaluated through `UnsafeProjection.create(...)`, with array and map decimal
rows added and the conf column dropped?
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/exchange/EnsureRequirements.scala:
##########
@@ -356,27 +356,57 @@ case class EnsureRequirements(
// consider the minimum parallelism iff ALL children need to be
re-shuffled.
//
// A child needs to be re-shuffled iff either of the following is true:
- // 1. It can't create partitioning by itself, i.e.,
`canCreatePartitioning` returns false
- // (as for the case of `RangePartitioning`), therefore it needs to be
re-shuffled
- // according to another child's shuffle spec.
+ // 1. It can't create partitioning by itself, i.e., no member of its
spec can serve as the
+ // layout (`canServe` below, as for the case of `RangePartitioning`),
therefore it needs
+ // to be re-shuffled according to another child's shuffle spec.
// 2. It already has `ShuffleExchangeLike`, so we can re-use existing
shuffle without
// introducing extra shuffle.
//
// On the other hand, in scenarios such as:
// HashPartitioning(5) <-> HashPartitioning(6)
// while `spark.sql.shuffle.partitions` is 10, we'll only re-shuffle the
left side and make it
// HashPartitioning(6).
- val shouldConsiderMinParallelism = children.zip(specs).forall { case
(child, spec) =>
- spec.forall(!_.canCreatePartitioning) ||
child.isInstanceOf[ShuffleExchangeLike]
+ //
+ // A member of a child serves as the layout when it can create a
partitioning, and either Spark
+ // can evaluate that partitioning or no child has to be shuffled onto it:
every other
+ // co-partitioned child pairs with it as it stands, which evaluates
nothing. One that cannot
+ // serve does not keep its child from being re-shuffled, so it counts
towards neither the
+ // parallelism guard nor the ranking below.
+ //
+ // A `ShuffleSpecCollection` answers `isCompatibleWith` if *any* of its
members does, so the
+ // winner alone does not say which member the sides agreed on. The
projection pushed into a
+ // paired child and the partitioning built for a re-shuffled child both
have to come from one
+ // member, otherwise the sides end up grouped on different keys, or on a
key set the child does
+ // not even have. So what is picked here is the pairing, not a member per
side.
+ //
+ // Per child, each member of its spec with the member each child pairs
with it, in child order.
+ // The pairing compares partition keys, so it is worked out when first
asked for, and once: for
+ // a member that cannot be evaluated by `canServe`, and for the winner's
members below.
+ val childMembers: Seq[Seq[LeafShuffleSpec]] =
specs.map(_.toSeq.flatMap(_.flatten))
+ final class Pairing(val member: LeafShuffleSpec, index: Int) {
+ lazy val paired: Seq[Option[LeafShuffleSpec]] =
+ childMembers.map(_.find(member.isCompatibleWith))
+ // `canEvaluatePartitioning` first: it is kept by the spec, while the
pairing compares keys.
+ lazy val canServe: Boolean =
+ member.canCreatePartitioning && (member.canEvaluatePartitioning ||
+ paired.zip(specs).zipWithIndex.forall { case ((pairedMember, spec),
j) =>
+ j == index || spec.isEmpty || pairedMember.isDefined
+ })
+ }
+ val pairings: Seq[Seq[Pairing]] = childMembers.zipWithIndex.map { case
(members, index) =>
+ members.map(new Pairing(_, index))
+ }
+ val shouldConsiderMinParallelism = children.zipWithIndex.forall { case
(child, index) =>
Review Comment:
**[11] Nit: `shouldConsiderMinParallelism` and `candidateSpecs` (L406)
switched to `zipWithIndex` only to look up `pairings(index)`.**
`pairings(i)` is empty exactly when `specs(i)` is `None`, since `flatten` is
never empty, so `children.zip(pairings).forall { case (child, ps) =>
!ps.exists(_.canServe) || child.isInstanceOf[ShuffleExchangeLike] }`, and a
`children.zip(pairings).flatMap` with the same guard and a `case _ => None`,
give the same results. They force the lazy `canServe` in the same order and
hand the same `Pairing`s to `bestPairings`. Only the `Pairing` construction at
L396 needs the index. Could these keep the `zip` shape they had before this PR?
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala:
##########
@@ -1599,6 +1599,80 @@ class EnsureRequirementsSuite extends SharedSparkSession
{
}
}
+ test("SPARK-60522: a cogroup over a transform Spark cannot evaluate still
pairs as it stands") {
+ // `pickCoPartitionTarget` decides both which layout the other children
are shuffled onto and
+ // which a child is read against as it stands. Only the first evaluates
the transform, so a
+ // transform Spark cannot evaluate is refused only for that. A cogroup has
only this path,
+ // since `checkKeyGroupCompatible` declines for anything that is not a
`ShuffledJoin`. So two
+ // sides on the same such transform and keys, which nothing has to
evaluate, keep their
+ // layout, while a plain side is not shuffled onto it.
+ val nonScalar = new BoundFunction {
+ override def inputTypes(): Array[DataType] = Array(IntegerType, LongType)
+ override def resultType(): DataType = IntegerType
+ override def name(): String = "bucket"
+ override def canonicalName(): String = "bucket"
+ }
+ val idL = AttributeReference("idL", LongType)()
+ val idR = AttributeReference("idR", LongType)()
+ val keys = Seq(InternalRow(0), InternalRow(1), InternalRow(2),
InternalRow(3))
+ val left = DummySparkPlan(outputPartitioning = KeyedPartitioning(
+ Seq(TransformExpression(nonScalar, Seq(idL), Some(4))), keys))
+ val right = DummySparkPlan(outputPartitioning = KeyedPartitioning(
+ Seq(TransformExpression(nonScalar, Seq(idR), Some(4))), keys))
+ val pythonUdf = PythonUDF("pyUDF", null,
+ StructType(Seq(StructField("value", IntegerType))),
+ Seq.empty,
+ PythonEvalType.SQL_COGROUPED_MAP_PANDAS_UDF,
+ true)
+ val cogroup = FlatMapCoGroupsInPandasExec(Seq(idL), Seq(idR), pythonUdf,
+ AttributeReference("value", IntegerType)() :: Nil, left, right)
+
+ withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ assert(EnsureRequirements.apply(cogroup).collect {
+ case s: ShuffleExchangeExec => s
+ }.isEmpty, "the two sides are read as they stand")
+ val plain = cogroup.copy(right = DummySparkPlan(outputPartitioning =
UnknownPartitioning(0)))
+ assert(!EnsureRequirements.apply(plain).collect {
+ case s: ShuffleExchangeExec => s.outputPartitioning
+ }.exists(_.isInstanceOf[KeyedPartitioning]), "a plain side is not
shuffled onto it")
+ // A sort-merge join over the same children is paired by
`checkKeyGroupCompatible` first.
+ val join = SortMergeJoinExec(Seq(idL), Seq(idR), Inner, None, left,
right)
+ assert(EnsureRequirements.apply(join).collect {
+ case s: ShuffleExchangeExec => s
+ }.isEmpty, "a join still pairs them as they stand")
+ }
+ }
+
+ test("SPARK-60522: a layout Spark cannot evaluate does not lower the shuffle
parallelism") {
Review Comment:
**[5] This PR quietly narrows the SPARK-59272 sweep "every plan the rule
produces for a generated pair of shapes validates" (L2905).**
Its `years` shape (L2930) builds `years(a)` (L1447-1448) over `exprA` and
`exprB`, both `INT` (L45-46), and `YearsFunction` declares `TIMESTAMP`.
`atDeclaredType` refuses that pair, so every `years` cell is now a layout Spark
cannot evaluate, and with `keyedShuffle = true` it can no longer be the layout
another child is shuffled onto. For `years/123/unmarked` against
`id/12/unmarked`, for example, master shuffled the right side onto
`KeyedPartitioning([years(b)], ...)`, and this PR picks the right side's `[b]`
layout instead. The sweep asserts only
`ValidateRequirements.validate(planned)`, so it still passes, but the cells
that exercised `createPartitioning` with a non-bucket template (marked, the
reordered `312`, the interleaved keys) are gone without a failure. The sweep is
outside this diff, so I left this here. Could the `years` shape use a fixture
that declares the attribute's type, or a `DATE` or `TIMESTAMP` attribute?
##########
sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/ShuffleSpecSuite.scala:
##########
@@ -705,6 +713,249 @@ class ShuffleSpecSuite extends SparkFunSuite with
SQLHelper {
}
}
+ test("SPARK-60522: KeyedShuffleSpec refuses a transform Spark cannot
evaluate") {
+ // A connector may bind its partition transform to a `BoundFunction` that
is not a
+ // `ScalarFunction`, or to a `ScalarFunction` whose `produceResult` is
inherited, which
+ // resolving it does not find. `TransformExpression.eval` cannot run
either, so an identity
+ // side is not reduced onto it and no other child is shuffled onto it. The
gates answer that
+ // without throwing, since they are asked of candidates that may not be
chosen.
+ val nonScalarBucket = new BoundFunction {
+ override def inputTypes(): Array[DataType] = Array(IntegerType,
IntegerType)
+ override def resultType(): DataType = IntegerType
+ override def name(): String = "bucket"
+ override def canonicalName(): String = "test.bucket"
+ }
+ abstract class BucketBase extends ScalarFunction[Int] {
+ override def produceResult(input: InternalRow): Int = input.getInt(1)
+ }
+ // Declares the bucket count too, so that only the inherited
`produceResult` refuses it.
+ val inheritedBucket = new BucketBase {
+ override def inputTypes(): Array[DataType] = Array(IntegerType,
IntegerType)
+ override def resultType(): DataType = IntegerType
+ override def name(): String = "bucket"
+ override def canonicalName(): String = "test.bucket"
+ }
+ val a = $"a".int
+ def spec(expression: Expression): KeyedShuffleSpec = KeyedShuffleSpec(
+ KeyedPartitioning(Seq(expression), Seq(InternalRow(1))),
ClusteredDistribution(Seq(a)))
+ val identity = spec(a)
+
+ Seq(
+ ("scalar", TestBucketFunction, true),
+ ("non-scalar", nonScalarBucket, false),
+ ("inherited produceResult", inheritedBucket, false)
+ ).foreach {
+ case (name, fn, evaluable) =>
+ val bucket = spec(TransformExpression(fn, Seq(a), Some(4)))
+ 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(identity.areKeysCompatible(bucket, allowReduce = true) ==
evaluable, name)
+ assert(bucket.areKeysCompatible(identity, allowReduce = true) ==
evaluable, name)
+ }
+ withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ assert(bucket.canCreatePartitioning, s"$name: a layout as it stands")
+ assert(bucket.canEvaluatePartitioning == evaluable, name)
+ }
+ }
+ }
+
+ test("SPARK-60027: KeyedShuffleSpec refuses a transform whose arguments miss
its input types") {
+ // Nothing casts a transform's arguments, the bucket count included, to
its function's
+ // `inputTypes()` after analysis, so the call would get the raw values. So
the gates require
+ // the arguments to match the input types in number and type. A narrower
integral argument is
+ // widened, as the cast keeps its value, and anything else is refused.
+ def bucketFunction(declared: DataType*): ScalarFunction[Int] = new
ScalarFunction[Int] {
+ override def inputTypes(): Array[DataType] = declared.toArray
+ override def resultType(): DataType = IntegerType
+ override def name(): String = "bucket"
+ override def canonicalName(): String = "test.bucket"
+ override def produceResult(input: InternalRow): Int = 0
+ }
+ def spec(expression: Expression, key: InternalRow): KeyedShuffleSpec =
KeyedShuffleSpec(
+ KeyedPartitioning(Seq(expression), Seq(key)),
+ ClusteredDistribution(expression.references.toSeq))
+
+ // (the function's input types, the column's type, whether Spark can
evaluate the transform)
+ Seq(
+ (Seq(IntegerType, LongType), LongType, true),
+ (Seq(IntegerType, LongType), IntegerType, true),
+ (Seq(IntegerType, LongType), ShortType, true),
+ (Seq(IntegerType, IntegerType), LongType, false),
+ (Seq(IntegerType, StringType), IntegerType, false),
+ (Seq(LongType, LongType), LongType, true),
+ (Seq(ShortType, LongType), LongType, false),
+ (Seq(LongType), LongType, false),
+ (Seq(IntegerType, LongType, LongType), LongType, false),
+ // A `CHAR`, `VARCHAR` or `STRING` is held as a string of its collation,
so it suits any of
+ // them of that collation. Not another collation, which compares strings
differently.
+ (Seq(IntegerType, StringType), VarcharType(10), true),
+ (Seq(IntegerType, StringType), CharType(5), true),
+ (Seq(IntegerType, VarcharType(5)), StringType, true),
+ (Seq(IntegerType, StringType), StringType("UTF8_LCASE"), false),
+ (Seq(IntegerType, StringType("UTF8_LCASE")), VarcharType(10), false),
+ // A decimal suits a declared decimal of its scale that can hold its
value, so that its
+ // unscaled digits are what a function reads. Not one of another scale.
+ (Seq(IntegerType, DecimalType(38, 2)), DecimalType(10, 2), true),
+ (Seq(IntegerType, DecimalType(38, 18)), DecimalType(10, 2), false),
+ (Seq(IntegerType, DecimalType(10, 2)), DecimalType(5, 0), false),
+ // A `DATE` and a `TIMESTAMP` are held as an `INT` and a `BIGINT`, so
they suit those.
+ (Seq(IntegerType, IntegerType), DateType, true),
+ (Seq(IntegerType, LongType), TimestampType, true),
+ (Seq(IntegerType, DecimalType(10, 2)), DecimalType(38, 18), false),
+ (Seq(IntegerType, DecimalType(10, 4)), DecimalType(10, 2), false)
+ ).foreach { case (declared, columnType, evaluable) =>
+ val name = s"${declared.mkString("(", ", ", ")")} over $columnType"
+ val column = AttributeReference("k", columnType)()
+ val bucket = spec(
+ TransformExpression(bucketFunction(declared: _*), Seq(column),
Some(4)), InternalRow(1))
+ val identityKey = columnType match {
+ case LongType | TimestampType => InternalRow(1L)
+ case _: StringType => InternalRow(UTF8String.fromString("a"))
+ case _: DecimalType => InternalRow(Decimal(1))
+ case _ => InternalRow(1)
+ }
+ val identity = spec(column, identityKey)
+ 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(identity.areKeysCompatible(bucket, allowReduce = true) ==
evaluable, name)
+ assert(bucket.areKeysCompatible(identity, allowReduce = true) ==
evaluable, name)
+ }
+ withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ assert(bucket.canCreatePartitioning, s"$name: a layout as it stands")
+ assert(bucket.canEvaluatePartitioning == evaluable, name)
+ }
+ }
+
+ // A `years`-like transform over a `DATE` column. A `DATE` matches only a
declared `DATE`, as
+ // Iceberg's `years` declares, and an `INT`, which a `DATE` is held as.
Not a `TIMESTAMP`, as
+ // the test catalog's `years` declares, which is held as a `BIGINT`, since
its cast takes the
+ // session time zone. Nor a `TIMESTAMP_NTZ`, also held as a `BIGINT`.
+ val d = AttributeReference("d", DateType)()
+ Seq(
+ DateType -> true,
+ TimestampType -> false,
+ IntegerType -> true,
+ TimestampNTZType -> false
+ ).foreach { case (declared, evaluable) =>
+ val yearsSpec = spec(TransformExpression(bucketFunction(declared),
Seq(d)), InternalRow(1))
+ val identity = spec(d, InternalRow(1))
+ 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(identity.areKeysCompatible(yearsSpec, allowReduce = true) ==
evaluable,
+ s"DATE against $declared")
+ }
+ withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ assert(yearsSpec.canEvaluatePartitioning == evaluable, s"DATE against
$declared")
+ }
+ }
+ }
+
+ test("SPARK-60522: createPartitioning refuses a layout Spark cannot
evaluate") {
+ // `canCreatePartitioning` answers for a layout read as it stands, which
evaluates nothing, so
+ // a caller that shuffles a child onto the layout has to ask
`canEvaluatePartitioning` too.
+ // One that does not fails while planning, not on an executor.
+ val nonScalar = new BoundFunction {
+ override def inputTypes(): Array[DataType] = Array(IntegerType,
IntegerType)
+ override def resultType(): DataType = IntegerType
+ override def name(): String = "bucket"
+ override def canonicalName(): String = "test.bucket"
+ }
+ val a = $"a".int
+ val spec = KeyedShuffleSpec(KeyedPartitioning(
+ Seq(TransformExpression(nonScalar, Seq(a), Some(4))),
Seq(InternalRow(1))),
+ ClusteredDistribution(Seq(a)))
+ withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ assert(spec.canCreatePartitioning && !spec.canEvaluatePartitioning)
+ val e = intercept[AssertionError](spec.createPartitioning(Seq($"b".int)))
+ assert(e.getMessage.contains("Cannot shuffle a child onto a layout Spark
cannot evaluate"))
+ }
+ }
+
+ test("SPARK-60027: a layout over a nested column places the rows of a more
nullable key") {
+ // A join compares two keys that differ only in nullability without a
cast, so the side
+ // shuffled onto a layout can have a more nullable key than the column the
layout's transform
+ // is over. The partitioning built over that key is evaluated on each of
its rows, so it has to
+ // take the key as the function did before the input type check.
+ def nestedBucket(declared: DataType)(value: InternalRow => Int):
ScalarFunction[Int] =
+ new ScalarFunction[Int] {
+ override def inputTypes(): Array[DataType] = Array(IntegerType,
declared)
+ override def resultType(): DataType = IntegerType
+ override def name(): String = "bucket"
+ override def canonicalName(): String = "test.bucket"
+ override def produceResult(input: InternalRow): Int =
+ Math.floorMod(value(input), input.getInt(0))
+ }
+ val notNullStruct = new StructType().add("a", IntegerType, nullable =
false)
+ Seq(
+ (notNullStruct, new StructType().add("a", IntegerType, nullable = true),
+ nestedBucket(notNullStruct)(_.getStruct(1, 1).getInt(0)),
InternalRow(5)),
+ (ArrayType(IntegerType, containsNull = false), ArrayType(IntegerType,
containsNull = true),
+ nestedBucket(ArrayType(IntegerType, containsNull =
false))(_.getArray(1).getInt(0)),
+ new GenericArrayData(Array[Any](5))),
+ (MapType(IntegerType, IntegerType, valueContainsNull = false),
+ MapType(IntegerType, IntegerType, valueContainsNull = true),
+ nestedBucket(MapType(IntegerType, IntegerType, valueContainsNull =
false))(
+ _.getMap(1).valueArray().getInt(0)),
+ ArrayBasedMapData(Map(1 -> 5)))
+ ).foreach { case (columnType, keyType, fn, value) =>
+ val k = AttributeReference("k", columnType)()
+ val r = AttributeReference("r", keyType)()
+ val spec = KeyedShuffleSpec(KeyedPartitioning(
+ Seq(TransformExpression(fn, Seq(k), Some(4))), Seq(InternalRow(1))),
+ ClusteredDistribution(Seq(k)))
+ withSQLConf(SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key -> "true") {
+ assert(spec.canCreatePartitioning && spec.canEvaluatePartitioning,
s"$columnType")
+ val built =
spec.createPartitioning(Seq(r)).asInstanceOf[KeyedPartitioning]
+ val bound = BindReferences.bindReference(built.expressions.head,
Seq(r))
+ assert(bound.eval(InternalRow(value)) == 1, s"$columnType placing a
$keyType key")
+ }
+ }
+ }
+
+ test("SPARK-60027: the identity reduce asks of the transform over the
identity column") {
Review Comment:
**[7] Minor: this test cannot tell which transform the gate asks.**
The gate asks the reported transform (`canEvaluate(t)` at
partitioning.scala:2405), and both rows expect `true` (L936), so `== reduced`
and `if (reduced)` are constant. Changing partitioning.scala:2405 and 2407 to
`canEvaluate(rebuiltOver(t, a))`, dropping `&& canEvaluate(t)`, or building
`IdentityReducer(t)` without the rebuild at partitioning.scala:2562 all keep it
green, and it passes on the merge-base too. What it pins is that a struct whose
field nullability differs is still reduced, and the symmetric `(t:
TransformExpression, a: AttributeReference)` arm never runs. Could you name it
after what it checks, drop the constant column, and add a row where the two
answers differ, e.g. an identity `a: LongType` against a transform over an
`IntegerType` column that declares `IntegerType`?
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -167,30 +175,153 @@ case class TransformExpression(
copy(children = newChildren)
/**
- * The scalar function call this transform stands for, with the bucket count
prepended as a
- * literal argument. `None` when the bound function is not a
[[ScalarFunction]], and also when a
- * join reduced this expression's keys, since then the call no longer
computes them. Evaluating it
- * to place a row would send the row to a partition it does not belong to.
That second arm is a
- * local gate. No caller reaches it today, because every consumer of a
reduced partitioning
- * refuses it first, and the write path never sees one.
+ * The bound function, when Spark could call it for this transform: a
[[ScalarFunction]], over
+ * keys no join reduced. See [[resolvedFunction]].
*/
- lazy val resolvedFunction: Option[Expression] = function match {
- case scalarFunc: ScalarFunction[_] if reducedWith.isEmpty =>
- val arguments = numBucketsOpt.fold(children)(n => Literal(n) +: children)
- Some(V2ExpressionUtils.resolveScalarFunction(scalarFunc, arguments))
+ private def scalarFunction: Option[ScalarFunction[_]] = function match {
+ case scalarFunc: ScalarFunction[_] if reducedWith.isEmpty =>
Some(scalarFunc)
case _ => None
}
- override def eval(input: InternalRow): Any = resolvedFunction match {
+ /** The arguments of the call: the bucket count, if any, as a literal, then
the children. */
+ private def callArguments: Seq[Expression] =
+ numBucketsOpt.fold(children)(n => Literal(n) +: children)
+
+ /**
+ * The scalar function call this transform stands for, over its arguments as
they are. `None`
+ * when the bound function is not a [[ScalarFunction]], and also when a join
reduced this
+ * expression's keys, since then the call no longer computes them.
Evaluating it to place a row
+ * would send the row to a partition it does not belong to. That second arm
is a local gate. No
+ * caller reaches it today, because every consumer of a reduced partitioning
refuses it first,
+ * and the write path never sees one. The write path
(`DistributionAndOrderingUtils`) casts the
+ * arguments to the function's input types itself. `eval` uses
[[evaluableFunction]] instead.
+ */
+ lazy val resolvedFunction: Option[Expression] =
+ scalarFunction.map(V2ExpressionUtils.resolveScalarFunction(_,
callArguments))
+
+ /**
+ * The call `eval` runs: [[resolvedFunction]], with each argument at the
type the function
+ * declares for it in `inputTypes()`. Nothing casts the arguments after
analysis, so the call
+ * gets the values as the argument holds them. Most arguments go as they
are, when they are held
+ * the way the declared type is, which is what a function read before this
check. One kind is
+ * cast: an integral argument narrower than declared is widened with a
[[Cast]], which keeps its
+ * value and needs no time zone. A function bound to a `SMALLINT` column but
declaring `INT`
+ * stays evaluable that way. Any other mismatch, or a different number of
arguments, gives
+ * `None`, as does everything `resolvedFunction` is `None` for. See
`atDeclaredType`. A mismatch
+ * is logged at info level once per function class and argument types, since
it makes a
+ * storage-partitioned join shuffle where it would not otherwise.
+ *
+ * Resolving the call can throw, e.g. when Spark finds no method to call on
a `ScalarFunction`
+ * that inherits `produceResult` from a base class, or when connector code
it calls throws. That
+ * gives `None` too, and `eval` rethrows it. So does a `LinkageError`, e.g.
a connector class
+ * whose dependency is missing, which the planner would otherwise hit for a
layout it does not
+ * pick. The planner asks this of such layouts on every query, so a failure
is logged as a
+ * warning once per function class and kind of failure, and with its stack
trace at debug level.
+ * Not serialized, so an executor resolves it again rather than receiving
the call (or the
+ * failure) the driver resolved while planning.
+ */
+ def evaluableFunction: Option[Expression] =
evaluableFunctionAttempt.toOption.flatten
+
+ /**
+ * The call, or why it failed to resolve, computed and logged once, since a
`lazy val` whose
+ * initializer throws runs it again on each access.
+ */
+ @transient private lazy val evaluableFunctionAttempt:
Try[Option[Expression]] = {
+ val attempt = try Success(resolveEvaluableFunction()) catch {
+ case NonFatal(e) => Failure(e)
+ case e: LinkageError => Failure(e)
+ }
+ attempt.failed.foreach(TransformExpression.logCannotResolve(function, _))
+ attempt
+ }
+
+ private def resolveEvaluableFunction(): Option[Expression] =
+ scalarFunction.flatMap { scalarFunc =>
+ val arguments = callArguments
+ val declaredTypes = scalarFunc.inputTypes().toSeq
+ val typed = if (arguments.length == declaredTypes.length) {
+ arguments.zip(declaredTypes).map { case (argument, declaredType) =>
+ TransformExpression.atDeclaredType(argument, declaredType)
+ }
+ } else {
+ Seq(None)
Review Comment:
**[10] Nit: `Seq(None)` is only a sentinel that makes
`typed.forall(_.isDefined)` fail on an arity mismatch, which splits one
decision across two `if`s.**
Mapping `arguments.zip(declaredTypes)` unconditionally and testing
`arguments.length == declaredTypes.length && typed.forall(_.isDefined)` behaves
the same, with one `logArgumentsMismatch` call and the same `None`. The length
check is still needed, e.g. for the three call arguments at
TransformExpressionSuite.scala:150. Could you test the length directly?
--
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]
