fix: have the native Iceberg write read a converted child's Arrow stream

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

@@ -21,12 +21,13 @@ package org.apache.spark.sql.comet
 
 import java.util.concurrent.atomic.AtomicReference
 
+import org.apache.arrow.c.ArrowArrayStream
 import org.apache.spark.TaskContext
 import org.apache.spark.rdd.RDD
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.Attribute
 import org.apache.spark.sql.catalyst.expressions.UnsafeProjection
-import org.apache.spark.sql.comet.execution.arrow.CometArrowStream
+import org.apache.spark.sql.comet.execution.arrow.{CometArrowStream, CometNativeArrowSource}
 import org.apache.spark.sql.comet.util.{Utils => CometUtils}
 import org.apache.spark.sql.connector.write.{BatchWrite, WriterCommitMessage}
 import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode}
@@ -230,29 +231,41 @@ case class CometIcebergWriteExec(
 
   override def doExecuteColumnar(): RDD[ColumnarBatch] = {
     val schedulerProbe = IcebergSchedulerTestProbe.current
-    val inputRDD = if (child.supportsColumnar) {
-      child.executeColumnar()
-    } else {
-      throw new UnsupportedOperationException(
-        "CometIcebergWriteExec requires a columnar (Comet native) child; got " +
-          child.getClass.getName)
+    // One Arrow stream per partition, which native consumes through the C Stream Interface.
+    val streamRDD: RDD[ArrowArrayStream] = child match {
+      // Native reads the source's Arrow stream, as it does for a native operator's input. The
+      // source reuses its vectors for the next batch, and wrapping its `executeColumnar` batches
+      // below closes each batch once native has it; closing a struct vector drops its children,
+      // so the next batch would lose them (#7239). The scheduler probe counts rows as batches.
+      case source: CometNativeArrowSource if schedulerProbe.isEmpty =>
+        source.doExecuteAsArrowStream()
+      case _ if child.supportsColumnar =>
+        val inputRDD = child.executeColumnar()
+        val childRDD = schedulerProbe
+          .map(probe =>
+            new IcebergSchedulerProbeRDD(
+              inputRDD,
+              probe,
+              nativeOp.getIcebergWrite.getCommon.getDataLocation,
+              nativeOp.getIcebergWrite.getCommon.getOperationId))
+          .getOrElse(inputRDD)
+        CometArrowStream.wrapColumnarBatchRDD(
+          childRDD,
+          CometUtils.fromAttributes(child.output),
+          CometArrowStream.NATIVE_TIMEZONE,
+          "CometIcebergWriteExec")
+      case _ =>
+        throw new UnsupportedOperationException(
+          "CometIcebergWriteExec requires a columnar (Comet native) child; got " +
+            child.getClass.getName)
     }
-
-    val childRDD = schedulerProbe
-      .map(probe =>
-        new IcebergSchedulerProbeRDD(
-          inputRDD,
-          probe,
-          nativeOp.getIcebergWrite.getCommon.getDataLocation,
-          nativeOp.getIcebergWrite.getCommon.getOperationId))
-      .getOrElse(inputRDD)
-    val numPartitions = childRDD.getNumPartitions
+    val numPartitions = streamRDD.getNumPartitions
     // Native side emits the per-task Avro `DataFile` blob and the locations it wrote, one Binary
     // column each (see `build_output_schema` in `iceberg_write.rs`).
     val numOutputCols = 2
     val capturedNativeOp = nativeOp
 
-    childRDD.mapPartitionsInternal { iter =>
+    streamRDD.mapPartitionsInternal { iter =>
       val partitionId = TaskContext.getPartitionId()
       val taskAttemptId = TaskContext.get().taskAttemptId()
 
@@ -280,10 +293,7 @@ case class CometIcebergWriteExec(
 
       val execIterator = new CometExecIterator(
         CometExec.newIterId,
-        CometArrowStream.inputObjects(
-          iter,
-          CometUtils.fromAttributes(child.output),
-          "CometIcebergWriteExec"),
+        Array[Object](iter.next()),
         numOutputCols,
         planBytes,
         nativeMetrics,

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

@@ -47,6 +47,7 @@ import org.apache.spark.sql.Row
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.Attribute
 import org.apache.spark.sql.comet.{CometIcebergWriteExec, CometMergeRowsExec, IcebergCommitExec, IcebergSchedulerTestProbe, IcebergWriteExec}
+import org.apache.spark.sql.comet.execution.arrow.CometNativeArrowSource
 import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog
 import org.apache.spark.sql.connector.write.{BatchWrite, DataWriterFactory, PhysicalWriteInfo, Write, WriterCommitMessage}
 import org.apache.spark.sql.execution.{ColumnarToRowTransition, LeafExecNode, SparkPlan}
@@ -412,6 +413,45 @@ class CometIcebergWriteActionSuite
     }
   }
 
+  // https://github.com/apache/datafusion-comet/issues/7239. A conversion of Spark rows writes
+  // every batch into the same vectors, and closing a struct vector drops its children. The write
+  // reads the conversion's Arrow stream, as native shuffle does, rather than batches it would
+  // close once native had them, which left the second batch of a partition without them.
+  test("native write reads a struct column that a conversion writes over several batches") {
+    assume(icebergAvailable, "Iceberg not available in classpath")
+    withIcebergCatalog { _ =>
+      spark.sql(s"CREATE TABLE $catalog.$ns.struct_rdd (id INT, s STRUCT<a: INT>) USING iceberg")
+      val schema = new StructType()
+        .add("id", IntegerType)
+        .add("s", new StructType().add("a", IntegerType))
+      val rows = spark.sparkContext.parallelize((1 to 10).map(i => Row(i, Row(i))), 1)
+      val snapshot = captureWrite("struct_rdd") {
+        withNativeEnabled {
+          withSQLConf(
+            CometConf.COMET_BATCH_SIZE.key -> "2",
+            CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> "true") {
+            spark.createDataFrame(rows, schema).writeTo(s"$catalog.$ns.struct_rdd").append()
+          }
+        }
+      }
+      assert(snapshot.snapshotDelta == 1L, s"expected 1 commit, got ${snapshot.snapshotDelta}")
+      val nativeWrites = snapshot.plans.flatMap { plan =>
+        collectWithSubqueries(plan) { case e: CometIcebergWriteExec => e }
+      }
+      assert(
+        nativeWrites.exists(_.child.isInstanceOf[CometNativeArrowSource]),
+        "expected a CometIcebergWriteExec over the conversion. Plans:\n" +
+          snapshot.plans.mkString("\n--\n"))
+      assertRows("struct_rdd", expectedIds = 1 to 10)
+      val fields = spark
+        .sql(s"SELECT s.a FROM $catalog.$ns.struct_rdd ORDER BY id")
+        .collect()
+        .map(_.getInt(0))
+        .toSeq
+      assert(fields == (1 to 10), s"expected ${1 to 10}, got $fields")
+    }
+  }
+
   test("AQE re-plan of the writer subtree writes and commits exactly once") {
     assume(icebergAvailable, "Iceberg not available in classpath")
     withIcebergCatalog { _ =>