fix: restore columnar transitions under the native Iceberg write

spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala

@@ -22,7 +22,7 @@ package org.apache.comet.rules
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.util.sideBySide
-import org.apache.spark.sql.comet.{CometCollectLimitExec, CometColumnarToRowExec, CometMapInBatchExec, CometNativeColumnarToRowExec, CometNativeWriteExec, CometPlan, CometSparkToColumnarExec}
+import org.apache.spark.sql.comet.{CometCollectLimitExec, CometColumnarToRowExec, CometIcebergWriteExec, CometMapInBatchExec, CometNativeColumnarToRowExec, CometNativeWriteExec, CometPlan, CometSparkToColumnarExec}
 import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometShuffleExchangeExec}
 import org.apache.spark.sql.comet.shims.{MapInBatchInfo, ShimCometMapInBatch}
 import org.apache.spark.sql.execution.{ColumnarToRowExec, RowToColumnarExec, SparkPlan}
@@ -74,7 +74,7 @@ case class EliminateRedundantTransitions(session: SparkSession)
   }
 
   private def _apply(plan: SparkPlan): SparkPlan = {
-    val eliminatedPlan = plan transformUp {
+    val eliminatedPlan = stripIcebergWriteInputTransition(plan) transformUp {
       case ColumnarToRowExec(shuffleExchangeExec: CometShuffleExchangeExec)
           if plan.conf.adaptiveExecutionEnabled =>
         shuffleExchangeExec
@@ -169,6 +169,57 @@ case class EliminateRedundantTransitions(session: SparkSession)
     }
   }
 
+  /**
+   * `CometIcebergWriteExec` is row-based (it emits the serialised Iceberg commit message) but
+   * consumes Arrow batches from its child over FFI, so Spark's
+   * `ApplyColumnarRulesAndInsertTransitions` inserts a columnar-to-row transition *underneath*
+   * it. Strip that transition so `doExecuteColumnar` sees the columnar child directly;
+   * `CometIcebergNativeWrite.requiresNativeChildren` already guarantees the child was
+   * Comet-native when the write was converted.
+   *
+   * The write deliberately does not tag itself as a `ColumnarToRowTransition` to suppress the
+   * insertion: Spark leaves such a node untouched, so the whole subtree below the write is never
+   * visited and the transitions the rest of that subtree needs are never inserted
+   * (https://github.com/apache/datafusion-comet/issues/5689).
+   *
+   * This runs as a separate pass *before* the main `transformUp`, not as an arm inside it,
+   * because the write's input transition has to be removed before the generic transition
+   * cancellation can consume it. `transformUp` visits children first, so the
+   * `ColumnarToRowExec(CometSparkToColumnarExec)` arm would otherwise rewrite the write's child
+   * before the write itself is visited, and that arm is destructive at this boundary:
+   *   - over a row source it drops the `CometSparkToColumnarExec` as well, leaving the write with
+   *     a row child and no Arrow producer at all;
+   *   - over a Spark-columnar source it keeps a `ColumnarToRowExec` but drops the Arrow bridge,
+   *     so the write would be handed Spark `ColumnarVector`s where the FFI adapter requires
+   *     `CometVector`s. That shape is reachable whenever `spark.comet.sparkToColumnar.enabled`
+   *     admits the write's source: `CometSparkToColumnarExec.createExec` wraps it in a
+   *     `CometScanWrapper` (a `CometNativeExec`, so `requiresNativeChildren` accepts it), and
+   *     `CometExecRule` then unwraps the placeholder, leaving the bridge directly beneath the
+   *     write.
+   */
+  private def stripIcebergWriteInputTransition(plan: SparkPlan): SparkPlan = plan.transform {
+    case w: CometIcebergWriteExec =>
+      stripColumnarToRow(w.child).map(child => w.withNewChildren(Seq(child))).getOrElse(w)
+  }
+
+  /**
+   * Unwraps a columnar-to-row transition, returning the columnar child underneath it, or `None`
+   * when `plan` is not such a transition.
+   *
+   * The plain `ColumnarToRowExec` is what Spark's insertion pass produces and is the only form
+   * observed in the suites. The two Comet variants are handled as well because this rule is part
+   * of `postColumnarTransitions` and AQE applies those rules once per materialised stage plus
+   * once for the final plan, so it can be handed a plan whose transitions an earlier pass already
+   * rewrote. Missing a variant would not be caught at planning time -- the write would fail at
+   * runtime with `requires a columnar (Comet native) child`.
+   */
+  private def stripColumnarToRow(plan: SparkPlan): Option[SparkPlan] = plan match {
+    case CometNativeColumnarToRowExec(child) => Some(child)
+    case CometColumnarToRowExec(child) => Some(child)
+    case ColumnarToRowExec(child) => Some(child)
+    case _ => None
+  }
+
   /**
    * If the given plan is a Comet ColumnarToRow transition, returns the columnar child the Python
    * UDF operator can consume directly. By the time this rule runs the earlier

spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala

@@ -27,7 +27,7 @@ import org.apache.spark.sql.catalyst.expressions.UnsafeProjection
 import org.apache.spark.sql.comet.execution.arrow.CometArrowStream
 import org.apache.spark.sql.comet.util.{Utils => CometUtils}
 import org.apache.spark.sql.connector.write.{BatchWrite, WriterCommitMessage}
-import org.apache.spark.sql.execution.{ColumnarToRowTransition, SparkPlan, UnaryExecNode}
+import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode}
 import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
 import org.apache.spark.sql.types.BinaryType
 import org.apache.spark.sql.vectorized.ColumnarBatch
@@ -68,13 +68,7 @@ case class CometIcebergWriteExec(
     @transient table: AnyRef,
     partitionSpecId: Int)
     extends CometNativeExec
-    with UnaryExecNode
-    // We consume Arrow batches (via FFI) and emit row-shaped commit messages, so we are a
-    // columnar-to-row transition. Without this trait Spark's
-    // `ApplyColumnarRulesAndInsertTransitions` wedges a `CometNativeColumnarToRowExec` between
-    // us and the Comet-native child, which would then fail `child.executeColumnar()` in
-    // `doExecuteColumnar`.
-    with ColumnarToRowTransition {
+    with UnaryExecNode {
 
   override def originalPlan: SparkPlan = child
 
@@ -86,6 +80,16 @@ case class CometIcebergWriteExec(
   // Native exec emits a single Binary column; the surrounding command framework expects rows, so
   // the outer commit exec calls executeCollect on us. supportsColumnar = false keeps Spark from
   // inserting a ColumnarToRow that would clash with our (Nil-output-like) row contract.
+  //
+  // We do consume Arrow batches from `child` over FFI, so Spark's
+  // `ApplyColumnarRulesAndInsertTransitions` inserts a columnar-to-row transition *below* us
+  // (our child is columnar-only, we are row-based). `EliminateRedundantTransitions` strips that
+  // transition back off so `doExecuteColumnar` sees the Comet-native child directly.
+  //
+  // Tagging this node as a `ColumnarToRowTransition` would also stop Spark inserting it, but at
+  // the cost of correctness: `ensureOutputsRowBased` returns such a node untouched, so nothing
+  // below the write is visited and the transitions the rest of that subtree needs are never
+  // inserted. See https://github.com/apache/datafusion-comet/issues/5689.
   override def supportsColumnar: Boolean = false
 
   override def executeCollect(): Array[InternalRow] = {

spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala

@@ -32,8 +32,9 @@ import org.apache.spark.sql.CometTestBase
 import org.apache.spark.sql.Row
 import org.apache.spark.sql.comet.{CometIcebergWriteExec, IcebergCommitExec, IcebergWriteExec}
 import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog
-import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.{ColumnarToRowTransition, SparkPlan}
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types.{DoubleType, IntegerType, StringType, StructField, StructType}
 
 import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus, isSpark41Plus}
@@ -696,6 +697,57 @@ class CometIcebergWriteActionSuite
     }
   }
 
+  // https://github.com/apache/datafusion-comet/issues/5689: the Iceberg CoW rewrite plan mixes a
+  // Spark-columnar `BatchScan` with row-based joins/filters underneath the write, so the write's
+  // subtree needs Spark to insert columnar-to-row transitions inside it. With AQE on those
+  // transitions are re-inserted when each stage materialises, which is why every other Iceberg
+  // suite in this file misses the problem; with AQE off the plan gets exactly one insertion pass
+  // and a missed transition fails the task with
+  // `ColumnarBatch cannot be cast to InternalRow`.
+  test("native acceleration: ReplaceData (CoW DELETE) with AQE disabled") {
+    assumeNativeAcceleration()
+    withIcebergCatalog { warehouseDir =>
+      createTable(
+        warehouseDir,
+        "native_cow_delete_no_aqe",
+        partitionSpec = "PARTITIONED BY (region)",
+        properties = Some("'write.delete.mode'='copy-on-write'"))
+      withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") {
+        coalesceInsert(
+          "native_cow_delete_no_aqe",
+          Seq((1, "us-east", 10.0), (2, "us-west", 20.0), (3, "eu", 30.0), (4, "us-east", 40.0)))
+      }
+      // The IN-subquery is what puts a row-based join between the columnar CoW scan and the
+      // write, so the subtree genuinely needs a transition rather than being Comet-native
+      // end to end.
+      withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+        val snapshot = withNativeEnabled {
+          captureWrite("native_cow_delete_no_aqe") {
+            spark.sql("DELETE FROM cat.db.native_cow_delete_no_aqe WHERE id IN " +
+              "(SELECT col1 FROM VALUES (2) AS t(col1))")
+          }
+        }
+        assert(
+          snapshot.snapshotDelta == 1L,
+          s"expected 1 commit via native path, got ${snapshot.snapshotDelta}")
+        val nativeExecs = snapshot.plans.flatMap { p =>
+          collectWithSubqueries(p) { case e: CometIcebergWriteExec => e }
+        }
+        assert(
+          nativeExecs.nonEmpty,
+          "expected >= 1 CometIcebergWriteExec in captured plans, got 0. Plans:\n" +
+            snapshot.plans.mkString("\n--\n"))
+        // The transition Spark inserts below the write must have been stripped again, otherwise
+        // `doExecuteColumnar` has no columnar child to pull Arrow batches from.
+        assert(
+          nativeExecs.forall(_.child.supportsColumnar),
+          "the native write must sit directly on its columnar child, got " +
+            nativeExecs.map(_.child.nodeName).mkString(", "))
+      }
+      assertRows("native_cow_delete_no_aqe", Seq(1, 3, 4))
+    }
+  }
+
   test("native acceleration: ReplaceData (CoW UPDATE)") {
     assumeNativeAcceleration()
     withIcebergCatalog { warehouseDir =>
@@ -752,6 +804,73 @@ class CometIcebergWriteActionSuite
     }
   }
 
+  // The remaining two CoW shapes from issue #5689, with AQE off so the plan gets exactly one
+  // transition-insertion pass. `captureWrite` runs `assertColumnarContract` on every captured
+  // plan, so these pin the same invariant as the DELETE case over the UPDATE and MERGE rewrite
+  // shapes, which put different operators between the columnar CoW scan and the write.
+  test("native acceleration: ReplaceData (CoW UPDATE) with AQE disabled") {
+    assumeNativeAcceleration()
+    withIcebergCatalog { warehouseDir =>
+      createTable(
+        warehouseDir,
+        "native_cow_update_no_aqe",
+        partitionSpec = "PARTITIONED BY (region)",
+        properties = Some("'write.update.mode'='copy-on-write'"))
+      withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") {
+        coalesceInsert(
+          "native_cow_update_no_aqe",
+          Seq((1, "us-east", 10.0), (2, "us-west", 20.0), (3, "eu", 30.0)))
+      }
+      // As in the DELETE case, the IN-subquery is what puts a row-based join between the
+      // Spark-columnar CoW scan and the write.
+      withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+        assertNativeWriteEngages("native_cow_update_no_aqe", Seq(1, 2, 3)) {
+          spark.sql(
+            "UPDATE cat.db.native_cow_update_no_aqe SET amount = amount * 2 " +
+              "WHERE id IN (SELECT col1 FROM VALUES (2) AS t(col1))")
+        }
+      }
+      val r = spark
+        .sql("SELECT amount FROM cat.db.native_cow_update_no_aqe WHERE id = 2")
+        .collect()
+      assert(r.length == 1 && r(0).getDouble(0) == 40.0, s"got ${r.toSeq}")
+    }
+  }
+
+  test("native acceleration: ReplaceData (CoW MERGE) with AQE disabled") {
+    // Unlike the unpartitioned MERGE above, this one *does* engage natively: partitioning the
+    // table puts a `CometColumnarExchange` (REBALANCE_PARTITIONS_BY_COL) between the JVM
+    // `MergeRowsExec` and the write, so the write's own child is Comet-native and
+    // `requiresNativeChildren` is satisfied even though `MergeRowsExec` itself is not.
+    //
+    // That makes this the strongest of the three CoW shapes for issue #5689: the subtree below
+    // the write mixes a Spark-columnar `BatchScan`, a row-based `MergeRowsExec` and Comet
+    // operators, so it needs transitions inserted in three different places.
+    assumeNativeAcceleration()
+    withIcebergCatalog { warehouseDir =>
+      createTable(
+        warehouseDir,
+        "native_cow_merge_no_aqe",
+        partitionSpec = "PARTITIONED BY (region)",
+        properties = Some("'write.merge.mode'='copy-on-write'"))
+      withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "false") {
+        coalesceInsert("native_cow_merge_no_aqe", Seq((1, "us-east", 10.0), (2, "us-west", 20.0)))
+      }
+      withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+        assertNativeWriteEngages("native_cow_merge_no_aqe", Seq(1, 2, 3)) {
+          spark.sql("""
+            |MERGE INTO cat.db.native_cow_merge_no_aqe t
+            |USING (SELECT 2 AS id, 'us-west' AS region, 200.0 AS amount UNION ALL
+            |       SELECT 3 AS id, 'eu' AS region, 30.0 AS amount) s
+            |ON t.id = s.id
+            |WHEN MATCHED THEN UPDATE SET t.amount = s.amount
+            |WHEN NOT MATCHED THEN INSERT (id, region, amount) VALUES (s.id, s.region, s.amount)
+            |""".stripMargin)
+        }
+      }
+    }
+  }
+
   test("native acceleration: complex types (struct, array, map) round-trip with field IDs") {
     assumeNativeAcceleration()
     withIcebergCatalog { _ =>
@@ -1633,6 +1752,10 @@ class CometIcebergWriteActionSuite
   private def captureWrite(tableName: String)(action: => Unit): WriteSnapshot = {
     val before = countSnapshots(tableName)
     val plans = capturePlans(spark)(action)
+    // Every write in this suite gets the columnar contract checked, not just the ones written to
+    // exercise it: a Comet rewrite that hides part of a subtree from Spark's transition-insertion
+    // pass is a whole class of bug (issue #5689) and the check is free once the plans are here.
+    plans.foreach(assertColumnarContract)
     WriteSnapshot(countSnapshots(tableName) - before, plans)
   }
 
@@ -1688,6 +1811,31 @@ class CometIcebergWriteActionSuite
     assume(icebergAvailable, "Iceberg not available in classpath")
   }
 
+  /**
+   * Every row-consuming operator must receive row-based input. Spark guarantees that by inserting
+   * `ColumnarToRow` transitions in `ApplyColumnarRulesAndInsertTransitions`; an operator that
+   * Comet rewrote in a way that skips the insertion pass shows up here as a row-based node with a
+   * columnar-only child, and would fail at runtime with a `ColumnarBatch cannot be cast to
+   * InternalRow` `ClassCastException` rather than at planning time.
+   *
+   * `CometIcebergWriteExec` is the one legitimate exception: it is row-based on the outside but
+   * pulls Arrow batches from its Comet-native child over FFI (see the class docstring), so the
+   * transition below it is deliberately stripped again by `EliminateRedundantTransitions`.
+   */
+  private def assertColumnarContract(plan: SparkPlan): Unit = {
+    val violations = collectWithSubqueries(plan) {
+      case p
+          if !p.supportsColumnar && !p.isInstanceOf[ColumnarToRowTransition] &&
+            !p.isInstanceOf[CometIcebergWriteExec] &&
+            p.children.exists(c => c.supportsColumnar && !c.supportsRowBased) =>
+        p
+    }
+    assert(
+      violations.isEmpty,
+      "row-based operators consuming columnar-only children (missing ColumnarToRow): " +
+        s"${violations.map(_.nodeName).mkString(", ")}. Plan:\n$plan")
+  }
+
   /**
    * Flip [[CometConf.COMET_ICEBERG_NATIVE_WRITE_ENABLED]] for the duration of `action`.
    *