Port SPARK-48949 scan-side partition filter to GpuBatchScanExec [fast-ut][databricks][reduced-it]

[amahussein]

Fixes #11453.

Description

The problem.

Apache Spark 4.0.0 added planning-time partition filtering for Storage-Partitioned Join (SPARK-48949, dbba92a1b4e). With spark.sql.sources.v2.bucketing.partition.filter.enabled=true, EnsureRequirements pushes down the intersection of the two join sides' partition values instead of their union. That planner half is untouched Spark code and runs whether or not the plugin is loaded — the plugin cannot opt out of it. The scan half lives inside BatchScanExec, which the plugin hand-copies into GpuBatchScanExec, and that half was never ported.

The consequence is not a missed optimization, it is a hard failure. GpuBatchScanExec.inputRDD walks the scan's own partition groups and asserts each one is present in commonPartitionValues. That assertion was safe when the value set was a union; against an intersection it is not, and any group the planner pruned trips it. With both ...partition.filter.enabled=true and ...partiallyClusteredDistribution.enabled=true, a GPU Iceberg SPJ inner join whose left side has a partition key the right side lacks dies on the driver during RDD construction:

java.lang.AssertionError: assertion failed: Partition value ... does not exist in
common partition values from Spark plan

CPU Spark 4.0.2 runs the same query correctly. Both configs are user-facing, documented and independently settable, so this is reachable in production; it fails before any task launches, so no wrong answer is possible.

The fix.

Transcribe the four-line filter Spark already ships (v4.0.2:BatchScanExec.scala:203) into both GpuBatchScanExec shims that carry the full SPJ logic, dropping the scan's own groups that are absent from commonPartitionValues before the assertion runs. No new configuration, no behavioural change when the config is off.

Why this cannot skew partition counts between the two join sides.

This is the usual risk with changes in this area — Can't zip RDDs with unequal numbers of partitions — and the structure of the branch rules it out. The partially-clustered branch does not build its output from the collection being filtered; it builds it by flatMap-ing over commonPartitionValues, using the filtered groups only as a lookup table for how to populate each key. EnsureRequirements.populateCommonPartitionInfo pushes the identical merged value list to both sides of the join, so both scans iterate the same list and produce the same number of partitions. Filtering changes which keys the lookup table can satisfy; it cannot change how many partitions come out. The groups it removes are precisely those whose key is absent from that list, so they could never have been looked up during output construction — before this change they simply reached the assert first.

After Change

the plugin's partition.filter × partiallyClusteredDistribution behaviour matches CPU Spark 4.0.2 cell for cell, including RDD partition counts, on the four-way matrix.

Testing

New test_iceberg_spj_partition_filter in integration_tests/src/main/python/iceberg/iceberg_test.py, parameterised over partition.filter.enabled × partiallyClusteredDistribution.enabled. It compares GPU against CPU on join output — CPU is the oracle, and that row comparison is what proves the filter correct. Before the fix exactly one of the four cells fails.

It is gated on is_spark_400_or_later(), with no patch-level condition — deliberately, so it runs on 4.0.x and 4.1.x rather than inheriting the 3.5.9+/4.0.3+/4.1.2+ skip that test_iceberg_spj_partial_clustering_distinct carries for unrelated reasons (see below). That existing test is not modified.

Performance

1. The change cannot slow down anything that works today. With spark.sql.sources.v2.bucketing.partition.filter.enabled=false — the default — commonPartitionValues is the union of both join sides and is therefore a superset of each scan's own partition groups, so the new .filter is provably unable to remove anything. Measured rather than argued: the four-cell partition.filter × partiallyClusteredDistribution matrix was run before and after the change on the same machine, and the three cells that work today return identical join-stage RDD partition counts — 3, 4 and 2 — in both runs.

2. Its own cost is one hash-set lookup per partition group, on the driver, at plan time. O(g) in the number of groups the scan enumerates, over the same collection the next line already maps across to build nestGroupedPartitions. No per-row, per-batch or executor-side cost. Not worth benchmarking.

3. The pruning it unlocks on the path that previously crashed. Iceberg fact table with 64 partitions × 20,000 rows, dimension table with 4 of those keys, inner join on the partition key, partiallyClusteredDistribution.enabled=true, AQE and broadcast join disabled, single-node local[4], one RTX A5000, Iceberg 1.10.1:

	partition.filter=false	partition.filter=true
GPU scan RDD partitions	64, 64	4, 4
GPU elapsed	0.65 s / 0.55 s	0.17 s / 0.16 s
CPU (Spark 4.0.2) scan RDD partitions	64, 64	4, 4
CPU elapsed	0.22 s / 0.16 s	0.05 s / 0.03 s

Both configurations return the same result (count=80000, sum=51197560000). Two repeats after a warm-up.

These are after numbers only. There is no before number for the partition.filter=true column, because before this change that configuration does not run — it dies on the driver with the AssertionError this PR fixes. The partition.filter=false column is unchanged by the PR (64 partitions, 0.56 s pre-fix against 0.65 s / 0.55 s post-fix), which is claim 1 restated on this workload.

The claim to take from the table is "GPU now prunes identically to CPU on this path", not "GPU got faster". Two things this is deliberately not claiming. It is not claiming new pruning in general: with partiallyClusteredDistribution disabled the plugin already prunes 64 → 4 with no code change, and that is not a benefit of this PR. And it is not claiming a GPU-specific advantage in ratio — the speed-ups above are 3.6× on GPU against 4.8× on CPU, so if anything the ratio favours CPU. What the numbers do support, as an indicative absolute statement, is that each eliminated partition is worth more wall-clock on GPU than on CPU: dividing the elapsed delta by the 60 removed partitions gives roughly 7 ms per partition on GPU against roughly 2.5 ms on CPU. This is a small single-node workload and those per-partition figures should be read as indicative, not as a benchmark result.

