This is an automated email from the ASF dual-hosted git repository.
agrove 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 9b7e26d6 fix: withInfo was overwriting information in some cases (#780)
9b7e26d6 is described below
commit 9b7e26d626454fd5411faee9032205d703a9519b
Author: Andy Grove <[email protected]>
AuthorDate: Mon Aug 5 15:24:21 2024 -0600
fix: withInfo was overwriting information in some cases (#780)
* add test for withInfo
* improve test
* Revert whitespace changes
* update test
* format
---
.../apache/comet/CometSparkSessionExtensions.scala | 6 ++--
.../org/apache/comet/CometExpressionSuite.scala | 35 ++++++++++++++++++++++
2 files changed, 39 insertions(+), 2 deletions(-)
diff --git
a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
index 78477486..9d2b5651 100644
--- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
+++ b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
@@ -1235,8 +1235,10 @@ object CometSparkSessionExtensions extends Logging {
* The node with information (if any) attached
*/
def withInfos[T <: TreeNode[_]](node: T, info: Set[String], exprs: T*): T = {
- val exprInfo =
exprs.flatMap(_.getTagValue(CometExplainInfo.EXTENSION_INFO)).flatten.toSet
- node.setTagValue(CometExplainInfo.EXTENSION_INFO, exprInfo ++ info)
+ val existingNodeInfos = node.getTagValue(CometExplainInfo.EXTENSION_INFO)
+ val newNodeInfo = (existingNodeInfos ++ exprs
+ .flatMap(_.getTagValue(CometExplainInfo.EXTENSION_INFO))).flatten.toSet
+ node.setTagValue(CometExplainInfo.EXTENSION_INFO, newNodeInfo ++ info)
node
}
diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
index e787eb08..cce48719 100644
--- a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
@@ -27,6 +27,8 @@ import scala.util.Random
import org.apache.hadoop.fs.Path
import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
+import org.apache.spark.sql.comet.CometProjectExec
+import org.apache.spark.sql.execution.{ColumnarToRowExec, InputAdapter,
WholeStageCodegenExec}
import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
import org.apache.spark.sql.functions._
import org.apache.spark.sql.internal.SQLConf
@@ -634,6 +636,37 @@ class CometExpressionSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
}
}
+ test("withInfo") {
+ val table = "with_info"
+ withTable(table) {
+ sql(s"create table $table(id int, name varchar(20)) using parquet")
+ sql(s"insert into $table values(1,'James Smith')")
+ val query = sql(s"select cast(id as string) from $table")
+ val (_, cometPlan) = checkSparkAnswer(query)
+ val project = cometPlan
+ .asInstanceOf[WholeStageCodegenExec]
+ .child
+ .asInstanceOf[ColumnarToRowExec]
+ .child
+ .asInstanceOf[InputAdapter]
+ .child
+ .asInstanceOf[CometProjectExec]
+ val id = project.expressions.head
+ CometSparkSessionExtensions.withInfo(id, "reason 1")
+ CometSparkSessionExtensions.withInfo(project, "reason 2")
+ CometSparkSessionExtensions.withInfo(project, "reason 3", id)
+ CometSparkSessionExtensions.withInfo(project, id)
+ CometSparkSessionExtensions.withInfo(project, "reason 4")
+ CometSparkSessionExtensions.withInfo(project, "reason 5", id)
+ CometSparkSessionExtensions.withInfo(project, id)
+ CometSparkSessionExtensions.withInfo(project, "reason 6")
+ val explain = new ExtendedExplainInfo().generateExtendedInfo(project)
+ for (i <- 1 until 7) {
+ assert(explain.contains(s"reason $i"))
+ }
+ }
+ }
+
test("rlike fallback for non scalar pattern") {
val table = "rlike_fallback"
withTable(table) {
@@ -1697,6 +1730,7 @@ class CometExpressionSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
s"SELECT sum(c0), sum(c2) from $table group by c1",
Set(
"HashAggregate is not native because the following children are
not native (AQEShuffleRead)",
+ "HashAggregate is not native because the following children are
not native (Exchange)",
"Comet shuffle is not enabled: spark.comet.exec.shuffle.enabled
is not enabled",
"AQEShuffleRead is not supported")),
(
@@ -1709,6 +1743,7 @@ class CometExpressionSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
"AQEShuffleRead is not supported",
"make_interval is not supported",
"HashAggregate is not native because the following children are
not native (AQEShuffleRead)",
+ "HashAggregate is not native because the following children are
not native (Exchange)",
"Project is not native because the following children are not
native (BroadcastHashJoin)",
"BroadcastHashJoin is not enabled because the following children
are not native" +
" (BroadcastExchange, Project)")))
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
