diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index 07c1086624..a1f757b83e 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -84,15 +84,16 @@ 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.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. | +| 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.explainFallback.enabled` @@ -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..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,12 +19,14 @@ 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 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 +72,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 +100,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 +113,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 +142,17 @@ 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(Locale.ROOT)}: 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..51665a3ff1 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -131,6 +131,42 @@ 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) + 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