[SPARK-50593][SQL] SPJ: Support truncate transform via generalized ReducibleFunction API

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<I, O> {
 
+  /**
+   * Generic reducer for parameterized functions (bucket, truncate, etc.).
+   *
+   * If this function is 'reducible' on another function, return the {@link Reducer}.
+   * <p>
+   * Each parameter is a non-complex {@link Literal} carrying both its value and data type:
+   * array/map/struct/UDT-typed values are filtered out by Spark and not passed here, but other
+   * scalar values (e.g. bucket numBuckets, truncate width, or a
+   * {@code CalendarInterval}) may be. {@link Literal#value()} is Spark's internal representation
+   * (e.g. {@code UTF8String} for strings, {@code Decimal} for decimals); use
+   * {@link Literal#dataType()} to interpret it rather than assuming a JVM type.
+   * <p>
+   * {@code thisParams} and {@code otherParams} hold each side's own literal parameters and may have
+   * different lengths -- for example a zero-parameter transform reducing onto a one-parameter one.
+   * Implementations must check each array's length before indexing into it.
+   * <p>
+   * Returning {@code null} means "not reducible for these parameters" 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).
+   * <p>
+   * Examples:
+   * <ul>
+   *     <li>bucket(4, x) and bucket(2, x): thisParams = [4], otherParams = [2]</li>
+   *     <li>truncate(x, 3) and truncate(x, 5): thisParams = [3], otherParams = [5]</li>
+   *     <li>hypothetical range_bucket(x, 0L, 100L, 4): thisParams = [0L, 100L, 4]</li>
+   * </ul>
+   *
+   * @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<I, O> reducer(
+          Literal<?>[] thisParams,
+          ReducibleFunction<?, ?> otherFunction,
+          Literal<?>[] otherParams) {
+    throw new UnsupportedOperationException();
+  }
+
   /**
    * This method is for the bucket function.
    *
@@ -78,7 +125,12 @@ public interface ReducibleFunction<I, O> {
    * @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<I, O> reducer(
       int thisNumBuckets,
       ReducibleFunction<?, ?> otherBucketFunction,
@@ -101,6 +153,6 @@ default Reducer<I, O> reducer(
    * @return a reduction function if it is reducible, null if not.
    */
   default Reducer<I, O> reducer(ReducibleFunction<?, ?> otherFunction) {
-    throw new UnsupportedOperationException();
+    return reducer(new Literal<?>[0], otherFunction, new Literal<?>[0]);
   }
 }

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 =>

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)

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)
   }