metanil commented on code in PR #55885:
URL: https://github.com/apache/spark/pull/55885#discussion_r4418405869
##########
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:
Thanks @peter-toth
I thought about this layout based check, but it's still would be a guess
about which function is a bucket. Example, `xyz(2, data)` vs `xyz(4, data)`
gets through it the same way bucket does.
I was also initially exploring using `name()`, but turns out it'll have
different issue, since name() has no contract (your finding 5).
So only way to completely resolve this AFAIK (which I don't prefer) is
adding a flag (say `fromBucketTransform`) on `TransformExpression` which will
get set in `V2ExpressionUtils` when the V2 transform is `bucket(n, col)`. The
flag isn't part of `functionId` and doesn't show up in plan strings. I don't
prefer this as this brings back same style as that of the `numBucketsOpt` as a
marker.
Happy to go with your layout check instead if you think the flag isn't
worth it.
--
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]
