Keep the TableCacheQueryStageExec on the CPU but convert the child plan by razajafri · Pull Request #13944 · NVIDIA/cudf-spark

razajafri · 2025-12-04T04:54:34Z

Description

A customer has reported a bug where writing to a Delta table after caching Parquet scan results in an exception. The fix keeps the TableCacheQueryStageExec on the CPU but brings the child plan InMemoryTableScanExec on the GPU.

Checklists

This PR has added documentation for new or modified features or behaviors.
This PR has added new tests or modified existing tests to cover new code paths.
(Please explain in the PR description how the new code paths are tested, such as names of the new/existing tests that cover them.)
Performance testing has been performed and its results are added in the PR description. Or, an issue has been filed with a link in the PR description.

…ting partition count in GpuOptimizeWriteExchangeExec (NVIDIA#13804) Backport of NVIDIA#13780. Cherry picked from commit f02e77d. ### Description `GpuOptimizeWriteExchangeExec` computes the `actualNumPartitions` based on the partition count of its input. The computed `actualNumPartitions` is used to dynamically optimize the partitioning. This logic is currently missing handling of the case when the input partition count is 0, which can cause the ArithmeticException error. This PR adds a proper handling of the case. ### Checklists - [ ] This PR has added documentation for new or modified features or behaviors. - [x] This PR has added new tests or modified existing tests to cover new code paths. (Please explain in the PR description how the new code paths are tested, such as names of the new/existing tests that cover them.) - [ ] Performance testing has been performed and its results are added in the PR description. Or, an issue has been filed with a link in the PR description. Signed-off-by: Jihoon Son <ghoonson@gmail.com>

…VIDIA#13782) This is a cherry-pick of NVIDIA#13775 for 25.10. ### Description The PR is a clean cherry-pick of a commit we are adding in 25.12 to circumvent an exception in the state machine. See NVIDIA#13774 for more info. ### Checklists - [x] This PR has added documentation for new or modified features or behaviors. - [ ] This PR has added new tests or modified existing tests to cover new code paths. (Please explain in the PR description how the new code paths are tested, such as names of the new/existing tests that cover them.) - [ ] Performance testing has been performed and its results are added in the PR description. Or, an issue has been filed with a link in the PR description. Signed-off-by: Alessandro Bellina <abellina@nvidia.com>

…DIA#13831) This cherry picks NVIDIA#13806 to 25.10 - [ ] This PR has added documentation for new or modified features or behaviors. - [ ] This PR has added new tests or modified existing tests to cover new code paths. (Please explain in the PR description how the new code paths are tested, such as names of the new/existing tests that cover them.) - [ ] Performance testing has been performed and its results are added in the PR description. Or, an issue has been filed with a link in the PR description. --------- Signed-off-by: Alessandro Bellina <abellina@nvidia.com>

Signed-off-by: Raza Jafri <raza.jafri@gmail.com>

greptile-apps · 2025-12-04T04:58:56Z

Greptile Overview

Greptile Summary

This PR attempts to fix a bug where writing to a Delta table after caching Parquet scan results throws a RapidsHostColumnVector cannot be cast to GpuColumnVector exception. The approach adds a new shim method handleTableCacheInOptimizeAdaptiveTransitions intended to keep TableCacheQueryStageExec on the CPU without recursively processing its children.

Added handleTableCacheInOptimizeAdaptiveTransitions method to SparkShims trait with default returning None
Overrode in Spark350PlusNonDBShims to return Some(tcqs) for TableCacheQueryStageExec, which would prevent child processing
Modified default case in optimizeAdaptiveTransitions to check the new handler before processing children

Critical Issue: The new handler appears to be unreachable for TableCacheQueryStageExec because an earlier case at line 177-179 already matches it via getTableCacheNonQueryStagePlan(p).nonEmpty. Both shim methods use identical pattern matching for TableCacheQueryStageExec, so the earlier case always wins and recursively processes children regardless of the new handler.

Confidence Score: 2/5

This PR has a potential logic issue where the new code path may never be executed due to pattern matching order
The new handleTableCacheInOptimizeAdaptiveTransitions method is placed in the default case but TableCacheQueryStageExec is already matched by an earlier case at line 177-179. This means the fix may not work as intended. The shim additions are clean but the integration point in GpuTransitionOverrides appears to be dead code for the target use case.
GpuTransitionOverrides.scala - the new handler placement appears incorrect; it should either be moved before line 177 or line 177-179 should be modified

Important Files Changed

File Analysis

Filename	Score	Overview
sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuTransitionOverrides.scala	2/5	Added handler for TableCacheQueryStageExec in default case, but this code appears unreachable since line 177-179 already matches TableCacheQueryStageExec first
sql-plugin/src/main/scala/com/nvidia/spark/rapids/SparkShims.scala	4/5	Added new shim method `handleTableCacheInOptimizeAdaptiveTransitions` with default returning None; clean addition to trait
sql-plugin/src/main/spark350/scala/com/nvidia/spark/rapids/shims/Spark350PlusNonDBShims.scala	4/5	Override returns `Some(tcqs)` for TableCacheQueryStageExec to prevent child processing; implementation is correct

Sequence Diagram

sequenceDiagram
    participant AQE as AdaptiveQueryExec
    participant OAT as optimizeAdaptiveTransitions
    participant TCQS as TableCacheQueryStageExec
    participant IMTS as InMemoryTableScanExec

    AQE->>OAT: Process plan with TableCacheQueryStageExec
    OAT->>OAT: Match case at line 177-179<br/>(getTableCacheNonQueryStagePlan.nonEmpty)
    OAT->>TCQS: Get children (returns wrapped plan)
    OAT->>OAT: Recursively process child<br/>InMemoryTableScanExec
    Note over OAT: New handler at line 245-249<br/>is never reached for TCQS
    OAT->>TCQS: withNewChildren(optimized children)
    TCQS-->>AQE: Return modified TCQS

greptile-apps

_{10 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

razajafri · 2025-12-04T16:43:53Z

build

revans2

I still don't totally understand what is happening that caused this to be a problem, but making sure that the cached table wrapper stays on the CPU is confusing. Do we end up converting the data to and from the CPU around it?

Could we make sure that we test this on databricks, because you did not update the shims for databricks.

revans2 · 2025-12-04T19:16:48Z

+        df2 = unary_op_df(spark, data_gen).withColumnRenamed("a", "r_a").cache()
+        df1.count()
+        df2.count()
+        return df1.join(df2, df1.a == df2.r_a, 'Outer')


nit: I don't think that a and r_a will ever be equal. At least statistically it should be fairly rare.

revans2 · 2025-12-04T19:18:45Z

+@ignore_order(local=True)
+@allow_non_gpu("CollectLimitExec", "ShuffleExchangeExec", "ColumnarToRowExec")
+@pytest.mark.parametrize('enable_vectorized_conf', enable_vectorized_confs, ids=idfn)
+@allow_non_gpu_conditional(is_spark_350_or_351(), "InMemoryTableScanExec")


Is there a follow on issue for the cache not being on the GPU?

revans2 · 2025-12-04T19:19:28Z

+        ).persist(StorageLevel.MEMORY_AND_DISK)
+        minNbrs1.count()
+
+        ee.join(minNbrs1, "src")


This is a noop and should be deleted

revans2 · 2025-12-04T20:59:00Z

+
+  override def handleTableCacheInOptimizeAdaptiveTransitions(plan: SparkPlan,
+      parent: Option[SparkPlan]): Option[SparkPlan] = {
+    plan match {


So effectively we are saying for Spark 3.5.0+ (non-DB) don't do into TableCacheQueryStageExec? It is already handled.

greptile-apps

Additional Comments (1)

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuTransitionOverrides.scala, line 245-249 (link)

logic: The new handleTableCacheInOptimizeAdaptiveTransitions call appears to be unreachable for TableCacheQueryStageExec because line 177-179 already matches it first:
```
// Line 177-179 (earlier in the match):
case p if SparkShimImpl.getTableCacheNonQueryStagePlan(p).nonEmpty =>
  p.withNewChildren(p.children.map(c => optimizeAdaptiveTransitions(c, Some(p))))
```
Both getTableCacheNonQueryStagePlan and handleTableCacheInOptimizeAdaptiveTransitions match TableCacheQueryStageExec, so the earlier case wins.

If the goal is to prevent child processing for TableCacheQueryStageExec, the handler should be placed before line 177, or line 177-179 should be modified to use the new shim method's behavior.

_{3 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

gerashegalov · 2025-12-05T04:59:46Z

Please please up-merge to HEAD of release/25.12

razajafri · 2025-12-05T22:50:43Z

This is not needed as the bug is not reproducible in 25.10+

jihoonson and others added 5 commits November 17, 2025 14:07

Fix for TableCacheQueryStageExec

a41ec65

Signing off

851b5b7

Signed-off-by: Raza Jafri <raza.jafri@gmail.com>

greptile-apps Bot reviewed Dec 4, 2025

View reviewed changes

sameerz requested a review from a team December 4, 2025 14:28

sameerz added the bug Something isn't working label Dec 4, 2025

revans2 reviewed Dec 4, 2025

View reviewed changes

razajafri changed the base branch from branch-25.08 to branch-25.10 December 4, 2025 22:20

Merge branch 'branch-25.10' into SP-13878-table-cache-query-stage-fix

4462ce1

greptile-apps Bot reviewed Dec 4, 2025

View reviewed changes

razajafri changed the base branch from branch-25.10 to release/25.12 December 5, 2025 04:45

razajafri closed this Dec 5, 2025

Navigation Menu

Uh oh!

Keep the TableCacheQueryStageExec on the CPU but convert the child plan - #13944

Keep the TableCacheQueryStageExec on the CPU but convert the child plan#13944
razajafri wants to merge 6 commits into
NVIDIA:release/25.12from
razajafri:SP-13878-table-cache-query-stage-fix

razajafri commented Dec 4, 2025 •

edited

Loading

Uh oh!

greptile-apps Bot commented Dec 4, 2025 •

edited

Loading

Uh oh!

greptile-apps Bot left a comment

Uh oh!

razajafri commented Dec 4, 2025

Uh oh!

revans2 left a comment

Uh oh!

revans2 Dec 4, 2025

Uh oh!

revans2 Dec 4, 2025

Uh oh!

revans2 Dec 4, 2025

Uh oh!

revans2 Dec 4, 2025

Uh oh!

greptile-apps Bot left a comment •

edited

Loading

Uh oh!

gerashegalov commented Dec 5, 2025

Uh oh!

razajafri commented Dec 5, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

Uh oh!

Conversation

razajafri commented Dec 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Checklists

Uh oh!

greptile-apps Bot commented Dec 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Overview

Greptile Summary

Confidence Score: 2/5

Important Files Changed

Sequence Diagram

Uh oh!

greptile-apps Bot left a comment

Choose a reason for hiding this comment

Uh oh!

razajafri commented Dec 4, 2025

Uh oh!

revans2 left a comment

Choose a reason for hiding this comment

Uh oh!

revans2 Dec 4, 2025

Choose a reason for hiding this comment

Uh oh!

revans2 Dec 4, 2025

Choose a reason for hiding this comment

Uh oh!

revans2 Dec 4, 2025

Choose a reason for hiding this comment

Uh oh!

revans2 Dec 4, 2025

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Additional Comments (1)

Uh oh!

gerashegalov commented Dec 5, 2025

Uh oh!

razajafri commented Dec 5, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

razajafri commented Dec 4, 2025 •

edited

Loading

greptile-apps Bot commented Dec 4, 2025 •

edited

Loading

greptile-apps Bot left a comment •

edited

Loading