This is an automated email from the ASF dual-hosted git repository.
paleolimbot pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-nanoarrow.git
The following commit(s) were added to refs/heads/main by this push:
new 9dd6c4b0 feat: Add DictionaryBatch write support to IPC encoder and
writer (#926)
9dd6c4b0 is described below
commit 9dd6c4b0b84a347039e109d24f67c3692269a0c9
Author: NIEK VERWEIJ <[email protected]>
AuthorDate: Tue Sep 8 13:35:22 2026 -0400
feat: Add DictionaryBatch write support to IPC encoder and writer (#926)
Hi, I had implemented dictionary batch writing for a project that I was
working on, I noticed a thread
(https://github.com/apache/arrow-nanoarrow/issues/622) about it, so I
just wanted to share what I had in case it is useful. I needed this for
a project that emits dictionary-encoded columns over IPC. I use
ArrowIpcWriterWriteDictionaryBatch directly, for the PR I wired it into
WriteArrayStream so dictionaries get emitted automatically on the
high-level path and added a write→read round-trip test. Happy to also
adjust.
---------
Co-authored-by: niekverw <[email protected]>
Co-authored-by: Dewey Dunnington <[email protected]>
Co-authored-by: Dewey Dunnington <[email protected]>
---
src/nanoarrow/ipc/encoder.c | 116 +++++++++++++++++++++-
src/nanoarrow/ipc/encoder_test.cc | 71 ++++++++++++++
src/nanoarrow/ipc/writer.c | 85 ++++++++++++++++
src/nanoarrow/ipc/writer_test.cc | 201 ++++++++++++++++++++++++++++++++++++++
src/nanoarrow/nanoarrow_ipc.h | 29 ++++++
5 files changed, 500 insertions(+), 2 deletions(-)
diff --git a/src/nanoarrow/ipc/encoder.c b/src/nanoarrow/ipc/encoder.c
index 2c12a1b5..d1435f49 100644
--- a/src/nanoarrow/ipc/encoder.c
+++ b/src/nanoarrow/ipc/encoder.c
@@ -445,6 +445,13 @@ static ArrowErrorCode ArrowIpcEncodeField(
flatcc_builder_t* builder, const struct ArrowSchema* schema,
const struct ArrowIpcDictionaryEncodings* dictionary_encodings,
struct ArrowError* error) {
+ // Check before ArrowSchemaViewInit(), which assumes dictionary values are
not
+ // themselves dictionary-encoded.
+ if (schema->dictionary != NULL && schema->dictionary->dictionary != NULL) {
+ ArrowErrorSet(error, "IPC encoding of nested dictionary values
unsupported");
+ return ENOTSUP;
+ }
+
FLATCC_RETURN_UNLESS_0(Field_name_create_str(builder, schema->name), error);
FLATCC_RETURN_UNLESS_0(
Field_nullable_add(builder, (schema->flags & ARROW_FLAG_NULLABLE) != 0),
error);
@@ -519,6 +526,15 @@ static ArrowErrorCode ArrowIpcEncodeField(
// Add the dictionary encoding to the field
FLATCC_RETURN_UNLESS_0(Field_dictionary_add(builder, dict_encoding_ref),
error);
+ // Support dictionary values with children by encoding children from
+ // schema->dictionary (and add a roundtrip test for a nested value type).
+ // Using schema below would encode the index type's children instead and
+ // produce a Field whose type and children do not agree.
+ if (schema->dictionary->n_children != 0) {
+ ArrowErrorSet(error, "IPC encoding of dictionary values with children
unsupported");
+ return ENOTSUP;
+ }
+
NANOARROW_RETURN_NOT_OK(ArrowSchemaViewInit(&schema_view,
schema->dictionary, error));
}
@@ -689,8 +705,10 @@ static ArrowErrorCode ArrowIpcEncoderEncodeRecordBatchImpl(
}
if (array_view->dictionary != NULL) {
- ArrowErrorSet(error, "Cannot encode dictionary arrays");
- return ENOTSUP;
+ // Values live in a separate DictionaryBatch message per the Arrow IPC
spec;
+ // the parent's index node + buffers were already emitted by the caller
loop,
+ // so stop recursing here.
+ return NANOARROW_OK;
}
for (int64_t c = 0; c < array_view->n_children; ++c) {
@@ -783,6 +801,81 @@ ArrowErrorCode ArrowIpcEncoderEncodeSimpleRecordBatch(
return ArrowIpcEncoderEncodeRecordBatch(encoder, &buffer_encoder,
array_view, error);
}
+static ArrowErrorCode ArrowIpcEncoderEncodeDictionaryBatch(
+ struct ArrowIpcEncoder* encoder, struct ArrowIpcBufferEncoder*
buffer_encoder,
+ int64_t dictionary_id, char is_delta, const struct ArrowArrayView*
values_view,
+ struct ArrowError* error) {
+ NANOARROW_DCHECK(encoder != NULL && encoder->private_data != NULL &&
+ buffer_encoder != NULL && buffer_encoder->encode_buffer !=
NULL);
+ if (values_view->dictionary != NULL) {
+ ArrowErrorSet(error,
+ "DictionaryBatch values array must not itself be
dictionary-encoded");
+ return EINVAL;
+ }
+
+ struct ArrowIpcEncoderPrivate* private =
+ (struct ArrowIpcEncoderPrivate*)encoder->private_data;
+ flatcc_builder_t* builder = &private->builder;
+
+ FLATCC_RETURN_UNLESS_0(Message_start_as_root(builder), error);
+ FLATCC_RETURN_UNLESS_0(Message_version_add(builder, ns(MetadataVersion_V5)),
error);
+
+ FLATCC_RETURN_UNLESS_0(Message_header_DictionaryBatch_start(builder), error);
+ FLATCC_RETURN_UNLESS_0(DictionaryBatch_id_add(builder, dictionary_id),
error);
+ FLATCC_RETURN_UNLESS_0(DictionaryBatch_data_start(builder), error);
+ FLATCC_RETURN_UNLESS_0(RecordBatch_length_add(builder, values_view->length),
error);
+
+ NANOARROW_ASSERT_OK(ArrowBufferResize(&private->buffers, 0, 0));
+ NANOARROW_ASSERT_OK(ArrowBufferResize(&private->nodes, 0, 0));
+
+ // The values array is a single top-level column. Emit the top-level node +
+ // buffers here, then descend into any nested children.
+ struct ns(FieldNode) top_node = {values_view->length,
values_view->null_count};
+ NANOARROW_RETURN_NOT_OK_WITH_ERROR(
+ ArrowBufferAppend(&private->nodes, &top_node, sizeof(top_node)), error);
+ for (int64_t b = 0; b < values_view->array->n_buffers; ++b) {
+ struct ns(Buffer) buffer;
+ NANOARROW_RETURN_NOT_OK(buffer_encoder->encode_buffer(
+ values_view->buffer_views[b], encoder, buffer_encoder, &buffer.offset,
+ &buffer.length, error));
+ NANOARROW_RETURN_NOT_OK_WITH_ERROR(
+ ArrowBufferAppend(&private->buffers, &buffer, sizeof(buffer)), error);
+ }
+ NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderEncodeRecordBatchImpl(
+ encoder, buffer_encoder, values_view, &private->buffers,
&private->nodes, error));
+
+ FLATCC_RETURN_UNLESS_0(
+ RecordBatch_nodes_create(builder, (struct
ns(FieldNode)*)private->nodes.data,
+ private->nodes.size_bytes / sizeof(struct
ns(FieldNode))),
+ error);
+ FLATCC_RETURN_UNLESS_0(
+ RecordBatch_buffers_create(builder, (struct
ns(Buffer)*)private->buffers.data,
+ private->buffers.size_bytes / sizeof(struct
ns(Buffer))),
+ error);
+ FLATCC_RETURN_UNLESS_0(DictionaryBatch_data_end(builder), error);
+ FLATCC_RETURN_UNLESS_0(DictionaryBatch_isDelta_add(builder, is_delta ? 1 :
0), error);
+ FLATCC_RETURN_UNLESS_0(Message_header_DictionaryBatch_end(builder), error);
+ FLATCC_RETURN_UNLESS_0(Message_bodyLength_add(builder,
buffer_encoder->body_length),
+ error);
+ FLATCC_RETURN_IF_NULL(ns(Message_end_as_root(builder)), error);
+ return NANOARROW_OK;
+}
+
+ArrowErrorCode ArrowIpcEncoderEncodeSimpleDictionaryBatch(
+ struct ArrowIpcEncoder* encoder, int64_t dictionary_id, char is_delta,
+ const struct ArrowArrayView* values_view, struct ArrowBuffer* body_buffer,
+ struct ArrowError* error) {
+ NANOARROW_DCHECK(encoder != NULL && encoder->private_data != NULL &&
+ body_buffer != NULL);
+ struct ArrowIpcBufferEncoder buffer_encoder = {
+ .encode_buffer = &ArrowIpcEncoderBuildContiguousBodyBufferCallback,
+ .encode_buffer_state = body_buffer,
+ .body_length = 0,
+ };
+ return ArrowIpcEncoderEncodeDictionaryBatch(encoder, &buffer_encoder,
dictionary_id,
+ is_delta, values_view, error);
+}
+
void ArrowIpcFooterInit(struct ArrowIpcFooter* footer) {
footer->schema.release = NULL;
ArrowBufferInit(&footer->record_batch_blocks);
@@ -837,6 +930,25 @@ ArrowErrorCode ArrowIpcEncoderEncodeFooter(struct
ArrowIpcEncoder* encoder,
}
FLATCC_RETURN_UNLESS_0(Footer_recordBatches_end(builder), error);
+ const struct ArrowIpcFileBlock* dict_blocks =
+ (struct ArrowIpcFileBlock*)footer->dictionary_blocks.data;
+ int64_t n_dict_blocks =
+ footer->dictionary_blocks.size_bytes / sizeof(struct ArrowIpcFileBlock);
+
+ FLATCC_RETURN_UNLESS_0(Footer_dictionaries_start(builder), error);
+ struct ns(Block)* flatcc_dict_blocks =
+ ns(Footer_dictionaries_extend(builder, n_dict_blocks));
+ FLATCC_RETURN_IF_NULL(flatcc_dict_blocks, error);
+ for (int64_t i = 0; i < n_dict_blocks; i++) {
+ struct ns(Block) block = {
+ dict_blocks[i].offset,
+ dict_blocks[i].metadata_length,
+ dict_blocks[i].body_length,
+ };
+ flatcc_dict_blocks[i] = block;
+ }
+ FLATCC_RETURN_UNLESS_0(Footer_dictionaries_end(builder), error);
+
FLATCC_RETURN_IF_NULL(ns(Footer_end_as_root(builder)), error);
return NANOARROW_OK;
}
diff --git a/src/nanoarrow/ipc/encoder_test.cc
b/src/nanoarrow/ipc/encoder_test.cc
index 3eeff1ca..4fa0e4fe 100644
--- a/src/nanoarrow/ipc/encoder_test.cc
+++ b/src/nanoarrow/ipc/encoder_test.cc
@@ -106,6 +106,29 @@ TEST(NanoarrowIpcTest, NanoarrowIpcFooterEncoding) {
EXPECT_GT(footer_buffer->size_bytes, raw_schema_buffer->size_bytes);
}
+TEST(NanoarrowIpcTest, NanoarrowIpcEncoderRejectsNestedDictionary) {
+ nanoarrow::UniqueSchema schema;
+ ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_STRUCT),
NANOARROW_OK);
+ ASSERT_EQ(ArrowSchemaAllocateChildren(schema.get(), 1), NANOARROW_OK);
+ ASSERT_EQ(ArrowSchemaInitFromType(schema->children[0], NANOARROW_TYPE_INT32),
+ NANOARROW_OK);
+ ASSERT_EQ(ArrowSchemaAllocateDictionary(schema->children[0]), NANOARROW_OK);
+ ASSERT_EQ(
+ ArrowSchemaInitFromType(schema->children[0]->dictionary,
NANOARROW_TYPE_INT32),
+ NANOARROW_OK);
+ ASSERT_EQ(ArrowSchemaAllocateDictionary(schema->children[0]->dictionary),
NANOARROW_OK);
+
ASSERT_EQ(ArrowSchemaInitFromType(schema->children[0]->dictionary->dictionary,
+ NANOARROW_TYPE_STRING),
+ NANOARROW_OK);
+
+ nanoarrow::ipc::UniqueEncoder encoder;
+ ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK);
+
+ struct ArrowError error;
+ EXPECT_EQ(ArrowIpcEncoderEncodeSchema(encoder.get(), schema.get(), &error),
ENOTSUP);
+ EXPECT_STREQ(error.message, "IPC encoding of nested dictionary values
unsupported");
+}
+
using KeyValues = std::vector<std::pair<std::string, std::string>>;
// Unpack nanoarrow's metadata representation into something comparable
@@ -430,3 +453,51 @@ TEST(NanoarrowIpcTest,
NanoarrowIpcVisitMessageMetadataError) {
EXPECT_EQ(visited, (KeyValues{{"key1", "value1"}}));
EXPECT_STREQ(error.message, "visitor stopped at key1");
}
+
+TEST(NanoarrowIpcTest, NanoarrowIpcEncoderDictionaryBatch) {
+ nanoarrow::ipc::UniqueEncoder encoder;
+ ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK);
+
+ // Build a simple Utf8 values array
+ nanoarrow::UniqueSchema values_schema;
+ ASSERT_EQ(ArrowSchemaInitFromType(values_schema.get(),
NANOARROW_TYPE_STRING),
+ NANOARROW_OK);
+
+ nanoarrow::UniqueArray values_array;
+ ASSERT_EQ(ArrowArrayInitFromSchema(values_array.get(), values_schema.get(),
nullptr),
+ NANOARROW_OK);
+
+ struct ArrowError error;
+ ASSERT_EQ(ArrowArrayStartAppending(values_array.get()), NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayAppendString(values_array.get(), ArrowCharView("foo")),
+ NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayAppendString(values_array.get(), ArrowCharView("bar")),
+ NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayFinishBuildingDefault(values_array.get(), &error),
NANOARROW_OK)
+ << error.message;
+
+ nanoarrow::UniqueArrayView values_view;
+ ASSERT_EQ(ArrowArrayViewInitFromSchema(values_view.get(),
values_schema.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(ArrowArrayViewSetArray(values_view.get(), values_array.get(),
&error),
+ NANOARROW_OK)
+ << error.message;
+
+ // Encode a non-delta DictionaryBatch with dictionary_id=0
+ nanoarrow::UniqueBuffer body_buffer;
+ EXPECT_EQ(ArrowIpcEncoderEncodeSimpleDictionaryBatch(encoder.get(),
/*dictionary_id=*/0,
+ /*is_delta=*/0,
values_view.get(),
+ body_buffer.get(),
&error),
+ NANOARROW_OK)
+ << error.message;
+
+ nanoarrow::UniqueBuffer message_buffer;
+ EXPECT_EQ(ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/1,
+ message_buffer.get()),
+ NANOARROW_OK);
+
+ // The encapsulated message must be non-empty and 8-byte aligned
+ EXPECT_GT(message_buffer->size_bytes, 8);
+ EXPECT_EQ(message_buffer->size_bytes % 8, 0);
+}
diff --git a/src/nanoarrow/ipc/writer.c b/src/nanoarrow/ipc/writer.c
index 1f100865..00ea93ec 100644
--- a/src/nanoarrow/ipc/writer.c
+++ b/src/nanoarrow/ipc/writer.c
@@ -317,6 +317,86 @@ ArrowErrorCode ArrowIpcWriterWriteArrayView(struct
ArrowIpcWriter* writer,
return NANOARROW_OK;
}
+ArrowErrorCode ArrowIpcWriterWriteDictionaryBatch(
+ struct ArrowIpcWriter* writer, int64_t dictionary_id, char is_delta,
+ const struct ArrowArrayView* values_view, struct ArrowError* error) {
+ NANOARROW_DCHECK(writer != NULL && writer->private_data != NULL &&
values_view != NULL);
+ struct ArrowIpcWriterPrivate* private =
+ (struct ArrowIpcWriterPrivate*)writer->private_data;
+
+ // This check is intentionally minimal: we're allowed to write one dictionary
+ // batch per ID in a file but we would need to add bookkeeping to keep track
+ // of written IDs (and usefully a fingerprint or reference to the dictionary
+ // so we can check if we need to emit it again).
+ if (private->writing_file &&
+ (is_delta || private->footer.dictionary_blocks.size_bytes != 0)) {
+ ArrowErrorSet(error,
+ "IPC file writing supports exactly one non-delta dictionary
batch");
+ return ENOTSUP;
+ }
+
+ NANOARROW_ASSERT_OK(ArrowBufferResize(&private->buffer, 0, 0));
+ NANOARROW_ASSERT_OK(ArrowBufferResize(&private->body_buffer, 0, 0));
+
+ NANOARROW_RETURN_NOT_OK(ArrowIpcEncoderEncodeSimpleDictionaryBatch(
+ &private->encoder, dictionary_id, is_delta, values_view,
&private->body_buffer,
+ error));
+ NANOARROW_RETURN_NOT_OK_WITH_ERROR(
+ ArrowIpcEncoderFinalizeBuffer(&private->encoder, /*encapsulate=*/1,
+ &private->buffer),
+ error);
+
+ if (private->writing_file) {
+ _NANOARROW_CHECK_RANGE(private->buffer.size_bytes, 0, INT32_MAX);
+ struct ArrowIpcFileBlock block = {
+ .offset = private->bytes_written,
+ .metadata_length = (int32_t) private->buffer.size_bytes,
+ .body_length = private->body_buffer.size_bytes,
+ };
+ NANOARROW_RETURN_NOT_OK_WITH_ERROR(
+ ArrowBufferAppend(&private->footer.dictionary_blocks, &block,
sizeof(block)),
+ error);
+ }
+ private->bytes_written += private->buffer.size_bytes;
+ private->bytes_written += private->body_buffer.size_bytes;
+
+ NANOARROW_RETURN_NOT_OK(ArrowIpcOutputStreamWrite(
+ &private->output_stream, ArrowBufferToBufferView(&private->buffer),
error));
+ NANOARROW_RETURN_NOT_OK(ArrowIpcOutputStreamWrite(
+ &private->output_stream, ArrowBufferToBufferView(&private->body_buffer),
error));
+ return NANOARROW_OK;
+}
+
+// Walk the array in the same depth-first order the schema encoder uses to
assign
+// dictionary ids (see ArrowIpcDictionaryEncodingsAppendSchema): a
dictionary-encoded
+// node claims the next id before descending into its children and then its
values.
+// Emitting a full (non-delta) DictionaryBatch for each dictionary before every
+// RecordBatch keeps each batch's indices valid against the dictionary that
precedes
+// it, which is required because each array in the stream carries its own
dictionary.
+// In the future we can reduce the number of dictionaries emitted by checking
for
+// identical dictionary arrays.
+static ArrowErrorCode ArrowIpcWriterWriteDictionariesForArrayView(
+ struct ArrowIpcWriter* writer, const struct ArrowArrayView* array_view,
+ int64_t* next_id, struct ArrowError* error) {
+ if (array_view->dictionary != NULL) {
+ int64_t dictionary_id = (*next_id)++;
+ NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionaryBatch(
+ writer, dictionary_id, /*is_delta=*/0, array_view->dictionary, error));
+ }
+
+ for (int64_t i = 0; i < array_view->n_children; i++) {
+ NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionariesForArrayView(
+ writer, array_view->children[i], next_id, error));
+ }
+
+ if (array_view->dictionary != NULL) {
+ NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteDictionariesForArrayView(
+ writer, array_view->dictionary, next_id, error));
+ }
+
+ return NANOARROW_OK;
+}
+
static ArrowErrorCode ArrowIpcWriterWriteArrayStreamImpl(
struct ArrowIpcWriter* writer, struct ArrowArrayStream* in,
struct ArrowSchema* schema, struct ArrowArray* array,
@@ -332,6 +412,11 @@ 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));
ArrowArrayRelease(array);
}
diff --git a/src/nanoarrow/ipc/writer_test.cc b/src/nanoarrow/ipc/writer_test.cc
index a07ae516..c30dd603 100644
--- a/src/nanoarrow/ipc/writer_test.cc
+++ b/src/nanoarrow/ipc/writer_test.cc
@@ -204,3 +204,204 @@ TEST(NanoarrowIpcWriter, FileWriting) {
auto after_footer = p->bytes_written;
EXPECT_GT(after_footer, after_eos);
}
+
+TEST(NanoarrowIpcWriter, WriteDictionaryBatch) {
+ struct ArrowError error;
+
+ 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);
+
+ auto* p = static_cast<struct ArrowIpcWriterPrivate*>(writer->private_data);
+
+ // Build a simple Utf8 values array
+ nanoarrow::UniqueSchema values_schema;
+ ASSERT_EQ(ArrowSchemaInitFromType(values_schema.get(),
NANOARROW_TYPE_STRING),
+ NANOARROW_OK);
+
+ nanoarrow::UniqueArray values_array;
+ ASSERT_EQ(ArrowArrayInitFromSchema(values_array.get(), values_schema.get(),
nullptr),
+ NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayStartAppending(values_array.get()), NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayAppendString(values_array.get(), ArrowCharView("foo")),
+ NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayAppendString(values_array.get(), ArrowCharView("bar")),
+ NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayFinishBuildingDefault(values_array.get(), &error),
NANOARROW_OK)
+ << error.message;
+
+ nanoarrow::UniqueArrayView values_view;
+ ASSERT_EQ(ArrowArrayViewInitFromSchema(values_view.get(),
values_schema.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(ArrowArrayViewSetArray(values_view.get(), values_array.get(),
&error),
+ NANOARROW_OK)
+ << error.message;
+
+ // stream mode: write a DictionaryBatch — bytes are emitted but no block is
tracked
+ EXPECT_EQ(p->bytes_written, 0);
+ EXPECT_EQ(p->footer.dictionary_blocks.size_bytes, 0);
+
+ EXPECT_EQ(ArrowIpcWriterWriteDictionaryBatch(writer.get(),
/*dictionary_id=*/0,
+ /*is_delta=*/0,
values_view.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+
+ auto after_dict_stream = p->bytes_written;
+ EXPECT_GT(after_dict_stream, 0);
+ // no block tracked in stream mode
+ EXPECT_EQ(p->footer.dictionary_blocks.size_bytes, 0);
+
+ // file mode: the block is tracked in the footer
+ nanoarrow::ipc::UniqueOutputStream stream2;
+ nanoarrow::UniqueBuffer output2;
+ ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream2.get(), output2.get()),
NANOARROW_OK);
+
+ nanoarrow::ipc::UniqueWriter writer2;
+ ASSERT_EQ(ArrowIpcWriterInit(writer2.get(), stream2.get()), NANOARROW_OK);
+
+ auto* p2 = static_cast<struct ArrowIpcWriterPrivate*>(writer2->private_data);
+
+ ASSERT_EQ(ArrowIpcWriterStartFile(writer2.get(), &error), NANOARROW_OK)
+ << error.message;
+ EXPECT_EQ(p2->footer.dictionary_blocks.size_bytes, 0);
+
+ EXPECT_EQ(ArrowIpcWriterWriteDictionaryBatch(writer2.get(),
/*dictionary_id=*/0,
+ /*is_delta=*/0,
values_view.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+
+ // one block tracked in file mode
+ EXPECT_EQ(p2->footer.dictionary_blocks.size_bytes, sizeof(struct
ArrowIpcFileBlock));
+
+ int64_t bytes_written = p2->bytes_written;
+ EXPECT_EQ(ArrowIpcWriterWriteDictionaryBatch(writer2.get(),
/*dictionary_id=*/0,
+ /*is_delta=*/0,
values_view.get(), &error),
+ ENOTSUP);
+ EXPECT_STREQ(error.message,
+ "IPC file writing supports exactly one non-delta dictionary
batch");
+ EXPECT_EQ(p2->bytes_written, bytes_written);
+ EXPECT_EQ(p2->footer.dictionary_blocks.size_bytes, sizeof(struct
ArrowIpcFileBlock));
+
+ nanoarrow::ipc::UniqueOutputStream stream3;
+ nanoarrow::UniqueBuffer output3;
+ ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(stream3.get(), output3.get()),
NANOARROW_OK);
+
+ nanoarrow::ipc::UniqueWriter writer3;
+ ASSERT_EQ(ArrowIpcWriterInit(writer3.get(), stream3.get()), NANOARROW_OK);
+ ASSERT_EQ(ArrowIpcWriterStartFile(writer3.get(), &error), NANOARROW_OK)
+ << error.message;
+ EXPECT_EQ(ArrowIpcWriterWriteDictionaryBatch(writer3.get(),
/*dictionary_id=*/0,
+ /*is_delta=*/1,
values_view.get(), &error),
+ ENOTSUP);
+ EXPECT_STREQ(error.message,
+ "IPC file writing supports exactly one non-delta dictionary
batch");
+}
+
+// Build a struct array with a single dictionary-encoded (int32 -> utf8) child.
+static void MakeDictionaryStructArray(struct ArrowArray* array,
+ struct ArrowSchema* schema) {
+ 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),
+ NANOARROW_OK);
+ ASSERT_EQ(ArrowSchemaSetName(schema->children[0], "dict_col"), NANOARROW_OK);
+ ASSERT_EQ(ArrowSchemaAllocateDictionary(schema->children[0]), NANOARROW_OK);
+ ASSERT_EQ(
+ ArrowSchemaInitFromType(schema->children[0]->dictionary,
NANOARROW_TYPE_STRING),
+ NANOARROW_OK);
+
+ ASSERT_EQ(ArrowArrayInitFromSchema(array, schema, nullptr), NANOARROW_OK);
+ struct ArrowArray* indices = array->children[0];
+ struct ArrowArray* values = indices->dictionary;
+
+ ASSERT_EQ(ArrowArrayStartAppending(array), NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayAppendString(values, ArrowCharView("foo")),
NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayAppendString(values, ArrowCharView("bar")),
NANOARROW_OK);
+
+ ASSERT_EQ(ArrowArrayAppendInt(indices, 0), NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayAppendInt(indices, 1), NANOARROW_OK);
+ ASSERT_EQ(ArrowArrayAppendInt(indices, 0), NANOARROW_OK);
+ array->length = 3;
+
+ ASSERT_EQ(ArrowArrayFinishBuildingDefault(array, nullptr), NANOARROW_OK);
+}
+
+// Write a dictionary-encoded stream through the high-level WriteArrayStream
path
+// and read it back through the IPC reader, confirming the DictionaryBatch is
+// emitted automatically and the decoded values match.
+TEST(NanoarrowIpcWriter, RoundtripDictionaryStream) {
+ struct ArrowError error;
+
+ nanoarrow::UniqueSchema schema;
+ nanoarrow::UniqueArray array;
+ MakeDictionaryStructArray(array.get(), schema.get());
+
+ nanoarrow::UniqueArrayStream array_stream;
+ ASSERT_EQ(ArrowBasicArrayStreamInit(array_stream.get(), schema.get(), 1),
NANOARROW_OK);
+ ArrowBasicArrayStreamSetArray(array_stream.get(), 0, array.get());
+
+ nanoarrow::UniqueBuffer output;
+ nanoarrow::ipc::UniqueOutputStream out_stream;
+ ASSERT_EQ(ArrowIpcOutputStreamInitBuffer(out_stream.get(), output.get()),
NANOARROW_OK);
+
+ nanoarrow::ipc::UniqueWriter writer;
+ ASSERT_EQ(ArrowIpcWriterInit(writer.get(), out_stream.get()), NANOARROW_OK);
+ ASSERT_EQ(ArrowIpcWriterWriteArrayStream(writer.get(), array_stream.get(),
&error),
+ NANOARROW_OK)
+ << error.message;
+
+ // Read the encoded bytes back
+ 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;
+ ASSERT_EQ(roundtrip_schema->n_children, 1);
+ ASSERT_NE(roundtrip_schema->children[0]->dictionary, nullptr);
+ EXPECT_STREQ(roundtrip_schema->children[0]->dictionary->format, "u");
+
+ nanoarrow::UniqueArray roundtrip_array;
+ ASSERT_EQ(ArrowArrayStreamGetNext(reader.get(), roundtrip_array.get(),
&error),
+ NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(roundtrip_array->length, 3);
+ ASSERT_EQ(roundtrip_array->n_children, 1);
+ ASSERT_NE(roundtrip_array->children[0]->dictionary, nullptr);
+ EXPECT_EQ(roundtrip_array->children[0]->dictionary->length, 2);
+
+ // Validate the decoded indices resolve to the original values
+ nanoarrow::UniqueArrayView view;
+ ASSERT_EQ(ArrowArrayViewInitFromSchema(view.get(), roundtrip_schema.get(),
&error),
+ NANOARROW_OK)
+ << error.message;
+ ASSERT_EQ(ArrowArrayViewSetArray(view.get(), roundtrip_array.get(), &error),
+ NANOARROW_OK)
+ << error.message;
+
+ struct ArrowArrayView* indices_view = view->children[0];
+ struct ArrowArrayView* values_view = indices_view->dictionary;
+ ASSERT_NE(values_view, nullptr);
+ EXPECT_EQ(ArrowArrayViewGetIntUnsafe(indices_view, 0), 0);
+ EXPECT_EQ(ArrowArrayViewGetIntUnsafe(indices_view, 1), 1);
+ EXPECT_EQ(ArrowArrayViewGetIntUnsafe(indices_view, 2), 0);
+
+ struct ArrowStringView v0 = ArrowArrayViewGetStringUnsafe(values_view, 0);
+ struct ArrowStringView v1 = ArrowArrayViewGetStringUnsafe(values_view, 1);
+ EXPECT_EQ(std::string(v0.data, v0.size_bytes), "foo");
+ EXPECT_EQ(std::string(v1.data, v1.size_bytes), "bar");
+
+ roundtrip_array.reset();
+ ASSERT_EQ(ArrowArrayStreamGetNext(reader.get(), roundtrip_array.get(),
&error),
+ NANOARROW_OK)
+ << error.message;
+ EXPECT_EQ(roundtrip_array->release, nullptr);
+}
diff --git a/src/nanoarrow/nanoarrow_ipc.h b/src/nanoarrow/nanoarrow_ipc.h
index fb38750c..0523ac38 100644
--- a/src/nanoarrow/nanoarrow_ipc.h
+++ b/src/nanoarrow/nanoarrow_ipc.h
@@ -96,6 +96,8 @@
NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcEncoderEncodeSchema)
#define ArrowIpcEncoderEncodeSimpleRecordBatch \
NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcEncoderEncodeSimpleRecordBatch)
+#define ArrowIpcEncoderEncodeSimpleDictionaryBatch \
+ NANOARROW_SYMBOL(NANOARROW_NAMESPACE,
ArrowIpcEncoderEncodeSimpleDictionaryBatch)
#define ArrowIpcOutputStreamInitBuffer \
NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcOutputStreamInitBuffer)
#define ArrowIpcOutputStreamInitFile \
@@ -110,6 +112,8 @@
NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcWriterWriteSchema)
#define ArrowIpcWriterWriteArrayView \
NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcWriterWriteArrayView)
+#define ArrowIpcWriterWriteDictionaryBatch \
+ NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcWriterWriteDictionaryBatch)
#define ArrowIpcWriterWriteArrayStream \
NANOARROW_SYMBOL(NANOARROW_NAMESPACE, ArrowIpcWriterWriteArrayStream)
#define ArrowIpcWriterStartFile \
@@ -881,6 +885,20 @@ NANOARROW_DLL ArrowErrorCode
ArrowIpcEncoderEncodeSimpleRecordBatch(
struct ArrowIpcEncoder* encoder, const struct ArrowArrayView* array_view,
struct ArrowBuffer* body_buffer, struct ArrowError* error);
+/// \brief Encode an ArrayView as a DictionaryBatch flatbuffer, embedded in a
Message.
+///
+/// dictionary_id must match the id assigned to the dictionary-encoded field
in the
+/// schema. is_delta selects DictionaryBatch.isDelta. values_view must not
itself be
+/// dictionary-encoded. Body buffers are concatenated into a contiguous, padded
+/// body_buffer.
+///
+/// Returns ENOMEM if allocation fails, EINVAL if values_view is
dictionary-encoded,
+/// NANOARROW_OK otherwise.
+NANOARROW_DLL ArrowErrorCode ArrowIpcEncoderEncodeSimpleDictionaryBatch(
+ struct ArrowIpcEncoder* encoder, int64_t dictionary_id, char is_delta,
+ const struct ArrowArrayView* values_view, struct ArrowBuffer* body_buffer,
+ struct ArrowError* error);
+
/// \brief An user-extensible output data sink
struct ArrowIpcOutputStream {
/// \brief Write up to buf_size_bytes from stream into buf
@@ -964,6 +982,17 @@ NANOARROW_DLL ArrowErrorCode
ArrowIpcWriterWriteArrayView(struct ArrowIpcWriter*
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
+/// schema. is_delta selects DictionaryBatch.isDelta. values_view must not
itself be
+/// dictionary-encoded. The writer does not check that a schema was already
written.
+///
+/// Errors are propagated from the underlying encoder and output byte stream.
+NANOARROW_DLL ArrowErrorCode ArrowIpcWriterWriteDictionaryBatch(
+ struct ArrowIpcWriter* writer, int64_t dictionary_id, char is_delta,
+ const struct ArrowArrayView* values_view, struct ArrowError* error);
+
/// \brief Write an entire stream (including EOS) to the output byte stream
///
/// Errors are propagated from the underlying encoder, array stream, and
output byte
