Skip to content

feat(ipc): add dictionary batch write support #928

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

Open
rustyconover wants to merge 11 commits into
base: main
Choose a base branch
from
Open

feat(ipc): add dictionary batch write support #928

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
5e364cd
feat(ipc): Add DictionaryBatch write support to IPC encoder and writer
niekverw Aug 25, 2026
251577a
fix(ipc): attach metadata to dictionary batches
rustyconover Sep 1, 2026
717cf57
fix(ipc): avoid repeating unchanged dictionaries
rustyconover Sep 1, 2026
ba6c18d
fix(ipc): advertise dictionary replacements
rustyconover Sep 1, 2026
72dd083
fix(ipc): encode nested dictionary value fields
rustyconover Sep 2, 2026
b5d41fd
fix(ipc): write nested dictionaries dependency first
rustyconover Sep 2, 2026
1fb1d40
fix(ipc): invalidate nested dictionary dependents
rustyconover Sep 2, 2026
ce60996
fix(ipc): reject dictionary mutation in files
rustyconover Sep 2, 2026
522f9ac
test(ipc): enable dictionary write integration cases
rustyconover Sep 2, 2026
a8268d8
perf(ipc): reduce dictionary cache copies
rustyconover Sep 2, 2026
0951095
test(ipc): finish dictionary writer coverage
rustyconover Sep 2, 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
142 changes: 130 additions & 12 deletions src/nanoarrow/ipc/encoder.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ struct ArrowIpcEncoderPrivate {
struct ArrowBuffer buffers;
struct ArrowBuffer nodes;
int encoding_footer;
int dictionary_replacement;
struct ArrowIpcDictionaryEncodings dictionary_encodings;
// Metadata to attach to the next encoded Message (in nanoarrow's packed
// representation), or an empty buffer if the next Message has no metadata.
Expand All @@ -65,6 +66,7 @@ ArrowErrorCode ArrowIpcEncoderInit(struct ArrowIpcEncoder* encoder) {
return ESPIPE;
}
private->encoding_footer = 0;
private->dictionary_replacement = 0;
ArrowBufferInit(&private->buffers);
ArrowBufferInit(&private->nodes);
ArrowIpcDictionaryEncodingsInit(&private->dictionary_encodings);
Expand Down Expand Up @@ -119,6 +121,14 @@ ArrowErrorCode ArrowIpcEncoderSetMessageMetadata(struct ArrowIpcEncoder* encoder
return NANOARROW_OK;
}

void ArrowIpcEncoderSetDictionaryReplacement(struct ArrowIpcEncoder* encoder,
char enabled) {
NANOARROW_DCHECK(encoder != NULL && encoder->private_data != NULL);
struct ArrowIpcEncoderPrivate* private =
(struct ArrowIpcEncoderPrivate*)encoder->private_data;
private->dictionary_replacement = enabled != 0;
}

static ArrowErrorCode ArrowIpcEncoderWriteContinuationAndSize(struct ArrowBuffer* out,
size_t size) {
_NANOARROW_CHECK_UPPER_LIMIT(size, INT32_MAX);
Expand Down Expand Up @@ -451,6 +461,7 @@ static ArrowErrorCode ArrowIpcEncodeField(

struct ArrowSchemaView schema_view;
NANOARROW_RETURN_NOT_OK(ArrowSchemaViewInit(&schema_view, schema, error));
const struct ArrowSchema* value_schema = schema;

if (schema_view.type == NANOARROW_TYPE_DICTIONARY) {
const struct ArrowIpcDictionaryEncoding* encoding =
Expand Down Expand Up @@ -519,15 +530,16 @@ static ArrowErrorCode ArrowIpcEncodeField(
// Add the dictionary encoding to the field
FLATCC_RETURN_UNLESS_0(Field_dictionary_add(builder, dict_encoding_ref), error);

NANOARROW_RETURN_NOT_OK(ArrowSchemaViewInit(&schema_view, schema->dictionary, error));
value_schema = schema->dictionary;
NANOARROW_RETURN_NOT_OK(ArrowSchemaViewInit(&schema_view, value_schema, error));
}

NANOARROW_RETURN_NOT_OK(ArrowIpcEncodeFieldType(builder, &schema_view, error));

if (schema->n_children != 0) {
if (value_schema->n_children != 0) {
FLATCC_RETURN_UNLESS_0(Field_children_start(builder), error);
NANOARROW_RETURN_NOT_OK(
ArrowIpcEncodeFields(builder, schema, &ns(Field_children_push_start),
ArrowIpcEncodeFields(builder, value_schema, &ns(Field_children_push_start),
&ns(Field_children_push_end), dictionary_encodings, error));
FLATCC_RETURN_UNLESS_0(Field_children_end(builder), error);
}
Expand All @@ -545,7 +557,7 @@ static ArrowErrorCode ArrowIpcEncodeField(
static ArrowErrorCode ArrowIpcEncodeSchema(
flatcc_builder_t* builder, const struct ArrowSchema* schema,
const struct ArrowIpcDictionaryEncodings* dictionary_encodings,
struct ArrowError* error) {
int dictionary_replacement, struct ArrowError* error) {
NANOARROW_DCHECK(schema->release != NULL);

if (strcmp(schema->format, "+s") != 0) {
Expand Down Expand Up @@ -576,8 +588,13 @@ static ArrowErrorCode ArrowIpcEncodeSchema(
}
FLATCC_RETURN_UNLESS_0(Schema_custom_metadata_end(builder), error);

FLATCC_RETURN_UNLESS_0(Schema_features_start(builder), error);
FLATCC_RETURN_UNLESS_0(Schema_features_end(builder), error);
if (dictionary_replacement && dictionary_encodings->encodings.size_bytes > 0) {
ns(Feature_enum_t) feature = ns(Feature_DICTIONARY_REPLACEMENT);
FLATCC_RETURN_UNLESS_0(Schema_features_create(builder, &feature, 1), error);
} else {
FLATCC_RETURN_UNLESS_0(Schema_features_start(builder), error);
FLATCC_RETURN_UNLESS_0(Schema_features_end(builder), error);
}

return NANOARROW_OK;
}
Expand Down Expand Up @@ -607,8 +624,9 @@ ArrowErrorCode ArrowIpcEncoderEncodeSchema(struct ArrowIpcEncoder* encoder,
ArrowIpcDictionaryEncodingsAppendSchema(&private->dictionary_encodings, schema),
error);

NANOARROW_RETURN_NOT_OK(
ArrowIpcEncodeSchema(builder, schema, &private->dictionary_encodings, error));
NANOARROW_RETURN_NOT_OK(ArrowIpcEncodeSchema(builder, schema,
&private->dictionary_encodings,
private->dictionary_replacement, error));

FLATCC_RETURN_UNLESS_0(Message_header_Schema_end(builder), error);

Expand Down Expand Up @@ -689,8 +707,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 +803,84 @@ 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);

NANOARROW_RETURN_NOT_OK(ArrowIpcEncodeMessageMetadata(private, 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 @@ -814,8 +912,9 @@ ArrowErrorCode ArrowIpcEncoderEncodeFooter(struct ArrowIpcEncoder* encoder,
FLATCC_RETURN_UNLESS_0(Footer_version_add(builder, ns(MetadataVersion_V5)), error);

FLATCC_RETURN_UNLESS_0(Footer_schema_start(builder), error);
NANOARROW_RETURN_NOT_OK(
ArrowIpcEncodeSchema(builder, &footer->schema, &footer->dictionaries, error));
NANOARROW_RETURN_NOT_OK(ArrowIpcEncodeSchema(builder, &footer->schema,
&footer->dictionaries,
/*dictionary_replacement=*/0, error));
FLATCC_RETURN_UNLESS_0(Footer_schema_end(builder), error);

const struct ArrowIpcFileBlock* blocks =
Expand All @@ -837,6 +936,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;
}
Loading
Loading