You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/06 16:33:43 UTC
(datafusion-comet) 01/01: perf: skip Spark's encoders for boxed primitive Scala UDF values in the codegen dispatcher (#6714)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6714-6940c2e584b96ae9e6666228ab14d5d8bcc693f6
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 62c5d44921a25f823c53d8fc58566874fcda3c25
Author: Andy Grove <ag...@apache.org>
AuthorDate: Tue Oct 6 16:33:20 2026 +0000
perf: skip Spark's encoders for boxed primitive Scala UDF values in the codegen dispatcher (#6714)
* perf: skip Spark's encoders for boxed primitive Scala UDF values in the codegen dispatcher
A typed Scala UDF converts a boxed primitive parameter (`java.lang.Long`, ...) through its
input encoder and a boxed primitive result through its output encoder, each a projection per
row. The kernel now drops those encoders when their expression is exactly the boxing or
unboxing call, so `ScalaUDF` passes the value through `CatalystTypeConverters` unchanged, as
Spark 4 does for a Java UDF. Every other encoder is left alone.
Closes #6706.
* Compare boxed encoders with the trees Spark's builders produce
Check each encoder by building the expected tree with Spark's own
`DeserializerBuildHelper` and `SerializerBuildHelper` builders and comparing with `==`, in
place of field-by-field matching. Drop parameter encoders with `filterNot` like the result's,
and always copy the `ScalaUDF`, since an equal copy leaves its parents untouched.
In the tests, register the shared UDFs once, pass vector classes directly, drop two no-op
casts, and run the mixed-parameter calls as one query.
---
.../comet/codegen/CometBatchKernelCodegen.scala | 86 +++++++++++--
.../scala/org/apache/comet/CometCodegenSuite.scala | 133 ++++++++++++++++++++-
2 files changed, 211 insertions(+), 8 deletions(-)
diff --git a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala
index e478057cee..e2edb3c989 100644
--- a/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala
+++ b/spark/src/main/scala/org/apache/comet/codegen/CometBatchKernelCodegen.scala
@@ -24,7 +24,10 @@ import org.apache.arrow.vector._
import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector}
import org.apache.arrow.vector.types.pojo.Field
import org.apache.spark.internal.Logging
-import org.apache.spark.sql.catalyst.expressions.{BoundReference, Expression, Literal, Unevaluable}
+import org.apache.spark.sql.catalyst.DeserializerBuildHelper.createDeserializerForTypesSupportValueOf
+import org.apache.spark.sql.catalyst.SerializerBuildHelper.{createSerializerForBoolean, createSerializerForByte, createSerializerForDouble, createSerializerForFloat, createSerializerForInteger, createSerializerForLong, createSerializerForShort}
+import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
+import org.apache.spark.sql.catalyst.expressions.{BoundReference, Expression, Literal, ScalaUDF, Unevaluable}
import org.apache.spark.sql.catalyst.expressions.codegen._
import org.apache.spark.sql.execution.ExecSubqueryExpression
import org.apache.spark.sql.internal.SQLConf
@@ -47,7 +50,8 @@ import org.apache.comet.shims.{CometExprTraitShim, CometTypeShim}
* Input- and output-side emission live in [[CometBatchKernelCodegenInput]] and
* [[CometBatchKernelCodegenOutput]]. This file owns the [[ArrowColumnSpec]] vocabulary, the
* [[canHandle]] / [[allocateOutput]] / [[compile]] / [[generateSource]] entry points, and
- * cross-cutting kernel-shape decisions (NullIntolerant short-circuit, CSE variant).
+ * cross-cutting kernel-shape decisions (NullIntolerant short-circuit, CSE variant, boxed
+ * primitive `ScalaUDF` values without their encoders).
*
* The generated kernel is the `InternalRow` that Spark's `BoundReference.genCode` reads from. See
* [[generateSource]] for how the wiring is set up.
@@ -263,6 +267,7 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come
inputSchema: Seq[ArrowColumnSpec]): GeneratedSource = {
canHandle(boundExpr).foreach(reason =>
throw new IllegalArgumentException(s"CometBatchKernelCodegen.generateSource: $reason"))
+ val expr = withoutBoxedPrimitiveEncoders(boundExpr)
val ctx = new CodegenContext
// `BoundReference.genCode` emits `${ctx.INPUT_ROW}.getUTF8String(ord)`. Aliasing `row` to
// `this` at the top of `process` routes those reads to the kernel's typed getters (final
@@ -293,19 +298,19 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come
// `subexprFunctionsCode` is the concatenated helper invocation block, spliced into the
// per-row body by `defaultBody`.
val ev = if (SQLConf.get.subexpressionEliminationEnabled) {
- ctx.generateExpressions(Seq(boundExpr), doSubexpressionElimination = true).head
+ ctx.generateExpressions(Seq(expr), doSubexpressionElimination = true).head
} else {
- boundExpr.genCode(ctx)
+ expr.genCode(ctx)
}
val subExprsCode = ctx.subexprFunctionsCode
val (cls, setup, snippet) =
- CometBatchKernelCodegenOutput.emitOutputWriter(boundExpr.dataType, ev.value, ctx)
- (cls, setup, defaultBody(boundExpr, inputSchema, ev, snippet, subExprsCode))
+ CometBatchKernelCodegenOutput.emitOutputWriter(expr.dataType, ev.value, ctx)
+ (cls, setup, defaultBody(expr, inputSchema, ev, snippet, subExprsCode))
}
val typedFieldDecls = CometBatchKernelCodegenInput.emitInputFieldDecls(inputSchema)
val typedInputCasts = CometBatchKernelCodegenInput.emitInputCasts(inputSchema)
- val decimalTypeByOrdinal = CometBatchKernelCodegenInput.decimalPrecisionByOrdinal(boundExpr)
+ val decimalTypeByOrdinal = CometBatchKernelCodegenInput.decimalPrecisionByOrdinal(expr)
val getters =
CometBatchKernelCodegenInput.emitTypedGetters(inputSchema, decimalTypeByOrdinal)
val nested = CometBatchKernelCodegenInput.emitNestedClasses(inputSchema)
@@ -371,6 +376,73 @@ object CometBatchKernelCodegen extends Logging with CometExprTraitShim with Come
GeneratedSource(code.body, code, ctx.references.toArray)
}
+ /**
+ * Drops the encoder from each boxed primitive parameter and result of every `ScalaUDF` in the
+ * tree, so the kernel passes those values to and from the user function directly (#6706).
+ *
+ * A typed Scala UDF carries an `ExpressionEncoder` for each parameter and one for its result.
+ * For a boxed primitive (`(x: java.lang.Long) => ...`) the parameter's deserializer is
+ * `java.lang.Long.valueOf(input[0])`, which yields null for a null input, and the result's
+ * serializer is `input[0].longValue()`. `ScalaUDF` runs each one on every row through a
+ * projection: a `SafeProjection` into a `GenericInternalRow` going in, and an
+ * `UnsafeProjection` into an `UnsafeRow` coming out. Spark's whole-stage codegen runs the same
+ * projections.
+ *
+ * Without an encoder, `ScalaUDF` converts the value with `CatalystTypeConverters` instead,
+ * which passes a primitive through in its boxed form: `identity` going in, and coming out an
+ * `Option` unwrap that a boxed value never takes. The function receives the same value or null,
+ * and Catalyst gets the same result. Spark 4 builds a Java UDF the same way, with no encoders.
+ *
+ * Only an encoder whose expression is exactly the one Spark builds for a boxed primitive is
+ * dropped. Every other encoder (`String`, `Option`, case classes, collections) stays.
+ */
+ private[comet] def withoutBoxedPrimitiveEncoders(expr: Expression): Expression =
+ expr.transformUp { case udf: ScalaUDF =>
+ udf.copy(
+ inputEncoders = udf.inputEncoders.zip(udf.children).map { case (enc, child) =>
+ enc.filterNot(isBoxedPrimitiveDeserializer(_, child.dataType))
+ },
+ outputEncoder = udf.outputEncoder.filterNot(isBoxedPrimitiveSerializer(_, udf.dataType)))
+ }
+
+ /**
+ * Boxed class of each type `CatalystTypeConverters` passes through, and Spark's builder for the
+ * serializer of that class's encoder.
+ */
+ private def boxedPrimitive(dt: DataType): Option[(Class[_], Expression => Expression)] =
+ dt match {
+ case BooleanType => Some((classOf[java.lang.Boolean], createSerializerForBoolean))
+ case ByteType => Some((classOf[java.lang.Byte], createSerializerForByte))
+ case ShortType => Some((classOf[java.lang.Short], createSerializerForShort))
+ case IntegerType => Some((classOf[java.lang.Integer], createSerializerForInteger))
+ case LongType => Some((classOf[java.lang.Long], createSerializerForLong))
+ case FloatType => Some((classOf[java.lang.Float], createSerializerForFloat))
+ case DoubleType => Some((classOf[java.lang.Double], createSerializerForDouble))
+ case _ => None
+ }
+
+ /**
+ * True iff `enc` deserializes a `dt` value exactly as the encoder of its boxed class does, with
+ * `valueOf` over a nullable reference. The reference must be nullable because `StaticInvoke`
+ * propagates null only from a nullable argument: over a non-nullable one, Spark would box the 0
+ * that a row holds for a null.
+ */
+ private def isBoxedPrimitiveDeserializer(enc: ExpressionEncoder[_], dt: DataType): Boolean =
+ boxedPrimitive(dt).exists { case (boxed, _) =>
+ enc.objDeserializer ==
+ createDeserializerForTypesSupportValueOf(BoundReference(0, dt, nullable = true), boxed)
+ }
+
+ /**
+ * True iff `enc` serializes a boxed `dt` value exactly as the encoder of the boxed class does,
+ * by unboxing it. `ScalaUDF` turns a null result into null before it calls the serializer, so
+ * the serializer only ever sees a boxed value.
+ */
+ private def isBoxedPrimitiveSerializer(enc: ExpressionEncoder[_], dt: DataType): Boolean =
+ boxedPrimitive(dt).exists { case (boxed, serializerFor) =>
+ enc.objSerializer == serializerFor(BoundReference(0, ObjectType(boxed), nullable = true))
+ }
+
/**
* Per-row body. For `NullIntolerant` expressions whose input nulls fully determine a null
* output (see [[canShortCircuitNulls]]), prepends a short-circuit on those ordinals so the
diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
index e9105230d7..8f6d7c4a8d 100644
--- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala
@@ -28,7 +28,7 @@ import org.apache.spark.{SparkConf, SparkEnv, TaskContext}
import org.apache.spark.sql.{CometTestBase, Row}
import org.apache.spark.sql.api.java.UDF1
import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{Add, Alias, ApplyFunctionExpression, AttributeReference, BoundReference, Cast, CreateArray, CreateMap, CreateNamedStruct, Expression, GenericInternalRow, Hypot, Literal, MapConcat, ScalaUDF}
+import org.apache.spark.sql.catalyst.expressions.{Add, Alias, ApplyFunctionExpression, AttributeReference, AttributeSeq, BindReferences, BoundReference, Cast, CreateArray, CreateMap, CreateNamedStruct, Expression, GenericInternalRow, Hypot, Literal, MapConcat, ScalaUDF}
import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial}
import org.apache.spark.sql.catalyst.expressions.objects.{Invoke, StaticInvoke}
import org.apache.spark.sql.catalyst.util.GenericArrayData
@@ -1304,6 +1304,137 @@ class CometCodegenSuite
}
}
+ private val boxedPrimitiveTypes = Seq("bool", "byte", "short", "int", "long", "float", "double")
+
+ /**
+ * Boxed primitive UDFs for the #6706 tests, where the kernel drops Spark's encoders. Each
+ * `<type>_str` function names the exact value it receives, null included, so a value converted
+ * differently from Spark's encoder (a null read as 0, a -0.0 read as 0.0) fails the comparison
+ * with Spark. Each `<type>_id` function returns the boxed value it receives, which covers the
+ * result side. `long_tag` and `opt_long` pair a boxed parameter with ones that keep their
+ * encoders.
+ */
+ private def registerBoxedPrimitiveUdfs(): Unit = {
+ spark.udf.register("bool_str", (x: java.lang.Boolean) => String.valueOf(x))
+ spark.udf.register("byte_str", (x: java.lang.Byte) => String.valueOf(x))
+ spark.udf.register("short_str", (x: java.lang.Short) => String.valueOf(x))
+ spark.udf.register("int_str", (x: java.lang.Integer) => String.valueOf(x))
+ spark.udf.register("long_str", (x: java.lang.Long) => String.valueOf(x))
+ // Raw bits, so the sign of a zero or a NaN payload would show.
+ spark.udf.register(
+ "float_str",
+ (x: java.lang.Float) =>
+ if (x == null) "null" else java.lang.Float.floatToRawIntBits(x).toString)
+ spark.udf.register(
+ "double_str",
+ (x: java.lang.Double) =>
+ if (x == null) "null" else java.lang.Double.doubleToRawLongBits(x).toString)
+ spark.udf.register("bool_id", (x: java.lang.Boolean) => x)
+ spark.udf.register("byte_id", (x: java.lang.Byte) => x)
+ spark.udf.register("short_id", (x: java.lang.Short) => x)
+ spark.udf.register("int_id", (x: java.lang.Integer) => x)
+ spark.udf.register("long_id", (x: java.lang.Long) => x)
+ spark.udf.register("float_id", (x: java.lang.Float) => x)
+ spark.udf.register("double_id", (x: java.lang.Double) => x)
+ spark.udf.register("long_tag", (x: java.lang.Long, s: String) => s"$x:$s")
+ spark.udf.register("opt_long", (x: Option[Long]) => x.map(_ + 1).getOrElse(-1L))
+ }
+
+ private def withBoxedPrimitiveTable(f: => Unit): Unit = {
+ withTable("t") {
+ sql(
+ "CREATE TABLE t (c_bool BOOLEAN, c_byte TINYINT, c_short SMALLINT, c_int INT, " +
+ "c_long BIGINT, c_float FLOAT, c_double DOUBLE, s STRING) USING parquet")
+ sql(
+ "INSERT INTO t VALUES " +
+ "(true, CAST(-128 AS TINYINT), CAST(-32768 AS SMALLINT), -2147483648, " +
+ "-9223372036854775808, CAST('-0.0' AS FLOAT), CAST('-0.0' AS DOUBLE), 'a'), " +
+ "(false, CAST(127 AS TINYINT), CAST(32767 AS SMALLINT), 2147483647, " +
+ "9223372036854775807, CAST('NaN' AS FLOAT), CAST('NaN' AS DOUBLE), NULL), " +
+ "(NULL, NULL, NULL, NULL, NULL, NULL, NULL, 'c'), " +
+ "(true, CAST(0 AS TINYINT), CAST(1000 AS SMALLINT), 1000, 1000, " +
+ "CAST('Infinity' AS FLOAT), CAST(0.0 AS DOUBLE), 'd'), " +
+ "(false, CAST(-1 AS TINYINT), CAST(-1000 AS SMALLINT), -1000, -1000, " +
+ "CAST(1.5 AS FLOAT), CAST('-Infinity' AS DOUBLE), NULL)")
+ f
+ }
+ }
+
+ test("the kernel drops only boxed primitive encoders from a ScalaUDF (#6706)") {
+ registerBoxedPrimitiveUdfs()
+ spark.udf.register("str_id", (x: String) => x)
+ spark.udf.register("long_in_prim_out", (x: java.lang.Long) => if (x == null) -1L else x + 1L)
+ spark.udf.register("prim_add_one", (x: Long) => x + 1)
+ withBoxedPrimitiveTable {
+ def boundUdf(call: String): ScalaUDF = {
+ val udf = sql(s"SELECT $call FROM t").queryExecution.optimizedPlan.expressions
+ .flatMap(_.collect { case u: ScalaUDF => u })
+ .head
+ val attrs = udf.collect { case a: AttributeReference => a }.distinct
+ BindReferences.bindReference(udf, AttributeSeq(attrs))
+ }
+ // Whether the kernel drops the encoder of each parameter and then of the result.
+ def dropped(call: String): Seq[Boolean] = {
+ val udf = boundUdf(call)
+ val kernelUdf =
+ CometBatchKernelCodegen.withoutBoxedPrimitiveEncoders(udf).asInstanceOf[ScalaUDF]
+ (udf.inputEncoders :+ udf.outputEncoder)
+ .zip(kernelUdf.inputEncoders :+ kernelUdf.outputEncoder)
+ .map { case (before, after) => before.isDefined && after.isEmpty }
+ }
+ boxedPrimitiveTypes.foreach { t =>
+ assert(dropped(s"${t}_id(c_$t)") === Seq(true, true), t)
+ }
+ assert(dropped("long_in_prim_out(c_long)") === Seq(true, false))
+ assert(dropped("long_tag(c_long, s)") === Seq(true, false, false))
+ assert(dropped("prim_add_one(c_long)") === Seq(false, false))
+ assert(dropped("opt_long(c_long)") === Seq(false, false))
+ assert(dropped("str_id(s)") === Seq(false, false))
+
+ // Spark's encoder path calls the parameter's converter even for a null input. Without the
+ // encoder, `ScalaUDF` tests for null itself, which shows the kernel compiles the rewrite.
+ def kernelSource(call: String, vectorClass: Class[_ <: ValueVector]): String =
+ CometBatchKernelCodegen
+ .generateSource(
+ boundUdf(call),
+ IndexedSeq(ArrowColumnSpec(vectorClass, nullable = true)))
+ .body
+ assert(kernelSource("str_id(s)", classOf[VarCharVector]).contains(".apply(null)"))
+ assert(!kernelSource("long_id(c_long)", classOf[BigIntVector]).contains(".apply(null)"))
+ }
+ }
+
+ test("boxed primitive UDF parameters and results match Spark (#6706)") {
+ registerBoxedPrimitiveUdfs()
+ withBoxedPrimitiveTable {
+ val calls = boxedPrimitiveTypes.flatMap(t => Seq(s"${t}_str(c_$t)", s"${t}_id(c_$t)"))
+ assertCodegenRan {
+ checkSparkAnswerAndOperator(sql(s"SELECT ${calls.mkString(", ")} FROM t"))
+ }
+ }
+ }
+
+ test("boxed primitive UDF parameters next to other parameters match Spark (#6706)") {
+ registerBoxedPrimitiveUdfs()
+ // Spark guards the call with a null check for the primitive parameter, not the boxed one.
+ spark.udf.register("long_or", (x: java.lang.Long, y: Long) => if (x == null) y else x + y)
+ withBoxedPrimitiveTable {
+ val calls = Seq(
+ "long_tag(c_long, s)",
+ "opt_long(c_long)",
+ "long_or(c_long, c_int)",
+ // Spark casts the int column to the `java.lang.Long` parameter's type.
+ "long_str(c_int)",
+ "long_id(long_id(c_long))",
+ // Inside a lambda, Spark evaluates the UDF through `eval` rather than generated code,
+ // with the same converters.
+ "transform(array(c_int, NULL), x -> int_str(x))")
+ assertCodegenRan {
+ checkSparkAnswerAndOperator(sql(s"SELECT ${calls.mkString(", ")} FROM t"))
+ }
+ }
+ }
+
test("ScalaUDF returning a different type than its input") {
// String -> Int output transition. Identity-loop above keeps input == output. This asserts
// the writer can switch types per the UDF's declared return.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org