parthchandra opened a new issue, #7271:
URL: https://github.com/apache/datafusion-comet/issues/7271
Part of #7267.
## Problem
arrow's `BatchCoalescer` (`arrow-select/src/coalesce.rs`, 59.3.1) builds
output
batches up to a target row count. It has no byte limit.
`biggest_coalesce_batch_size` (`coalesce.rs:164, 210`, applied at
`:452-470`) is a
row-count pass-through threshold, not a cap: larger input batches skip
merging.
For `Utf8`/`Binary` columns it buffers input slices in
`GenericInProgressArray` and
calls `concat` when the target is reached (`coalesce/generic.rs:94-106`). If
those
rows hold more than 2 GiB, `concat` fails with `OffsetOverflowError`. Before
failing,
`concat` also reserves the full output buffer, so peak memory is about twice
the
buffered data.
DataFusion's `LimitedBatchCoalescer` wraps it with the pass-through
threshold set
to half the target (`coalesce/mod.rs:63-64`). It is used by `FilterExec`,
`CoalesceBatchesExec`, `RepartitionExec` and the hash join output stream
(`joins/hash_join/stream.rs:493`). The raw `BatchCoalescer` is used by
sort-merge
join (`materializing_stream.rs:417, 520, 574, 578`), nested loop join
(`:1499`) and
TopK. Comet's shuffle writer uses it directly (`buf_batch_writer.rs:133`).
Nested types overflow the same way: list, map and struct `concat` recurse
into
`concat_bytes` (`concat.rs:399, 464, 553-558`). The over-2 GiB buffer is
reserved
before the error and is not tracked by the DataFusion memory pool.
Even with the correct batch size (sub-issue 1), the default of 8192 rows
overflows
when values average more than ~256 KB.
## Proposed fix
arrow-rs:
- Add an optional byte target, for example
`with_target_batch_bytes(Option<usize>)`.
- Before buffering a batch, if the buffered bytes of any `Utf8`/`Binary`
column,
including ones nested inside lists, maps and structs, plus the incoming
bytes would
pass the target (capped at i32::MAX), finish the current output batch
first. If one
input batch alone is over the target, slice it.
- View types and primitive types are unaffected, or can use the same budget
for
memory reasons.
DataFusion:
- Pass the byte target through `LimitedBatchCoalescer` and expose it on
`FilterExec` (and any other operator using the coalescer), with a session
config
default.
Comet:
- Set the byte target from a Comet config, default about 1 GiB.
- Use it in the shuffle writer's coalescer (sub-issue 2).
## Tests
- arrow-rs unit tests with a small byte target: outputs never exceed it; row
order
and content are unchanged; an input batch over the target is sliced.
- DataFusion `FilterExec` test with a small byte target.
- Comet test with a small configured target and a binary column, compared
with Spark.
## Notes
- Before filing upstream, check arrow-rs #7764 and #10799. Recent benchmark
work
(#10797) points to a specialized Utf8/Binary in-progress array that would
replace
the generic concat path; a byte target may fit there.
- Upstream first, then bump arrow and DataFusion in Comet. Comet's interim
workaround is sub-issue 1 plus a lower `spark.comet.batchSize`.
--
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]
