From a0915fe2993bf6d806cee9cdc0b4316cacc01aec Mon Sep 17 00:00:00 2001 From: zhangfengcdt Date: Fri, 6 Nov 2026 06:05:52 -0800 Subject: [PATCH] fix: have the native Iceberg write read a converted child's Arrow stream A native Iceberg write over a Spark-to-Arrow conversion wrapped the conversion's executeColumnar batches in a reader that closes each batch once native has it. The conversion reuses its vectors for the next batch, and closing a struct vector drops its children, so the second batch of a partition failed to import with "no more field nodes". The write now reads such a child's Arrow stream, as native shuffle does since #6689. --- .../sql/comet/CometIcebergWriteExec.scala | 54 +++++++++++-------- .../comet/CometIcebergWriteActionSuite.scala | 40 ++++++++++++++ 2 files changed, 72 insertions(+), 22 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala index 98ba9babe77..66c8c9f0895 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometIcebergWriteExec.scala +++ b/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, diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala index 6f752267a63..e0a2de8d2c4 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala +++ b/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) 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 { _ =>