Fix SPJ runtime partition validation [databricks][reduced-it] by amahussein · Pull Request #15924 · NVIDIA/cudf-spark

amahussein · 2026-09-08T21:49:48Z

Description

During runtime filtering, GpuBatchScanExec.filteredPartitions validates full scan keys against SPJ-facing outputPartitioning, whose values may be intersected or projected. Valid scan-owned keys can consequently be rejected as new.

This ports Spark's SPARK-58783 receiver correction to shared spark340 and spark350db143: validation now uses super.outputPartitioning, preserving scan-native keys while later SPJ projection and grouping remain unchanged. Spark 3.4 and 3.5 cannot reach the false rejection; aligning spark340 does not change conforming prune-only behavior.

On affected Spark 4.0.x and 4.1.x releases, test_iceberg_spj_partition_filter_with_runtime_filter_cpu_fails_gpu_succeeds asserts the upstream CPU exception, then verifies GPU results, evaluated DPP values, and an exchange-free GPU SPJ plan. The fixed-release parity test remains gated to Spark 4.0.5+ or 4.1.4+.

test_iceberg_spj_runtime_filter_with_trailing_join_key adds an (INT, STRING) identity-partitioned DPP case joining on the trailing string key. It asserts projected joinKeyPositions, pruning and common partition values, expected GPU rows, and exchange-free SPJ execution. Affected releases assert the observed CPU-side Iceberg partition-data type-check failure; fixed releases require CPU/GPU parity. Both regressions are selected in the only premerge Iceberg job, its extra-classpath smoke list.

Affected Apache Spark profiles served by the shared source intentionally differ from their buggy CPU behavior. Mapped Databricks builds receive the same correction, but their CPU/GPU parity remains unverified.

Validation passed a Scala 2.13 buildver 402 install, RAT, Scalastyle, resource-nesting lint, and static checks. Spark 4.0.2 with Iceberg 1.10.1 passed both regressions in normal and CI extra-classpath layouts. Reverting the receiver fix made both GPU legs fail; restoring it made both pass. Premerge itself runs Spark 4.0.1.

No public API or configuration changes. Performance is unaffected: validation only selects different existing metadata.

AI assistance was used to implement and review this change.

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
(Please provide the names of the existing tests in the PR description.)
Not required

Performance

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

Fixes NVIDIA#15839 Validate runtime-filtered scan partition keys against the scan's original full partitioning instead of the SPJ-facing projected view. Add a release-gated Iceberg regression that proves the partition-filter and dynamic-pruning plan shape. Signed-off-by: Ahmed Hussein <ahussein@nvidia.com>

amahussein · 2026-09-08T21:50:35Z

build

greptile-apps · 2026-09-08T21:53:13Z

Confidence Score: 5/5

The PR appears safe to merge; no actionable correctness, test-coverage, security, or repository-rule failures remain.

Summary

Applies the receiver correction to the shared Spark 3.4 and Databricks 14.3 shim implementations.
Adds Iceberg regressions for partition-filter and trailing-join-key runtime-filter paths.
Adds affected-release regression coverage to the Iceberg extra-classpath pre-merge job.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Scan-native partitioning] --> B[Runtime filtering]
  B --> C[Filtered data-source partitions]
  A --> D[Validate filtered partition keys]
  C --> D
  D --> E[SPJ projection and intersection]
  E --> F[GPU storage-partitioned join]

_{Reviews (3) · Last reviewed commit: "Merge branch 'main' into rapids-15839-sp..."}

amahussein · 2026-09-09T12:08:29Z

build

gerashegalov · 2026-09-10T22:21:16Z

+@pytest.mark.skipif(
+    not (
+        (is_spark_40x() and _is_spark_patch_at_least(spark_version(), 5))
+        or (is_spark_41x() and _is_spark_patch_at_least(spark_version(), 4))
+    ),
+    reason="SPARK-58783 was fixed in Spark 4.0.5 and 4.1.4; Spark 4.2+ is unaffected")


As written, CI compiles the change but never exercises the production code change

Could we add a current-release test that separately expects the CPU exception and verifies the successful GPU result and plan?

