You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "andygrove (via GitHub)" <gi...@apache.org> on 2026/10/14 15:08:42 UTC
[PR] perf: combine small columnar batches before caching them [datafusion-comet]
andygrove opened a new pull request, #6988:
URL: https://github.com/apache/datafusion-comet/pull/6988
## Which issue does this PR close?
Closes #6950.
## Rationale for this change
The columnar cache write stores each input batch as one cached batch, and nothing bounds how small those batches are. A shuffle read returns at least one batch per map output block, so a relation cached straight from a shuffle over many map tasks holds thousands of small batches. Each carries its own message metadata, compression framing and statistics row, and every read of the cache pays for the batch count. #6950 has the details.
## What changes are included in this PR?
- A new `CacheBatchCoalescer` sits in front of the columnar write in `ArrowCachedBatchSerializer.convertColumnarBatchToCachedBatch`. It passes a batch through when two things hold: the batch is Arrow-backed and laid out the way the reader reads it, which is the test the write already used to store a batch as it stands, and it holds at least half of `spark.sql.inMemoryColumnarStorage.batchSize` rows. Full native batches of 8192 rows therefore still reach the writer without a copy.
- Every other batch is copied with `ArrowWriter.writeColumns` into batches of `batchSize` rows, the size the row path cuts at, and is split where it crosses that size. A small batch is worth the copy. A batch of Spark's own column vectors, or one laid out differently from the reader, was already being copied into Arrow one batch at a time, so combining it adds no copy. The copy reads each column through the `ColumnVector` interface, so it handles dictionary-encoded and nested columns, and offsets that don't start at zero, as a native producer's sliced arrays can have. Arrow's `VectorAppender` assumes offsets start at zero, which is why this doesn't use it.
- Row order is kept, so a sorted cache stays sorted: the rows copied so far are emitted before a batch that passes through. Each batch is copied before the input advances, because a native producer such as `CometExecIterator` releases the batch it returned last when it is asked for the next. The coalescer owns the batches it builds and releases each one when the next is requested, and anything still open when the task completes.
- The row path is unchanged, since it already cuts at `batchSize`.
- The in-memory cache guide's "Storage format" section describes the batch sizes.
## How are these changes tested?
Three new tests in `CometInMemoryCacheSuite`:
- "combines small columnar batches in order" feeds the serializer a mix of Arrow and Spark-vector batches at a batch size of 1000. The input releases each batch when it is asked for the next, as a native producer does. The test checks the exact cached batch sizes (600, 9000, 700, 1000, 1000, 1000 and 100) and every row in order.
- "combines the small batches of a shuffle read" caches a 20,000-row relation with nested columns straight from a shuffle over 20 map tasks into 2 partitions. Without this change it caches 40 batches of about 500 rows. With it there are at most 4, and the rows match Spark's.
- "keeps a sorted cache sorted when it combines its batches" caches a sort whose native batches hold 100 rows and checks that each partition reads back in order. Without this change it caches 201 batches of 100 rows.
Each of five mutants fails at least one of these tests: removing the coalescer, passing through at the full batch size instead of half of it, not splitting at the batch size, emitting a pass-through batch ahead of the rows combined before it, and passing Spark-vector batches through.
`CometInMemoryCacheSuite`, `CometInMemoryCachePruningSuite`, the three Kryo suites and `CachedBatchRowIteratorSuite` pass locally on Spark 3.4, 3.5, 4.0, 4.1 and 4.2. Scalafix and the strict-warnings build pass on 3.5.
I measured #6950's shape with a probe that isn't part of this PR. It reads 5M rows of four columns (bigint, int, double and string) from 200 Parquet files as 200 splits, and caches them with `repartition(50, $"a")`. Both caches were built and read in the same JVM, interleaved. The first was built at a batch size of 2, so every one of these batches passes through, which is the layout `main` produces. The figures are medians of 3 builds and 9 reads, on an Apple M3 Max with JDK 17, Spark 4.1 and a debug native library:
| | `main`'s layout | This PR |
| --- | ---: | ---: |
| Cached batches | 10000, of 250 to 775 rows | 510, of 5000 to 10000 rows |
| Payload | 36.1 MiB | 32.8 MiB |
| Storage memory | 44.0 MiB | 33.2 MiB |
| Build | 3961 ms | 3815 ms |
| Aggregate through `CometInMemoryTableScan` | 369 ms | 289 ms |
| `noop` write of every row | 141 ms | 113 ms |
| Aggregate by Spark operators | 315 ms | 299 ms |
On a copy of the same rows in random order, which compresses less well, storage memory went from 77.3 to 66.0 MiB, the build from 4.4 to 3.9 seconds, and the aggregate from 394 to 322 ms. Separate runs of `main` and this PR gave the same layouts. From a single map task, the same relation caches the same 633 full batches with or without this PR.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: combine small columnar batches before caching them [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #6988:
URL: https://github.com/apache/datafusion-comet/pull/6988#discussion_r4402105113
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CacheBatchCoalescer.scala:
##########
@@ -0,0 +1,216 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.comet.execution.arrow
+
+import scala.collection.JavaConverters._
+import scala.collection.mutable
+import scala.util.control.NonFatal
+
+import org.apache.arrow.memory.BufferAllocator
+import org.apache.arrow.vector.VectorSchemaRoot
+import org.apache.spark.TaskContext
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+import org.apache.comet.vector.NativeUtil
+
+/**
+ * The columnar input of a cache build, with its small batches combined.
+ *
+ * The columnar write stores each batch it is handed as one cached batch, and nothing upstream
+ * bounds how small a batch can be. A shuffle read returns at least one batch per map output
+ * block, so a relation cached from a shuffle over many map tasks would otherwise hold thousands
+ * of batches of a few hundred rows. Each carries its own message metadata, compression framing
+ * and statistics row, and every read of the cache pays for the batch count.
+ *
+ * A batch the writer can store as it stands, Arrow-backed and laid out the way the reader reads
+ * it, is passed through when it holds at least half of `targetRows`, so full native batches are
+ * still written without a copy. Every other batch is copied into batches of `targetRows` rows,
+ * the size the row path cuts at, and split where it crosses that size. A small batch is worth the
+ * copy, and a batch the writer cannot store as it stands has to be copied into Arrow anyway, so
+ * combining it costs nothing more. The copy reads each column through the `ColumnVector`
+ * interface, as `ArrowWriter.writeColumns` does for any input, so it handles dictionary-encoded
+ * and nested columns, and offsets that do not start at zero, as a native producer's sliced arrays
+ * can have. Arrow's `VectorAppender` assumes they start at zero.
+ *
+ * Row order is kept: the rows copied so far are emitted before a batch that passes through. Each
+ * batch is copied before the input is advanced, because a native producer releases the batch it
+ * returned last when it is asked for the next one.
+ *
+ * The batches built here are owned here. Each is released when the next batch is requested, and
+ * whatever is still open when the task completes is released then. A batch passed through still
+ * belongs to its producer.
+ */
+private[arrow] class CacheBatchCoalescer(
+ input: Iterator[ColumnarBatch],
+ attrs: Seq[Attribute],
+ targetRows: Int,
+ allocator: BufferAllocator)
+ extends Iterator[ColumnarBatch] {
+
+ require(targetRows > 0, s"The target batch size must be positive, but was $targetRows")
+
+ private val arrowSchema =
+ Utils.toArrowSchema(Utils.fromAttributes(attrs), CometArrowStream.NATIVE_TIMEZONE)
+ private val readerFields = arrowSchema.getFields.asScala.toIndexedSeq
+
+ // A batch at least this large that the writer can store as it stands is not copied.
+ private val passThroughRows = math.max(1, targetRows / 2)
+
+ // The batch being filled.
+ private var pending: VectorSchemaRoot = null
+ private var pendingWriter: ArrowWriter = null
+ private var pendingRows = 0
+
+ // Batches built here and filled, in input order, ahead of `held`.
+ private val filled = mutable.Queue.empty[VectorSchemaRoot]
+
+ // A batch that passes through, waiting for the rows copied ahead of it to be emitted.
+ private var held: ColumnarBatch = null
+
+ // The batch the next call to `next` returns, and the root behind it when it was built here.
+ private var ready: ColumnarBatch = null
+ private var readyRoot: VectorSchemaRoot = null
+
+ // The root behind the batch `next` returned last, released once the consumer asks for more.
+ private var emitted: VectorSchemaRoot = null
+
+ private var closed = false
+
+ Option(TaskContext.get()).foreach(_.addTaskCompletionListener[Unit](_ => close()))
+
+ override def hasNext: Boolean = {
+ if (ready == null && !closed) {
+ release(emitted)
+ emitted = null
+ advance()
+ }
+ ready != null
+ }
+
+ override def next(): ColumnarBatch = {
+ if (!hasNext) {
+ throw new NoSuchElementException("No more cache input batches")
+ }
+ val batch = ready
+ emitted = readyRoot
+ ready = null
+ readyRoot = null
+ batch
+ }
+
+ /** Releases every batch built here that is still open. Idempotent. */
+ def close(): Unit = {
+ if (!closed) {
+ closed = true
+ release(emitted)
+ release(readyRoot)
+ release(pending)
+ filled.foreach(release)
+ filled.clear()
+ emitted = null
+ ready = null
+ readyRoot = null
+ pending = null
+ pendingWriter = null
+ held = null
+ }
+ }
+
+ private def advance(): Unit = {
+ while (filled.isEmpty && held == null && input.hasNext) {
+ val batch = input.next()
+ if (passesThrough(batch)) {
+ finishPending()
+ held = batch
+ } else {
+ append(batch)
+ }
+ }
+ if (filled.isEmpty && held == null) {
+ // The input is exhausted.
+ finishPending()
+ }
+ if (filled.nonEmpty) {
+ readyRoot = filled.dequeue()
+ ready = NativeUtil.rootAsBatch(readyRoot)
+ } else if (held != null) {
+ ready = held
+ held = null
+ }
+ }
+
+ private def passesThrough(batch: ColumnarBatch): Boolean =
+ batch.numRows() >= passThroughRows &&
+ Utils.isArrowBacked(batch) &&
+ CachedBatchIpc.matchesReaderLayout(batch, readerFields)
+
+ // Copies every row of `batch`, filling batches of targetRows rows.
+ private def append(batch: ColumnarBatch): Unit = {
+ val numRows = batch.numRows()
+ var start = 0
+ while (start < numRows) {
Review Comment:
You were right about this. With four string columns and a 4 MiB limit on the allocator, one 4096-row input at a target of 10 ran out of memory before the first output batch came back. The coalescer now keeps the input it is copying and a row offset, and copies only as far as one output batch each time it is asked for one, so the producer is not advanced until that input is used up. The new `CacheBatchCoalescerSuite` has the bounded-allocator case, and it fails on the previous code with the same out-of-memory error. Fixed in 4220514d6f.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: combine small columnar batches before caching them [datafusion-comet]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #6988:
URL: https://github.com/apache/datafusion-comet/pull/6988#discussion_r4272822039
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CacheBatchCoalescer.scala:
##########
@@ -0,0 +1,216 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.comet.execution.arrow
+
+import scala.collection.JavaConverters._
+import scala.collection.mutable
+import scala.util.control.NonFatal
+
+import org.apache.arrow.memory.BufferAllocator
+import org.apache.arrow.vector.VectorSchemaRoot
+import org.apache.spark.TaskContext
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+import org.apache.comet.vector.NativeUtil
+
+/**
+ * The columnar input of a cache build, with its small batches combined.
+ *
+ * The columnar write stores each batch it is handed as one cached batch, and nothing upstream
+ * bounds how small a batch can be. A shuffle read returns at least one batch per map output
+ * block, so a relation cached from a shuffle over many map tasks would otherwise hold thousands
+ * of batches of a few hundred rows. Each carries its own message metadata, compression framing
+ * and statistics row, and every read of the cache pays for the batch count.
+ *
+ * A batch the writer can store as it stands, Arrow-backed and laid out the way the reader reads
+ * it, is passed through when it holds at least half of `targetRows`, so full native batches are
+ * still written without a copy. Every other batch is copied into batches of `targetRows` rows,
+ * the size the row path cuts at, and split where it crosses that size. A small batch is worth the
+ * copy, and a batch the writer cannot store as it stands has to be copied into Arrow anyway, so
+ * combining it costs nothing more. The copy reads each column through the `ColumnVector`
+ * interface, as `ArrowWriter.writeColumns` does for any input, so it handles dictionary-encoded
+ * and nested columns, and offsets that do not start at zero, as a native producer's sliced arrays
+ * can have. Arrow's `VectorAppender` assumes they start at zero.
+ *
+ * Row order is kept: the rows copied so far are emitted before a batch that passes through. Each
+ * batch is copied before the input is advanced, because a native producer releases the batch it
+ * returned last when it is asked for the next one.
+ *
+ * The batches built here are owned here. Each is released when the next batch is requested, and
+ * whatever is still open when the task completes is released then. A batch passed through still
+ * belongs to its producer.
+ */
+private[arrow] class CacheBatchCoalescer(
+ input: Iterator[ColumnarBatch],
+ attrs: Seq[Attribute],
+ targetRows: Int,
+ allocator: BufferAllocator)
+ extends Iterator[ColumnarBatch] {
+
+ require(targetRows > 0, s"The target batch size must be positive, but was $targetRows")
+
+ private val arrowSchema =
+ Utils.toArrowSchema(Utils.fromAttributes(attrs), CometArrowStream.NATIVE_TIMEZONE)
+ private val readerFields = arrowSchema.getFields.asScala.toIndexedSeq
+
+ // A batch at least this large that the writer can store as it stands is not copied.
+ private val passThroughRows = math.max(1, targetRows / 2)
+
+ // The batch being filled.
+ private var pending: VectorSchemaRoot = null
+ private var pendingWriter: ArrowWriter = null
+ private var pendingRows = 0
+
+ // Batches built here and filled, in input order, ahead of `held`.
+ private val filled = mutable.Queue.empty[VectorSchemaRoot]
+
+ // A batch that passes through, waiting for the rows copied ahead of it to be emitted.
+ private var held: ColumnarBatch = null
+
+ // The batch the next call to `next` returns, and the root behind it when it was built here.
+ private var ready: ColumnarBatch = null
+ private var readyRoot: VectorSchemaRoot = null
+
+ // The root behind the batch `next` returned last, released once the consumer asks for more.
+ private var emitted: VectorSchemaRoot = null
+
+ private var closed = false
+
+ Option(TaskContext.get()).foreach(_.addTaskCompletionListener[Unit](_ => close()))
+
+ override def hasNext: Boolean = {
+ if (ready == null && !closed) {
+ release(emitted)
+ emitted = null
+ advance()
+ }
+ ready != null
+ }
+
+ override def next(): ColumnarBatch = {
+ if (!hasNext) {
+ throw new NoSuchElementException("No more cache input batches")
+ }
+ val batch = ready
+ emitted = readyRoot
+ ready = null
+ readyRoot = null
+ batch
+ }
+
+ /** Releases every batch built here that is still open. Idempotent. */
+ def close(): Unit = {
+ if (!closed) {
+ closed = true
+ release(emitted)
+ release(readyRoot)
+ release(pending)
+ filled.foreach(release)
+ filled.clear()
+ emitted = null
+ ready = null
+ readyRoot = null
+ pending = null
+ pendingWriter = null
+ held = null
+ }
+ }
+
+ private def advance(): Unit = {
+ while (filled.isEmpty && held == null && input.hasNext) {
+ val batch = input.next()
+ if (passesThrough(batch)) {
+ finishPending()
+ held = batch
+ } else {
+ append(batch)
+ }
+ }
+ if (filled.isEmpty && held == null) {
+ // The input is exhausted.
+ finishPending()
+ }
+ if (filled.nonEmpty) {
+ readyRoot = filled.dequeue()
+ ready = NativeUtil.rootAsBatch(readyRoot)
+ } else if (held != null) {
+ ready = held
+ held = null
+ }
+ }
+
+ private def passesThrough(batch: ColumnarBatch): Boolean =
+ batch.numRows() >= passThroughRows &&
+ Utils.isArrowBacked(batch) &&
+ CachedBatchIpc.matchesReaderLayout(batch, readerFields)
+
+ // Copies every row of `batch`, filling batches of targetRows rows.
+ private def append(batch: ColumnarBatch): Unit = {
+ val numRows = batch.numRows()
+ var start = 0
+ while (start < numRows) {
Review Comment:
[P2] Emit split fragments incrementally instead of queuing the entire input. When a Spark-vector batch exceeds `targetRows`, this loop allocates and retains every output root before `advance()` returns the first one. Each variable-width vector receives Arrow's initial buffer allocation, so reducing `spark.sql.inMemoryColumnarStorage.batchSize` can greatly increase memory usage. For a normal-sized 4096-row input with four one-byte string columns, the previous conversion held 256 KiB, while a target of 100 retained 7.69 MiB and a target of 10 exhausted a bounded 64 MiB allocator before returning any output. Production uses the unbounded `CometArrowAllocator`, so this additional off-heap memory bypasses Spark's memory budget. Please retain the current input and row offset, emit one completed fragment, and continue copying only when the consumer requests another. The producer can remain unadvanced until that input is exhausted.
Evidence: Reproduced with the exact-head `CacheBatchCoalescer` and `ArrowWriters`, Spark 4.1.3, Arrow 18.3.0, and `-Darrow.allocation.manager.type=Unsafe`. Construct four `OnHeapColumnVector(4096, StringType)` columns containing the single byte 'x', wrap them in a `ColumnarBatch`, and call `new CacheBatchCoalescer(Iterator(batch), attrs, targetRows, new RootAllocator(64L * 1024 * 1024)).next()`. The unchanged base conversion sequence (`ArrowWriter.create(root, 4096)`, `writeColumns(batch, 0, 4096)`, `finish()`) retained 262144 bytes. Head retained 8060928 bytes at target 100. At target 10 it threw `OutOfMemoryException` before the first output, with 67092480 bytes allocated. Explicit cleanup returned allocation to zero. Reproduction and logs are under `/tmp/pr6988-review/`.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: combine small columnar batches before caching them [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #6988:
URL: https://github.com/apache/datafusion-comet/pull/6988#discussion_r4402105113
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CacheBatchCoalescer.scala:
##########
@@ -0,0 +1,216 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.comet.execution.arrow
+
+import scala.collection.JavaConverters._
+import scala.collection.mutable
+import scala.util.control.NonFatal
+
+import org.apache.arrow.memory.BufferAllocator
+import org.apache.arrow.vector.VectorSchemaRoot
+import org.apache.spark.TaskContext
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+import org.apache.comet.vector.NativeUtil
+
+/**
+ * The columnar input of a cache build, with its small batches combined.
+ *
+ * The columnar write stores each batch it is handed as one cached batch, and nothing upstream
+ * bounds how small a batch can be. A shuffle read returns at least one batch per map output
+ * block, so a relation cached from a shuffle over many map tasks would otherwise hold thousands
+ * of batches of a few hundred rows. Each carries its own message metadata, compression framing
+ * and statistics row, and every read of the cache pays for the batch count.
+ *
+ * A batch the writer can store as it stands, Arrow-backed and laid out the way the reader reads
+ * it, is passed through when it holds at least half of `targetRows`, so full native batches are
+ * still written without a copy. Every other batch is copied into batches of `targetRows` rows,
+ * the size the row path cuts at, and split where it crosses that size. A small batch is worth the
+ * copy, and a batch the writer cannot store as it stands has to be copied into Arrow anyway, so
+ * combining it costs nothing more. The copy reads each column through the `ColumnVector`
+ * interface, as `ArrowWriter.writeColumns` does for any input, so it handles dictionary-encoded
+ * and nested columns, and offsets that do not start at zero, as a native producer's sliced arrays
+ * can have. Arrow's `VectorAppender` assumes they start at zero.
+ *
+ * Row order is kept: the rows copied so far are emitted before a batch that passes through. Each
+ * batch is copied before the input is advanced, because a native producer releases the batch it
+ * returned last when it is asked for the next one.
+ *
+ * The batches built here are owned here. Each is released when the next batch is requested, and
+ * whatever is still open when the task completes is released then. A batch passed through still
+ * belongs to its producer.
+ */
+private[arrow] class CacheBatchCoalescer(
+ input: Iterator[ColumnarBatch],
+ attrs: Seq[Attribute],
+ targetRows: Int,
+ allocator: BufferAllocator)
+ extends Iterator[ColumnarBatch] {
+
+ require(targetRows > 0, s"The target batch size must be positive, but was $targetRows")
+
+ private val arrowSchema =
+ Utils.toArrowSchema(Utils.fromAttributes(attrs), CometArrowStream.NATIVE_TIMEZONE)
+ private val readerFields = arrowSchema.getFields.asScala.toIndexedSeq
+
+ // A batch at least this large that the writer can store as it stands is not copied.
+ private val passThroughRows = math.max(1, targetRows / 2)
+
+ // The batch being filled.
+ private var pending: VectorSchemaRoot = null
+ private var pendingWriter: ArrowWriter = null
+ private var pendingRows = 0
+
+ // Batches built here and filled, in input order, ahead of `held`.
+ private val filled = mutable.Queue.empty[VectorSchemaRoot]
+
+ // A batch that passes through, waiting for the rows copied ahead of it to be emitted.
+ private var held: ColumnarBatch = null
+
+ // The batch the next call to `next` returns, and the root behind it when it was built here.
+ private var ready: ColumnarBatch = null
+ private var readyRoot: VectorSchemaRoot = null
+
+ // The root behind the batch `next` returned last, released once the consumer asks for more.
+ private var emitted: VectorSchemaRoot = null
+
+ private var closed = false
+
+ Option(TaskContext.get()).foreach(_.addTaskCompletionListener[Unit](_ => close()))
+
+ override def hasNext: Boolean = {
+ if (ready == null && !closed) {
+ release(emitted)
+ emitted = null
+ advance()
+ }
+ ready != null
+ }
+
+ override def next(): ColumnarBatch = {
+ if (!hasNext) {
+ throw new NoSuchElementException("No more cache input batches")
+ }
+ val batch = ready
+ emitted = readyRoot
+ ready = null
+ readyRoot = null
+ batch
+ }
+
+ /** Releases every batch built here that is still open. Idempotent. */
+ def close(): Unit = {
+ if (!closed) {
+ closed = true
+ release(emitted)
+ release(readyRoot)
+ release(pending)
+ filled.foreach(release)
+ filled.clear()
+ emitted = null
+ ready = null
+ readyRoot = null
+ pending = null
+ pendingWriter = null
+ held = null
+ }
+ }
+
+ private def advance(): Unit = {
+ while (filled.isEmpty && held == null && input.hasNext) {
+ val batch = input.next()
+ if (passesThrough(batch)) {
+ finishPending()
+ held = batch
+ } else {
+ append(batch)
+ }
+ }
+ if (filled.isEmpty && held == null) {
+ // The input is exhausted.
+ finishPending()
+ }
+ if (filled.nonEmpty) {
+ readyRoot = filled.dequeue()
+ ready = NativeUtil.rootAsBatch(readyRoot)
+ } else if (held != null) {
+ ready = held
+ held = null
+ }
+ }
+
+ private def passesThrough(batch: ColumnarBatch): Boolean =
+ batch.numRows() >= passThroughRows &&
+ Utils.isArrowBacked(batch) &&
+ CachedBatchIpc.matchesReaderLayout(batch, readerFields)
+
+ // Copies every row of `batch`, filling batches of targetRows rows.
+ private def append(batch: ColumnarBatch): Unit = {
+ val numRows = batch.numRows()
+ var start = 0
+ while (start < numRows) {
Review Comment:
You were right about this. With four string columns and a 4 MiB limit on the allocator, one 4096-row input at a target of 10 ran out of memory before the first output batch came back. The coalescer now keeps the input it is copying and a row offset, and copies only as far as one output batch each time it is asked for one, so the producer is not advanced until that input is used up. The new `CacheBatchCoalescerSuite` has the bounded-allocator case, and it fails on the previous code with the same out-of-memory error. Fixed in 59c9b6d748.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] perf: combine small columnar batches before caching them [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #6988:
URL: https://github.com/apache/datafusion-comet/pull/6988#discussion_r4414768011
##########
spark/src/test/scala/org/apache/spark/sql/comet/execution/arrow/CacheBatchCoalescerSuite.scala:
##########
@@ -0,0 +1,116 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.comet.execution.arrow
+
+import java.nio.charset.StandardCharsets.UTF_8
+
+import org.scalatest.funsuite.AnyFunSuite
+
+import org.apache.arrow.memory.RootAllocator
+import org.apache.spark.sql.catalyst.expressions.AttributeReference
+import org.apache.spark.sql.execution.vectorized.OnHeapColumnVector
+import org.apache.spark.sql.types.StringType
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
+
+class CacheBatchCoalescerSuite extends AnyFunSuite {
+
+ private val numColumns = 4
+ private val attrs = (0 until numColumns).map(i => AttributeReference(s"s$i", StringType)())
+
+ // A batch of Spark's own on-heap vectors, which the coalescer always copies. Row `i` of column
+ // `c` holds "c<c>r<first + i>", so the output order can be checked.
+ private def sparkBatch(first: Int, rows: Int): ColumnarBatch = {
+ val columns = (0 until numColumns).map { c =>
+ val vector = new OnHeapColumnVector(rows, StringType)
+ (0 until rows).foreach(i => vector.putByteArray(i, s"c${c}r${first + i}".getBytes(UTF_8)))
+ vector: ColumnVector
+ }
+ new ColumnarBatch(columns.toArray, rows)
+ }
+
+ // Runs `f` on a coalescer over `inputs`, checks that closing it returns every byte it took, and
+ // returns the most the allocator held at once.
+ private def withCoalescer(limit: Long, inputs: Seq[ColumnarBatch], target: Int)(
Review Comment:
Every test here, and the new `CometInMemoryCacheSuite` cases, drains the coalescer before `close()` runs, so the release paths for a partly consumed one have no test: `pending`, `filled`, a held `readyRoot` and `emitted` at close time, and the `NonFatal` cleanup in `startPending`. A change that made `close()` skip `pending` or `filled` would still pass. Could one case take a single batch from a multi-batch input (a 25-row `sparkBatch` at target 10 leaves a partial `pending` and a non-empty `filled`) and then close, and another use an allocator limit small enough that `startPending` or `writeColumns` throws, both checking the allocator returns to zero?
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/CacheBatchCoalescer.scala:
##########
@@ -0,0 +1,227 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.comet.execution.arrow
+
+import scala.collection.JavaConverters._
+import scala.collection.mutable
+import scala.util.control.NonFatal
+
+import org.apache.arrow.memory.BufferAllocator
+import org.apache.arrow.vector.VectorSchemaRoot
+import org.apache.spark.TaskContext
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.comet.util.Utils
+import org.apache.spark.sql.vectorized.ColumnarBatch
+
+import org.apache.comet.vector.NativeUtil
+
+/**
+ * The columnar input of a cache build, with its small batches combined.
+ *
+ * The columnar write stores each batch it is handed as one cached batch, and nothing upstream
+ * bounds how small a batch can be. A shuffle read returns at least one batch per map output
+ * block, so a relation cached from a shuffle over many map tasks would otherwise hold thousands
+ * of batches of a few hundred rows. Each carries its own message metadata, compression framing
+ * and statistics row, and every read of the cache pays for the batch count.
+ *
+ * A batch the writer can store as it stands, Arrow-backed and laid out the way the reader reads
+ * it, is passed through when it holds at least half of `targetRows`, so full native batches are
+ * still written without a copy. Every other batch is copied into batches of `targetRows` rows,
+ * the size the row path cuts at, and split where it crosses that size. A small batch is worth the
+ * copy, and a batch the writer cannot store as it stands has to be copied into Arrow anyway, so
+ * combining it costs nothing more. The copy reads each column through the `ColumnVector`
+ * interface, as `ArrowWriter.writeColumns` does for any input, so it handles dictionary-encoded
+ * and nested columns, and offsets that do not start at zero, as a native producer's sliced arrays
+ * can have. Arrow's `VectorAppender` assumes they start at zero.
+ *
+ * Row order is kept: the rows copied so far are emitted before a batch that passes through. A
+ * large input is copied one output batch at a time, as the consumer asks for them, so the batches
+ * open at once do not grow with its size. The input is not advanced until it has been copied in
+ * full, because a native producer releases the batch it returned last when it is asked for the
+ * next one.
+ *
+ * The batches built here are owned here. Each is released when the next batch is requested, and
+ * whatever is still open when the task completes is released then. A batch passed through still
+ * belongs to its producer.
+ */
+private[arrow] class CacheBatchCoalescer(
+ input: Iterator[ColumnarBatch],
+ attrs: Seq[Attribute],
+ targetRows: Int,
+ allocator: BufferAllocator)
+ extends Iterator[ColumnarBatch] {
+
+ require(targetRows > 0, s"The target batch size must be positive, but was $targetRows")
+
+ private val arrowSchema =
+ Utils.toArrowSchema(Utils.fromAttributes(attrs), CometArrowStream.NATIVE_TIMEZONE)
+ private val readerFields = arrowSchema.getFields.asScala.toIndexedSeq
+
+ // A batch at least this large that the writer can store as it stands is not copied.
+ private val passThroughRows = math.max(1, targetRows / 2)
Review Comment:
The threshold follows `spark.sql.inMemoryColumnarStorage.batchSize`, so it is 5000 rows by default, and default 8192-row native batches pass through. With `spark.comet.batchSize` at 4096 or lower, or the cache batch size above 16384, every native Arrow batch now takes the `ArrowWriter.writeColumns` copy, where before it was written without one. The table in the description measures inputs of 250 to 775 rows, where the copy clearly pays. Could you measure the build at 4096-row native batches against main, or tie the threshold to Comet's batch size? I have not measured the copy cost.
##########
docs/source/user-guide/latest/in-memory-cache.md:
##########
@@ -77,6 +77,14 @@ metrics only through that scan.
Each cached batch is stored as a single Arrow IPC record batch message and its body.
+A cache built from rows cuts its batches at `spark.sql.inMemoryColumnarStorage.batchSize` rows,
+10000 by default, as Spark's own format does. A cache built from the Arrow batches of Comet
Review Comment:
Spark's row path also ends a batch once it reaches about 4 MiB (`ColumnBuilder.MAX_BATCH_SIZE_IN_BYTE` in `InMemoryRelation`, in 3.4, 3.5 and 4.1), while Comet's row path and this coalescer cut by rows only. Could the sentence say that, so readers know wide rows give larger cached batches than Spark's format?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org