Addition Context

One observation that is not a regression. With the assertion no longer firing, SELECT DISTINCT over a partially-clustered SPJ at partition.filter=true now returns duplicates on Spark 4.0.2 instead of crashing. CPU Spark 4.0.2 returns the same duplicates at the same configs — this is SPARK-55848, fixed upstream in Spark 3.5.9 / 4.0.3 / 4.1.2, which the crash was previously masking. It is why test_iceberg_spj_partial_clustering_distinct carries a patch-level skip and why the new test asserts join output instead.

Not in scope, tracked separately: SPARK-58783's originalPartitioning one-liner in filteredPartitions (#15839 — CPU Spark 4.0.2 fails identically today, so there is no plugin-specific exposure); delegating to KeyGroupedPartitionedScan on Spark 4.1 to delete ~125 lines of hand-copied Spark internals (#15840, sequenced after this so the refactor lands with a test protecting it); and outputPartitioning not projecting on joinKeyPositions (#15338, which cannot co-occur with this issue's configs).

Checklists

Documentation

• Updated for new or modified user-facing features or behaviors
• No user-facing change

Testing

• Added or modified tests to cover new code paths
• Covered by existing tests
• Not required

Performance

• Tests ran and results are added in the PR description
• Issue filed with a link in the PR description
• Not required

[amahussein]

build

[greptile-apps]

Greptile Summary

The PR ports Spark's scan-side partition filtering into the two GPU batch-scan shims that implement partially clustered storage-partitioned joins.

• Filters scan partition groups against planner-provided common partition values before split-count lookup.
• Adds an Iceberg CPU/GPU parity test covering the partition-filtering and partially-clustered-distribution configuration matrix.
• Preserves the pre-Spark-4 shim copy for consistency while documenting that the filter is inert there.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
integration_tests/src/main/python/iceberg/iceberg_test.py	Adds deterministic CPU/GPU parity and GPU-plan validation across the four relevant SPJ configuration combinations.
sql-plugin/src/main/spark340/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala	Adds an effectively inert pre-Spark-4 filter to keep the parallel SPJ shim implementations aligned.
sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala	Removes planner-pruned partition groups before the split-count lookup, matching Spark's scan-side filtering behavior.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Grouped scan partitions] --> B{Present in common partition values?}
  B -- No --> C[Drop pruned group]
  B -- Yes --> D[Look up split count]
  D --> E[Build nested grouped partitions]
  E --> F[Create scan RDD partitions]

Loading

_{Reviews (4): Last reviewed commit: "Merge branch 'main' into rapids-11453" | Re-trigger Greptile}

[firestarman]

NIT:
The extra key that partition filtering should prune (id=3) only exists on the left table, and the comments above say the two id=1 files are what make the left scan partially clustered.

GpuBatchScanExec applies the SPARK-48949 filter to both scans, then branches on replicatePartitions (pad vs replicate). A regression that applied the filter only in the pad (!replicatePartitions) branch would still pass this matrix, because the replicated right scan never enumerates a group that is absent from commonPartitionValues.

Could we add a right-only partition key (e.g. id=4 on the right table only) in the filtered + partially clustered cell, or parametrize which side holds the extra key? That would force the replicated scan through the same commonPartValuesMap filter before Seq.fill(numSplits.get)(splits). Inner-join output would be unchanged — id=4 would not appear in the result.

Suggested change

	spark.sql(f"INSERT INTO {left_table} VALUES (1, 40.0), (2, 10.0), (3, 15.5)")
	spark.sql(f"INSERT INTO {left_table} VALUES (1, 41.0)")
	spark.sql(f"INSERT INTO {right_table} VALUES (1, 'a'), (2, 'b')")
	spark.sql(f"INSERT INTO {left_table} VALUES (1, 40.0), (2, 10.0), (3, 15.5)")
	spark.sql(f"INSERT INTO {left_table} VALUES (1, 41.0)")
	spark.sql(f"INSERT INTO {right_table} VALUES (1, 'a'), (2, 'b')")
	# Right-only key so the replicated scan also enumerates a group absent from
	# the intersection (mirror of left-only id=3 on the pad side).
	spark.sql(f"INSERT INTO {right_table} VALUES (4, 'd')")

[amahussein]

Good catch, real gap. With the old fixture the right scan is the replicate side and its keys {1, 2} are exactly the intersection, so it never had a group to drop.

Verified it bites: I built with the filter skipped when replicatePartitions=true. The old fixture passed, the new one fails on partially_clustered-filtered with the expected AssertionError.

One wrinkle with adding only id=4: it swaps the pad and replicate roles, and the new pad side has one file per key, so every numSplits is 1 and the counts collapse to 4/2 for both partially_clustered and clustered. That quietly disables the assertion covering partial clustering.

Made it symmetric instead: a key each side lacks (id=3 left, id=4 right), plus a two-file key inside the intersection on each side. Both branches are then covered whichever side Spark picks, counts are 3/5/2/4, and all four cells pass.

[firestarman]

LGTM, and the comment above is not a blocker. Better have more eyes on this.

[amahussein]

build

[firestarman]

build

[firestarman]

build

[amahussein]

build