From 123bcae40f8a3cc20b2138da4f11f956d5bc4208 Mon Sep 17 00:00:00 2001
From: advancedxy <xianjin@apache.org>
Date: Sat, 20 Apr 2024 02:45:12 +0800
Subject: [PATCH] chore: Fix NPE when running CometTPCHQueriesList directly
 (#285)

* chore: Fix NPE when running CometTPCHQueriesList directly

* fix match error.
---
 .../scala/org/apache/spark/sql/CometTPCHQueriesList.scala   | 5 ++++-
 .../scala/org/apache/spark/sql/CometTPCQueryListBase.scala  | 6 ++++--
 2 files changed, 8 insertions(+), 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
index 3d8e1ed02b..3586e240a0 100644
--- 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
index 1f28b76a1e..ac3c047023 100644
--- 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) {