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 <[email protected]> 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: [email protected] For additional commands, e-mail: [email protected]
