Skip to content

[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
apache:masterfrom
peter-toth:SPARK-59252-memoize-scan-output-partitioning
Open

[SPARK-59252][SQL] Fix an SPJ scan reporting a different partitioning at execution than at planning#58529
peter-toth wants to merge 1 commit into
apache:masterfrom
peter-toth:SPARK-59252-memoize-scan-output-partitioning

Conversation

@peter-toth

Copy link
Copy Markdown
Contributor

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)

… 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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant