viirya commented on code in PR #7225:
URL: https://github.com/apache/datafusion-comet/pull/7225#discussion_r4372554728
##########
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:
"drops the rows before a sliced aggregate output" reads as if the cache
loses rows. What the test checks is that the bytes of the earlier rows are not
stored. How about "does not store the bytes before a sliced aggregate output"?
##########
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:
There is one more writer that doesn't follow this yet. In
`CometArrowPythonRunnerBase.withInputBatchRange`, the `offset == 0 && length ==
numRows` branch passes the batch's own vectors to `serializeBatch`, which
unloads them with `VectorUnloader`. So when a sliced aggregate output feeds
`mapInArrow` or `mapInPandas`, the bytes of the earlier slices go to the Python
worker too. The values still come out right because pyarrow honors the offsets,
but it is the same inflation this PR fixes for the cache. The partial-range
branch is fine, since `column.slice` goes through `TransferPair` and rebases.
Could that branch use `withRebasedOffsets` as well? It would need wider
visibility than `private[comet]`, because the runner is in
`org.apache.spark.sql.execution.python`. If you'd rather keep this PR to the
cache, could you file an issue and link it here?
##########
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:
I think this test can go. The success half is covered end to end by the
aggregate test above. The failure half exercises the `finally` in
`withRebasedOffsets`, which is now shared and already covered by `UtilsSuite`'s
"serializeBatches closes the slice it writes when the write fails". Removing
it, plus moving the list test, would also let `nativeSlice`,
`slicedStringVector`, `stringVector` and `assertSliceWritesAsOwnValues` go,
which removes about half of the new test code.
##########
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:
This is the test that adds new coverage, the rebase recursing into the
list's child. That is a property of `withRebasedOffsets` rather than of the
cache, though, so it would fit better in `UtilsSuite`, tested against the
helper directly. That would cover the broadcast path too, and `slicedStrings`
already lives there, so `nativeSlice` wouldn't need to become a third copy of
the native-slice construction.
Also, the comment on L2883 seems to be missing a word. I think it means
"Only the list's offsets start past zero here".
--
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]
