A cached aggregate stores each output slice with the bytes of the slices before it, inflating its size statistics and memory #6995

What is the problem the feature request solves?

Part of #6834.

A grouped aggregate (a GROUP BY or DISTINCT) emits a partition's output in spark.comet.batchSize-row slices of one batch. As #6779 describes, a slice of a string, binary, list or map array reaches the JVM with offsets that start past zero, and the rows before the slice are still at the front of its data buffer. CachedBatchIpc.serialize unloads each batch with VectorUnloader (CachedBatchIpc.scala:224), which writes a data buffer from byte 0 up to its last offset. So the k-th slice of a partition is cached with the bytes of every slice before it, and a partition that emits K slices stores about (K + 1) / 2 times its variable-width data. The values read back correctly.

columnSizes (:623) measures the same buffers, so sizeInBytes grows by the same factor. statsRow (ArrowCachedBatchSerializer.scala:340) records those sizes so that the planner sees what Spark's own cache format would report, and for these batches it sees several times more.

Measured on main (7d8d9b1) with Spark 4.1 and local[5] on an Apple M3 Max. The plugin installs Comet's cache serializer, as it does by default, and spark.comet.batchSize is at its default of 8192. k is a distinct 11-character string per row. Each SELECT DISTINCT k is compared with the same rows, in the same order, cached from a Parquet scan, whose batches start at offset 0:

Storage memory is the cached RDD's memSize, as the Storage tab shows it. At Spark's defaults each of the 200 partitions holds about 50,000 keys and emits 7 slices, so the cache reports 3.5 times its size and holds 3.1 times the memory. In one partition the 200K keys come out as 25 slices, and the reported size is 9.9 times the real one.

That is enough to change the plan. Joined against range(0, 10000000) with the default 10 MiB spark.sql.autoBroadcastJoinThreshold and AQE off, the 200K-key DISTINCT cache becomes a CometSortMergeJoin, while the scan-built cache and Spark's own cache format are both broadcast with CometBroadcastHashJoin. Reads pay for the extra bytes too: max(k), count(*) over the 200K-row cache took a median of 71 ms against 55 ms for the scan-built cache (nine interleaved runs, with a debug libcomet, so the absolute times are high).

This is new relative to 1.1.0, where the cache used Spark's format unless spark.comet.exec.inMemoryCache.enabled was set. #5634 made Comet's format the default from Spark 3.5.

Describe the potential solution

Rebase the offsets before unloading, as #6900 does for broadcasts in Utils.serializeBatches. When an offset buffer in the batch, children included, starts past zero, unload root.slice(0, numRows) instead of root, and close the slice in a finally. The slice rebases the offsets of the variable-width and list vectors and shares the other buffers, so the cost is a copy of their offsets, and columnSizes then measures the slice with no further change. #6900's offsetsStartAtZero is private to Utils and could be shared.

Rebasing in zero_offsets on the native side instead would fix every JVM consumer of a sliced batch at once, at the cost of rebasing on every sliced export, including the ones that are never serialized.

Additional context

The factor depends on how many groups a partition holds. With the defaults, the 200K keys spread over 200 partitions of about 1,000 groups, so each partition emitted one batch and sizeInBytes was right.

#6988 combines small batches before the write, but passes batches of at least half of spark.comet.batchSize rows through unchanged, so full-size slices still carry the bytes before them. #6985 changes the same method.