ranflarion opened a new issue, #11647:
URL: https://github.com/apache/arrow-rs/issues/11647
**Describe the bug**
When a message arrives split across input buffers, `StreamDecoder` assembles
it in its scratch `MutableBuffer`, which starts empty and grows through
`try_extend_from_slice`, so its capacity doubles as the pieces arrive. The
decoded `RecordBatch` is zero-copy over that buffer and keeps its whole
capacity alive: anywhere from 1x to 2x the body length, depending on where the
first piece of the body falls in the input.
In `StreamDecoder::decode` (`arrow-ipc/src/reader/stream.rs`):
```rust
let to_read = buffer.len().min(body_length - self.buf.len());
self.buf
.try_extend_from_slice(&buffer[..to_read])
.map_err(|e| ArrowError::MemoryError(e.to_string()))?;
buffer.advance(to_read);
if self.buf.len() != body_length {
continue;
}
std::mem::take(&mut self.buf).into()
```
`StreamReader` no longer has this shape: since #10522 it reads a body into
`min(body_length, MAX_PREALLOC_BYTES)` up front, so a body up to 64 MiB is
allocated once at its size.
**To Reproduce**
```rust
let schema = Arc::new(Schema::new(vec![Field::new("s", DataType::Utf8,
false)]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(StringArray::from_iter_values(
(0..7000).map(|i| format!("value{i:05}")),
)) as _],
)
.unwrap();
let mut stream = Vec::new();
let mut writer = StreamWriter::try_new(&mut stream, &schema).unwrap();
writer.write(&batch).unwrap();
writer.finish().unwrap();
drop(writer);
let mut decoder = StreamDecoder::new();
let mut decoded = None;
for byte in stream.chunks(1) {
if let Some(batch) = decoder.decode(&mut Buffer::from(byte)).unwrap() {
decoded = Some(batch);
}
}
let held = decoded.unwrap().column(0).to_data().buffers()[0].capacity();
// holds 131072 bytes for a 99272 byte stream
assert!(held <= stream.len());
```
Feeding one byte at a time just starts the doubling from the smallest seed.
Any reader that hands the decoder fixed-size chunks smaller than its messages
gets the same result, and that is how DataFusion 55 reads spill files back, in
128 KiB chunks. In a spilling sort with 8192-row batches of 4 KB rows, the 152
batches read back from spill held 2,969 MiB for 2,198 MiB of data (1.39x per
batch at the median, 1.97x at worst), and none of the excess was in the sort's
memory reservation, which is sized from the batches as they were written.
DataFusion has since worked around it on its side in apache/datafusion#24592,
which traced apache/datafusion#17340 to this, by copying each message into an
exact-size buffer before handing it to the decoder.
**Expected behavior**
A split message should be held at its declared size, with the same bound
`StreamReader` puts on how much an untrusted length can reserve before the
bytes behind it arrive.
**Additional context**
The message flatbuffer goes through the same path in
`DecoderState::Message`, at much smaller sizes. I have a fix that reserves
`min(len, MAX_PREALLOC_BYTES)` when the first bytes of a message or body are
copied into the empty scratch buffer, and will post it shortly.
--
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]
