andygrove opened a new pull request, #7225: URL: https://github.com/apache/datafusion-comet/pull/7225
## 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 applies the same fix to the cache. ## What changes are included in this PR? - `CachedBatchIpc.serialize` writes `root.slice(0, numRows)` instead of `root` when any offset buffer in the batch, children included, starts past zero, and closes 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 the offsets. `columnSizes` then measures the slice with no further change. - `Utils.offsetsStartAtZero` is `private[comet]` instead of `private`, so the cache and the broadcast path share one check. - `ffi.md` no longer says the cache serializer writes the earlier rows. 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? Two new tests in `CometInMemoryCacheSuite`: - A native `DISTINCT` over 4000 rows at `spark.comet.batchSize=100` (40 slices) is cached, and its `sizeInBytes` is compared with the same rows cached from a Parquet scan. Without the fix it is 2,657,360 bytes against 161,360. The test also checks the answer. - A sliced `VarCharVector`, built the way the C Data import builds one, serializes to the same number of bytes as a fresh vector with the same values. A write that fails part way (an injected compression failure) leaves nothing allocated in the child allocator, which throws on close if it does. I checked each with a mutant. With the slice removed, both fail. With the slice never closed, only the allocation test fails. `CometInMemoryCacheSuite` and `UtilsSuite` pass on Spark 4.1 (91 tests), including Spotless and Scalastyle. I have not run the other Spark profiles, Scalafix or the strict-warnings build locally. The change uses `VectorSchemaRoot.slice`, which `Utils.serializeBatches` already uses on every profile. -- 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]
