Fix incorrect GPU partitioning and ordering metadata [databricks] by revans2 · Pull Request #16620 · NVIDIA/cudf-spark

revans2 · 2026-10-27T01:35:02Z

GPU projections and aggregations could report partitioning by output position and pass ordering through unchanged. For a query like

spark.range(100).selectExpr('id % 10 as k', 'id as v').groupBy('k').agg(max(col('v')).collect()

The result is a correct 10 rows, but if you persist the result in between the selectExpr and the groupBy like

val data = spark.range(100).selectExpr('id % 10 as k', 'id as v').cache()
data.groupBy('k').agg(max(col('v')).collect()

on the GPU, this causes the ordering and partitioning from the GpuProject to be used and the shuffle before the partial aggregation is not inserted, resulting in 100 rows being returned. This is because the existing code only goes off of position, not inspecting the actual expressions.

This fixes the issue by reusing the installed Spark version’s alias-aware metadata logic for GPU projections and aggregates. It keeps CPU expressions for metadata matching separate from GPU execution expressions, and preserves them through plan copies and projection pruning. If those are not available a conservative fallback happens. This will not be perfect, because the ordering and partitioning can include CPU expressions in them that will not match the GPU expressions, but it will not corrupt the data. In the worst case we may end up with more shuffle operations than before.

Validation:

Standalone Spark 3.5.5 build: 27 Scala and 9 Java tests passed.
Cached aggregate regression passed with AQE on/off under both PCBS and the default cache serializer.
The same metadata implementation passed focused suites on Spark 3.3.4 and 4.1.3.

This appears to have been caused by #14164 and #14264. The first one put in the logic and the second one made it so more operators adopted it. The original logic really only worked for union, and even there I am not 100% sure that it always did.

I will try to run some performance tests, but as this is a data corruption fix it is not strictly required, and if the issue is hit, we will take a performance hit in order to produce the correct result (as a shuffle must be inserted to make the plan correct).

AI-assisted implementation and testing.

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

Map partitioning attributes through direct aliases rather than projection output positions. Projecting id % 10 AS k before id AS v otherwise falsely advertised range partitioning on k. Cached scans preserved that claim, so Spark omitted the grouping shuffle and returned per-partition results. Preserve identity and direct CPU/GPU attribute aliases, and let existing reference filtering discard unsupported computed-column partitioning. Keep aggregate and other operator implementations unchanged. Add a compact projection metadata regression and the 100-row cached aggregation reproduction with AQE both enabled and disabled. The new integration regression fails the original jar with CPU 10 / GPU 100 rows. Validation: 25 Scala and 9 Java tests passed on Spark 3.5.5. The new integration regression passed with default cache and PCBS, AQE on/off; full PCBS cache suite: 625 passed, 4 large-data skips. Standalone shell reproductions and controls returned exact results with both serializers. SQL compile/package checks passed on Spark 3.3.4 and 4.0.1/Scala 2.13. Status: READY FOR REVIEW. Verification evidence is retained outside the repository under /data/pcbs-oom/project-partitioning-fix, as requested. Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

Use the installed Spark version's partitioning and ordering traits instead of maintaining a separate attribute remapping algorithm. Keep original CPU expressions for metadata matching in a second constructor parameter list, following GPU sort and shuffle, so Spark does not treat them as executable expressions. Preserve those expressions across child replacement and projection pruning. Compare metadata directly with native Spark for aliases, computed expressions, ordering prefixes, null ordering and candidate limits. Keep the cached aggregate regression for both AQE settings. Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

Remove canonicalized-equality assertions that cannot establish anything about second-parameter-list CPU metadata. Keep the direct expression and copy-preservation checks. Collect cached query results before unpersisting in a finally block for each CPU and GPU session, including when query execution fails. Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

greptile-apps · 2026-10-27T01:38:58Z

Confidence Score: 2/5

[High impact] The PR changes shuffle and sort planning across GPU SQL queries.

The PR still needs the performance evidence required by the repository before merging.

Summary

The PR uses the installed Spark version’s alias-aware logic to report partitioning and ordering for GPU projections and aggregates.

GPU projections report partitioning and ordering from their actual outputs.
GPU aggregates match output partitioning through their CPU result expressions.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  CPU[Original CPU expressions] --> Metadata[Spark alias-aware metadata]
  GPU[GPU expressions] --> Execution[GPU execution]
  Metadata --> Planning[Partitioning and ordering]
  Copies[Plan copies and projection pruning] --> CPU
  Copies --> GPU

_{Reviews (4) · Last reviewed commit: "Use standard CPU GPU comparison helper f..." · Reviewed by Greptile}

Copilot

🟡 Changes recommended

Required performance results or a linked validation issue are missing for changes that can alter shuffle and sort execution.

1 open issue

Add performance validation for Spark metadata change · New

What changed in this PR

Fixes incorrect partitioning and ordering metadata in GPU projections and aggregates, preventing Spark from skipping required shuffles and returning incorrect results.

Changes:

Reuses Spark’s version-specific alias-aware metadata logic while keeping CPU metadata expressions separate from GPU execution expressions.
Preserves metadata through plan copies, projection pruning, and Delta rewrites.
Adds metadata parity tests and a cached-aggregation regression with adaptive execution enabled and disabled.

Shim validation remains INCOMPLETE: the repository snapshot was unavailable, and current-head/base shim-compatibility and shim-related-prs reports could not be verified. The reported focused tests do not establish full affected-profile coverage.

File	Description
tests/src/test/scala/org/apache/spark/sql/rapids/ProjectExprSuite.scala	Tests alias metadata, ordering, candidate limits, and rewrites.
tests/src/test/scala/com/nvidia/spark/rapids/ProjectSplitRetrySuite.scala	Updates projection construction.
tests/src/test/scala/com/nvidia/spark/rapids/HashAggregatesSuite.scala	Updates aggregate construction.
sql-plugin/src/main/spark340/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinExec.scala	Updates join projection construction.
sql-plugin/src/main/spark340/scala/com/nvidia/spark/rapids/shims/GpuAliasAwareExecNodeShims.scala	Adds Spark 3.4+ alias-aware metadata adapters.
sql-plugin/src/main/spark332db/scala/org/apache/spark/sql/rapids/GpuFileFormatWriter.scala	Updates write projection construction.
sql-plugin/src/main/spark330db/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinExec.scala	Updates Databricks join projections.
sql-plugin/src/main/spark330/scala/org/apache/spark/sql/rapids/GpuFileFormatWriter.scala	Updates write projection construction.
sql-plugin/src/main/spark330/scala/org/apache/spark/sql/rapids/execution/GpuBroadcastNestedLoopJoinExec.scala	Updates Spark 3.3 join projections.
sql-plugin/src/main/spark330/scala/com/nvidia/spark/rapids/shims/GpuAliasAwareExecNodeShims.scala	Adds Spark 3.3 alias-aware metadata adapters.
sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuRangeBoundaryPlan.scala	Preserves metadata during projection pruning.
sql-plugin/src/main/scala/org/apache/spark/sql/rapids/execution/GpuHashJoin.scala	Updates generated join projections.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/window/GpuWindowExecMeta.scala	Updates generated window projections.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAliasAwareOutputExpressions.scala	Converts GPU aliases for metadata matching.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAggregateExec.scala	Uses Spark metadata logic with retained CPU expressions.
sql-plugin/src/main/scala/com/nvidia/spark/rapids/basicPhysicalOperators.scala	Replaces positional metadata mapping and preserves projection metadata.
integration_tests/src/main/python/cache_test.py	Adds cached-projection aggregation regression.
delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/DeltaSpark400DB173Provider.scala	Preserves metadata during Databricks Delta rewrites.
delta-lake/delta-spark332db/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuOptimisticTransactionBase.scala	Updates transaction projection construction.
delta-lake/delta-spark330db/src/main/scala/com/databricks/sql/transaction/tahoe/rapids/GpuOptimisticTransactionBase.scala	Updates transaction projection construction.
delta-lake/common/src/main/delta-io/scala/org/apache/spark/sql/delta/rapids/AbstractGpuOptimisticTransactionBase.scala	Updates shared Delta projection construction.
delta-lake/common/src/main/delta-33x-42x/scala/com/nvidia/spark/rapids/delta/common/DeltaProviderBase.scala	Preserves metadata through Delta pruning and merging.
delta-lake/common/src/main/db-350db143-410db183/scala/com/databricks/sql/transaction/tahoe/rapids/GpuOptimisticTransactionBaseCommon.scala	Updates shared Databricks projection construction.

_{Review effort: Balanced}

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

revans2 · 2026-10-27T02:05:12Z

build

nartal1

Thanks @revans2 for tracking this down and fixing it. Reusing Spark’s alias-aware logic makes sense. Our earlier tests focused on union behavior and missed this completely.

amahussein

LGTME

A couple of small inline requests.

A test could cover beyond range partitioning.
add a couple of comments to define the contract for the new cpuProjectList argument in the GpuProjectExec.

amahussein · 2026-10-27T21:42:36Z


+@pytest.mark.parametrize('aqe_enabled', ['true', 'false'])
+@allow_non_gpu('InMemoryTableScanExec', 'ColumnarToRowExec')
+def test_cache_project_partitioning(aqe_enabled):


This test feels narrower than the bug. It covers spark.range which is ranged-partition on id which it is named later to v. It misses the hash-partitioned; join followed by projectio; and a join whose condition the GPU can't run as one AST expression.

Can we add a different test that will cover the 3 different cases with both AQE on/off?

@pytest.mark.parametrize('shape', ['repartition', 'join', 'join_split_condition']) ... # set the condition bawed on the condition parameter. ...

amahussein · 2026-10-27T21:58:22Z

   child: SparkPlan,
-   enablePreSplit: Boolean = true) extends GpuProjectExecLike {
+   enablePreSplit: Boolean = true)(
+   val cpuProjectList: Option[List[NamedExpression]] = None) extends GpuProjectExecLike {


It took me a while to understand the change in the construction and why the PR used to replace all the call sites with withProjectList.
My concern is that someone later unfamiliar does not maintain the CPU/GPU expression lists mapped correctly.
I don't believe there is a bug in there so far.

Either we can add some rule to auto-detect that, or just add a comment to record the rule.

// The CPU form of each projectList entry, read by Spark's partitioning // and ordering logic. // Change projectList through withProjectList, which updates this list // to match.

amahussein · 2026-10-27T22:33:08Z

@revans2 did you check the [AI-AUDIT] GpuUnionExec re-derives partitioning-aware vs concatenating union on every call (SPARK-59122) #16609 ? It sounds like spark-4.1+ GpuUnionExec re-derives its plan at execution.
@nartal1 do you need to update Refactor outputPartitioning for few execs and add tests #14229 striking off item number 2 off the list?

revans2 added 3 commits October 27, 2026 01:18

revans2 requested review from jihoonson and nartal1 and a balanced review from Copilot October 27, 2026 01:35

Copilot started reviewing on behalf of revans2 October 27, 2026 01:35

View session

Merge main into fix-gpu-output-metadata

57bee97

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

Copilot AI reviewed Oct 27, 2026

View reviewed changes

Comment thread sql-plugin/src/main/scala/com/nvidia/spark/rapids/basicPhysicalOperators.scala

greptile-apps Bot reviewed Oct 27, 2026

View reviewed changes

Comment thread integration_tests/src/main/python/cache_test.py Outdated

Update copyright years in modified Delta sources

7ef09c6

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

greptile-apps Bot reviewed Oct 27, 2026

View reviewed changes

Comment thread sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuAggregateExec.scala

Use standard CPU GPU comparison helper for cache regression

1784e2e

Signed-off-by: Robert (Bobby) Evans <bobby@apache.org>

revans2 mentioned this pull request Oct 27, 2026

Reduce Parquet cache heap pressure with smaller row-group records [databricks] #16633

Draft

3 of 8 tasks

sameerz added the bug Something isn't working label Oct 27, 2026

amahussein self-requested a review October 27, 2026 16:40

nartal1 approved these changes Oct 27, 2026

View reviewed changes

amahussein requested changes Oct 27, 2026

View reviewed changes

nartal1 mentioned this pull request Oct 28, 2026

Refactor outputPartitioning for few execs and add tests #14229

Open

Navigation Menu

Repository navigation

Fix incorrect GPU partitioning and ordering metadata [databricks] - #16620

Fix incorrect GPU partitioning and ordering metadata [databricks]#16620
revans2 wants to merge 6 commits into
NVIDIA:mainfrom
revans2:fix-gpu-output-metadata

revans2 commented Oct 27, 2026 •

edited

Loading

Uh oh!

greptile-apps Bot commented Oct 27, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

revans2 commented Oct 27, 2026

Uh oh!

nartal1 left a comment

Uh oh!

amahussein left a comment

Uh oh!

amahussein Oct 27, 2026

Uh oh!

amahussein Oct 27, 2026

Uh oh!

amahussein commented Oct 27, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Conversation

revans2 commented Oct 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Checklists

Uh oh!

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

Uh oh!

Confidence Score: 2/5

Summary

Diagram

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

Uh oh!

Uh oh!

Uh oh!

revans2 commented Oct 27, 2026

Uh oh!

nartal1 left a comment

Choose a reason for hiding this comment

Uh oh!

amahussein left a comment

Choose a reason for hiding this comment

Uh oh!

amahussein Oct 27, 2026

Choose a reason for hiding this comment

Uh oh!

amahussein Oct 27, 2026

Choose a reason for hiding this comment

Uh oh!

amahussein commented Oct 27, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

revans2 commented Oct 27, 2026 •

edited

Loading

greptile-apps Bot commented Oct 27, 2026 •

edited

Loading