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);