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/19 17:54:52 UTC
(datafusion-comet) branch gh-readonly-queue/main/pr-6371-2a2aa05404337919a214598f2873dba08e991ddb created (now 0c7c5d69db)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-6371-2a2aa05404337919a214598f2873dba08e991ddb
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
at 0c7c5d69db fix: serialize large-offset string and binary vectors with 32-bit offsets (#6371)
This branch includes the following new commits:
new 0c7c5d69db fix: serialize large-offset string and binary vectors with 32-bit offsets (#6371)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion-comet) 01/01: fix: serialize large-offset string and binary vectors with 32-bit offsets (#6371)
Posted by gi...@apache.org.
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-6371-2a2aa05404337919a214598f2873dba08e991ddb
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 0c7c5d69db8a4b491088a1e8c862e3d4cccedb20
Author: Andy Grove <ag...@apache.org>
AuthorDate: Mon Oct 19 16:19:01 2026 +0000
fix: serialize large-offset string and binary vectors with 32-bit offsets (#6371)
* fix: serialize large-offset string and binary vectors with 32-bit offsets
Utils.getFieldVector rejects LargeVarCharVector and LargeVarBinaryVector,
so Utils.serializeBatches failed on a batch holding either, although a
PyArrow UDF can produce them and Comet reads them elsewhere.
getBatchFieldVectorsWithProviders now writes a VarCharVector or
VarBinaryVector copy of such a column, failing only when its data does
not fit 32-bit offsets. Every reader of the serialized bytes then gets
the Utf8 or Binary type the column's Spark type maps to, and the batches
of one broadcast share a schema whichever offset width each arrived with.
getFieldVector, isSupportedFieldVector and the FFI export are unchanged.
* fix: release the vectors a failed batch assembly made, and test offsets past zero
getBatchFieldVectorsWithProviders keeps the vectors it makes, the 32-bit
copies of large-offset columns and the materialized ConstantColumnVectors,
and closes them if a later column fails. Only the root serializeBatches
builds afterwards would have released them, so a failure left them
allocated from CometArrowAllocator. The batch's own vectors stay with their
owner.
narrowOffsets passes the validity length as a Long, which -Pstrict-warnings
requires, explains why the validity bitmap is copied from bit 0, and points
at #6933 for large-offset children of nested columns.
UtilsSuite covers offsets that start past zero, empty large-offset vectors
with no offset buffer, and the cleanup after a later column fails, and
builds its vectors with fill throughout.
* address review: release every vector a failed write made, and tidy the large-offset tests
---
.../sql/comet/execution/arrow/CachedBatchIpc.scala | 26 +-
.../org/apache/spark/sql/comet/util/Utils.scala | 258 ++++++++++++++-----
.../comet/exec/CometInMemoryCacheSuite.scala | 53 +++-
.../apache/spark/sql/comet/util/UtilsSuite.scala | 279 +++++++++++++++++++--
4 files changed, 520 insertions(+), 96 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 16dc681389..2526d6b651 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
@@ -209,7 +209,7 @@ private[comet] object CachedBatchIpc {
codec: CompressionCodec,
allocator: BufferAllocator,
chunkSize: Int): (ChunkedByteBuffer, Array[Long]) = {
- val (vectors, decoded) = decodeDictionaries(batch, allocator)
+ val (vectors, allocated) = decodeDictionaries(batch, allocator)
try {
val root = new VectorSchemaRoot(vectors.asJava)
// A batch of zero columns carries only a row count, which a VectorSchemaRoot cannot infer
@@ -246,8 +246,8 @@ private[comet] object CachedBatchIpc {
recordBatch.close()
}
} finally {
- // Only the vectors this method allocated. The rest belong to the input batch.
- decoded.foreach(v =>
+ // Only the vectors allocated for this write. The rest belong to the input batch.
+ allocated.foreach(v =>
try v.close()
catch { case NonFatal(_) => () })
}
@@ -641,32 +641,34 @@ private[comet] object CachedBatchIpc {
/**
* Replace every dictionary-encoded column of `batch` with its decoded form.
*
- * Returns the vectors to write and, separately, the ones allocated here so the caller can close
- * exactly those. Columns that needed no decoding are returned as they are and stay owned by
+ * Returns the vectors to write and, separately, the ones allocated for the write, so the caller
+ * can close exactly those: the decoded columns, and any that
+ * [[Utils.getBatchFieldVectorsWithProviders]] allocated, such as a 32-bit copy of a
+ * large-offset column. Columns that needed neither are returned as they are and stay owned by
* `batch`.
*/
private def decodeDictionaries(
batch: ColumnarBatch,
allocator: BufferAllocator): (Seq[FieldVector], Seq[ValueVector]) = {
- val decoded = mutable.ArrayBuffer.empty[ValueVector]
+ val columns = Utils.getBatchFieldVectorsWithProviders(batch)
+ val allocated = mutable.ArrayBuffer.empty[ValueVector]
+ allocated ++= Utils.allocatedFor(batch, columns.map(_._1))
try {
val vectors =
- Utils.getBatchFieldVectorsWithProviders(batch).map { case (vector, providerOpt) =>
+ columns.map { case (vector, providerOpt) =>
if (vector.getField.getDictionary == null) {
vector
} else {
val dictionary = Utils.lookupDictionary(vector, providerOpt)
val plain = DictionaryEncoder.decode(vector, dictionary, allocator)
- decoded += plain
+ allocated += plain
plain.asInstanceOf[FieldVector]
}
}
- (vectors, decoded.toSeq)
+ (vectors, allocated.toSeq)
} catch {
case NonFatal(e) =>
- decoded.foreach(v =>
- try v.close()
- catch { case NonFatal(closeError) => e.addSuppressed(closeError) })
+ Utils.closeAfterFailure(allocated.toSeq, e)
throw e
}
}
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 98b70a1da7..83bb26d09d 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
@@ -24,6 +24,7 @@ import java.nio.ByteBuffer
import java.nio.channels.Channels
import scala.jdk.CollectionConverters._
+import scala.util.control.NonFatal
import org.apache.arrow.c.CDataDictionaryProvider
import org.apache.arrow.vector._
@@ -260,28 +261,36 @@ object Utils extends CometTypeShim with Logging {
val out = new DataOutputStream(codec.compressedOutputStream(cbbos))
val (fieldVectors, batchProviderOpt) = getBatchFieldVectors(batch)
- val root = new VectorSchemaRoot(fieldVectors.asJava)
- if (fieldVectors.isEmpty) {
- // VSR cannot infer rowCount without field vectors
- root.setRowCount(batch.numRows())
- }
- 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 =
- 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()
+ val root = new VectorSchemaRoot(fieldVectors.asJava)
+ if (fieldVectors.isEmpty) {
+ // VSR cannot infer rowCount without field vectors
+ root.setRowCount(batch.numRows())
+ }
+ 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 =
+ 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()
+ }
+ } catch {
+ case NonFatal(e) =>
+ // The batch's own vectors stay with their owner. The ones made for this write are held
+ // by nothing else, and a narrowed copy can be close to 2 GiB.
+ closeAfterFailure(allocatedFor(batch, fieldVectors), e)
+ throw e
}
if (out.size() > 0) {
@@ -479,22 +488,59 @@ object Utils extends CometTypeShim with Logging {
def isArrowBacked(batch: ColumnarBatch): Boolean =
(0 until batch.numCols()).forall { i =>
batch.column(i) match {
- // Not every CometVector can be handed to getFieldVector: a CometPlainVector can wrap a
+ // Not every CometVector is handed out as it stands: a CometPlainVector can wrap a
// LargeVarCharVector or LargeVarBinaryVector (an accelerated mapInArrow returning
- // pa.large_string(), for instance), which it rejects. Answering true for those would
- // send a batch down the direct write path that then fails, so check the vector itself
- // and let the caller convert instead.
+ // pa.large_string(), for instance), which getFieldVector rejects and
+ // getBatchFieldVectorsWithProviders replaces with a 32-bit copy. So check the vector
+ // itself and let the caller convert such a batch.
case v: CometVector => isSupportedFieldVector(v.getValueVector)
case _ => false
}
}
+ /**
+ * Field vectors of `batch` and the dictionaries they refer to, as one provider.
+ *
+ * Some of the vectors may have been allocated here rather than taken from `batch`; see
+ * [[getBatchFieldVectorsWithProviders]]. They are released with the rest when the caller clears
+ * the vectors it wrote, or by [[closeAfterFailure]] if the write fails.
+ */
def getBatchFieldVectors(
batch: ColumnarBatch): (Seq[FieldVector], Option[DictionaryProvider]) = {
val columns = getBatchFieldVectorsWithProviders(batch)
- (columns.map(_._1), combineDictionaryProviders(columns))
+ val vectors = columns.map(_._1)
+ try {
+ (vectors, combineDictionaryProviders(columns))
+ } catch {
+ case NonFatal(e) =>
+ closeAfterFailure(allocatedFor(batch, vectors), e)
+ throw e
+ }
+ }
+
+ /**
+ * The vectors of `vectors`, which were returned for `batch`, that were allocated for it rather
+ * than taken from it: a 32-bit copy of a large-offset column, or a materialized
+ * `ConstantColumnVector`. Whoever holds the returned vectors is the only one who can release
+ * them, and `batch` still owns the rest.
+ */
+ def allocatedFor(batch: ColumnarBatch, vectors: Seq[FieldVector]): Seq[FieldVector] = {
+ val borrowed = (0 until batch.numCols()).flatMap { i =>
+ batch.column(i) match {
+ case v: CometVector => Some(v.getValueVector)
+ case _ => None
+ }
+ }
+ vectors.filterNot(v => borrowed.exists(_ eq v))
}
+ /** Closes `vectors` for a caller that is failing with `failure`, which keeps any error. */
+ def closeAfterFailure(vectors: Seq[ValueVector], failure: Throwable): Unit =
+ vectors.foreach { v =>
+ try v.close()
+ catch { case NonFatal(closeError) => failure.addSuppressed(closeError) }
+ }
+
/**
* The dictionaries every dictionary-encoded column of `columns` refers to, as one provider.
*
@@ -557,44 +603,138 @@ object Utils extends CometTypeShim with Logging {
* what a single stream over every column needs. Comet's cache serializer keeps the pairing
* instead: its payload has no schema message to describe a dictionary encoding, so it decodes
* each such column against the provider that column arrived with.
+ *
+ * Not every returned vector is `batch`'s own. A large-offset column is replaced by a 32-bit
+ * copy and a `ConstantColumnVector` by a materialized one, and the caller releases those, as
+ * [[allocatedFor]] tells them apart. A caller that only unloads or closes the vectors it
+ * decoded itself would leak them. If a later column fails, the ones already made are closed
+ * here before the failure is rethrown.
*/
def getBatchFieldVectorsWithProviders(
batch: ColumnarBatch): Seq[(FieldVector, Option[DictionaryProvider])] = {
val rows = batch.numRows()
- (0 until batch.numCols()).map { index =>
- batch.column(index) match {
- case a: CometVector =>
- val valueVector = a.getValueVector
- val provider =
- if (valueVector.getField.getDictionary != null) Some(a.getDictionaryProvider)
- else None
-
- (getFieldVector(valueVector, "serialize"), provider)
-
- case cv: ConstantColumnVector =>
- // Spark wraps file-source partition columns and other per-batch constants in
- // `ConstantColumnVector`. Materialise to an Arrow vector so the serialisation path
- // doesn't reject the batch. "UTC" is intentional -- see `ConstantColumnVectors`.
- val materialized = ConstantColumnVectors.materialize(
- cv,
- cv.dataType(),
- rows,
- s"_const_$index",
- org.apache.comet.CometArrowAllocator,
- "UTC")
- (materialized, None)
-
- case c =>
- throw new SparkException(
- s"Comet execution only takes Arrow Arrays, but got ${c.getClass}. " +
- "This typically happens when a Comet scan falls back to Spark due to unsupported " +
- "data types (e.g., complex types like structs, arrays, or maps). " +
- "To resolve this, you can: " +
- "(1) enable spark.comet.scan.allowIncompatible=true to use a compatible native " +
- "scan variant, or " +
- "(2) enable spark.comet.convert.parquet.enabled=true to convert Spark Parquet " +
- "data to Arrow format automatically.")
+ // Vectors made here rather than taken from `batch`. The caller releases them with the rest of
+ // what it writes, but if a later column fails it never gets them, so they are closed here.
+ val made = scala.collection.mutable.ArrayBuffer.empty[FieldVector]
+ def keep(vector: FieldVector): FieldVector = {
+ made += vector
+ vector
+ }
+
+ try {
+ (0 until batch.numCols()).map { index =>
+ batch.column(index) match {
+ case a: CometVector =>
+ val valueVector = a.getValueVector
+ val provider =
+ if (valueVector.getField.getDictionary != null) Some(a.getDictionaryProvider)
+ else None
+
+ val fieldVector = valueVector match {
+ case v: LargeVarCharVector => keep(narrowOffsets(v, ArrowType.Utf8.INSTANCE))
+ case v: LargeVarBinaryVector => keep(narrowOffsets(v, ArrowType.Binary.INSTANCE))
+ case v => getFieldVector(v, "serialize")
+ }
+ (fieldVector, provider)
+
+ case cv: ConstantColumnVector =>
+ // Spark wraps file-source partition columns and other per-batch constants in
+ // `ConstantColumnVector`. Materialise to an Arrow vector so the serialisation path
+ // doesn't reject the batch. "UTC" is intentional -- see `ConstantColumnVectors`.
+ val materialized = keep(
+ ConstantColumnVectors.materialize(
+ cv,
+ cv.dataType(),
+ rows,
+ s"_const_$index",
+ org.apache.comet.CometArrowAllocator,
+ "UTC"))
+ (materialized, None)
+
+ case c =>
+ throw new SparkException(
+ s"Comet execution only takes Arrow Arrays, but got ${c.getClass}. " +
+ "This typically happens when a Comet scan falls back to Spark due to " +
+ "unsupported data types (e.g., complex types like structs, arrays, or maps). " +
+ "To resolve this, you can: " +
+ "(1) enable spark.comet.scan.allowIncompatible=true to use a compatible native " +
+ "scan variant, or " +
+ "(2) enable spark.comet.convert.parquet.enabled=true to convert Spark Parquet " +
+ "data to Arrow format automatically.")
+ }
+ }
+ } catch {
+ case NonFatal(e) =>
+ closeAfterFailure(made.toSeq, e)
+ throw e
+ }
+ }
+
+ /**
+ * A copy of `large` with 32-bit offsets, as a vector of `narrowType`.
+ *
+ * A PyArrow UDF can hand back a `large_string` or `large_binary` column, which Comet reads but
+ * [[getFieldVector]] does not accept. The column is still `StringType` or `BinaryType`, which
+ * Comet otherwise carries as `Utf8` or `Binary`, so it is written in that form. Whatever reads
+ * the bytes back then gets the type it plans for, and the batches of one broadcast share a
+ * schema whichever offset width each arrived with: [[coalesceBroadcastBatches]] appends them to
+ * a root built from the first one's schema, and a native plan reading them takes its input
+ * stream's schema from the first one too. Only a column holding more data than 32-bit offsets
+ * can address is refused.
+ *
+ * Only top-level columns are narrowed. A `large_string` or `large_binary` child of a struct,
+ * list or map column passes [[getFieldVector]] with its parent and keeps its 64-bit offsets
+ * (#6933).
+ *
+ * Like a materialized `ConstantColumnVector`, the copy is released when the caller clears the
+ * vectors it wrote. `large` itself stays with its owner.
+ */
+ private def narrowOffsets(
+ large: BaseLargeVariableWidthVector,
+ narrowType: ArrowType): FieldVector = {
+ val numValues = large.getValueCount
+ val largeOffsets = large.getOffsetBuffer
+ val largeWidth = BaseLargeVariableWidthVector.OFFSET_WIDTH.toLong
+ // An empty vector need not have an offset buffer to read.
+ val start = if (numValues == 0) 0L else largeOffsets.getLong(0)
+ val dataLength =
+ if (numValues == 0) 0L else largeOffsets.getLong(numValues * largeWidth) - start
+ if (dataLength > Int.MaxValue) {
+ throw new SparkException(
+ s"Column ${large.getField.getName} holds $dataLength bytes of data, more than 32-bit " +
+ "offsets can address")
+ }
+
+ val fieldType = large.getField.getFieldType
+ val narrow = new Field(
+ large.getField.getName,
+ new FieldType(fieldType.isNullable, narrowType, null, fieldType.getMetadata),
+ null)
+ .createVector(org.apache.comet.CometArrowAllocator)
+ .asInstanceOf[BaseVariableWidthVector]
+ try {
+ narrow.allocateNew(dataLength, numValues)
+ if (numValues > 0) {
+ // An Arrow Java vector has no array offset, so validity bit i belongs to value i even when
+ // the offsets start past zero, and the bitmap is copied from bit 0.
+ val validityBytes = BitVectorHelper.getValidityBufferSize(numValues).toLong
+ narrow.getValidityBuffer.setBytes(0, large.getValidityBuffer, 0, validityBytes)
+ val offsets = narrow.getOffsetBuffer
+ (0 to numValues).foreach { i =>
+ val offset = largeOffsets.getLong(i * largeWidth) - start
+ offsets.setInt(i.toLong * BaseVariableWidthVector.OFFSET_WIDTH, offset.toInt)
+ }
+ narrow.getDataBuffer.setBytes(0, large.getDataBuffer, start, dataLength)
}
+ // setValueCount overwrites every value after lastSet with an empty one, so mark them all
+ // as written.
+ narrow.setLastSet(numValues - 1)
+ narrow.setValueCount(numValues)
+ narrow
+ } catch {
+ case NonFatal(e) =>
+ narrow.close()
+ throw e
}
}
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 4dd68e6b5f..e25c24341c 100644
--- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala
@@ -30,7 +30,7 @@ import scala.jdk.CollectionConverters._
import org.apache.arrow.compression.ZstdCompressionCodec
import org.apache.arrow.memory.{ArrowBuf, BufferAllocator}
-import org.apache.arrow.vector.{BitVector, FixedSizeBinaryVector, IntVector, VarBinaryVector, VarCharVector}
+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.types.pojo.ArrowType
import org.apache.spark.{SparkConf, SparkException}
@@ -3039,6 +3039,57 @@ class CometInMemoryCacheSuite extends CometTestBase {
}
}
+ test("Comet in-memory cache releases the 32-bit copy it makes of a large-offset column") {
+ // A large-offset column is written as a copy with 32-bit offsets, which nothing outside the
+ // writer ever holds. It has to be closed whether the write succeeds or fails, and a copy can
+ // be close to 2 GiB.
+ val rows = 16
+ // Returns the failure of the write, if any, once the batch's own vectors are closed and the
+ // allocator is back where it started.
+ def write(codec: CompressionCodec): Option[Throwable] = {
+ val baseline = CometArrowAllocator.getAllocatedMemory
+ val ints = new IntVector("i", CometArrowAllocator)
+ val large = new LargeVarCharVector("s", CometArrowAllocator)
+ val failure =
+ try {
+ ints.allocateNew(rows)
+ (0 until rows).foreach(i => ints.set(i, i))
+ ints.setValueCount(rows)
+ large.allocateNew(rows)
+ (0 until rows).foreach(i => large.setSafe(i, s"value_$i".getBytes("UTF-8")))
+ large.setValueCount(rows)
+
+ val batch = new ColumnarBatch(
+ Array[ColumnVector](new CometPlainVector(ints), new CometPlainVector(large)),
+ rows)
+ try {
+ CometCachedBatchHelper.serialize(batch, codec, CometArrowAllocator)
+ None
+ } catch {
+ case e: Exception => Some(e)
+ }
+ } finally {
+ ints.close()
+ large.close()
+ }
+ assert(
+ CometArrowAllocator.getAllocatedMemory == baseline,
+ "the copy of the large-offset column must be released")
+ failure
+ }
+
+ assert(write(new ZstdCompressionCodec(1)).isEmpty)
+
+ // The int column's two buffers compress, then the copy's first one fails.
+ val codec = new FailAfterCompressionCodec(succeedFor = 2)
+ val thrown = write(codec).getOrElse(fail("the injected failure must surface"))
+ assert(
+ causeChain(thrown).exists(t =>
+ Option(t.getMessage).contains(FailAfterCompressionCodec.Message)),
+ s"a compression failure must surface as itself: $thrown")
+ assert(codec.compressed == 2, "the failure must come after some buffers were compressed")
+ }
+
test("Comet in-memory cache keeps no compressed bytes alive behind a buffer stored raw") {
// Arrow stores a buffer verbatim when compressing it would not make it smaller -- a bitmap of
// random bits, say -- and decompressing hands it back as a slice of whatever holds it. Copied
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 54a55ccde8..483309b0a9 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,11 +26,14 @@ 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.{BaseLargeVariableWidthVector, BitVectorHelper, IntVector, LargeVarBinaryVector, LargeVarCharVector, ValueVector, VarBinaryVector, VarCharVector, VariableWidthFieldVector}
+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}
+import org.apache.spark.SparkException
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.types.{BinaryType, IntegerType, StringType, StructField, StructType, TimestampType}
import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
import org.apache.spark.util.io.ChunkedByteBuffer
@@ -233,29 +236,18 @@ class UtilsSuite extends CometTestBase {
test("isArrowBacked rejects large-offset Arrow vectors") {
// A CometPlainVector can wrap a LargeVarCharVector or LargeVarBinaryVector -- an accelerated
- // mapInArrow returning pa.large_string() produces one -- but getFieldVector rejects both. If
- // isArrowBacked accepted them, a caller would take the direct write path and then fail, so it
- // must report false and let the caller convert the batch instead.
- val numRows = 2
- Seq[org.apache.arrow.vector.FieldVector](
- {
- val v = new org.apache.arrow.vector.LargeVarCharVector("s", CometArrowAllocator)
- v.allocateNew(numRows)
- v.setSafe(0, "hello".getBytes("UTF-8"))
- v.setSafe(1, "world".getBytes("UTF-8"))
- v.setValueCount(numRows)
- v
- }, {
- val v = new org.apache.arrow.vector.LargeVarBinaryVector("b", CometArrowAllocator)
- v.allocateNew(numRows)
- v.setSafe(0, "hello".getBytes("UTF-8"))
- v.setSafe(1, "world".getBytes("UTF-8"))
- v.setValueCount(numRows)
- v
- }).foreach { vector =>
+ // mapInArrow returning pa.large_string() produces one -- but getFieldVector rejects both, and
+ // getBatchFieldVectors writes a 32-bit copy of them rather than the vector itself. A caller
+ // taking the direct write path expects the batch's own vectors, so isArrowBacked must report
+ // false and let the caller convert the batch instead. The cache's direct write path also asks
+ // CachedBatchIpc.matchesReaderLayout, which refuses LargeUtf8 against the reader's Utf8, so
+ // either check alone keeps such a batch off it.
+ val values = Seq("hello", "world")
+ Seq[ValueVector](
+ fill(new LargeVarCharVector("s", CometArrowAllocator), values),
+ fill(new LargeVarBinaryVector("b", CometArrowAllocator), values)).foreach { vector =>
try {
- val col = CometVector.getVector(vector, new CDataDictionaryProvider)
- val batch = new ColumnarBatch(Array[ColumnVector](col), numRows)
+ val batch = new ColumnarBatch(Array(cometVector(vector)), values.length)
assert(
!Utils.isArrowBacked(batch),
s"${vector.getClass.getSimpleName} must not be reported as directly writable")
@@ -265,6 +257,245 @@ class UtilsSuite extends CometTestBase {
}
}
+ test("serializeBatches writes large-offset string and binary vectors with 32-bit offsets") {
+ // A PyArrow UDF returning pa.large_string() or pa.large_binary() hands Comet a
+ // LargeVarCharVector or LargeVarBinaryVector. The column is still StringType or BinaryType,
+ // so it is written as the Utf8 or Binary vector every reader of the bytes expects for those.
+ val values = Seq("hello", null, "", "wörld")
+ val strings = fill(new LargeVarCharVector("s", CometArrowAllocator), values)
+ val binaries = fill(new LargeVarBinaryVector("b", CometArrowAllocator), values)
+ try {
+ val batch =
+ new ColumnarBatch(Array(cometVector(strings), cometVector(binaries)), values.length)
+ assert(
+ roundTrip(batch) == Seq(
+ classOf[VarCharVector] -> values,
+ classOf[VarBinaryVector] -> values))
+ } finally {
+ strings.close()
+ binaries.close()
+ }
+ }
+
+ test("serializeBatches rebases large-offset vectors whose offsets start past zero") {
+ // Offsets need not start at zero. A slice imported over the C Data interface keeps its
+ // parent's offsets and data buffer, so they start past bytes that no value covers. The copy
+ // holds only the bytes the values cover, with offsets counted from the first of them.
+ val values = Seq("hello", null, "wörld")
+ val strings = startingPast("xx", new LargeVarCharVector("s", CometArrowAllocator), values)
+ val binaries = startingPast("xx", new LargeVarBinaryVector("b", CometArrowAllocator), values)
+ try {
+ assert(strings.getOffsetBuffer.getLong(0) == 2)
+ assert(binaries.getOffsetBuffer.getLong(0) == 2)
+ val batch =
+ new ColumnarBatch(Array(cometVector(strings), cometVector(binaries)), values.length)
+ assert(
+ roundTrip(batch) == Seq(
+ classOf[VarCharVector] -> values,
+ classOf[VarBinaryVector] -> values))
+ } finally {
+ strings.close()
+ binaries.close()
+ }
+ }
+
+ test("serializeBatches writes empty large-offset vectors with 32-bit offsets") {
+ // A vector that holds no values need not have an offset buffer, so none is read.
+ val strings = new LargeVarCharVector("s", CometArrowAllocator)
+ val binaries = new LargeVarBinaryVector("b", CometArrowAllocator)
+ try {
+ strings.setValueCount(0)
+ binaries.setValueCount(0)
+ assert(strings.getOffsetBuffer.capacity == 0)
+ val batch = new ColumnarBatch(Array(cometVector(strings), cometVector(binaries)), 0)
+ assert(
+ roundTrip(batch) == Seq(
+ classOf[VarCharVector] -> Seq.empty,
+ classOf[VarBinaryVector] -> Seq.empty))
+ } finally {
+ strings.close()
+ binaries.close()
+ }
+ }
+
+ test("coalesceBroadcastBatches coalesces batches of mixed offset widths") {
+ // The coalescer appends every batch to a root built from the first one's schema, and leaves
+ // the broadcast uncoalesced when a batch does not match it. A producer can hand back
+ // large_string for one batch and string for the next, so both must serialize alike.
+ val large = fill(new LargeVarCharVector("s", CometArrowAllocator), Seq("a", null))
+ val regular = fill(new VarCharVector("s", CometArrowAllocator), Seq("b", "c"))
+ try {
+ val batches = Seq(large, regular).map(v => new ColumnarBatch(Array(cometVector(v)), 2))
+ val bufs = Utils.serializeBatches(batches.iterator).map(_._2).toSeq.iterator
+ val (coalesced, batchCount, totalRows) = Utils.coalesceBroadcastBatches(bufs)
+ // A batch count of zero is the fallback to the uncoalesced buffers.
+ assert(batchCount == 2)
+ assert(totalRows == 4)
+
+ val it = coalesced.iterator.flatMap(Utils.decodeBatches(_, "test"))
+ val out = it.next()
+ val got = stringValues(out.column(0), out.numRows())
+ assert(!it.hasNext)
+ assert(got == Seq("a", null, "b", "c"))
+ } finally {
+ large.close()
+ regular.close()
+ }
+ }
+
+ test("serializeBatches releases the vectors it made when a later column fails") {
+ // A materialized constant and a 32-bit copy of a large-offset column are both made before the
+ // last column is refused. Nothing outside serializeBatches ever holds them, so it has to
+ // close them itself, while the batch's own vectors stay with their owner.
+ val before = CometArrowAllocator.getAllocatedMemory
+ val constant = new ConstantColumnVector(1, IntegerType)
+ constant.setInt(42)
+ val small = fill(new LargeVarCharVector("small", CometArrowAllocator), Seq("a"))
+ val tooBig = tooBigForIntOffsets("big")
+ try {
+ val batch = new ColumnarBatch(Array(constant, cometVector(small), cometVector(tooBig)), 1)
+ val e = intercept[SparkException](Utils.serializeBatches(Iterator(batch)).next())
+ assert(e.getMessage.contains("more than 32-bit offsets can address"), e.getMessage)
+ assert(small.getValueCount == 1 && small.getObject(0).toString == "a")
+ } finally {
+ constant.close()
+ small.close()
+ tooBig.close()
+ }
+ assert(CometArrowAllocator.getAllocatedMemory == before)
+ }
+
+ test("serializeBatches releases the vectors it made when the dictionaries of a batch clash") {
+ // The vectors are all made before the dictionaries of the columns are combined, which is where
+ // two columns sharing an ID without sharing a dictionary are refused. By then the caller has
+ // not got them either, and a copy of a large-offset column can be close to 2 GiB.
+ val before = CometArrowAllocator.getAllocatedMemory
+ val large = fill(new LargeVarCharVector("large", CometArrowAllocator), Seq("a"))
+ val (first, firstProvider) = dictionaryEncoded("first", 7L, "x")
+ val (second, secondProvider) = dictionaryEncoded("second", 7L, "y")
+ try {
+ val batch = new ColumnarBatch(
+ Array(
+ cometVector(large),
+ CometVector.getVector(first, firstProvider),
+ CometVector.getVector(second, secondProvider)),
+ 1)
+ val e = intercept[SparkException](Utils.serializeBatches(Iterator(batch)).next())
+ assert(e.getMessage.contains("different dictionaries under ID 7"), e.getMessage)
+ assert(large.getValueCount == 1 && large.getObject(0).toString == "a")
+ } finally {
+ large.close()
+ Seq(first, second).foreach(_.close())
+ Seq(firstProvider, secondProvider).foreach(_.lookup(7L).getVector.close())
+ }
+ assert(CometArrowAllocator.getAllocatedMemory == before)
+ }
+
+ /** A dictionary-encoded index column holding the one entry of a dictionary holding `value`. */
+ private def dictionaryEncoded(
+ name: String,
+ id: Long,
+ value: String): (IntVector, DictionaryProvider) = {
+ val encoding = new DictionaryEncoding(id, false, null)
+ val dictionary =
+ new Dictionary(fill(new VarCharVector("dict", CometArrowAllocator), Seq(value)), encoding)
+ val indices = new IntVector(
+ new Field(name, new FieldType(true, new ArrowType.Int(32, true), encoding), null),
+ CometArrowAllocator)
+ indices.allocateNew(1)
+ indices.set(0, 0)
+ indices.setValueCount(1)
+ (indices, new DictionaryProvider.MapDictionaryProvider(dictionary))
+ }
+
+ /** Writes `values` into `vector`, leaving a null entry unset so that it reads back as null. */
+ private def fill[V <: VariableWidthFieldVector](vector: V, values: Seq[String]): V = {
+ vector.allocateNew(values.length)
+ values.zipWithIndex.foreach { case (value, i) =>
+ if (value != null) vector.setSafe(i, value.getBytes(StandardCharsets.UTF_8))
+ }
+ vector.setValueCount(values.length)
+ vector
+ }
+
+ /**
+ * Loads `values` into `vector` the way an imported slice holds them: the data buffer starts
+ * with the bytes of `skipped`, which no value covers, so the offsets start past zero.
+ */
+ private def startingPast[V <: BaseLargeVariableWidthVector](
+ skipped: String,
+ vector: V,
+ values: Seq[String]): V = {
+ val skippedBytes = skipped.getBytes(StandardCharsets.UTF_8)
+ val valueBytes =
+ values.map(v => if (v == null) Array.emptyByteArray else v.getBytes(StandardCharsets.UTF_8))
+ val offsets = valueBytes.scanLeft(skippedBytes.length.toLong)(_ + _.length)
+ val width = BaseLargeVariableWidthVector.OFFSET_WIDTH.toLong
+
+ val validity =
+ CometArrowAllocator.buffer(BitVectorHelper.getValidityBufferSize(values.length).toLong)
+ val offsetBuffer = CometArrowAllocator.buffer(offsets.length * width)
+ val data = CometArrowAllocator.buffer(offsets.last)
+ val buffers = Seq(validity, offsetBuffer, data)
+ try {
+ validity.setZero(0, validity.capacity)
+ values.indices.filter(values(_) != null).foreach { i =>
+ BitVectorHelper.setBit(validity, i.toLong)
+ }
+ offsets.zipWithIndex.foreach { case (offset, i) => offsetBuffer.setLong(i * width, offset) }
+ data.setBytes(0, skippedBytes ++ valueBytes.flatten)
+ vector.loadFieldBuffers(
+ new ArrowFieldNode(values.length.toLong, values.count(_ == null).toLong),
+ buffers.asJava)
+ vector
+ } finally {
+ // The vector takes references of its own.
+ buffers.foreach(_.close())
+ }
+ }
+
+ /** A one-value vector whose offsets claim more data than 32-bit offsets can address. */
+ private def tooBigForIntOffsets(name: String): LargeVarCharVector = {
+ val vector = fill(new LargeVarCharVector(name, CometArrowAllocator), Seq("a"))
+ vector.getOffsetBuffer.setLong(BaseLargeVariableWidthVector.OFFSET_WIDTH, Int.MaxValue + 1L)
+ vector
+ }
+
+ private def cometVector(vector: ValueVector): ColumnVector =
+ CometVector.getVector(vector, new CDataDictionaryProvider)
+
+ /**
+ * Serializes `batch` on its own and decodes it as the broadcast consumers do, returning each
+ * column's vector class and values.
+ */
+ private def roundTrip(batch: ColumnarBatch): Seq[(Class[_], Seq[String])] = {
+ val (rowCount, buf) = Utils.serializeBatches(Iterator(batch)).next()
+ assert(rowCount == batch.numRows())
+
+ // Everything is read before the iterator moves on, which releases the batch it returned.
+ val it = Utils.decodeBatches(buf, "test")
+ val out = it.next()
+ val columns = (0 until out.numCols()).map { i =>
+ val column = out.column(i)
+ val vectorClass = column.asInstanceOf[CometVector].getValueVector.getClass
+ (vectorClass, stringValues(column, out.numRows()))
+ }
+ assert(!it.hasNext)
+ columns
+ }
+
+ /** The values of a string or binary column, with binary values read as UTF-8. */
+ private def stringValues(column: ColumnVector, numRows: Int): Seq[String] =
+ (0 until numRows).map { i =>
+ if (column.isNullAt(i)) {
+ null
+ } else if (column.dataType == BinaryType) {
+ new String(column.getBinary(i), StandardCharsets.UTF_8)
+ } else {
+ column.getUTF8String(i).toString
+ }
+ }
+
private def strings(
values: Seq[String],
allocator: BufferAllocator = CometArrowAllocator): VarCharVector = {
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org