metanil commented on code in PR #55885:
URL: https://github.com/apache/spark/pull/55885#discussion_r4139433665
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -133,18 +182,178 @@ case class TransformExpression(
def withReference(attr: Attribute): TransformExpression =
transform { case _: AttributeReference => attr
}.asInstanceOf[TransformExpression]
- // 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)
+ /**
+ * Rewrites this transform's column arguments with `rewriteColumn`, leaving
its literal parameters
+ * untouched. Unlike [[withReference]], which swaps the attribute *inside* a
column argument (so a
+ * `GetStructField` path above it is kept), this replaces the whole argument.
+ *
+ * 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. Every such path should go through here rather than reimplement the
skip.
+ *
+ * Callers: `KeyedShuffleSpec.createPartitioning`, which retargets a
transform at the other side's
+ * clustering key, and `PartitioningPreservingUnaryExecNode`, which
retargets it at an aliased
+ * output attribute.
+ */
+ 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. [[withReference]]
retargets a transform at
+ * another key, and a [[GetStructField]] column slot retargeted at a
non-struct attribute -- an
+ * `identity` column on the other side of a join, say -- leaves
`GetStructField(intCol, 0)`, whose
+ * `dataType` raises `INTERNAL_ERROR: GetStructField requires a StructType
child`. 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))
+ }
+ }
+
+ /**
+ * 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 or a column reference
on both sides.
+ * Literal values and arity may differ; the connector reducer decides those.
+ */
+ private def sameArgumentLayout(other: TransformExpression): Boolean =
+ children.zip(other.children).forall {
+ case (_: Literal, _: Literal) => true
+ case (c1, c2) => TransformExpression.isColumnRef(c1) &&
TransformExpression.isColumnRef(c2)
+ }
+
+ /**
+ * Whether no literal parameter has an array, map, struct or UDT type; those
are not passed to a
+ * connector reducer. Other scalars, such as intervals, are.
+ */
+ private def noComplexLiteralParams: Boolean =
+ literalChildren.forall(_.dataType match {
+ case _: ArrayType | _: MapType | _: StructType | _: UserDefinedType[_]
=> false
+ case _ => true
+ })
+
+ /**
+ * Return a [[Reducer]] that maps this transform's partition keys onto
`other`'s, or None if there
+ * is none: when either function is not a [[ReducibleFunction]], when the
argument layouts or
+ * literal parameters rule a reducer out, or when the connector reports it
is not reducible.
+ * Handles both parameterized (bucket, truncate) and non-parameterized
(days, hours) functions.
+ */
+ def reducers(other: TransformExpression): Option[Reducer[_, _]] = {
+ import TransformExpression._
+ val thisFunction: ReducibleFunction[_, _] = function match {
+ case f: ReducibleFunction[_, _] => f
+ case _ => return None
+ }
+ val otherFunction: ReducibleFunction[_, _] = other.function match {
+ case o: ReducibleFunction[_, _] => o
+ case _ => return None
+ }
+ if (!sameArgumentLayout(other) ||
+ !literalParamsMatchInputTypes || !other.literalParamsMatchInputTypes ||
+ !noComplexLiteralParams || !other.noComplexLiteralParams) {
+ return None
+ }
+
+ val thisParams = extractParameters
+ val otherParams = other.extractParameters
+ val thisName = function.canonicalName()
+
+ // Only a single non-null IntegerType parameter per side may use the
deprecated int overload; a
+ // typed null would otherwise be read as 0.
+ def isSingleInt(p: Array[V2Literal[_]]): Boolean = {
+ p.length == 1 && p(0).dataType == IntegerType && p(0).value() != null
+ }
+
+ // 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
Review Comment:
Yes, this was the known issue, rare but not impossible.
I didn't want to use reflection for it, but as you pointed out, there's
already a history of using it, I might rethink this.
--
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]
