Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
41 changes: 26 additions & 15 deletions src/nanoarrow/ipc/encoder.c
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,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 Down Expand Up @@ -76,6 +77,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 @@ -216,6 +218,14 @@ ArrowErrorCode ArrowIpcEncoderSetCompression(
return ArrowIpcEncoderSetCompressor(encoder, &compressor);
}

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;
}
Comment on lines +221 to +227

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.

This is a great and super useful change, but a separate one from the decoding or encoding of delta dictionaries

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Agreed. It stays in this PR, which is now only the non-delta writer and encoder changes.


static ArrowErrorCode ArrowIpcEncoderWriteContinuationAndSize(struct ArrowBuffer* out,
size_t size) {
_NANOARROW_CHECK_UPPER_LIMIT(size, INT32_MAX);
Expand Down Expand Up @@ -555,6 +565,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 @@ -623,24 +634,16 @@ 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));
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 @@ -658,7 +661,7 @@ static ArrowErrorCode ArrowIpcEncodeField(
static ArrowErrorCode ArrowIpcEncodeSchema(
flatcc_builder_t* builder, const struct ArrowSchema* schema,
const struct ArrowIpcDictionaryEncodings* dictionary_encodings, int compressed_body,
struct ArrowError* error) {
int dictionary_replacement, struct ArrowError* error) {
NANOARROW_DCHECK(schema->release != NULL);

if (strcmp(schema->format, "+s") != 0) {
Expand Down Expand Up @@ -695,6 +698,10 @@ static ArrowErrorCode ArrowIpcEncodeSchema(
ns(Feature_enum_t) feature = ns(Feature_COMPRESSED_BODY);
FLATCC_RETURN_IF_NULL(ns(Feature_vec_push(builder, &feature)), error);
}
if (dictionary_replacement && dictionary_encodings->encodings.size_bytes > 0) {
ns(Feature_enum_t) feature = ns(Feature_DICTIONARY_REPLACEMENT);
FLATCC_RETURN_IF_NULL(ns(Feature_vec_push(builder, &feature)), error);
}
FLATCC_RETURN_UNLESS_0(Schema_features_end(builder), error);

return NANOARROW_OK;
Expand Down Expand Up @@ -727,7 +734,8 @@ ArrowErrorCode ArrowIpcEncoderEncodeSchema(struct ArrowIpcEncoder* encoder,

NANOARROW_RETURN_NOT_OK(ArrowIpcEncodeSchema(
builder, schema, &private->dictionary_encodings,
ArrowIpcEncoderCodec(private) != NANOARROW_IPC_COMPRESSION_TYPE_NONE, error));
ArrowIpcEncoderCodec(private) != NANOARROW_IPC_COMPRESSION_TYPE_NONE,
private->dictionary_replacement, error));

FLATCC_RETURN_UNLESS_0(Message_header_Schema_end(builder), error);

Expand Down Expand Up @@ -1124,6 +1132,9 @@ static ArrowErrorCode ArrowIpcEncoderEncodeDictionaryBatch(
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);
Expand Down Expand Up @@ -1180,7 +1191,7 @@ ArrowErrorCode ArrowIpcEncoderEncodeFooter(struct ArrowIpcEncoder* encoder,
builder, &footer->schema, &footer->dictionaries,
private->has_compressed_body ||
ArrowIpcEncoderCodec(private) != NANOARROW_IPC_COMPRESSION_TYPE_NONE,
error));
/*dictionary_replacement=*/0, error));
FLATCC_RETURN_UNLESS_0(Footer_schema_end(builder), error);

const struct ArrowIpcFileBlock* blocks =
Expand Down
107 changes: 107 additions & 0 deletions src/nanoarrow/ipc/encoder_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -326,6 +326,89 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderSchemaMessageMetadata) {
(KeyValues{{"schema_key", "schema_value"}}));
}

TEST(NanoarrowIpcTest, NanoarrowIpcEncoderDictionaryReplacementFeature) {
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(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);

nanoarrow::ipc::UniqueEncoder encoder;
ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK);
ArrowIpcEncoderSetDictionaryReplacement(encoder.get(), /*enabled=*/1);