Added an affected-release regression to the existing Iceberg extra-classpath premerge job, the only premerge path that executes Iceberg tests. It asserts the CPU exception, then verifies the GPU result, DPP values, and exchange-free SPJ plan. On Spark 4.0.2, reverting this fix made it fail and restoring it passed; premerge will run Spark 4.0.1. Fixed releases use parity.

gerashegalov · 2026-09-10T22:21:58Z

+        "spark.sql.optimizer.dynamicPartitionPruning.fallbackFilterRatio": "10",
+        "spark.sql.sources.v2.bucketing.enabled": "true",
+        "spark.sql.sources.v2.bucketing.pushPartValues.enabled": "true",
+        "spark.sql.sources.v2.bucketing.partition.filter.enabled": "true",


Could we also cover the projected trailing-join-key path? This case uses one INT partition key and exercises only the common-value intersection. SPARK-58783 also fixes joinKeyPositions projections; with partition keys such as
(INT, STRING) and a join on the trailing string, the old code can interpret the first raw field using the projected StringType and fail with a ClassCastException. A DPP test using a non-leading join key and asserting
joinKeyPositions would cover that distinct part of the fix.

Added an (INT, STRING) identity-partitioned DPP case to that premerge job. It joins on the trailing string key and asserts joinKeyPositions == [1], pruning and common partition values, GPU rows, and exchange-free SPJ. Reverting this fix makes its GPU leg fail. The measured CPU path fails in Iceberg's partition-data check before Spark's equivalent cast; fixed releases use parity.

Review feedback on NVIDIA#15924. The parity test added with the receiver fix is gated to Spark 4.0.5+ or 4.1.4+, neither of which exists, so CI compiled the production change without ever executing it. Separately, the joinKeyPositions projection route that SPARK-58783 also repairs had no plugin coverage of any kind. test_iceberg_spj_partition_filter_with_runtime_filter_cpu_fails_gpu_succeeds runs on affected Spark 4.0.x and 4.1.x releases. It asserts the CPU runtime-filtering exception, then requires the GPU to return the expected row with evaluated DPP values and an exchange-free SPJ plan. It retires itself once a fixed release ships, at which point the existing parity test takes over, so the asymmetric assertion never needs to be undone by hand. test_iceberg_spj_runtime_filter_with_trailing_join_key partitions the left table by (dept_id, data) and joins on the trailing string key, so joinKeyPositions is [1] and the pre-fix receiver wraps a raw INT key with the projected StringType. Asserting commonPartitionValues keeps the case on the two-sided route and away from open issue NVIDIA#15338. On affected releases the CPU leg fails inside Iceberg's PartitionData type check, which fires before Spark's own cast, so the matcher accepts either message. Both regressions are added to the Iceberg extra-classpath premerge job, which is the only premerge path that executes Iceberg tests. Verified on Spark 4.0.2 with Iceberg 1.10.1 in both the normal and extra-classpath layouts. Reverting the receiver fix makes both GPU legs fail and restoring it makes both pass, so these exercise the production change rather than only its plan shape. Premerge itself runs Spark 4.0.1. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Ahmed Hussein (amahussein) <a@ahussein.me>

amahussein · 2026-09-11T03:42:59Z

build

amahussein

Thanks @gerashegalov
Addressed review comments

amahussein · 2026-09-11T03:45:32Z

+        "spark.sql.optimizer.dynamicPartitionPruning.fallbackFilterRatio": "10",
+        "spark.sql.sources.v2.bucketing.enabled": "true",
+        "spark.sql.sources.v2.bucketing.pushPartValues.enabled": "true",
+        "spark.sql.sources.v2.bucketing.partition.filter.enabled": "true",


Added an (INT, STRING) identity-partitioned DPP case to that premerge job. It joins on the trailing string key and asserts joinKeyPositions == [1], pruning and common partition values, GPU rows, and exchange-free SPJ. Reverting this fix makes its GPU leg fail. The measured CPU path fails in Iceberg's partition-data check before Spark's equivalent cast; fixed releases use parity.

amahussein · 2026-09-11T03:46:11Z

