fix: Native Scan assert is irrelevant once in a different Comet native plan

Author: EmilyMatt

spark/src/main/scala/org/apache/spark/sql/comet/operators.scala

@@ -219,10 +219,12 @@ abstract class CometNativeExec extends CometExec {
 
         // For each relation in a CometNativeScan generate a hadoopConf,
         // for each file path in a relation associate with hadoopConf
-        val cometNativeScans: Seq[CometNativeScanExec] = this
-          .collectLeaves()
-          .filter(_.isInstanceOf[CometNativeScanExec])
-          .map(_.asInstanceOf[CometNativeScanExec])
+        // This is done per native plan, so only count scans until a comet input is reached.
+        val cometNativeScans = mutable.ArrayBuffer.empty[CometNativeScanExec]
+        foreachUntilCometInput(this) {
+          case scan: CometNativeScanExec => cometNativeScans += scan
+          case _ => // no-op
+        }
         assert(
           cometNativeScans.size <= 1,
           "We expect one native scan in a Comet plan since we will broadcast one hadoopConf.")

spark/src/test/scala/org/apache/spark/sql/comet/ParquetEncryptionITCase.scala

@@ -22,20 +22,17 @@ package org.apache.spark.sql.comet
 import java.io.{File, RandomAccessFile}
 import java.nio.charset.StandardCharsets
 import java.util.Base64
-
 import org.junit.runner.RunWith
 import org.scalactic.source.Position
 import org.scalatest.Tag
 import org.scalatestplus.junit.JUnitRunner
-
 import org.apache.parquet.crypto.DecryptionPropertiesFactory
 import org.apache.parquet.crypto.keytools.{KeyToolkit, PropertiesDrivenCryptoFactory}
 import org.apache.parquet.crypto.keytools.mocks.InMemoryKMS
 import org.apache.spark.{DebugFilesystem, SparkConf}
-import org.apache.spark.sql.{CometTestBase, SQLContext}
+import org.apache.spark.sql.{CometTestBase, SQLContext, functions}
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.test.SQLTestUtils
-
 import org.apache.comet.{CometConf, IntegrationTestSuite}
 import org.apache.comet.CometConf.{SCAN_NATIVE_COMET, SCAN_NATIVE_DATAFUSION, SCAN_NATIVE_ICEBERG_COMPAT}
 
@@ -394,11 +391,16 @@ class ParquetEncryptionITCase extends CometTestBase with SQLTestUtils {
         val parquetDF2 = spark.read.parquet(parquetDir2)
 
         val unionDF = parquetDF1.union(parquetDF2)
+        // Since the union has its own executeColumnar, problems would not surface if it is the last operator
+        // If we add another comet aggregate after the union, we see the need for the
+        // foreachUntilCometInput() in operator.scala
+        // as we would error on multiple native scan execs despite no longer being in the same plan at all
+        val aggDf = unionDF.agg(functions.sum("id"))
 
         if (CometConf.COMET_ENABLED.get(conf)) {
-          checkSparkAnswerAndOperator(unionDF)
+          checkSparkAnswerAndOperator(aggDf)
         } else {
-          checkSparkAnswer(unionDF)
+          checkSparkAnswer(aggDf)
         }
       }
     }