Skip to content

feat(ipc): Add ArrowIpcWriterWriteArrayViewWithDictionaries #968

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): Add ArrowIpcWriterWriteArrayViewWithDictionaries #968

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
17 changes: 12 additions & 5 deletions src/nanoarrow/ipc/writer.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -408,6 +408,16 @@ static ArrowErrorCode ArrowIpcWriterWriteDictionariesForArrayView(
return NANOARROW_OK;
}

ArrowErrorCode ArrowIpcWriterWriteArrayViewWithDictionaries(
struct ArrowIpcWriter* writer, const struct ArrowArrayView* in,
struct ArrowError* error) {
NANOARROW_DCHECK(writer != NULL && writer->private_data != NULL && in != NULL);
int64_t next_dictionary_id = 0;
NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionariesForArrayView(
writer, in, &next_dictionary_id, error));
return ArrowIpcWriterWriteArrayView(writer, in, error);
}

static ArrowErrorCode ArrowIpcWriterWriteArrayStreamImpl(
struct ArrowIpcWriter* writer, struct ArrowArrayStream* in,
struct ArrowSchema* schema, struct ArrowArray* array,
Expand All @@ -424,11 +434,8 @@ static ArrowErrorCode ArrowIpcWriterWriteArrayStreamImpl(

NANOARROW_RETURN_NOT_OK(ArrowArrayViewSetArray(array_view, array, error));

int64_t next_dictionary_id = 0;
NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionariesForArrayView(
writer, array_view, &next_dictionary_id, error));

NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteArrayView(writer, array_view, error));
NANOARROW_RETURN_NOT_OK(
ArrowIpcWriterWriteArrayViewWithDictionaries(writer, array_view, error));
ArrowArrayRelease(array);
}

Expand Down
165 changes: 163 additions & 2 deletions src/nanoarrow/ipc/writer_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -303,7 +303,8 @@ TEST(NanoarrowIpcWriter, WriteDictionaryBatch) {

// Build a struct array with a single dictionary-encoded (int32 -> utf8) child.
static void MakeDictionaryStructArray(struct ArrowArray* array,
struct ArrowSchema* schema) {
struct ArrowSchema* schema,
const char* value1 = "bar") {
ASSERT_EQ(ArrowSchemaInitFromType(schema, NANOARROW_TYPE_STRUCT), NANOARROW_OK);
ASSERT_EQ(ArrowSchemaAllocateChildren(schema, 1), NANOARROW_OK);
ASSERT_EQ(ArrowSchemaInitFromType(schema->children[0], NANOARROW_TYPE_INT32),
Expand All @@ -320,7 +321,7 @@ static void MakeDictionaryStructArray(struct ArrowArray* array,

ASSERT_EQ(ArrowArrayStartAppending(array), NANOARROW_OK);
ASSERT_EQ(ArrowArrayAppendString(values, ArrowCharView("foo")), NANOARROW_OK);
ASSERT_EQ(ArrowArrayAppendString(values, ArrowCharView("bar")), NANOARROW_OK);
ASSERT_EQ(ArrowArrayAppendString(values, ArrowCharView(value1)), NANOARROW_OK);

ASSERT_EQ(ArrowArrayAppendInt(indices, 0), NANOARROW_OK);
ASSERT_EQ(ArrowArrayAppendInt(indices, 1), NANOARROW_OK);
Expand Down Expand Up @@ -726,3 +727,163 @@ TEST(NanoarrowIpcWriter, CompressedDictionaryStream) {
EXPECT_EQ(std::string(v0.data, v0.size_bytes), "foo");
EXPECT_EQ(std::string(v1.data, v1.size_bytes), "bar");
}

// The type of each message in an encapsulated IPC stream, up to the end of the stream
static std::vector<enum ArrowIpcMessageType> ReadMessageTypes(
const struct ArrowBuffer* stream_bytes) {
std::vector<enum ArrowIpcMessageType> types;
nanoarrow::ipc::UniqueDecoder decoder;
EXPECT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK);

struct ArrowError error;
struct ArrowBufferView view;
view.data.as_uint8 = stream_bytes->data;
view.size_bytes = stream_bytes->size_bytes;
while (view.size_bytes > 0) {
ArrowErrorCode result = ArrowIpcDecoderVerifyHeader(decoder.get(), view, &error);
if (result == ENODATA) {
break;
}
EXPECT_EQ(result, NANOARROW_OK) << error.message;
if (result != NANOARROW_OK) {
break;
}
types.push_back(decoder->message_type);
int64_t advance = decoder->header_size_bytes + decoder->body_size_bytes;
view.data.as_uint8 += advance;
view.size_bytes -= advance;
}
return types;
}

TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesMatchesArrayStream) {
// Writing a batch with ArrowIpcWriterWriteArrayViewWithDictionaries() produces the
// same bytes as writing it through ArrowIpcWriterWriteArrayStream()
struct ArrowError error;

nanoarrow::UniqueSchema schema;
nanoarrow::UniqueArray array;
ASSERT_NO_FATAL_FAILURE(MakeDictionaryStructArray(array.get(), schema.get()));

nanoarrow::UniqueArrayView array_view;
ASSERT_EQ(ArrowArrayViewInitFromSchema(array_view.get(), schema.get(), &error),
NANOARROW_OK)
<< error.message;
ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), array.get(), &error), NANOARROW_OK)
<< error.message;

