Skip to content

feat(ipc): Reference message bodies in place when reading from a buffer - #965

Open
rustyconover wants to merge 1 commit into
apache:mainfrom
rustyconover:ipc-reader-zero-copy-buffer
Open

rustyconover wants to merge 1 commit into
apache:mainfrom
rustyconover:ipc-reader-zero-copy-buffer

Conversation

@rustyconover

Copy link
Copy Markdown
Contributor

When ArrowIpcArrayStreamReader reads from an input stream created with ArrowIpcInputStreamInitBuffer(), it copies every message body out of the input buffer into a newly allocated buffer before decoding it. With shared buffers, which are the default wherever C11 atomics are available, the decoded arrays then reference that copy. The bytes were already in memory and owned by the stream, so the copy isn't needed.

With this PR, the buffer input stream holds its input as a shared buffer, or takes it over directly if it already is one. When the reader uses shared buffers, it clones and slices each message body out of that buffer. This is the same way ArrowIpcDecoderDecodeArrayFromShared() already slices a body into array buffers. There's no API change. File input, other input streams and use_shared_buffers = 0 behave exactly as before.

Things worth a look in review:

  • Decoded arrays now keep the whole input buffer alive, not just their own message body. That's the tradeoff use_shared_buffers already makes within a message, and the input was already resident. I've noted it in the ArrowIpcInputStreamInitBuffer() docs.
  • A body is only referenced in place if it starts at an 8-byte-aligned address. Bodies are 8-byte aligned within a valid stream, so this only fails when the caller's buffer itself starts at an odd address. In that case the reader falls back to copying, so decoded arrays stay as aligned as they are today.
  • R's read_nanoarrow() on a raw vector goes through this path with default options, so it benefits too. Its buffers are released through the preserve registry, which already handles release from other threads. I don't have R set up locally, so I'm relying on CI for the R tests.

Timings come from a harness that does what BaseBenchmarIpcFixtureBuffer() in dev/benchmarks does: it reads the whole stream from a borrowed in-memory buffer. Fixtures are from generate-fixtures.py, plus a 23 MB stream of 1M rows (int32, double, string). Release build, Apple M3, median of 50 runs:

fixture main this PR
float64_basic (10.5 MB, 2 batches) 0.52 ms 0.004 ms
float64_long (10.5 MB, 20 batches) 0.16 ms 0.006 ms
float64_wide (10.6 MB, 1280 columns) 0.58 ms 0.35 ms
1M rows int32/double/string 0.66 ms 0.28 ms

For the float64 fixtures, the copy was nearly the entire cost of reading. What's left on the wide fixture is per-column decoding, and what's left on the mixed fixture is full validation of the string offsets. Together with #964, reading the mixed fixture at NANOARROW_VALIDATION_LEVEL_DEFAULT takes 0.010 ms. That's the case I care about: a WebAssembly build of DuckDB that I'm moving from Arrow C++ to nanoarrow, where Arrow JS hands it IPC streams that are already in memory. I'm happy to add the mixed fixture to dev/benchmarks if that's useful.

Tests: StreamReaderBufferInputSharedBuffers covers these cases:

  • With shared buffers, decoded values point into the input buffer.
  • Without shared buffers, or with a misaligned input, they are copied (and still aligned).
  • An input that is already a shared buffer works.
  • Arrays stay valid after the stream has been released.

The test fails on main. ctest passes locally (macOS arm64, Debug, extra warnings). Under ASan/UBSan, the IPC and integration tests passed 10 repeated runs. The only failures in the full ASan suite are the six allocation tests that also fail on main under ASan.

🤖 Generated with Claude Code

When an ArrowIpcArrayStreamReader using shared buffers reads from an
input stream created with ArrowIpcInputStreamInitBuffer(), slice each
message body out of the input buffer instead of copying it. The input
buffer is held as a shared buffer, so arrays decoded from it keep it
alive after the stream has been released.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@codecov-commenter

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 75.00000% with 10 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.24%. Comparing base (ea215b0) to head (02c9d14).

Files with missing lines Patch % Lines
src/nanoarrow/ipc/reader.c 75.00% 3 Missing and 7 partials ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #965      +/-   ##
==========================================
- Coverage   78.24%   78.24%   -0.01%     
==========================================
  Files         106      106              
  Lines       16867    16902      +35     
  Branches     1993     2002       +9     
==========================================
+ Hits        13198    13225      +27     
- Misses       2437     2440       +3     
- Partials     1232     1237       +5     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@paleolimbot paleolimbot left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a great idea; however, I think we should add a callback to the ArrowIpcInputStream (or update the existing one) to handle this case, which comes up when reading files. Maybe something like

    ArrowErrorCode (*peek_at)(struct ArrowIpcInputStream* stream,
                         int64_t offset_from_current,
                         int64_t buf_size_bytes,
                         const uint8_t** buf,
                         struct ArrowError* error);
    ArrowErrorCode (*read_at)(struct ArrowIpcInputStream* stream,
                         int64_t offset_from_current,
                         int64_t buf_size_bytes, struct ArrowBuffer* out,
                         struct ArrowError* error);

Those callbacks can return ENOTSUP if IO would otherwise be required but the reader could try them (the first variant for metadata to avoid a copy, and the second variant for buffers). There's one other case we eventually have to handle which is firing off a bunch of concurrentread_at and waiting for the result; however, I don't think we need to handle that here.

Also, while bodies usually contain aligned buffers when the bodies themselves are aligned, there's no guarantee of this. We can provide a separate ArrowArrayEnsureAligned() that recursively walks the array and ensures buffers are aligned to at least the fixed width of their type.

@rustyconover

Copy link
Copy Markdown
Contributor Author

Thanks, that makes sense: callbacks on the input stream are a better fit than special-casing the buffer stream, and they cover files too. Two questions before I rework it.

How should the struct gain the callbacks? ArrowIpcInputStream is allocated by callers who fill in only read, release and private_data. The R package does this: r/src/ipc.c ArrowMallocs the struct and sets read/release one at a time. If peek_at/read_at are appended to the struct, those streams (and any out-of-tree ones) will have garbage in the new members, and the reader would call them. Options I can see:

  1. Append the members, document that unsupported callbacks must be NULL, and update R/Python in this PR. That's simple, but out-of-tree implementations silently get undefined behaviour until they're updated.
  2. Keep ArrowIpcInputStream as it is and add a separate struct for random access (peek_at, read_at, release, private_data). The buffer and file inputs implement both, and a reader init function takes the new one. That's more API, but nothing existing changes.

My preference is (2), unless nanoarrow is fine with breaking out-of-tree streams in a release.

Semantics. For peek_at, I'd make the returned pointer valid until the next call on the stream. For read_at, the ArrowBuffer would be a shared slice when the bytes are already in memory, and a newly read buffer otherwise. Both would advance the stream position only when offset_from_current is 0 and the full length is consumed. Does that match what you had in mind, or should offsets be absolute?

Alignment: agreed. I'd write ArrowArrayEnsureAligned() first as its own small core PR: walk the array and copy any buffer that isn't aligned to its type's fixed width. The reader would then call it on arrays that reference a body in place, replacing the 8-byte body-alignment check here. Does that split work for you?

@paleolimbot

Copy link
Copy Markdown
Member

It's a good point, and (2) does seem safer since it's not an obvious compilation break. I am not sure how that interacts with the existing reader in terms of what needs to be rewritten specifically for the random access reader. As long as this doesn't add a fully duplicate path to the array stream reader (which I think can be done), go for it!

Both would advance the stream position

advancing the stream position I think is separate (probably need a seek)

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants