andygrove commented on code in PR #6455:
URL: https://github.com/apache/datafusion-comet/pull/6455#discussion_r4162151865
##########
spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegenOutput.scala:
##########
@@ -231,15 +231,39 @@ 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.
+ //
+ // 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);
Review Comment:
@sunchao @parthchandra Fixed the transitive case in
5d07da30e347909c38e288a951e839fdceaecb9e. Both checks now search the full
argument trees: scalar consumers run with the DSv2 decimal call in one kernel,
and aggregates containing such a call fall back to Spark. This keeps the
intermediate decimal from being normalized before its consumer reads it.
The regression covers nested `abs`, instance/static calls, array/struct
access, nullness and string consumers, filtering, and `count`/`max`/`sum`, in
both ANSI modes. It also checks dispatcher-disabled fallback and a parquet
write/read boundary, where normalization really should occur. The new
regression failed before the correction on both branches while the five
existing decimal tests passed.
Local validation: source Spark 4.1 codegen suites 170/170; branch-1.1 Spark
4.1 169/169; branch-1.1 Spark 3.5 targeted decimal/writer tests 66/66. Native
builds and release `cargo check --all-targets` passed. The tested backport is
66ec6969645276d9d390a8bcfecc89d805b1a07c in #6490. Fresh CI is running; full
Spark SQL CI is not yet verified.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
