This is an automated email from the ASF dual-hosted git repository.
sunchao pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 26a70bddca fix: ignore structural tags when lifting expression
coverage (#5471)
26a70bddca is described below
commit 26a70bddca1ae2bdca614677b37cdb9fbdb63c7f
Author: Chao Sun <[email protected]>
AuthorDate: Wed Sep 9 16:03:42 2026 -0700
fix: ignore structural tags when lifting expression coverage (#5471)
* fix: ignore structural tags when lifting expression coverage
* test: preserve dispatcher coverage across decimal promotion
---
.../org/apache/comet/ExtendedExplainInfo.scala | 11 ++++----
.../org/apache/comet/serde/QueryPlanSerde.scala | 15 +++++-----
.../scala/org/apache/comet/CometCodegenSuite.scala | 33 +++++++++++++++++++++-
3 files changed, 46 insertions(+), 13 deletions(-)
diff --git a/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala
b/spark/src/main/scala/org/apache/comet/ExtendedExplainInfo.scala
index b913705155..72e168d7b0 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 ff6ecc471f..1fa43f07de 100644
--- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
+++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -885,14 +885,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)
}
/**
@@ -1047,6 +1043,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 e0175cf756..247f69acdd 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, 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
@@ -358,7 +374,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 +410,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 +418,7 @@ class CometCodegenSuite
} finally {
planted.unsetTagValue(CometExplainInfo.EXTENSION_INFO)
planted.unsetTagValue(CometExplainInfo.NATIVE_EXPRS)
+ planted.unsetTagValue(CometExplainInfo.CODEGEN_DISPATCH_EXPRS)
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
