chore: Fix NPE when running CometTPCHQueriesList directly (apache#285)

* chore: Fix NPE when running CometTPCHQueriesList directly * fix match error.
himadripal · Apr 21, 2024 · 123bcae · 123bcae
1 parent 32b32a1
commit 123bcae
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTPCHQueriesList.scala b/spark/src/test/scala/org/apache/spark/sql/CometTPCHQueriesList.scala
@@ -49,9 +49,12 @@ object CometTPCHQueriesList extends CometTPCQueryListBase with CometTPCQueryBase
 
     // List of all TPC-H queries
     val tpchQueries = (1 to 22).map(n => s"q$n")
+    // Only q1 in the extended queries
+    val tpchExtendedQueries = Seq("q1")
 
     // If `--query-filter` defined, filters the queries that this option selects
     val queries = filterQueries(tpchQueries, benchmarkArgs.queryFilter)
+    val extendedQueries = filterQueries(tpchExtendedQueries, benchmarkArgs.queryFilter)
 
     if (queries.isEmpty) {
       throw new RuntimeException(
@@ -63,6 +66,6 @@ object CometTPCHQueriesList extends CometTPCQueryListBase with CometTPCQueryBase
     setupCBO(cometSpark, benchmarkArgs.cboEnabled, tables)
 
     runQueries("tpch", queries, " TPCH Snappy")
-    runQueries("tpch-extended", queries, " TPCH Extended Snappy")
+    runQueries("tpch-extended", extendedQueries, " TPCH Extended Snappy")
   }
 }
diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTPCQueryListBase.scala b/spark/src/test/scala/org/apache/spark/sql/CometTPCQueryListBase.scala
@@ -88,8 +88,10 @@ trait CometTPCQueryListBase
 
         val df = cometSpark.sql(queryString)
         val cometPlans = mutable.HashSet.empty[String]
-        stripAQEPlan(df.queryExecution.executedPlan).foreach { case op: CometExec =>
-          cometPlans += s"${op.nodeName}"
+        stripAQEPlan(df.queryExecution.executedPlan).foreach {
+          case op: CometExec =>
+            cometPlans += s"${op.nodeName}"
+          case _ =>
         }
 
         if (cometPlans.nonEmpty) {