ranflarion opened a new pull request, #11648:
URL: https://github.com/apache/arrow-rs/pull/11648
# Which issue does this PR close?
- Closes #11647.
# Rationale for this change
`StreamDecoder` assembles a message that arrives split across input buffers
in a scratch `MutableBuffer` that starts empty and doubles as the pieces are
copied in, and the decoded batch is zero-copy over that buffer, so it keeps
between 1x and 2x its body alive. A reader that feeds fixed-size chunks hits
this on nearly every message large enough to span one: DataFusion 55 reads
spill files back in 128 KiB chunks, and a spilling sort's read-back batches
held 2,969 MiB for 2,198 MiB of data. `StreamReader` avoids it since #10522 by
reading a body into `min(body_length, MAX_PREALLOC_BYTES)` up front.
# What changes are included in this PR?
When the first bytes of a message flatbuffer or body are copied into the
empty scratch buffer, the decoder now reserves the declared length, capped at
`MAX_PREALLOC_BYTES` the same way `read_body_bounded` caps it. A message up to
64 MiB is therefore allocated once at its size, while an implausible declared
length still reserves at most 64 MiB and leaves the decoder waiting for bytes
that never arrive. Larger messages keep growing geometrically past the cap, as
they do in `StreamReader`. The path that slices a message straight out of a
single input buffer is unchanged.
# Are these changes tested?
Yes, two tests in `arrow-ipc/src/reader/stream.rs`.
`test_split_body_is_buffered_at_its_declared_size` feeds a 99,272 byte stream
one byte at a time and checks the decoded batch holds no more than the stream:
without the change it holds 131,072 bytes.
`test_implausible_body_length_reserves_a_bounded_buffer` declares body lengths
of `i64::MAX` and `1 << 50` and checks the decoder keeps waiting for the body:
reserving the declared length without the cap fails it with
`MemoryError("invalid allocation layout for requested capacity")`.
In the DataFusion sort from the issue, built against arrow 59.3 with this
change, peak native allocation went from 201.8 MiB to 161.0 MiB.
The `arrow-ipc` suite passes (138 library tests, including the two above),
along with `cargo fmt --all -- --check` and `cargo clippy -p arrow-ipc
--all-targets --all-features -- -D warnings`.
# Are there any user-facing changes?
No API changes. Batches decoded from split messages hold less memory. No
breaking changes.
--
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]
