[SPARK-59248][SQL] Keep storage-partitioned join when partition keys are pruned from the scan output

[ulysses-you]

What changes were proposed in this pull request?

Keep the reported key-grouped partitioning when some of its keys are column-pruned out of the scan
output, instead of dropping the whole partitioning, so a storage-partitioned join still applies.

• V2ScanPartitioningAndOrdering keeps the reported keys as long as at least one of them is still in
  the scan output. The full key list is kept, not the surviving subset, to stay positionally aligned
  with the raw partition-key rows.
• DataSourceV2ScanExecBase splits the two views of that report. reportedKeyedPartitioning is the
  source's own, at full key width, and is what reads and orders the raw HasPartitionKey rows
  (BatchScanExec.filteredPartitions). outputPartitioning projects the pruned key positions away,
  so the partitioning Spark plans against only references output columns. replanWithRuntimeFilters
  takes Option[KeyedPartitioning], so the projected view cannot reach the raw key rows.
• DataSourceV2ScanRelation.doCanonicalize and BatchScanExec (doCanonicalize / equals /
  hashCode, through a shared prunedKeyGroupedPartitioning) ignore partitioning keys that
  reference pruned columns, and the canonicalized ordering keeps only its leading run of sort orders
  over output columns. Otherwise a dangling attribute's exprId keeps otherwise-equivalent scans
  unequal and defeats subplan merging and exchange reuse.

Whether the surviving projection is usable is decided one layer down, by KeyedPartitioning:
pruning a key can map keys that were distinct onto the same key, and such a projection is marked
isCollapsed, so grouping it stays gated on
spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled in mayGroupToSatisfy. A
projection that collapses nothing leaves the keys unique and needs no opt-in. No config doc changes:
the config keeps the roles it already has.

Why are the changes needed?

An extra partition key that is column-pruned out of the scan output should not disable a
storage-partitioned join. Dropping the whole report also lost it for the case where the surviving
keys are still unique, which is sound with no opt-in at all.

