sunchao commented on code in PR #55885:
URL: https://github.com/apache/spark/pull/55885#discussion_r3503138262
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala:
##########
@@ -92,24 +126,145 @@ case class TransformExpression(
*/
def reducers(other: TransformExpression): Option[Reducer[_, _]] = {
(function, other.function) match {
- case(e1: ReducibleFunction[_, _], e2: ReducibleFunction[_, _]) =>
- reducer(e1, numBucketsOpt, e2, other.numBucketsOpt)
+ case (e1: ReducibleFunction[_, _], e2: ReducibleFunction[_, _]) =>
+ reducer(e1, this, e2, other)
case _ => None
}
}
- // Return a Reducer for a reducible function on another reducible function
+ /**
+ * 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 every literal parameter already matches the bound function's
declared input type at its
+ * position. The SPJ reducer paths hand literal values to the connector
reducer, or evaluate the
+ * transform directly, without Analyzer type coercion. A literal whose type
differs from the
+ * declared input type (a legal implicit cast under [[BoundFunction]]) is
therefore treated as not
+ * reducible: the join falls back to a shuffle rather than coercing a value
the partitions were
+ * not built on, or crashing when the connector casts it to the declared
type.
+ */
+ lazy val literalParamsMatchInputTypes: Boolean = {
+ val declaredTypes = function.inputTypes()
+ children.zipWithIndex.forall {
+ // Reject only a genuine type mismatch at a declared position. A literal
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 / the other guards rather than rejected
here.
+ case (l: Literal, i) => i >= declaredTypes.length || l.dataType ==
declaredTypes(i)
+ case _ => true
+ }
+ }
+
+ /**
+ * Reducer precondition: positionally-aligned argument structure with
`other` -- at each zipped
+ * position a literal aligns with a literal, nested transforms are
recursively the same function,
+ * and any other slot is a column reference on both sides. Only literal
*values* may differ. Arity
+ * is NOT required to match: children are zipped (a shorter side truncates),
so a zero-vs-one
+ * parameter pair is admitted and left to the connector reducer. Unlike
[[isSameFunction]] the
+ * function name is not compared.
+ */
+ private def sameArgumentLayout(other: TransformExpression): Boolean =
+ childrenMatch(other)((_, _) => true)
+
+ /**
+ * Whether no literal parameter has a complex type. A literal is rejected if
its [[DataType]] is
+ * [[ArrayType]] / [[MapType]] / [[StructType]] / [[UserDefinedType]]. Such
params (whose value is
+ * a Catalyst-internal container, or -- for a UDT -- whatever its `sqlType`
serializes to) must
+ * not cross the public reducer boundary, so the transform is treated as not
reducible. Keying off
+ * the type (not the value) also rejects a null-valued complex literal, and
rejecting all UDTs is
+ * a safe over-approximation (a UDT transform parameter is exotic; the cost
is a shuffle). Scalar
+ * types such as `CalendarIntervalType` are admitted (the connector
interprets them via the type).
+ */
+ private def noComplexLiteralParams: Boolean =
+ literalChildren.forall(_.dataType match {
+ case _: ArrayType | _: MapType | _: StructType | _: UserDefinedType[_]
=> false
+ case _ => true
+ })
+
+ /**
+ * Return a Reducer for a reducible function on another reducible function
+ * Handles both parameterized (bucket, truncate) and non-parameterized
(days, hours) functions.
+ */
private def reducer(
thisFunction: ReducibleFunction[_, _],
- thisNumBucketsOpt: Option[Int],
+ thisExpr: TransformExpression,
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)
+ otherExpr: TransformExpression): Option[Reducer[_, _]] = {
+ if (!thisExpr.sameArgumentLayout(otherExpr) ||
+ !thisExpr.literalParamsMatchInputTypes ||
!otherExpr.literalParamsMatchInputTypes ||
+ !thisExpr.noComplexLiteralParams || !otherExpr.noComplexLiteralParams)
{
+ return None
}
- Option(res)
+
+ val thisParams = thisExpr.extractParameters
+ val otherParams = otherExpr.extractParameters
+ val thisName = thisExpr.function.canonicalName()
+
+ // Gate on DataType, not the boxed runtime class
(DateType/YearMonthInterval box to Int).
+ def isSingleInt(p: Array[V2Literal[_]]): Boolean = {
+ p.length == 1 && p(0).dataType == IntegerType
Review Comment:
[P2] Please exclude typed-null integer literals from the deprecated
primitive-int probe. The public `Literal` contract does not prohibit `(value =
null, dataType = IntegerType)`, and `isSingleInt` accepts it; Scala's
`null.asInstanceOf[Int]` then becomes `0`. A legacy reducer can accept that
fabricated zero and return a reducer, so the generalized overload never
receives the actual typed null. If the transform distinguishes null from zero,
Spark can falsely declare the pair compatible and skip a required shuffle.
Require both literal values to be non-null before trying the deprecated
overload; otherwise invoke the generalized overload directly. Please add a
typed-null-vs-int regression where the legacy overload succeeds but the
generalized overload rejects the actual null.
##########
sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/functions/ReducibleFunction.java:
##########
@@ -60,6 +61,48 @@
@Evolving
public interface ReducibleFunction<I, O> {
+ /**
+ * Generic reducer for parameterized functions (bucket, truncate, etc.).
+ *
+ * If this function is 'reducible' on another function, return the {@link
Reducer}.
+ * <p>
+ * Each parameter is a non-complex {@link Literal} carrying both its value
and data type:
+ * array/map/struct/UDT-typed values are filtered out by Spark and not
passed here, but other
+ * scalar values (e.g. bucket numBuckets, truncate width, or a
+ * {@code CalendarInterval}) may be. {@link Literal#value()} is Spark's
internal representation
+ * (e.g. {@code UTF8String} for strings, {@code Decimal} for decimals); use
+ * {@link Literal#dataType()} to interpret it rather than assuming a JVM
type.
+ * <p>
+ * {@code thisParams} and {@code otherParams} hold each side's own literal
parameters and may have
+ * different lengths -- for example a zero-parameter transform reducing onto
a one-parameter one.
+ * Implementations must check each array's length before indexing into it.
+ * <p>
+ * Returning {@code null} means "not reducible for these parameters".
Throwing
+ * {@link UnsupportedOperationException} signals that this overload is not
implemented (Spark then
Review Comment:
[P2] Please correct this public fallback contract. For a single integer
parameter on each side, Spark probes the deprecated `reducer(int, ..., int)`
overload first and invokes this generalized overload only if that probe did not
return a reducer; an `UnsupportedOperationException` from this method never
triggers a later fallback to the deprecated method. For other parameterized
shapes there is no deprecated-int probe at all. This exception is not always
silent either: when every eligible overload throws
`UnsupportedOperationException`, Spark logs the `implements no reducer`
warning. Since connector authors will rely on this Javadoc to understand
dispatch, please document the actual deprecated-first order and terminal
behavior.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala:
##########
@@ -641,19 +640,15 @@ object KeyedPartitioning {
def supportsExpressions(expressions: Seq[Expression]): Boolean = {
def isSupportedTransform(transform: TransformExpression): Boolean = {
- transform.children.size == 1 && isReference(transform.children.head)
- }
-
- @tailrec
- def isReference(e: Expression): Boolean = e match {
- case _: Attribute => true
- case g: GetStructField => isReference(g.child)
- case _ => false
+ // Should only consider column references, not literals.
+ val nonLiteralChildren =
transform.children.filterNot(_.isInstanceOf[Literal])
+ // We need exactly one column reference per transform.
+ nonLiteralChildren.size == 1 &&
TransformExpression.isColumnRef(nonLiteralChildren.head)
Review Comment:
[P1] Thanks—rejecting the mismatched literal fixes that reproducer, but this
still checks only `Literal` children. `BoundFunction.inputTypes()` applies to
every argument, and `identityReducer` directly evaluates every raw child
through the same unanalyzed path. For example, a connector can legally bind
`(ShortType, IntegerType)` and return a scalar function whose declared inputs
are `(IntegerType, IntegerType)`. Because the literal is already an exact
`IntegerType`, `literalParamsMatchInputTypes` passes; then
`identityReducer.reduce` feeds the boxed `Short` through
`ApplyFunctionExpression` into `SpecificInternalRow(IntegerType)`, which throws
`ClassCastException`. Please validate every child against its corresponding
declared input type before this direct evaluation (or insert the required
casts), and add an identity-vs-transform regression with a `ShortType` column
and exact-typed integer literal. The current tests cover only the literal slot.
--
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]
