Skip to content

feat: Add DictionaryBatch write support to IPC encoder and writer #926

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’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
paleolimbot merged 3 commits into from
Sep 8, 2026
Merged

feat: Add DictionaryBatch write support to IPC encoder and writer #926

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
b1a66b6
feat(ipc): Add DictionaryBatch write support to IPC encoder and writer
niekverw Aug 25, 2026
6938ef4
Update src/nanoarrow/ipc/writer_test.cc
paleolimbot Sep 8, 2026
840757e
ensure unsupported items fail
paleolimbot Sep 8, 2026
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
116 changes: 114 additions & 2 deletions src/nanoarrow/ipc/encoder.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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));
}

Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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;
}
71 changes: 71 additions & 0 deletions src/nanoarrow/ipc/encoder_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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);
}
85 changes: 85 additions & 0 deletions src/nanoarrow/ipc/writer.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Comment on lines +373 to +375

@paleolimbot paleolimbot Sep 5, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is a good scope for this PR, but we will want to ensure we have some way to avoid emitting duplicate dictionaries as tracked follow-up.

// 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));
}
Comment on lines +392 to +395

@paleolimbot paleolimbot Sep 5, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The encoder rejects dictionary values that have a dictionary. Should this error instead? (Or if it works, can we add a test?)


return NANOARROW_OK;
}

static ArrowErrorCode ArrowIpcWriterWriteArrayStreamImpl(
struct ArrowIpcWriter* writer, struct ArrowArrayStream* in,
struct ArrowSchema* schema, struct ArrowArray* array,
Expand All @@ -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);
}
Expand Down
Loading
Loading