Apply the change to BenchmarkQueryTest too.

Author: sarutak

sql/core/src/main/scala/org/apache/spark/sql/execution/debug/package.scala

@@ -109,13 +109,11 @@ package object debug {
     val codegenSubtrees = new collection.mutable.HashSet[WholeStageCodegenExec]()
 
     def findSubtrees(plan: SparkPlan): Unit = {
-      plan transform {
+      plan foreach {
         case s: WholeStageCodegenExec =>
           codegenSubtrees += s
-          s
         case s =>
           s.subqueries.foreach(findSubtrees)
-          s
       }
     }
 

sql/core/src/test/scala/org/apache/spark/sql/BenchmarkQueryTest.scala

@@ -63,11 +63,17 @@ abstract class BenchmarkQueryTest extends QueryTest with SharedSparkSession {
 
   protected def checkGeneratedCode(plan: SparkPlan, checkMethodCodeSize: Boolean = true): Unit = {
     val codegenSubtrees = new collection.mutable.HashSet[WholeStageCodegenExec]()
-    plan foreach {
-      case s: WholeStageCodegenExec =>
-        codegenSubtrees += s
-      case _ =>
+
+    def findSubtrees(paln: SparkPlan): Unit = {
+      plan foreach {
+        case s: WholeStageCodegenExec =>
+          codegenSubtrees += s
+        case s =>
+          s.subqueries.foreach(findSubtrees)
+      }
     }
+
+    findSubtrees(plan)
     codegenSubtrees.toSeq.foreach { subtree =>
       val code = subtree.doCodeGen()._2
       val (_, ByteCodeStats(maxMethodCodeSize, _, _)) = try {

sql/core/src/test/scala/org/apache/spark/sql/ExplainSuite.scala

@@ -231,11 +231,8 @@ class ExplainSuite extends ExplainSuiteHelper with DisableAdaptiveExecutionSuite
   test("SPARK-33853: explain codegen - check presence of subquery") {
     withSQLConf(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true") {
       withTable("df1") {
-        spark.range(1, 100)
-          .write
-          .format("parquet")
-          .mode("overwrite")
-          .saveAsTable("df1")
+        val df1 = spark.range(1, 100)
+        df1.createTempView("df1")
 
         val sqlText = "EXPLAIN CODEGEN SELECT (SELECT min(id) FROM df1)"
         val expectedText = "Found 3 WholeStageCodegen subtrees."