diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 395b5c45909..03a9c3a0b9d 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -243,13 +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 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)). 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/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 81f8e56918a..615df85e6f9 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 @@ -224,32 +224,37 @@ 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) = + // 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) + 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 { - (columnSizes(fields, plainBatch), compressed(plainBatch, codec, allocator)) - } finally { - plainBatch.close() - } - 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() + // Leaves the batch in the state serializeBatches leaves one. The record batch holds its + // own buffers by now, and so does the slice, so this does not touch either. + root.clear() - val out = new ChunkedByteBufferOutputStream(chunkSize, ByteBuffer.allocate) - try { - val channel = new WriteChannel(Channels.newChannel(out)) - MessageSerializer.serialize(channel, recordBatch) + 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() } } 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 975d4b05823..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 @@ -277,20 +277,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 = - if (fieldVectors.forall(offsetsStartAtZero)) root else root.slice(0, root.getRowCount) - 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) => @@ -309,20 +302,44 @@ object Utils extends CometTypeShim with Logging { } /** - * Whether every offset buffer in `vector`, its children included, starts at zero. + * 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[sql] def withRebasedOffsets[T](root: VectorSchemaRoot)( + write: VectorSchemaRoot => T): T = { + val written = + 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 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/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/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala index 2e1f4d56d33..59cede8d87e 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala @@ -2993,6 +2993,63 @@ class CometInMemoryCacheSuite extends CometTestBase { assert(zstdSize == plainSize, s"the relation size should not depend on the codec: $sizes") } + 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 + // bytes of the k - 1 slices before it, and the relation reports them in its size. + val rows = 1000L + 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() + spark + .range(0, rows, 1, 1) + .selectExpr("concat('key-', lpad(cast(id AS string), 12, '0')) AS k") + .write + .parquet(dir.getCanonicalPath) + + def cachedSize(name: String, df: DataFrame): (Long, 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 + 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( + 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. + 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 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") + checkAnswer(spark.table("distinct_cache"), scanned) + } finally { + spark.catalog.clearCache() + } + } + } + } + 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 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 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"),