From b344dea2defa9f7f9bb9aaf41413f88feef3268b Mon Sep 17 00:00:00 2001 From: comphead Date: Fri, 10 Jul 2026 13:43:30 -0700 Subject: [PATCH 1/6] chore: Codegen fallback explain for plan --- .../latest/understanding-comet-plans.md | 41 +++++++++++++++++- .../scala/org/apache/comet/CometConf.scala | 11 +++++ .../apache/comet/serde/CometScalaUDF.scala | 42 +++++++++++++++---- .../org/apache/comet/CometCodegenSuite.scala | 40 ++++++++++++++++++ 4 files changed, 124 insertions(+), 10 deletions(-) diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index 07c1086624..90f2cb332a 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -84,13 +84,14 @@ incompatibility details. ## Configs for Inspecting Plans and Fallback -Comet provides four configs for understanding what is happening in a plan. +Comet provides five configs for understanding what is happening in a plan. They serve different purposes and produce output in different places. | Config | Output destination | What you see | | ---------------------------------------- | ---------------------------------- | --------------------------------------------------------------------------------------------- | | `spark.comet.explainFallback.enabled` | Driver log (only when fallback) | A WARN with the list of reasons each query stage could not run in Comet. | | `spark.comet.logFallbackReasons.enabled` | Driver log | One WARN per fallback reason as it is encountered, without surrounding plan context. | +| `spark.comet.explainCodegen.enabled` | `ExtendedExplainInfo` / SQL UI | A `[COMET-INFO: : routes through JVM codegen dispatcher]` annotation per routed expression on the surrounding Comet operator. | | `spark.comet.explain.format` | Spark SQL UI (Spark 4.0 and newer) | Annotated plan or fallback-reason list, depending on `verbose` (default) or `fallback` value. | | `spark.comet.explain.native.enabled` | Executor logs, per task | The DataFusion plan with metrics, useful for inspecting Rust execution. | @@ -108,6 +109,44 @@ when you want to see all reasons, including ones that not include the surrounding plan, so it is best for accumulating diagnostics across many queries. +### `spark.comet.explainCodegen.enabled` + +Disabled by default. When enabled, every Spark expression that Comet routes +through the JVM codegen dispatcher (Spark's own `doGenCode` running inside a +Comet kernel — the path used for expressions with no native DataFusion +implementation, gated by `spark.comet.exec.scalaUDF.codegen.enabled=true`) is +annotated with `: routes through JVM codegen dispatcher`. The +annotation rolls up onto the surrounding `CometProject` / `CometFilter` / +etc. and is rendered inside the `[COMET-INFO: ...]` segment of extended +explain output (verbose format). The projection stays Comet-native — this is +informational only, useful for distinguishing "runs natively in DataFusion" +from "runs Spark's generated code inside a Comet kernel". + +Example: + +```scala +spark.conf.set("spark.comet.exec.scalaUDF.codegen.enabled", "true") +spark.conf.set("spark.comet.explainCodegen.enabled", "true") + +val df = spark.sql("SELECT hypot(a, b), levenshtein(s1, s2) FROM t") +println(new org.apache.comet.ExtendedExplainInfo() + .generateExtendedInfo(df.queryExecution.executedPlan)) +``` + +Output: + +``` +CometNativeColumnarToRow ++- CometProject [COMET-INFO: hypot: routes through JVM codegen dispatcher, levenshtein: routes through JVM codegen dispatcher] + +- CometNativeScan parquet spark_catalog.default.t + +Comet accelerated 2 out of 2 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. +``` + +Note that the operator is still `CometProject` (Comet-accelerated); only the +per-expression evaluation for `hypot` and `levenshtein` takes the JVM codegen +path. + ### `spark.comet.explain.format` This config is read by `org.apache.comet.ExtendedExplainInfo`, which Spark diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 5c130d457e..6288a3ba10 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -686,6 +686,17 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_EXPLAIN_CODEGEN_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.explainCodegen.enabled") + .category(CATEGORY_EXEC_EXPLAIN) + .doc( + "When this setting is enabled, Comet will annotate expressions that route through the " + + "JVM codegen dispatcher (Spark's `doGenCode` running inside a Comet kernel) with a " + + "`[COMET-INFO: ...]` message in extended explain output. Disabled by default to keep " + + "the plan display quiet unless the user is investigating dispatch coverage.") + .booleanConf + .createWithDefault(false) + val COMET_ONHEAP_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.exec.onHeap.enabled") .category(CATEGORY_TESTING) 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 3df275f057..267e7cc92a 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala @@ -24,7 +24,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, import org.apache.spark.sql.types.BinaryType import org.apache.comet.CometConf -import org.apache.comet.CometSparkSessionExtensions.withFallbackReason +import org.apache.comet.CometSparkSessionExtensions.{withFallbackReason, withInfo} import org.apache.comet.codegen.CometBatchKernelCodegen import org.apache.comet.serde.ExprOuterClass.Expr import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, serializeDataType} @@ -70,11 +70,12 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { expr: Expression, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = { + val exprName = expr.prettyName if (!CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.get()) { withFallbackReason( expr, - s"${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false; expression has no native " + - "path so the plan falls back to Spark") + s"$exprName: ${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false; expression has " + + "no native path so the plan falls back to Spark") return None } @@ -97,7 +98,7 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { // at execute. CometBatchKernelCodegen.canHandle(boundExpr) match { case Some(reason) => - withFallbackReason(expr, reason) + withFallbackReason(expr, s"$exprName: $reason") return None case None => } @@ -110,12 +111,26 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { val buffer = serializer.serialize(boundExpr) val bytes = new Array[Byte](buffer.remaining()) buffer.get(bytes) - val exprArg = exprToProtoInternal(Literal(bytes, BinaryType), inputs, binding) - .getOrElse(return None) + val exprArg = exprToProtoInternal(Literal(bytes, BinaryType), inputs, binding).getOrElse { + withFallbackReason( + expr, + s"$exprName: codegen dispatch: could not serialize closure-serialized bound " + + "expression payload") + return None + } - val dataArgs = - attrs.map(a => exprToProtoInternal(a, inputs, binding).getOrElse(return None)) - val returnTypeProto = serializeDataType(expr.dataType).getOrElse(return None) + val dataArgs = attrs.map { a => + exprToProtoInternal(a, inputs, binding).getOrElse { + withFallbackReason(expr, s"$exprName: codegen dispatch: could not serialize data arg $a") + return None + } + } + val returnTypeProto = serializeDataType(expr.dataType).getOrElse { + withFallbackReason( + expr, + s"$exprName: codegen dispatch: unsupported return type ${expr.dataType}") + return None + } val udfBuilder = ExprOuterClass.JvmScalarUdf .newBuilder() @@ -125,6 +140,15 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { udfBuilder .setReturnType(returnTypeProto) .setReturnNullable(expr.nullable) + // Surface the dispatch route in extended explain output when the user opts in via + // `spark.comet.explainCodegen.enabled=true`. `rollUpInfoMessages` in `CometExecRule` walks + // `op.expressions` via `TreeNode.collect` and copies each expression's `EXTENSION_INFO` tag + // onto the resulting Comet plan node, where `ExtendedExplainInfo` renders it as + // `[COMET-INFO: ...]`. Informational only: `withInfo` does not trigger fallback, so the + // projection remains a Comet operator regardless of this setting. + if (CometConf.COMET_EXPLAIN_CODEGEN_ENABLED.get()) { + withInfo(expr, s"${exprName.toLowerCase}: routes through JVM codegen dispatcher") + } Some( ExprOuterClass.Expr .newBuilder() diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index 76ffbc8c99..9b029da687 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -131,6 +131,46 @@ class CometCodegenSuite } } + test("explainCodegen.enabled surfaces routed expressions in COMET-INFO") { + // Two built-in expressions in the same SELECT that both route through the JVM codegen + // dispatcher (hypot -> CometHypot, levenshtein -> CometLevenshtein, both extend + // CometCodegenDispatch). With codegen dispatch enabled AND + // `spark.comet.explainCodegen.enabled=true`, `emitJvmCodegenDispatch` tags each expression + // with a withInfo message on its EXTENSION_INFO tag. `CometExecRule.rollUpInfoMessages` + // walks the projection's expression trees and copies those tags onto the CometProjectExec + // node, where `ExtendedExplainInfo` renders them as `[COMET-INFO: ...]`. This test verifies + // both named expressions surface in the info block while the projection stays Comet-native. + withTable("t") { + sql("CREATE TABLE t (a DOUBLE, b DOUBLE, s1 STRING, s2 STRING) USING parquet") + sql("INSERT INTO t VALUES (3.0, 4.0, 'kitten', 'sitting')") + + withSQLConf( + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true", + CometConf.COMET_EXPLAIN_CODEGEN_ENABLED.key -> "true", + CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "true", + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT.key -> + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_VERBOSE) { + val df = sql("SELECT hypot(a, b), levenshtein(s1, s2) FROM t") + checkSparkAnswerAndOperator(df) + val explain = + new ExtendedExplainInfo().generateExtendedInfo(df.queryExecution.executedPlan) + // Print so the test log shows the shape of the info annotation. + // scalastyle:off println + println(s"[codegen-fallback explain]\n$explain") + // scalastyle:on println + assert( + explain.contains("[COMET-INFO:"), + s"expected a [COMET-INFO: segment in explain output, got:\n$explain") + assert( + explain.contains("hypot: routes through JVM codegen dispatcher"), + s"expected 'hypot' codegen-dispatch info in explain output, got:\n$explain") + assert( + explain.contains("levenshtein: routes through JVM codegen dispatcher"), + s"expected 'levenshtein' codegen-dispatch info in explain output, got:\n$explain") + } + } + } + test("dispatcher caches the compiled kernel across batches of one query") { // Within a single query, the dispatcher compiles a kernel for the (expression, schema) pair // once and reuses it across every subsequent batch of the same shape. Force multiple batches From f45e48682374d160f0117ef65bae4b872fff9df8 Mon Sep 17 00:00:00 2001 From: comphead Date: Fri, 10 Jul 2026 14:02:31 -0700 Subject: [PATCH 2/6] chore: Codegen fallback explain for plan --- .../user-guide/latest/understanding-comet-plans.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index 90f2cb332a..a1f757b83e 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -87,13 +87,13 @@ incompatibility details. Comet provides five configs for understanding what is happening in a plan. They serve different purposes and produce output in different places. -| Config | Output destination | What you see | -| ---------------------------------------- | ---------------------------------- | --------------------------------------------------------------------------------------------- | -| `spark.comet.explainFallback.enabled` | Driver log (only when fallback) | A WARN with the list of reasons each query stage could not run in Comet. | -| `spark.comet.logFallbackReasons.enabled` | Driver log | One WARN per fallback reason as it is encountered, without surrounding plan context. | +| Config | Output destination | What you see | +| ---------------------------------------- | ---------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------- | +| `spark.comet.explainFallback.enabled` | Driver log (only when fallback) | A WARN with the list of reasons each query stage could not run in Comet. | +| `spark.comet.logFallbackReasons.enabled` | Driver log | One WARN per fallback reason as it is encountered, without surrounding plan context. | | `spark.comet.explainCodegen.enabled` | `ExtendedExplainInfo` / SQL UI | A `[COMET-INFO: : routes through JVM codegen dispatcher]` annotation per routed expression on the surrounding Comet operator. | -| `spark.comet.explain.format` | Spark SQL UI (Spark 4.0 and newer) | Annotated plan or fallback-reason list, depending on `verbose` (default) or `fallback` value. | -| `spark.comet.explain.native.enabled` | Executor logs, per task | The DataFusion plan with metrics, useful for inspecting Rust execution. | +| `spark.comet.explain.format` | Spark SQL UI (Spark 4.0 and newer) | Annotated plan or fallback-reason list, depending on `verbose` (default) or `fallback` value. | +| `spark.comet.explain.native.enabled` | Executor logs, per task | The DataFusion plan with metrics, useful for inspecting Rust execution. | ### `spark.comet.explainFallback.enabled` From e0889c558e7fe627ee659da9f2d6f12bf77f4741 Mon Sep 17 00:00:00 2001 From: comphead Date: Sat, 11 Jul 2026 09:49:14 -0700 Subject: [PATCH 3/6] chore: Codegen fallback explain for plan --- .../main/scala/org/apache/comet/serde/CometScalaUDF.scala | 6 +++++- .../src/test/scala/org/apache/comet/CometCodegenSuite.scala | 4 ---- 2 files changed, 5 insertions(+), 5 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 267e7cc92a..c99265bab2 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala @@ -19,6 +19,8 @@ package org.apache.comet.serde +import java.util.Locale + import org.apache.spark.SparkEnv import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSeq, BindReferences, Expression, Literal, RuntimeReplaceable, ScalaUDF} import org.apache.spark.sql.types.BinaryType @@ -147,7 +149,9 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { // `[COMET-INFO: ...]`. Informational only: `withInfo` does not trigger fallback, so the // projection remains a Comet operator regardless of this setting. if (CometConf.COMET_EXPLAIN_CODEGEN_ENABLED.get()) { - withInfo(expr, s"${exprName.toLowerCase}: routes through JVM codegen dispatcher") + withInfo( + expr, + s"${exprName.toLowerCase(Locale.ROOT)}: routes through JVM codegen dispatcher") } Some( ExprOuterClass.Expr diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index 9b029da687..51665a3ff1 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -154,10 +154,6 @@ class CometCodegenSuite checkSparkAnswerAndOperator(df) val explain = new ExtendedExplainInfo().generateExtendedInfo(df.queryExecution.executedPlan) - // Print so the test log shows the shape of the info annotation. - // scalastyle:off println - println(s"[codegen-fallback explain]\n$explain") - // scalastyle:on println assert( explain.contains("[COMET-INFO:"), s"expected a [COMET-INFO: segment in explain output, got:\n$explain") From 125e07b7fd96bc5d0c7fc19c3a634354abf2e5af Mon Sep 17 00:00:00 2001 From: comphead Date: Mon, 13 Jul 2026 11:39:10 -0700 Subject: [PATCH 4/6] chore: Codegen fallback explain for plan --- .../latest/understanding-comet-plans.md | 28 ++++++---- .../scala/org/apache/comet/CometConf.scala | 8 +-- .../comet/CometSparkSessionExtensions.scala | 15 +++++ .../apache/comet/ExtendedExplainInfo.scala | 5 ++ .../apache/comet/rules/CometExecRule.scala | 23 ++++++-- .../apache/comet/serde/CometScalaUDF.scala | 30 ++++++---- .../org/apache/comet/CometCodegenSuite.scala | 56 ++++++++++++++----- 7 files changed, 118 insertions(+), 47 deletions(-) diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index a1f757b83e..4062945d52 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -91,7 +91,7 @@ They serve different purposes and produce output in different places. | ---------------------------------------- | ---------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------- | | `spark.comet.explainFallback.enabled` | Driver log (only when fallback) | A WARN with the list of reasons each query stage could not run in Comet. | | `spark.comet.logFallbackReasons.enabled` | Driver log | One WARN per fallback reason as it is encountered, without surrounding plan context. | -| `spark.comet.explainCodegen.enabled` | `ExtendedExplainInfo` / SQL UI | A `[COMET-INFO: : routes through JVM codegen dispatcher]` annotation per routed expression on the surrounding Comet operator. | +| `spark.comet.explainCodegen.enabled` | `ExtendedExplainInfo` / SQL UI | A single `[COMET-INFO: JVM codegen dispatcher: , , ...]` segment naming each expression on the operator that was routed through the JVM codegen dispatcher. | | `spark.comet.explain.format` | Spark SQL UI (Spark 4.0 and newer) | Annotated plan or fallback-reason list, depending on `verbose` (default) or `fallback` value. | | `spark.comet.explain.native.enabled` | Executor logs, per task | The DataFusion plan with metrics, useful for inspecting Rust execution. | @@ -111,16 +111,20 @@ across many queries. ### `spark.comet.explainCodegen.enabled` -Disabled by default. When enabled, every Spark expression that Comet routes -through the JVM codegen dispatcher (Spark's own `doGenCode` running inside a -Comet kernel — the path used for expressions with no native DataFusion -implementation, gated by `spark.comet.exec.scalaUDF.codegen.enabled=true`) is -annotated with `: routes through JVM codegen dispatcher`. The -annotation rolls up onto the surrounding `CometProject` / `CometFilter` / -etc. and is rendered inside the `[COMET-INFO: ...]` segment of extended -explain output (verbose format). The projection stays Comet-native — this is -informational only, useful for distinguishing "runs natively in DataFusion" -from "runs Spark's generated code inside a Comet kernel". +Disabled by default. When enabled, Comet annotates the surrounding Comet +operator (`CometProject`, `CometFilter`, etc.) with a single combined +`[COMET-INFO: JVM codegen dispatcher: , , ...]` segment listing +every expression on that operator that was routed through the JVM codegen +dispatcher (Spark's own `doGenCode` running inside a Comet kernel), gated by +`spark.comet.exec.scalaUDF.codegen.enabled=true`. The projection stays +Comet-native — this is informational only, useful for distinguishing "runs +natively in DataFusion" from "runs Spark's generated code inside a Comet +kernel". The annotation only appears for expressions Comet actually routed +through the dispatcher on this plan — either because no native DataFusion +implementation exists (`CometCodegenDispatch` serdes such as `hypot`, +`levenshtein`, `pmod`), or because a native path exists but declined this +specific input (`CodegenDispatchFallback` mixins). It is not a general "this +expression ran here" marker. Example: @@ -137,7 +141,7 @@ Output: ``` CometNativeColumnarToRow -+- CometProject [COMET-INFO: hypot: routes through JVM codegen dispatcher, levenshtein: routes through JVM codegen dispatcher] ++- CometProject [COMET-INFO: JVM codegen dispatcher: hypot, levenshtein] +- CometNativeScan parquet spark_catalog.default.t Comet accelerated 2 out of 2 eligible operators (100%). Final plan contains 1 transitions between Spark and Comet. diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 6288a3ba10..6194ee7805 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -689,11 +689,9 @@ object CometConf extends ShimCometConf { val COMET_EXPLAIN_CODEGEN_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.explainCodegen.enabled") .category(CATEGORY_EXEC_EXPLAIN) - .doc( - "When this setting is enabled, Comet will annotate expressions that route through the " + - "JVM codegen dispatcher (Spark's `doGenCode` running inside a Comet kernel) with a " + - "`[COMET-INFO: ...]` message in extended explain output. Disabled by default to keep " + - "the plan display quiet unless the user is investigating dispatch coverage.") + .doc("When enabled, Comet annotates the surrounding Comet operator with a `[COMET-INFO: " + + "JVM codegen dispatcher: ]` segment listing every expression it routed through " + + "the JVM codegen dispatcher. Disabled by default.") .booleanConf .createWithDefault(false) diff --git a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala index f117372b3c..202c52b158 100644 --- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala +++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala @@ -383,4 +383,19 @@ object CometSparkSessionExtensions extends Logging { node } + /** + * Record that `node` (typically an `Expression`) is routing through the JVM codegen dispatcher. + * `CometExecRule.rollUpInfoMessages` collects the names across an operator's expression trees + * and emits one combined `[COMET-INFO: ...]` segment. + */ + def withCodegenDispatchExpr[T <: TreeNode[_]](node: T, name: String): T = { + if (name != null && name.nonEmpty) { + val existing = node + .getTagValue(CometExplainInfo.CODEGEN_DISPATCH_EXPRS) + .getOrElse(Set.empty[String]) + node.setTagValue(CometExplainInfo.CODEGEN_DISPATCH_EXPRS, existing + name) + } + node + } + } diff --git a/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala b/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala index 5318c3f87c..0a7a8f4e30 100644 --- a/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala +++ b/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala @@ -224,6 +224,11 @@ object CometExplainInfo { val FALLBACK_REASONS = new TreeNodeTag[Set[String]]("CometFallbackReasons") val EXTENSION_INFO = new TreeNodeTag[Set[String]]("CometExtensionInfo") + // Expression names the serde routed through the JVM codegen dispatcher. Rolled up per + // operator by `CometExecRule.rollUpInfoMessages` into one combined `[COMET-INFO: ...]`. + val CODEGEN_DISPATCH_EXPRS = + new TreeNodeTag[Set[String]]("CometCodegenDispatchExprs") + def getActualPlan(node: TreeNode[_]): TreeNode[_] = { node match { case p: AdaptiveSparkPlanExec => getActualPlan(p.executedPlan) diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index dc8268950e..70982bdb26 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -729,14 +729,25 @@ case class CometExecRule(session: SparkSession) * Lift informational (non-fallback) messages tagged on an operator and its expressions onto the * converted Comet plan node so they appear in verbose extended explain output. Expression-level * hints would otherwise be invisible because explain only traverses plan nodes, not - * expressions. + * expressions. `CODEGEN_DISPATCH_EXPRS` names across the tree are aggregated into one combined + * info line. */ private def rollUpInfoMessages(op: SparkPlan, exec: SparkPlan): Unit = { - val fromOp = op.getTagValue(CometExplainInfo.EXTENSION_INFO).getOrElse(Set.empty[String]) - val fromExprs = op.expressions - .flatMap(_.collect { case e: Expression => e }) - .flatMap(_.getTagValue(CometExplainInfo.EXTENSION_INFO).getOrElse(Set.empty[String])) - (fromOp ++ fromExprs).foreach(msg => withInfo(exec, msg)) + val allExprs = op.expressions.flatMap(_.collect { case e: Expression => e }) + + val infos = + op.getTagValue(CometExplainInfo.EXTENSION_INFO).getOrElse(Set.empty[String]) ++ + allExprs.flatMap(_.getTagValue(CometExplainInfo.EXTENSION_INFO)).flatten + infos.foreach(msg => withInfo(exec, msg)) + + val routedNames = allExprs + .flatMap(_.getTagValue(CometExplainInfo.CODEGEN_DISPATCH_EXPRS)) + .flatten + .distinct + .sorted + if (routedNames.nonEmpty) { + withInfo(exec, s"JVM codegen dispatcher: ${routedNames.mkString(", ")}") + } } private def isOperatorEnabled( 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 c99265bab2..8013342f51 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala @@ -26,7 +26,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, import org.apache.spark.sql.types.BinaryType import org.apache.comet.CometConf -import org.apache.comet.CometSparkSessionExtensions.{withFallbackReason, withInfo} +import org.apache.comet.CometSparkSessionExtensions.{withCodegenDispatchExpr, withFallbackReason} import org.apache.comet.codegen.CometBatchKernelCodegen import org.apache.comet.serde.ExprOuterClass.Expr import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, serializeDataType} @@ -57,6 +57,19 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { override def convert(expr: ScalaUDF, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = emitJvmCodegenDispatch(expr, inputs, binding) + /** + * Canonical name for annotation. `BinaryMathExpression` (Hypot, Pow, ...) overrides + * `prettyName` to raw uppercase; `ScalaUDF.prettyName` collapses to `"scalaudf"` for every user + * UDF. Prefer `ScalaUDF.udfName` when set, then lowercase to normalize. + */ + private def displayName(expr: Expression): String = { + val raw = expr match { + case s: ScalaUDF => s.udfName.getOrElse(s.prettyName) + case other => other.prettyName + } + raw.toLowerCase(Locale.ROOT) + } + /** * Bind `expr`, closure-serialize it, and emit a `JvmScalarUdf` proto routed through * [[CometScalaUDFCodegen]] so that native execution evaluates the expression inside the @@ -72,7 +85,7 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { expr: Expression, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = { - val exprName = expr.prettyName + val exprName = displayName(expr) if (!CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.get()) { withFallbackReason( expr, @@ -142,16 +155,11 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { udfBuilder .setReturnType(returnTypeProto) .setReturnNullable(expr.nullable) - // Surface the dispatch route in extended explain output when the user opts in via - // `spark.comet.explainCodegen.enabled=true`. `rollUpInfoMessages` in `CometExecRule` walks - // `op.expressions` via `TreeNode.collect` and copies each expression's `EXTENSION_INFO` tag - // onto the resulting Comet plan node, where `ExtendedExplainInfo` renders it as - // `[COMET-INFO: ...]`. Informational only: `withInfo` does not trigger fallback, so the - // projection remains a Comet operator regardless of this setting. + // Opt-in dispatch annotation for extended explain. Rolled up per operator by + // `CometExecRule.rollUpInfoMessages` into a single `[COMET-INFO: JVM codegen dispatcher: + // ...]` line. Informational only — does not trigger fallback. if (CometConf.COMET_EXPLAIN_CODEGEN_ENABLED.get()) { - withInfo( - expr, - s"${exprName.toLowerCase(Locale.ROOT)}: routes through JVM codegen dispatcher") + withCodegenDispatchExpr(expr, exprName) } Some( ExprOuterClass.Expr diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index 51665a3ff1..aa610d0751 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -132,14 +132,9 @@ class CometCodegenSuite } test("explainCodegen.enabled surfaces routed expressions in COMET-INFO") { - // Two built-in expressions in the same SELECT that both route through the JVM codegen - // dispatcher (hypot -> CometHypot, levenshtein -> CometLevenshtein, both extend - // CometCodegenDispatch). With codegen dispatch enabled AND - // `spark.comet.explainCodegen.enabled=true`, `emitJvmCodegenDispatch` tags each expression - // with a withInfo message on its EXTENSION_INFO tag. `CometExecRule.rollUpInfoMessages` - // walks the projection's expression trees and copies those tags onto the CometProjectExec - // node, where `ExtendedExplainInfo` renders them as `[COMET-INFO: ...]`. This test verifies - // both named expressions surface in the info block while the projection stays Comet-native. + // With the opt-in flag on, `hypot` and `levenshtein` (both `CometCodegenDispatch`) roll up + // into one `[COMET-INFO: JVM codegen dispatcher: hypot, levenshtein]` line on the + // `CometProject`. With the flag off (default), no such line appears. withTable("t") { sql("CREATE TABLE t (a DOUBLE, b DOUBLE, s1 STRING, s2 STRING) USING parquet") sql("INSERT INTO t VALUES (3.0, 4.0, 'kitten', 'sitting')") @@ -156,13 +151,48 @@ class CometCodegenSuite new ExtendedExplainInfo().generateExtendedInfo(df.queryExecution.executedPlan) assert( explain.contains("[COMET-INFO:"), - s"expected a [COMET-INFO: segment in explain output, got:\n$explain") + s"expected a [COMET-INFO: segment, got:\n$explain") + // Names appear alphabetically via `.distinct.sorted` in rollUpInfoMessages. assert( - explain.contains("hypot: routes through JVM codegen dispatcher"), - s"expected 'hypot' codegen-dispatch info in explain output, got:\n$explain") + explain.contains("JVM codegen dispatcher: hypot, levenshtein"), + s"expected combined codegen-dispatch info, got:\n$explain") + } + + withSQLConf( + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true", + CometConf.COMET_EXPLAIN_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_EXEC_PROJECT_ENABLED.key -> "true", + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT.key -> + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_VERBOSE) { + val df = sql("SELECT hypot(a, b), levenshtein(s1, s2) FROM t") + checkSparkAnswerAndOperator(df) + val explain = + new ExtendedExplainInfo().generateExtendedInfo(df.queryExecution.executedPlan) + assert( + !explain.contains("JVM codegen dispatcher"), + s"expected NO codegen-dispatch info with the flag off, got:\n$explain") + } + } + } + + test("codegen dispatch fallback reasons name the expression") { + // Flag-off short-circuit tags the expression `: ` so distinct expressions + // don't collapse in the `Set[String]` roll-up. + withTable("t") { + sql("CREATE TABLE t (a DOUBLE, b DOUBLE) USING parquet") + sql("INSERT INTO t VALUES (3.0, 4.0)") + withSQLConf( + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT.key -> + CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_VERBOSE) { + val df = sql("SELECT hypot(a, b) FROM t") + checkSparkAnswer(df) + val explain = + new ExtendedExplainInfo().generateExtendedInfo(df.queryExecution.executedPlan) assert( - explain.contains("levenshtein: routes through JVM codegen dispatcher"), - s"expected 'levenshtein' codegen-dispatch info in explain output, got:\n$explain") + explain.contains("hypot:") && + explain.contains(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key + "=false"), + s"expected 'hypot:' prefix and disabled-flag reason, got:\n$explain") } } } From 18121eb505b3305cfa208ab09da6070576707f28 Mon Sep 17 00:00:00 2001 From: comphead Date: Mon, 13 Jul 2026 12:10:15 -0700 Subject: [PATCH 5/6] chore: Codegen fallback explain for plan --- .../user-guide/latest/understanding-comet-plans.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index 4062945d52..8eac031540 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -87,13 +87,13 @@ incompatibility details. Comet provides five configs for understanding what is happening in a plan. They serve different purposes and produce output in different places. -| Config | Output destination | What you see | -| ---------------------------------------- | ---------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------- | -| `spark.comet.explainFallback.enabled` | Driver log (only when fallback) | A WARN with the list of reasons each query stage could not run in Comet. | -| `spark.comet.logFallbackReasons.enabled` | Driver log | One WARN per fallback reason as it is encountered, without surrounding plan context. | +| Config | Output destination | What you see | +| ---------------------------------------- | ---------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `spark.comet.explainFallback.enabled` | Driver log (only when fallback) | A WARN with the list of reasons each query stage could not run in Comet. | +| `spark.comet.logFallbackReasons.enabled` | Driver log | One WARN per fallback reason as it is encountered, without surrounding plan context. | | `spark.comet.explainCodegen.enabled` | `ExtendedExplainInfo` / SQL UI | A single `[COMET-INFO: JVM codegen dispatcher: , , ...]` segment naming each expression on the operator that was routed through the JVM codegen dispatcher. | -| `spark.comet.explain.format` | Spark SQL UI (Spark 4.0 and newer) | Annotated plan or fallback-reason list, depending on `verbose` (default) or `fallback` value. | -| `spark.comet.explain.native.enabled` | Executor logs, per task | The DataFusion plan with metrics, useful for inspecting Rust execution. | +| `spark.comet.explain.format` | Spark SQL UI (Spark 4.0 and newer) | Annotated plan or fallback-reason list, depending on `verbose` (default) or `fallback` value. | +| `spark.comet.explain.native.enabled` | Executor logs, per task | The DataFusion plan with metrics, useful for inspecting Rust execution. | ### `spark.comet.explainFallback.enabled` From 1dfc5df9cf9e4e7b8ea6c6090c7074612c88ac69 Mon Sep 17 00:00:00 2001 From: comphead Date: Mon, 13 Jul 2026 14:59:32 -0700 Subject: [PATCH 6/6] chore: Codegen fallback explain for plan --- spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 8013342f51..e04f680b3b 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala @@ -157,7 +157,7 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { .setReturnNullable(expr.nullable) // Opt-in dispatch annotation for extended explain. Rolled up per operator by // `CometExecRule.rollUpInfoMessages` into a single `[COMET-INFO: JVM codegen dispatcher: - // ...]` line. Informational only — does not trigger fallback. + // ...]` line. Informational only - does not trigger fallback. if (CometConf.COMET_EXPLAIN_CODEGEN_ENABLED.get()) { withCodegenDispatchExpr(expr, exprName) }