From de1f04a66764bbd5a627bd6051d187834a85eecc Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 6 Sep 2026 00:23:43 +0800 Subject: [PATCH 1/2] refactor: make Spark-to-Arrow batch row counts explicit --- .../comet/execution/arrow/ArrowWriters.scala | 12 +- .../arrow/CometArrowConverters.scala | 4 +- .../execution/arrow/RowArrowReader.scala | 2 +- .../arrow/SparkColumnarArrowReader.scala | 2 +- .../benchmark/CometArrowWriterBenchmark.scala | 2 +- .../arrow/CometArrowStreamSuite.scala | 145 +++++++++++++++++- 6 files changed, 147 insertions(+), 20 deletions(-) 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 e2632f563e3..67882bdb357 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 @@ -144,25 +144,20 @@ class ArrowWriter(val root: VectorSchemaRoot, fields: Array[ArrowFieldWriter]) { def schema: StructType = Utils.fromArrowSchema(root.getSchema()) - private var count: Int = 0 - def write(row: InternalRow): Unit = { var i = 0 while (i < fields.length) { fields(i).writeUnsafe(row, i) i += 1 } - count += 1 } def writeCol(input: ColumnarArray, columnIndex: Int): Unit = { fields(columnIndex).writeCol(input) - count = input.numElements() } def writeColNoNull(input: ColumnarArray, columnIndex: Int): Unit = { fields(columnIndex).writeColNoNull(input) - count = input.numElements() } def writeColumns(input: ColumnarBatch, startRow: Int, numRows: Int): Unit = { @@ -171,17 +166,16 @@ class ArrowWriter(val root: VectorSchemaRoot, fields: Array[ArrowFieldWriter]) { fields(columnIndex).writeColumnSlice(input.column(columnIndex), startRow, numRows) columnIndex += 1 } - count = numRows } - def finish(): Unit = { - root.setRowCount(count) + /** Finish with the caller's logical batch size, including batches with no columns. */ + def finish(rowCount: Int): Unit = { + root.setRowCount(rowCount) fields.foreach(_.finish()) } def reset(): Unit = { root.setRowCount(0) - count = 0 fields.foreach(_.reset()) } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala index e68eee6b79b..6ccdf1370e5 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala @@ -76,7 +76,7 @@ object CometArrowConverters extends Logging { writer.write(rowIter.next()) rowCount += 1 } - writer.finish() + writer.finish(rowCount) NativeUtil.rootAsBatch(root) } } @@ -102,7 +102,7 @@ object CometArrowConverters extends Logging { closingRootOnFailure(root) { val writer = ArrowWriter.create(root, numRows) writer.writeColumns(batch, 0, numRows) - writer.finish() + writer.finish(numRows) NativeUtil.rootAsBatch(root) } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/RowArrowReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/RowArrowReader.scala index e4f484aa1e2..096082f8117 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/RowArrowReader.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/RowArrowReader.scala @@ -63,7 +63,7 @@ private[comet] class RowArrowReader( writer.write(rowIter.next()) rowCount += 1 } - writer.finish() + writer.finish(rowCount) onConversionNs(System.nanoTime() - startNs) true } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/SparkColumnarArrowReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/SparkColumnarArrowReader.scala index 90c8be2a613..f63bbf469b0 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/SparkColumnarArrowReader.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/SparkColumnarArrowReader.scala @@ -80,7 +80,7 @@ private[comet] class SparkColumnarArrowReader( writer.writeColumns(current, rowsConsumedInCurrent, rowsToProduce) rowsConsumedInCurrent += rowsToProduce - writer.finish() + writer.finish(rowsToProduce) onConversionNs(System.nanoTime() - startNs) true } diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometArrowWriterBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometArrowWriterBenchmark.scala index 11033703afc..f4bd3b5f323 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometArrowWriterBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometArrowWriterBenchmark.scala @@ -86,7 +86,7 @@ object CometArrowWriterBenchmark extends BenchmarkBase { col += 1 } } - writer.finish() + writer.finish(input.numRows()) } try { 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 1b820fdddf5..71878544437 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 @@ -31,12 +31,12 @@ import org.apache.arrow.vector.{BaseFixedWidthVector, BaseValueVector, BigIntVec import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType, Schema} import org.apache.spark.sql.catalyst.expressions.{GenericInternalRow, SpecializedGetters} import org.apache.spark.sql.comet.util.Utils -import org.apache.spark.sql.execution.vectorized.{Dictionary, OffHeapColumnVector, OnHeapColumnVector} -import org.apache.spark.sql.types.{BooleanType, ByteType, CalendarIntervalType, DataType, DateType, DayTimeIntervalType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, LongType, ShortType, StructField, StructType, TimestampNTZType, TimestampType, YearMonthIntervalType} +import org.apache.spark.sql.execution.vectorized.{ConstantColumnVector, Dictionary, OffHeapColumnVector, OnHeapColumnVector} +import org.apache.spark.sql.types.{ArrayType, BooleanType, ByteType, CalendarIntervalType, DataType, DateType, DayTimeIntervalType, Decimal, DecimalType, DoubleType, FloatType, IntegerType, LongType, ShortType, StringType, StructField, StructType, TimestampNTZType, TimestampType, YearMonthIntervalType} import org.apache.spark.sql.vectorized.{ColumnarArray, ColumnarBatch, ColumnVector} import org.apache.spark.unsafe.types.CalendarInterval -import org.apache.comet.vector.{CometPlainVector, CometVector} +import org.apache.comet.vector.{CometPlainVector, CometVector, NativeUtil} /** * Direct tests for [[CometArrowStream.reconcileStreamSchema]]. The end-to-end regression that @@ -68,7 +68,7 @@ class CometArrowStreamSuite extends AnyFunSuite with Matchers { val writer = ArrowWriter.create(root, 2) writer.write(new GenericInternalRow(Array[Any](expected))) writer.write(new GenericInternalRow(Array[Any](null))) - writer.finish() + writer.finish(2) val arrow = root.getVector(0).asInstanceOf[IntervalMonthDayNanoVector] IntervalMonthDayNanoVector.getMonths(arrow.getDataBuffer, 0) shouldBe expected.months @@ -297,9 +297,16 @@ class CometArrowStreamSuite extends AnyFunSuite with Matchers { try { val writer = ArrowWriter.create(root, numRows) writer.writeColumns(input, 0, numRows) - writer.finish() + writer.finish(numRows) root.getRowCount shouldBe numRows + writer.reset() + root.getRowCount shouldBe 0 + // The logical count comes from finish, even when there are no values to encode. + writer.finish(3) + root.getRowCount shouldBe 3 + writer.finish(0) + root.getRowCount shouldBe 0 } finally { input.close() root.close() @@ -307,6 +314,132 @@ class CometArrowStreamSuite extends AnyFunSuite with Matchers { } } + test("row and columnar conversion preserve zero-column batch sizes") { + val allocator = new RootAllocator(Long.MaxValue) + val schema = StructType(Seq.empty[StructField]) + val arrowSchema = Utils.toArrowSchema(schema, "UTC") + val input = new ColumnarBatch(Array.empty[ColumnVector], 5) + def rows = Iterator.fill(5)(new GenericInternalRow(0)) + val rowReader = new RowArrowReader(allocator, arrowSchema, rows, 2) + val columnReader = + new SparkColumnarArrowReader(allocator, arrowSchema, Iterator.single(input), 2) + try { + Seq(rowReader, columnReader).foreach { reader => + Seq(2, 2, 1).foreach { expected => + reader.loadNextBatch() shouldBe true + reader.getVectorSchemaRoot.getRowCount shouldBe expected + } + reader.loadNextBatch() shouldBe false + } + val batches = CometArrowConverters.rowToArrowBatchIter(rows, schema, 2, "UTC", allocator) + batches.map { batch => + try batch.numRows() + finally batch.close() + }.toList shouldBe List(2, 2, 1) + Seq(5, 0).foreach { numRows => + input.setNumRows(numRows) + val batch = CometArrowConverters.columnarBatchToArrowBatch(input, arrowSchema, allocator) + try batch.numRows() shouldBe numRows + finally batch.close() + } + } finally { + rowReader.close() + columnReader.close() + input.close() + allocator.close() + } + } + + test("nested and foreign-vector fallback preserves slices and independently owned batches") { + val allocator = new RootAllocator(Long.MaxValue) + val arrayType = ArrayType(IntegerType, containsNull = true) + val schema = StructType( + Seq( + StructField("array", arrayType), + StructField("string", StringType), + StructField("constant", LongType, nullable = false))) + val arrowSchema = Utils.toArrowSchema(schema, "UTC") + val arrays = new OnHeapColumnVector(7, arrayType) + val strings = new OnHeapColumnVector(7, StringType) + val constant = new ConstantColumnVector(7, LongType) + constant.setLong(42L) + val input = new ColumnarBatch(Array[ColumnVector](arrays, strings, constant), 7) + val reader = + new SparkColumnarArrowReader(allocator, arrowSchema, Iterator.single(input), 3) + try { + (0 until 7).foreach { i => + arrays.putArray(i, i, 1) + if (i % 2 == 0) arrays.getChild(0).putNull(i) + else arrays.getChild(0).putInt(i, i) + if (i % 2 == 0) strings.putNull(i) + else + strings.putByteArray(i, s"value-$i".getBytes(java.nio.charset.StandardCharsets.UTF_8)) + } + def check(batch: ColumnarBatch, startRow: Int): Unit = { + (0 until batch.numRows()).foreach { i => + val sourceRow = startRow + i + val array = batch.column(0).getArray(i) + array.numElements() shouldBe 1 + array.isNullAt(0) shouldBe (sourceRow % 2 == 0) + batch.column(1).isNullAt(i) shouldBe (sourceRow % 2 == 0) + if (sourceRow % 2 != 0) { + array.getInt(0) shouldBe sourceRow + batch.column(1).getUTF8String(i).toString shouldBe s"value-$sourceRow" + } + batch.column(2).getLong(i) shouldBe 42L + } + } + val first = CometArrowConverters.columnarBatchToArrowBatch(input, arrowSchema, allocator) + try { + val second = CometArrowConverters.columnarBatchToArrowBatch(input, arrowSchema, allocator) + try check(second, 0) + finally second.close() + val root = reader.getVectorSchemaRoot + Seq(0, 3, 6).foreach { startRow => + reader.loadNextBatch() shouldBe true + reader.getVectorSchemaRoot should be theSameInstanceAs root + root.getRowCount shouldBe math.min(3, 7 - startRow) + check(NativeUtil.rootAsBatch(root), startRow) + } + reader.loadNextBatch() shouldBe false + check(first, 0) + // Closing converted batches must not close the producer's input. + strings.getUTF8String(1).toString shouldBe "value-1" + } finally first.close() + } finally { + reader.close() + input.close() + allocator.close() + } + } + + test("fresh row and columnar conversion release allocations when encoding throws") { + val allocator = new RootAllocator(Long.MaxValue) + val schema = StructType(Seq(StructField("int", IntegerType))) + val arrowSchema = Utils.toArrowSchema(schema, "UTC") + val failure = new IllegalStateException("conversion failed") + val column = new ConstantColumnVector(2, IntegerType) { + override def getInt(rowId: Int): Int = throw failure + } + val input = new ColumnarBatch(Array[ColumnVector](column), 2) + try { + intercept[IllegalStateException] { + CometArrowConverters.columnarBatchToArrowBatch(input, arrowSchema, allocator) + } should be theSameInstanceAs failure + allocator.getAllocatedMemory shouldBe 0L + val rows = Iterator(new GenericInternalRow(Array[Any](1)) { + override def getInt(ordinal: Int): Int = throw failure + }) + intercept[IllegalStateException] { + CometArrowConverters.rowToArrowBatchIter(rows, schema, 2, "UTC", allocator).next() + } should be theSameInstanceAs failure + allocator.getAllocatedMemory shouldBe 0L + } finally { + input.close() + allocator.close() + } + } + test("Spark columnar reader preserves split input slice offsets") { val allocator = new RootAllocator(Long.MaxValue) val numRows = 12 @@ -390,7 +523,7 @@ class CometArrowStreamSuite extends AnyFunSuite with Matchers { vector.getValueCapacity should be < numRows writer.writeColNoNull(new ColumnarArray(input, 0, numRows), 0) - writer.finish() + writer.finish(numRows) vector.getValueCapacity should be >= numRows vector.get(numRows - 1) shouldBe numRows - 1 From ab6af9454fe0897fc7880e2be50888548146b5e4 Mon Sep 17 00:00:00 2001 From: peterxcli Date: Sun, 6 Sep 2026 23:36:30 +0800 Subject: [PATCH 2/2] test: retain conversion coverage with the existing writer API --- .../sql/comet/execution/arrow/ArrowWriters.scala | 12 +++++++++--- .../execution/arrow/CometArrowConverters.scala | 4 ++-- .../sql/comet/execution/arrow/RowArrowReader.scala | 2 +- .../execution/arrow/SparkColumnarArrowReader.scala | 2 +- .../sql/benchmark/CometArrowWriterBenchmark.scala | 2 +- .../execution/arrow/CometArrowStreamSuite.scala | 13 +++---------- 6 files changed, 17 insertions(+), 18 deletions(-) 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 67882bdb357..e2632f563e3 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 @@ -144,20 +144,25 @@ class ArrowWriter(val root: VectorSchemaRoot, fields: Array[ArrowFieldWriter]) { def schema: StructType = Utils.fromArrowSchema(root.getSchema()) + private var count: Int = 0 + def write(row: InternalRow): Unit = { var i = 0 while (i < fields.length) { fields(i).writeUnsafe(row, i) i += 1 } + count += 1 } def writeCol(input: ColumnarArray, columnIndex: Int): Unit = { fields(columnIndex).writeCol(input) + count = input.numElements() } def writeColNoNull(input: ColumnarArray, columnIndex: Int): Unit = { fields(columnIndex).writeColNoNull(input) + count = input.numElements() } def writeColumns(input: ColumnarBatch, startRow: Int, numRows: Int): Unit = { @@ -166,16 +171,17 @@ class ArrowWriter(val root: VectorSchemaRoot, fields: Array[ArrowFieldWriter]) { fields(columnIndex).writeColumnSlice(input.column(columnIndex), startRow, numRows) columnIndex += 1 } + count = numRows } - /** Finish with the caller's logical batch size, including batches with no columns. */ - def finish(rowCount: Int): Unit = { - root.setRowCount(rowCount) + def finish(): Unit = { + root.setRowCount(count) fields.foreach(_.finish()) } def reset(): Unit = { root.setRowCount(0) + count = 0 fields.foreach(_.reset()) } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala index 6ccdf1370e5..e68eee6b79b 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CometArrowConverters.scala @@ -76,7 +76,7 @@ object CometArrowConverters extends Logging { writer.write(rowIter.next()) rowCount += 1 } - writer.finish(rowCount) + writer.finish() NativeUtil.rootAsBatch(root) } } @@ -102,7 +102,7 @@ object CometArrowConverters extends Logging { closingRootOnFailure(root) { val writer = ArrowWriter.create(root, numRows) writer.writeColumns(batch, 0, numRows) - writer.finish(numRows) + writer.finish() NativeUtil.rootAsBatch(root) } } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/RowArrowReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/RowArrowReader.scala index 096082f8117..e4f484aa1e2 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/RowArrowReader.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/RowArrowReader.scala @@ -63,7 +63,7 @@ private[comet] class RowArrowReader( writer.write(rowIter.next()) rowCount += 1 } - writer.finish(rowCount) + writer.finish() onConversionNs(System.nanoTime() - startNs) true } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/SparkColumnarArrowReader.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/SparkColumnarArrowReader.scala index f63bbf469b0..90c8be2a613 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/SparkColumnarArrowReader.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/SparkColumnarArrowReader.scala @@ -80,7 +80,7 @@ private[comet] class SparkColumnarArrowReader( writer.writeColumns(current, rowsConsumedInCurrent, rowsToProduce) rowsConsumedInCurrent += rowsToProduce - writer.finish(rowsToProduce) + writer.finish() onConversionNs(System.nanoTime() - startNs) true } diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometArrowWriterBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometArrowWriterBenchmark.scala index f4bd3b5f323..11033703afc 100644 --- a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometArrowWriterBenchmark.scala +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometArrowWriterBenchmark.scala @@ -86,7 +86,7 @@ object CometArrowWriterBenchmark extends BenchmarkBase { col += 1 } } - writer.finish(input.numRows()) + writer.finish() } try { 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 71878544437..d95a4c718e5 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 @@ -68,7 +68,7 @@ class CometArrowStreamSuite extends AnyFunSuite with Matchers { val writer = ArrowWriter.create(root, 2) writer.write(new GenericInternalRow(Array[Any](expected))) writer.write(new GenericInternalRow(Array[Any](null))) - writer.finish(2) + writer.finish() val arrow = root.getVector(0).asInstanceOf[IntervalMonthDayNanoVector] IntervalMonthDayNanoVector.getMonths(arrow.getDataBuffer, 0) shouldBe expected.months @@ -297,16 +297,9 @@ class CometArrowStreamSuite extends AnyFunSuite with Matchers { try { val writer = ArrowWriter.create(root, numRows) writer.writeColumns(input, 0, numRows) - writer.finish(numRows) + writer.finish() root.getRowCount shouldBe numRows - writer.reset() - root.getRowCount shouldBe 0 - // The logical count comes from finish, even when there are no values to encode. - writer.finish(3) - root.getRowCount shouldBe 3 - writer.finish(0) - root.getRowCount shouldBe 0 } finally { input.close() root.close() @@ -523,7 +516,7 @@ class CometArrowStreamSuite extends AnyFunSuite with Matchers { vector.getValueCapacity should be < numRows writer.writeColNoNull(new ColumnarArray(input, 0, numRows), 0) - writer.finish(numRows) + writer.finish() vector.getValueCapacity should be >= numRows vector.get(numRows - 1) shouldBe numRows - 1