From 9c4b983064076fe512e0e61aad79dcc288e243eb Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Mon, 31 Aug 2026 10:11:29 -0600 Subject: [PATCH] chore: report the codegen-dispatch gate from getSupportLevel Serdes whose only path is the JVM codegen dispatcher reported Compatible from getSupportLevel and then discovered inside convert that the dispatcher would not run the expression, either because spark.comet.exec.scalaUDF.codegen.enabled is off or because CometBatchKernelCodegen.canHandle rejected the bound tree. That breaks the serde invariant, and it is the one shape the dispatcher cannot see: dispatchIfFallback is reached only from the Unsupported and Incompatible arms of exprToProtoInternal. Extract the two checks into CometScalaUDF.dispatchSupportLevel and report them from getSupportLevel on CometCodegenDispatch (62 serdes inherit it) and CometScalaUDF. Behaviour is unchanged: CometCodegenDispatch does not mix in CodegenDispatchFallback, so an Unsupported result tags the same reason text and falls the operator back to Spark exactly as before. Six CometCodegenDispatch subclasses override getSupportLevel (CometGetJsonObject, CometLengthOfJsonArray, CometStructsToJson, CometJsonToStructs, CometMakeTimestamp, CometToUnixTimestamp) and are not covered here, since each has to compose the gate with its own native opt-in branch. Same for the serdes that call emitJvmCodegenDispatch inline from convert. Both are follow-ups. Related to #5574. --- .../apache/comet/serde/CometScalaUDF.scala | 89 ++++++++++++++----- .../org/apache/comet/CometCodegenSuite.scala | 33 ++++++- 2 files changed, 100 insertions(+), 22 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala index 26fe0c591ab..853245a61ad 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala @@ -53,9 +53,61 @@ import org.apache.comet.udf.codegen.CometScalaUDFCodegen */ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { + override def getSupportLevel(expr: ScalaUDF): SupportLevel = dispatchSupportLevel(expr) + override def convert(expr: ScalaUDF, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = emitJvmCodegenDispatch(expr, inputs, binding) + /** + * Bind `expr` the way [[emitJvmCodegenDispatch]] will. + * + * `RuntimeReplaceable` expressions (e.g. Spark 4's `StructsToJson`) have a `doGenCode` that + * always throws "Cannot generate code for expression". Catalyst's `ReplaceExpressions` rule + * normally rewrites them to their `replacement` form before codegen runs. Comet's serde + * sometimes works with the pre-rewrite form (via shim reconstruction) for matching purposes, so + * unwrap to the replacement here before binding so the kernel compiles. + * + * Binding is against only the `AttributeReference`s the tree actually reads, so ordinals align + * with the data args shipped alongside the closure. Those attributes are returned too, since + * [[emitJvmCodegenDispatch]] needs them in the same order to build the data args. + */ + private def bindForDispatch(expr: Expression): (Expression, Seq[AttributeReference]) = { + val target = expr match { + case rr: RuntimeReplaceable => rr.replacement + case other => other + } + val attrs = target.collect { case a: AttributeReference => a }.distinct + (BindReferences.bindReference(target, AttributeSeq(attrs)), attrs) + } + + /** + * `SupportLevel` for a serde whose only path is the codegen dispatcher. `Compatible` when the + * dispatcher will accept the expression, `Unsupported` (with the same reason + * [[emitJvmCodegenDispatch]] would have tagged) when it will not. + * + * Reporting this from `getSupportLevel` rather than discovering it inside `convert` keeps the + * serde invariant intact and lets `exprToProtoInternal` handle the decline on its normal + * `Unsupported` path. Behaviour is unchanged: `CometCodegenDispatch` does not mix in + * `CodegenDispatchFallback`, so an `Unsupported` result still tags the reason and falls the + * operator back to Spark, exactly as the `convert`-side decline did. + * + * This is a pure predicate -- it records no fallback reason of its own, because the + * `Unsupported` arm of `exprToProtoInternal` already tags the notes returned here. + */ + def dispatchSupportLevel(expr: Expression): SupportLevel = { + val exprName = CometExplainInfo.exprDisplayName(expr) + if (!CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.get()) { + return Unsupported( + Some( + s"$exprName: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false; expression has " + + "no native path so the plan falls back to Spark")) + } + CometBatchKernelCodegen.canHandle(bindForDispatch(expr)._1) match { + case Some(reason) => Unsupported(Some(s"$exprName: $reason")) + case None => Compatible() + } + } + /** * Bind `expr`, closure-serialize it, and emit a `JvmScalarUdf` proto routed through * [[CometScalaUDFCodegen]] so that native execution evaluates the expression inside the @@ -66,6 +118,10 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { * via [[CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED]] or when * [[CometBatchKernelCodegen.canHandle]] refuses the expression tree. Callers should treat * `None` as a clean Spark-fallback signal. + * + * Serdes that gate on [[dispatchSupportLevel]] have already screened both of those conditions, + * so for them the checks below are a cheap re-verification. They are kept because several + * serdes call this directly from `convert` without gating first. */ def emitJvmCodegenDispatch( expr: Expression, @@ -80,20 +136,7 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { return None } - // `RuntimeReplaceable` expressions (e.g. Spark 4's `StructsToJson`) have a `doGenCode` that - // always throws "Cannot generate code for expression". Catalyst's `ReplaceExpressions` rule - // normally rewrites them to their `replacement` form before codegen runs. Comet's serde - // sometimes works with the pre-rewrite form (via shim reconstruction) for matching purposes, - // so unwrap to the replacement here before binding so the kernel compiles. - val target = expr match { - case rr: RuntimeReplaceable => rr.replacement - case other => other - } - - // Bind against only the AttributeReferences the tree actually reads, so ordinals align with - // the data args we ship. - val attrs = target.collect { case a: AttributeReference => a }.distinct - val boundExpr = BindReferences.bindReference(target, AttributeSeq(attrs)) + val (boundExpr, attrs) = bindForDispatch(expr) // Gate at plan time. Surface the reason via withFallbackReason rather than crashing Janino // at execute. @@ -159,18 +202,22 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { /** * Convenience base for serdes that route a non-ScalaUDF Spark expression through the codegen - * dispatcher. Delegates `convert` to [[CometScalaUDF.emitJvmCodegenDispatch]] and marks the - * expression `Compatible()` because the dispatcher runs Spark's own `doGenCode` inside the - * kernel: behavior matches Spark exactly when [[CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED]] is - * enabled, and the operator falls back to Spark cleanly when it is not. + * dispatcher. Delegates `convert` to [[CometScalaUDF.emitJvmCodegenDispatch]], and reports + * [[CometScalaUDF.dispatchSupportLevel]] so that the two conditions the dispatcher can refuse on + * -- the global [[CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED]] flag being off, and + * `CometBatchKernelCodegen.canHandle` rejecting the bound tree -- are reported from + * `getSupportLevel` instead of surfacing as a `Compatible` serde that then declines in `convert`. + * + * When the dispatcher will run the expression this is `Compatible()`: behavior then matches Spark + * exactly, because the kernel runs Spark's own `doGenCode`. */ class CometCodegenDispatch[T <: Expression] extends CometExpressionSerde[T] { - override def getSupportLevel(expr: T): SupportLevel = Compatible() + override def getSupportLevel(expr: T): SupportLevel = CometScalaUDF.dispatchSupportLevel(expr) // Intentionally no getCompatibleNotes override: the docs generator emits compat notes under // a heading that promises "no additional configuration required". The dispatcher flag is a // global concern documented elsewhere; tagging each expression here would contradict the - // heading. When the flag is off, `convert` returns None with a clear fallback reason that - // shows up in EXPLAIN, which is the right place for that signal. + // heading. When the flag is off, `getSupportLevel` reports Unsupported with a clear reason + // that shows up in EXPLAIN, which is the right place for that signal. override def convert(expr: T, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = CometScalaUDF.emitJvmCodegenDispatch(expr, inputs, binding) } diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index 5806cb35015..ca091e924db 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -25,7 +25,7 @@ import org.apache.arrow.vector._ import org.apache.spark.{SparkConf, SparkEnv, TaskContext} import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.api.java.UDF1 -import org.apache.spark.sql.catalyst.expressions.{BoundReference, CreateArray, CreateMap, CreateNamedStruct, Expression, Literal, MapConcat} +import org.apache.spark.sql.catalyst.expressions.{BoundReference, CreateArray, CreateMap, CreateNamedStruct, Expression, FindInSet, Literal, MapConcat} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types._ @@ -34,6 +34,7 @@ import org.apache.spark.unsafe.types.UTF8String import org.apache.comet.CometSparkSessionExtensions.isSpark41Plus import org.apache.comet.codegen.CometBatchKernelCodegen import org.apache.comet.codegen.CometBatchKernelCodegen.ArrowColumnSpec +import org.apache.comet.serde.{CometFindInSet, CometScalaUDF, Compatible, Unsupported} import org.apache.comet.udf.codegen.CometScalaUDFCodegen import org.apache.comet.vector.CometVector @@ -1754,6 +1755,36 @@ class CometCodegenSuite } } } + + // The two conditions the dispatcher can refuse on are reported from `getSupportLevel` rather + // than discovered inside `convert`, so a dispatch-only serde never claims `Compatible` and then + // declines. `find_in_set` stands in for the ~62 plain `CometCodegenDispatch` serdes. + private def findInSet = FindInSet(Literal("b"), Literal("a,b,c")) + + test("dispatch-only serdes report Compatible when the dispatcher will run the expression") { + assert(CometScalaUDF.dispatchSupportLevel(findInSet).isInstanceOf[Compatible]) + assert(CometFindInSet.getSupportLevel(findInSet).isInstanceOf[Compatible]) + } + + test("dispatch-only serdes report Unsupported when the dispatcher is disabled") { + withSQLConf(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") { + Seq( + CometScalaUDF.dispatchSupportLevel(findInSet), + CometFindInSet.getSupportLevel(findInSet)).foreach { + case Unsupported(Some(reason)) => + assert(reason.contains(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key)) + case other => fail(s"expected Unsupported, got $other") + } + } + } + + test("dispatch-only serdes report Unsupported when canHandle refuses the tree") { + // NullType is outside CometBatchKernelCodegen.isSupportedDataType. + CometScalaUDF.dispatchSupportLevel(Literal(null, NullType)) match { + case Unsupported(Some(reason)) => assert(reason.contains("unsupported output type")) + case other => fail(s"expected Unsupported, got $other") + } + } } /**