fix: rescale decimal results from dispatched functions to the declared type

spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala

@@ -231,15 +231,43 @@ private[codegen] object CometBatchKernelCodegenOutput extends CometTypeShim {
       val set = if (nested) "setSafe" else "set"
       OutputEmit("", s"$targetVec.$set($idx, $source);")
     case dt: DecimalType =>
+      // Rescale to the declared type, and write null when the value does not fit, as Spark's
+      // `UnsafeRowWriter` and `UnsafeArrayWriter` do in `write(ordinal, Decimal, precision,
+      // scale)`. A Spark expression already produces its declared precision and scale, but a
+      // DSv2 function called through `Invoke` / `StaticInvoke` can return a `Decimal` of any
+      // scale (#6425). Like Spark's writers, this rescales the value in place, and leaves it
+      // untouched when it does not fit. Unlike them, it does not test `source` for null: the
+      // callers write null values themselves, and skip that test only for a type that is not
+      // nullable.
+      //
+      // Only the kernel's own output sees the null. Spark rescales such a value when it writes a
+      // row, and an expression around the call reads the value the function returned, so
+      // `QueryPlanSerde` dispatches that expression with the call, and falls an aggregate back.
+      //
+      // The precision and scale test repeats `changePrecision`'s own fast path. It keeps the call
+      // off the common path, so the JIT can still scalar-replace the `Decimal` that an input
+      // getter allocates. With the bare call, passing a `DECIMAL(18, 2)` column through took
+      // about half as long again per row.
+      //
       // DecimalOutputShortFastPath: precision <= 18 fits in a signed long, so pass the unscaled
       // value to `setSafe(int, long)` and skip the BigDecimal allocation.
+      val dec = ctx.freshName("dec")
+      val (precision, scale) = (dt.precision, dt.scale)
       val write =
-        if (dt.precision <= Decimal.MAX_LONG_DIGITS) {
-          s"$targetVec.setSafe($idx, $source.toUnscaledLong());"
+        if (precision <= Decimal.MAX_LONG_DIGITS) {
+          s"$targetVec.setSafe($idx, $dec.toUnscaledLong());"
         } else {
-          s"$targetVec.setSafe($idx, $source.toJavaBigDecimal());"
+          s"$targetVec.setSafe($idx, $dec.toJavaBigDecimal());"
         }
-      OutputEmit("", write)
+      OutputEmit(
+        "",
+        s"""org.apache.spark.sql.types.Decimal $dec = $source;
+           |if (($dec.precision() == $precision && $dec.scale() == $scale) ||
+           |    $dec.changePrecision($precision, $scale)) {
+           |  $write
+           |} else {
+           |  $targetVec.setNull($idx);
+           |}""".stripMargin)
     case _: StringType =>
       // Utf8OutputOnHeapShortcut: when the UTF8String is on-heap (Spark's string functions
       // allocate results on-heap), pass its backing byte[] directly to `setSafe`, skipping the

spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala

@@ -22,7 +22,7 @@ package org.apache.comet.rules
 import scala.collection.mutable.ListBuffer
 
 import org.apache.spark.sql.SparkSession
-import org.apache.spark.sql.catalyst.expressions.{Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, NamedExpression, Remainder}
+import org.apache.spark.sql.catalyst.expressions.{Divide, DoubleLiteral, EqualNullSafe, EqualTo, Expression, ExprId, FloatLiteral, GreaterThan, GreaterThanOrEqual, KnownFloatingPointNormalized, LessThan, LessThanOrEqual, NamedExpression, Remainder}
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateMode, Final, Partial, PartialMerge}
 import org.apache.spark.sql.catalyst.optimizer.NormalizeNaNAndZero
 import org.apache.spark.sql.catalyst.rules.Rule
@@ -79,6 +79,10 @@ object CometExecRule {
   val COMET_UNSAFE_PARTIAL: TreeNodeTag[String] =
     TreeNodeTag[String]("comet.unsafePartialAgg")
 
+  /** Keep a DSv2 decimal producer and its consumers within Spark's row-materialization path. */
+  val COMET_UNMATERIALIZED_DECIMAL: TreeNodeTag[Unit] =
+    TreeNodeTag[Unit]("comet.unmaterializedDecimal")
+
   /**
    * Fully native operators.
    */
@@ -351,6 +355,12 @@ case class CometExecRule(session: SparkSession)
   // spotless:on
   private def transform(plan: SparkPlan): SparkPlan = {
     def convertNode(op: SparkPlan): SparkPlan = op match {
+      case op if op.getTagValue(CometExecRule.COMET_UNMATERIALIZED_DECIMAL).isDefined =>
+        withFallbackReason(
+          op,
+          "DSv2 decimal result and its consumers must share Spark materialization")
+        op
+
       // Scan marker produced by an optional, out-of-tree scan contrib (e.g. contrib/delta).
       // Matched by trait (no compile-time dependency on the contrib) and present only when that
       // contrib is on the classpath. The marker carries its own serde handler and typically wraps
@@ -704,6 +714,68 @@ case class CometExecRule(session: SparkSession)
     }
   }
 
+  /**
+   * Spark can pass a Decimal between fused operators without writing a row. An operator that
+   * defines an attribute from a dispatched DSv2 decimal call therefore cannot write it to Arrow
+   * before its consumers run: that would rescale/null the value early. The operator can be a
+   * retained Project, a generator (`explode` passes the array's elements on as they are), an
+   * Expand, or an aggregate whose result expression makes the call. Follow those attributes and
+   * keep both ends, including intervening operators, in Spark. Declining only the consumer would
+   * leave the early write.
+   *
+   * Paths are propagated conservatively through decimal outputs, including row boundaries. This
+   * may retain more Spark operators after a real materialization, but never introduces a new one
+   * before a consumer. Unrelated inputs, and trees that consume the call within one expression,
+   * remain eligible for native execution.
+   */
+  private def tagUnmaterializedDecimals(plan: SparkPlan): Unit = {
+    type Paths = Map[ExprId, Set[SparkPlan]]
+    def visit(op: SparkPlan): Paths = {
+      val inputs = op.children.flatMap(visit).groupBy(_._1).map { case (id, paths) =>
+        id -> paths.flatMap(_._2).toSet
+      }
+      val consumed = op.references.toSeq.flatMap(a => inputs.getOrElse(a.exprId, Set.empty)).toSet
+      if (consumed.nonEmpty) {
+        (consumed + op).foreach(_.setTagValue(CometExecRule.COMET_UNMATERIALIZED_DECIMAL, ()))
+      }
+      val sources = op match {
+        case p: ProjectExec =>
+          p.projectList
+            .filter(QueryPlanSerde.producesUnmaterializedDsv2Decimal)
+            .map(_.exprId)
+            .toSet
+        // A result expression such as `fn(max(i))`, which HAVING can read. QueryPlanSerde already
+        // falls back an aggregate whose function makes the call.
+        case a: BaseAggregateExec =>
+          a.resultExpressions
+            .filter(QueryPlanSerde.producesUnmaterializedDsv2Decimal)
+            .map(_.exprId)
+            .toSet
+        // Other operators, such as a generator or an Expand, do not tie each new attribute to one
+        // expression, so take every decimal attribute that one defines. A Comet operator from an
+        // earlier pass is already decided.
+        case _
+            if !op.isInstanceOf[CometPlan] &&
+              op.expressions.exists(QueryPlanSerde.evaluatesDispatchedDsv2Decimal) =>
+          op.output
+            .filter(a => !op.inputSet.contains(a) && QueryPlanSerde.containsDecimal(a.dataType))
+            .map(_.exprId)
+            .toSet
+        case _ => Set.empty[ExprId]
+      }
+      op.output.flatMap { attr =>
+        val inherited = inputs.getOrElse(attr.exprId, Set.empty)
+        // Computed decimal outputs can retain the raw value as well (e.g. abs(d), or max(d)).
+        val derived =
+          if (QueryPlanSerde.containsDecimal(attr.dataType)) consumed else Set.empty[SparkPlan]
+        val path = inherited ++ derived
+        if (sources.contains(attr.exprId) || path.nonEmpty) Some(attr.exprId -> (path + op))
+        else None
+      }.toMap
+    }
+    val _ = visit(plan)
+  }
+
   private def normalizePlan(plan: SparkPlan): SparkPlan = {
     plan.transformUp {
       case p: ProjectExec =>
@@ -797,6 +869,7 @@ case class CometExecRule(session: SparkSession)
       // corresponding Final or PartialMerge cannot be converted and the intermediate buffer
       // formats are incompatible. This runs before transform() so the tags are checked
       // during the bottom-up conversion. Tags persist through AQE stage creation.
+      tagUnmaterializedDecimals(planWithJoinRewritten)
       tagUnsafePartialAggregates(planWithJoinRewritten)
 
       var newPlan = revertUnsafePartialAggregates(transform(planWithJoinRewritten))
@@ -1370,6 +1443,8 @@ case class CometExecRule(session: SparkSession)
     if (handler.isEmpty) return false
     val serde = handler.get.asInstanceOf[CometOperatorSerde[SparkPlan]]
     if (!isOperatorEnabled(serde, agg.asInstanceOf[SparkPlan])) return false
+    // tagUnmaterializedDecimals runs first, and transform() leaves a tagged operator in Spark.
+    if (agg.getTagValue(CometExecRule.COMET_UNMATERIALIZED_DECIMAL).isDefined) return false
 
     val aggregateExpressions = agg.aggregateExpressions
     val groupingExpressions = agg.groupingExpressions

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -31,6 +31,7 @@ import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke}
 import org.apache.spark.sql.catalyst.expressions.xml.{XPathBoolean, XPathDouble, XPathFloat, XPathInt, XPathList, XPathLong, XPathShort, XPathString}
 import org.apache.spark.sql.comet.DecimalPrecision
+import org.apache.spark.sql.connector.catalog.functions.ScalarFunction
 import org.apache.spark.sql.execution.{ScalarSubquery, SparkPlan}
 import org.apache.spark.sql.execution.datasources.parquet.ParquetUtils
 import org.apache.spark.sql.internal.SQLConf
@@ -836,6 +837,17 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
     }
 
     val fn = aggExpr.aggregateFunction
+    // An aggregate reads the `Decimal` a DSv2 function returned, whether or not it fits the
+    // declared type, and the kernel has nulled such a value at its own output: Spark counts it
+    // and takes it as the maximum, and the row writer nulls the result afterwards. An aggregate
+    // cannot run in the kernel, so it falls back. See [[readsDispatchedDsv2Decimal]].
+    if (fn.children.exists(_.exists(isDispatchedDsv2DecimalCall))) {
+      withFallbackReason(
+        aggExpr,
+        s"${fn.prettyName} aggregates the decimal result of a DSv2 function, which Spark " +
+          "rescales to the declared type only when it writes a row")
+      return None
+    }
     val cometExpr = aggrSerdeMap.get(fn.getClass)
     val protoAggExprOpt = cometExpr match {
       case Some(handler) =>
@@ -1032,6 +1044,12 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
             s"${CometConf.getExprEnabledConfigKey(exprConfName)}=true to enable it.")
         return None
       }
+      if (readsDispatchedDsv2Decimal(expr)) {
+        // Returns None, tagged with the reason, when the dispatcher declines, which falls the
+        // operator back to Spark. Converting `expr` natively is not an alternative: it would read
+        // the call's output rather than the value the function returned.
+        return CometScalaUDF.emitJvmCodegenDispatch(expr, inputs, binding)
+      }
       handler.getSupportLevel(expr) match {
         case Unsupported(notes) =>
           // `CodegenDispatchFallback` serdes have no native path for these cases either, but the
@@ -1141,6 +1159,58 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim {
     case _ => false
   }
 
+  /**
+   * Whether `expr` consumes a decimal result of a dispatched DSv2 scalar function anywhere in its
+   * argument trees, including through intermediate expressions and container access.
+   *
+   * Spark does not rescale such a result to the type the function declares, or write null when it
+   * does not fit, until it writes a row. An expression around the call reads the `Decimal` the
+   * function returned. The dispatcher has to write an Arrow vector of the declared type, so it
+   * rescales and nulls at its own output, and a native expression over that output would read
+   * something else. `IS NULL` of a value that does not fit is false in Spark, a cast to string
+   * keeps the function's scale, and `hash` reads the unscaled value at that scale (#6425). So
+   * such an expression runs in the same kernel as the call, where Spark's own code reads the
+   * value the function returned. Checking only immediate children would let an intermediate
+   * expression, such as `abs(call)` or `call[0]`, normalize the decimal before its parent reads
+   * it. `Alias` is skipped because it computes nothing. An operator's output is not necessarily a
+   * row-materialization boundary: CometExecRule keeps an operator that defines a decimal
+   * attribute from such a call, such as a Project, a generator or an aggregate's result, in Spark
+   * with that attribute's consumers until they have read the original value.
+   */
+  private def readsDispatchedDsv2Decimal(expr: Expression): Boolean =
+    !isStructuralExpr(expr) && expr.children.exists(_.exists(isDispatchedDsv2DecimalCall))
+
+  /** Whether `expr` produces a decimal that holds the value a dispatched DSv2 call returned. */
+  def producesUnmaterializedDsv2Decimal(expr: Expression): Boolean =
+    containsDecimal(expr.dataType) && evaluatesDispatchedDsv2Decimal(expr)
+
+  /** Whether `expr` calls a dispatched DSv2 scalar function that returns a decimal. */
+  private[comet] def evaluatesDispatchedDsv2Decimal(expr: Expression): Boolean =
+    expr.exists(isDispatchedDsv2DecimalCall)
+
+  private def isDispatchedDsv2DecimalCall(expr: Expression): Boolean = {
+    val dispatchedDsv2Call = expr match {
+      case i: Invoke =>
+        i.targetObject match {
+          case Literal(_: ScalarFunction[_], _) => true
+          case _ => false
+        }
+      case s: StaticInvoke =>
+        classOf[ScalarFunction[_]].isAssignableFrom(s.staticObject) &&
+        CometStaticInvoke.runsInDispatcher(s)
+      case _ => false
+    }
+    dispatchedDsv2Call && containsDecimal(expr.dataType)
+  }
+
+  private[comet] def containsDecimal(dataType: DataType): Boolean = dataType match {
+    case _: DecimalType => true
+    case ArrayType(elementType, _) => containsDecimal(elementType)
+    case MapType(keyType, valueType, _) => containsDecimal(keyType) || containsDecimal(valueType)
+    case StructType(fields) => fields.exists(f => containsDecimal(f.dataType))
+    case _ => false
+  }
+
   /**
    * Creates a UnaryExpr by calling exprToProtoInternal for the provided child expression and then
    * invokes the supplied function to wrap this UnaryExpr in a top-level Expr.

spark/src/main/scala/org/apache/comet/serde/statics.scala

@@ -58,6 +58,10 @@ object CometStaticInvoke extends CometExpressionSerde[StaticInvoke] {
   private def handlerFor(expr: StaticInvoke): Option[CometExpressionSerde[StaticInvoke]] =
     staticInvokeExpressions.get((expr.functionName, expr.staticObject.getName))
 
+  /** Whether [[convert]] hands `expr` to the codegen dispatcher, because no handler claims it. */
+  def runsInDispatcher(expr: StaticInvoke): Boolean =
+    handlerFor(expr).isEmpty && icebergHandlerFor(expr).isEmpty
+
   /**
    * Iceberg's system functions are keyed by class name only because both the `StaticInvoke` and
    * `ApplyFunctionExpression` lowerings share the same identity class. Consulted after the
@@ -80,11 +84,13 @@ object CometStaticInvoke extends CometExpressionSerde[StaticInvoke] {
    * [[CodegenDispatchFallback]]: that would also route a *handler's* `Unsupported` through the
    * dispatcher, and at least one of those is not dispatchable. `CometIcebergTruncate` declines a
    * decimal because Iceberg's `truncate` can return a value wider than the column's declared
-   * precision, which Spark nulls only when the row is materialized; the dispatcher writes into an
-   * Arrow `Decimal128(precision, scale)` vector just like a native kernel does, so it produces
-   * the out-of-range value instead of a null. The mixin's contract ("the case must be something
-   * `doGenCode` can compile") does not cover a limit that lives at the Arrow output boundary, so
-   * enrollment stays with the individual handlers.
+   * precision, which Spark nulls only when the row is materialized. The dispatcher nulls such a
+   * value at its own output, as the row writer does, and an expression or aggregate that reads a
+   * dispatched call either runs in the same kernel or falls back (see
+   * `QueryPlanSerde.readsDispatchedDsv2Decimal`), but the decline predates that and stays. The
+   * mixin's contract ("the case must be something `doGenCode` can compile") does not cover a
+   * limit that lives at the Arrow output boundary, so enrollment stays with the individual
+   * handlers.
    */
   override def getSupportLevel(expr: StaticInvoke): SupportLevel =
     handlerFor(expr)

spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala

@@ -616,6 +616,13 @@ object CometShuffleExchangeExec
           if (QueryPlanSerde.exprToProto(expr, inputs).isEmpty) {
             reasons += s"unsupported hash partitioning expression: $expr"
           }
+          // Spark hashes the `Decimal` a DSv2 function returned, at the scale the function chose,
+          // while the native shuffle would hash the dispatcher's rescaled output, and so place
+          // rows in other partitions (#6425). The JVM columnar shuffle computes the partition with
+          // Spark's own projection.
+          if (QueryPlanSerde.producesUnmaterializedDsv2Decimal(expr)) {
+            reasons += s"hash partitioning expression reads a DSv2 decimal result: $expr"
+          }
         }
         for (dt <- expressions.map(_.dataType).distinct) {
           if (!supportedHashPartitioningDataType(dt)) {

spark/src/test/scala/org/apache/comet/CometCodegenSourceSuite.scala

@@ -401,6 +401,10 @@ class CometCodegenSourceSuite extends AnyFunSuite {
       !result.body.contains(".toJavaBigDecimal("),
       "expected no BigDecimal allocation for short-precision output; got:\n" +
         CodeFormatter.format(result.code))
+    // Rescaled to the declared type first, as Spark's row writer does (#6425).
+    assert(
+      result.body.contains(".changePrecision(18, 2)"),
+      s"expected changePrecision before the write; got:\n${CodeFormatter.format(result.code)}")
   }
 
   test("DecimalVector setSafe uses BigDecimal slow path for long-precision output") {
@@ -417,6 +421,9 @@ class CometCodegenSourceSuite extends AnyFunSuite {
       !result.body.contains(".toUnscaledLong()"),
       "expected no unscaled-long write for long-precision output; got:\n" +
         CodeFormatter.format(result.code))
+    assert(
+      result.body.contains(".changePrecision(38, 10)"),
+      s"expected changePrecision before the write; got:\n${CodeFormatter.format(result.code)}")
   }
 
   test("VarCharVector setSafe uses on-heap UTF8String shortcut") {