[andygrove]

Which issue does this PR close?

Closes #6995.

Part of #6834.

Rationale for this change

A grouped aggregate emits a partition's output as spark.comet.batchSize-row slices of one batch. 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 the batch with VectorUnloader, which writes a data buffer from byte 0, so the k-th slice of a partition is cached with the bytes of every slice before it. A partition that emits K slices stores about (K + 1) / 2 times its variable-width data, and columnSizes, which measures the same buffers, reports the same inflated size in the relation's statistics. The values read back correctly, so only memory and the planner's size estimate are affected.

On the issue's measurements, 200K distinct 11-character keys in one partition (25 slices) report 28.7 MiB against 2.9 MiB for the same rows cached from a Parquet scan. That is enough to turn a broadcast join into a sort-merge join at the default 10 MiB threshold.

#6900 fixed the same layout for broadcasts in Utils.serializeBatches, which rebases the offsets with VectorSchemaRoot.slice before writing. This shares that fix with the cache, and with the one other writer that unloads imported vectors as they are, the Arrow Python runner.

What changes are included in this PR?

• Utils.withRebasedOffsets(root)(write) runs write on root when none of its vectors is a slice of a larger one. Otherwise it runs write on root.slice(0, numRows), which rebases the offsets of the variable-width and list vectors and shares the other buffers, and it closes the slice whether write returns or fails. Utils.serializeBatches uses it in place of the inline slice and close from fix: keep the values of sliced batches when coalescing a broadcast #6900.
• A list or map counts as a slice when its first offset is past zero, and also when its last offset is short of its child's length. The first slice of an array starts at zero, but its child still holds the elements of every later slice, so a check of the first offset alone writes them all. A string or binary vector stops writing its data at the last offset, so its first offset is enough.
• CachedBatchIpc.serialize unloads through withRebasedOffsets, so a cached slice no longer stores the rows before it, and columnSizes measures the rebased slice with no further change.
• CometArrowPythonRunnerBase.serializeBatch unloads through withRebasedOffsets as well. The whole-batch range of withInputBatchRange handed an imported slice's vectors to VectorUnloader as they were, so a sliced aggregate output feeding mapInArrow or mapInPandas sent the bytes of the earlier slices to the Python worker. The partial-range branch already rebased through column.slice. withRebasedOffsets is private[sql] for this.
• ffi.md names the three writers that rebase and tells new code that writes imported vectors to use withRebasedOffsets.

Rebasing in zero_offsets on the native side would fix every JVM consumer of a sliced batch at once, but it would rebase every sliced export, including the ones that are never serialized, so this stays on the write path as in #6900.

How are these changes tested?

• CometInMemoryCacheSuite has an end-to-end test. A native DISTINCT over 1000 rows at spark.comet.batchSize=100 (10 slices) is cached before its Parquet source, so that it reads the scan and not a cached copy of it. Its sizeInBytes is compared with the same rows cached from the scan, and the test asserts that the aggregate emitted more than one batch and that the answer matches. Without the fix the aggregate's cache is several times larger than the scan's.
• UtilsSuite tests withRebasedOffsets directly on slices of a list<string> built the way the C Data import builds them, including the first slice. Each slice must write the same number of bytes and the same values as a vector of its own, and the child allocator must be empty afterwards. Without the list check on the last offset, the first slice of 8 lists out of 24 writes 536 bytes against 232.
• CometArrowPythonRunnerSuite (Spark 4.x) writes a whole-batch range of a sliced string vector and compares the stream's length with the same values in a vector of their own. Without the change it is 688 bytes against 632.

CometInMemoryCacheSuite, UtilsSuite and CometArrowPythonRunnerSuite pass on Spark 4.1, and UtilsSuite and CometInMemoryCacheSuite pass on Spark 3.5. Spotless, Scalafix and the strict-warnings build (Spark 3.5) are clean.

[viirya]

Thanks for picking this up. Sharing the rebase between the broadcast and cache paths through withRebasedOffsets is a nice cleanup of the inline slice and close from #6900. It also means columnSizes measures what is actually written, with no extra change. I traced the slice through Arrow 18.3. VectorSchemaRoot.slice always goes through splitAndTransfer, the variable-width and list vectors get a fresh rebased offset buffer, and the data buffer is shared via slice. Calling root.clear() while the slice is live is safe because the slice holds its own references. The 32-bit-only match in offsetsStartAtZero doesn't leave a gap on this path either, because large-offset types fail matchesReaderLayout and take the convert path. The end-to-end test is well built. Caching the aggregate before the scan and asserting more than one batch makes sure it really exercises the sliced layout.

I left a few inline comments. One is about the Python runner, which has the same issue on a path this PR doesn't touch. I anchored it on the new ffi.md sentence, since that sentence now describes it as a rule. The others are about trimming the new unit tests, which mostly repeat existing coverage, and a test name that reads as data loss.

[viirya]

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?

[andygrove]

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 a3a85a4.

[viirya]

"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"?

[andygrove]

Renamed to "does not store the bytes before a sliced aggregate output" in c41846e.

[viirya]

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".

[andygrove]

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 b4a0ae1 and the removal in c41846e.

[viirya]

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.

[andygrove]

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 c41846e.

[comphead]

One gap beyond what viirya covered, on the first slice of a list or map column.

[comphead]

This only checks where the offsets start, so the first slice of a list or map column passes as is. DataFusion's grouped aggregate emits batch.slice(0, size) first, and arrow-rs GenericListArray::slice keeps the whole values array, so that slice arrives with offsets from zero and a child vector that the C Data import sizes to every later slice's elements too. VectorUnloader then writes the child with its own getValueCount, so the first cached slice stores, and columnSizes reports, the child of the whole batch. Strings are not affected, because BaseVariableWidthVector stops the data buffer at the last offset.

Could the BaseRepeatedValueVector case also require the last offset to equal getDataVector.getValueCount? ListVector's splitAndTransfer already cuts the child to [first, last), so the existing slice would trim both ends, which is what arrow-rs's own IPC writer does in get_list_array_buffers. The broadcast path shares the check, so its bytes would shrink too. A head slice (offset 0, fewer rows than the batch) in the list test would cover it, and the offsetsStartAtZero doc could mention the elements after the slice. I traced this through the Arrow Java 18.3 and arrow-rs sources and have not run it.

[andygrove]

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 b4a0ae1.

[viirya]

Thanks for the quick turnaround. The holdsOnlyItsOwnValues check now covers both ends of a list or map slice, and the new UtilsSuite test pins the head, middle and tail cases with both values and bytes. Routing the Python runner through the same helper and naming all three writers in ffi.md makes the rule easy to find for the next writer. LGTM.

[rich7420]

LGTM. Three regressions passed CI; Arrow component probes passed locally.

[andygrove]

Thanks @viirya @rich7420 @comphead