comphead opened a new pull request, #6900:
URL: https://github.com/apache/datafusion-comet/pull/6900
## Which issue does this PR close?
Closes #6779.
## Rationale for this change
A Comet broadcast can return wrong values, or fail with
`OversizedAllocationException`, when its
build side emits sliced batches with string, binary, array or map columns,
at any nesting level.
arrow-rs slices these arrays by moving the start of their offsets, so the
JVM imports a slice
whose offsets start past zero and whose data buffer still begins with the
rows before the slice.
`zero_offsets` (#6339) only handles `ArrayData::offset`, which arrow-rs
keeps for booleans.
DataFusion produces such slices whenever an operator emits one batch in
`batch_size` pieces.
Grouped aggregates do this for every `GROUP BY` and `DISTINCT`, and so do
`LIMIT ... OFFSET`,
TopK and window aggregates.
`Utils.serializeBatches` writes these vectors with Arrow Java's
`VectorUnloader`, which keeps the
offsets as they are and writes the data buffer from byte 0.
`coalesceBroadcastBatches` (#3703)
then appends the decoded batches with `VectorSchemaRootAppender`. In Arrow
Java 18.3 that assumes
offsets start at zero, so the first value of every slice absorbs all the
bytes before it. The
issue's query returns 199976 rows instead of 200000, and the appended data
grows with every
slice, which is where the `OversizedAllocationException` comes from. Every
release since 0.15.0
has this.
Arrow Java fixed `VectorAppender` in apache/arrow-java#1180, but no release
includes it yet.
19.0.0 is the latest.
## What changes are included in this PR?
- `serializeBatches` checks each column, children included, for an offset
buffer that starts past
zero. When one does, it writes `root.slice(0, rowCount)` instead of the
root. The slice rebases
those offsets to zero and drops the bytes before them, and shares every
other buffer. Batches
whose offsets already start at zero are written as before.
- The bytes before a slice are no longer written. Before this change, each
slice carried its
buffer from the start up to its own end, so the bytes collected to the
driver grew with the
square of the number of slices per partition.
- A comment at the `VectorSchemaRootAppender.append` call notes that
coalescing relies on this.
The in-memory cache serializer (`CachedBatchIpc.serialize`) also writes the
bytes before the first
offset. There they only inflate the payload and its size statistics, and the
values read back
correctly, so this PR leaves it unchanged.
## How are these changes tested?
- `CometJoinSuite` "Broadcast coalescing keeps the values of sliced batches"
broadcasts
`SELECT DISTINCT *` over string, binary, `array<string>` and struct
columns with nulls at batch
size 100. It compares every value with Spark and checks that the broadcast
still coalesced.
- `UtilsSuite` "serializeBatches rebases offsets that start past zero"
coalesces string vectors
shaped like imported slices and checks the values. It also checks that
each slice takes as many
bytes as the same values in a vector of their own.
Both tests fail without the fix. Locally, on Spark 4.1 against a build of
recent `main`, the new
tests, the issue's reproducer, the rest of `UtilsSuite` and 25 existing
broadcast tests in
`CometJoinSuite`, `CometExecSuite` and `CometInMemoryCacheSuite` pass. The
Spark 3.x profiles were
not run locally.
--
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]