Background: SPARK-40429 (#37886) introduced the drop-to-None guard to avoid a "missing inputs"
plan-validation error when a partitioning key referenced a pruned column. That failure mode no
longer applies: SPARK-40259 made DataSourceV2ScanRelation.references return empty (partitioning,
ordering and pushed filters are scan metadata, not references to resolve), so a dangling
partitioning key no longer trips plan validation. This change relies on that. SPARK-40259 is on
master only, so this change is master-only too and is not backportable to branch-4.x as it
stands.

Does this PR introduce any user-facing change?

Yes. When a partition key is pruned out of the scan output, Spark now reports the projected
partitioning instead of no partitioning, so a query that previously shuffled can plan a
storage-partitioned join:

• with spark.sql.sources.v2.bucketing.allowKeysSubsetOfPartitionKeys.enabled off (default), when
  the surviving keys are still unique;
• with it on, also when pruning collapses distinct keys onto the same key.

Results are unchanged either way.

How was this patch tested?

New tests in KeyGroupedPartitioningSuite:

• a join whose key is a subset of the partition keys, with the extra key pruned, triggers SPJ with
  allowKeysSubsetOfPartitionKeys on and shuffles with it off;
• a pruned key that collapses nothing keeps SPJ with the config off;
• a pruned leading key keeps the raw key rows aligned: one test on row loss, one on the
  ClassCastException from reading a leading String key at the projected Integer type. Pointing
  filteredPartitions at the projected partitioning fails exactly those two;
• all partition keys pruned reports no partitioning, even with the config on;
• a self-join with a pruned key stays canonicalizable, and a pruned key does not defeat plan reuse.

MergeSubplansSuite: identical scans whose reported ordering is on a pruned column are
deduplicated, not fused. Verified this test fails without the ordering half of the canonicalization
change.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (qwen3.8-max, Claude Opus 5)

[ulysses-you]

cc @peter-toth@cloud-fan thank you

[peter-toth]

Thanks for the PR, @ulysses-you!

This closes the gap SPARK-46367 deferred. V2ScanPartitioningAndOrdering dropped the whole reported partitioning as soon as one key referenced a pruned column. An unselected extra partition key therefore cost the storage-partitioned join. Keeping the full key list on the logical node and projecting the pruned positions away in the physical outputPartitioning is the right shape.

The projection has one blocking problem. It changes what the partitioning's keys and key types describe. BatchScanExec.filteredPartitions still reads the raw HasPartitionKey.partitionKey() rows with them. A pruned key that is not the trailing one then silently drops join rows. Every existing test prunes the trailing key, which is why the suite stays green.

Blocking

• 1.Projected partitioning read against full-width partition-key rows:outputPartitioning returns projected keys and projected key types, and filteredPartitions applies them to the unprojected partitionKey() rows. Measured: a join silently loses rows, and a differently-typed leading pruned key throws ClassCastException. inline

Non-blocking

• 2.The config gate on the report is stricter than the safety property: projecting a key away needs the opt-in only when it collapses keys, and KeyedPartitioning.mayGroupToSatisfy already gates exactly that. Measured with the conjunct dropped: the unique-projected-keys case gets SPJ with the config off, and the collapsing case still falls back to a shuffle. inline
• 3.a pruned source-reported ordering must not defeat exchange reuse passes on master:BatchScanExec never compares ordering, so a dangling ordering could not block physical reuse. The MergeSubplansSuite test is the one that covers the takeWhile. inline
• 4.The config's doc() does not mention its new role:allowKeysSubsetOfPartitionKeys now also decides whether a partitioning is reported at all. inline

Minor

• 5.Stale citation:PlanMerger.scala:985 says the partitioning pass is "reference-subset guarded". After this PR it drops the report only when every key is pruned, or when the config is off. The conclusion around it still holds, only the parenthetical needs updating.

[peter-toth]

Finding 1.project rebuilds the key rows onto resolvablePositions, so from here the partitioning's partitionKeys are projected rows and its keyDataTypes are the projected types. BatchScanExec.filteredPartitions passes this same partitioning to PushDownUtils.replanWithRuntimeFilters, which reads the raw, full-width HasPartitionKey.partitionKey() rows with it:

• parts.sortBy(_.asInstanceOf[HasPartitionKey].partitionKey())(k.keyRowOrdering) at sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala:357
• getInternalRowComparableWrapperFactory(k.keyDataTypes) applied to partitionKey() at sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala:316

Those line up only when the surviving positions are the leading prefix. Every existing test prunes the trailing key, so nothing catches it.

Wrong results.GroupPartitionsExec.grouping labels child RDD partition i with childKp.partitionKeys(i). partitionKeys is ordered by the full key, while filteredPartitions is ordered by the projected prefix. The two disagree wherever the full-key sort reorders inside a prefix group. Measured on a 7-split table partitioned by (store_id, dept_id) with store_id pruned:

rdd partition keys (store,dept): (1,10), (1,20), (1,30), (2,5), (2,40), (3,7), (3,1)
partitioning keys (dept): 10, 20, 30, 5, 40, 1, 7

Indices 5 and 6 carry each other's label. A join that lays the other side out on those keys loses exactly those two:

valcols=Array(
Column.create("store_id", IntegerType),
Column.create("dept_id", IntegerType))
createTable("t", cols, Array(identity("store_id"), identity("dept_id")))
sql("INSERT INTO testcat.ns.t VALUES (1, 20), (1, 10), (1, 30), (2, 5), (2, 40), (3, 7), (3, 1)")
withTempView("other") {
spark.range(1, 41).selectExpr("cast(id as int) as dept_id").createOrReplaceTempView("other")
withSQLConf(
SQLConf.ADAPTIVE_EXECUTION_ENABLED.key ->"false",
SQLConf.V2_BUCKETING_SHUFFLE_ENABLED.key ->"true",
SQLConf.V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.key ->"true") {
valdf= sql("SELECT /*+ MERGE(t, o) */ t.dept_id "+"FROM testcat.ns.t t JOIN other o ON t.dept_id = o.dept_id")
assert(df.collect().map(_.getInt(0)).sorted.toSeq ===Seq(1, 5, 7, 10, 20, 30, 40))
}
}

This branch returns 5, 10, 20, 30, 40. With allowKeysSubsetOfPartitionKeys off the same query is correct, so it is this path.

ClassCastException when the pruned leading key has a different type. Two tables partitioned by (identity("data"), identity("id")), joined on id with data pruned, throw java.lang.ClassCastException: class org.apache.spark.unsafe.types.UTF8String cannot be cast to class java.lang.Integer at PushDownUtils.scala:357. keyRowOrdering comes from the projected [IntegerType] and is applied to a (String, Int) row.

Fix. Give replanWithRuntimeFilters the unprojected partitioning: its keys are the raw rows, in the order it sorts the input partitions into. Splitting the construction out of outputPartitioning is enough:

/** * The partitioning as the source reported it: one key per input partition, holding every reported * key position, in the order a consumer must sort the input partitions into.*/protecteddefreportedKeyedPartitioning:Option[KeyedPartitioning] = {
keyGroupedPartitioning match {
caseSome(exprs) if conf.v2BucketingEnabled &&KeyedPartitioning.supportsExpressions(exprs) &&
inputPartitions.nonEmpty && inputPartitions.forall(_.isInstanceOf[HasPartitionKey]) =>valdataTypes= exprs.map(_.dataType)
valrowOrdering=RowOrdering.createNaturalAscendingOrdering(dataTypes)
valpartitionKeys=
inputPartitions.map(_.asInstanceOf[HasPartitionKey].partitionKey()).sorted(rowOrdering)
Some(KeyedPartitioning(exprs, partitionKeys))
case _ =>None
}
}
/** * What orders the input partitions. `outputPartitioning` may have projected key positions away, * and its keys and key types then no longer describe the raw `HasPartitionKey.partitionKey()` * rows, so a consumer reading those rows must take the types and the order from here.*/protecteddefinputPartitionOrdering: physical.Partitioning=if (outputPartitioning.isInstanceOf[KeyedPartitioning]) {
reportedKeyedPartitioning.getOrElse(super.outputPartitioning)
} else {
super.outputPartitioning
}
overridedefoutputPartitioning: physical.Partitioning= reportedKeyedPartitioning match {
caseSome(partitioning) =>valexprs= partitioning.expressions
valresolvablePositions= exprs.indices.filter(i => exprs(i).references.subsetOf(outputSet))
if (resolvablePositions.isEmpty) super.outputPartitioning
else partitioning.project(resolvablePositions)
case _ =>super.outputPartitioning
}

with BatchScanExec.filteredPartitions passing inputPartitionOrdering instead of outputPartitioning. I ran KeyGroupedPartitioningSuite, DataSourceV2CatalystRuntimeFilterSuite and MergeSubplansSuite with that applied plus the two cases above: 245 tests, all green.

[ulysses-you]

Same issue @sunchao raised; fixed in 973bde7. reportedKeyedPartitioning carries the source's
report at full key width and is what filteredPartitions passes down; outputPartitioning stays
projected for planning. replanWithRuntimeFilters now takes Option[KeyedPartitioning], so the
projected view cannot reach the raw key rows.

[peter-toth]

Finding 2. The opt-in is needed for the grouping, not for the report. Projecting a key position away is a collapse only when two distinct source keys land on the same projected key. KeyedPartitioning.project sets isCollapsed exactly then, and mayGroupToSatisfy refuses to group a collapsed partitioning without this config. The safety property is therefore already enforced one layer down, on the case that needs it.

What this conjunct costs is the other case: a pruned key whose removal collapses nothing. The projected keys stay unique, isGrouped is true, and satisfies(ClusteredDistribution) holds whatever the config says.

That split is the one SPARK-46367 (e656d04c157) already drew for the sibling narrowing in PartitioningPreservingUnaryExecNode: distinct projected keys need no config, duplicate projected keys require allowKeysSubsetOfPartitionKeys. KeyedPartitioning's class doc argues the same for not gating the report - "Reporting UnknownPartitioning would give up all of them, and make the plan shape depend on a config."

I measured both halves with the conjunct dropped, each with allowKeysSubsetOfPartitionKeys=false:

• partitioned by (dept_id, store_id), one store_id per dept_id, store_id pruned, join on dept_id: no shuffle, correct answer.
• partitioned by (id, data), two data values per id, data pruned, join on id: two shuffles, correct answer. mayGroupToSatisfy refuses, as it should.

So else if (inOutput.exists(identity)) looks like the right gate.

[ulysses-you]

Agreed, gate removed in 973bde7.

Your mechanism holds: project marks isCollapsed exactly on the collapsing case and
mayGroupToSatisfy keeps the config over it, so the rule only has to decide whether any key
survives. The rule also has no partition key values in hand, so it cannot tell the two cases apart
even in principle.

The case the conjunct was costing is now pinned by a pruned key that collapses nothing keeps SPJ without the config, with allowKeysSubsetOfPartitionKeys set to false explicitly.

[peter-toth]

Finding 3. This one passes on master. BatchScanExec never looks at ordering - equals, hashCode and doCanonicalize all leave it out - so a dangling reported ordering could not block physical exchange reuse in the first place.

Measured two ways. On e261626f152 with only the test files applied, it passes. With the whole PR applied but the takeWhile in DataSourceV2ScanRelation.doCanonicalize reverted, it still passes. The other tests behave as expected under the same treatment: "join key subset ..." and "self-join ..." fail on base, and "a pruned partition key must not defeat plan reuse" fails when only BatchScanExec.scala is reverted.

The MergeSubplansSuite test is the one that covers the takeWhile - it fails with that hunk alone reverted, matching your description. So either drop this test, or point it at the logical path (sameResult over two DataSourceV2ScanRelations), where the ordering really is compared.

[ulysses-you]

Confirmed, and the test is removed in 973bde7.

The ordering half of doCanonicalize is pinned instead by MergeSubplansSuite's identical DSv2 scans whose reported ordering is on a pruned column are deduplicated, not fused, which does fail
with the takeWhile reverted.

[peter-toth]

Finding 4. If the gate stays (see finding 2), V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.doc() needs a sentence for it. The doc enumerates the config's roles today - allowing the operation keys to be a subset of the partition keys, and gating the grouping of a collapsed partitioning. This adds a third that differs in kind: whether the scan reports a partitioning at all. That is the one a user meets as "the plan shape changed", so it is worth naming.

[ulysses-you]

Moot now that the gate is gone (finding 2): the config keeps the two roles it already documents, so
V2_BUCKETING_ALLOW_KEYS_SUBSET_OF_PARTITION_KEYS.doc() is unchanged.

[sunchao]

Request changes: one P1 correctness issue. Reviewed all six changed files at 78e2e312f272.

[P1] Preserve full partition keys for scan preparation

At DataSourceV2ScanExecBase.scala:108, the projection shortens partition-key rows and their types. However, BatchScanExec.filteredPartitions passes that projected partitioning into PushDownUtils.replanWithRuntimeFilters, which still reads full source keys.

This has three consequences:

• Wrong results: pruning a leading or middle key can make the actual partition order disagree with its advertised keys, causing joins to lose rows.
• Type errors: pruning a leading string key while retaining an integer key makes the sorter read a string as an integer.
• Runtime-filtering failures, even with trailing pruning: hashing iterates over the full key’s fields but indexes the shorter projected type array, causing ArrayIndexOutOfBoundsException when partitions survive filtering.

Keep the full source-key layout for input sorting and runtime filtering; expose the projected partitioning downstream. Add tests for leading/middle pruning and runtime filtering with surviving partitions.

This confirms the existing review’s diagnosis and extends its runtime-filtering impact. I found no additional actionable canonicalization/equality issues.

Validation: source tracing and an executable model reproduced the ordering mismatch and schema-index failure. Full Spark tests were not run; the checkout lacks the required build dependencies.

[ulysses-you]

Thank you @peter-toth and @sunchao . Fixed in 973bde7. This was a regression of this change.

DataSourceV2ScanExecBase now reports two separate views: reportedKeyedPartitioning, the source's
own report at full key width, and outputPartitioning, with the pruned key positions projected
away. filteredPartitions takes the former, and replanWithRuntimeFilters now takes
Option[KeyedPartitioning] so the projected view is not passable there at all.

Pinned by two new tests: a pruned leading key losing join rows, and a pruned leading String key
throwing ClassCastException when read at the projected Integer type. Pointing the call site back
at outputPartitioning fails exactly those two and nothing else.