Fix columnar mismatch bug in iceberg dml when aqe enabled.

[liurenjie1024]

Fixes nvbugs 5689547.

Description

When aqe is enabled, we need to change it to support columnar for iceberg dml operations. This pr also added necessary tests.

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 Overview

Greptile Summary

Fixed columnar mismatch bug in Iceberg DML operations when Adaptive Query Execution (AQE) is enabled by properly handling AdaptiveSparkPlanExec in the query plan.

Key Changes

• IcebergProviderImpl.scala: Added checkChildPlan validation method that allows AdaptiveSparkPlanExec children to proceed without requiring immediate GPU replacement check. Removed manual GpuRowToColumnarExec wrapping logic from all convertToGpu methods across write operations (AppendData, OverwritePartitions, OverwriteByExpression, ReplaceData, WriteDelta).

• WriteToDataSourceV2Exec.scala: Centralized columnar handling by introducing finalQuery that sets supportsColumnar=true on AdaptiveSparkPlanExec and unwraps GpuColumnarToRowExec. Removed redundant query override methods that were unwrapping columnar-to-row conversions.

• Test Coverage: Added comprehensive AQE-enabled tests for all DML operations (INSERT, CTAS, RTAS, DELETE, UPDATE, MERGE, dynamic/static overwrite) with both partitioned and unpartitioned tables, covering copy-on-write and merge-on-read modes.

Root Cause

When AQE is enabled, GpuShuffleCoalesceExec would end up as a child without proper columnar support flagging, causing "ROW BASED PROCESSING IS NOT SUPPORTED" errors. The fix ensures AQE plans are marked as columnar-capable and validation happens at the right stage.

Confidence Score: 5/5

• This PR is safe to merge with minimal risk
• The fix correctly addresses the root cause by properly handling AQE plans, removes redundant code, and includes comprehensive test coverage across all affected DML operations. The logic is sound: validation is moved to tagging phase and columnar handling is centralized, preventing the mismatch that caused the bug.
• No files require special attention

Important Files Changed

File Analysis

Filename	Score	Overview
iceberg/src/main/scala/com/nvidia/spark/rapids/iceberg/IcebergProviderImpl.scala	5/5	Added checkChildPlan validation to all write operations and removed manual GpuRowToColumnarExec wrapping from convert methods
sql-plugin/src/main/spark350/scala/org/apache/spark/sql/execution/datasources/v2/WriteToDataSourceV2Exec.scala	5/5	Centralized columnar handling in finalQuery to properly support AQE by setting supportsColumnar=true on AdaptiveSparkPlanExec
integration_tests/src/main/python/iceberg/iceberg_ctas_test.py	5/5	Added test_ctas_aqe test documenting NVBUGS-5689547 fix, plus @allow_non_gpu decorator to VALUES test

Sequence Diagram

sequenceDiagram
    participant User as User Query
    participant Meta as SparkPlanMeta
    participant IcebergProvider as IcebergProviderImpl
    participant WriteExec as GpuV2TableWriteExec
    participant AQE as AdaptiveSparkPlanExec
    participant FinalQuery as finalQuery

    Note over User,FinalQuery: When AQE is Enabled

    User->>Meta: tagForGpu()
    Meta->>IcebergProvider: tagForGpu(cpuExec, meta)
    IcebergProvider->>IcebergProvider: checkChildPlan(meta)
    alt Child is AdaptiveSparkPlanExec
        IcebergProvider-->>Meta: Allow (AQE handles columnar)
    else Child cannot run on GPU
        IcebergProvider-->>Meta: willNotWorkOnGpu()
    end

    User->>Meta: convertToGpu()
    Meta->>IcebergProvider: convertToGpu(cpuExec, meta)
    IcebergProvider->>WriteExec: Create GpuAppendDataExec(inner, ...)
    Note over WriteExec: inner = meta.childPlans.head.convertIfNeeded()

    User->>WriteExec: execute()
    WriteExec->>FinalQuery: Determine finalQuery
    alt query is AdaptiveSparkPlanExec
        FinalQuery->>AQE: aqe.copy(supportsColumnar = true)
        Note over AQE: AQE now supports columnar output
    else query is GpuColumnarToRowExec
        FinalQuery->>FinalQuery: Unwrap to get inner plan
    else
        FinalQuery->>FinalQuery: Use query as-is
    end

    WriteExec->>FinalQuery: executeColumnar()
    FinalQuery-->>WriteExec: RDD[ColumnarBatch]
    WriteExec->>WriteExec: writeWithV2(batchWrite)
    WriteExec-->>User: Write complete

Loading

[greptile-apps]

_{12 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{10 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[liurenjie1024]

build

[res-life]

LGTM