Fix Spark master GpuBatchScanExec compile after replanWithRuntimeFilters API change[reduced-it] by firestarman · Pull Request #15932 · NVIDIA/cudf-spark

firestarman · 2026-09-09T05:34:32Z

Description

Spark master nightly builds fail in sql-plugin while compiling the Spark 5.0 shim: GpuBatchScanExec still passes outputPartitioning (Partitioning) into PushDownUtils.replanWithRuntimeFilters, which now takes Option[KeyedPartitioning]. Released Spark versions are unaffected. After this change, Spark 5 / master snapshot compilation succeeds again. No new configuration.

The helper must use the full-width keys the source reported, not the planner view. outputPartitioning may project away pruned key columns and is the wrong type. GPU GpuBatchScanExec now passes reportedKeyedPartitioning from DataSourceV2ScanExecBase, matching Spark's BatchScanExec.

The same Spark 5 BatchScanExec also treats pruned-out partition keys as planner metadata: equals, hashCode, and doCanonicalize use prunedKeyGroupedPartitioning, and canonicalize uses normalizeExpressions so multi-key order is preserved (SPARK-58120). Those were ported as well so GPU AQE reuse / sameResult does not diverge from CPU, and so a canonicalized keyed scan cannot reorder Integer/String keys.

Compile and the Spark 5 suite were checked with mvn -f scala2.13/pom.xml -Dbuildver=500 -Dcuda.version=cuda13 -pl sql-plugin,tests -am package -DwildcardSuites=com.nvidia.spark.rapids.shims.GpuBatchScanExecCanonicalizeSuite (BUILD SUCCESS, 4 tests). GpuBatchScanExecCanonicalizeSuite covers dangling-key equality, sameResult across different ExprIds, SPARK-58120 key order, and SPARK-59248-style runtime-filter replanning after a pruned leading key (filteredPartitions on a keyed GpuScan). Spark 5 unit tests use the Iceberg stub, so that last case is a GpuScan with HasPartitionKey partitions rather than an Iceberg query. Existing Spark 4 Iceberg SPJ / DSv2 scan tests still cover the Iceberg runtime-filter path.

This change was drafted with AI assistance. A human reviewed the diff and this description before updating the PR.

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

Driver-side plan identity and a compile-time argument swap when building scan partitions. No per-row or executor GPU work is added. reportedKeyedPartitioning is already the lazy value outputPartitioning is derived from in Spark's DataSourceV2ScanExecBase.

Spark master now takes Option[KeyedPartitioning] and requires the full-width reported keys, not the possibly projected outputPartitioning. Signed-off-by: Firestarman <firestarmanllc@gmail.com>

greptile-apps · 2026-09-09T05:36:38Z

Confidence Score: 5/5

The PR appears safe to merge, with no outstanding correctness or repository-rule issues identified.

Summary

Passes full-width reported keyed partitioning to runtime-filter replanning.
Excludes pruned partition keys from plan identity and canonicalization.
Preserves multi-key expression order during canonicalization.
Adds Spark 5 coverage for plan equality, canonicalization, and runtime-filter replanning.
Corrects the new test suite’s import ordering.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Source-reported full-width keys] --> B[GpuBatchScanExec]
  B --> C[Runtime-filter replanning]
  C --> D[Filtered partitions with original key slots]
  B --> E[Prune keys absent from output]
  E --> F[Plan equality and hash code]
  E --> G[Order-preserving canonicalization]

_{Reviews (4) · Last reviewed commit: "Fix scalastyle import order in GpuBatchS..."}

Ignore pruned-out partition keys in equals/hashCode and canonicalize with normalizeExpressions so key order is preserved (SPARK-58120). Signed-off-by: Firestarman <firestarmanllc@gmail.com>

firestarman · 2026-09-09T06:07:59Z

build

res-life · 2026-09-10T02:53:52Z

-      outputPartitioning,
+      // Full-width keys as the source reported them. outputPartitioning may project
+      // pruned key columns away and is a Partitioning, not Option[KeyedPartitioning].
+      reportedKeyedPartitioning,


[P1] Could we add an end-to-end regression test for the full-width key behavior exercised here? The new suite never evaluates filteredPartitions, and the existing Iceberg DPP case uses a single, unpruned partition key. An implementation that passes None or the pruned keys here would therefore still pass the current tests, while a runtime-filter replan after pruning a leading key can misalign keyed partitions, return wrong rows, or throw ClassCastException. Please port the SPARK-59248-style query with a multi-column key, a pruned leading key, and a runtime filter, and assert CPU/GPU result equality plus the GPU scan/SPJ plan.

Added a Spark 5 suite case that actually evaluates filteredPartitions. It uses a keyed GpuScan whose partitions implement HasPartitionKey with (store_id, dept_id), prunes the leading store_id from scan output, and applies a runtime filter on dept_id (SPARK-59248 shape). The test asserts padded None slots, remaining keys still two fields wide, and only dept_id=10 kept.

Spark 5 unit tests compile against the Iceberg stub, so this is not an Iceberg SQL IT. Passing None or the pruned one-column planner keys here fails the keyed replan (misaligned HasPartitionKey rows) rather than staying green.

res-life · 2026-09-10T02:53:52Z