nanoarrow::UniqueBuffer by_view;
{
nanoarrow::ipc::UniqueOutputStream stream;
ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream.get(), by_view.get()), NANOARROW_OK);
nanoarrow::ipc::UniqueWriter writer;
ASSERT_EQ(ArrowIpcWriterInit(writer.get(), stream.get()), NANOARROW_OK);
ASSERT_EQ(ArrowIpcWriterWriteSchema(writer.get(), schema.get(), &error), NANOARROW_OK)
<< error.message;
ASSERT_EQ(ArrowIpcWriterWriteArrayViewWithDictionaries(writer.get(), array_view.get(),
&error),
NANOARROW_OK)
<< error.message;
ASSERT_EQ(ArrowIpcWriterWriteArrayView(writer.get(), nullptr, &error), NANOARROW_OK);
}

nanoarrow::UniqueBuffer by_stream;
{
nanoarrow::UniqueSchema stream_schema;
ASSERT_EQ(ArrowSchemaDeepCopy(schema.get(), stream_schema.get()), NANOARROW_OK);
nanoarrow::UniqueArrayStream array_stream;
ASSERT_EQ(ArrowBasicArrayStreamInit(array_stream.get(), stream_schema.get(), 1),
NANOARROW_OK);
ArrowBasicArrayStreamSetArray(array_stream.get(), 0, array.get());

nanoarrow::ipc::UniqueOutputStream stream;
ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream.get(), by_stream.get()),
NANOARROW_OK);
nanoarrow::ipc::UniqueWriter writer;
ASSERT_EQ(ArrowIpcWriterInit(writer.get(), stream.get()), NANOARROW_OK);
ASSERT_EQ(ArrowIpcWriterWriteArrayStream(writer.get(), array_stream.get(), &error),
NANOARROW_OK)
<< error.message;
}

EXPECT_EQ(
ReadMessageTypes(by_view.get()),
(std::vector<enum ArrowIpcMessageType>{NANOARROW_IPC_MESSAGE_TYPE_SCHEMA,
NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH,
NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH}));
ASSERT_EQ(by_view->size_bytes, by_stream->size_bytes);
EXPECT_EQ(memcmp(by_view->data, by_stream->data, by_view->size_bytes), 0);
}

TEST(NanoarrowIpcWriter, WriteArrayViewWithDictionariesPerBatch) {
// Batches written one at a time each carry their own dictionary, which the reader
// resolves against the DictionaryBatch written just before it
struct ArrowError error;

nanoarrow::UniqueSchema schema;
nanoarrow::UniqueArray first;
ASSERT_NO_FATAL_FAILURE(MakeDictionaryStructArray(first.get(), schema.get(), "bar"));
nanoarrow::UniqueSchema second_schema;
nanoarrow::UniqueArray second;
ASSERT_NO_FATAL_FAILURE(
MakeDictionaryStructArray(second.get(), second_schema.get(), "baz"));

nanoarrow::UniqueArrayView array_view;
ASSERT_EQ(ArrowArrayViewInitFromSchema(array_view.get(), schema.get(), &error),
NANOARROW_OK)
<< error.message;

nanoarrow::UniqueBuffer output;
nanoarrow::ipc::UniqueOutputStream stream;
ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream.get(), output.get()), NANOARROW_OK);
nanoarrow::ipc::UniqueWriter writer;
ASSERT_EQ(ArrowIpcWriterInit(writer.get(), stream.get()), NANOARROW_OK);
ASSERT_EQ(ArrowIpcWriterWriteSchema(writer.get(), schema.get(), &error), NANOARROW_OK)
<< error.message;
for (struct ArrowArray* array : {first.get(), second.get()}) {
ASSERT_EQ(ArrowArrayViewSetArray(array_view.get(), array, &error), NANOARROW_OK)
<< error.message;
ASSERT_EQ(ArrowIpcWriterWriteArrayViewWithDictionaries(writer.get(), array_view.get(),
&error),
NANOARROW_OK)
<< error.message;
}
ASSERT_EQ(ArrowIpcWriterWriteArrayView(writer.get(), nullptr, &error), NANOARROW_OK);

