[25.10] Add missing handling for the 0 input partition count in computing partition count in GpuOptimizeWriteExchangeExec

[jihoonson]

Backport of #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.
• 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.

[greptile-apps]

Greptile Summary

• Fixes ArithmeticException by adding zero partition check in actualNumPartitions calculation (line 122-129)
• Improves consistency by marking childNumPartitions, actualNumPartitions, and actualPartitioning as @transient private lazy val
• Adds comprehensive test coverage for empty output scenario with 20 partitions to verify fix

Confidence Score: 5/5

• This PR is safe to merge with minimal risk
• Score reflects a simple, focused bug fix with proper zero-partition handling, comprehensive test coverage, and improved code consistency through @transient annotations. The fix directly addresses the ArithmeticException issue without introducing new complexity
• No files require special attention

Important Files Changed

Filename	Overview
delta-lake/common/src/main/delta-33x/scala/org/apache/spark/sql/delta/rapids/GpuOptimizeWriteExchangeExec.scala	Added zero partition check to prevent ArithmeticException in actualNumPartitions calculation and improved lazy val consistency

Sequence Diagram

sequenceDiagram
    participant Spark as "Spark Execution"
    participant Exchange as "GpuOptimizeWriteExchangeExec"
    participant Child as "Child SparkPlan"
    
    Spark->>Exchange: "executeColumnar()"
    Exchange->>Child: "executeColumnar()"
    Child-->>Exchange: "inputRDD"
    Exchange->>Exchange: "Check childNumPartitions == 0"
    alt childNumPartitions == 0
        Exchange->>Exchange: "Set actualNumPartitions = 0"
        Exchange->>Exchange: "Return Future.successful(null)"
        Exchange->>Spark: "Return ShuffledBatchRDD"
    else childNumPartitions > 0
        Exchange->>Exchange: "Calculate actualNumPartitions"
        Exchange->>Exchange: "Submit map stage"
        Exchange->>Exchange: "Collect statistics"
        Exchange->>Exchange: "Rebalance partitions"
        Exchange->>Spark: "Return ShuffledBatchRDD with partition specs"
    end

Loading

[greptile-apps]

_{2 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile
React with 👍 or 👎 to share your feedback on this new summary format}

[gerashegalov]

I usually do git cherry-pick -x

Can you manually add a line

"(cherry picked from commit ...)"

to the PR description to track lineage

[sameerz]

build

[gerashegalov]

LGTM

[jihoonson]

@gerashegalov updated the PR description.