feat(ipc): Reduce redundant zero-initialization in IPC reader#9778
feat(ipc): Reduce redundant zero-initialization in IPC reader#9778pchintar wants to merge 4 commits intoapache:mainfrom
Conversation
| let mut buf = MutableBuffer::from_len_zeroed(total_len); | ||
| reader.read_exact(&mut buf)?; | ||
| let mut buf = MutableBuffer::with_capacity(total_len); | ||
| // Buffer is immediately fully initialized by `read_exact` before any read occurs |
There was a problem hiding this comment.
Strictly speaking, this is not sound for arbitrary Read implementations. Unfortunately, nothing would prevent a Read impl from reading the uninitialized bytes in buf, causing undefined behavior. For a known implementation like File this pattern could be fine though.
An alternative could be to combine take and read_to_end to read into a Vec, but note that this changes the alignment of the resulting Buffer:
let mut buf = Vec::with_capacity(total_len);
reader.take(total_len as u64).read_to_end(&mut buf)?;
if buf.len() != total_len {
return Err(...)
}
Ok(buf.into())There was a problem hiding this comment.
Thanks for the reply @jhorstmann , so I dug into this a bit more and agree the set_len + read_exact approach isn’t suitable for generic Read.
While looking into alternatives, I came across the newer read_buf API using BorrowedBuf, which is designed specifically for reading into uninitialized memory safely. It seems like it would let us avoid the zero-fill while still preserving alignment via MutableBuffer.
Something like:
use std::io::{Read, BorrowedBuf};
let mut buf = MutableBuffer::with_capacity(total_len);
{
let spare = buf.spare_capacity_mut();
let mut borrowed = BorrowedBuf::from(spare);
reader.read_buf_exact(borrowed)?;
}
unsafe { buf.set_len(total_len) };There was a problem hiding this comment.
so it seems to be that the read_buf is a nightly api and so may not be compatible with stable rust, so I'm unsure about it.
There was a problem hiding this comment.
Thanks @jhorstmann and @alamb — I revisited this and reworked my approach to avoid unsafe entirely.
The earlier set_len + read_exact pattern is not sound for generic Read. A fully safe alternative must also account for alignment: Arrow buffers require 64-byte alignment (arrow_buffer::buffer::ALIGNMENT), while Vec<u8> does not guarantee this.
To handle both safely:
- read into
Vec<u8>viatake(...).read_to_end(...)(works for allRead) - reuse the allocation only if it is 64-byte aligned
- otherwise copy into
MutableBuffer::from_len_zeroed(...)to ensure proper alignment
To avoid overhead on small reads, my implementation retains the existing path for smaller buffer sizes and applies the optimized path only to larger buffers (8192 bytes or higher), which constitute the majority of cases in the benchmark.
Benchmark (official ipc_reader)
cargo bench -p arrow-ipc --bench ipc_reader --features zstdStreamReader/read_10 ~910 µs → ~772 µs (~15% faster)
StreamReader/no_validation ~441 µs → ~300 µs (~32% faster)
FileReader/read_10 ~900 µs → ~773 µs (~14% faster)
FileReader/no_validation ~576 µs → ~300 µs (~48% faster)
zstd cases: small but consistent improvements
mmap: unchanged or slightly improved
pchintar
force-pushed
the
reader-zero-init
branch
from
2a3d9bc to
b8bad7c
Compare
pchintar
changed the title
|
run benchmark ipc_reader |
alamb
left a comment
alamb
left a comment
There was a problem hiding this comment.
Thanks @pchintar -- using Vec and Vec::read_exact sounds like a great optimization to me
I don't really understand the reason to keep any of the other MutableBuffer based paths though -- i think we can simplify this PR significantly (and probably make it faster) if it always uses Vec
| let mut buf = MutableBuffer::from_len_zeroed(total_len); | ||
| reader.read_exact(&mut buf)?; | ||
| Ok(buf.into()) | ||
| if total_len < 8 * 1024 { |
There was a problem hiding this comment.
Did you benchmark always using Vec? Vec is pretty fast so it might be better to always use the Vec path 🤔
| if ((vec.as_ptr() as usize) & 63) == 0 { | ||
| Ok(Buffer::from_vec(vec)) | ||
| } else { | ||
| let mut buf = MutableBuffer::from_len_zeroed(total_len); |
There was a problem hiding this comment.
I don't understand why we would ever want to do a second copy? (why not always go directly to a Buffer from. a vec?)
|
🤖 Arrow criterion benchmark running (GKE) | trigger CPU Details (lscpu)Comparing reader-zero-init (8de1dd4) to b93240a (merge-base) diff File an issue against this benchmark runner |
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: CPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
Thanks @alamb — this is very helpful. My local results were positive, but the CI benchmark clearly shows the current hybrid approach regresses on this machine, so I simplified this and used the always-`Vec' approach now. i got these results below: 
|
4 participants
Which issue does this PR close?
Rationale for this change
In
arrow-ipc/src/reader.rs, buffers are currently allocated withMutableBuffer::from_len_zeroed(len)and then passed toread_exact, which fully overwrites the entire buffer contents.This results in an unnecessary full memory pass for all data bytes:
Since
read_exactguarantees that the provided slice is fully written on success, the initial zeroing step is redundant and can be safely avoided by allocating with capacity and setting the length before the read.What changes are included in this PR?
Affected locations
read_block(file reader path)MessageReader::maybe_next(stream reader path, message body buffer)Before/Current Approach
After Approach
Rationale
read_exactfully overwrites the provided slice on success, so the initial zero-fill is redundant.This change removes one full memory write pass per read without altering behavior.
Safety
The
unsafe set_lenusage is sound because:set_lenandread_exactread_exacteither fully initializes the buffer or returns an errorCrash or interruption scenarios (panic, early return, process termination) do not introduce unsoundness, as the buffer is never observed unless
read_exactcompletes successfully by returningOk(()).Are these changes tested?
Yes the changes were tested successfully by running:
cargo test -p arrow-ipc cargo fmt --all cargo clippy --all-targets --all-features -- -D warningsAre there any user-facing changes?
No, there are no changes made to any public api code.