Preserve partial clustering across Spark versions [databricks] - #15397
Preserve partial clustering across Spark versions [databricks]#15397res-life wants to merge 11 commits into
Conversation
|
build |
Signed-off-by: Chong Gao <chongg@nvidia.com>
Signed-off-by: Chong Gao <chongg@nvidia.com>
Signed-off-by: Chong Gao <chongg@nvidia.com>
5598e52 to
57c9565
Compare
|
build |
Greptile SummaryThe PR completes the partial-clustering compatibility fix across supported Spark patch releases.
Confidence Score: 5/5The PR appears safe to merge. The previously reported marker loss is fixed because the affected scan implementation now passes the rebuilt partitioning and Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart LR
A[GpuBatchScanExec] --> B[Compute final partition values]
B --> C[Marker-aware KeyGroupedPartitioning shim]
C --> D[Preserve isPartiallyClustered]
B --> E[Build RDD from finalPartitions]
D --> F[Shuffle-free storage-partitioned join]
E --> F
F --> G[Post-join DISTINCT shuffle]
Reviews (6): Last reviewed commit: "Cover padded partial clustering partitio..." | Re-trigger Greptile |
Signed-off-by: Chong Gao <chongg@nvidia.com>
|
build |
Signed-off-by: Chong Gao <chongg@nvidia.com>
Signed-off-by: Chong Gao <chongg@nvidia.com>
firestarman
left a comment
There was a problem hiding this comment.
Code Review: cudf-spark — (pr-15397)
1. [🟡 SHOULD FIX] Could you add a docstring for the newly added `gpu_plan_assertion` callback at integration_tests/src/main/python/asserts.py:507? The public reusable helper beginning at integration_tests/src/main/python/asserts.py:502 has no docstring, so callers cannot discover from its interface that the callback runs after GPU collection and receives the JVM executed plan from integration_tests/src/main/python/asserts.py:523
Line: 507 (diff)
Issue: Could you add a docstring for the newly added `gpu_plan_assertion` callback at integration_tests/src/main/python/asserts.py:507? The public reusable helper beginning at integration_tests/src/main/python/asserts.py:502 has no docstring, so callers cannot discover from its interface that the callback runs after GPU collection and receives the JVM executed plan from integration_tests/src/main/python/asserts.py:523. Documenting that contract would satisfy the public-function documentation convention and make future callback implementations safer to write.
Confidence: 🟣❗ CERTAIN
Code:
```python
|
def assert_cpu_and_gpu_are_equal_collect_with_capture(func,
exist_classes='',
non_exist_classes='',
conf={},
require_non_empty=False,
gpu_plan_assertion=None):
(bring_back, collect_type) = _prep_func_for_compare(func, 'COLLECT_WITH_DATAFRAME')```
Fix: Add a function docstring that summarizes the CPU/GPU collection behavior and explicitly documents that `gpu_plan_assertion`, when supplied, is called with the GPU dataframe's executed JVM plan after collection.
Diff:
```diff
|
@@
def assert_cpu_and_gpu_are_equal_collect_with_capture(func,
exist_classes='',
non_exist_classes='',
conf={},
require_non_empty=False,
gpu_plan_assertion=None):
+ """Compare collected CPU/GPU results and validate the executed GPU plan.
+
+ :param func: Function that creates the dataframe to collect in each Spark session.
+ :param exist_classes: Comma-separated class names required in the GPU plan.
+ :param non_exist_classes: Comma-separated class names forbidden in the GPU plan.
+ :param conf: Spark configuration used for both executions.
+ :param require_non_empty: Require the collected CPU result to contain at least one row.
+ :param gpu_plan_assertion: Optional callback invoked with the GPU dataframe's executed
+ JVM plan after collection.
+ """
(bring_back, collect_type) = _prep_func_for_compare(func, 'COLLECT_WITH_DATAFRAME')```
Reference: .review/conventions.md:82 explicitly requires docstrings on public functions and classes; integration_tests/src/main/python/asserts.py:696 shows the repository's established public-helper docstring pattern, including parameter documentation at integration_tests/src/main/python/asserts.py:701. PEP 257 also requires function docstrings to summarize behavior and document applicable arguments: https://peps.python.org/pep-0257/
(Reviewer: compliant)
2. [🔴 MUST FIX] Can this callback assert the rebuilt GPU scans' partial-clustering marker directly? The checks at integration_tests/src/main/python/iceberg/iceberg_test.py:71 through integration_tests/src/main/python/iceberg/iceberg_test.py:85 only count scan, join, and exchange nodes
Line: 71 (diff)
Issue: Can this callback assert the rebuilt GPU scans' partial-clustering marker directly? The checks at integration_tests/src/main/python/iceberg/iceberg_test.py:71 through integration_tests/src/main/python/iceberg/iceberg_test.py:85 only count scan, join, and exchange nodes. They do not exercise the behavior that the marker-preserving shim added at sql-plugin/src/main/spark359/scala/com/nvidia/spark/rapids/shims/KeyGroupedPartitioningShim.scala:37 is meant to guarantee. This leaves the regression test unable to identify the affected Spark 4.0.3/4.0.4/4.1.2/4.1.3 call site: those versions select sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala:23, while its reconstruction at sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala:140 copies only expressions, partition count, and values. A regression test for marker preservation needs to fail when that output marker is false.
Confidence: 🟣❗ CERTAIN
Code:
```python
`scans = nodes_of_class("GpuBatchScanExec")` followed by count and shuffle assertions at integration_tests/src/main/python/iceberg/iceberg_test.py:71, without an assertion on each scan's `outputPartitioning` marker.```
Fix: After verifying the two GPU scans at integration_tests/src/main/python/iceberg/iceberg_test.py:75, inspect their output partitioning and require at least one scan to report the partial-clustering marker. This directly covers the value propagated by sql-plugin/src/main/spark340/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala:143 and fails for the marker-dropping reconstruction at sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala:140.
Diff:
```diff
@@
assert len(scans) == 2, f"Expected two GPU batch scans, found {len(scans)}:\n{plan}"
+ assert any(
+ scan.outputPartitioning().isPartiallyClustered() for scan in scans
+ ), f"Expected a partially-clustered GPU batch scan:\n{plan}"
assert len(joins) == 1, f"Expected one GPU SPJ join, found {len(joins)}:\n{plan}"```
Reference: sql-plugin/src/test/spark359/scala/com/nvidia/spark/rapids/shims/KeyGroupedPartitioningShimSuite.scala:47 demonstrates the direct marker assertion; sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala:128 verifies that the affected node exposes `outputPartitioning`.
(Reviewer: test-coverage)
4. [🔴 MUST FIX] Can the Spark 4.0.3, 4.0.4, 4.1.2, and 4.1.3 registrations at sql-plugin/src/main/spark359/scala/com/nvidia/spark/rapids/shims/KeyGroupedPartitioningShim.scala:18-21 also route the selected GPU scan reconstruction through this helper? Those builds select sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala:23-27, whose outputPartitioning implementation at sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala:128-141 still calls KeyGroupedPartitioning.copy without spjParams.applyPartialClustering
Line: 18 (diff)
Issue: Can the Spark 4.0.3, 4.0.4, 4.1.2, and 4.1.3 registrations at sql-plugin/src/main/spark359/scala/com/nvidia/spark/rapids/shims/KeyGroupedPartitioningShim.scala:18-21 also route the selected GPU scan reconstruction through this helper? Those builds select sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala:23-27, whose outputPartitioning implementation at sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala:128-141 still calls KeyGroupedPartitioning.copy without spjParams.applyPartialClustering. Thus a concrete partial-clustering plan with commonPartitionValues defined and applyPartialClustering=true returns a GPU KeyGroupedPartitioning whose isPartiallyClustered marker was never set from the SPJ parameters. Apache Spark added this marker specifically so partially clustered data does not satisfy a downstream ClusteredDistribution; omitting it breaks the output-partitioning correctness contract and can suppress a required exchange, producing incorrect deduplication or aggregation results.
Confidence: 🟣❗ CERTAIN
Code:
```scala
|
sql-plugin/src/main/spark359/scala/com/nvidia/spark/rapids/shims/KeyGroupedPartitioningShim.scala:18-21
{"spark": "403"}
{"spark": "404"}
{"spark": "412"}
{"spark": "413"}
sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala:140-141
k.copy(expressions = expressions, numPartitions = newPartValues.length,
partitionValues = newPartValues)```
Fix: Route sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala:140 through the already-shimmed copyWithNewPartitionValues API, retaining the projected expressions and passing spjParams.applyPartialClustering. Also assert the marker directly in integration_tests/src/main/python/iceberg/iceberg_test.py:75 so the new regression test fails on the affected GPU implementation rather than only checking the CPU-planned post-join exchange.
Diff:
```diff
|
--- sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala
+++ sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala
@@ -137,5 +137,6 @@
case Some(projectionPositions) => projectionPositions.map(i => k.expressions(i))
case _ => k.expressions
}
- k.copy(expressions = expressions, numPartitions = newPartValues.length,
- partitionValues = newPartValues)
+ KeyGroupedPartitioningShim.copyWithNewPartitionValues(
+ k.copy(expressions = expressions), newPartValues,
+ spjParams.applyPartialClustering)
case p => p
--- integration_tests/src/main/python/iceberg/iceberg_test.py
+++ integration_tests/src/main/python/iceberg/iceberg_test.py
@@ -73,6 +73,9 @@
exchanges = nodes_of_class("GpuShuffleExchangeExec")
assert len(scans) == 2, f"Expected two GPU batch scans, found {len(scans)}:\n{plan}"
+ assert any(scan.outputPartitioning().isPartiallyClustered() for scan in scans), \
+ "Expected at least one GPU batch scan to retain the partial-clustering marker:" \
+ f"\n{plan}"
assert len(joins) == 1, f"Expected one GPU SPJ join, found {len(joins)}:\n{plan}"```
Reference: sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/BatchScanExecMeta.scala:41-44 carries the CPU SPJ parameters into the selected GPU scan; sql-plugin/src/main/spark340/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala:142-143 demonstrates the repository's correct marker-aware call pattern; and sql-plugin/src/main/spark359/scala/com/nvidia/spark/rapids/shims/KeyGroupedPartitioningShim.scala:34-37 is the implementation selected by the four newly registered builds. Apache Spark 4.0.3 explicitly passes isPartiallyClustered = spjParams.applyPartialClustering at the equivalent reconstruction and makes partially clustered partitioning reject ClusteredDistribution: https://github.com/apache/spark/blob/v4.0.3/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExec.scala#L117-L145 and https://github.com/apache/spark/pull/54851/files
(Reviewer: correctness)
5. [🔴 MUST FIX] Can the Spark 4.0.3, 4.0.4, 4.1.2, and 4.1.3 scan reconstruction be routed through the marker-aware shim boundary too? The changed build ownership at sql-plugin/src/main/spark359/scala/com/nvidia/spark/rapids/shims/KeyGroupedPartitioningShim.scala:18-21 makes those builds select the implementation that writes `isPartiallyClustered` at sql-plugin/src/main/spark359/scala/com/nvidia/spark/rapids/shims/KeyGroupedPartitioningShim.scala:34-37
Line: 18 (diff)
Issue: Can the Spark 4.0.3, 4.0.4, 4.1.2, and 4.1.3 scan reconstruction be routed through the marker-aware shim boundary too? The changed build ownership at sql-plugin/src/main/spark359/scala/com/nvidia/spark/rapids/shims/KeyGroupedPartitioningShim.scala:18-21 makes those builds select the implementation that writes `isPartiallyClustered` at sql-plugin/src/main/spark359/scala/com/nvidia/spark/rapids/shims/KeyGroupedPartitioningShim.scala:34-37. However, those same builds select the separate scan implementation declared at sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala:23-27, whose reconstruction at sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala:140-141 calls `k.copy` directly and never forwards `sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala:53`'s `spjParams.applyPartialClustering`. The project’s shimplify contract at docs/dev/shimplify.md:67-69 confirms that each build compiles every file listing its build version, so adding the versions to this helper does not redirect the scan call site. Apache Spark’s fixed reconstruction passes the flag explicitly; without it, partial clustering can satisfy the downstream distribution incorrectly and omit the correctness-required shuffle for post-join deduplication.
Confidence: 🟣❗ CERTAIN
Code:
```scala
sql-plugin/src/main/spark359/scala/com/nvidia/spark/rapids/shims/KeyGroupedPartitioningShim.scala:18-21
{"spark": "403"}
{"spark": "404"}
{"spark": "412"}
{"spark": "413"}
sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala:140-141
k.copy(expressions = expressions, numPartitions = newPartValues.length,
partitionValues = newPartValues)```
Fix: Route the reconstruction at sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala:140-141 through `copyWithNewPartitionValues`. Pass a copy containing the projected expressions plus `sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala:53`'s `spjParams.applyPartialClustering`. This keeps one scan implementation: the old-version helper at sql-plugin/src/main/spark350/scala/com/nvidia/spark/rapids/shims/KeyGroupedPartitioningShim.scala:38-42 accepts and intentionally ignores the unavailable marker, while the fixed-version helper at sql-plugin/src/main/spark359/scala/com/nvidia/spark/rapids/shims/KeyGroupedPartitioningShim.scala:30-37 writes it.
Diff:
```diff
--- sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala
+++ sql-plugin/src/main/spark350db143/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala
@@ -137,6 +137,8 @@
case Some(projectionPositions) => projectionPositions.map(i => k.expressions(i))
case _ => k.expressions
}
- k.copy(expressions = expressions, numPartitions = newPartValues.length,
- partitionValues = newPartValues)
+ KeyGroupedPartitioningShim.copyWithNewPartitionValues(
+ k.copy(expressions = expressions),
+ newPartValues,
+ spjParams.applyPartialClustering)
case p => p```
Reference: docs/dev/shimplify.md:67-69 defines build-version source selection; sql-plugin/src/main/spark340/scala/com/nvidia/spark/rapids/shims/GpuBatchScanExec.scala:142-143 demonstrates the project’s existing scan-to-helper boundary; sql-plugin/src/main/spark350/scala/com/nvidia/spark/rapids/shims/KeyGroupedPartitioningShim.scala:38-42 and sql-plugin/src/main/spark359/scala/com/nvidia/spark/rapids/shims/KeyGroupedPartitioningShim.scala:30-37 provide the compatible old/new helper implementations. Apache Spark’s official fix writes `isPartiallyClustered = spjParams.applyPartialClustering` in its scan reconstruction at https://github.com/apache/spark/pull/54851/files#diff-42b726bc796730a2ccbc17bf11a922033893751208db376835cd26f1a0e58ca7, and https://issues.apache.org/jira/browse/SPARK-55848 documents the wrong-result impact and fixed versions.
(Reviewer: architecture)
Signed-off-by: Allen Xu <allxu@nvidia.com>
|
build |
|
Thanks, updated. |
| */ | ||
| /*** spark-rapids-shim-json-lines | ||
| {"spark": "411"} | ||
| {"spark": "412"} |
There was a problem hiding this comment.
1. [🔴 MUST FIX] Can the Spark 4.1.2 and 4.1.3 release profiles be wired to the Iceberg 1.11 module before these changed shim registrations are enabled? (confirmed: still present from prior review) The annotations at iceberg/iceberg-1-11-x/src/main/spark411/scala/com/nvidia/spark/rapids/iceberg/iceberg111x/IcebergProviderImpl.scala:18 and iceberg/iceberg-1-11-x/src/main/spark411/scala/com/nvidia/spark/rapids/iceberg/iceberg111x/IcebergProviderImpl.scala:19 declare the Iceberg 1.11 implementation available to both builds, but pom.xml:789 through pom.xml:796 still configure Spark 4.1.2 with the Spark 4.0 Iceberg 1.10 runtime and the empty stub module, and pom.xml:827 through pom.xml:834 do the same for Spark 4.1.3
Line: 18 (diff)
Issue: Can the Spark 4.1.2 and 4.1.3 release profiles be wired to the Iceberg 1.11 module before these changed shim registrations are enabled? (confirmed: still present from prior review) The annotations at iceberg/iceberg-1-11-x/src/main/spark411/scala/com/nvidia/spark/rapids/iceberg/iceberg111x/IcebergProviderImpl.scala:18 and iceberg/iceberg-1-11-x/src/main/spark411/scala/com/nvidia/spark/rapids/iceberg/iceberg111x/IcebergProviderImpl.scala:19 declare the Iceberg 1.11 implementation available to both builds, but pom.xml:789 through pom.xml:796 still configure Spark 4.1.2 with the Spark 4.0 Iceberg 1.10 runtime and the empty stub module, and pom.xml:827 through pom.xml:834 do the same for Spark 4.1.3. The inherited artifact selection at pom.xml:1096 therefore makes aggregator/pom.xml:100 through aggregator/pom.xml:110 package the stub instead of the newly registered GPU Iceberg implementation. A normal Spark 4.1.2 or 4.1.3 distribution cannot exercise or ship the intended GPU Iceberg scan, so the newly enabled end-to-end case at integration_tests/src/main/python/iceberg/iceberg_test.py:94 through integration_tests/src/main/python/iceberg/iceberg_test.py:98 does not validate the released artifact.
Confidence: 🟣❗ CERTAIN
Code:
```scala
|
iceberg/iceberg-1-11-x/src/main/spark411/scala/com/nvidia/spark/rapids/iceberg/iceberg111x/IcebergProviderImpl.scala:18-19
{"spark": "412"}
{"spark": "413"}
pom.xml:789-796
<iceberg.artifact.suffix>${spark40x.iceberg.artifact.suffix}</iceberg.artifact.suffix>
<iceberg.runtime.version>${iceberg.110x.version}</iceberg.runtime.version>
...
<module>iceberg/iceberg-stub</module>```
Fix: Mirror the established Spark 4.1.1 Iceberg 1.11 profile at pom.xml:750 through pom.xml:758 for the Spark 4.1.2 and 4.1.3 profiles, and make the identical generated-profile updates in scala2.13/pom.xml:789 through scala2.13/pom.xml:834. That selects the Spark 4.1 Iceberg 1.11 runtime, builds iceberg/iceberg-1-11-x/pom.xml:29, and packages its artifact.
Diff:
```diff
|
--- pom.xml
+++ pom.xml
@@ release412
- <iceberg.artifact.suffix>${spark40x.iceberg.artifact.suffix}</iceberg.artifact.suffix>
- <iceberg.runtime.version>${iceberg.110x.version}</iceberg.runtime.version>
+ <iceberg.artifact.suffix>${spark41x.iceberg.artifact.suffix}</iceberg.artifact.suffix>
+ <iceberg.runtime.version>${iceberg.111x.version}</iceberg.runtime.version>
+ <rapids.iceberg.artifactId>rapids-4-spark-iceberg-1-11-x</rapids.iceberg.artifactId>
@@ release412 modules
- <module>iceberg/iceberg-stub</module>
+ <module>iceberg/iceberg-1-11-x</module>
@@ release413
- <iceberg.artifact.suffix>${spark40x.iceberg.artifact.suffix}</iceberg.artifact.suffix>
- <iceberg.runtime.version>${iceberg.110x.version}</iceberg.runtime.version>
+ <iceberg.artifact.suffix>${spark41x.iceberg.artifact.suffix}</iceberg.artifact.suffix>
+ <iceberg.runtime.version>${iceberg.111x.version}</iceberg.runtime.version>
+ <rapids.iceberg.artifactId>rapids-4-spark-iceberg-1-11-x</rapids.iceberg.artifactId>
@@ release413 modules
- <module>iceberg/iceberg-stub</module>
+ <module>iceberg/iceberg-1-11-x</module>
--- scala2.13/pom.xml
+++ scala2.13/pom.xml
@@ release412
- <iceberg.artifact.suffix>${spark40x.iceberg.artifact.suffix}</iceberg.artifact.suffix>
- <iceberg.runtime.version>${iceberg.110x.version}</iceberg.runtime.version>
+ <iceberg.artifact.suffix>${spark41x.iceberg.artifact.suffix}</iceberg.artifact.suffix>
+ <iceberg.runtime.version>${iceberg.111x.version}</iceberg.runtime.version>
+ <rapids.iceberg.artifactId>rapids-4-spark-iceberg-1-11-x</rapids.iceberg.artifactId>
@@ release412 modules
- <module>iceberg/iceberg-stub</module>
+ <module>iceberg/iceberg-1-11-x</module>
@@ release413
- <iceberg.artifact.suffix>${spark40x.iceberg.artifact.suffix}</iceberg.artifact.suffix>
- <iceberg.runtime.version>${iceberg.110x.version}</iceberg.runtime.version>
+ <iceberg.artifact.suffix>${spark41x.iceberg.artifact.suffix}</iceberg.artifact.suffix>
+ <iceberg.runtime.version>${iceberg.111x.version}</iceberg.runtime.version>
+ <rapids.iceberg.artifactId>rapids-4-spark-iceberg-1-11-x</rapids.iceberg.artifactId>
@@ release413 modules
- <module>iceberg/iceberg-stub</module>
+ <module>iceberg/iceberg-1-11-x</module>```
Reference: pom.xml:750 through pom.xml:758 are the repository's working Spark 4.1.1 pattern: they select the Spark 4.1 artifact suffix, Iceberg 1.11 runtime, Iceberg 1.11 RAPIDS artifact, and Iceberg 1.11 module. iceberg/iceberg-1-11-x/pom.xml:29 declares the artifact that must replace the stub, while iceberg/iceberg-stub/pom.xml:29 and iceberg/iceberg-stub/pom.xml:41 through iceberg/iceberg-stub/pom.xml:51 show that the currently selected artifact contains no Iceberg implementation dependency.
(Reviewer: correctness)
There was a problem hiding this comment.
Agreed. The Spark 4.1.2/4.1.3 Iceberg profile and packaging wiring is a separate issue outside the scope of this PR. It should be addressed in a follow-up PR, so I am leaving this thread unresolved here.
Signed-off-by: Chong Gao <chongg@nvidia.com>
|
build |
firestarman
left a comment
There was a problem hiding this comment.
Posted the remaining MUST FIX and SHOULD FIX findings from the refreshed quick review as inline comments. The other MUST FIX is already tracked on its code line: #15397 (comment)
Signed-off-by: Chong Gao <chongg@nvidia.com>
Signed-off-by: Chong Gao <chongg@nvidia.com>
Follow-up to #15286.
Description
Preserves Spark's partial-clustering correctness marker when
GpuBatchScanExecrebuilds aKeyGroupedPartitioningon Spark versions containing the SPARK-55848 fix:The marker-aware shim and its focused unit test cover these fixed version ranges. Earlier patch
releases remain on the old shim because their
KeyGroupedPartitioningdoes not contain thecorrectness marker.
On Spark 3.5.9,
GpuBatchScanExeccalculated the replicated, padded, and reorderedfinalPartitionsrequired by a partially clustered storage-partitioned join, but constructedGpuDataSourceRDDfrom the originalfilteredPartitions. This could fail execution with unequaljoin-side partition counts. The scan now constructs the RDD from
finalPartitions, matchingSpark's
BatchScanExec.The Iceberg end-to-end regression test enables
spark.sql.iceberg.planning.preserve-data-groupingso the scans actually participate in ashuffle-free storage-partitioned join. It compares non-empty CPU and GPU results and applies a
Python callback to the executed GPU plan. The callback asserts:
GpuBatchScanExecnodes;GpuShuffledSymmetricHashJoinExec;GpuShuffleExchangeExecin the complete plan, representing the post-joinDISTINCTshuffle; andGpuShuffleExchangeExecanywhere in the join subtree, proving both SPJ inputs areshuffle-free.
This structural assertion replaces the previous flat class-existence check, which could mistake
join-side shuffles for the required post-join shuffle. As a negative control, disabling Iceberg
data-grouping preservation produces two GPU shuffles and fails the new assertion with
Expected one post-join GPU shuffle, found 2.The Spark 4.1 Iceberg 1.11 shim sources are also enabled for Spark 4.1.2 and 4.1.3 so the
end-to-end test can exercise a GPU Iceberg scan on those versions.
Local validation with JDK 17:
KeyGroupedPartitioningShimSuite: passed on Spark 3.5.9, 4.0.3, 4.0.4, 4.1.2, and 4.1.3test_iceberg_spj_partial_clustering_distinct: passed on Spark 3.5.9 with Iceberg 1.9.2test_iceberg_spj_partial_clustering_distinct: passed on Spark 4.0.3 and 4.0.4 withIceberg 1.10.1
test_iceberg_spj_partial_clustering_distinct: passed on Spark 4.1.2 and 4.1.3 withIceberg 1.11.0
Can't zip RDDs with unequal numbers of partitions: List(4, 3)Spark APIs
python -m py_compile integration_tests/src/main/python/asserts.py integration_tests/src/main/python/iceberg/iceberg_test.pygit diff --checkChecklists
Documentation
Testing
(Please provide the names of the existing tests in the PR description.)
Performance