Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rustyconover wants to merge 1 commit into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
100 changes: 89 additions & 11 deletions src/nanoarrow/ipc/reader.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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,
Expand All @@ -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) {
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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);
Expand Down
87 changes: 87 additions & 0 deletions src/nanoarrow/ipc/reader_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@

#include <stdio.h>

#include <string>

#include "nanoarrow/nanoarrow_ipc.h"

static uint8_t kSimpleSchema[] = {
Expand Down Expand Up @@ -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<const uint8_t*>(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<uintptr_t>(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);
Expand Down
5 changes: 4 additions & 1 deletion src/nanoarrow/nanoarrow_ipc.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);

Expand Down
Loading