diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala index f00d8b9b82cb4..cde8b1041a952 100644 --- a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala +++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/DataSourceV2ScanExecBase.scala @@ -53,6 +53,7 @@ trait DataSourceV2ScanExecBase * `SupportsReportOrdering` */ def ordering: Option[Seq[SortOrder]] + /** Must be stable for the instance, since `outputPartitioning` memoizes over it. */ protected def inputPartitions: Seq[InputPartition] override def simpleString(maxFields: Int): String = { @@ -88,7 +89,9 @@ trait DataSourceV2ScanExecBase |""".stripMargin } - override def outputPartitioning: physical.Partitioning = { + // A `lazy val` because the planner asks a node for its partitioning many times, and the + // key-grouped arm sorts every partition key, wraps each one and runs a `distinct` over them. + @transient override lazy val outputPartitioning: physical.Partitioning = { keyGroupedPartitioning match { case Some(exprs) if conf.v2BucketingEnabled && KeyedPartitioning.supportsExpressions(exprs) && inputPartitions.nonEmpty && inputPartitions.forall(_.isInstanceOf[HasPartitionKey]) => diff --git a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala index 6a911ce652615..1fea555de92d7 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/connector/KeyGroupedPartitioningSuite.scala @@ -1186,6 +1186,23 @@ class KeyGroupedPartitioningSuite } } + test("SPARK-59252: a planned scan keeps the partitioning it was planned with") { + createBucketedIdTable("l4", 4) + createBucketedIdTable("r4", 4) + + val df = sql("SELECT l.id FROM testcat.ns.l4 l JOIN testcat.ns.r4 r ON l.id = r.id") + val plan = stripAQEPlan(df.queryExecution.executedPlan) + // The planner committed to the key-grouped layout: it dropped both shuffles and put a + // `GroupPartitionsExec` on each side. Those nodes ask their child for the partitioning again at + // execution, so the scan has to keep answering what it was planned with. + assert(collectShuffles(plan).isEmpty) + assert(collectGroupPartitions(plan).size == 2) + + withSQLConf(SQLConf.V2_BUCKETING_ENABLED.key -> "false") { + checkAnswer(df, (0 until 12).map(i => Row(i.toLong))) + } + } + test("partitioned join: join with two partition keys and matching & sorted partitions") { val items_partitions = Array(bucket(8, "id"), days("arrive_time")) createTable(items, itemsColumns, items_partitions)