[SPARK-59252][SQL] Fix an SPJ scan reporting a different partitioning at execution than at planning - #58529
Open
peter-toth wants to merge 1 commit into
Conversation
… 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 conf-derived `bucketedScan` `lazy val`, so this aligns V2 with V1.
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 us at 100 partitions, 124 us at 1,000 and 1,515 us at 10,000.
### 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)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changes were proposed in this pull request?
DataSourceV2ScanExecBase.outputPartitioningbecomes alazy val, so a scan node answers for its whole life what it answered the planner.Why are the changes needed?
It reads
conf.v2BucketingEnabledoff the live session conf, so as adefit 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:
The planner dropped both shuffles on the strength of the key-grouped layout and put a
GroupPartitionsExecon each side. Those nodes ask the child for its partitioning again at execution,GroupPartitionsExec.groupingcasts it toPartitioning with Expression, and by then the scan reportsUnknownPartitioning.Measured on master,
branch-4.3andbranch-4.2: the same exception on all three, with the same plan shape.branch-4.1and older have noGroupPartitionsExec, so the crash site does not exist there.Everything else the body reads is fixed for the instance.
keyGroupedPartitioningis a constructor field, every implementation ofinputPartitionsis already alazy val(BatchScanExec,MicroBatchScanExec,ContinuousScanExec,RealTimeStreamScanExec), andBatchScanExec.filteredPartitionsderives a new sequence rather than replacinginputPartitions.FileSourceScanExec, the V1 twin of this node, is alreadyoverride lazy val (outputPartitioning, outputOrdering)over a conf-derivedbucketedScanlazy val, so this aligns V2 with V1.The fix also removes repeated work, which is not the reason for it but is worth recording. As a
defthe key-grouped arm sorted every partition key and handed them toKeyedPartitioning.apply, which wraps each one and runs adistinct. Instrumenting the body and runningKeyGroupedPartitioningSuitecounted 33,051 executions as adefagainst 1,262 as alazy val, so 26 per scan instance instead of one. Separately, in a microbenchmark over a two-column key, one call measured 35 us at 100 partitions, 124 us at 1,000 and 1,515 us at 10,000.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:
ValidateRequirementsreads one child's partitioning twice and compares specs built from the two reads,EnsureRequirementsreads it at five points and threads the results between them, and atEnsureRequirements.scala:335it is the pruning predicate of amultiTransformDownWithPruning, re-evaluated per generated alternative.One source-level note for out-of-tree code: Scala forbids a
defoverriding alazy val, so a subclass of this trait that overridesoutputPartitioningwith adefneeds to become alazy valtoo. Binary compatibility is unaffected, since the trait still emits a default method carrying the initializer.How was this patch tested?
A new
KeyGroupedPartitioningSuitetest builds the plan above, asserts the planner really did commit to the layout (no shuffles, twoGroupPartitionsExec), then flips the config and checks the answer. It fails with theClassCastExceptionwithout the fix.Existing suites:
KeyGroupedPartitioningSuite,DataSourceV2CatalystRuntimeFilterSuitefor the runtime-filter re-plan path,DataSourceV2Suite,GroupPartitionsExecSuiteandEnsureRequirementsSuite, plus the continuous and micro-batch scan suites, which reach this trait through the otherinputPartitionsimplementations.Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Opus 5)