From 02c9d14d4270ee037ffb7c1ac83ede2c38d3af12 Mon Sep 17 00:00:00 2001 From: Rusty Conover Date: Thu, 15 Oct 2026 21:00:00 -0400 Subject: [PATCH] feat(ipc): Reference message bodies in place when reading from a buffer 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) --- src/nanoarrow/ipc/reader.c | 100 +++++++++++++++++++++++++++---- src/nanoarrow/ipc/reader_test.cc | 87 +++++++++++++++++++++++++++ src/nanoarrow/nanoarrow_ipc.h | 5 +- 3 files changed, 180 insertions(+), 12 deletions(-) diff --git a/src/nanoarrow/ipc/reader.c b/src/nanoarrow/ipc/reader.c index c699ca331..240b2f0ab 100644 --- a/src/nanoarrow/ipc/reader.c +++ b/src/nanoarrow/ipc/reader.c @@ -74,6 +74,51 @@ static ArrowErrorCode ArrowIpcInputStreamBufferRead(struct ArrowIpcInputStream* return NANOARROW_OK; } +// Check whether the next bytes of a buffer input stream can be referenced in place. +// Arrays decoded from a message body are only aligned if the body is, so a body +// that starts at a misaligned address (e.g., because the input buffer does) is +// copied instead. +static int ArrowIpcInputStreamBufferCanReadShared(struct ArrowIpcInputStream* stream) { + if (stream->read != &ArrowIpcInputStreamBufferRead) { + return 0; + } + + struct ArrowIpcInputStreamBufferPrivate* private_data = + (struct ArrowIpcInputStreamBufferPrivate*)stream->private_data; + uintptr_t next_address = + (uintptr_t)(private_data->input.data + private_data->cursor_bytes); + return next_address % 8 == 0; +} + +// Like ArrowIpcInputStreamBufferRead(), but sets out to a reference to the input +// buffer instead of copying the bytes out of it. The input buffer remains valid +// until out (and the input stream) have been released. +static ArrowErrorCode ArrowIpcInputStreamBufferReadShared( + struct ArrowIpcInputStream* stream, int64_t buf_size_bytes, struct ArrowBuffer* out, + int64_t* size_read_out) { + struct ArrowIpcInputStreamBufferPrivate* private_data = + (struct ArrowIpcInputStreamBufferPrivate*)stream->private_data; + int64_t bytes_remaining = private_data->input.size_bytes - private_data->cursor_bytes; + int64_t bytes_to_read; + if (bytes_remaining > buf_size_bytes) { + bytes_to_read = buf_size_bytes; + } else { + bytes_to_read = bytes_remaining; + } + + ArrowBufferReset(out); + if (bytes_to_read > 0) { + NANOARROW_RETURN_NOT_OK(ArrowSharedBufferClone(&private_data->input, out)); + out->data += private_data->cursor_bytes; + out->size_bytes = bytes_to_read; + out->capacity_bytes = bytes_to_read; + } + + *size_read_out = bytes_to_read; + private_data->cursor_bytes += bytes_to_read; + return NANOARROW_OK; +} + static void ArrowIpcInputStreamBufferRelease(struct ArrowIpcInputStream* stream) { struct ArrowIpcInputStreamBufferPrivate* private_data = (struct ArrowIpcInputStreamBufferPrivate*)stream->private_data; @@ -93,7 +138,18 @@ ArrowErrorCode ArrowIpcInputStreamInitBuffer(struct ArrowIpcInputStream* stream, return ENOMEM; } - ArrowBufferMove(input, &private_data->input); + // Hold the input as a shared buffer so that the array stream reader can + // reference message bodies in place instead of copying them + if (ArrowIsSharedBuffer(input)) { + ArrowBufferMove(input, &private_data->input); + } else { + ArrowErrorCode result = ArrowSharedBufferInit(&private_data->input, input); + if (result != NANOARROW_OK) { + ArrowFree(private_data); + return result; + } + } + private_data->cursor_bytes = 0; stream->read = &ArrowIpcInputStreamBufferRead; stream->release = &ArrowIpcInputStreamBufferRelease; @@ -354,14 +410,23 @@ static int ArrowIpcArrayStreamReaderNextBody( int64_t bytes_read; int64_t bytes_to_read = private_data->decoder.body_size_bytes; - // Read the body bytes - private_data->body.size_bytes = 0; - NANOARROW_RETURN_NOT_OK_WITH_ERROR( - ArrowBufferReserve(&private_data->body, bytes_to_read), &private_data->error); - NANOARROW_RETURN_NOT_OK(private_data->input.read(&private_data->input, - private_data->body.data, bytes_to_read, - &bytes_read, &private_data->error)); - private_data->body.size_bytes += bytes_read; + if (private_data->use_shared_buffers && + ArrowIpcInputStreamBufferCanReadShared(&private_data->input)) { + // The input is already in memory: reference the body instead of copying it + NANOARROW_RETURN_NOT_OK_WITH_ERROR( + ArrowIpcInputStreamBufferReadShared(&private_data->input, bytes_to_read, + &private_data->body, &bytes_read), + &private_data->error); + } else { + // Read the body bytes + private_data->body.size_bytes = 0; + NANOARROW_RETURN_NOT_OK_WITH_ERROR( + ArrowBufferReserve(&private_data->body, bytes_to_read), &private_data->error); + NANOARROW_RETURN_NOT_OK( + private_data->input.read(&private_data->input, private_data->body.data, + bytes_to_read, &bytes_read, &private_data->error)); + private_data->body.size_bytes += bytes_read; + } if (bytes_read != bytes_to_read) { ArrowErrorSet(&private_data->error, @@ -374,6 +439,17 @@ static int ArrowIpcArrayStreamReaderNextBody( } } +// Transfer the body of the current message to a shared buffer +static ArrowErrorCode ArrowIpcArrayStreamReaderTakeSharedBody( + struct ArrowIpcArrayStreamReaderPrivate* private_data, struct ArrowBuffer* out) { + if (ArrowIsSharedBuffer(&private_data->body)) { + ArrowBufferMove(&private_data->body, out); + return NANOARROW_OK; + } + + return ArrowSharedBufferInit(out, &private_data->body); +} + static int ArrowIpcArrayStreamReaderReadSchemaIfNeeded( struct ArrowIpcArrayStreamReaderPrivate* private_data) { if (private_data->out_schema.release != NULL) { @@ -456,7 +532,8 @@ static int ArrowIpcArrayStreamReaderProcessRecordBatch( if (private_data->use_shared_buffers) { struct ArrowBuffer shared; NANOARROW_RETURN_NOT_OK_WITH_ERROR( - ArrowSharedBufferInit(&shared, &private_data->body), &private_data->error); + ArrowIpcArrayStreamReaderTakeSharedBody(private_data, &shared), + &private_data->error); ArrowErrorCode result = ArrowIpcDecoderDecodeArrayFromSharedWithDictionaries( &private_data->decoder, &shared, private_data->field_index, &private_data->dictionaries, out, NANOARROW_VALIDATION_LEVEL_FULL, @@ -486,7 +563,8 @@ static int ArrowIpcArrayStreamReaderProcessDictionary( // Decode the dictionary struct ArrowBuffer shared; NANOARROW_RETURN_NOT_OK_WITH_ERROR( - ArrowSharedBufferInit(&shared, &private_data->body), &private_data->error); + ArrowIpcArrayStreamReaderTakeSharedBody(private_data, &shared), + &private_data->error); int result = ArrowIpcDecoderDecodeDictionaryFromShared( &private_data->decoder, &shared, NANOARROW_VALIDATION_LEVEL_FULL, &private_data->dictionaries, &private_data->error); diff --git a/src/nanoarrow/ipc/reader_test.cc b/src/nanoarrow/ipc/reader_test.cc index d257e111f..129cb0ba0 100644 --- a/src/nanoarrow/ipc/reader_test.cc +++ b/src/nanoarrow/ipc/reader_test.cc @@ -19,6 +19,8 @@ #include +#include + #include "nanoarrow/nanoarrow_ipc.h" static uint8_t kSimpleSchema[] = { @@ -293,6 +295,91 @@ TEST(NanoarrowIpcReader, StreamReaderBasicNoSharedBuffers) { ArrowArrayStreamRelease(&stream); } +TEST(NanoarrowIpcReader, StreamReaderBufferInputSharedBuffers) { + // When the input is a buffer, decoded arrays should reference it in place when + // shared buffers are requested (and the input is aligned), and keep it alive + // after the stream has been released. + struct TestCase { + int use_shared_buffers; + bool input_is_shared; + bool input_is_misaligned; + bool expect_values_in_input; + }; + + alignas(8) + uint8_t misaligned_storage[1 + sizeof(kSimpleSchema) + sizeof(kSimpleRecordBatch)]; + + for (const auto& test_case : + {TestCase{1, false, false, true}, TestCase{1, true, false, true}, + TestCase{0, false, false, false}, TestCase{0, true, false, false}, + TestCase{1, false, true, false}}) { + SCOPED_TRACE( + "use_shared_buffers: " + std::to_string(test_case.use_shared_buffers) + + ", input_is_shared: " + std::to_string(test_case.input_is_shared) + + ", input_is_misaligned: " + std::to_string(test_case.input_is_misaligned)); + + struct ArrowBuffer input_buffer; + ArrowBufferInit(&input_buffer); + if (test_case.input_is_misaligned) { + // Borrow storage that starts one byte past an aligned address + uint8_t* data = misaligned_storage + 1; + memcpy(data, kSimpleSchema, sizeof(kSimpleSchema)); + memcpy(data + sizeof(kSimpleSchema), kSimpleRecordBatch, + sizeof(kSimpleRecordBatch)); + ASSERT_EQ(ArrowBufferSetAllocator( + &input_buffer, + ArrowBufferDeallocator( + [](ArrowBufferAllocator*, uint8_t*, int64_t) {}, nullptr)), + NANOARROW_OK); + input_buffer.data = data; + input_buffer.size_bytes = sizeof(kSimpleSchema) + sizeof(kSimpleRecordBatch); + } else { + ASSERT_EQ(ArrowBufferAppend(&input_buffer, kSimpleSchema, sizeof(kSimpleSchema)), + NANOARROW_OK); + ASSERT_EQ(ArrowBufferAppend(&input_buffer, kSimpleRecordBatch, + sizeof(kSimpleRecordBatch)), + NANOARROW_OK); + } + + if (test_case.input_is_shared) { + struct ArrowBuffer shared; + ASSERT_EQ(ArrowSharedBufferInit(&shared, &input_buffer), NANOARROW_OK); + ArrowBufferMove(&shared, &input_buffer); + } + + const uint8_t* input_begin = input_buffer.data; + const uint8_t* input_end = input_buffer.data + input_buffer.size_bytes; + + struct ArrowIpcInputStream input; + ASSERT_EQ(ArrowIpcInputStreamInitBuffer(&input, &input_buffer), NANOARROW_OK); + + struct ArrowArrayStream stream; + struct ArrowIpcArrayStreamReaderOptions options; + options.field_index = -1; + options.use_shared_buffers = test_case.use_shared_buffers; + ASSERT_EQ(ArrowIpcArrayStreamReaderInit(&stream, &input, &options), NANOARROW_OK); + + struct ArrowArray array; + ASSERT_EQ(ArrowArrayStreamGetNext(&stream, &array, nullptr), NANOARROW_OK); + ArrowArrayStreamRelease(&stream); + + ASSERT_EQ(array.n_children, 1); + ASSERT_EQ(array.children[0]->length, 3); + auto values = reinterpret_cast(array.children[0]->buffers[1]); + bool values_in_input = values >= input_begin && values < input_end; + EXPECT_EQ(values_in_input, test_case.expect_values_in_input); + EXPECT_EQ(reinterpret_cast(values) % 8, 0); + + int32_t int_values[3]; + memcpy(int_values, values, sizeof(int_values)); + EXPECT_EQ(int_values[0], 1); + EXPECT_EQ(int_values[1], 2); + EXPECT_EQ(int_values[2], 3); + + ArrowArrayRelease(&array); + } +} + TEST(NanoarrowIpcReader, StreamReaderBasicWithEndOfStream) { struct ArrowBuffer input_buffer; ArrowBufferInit(&input_buffer); diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h index 7f5559759..cbfabf3af 100644 --- a/src/nanoarrow/nanoarrow_ipc.h +++ b/src/nanoarrow/nanoarrow_ipc.h @@ -908,7 +908,10 @@ NANOARROW_DLL void ArrowIpcInputStreamMove(struct ArrowIpcInputStream* src, /// \brief Create an input stream from an ArrowBuffer /// -/// The stream takes ownership of the buffer and reads bytes from it. +/// The stream takes ownership of the buffer and reads bytes from it. When this +/// stream is read by an ArrowIpcArrayStreamReader that uses shared buffers, the +/// decoded arrays reference the input buffer directly instead of a copy of each +/// message body, and keep the input buffer alive until they have been released. NANOARROW_DLL ArrowErrorCode ArrowIpcInputStreamInitBuffer( struct ArrowIpcInputStream* stream, struct ArrowBuffer* input);