(datafusion-comet) 01/01: fix: write cached batches to the schema width, not the batch width (#6090)

[gi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6090-57ef3275d3d4e6e3c4d5f3de27b9b71daf736164
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git

commit 30c24c078d863e6aae9e2bef0919d97d6105ad9c
Author: Andy Grove <ag...@apache.org>
AuthorDate: Mon Sep 21 22:59:20 2026 +0000

    fix: write cached batches to the schema width, not the batch width (#6090)
    
    ArrowWriter.writeColumns drove its loop from the input ColumnarBatch's width while indexing the writer's fields, which come from the schema the batch is written under. That assumed every producer hands over a batch exactly as wide as the schema.
    
    Iceberg's vectorized reader does not. BatchDeleteFilter.filterBatch reads with the delete filter's requiredSchema, which carries _pos after the projected columns when a data file has position deletes, and trims the extras back only when the file also has equality deletes. A merge-on-read UPDATE writes position deletes and no equality deletes, so the extra column survives into the batch, and caching such a relation failed with ArrayIndexOutOfBoundsException inside the write loop.
    
    Drive the loop from the writer's fields instead, which writes exactly the columns the schema describes: the extras are trailing, the same prefix Iceberg keeps when it does trim. A batch narrower than the schema is a genuine contract violation and is now refused with a message naming both widths.
    
    Closes #6087.
---
 .../sql/comet/execution/arrow/ArrowWriters.scala   | 14 +++++-
 .../execution/arrow/CometArrowStreamSuite.scala    | 56 ++++++++++++++++++++++
 2 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala
index e2632f563e..bb883f8d2a 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala
@@ -165,9 +165,21 @@ class ArrowWriter(val root: VectorSchemaRoot, fields: Array[ArrowFieldWriter]) {
     count = input.numElements()
   }
 
+  // Driven by the writer's fields rather than by the input's width, because a producer may hand
+  // over a batch wider than the schema it is written under. Iceberg's vectorized reader does:
+  // it reads with the schema its delete filter required, which carries `_pos` after the projected
+  // columns when a data file has position deletes, and trims the extras back only when the file
+  // also has equality deletes. Those extras are trailing -- `removeExtraColumns` keeps the leading
+  // `expectedSchema` prefix when it does trim -- so writing the first `fields.length` columns
+  // writes exactly the columns the schema describes. A batch with fewer columns than the schema
+  // has no such reading and is refused rather than written short.
   def writeColumns(input: ColumnarBatch, startRow: Int, numRows: Int): Unit = {
+    require(
+      input.numCols() >= fields.length,
+      s"Cannot write ${fields.length} columns from a batch of ${input.numCols()} " +
+        (if (input.numCols() == 1) "column" else "columns"))
     var columnIndex = 0
-    while (columnIndex < input.numCols()) {
+    while (columnIndex < fields.length) {
       fields(columnIndex).writeColumnSlice(input.column(columnIndex), startRow, numRows)
       columnIndex += 1
     }
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala
index cfdefedd6e..3ad59e96ef 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowStreamSuite.scala
@@ -872,4 +872,60 @@ class CometArrowStreamSuite extends AnyFunSuite with Matchers {
       allocator.close()
     }
   }
+
+  test("columnar writes are driven by the schema, not by the input batch width") {
+    val allocator = new RootAllocator(Long.MaxValue)
+    val numRows = 3
+    val schema = StructType(Seq(StructField("name", StringType)))
+    val arrowSchema = Utils.toArrowSchema(schema, "UTC")
+    // A connector may hand over a batch wider than the schema it is read under. Iceberg's
+    // vectorized reader is the case that found this: it reads with the schema its delete filter
+    // required, which carries `_pos` after the projected columns when a data file has position
+    // deletes, and only trims the extras back when the file also has equality deletes. The
+    // trailing columns are the extras, so the schema's fields line up with the leading ones.
+    val names = new OnHeapColumnVector(numRows, StringType)
+    val positions = new OnHeapColumnVector(numRows, LongType)
+    val input = new ColumnarBatch(Array[ColumnVector](names, positions), numRows)
+    try {
+      (0 until numRows).foreach { i =>
+        names.putByteArray(i, s"n$i".getBytes(StandardCharsets.UTF_8))
+        positions.putLong(i, i.toLong)
+      }
+      val batch = CometArrowConverters.columnarBatchToArrowBatch(input, arrowSchema, allocator)
+      try {
+        batch.numCols() shouldBe 1
+        batch.numRows() shouldBe numRows
+        (0 until numRows).foreach { i =>
+          batch.column(0).getUTF8String(i).toString shouldBe s"n$i"
+        }
+      } finally batch.close()
+    } finally {
+      input.close()
+      allocator.close()
+    }
+  }
+
+  test("a batch narrower than the schema is refused rather than written short") {
+    val allocator = new RootAllocator(Long.MaxValue)
+    val numRows = 2
+    val schema =
+      StructType(Seq(StructField("name", StringType), StructField("id", LongType)))
+    val arrowSchema = Utils.toArrowSchema(schema, "UTC")
+    val names = new OnHeapColumnVector(numRows, StringType)
+    val input = new ColumnarBatch(Array[ColumnVector](names), numRows)
+    try {
+      (0 until numRows).foreach { i =>
+        names.putByteArray(i, s"n$i".getBytes(StandardCharsets.UTF_8))
+      }
+      val failure = intercept[IllegalArgumentException] {
+        CometArrowConverters.columnarBatchToArrowBatch(input, arrowSchema, allocator)
+      }
+      failure.getMessage should include("1 column")
+      failure.getMessage should include("2")
+      allocator.getAllocatedMemory shouldBe 0L
+    } finally {
+      input.close()
+      allocator.close()
+    }
+  }
 }


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org