fix: do not resolve DPP subqueries when AQE reads a native scan's partitioning

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

@@ -128,7 +128,14 @@ case class CometIcebergNativeScanExec(
   // Only accessed during execution, not planning
   def numPartitions: Int = perPartitionData.length
 
-  override lazy val outputPartitioning: Partitioning = UnknownPartitioning(numPartitions)
+  // Like BatchScanExec on Spark 3.5 and later without a key-grouped partitioning, the scan does
+  // not know its partition count while planning, because dynamic partition pruning decides which
+  // tasks survive and its subqueries resolve at execution. AQE reads outputPartitioning while it
+  // optimizes a stage (ValidateRequirements after CoalesceShufflePartitions), which is before
+  // CometPlanAdaptiveDynamicPruningFilters converts the DPP placeholders, so serializing the
+  // tasks here would run a placeholder subquery. At execution, CometNativeExec sizes the native
+  // RDD from perPartitionData instead.
+  override lazy val outputPartitioning: Partitioning = UnknownPartitioning(0)
 
   override lazy val outputOrdering: Seq[SortOrder] = Nil
 

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

@@ -120,27 +120,15 @@ case class CometNativeScanExec(
   // exposed for testing
   lazy val bucketedScan: Boolean = originalPlan.bucketedScan && !disableBucketedScan
 
-  override lazy val outputPartitioning: Partitioning = {
-    if (bucketedScan) {
-      originalPlan.outputPartitioning
-    } else {
-      // Use perPartitionData.length instead of originalPlan.inputRDD.getNumPartitions.
-      //
-      // originalPlan.inputRDD triggers FileSourceScanExec's full scan pipeline including
-      // codegen on partition filter expressions. With DPP, this calls
-      // InSubqueryExec.doGenCode which requires the subquery to have finished - but
-      // outputPartitioning can be accessed before prepare() runs (e.g., by
-      // ValidateRequirements during plan validation).
-      //
-      // perPartitionData goes through serializedPartitionData, which explicitly resolves
-      // DPP subqueries (via updateResult()) before accessing file partitions. This is the
-      // same pattern CometIcebergNativeScanExec uses.
-      //
-      // This is also more correct: perPartitionData.length reflects the post-DPP partition
-      // count, matching what CometExecRDD actually uses in doExecuteColumnar().
-      UnknownPartitioning(perPartitionData.length)
-    }
-  }
+  // Like FileSourceScanExec, only a bucketed scan knows its partitioning while planning. The
+  // file-split count of any other scan depends on dynamic partition pruning, whose subqueries
+  // resolve at execution. AQE reads outputPartitioning while it optimizes a stage
+  // (ValidateRequirements after CoalesceShufflePartitions), which is before
+  // CometPlanAdaptiveDynamicPruningFilters converts the DPP placeholders, so computing the file
+  // partitions here would run a placeholder subquery. At execution, CometNativeExec sizes the
+  // native RDD from perPartitionData instead.
+  override lazy val outputPartitioning: Partitioning =
+    if (bucketedScan) originalPlan.outputPartitioning else UnknownPartitioning(0)
 
   override lazy val outputOrdering: Seq[SortOrder] = originalPlan.outputOrdering
 

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

@@ -1107,6 +1107,13 @@ abstract class CometNativeExec extends CometExec {
     // broadcast plan.
     val (firstNonBroadcastPlanRDD, firstNonBroadcastPlanNumPartitions) =
       firstNonBroadcastPlan.get._1 match {
+        // Plan-data scans, Iceberg included, are sized from perPartitionData, which
+        // findAllPlanData above has resolved. A bucketed scan has one entry per bucket, after any
+        // bucket coalescing, so the count still matches its HashPartitioning.
+        case scan: CometScanWithPlanData =>
+          (null.asInstanceOf[RDD[Any]], scan.perPartitionData.length)
+        case scan: CometIcebergNativeScanExec =>
+          (null.asInstanceOf[RDD[Any]], scan.numPartitions)
         case plan: CometNativeExec =>
           (null.asInstanceOf[RDD[Any]], plan.outputPartitioning.numPartitions)
         case plan =>
@@ -1346,6 +1353,17 @@ abstract class CometLeafExec extends CometNativeExec with LeafExecNode {
  * parent's native execution receives an empty input. (`CometIcebergNativeScanExec` does NOT use
  * this trait; it has a dedicated `findAllPlanData` case.)
  *
+ * `perPartitionData.length` is the partition count the native block runs with:
+ * `CometNativeExec.buildNativeContext` sizes the native RDD from it at execution, and the count
+ * in `outputPartitioning` is not used for that.
+ *
+ * `outputPartitioning` must not read `perPartitionData`, or anything else that evaluates the
+ * scan's DPP subqueries. AQE calls it while optimizing a stage, before
+ * `CometPlanAdaptiveDynamicPruningFilters` has converted the adaptive DPP placeholders, so
+ * computing it from those runs a placeholder subquery. A scan that cannot report a real
+ * partitioning without them should report `UnknownPartitioning(0)`, as `CometNativeScanExec` does
+ * for a non-bucketed scan.
+ *
  * Each implementation also resolves its own DPP subqueries via `ensureSubqueriesResolved` before
  * `commonData`/`perPartitionData` are read. That method lives on [[CometLeafExec]], so the `self:
  * CometLeafExec` self-type below makes "is a leaf scan" a compile-time requirement: an

spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala

@@ -36,7 +36,7 @@ import org.apache.spark.sql.catalyst.expressions.DynamicPruningExpression
 import org.apache.spark.sql.comet._
 import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec}
 import org.apache.spark.sql.execution.{InSubqueryExec, ReusedSubqueryExec, SparkPlan, SubqueryAdaptiveBroadcastExec, SubqueryBroadcastExec, SubqueryExec}
-import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, BroadcastQueryStageExec}
+import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, AdaptiveSparkPlanHelper, AQEShuffleReadExec, BroadcastQueryStageExec}
 import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
 import org.apache.spark.sql.execution.exchange.{ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike}
 import org.apache.spark.sql.functions.col
@@ -5359,6 +5359,102 @@ class CometIcebergNativeSuite
     }
   }
 
