diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index 07c1086624..8eac031540 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 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.explainFallback.enabled` @@ -108,6 +109,48 @@ 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, 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: + +```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: 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. +``` + +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..6194ee7805 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -686,6 +686,15 @@ 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 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) + 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/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 3df275f057..e04f680b3b 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.{withCodegenDispatchExpr, withFallbackReason} import org.apache.comet.codegen.CometBatchKernelCodegen import org.apache.comet.serde.ExprOuterClass.Expr import org.apache.comet.serde.QueryPlanSerde.{exprToProtoInternal, serializeDataType} @@ -55,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 @@ -70,11 +85,12 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { expr: Expression, inputs: Seq[Attribute], binding: Boolean): Option[Expr] = { + val exprName = displayName(expr) 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 +113,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 +126,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 +155,12 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { udfBuilder .setReturnType(returnTypeProto) .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. + if (CometConf.COMET_EXPLAIN_CODEGEN_ENABLED.get()) { + withCodegenDispatchExpr(expr, exprName) + } 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..aa610d0751 100644 --- a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala @@ -131,6 +131,72 @@ class CometCodegenSuite } } + test("explainCodegen.enabled surfaces routed expressions in COMET-INFO") { + // 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')") + + 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, got:\n$explain") + // Names appear alphabetically via `.distinct.sorted` in rollUpInfoMessages. + assert( + 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("hypot:") && + explain.contains(CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key + "=false"), + s"expected 'hypot:' prefix and disabled-flag reason, 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