You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/15 03:34:06 UTC
(datafusion-comet) 01/01: fix: keep the values of sliced batches when coalescing a broadcast (#6900)
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-6900-a230ae54d223e0b60097c12e470a87d54e802b0e
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 66c63bd307b1b11fb684e6b89e9bb220f444de7c
Author: Oleks V <co...@users.noreply.github.com>
AuthorDate: Thu Oct 15 03:33:44 2026 +0000
fix: keep the values of sliced batches when coalescing a broadcast (#6900)
* 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.
* 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.
* 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.
* 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.
---
docs/source/contributor-guide/ffi.md | 10 ++
.../org/apache/spark/sql/comet/util/Utils.scala | 41 ++++++--
.../org/apache/comet/exec/CometJoinSuite.scala | 30 ++++++
.../apache/spark/sql/comet/util/UtilsSuite.scala | 106 +++++++++++++++++++++
4 files changed, 181 insertions(+), 6 deletions(-)
diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md
index 36328468f3..395b5c4590 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:
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 c884344ef1..98b70a1da7 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,11 +267,22 @@ object Utils extends CometTypeShim with Logging {
}
val provider = batchProviderOpt.getOrElse(dictionaryProvider)
- val writer = new ArrowStreamWriter(root, provider, Channels.newChannel(out))
- writer.start()
- writer.writeBatch()
- root.clear()
- writer.close()
+ // 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)
+ 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)
@@ -281,6 +292,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 +408,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 36e7c5c4f4..d96d4f7b8e 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 4510f9d0ac..54a55ccde8 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,20 @@
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
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 +67,64 @@ 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 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
@@ -197,4 +264,43 @@ class UtilsSuite extends CometTestBase {
}
}
}
+
+ 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)
+ 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,
+ allocator: BufferAllocator = CometArrowAllocator): VarCharVector = {
+ val full = strings(values, allocator)
+ try {
+ val length = values.length - skip
+ val sliced = new VarCharVector("s", allocator)
+ sliced.loadFieldBuffers(
+ new ArrowFieldNode(length.toLong, values.drop(skip).count(_ == null).toLong),
+ Seq(
+ full.getValidityBuffer.slice(skip / 8L, (length + 7) / 8L),
+ full.getOffsetBuffer.slice(skip * 4L, (length + 1) * 4L),
+ full.getDataBuffer).asJava)
+ sliced
+ } finally {
+ full.close()
+ }
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org