diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala index e058a5c0916..dd644cbc132 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergNativeScanExec.scala +++ b/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 diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala index 37694992211..8cb580db812 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala +++ b/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 diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index bcae614556b..d033353daaa 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/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 diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index f271f7d6ec2..36ec57fbbaf 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/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 => diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index 478c62f4eda..79438775b25 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/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 { diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala index 318a416ce47..9bf5e955954 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/CometTaskMetricsSuite.scala +++ b/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")