peter-toth commented on code in PR #59760:
URL: https://github.com/apache/spark/pull/59760#discussion_r4457220509
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -2570,24 +2606,38 @@ 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 {
+ case t: TransformExpression => canEvaluate(t)
+ case _ => true
+ }
+
override def createPartitioning(clustering: Seq[Expression]): Partitioning =
{
assert(clustering.size == distribution.clustering.size,
"Required distributions of join legs should be the same size.")
-
val newExpressions = partitioning.expressions.zip(keyPositions).map {
case (te: TransformExpression, positionSet) =>
// `canCreatePartitioning` admits only a template over bare columns,
see there.
rebuiltOver(te, clustering(positionSet.head))
case (_, positionSet) => clustering(positionSet.head)
}
+ // The child shuffled onto this layout is placed by evaluating what is
built here, which would
+ // fail on an executor. A caller has to ask `canEvaluatePartitioning`
first, which answers for
+ // this layout's own keys. A join's keys share their types, so that
answers for what is built
+ // too. A cogroup's need not (SPARK-60782), and falling back to a hash
shuffle would pair its
+ // groups wrongly, so such a child fails here, while planning.
+ assert(newExpressions.forall {
Review Comment:
**Finding 13.** This is the assert finding 5 asked for, and it now fails
while planning, as intended. But it fires on user input: a cogroup key whose
type the layout's function doesn't take. `QueryExecution.toInternalError` turns
an `AssertionError` into an `INTERNAL_ERROR` that says "You hit a bug in Spark
... Please, report this bug", for a failure the description lists as expected.
A user-facing error that names the key's type and the function's declared types
would fit better, until SPARK-60782 rejects such keys up front.
It also fires for a key of another collation, which master did not reject. I
planned a `FlatMapCoGroupsInPandasExec` in `EnsureRequirementsSuite` under
`spark.sql.sources.v2.bucketing.shuffle.enabled`. The left child is on a
`bucket(4)` over a `STRING`, with a function declaring `(INT, STRING)`, and the
right key is a `STRING COLLATE UTF8_LCASE`:
- master (`fce7d51`) plans the shuffle onto `bucket(4)` over the right key;
- this head fails here, since a string of another collation is refused.
Failing is the better outcome, since master pairs the groups wrongly there.
The cogroup sorts the right side by `UTF8_LCASE` and `CoGroupedIterator` then
merges the two sides with the left key's binary ordering. I fed
`GroupedIterator` and `CoGroupedIterator` the keys `B` and `a` on both sides,
each side sorted as the cogroup sorts it. `B` came out twice, once with only
its left rows and once with only its right rows. Lowercase keys pair correctly.
So the user-facing line "still fails, as on master" does not hold for it. That
case returned wrong groups on master and fails now. It is worth listing under
the user-facing changes.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -167,30 +175,164 @@ 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 {
- case Some(fn) => fn.eval(input)
- case None => throw QueryExecutionErrors.cannotEvaluateExpressionError(this)
+ /** 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()`, see `atDeclaredType` for which
arguments suit. `None` when
+ * an argument does not, when there is another number of arguments, when the
call fails to
+ * resolve, and wherever `resolvedFunction` is `None`. `eval` rethrows a
failure to resolve. A
+ * storage-partitioned join shuffles instead of evaluating such a transform,
so each case but
+ * reduced keys is logged once, see `logNotScalar`, `logArgumentsMismatch`
and
+ * `logCannotResolve`.
+ */
+ def evaluableFunction: Option[Expression] = evaluableFunctionAttempt match {
+ case Success(call) => call
+ case _ => None
+ }
+
+ /**
+ * The call, or why it failed to resolve, e.g. a `ScalarFunction` that
inherits `produceResult`
+ * from a base class, connector code that throws, or a `LinkageError` for a
connector class whose
+ * dependency is missing, which the planner would otherwise hit for a layout
it does not pick.
+ * Computed and logged once, since a `lazy val` whose initializer throws
runs it again on each
+ * access. Not serialized, so an executor resolves it again rather than
receiving the call (or
+ * the failure) the driver resolved while planning.
+ */
+ @transient private lazy val evaluableFunctionAttempt:
Try[Option[Expression]] =
+ try Success(resolveEvaluableFunction()) catch {
+ case NonFatal(e) => cannotResolve(e)
+ case e: LinkageError => cannotResolve(e)
+ }
+
+ private def cannotResolve(e: Throwable): Try[Option[Expression]] = {
+ TransformExpression.logCannotResolve(function, e)
+ Failure(e)
+ }
+
+ private def resolveEvaluableFunction(): Option[Expression] = scalarFunction
match {
+ case None =>
+ // Reduced keys are the planner's own doing, so only a function Spark
cannot call is logged.
+ if (!function.isInstanceOf[ScalarFunction[_]])
TransformExpression.logNotScalar(function)
+ None
+ case Some(scalarFunc) =>
+ val arguments = callArguments
+ val declaredTypes = scalarFunc.inputTypes().toSeq
+ val typed = arguments.zip(declaredTypes).map { case (argument,
declaredType) =>
+ TransformExpression.atDeclaredType(argument, declaredType)
+ }
+ if (arguments.length == declaredTypes.length &&
typed.forall(_.isDefined)) {
+ Some(V2ExpressionUtils.resolveScalarFunction(scalarFunc,
typed.flatten))
Review Comment:
**Finding 14.** A user-defined type goes as its `sqlType` here, and
`resolveScalarFunction` finds a magic `invoke` from the declared types. But
`Invoke` and `StaticInvoke` look the method up again from the arguments'
classes (`EncoderUtils.expressionJavaClasses`, `objects.scala:346` and `:460`),
and `dataTypeJavaClass` maps a UDT to `Object`. Measured with a UDT whose
`sqlType` is `INT`, and a `bucket` declaring `(INT, INT)`:
- with a magic `invoke(int, int)`, `evaluableFunction` is defined, so both
gates admit the transform, and `eval` throws `[INTERNAL_ERROR] Couldn't find
method invoke with arguments (int, class java.lang.Object)`;
- with `produceResult` instead, it evaluates, as the UDT row in
`TransformExpressionSuite` does.
Master fails the same way, but this check is there to refuse such a call.
Looking the method up here makes the failure land in
`evaluableFunctionAttempt`, which logs it and refuses the transform. With the
change below the magic case is refused, the `produceResult` case still
evaluates, and `TransformExpressionSuite` and `ShuffleSpecSuite` pass:
```suggestion
val call = V2ExpressionUtils.resolveScalarFunction(scalarFunc,
typed.flatten)
// `Invoke` and `StaticInvoke` look the method up again from the
arguments' classes, a
// UDT's being `Object`, so look it up here, where a failure is
caught and logged.
call match {
case i: objects.Invoke => i.method
case s: objects.StaticInvoke => s.method
case _ =>
}
Some(call)
```
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -2415,6 +2430,26 @@ case class KeyedShuffleSpec(
private def argumentsAreColumns(transform: TransformExpression): Boolean =
transform.children.forall(_.isInstanceOf[Attribute])
+ /**
+ * Whether Spark can evaluate `transform`, see
`TransformExpression.evaluableFunction`. An
+ * identity reduce evaluates it on the identity side's partition keys, and a
one-side shuffle on
+ * each row of the side it shuffles, both after rebuilding it over the other
side's key. An
+ * equi-join pairs two keys only when their types are equal apart from
nullability and field
+ * names (`BinaryComparison.sameType`), which `evaluableFunction` does not
tell apart either. So
+ * for a join, asking this of the reported transform answers it for the
rebuilt one, and the
+ * reported one keeps the answer. A cogroup does not coerce its keys, so a
key of another type,
+ * e.g. a `BIGINT` against a transform over an `INT` column that declares
`INT`, can still fail
+ * where its rows are placed, as it did before this check. A decimal key of
the column's scale
+ * passes, since its precision does not change what the function reads.
Review Comment:
**Finding 15.** Since `3071d6d`, `createPartitioning` asserts the
partitioning it builds over the shuffled child's key. So a cogroup child
shuffled from a key of another type now fails there, while planning, and the
new `EnsureRequirementsSuite` test pins that. This sentence still says it fails
where its rows are placed, as before this check.
The description has the same sentence in the paragraph that starts "Asking
this of the reported transform", while its user-facing list says "while
planning rather than in a task".
```suggestion
* e.g. a `BIGINT` against a transform over an `INT` column that declares
`INT`, can still fail
* as it did before this check. A one-side shuffle fails while planning,
in `createPartitioning`.
* A decimal key of the column's scale passes, since its precision does
not change what the
* function reads.
```
--
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]
