From 0564f2bbb1719243afa1a59444417a54279bee1a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 25 Oct 2026 13:25:31 -0600 Subject: [PATCH 1/8] fix: write a sliced batch to the in-memory cache without the rows before it --- docs/source/contributor-guide/ffi.md | 7 +- .../execution/arrow/CachedBatchIpc.scala | 55 +++++--- .../apache/spark/sql/comet/util/Utils.scala | 2 +- .../comet/exec/CometInMemoryCacheSuite.scala | 125 ++++++++++++++++++ 4 files changed, 164 insertions(+), 25 deletions(-) diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 395b5c45909..eae5f786ab7 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -247,9 +247,10 @@ rows before the slice still at the front of the data buffer, or of the child vec 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)). +`coalesceBroadcastBatches` appends it ([#6779](https://github.com/apache/datafusion-comet/issues/6779)), and +`CachedBatchIpc.serialize` does the same before the in-memory cache writes a batch, so that a cached slice neither +stores nor reports the rows before it ([#6995](https://github.com/apache/datafusion-comet/issues/6995)). Code that +appends or writes imported vectors has to rebase them too. ### Ownership and Lifecycle diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala index 2526d6b6511..17c16d190fd 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala @@ -218,32 +218,45 @@ private[comet] object CachedBatchIpc { root.setRowCount(batch.numRows()) } - // Unloaded plain and compressed afterwards rather than by handing the codec to the unloader; - // see compressed for why. - val fields = vectors.map(_.getField) - val unloader = new VectorUnloader(root, true, NoCompressionCodec.INSTANCE, true) - val plainBatch = unloader.getRecordBatch - val (sizes, recordBatch) = - try { - (columnSizes(fields, plainBatch), compressed(plainBatch, codec, allocator)) - } finally { - plainBatch.close() - } + // A column sliced on the native side reaches the JVM with offsets that start past zero, and + // VectorUnloader writes the bytes before the first offset too. A grouped aggregate emits its + // groups as slices of one batch, so each slice would be cached with the bytes of the slices + // before it, and columnSizes would report them. A slice of the root rebases the offsets to + // zero and shares the other buffers. + val written = + if (vectors.forall(Utils.offsetsStartAtZero)) root else root.slice(0, root.getRowCount) try { - // Leaves the batch in the state serializeBatches leaves one. The record batch holds its - // own buffers by now, so this does not touch it. - root.clear() - - val out = new ChunkedByteBufferOutputStream(chunkSize, ByteBuffer.allocate) + // Unloaded plain and compressed afterwards rather than by handing the codec to the + // unloader; see compressed for why. + val fields = vectors.map(_.getField) + val unloader = new VectorUnloader(written, true, NoCompressionCodec.INSTANCE, true) + val plainBatch = unloader.getRecordBatch + val (sizes, recordBatch) = + try { + (columnSizes(fields, plainBatch), compressed(plainBatch, codec, allocator)) + } finally { + plainBatch.close() + } try { - val channel = new WriteChannel(Channels.newChannel(out)) - MessageSerializer.serialize(channel, recordBatch) + // Leaves the batch in the state serializeBatches leaves one. The record batch holds its + // own buffers by now, so this does not touch it. + root.clear() + + val out = new ChunkedByteBufferOutputStream(chunkSize, ByteBuffer.allocate) + try { + val channel = new WriteChannel(Channels.newChannel(out)) + MessageSerializer.serialize(channel, recordBatch) + } finally { + out.close() + } + (out.toChunkedByteBuffer, sizes) } finally { - out.close() + recordBatch.close() } - (out.toChunkedByteBuffer, sizes) } finally { - recordBatch.close() + // The slice holds its own references on the batch's buffers, so it is closed even when the + // write fails. + if (written ne root) written.close() } } finally { // Only the vectors allocated for this write. The rest belong to the input batch. 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 2fb1dfb5569..cfd22caff36 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 @@ -310,7 +310,7 @@ object Utils extends CometTypeShim with Logging { * 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 = { + private[comet] 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) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 3273b2af7bd..1c488e303e5 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -32,6 +32,7 @@ import org.apache.arrow.compression.ZstdCompressionCodec import org.apache.arrow.memory.{ArrowBuf, BufferAllocator} import org.apache.arrow.vector.{BitVector, FixedSizeBinaryVector, IntVector, LargeVarCharVector, VarBinaryVector, VarCharVector} import org.apache.arrow.vector.compression.{CompressionCodec, CompressionUtil, NoCompressionCodec} +import org.apache.arrow.vector.ipc.message.ArrowFieldNode import org.apache.arrow.vector.types.pojo.ArrowType import org.apache.spark.{SparkConf, SparkException} import org.apache.spark.CometDriverPlugin @@ -2758,6 +2759,130 @@ class CometInMemoryCacheSuite extends CometTestBase { assert(zstdSize == plainSize, s"the relation size should not depend on the codec: $sizes") } + test("Comet in-memory cache drops the rows before a sliced aggregate output") { + // A grouped aggregate emits its groups as batch-size slices of one batch, and a slice of a + // string or binary array reaches the JVM with offsets that start past zero and the rows before + // it still at the front of its data buffer. Written as it stands, the k-th slice is cached with + // the bytes of the k - 1 slices before it, and the relation reports them in its size. + val rows = 4000 + withTempPath { dir => + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_BATCH_SIZE.key -> "100", + CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { + spark.catalog.clearCache() + val keys = "concat('key-', lpad(cast(id AS string), 12, '0'))" + spark + .range(0, rows, 1, 1) + .selectExpr(s"$keys AS k", s"cast($keys AS binary) AS b") + .write + .parquet(dir.getCanonicalPath) + + def cachedSize(name: String, df: DataFrame): Long = { + df.createOrReplaceTempView(name) + spark.catalog.cacheTable(name) + assert(spark.table(name).count() == rows) + val relation = spark.sharedState.cacheManager + .lookupCachedData(spark.table(name)) + .get + .cachedRepresentation + val batches = relation.cacheBuilder.cachedColumnBuffers.collect() + assert(batches.forall( + _.getClass.getName == "org.apache.spark.sql.comet.execution.arrow.CometCachedBatch")) + relation.computeStats().sizeInBytes.toLong + } + + try { + val scanned = spark.read.parquet(dir.getCanonicalPath) + val distinct = scanned.distinct() + // The aggregate has to emit more than one slice for the bytes before a slice to matter. + val plan = distinct.queryExecution.executedPlan + assert(collect(plan) { case a: CometHashAggregateExec => a }.nonEmpty) + + val scannedSize = cachedSize("scanned_cache", scanned) + val distinctSize = cachedSize("distinct_cache", distinct) + // The same rows, so the same bytes of string data. A slice that carried the bytes before + // it would make this about 20 times larger, one for each of the 40 slices. + assert( + distinctSize < scannedSize * 1.5, + s"the aggregate's cache is $distinctSize bytes, a scan's of the same rows $scannedSize") + checkAnswer(spark.table("distinct_cache"), scanned) + } finally { + spark.catalog.clearCache() + } + } + } + } + + test("Comet in-memory cache writes a sliced string vector as its own values and releases it") { + // The rebased copy of a sliced vector shares the batch's buffers, so it has to be closed + // whether the write succeeds or fails, or the batch's data outlives the batch. + val values = (0 until 24).map(i => if (i % 5 == 0) null else s"value-$i") + val skip = 8 + def strings(allocator: BufferAllocator, vs: Seq[String]): VarCharVector = { + val vector = new VarCharVector("s", allocator) + vector.allocateNew() + vs.zipWithIndex.foreach { + case (null, i) => vector.setNull(i) + case (v, i) => vector.setSafe(i, v.getBytes(StandardCharsets.UTF_8)) + } + vector.setValueCount(vs.length) + vector + } + + // The slice of `values` from `skip` on, imported as a native slice is: offsets that start + // past zero over the whole data buffer. + def sliced(allocator: BufferAllocator): VarCharVector = { + val full = strings(allocator, values) + try { + val length = values.length - skip + val vector = new VarCharVector("s", allocator) + vector.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) + vector + } finally { + full.close() + } + } + + def payload(vector: VarCharVector, codec: CompressionCodec): Long = { + val batch = + new ColumnarBatch(Array[ColumnVector](new CometPlainVector(vector)), vector.getValueCount) + CometCachedBatchHelper.payloadSize( + CometCachedBatchHelper.cachedBatch( + CometCachedBatchHelper.serialize(batch, codec, CometArrowAllocator), + batch.numRows())) + } + + val codec = CometCachedBatchHelper.compressionCodec("none") + // Closing a child allocator that still holds a buffer throws, which is the leak check. + val allocator = CometArrowAllocator.newChildAllocator("sliced-write", 0, Long.MaxValue) + try { + val vector = sliced(allocator) + val fresh = strings(allocator, values.drop(skip)) + val failed = sliced(allocator) + try { + assert(vector.getOffsetBuffer.getInt(0) > 0, "the test needs offsets past zero") + assert(payload(vector, codec) == payload(fresh, codec)) + + // The write fails after the slice has rebased its offsets into a buffer of its own. + val failing = new FailAfterCompressionCodec(succeedFor = 1) + intercept[Exception](payload(failed, failing)) + } finally { + fresh.close() + vector.close() + failed.close() + } + } finally { + allocator.close() + } + } + test("Comet in-memory cache scans no columns for a row-count-only query") { // SELECT count(*) selects no columns, and the scan must keep it that way. Widening it -- to // the whole cache schema, or to a single placeholder column -- makes the emitted batches From cdd555befae93140035c064e6a648ab0cf0fe8ff Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 25 Oct 2026 13:29:33 -0600 Subject: [PATCH 2/8] refactor: share the offset rebase and make the cache test read the scan --- .../execution/arrow/CachedBatchIpc.scala | 10 +++------ .../apache/spark/sql/comet/util/Utils.scala | 15 ++++++++++--- .../comet/exec/CometInMemoryCacheSuite.scala | 21 +++++++++++-------- 3 files changed, 27 insertions(+), 19 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala index 17c16d190fd..b7dddddaba6 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala @@ -221,10 +221,8 @@ private[comet] object CachedBatchIpc { // A column sliced on the native side reaches the JVM with offsets that start past zero, and // VectorUnloader writes the bytes before the first offset too. A grouped aggregate emits its // groups as slices of one batch, so each slice would be cached with the bytes of the slices - // before it, and columnSizes would report them. A slice of the root rebases the offsets to - // zero and shares the other buffers. - val written = - if (vectors.forall(Utils.offsetsStartAtZero)) root else root.slice(0, root.getRowCount) + // before it, and columnSizes would report them. + val written = Utils.rebasedForWrite(root) try { // Unloaded plain and compressed afterwards rather than by handing the codec to the // unloader; see compressed for why. @@ -239,7 +237,7 @@ private[comet] object CachedBatchIpc { } try { // Leaves the batch in the state serializeBatches leaves one. The record batch holds its - // own buffers by now, so this does not touch it. + // own buffers by now, and so does the slice, so this does not touch either. root.clear() val out = new ChunkedByteBufferOutputStream(chunkSize, ByteBuffer.allocate) @@ -254,8 +252,6 @@ private[comet] object CachedBatchIpc { recordBatch.close() } } finally { - // The slice holds its own references on the batch's buffers, so it is closed even when the - // write fails. if (written ne root) written.close() } } finally { 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 cfd22caff36..35cce5055c9 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 @@ -273,8 +273,7 @@ object Utils extends CometTypeShim with Logging { // 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 written = rebasedForWrite(root) try { val writer = new ArrowStreamWriter(written, provider, Channels.newChannel(out)) writer.start() @@ -302,6 +301,16 @@ object Utils extends CometTypeShim with Logging { } } + /** + * `root` itself when every offset buffer in it starts at zero, otherwise a slice of it with the + * offsets rebased to zero and the bytes before the first of them left out. The slice shares the + * other buffers and holds its own references on them, so the caller closes it whenever it is + * not `root`, even when the write fails. + */ + private[comet] def rebasedForWrite(root: VectorSchemaRoot): VectorSchemaRoot = + if (root.getFieldVectors.asScala.forall(offsetsStartAtZero)) root + else root.slice(0, root.getRowCount) + /** * Whether every offset buffer in `vector`, its children included, starts at zero. * @@ -310,7 +319,7 @@ object Utils extends CometTypeShim with Logging { * 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[comet] def offsetsStartAtZero(vector: FieldVector): Boolean = { + 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) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 1c488e303e5..4dbcc7aed2d 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -2779,7 +2779,7 @@ class CometInMemoryCacheSuite extends CometTestBase { .write .parquet(dir.getCanonicalPath) - def cachedSize(name: String, df: DataFrame): Long = { + def cachedSize(name: String, df: DataFrame): (Long, Int) = { df.createOrReplaceTempView(name) spark.catalog.cacheTable(name) assert(spark.table(name).count() == rows) @@ -2790,22 +2790,25 @@ class CometInMemoryCacheSuite extends CometTestBase { val batches = relation.cacheBuilder.cachedColumnBuffers.collect() assert(batches.forall( _.getClass.getName == "org.apache.spark.sql.comet.execution.arrow.CometCachedBatch")) - relation.computeStats().sizeInBytes.toLong + (relation.computeStats().sizeInBytes.toLong, batches.length) } try { val scanned = spark.read.parquet(dir.getCanonicalPath) val distinct = scanned.distinct() - // The aggregate has to emit more than one slice for the bytes before a slice to matter. - val plan = distinct.queryExecution.executedPlan - assert(collect(plan) { case a: CometHashAggregateExec => a }.nonEmpty) - - val scannedSize = cachedSize("scanned_cache", scanned) - val distinctSize = cachedSize("distinct_cache", distinct) + assert(collect(distinct.queryExecution.executedPlan) { case a: CometHashAggregateExec => + a + }.nonEmpty) + + // The aggregate is cached first, so that it reads the Parquet scan and not a cached copy + // of it. It has to emit more than one slice for the bytes before a slice to matter. + val (distinctSize, distinctBatches) = cachedSize("distinct_cache", distinct) + assert(distinctBatches > 1, s"the aggregate emitted $distinctBatches batches") + val (scannedSize, _) = cachedSize("scanned_cache", scanned) // The same rows, so the same bytes of string data. A slice that carried the bytes before // it would make this about 20 times larger, one for each of the 40 slices. assert( - distinctSize < scannedSize * 1.5, + distinctSize < scannedSize * 1.1, s"the aggregate's cache is $distinctSize bytes, a scan's of the same rows $scannedSize") checkAnswer(spark.table("distinct_cache"), scanned) } finally { From ca710474473dfe4c8ff9bf530ffeb63ebdaaccfa Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 25 Oct 2026 13:32:44 -0600 Subject: [PATCH 3/8] test: cover a sliced list and assert the injected failure in the cache write tests --- .../comet/exec/CometInMemoryCacheSuite.scala | 108 +++++++++++++++--- 1 file changed, 95 insertions(+), 13 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 4dbcc7aed2d..7cab3b63be6 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -28,12 +28,14 @@ import scala.concurrent.ExecutionContext.Implicits.global import scala.concurrent.duration.DurationInt import scala.jdk.CollectionConverters._ +import org.apache.arrow.c.CDataDictionaryProvider import org.apache.arrow.compression.ZstdCompressionCodec import org.apache.arrow.memory.{ArrowBuf, BufferAllocator} -import org.apache.arrow.vector.{BitVector, FixedSizeBinaryVector, IntVector, LargeVarCharVector, VarBinaryVector, VarCharVector} +import org.apache.arrow.vector.{BitVector, FieldVector, FixedSizeBinaryVector, IntVector, LargeVarCharVector, VarBinaryVector, VarCharVector} +import org.apache.arrow.vector.complex.ListVector import org.apache.arrow.vector.compression.{CompressionCodec, CompressionUtil, NoCompressionCodec} import org.apache.arrow.vector.ipc.message.ArrowFieldNode -import org.apache.arrow.vector.types.pojo.ArrowType +import org.apache.arrow.vector.types.pojo.{ArrowType, FieldType} import org.apache.spark.{SparkConf, SparkException} import org.apache.spark.CometDriverPlugin import org.apache.spark.sql.{CometTestBase, DataFrame, Observation, QueryTest, Row} @@ -2818,6 +2820,91 @@ class CometInMemoryCacheSuite extends CometTestBase { } } + /** The bytes the cache writer stores for a batch of `vector` alone. */ + private def payloadSize(vector: FieldVector, codec: CompressionCodec): Long = { + val column = CometVector.getVector(vector, new CDataDictionaryProvider) + val batch = new ColumnarBatch(Array[ColumnVector](column), vector.getValueCount) + CometCachedBatchHelper.payloadSize( + CometCachedBatchHelper.cachedBatch( + CometCachedBatchHelper.serialize(batch, codec, CometArrowAllocator), + batch.numRows())) + } + + test("Comet in-memory cache writes a sliced list of strings as its own values") { + // A list slice has offsets past zero in two places: its own, and its child's once the child + // is sliced too. Only the list's start past zero here, over a child that holds every element, + // which is how a native slice of a list reaches the JVM. + val values: Seq[Seq[String]] = (0 until 24).map { i => + if (i % 5 == 0) null else (0 until i % 4).map(j => s"element-$i-$j") + } + val skip = 8 + + def lists(allocator: BufferAllocator, vs: Seq[Seq[String]]): ListVector = { + val list = ListVector.empty("l", allocator) + val child = list + .addOrGetVector[VarCharVector](FieldType.nullable(ArrowType.Utf8.INSTANCE)) + .getVector + child.allocateNew() + var position = 0 + vs.zipWithIndex.foreach { case (elements, i) => + if (elements == null) { + list.setNull(i) + } else { + list.startNewValue(i) + elements.foreach { e => + child.setSafe(position, e.getBytes(StandardCharsets.UTF_8)) + position += 1 + } + list.endValue(i, elements.length) + } + } + child.setValueCount(position) + list.setValueCount(vs.length) + list + } + + def sliced(allocator: BufferAllocator): ListVector = { + val full = lists(allocator, values) + try { + val length = values.length - skip + val list = + new ListVector("l", allocator, FieldType.nullable(ArrowType.List.INSTANCE), null) + list.initializeChildrenFromFields(full.getField.getChildren) + list.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)).asJava) + val fullChild = full.getDataVector.asInstanceOf[VarCharVector] + list.getDataVector.loadFieldBuffers( + new ArrowFieldNode(fullChild.getValueCount.toLong, 0L), + Seq( + fullChild.getValidityBuffer, + fullChild.getOffsetBuffer, + fullChild.getDataBuffer).asJava) + list + } finally { + full.close() + } + } + + val codec = CometCachedBatchHelper.compressionCodec("none") + val allocator = CometArrowAllocator.newChildAllocator("sliced-list-write", 0, Long.MaxValue) + try { + val vector = sliced(allocator) + val fresh = lists(allocator, values.drop(skip)) + try { + assert(vector.getOffsetBuffer.getInt(0) > 0, "the test needs offsets past zero") + assert(payloadSize(vector, codec) == payloadSize(fresh, codec)) + } finally { + fresh.close() + vector.close() + } + } finally { + allocator.close() + } + } + test("Comet in-memory cache writes a sliced string vector as its own values and releases it") { // The rebased copy of a sliced vector shares the batch's buffers, so it has to be closed // whether the write succeeds or fails, or the batch's data outlives the batch. @@ -2853,15 +2940,6 @@ class CometInMemoryCacheSuite extends CometTestBase { } } - def payload(vector: VarCharVector, codec: CompressionCodec): Long = { - val batch = - new ColumnarBatch(Array[ColumnVector](new CometPlainVector(vector)), vector.getValueCount) - CometCachedBatchHelper.payloadSize( - CometCachedBatchHelper.cachedBatch( - CometCachedBatchHelper.serialize(batch, codec, CometArrowAllocator), - batch.numRows())) - } - val codec = CometCachedBatchHelper.compressionCodec("none") // Closing a child allocator that still holds a buffer throws, which is the leak check. val allocator = CometArrowAllocator.newChildAllocator("sliced-write", 0, Long.MaxValue) @@ -2871,11 +2949,15 @@ class CometInMemoryCacheSuite extends CometTestBase { val failed = sliced(allocator) try { assert(vector.getOffsetBuffer.getInt(0) > 0, "the test needs offsets past zero") - assert(payload(vector, codec) == payload(fresh, codec)) + assert(payloadSize(vector, codec) == payloadSize(fresh, codec)) // The write fails after the slice has rebased its offsets into a buffer of its own. val failing = new FailAfterCompressionCodec(succeedFor = 1) - intercept[Exception](payload(failed, failing)) + val thrown = intercept[Exception](payloadSize(failed, failing)) + assert( + causeChain(thrown).exists(t => + Option(t.getMessage).contains(FailAfterCompressionCodec.Message)), + s"a compression failure must surface as itself: $thrown") } finally { fresh.close() vector.close() From 2233891bc1740aadf97e2ed1469dac14c3a3438d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 25 Oct 2026 13:38:03 -0600 Subject: [PATCH 4/8] refactor: own the sliced root's close in one helper and slim the cache tests --- .../execution/arrow/CachedBatchIpc.scala | 12 +- .../apache/spark/sql/comet/util/Utils.scala | 30 ++-- .../comet/exec/CometInMemoryCacheSuite.scala | 137 ++++++++++-------- 3 files changed, 92 insertions(+), 87 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala index b7dddddaba6..6de3e1b2ddb 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CachedBatchIpc.scala @@ -218,12 +218,10 @@ private[comet] object CachedBatchIpc { root.setRowCount(batch.numRows()) } - // A column sliced on the native side reaches the JVM with offsets that start past zero, and - // VectorUnloader writes the bytes before the first offset too. A grouped aggregate emits its - // groups as slices of one batch, so each slice would be cached with the bytes of the slices - // before it, and columnSizes would report them. - val written = Utils.rebasedForWrite(root) - try { + // A grouped aggregate emits its groups as slices of one batch, which reach the JVM with + // offsets that start past zero. Unloaded as they are, each slice would be cached, and + // reported by columnSizes, with the bytes of the slices before it. + Utils.withRebasedOffsets(root) { written => // Unloaded plain and compressed afterwards rather than by handing the codec to the // unloader; see compressed for why. val fields = vectors.map(_.getField) @@ -251,8 +249,6 @@ private[comet] object CachedBatchIpc { } finally { recordBatch.close() } - } finally { - if (written ne root) written.close() } } finally { // Only the vectors allocated for this write. The rest belong to the input batch. 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 35cce5055c9..b5932602ef7 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 @@ -271,19 +271,13 @@ object Utils extends CometTypeShim with Logging { val provider = batchProviderOpt.getOrElse(dictionaryProvider) // 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 = rebasedForWrite(root) - try { + // at zero. + withRebasedOffsets(root) { written => 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() } } catch { case NonFatal(e) => @@ -302,14 +296,20 @@ object Utils extends CometTypeShim with Logging { } /** - * `root` itself when every offset buffer in it starts at zero, otherwise a slice of it with the - * offsets rebased to zero and the bytes before the first of them left out. The slice shares the - * other buffers and holds its own references on them, so the caller closes it whenever it is - * not `root`, even when the write fails. + * Runs `write` on `root` itself when every offset buffer in it starts at zero, and otherwise on + * a slice of it that rebases the offsets to zero and leaves out the bytes before the first of + * them, which Arrow Java would write. The slice shares the other buffers and holds its own + * references on them, so it is closed whether `write` returns or fails, as it does when a + * cancelled broadcast interrupts the task. */ - private[comet] def rebasedForWrite(root: VectorSchemaRoot): VectorSchemaRoot = - if (root.getFieldVectors.asScala.forall(offsetsStartAtZero)) root - else root.slice(0, root.getRowCount) + private[comet] def withRebasedOffsets[T](root: VectorSchemaRoot)( + write: VectorSchemaRoot => T): T = { + val written = + if (root.getFieldVectors.asScala.forall(offsetsStartAtZero)) root + else root.slice(0, root.getRowCount) + try write(written) + finally if (written ne root) written.close() + } /** * Whether every offset buffer in `vector`, its children included, starts at zero. diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 7cab3b63be6..332acb8e2b9 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -2763,10 +2763,10 @@ class CometInMemoryCacheSuite extends CometTestBase { test("Comet in-memory cache drops the rows before a sliced aggregate output") { // A grouped aggregate emits its groups as batch-size slices of one batch, and a slice of a - // string or binary array reaches the JVM with offsets that start past zero and the rows before - // it still at the front of its data buffer. Written as it stands, the k-th slice is cached with - // the bytes of the k - 1 slices before it, and the relation reports them in its size. - val rows = 4000 + // string array reaches the JVM with offsets that start past zero and the rows before it still + // at the front of its data buffer. Written as it stands, the k-th slice is cached with the + // bytes of the k - 1 slices before it, and the relation reports them in its size. + val rows = 1000 withTempPath { dir => withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", @@ -2774,14 +2774,13 @@ class CometInMemoryCacheSuite extends CometTestBase { CometConf.COMET_BATCH_SIZE.key -> "100", CometConf.COMET_EXEC_IN_MEMORY_CACHE_ENABLED.key -> "true") { spark.catalog.clearCache() - val keys = "concat('key-', lpad(cast(id AS string), 12, '0'))" spark .range(0, rows, 1, 1) - .selectExpr(s"$keys AS k", s"cast($keys AS binary) AS b") + .selectExpr("concat('key-', lpad(cast(id AS string), 12, '0')) AS k") .write .parquet(dir.getCanonicalPath) - def cachedSize(name: String, df: DataFrame): (Long, Int) = { + def cachedSize(name: String, df: DataFrame): (Long, Long) = { df.createOrReplaceTempView(name) spark.catalog.cacheTable(name) assert(spark.table(name).count() == rows) @@ -2789,18 +2788,17 @@ class CometInMemoryCacheSuite extends CometTestBase { .lookupCachedData(spark.table(name)) .get .cachedRepresentation - val batches = relation.cacheBuilder.cachedColumnBuffers.collect() - assert(batches.forall( - _.getClass.getName == "org.apache.spark.sql.comet.execution.arrow.CometCachedBatch")) - (relation.computeStats().sizeInBytes.toLong, batches.length) + val batches = relation.cacheBuilder.cachedColumnBuffers + val classes = batches.map(_.getClass.getSimpleName).distinct().collect() + assert(classes.sameElements(Array("CometCachedBatch"))) + (relation.computeStats().sizeInBytes.toLong, batches.count()) } try { val scanned = spark.read.parquet(dir.getCanonicalPath) val distinct = scanned.distinct() - assert(collect(distinct.queryExecution.executedPlan) { case a: CometHashAggregateExec => - a - }.nonEmpty) + assert( + distinct.queryExecution.executedPlan.exists(_.isInstanceOf[CometHashAggregateExec])) // The aggregate is cached first, so that it reads the Parquet scan and not a cached copy // of it. It has to emit more than one slice for the bytes before a slice to matter. @@ -2808,7 +2806,7 @@ class CometInMemoryCacheSuite extends CometTestBase { assert(distinctBatches > 1, s"the aggregate emitted $distinctBatches batches") val (scannedSize, _) = cachedSize("scanned_cache", scanned) // The same rows, so the same bytes of string data. A slice that carried the bytes before - // it would make this about 20 times larger, one for each of the 40 slices. + // it would make this about 5 times larger, as the 10 slices average 5.5 slices of data. assert( distinctSize < scannedSize * 1.1, s"the aggregate's cache is $distinctSize bytes, a scan's of the same rows $scannedSize") @@ -2830,6 +2828,56 @@ class CometInMemoryCacheSuite extends CometTestBase { batch.numRows())) } + /** + * The field node, validity and offset buffers of the rows of `full` from `skip` on, which the C + * Data import hands the JVM for a native slice: the offsets start past zero, over the buffers + * of all of `full`'s rows. + */ + private def nativeSlice( + full: FieldVector, + skip: Int, + nulls: Int): (ArrowFieldNode, Seq[ArrowBuf]) = { + val length = full.getValueCount - skip + ( + new ArrowFieldNode(length.toLong, nulls.toLong), + Seq( + full.getValidityBuffer.slice(skip / 8L, (length + 7) / 8L), + full.getOffsetBuffer.slice(skip * 4L, (length + 1) * 4L))) + } + + private def assertSliceWritesAsOwnValues(sliced: FieldVector, fresh: FieldVector): Unit = { + assert(sliced.getOffsetBuffer.getInt(0) > 0, "the test needs offsets past zero") + val codec = CometCachedBatchHelper.compressionCodec("none") + assert(payloadSize(sliced, codec) == payloadSize(fresh, codec)) + } + + private val sliceSkip = 8 + + private def stringVector(allocator: BufferAllocator, values: Seq[String]): VarCharVector = { + val vector = new VarCharVector("s", allocator) + vector.allocateNew() + values.zipWithIndex.foreach { + case (null, i) => vector.setNull(i) + case (v, i) => vector.setSafe(i, v.getBytes(StandardCharsets.UTF_8)) + } + vector.setValueCount(values.length) + vector + } + + private def slicedStringVector( + allocator: BufferAllocator, + values: Seq[String]): VarCharVector = { + val full = stringVector(allocator, values) + try { + val (node, buffers) = nativeSlice(full, sliceSkip, values.drop(sliceSkip).count(_ == null)) + val vector = new VarCharVector("s", allocator) + vector.loadFieldBuffers(node, (buffers :+ full.getDataBuffer).asJava) + vector + } finally { + full.close() + } + } + test("Comet in-memory cache writes a sliced list of strings as its own values") { // A list slice has offsets past zero in two places: its own, and its child's once the child // is sliced too. Only the list's start past zero here, over a child that holds every element, @@ -2837,7 +2885,6 @@ class CometInMemoryCacheSuite extends CometTestBase { val values: Seq[Seq[String]] = (0 until 24).map { i => if (i % 5 == 0) null else (0 until i % 4).map(j => s"element-$i-$j") } - val skip = 8 def lists(allocator: BufferAllocator, vs: Seq[Seq[String]]): ListVector = { val list = ListVector.empty("l", allocator) @@ -2866,15 +2913,12 @@ class CometInMemoryCacheSuite extends CometTestBase { def sliced(allocator: BufferAllocator): ListVector = { val full = lists(allocator, values) try { - val length = values.length - skip + val (node, buffers) = + nativeSlice(full, sliceSkip, values.drop(sliceSkip).count(_ == null)) val list = new ListVector("l", allocator, FieldType.nullable(ArrowType.List.INSTANCE), null) list.initializeChildrenFromFields(full.getField.getChildren) - list.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)).asJava) + list.loadFieldBuffers(node, buffers.asJava) val fullChild = full.getDataVector.asInstanceOf[VarCharVector] list.getDataVector.loadFieldBuffers( new ArrowFieldNode(fullChild.getValueCount.toLong, 0L), @@ -2888,15 +2932,12 @@ class CometInMemoryCacheSuite extends CometTestBase { } } - val codec = CometCachedBatchHelper.compressionCodec("none") val allocator = CometArrowAllocator.newChildAllocator("sliced-list-write", 0, Long.MaxValue) try { val vector = sliced(allocator) - val fresh = lists(allocator, values.drop(skip)) - try { - assert(vector.getOffsetBuffer.getInt(0) > 0, "the test needs offsets past zero") - assert(payloadSize(vector, codec) == payloadSize(fresh, codec)) - } finally { + val fresh = lists(allocator, values.drop(sliceSkip)) + try assertSliceWritesAsOwnValues(vector, fresh) + finally { fresh.close() vector.close() } @@ -2909,47 +2950,15 @@ class CometInMemoryCacheSuite extends CometTestBase { // The rebased copy of a sliced vector shares the batch's buffers, so it has to be closed // whether the write succeeds or fails, or the batch's data outlives the batch. val values = (0 until 24).map(i => if (i % 5 == 0) null else s"value-$i") - val skip = 8 - def strings(allocator: BufferAllocator, vs: Seq[String]): VarCharVector = { - val vector = new VarCharVector("s", allocator) - vector.allocateNew() - vs.zipWithIndex.foreach { - case (null, i) => vector.setNull(i) - case (v, i) => vector.setSafe(i, v.getBytes(StandardCharsets.UTF_8)) - } - vector.setValueCount(vs.length) - vector - } - - // The slice of `values` from `skip` on, imported as a native slice is: offsets that start - // past zero over the whole data buffer. - def sliced(allocator: BufferAllocator): VarCharVector = { - val full = strings(allocator, values) - try { - val length = values.length - skip - val vector = new VarCharVector("s", allocator) - vector.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) - vector - } finally { - full.close() - } - } - val codec = CometCachedBatchHelper.compressionCodec("none") // Closing a child allocator that still holds a buffer throws, which is the leak check. val allocator = CometArrowAllocator.newChildAllocator("sliced-write", 0, Long.MaxValue) try { - val vector = sliced(allocator) - val fresh = strings(allocator, values.drop(skip)) - val failed = sliced(allocator) + val vector = slicedStringVector(allocator, values) + val fresh = stringVector(allocator, values.drop(sliceSkip)) + val failed = slicedStringVector(allocator, values) try { - assert(vector.getOffsetBuffer.getInt(0) > 0, "the test needs offsets past zero") - assert(payloadSize(vector, codec) == payloadSize(fresh, codec)) + assertSliceWritesAsOwnValues(vector, fresh) // The write fails after the slice has rebased its offsets into a buffer of its own. val failing = new FailAfterCompressionCodec(succeedFor = 1) From 72fd2d7c5af190cfb1373da273dece19ea85aec6 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Sun, 25 Oct 2026 15:07:59 -0600 Subject: [PATCH 5/8] test: use a Long row count for the sliced aggregate cache test --- .../scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 332acb8e2b9..0ca12b007a6 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -2766,7 +2766,7 @@ class CometInMemoryCacheSuite extends CometTestBase { // string array reaches the JVM with offsets that start past zero and the rows before it still // at the front of its data buffer. Written as it stands, the k-th slice is cached with the // bytes of the k - 1 slices before it, and the relation reports them in its size. - val rows = 1000 + val rows = 1000L withTempPath { dir => withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", From b4a0ae1d000c77f48b8b4256fb66f244049d678f Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 28 Oct 2026 08:54:02 -0600 Subject: [PATCH 6/8] fix: rebase the first slice of a list or map before writing it The first slice of a native list or map array starts its offsets at zero, but its child still holds the elements of the later slices, so it was written with all of them. Rebase a list or map whose last offset is short of its child's length too, and test the slices of a list in UtilsSuite against withRebasedOffsets. Widen withRebasedOffsets to private[sql] so the Python runner can use it. --- .../apache/spark/sql/comet/util/Utils.scala | 38 +++--- .../spark/sql/comet/util/UtilsSuite.scala | 113 +++++++++++++++++- 2 files changed, 134 insertions(+), 17 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 fdd1d130e55..3561a5ef570 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 @@ -302,36 +302,44 @@ object Utils extends CometTypeShim with Logging { } /** - * Runs `write` on `root` itself when every offset buffer in it starts at zero, and otherwise on - * a slice of it that rebases the offsets to zero and leaves out the bytes before the first of - * them, which Arrow Java would write. The slice shares the other buffers and holds its own - * references on them, so it is closed whether `write` returns or fails, as it does when a - * cancelled broadcast interrupts the task. + * Runs `write` on `root` itself when none of its vectors is a slice of a larger one, and + * otherwise on a slice of it that rebases the offsets to zero and leaves out the bytes or + * elements outside them, which Arrow Java would write. The slice shares the other buffers and + * holds its own references on them, so it is closed whether `write` returns or fails, as it + * does when a cancelled broadcast interrupts the task. */ - private[comet] def withRebasedOffsets[T](root: VectorSchemaRoot)( + private[sql] def withRebasedOffsets[T](root: VectorSchemaRoot)( write: VectorSchemaRoot => T): T = { val written = - if (root.getFieldVectors.asScala.forall(offsetsStartAtZero)) root + if (root.getFieldVectors.asScala.forall(holdsOnlyItsOwnValues)) root else root.slice(0, root.getRowCount) try write(written) finally if (written ne root) written.close() } /** - * Whether every offset buffer in `vector`, its children included, starts at zero. + * Whether every vector in `vector`, its children included, holds only the bytes or elements its + * own offsets cover. * * 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. + * its child vector for a list or map. The first slice of an array starts at zero instead, and a + * list or map still has the elements of the later slices in its child, which is longer than the + * last offset. A string or binary vector stops writing its data at the last offset, so only its + * first offset can show a slice. */ - 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 + private def holdsOnlyItsOwnValues(vector: FieldVector): Boolean = { + val ownsItsValues = vector match { + case v: BaseVariableWidthVector if v.getValueCount > 0 => + v.getOffsetBuffer.getInt(0) == 0 + case v: BaseRepeatedValueVector if v.getValueCount > 0 => + v.getOffsetBuffer.getInt(0) == 0 && + v.getOffsetBuffer.getInt(v.getValueCount.toLong * BaseRepeatedValueVector.OFFSET_WIDTH) == + v.getDataVector.getValueCount + case _ => true } - firstOffset == 0 && vector.getChildrenFromFields.asScala.forall(offsetsStartAtZero) + ownsItsValues && vector.getChildrenFromFields.asScala.forall(holdsOnlyItsOwnValues) } /** 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 29c4547ccd3..b23f40684ee 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 @@ -26,8 +26,8 @@ import scala.jdk.CollectionConverters._ import org.apache.arrow.c.CDataDictionaryProvider import org.apache.arrow.memory.BufferAllocator -import org.apache.arrow.vector.{BaseLargeVariableWidthVector, BitVectorHelper, IntVector, LargeVarBinaryVector, LargeVarCharVector, ValueVector, VarBinaryVector, VarCharVector, VariableWidthFieldVector} -import org.apache.arrow.vector.complex.StructVector +import org.apache.arrow.vector.{BaseLargeVariableWidthVector, BitVectorHelper, FieldVector, IntVector, LargeVarBinaryVector, LargeVarCharVector, ValueVector, VarBinaryVector, VarCharVector, VariableWidthFieldVector, VectorSchemaRoot, VectorUnloader} +import org.apache.arrow.vector.complex.{ListVector, StructVector} import org.apache.arrow.vector.dictionary.{Dictionary, DictionaryProvider} import org.apache.arrow.vector.ipc.message.ArrowFieldNode import org.apache.arrow.vector.types.pojo.{ArrowType, DictionaryEncoding, Field, FieldType} @@ -105,6 +105,41 @@ class UtilsSuite extends CometTestBase { assert(sliced.map(_.size) == fresh.map(_.size)) } + test("withRebasedOffsets writes a sliced list as its own values") { + // A native slice of a list array has the offsets of the list and a child that both still hold + // the rows around the slice. The first slice starts at zero, so only the child shows that it + // is a slice: it holds the elements of every later one. Written as it stands, a slice carries + // those elements too. + val values: Seq[Seq[String]] = (0 until 24).map { i => + if (i % 5 == 0) null else (0 until i % 4).map(j => s"element-$i-$j") + } + val allocator = CometArrowAllocator.newChildAllocator("sliced-lists", 0, Long.MaxValue) + try { + Seq((0, 8), (8, 8), (8, 16), (16, 8)).foreach { case (skip, take) => + val sliced = slicedLists(values, skip, take, allocator) + val fresh = lists(values.slice(skip, skip + take), allocator) + try { + val offsets = sliced.getOffsetBuffer + assert( + offsets.getInt(0) > 0 || offsets.getInt( + take * 4L) < sliced.getDataVector.getValueCount, + s"the test needs a slice: skip=$skip, take=$take") + val (slicedBytes, slicedValues) = writtenBytesAndValues(sliced) + val (freshBytes, freshValues) = writtenBytesAndValues(fresh) + assert(slicedValues == freshValues, s"values of skip=$skip, take=$take") + assert(slicedBytes == freshBytes, s"bytes of skip=$skip, take=$take") + } finally { + fresh.close() + sliced.close() + } + } + // The slice withRebasedOffsets made is released with the rest. + assert(allocator.getAllocatedMemory == 0) + } finally { + allocator.close() + } + } + 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 @@ -563,6 +598,80 @@ class UtilsSuite extends CometTestBase { vector } + private def lists(values: Seq[Seq[String]], allocator: BufferAllocator): ListVector = { + val list = ListVector.empty("l", allocator) + val child = list + .addOrGetVector[VarCharVector](FieldType.nullable(ArrowType.Utf8.INSTANCE)) + .getVector + child.allocateNew() + var position = 0 + values.zipWithIndex.foreach { case (elements, i) => + if (elements == null) { + list.setNull(i) + } else { + list.startNewValue(i) + elements.foreach { e => + child.setSafe(position, e.getBytes(StandardCharsets.UTF_8)) + position += 1 + } + list.endValue(i, elements.length) + } + } + child.setValueCount(position) + list.setValueCount(values.length) + list + } + + /** + * `values.slice(skip, skip + take)` in the shape the C Data import gives a slice of a native + * list array: the offsets are the original ones from `skip` on, over a child that still holds + * the elements of every list. `skip` is a multiple of 8, so the validity bitmap slices on a + * byte boundary. + */ + private def slicedLists( + values: Seq[Seq[String]], + skip: Int, + take: Int, + allocator: BufferAllocator): ListVector = { + val full = lists(values, allocator) + try { + val sliced = + new ListVector("l", allocator, FieldType.nullable(ArrowType.List.INSTANCE), null) + sliced.initializeChildrenFromFields(full.getField.getChildren) + sliced.loadFieldBuffers( + new ArrowFieldNode(take.toLong, values.slice(skip, skip + take).count(_ == null).toLong), + Seq( + full.getValidityBuffer.slice(skip / 8L, (take + 7) / 8L), + full.getOffsetBuffer.slice(skip * 4L, (take + 1) * 4L)).asJava) + val child = full.getDataVector + sliced.getDataVector.loadFieldBuffers( + new ArrowFieldNode(child.getValueCount.toLong, child.getNullCount.toLong), + child.getFieldBuffers) + sliced + } finally { + full.close() + } + } + + /** + * The size of the body `vector` is written with after `withRebasedOffsets`, and the values that + * body holds. + */ + private def writtenBytesAndValues(vector: FieldVector): (Long, Seq[String]) = { + val root = + new VectorSchemaRoot(Seq(vector.getField).asJava, Seq(vector).asJava, vector.getValueCount) + Utils.withRebasedOffsets(root) { written => + val batch = new VectorUnloader(written).getRecordBatch + try { + val column = written.getVector(0) + val values = (0 until column.getValueCount).map(i => String.valueOf(column.getObject(i))) + (batch.computeBodyLength(), values) + } finally { + batch.close() + } + } + } + /** * `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 From a3a85a499b41eee24c067a7cae78b1b7251a2b8d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 28 Oct 2026 08:54:03 -0600 Subject: [PATCH 7/8] fix: rebase sliced vectors before sending them to a Python worker The whole-batch range of withInputBatchRange handed an imported slice's vectors to VectorUnloader as they were, so the worker received the bytes of the rows before it. Rebase them with withRebasedOffsets, as the broadcast and cache writers do. --- docs/source/contributor-guide/ffi.md | 8 +-- .../python/CometArrowPythonRunnerBase.scala | 11 +++++ .../python/CometArrowPythonRunnerSuite.scala | 49 +++++++++++++++++++ 3 files changed, 65 insertions(+), 3 deletions(-) diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index eae5f786ab7..03a9c3a0b9d 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -243,14 +243,16 @@ or sliced booleans reach the JVM misaligned `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 +rows before the slice still at the front of the data buffer, or of the child vector for a list or map. The first slice +of an array starts at zero, but a list or map still has the elements of the later slices in its child. 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 `CachedBatchIpc.serialize` does the same before the in-memory cache writes a batch, so that a cached slice neither -stores nor reports the rows before it ([#6995](https://github.com/apache/datafusion-comet/issues/6995)). Code that -appends or writes imported vectors has to rebase them too. +stores nor reports the rows before it ([#6995](https://github.com/apache/datafusion-comet/issues/6995)). The Arrow +Python runner rebases the same way before it sends a batch to the worker. All three go through +`Utils.withRebasedOffsets`, and code that appends or writes imported vectors has to use it too. ### Ownership and Lifecycle diff --git a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala index cf2ac1ab0c7..ef76fe54a7d 100644 --- a/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala +++ b/spark/src/main/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerBase.scala @@ -730,6 +730,17 @@ private[python] object CometArrowPythonRunnerBase { allocator: BufferAllocator): Unit = { val sourceRoot = new VectorSchemaRoot(sourceVectors.map(_.getField).asJava, sourceVectors.asJava, numRows) + // A slice of a native array reaches the JVM with offsets that start past zero, and the bytes + // before the first of them would be written with the batch. The slice this makes holds its + // own references, so it is closed here and the source root still is not. + Utils.withRebasedOffsets(sourceRoot)(serializeRoot(writeChannel, _, numRows, allocator)) + } + + private def serializeRoot( + writeChannel: WriteChannel, + sourceRoot: VectorSchemaRoot, + numRows: Int, + allocator: BufferAllocator): Unit = { val sourceBatch = new VectorUnloader(sourceRoot).getRecordBatch try { val validityBytes = (numRows.toLong + 7L) / 8L diff --git a/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala b/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala index e2f4c9e788d..cda8fb3834a 100644 --- a/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala +++ b/spark/src/test/spark-4.x/org/apache/spark/sql/execution/python/CometArrowPythonRunnerSuite.scala @@ -39,6 +39,7 @@ import org.apache.arrow.vector.{BaseVariableWidthVector, BigIntVector, FieldVect import org.apache.arrow.vector.complex.{ListVector, MapVector, StructVector} import org.apache.arrow.vector.dictionary.{Dictionary, DictionaryProvider} import org.apache.arrow.vector.ipc.{ArrowStreamReader, ArrowStreamWriter, WriteChannel} +import org.apache.arrow.vector.ipc.message.ArrowFieldNode import org.apache.arrow.vector.types.{IntervalUnit, TimeUnit} import org.apache.arrow.vector.types.pojo.{ArrowType, DictionaryEncoding, Field, FieldType, Schema} import org.apache.spark.{SparkConf, SparkEnv, SparkException, TaskContext, TaskContextImpl} @@ -1152,6 +1153,54 @@ class CometArrowPythonRunnerSuite extends AnyFunSuite with Matchers { } } + test("a whole-batch range of a sliced string vector writes only the sliced values") { + // A slice of a native string array reaches the JVM with offsets that start past zero and the + // bytes of the rows before it at the front of its data buffer. Written as it stands, the + // worker would receive those bytes too, however the range is cut. + val allocator = new RootAllocator(Long.MaxValue) + val values = (0 until 24).map(i => s"value-$i") + val skip = 8 + val full = new VarCharVector("text", allocator) + val sliced = new VarCharVector("text", allocator) + val fresh = new VarCharVector("text", allocator) + try { + full.allocateNew() + values.zipWithIndex.foreach { case (v, i) => + full.setSafe(i, v.getBytes(StandardCharsets.UTF_8)) + } + full.setValueCount(values.length) + val length = values.length - skip + sliced.loadFieldBuffers( + new ArrowFieldNode(length.toLong, 0L), + Seq( + full.getValidityBuffer.slice(skip / 8L, (length + 7) / 8L), + full.getOffsetBuffer.slice(skip * 4L, (length + 1) * 4L), + full.getDataBuffer).asJava) + sliced.getOffsetBuffer.getInt(0) should be > 0 + + fresh.allocateNew() + values.drop(skip).zipWithIndex.foreach { case (v, i) => + fresh.setSafe(i, v.getBytes(StandardCharsets.UTF_8)) + } + fresh.setValueCount(length) + + // The writer allocator is separate because writeInput checks that nothing is left in it. + val writerAllocator = new RootAllocator(Long.MaxValue) + try { + val written = + writeInput(Seq(new CometPlainVector(sliced)), length, 0, 0L, writerAllocator) + val expected = + writeInput(Seq(new CometPlainVector(fresh)), length, 0, 0L, writerAllocator) + written.length shouldBe expected.length + } finally { + writerAllocator.close() + } + } finally { + Seq(fresh, sliced, full).foreach(_.close()) + allocator.close() + } + } + test("the hard Arrow ceiling splits gigabyte values with soft limits raised or disabled") { val input = dictionaryInput( Seq("x"), From c41846e92403308f85cda12cbc55f12bcffdc34a Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 28 Oct 2026 08:54:03 -0600 Subject: [PATCH 8/8] test: keep only the end-to-end test of a sliced aggregate in the cache suite The list and failure-path tests exercised withRebasedOffsets, which UtilsSuite covers directly, so drop them and their helpers. Rename the remaining test so that it does not read as data loss. --- .../comet/exec/CometInMemoryCacheSuite.scala | 168 +----------------- 1 file changed, 3 insertions(+), 165 deletions(-) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 0f45f37f633..59cede8d87e 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -28,14 +28,11 @@ import scala.concurrent.ExecutionContext.Implicits.global import scala.concurrent.duration.DurationInt import scala.jdk.CollectionConverters._ -import org.apache.arrow.c.CDataDictionaryProvider import org.apache.arrow.compression.ZstdCompressionCodec import org.apache.arrow.memory.{ArrowBuf, BufferAllocator} -import org.apache.arrow.vector.{BitVector, FieldVector, FixedSizeBinaryVector, IntVector, LargeVarCharVector, VarBinaryVector, VarCharVector} -import org.apache.arrow.vector.complex.ListVector +import org.apache.arrow.vector.{BitVector, FixedSizeBinaryVector, IntVector, LargeVarCharVector, VarBinaryVector, VarCharVector} import org.apache.arrow.vector.compression.{CompressionCodec, CompressionUtil, NoCompressionCodec} -import org.apache.arrow.vector.ipc.message.ArrowFieldNode -import org.apache.arrow.vector.types.pojo.{ArrowType, FieldType} +import org.apache.arrow.vector.types.pojo.ArrowType import org.apache.spark.{SparkConf, SparkException} import org.apache.spark.CometDriverPlugin import org.apache.spark.sql.{CometTestBase, DataFrame, Observation, QueryTest, Row} @@ -2996,7 +2993,7 @@ class CometInMemoryCacheSuite extends CometTestBase { assert(zstdSize == plainSize, s"the relation size should not depend on the codec: $sizes") } - test("Comet in-memory cache drops the rows before a sliced aggregate output") { + test("Comet in-memory cache does not store the bytes before a sliced aggregate output") { // A grouped aggregate emits its groups as batch-size slices of one batch, and a slice of a // string array reaches the JVM with offsets that start past zero and the rows before it still // at the front of its data buffer. Written as it stands, the k-th slice is cached with the @@ -3053,165 +3050,6 @@ class CometInMemoryCacheSuite extends CometTestBase { } } - /** The bytes the cache writer stores for a batch of `vector` alone. */ - private def payloadSize(vector: FieldVector, codec: CompressionCodec): Long = { - val column = CometVector.getVector(vector, new CDataDictionaryProvider) - val batch = new ColumnarBatch(Array[ColumnVector](column), vector.getValueCount) - CometCachedBatchHelper.payloadSize( - CometCachedBatchHelper.cachedBatch( - CometCachedBatchHelper.serialize(batch, codec, CometArrowAllocator), - batch.numRows())) - } - - /** - * The field node, validity and offset buffers of the rows of `full` from `skip` on, which the C - * Data import hands the JVM for a native slice: the offsets start past zero, over the buffers - * of all of `full`'s rows. - */ - private def nativeSlice( - full: FieldVector, - skip: Int, - nulls: Int): (ArrowFieldNode, Seq[ArrowBuf]) = { - val length = full.getValueCount - skip - ( - new ArrowFieldNode(length.toLong, nulls.toLong), - Seq( - full.getValidityBuffer.slice(skip / 8L, (length + 7) / 8L), - full.getOffsetBuffer.slice(skip * 4L, (length + 1) * 4L))) - } - - private def assertSliceWritesAsOwnValues(sliced: FieldVector, fresh: FieldVector): Unit = { - assert(sliced.getOffsetBuffer.getInt(0) > 0, "the test needs offsets past zero") - val codec = CometCachedBatchHelper.compressionCodec("none") - assert(payloadSize(sliced, codec) == payloadSize(fresh, codec)) - } - - private val sliceSkip = 8 - - private def stringVector(allocator: BufferAllocator, values: Seq[String]): VarCharVector = { - val vector = new VarCharVector("s", allocator) - vector.allocateNew() - values.zipWithIndex.foreach { - case (null, i) => vector.setNull(i) - case (v, i) => vector.setSafe(i, v.getBytes(StandardCharsets.UTF_8)) - } - vector.setValueCount(values.length) - vector - } - - private def slicedStringVector( - allocator: BufferAllocator, - values: Seq[String]): VarCharVector = { - val full = stringVector(allocator, values) - try { - val (node, buffers) = nativeSlice(full, sliceSkip, values.drop(sliceSkip).count(_ == null)) - val vector = new VarCharVector("s", allocator) - vector.loadFieldBuffers(node, (buffers :+ full.getDataBuffer).asJava) - vector - } finally { - full.close() - } - } - - test("Comet in-memory cache writes a sliced list of strings as its own values") { - // A list slice has offsets past zero in two places: its own, and its child's once the child - // is sliced too. Only the list's start past zero here, over a child that holds every element, - // which is how a native slice of a list reaches the JVM. - val values: Seq[Seq[String]] = (0 until 24).map { i => - if (i % 5 == 0) null else (0 until i % 4).map(j => s"element-$i-$j") - } - - def lists(allocator: BufferAllocator, vs: Seq[Seq[String]]): ListVector = { - val list = ListVector.empty("l", allocator) - val child = list - .addOrGetVector[VarCharVector](FieldType.nullable(ArrowType.Utf8.INSTANCE)) - .getVector - child.allocateNew() - var position = 0 - vs.zipWithIndex.foreach { case (elements, i) => - if (elements == null) { - list.setNull(i) - } else { - list.startNewValue(i) - elements.foreach { e => - child.setSafe(position, e.getBytes(StandardCharsets.UTF_8)) - position += 1 - } - list.endValue(i, elements.length) - } - } - child.setValueCount(position) - list.setValueCount(vs.length) - list - } - - def sliced(allocator: BufferAllocator): ListVector = { - val full = lists(allocator, values) - try { - val (node, buffers) = - nativeSlice(full, sliceSkip, values.drop(sliceSkip).count(_ == null)) - val list = - new ListVector("l", allocator, FieldType.nullable(ArrowType.List.INSTANCE), null) - list.initializeChildrenFromFields(full.getField.getChildren) - list.loadFieldBuffers(node, buffers.asJava) - val fullChild = full.getDataVector.asInstanceOf[VarCharVector] - list.getDataVector.loadFieldBuffers( - new ArrowFieldNode(fullChild.getValueCount.toLong, 0L), - Seq( - fullChild.getValidityBuffer, - fullChild.getOffsetBuffer, - fullChild.getDataBuffer).asJava) - list - } finally { - full.close() - } - } - - val allocator = CometArrowAllocator.newChildAllocator("sliced-list-write", 0, Long.MaxValue) - try { - val vector = sliced(allocator) - val fresh = lists(allocator, values.drop(sliceSkip)) - try assertSliceWritesAsOwnValues(vector, fresh) - finally { - fresh.close() - vector.close() - } - } finally { - allocator.close() - } - } - - test("Comet in-memory cache writes a sliced string vector as its own values and releases it") { - // The rebased copy of a sliced vector shares the batch's buffers, so it has to be closed - // whether the write succeeds or fails, or the batch's data outlives the batch. - val values = (0 until 24).map(i => if (i % 5 == 0) null else s"value-$i") - - // Closing a child allocator that still holds a buffer throws, which is the leak check. - val allocator = CometArrowAllocator.newChildAllocator("sliced-write", 0, Long.MaxValue) - try { - val vector = slicedStringVector(allocator, values) - val fresh = stringVector(allocator, values.drop(sliceSkip)) - val failed = slicedStringVector(allocator, values) - try { - assertSliceWritesAsOwnValues(vector, fresh) - - // The write fails after the slice has rebased its offsets into a buffer of its own. - val failing = new FailAfterCompressionCodec(succeedFor = 1) - val thrown = intercept[Exception](payloadSize(failed, failing)) - assert( - causeChain(thrown).exists(t => - Option(t.getMessage).contains(FailAfterCompressionCodec.Message)), - s"a compression failure must surface as itself: $thrown") - } finally { - fresh.close() - vector.close() - failed.close() - } - } finally { - allocator.close() - } - } - test("Comet in-memory cache scans no columns for a row-count-only query") { // SELECT count(*) selects no columns, and the scan must keep it that way. Widening it -- to // the whole cache schema, or to a single placeholder column -- makes the emitted batches