+@pytest.mark.skipif(
+    not (
+        (is_spark_40x() and _is_spark_patch_at_least(spark_version(), 5))
+        or (is_spark_41x() and _is_spark_patch_at_least(spark_version(), 4))
+    ),
+    reason="SPARK-58783 was fixed in Spark 4.0.5 and 4.1.4; Spark 4.2+ is unaffected")


Added an affected-release regression to the existing Iceberg extra-classpath premerge job, the only premerge path that executes Iceberg tests. It asserts the CPU exception, then verifies the GPU result, DPP values, and exchange-free SPJ plan. On Spark 4.0.2, reverting this fix made it fail and restoring it passed; premerge will run Spark 4.0.1. Fixed releases use parity.

gerashegalov · 2026-09-11T06:49:10Z

+        return _partition_filter_runtime_filter_query(
+            spark, left_table, right_table, dim_table)
+
+    def assert_plan(plan):


After the merge from main, assert_cpu_and_gpu_are_equal_collect_with_capture calls this callback with (cpu_plan, gpu_plan). Once 4.0.5 or 4.1.4 enables this test, the one-argument function will fail with TypeError before checking the plan. Could this be def assert_plan(_cpu_plan, plan)?

gerashegalov · 2026-09-11T06:50:00Z

+            join_with_runtime_filter,
+            conf=conf,
+            require_non_empty=True,
+            gpu_plan_assertion=assert_plan)


After the merge from main, assert_cpu_and_gpu_are_equal_collect_with_capture calls this callback with (cpu_plan, gpu_plan). Once 4.0.5 or 4.1.4 enables this test, the one-argument function will fail with TypeError before checking the plan. Could this be def assert_plan(_cpu_plan, plan)?

amahussein self-assigned this Sep 8, 2026

amahussein added Spark 4.0+ Spark 4.0+ issues Spark 4.1+ Spark 4.1+ issues labels Sep 8, 2026

amahussein requested a review from a team September 8, 2026 22:06

amahussein mentioned this pull request Sep 9, 2026

BUG] RapidsCachedTableSuite accumulator cleanup test is flaky #15928

Open

Merge branch 'main' into rapids-15839-spj-filtered-partitions

240db1b

amahussein requested a review from nartal1 September 9, 2026 16:07

gerashegalov requested changes Sep 10, 2026

View reviewed changes

amahussein requested a review from a team as a code owner September 11, 2026 03:41

Merge branch 'main' into rapids-15839-spj-filtered-partitions

51cdff9

amahussein requested a review from gerashegalov September 11, 2026 03:47

amahussein commented Sep 11, 2026

View reviewed changes

gerashegalov reviewed Sep 11, 2026

View reviewed changes

Navigation Menu

Uh oh!

Fix SPJ runtime partition validation [databricks][reduced-it] - #15924

Fix SPJ runtime partition validation [databricks][reduced-it]#15924
amahussein wants to merge 4 commits into
NVIDIA:mainfrom
amahussein:rapids-15839-spj-filtered-partitions

amahussein commented Sep 8, 2026 •

edited

Loading

Uh oh!

amahussein commented Sep 8, 2026

Uh oh!

greptile-apps Bot commented Sep 8, 2026 •

edited

Loading

Uh oh!

amahussein commented Sep 9, 2026

Uh oh!

gerashegalov Sep 10, 2026

Uh oh!

amahussein Sep 11, 2026

Uh oh!

gerashegalov Sep 10, 2026

Uh oh!

amahussein Sep 11, 2026

Uh oh!

amahussein commented Sep 11, 2026

Uh oh!

amahussein left a comment

Uh oh!

amahussein Sep 11, 2026

Uh oh!

amahussein Sep 11, 2026

Uh oh!

gerashegalov Sep 11, 2026

Uh oh!

gerashegalov Sep 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

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

Uh oh!

Description

Checklists

Uh oh!

amahussein commented Sep 8, 2026

Uh oh!

greptile-apps Bot commented Sep 8, 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!

amahussein commented Sep 9, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

amahussein commented Sep 11, 2026

Uh oh!

amahussein left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

amahussein commented Sep 8, 2026 •

edited

Loading

greptile-apps Bot commented Sep 8, 2026 •

edited

Loading