andygrove commented on code in PR #7225: URL: https://github.com/apache/datafusion-comet/pull/7225#discussion_r4402129894
########## 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. Review Comment: You're right, the whole-batch branch of `withInputBatchRange` handed an imported slice's vectors to `VectorUnloader` as they were. I confirmed it with a test before changing anything: a sliced string vector wrote 688 bytes against 632 for the same values in a vector of its own. `serializeBatch` now goes through `withRebasedOffsets`, which is `private[sql]` so the Python runner can call it, and this paragraph names the runner as the third writer. The fix and the test are in a3a85a499b. ########## spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala: ########## @@ -2758,6 +2761,222 @@ 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") { Review Comment: Renamed to "does not store the bytes before a sliced aggregate output" in c41846e924. ########## spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala: ########## @@ -2758,6 +2761,222 @@ 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 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() + } + } + } + } + + /** 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") { Review Comment: Moved to `UtilsSuite` as "withRebasedOffsets writes a sliced list as its own values", which checks the written bytes and values of list slices against `withRebasedOffsets` directly. `nativeSlice`, `slicedStringVector`, `stringVector` and `assertSliceWritesAsOwnValues` are gone from `CometInMemoryCacheSuite`, and the comment with the missing word went with the test. The move is in b4a0ae1d00 and the removal in c41846e924. ########## spark/src/test/scala/org/apache/comet/exec/CometInMemoryCacheSuite.scala: ########## @@ -2758,6 +2761,222 @@ 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 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() + } + } + } + } + + /** 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") { Review Comment: Dropped, with the helpers it needed. The `finally` in `withRebasedOffsets` is covered by the `UtilsSuite` test you pointed at, and the end-to-end test covers the success half. Removed in c41846e924. ########## spark/src/main/scala/org/apache/spark/sql/comet/util/Utils.scala: ########## @@ -302,6 +295,22 @@ 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. + */ + private[comet] def withRebasedOffsets[T](root: VectorSchemaRoot)( + write: VectorSchemaRoot => T): T = { + val written = + if (root.getFieldVectors.asScala.forall(offsetsStartAtZero)) root Review Comment: Confirmed, and it was easy to miss because the first slice looks like an ordinary list. In a new `UtilsSuite` test, a slice of the first 8 of 24 lists wrote 536 bytes against 232 for the same lists in a vector of their own, because its child still held the elements of the later slices. A list or map is now rebased when its last offset is short of its child's value count as well as when its first offset is past zero, and that test passes for the first slice and for the later ones. Strings keep the first-offset check, since their data buffer stops at the last offset. I renamed `offsetsStartAtZero` to `holdsOnlyItsOwnValues` because it no longer only looks at the first offset. The change is in b4a0ae1d00. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