+      keyGroupedPartitioning = keys)
+  }
+
+  test("equals and hashCode ignore partition keys pruned out of output") {


[P2] Could this be covered through Spark's actual plan-comparison path? These assertions call equals, hashCode, and doCanonicalize directly, so they do not verify the AQE reuse / sameResult behavior claimed in the PR description. Please add a query that produces semantically equivalent scans with different ExprIds and assert sameResult or ReusedExchangeExec; that would catch integration issues between canonicalization and physical-plan reuse.

The dangling-key case now also asserts sameResult, which goes through SparkPlan.canonicalized (the path AQE reuse uses). There is a second test that builds two GpuBatchScanExecs that differ only by ExprId and asserts sameResult between them.

A ReusedExchangeExec SQL query is not added here: Spark 5 UTs do not have a GPU DSv2 source that reports HasPartitionKey (parquet file partitions do not implement it; Iceberg is stubbed). sameResult on the GPU scan nodes is what would decide reuse for those plans.

res-life · 2026-09-10T02:53:52Z

+import org.apache.spark.sql.connector.read.{Batch, InputPartition, PartitionReaderFactory}
+import org.apache.spark.sql.types.{IntegerType, StringType, StructType}
+
+class GpuBatchScanExecCanonicalizeSuite extends SparkQueryCompareTestSuite with MockitoSugar {


[P2] PR description issue (anchored here because it describes validation for this suite): mvn ... validate stops before compile and test, so the listed command does not support "Compile was checked." Please replace it with an actual compile, package, install, or test command and its result. For example, a Spark 5 / Scala 2.13 reactor install followed by this targeted suite would substantiate both compilation and test coverage.

Updated the PR description. Compile and this suite were checked with:

mvn -f scala2.13/pom.xml -Dbuildver=500 -Dcuda.version=cuda13 -pl sql-plugin,tests -am package -DwildcardSuites=com.nvidia.spark.rapids.shims.GpuBatchScanExecCanonicalizeSuite

That package run compiled sql-plugin + tests and reported BUILD SUCCESS with 4 tests in GpuBatchScanExecCanonicalizeSuite.

Evaluate filteredPartitions after a pruned leading key plus runtime filter, and assert sameResult through SparkPlan canonicalization. Signed-off-by: Firestarman <firestarmanllc@gmail.com>

Signed-off-by: Firestarman <firestarmanllc@gmail.com>

firestarman · 2026-09-10T07:05:36Z

build

res-life

LGTM

firestarman · 2026-09-11T01:58:19Z

build

Fix Spark 5 GpuBatchScanExec after replanWithRuntimeFilters API change

6c47e44

Spark master now takes Option[KeyedPartitioning] and requires the full-width reported keys, not the possibly projected outputPartitioning. Signed-off-by: Firestarman <firestarmanllc@gmail.com>

firestarman added bug Something isn't working SQL part of the SQL/Dataframe plugin labels Sep 9, 2026

Align Spark 5 GpuBatchScanExec canonicalize with CPU BatchScanExec

1300a22

Ignore pruned-out partition keys in equals/hashCode and canonicalize with normalizeExpressions so key order is preserved (SPARK-58120). Signed-off-by: Firestarman <firestarmanllc@gmail.com>

firestarman changed the title ~~Fix Spark master GpuBatchScanExec compile after replanWithRuntimeFilters API change[fast-ut][reduced-it]~~ Fix Spark master GpuBatchScanExec compile after replanWithRuntimeFilters API change[reduced-it] Sep 9, 2026

firestarman self-assigned this Sep 9, 2026

firestarman requested a review from a team September 9, 2026 06:08

res-life reviewed Sep 10, 2026

View reviewed changes

firestarman added 2 commits September 10, 2026 11:21

Cover Spark 5 GpuBatchScanExec runtime-filter keys and sameResult

21b5e1b

Evaluate filteredPartitions after a pruned leading key plus runtime filter, and assert sameResult through SparkPlan canonicalization. Signed-off-by: Firestarman <firestarmanllc@gmail.com>

Fix scalastyle import order in GpuBatchScanExecCanonicalizeSuite

13968de

Signed-off-by: Firestarman <firestarmanllc@gmail.com>

res-life approved these changes Sep 10, 2026

View reviewed changes

firestarman merged commit 3f731d7 into NVIDIA:main Sep 11, 2026
57 checks passed

firestarman deleted the fix/15930-spark500-batchscan-keyed-partitioning branch September 11, 2026 04:52

Navigation Menu

Uh oh!

Fix Spark master GpuBatchScanExec compile after replanWithRuntimeFilters API change[reduced-it] - #15932

Fix Spark master GpuBatchScanExec compile after replanWithRuntimeFilters API change[reduced-it]#15932
firestarman merged 4 commits into
NVIDIA:mainfrom
firestarman:fix/15930-spark500-batchscan-keyed-partitioning

firestarman commented Sep 9, 2026 •

edited

Loading

Uh oh!

greptile-apps Bot commented Sep 9, 2026 •

edited

Loading

Uh oh!

firestarman commented Sep 9, 2026

Uh oh!

res-life Sep 10, 2026

Uh oh!

firestarman Sep 10, 2026

Uh oh!

res-life Sep 10, 2026

Uh oh!

firestarman Sep 10, 2026

Uh oh!

res-life Sep 10, 2026

Uh oh!

firestarman Sep 10, 2026

Uh oh!

firestarman commented Sep 10, 2026

Uh oh!

res-life left a comment

Uh oh!

firestarman commented Sep 11, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

firestarman commented Sep 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Checklists

Uh oh!

greptile-apps Bot commented Sep 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Confidence Score: 5/5

Summary

Diagram

Uh oh!

firestarman commented Sep 9, 2026

Uh oh!

res-life Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

firestarman Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

res-life Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

firestarman Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

res-life Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

firestarman Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

firestarman commented Sep 10, 2026

Uh oh!

res-life left a comment

Choose a reason for hiding this comment

Uh oh!

firestarman commented Sep 11, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

firestarman commented Sep 9, 2026 •

edited

Loading

greptile-apps Bot commented Sep 9, 2026 •

edited

Loading