From ad03d40b3b8f9b24ca60cd9b8a3e9eb3d083f823 Mon Sep 17 00:00:00 2001 From: comphead Date: Mon, 12 Oct 2026 09:46:14 -0700 Subject: [PATCH 1/4] fix: keep the values of sliced batches when coalescing a broadcast A slice of a native string, binary, list or map array reaches the JVM with offsets that start past zero, with the rows before the slice still at the front of its data. Arrow Java writes those bytes into the IPC stream, and the VectorSchemaRootAppender that coalesces broadcast batches assumes offsets start at zero, so it prepended them to the first value of each slice. serializeBatches now writes such a batch through a slice of its root, which rebases the offsets to zero and leaves the earlier bytes out. --- .../apache/spark/sql/comet/util/Utils.scala | 28 ++++++- .../apache/comet/exec/CometJoinSuite.scala | 30 ++++++++ .../spark/sql/comet/util/UtilsSuite.scala | 75 +++++++++++++++++++ 3 files changed, 131 insertions(+), 2 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index c884344ef15..1fc68aac6dc 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -27,7 +27,7 @@ import scala.jdk.CollectionConverters._ import org.apache.arrow.c.CDataDictionaryProvider import org.apache.arrow.vector._ -import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector} +import org.apache.arrow.vector.complex.{BaseRepeatedValueVector, ListVector, MapVector, StructVector} import org.apache.arrow.vector.dictionary.{Dictionary, DictionaryProvider} import org.apache.arrow.vector.dictionary.DictionaryProvider.MapDictionaryProvider import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter} @@ -267,9 +267,15 @@ object Utils extends CometTypeShim with Logging { } val provider = batchProviderOpt.getOrElse(dictionaryProvider) - val writer = new ArrowStreamWriter(root, provider, Channels.newChannel(out)) + // VectorSchemaRootAppender, which coalesces broadcast batches, assumes offsets that start + // at zero, and Arrow Java writes the bytes before the first offset too. A slice of the root + // rebases the offsets to zero and leaves those bytes out. + val written = + if (fieldVectors.forall(offsetsStartAtZero)) root else root.slice(0, root.getRowCount) + val writer = new ArrowStreamWriter(written, provider, Channels.newChannel(out)) writer.start() writer.writeBatch() + if (written ne root) written.close() root.clear() writer.close() @@ -281,6 +287,23 @@ object Utils extends CometTypeShim with Logging { } } + /** + * Whether every offset buffer in `vector`, its children included, starts at zero. + * + * arrow-rs slices a string, binary, list or map array by moving the start of its offsets, and + * the C Data import keeps them, so a vector imported from a sliced native array has offsets + * that start past zero. The rows before the slice stay at the front of its data buffer, or of + * its child vector for a list or map. + */ + private def offsetsStartAtZero(vector: FieldVector): Boolean = { + val firstOffset = vector match { + case v: BaseVariableWidthVector if v.getValueCount > 0 => v.getOffsetBuffer.getInt(0) + case v: BaseRepeatedValueVector if v.getValueCount > 0 => v.getOffsetBuffer.getInt(0) + case _ => 0 + } + firstOffset == 0 && vector.getChildrenFromFields.asScala.forall(offsetsStartAtZero) + } + /** * The classes that carry the output of [[serializeBatches]] out of Comet, for Kryo registration * by [[org.apache.comet.CometKryoRegistrator]]. @@ -380,6 +403,7 @@ object Utils extends CometTypeShim with Logging { targetRoot.allocateNew() } try { + // Appending needs offsets that start at zero, which serializeBatches writes. VectorSchemaRootAppender.append(targetRoot, sourceRoot) } catch { case e: IllegalArgumentException => diff --git a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala index 36e7c5c4f4a..d96d4f7b8e7 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometJoinSuite.scala @@ -19,6 +19,8 @@ package org.apache.comet.exec +import java.nio.charset.StandardCharsets + import org.scalactic.source.Position import org.scalatest.Tag @@ -1716,6 +1718,34 @@ class CometJoinSuite extends CometTestBase { } } + test("Broadcast coalescing keeps the values of sliced batches") { + // A grouped aggregate emits its groups as batch-size slices of one batch, so every slice after + // the first reaches the JVM with offsets that start past zero. Coalescing those slices used to + // prepend the bytes of the rows before each slice to its first value (#6779). + withSQLConf( + CometConf.COMET_BATCH_SIZE.key -> "100", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "2") { + val rows = (0 until 1000).map { i => + val s = if (i % 7 == 0) null else s"value-$i" + val bytes = if (s == null) null else s.getBytes(StandardCharsets.UTF_8) + val array = if (i % 11 == 0) null else Seq(s, s"element-$i") + (i, s, bytes, array, (s, i)) + } + withParquetTable(rows, "t") { + val (_, cometPlan) = checkSparkAnswerAndOperator( + sql(""" + |SELECT /*+ BROADCAST(d) */ t._1, d._2, d._3, d._4, d._5 + |FROM t JOIN (SELECT DISTINCT * FROM t) d ON t._1 = d._1 + |""".stripMargin), + Seq(classOf[CometBroadcastExchangeExec], classOf[CometBroadcastHashJoinExec])) + + val broadcast = collect(cometPlan) { case b: CometBroadcastExchangeExec => b }.head + assert(broadcast.metrics("numCoalescedRows").value == rows.length) + } + } + } + // Reproducer for SPARK-43113: full outer SMJ with a join filter that references // a nullable column should not match when the filter evaluates to NULL. test("SPARK-43113: Full outer SMJ with NULL in join filter") { diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala index 4510f9d0ac1..53134264cd8 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala @@ -19,11 +19,18 @@ package org.apache.spark.sql.comet.util +import java.nio.charset.StandardCharsets + +import scala.jdk.CollectionConverters._ + import org.apache.arrow.c.CDataDictionaryProvider +import org.apache.arrow.vector.{BitVectorHelper, VarCharVector} +import org.apache.arrow.vector.ipc.message.ArrowFieldNode import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.execution.vectorized.ConstantColumnVector import org.apache.spark.sql.types.{IntegerType, StringType, StructField, StructType, TimestampType} import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} +import org.apache.spark.util.io.ChunkedByteBuffer import org.apache.comet.CometArrowAllocator import org.apache.comet.vector.CometVector @@ -58,6 +65,40 @@ class UtilsSuite extends CometTestBase { assert(decoded.map(_.numRows()).sum == expected) } + test("serializeBatches rebases offsets that start past zero") { + // A slice of a native string array reaches the JVM with its original offsets, which start + // past zero, and with the bytes of the rows before the slice at the front of its data buffer. + // Coalescing such batches for a broadcast used to fold those bytes into each batch's first + // value (#6779). + val values = (0 until 24).map(i => if (i % 5 == 0) null else s"value-$i") + val skips = Seq(8, 16) + + def serialize(vectors: Seq[VarCharVector]): Seq[ChunkedByteBuffer] = { + val batches = vectors.map { v => + val column = CometVector.getVector(v, new CDataDictionaryProvider) + new ColumnarBatch(Array[ColumnVector](column), v.getValueCount) + } + try Utils.serializeBatches(batches.iterator).map(_._2).toList + finally batches.foreach(_.close()) + } + + val sliced = serialize(skips.map(slicedStrings(values, _))) + val (coalesced, batchCount, _) = Utils.coalesceBroadcastBatches(sliced.iterator) + assert(batchCount == skips.length) + val it = Utils.decodeBatches(coalesced.head, "test") + val out = it.next() + val got = (0 until out.numRows()).map { i => + if (out.column(0).isNullAt(i)) null else out.column(0).getUTF8String(i).toString + } + assert(!it.hasNext) + assert(got == skips.flatMap(values.drop)) + + // The bytes before a slice are not written either: it takes as many bytes as the same values + // in a vector of their own. + val fresh = serialize(skips.map(skip => strings(values.drop(skip)))) + assert(sliced.map(_.size) == fresh.map(_.size)) + } + test("serializeBatches materializes ConstantColumnVector columns") { // Spark wraps file-source partition columns and other per-batch constants in // ConstantColumnVector. When such a batch reaches Comet's serialization/export path @@ -197,4 +238,38 @@ class UtilsSuite extends CometTestBase { } } } + + private def strings(values: Seq[String]): VarCharVector = { + val vector = new VarCharVector("s", CometArrowAllocator) + vector.allocateNew() + values.zipWithIndex.foreach { + case (null, i) => vector.setNull(i) + case (value, i) => vector.setSafe(i, value.getBytes(StandardCharsets.UTF_8)) + } + vector.setValueCount(values.length) + vector + } + + /** + * `values.drop(skip)` in the shape the C Data import gives a slice of a native string array: + * the offsets are the original ones from `skip` on, and the data buffer still starts with the + * bytes of the first `skip` values. `skip` is a multiple of 8, so the validity bitmap slices on + * a byte boundary. + */ + private def slicedStrings(values: Seq[String], skip: Int): VarCharVector = { + val full = strings(values) + try { + val length = values.length - skip + val sliced = new VarCharVector("s", CometArrowAllocator) + sliced.loadFieldBuffers( + new ArrowFieldNode(length, values.drop(skip).count(_ == null)), + Seq( + full.getValidityBuffer.slice(skip / 8, BitVectorHelper.getValidityBufferSize(length)), + full.getOffsetBuffer.slice(skip * 4L, (length + 1) * 4L), + full.getDataBuffer).asJava) + sliced + } finally { + full.close() + } + } } From e1e392a918d2108af7f78987960ddf513207082c Mon Sep 17 00:00:00 2001 From: comphead Date: Tue, 13 Oct 2026 13:45:58 -0700 Subject: [PATCH 2/4] Pass Long arguments explicitly in the UtilsSuite slice helper The strict-warnings build rejects the implicit Int to Long widening in the `ArrowFieldNode` and `ArrowBuf.slice` calls, so convert the arguments explicitly. --- .../scala/org/apache/spark/sql/comet/util/UtilsSuite.scala | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala index 53134264cd8..e745647714e 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala @@ -24,7 +24,7 @@ import java.nio.charset.StandardCharsets import scala.jdk.CollectionConverters._ import org.apache.arrow.c.CDataDictionaryProvider -import org.apache.arrow.vector.{BitVectorHelper, VarCharVector} +import org.apache.arrow.vector.VarCharVector import org.apache.arrow.vector.ipc.message.ArrowFieldNode import org.apache.spark.sql.CometTestBase import org.apache.spark.sql.execution.vectorized.ConstantColumnVector @@ -262,9 +262,9 @@ class UtilsSuite extends CometTestBase { val length = values.length - skip val sliced = new VarCharVector("s", CometArrowAllocator) sliced.loadFieldBuffers( - new ArrowFieldNode(length, values.drop(skip).count(_ == null)), + new ArrowFieldNode(length.toLong, values.drop(skip).count(_ == null).toLong), Seq( - full.getValidityBuffer.slice(skip / 8, BitVectorHelper.getValidityBufferSize(length)), + full.getValidityBuffer.slice(skip / 8L, (length + 7) / 8L), full.getOffsetBuffer.slice(skip * 4L, (length + 1) * 4L), full.getDataBuffer).asJava) sliced From d0e841ca4b416563be63527d5ab756d43b6d9af0 Mon Sep 17 00:00:00 2001 From: comphead Date: Wed, 14 Oct 2026 16:04:32 -0700 Subject: [PATCH 3/4] Document that sliced variable-width arrays reach the JVM with offsets past zero `zero_offsets` handles `ArrayData::offset`, but a sliced string, binary, list or map array still reaches the JVM with offsets that start past zero and the rows before the slice at the front of its data. Say so in the "Array Offsets" section, and that `serializeBatches` rebases broadcast batches because `VectorSchemaRootAppender` needs offsets that start at zero. --- docs/source/contributor-guide/ffi.md | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 36328468f34..395b5c45909 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -241,6 +241,16 @@ level and shares every other buffer. `move_to_spark` applies it to executed batc or sliced booleans reach the JVM misaligned ([#6288](https://github.com/apache/datafusion-comet/issues/6288)). +`zero_offsets` leaves the offsets of a sliced string, binary, list or map array as they are. arrow-rs slices those +arrays by moving the start of their offsets, so the JVM imports them with offsets that start past zero and with the +rows before the slice still at the front of the data buffer, or of the child vector for a list or map. Reading values +through the offsets is correct, but Arrow Java 18.3's `VectorSchemaRootAppender` assumes offsets that start at zero and +folds those rows into the first value it appends, and `VectorUnloader` writes them into an IPC stream. So +`Utils.serializeBatches` rebases a broadcast batch's offsets with `VectorSchemaRoot.slice` before +`coalesceBroadcastBatches` appends it ([#6779](https://github.com/apache/datafusion-comet/issues/6779)), and code that +appends imported vectors has to rebase them too. The in-memory cache serializer still writes the earlier rows +([#6995](https://github.com/apache/datafusion-comet/issues/6995)). + ### Ownership and Lifecycle Native allocates the data, and the JVM references it without copying: From 63e83b0ac9bd27598ebd63e353b5052bb5e312ed Mon Sep 17 00:00:00 2001 From: comphead Date: Wed, 14 Oct 2026 16:15:02 -0700 Subject: [PATCH 4/4] Close the rebased slice when writing a broadcast batch fails `splitAndTransfer` takes its own references on the buffers the slice shares with the batch, so a write that failed, as it does when a cancelled broadcast interrupts the task, left them pinned after the batch closed. Close the slice in a `finally`, and test that the slice's allocator drains to zero after an interrupted write. --- .../apache/spark/sql/comet/util/Utils.scala | 17 +++++--- .../spark/sql/comet/util/UtilsSuite.scala | 41 ++++++++++++++++--- 2 files changed, 47 insertions(+), 11 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala index 1fc68aac6dc..98b70a1da78 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala @@ -272,12 +272,17 @@ object Utils extends CometTypeShim with Logging { // rebases the offsets to zero and leaves those bytes out. val written = if (fieldVectors.forall(offsetsStartAtZero)) root else root.slice(0, root.getRowCount) - val writer = new ArrowStreamWriter(written, provider, Channels.newChannel(out)) - writer.start() - writer.writeBatch() - if (written ne root) written.close() - root.clear() - writer.close() + try { + val writer = new ArrowStreamWriter(written, provider, Channels.newChannel(out)) + writer.start() + writer.writeBatch() + root.clear() + writer.close() + } finally { + // The slice holds its own references on the batch's buffers, so it is closed even when the + // write fails, as it does when a cancelled broadcast interrupts the task. + if (written ne root) written.close() + } if (out.size() > 0) { (batch.numRows().toLong, cbbos.toChunkedByteBuffer) diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala index e745647714e..54a55ccde82 100644 --- a/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala +++ b/spark/src/test/scala/org/apache/spark/sql/comet/util/UtilsSuite.scala @@ -19,11 +19,13 @@ package org.apache.spark.sql.comet.util +import java.nio.channels.ClosedByInterruptException import java.nio.charset.StandardCharsets import scala.jdk.CollectionConverters._ import org.apache.arrow.c.CDataDictionaryProvider +import org.apache.arrow.memory.BufferAllocator import org.apache.arrow.vector.VarCharVector import org.apache.arrow.vector.ipc.message.ArrowFieldNode import org.apache.spark.sql.CometTestBase @@ -99,6 +101,30 @@ class UtilsSuite extends CometTestBase { assert(sliced.map(_.size) == fresh.map(_.size)) } + test("serializeBatches closes the slice it writes when the write fails") { + // A cancelled broadcast interrupts its task, so the next write to the interruptible channel + // serializeBatches writes through fails. The slice holds its own references on the batch's + // buffers, and they would outlive the batch unless the slice is closed on that path too. + val allocator = CometArrowAllocator.newChildAllocator("sliced-batch", 0, Long.MaxValue) + try { + val vector = slicedStrings((0 until 24).map(i => s"value-$i"), 8, allocator) + val column = CometVector.getVector(vector, new CDataDictionaryProvider) + val batch = new ColumnarBatch(Array[ColumnVector](column), vector.getValueCount) + Thread.currentThread().interrupt() + try { + intercept[ClosedByInterruptException] { + Utils.serializeBatches(Iterator(batch)).next() + } + } finally { + Thread.interrupted() + batch.close() + } + assert(allocator.getAllocatedMemory == 0) + } finally { + allocator.close() + } + } + test("serializeBatches materializes ConstantColumnVector columns") { // Spark wraps file-source partition columns and other per-batch constants in // ConstantColumnVector. When such a batch reaches Comet's serialization/export path @@ -239,8 +265,10 @@ class UtilsSuite extends CometTestBase { } } - private def strings(values: Seq[String]): VarCharVector = { - val vector = new VarCharVector("s", CometArrowAllocator) + private def strings( + values: Seq[String], + allocator: BufferAllocator = CometArrowAllocator): VarCharVector = { + val vector = new VarCharVector("s", allocator) vector.allocateNew() values.zipWithIndex.foreach { case (null, i) => vector.setNull(i) @@ -256,11 +284,14 @@ class UtilsSuite extends CometTestBase { * bytes of the first `skip` values. `skip` is a multiple of 8, so the validity bitmap slices on * a byte boundary. */ - private def slicedStrings(values: Seq[String], skip: Int): VarCharVector = { - val full = strings(values) + private def slicedStrings( + values: Seq[String], + skip: Int, + allocator: BufferAllocator = CometArrowAllocator): VarCharVector = { + val full = strings(values, allocator) try { val length = values.length - skip - val sliced = new VarCharVector("s", CometArrowAllocator) + val sliced = new VarCharVector("s", allocator) sliced.loadFieldBuffers( new ArrowFieldNode(length.toLong, values.drop(skip).count(_ == null).toLong), Seq(