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 31e518cfcf8..e3466c71d27 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 @@ -26,7 +26,7 @@ import org.apache.spark.rdd.RDD import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, SortOrder} import org.apache.spark.sql.catalyst.plans.QueryPlan -import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, UnknownPartitioning} +import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, SinglePartition, UnknownPartitioning} import org.apache.spark.sql.execution.{SparkPlan, SQLExecution} import org.apache.spark.sql.execution.datasources.v2.BatchScanExec import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} @@ -135,7 +135,18 @@ case class CometIcebergNativeScanExec( // 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) + // + // Spark 3.4's BatchScanExec reports SinglePartition for one input partition, which it counts + // before any DPP filtering, and Spark plans no exchange above it on that basis. The scan reports + // the same. Otherwise EnsureRequirements, which AQE runs again over a DPP subquery's converted + // plan, inserts an exchange inside a native block that is already serialized (#7422). + // serializePartitions keeps that one partition when DPP prunes it, as Spark does. + override lazy val outputPartitioning: Partitioning = + if (originalPlan != null && originalPlan.outputPartitioning == SinglePartition) { + SinglePartition + } else { + UnknownPartitioning(0) + } override lazy val outputOrdering: Seq[SortOrder] = Nil diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index 4153498516f..b8835212129 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -6743,8 +6743,6 @@ class CometIcebergNativeSuite // Spark's V2 path. test("AQE DPP - V2 broadcast query stage creation order (SPARK-34637)") { assume(icebergAvailable, "Iceberg not available") - // https://github.com/apache/datafusion-comet/issues/7422 - assume(isSpark35Plus, "Spark 3.4 fails the DPP build side, see #7422") withTempIcebergDir { warehouseDir => withSQLConf( "spark.sql.catalog.aqe_cat" -> "org.apache.iceberg.spark.SparkCatalog", @@ -6780,12 +6778,14 @@ class CometIcebergNativeSuite assertNoLeftoverCSAB(cometPlan) - val cometSubqueries = collectWithSubqueries(cometPlan) { - case s: CometSubqueryBroadcastExec => s + if (isSpark35Plus) { + val cometSubqueries = collectWithSubqueries(cometPlan) { + case s: CometSubqueryBroadcastExec => s + } + assert( + cometSubqueries.nonEmpty, + s"Expected CometSubqueryBroadcastExec for DPP, got 0:\n${cometPlan.treeString}") } - assert( - cometSubqueries.nonEmpty, - s"Expected CometSubqueryBroadcastExec for DPP, got 0:\n${cometPlan.treeString}") spark.sql("DROP TABLE aqe_cat.db.q34637_fact") }