diff --git a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/functions/ReducibleFunction.java b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/functions/ReducibleFunction.java index ef1a14e50cdad..1ac7b8c3776e3 100644 --- a/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/functions/ReducibleFunction.java +++ b/sql/catalyst/src/main/java/org/apache/spark/sql/connector/catalog/functions/ReducibleFunction.java @@ -17,6 +17,7 @@ package org.apache.spark.sql.connector.catalog.functions; import org.apache.spark.annotation.Evolving; +import org.apache.spark.sql.connector.expressions.Literal; /** * Base class for user-defined functions that can be 'reduced' on another function. @@ -60,6 +61,52 @@ @Evolving public interface ReducibleFunction { + /** + * Generic reducer for parameterized functions (bucket, truncate, etc.). + * + * If this function is 'reducible' on another function, return the {@link Reducer}. + *

+ * 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. + *

+ * {@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. + *

+ * Returning {@code null} means "not reducible for these parameters" and is authoritative: + * Spark consults no other overload. Dispatch order: Spark tries this generalized overload + * first; only if it is not implemented (throws {@link UnsupportedOperationException}) and each + * side has a single non-null integer parameter does Spark fall back to the deprecated + * {@code reducer(int, ReducibleFunction, int)} overload. If every eligible overload throws + * {@link UnsupportedOperationException}, Spark logs an "implements no reducer" warning; any + * other exception is logged and the pair is treated as not reducible (the join falls back to a + * shuffle). + *

+ * Examples: + *

+ * + * @param thisParams literal parameters for this function (may differ in length from otherParams) + * @param otherFunction the other parameterized function + * @param otherParams literal parameters for the other function (may differ in length from + * thisParams) + * @return a reduction function if reducible, null otherwise + * @since 4.3.0 + */ + default Reducer reducer( + Literal[] thisParams, + ReducibleFunction otherFunction, + Literal[] otherParams) { + throw new UnsupportedOperationException(); + } + /** * This method is for the bucket function. * @@ -78,7 +125,12 @@ public interface ReducibleFunction { * @param otherBucketFunction the other parameterized function * @param otherNumBuckets parameter for the other function * @return a reduction function if it is reducible, null if not + * @deprecated as of 4.3.0. Please override + * {@link #reducer(Literal[], ReducibleFunction, Literal[])} instead. + * The new overload supports transforms with any number of parameters of any type + * (e.g. truncate width, multi-arg range buckets), not just a single int. */ + @Deprecated(since = "4.3.0") default Reducer reducer( int thisNumBuckets, ReducibleFunction otherBucketFunction, @@ -101,6 +153,6 @@ default Reducer reducer( * @return a reduction function if it is reducible, null if not. */ default Reducer reducer(ReducibleFunction otherFunction) { - throw new UnsupportedOperationException(); + return reducer(new Literal[0], otherFunction, new Literal[0]); } } diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala index 9041ed15fc501..b20fe5a18b841 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/TransformExpression.scala @@ -17,45 +17,79 @@ package org.apache.spark.sql.catalyst.expressions +import scala.annotation.tailrec +import scala.util.{Failure, Success, Try} + +import org.apache.spark.internal.Logging +import org.apache.spark.internal.LogKeys.FUNCTION_NAME import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode} import org.apache.spark.sql.connector.catalog.functions.{BoundFunction, Reducer, ReducibleFunction, ScalarFunction} +import org.apache.spark.sql.connector.expressions.{Literal => V2Literal, LiteralValue} import org.apache.spark.sql.errors.QueryExecutionErrors -import org.apache.spark.sql.types.DataType +import org.apache.spark.sql.types.{ArrayType, DataType, IntegerType, MapType, StructType, UserDefinedType} /** * Represents a partition transform expression, for instance, `bucket`, `days`, `years`, etc. * * @param function the transform function itself. Spark will use it to decide whether two * partition transform expressions are compatible. - * @param numBucketsOpt the number of buckets if the transform is `bucket`. Unset otherwise. */ -case class TransformExpression( - function: BoundFunction, - children: Seq[Expression], - numBucketsOpt: Option[Int] = None) extends Expression { +case class TransformExpression(function: BoundFunction, children: Seq[Expression]) + extends Expression with Logging { override def nullable: Boolean = true /** - * Whether this [[TransformExpression]] has the same semantics as `other`. - * For instance, `bucket(32, c)` is equal to `bucket(32, d)`, but not to `bucket(16, d)` or - * `year(c)`. + * Extract literal children (constant parameters) from this transform. These are constant + * arguments like width in truncate(col, width). Literals are compared when checking if two + * transforms are the same. + */ + private lazy val literalChildren: Seq[Literal] = + children.collect { case l: Literal => l } + + /** + * Whether this [[TransformExpression]] has the same semantics as `other`. For instance, + * `bucket(32, c)` is equal to `bucket(32, d)`, but not to `bucket(16, d)` or `year(c)`. + * Similarly, `truncate(c, 2)` is equal to `truncate(d, 2)`, but may not to `truncate(c, 4)`. * * This will be used, for instance, by Spark to determine whether storage-partitioned join can * be triggered, by comparing partition transforms from both sides of the join and checking * whether they are compatible. * - * @param other the transform expression to compare to - * @return true if this and `other` has the same semantics w.r.t to transform, false otherwise. + * Two transforms are considered the same when they have the same function name, the same arity, + * and each pair of corresponding children matches: + * - literal arguments must be equal (e.g. numBuckets for bucket, width for truncate), so that + * `bucket(32, c)` is not the same as `bucket(16, c)`; + * - nested transform arguments must recursively be the same function, so that + * `bucket(4, years(c))` is not the same as `bucket(4, days(c))`; + * - everything else must be a plain column reference on both sides. Column identity is + * intentionally ignored (it is reconciled separately via positional matching), but a + * non-reference slot such as `c + 1` or `cast(c)`, or a literal/transform-vs-reference + * mismatch, is treated as not the same. + * + * @param other + * the transform expression to compare to + * @return + * true if this and `other` has the same semantics w.r.t to transform, false otherwise. */ - def isSameFunction(other: TransformExpression): Boolean = other match { - case TransformExpression(otherFunction, _, otherNumBucketsOpt) => - function.canonicalName() == otherFunction.canonicalName() && - numBucketsOpt == otherNumBucketsOpt - case _ => - false - } + def isSameFunction(other: TransformExpression): Boolean = + function.canonicalName() == other.function.canonicalName() && + children.length == other.children.length && + childrenMatch(other)(_ == _) + + /** + * Per-position match of the zipped children (callers enforce arity where needed). Literal slots + * are compared by the caller-supplied `literalsMatch`; nested transform slots must recursively be + * the same function; any other slot must be a plain column reference on both sides. + */ + private def childrenMatch(other: TransformExpression) + (literalsMatch: (Literal, Literal) => Boolean): Boolean = + children.zip(other.children).forall { + case (l1: Literal, l2: Literal) => literalsMatch(l1, l2) + case (t1: TransformExpression, t2: TransformExpression) => t1.isSameFunction(t2) + case (c1, c2) => TransformExpression.isColumnRef(c1) && TransformExpression.isColumnRef(c2) + } /** * Whether this [[TransformExpression]]'s function is compatible with the `other` @@ -73,8 +107,8 @@ case class TransformExpression( } 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) + val thisReducer = reducer(f, this, o, other) + val otherReducer = reducer(o, other, f, this) thisReducer.isDefined || otherReducer.isDefined case _ => false } @@ -92,24 +126,166 @@ 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 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. + */ + private def inputTypesMatch(select: Expression => Boolean): Boolean = { + val declaredTypes = function.inputTypes() + children.zipWithIndex.forall { + case (c, i) => !select(c) || i >= declaredTypes.length || c.dataType == 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 -- columns AND literals -- matches its declared input type, at exactly + * the declared arity. Required before directly evaluating the transform (the + * identity-vs-transform reducer in `KeyedShuffleSpec`), which feeds every child through the + * function's `SpecificInternalRow(inputTypes())`: a child whose type differs would raise a + * `ClassCastException`, and a child *beyond* the declared arity would raise an + * `ArrayIndexOutOfBoundsException` (the row is sized to `inputTypes().length`). The exact-arity + * requirement is specific to this eval path -- the transform-vs-transform path passes literals to + * the connector's reducer (no eval) and deliberately allows mixed arity, so it uses + * [[literalParamsMatchInputTypes]], which keeps the beyond-arity short-circuit. The check is one + * level and reads each child's `dataType`: the non-literal child is a column reference -- an + * [[Attribute]] or [[GetStructField]] chain, never a nested transform (rejected by the scan gate + * `supportsExpressions`/`isColumnRef`) -- so there is no inner-transform column to recurse into, + * and a gate-admitted child always has a resolvable `dataType`, so the eager read is safe. + * Stronger than [[literalParamsMatchInputTypes]]. + */ + lazy val argsMatchInputTypes: Boolean = + children.length == function.inputTypes().length && inputTypesMatch(_ => 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[_, _]] = { + import TransformExpression._ + if (!thisExpr.sameArgumentLayout(otherExpr) || + !thisExpr.literalParamsMatchInputTypes || !otherExpr.literalParamsMatchInputTypes || + !thisExpr.noComplexLiteralParams || !otherExpr.noComplexLiteralParams) { + return None + } + + val thisParams = thisExpr.extractParameters + val otherParams = otherExpr.extractParameters + val thisName = thisExpr.function.canonicalName() + + // A single non-null IntegerType param on each side is the shape the deprecated + // reducer(int, ..., int) fallback accepts. Gate on the DataType, not the boxed runtime class + // (DateType / YearMonthInterval also box to Int). A typed null (Literal(null, IntegerType)) is + // excluded: null.asInstanceOf[Int] would fabricate a 0 a legacy reducer might accept, so a + // typed null must not reach the deprecated fallback (the generalized overload sees the real + // null). + 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 + case Failure(e) => Threw(e) + } + + // Prefer the generalized Literal[] overload; fall back to the deprecated int overload only for + // a single-int pair, and only when the generalized one is not implemented. + // Modern connectors never touch the deprecated path; deprecated-only connectors still reduce. + val outcome = + if (thisParams.isEmpty && otherParams.isEmpty) { + probe(thisFunction.reducer(otherFunction)) + } else { + probe(thisFunction.reducer(thisParams, otherFunction, otherParams)) match { + // Generalized overload not implemented: fall back to the deprecated int overload, but + // only for a single-int pair. Any other generalized outcome (reducible, deliberately not + // reducible, or a thrown bug) is authoritative. + case Unimplemented if isSingleInt(thisParams) && isSingleInt(otherParams) => + probe(thisFunction.reducer( + thisParams(0).value().asInstanceOf[Int], otherFunction, + otherParams(0).value().asInstanceOf[Int])) + case other => other + } + } + + outcome match { + case Reducible(r) => Some(r) + case NotReducible => None + case Threw(e) => + logWarning(log"V2 function ${MDC(FUNCTION_NAME, thisName)} reducer threw an exception; " + + log"treating as not reducible.", e) + None + case Unimplemented => + logWarning(log"V2 function ${MDC(FUNCTION_NAME, thisName)} implements no reducer; " + + log"treating as not reducible. Override " + + log"reducer(Literal[], ReducibleFunction, Literal[]) to enable SPJ.") + None } - Option(res) } override def dataType: DataType = function.resultType() @@ -118,10 +294,7 @@ case class TransformExpression( copy(children = newChildren) private lazy val resolvedFunction: Option[Expression] = this match { - case TransformExpression(scalarFunc: ScalarFunction[_], arguments, Some(numBuckets)) => - Some(V2ExpressionUtils.resolveScalarFunction(scalarFunc, - Seq(Literal(numBuckets)) ++ arguments)) - case TransformExpression(scalarFunc: ScalarFunction[_], arguments, None) => + case TransformExpression(scalarFunc: ScalarFunction[_], arguments) => Some(V2ExpressionUtils.resolveScalarFunction(scalarFunc, arguments)) case _ => None } @@ -136,3 +309,25 @@ case class TransformExpression( override protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = throw QueryExecutionErrors.cannotGenerateCodeForExpressionError(this) } + +object TransformExpression { + /** + * Whether `e` is a bare column reference: an [[Attribute]] or a [[GetStructField]] chain + * (struct-field access on a column). Shared by [[TransformExpression.isSameFunction]] and by + * `KeyedPartitioning.supportsExpressions`, which both decide whether a transform's single + * non-literal argument is a plain column. + */ + @tailrec + private[sql] def isColumnRef(e: Expression): Boolean = e match { + case _: Attribute => true + case g: GetStructField => isColumnRef(g.child) + case _ => false + } + + /** The result of probing one reducer overload, for the dispatch in [[TransformExpression]]. */ + private sealed trait Outcome + private case class Reducible(reducer: Reducer[_, _]) extends Outcome + private case object NotReducible extends Outcome + private case class Threw(e: Throwable) extends Outcome + private case object Unimplemented extends Outcome +} diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala index 702255e075743..dee76a8588e27 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/V2ExpressionUtils.scala @@ -30,7 +30,7 @@ import org.apache.spark.sql.catalyst.plans.logical.{LocalRelation, LogicalPlan, import org.apache.spark.sql.connector.catalog.{FunctionCatalog, Identifier} import org.apache.spark.sql.connector.catalog.functions._ import org.apache.spark.sql.connector.catalog.functions.ScalarFunction.MAGIC_METHOD_NAME -import org.apache.spark.sql.connector.expressions.{BucketTransform, Cast => V2Cast, Expression => V2Expression, FieldReference, GeneralScalarExpression, IdentityTransform, Literal => V2Literal, NamedReference, NamedTransform, NullOrdering => V2NullOrdering, SortDirection => V2SortDirection, SortOrder => V2SortOrder, SortValue, Transform} +import org.apache.spark.sql.connector.expressions.{Cast => V2Cast, Expression => V2Expression, FieldReference, GeneralScalarExpression, IdentityTransform, Literal => V2Literal, NamedReference, NamedTransform, NullOrdering => V2NullOrdering, SortDirection => V2SortDirection, SortOrder => V2SortOrder, SortValue, Transform} import org.apache.spark.sql.connector.expressions.filter.{AlwaysFalse, AlwaysTrue} import org.apache.spark.sql.connector.read.{SampleMethod => V2SampleMethod} import org.apache.spark.sql.errors.DataTypeErrors.toSQLId @@ -116,17 +116,6 @@ object V2ExpressionUtils extends SQLConfHelper with Logging { funCatalogOpt: Option[FunctionCatalog] = None): Option[Expression] = trans match { case IdentityTransform(ref) => Some(resolveRef[NamedExpression](ref, query)) - case BucketTransform(numBuckets, refs, sorted) - if sorted.isEmpty && refs.length == 1 && refs.forall(_.isInstanceOf[NamedReference]) => - val resolvedRefs = refs.map(r => resolveRef[NamedExpression](r, query)) - // Create a dummy reference for `numBuckets` here and use that, together with `refs`, to - // look up the V2 function. - val numBucketsRef = AttributeReference("numBuckets", IntegerType, nullable = false)() - funCatalogOpt.flatMap { catalog => - loadV2FunctionOpt(catalog, "bucket", Seq(numBucketsRef) ++ resolvedRefs).map { bound => - TransformExpression(bound, resolvedRefs, Some(numBuckets)) - } - } case NamedTransform(name, args) => val catalystArgs = args.map(toCatalyst(_, query, funCatalogOpt)) funCatalogOpt.flatMap { catalog => diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala index d2bb12d2053aa..111532f1cabd0 100644 --- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala +++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/physical/partitioning.scala @@ -17,7 +17,6 @@ package org.apache.spark.sql.catalyst.plans.physical -import scala.annotation.tailrec import scala.collection.mutable import org.apache.spark.{SparkException, SparkUnsupportedOperationException} @@ -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) } expressions.forall { case t: TransformExpression if isSupportedTransform(t) => true - case e: Expression if isReference(e) => true + case e: Expression if TransformExpression.isColumnRef(e) => true case _ => false } } @@ -1302,24 +1297,58 @@ case class KeyedShuffleSpec( } } - private def isExpressionCompatible(left: Expression, right: Expression): Boolean = + /** + * The reducer mapping a raw identity column `col` onto transform `t` (it applies `t` to the + * identity values), or None if not reducible. Single source of the identity-vs-transform + * decision: [[isExpressionCompatible]] derives the gate from it (`.isDefined`) and [[reducers]] + * returns it, so the two cannot drift (a divergence would keep raw keys -> mis-join). + * + * The reducer evals the transform with `col` substituted for its column, so it validates the + * SUBSTITUTED expression's arg types ([[TransformExpression.argsMatchInputTypes]]) -- not t's own + * -- keeping a type mismatch (e.g. a `ShortType` `col` at an `IntegerType` slot) from reaching + * eval and raising a `ClassCastException`. + */ + private def identityReducer( + col: AttributeReference, t: TransformExpression): Option[Reducer[_, _]] = { + // `transform` preserves the root node type, so this is always a TransformExpression; the only + // real gate is argsMatchInputTypes on the substituted (identity) column -- see the doc above. + val reducerExpr = + t.transform { case _: AttributeReference => col }.asInstanceOf[TransformExpression] + if (reducerExpr.argsMatchInputTypes) { + val boundExpr = BindReferences.bindReference(reducerExpr, AttributeSeq(Seq(col))) + Some(new Reducer[Any, Any] { + override def reduce(v: Any): Any = boundExpr.eval(new GenericInternalRow(Array[Any](v))) + override def resultType(): DataType = reducerExpr.dataType + override def displayName(): String = reducerExpr.toString + }) + } else { + None + } + } + + private def isExpressionCompatible(left: Expression, right: Expression): Boolean = { + def compatibleTransformsAllowed: Boolean = + SQLConf.get.v2BucketingPushPartValuesEnabled && + !SQLConf.get.v2BucketingPartiallyClusteredDistributionEnabled && + SQLConf.get.v2BucketingAllowCompatibleTransforms (left, right) match { case (_: LeafExpression, _: LeafExpression) => true case (left: TransformExpression, right: TransformExpression) => - if (SQLConf.get.v2BucketingPushPartValuesEnabled && - !SQLConf.get.v2BucketingPartiallyClusteredDistributionEnabled && - SQLConf.get.v2BucketingAllowCompatibleTransforms) { + if (compatibleTransformsAllowed) { left.isCompatible(right) } else { left.isSameFunction(right) } - case (_: AttributeReference, _: TransformExpression) | - (_: TransformExpression, _: AttributeReference) => - SQLConf.get.v2BucketingPushPartValuesEnabled && - !SQLConf.get.v2BucketingPartiallyClusteredDistributionEnabled && - SQLConf.get.v2BucketingAllowCompatibleTransforms + // Identity transform on one side, arbitrary transform on the other. Derive the gate from the + // producer (identityReducer): the pair is compatible only if a reducer can actually be built, + // so the gate and reducers cannot drift (a divergence would keep raw keys -> mis-join). + case (col: AttributeReference, t: TransformExpression) => + compatibleTransformsAllowed && identityReducer(col, t).isDefined + case (t: TransformExpression, col: AttributeReference) => + compatibleTransformsAllowed && identityReducer(col, t).isDefined case _ => false } + } /** * Return a set of [[Reducer]] for the partition expressions of this shuffle spec, @@ -1341,19 +1370,11 @@ case class KeyedShuffleSpec( val results = partitioning.expressions.zip(other.partitioning.expressions).map { case (e1: TransformExpression, e2: TransformExpression) => e1.reducers(e2) - // Identity transform on this side, arbitrary transform on the other side: create a reducer - // that applies the other's transform to the raw identity values. The symmetric case + // Identity transform on this side, arbitrary transform on the other side. The symmetric case // (TransformExpression, AttributeReference) is handled when the other side calls reducers. - // Each partition expression is guaranteed to have exactly one leaf child (asserted in - // keyPositions), so `a` lives at position 0 in the row we construct. - case (a: AttributeReference, t: TransformExpression) => - val reducerExpr = t.transform { case _: AttributeReference => a } - val boundExpr = BindReferences.bindReference(reducerExpr, AttributeSeq(Seq(a))) - Some(new Reducer[Any, Any] { - override def reduce(v: Any): Any = boundExpr.eval(new GenericInternalRow(Array[Any](v))) - override def resultType(): DataType = reducerExpr.dataType - override def displayName(): String = reducerExpr.toString - }) + // identityReducer is the shared decision the compatibility gate also consults, so the two + // cannot drift. + case (col: AttributeReference, t: TransformExpression) => identityReducer(col, t) case (_, _) => None } @@ -1375,7 +1396,13 @@ case class KeyedShuffleSpec( val newExpressions = partitioning.expressions.zip(keyPositions).map { case (te: TransformExpression, positionSet) => - te.copy(children = te.children.map(_ => clustering(positionSet.head))) + // Preserve literal parameters (e.g., numBuckets, truncate width) + // while replacing only column references with the new clustering expression + val newChildren = te.children.map { + case l: Literal => l // Keep literals as-is + case _ => clustering(positionSet.head) // Replace column references + } + te.copy(children = newChildren) case (_, positionSet) => clustering(positionSet.head) } KeyedPartitioning(newExpressions, partitioning.partitionKeys, partitioning.isGrouped) diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DistributionAndOrderingUtils.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DistributionAndOrderingUtils.scala index 02e19dd053f29..9f3c4daa7a6f3 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DistributionAndOrderingUtils.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DistributionAndOrderingUtils.scala @@ -18,7 +18,7 @@ package org.apache.spark.sql.execution.datasources.v2 import org.apache.spark.sql.catalyst.analysis.{AnsiTypeCoercion, ResolveTimeZone, TypeCoercion} -import org.apache.spark.sql.catalyst.expressions.{Expression, Literal, SortOrder, TransformExpression, V2ExpressionUtils} +import org.apache.spark.sql.catalyst.expressions.{Expression, SortOrder, TransformExpression, V2ExpressionUtils} import org.apache.spark.sql.catalyst.expressions.V2ExpressionUtils._ import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, RebalancePartitions, RepartitionByExpression, Sort} import org.apache.spark.sql.catalyst.rules.{Rule, RuleExecutor} @@ -96,9 +96,7 @@ object DistributionAndOrderingUtils { } private def resolveTransformExpression(expr: Expression): Expression = expr.transform { - case TransformExpression(scalarFunc: ScalarFunction[_], arguments, Some(numBuckets)) => - V2ExpressionUtils.resolveScalarFunction(scalarFunc, Seq(Literal(numBuckets)) ++ arguments) - case TransformExpression(scalarFunc: ScalarFunction[_], arguments, None) => + case TransformExpression(scalarFunc: ScalarFunction[_], arguments) => V2ExpressionUtils.resolveScalarFunction(scalarFunc, arguments) } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala index e765b86301892..c4873e6494abb 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala @@ -45,6 +45,7 @@ import org.apache.spark.sql.functions.{col, max} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.internal.SQLConf._ import org.apache.spark.sql.types._ +import org.apache.spark.unsafe.types.{CalendarInterval, UTF8String} class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with ExplainSuiteHelper { private val functions = Seq( @@ -133,7 +134,7 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with val df = sql(s"SELECT * FROM testcat.ns.$table") val distribution = physical.ClusteredDistribution( - Seq(TransformExpression(BucketFunction, Seq(attr("ts")), Some(32)))) + Seq(TransformExpression(BucketFunction, Seq(Literal(32), attr("ts"))))) checkQueryPlan(df, distribution, physical.UnknownPartitioning(0)) } @@ -145,7 +146,7 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with val df = sql(s"SELECT * FROM testcat.ns.$table") val distribution = physical.ClusteredDistribution( - Seq(TransformExpression(BucketFunction, Seq(attr("ts")), Some(32)))) + Seq(TransformExpression(BucketFunction, Seq(Literal(32), attr("ts"))))) // Has exactly one partition. val partitionKeys = Seq(0).map(v => InternalRow.fromSeq(Seq(v))) @@ -201,13 +202,13 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with val df = sql(s"SELECT * FROM testcat.ns.$table") val distribution = physical.ClusteredDistribution( - Seq(TransformExpression(BucketFunction, Seq(attr("ts")), Some(32)))) + Seq(TransformExpression(BucketFunction, Seq(Literal(32), attr("ts"))))) checkQueryPlan(df, distribution, physical.UnknownPartitioning(0)) } } - test("non-clustered distribution: V2 function with multiple args") { + test("clustered distribution: V2 function with multiple args") { val partitions: Array[Transform] = Array( Expressions.apply("truncate", Expressions.column("data"), Expressions.literal(2)) ) @@ -223,7 +224,11 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with val distribution = physical.ClusteredDistribution( Seq(TransformExpression(TruncateFunction, Seq(attr("data"), Literal(2))))) - checkQueryPlan(df, distribution, physical.UnknownPartitioning(0)) + // With truncate transform support, KeyedPartitioning should now work + val partitionKeys = Seq("aa", "bb", "cc").map(v => + InternalRow(UTF8String.fromString(v))) + checkQueryPlan(df, distribution, + physical.KeyedPartitioning(distribution.clustering, partitionKeys)) } /** @@ -4191,6 +4196,606 @@ class KeyGroupedPartitioningSuite extends DistributionAndOrderingSuiteBase with } } + test("SPARK-50593: cross-function truncate vs bucket should NOT trigger SPJ") { + val partitions1 = Array( + Expressions.apply("truncate", Expressions.column("data"), Expressions.literal(3)) + ) + val partitions2 = Array( + Expressions.bucket(4, "data") + ) + + createTable("trunc_cross1", columns, partitions1) + sql("INSERT INTO testcat.ns.trunc_cross1 VALUES " + + "(0, 'aaa', CAST('2022-01-01' AS timestamp)), " + + "(1, 'bbb', CAST('2021-01-01' AS timestamp))") + + createTable("trunc_cross2", columns2, partitions2) + sql("INSERT INTO testcat.ns.trunc_cross2 VALUES " + + "(1, 5, 'aaa'), " + + "(5, 10, 'bbb')") + + withSQLConf( + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") { + + val df = sql( + s""" + |${selectWithMergeJoinHint("trunc_cross1", "trunc_cross2")} + |trunc_cross1.id, trunc_cross2.store_id + |FROM testcat.ns.trunc_cross1 JOIN testcat.ns.trunc_cross2 + |ON trunc_cross1.data = trunc_cross2.data + |ORDER BY trunc_cross1.id + |""".stripMargin) + + // Different functions (truncate vs bucket) are not mutually reducible, so a shuffle + // must still be planned. + val shuffles = collectShuffles(df.queryExecution.executedPlan) + assert(shuffles.nonEmpty, + "truncate vs bucket are not compatible - a shuffle should be present, " + + "but none was planned") + checkAnswer(df, Seq(Row(0, 1), Row(1, 5))) + } + } + + test("SPARK-50593: truncate(3) vs truncate(5) triggers SPJ via width reducer") { + // Exercises the Literal[]-based reducer path end-to-end: truncate widths 3 and 5 + // are mutually reducible (reduce the larger to the smaller), so SPJ must avoid the shuffle. + val table1 = "trunc_three" + val table2 = "trunc_five" + + val partitions1 = Array( + Expressions.apply("truncate", Expressions.column("data"), Expressions.literal(3))) + val partitions2 = Array( + Expressions.apply("truncate", Expressions.column("data"), Expressions.literal(5))) + + createTable(table1, columns, partitions1) + sql(s"INSERT INTO testcat.ns.$table1 VALUES " + + "(0, 'apple', CAST('2022-01-01' AS timestamp)), " + + "(1, 'grape', CAST('2021-01-01' AS timestamp)), " + + "(2, 'orange', CAST('2020-01-01' AS timestamp))") + + createTable(table2, columns, partitions2) + sql(s"INSERT INTO testcat.ns.$table2 VALUES " + + "(10, 'apple', CAST('2022-01-01' AS timestamp)), " + + "(20, 'grape', CAST('2021-01-01' AS timestamp)), " + + "(30, 'orange', CAST('2020-01-01' AS timestamp))") + + withSQLConf( + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") { + + val df = sql( + s""" + |${selectWithMergeJoinHint(table1, table2)} + |$table1.id AS left_id, $table2.id AS right_id + |FROM testcat.ns.$table1 JOIN testcat.ns.$table2 + |ON $table1.data = $table2.data + |ORDER BY $table1.id + |""".stripMargin) + + val shuffles = collectShuffles(df.queryExecution.executedPlan) + assert(shuffles.isEmpty, + "truncate(3) vs truncate(5) should avoid shuffle via the width reducer, " + + "but a shuffle was planned") + checkAnswer(df, Seq(Row(0, 10), Row(1, 20), Row(2, 30))) + } + } + + test("SPARK-50593: existing bucket SPJ still works with Literal[] API") { + // Exercises the new Literal[]-based reducer path end-to-end: bucket(4) and + // bucket(2) differ, so SPJ can only avoid the shuffle if BucketFunction's reducer + // (now implemented via Literal[] params) correctly returns a GCD-based Reducer. + // BucketFunction overrides only the new API, so this also covers the deprecated->new + // fallback: the single-int dispatch tries reducer(int, ...) first (UOE), then the Literal[]. + val table1 = "bucket_compat1" + val table2 = "bucket_compat2" + + val partitions1 = Array(Expressions.bucket(4, "id")) + val partitions2 = Array(Expressions.bucket(2, "store_id")) + + createTable(table1, columns, partitions1) + sql(s"INSERT INTO testcat.ns.$table1 VALUES " + + "(0, 'aaa', CAST('2022-01-01' AS timestamp)), " + + "(1, 'bbb', CAST('2021-01-01' AS timestamp)), " + + "(2, 'ccc', CAST('2020-01-01' AS timestamp)), " + + "(3, 'ddd', CAST('2019-01-01' AS timestamp))") + + createTable(table2, columns2, partitions2) + sql(s"INSERT INTO testcat.ns.$table2 VALUES " + + "(0, 5, 'aaa'), " + + "(1, 10, 'bbb'), " + + "(2, 15, 'ccc'), " + + "(3, 20, 'ddd')") + + withSQLConf( + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") { + val df = sql( + s""" + |${selectWithMergeJoinHint(table1, table2)} + |$table1.id, $table2.store_id + |FROM testcat.ns.$table1 JOIN testcat.ns.$table2 + |ON $table1.id = $table2.store_id + |ORDER BY $table1.id + |""".stripMargin) + + val shuffles = collectShuffles(df.queryExecution.executedPlan) + assert(shuffles.isEmpty, + "bucket(4) vs bucket(2) should avoid shuffle via the GCD reducer, " + + "but a shuffle was planned") + checkAnswer(df, Seq(Row(0, 0), Row(1, 1), Row(2, 2), Row(3, 3))) + } + } + + test("SPARK-50593: bucket(4) vs bucket(3) - no common divisor, must shuffle") { + // GCD(4, 3) = 1 -- BucketFunction.reducer returns null. Spark must NOT enable SPJ. + // Regression guard: a buggy null-handling in TransformExpression.reducer (e.g., + // Try(...).toOption instead of Try(Option(...))) would treat null as Some(null), + // enable SPJ, and produce wrong join results for incompatible bucket layouts. + val table1 = "bucket_gcd1_a" + val table2 = "bucket_gcd1_b" + + val partitions1 = Array(Expressions.bucket(4, "id")) + val partitions2 = Array(Expressions.bucket(3, "store_id")) + + createTable(table1, columns, partitions1) + sql(s"INSERT INTO testcat.ns.$table1 VALUES " + + "(0, 'aaa', CAST('2022-01-01' AS timestamp)), " + + "(1, 'bbb', CAST('2021-01-01' AS timestamp)), " + + "(2, 'ccc', CAST('2020-01-01' AS timestamp))") + + createTable(table2, columns2, partitions2) + sql(s"INSERT INTO testcat.ns.$table2 VALUES " + + "(0, 5, 'aaa'), " + + "(1, 10, 'bbb'), " + + "(2, 15, 'ccc')") + + withSQLConf( + SQLConf.V2_BUCKETING_PUSH_PART_VALUES_ENABLED.key -> "true", + SQLConf.V2_BUCKETING_ALLOW_COMPATIBLE_TRANSFORMS.key -> "true") { + val df = sql( + s""" + |${selectWithMergeJoinHint(table1, table2)} + |$table1.id, $table2.store_id + |FROM testcat.ns.$table1 JOIN testcat.ns.$table2 + |ON $table1.id = $table2.store_id + |ORDER BY $table1.id + |""".stripMargin) + + val shuffles = collectShuffles(df.queryExecution.executedPlan) + assert(shuffles.nonEmpty, + "bucket(4) vs bucket(3) have no common divisor (GCD=1), so the reducer " + + "returns null. SPJ must NOT be enabled; a shuffle is required.") + checkAnswer(df, Seq(Row(0, 0), Row(1, 1), Row(2, 2))) + } + } + + test("SPARK-50593: isSameFunction recurses into nested transforms, respects column-ref slots") { + import org.apache.spark.sql.catalyst.expressions.{Add, Expression, GetStructField} + val a = attr("a") + val b = attr("b") + def bucket(n: Int, e: Expression): TransformExpression = + TransformExpression(BucketFunction, Seq(Literal(n), e)) + def years(e: Expression): TransformExpression = TransformExpression(YearsFunction, Seq(e)) + def days(e: Expression): TransformExpression = TransformExpression(DaysFunction, Seq(e)) + + // Nested identical -> same (recursing into the inner transform), with column identity ignored. + // isSameFunction stays correct for nested shapes even though the SPJ gate currently rejects + // them; keeping this behavior is the right shape for any future nested support. + assert(bucket(4, years(a)).isSameFunction(bucket(4, years(a)))) + assert(bucket(4, years(a)).isSameFunction(bucket(4, years(b))), "column identity is ignored") + // Nested different inner -> not same. + assert(!bucket(4, years(a)).isSameFunction(bucket(4, days(a)))) + // Different outer literal -> not same. + assert(!bucket(4, years(a)).isSameFunction(bucket(2, years(a)))) + // Flat sanity (no nesting). + assert(bucket(4, a).isSameFunction(bucket(4, b))) + assert(!bucket(4, a).isSameFunction(bucket(2, b))) + + // A non-reference column slot (a + 1) carries value-changing semantics, so it is conservatively + // treated as not-same -- even compared to itself. + val add = bucket(4, Add(a, Literal(1))) + assert(!add.isSameFunction(bucket(4, Add(b, Literal(1))))) + assert(!add.isSameFunction(add), "a non-reference slot is treated as not-same by design") + + // Struct-field column references are recognized (reflexivity preserved for genuine refs). + val s = AttributeReference("s", StructType(Seq(StructField("f", IntegerType))))() + val sf = GetStructField(s, 0) + assert(bucket(4, sf).isSameFunction(bucket(4, sf))) + } + + test("SPARK-50593: supportsExpressions admits flat parameterized transforms, " + + "rejects nested and non-reference slots") { + import org.apache.spark.sql.catalyst.expressions.{Add, Expression} + val a = AttributeReference("a", IntegerType)() + val b = AttributeReference("b", IntegerType)() + def bucket(n: Int, e: Expression): TransformExpression = + TransformExpression(BucketFunction, Seq(Literal(n), e)) + + // Flat parameterized transform over a bare column -> admitted (one non-literal child = column). + assert(physical.KeyedPartitioning.supportsExpressions(Seq(bucket(4, a)))) + // Bare identity column -> admitted. + assert(physical.KeyedPartitioning.supportsExpressions(Seq(a))) + + // Nested transform -> rejected: the non-literal child is a transform, not a column reference. + // SPJ reasons about a transform via its function and literal params alone, which is unsound + // when the remaining argument is itself a transform. + val nested = bucket(4, TransformExpression(YearsFunction, Seq(a))) + assert(!physical.KeyedPartitioning.supportsExpressions(Seq(nested))) + + // Value-changing slot (a + 1) -> rejected: not a plain column reference. + assert(!physical.KeyedPartitioning.supportsExpressions(Seq(bucket(4, Add(a, Literal(1)))))) + + // Two non-literal column references -> rejected: a partition expression must map to exactly one + // clustering column (the positional keyPositions model needs a single column per transform). + assert(!physical.KeyedPartitioning.supportsExpressions( + Seq(TransformExpression(BucketFunction, Seq(Literal(4), a, b))))) + } + + test("SPARK-50593: integer truncate is reducible via lcm (generalized reducer, non-bucket)") { + // A second reducible transform exercising the generalized Literal[] reducer API with reducer + // math distinct from bucket (GCD) and string truncate (prefix-min): integer truncate snaps to + // a coarser grid, so truncate(v, W1) and truncate(v, W2) reduce onto multiples of lcm(W1, W2). + import org.apache.spark.sql.catalyst.expressions.Expression + val id = attr("id") + def itrunc(e: Expression, w: Int): TransformExpression = + TransformExpression(IntegerTruncateFunction, Seq(e, Literal(w))) + + // Same width -> same function (no reduction needed). + assert(itrunc(id, 4).isSameFunction(itrunc(id, 4))) + + // W2 is a multiple of W1: the finer side (W1=2) reduces onto the coarser grid (W2=4). + assert(itrunc(id, 2).isCompatible(itrunc(id, 4))) + val r = itrunc(id, 2).reducers(itrunc(id, 4)) + assert(r.isDefined, "truncate(2) must reduce onto truncate(4)") + val red = r.get.asInstanceOf[Reducer[Integer, Integer]] + // truncate(.,2) values snapped to multiples of 4: 6 -> 4, 2 -> 0, 8 -> 8 + assert(red.reduce(6) == 4 && red.reduce(2) == 0 && red.reduce(8) == 8) + // The coarser side (4) is already the common grid -> no reducer. + assert(itrunc(id, 4).reducers(itrunc(id, 2)).isEmpty) + + // Neither divides the other: both sides reduce to the lcm grid. + assert(itrunc(id, 6).isCompatible(itrunc(id, 4))) // lcm(6, 4) = 12 + assert(itrunc(id, 6).reducers(itrunc(id, 4)).isDefined) + assert(itrunc(id, 4).reducers(itrunc(id, 6)).isDefined) + assert(itrunc(id, 3).isCompatible(itrunc(id, 5))) // coprime -> lcm(3, 5) = 15 + } + + test("SPARK-50593: deprecated int reducer API still works (legacy connector backward compat)") { + // The reducer dispatch attempts the deprecated reducer(int, func, int) first for single-int + // params, so a ReducibleFunction that overrides ONLY the deprecated method still reduces. + // This mirrors how Iceberg 1.10.0 (and earlier) ship -- they predate the Literal[] API. + val bucketExpr4 = TransformExpression(LegacyBucketFunction, Seq(Literal(4), attr("id"))) + val bucketExpr2 = TransformExpression(LegacyBucketFunction, Seq(Literal(2), attr("id"))) + + val reducer = bucketExpr4.reducers(bucketExpr2) + assert(reducer.isDefined, "Expected a reducer for legacy_bucket(4) on legacy_bucket(2)") + + // Verify the returned Reducer actually reduces bucket 4 -> bucket 2 (GCD = 2). + // bucket(4, x) produces values in [0, 4); reducing by GCD=2 gives v % 2. + val r = reducer.get.asInstanceOf[Reducer[Integer, Integer]] + assert(r.reduce(3) == 1, s"Expected reduce(3) == 1, got ${r.reduce(3)}") + assert(r.reduce(2) == 0, s"Expected reduce(2) == 0, got ${r.reduce(2)}") + } + + test("SPARK-50593: a non-IntegerType param (DateType) does not reach the deprecated " + + "int reducer") { + // DateType is stored as a boxed Integer (epoch days) internally, so the reducer dispatch must + // key off the DataType, not the runtime class -- otherwise a DateType param is mistaken for the + // bucket-style int param and routed to the deprecated reducer(int, ...). LegacyBucketFunction + // overrides ONLY that deprecated method, so with a DateType param it must be unreachable, + // leaving the pair not reducible (rather than producing a bogus GCD reducer over epoch-days). + val l = TransformExpression(LegacyBucketFunction, Seq(Literal(8, DateType), attr("id"))) + val r = TransformExpression(LegacyBucketFunction, Seq(Literal(4, DateType), attr("id"))) + assert(!l.isSameFunction(r)) + assert(!l.isCompatible(r), "a DateType param must not reach the deprecated int reducer") + assert(l.reducers(r).isEmpty && r.reducers(l).isEmpty) + } + + test("SPARK-50593: mismatched column/literal argument layout is not reducible") { + // Both transforms pass the strict gate (one column-reference non-literal child), but the column + // and literal sit in swapped positions: truncate(id, 2) is (col, lit) while truncate(4, sid) is + // (lit, col). The reducer only sees the literal positions ([2] vs [4]), so without an + // argument-layout check it would wrongly reduce these and co-locate non-matching rows. + // IntegerTruncateFunction has two same-typed (Int) args, which makes this layout reachable. + val l = TransformExpression(IntegerTruncateFunction, Seq(attr("id"), Literal(2))) + val r = TransformExpression(IntegerTruncateFunction, Seq(Literal(4), attr("store_id"))) + assert(!l.isSameFunction(r)) + assert(!l.isCompatible(r), "swapped column/literal layout must not be reducible") + assert(l.reducers(r).isEmpty && r.reducers(l).isEmpty) + + // Control: same layout (col, lit) on both sides remains reducible via lcm(2, 4). + val a = TransformExpression(IntegerTruncateFunction, Seq(attr("id"), Literal(2))) + val b = TransformExpression(IntegerTruncateFunction, Seq(attr("store_id"), Literal(4))) + assert(a.isCompatible(b), "aligned (col, lit) layout must remain reducible") + } + + test("SPARK-50593: a dual-API connector reduces via the generalized overload") { + // DualApiBucketFunction implements both overloads: the deprecated reducer(int, ...) returns + // null, the generalized reducer(Literal[], ...) returns a valid GCD reducer. Generalized-first + // dispatch reduces via the generalized overload directly; the deprecated overload is not + // consulted (its null is irrelevant). + val l = TransformExpression(DualApiBucketFunction, Seq(Literal(4), attr("id"))) + val r = TransformExpression(DualApiBucketFunction, Seq(Literal(2), attr("store_id"))) + assert(l.isCompatible(r), "the generalized overload must produce a reducer") + val red = l.reducers(r) + assert(red.isDefined, "generalized reducer must be reached") + assert(red.get.asInstanceOf[Reducer[Integer, Integer]].reduce(3) == 1) + } + + test("SPARK-50593: the generalized overload's null is authoritative, no deprecated fallback") { + // DualApiGeneralizedNullFunction's generalized reducer returns null (not reducible) for a + // single-int pair its deprecated reducer WOULD reduce (gcd). Under generalized-first dispatch + // the generalized null is authoritative: Spark must not fall back to the deprecated overload, + // so the pair is not reducible. (Deprecated-first would instead co-partition via gcd(4,2)=2.) + val l = TransformExpression(DualApiGeneralizedNullFunction, Seq(Literal(4), attr("id"))) + val r = TransformExpression(DualApiGeneralizedNullFunction, Seq(Literal(2), attr("store_id"))) + assert(l.reducers(r).isEmpty, + "a generalized null must not fall back to the deprecated overload") + assert(r.reducers(l).isEmpty, "symmetric") + } + + test("SPARK-50593: a complex (non-scalar) literal param is not reducible") { + // Reducer parameters must not carry Catalyst-internal containers. ArrayParamFunction's + // generalized reducer returns a reducer unconditionally, so reaching it at all is the leak; the + // guard must refuse the ArrayData-backed literal param first. Different array values keep + // isSameFunction false, forcing the reducer path where the guard applies. + val l = TransformExpression(ArrayParamFunction, + Seq(Literal.create(Array(1, 2, 3), ArrayType(IntegerType)), attr("id"))) + val r = TransformExpression(ArrayParamFunction, + Seq(Literal.create(Array(4, 5, 6), ArrayType(IntegerType)), attr("store_id"))) + assert(!l.isSameFunction(r)) + assert(!l.isCompatible(r), "a complex literal param must not be reducible") + assert(l.reducers(r).isEmpty && r.reducers(l).isEmpty) + } + + test("SPARK-50593: a UDT-typed literal param is not reducible (complex-type guard)") { + // noComplexLiteralParams rejects a UDT-typed literal param by its DataType. UdtParamFunction + // declares the UDT as its literal input type, so literalParamsMatchInputTypes passes (UDT == + // UDT) and only the complex-type guard can reject it. UdtParamFunction reduces unconditionally, + // so reaching its reducer is the leak (here a StructBackedUDT, whose value is an InternalRow). + val udt = new StructBackedUDT + val l = TransformExpression(UdtParamFunction, + Seq(Literal(udt.serialize(new StructBacked(1)), udt), attr("id"))) + val r = TransformExpression(UdtParamFunction, + Seq(Literal(udt.serialize(new StructBacked(2)), udt), attr("store_id"))) + assert(!l.isSameFunction(r)) + assert(!l.isCompatible(r), "a UDT-typed literal param must not be reducible") + assert(l.reducers(r).isEmpty && r.reducers(l).isEmpty) + } + + test("SPARK-50593: bundled reducers tolerate a length-mismatched params call (no AIOOBE)") { + // sameArgumentLayout is arity-less, so a 0-vs-1-parameter pair (e.g. truncate(col) vs + // truncate(col, w), whose column slots align) reaches the connector reducer with + // mismatched-length param arrays. The bundled reducers model the documented contract by + // length-checking before indexing -- returning null rather than throwing ArrayIndexOutOfBounds + // (which attempt()'s Try would swallow into a silent missed SPJ + a misleading warning). + val empty = Array.empty[org.apache.spark.sql.connector.expressions.Literal[_]] + val one = Array[org.apache.spark.sql.connector.expressions.Literal[_]](literal(3)) + assert(BucketFunction.reducer(empty, BucketFunction, one) == null) + assert(TruncateFunction.reducer(empty, TruncateFunction, one) == null) + assert(IntegerTruncateFunction.reducer(empty, IntegerTruncateFunction, one) == null) + } + + test("SPARK-50593: a non-UOE reducer exception is logged and treated as not reducible") { + // An UnsupportedOperationException means "overload not implemented" (silent). Any other + // throwable is a bug in an implemented reducer: the dispatch logs it (not the misleading + // "implements no reducer" hint) and treats the pair as not reducible -- it falls back to a + // shuffle. + val id = attr("id") + val l = TransformExpression(ThrowingReducerFunction, Seq(id, Literal(2))) + val r = TransformExpression(ThrowingReducerFunction, Seq(id, Literal(4))) + val appender = new LogAppender("non-UOE reducer exception") + withLogAppender(appender) { + assert(l.reducers(r).isEmpty, "a throwing reducer must be treated as not reducible") + } + val messages = appender.loggingEvents.map(_.getMessage.getFormattedMessage) + assert(messages.exists(_.contains("reducer threw an exception")), + "the non-UOE exception must be logged") + assert(!messages.exists(_.contains("implements no reducer")), + "must not emit the 'implements no reducer' hint for an implemented-but-throwing reducer") + } + + test("SPARK-50593: deprecated overload is not probed once the generalized one reduces") { + // Generalized-first dispatch: the generalized reducer is tried first, and the deprecated int + // overload must not be probed once it reduced. Here the generalized reducer succeeds and the + // deprecated one throws; SPJ must succeed via the generalized path with NO "reducer threw" + // warning (which an eager probe of the deprecated overload would spuriously log). + val l = TransformExpression(GeneralizedOkDeprecatedThrowsFunction, Seq(Literal(4), attr("id"))) + val r = TransformExpression( + GeneralizedOkDeprecatedThrowsFunction, Seq(Literal(2), attr("store_id"))) + val appender = new LogAppender("deprecated overload probed eagerly") + withLogAppender(appender) { + assert(l.reducers(r).isDefined, "the generalized reducer must produce a reducer") + } + assert(!appender.loggingEvents.map(_.getMessage.getFormattedMessage) + .exists(_.contains("reducer threw an exception")), + "the deprecated overload must not be probed (and throw) once the generalized one reduced") + } + + test("SPARK-50593: a throwing generalized overload is surfaced, not masked by the deprecated " + + "one") { + // DeprecatedOkGeneralizedThrowsFunction implements both: the generalized overload throws, the + // deprecated one would reduce. Generalized-first dispatch treats the generalized bug as + // authoritative -- it logs the exception and does NOT fall back to the deprecated overload, so + // the pair shuffles. (Deprecated-first would have silently reduced via the deprecated overload, + // masking the new-API bug.) + val l = TransformExpression(DeprecatedOkGeneralizedThrowsFunction, Seq(Literal(4), attr("id"))) + val r = TransformExpression( + DeprecatedOkGeneralizedThrowsFunction, Seq(Literal(2), attr("store_id"))) + val appender = new LogAppender("generalized bug masked") + withLogAppender(appender) { + assert(l.reducers(r).isEmpty, "a throwing generalized overload must not fall back and reduce") + } + assert(appender.loggingEvents.map(_.getMessage.getFormattedMessage) + .exists(_.contains("reducer threw an exception")), "the generalized bug must be surfaced") + } + + test("SPARK-50593: a typed-null integer param is not routed to the deprecated int reducer") { + // LegacyIntReducerFunction implements ONLY the deprecated int reducer (accepts any int), so the + // generalized probe is Unimplemented and dispatch falls back to the deprecated overload. A + // typed-null IntegerType param must be excluded by isSingleInt from that fallback -- otherwise + // null.asInstanceOf[Int] fabricates a 0 the legacy reducer accepts, falsely co-partitioning + // null vs 0. So the pair must NOT be reducible. + val col = AttributeReference("id", IntegerType)() + val l = TransformExpression(LegacyIntReducerFunction, Seq(Literal(null, IntegerType), col)) + val r = TransformExpression(LegacyIntReducerFunction, Seq(Literal(null, IntegerType), col)) + assert(l.reducers(r).isEmpty, + "a typed-null int param must not reach the deprecated int fallback") + } + + test("SPARK-50593: a column whose type differs from the declared input type is not reducible " + + "(identity-vs-transform, exact-typed literal)") { + // Models a cross-side join `a = b` on ShortType keys: left is identity(a), right is + // truncate(b, 4). The literal slot matches the declared input type exactly, but the column does + // not -- identityReducer evals the transform (with the identity column substituted in), feeding + // it through SpecificInternalRow(inputTypes), so a ShortType column at an IntegerType-declared + // position would ClassCastException. argsMatchInputTypes checks the column too, so the pair is + // not reducible (shuffle). IntegerTruncateFunction declares (IntegerType, IntegerType). + val a = AttributeReference("a", ShortType)() // left, identity side + val b = AttributeReference("b", ShortType)() // right, transform's value column + val identity = physical.KeyedPartitioning(Seq(a), Seq.empty) + val truncated = physical.KeyedPartitioning( + Seq(TransformExpression(IntegerTruncateFunction, Seq(b, Literal(4)))), Seq.empty) + val idSpec = physical.KeyedShuffleSpec(identity, physical.ClusteredDistribution(Seq(a))) + val trSpec = physical.KeyedShuffleSpec(truncated, physical.ClusteredDistribution(Seq(b))) + assert(idSpec.reducers(trSpec).isEmpty, + "a mismatched-type column must not be reducible via the identity-vs-transform eval path") + } + + test("SPARK-50593: identityReducer validates the substituted identity column, not the " + + "transform's own column") { + // The transform's own column matches its declared input type, but the identity column actually + // evaluated (substituted in) does not. identityReducer must reject based on the substituted + // column -- what it evals -- else it builds a reducer that ClassCastExceptions at eval. (The + // planner's keyPositions normally forces the two columns to share a type; this builds the specs + // directly to pin identityReducer's own correctness.) + val idCol = AttributeReference("a", ShortType)() // evaluated column: ShortType + val tCol = AttributeReference("b", StringType)() // transform's column: matches declared type + val identity = physical.KeyedPartitioning(Seq(idCol), Seq.empty) + val truncated = physical.KeyedPartitioning( + Seq(TransformExpression(TruncateFunction, Seq(tCol, Literal(3)))), Seq.empty) + val idSpec = physical.KeyedShuffleSpec(identity, physical.ClusteredDistribution(Seq(idCol))) + val trSpec = physical.KeyedShuffleSpec(truncated, physical.ClusteredDistribution(Seq(tCol))) + assert(idSpec.reducers(trSpec).isEmpty, + "must validate the substituted identity column (ShortType), not the transform's own column") + } + + test("SPARK-50593: an arity-flexible transform (children > declared inputTypes) is not " + + "reducible via the identity-vs-transform eval path") { + // ZeroOrOneParamFunction declares inputTypes() of length 1 but admits transforms with 2 + // children (col + one literal). The eval path (identityReducer) feeds every child through + // ApplyFunctionExpression's SpecificInternalRow(inputTypes()) -- sized 1 -- so evaluating a + // 2-child transform would ArrayIndexOutOfBoundsException at reduce time. argsMatchInputTypes' + // exact-arity check rejects it, so the pair is not reducible (shuffle) instead of crashing. + // (The transform-vs-transform path keeps its mixed-arity flexibility -- see the zero-vs-one + // test below -- because it passes literals to the connector reducer and never evals.) + val col = AttributeReference("id", IntegerType)() + val identity = physical.KeyedPartitioning(Seq(col), Seq.empty) + val arityFlexible = physical.KeyedPartitioning( + Seq(TransformExpression(ZeroOrOneParamFunction, Seq(col, Literal(2)))), Seq.empty) + val idSpec = physical.KeyedShuffleSpec(identity, physical.ClusteredDistribution(Seq(col))) + val afSpec = physical.KeyedShuffleSpec(arityFlexible, physical.ClusteredDistribution(Seq(col))) + assert(idSpec.reducers(afSpec).isEmpty, + "a child beyond the declared arity must not reach the eval path (would AIOOBE)") + } + + test("SPARK-50593: zero-param vs one-param transforms reach the reducer (no arity block)") { + // raw(id) has children [id]; withParam(id, 2) has children [id, 2]. Both pass + // supportsExpressions (one column ref each). The dispatch must not require equal child counts + // before the reducer: ZeroOrOneParamFunction is reducible across the 0-vs-1-parameter shape, so + // the pair must reach it (the column slots align under zip; the extra parameter is reconciled + // by the reducer). + val raw = TransformExpression(ZeroOrOneParamFunction, Seq(attr("id"))) + val withParam = TransformExpression(ZeroOrOneParamFunction, Seq(attr("id"), Literal(2))) + assert(!raw.isSameFunction(withParam)) // different arity -> not the "same" transform + assert(raw.isCompatible(withParam), + "zero-param vs one-param must reach the connector reducer, not be blocked by arity") + assert(raw.reducers(withParam).isDefined && withParam.reducers(raw).isDefined) + } + + test("SPARK-50593: CalendarIntervalType literal param is reducible (not treated as complex)") { + // CalendarIntervalType is non-complex but not an AtomicType; its literal param must not be + // rejected as a complex container before the reducer is consulted. IntervalParamFunction is + // reducible; differing interval params keep isSameFunction false, forcing the reducer path. + val l = TransformExpression(IntervalParamFunction, + Seq(attr("id"), Literal(new CalendarInterval(1, 0, 0), CalendarIntervalType))) + val r = TransformExpression(IntervalParamFunction, + Seq(attr("id"), Literal(new CalendarInterval(2, 0, 0), CalendarIntervalType))) + assert(!l.isSameFunction(r)) + assert(l.isCompatible(r), "a CalendarIntervalType param must reach the connector reducer") + assert(l.reducers(r).isDefined) + } + + // Builds (identity spec, truncate(col, ) spec) on the same column. The ShortType + // width mismatches truncate's declared IntegerType input, so the pair must be treated as not + // reducible. Shared by the reducers and gate tests below, which must agree on that decision. + private def mismatchedIdVsTransformSpecs() + : (physical.KeyedShuffleSpec, physical.KeyedShuffleSpec) = { + val data = AttributeReference("data", StringType)() + val identity = physical.KeyedPartitioning(Seq(data), Seq.empty) + val truncated = physical.KeyedPartitioning( + Seq(TransformExpression(TruncateFunction, Seq(data, Literal(2.toShort, ShortType)))), + Seq.empty) + val dist = physical.ClusteredDistribution(Seq(data)) + (physical.KeyedShuffleSpec(identity, dist), physical.KeyedShuffleSpec(truncated, dist)) + } + + test("SPARK-50593: a literal param whose type differs from the declared input type is not " + + "reducible (identity-vs-transform)") { + // A bound function may declare inputTypes() that differ from the literal's actual type (a legal + // implicit cast). truncate declares (StringType, IntegerType); a connector can report a Short + // width literal. The identity-vs-transform reducer binds and directly evals the transform, + // skipping Analyzer coercion -- a raw Short into an IntegerType slot would throw. Rather than + // coerce a value the partitions were not built on, this pair is not reducible (reducers => + // None); the companion gate test asserts areKeysCompatible also rejects it, so it shuffles. + val (idSpec, trSpec) = mismatchedIdVsTransformSpecs() + assert(idSpec.reducers(trSpec).isEmpty, + "a mismatched-type literal param must not be reducible via the identity-vs-transform path") + } + + test("SPARK-50593: identity-vs-transform compatibility gate agrees with reducers on a " + + "mismatched-type literal param") { + // The compatibility gate (areKeysCompatible -> isExpressionCompatible) and reducers MUST agree: + // if the gate says compatible but reducers returns None, EnsureRequirements keeps the identity + // side's raw keys (the reducedDataTypes check can't catch it -- both StringType) and SPJ joins + // raw-vs-transformed keys -> silent wrong results. So a mismatched-type literal must make the + // gate return false (force a shuffle), consistent with reducers returning None above. + val (idSpec, trSpec) = mismatchedIdVsTransformSpecs() + 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(!idSpec.areKeysCompatible(trSpec), + "gate must reject a mismatched-type literal so it agrees with reducers (no mis-join)") + assert(!trSpec.areKeysCompatible(idSpec), "symmetric") + } + } + + test("SPARK-50593: a literal param whose type differs from the declared input type is not " + + "reducible (transform-vs-transform)") { + // Both sides are truncate transforms whose width literal is ShortType, while the function + // declares (StringType, IntegerType). A literal whose type differs from the declared input type + // is treated as not reducible (no coercion), so the pair falls back to a shuffle. Uses a + // type-tolerant reducer (reads the width via Number) so emptiness is attributable to the gate, + // not to an incidental ClassCastException in the connector. The same widths typed as + // IntegerType remain reducible (control). + val data = AttributeReference("data", StringType)() + def trunc(w: Short): TransformExpression = + TransformExpression(TypeTolerantTruncateFunction, Seq(data, Literal(w, ShortType))) + assert(trunc(4).reducers(trunc(3)).isEmpty, + "mismatched-type (Short) width params must not be reducible") + assert(trunc(3).reducers(trunc(4)).isEmpty) + + // Control: IntegerType widths (matching the declared input type) still reduce. + def itrunc(w: Int): TransformExpression = + TransformExpression(TypeTolerantTruncateFunction, Seq(data, Literal(w))) + val reduced = itrunc(4).reducers(itrunc(3)) + assert(reduced.isDefined, "IntegerType widths must remain reducible") + assert(reduced.get.asInstanceOf[Reducer[Any, Any]] + .reduce(UTF8String.fromString("abcd")) == UTF8String.fromString("abc")) + } + test("SPARK-57881: storage-partitioned join leverages union output KeyedPartitioning to " + "avoid shuffle") { val cols = Array( diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/catalog/functions/transformFunctions.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/catalog/functions/transformFunctions.scala index 35102c6893d3b..d51fcffb78165 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/catalog/functions/transformFunctions.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/catalog/functions/transformFunctions.scala @@ -20,7 +20,9 @@ import java.time.{Instant, LocalDate, ZoneId} import java.time.temporal.ChronoUnit import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.GenericInternalRow import org.apache.spark.sql.catalyst.util.DateTimeUtils +import org.apache.spark.sql.connector.expressions.Literal import org.apache.spark.sql.types._ import org.apache.spark.unsafe.types.UTF8String @@ -213,11 +215,14 @@ object BucketFunction extends ScalarFunction[Int] with ReducibleFunction[Int, In } override def reducer( - thisNumBuckets: Int, + thisParams: Array[Literal[_]], otherFunc: ReducibleFunction[_, _], - otherNumBuckets: Int): Reducer[Int, Int] = { + otherParams: Array[Literal[_]]): Reducer[Int, Int] = { + + if (otherFunc == BucketFunction && thisParams.length == 1 && otherParams.length == 1) { + val thisNumBuckets = thisParams(0).value().asInstanceOf[Int] + val otherNumBuckets = otherParams(0).value().asInstanceOf[Int] - if (otherFunc == BucketFunction) { val gcd = this.gcd(thisNumBuckets, otherNumBuckets) if (gcd > 1 && gcd != thisNumBuckets) { return BucketReducer(gcd) @@ -235,6 +240,95 @@ case class BucketReducer(divisor: Int) extends Reducer[Int, Int] { override def displayName(): String = toString } +/** + * A bucket function that only overrides the deprecated `reducer(int, func, int)` method, + * not the new `reducer(Literal[], func, Literal[])` method. + * + * Used to verify that the default implementation of the new method correctly falls back + * to the deprecated int-based API, so legacy implementations continue to work. + */ +object LegacyBucketFunction extends ScalarFunction[Int] with ReducibleFunction[Int, Int] { + override def inputTypes(): Array[DataType] = Array(IntegerType, LongType) + override def resultType(): DataType = IntegerType + override def name(): String = "legacy_bucket" + override def canonicalName(): String = name() + override def toString: String = name() + override def produceResult(input: InternalRow): Int = { + Math.floorMod(input.getLong(1), input.getInt(0)) + } + + override def reducer( + thisNumBuckets: Int, + otherFunc: ReducibleFunction[_, _], + otherNumBuckets: Int): Reducer[Int, Int] = { + if (otherFunc == LegacyBucketFunction) { + val gcd = BigInt(thisNumBuckets).gcd(BigInt(otherNumBuckets)).toInt + if (gcd > 1 && gcd != thisNumBuckets) { + return BucketReducer(gcd) + } + } + null + } +} + +/** + * A bucket function that implements BOTH reducer overloads: the deprecated `reducer(int, ..., int)` + * always returns null (not reducible via the old API), while the new `reducer(Literal[], ...)` + * returns a GCD-based reducer. Used to verify that the dispatch falls back to the generalized + * overload when the deprecated one returns null (not only when it throws). + */ +object DualApiBucketFunction extends ScalarFunction[Int] with ReducibleFunction[Int, Int] { + override def inputTypes(): Array[DataType] = Array(IntegerType, LongType) + override def resultType(): DataType = IntegerType + override def name(): String = "dual_bucket" + override def canonicalName(): String = name() + override def toString: String = name() + override def produceResult(input: InternalRow): Int = { + Math.floorMod(input.getLong(1), input.getInt(0)) + } + + // Deprecated API: intentionally signals "not reducible" via null (not via an exception). + override def reducer( + thisNumBuckets: Int, + otherFunc: ReducibleFunction[_, _], + otherNumBuckets: Int): Reducer[Int, Int] = null + + // New API: a real GCD-based reducer. + override def reducer( + thisParams: Array[Literal[_]], + otherFunc: ReducibleFunction[_, _], + otherParams: Array[Literal[_]]): Reducer[Int, Int] = { + if (otherFunc == DualApiBucketFunction) { + val thisNumBuckets = thisParams(0).value().asInstanceOf[Int] + val otherNumBuckets = otherParams(0).value().asInstanceOf[Int] + val gcd = BigInt(thisNumBuckets).gcd(BigInt(otherNumBuckets)).toInt + if (gcd > 1 && gcd != thisNumBuckets) { + return BucketReducer(gcd) + } + } + null + } +} + +/** + * A function with a complex (ArrayType) literal parameter. Its generalized reducer returns a valid + * reducer unconditionally, so a test can prove the dispatch refuses to invoke it for a non-scalar + * literal param (rather than the call happening to fail on a cast). + */ +object ArrayParamFunction extends ScalarFunction[Int] with ReducibleFunction[Int, Int] { + override def inputTypes(): Array[DataType] = Array(ArrayType(IntegerType), LongType) + override def resultType(): DataType = IntegerType + override def name(): String = "array_param" + override def canonicalName(): String = name() + override def toString: String = name() + override def produceResult(input: InternalRow): Int = input.getInt(1) + + override def reducer( + thisParams: Array[Literal[_]], + otherFunc: ReducibleFunction[_, _], + otherParams: Array[Literal[_]]): Reducer[Int, Int] = BucketReducer(1) +} + object UnboundStringSelfFunction extends UnboundFunction { override def bind(inputType: StructType): BoundFunction = StringSelfFunction override def description(): String = name() @@ -253,12 +347,35 @@ object StringSelfFunction extends ScalarFunction[UTF8String] { } object UnboundTruncateFunction extends UnboundFunction { - override def bind(inputType: StructType): BoundFunction = TruncateFunction + override def bind(inputType: StructType): BoundFunction = { + if (inputType.size == 2) { + inputType.head.dataType match { + case StringType => TruncateFunction + case IntegerType => IntegerTruncateFunction + case _ => + throw new UnsupportedOperationException( + s"'truncate' does not support data type: ${inputType.head.dataType}") + } + } else { + throw new UnsupportedOperationException( + "'truncate' requires exactly 2 arguments: (column, width)") + } + } + override def description(): String = name() override def name(): String = "truncate" } -object TruncateFunction extends ScalarFunction[UTF8String] { +/** + * Truncate transform for String type. + * Follows Iceberg spec: truncate(str, L) = str[0:L] + * + * Implements ReducibleFunction: ANY two different widths are compatible. + * The reducer uses the smaller width. + */ +object TruncateFunction + extends ScalarFunction[UTF8String] + with ReducibleFunction[UTF8String, UTF8String] { override def inputTypes(): Array[DataType] = Array(StringType, IntegerType) override def resultType(): DataType = StringType override def name(): String = "truncate" @@ -266,7 +383,336 @@ object TruncateFunction extends ScalarFunction[UTF8String] { override def toString: String = name() override def produceResult(input: InternalRow): UTF8String = { val str = input.getUTF8String(0) - val length = input.getInt(1) - str.substring(0, length) + val width = input.getInt(1) + str.substring(0, width) + } + + override def reducer( + thisParams: Array[Literal[_]], + otherFunc: ReducibleFunction[_, _], + otherParams: Array[Literal[_]]): Reducer[UTF8String, UTF8String] = { + + if (otherFunc == TruncateFunction && thisParams.length == 1 && otherParams.length == 1) { + val thisWidth = thisParams(0).value().asInstanceOf[Int] + val otherWidth = otherParams(0).value().asInstanceOf[Int] + val smallerWidth = math.min(thisWidth, otherWidth) + + if (smallerWidth != thisWidth) { + return TruncateReducer(smallerWidth) + } + } + null + } +} + +case class TruncateReducer(width: Int) extends Reducer[UTF8String, UTF8String] { + override def reduce(value: UTF8String): UTF8String = { + value.substring(0, width) + } + override def resultType(): DataType = StringType + override def displayName(): String = s"truncate($width)" +} + +/** + * Truncate transform for Integer type. + * Follows Iceberg spec: truncate(value, W) = value - (((value % W) + W) % W), which snaps `value` + * down to a multiple of `W`. + * + * Implements ReducibleFunction: truncate(v, W1) and truncate(v, W2) are always reducible onto a + * common coarser grid of multiples of lcm(W1, W2). The finer side (whose width does not already + * equal the lcm) reduces by snapping to that grid; when W2 is a multiple of W1 the lcm is simply + * the coarser width W2. + */ +object IntegerTruncateFunction + extends ScalarFunction[Int] + with ReducibleFunction[Int, Int] { + override def inputTypes(): Array[DataType] = Array(IntegerType, IntegerType) + override def resultType(): DataType = IntegerType + override def name(): String = "truncate" + override def canonicalName(): String = name() + override def toString: String = name() + override def produceResult(input: InternalRow): Int = { + val value = input.getInt(0) + val width = input.getInt(1) + value - (((value % width) + width) % width) + } + + override def reducer( + thisParams: Array[Literal[_]], + otherFunc: ReducibleFunction[_, _], + otherParams: Array[Literal[_]]): Reducer[Int, Int] = { + if (otherFunc == IntegerTruncateFunction && thisParams.length == 1 && otherParams.length == 1) { + val thisWidth = thisParams(0).value().asInstanceOf[Int] + val otherWidth = otherParams(0).value().asInstanceOf[Int] + val common = lcm(thisWidth, otherWidth) + // Only the finer side reduces; if `common == thisWidth` this side is already the common grid. + if (common != thisWidth) { + return IntTruncateReducer(common) + } + } + null + } + + private def lcm(a: Int, b: Int): Int = { + val g = BigInt(a).gcd(BigInt(b)) + (BigInt(a) / g * BigInt(b)).toInt + } +} + +case class IntTruncateReducer(width: Int) extends Reducer[Int, Int] { + override def reduce(value: Int): Int = value - (((value % width) + width) % width) + override def resultType(): DataType = IntegerType + override def displayName(): String = s"truncate($width)" +} + +/** + * A transform whose reducer is defined across a zero-parameter vs one-parameter shape, e.g. + * `zero_or_one(col)` reducing onto `zero_or_one(col, 2)`. Used to verify the dispatch does not + * globally require equal child counts before invoking the reducer. + */ +object ZeroOrOneParamFunction extends ScalarFunction[Int] with ReducibleFunction[Int, Int] { + override def inputTypes(): Array[DataType] = Array(IntegerType) + override def resultType(): DataType = IntegerType + override def name(): String = "zero_or_one" + override def canonicalName(): String = name() + override def toString: String = name() + override def produceResult(input: InternalRow): Int = input.getInt(0) + + override def reducer( + thisParams: Array[Literal[_]], + otherFunc: ReducibleFunction[_, _], + otherParams: Array[Literal[_]]): Reducer[Int, Int] = { + if (otherFunc == ZeroOrOneParamFunction && thisParams.length != otherParams.length) { + BucketReducer(1) + } else { + null + } + } +} + +/** + * A transform with a `CalendarIntervalType` literal parameter (which is non-complex but not an + * `AtomicType`). Used to verify such a parameter is not rejected as a complex container before its + * reducer is consulted. + */ +object IntervalParamFunction extends ScalarFunction[Int] with ReducibleFunction[Int, Int] { + override def inputTypes(): Array[DataType] = Array(IntegerType, CalendarIntervalType) + override def resultType(): DataType = IntegerType + override def name(): String = "interval_param" + override def canonicalName(): String = name() + override def toString: String = name() + override def produceResult(input: InternalRow): Int = input.getInt(0) + + override def reducer( + thisParams: Array[Literal[_]], + otherFunc: ReducibleFunction[_, _], + otherParams: Array[Literal[_]]): Reducer[Int, Int] = { + if (otherFunc == IntervalParamFunction) BucketReducer(1) else null + } +} + +/** A user type whose UDT serializes to a struct (an [[InternalRow]]). */ +class StructBacked(val n: Int) extends Serializable + +/** + * A UDT whose `sqlType` is a [[StructType]]: a literal of this type carries an [[InternalRow]] + * value, yet its `dataType` is the UDT, not `StructType`. This is the case a `DataType`-based + * container check misses but a value-based one catches. + */ +class StructBackedUDT extends UserDefinedType[StructBacked] { + override def sqlType: DataType = StructType(Seq(StructField("n", IntegerType, nullable = false))) + override def serialize(obj: StructBacked): InternalRow = new GenericInternalRow(Array[Any](obj.n)) + override def deserialize(datum: Any): StructBacked = datum match { + case row: InternalRow => new StructBacked(row.getInt(0)) + } + override def userClass: Class[StructBacked] = classOf[StructBacked] +} + +/** + * A transform whose declared input type at the literal position is a UDT ([[StructBackedUDT]]). + * Used to make the value-based `noComplexLiteralParams` guard load-bearing: a UDT-over-struct + * literal matches the declared input type (so `literalParamsMatchInputTypes` passes), yet its value + * is an [[InternalRow]], so only the value-based guard can reject it. The reducer returns + * unconditionally, so reaching it at all is the leak. + */ +object UdtParamFunction extends ScalarFunction[Int] with ReducibleFunction[Int, Int] { + override def inputTypes(): Array[DataType] = Array(new StructBackedUDT, LongType) + override def resultType(): DataType = IntegerType + override def name(): String = "udt_param" + override def canonicalName(): String = name() + override def toString: String = name() + override def produceResult(input: InternalRow): Int = input.getInt(1) + + override def reducer( + thisParams: Array[Literal[_]], + otherFunc: ReducibleFunction[_, _], + otherParams: Array[Literal[_]]): Reducer[Int, Int] = BucketReducer(1) +} + +/** + * A string truncate whose reducer reads its width type-tolerantly (via [[Number]], so it accepts a + * boxed Short or Integer). Used to verify that the literal-param-type gate -- not an incidental + * ClassCastException in the connector -- is what makes a mismatched-type (e.g. ShortType) width + * non-reducible. With the gate removed, this reducer WOULD reduce a ShortType-width pair. + */ +object TypeTolerantTruncateFunction + extends ScalarFunction[UTF8String] + with ReducibleFunction[UTF8String, UTF8String] { + override def inputTypes(): Array[DataType] = Array(StringType, IntegerType) + override def resultType(): DataType = StringType + override def name(): String = "tolerant_truncate" + override def canonicalName(): String = name() + override def toString: String = name() + override def produceResult(input: InternalRow): UTF8String = + input.getUTF8String(0).substring(0, input.getInt(1)) + + override def reducer( + thisParams: Array[Literal[_]], + otherFunc: ReducibleFunction[_, _], + otherParams: Array[Literal[_]]): Reducer[UTF8String, UTF8String] = { + if (otherFunc == TypeTolerantTruncateFunction && + thisParams.length == 1 && otherParams.length == 1) { + val thisWidth = thisParams(0).value().asInstanceOf[Number].intValue() + val otherWidth = otherParams(0).value().asInstanceOf[Number].intValue() + val smaller = math.min(thisWidth, otherWidth) + if (smaller != thisWidth) return TruncateReducer(smaller) + } + null + } +} + +/** + * A function whose generalized reducer throws an unexpected (non-UnsupportedOperationException) + * exception. Used to verify the dispatch logs it and treats the pair as not reducible (a shuffle), + * rather than crashing or emitting the misleading "implements no reducer" hint (it does implement + * the overload -- it threw a bug, which is a different signal than UOE-means-unimplemented). + */ +object ThrowingReducerFunction extends ScalarFunction[Int] with ReducibleFunction[Int, Int] { + override def inputTypes(): Array[DataType] = Array(IntegerType, IntegerType) + override def resultType(): DataType = IntegerType + override def name(): String = "throwing_reducer" + override def canonicalName(): String = name() + override def toString: String = name() + override def produceResult(input: InternalRow): Int = input.getInt(1) + + override def reducer( + thisParams: Array[Literal[_]], + otherFunc: ReducibleFunction[_, _], + otherParams: Array[Literal[_]]): Reducer[Int, Int] = + throw new RuntimeException("boom from reducer") +} + +/** + * A bucket-like function implementing BOTH reducer overloads: the deprecated int overload succeeds + * (returns a reducer), while the generalized Literal[] overload throws. Used to verify the + * single-int dispatch is lazy -- it must not invoke (and log the throw from) the generalized + * overload once the deprecated one already produced a reducer. + */ +object DeprecatedOkGeneralizedThrowsFunction + extends ScalarFunction[Int] with ReducibleFunction[Int, Int] { + override def inputTypes(): Array[DataType] = Array(IntegerType, LongType) + override def resultType(): DataType = IntegerType + override def name(): String = "deprecated_ok_generalized_throws" + override def canonicalName(): String = name() + override def toString: String = name() + override def produceResult(input: InternalRow): Int = + Math.floorMod(input.getLong(1), input.getInt(0)) + + override def reducer( + thisNumBuckets: Int, + otherFunc: ReducibleFunction[_, _], + otherNumBuckets: Int): Reducer[Int, Int] = + if (otherFunc == DeprecatedOkGeneralizedThrowsFunction) BucketReducer(1) else null + + override def reducer( + thisParams: Array[Literal[_]], + otherFunc: ReducibleFunction[_, _], + otherParams: Array[Literal[_]]): Reducer[Int, Int] = + throw new RuntimeException("boom from generalized overload") +} + +/** + * The mirror of [[DeprecatedOkGeneralizedThrowsFunction]]: the generalized overload returns a + * reducer, the deprecated overload throws. Used to verify that under generalized-first dispatch the + * deprecated overload is NOT probed once the generalized one reduced (no "reducer threw" warning). + */ +object GeneralizedOkDeprecatedThrowsFunction + extends ScalarFunction[Int] with ReducibleFunction[Int, Int] { + override def inputTypes(): Array[DataType] = Array(IntegerType, LongType) + override def resultType(): DataType = IntegerType + override def name(): String = "generalized_ok_deprecated_throws" + override def canonicalName(): String = name() + override def toString: String = name() + override def produceResult(input: InternalRow): Int = + Math.floorMod(input.getLong(1), input.getInt(0)) + + override def reducer( + thisNumBuckets: Int, + otherFunc: ReducibleFunction[_, _], + otherNumBuckets: Int): Reducer[Int, Int] = + throw new RuntimeException("boom from deprecated overload") + + override def reducer( + thisParams: Array[Literal[_]], + otherFunc: ReducibleFunction[_, _], + otherParams: Array[Literal[_]]): Reducer[Int, Int] = + if (otherFunc == GeneralizedOkDeprecatedThrowsFunction) BucketReducer(1) else null +} + +/** + * A legacy connector: implements ONLY the deprecated int reducer (returns a reducer for any int) + * and not the generalized overload. Used to verify the isSingleInt null-guard under generalized- + * first dispatch: a typed-null int param must not reach the deprecated fallback, where + * null.asInstanceOf[Int] would fabricate a 0 this reducer would accept. + */ +object LegacyIntReducerFunction extends ScalarFunction[Int] with ReducibleFunction[Int, Int] { + override def inputTypes(): Array[DataType] = Array(IntegerType, IntegerType) + override def resultType(): DataType = IntegerType + override def name(): String = "legacy_int_reducer" + override def canonicalName(): String = name() + override def toString: String = name() + override def produceResult(input: InternalRow): Int = input.getInt(1) + + override def reducer( + thisNumBuckets: Int, + otherFunc: ReducibleFunction[_, _], + otherNumBuckets: Int): Reducer[Int, Int] = + if (otherFunc == LegacyIntReducerFunction) BucketReducer(1) else null +} + +/** + * A dual-API connector whose GENERALIZED reducer returns null (deliberately not reducible) while + * its DEPRECATED int reducer WOULD reduce a single-int pair (gcd). Used to pin generalized-first + * dispatch: the generalized null is authoritative, so Spark must NOT fall back to the deprecated + * overload (which would otherwise co-partition). + */ +object DualApiGeneralizedNullFunction extends ScalarFunction[Int] with ReducibleFunction[Int, Int] { + override def inputTypes(): Array[DataType] = Array(IntegerType, LongType) + override def resultType(): DataType = IntegerType + override def name(): String = "dual_api_generalized_null" + override def canonicalName(): String = name() + override def toString: String = name() + override def produceResult(input: InternalRow): Int = { + Math.floorMod(input.getLong(1), input.getInt(0)) + } + + // Generalized API: deliberately not reducible (returns null, not an exception). + override def reducer( + thisParams: Array[Literal[_]], + otherFunc: ReducibleFunction[_, _], + otherParams: Array[Literal[_]]): Reducer[Int, Int] = null + + // Deprecated int API: WOULD reduce via gcd -- a deprecated-first order would co-partition. + override def reducer( + thisNumBuckets: Int, + otherFunc: ReducibleFunction[_, _], + otherNumBuckets: Int): Reducer[Int, Int] = { + if (otherFunc == DualApiGeneralizedNullFunction) { + val gcd = BigInt(thisNumBuckets).gcd(BigInt(otherNumBuckets)).toInt + if (gcd > 1 && gcd != thisNumBuckets) { + return BucketReducer(gcd) + } + } + null } } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala index a70baece77844..629d65bb20c0b 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala @@ -19,7 +19,7 @@ package org.apache.spark.sql.execution import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, TransformExpression} +import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeReference, Literal, TransformExpression} import org.apache.spark.sql.catalyst.plans.physical.{ClusteredDistribution, HashPartitioning, KeyedPartitioning, Partitioning, PartitioningCollection, UnknownPartitioning} import org.apache.spark.sql.connector.catalog.functions.{BucketFunction, YearsFunction} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper @@ -492,7 +492,7 @@ class ProjectedOrderingAndPartitioningSuite // KP([bucket(32, id)], keys1d) through Project(id as pk) should produce // KP([bucket(32, pk)], keys1d): the alias is pushed into the bucket's column argument. val id = AttributeReference("id", IntegerType)() - val bucketExpr = TransformExpression(BucketFunction, Seq(id), Some(32)) + val bucketExpr = TransformExpression(BucketFunction, Seq(Literal(32), id)) val keys1d = Seq(InternalRow(0), InternalRow(1), InternalRow(2)) val child = DummyLeafExecWithPartitioning( output = Seq(id), @@ -507,7 +507,7 @@ class ProjectedOrderingAndPartitioningSuite case te: TransformExpression => assert(te.isSameFunction(bucketExpr), "bucket function and numBuckets must be preserved after alias substitution") - assert(te.children.head.asInstanceOf[Attribute].name === "pk", + assert(te.children.collectFirst { case a: Attribute => a }.get.name === "pk", "bucket's column argument must be rewritten to the aliased attribute") case other => fail(s"Expected TransformExpression, got $other") } @@ -524,7 +524,7 @@ class ProjectedOrderingAndPartitioningSuite // Result: KP([bucket(32, id)], keys1d, isNarrowed=true, isGrouped=false). val id = AttributeReference("id", IntegerType)() val ts = AttributeReference("ts", IntegerType)() - val bucketExpr = TransformExpression(BucketFunction, Seq(id), Some(32)) + val bucketExpr = TransformExpression(BucketFunction, Seq(Literal(32), id)) val yearsExpr = TransformExpression(YearsFunction, Seq(ts)) // Projected to position [0] (bucket): (0),(1),(0) -- bucket value 0 appears twice. val keys2d = Seq(InternalRow(0, 2020), InternalRow(1, 2020), InternalRow(0, 2021)) @@ -539,7 +539,7 @@ class ProjectedOrderingAndPartitioningSuite kp.expressions.head match { case te: TransformExpression => assert(te.isSameFunction(bucketExpr), "bucket must be the surviving expression") - assert(te.children.head.asInstanceOf[Attribute].name === "id") + assert(te.children.collectFirst { case a: Attribute => a }.get.name === "id") case other => fail(s"Expected TransformExpression, got $other") } assert(kp.isNarrowed, "dropping years(ts) position must mark the KP as narrowed") @@ -554,7 +554,7 @@ class ProjectedOrderingAndPartitioningSuite // Result: KP([bucket(32, id), years(ts_alias)], keys2d) -- not narrowed. val id = AttributeReference("id", IntegerType)() val ts = AttributeReference("ts", IntegerType)() - val bucketExpr = TransformExpression(BucketFunction, Seq(id), Some(32)) + val bucketExpr = TransformExpression(BucketFunction, Seq(Literal(32), id)) val yearsExpr = TransformExpression(YearsFunction, Seq(ts)) val keys2d = Seq(InternalRow(0, 2020), InternalRow(1, 2020), InternalRow(0, 2021)) val child = DummyLeafExecWithPartitioning( @@ -569,14 +569,14 @@ class ProjectedOrderingAndPartitioningSuite kp.expressions(0) match { case te: TransformExpression => assert(te.isSameFunction(bucketExpr)) - assert(te.children.head.asInstanceOf[Attribute].name === "id", + assert(te.children.collectFirst { case a: Attribute => a }.get.name === "id", "bucket's argument must remain id (no alias for id in this projection)") case other => fail(s"Expected TransformExpression at pos 0, got $other") } kp.expressions(1) match { case te: TransformExpression => assert(te.isSameFunction(yearsExpr)) - assert(te.children.head.asInstanceOf[Attribute].name === "ts_alias", + assert(te.children.collectFirst { case a: Attribute => a }.get.name === "ts_alias", "years() argument must be rewritten to ts_alias") case other => fail(s"Expected TransformExpression at pos 1, got $other") } diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala index 17d00ec055e07..84eee883aeeda 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/exchange/EnsureRequirementsSuite.scala @@ -1191,11 +1191,11 @@ class EnsureRequirementsSuite extends SharedSparkSession { } def bucket(numBuckets: Int, expr: Expression): TransformExpression = { - TransformExpression(BucketFunction, Seq(expr), Some(numBuckets)) + TransformExpression(BucketFunction, Seq(Literal(numBuckets), expr)) } def buckets(numBuckets: Int, expr: Seq[Expression]): TransformExpression = { - TransformExpression(BucketFunction, expr, Some(numBuckets)) + TransformExpression(BucketFunction, Seq(Literal(numBuckets)) ++ expr) } test("ShufflePartitionIdPassThrough - avoid unnecessary shuffle when children are compatible") {