struct ArrowError error;
nanoarrow::UniqueBuffer message;
ASSERT_EQ(ArrowIpcEncoderEncodeSchema(encoder.get(), schema.get(), &error),
NANOARROW_OK)
<< error.message;
ASSERT_EQ(
ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/true, message.get()),
NANOARROW_OK);

nanoarrow::ipc::UniqueDecoder decoder;
ASSERT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK);
struct ArrowBufferView view;
view.data.data = message->data;
view.size_bytes = message->size_bytes;
ASSERT_EQ(ArrowIpcDecoderDecodeHeader(decoder.get(), view, &error), NANOARROW_OK)
<< error.message;
EXPECT_EQ(decoder->feature_flags, NANOARROW_IPC_FEATURE_DICTIONARY_REPLACEMENT);
}

TEST(NanoarrowIpcTest, NanoarrowIpcEncoderNestedDictionaryValueSchema) {
nanoarrow::UniqueSchema schema;
ASSERT_EQ(ArrowSchemaInitFromType(schema.get(), NANOARROW_TYPE_STRUCT), NANOARROW_OK);
ASSERT_EQ(ArrowSchemaAllocateChildren(schema.get(), 1), NANOARROW_OK);

struct ArrowSchema* dictionary_field = schema->children[0];
ASSERT_EQ(ArrowSchemaInitFromType(dictionary_field, NANOARROW_TYPE_INT32),
NANOARROW_OK);
ASSERT_EQ(ArrowSchemaSetName(dictionary_field, "dict_col"), NANOARROW_OK);
ASSERT_EQ(ArrowSchemaAllocateDictionary(dictionary_field), NANOARROW_OK);

struct ArrowSchema* value_schema = dictionary_field->dictionary;
ASSERT_EQ(ArrowSchemaInitFromType(value_schema, NANOARROW_TYPE_STRUCT), NANOARROW_OK);
ASSERT_EQ(ArrowSchemaAllocateChildren(value_schema, 1), NANOARROW_OK);
ASSERT_EQ(ArrowSchemaInitFromType(value_schema->children[0], NANOARROW_TYPE_STRING),
NANOARROW_OK);
ASSERT_EQ(ArrowSchemaSetName(value_schema->children[0], "value"), NANOARROW_OK);

nanoarrow::ipc::UniqueEncoder encoder;
ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK);
struct ArrowError error;
nanoarrow::UniqueBuffer message;
ASSERT_EQ(ArrowIpcEncoderEncodeSchema(encoder.get(), schema.get(), &error),
NANOARROW_OK)
<< error.message;
ASSERT_EQ(
ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/true, message.get()),
NANOARROW_OK);

nanoarrow::ipc::UniqueDecoder decoder;
ASSERT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK);
struct ArrowBufferView view;
view.data.data = message->data;
view.size_bytes = message->size_bytes;
ASSERT_EQ(ArrowIpcDecoderDecodeHeader(decoder.get(), view, &error), NANOARROW_OK)
<< error.message;

nanoarrow::UniqueSchema roundtrip;
ASSERT_EQ(ArrowIpcDecoderDecodeSchema(decoder.get(), roundtrip.get(), &error),
NANOARROW_OK)
<< error.message;
ASSERT_NE(roundtrip->children[0]->dictionary, nullptr);
EXPECT_STREQ(roundtrip->children[0]->dictionary->format, "+s");
ASSERT_EQ(roundtrip->children[0]->dictionary->n_children, 1);
EXPECT_STREQ(roundtrip->children[0]->dictionary->children[0]->format, "u");
EXPECT_STREQ(roundtrip->children[0]->dictionary->children[0]->name, "value");
}

