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 975d4b05823..c8b34fc51c3 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 @@ -260,15 +260,25 @@ object Utils extends CometTypeShim with Logging { * the output stream */ def serializeBatches(batches: Iterator[ColumnarBatch]): Iterator[(Long, ChunkedByteBuffer)] = { + lazy val codec = CompressionCodec.createCodec(SparkEnv.get.conf) batches.map { batch => val dictionaryProvider: CDataDictionaryProvider = new CDataDictionaryProvider - val codec = CompressionCodec.createCodec(SparkEnv.get.conf) - val cbbos = new ChunkedByteBufferOutputStream(1024 * 1024, ByteBuffer.allocate) - val out = new DataOutputStream(codec.compressedOutputStream(cbbos)) - val (fieldVectors, batchProviderOpt) = getBatchFieldVectors(batch) try { + // Tiny shuffle batches otherwise allocate and zero a full 1 MiB heap chunk each. Keep + // larger chunks for large inputs, where a small fixed chunk would fragment the output. + val dictionaryBytes = batchProviderOpt + .map { provider => + provider.getDictionaryIds.asScala.foldLeft(0L) { (bytes, id) => + bytes + provider.lookup(id).getVector.getBufferSize + } + } + .getOrElse(0L) + val batchBytes = fieldVectors.foldLeft(dictionaryBytes)(_ + _.getBufferSize) + val chunkSize = math.max(4096L, math.min(1024 * 1024L, batchBytes)).toInt + val cbbos = new ChunkedByteBufferOutputStream(chunkSize, ByteBuffer.allocate) + val out = new DataOutputStream(codec.compressedOutputStream(cbbos)) val root = new VectorSchemaRoot(fieldVectors.asJava) if (fieldVectors.isEmpty) { // VSR cannot infer rowCount without field vectors @@ -292,6 +302,12 @@ object Utils extends CometTypeShim with Logging { // the write fails, as it does when a cancelled broadcast interrupts the task. if (written ne root) written.close() } + + if (out.size() > 0) { + (batch.numRows().toLong, cbbos.toChunkedByteBuffer) + } else { + (batch.numRows().toLong, new ChunkedByteBuffer(Array.empty[ByteBuffer])) + } } catch { case NonFatal(e) => // The batch's own vectors stay with their owner. The ones made for this write are held @@ -299,12 +315,6 @@ object Utils extends CometTypeShim with Logging { closeAfterFailure(allocatedFor(batch, fieldVectors), e) throw e } - - if (out.size() > 0) { - (batch.numRows().toLong, cbbos.toChunkedByteBuffer) - } else { - (batch.numRows().toLong, new ChunkedByteBuffer(Array.empty[ByteBuffer])) - } } } 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..6a4f4a32af6 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 @@ -43,6 +43,38 @@ import org.apache.comet.vector.CometVector class UtilsSuite extends CometTestBase { + test("serializeBatches round trips schema metadata larger than the output chunk estimate") { + val random = new scala.util.Random(42) + // Keep vector data below 4 KiB while the varied field names exceed it even after compression. + val vectors = (0 until 256).map { i => + val vector = new IntVector(random.alphanumeric.take(64).mkString, CometArrowAllocator) + vector.allocateNew(1) + vector.set(0, i) + vector.setValueCount(1) + vector + } + val columns = vectors.map { vector => + CometVector.getVector(vector, new CDataDictionaryProvider): ColumnVector + } + val batch = new ColumnarBatch(columns.toArray, 1) + try { + assert(vectors.map(_.getBufferSize).sum < 4096) + val (rows, buffer) = Utils.serializeBatches(Iterator(batch)).next() + try { + assert(rows == 1) + assert(buffer.size > 4096) + assert(buffer.getChunks.length > 1) + assert(buffer.getChunks.head.remaining() == 4096) + val decoded = Utils.decodeBatches(buffer, "schema-chunks") + val result = decoded.next() + assert(result.numRows() == 1) + assert(result.numCols() == vectors.length) + vectors.indices.foreach(i => assert(result.column(i).getInt(0) == i)) + assert(!decoded.hasNext) + } finally buffer.dispose() + } finally batch.close() + } + test("serializeBatches preserves row count for a zero-column batch") { val numRows = 5 val batch = new ColumnarBatch(Array.empty[ColumnVector], numRows)