+  // AQE runs CoalesceShufflePartitions before CometPlanAdaptiveDynamicPruningFilters. When that
+  // rule changes a stage, AQE validates the stage by reading outputPartitioning on every node in
+  // it. With Spark's UnionExec, the aggregate's shuffle read shares the final stage with the DPP
+  // scan and gets coalesced, so the scan is asked for its partitioning while its DPP subquery is
+  // still the adaptive placeholder. Answering must neither run that subquery nor lose the pruning.
+  test("AQE DPP - shuffle coalescing in the scan's stage keeps DPP working") {
+    assume(icebergAvailable, "Iceberg not available")
+    withTempIcebergDir { warehouseDir =>
+      val dimDir = new File(warehouseDir, "dim_parquet")
+      val otherDir = new File(warehouseDir, "other_parquet")
+      withSQLConf(
+        "spark.sql.catalog.aqe_cat" -> "org.apache.iceberg.spark.SparkCatalog",
+        "spark.sql.catalog.aqe_cat.type" -> "hadoop",
+        "spark.sql.catalog.aqe_cat.warehouse" -> warehouseDir.getAbsolutePath,
+        SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "1KB",
+        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+        SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true",
+        SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
+        SQLConf.COALESCE_PARTITIONS_PARALLELISM_FIRST.key -> "false",
+        CometConf.COMET_ENABLED.key -> "true",
+        CometConf.COMET_EXEC_ENABLED.key -> "true",
+        CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true",
+        CometConf.COMET_EXEC_UNION_ENABLED.key -> "false") {
+
+        spark.sql("""
+          CREATE TABLE aqe_cat.db.coalesce_fact (amount INT, store_id INT)
+          USING iceberg PARTITIONED BY (store_id)
+        """)
+        // One data file per store, so an unpruned scan reads 10 files and a pruned one reads 1.
+        spark
+          .range(100)
+          .selectExpr("cast(id as int) as amount", "cast(id % 10 as int) as store_id")
+          .coalesce(1)
+          .sortWithinPartitions("store_id")
+          .writeTo("aqe_cat.db.coalesce_fact")
+          .append()
+
+        spark
+          .range(10)
+          .selectExpr("cast(id as int) as store_id", "cast(id as string) as country")
+          .write
+          .parquet(dimDir.getAbsolutePath)
+        spark.read.parquet(dimDir.getAbsolutePath).createOrReplaceTempView("coalesce_dim")
+        spark
+          .range(100)
+          .selectExpr("cast(id % 7 as int) as k", "cast(id as int) as v")
+          .write
+          .parquet(otherDir.getAbsolutePath)
+        spark.read.parquet(otherDir.getAbsolutePath).createOrReplaceTempView("coalesce_other")
+
+        val query =
+          """SELECT /*+ BROADCAST(d) */ f.amount AS a, f.store_id AS b
+            |FROM aqe_cat.db.coalesce_fact f JOIN coalesce_dim d ON f.store_id = d.store_id
+            |WHERE d.country = '3'
+            |UNION ALL
+            |SELECT k AS a, CAST(count(*) AS INT) AS b
+            |FROM coalesce_other GROUP BY k""".stripMargin
+        val (_, cometPlan) = checkSparkAnswer(query)
+        assertIcebergNativeScanPresent(cometPlan)
+        assertNoLeftoverCSAB(cometPlan)
+
+        // On 3.4, Spark's own DPP rule plans the subquery before coalescing, so check results only.
+        if (isSpark35Plus) {
+          val dppScans = collectIcebergNativeScans(cometPlan).filter(
+            _.runtimeFilters.exists(_.isInstanceOf[DynamicPruningExpression]))
+          assert(dppScans.size == 1, s"Expected one DPP Iceberg scan:\n${cometPlan.treeString}")
+          val scan = dppScans.head
+          // The test only covers the bug if coalescing changed the stage holding the DPP scan.
+          // That is the final stage, and TreeNode.collect stops at its query stage leaves.
+          val finalStage = stripAQEPlan(cometPlan)
+          assert(
+            finalStage.find(_ eq scan).nonEmpty &&
+              finalStage.collect {
+                case r: AQEShuffleReadExec if r.isCoalescedRead => r
+              }.nonEmpty,
+            "Expected the DPP scan and a coalesced shuffle read in the final stage:\n" +
+              cometPlan.treeString)
+          assertCsbBroadcastReuse(collectIcebergDPPSubqueries(cometPlan), cometPlan)
+          // The dimension filter keeps store 3 only, so DPP reads 1 of the 10 data files.
+          // Iceberg filters its already planned tasks at runtime, so count the serialized tasks
+          // and the splits read rather than Iceberg's planning metrics.
+          val plannedTasks = scan.perPartitionData
+            .map(OperatorOuterClass.IcebergScan.parseFrom(_).getFileScanTasksCount)
+            .sum
+          val numSplits = scan.metrics("num_splits").value
+          assert(
+            plannedTasks == 1 && numSplits == 1,
+            s"Expected DPP to prune to 1 of 10 files, planned $plannedTasks tasks and read " +
+              s"$numSplits splits:\n${cometPlan.treeString}")
+        }
+
+        spark.sql("DROP TABLE aqe_cat.db.coalesce_fact")
+      }
+    }
+  }
+
   test("AQE DPP - cross-stage scalar subquery with broadcast reuse") {
     assume(icebergAvailable, "Iceberg not available")
     withTempIcebergDir { warehouseDir =>

spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala

@@ -1279,6 +1279,153 @@ class CometExecSuite extends CometTestBase {
     }
   }
 
+  // AQE runs CoalesceShufflePartitions before CometPlanAdaptiveDynamicPruningFilters. When that
+  // rule changes a stage, AQE validates the stage by reading outputPartitioning on every node in
+  // it. With Spark's UnionExec, the aggregate's shuffle read shares the final stage with the DPP
+  // scan and gets coalesced, so the scan is asked for its partitioning while its DPP subquery is
+  // still the adaptive placeholder. Answering must not run that subquery.
+  test("AQE DPP: shuffle coalescing in the scan's stage keeps DPP working") {
+    withTempDir { dir =>
+      val path = dir.getAbsolutePath
+      withSQLConf(CometConf.COMET_EXEC_ENABLED.key -> "false") {
+        spark
+          .range(100)
+          .selectExpr("cast(id % 10 as int) as store_id", "cast(id as int) as amount")
+          .write
+          .partitionBy("store_id")
+          .parquet(s"$path/fact")
+        spark
+          .range(10)
+          .selectExpr("cast(id as int) as store_id", "cast(id as string) as country")
+          .write
+          .parquet(s"$path/dim")
+        spark
+          .range(100)
+          .selectExpr("cast(id % 7 as int) as k", "cast(id as int) as v")
+          .write
+          .parquet(s"$path/other")
+      }
+
+      withSQLConf(
+        SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
+        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+        SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true",
+        SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "true",
+        SQLConf.COALESCE_PARTITIONS_PARALLELISM_FIRST.key -> "false",
+        CometConf.COMET_EXEC_UNION_ENABLED.key -> "false") {
+        spark.read.parquet(s"$path/fact").createOrReplaceTempView("dpp_coalesce_fact")
+        spark.read.parquet(s"$path/dim").createOrReplaceTempView("dpp_coalesce_dim")
+        spark.read.parquet(s"$path/other").createOrReplaceTempView("dpp_coalesce_other")
+
+        val df = spark.sql("""SELECT f.amount AS a, f.store_id AS b
+            |FROM dpp_coalesce_fact f JOIN dpp_coalesce_dim d ON f.store_id = d.store_id
+            |WHERE d.country = '3'
+            |UNION ALL
+            |SELECT k AS a, CAST(count(*) AS INT) AS b
+            |FROM dpp_coalesce_other GROUP BY k""".stripMargin)
+        val (_, cometPlan) = checkSparkAnswer(df)
+
+        if (isSpark35Plus) {
+          val dppScans = collect(cometPlan) {
+            case s: CometNativeScanExec
+                if s.partitionFilters.exists(_.isInstanceOf[DynamicPruningExpression]) =>
+              s
+          }
+          assert(dppScans.size == 1, s"Expected one DPP native scan:\n${cometPlan.treeString}")
+          // The test only covers the bug if coalescing changed the stage holding the DPP scan.
+          // That is the final stage, and TreeNode.collect stops at its query stage leaves.
+          val finalStage = stripAQEPlan(cometPlan)
+          assert(
+            finalStage.find(_ eq dppScans.head).nonEmpty &&
+              finalStage.collect {
+                case r: AQEShuffleReadExec if r.isCoalescedRead => r
+              }.nonEmpty,
+            "Expected the DPP scan and a coalesced shuffle read in the final stage:\n" +
+              cometPlan.treeString)
+          assertAqeDppShape(cometPlan, expectedCometSubqueryBroadcasts = Some(1))
+          // The dimension filter keeps store 3 only, so DPP reads 1 of the 10 partitions.
+          assert(
+            dppScans.head.metrics("numPartitions").value == 1,
+            s"Expected DPP to prune to 1 partition:\n${cometPlan.treeString}")
+        }
+      }
+    }
+  }
+
+  // A bucketed scan reports its bucket layout, which DPP does not change. The sort-merge join on
+  // the bucket column keeps relying on it, so neither join side is shuffled, while DPP from the
+  // broadcast dimension join still prunes the fact table's partitions.
+  test("AQE DPP: bucketed scan keeps its partitioning for a sort-merge join") {
+    withTable("dpp_bucket_fact", "dpp_bucket_dim", "dpp_bucket_other") {
+      withSQLConf(CometConf.COMET_EXEC_ENABLED.key -> "false") {
+        spark
+          .range(100)
+          .selectExpr(
+            "cast(id % 10 as int) as store_id",
+            "cast(id % 7 as int) as k",
+            "cast(id as int) as amount")
+          .write
+          .partitionBy("store_id")
+          .bucketBy(4, "k")
+          .format("parquet")
+          .saveAsTable("dpp_bucket_fact")
+        spark
+          .range(10)
+          .selectExpr("cast(id as int) as store_id", "cast(id as string) as country")
+          .write
+          .format("parquet")
+          .saveAsTable("dpp_bucket_dim")
+        spark
+          .range(50)
+          .selectExpr("cast(id % 7 as int) as k", "cast(id as int) as v")
+          .write
+          .bucketBy(4, "k")
+          .format("parquet")
+          .saveAsTable("dpp_bucket_other")
+      }
+
+      withSQLConf(
+        SQLConf.USE_V1_SOURCE_LIST.key -> "parquet",
+        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+        SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true",
+        SQLConf.BUCKETING_ENABLED.key -> "true") {
+        val df = sql("""SELECT /*+ BROADCAST(d), MERGE(o) */ f.amount, f.store_id, o.v
+            |FROM dpp_bucket_fact f
+            |JOIN dpp_bucket_dim d ON f.store_id = d.store_id
+            |JOIN dpp_bucket_other o ON f.k = o.k
+            |WHERE d.country = '3'""".stripMargin)
+        val (_, cometPlan) = checkSparkAnswer(df)
+
+        val smjs = collect(cometPlan) {
+          case j: SortMergeJoinExec => j
+          case j: CometSortMergeJoinExec => j
+        }
+        assert(smjs.size == 1, s"Expected one sort-merge join:\n${cometPlan.treeString}")
+        val shuffles = collect(cometPlan) { case s: ShuffleExchangeLike => s }
+        assert(
+          shuffles.isEmpty,
+          s"Expected no shuffle around the bucketed join:\n${cometPlan.treeString}")
+
+        if (isSpark35Plus) {
+          val dppScans = collect(cometPlan) {
+            case s: CometNativeScanExec
+                if s.partitionFilters.exists(_.isInstanceOf[DynamicPruningExpression]) =>
+              s
+          }
+          assert(dppScans.size == 1, s"Expected one DPP native scan:\n${cometPlan.treeString}")
+          assert(
+            dppScans.head.bucketedScan,
+            s"Expected a bucketed scan:\n${cometPlan.treeString}")
+          assertAqeDppShape(cometPlan, expectedCometSubqueryBroadcasts = Some(1))
+          // The dimension filter keeps store 3 only, so DPP reads 1 of the 10 partitions.
+          assert(
+            dppScans.head.metrics("numPartitions").value == 1,
+            s"Expected DPP to prune to 1 partition:\n${cometPlan.treeString}")
+        }
+      }
+    }
+  }
+
   // DPP subquery reuses the join's broadcast via AQE stageCache.
   test("AQE DPP: broadcast exchange reuse") {
     withDppTables {

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

@@ -1489,7 +1489,7 @@ class CometTaskMetricsSuite extends CometTestBase with AdaptiveSparkPlanHelper {
           run(plainQuery, confs)
         val plainScan = collectFirst(plainPlan) { case s: CometNativeScanExec => s }.get
         assert(
-          plainScan.outputPartitioning.numPartitions == 4,
+          plainScan.perPartitionData.length == 4,
           s"Expected one scan partition per file:\n${plainPlan.treeString}")
         assert(plainTasks == 4, s"Expected 4 tasks at interval $interval, got $plainTasks")
         assert(plainRows == totalRows, s"output_rows at interval $interval: $plainRows")