(spark) branch master updated: [SPARK-59252][SQL] Fix an SPJ scan reporting a different partitioning at execution than at planning

[pt...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

peter-toth pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/master by this push:
     new 8998b7ada107 [SPARK-59252][SQL] Fix an SPJ scan reporting a different partitioning at execution than at planning
8998b7ada107 is described below

commit 8998b7ada107df1e3eadc7937676b1fa1833b064
Author: Peter Toth <pe...@gmail.com>
AuthorDate: Mon Sep 7 09:29:10 2026 +0200

    [SPARK-59252][SQL] Fix an SPJ scan reporting a different partitioning at execution than at planning
    
    ### What changes were proposed in this pull request?
    
    `DataSourceV2ScanExecBase.outputPartitioning` becomes a `lazy val`, so a scan node answers for its whole life what it answered the planner.
    
    ### Why are the changes needed?
    
    It reads `conf.v2BucketingEnabled` off the live session conf, so as a `def` it can answer differently at execution time than it did at planning time. By then the plan is committed to the first answer.
    
    Two tables bucketed the same way, joined on the bucket column. Force the plan, turn the config off, run it:
    
        val df = sql("SELECT l.id FROM testcat.ns.l4 l JOIN testcat.ns.r4 r ON l.id = r.id")
        df.queryExecution.executedPlan       // 0 shuffles, 2 GroupPartitionsExec
        withSQLConf(SQLConf.V2_BUCKETING_ENABLED.key -> "false") {
          df.collect()
        }
    
        java.lang.ClassCastException: class org.apache.spark.sql.catalyst.plans.physical.UnknownPartitioning
          cannot be cast to class org.apache.spark.sql.catalyst.expressions.Expression
    
    The planner dropped both shuffles on the strength of the key-grouped layout and put a `GroupPartitionsExec` on each side. Those nodes ask the child for its partitioning again at execution, `GroupPartitionsExec.grouping` casts it to `Partitioning with Expression`, and by then the scan reports `UnknownPartitioning`.
    
    Measured on master, `branch-4.3` and `branch-4.2`: the same exception on all three, with the same plan shape. `branch-4.1` and older have no `GroupPartitionsExec`, so the crash site does not exist there.
    
    Everything else the body reads is fixed for the instance. `keyGroupedPartitioning` is a constructor field, every implementation of `inputPartitions` is already a `lazy val` (`BatchScanExec`, `MicroBatchScanExec`, `ContinuousScanExec`, `RealTimeStreamScanExec`), and `BatchScanExec.filteredPartitions` derives a new sequence rather than replacing `inputPartitions`. `FileSourceScanExec`, the V1 twin of this node, is already `override lazy val (outputPartitioning, outputOrdering)` over a c [...]
    
    The fix also removes repeated work, which is not the reason for it but is worth recording. As a `def` the key-grouped arm sorted every partition key and handed them to `KeyedPartitioning.apply`, which wraps each one and runs a `distinct`. Instrumenting the body and running `KeyGroupedPartitioningSuite` counted 33,051 executions as a `def` against 1,262 as a `lazy val`, so 26 per scan instance instead of one. Separately, in a microbenchmark over a two-column key, one call measured 35 u [...]
    
    ### Does this PR introduce _any_ user-facing change?
    
    Yes, it fixes the crash above. A plan keeps the bucketing setting it was planned under, which is what the planner already assumes: `ValidateRequirements` reads one child's partitioning twice and compares specs built from the two reads, `EnsureRequirements` reads it at five points and threads the results between them, and at `EnsureRequirements.scala:335` it is the pruning predicate of a `multiTransformDownWithPruning`, re-evaluated per generated alternative.
    
    One source-level note for out-of-tree code: Scala forbids a `def` overriding a `lazy val`, so a subclass of this trait that overrides `outputPartitioning` with a `def` needs to become a `lazy val` too. Binary compatibility is unaffected, since the trait still emits a default method carrying the initializer.
    
    ### How was this patch tested?
    
    A new `KeyGroupedPartitioningSuite` test builds the plan above, asserts the planner really did commit to the layout (no shuffles, two `GroupPartitionsExec`), then flips the config and checks the answer. It fails with the `ClassCastException` without the fix.
    
    Existing suites: `KeyGroupedPartitioningSuite`, `DataSourceV2CatalystRuntimeFilterSuite` for the runtime-filter re-plan path, `DataSourceV2Suite`, `GroupPartitionsExecSuite` and `EnsureRequirementsSuite`, plus the continuous and micro-batch scan suites, which reach this trait through the other `inputPartitions` implementations.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: Claude Code (Opus 5)
    
    Closes #58529 from peter-toth/SPARK-59252-memoize-scan-output-partitioning.
    
    Authored-by: Peter Toth <pe...@gmail.com>
    Signed-off-by: Peter Toth <pe...@gmail.com>
---
 .../datasources/v2/DataSourceV2ScanExecBase.scala       |  5 ++++-
 .../sql/connector/KeyGroupedPartitioningSuite.scala     | 17 +++++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

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 f00d8b9b82cb..cde8b1041a95 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 758e625c8389..02fc8afb7d9d 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
@@ -1240,6 +1240,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)


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org