TEST(NanoarrowIpcTest, NanoarrowIpcEncoderMessageMetadataEmpty) {
nanoarrow::ipc::UniqueEncoder encoder;
ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK);
Expand Down Expand Up @@ -468,6 +551,8 @@ TEST(NanoarrowIpcTest, NanoarrowIpcVisitMessageMetadataError) {
TEST(NanoarrowIpcTest, NanoarrowIpcEncoderDictionaryBatch) {
nanoarrow::ipc::UniqueEncoder encoder;
ASSERT_EQ(ArrowIpcEncoderInit(encoder.get()), NANOARROW_OK);
nanoarrow::ipc::UniqueDecoder decoder;
ASSERT_EQ(ArrowIpcDecoderInit(decoder.get()), NANOARROW_OK);

// Build a simple Utf8 values array
nanoarrow::UniqueSchema values_schema;
Expand Down Expand Up @@ -495,6 +580,13 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderDictionaryBatch) {
NANOARROW_OK)
<< error.message;

KeyValues message_key_values{{"dictionary_key", "dictionary_value"}};
auto message_metadata = PackMetadata(message_key_values);
ASSERT_EQ(
ArrowIpcEncoderSetMessageMetadata(encoder.get(), message_metadata.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,
Expand All @@ -511,6 +603,21 @@ TEST(NanoarrowIpcTest, NanoarrowIpcEncoderDictionaryBatch) {
// 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);
EXPECT_EQ(DecodeMessageMetadata(message_buffer.get(), decoder.get()),
message_key_values);

// The metadata applies to exactly one message: the next DictionaryBatch has none
message_buffer->size_bytes = 0;
body_buffer->size_bytes = 0;
ASSERT_EQ(ArrowIpcEncoderEncodeSimpleDictionaryBatch(encoder.get(), /*dictionary_id=*/0,
/*is_delta=*/0, values_view.get(),
body_buffer.get(), &error),
NANOARROW_OK)
<< error.message;
ASSERT_EQ(ArrowIpcEncoderFinalizeBuffer(encoder.get(), /*encapsulate=*/1,
message_buffer.get()),
NANOARROW_OK);
EXPECT_EQ(DecodeMessageMetadata(message_buffer.get(), decoder.get()), KeyValues{});
}

// A record batch whose columns exercise each path of the compressed body builder:
Expand Down
48 changes: 27 additions & 21 deletions src/nanoarrow/ipc/files_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -214,7 +214,7 @@ class TestFile {
}

ArrowErrorCode WriteNanoarrowStream(const nanoarrow::UniqueSchema& schema,
const std::vector<nanoarrow::UniqueArray>& arrays,
std::vector<nanoarrow::UniqueArray>& arrays,
enum ArrowIpcCompressionType codec,
struct ArrowBuffer* buffer,
struct ArrowError* error) {
Expand All @@ -226,19 +226,25 @@ class TestFile {
NANOARROW_RETURN_NOT_OK(ArrowIpcWriterSetCompression(
writer.get(), codec, NANOARROW_IPC_COMPRESSION_LEVEL_DEFAULT, error));

nanoarrow::UniqueArrayView array_view;
nanoarrow::UniqueSchema schema_copy;
NANOARROW_RETURN_NOT_OK(ArrowSchemaDeepCopy(schema.get(), schema_copy.get()));
nanoarrow::UniqueArrayStream array_stream;
NANOARROW_RETURN_NOT_OK(
ArrowArrayViewInitFromSchema(array_view.get(), schema.get(), error));

NANOARROW_RETURN_NOT_OK(ArrowIpcWriterWriteSchema(writer.get(), schema.get(), error));
for (const auto& array : arrays) {
NANOARROW_RETURN_NOT_OK(
ArrowArrayViewSetArray(array_view.get(), array.get(), error));

NANOARROW_RETURN_NOT_OK(
ArrowIpcWriterWriteArrayView(writer.get(), array_view.get(), error));
ArrowBasicArrayStreamInit(array_stream.get(), schema_copy.get(), arrays.size()));

for (size_t i = 0; i < arrays.size(); i++) {
// Preserve the decoded array for the subsequent Arrow C++ comparison while
// giving the basic stream an independently releasable shared clone.
nanoarrow::UniqueArray shared;
nanoarrow::UniqueArray clone;
NANOARROW_RETURN_NOT_OK(ArrowArrayMoveShared(arrays[i].get(), shared.get()));
ArrowErrorCode result = ArrowArrayCloneShared(shared.get(), clone.get());
ArrowArrayMove(shared.get(), arrays[i].get());
NANOARROW_RETURN_NOT_OK(result);
ArrowBasicArrayStreamSetArray(array_stream.get(), i, clone.get());
}
return ArrowIpcWriterWriteArrayView(writer.get(), nullptr, error);

return ArrowIpcWriterWriteArrayStream(writer.get(), array_stream.get(), error);
}

void TestEqualsArrowCpp(const std::string& dir_prefix,
Expand Down Expand Up @@ -529,10 +535,10 @@ INSTANTIATE_TEST_SUITE_P(
TestFile::OK("generated_primitive.stream"),
TestFile::OK("generated_recursive_nested.stream"),
TestFile::OK("generated_union.stream"),
TestFile::ReadOnly("generated_dictionary_unsigned.stream"),
TestFile::ReadOnly("generated_dictionary.stream"),
TestFile::ReadOnly("generated_nested_dictionary.stream"),
TestFile::ReadOnly("generated_extension.stream")
TestFile::OK("generated_dictionary_unsigned.stream"),
TestFile::OK("generated_dictionary.stream"),
TestFile::OK("generated_nested_dictionary.stream"),
TestFile::OK("generated_extension.stream")
// Comment to keep last line from wrapping
));

Expand Down Expand Up @@ -581,7 +587,7 @@ INSTANTIATE_TEST_SUITE_P(
TestFile::OK("0.14.1/generated_primitive.stream"),
TestFile::OK("0.14.1/generated_primitive_no_batches.stream"),
TestFile::OK("0.14.1/generated_primitive_zerolength.stream"),
TestFile::ReadOnly("4.0.0-shareddict/generated_shared_dict.stream"),
TestFile::OK("4.0.0-shareddict/generated_shared_dict.stream"),
// cpp-21.0.0 regenerated gold files
TestFile::OK("cpp-21.0.0/generated_binary.stream"),
TestFile::OK("cpp-21.0.0/generated_binary_no_batches.stream"),
Expand All @@ -608,10 +614,10 @@ INSTANTIATE_TEST_SUITE_P(
TestFile::OK("cpp-21.0.0/generated_primitive_zerolength.stream"),
TestFile::OK("cpp-21.0.0/generated_recursive_nested.stream"),
TestFile::OK("cpp-21.0.0/generated_union.stream"),
TestFile::ReadOnly("cpp-21.0.0/generated_dictionary.stream"),
TestFile::ReadOnly("cpp-21.0.0/generated_dictionary_unsigned.stream"),
TestFile::ReadOnly("cpp-21.0.0/generated_extension.stream"),
TestFile::ReadOnly("cpp-21.0.0/generated_nested_dictionary.stream"),
TestFile::OK("cpp-21.0.0/generated_dictionary.stream"),
TestFile::OK("cpp-21.0.0/generated_dictionary_unsigned.stream"),
TestFile::OK("cpp-21.0.0/generated_extension.stream"),
TestFile::OK("cpp-21.0.0/generated_nested_dictionary.stream"),
Comment on lines +617 to +620

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.

Just curious: what feature is needed to roundtrip the generated extension beyond what's here?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Nothing beyond this PR: the file's dict_exts column is a dictionary-encoded extension type, and the test previously wrote with ArrowIpcWriterWriteArrayView(), which doesn't write dictionaries. Now that the roundtrip goes through ArrowIpcWriterWriteArrayStream() it passes, so I switched it (and the shared-dictionary file) to OK.

TestFile::NotSupported("cpp-21.0.0/generated_list_view.stream"),
TestFile::NotSupported("cpp-21.0.0/generated_binary_view.stream"),
TestFile::NotSupported("cpp-21.0.0/generated_run_end_encoded.stream")
Expand Down
8 changes: 0 additions & 8 deletions src/nanoarrow/ipc/reader.c
Original file line number Diff line number Diff line change
Expand Up @@ -389,14 +389,6 @@ static int ArrowIpcArrayStreamReaderReadSchemaIfNeeded(
return EINVAL;
}

// ...or if it uses features we don't support
if (private_data->decoder.feature_flags &
NANOARROW_IPC_FEATURE_DICTIONARY_REPLACEMENT) {
ArrowErrorSet(&private_data->error,
"This stream uses unsupported feature DICTIONARY_REPLACEMENT");
return EINVAL;
}

// Notify the decoder of buffer endianness
NANOARROW_RETURN_NOT_OK_WITH_ERROR(
ArrowIpcDecoderSetEndianness(&private_data->decoder,
Expand Down
Loading
Loading