EXPECT_EQ(
ReadMessageTypes(output.get()),
(std::vector<enum ArrowIpcMessageType>{NANOARROW_IPC_MESSAGE_TYPE_SCHEMA,
NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH,
NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH,
NANOARROW_IPC_MESSAGE_TYPE_DICTIONARY_BATCH,
NANOARROW_IPC_MESSAGE_TYPE_RECORD_BATCH}));

struct ArrowIpcInputStream input;
ASSERT_EQ(ArrowIpcInputStreamInitBuffer(&input, output.get()), NANOARROW_OK);
nanoarrow::UniqueArrayStream reader;
ASSERT_EQ(ArrowIpcArrayStreamReaderInit(reader.get(), &input, nullptr), NANOARROW_OK);
nanoarrow::UniqueSchema roundtrip_schema;
ASSERT_EQ(ArrowArrayStreamGetSchema(reader.get(), roundtrip_schema.get(), &error),
NANOARROW_OK)
<< error.message;

for (const char* expected : {"bar", "baz"}) {
SCOPED_TRACE(expected);
nanoarrow::UniqueArray roundtrip;
ASSERT_EQ(ArrowArrayStreamGetNext(reader.get(), roundtrip.get(), &error),
NANOARROW_OK)
<< error.message;
ASSERT_NE(roundtrip->release, nullptr);

nanoarrow::UniqueArrayView view;
ASSERT_EQ(ArrowArrayViewInitFromSchema(view.get(), roundtrip_schema.get(), &error),
NANOARROW_OK)
<< error.message;
ASSERT_EQ(ArrowArrayViewSetArray(view.get(), roundtrip.get(), &error), NANOARROW_OK)
<< error.message;
struct ArrowArrayView* values = view->children[0]->dictionary;
ASSERT_NE(values, nullptr);
struct ArrowStringView value1 = ArrowArrayViewGetStringUnsafe(values, 1);
EXPECT_EQ(std::string(value1.data, value1.size_bytes), expected);
}
}
18 changes: 18 additions & 0 deletions src/nanoarrow/nanoarrow_ipc.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,8 @@
NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcWriterWriteSchema)
#define ArrowIpcWriterWriteArrayView \
NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcWriterWriteArrayView)
#define ArrowIpcWriterWriteArrayViewWithDictionaries \
NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcWriterWriteArrayViewWithDictionaries)
#define ArrowIpcWriterWriteDictionaryBatch \
NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcWriterWriteDictionaryBatch)
#define ArrowIpcWriterWriteArrayStream \
Expand Down Expand Up @@ -1163,6 +1165,22 @@ NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteArrayView(struct ArrowIpcWriter*
const struct ArrowArrayView* in,
struct ArrowError* error);

/// \brief Write an array view and the dictionaries it references
///
/// Writes a DictionaryBatch message for each dictionary-encoded field of in, followed
/// by the RecordBatch itself, exactly as ArrowIpcWriterWriteArrayStream() does for each
/// array of a stream. This lets a caller that produces batches one at a time (rather
/// than from an ArrowArrayStream) write dictionary-encoded data. Dictionary IDs are
/// assigned in the same order as when the schema was written. Each dictionary is
/// written in full (not as a delta) before every RecordBatch, and file mode has the
/// same limitations as ArrowIpcWriterWriteArrayStream(). in must not be NULL; use
/// ArrowIpcWriterWriteArrayView() to write EOS.
///
/// Errors are propagated from the underlying encoder and output byte stream.
NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteArrayViewWithDictionaries(
struct ArrowIpcWriter* writer, const struct ArrowArrayView* in,
struct ArrowError* error);

/// \brief Write a DictionaryBatch message to the output byte stream
///
/// dictionary_id must match the id assigned to the dictionary-encoded field in the
Expand Down
Loading