From ec19d6bd159c00bdebb7acff6a3c3a1f53ef757d Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Wed, 26 Aug 2026 05:16:21 +0000 Subject: [PATCH 1/2] fix: ignore structural tags when lifting expression coverage --- .../org/apache/comet/serde/QueryPlanSerde.scala | 10 +++------- .../org/apache/comet/CometCodegenSuite.scala | 17 ++++++++++++++++- 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index d6c9a496040..d36311fa845 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -879,14 +879,10 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { } private def liftCoverageTags(from: Expression, to: Expression): Unit = { - val native = mutable.Set.empty[String] - val dispatched = mutable.Set.empty[String] - from.foreach { e => - e.getTagValue(CometExplainInfo.NATIVE_EXPRS).foreach(native ++= _) - e.getTagValue(CometExplainInfo.CODEGEN_DISPATCH_EXPRS).foreach(dispatched ++= _) + val exprs = from.collect { case e: Expression => e } + Seq(CometExplainInfo.NATIVE_EXPRS, CometExplainInfo.CODEGEN_DISPATCH_EXPRS).foreach { tag => + appendTagValues(to, tag, CometExplainInfo.collectExprTagValues(exprs, tag)) } - appendTagValues(to, CometExplainInfo.NATIVE_EXPRS, native.toSet) - appendTagValues(to, CometExplainInfo.CODEGEN_DISPATCH_EXPRS, dispatched.toSet) } /** diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index e0175cf756c..1a9b15cd452 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.{AttributeReference, BoundReference, CreateArray, CreateMap, CreateNamedStruct, Expression, Literal, MapConcat} +import org.apache.spark.sql.catalyst.expressions.{Add, Alias, AttributeReference, BoundReference, CreateArray, CreateMap, CreateNamedStruct, Expression, Literal, MapConcat} import org.apache.spark.sql.catalyst.expressions.objects.Invoke import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.internal.SQLConf @@ -358,7 +358,20 @@ class CometCodegenSuite val planted = Literal.TrueLiteral planted.setTagValue(CometExplainInfo.EXTENSION_INFO, Set("PLANTED_INFO")) planted.setTagValue(CometExplainInfo.NATIVE_EXPRS, Set("plantedexpr")) + planted.setTagValue(CometExplainInfo.CODEGEN_DISPATCH_EXPRS, Set("planteddispatch")) try { + // Decimal promotion rebuilds this projection. Its coverage lift must not copy the + // singleton's stale tags onto the Alias, which is a legitimate coverage owner. + val decimal = AttributeReference("amount", DecimalType(10, 2), nullable = false)() + val projection = Alias( + CreateNamedStruct(Seq(Literal("flag"), planted, Literal("sum"), Add(decimal, decimal))), + "value")() + assert(QueryPlanSerde.exprToProto(projection, Seq(decimal)).isDefined) + val native = projection.getTagValue(CometExplainInfo.NATIVE_EXPRS).getOrElse(Set.empty) + assert(native.contains("checkoverflow"), s"expected lifted decimal coverage, got: $native") + assert(!native.contains("plantedexpr")) + assert(projection.getTagValue(CometExplainInfo.CODEGEN_DISPATCH_EXPRS).isEmpty) + withSQLConf( CometConf.COMET_EXTENDED_EXPLAIN_FORMAT.key -> CometConf.COMET_EXTENDED_EXPLAIN_FORMAT_VERBOSE, @@ -381,6 +394,7 @@ class CometCodegenSuite val info = new ExtendedExplainInfo() assert(!info.getNativeExpressions(plan).contains("plantedexpr")) + assert(!info.getCodegenDispatchExpressions(plan).contains("planteddispatch")) val explain = info.generateExtendedInfo(plan) assert(!explain.contains("PLANTED_INFO"), s"tag leaked into:\n$explain") } @@ -388,6 +402,7 @@ class CometCodegenSuite } finally { planted.unsetTagValue(CometExplainInfo.EXTENSION_INFO) planted.unsetTagValue(CometExplainInfo.NATIVE_EXPRS) + planted.unsetTagValue(CometExplainInfo.CODEGEN_DISPATCH_EXPRS) } } From 5674aeb60c57c25b849a5801eee769d66a90648a Mon Sep 17 00:00:00 2001 From: Chao Sun Date: Thu, 27 Aug 2026 16:47:07 +0000 Subject: [PATCH 2/2] test: preserve dispatcher coverage across decimal promotion --- .../org/apache/comet/ExtendedExplainInfo.scala | 11 ++++++----- .../apache/comet/serde/QueryPlanSerde.scala | 5 +++++ .../org/apache/comet/CometCodegenSuite.scala | 18 +++++++++++++++++- 3 files changed, 28 insertions(+), 6 deletions(-) diff --git a/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala b/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala index b913705155b..72e168d7b0b 100644 --- a/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala +++ b/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala @@ -311,7 +311,8 @@ object CometExplainInfo { } /** - * Union of a `Set`-valued tag over `exprs`, skipping nodes the serde never tags. + * Union of a coverage or info tag over `exprs`, skipping nodes the serde never tags for those + * purposes. This filter must not be used for `FALLBACK_REASONS`, which literals can carry. * * Catalyst copies a rewritten node's tags onto its replacement (`TreeNode.copyTagsFrom`, which * copies whenever the replacement has no tags of its own). Rewriting a tagged expression into a @@ -328,10 +329,10 @@ object CometExplainInfo { } /** - * Nodes that never carry a Comet tag of their own, so anything found on one arrived by the - * copying described in [[collectExprTagValues]]. `Literal` is the node that matters, being the - * only one with JVM-wide singletons (`Literal.TrueLiteral`, `Literal.FalseLiteral`); the other - * two are listed because nothing legitimate can live on them either. + * Nodes that never carry their own coverage or info tags, so those tags can only arrive by the + * copying described in [[collectExprTagValues]]. This set must match + * `QueryPlanSerde.isStructuralExpr` minus `Alias`; changing either set requires checking the + * other. This invariant does not apply to `FALLBACK_REASONS`. * * `Alias` is deliberately absent even though the serde does not tag one directly: * `QueryPlanSerde.liftCoverageTags` lands names on whichever node the operator holds, and for a diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index d36311fa845..ed5e501f4d4 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -1032,6 +1032,11 @@ object QueryPlanSerde extends Logging with CometExprShim with CometTypeShim { * Nodes that carry no computation of their own. They are excluded from the expression coverage * stats in extended explain because they appear in nearly every expression tree and would swamp * the names a user actually cares about. + * + * `CometExplainInfo.isNeverTagged` must be this set minus `Alias`: the read-side filter retains + * aliases because [[liftCoverageTags]] uses them to hold names from rewritten children. Keep + * both sets in sync. This coverage invariant does not exclude structural nodes from carrying + * `FALLBACK_REASONS`. */ private def isStructuralExpr(expr: Expression): Boolean = expr match { case _: Attribute | _: BoundReference | _: Literal | _: Alias => true diff --git a/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala b/spark/src/test/scala/org/apache/comet/CometCodegenSuite.scala index 1a9b15cd452..247f69acdd3 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.{Add, Alias, AttributeReference, BoundReference, CreateArray, CreateMap, CreateNamedStruct, Expression, Literal, MapConcat} +import org.apache.spark.sql.catalyst.expressions.{Add, Alias, AttributeReference, BoundReference, Cast, CreateArray, CreateMap, CreateNamedStruct, Expression, Hypot, Literal, MapConcat} import org.apache.spark.sql.catalyst.expressions.objects.Invoke import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.internal.SQLConf @@ -348,6 +348,22 @@ class CometCodegenSuite } } + test("codegen dispatch coverage survives the decimal promotion rewrite") { + val decimal = AttributeReference("amount", DecimalType(10, 2), nullable = false)() + val dispatched = Hypot(Cast(Add(decimal, decimal), DoubleType), Literal(4.0d)) + val projection = Alias(dispatched, "value")() + + // Promotion rebuilds Hypot as well as the Alias above it. Unlike the original Add, the + // dispatched copy is not reachable from the original tree, so only the coverage lift can + // bring its name back to the projection owner. + val proto = QueryPlanSerde.exprToProto(projection, Seq(decimal)).get + assert(proto.hasJvmScalarUdf) + assert(proto.getJvmScalarUdf.getClassName === classOf[CometScalaUDFCodegen].getName) + assert(dispatched.getTagValue(CometExplainInfo.DISPATCHED_SELF).isEmpty) + assert(dispatched.getTagValue(CometExplainInfo.CODEGEN_DISPATCH_EXPRS).isEmpty) + assert(projection.getTagValue(CometExplainInfo.CODEGEN_DISPATCH_EXPRS).contains(Set("hypot"))) + } + test("tags copied onto the shared TrueLiteral do not leak into unrelated plans") { // Catalyst copies a rewritten node's tags onto its replacement, so a tagged expression that an // earlier query rewrote into `Literal.TrueLiteral` brands that process-wide singleton for the