diff --git a/cpp/src/arrow/integration/json_integration_test.cc b/cpp/src/arrow/integration/json_integration_test.cc index 0f7e6c85dc55..4f70ac17893c 100644 --- a/cpp/src/arrow/integration/json_integration_test.cc +++ b/cpp/src/arrow/integration/json_integration_test.cc @@ -31,6 +31,7 @@ #include "arrow/array.h" #include "arrow/array/builder_binary.h" #include "arrow/array/builder_primitive.h" +#include "arrow/extension/uuid.h" #include "arrow/integration/json_integration.h" #include "arrow/integration/json_internal.h" #include "arrow/io/file.h" @@ -48,6 +49,7 @@ #include "arrow/type.h" #include "arrow/type_fwd.h" #include "arrow/util/io_util.h" +#include "arrow/util/key_value_metadata.h" #include "arrow/util/simdjson_internal.h" DEFINE_string(arrow, "", "Arrow file name"); @@ -725,6 +727,113 @@ static const char* json_example6 = R"example( } )example"; +// A batch with dictionaries whose value type is an extension type (GH-49704) +static const char* json_example7 = R"example( +{ + "schema": { + "fields": [ + { + "name": "uuid_dict", + "type": { + "name": "fixedsizebinary", + "byteWidth": 16 + }, + "nullable": true, + "children": [], + "dictionary": { + "id": 0, + "indexType": { + "name": "int", + "isSigned": true, + "bitWidth": 8 + }, + "isOrdered": false, + "metadata": [ + {"key": "ARROW:extension:name", "value": "uuid"}, + {"key": "ARROW:extension:metadata", "value": "uuid-serialized"} + ] + }, + "metadata": [ + {"key": "some_key", "value": "some_value"} + ] + }, + { + "name": "unknown_dict", + "type": { + "name": "fixedsizebinary", + "byteWidth": 16 + }, + "nullable": true, + "children": [], + "dictionary": { + "id": 1, + "indexType": { + "name": "int", + "isSigned": true, + "bitWidth": 8 + }, + "isOrdered": false, + "metadata": [ + {"key": "ARROW:extension:name", "value": "unknown.extension"}, + {"key": "ARROW:extension:metadata", "value": ""} + ] + } + } + ] + }, + "dictionaries": [ + { + "id": 0, + "data": { + "count": 2, + "columns": [ + { + "name": "DICT0", + "count": 2, + "VALIDITY": [1, 1], + "DATA": ["30313233343536373839616263646566", + "61626364656630313233343536373839"] + } + ] + } + }, + { + "id": 1, + "data": { + "count": 1, + "columns": [ + { + "name": "DICT1", + "count": 1, + "VALIDITY": [1], + "DATA": ["30313233343536373839616263646566"] + } + ] + } + } + ], + "batches": [ + { + "count": 4, + "columns": [ + { + "name": "uuid_dict", + "count": 4, + "DATA": [1, 0, 0, 1], + "VALIDITY": [1, 1, 0, 1] + }, + { + "name": "unknown_dict", + "count": 4, + "DATA": [0, 0, 0, 0], + "VALIDITY": [1, 0, 1, 1] + } + ] + } + ] +} +)example"; + void TestSchemaRoundTrip(const std::shared_ptr& schema) { arrow::internal::JsonWriter writer; @@ -1179,6 +1288,114 @@ void CheckRoundtrip(const RecordBatch& batch) { ApproxCompareBatch(batch, *result_batch); } +TEST(TestJsonFileReadWrite, JsonExample7) { + // Example 7: dictionaries whose value type is an extension type (GH-49704) + auto uuid_type = uuid(); + ExtensionTypeGuard ext_guard(uuid_type); + // The value type of the second dictionary is an unregistered extension type, + // and is read as its storage type + Schema ex_schema({field("uuid_dict", dictionary(int8(), uuid_type), /*nullable=*/true, + key_value_metadata({{"some_key", "some_value"}})), + field("unknown_dict", dictionary(int8(), fixed_size_binary(16)))}); + + std::shared_ptr batch; + ReadOneBatchJson(json_example7, ex_schema, &batch); + + auto uuid_values = ExtensionType::WrapArray( + uuid_type, ArrayFromJSON(fixed_size_binary(16), + R"(["0123456789abcdef", "abcdef0123456789"])")); + ASSERT_OK_AND_ASSIGN( + auto expected_uuid_dict, + DictionaryArray::FromArrays(dictionary(int8(), uuid_type), + ArrayFromJSON(int8(), "[1, 0, null, 1]"), uuid_values)); + AssertArraysEqual(*batch->column(0), *expected_uuid_dict, /*verbose=*/true); + + ASSERT_OK_AND_ASSIGN( + auto expected_unknown_dict, + DictionaryArray::FromArrays( + dictionary(int8(), fixed_size_binary(16)), + ArrayFromJSON(int8(), "[0, null, 0, 0]"), + ArrayFromJSON(fixed_size_binary(16), R"(["0123456789abcdef"])"))); + AssertArraysEqual(*batch->column(1), *expected_unknown_dict, /*verbose=*/true); +} + +TEST(TestJsonFileReadWrite, DictionaryOfExtension) { + // GH-49704: the extension metadata of a dictionary value type is written to + // the "metadata" member of the field's "dictionary" object + auto uuid_type = uuid(); + ExtensionTypeGuard ext_guard({uuid_type, dict_extension_type()}); + auto schema = ::arrow::schema({ + field("f0", dictionary(int16(), uuid_type), /*nullable=*/true, + key_value_metadata({{"some_key", "some_value"}})), + field("f1", list(dictionary(int32(), uuid_type))), + field("f2", dict_extension_type()), + }); + TestSchemaRoundTrip(schema); + + arrow::internal::JsonWriter writer; + DictionaryFieldMapper mapper(*schema); + writer.StartObject(); + ASSERT_OK(json::WriteSchema(*schema, mapper, &writer)); + writer.EndObject(); + ASSERT_OK_AND_ASSIGN(std::string_view json_schema, writer.GetString()); + + simdjson::dom::parser parser; + ASSERT_OK_AND_ASSIGN(auto doc, internal::ResolveSimdjsonResult( + parser.parse(json_schema.data(), json_schema.size()), + "Failed to parse JSON")); + using Metadata = std::vector>; + auto get_metadata = [](simdjson::dom::element obj) { + Metadata metadata; + simdjson::dom::array kvs; + if (obj["metadata"].get(kvs) == simdjson::SUCCESS) { + for (auto kv : kvs) { + metadata.emplace_back(std::string(kv["key"].get_string().value_unsafe()), + std::string(kv["value"].get_string().value_unsafe())); + } + } + return metadata; + }; + const Metadata uuid_metadata = {{"ARROW:extension:name", "uuid"}, + {"ARROW:extension:metadata", "uuid-serialized"}}; + + auto fields = doc["schema"]["fields"]; + auto f0 = fields.at(0).value_unsafe(); + ASSERT_EQ(get_metadata(f0), (Metadata{{"some_key", "some_value"}})); + ASSERT_EQ(get_metadata(f0["dictionary"].value_unsafe()), uuid_metadata); + + auto f1_child = fields.at(1).value_unsafe()["children"].at(0).value_unsafe(); + ASSERT_EQ(get_metadata(f1_child), Metadata{}); + ASSERT_EQ(get_metadata(f1_child["dictionary"].value_unsafe()), uuid_metadata); + + // An extension type with a dictionary storage type keeps its metadata on + // the field + auto f2 = fields.at(2).value_unsafe(); + ASSERT_EQ(get_metadata(f2), + (Metadata{{"ARROW:extension:name", "dict-extension"}, + {"ARROW:extension:metadata", "dict-extension-serialized"}})); + ASSERT_EQ(get_metadata(f2["dictionary"].value_unsafe()), Metadata{}); +} + +TEST(TestJsonFileReadWrite, DictionaryOfExtensionRoundTrip) { + // GH-49704 + for (const auto& value_type : {uuid(), extension::uuid()}) { + auto values = ExtensionType::WrapArray( + value_type, ArrayFromJSON(fixed_size_binary(16), + R"(["0123456789abcdef", "abcdef0123456789", null])")); + for (const auto& index_type : {int8(), int32()}) { + ARROW_SCOPED_TRACE("value_type = ", value_type->ToString(), + ", index_type = ", index_type->ToString()); + ASSERT_OK_AND_ASSIGN(auto dict_array, + DictionaryArray::FromArrays( + dictionary(index_type, value_type), + ArrayFromJSON(index_type, "[0, 1, null, 2, 1]"), values)); + auto batch = RecordBatch::Make(::arrow::schema({field("f0", dict_array->type())}), + dict_array->length(), {dict_array}); + CheckRoundtrip(*batch); + } + } +} + TEST_P(TestJsonRoundTrip, RoundTrip) { std::shared_ptr batch; ASSERT_OK((*GetParam())(&batch)); // NOLINT clang-tidy gtest issue diff --git a/cpp/src/arrow/integration/json_internal.cc b/cpp/src/arrow/integration/json_internal.cc index 9dd9963cd5d8..83c805283dd7 100644 --- a/cpp/src/arrow/integration/json_internal.cc +++ b/cpp/src/arrow/integration/json_internal.cc @@ -187,6 +187,14 @@ class SchemaWriter { writer_->Key("isOrdered"); writer_->Bool(type.ordered()); + + // Emulate DictionaryEncoding.custom_metadata for a dictionary whose value + // type is an extension type (GH-49704) + if (type.value_type()->id() == Type::EXTENSION) { + const auto& ext_type = checked_cast(*type.value_type()); + WriteKeyValueMetadata(nullptr, {{kExtensionTypeKeyName, ext_type.extension_name()}, + {kExtensionMetadataKeyName, ext_type.Serialize()}}); + } writer_->EndObject(); return Status::OK(); @@ -451,7 +459,14 @@ class SchemaWriter { return Status::OK(); } - Status Visit(const DictionaryType& type) { return VisitType(*type.value_type()); } + Status Visit(const DictionaryType& type) { + const DataType* value_type = type.value_type().get(); + if (value_type->id() == Type::EXTENSION) { + // The extension metadata is written to the "dictionary" object + value_type = checked_cast(*value_type).storage_type().get(); + } + return VisitType(*value_type); + } Status Visit(const RunEndEncodedType& type) { WriteName("runendencoded", type); @@ -1301,6 +1316,25 @@ Result> GetField(const JsonValue& obj, FieldPosition fiel std::shared_ptr index_type; RETURN_NOT_OK( ParseDictionary(dictionary_object, &dictionary_id, &is_ordered, &index_type)); + + // Emulate DictionaryEncoding.custom_metadata, which holds the extension + // metadata of the dictionary value type (GH-49704) + ARROW_ASSIGN_OR_RAISE(std::shared_ptr value_metadata, + GetKeyValueMetadata(dictionary_object)); + int value_ext_name_index = value_metadata->FindKey(kExtensionTypeKeyName); + if (value_ext_name_index != -1) { + auto value_ext_type = + GetExtensionType(value_metadata->value(value_ext_name_index)); + // As with IPC, an unregistered extension type leaves the storage type + if (value_ext_type != nullptr) { + int value_ext_data_index = value_metadata->FindKey(kExtensionMetadataKeyName); + ARROW_ASSIGN_OR_RAISE( + type, value_ext_type->Deserialize( + type, value_ext_data_index == -1 + ? "" + : value_metadata->value(value_ext_data_index))); + } + } dict_value_type = type; type = ::arrow::dictionary(index_type, type, is_ordered); } diff --git a/cpp/src/arrow/ipc/metadata_internal.cc b/cpp/src/arrow/ipc/metadata_internal.cc index 02134aa805ab..d1f33a54f736 100644 --- a/cpp/src/arrow/ipc/metadata_internal.cc +++ b/cpp/src/arrow/ipc/metadata_internal.cc @@ -478,6 +478,7 @@ Status TensorTypeToFlatbuffer(FBB& fbb, const DataType& type, flatbuf::Type* out static Status GetDictionaryEncoding(FBB& fbb, const std::shared_ptr& field, const DictionaryType& type, int64_t dictionary_id, + flatbuffers::Offset value_metadata, DictionaryOffset* out) { // We assume that the dictionary index type (as an integer) has already been // validated elsewhere, and can safely assume we are dealing with integers @@ -486,8 +487,9 @@ static Status GetDictionaryEncoding(FBB& fbb, const std::shared_ptr& fiel auto index_type_offset = flatbuf::CreateInt(fbb, index_type.bit_width(), index_type.is_signed()); - *out = flatbuf::CreateDictionaryEncoding(fbb, dictionary_id, index_type_offset, - type.ordered()); + *out = flatbuf::CreateDictionaryEncoding( + fbb, dictionary_id, index_type_offset, type.ordered(), + flatbuf::DictionaryKind::DictionaryKind_DenseArray, value_metadata); return Status::OK(); } @@ -783,7 +785,19 @@ class FieldToFlatbufferVisitor { // In this library, the dictionary "type" is a logical construct. Here we // pass through to the value type, as we've already captured the index // type in the DictionaryEncoding metadata in the parent field - return VisitType(*checked_cast(type).value_type()); + const DataType& value_type = *type.value_type(); + if (value_type.id() == Type::EXTENSION) { + // The extension metadata of the dictionary value type is written to + // DictionaryEncoding.custom_metadata, to distinguish it from an extension + // type whose storage type is a dictionary (GH-49704). + const auto& ext_type = checked_cast(value_type); + dictionary_value_metadata_.emplace_back(kExtensionTypeKeyName, + ext_type.extension_name()); + dictionary_value_metadata_.emplace_back(kExtensionMetadataKeyName, + ext_type.Serialize()); + return VisitType(*ext_type.storage_type()); + } + return VisitType(value_type); } Status Visit(const RunEndEncodedType& type) { @@ -822,9 +836,19 @@ class FieldToFlatbufferVisitor { if (storage_type->id() == Type::DICTIONARY) { ARROW_ASSIGN_OR_RAISE(const auto dictionary_id, mapper_.GetFieldId(field_pos_.path())); + // Omit DictionaryEncoding.custom_metadata (rather than writing an empty + // vector) if the dictionary value type has no metadata + flatbuffers::Offset fb_value_metadata; + if (!dictionary_value_metadata_.empty()) { + std::vector value_key_values; + for (const auto& [key, value] : dictionary_value_metadata_) { + value_key_values.push_back(AppendKeyValue(fbb_, key, value)); + } + fb_value_metadata = fbb_.CreateVector(value_key_values); + } RETURN_NOT_OK(GetDictionaryEncoding( fbb_, field, checked_cast(*storage_type), dictionary_id, - &dictionary)); + fb_value_metadata, &dictionary)); } auto metadata = field->metadata(); @@ -859,8 +883,37 @@ class FieldToFlatbufferVisitor { Offset type_offset_; std::vector children_; std::unordered_map extra_type_metadata_; + // Metadata of the dictionary value type, for DictionaryEncoding.custom_metadata + std::vector> dictionary_value_metadata_; }; +// If `metadata` describes a registered extension type, deserialize that +// extension type over `*type` and remove the extension keys from `metadata`, +// for faithful roundtripping. If the extension type is unknown, we do not raise +// and simply leave the storage type. +Status MaybeDeserializeExtensionType(KeyValueMetadata* metadata, + std::shared_ptr* type) { + if (metadata == nullptr) { + return Status::OK(); + } + int name_index = metadata->FindKey(kExtensionTypeKeyName); + if (name_index == -1) { + return Status::OK(); + } + std::shared_ptr ext_type = GetExtensionType(metadata->value(name_index)); + if (ext_type == nullptr) { + return Status::OK(); + } + int data_index = metadata->FindKey(kExtensionMetadataKeyName); + std::string type_data = data_index == -1 ? "" : metadata->value(data_index); + + ARROW_ASSIGN_OR_RAISE(*type, ext_type->Deserialize(*type, type_data)); + if (data_index != -1) { + return metadata->DeleteMany({name_index, data_index}); + } + return metadata->Delete(name_index); +} + Status FieldFromFlatbuffer(const flatbuf::Field* field, FieldPosition field_pos, DictionaryMemo* dictionary_memo, std::shared_ptr* out) { std::shared_ptr type; @@ -900,6 +953,13 @@ Status FieldFromFlatbuffer(const flatbuf::Field* field, FieldPosition field_pos, auto int_data = encoding->indexType(); CHECK_FLATBUFFERS_NOT_NULL(int_data, "DictionaryEncoding.indexType"); RETURN_NOT_OK(IntFromFlatbuffer(int_data, &index_type)); + // The dictionary value type may be an extension type, whose metadata is + // found in DictionaryEncoding.custom_metadata (GH-49704) + if (encoding->custom_metadata() != nullptr) { + ARROW_ASSIGN_OR_RAISE(auto value_metadata, + internal::GetKeyValueMetadata(encoding->custom_metadata())); + RETURN_NOT_OK(MaybeDeserializeExtensionType(value_metadata.get(), &type)); + } dict_value_type = type; ARROW_ASSIGN_OR_RAISE(type, DictionaryType::Make(index_type, type, encoding->isOrdered())); @@ -907,28 +967,7 @@ Status FieldFromFlatbuffer(const flatbuf::Field* field, FieldPosition field_pos, } // 4. Is it an extension type? - if (metadata != nullptr) { - // Look for extension metadata in custom_metadata field - int name_index = metadata->FindKey(kExtensionTypeKeyName); - if (name_index != -1) { - std::shared_ptr ext_type = - GetExtensionType(metadata->value(name_index)); - if (ext_type != nullptr) { - int data_index = metadata->FindKey(kExtensionMetadataKeyName); - std::string type_data = data_index == -1 ? "" : metadata->value(data_index); - - ARROW_ASSIGN_OR_RAISE(type, ext_type->Deserialize(type, type_data)); - // Remove the metadata, for faithful roundtripping - if (data_index != -1) { - RETURN_NOT_OK(metadata->DeleteMany({name_index, data_index})); - } else { - RETURN_NOT_OK(metadata->Delete(name_index)); - } - } - // NOTE: if extension type is unknown, we do not raise here and - // simply return the storage type. - } - } + RETURN_NOT_OK(MaybeDeserializeExtensionType(metadata.get(), &type)); // Reconstruct field auto field_name = StringFromFlatbuffers(field->name()); diff --git a/cpp/src/arrow/ipc/read_write_test.cc b/cpp/src/arrow/ipc/read_write_test.cc index cfa6a541d5cc..97ca2c32e6b2 100644 --- a/cpp/src/arrow/ipc/read_write_test.cc +++ b/cpp/src/arrow/ipc/read_write_test.cc @@ -17,6 +17,7 @@ #include #include +#include #include #include #include @@ -30,7 +31,9 @@ #include "arrow/array.h" #include "arrow/array/builder_primitive.h" +#include "arrow/array/concatenate.h" #include "arrow/buffer_builder.h" +#include "arrow/extension/uuid.h" #include "arrow/io/file.h" #include "arrow/io/memory.h" #include "arrow/io/test_common.h" @@ -51,10 +54,12 @@ #include "arrow/type_fwd.h" #include "arrow/util/bit_util.h" #include "arrow/util/checked_cast.h" +#include "arrow/util/endian.h" #include "arrow/util/io_util.h" #include "arrow/util/key_value_metadata.h" #include "arrow/util/ubsan.h" +#include "generated/File_generated.h" // IWYU pragma: keep #include "generated/Message_generated.h" // IWYU pragma: keep namespace arrow { @@ -370,6 +375,285 @@ TEST_F(TestSchemaMetadata, MetadataVersionForwardCompatibility) { ASSERT_RAISES(Invalid, ReadSchema(schema_v6_file.get(), &placeholder_memo)); } +// GH-49704: dictionary-encoded extension types + +// An extension type over an arbitrary storage type, e.g. a dictionary type +// whose value type is itself an extension type +class AnyStorageExtensionType : public ExtensionType { + public: + explicit AnyStorageExtensionType(std::shared_ptr storage_type) + : ExtensionType(std::move(storage_type)) {} + + std::string extension_name() const override { return "any-storage"; } + + bool ExtensionEquals(const ExtensionType& other) const override { + return other.extension_name() == extension_name() && + other.storage_type()->Equals(*storage_type()); + } + + std::shared_ptr MakeArray(std::shared_ptr data) const override { + return std::make_shared(std::move(data)); + } + + Result> Deserialize( + std::shared_ptr storage_type, + const std::string& serialized) const override { + if (serialized != "any-storage-serialized") { + return Status::Invalid("Type identifier did not match: '", serialized, "'"); + } + return std::make_shared(std::move(storage_type)); + } + + std::string Serialize() const override { return "any-storage-serialized"; } +}; + +// Return the first field of a schema serialized as an IPC message +const flatbuf::Field* GetFirstFlatbufferField(const Message& message) { + const auto* fb_message = flatbuf::GetMessage(message.metadata()->data()); + const auto* fb_schema = fb_message->header_as_Schema(); + EXPECT_NE(fb_schema, nullptr); + EXPECT_EQ(fb_schema->fields()->size(), 1); + return fb_schema->fields()->Get(0); +} + +void AssertKeyValues( + const KeyValueMetadata& expected, + const flatbuffers::Vector>* kv) { + ASSERT_NE(kv, nullptr); + ASSERT_OK_AND_ASSIGN(auto actual, internal::GetKeyValueMetadata(kv)); + ASSERT_TRUE(actual->Equals(expected)) + << "expected: " << expected.ToString() << " actual: " << actual->ToString(); +} + +// Serialize a schema of one field the way the IPC writer would, with explicit +// Field.custom_metadata and DictionaryEncoding.custom_metadata +Result> ReadSingleDictionaryFieldSchema( + flatbuf::Type fb_type, int index_bit_width, + const std::vector>& field_metadata, + const std::vector>& dictionary_metadata) { + flatbuffers::FlatBufferBuilder fbb; + auto make_key_values = + [&](const std::vector>& metadata) + -> flatbuffers::Offset { + if (metadata.empty()) { + return 0; + } + std::vector key_values; + for (const auto& [key, value] : metadata) { + key_values.push_back( + flatbuf::CreateKeyValue(fbb, fbb.CreateString(key), fbb.CreateString(value))); + } + return fbb.CreateVector(key_values); + }; + + flatbuffers::Offset type_offset; + switch (fb_type) { + case flatbuf::Type::Type_FixedSizeBinary: + type_offset = flatbuf::CreateFixedSizeBinary(fbb, 16).Union(); + break; + case flatbuf::Type::Type_Utf8: + type_offset = flatbuf::CreateUtf8(fbb).Union(); + break; + default: + return Status::NotImplemented("Unsupported test type"); + } + auto fb_dictionary_metadata = make_key_values(dictionary_metadata); + auto index_type = flatbuf::CreateInt(fbb, index_bit_width, /*is_signed=*/true); + auto encoding = flatbuf::CreateDictionaryEncoding( + fbb, /*id=*/0, index_type, /*isOrdered=*/false, + flatbuf::DictionaryKind::DictionaryKind_DenseArray, fb_dictionary_metadata); + auto fb_field_metadata = make_key_values(field_metadata); + auto fb_field = flatbuf::CreateField( + fbb, fbb.CreateString("f0"), /*nullable=*/true, fb_type, type_offset, encoding, + fbb.CreateVector(std::vector>{}), + fb_field_metadata); + auto fb_schema = flatbuf::CreateSchema( + fbb, flatbuf::Endianness::Endianness_Little, + fbb.CreateVector(std::vector>{fb_field})); + fbb.Finish(fb_schema); + + DictionaryMemo memo; + std::shared_ptr schema; + RETURN_NOT_OK(internal::GetSchema( + flatbuffers::GetRoot(fbb.GetBufferPointer()), &memo, &schema)); + return schema; +} + +TEST_F(TestSchemaMetadata, DictionaryOfExtensionFields) { + ExtensionTypeGuard guard(uuid()); + auto canonical_uuid = extension::uuid(); + CheckSchemaRoundtrip(Schema({ + field("f0", dictionary(int32(), canonical_uuid)), + field("f1", dictionary(int8(), uuid(), /*ordered=*/true)), + field("f2", list(dictionary(int16(), canonical_uuid))), + field("f3", struct_({field("s", dictionary(int32(), uuid()))})), + field("f4", dictionary(int32(), canonical_uuid), /*nullable=*/true, + key_value_metadata({{"key", "value"}})), + })); +} + +TEST_F(TestSchemaMetadata, ExtensionWithDictionaryStorage) { + auto any_storage_ext = + std::make_shared(dictionary(int8(), extension::uuid())); + ExtensionTypeGuard guard({dict_extension_type(), any_storage_ext}); + CheckSchemaRoundtrip(Schema({ + field("f0", dict_extension_type()), + // An extension type with a dictionary storage type whose value type is + // an extension type + field("f1", any_storage_ext), + })); +} + +TEST_F(TestSchemaMetadata, DictionaryOfExtensionFlatbuffer) { + auto check_serialized = [](const std::shared_ptr& field, + const KeyValueMetadata* expected_field_metadata, + const KeyValueMetadata* expected_dictionary_metadata) { + ASSERT_OK_AND_ASSIGN(auto buffer, SerializeSchema(Schema({field}))); + io::BufferReader reader(buffer); + ASSERT_OK_AND_ASSIGN(auto message, ReadMessage(&reader)); + const flatbuf::Field* fb_field = GetFirstFlatbufferField(*message); + ASSERT_NE(fb_field->dictionary(), nullptr); + if (expected_field_metadata == nullptr) { + ASSERT_EQ(fb_field->custom_metadata(), nullptr); + } else { + AssertKeyValues(*expected_field_metadata, fb_field->custom_metadata()); + } + if (expected_dictionary_metadata == nullptr) { + // The member is omitted, rather than written as an empty vector + ASSERT_EQ(fb_field->dictionary()->custom_metadata(), nullptr); + } else { + AssertKeyValues(*expected_dictionary_metadata, + fb_field->dictionary()->custom_metadata()); + } + }; + + // The extension metadata of the dictionary value type goes to + // DictionaryEncoding.custom_metadata, not Field.custom_metadata + KeyValueMetadata uuid_metadata({"ARROW:extension:name", "ARROW:extension:metadata"}, + {"arrow.uuid", ""}); + check_serialized(field("f0", dictionary(int32(), extension::uuid())), nullptr, + &uuid_metadata); + KeyValueMetadata field_metadata({"key"}, {"value"}); + check_serialized( + field("f0", dictionary(int32(), extension::uuid()), true, field_metadata.Copy()), + &field_metadata, &uuid_metadata); + + // A dictionary whose value type is not an extension type + check_serialized(field("f0", dictionary(int32(), utf8())), nullptr, nullptr); + + // An extension type whose storage type is a dictionary keeps its metadata + // in Field.custom_metadata + KeyValueMetadata dict_ext_metadata({"ARROW:extension:name", "ARROW:extension:metadata"}, + {"dict-extension", "dict-extension-serialized"}); + check_serialized(field("f0", dict_extension_type()), &dict_ext_metadata, nullptr); + + // Both + KeyValueMetadata any_storage_metadata( + {"ARROW:extension:name", "ARROW:extension:metadata"}, + {"any-storage", "any-storage-serialized"}); + check_serialized(field("f0", std::make_shared( + dictionary(int8(), extension::uuid()))), + &any_storage_metadata, &uuid_metadata); +} + +TEST_F(TestSchemaMetadata, DictionaryOfUnregisteredExtension) { + // The example "uuid" extension type is not registered: the dictionary value + // type is read as its storage type + ASSERT_OK_AND_ASSIGN( + auto buffer, SerializeSchema(Schema({field("f0", dictionary(int32(), uuid()))}))); + io::BufferReader reader(buffer); + DictionaryMemo memo; + ASSERT_OK_AND_ASSIGN(auto schema, ReadSchema(&reader, &memo)); + AssertSchemaEqual(Schema({field("f0", dictionary(int32(), fixed_size_binary(16)))}), + *schema, /*check_metadata=*/true); + + ASSERT_OK_AND_ASSIGN(schema, ReadSingleDictionaryFieldSchema( + flatbuf::Type::Type_FixedSizeBinary, 32, {}, + {{"ARROW:extension:name", "unknown.extension"}, + {"ARROW:extension:metadata", "unknown-serialized"}})); + AssertSchemaEqual(Schema({field("f0", dictionary(int32(), fixed_size_binary(16)))}), + *schema, /*check_metadata=*/true); +} + +TEST_F(TestSchemaMetadata, DictionaryOfExtensionManualFlatbuffer) { + ASSERT_OK_AND_ASSIGN( + auto schema, + ReadSingleDictionaryFieldSchema( + flatbuf::Type::Type_FixedSizeBinary, 32, {{"key", "value"}}, + {{"ARROW:extension:name", "arrow.uuid"}, {"ARROW:extension:metadata", ""}})); + AssertSchemaEqual(Schema({field("f0", dictionary(int32(), extension::uuid()), true, + key_value_metadata({{"key", "value"}}))}), + *schema, /*check_metadata=*/true); +} + +TEST_F(TestSchemaMetadata, DictionaryEncodingMetadataNotPreserved) { + // Pins a current C++ limitation: DictionaryType cannot hold metadata for its + // value type, so keys of DictionaryEncoding.custom_metadata other than those + // of a registered extension type are dropped on read + auto expected_storage = + Schema({field("f0", dictionary(int32(), fixed_size_binary(16)))}); + ASSERT_OK_AND_ASSIGN( + auto schema, ReadSingleDictionaryFieldSchema(flatbuf::Type::Type_FixedSizeBinary, + 32, {}, {{"some_key", "some_value"}})); + AssertSchemaEqual(expected_storage, *schema, /*check_metadata=*/true); + + ASSERT_OK_AND_ASSIGN(schema, ReadSingleDictionaryFieldSchema( + flatbuf::Type::Type_FixedSizeBinary, 32, {}, + {{"ARROW:extension:name", "unknown.extension"}, + {"ARROW:extension:metadata", ""}, + {"some_key", "some_value"}})); + AssertSchemaEqual(expected_storage, *schema, /*check_metadata=*/true); + + ASSERT_OK_AND_ASSIGN( + schema, ReadSingleDictionaryFieldSchema(flatbuf::Type::Type_FixedSizeBinary, 32, {}, + {{"ARROW:extension:name", "arrow.uuid"}, + {"ARROW:extension:metadata", ""}, + {"some_key", "some_value"}})); + AssertSchemaEqual(Schema({field("f0", dictionary(int32(), extension::uuid()))}), + *schema, /*check_metadata=*/true); +} + +TEST_F(TestSchemaMetadata, LegacyDictionaryExtensionMetadataOnField) { + // Before GH-49704, the extension metadata of a dictionary-encoded field was + // only ever written to Field.custom_metadata, which is still read as an + // extension type whose storage type is a dictionary + { + ExtensionTypeGuard guard(dict_extension_type()); + ASSERT_OK_AND_ASSIGN(auto schema, + ReadSingleDictionaryFieldSchema( + flatbuf::Type::Type_Utf8, 8, + {{"ARROW:extension:name", "dict-extension"}, + {"ARROW:extension:metadata", "dict-extension-serialized"}}, + {})); + AssertSchemaEqual(Schema({field("f0", dict_extension_type())}), *schema, + /*check_metadata=*/true); + } + { + // An unregistered extension type leaves the dictionary type and the metadata + auto metadata = key_value_metadata({{"ARROW:extension:name", "unknown.extension"}, + {"ARROW:extension:metadata", ""}}); + ASSERT_OK_AND_ASSIGN(auto schema, ReadSingleDictionaryFieldSchema( + flatbuf::Type::Type_FixedSizeBinary, 32, + {{"ARROW:extension:name", "unknown.extension"}, + {"ARROW:extension:metadata", ""}}, + {})); + AssertSchemaEqual( + Schema({field("f0", dictionary(int32(), fixed_size_binary(16)), true, metadata)}), + *schema, /*check_metadata=*/true); + } + { + // This is ambiguous, and arrow.uuid does not support a dictionary storage type + ASSERT_RAISES_WITH_MESSAGE( + Invalid, + "Invalid: Invalid storage type for UuidType: " + "dictionary", + ReadSingleDictionaryFieldSchema( + flatbuf::Type::Type_FixedSizeBinary, 32, + {{"ARROW:extension:name", "arrow.uuid"}, {"ARROW:extension:metadata", ""}}, + {})); + } +} + const std::vector kBatchCases = { &MakeIntRecordBatch, &MakeListRecordBatch, @@ -1044,6 +1328,62 @@ TEST_F(TestWriteRecordBatch, IntegerGetRecordBatchSize) { TestGetRecordBatchSize(options_, batch); } +TEST_F(TestWriteRecordBatch, DictionaryOfExtensionFileFooter) { + // GH-49704: the schema in the IPC file footer also carries the extension + // metadata of the dictionary value type in DictionaryEncoding.custom_metadata + auto type = dictionary(int32(), extension::uuid()); + auto values = ExtensionType::WrapArray( + extension::uuid(), ArrayFromJSON(fixed_size_binary(16), R"(["0123456789abcdef"])")); + ASSERT_OK_AND_ASSIGN( + auto array, + DictionaryArray::FromArrays(type, ArrayFromJSON(int32(), "[0, null]"), values)); + auto batch = RecordBatch::Make(schema({field("f0", type)}), 2, {array}); + + ASSERT_OK_AND_ASSIGN(auto sink, io::BufferOutputStream::Create()); + ASSERT_OK_AND_ASSIGN(auto writer, MakeFileWriter(sink, batch->schema())); + ASSERT_OK(writer->WriteRecordBatch(*batch)); + ASSERT_OK(writer->Close()); + ASSERT_OK_AND_ASSIGN(auto buffer, sink->Finish()); + + // The file ends with the footer, its int32 length and the "ARROW1" magic + const uint8_t* end = buffer->data() + buffer->size() - 6; + int32_t footer_length; + std::memcpy(&footer_length, end - sizeof(int32_t), sizeof(int32_t)); + footer_length = bit_util::FromLittleEndian(footer_length); + const uint8_t* footer_data = end - sizeof(int32_t) - footer_length; + flatbuffers::Verifier verifier(footer_data, footer_length); + ASSERT_TRUE(flatbuf::VerifyFooterBuffer(verifier)); + const auto* fb_fields = flatbuf::GetFooter(footer_data)->schema()->fields(); + ASSERT_EQ(fb_fields->size(), 1); + const flatbuf::Field* fb_field = fb_fields->Get(0); + ASSERT_EQ(fb_field->custom_metadata(), nullptr); + ASSERT_NE(fb_field->dictionary(), nullptr); + AssertKeyValues(KeyValueMetadata({"ARROW:extension:name", "ARROW:extension:metadata"}, + {"arrow.uuid", ""}), + fb_field->dictionary()->custom_metadata()); + + ASSERT_OK_AND_ASSIGN(auto reader, RecordBatchFileReader::Open( + std::make_shared(buffer))); + AssertSchemaEqual(*batch->schema(), *reader->schema()); + ASSERT_OK_AND_ASSIGN(auto read_batch, reader->ReadRecordBatch(0)); + AssertBatchesEqual(*batch, *read_batch); +} + +TEST_F(TestWriteRecordBatch, DictionaryOfExtension) { + // GH-49704: a dictionary whose value type is an extension type + auto storage = ArrayFromJSON(fixed_size_binary(16), + R"(["0123456789abcdef", "abcdef0123456789", null])"); + auto indices = ArrayFromJSON(int32(), "[0, 1, null, 1, 0, 2]"); + for (const auto& value_type : {extension::uuid(), uuid()}) { + ARROW_SCOPED_TRACE("value_type = ", value_type->ToString()); + auto values = ExtensionType::WrapArray(value_type, storage); + ASSERT_OK_AND_ASSIGN( + auto array, + DictionaryArray::FromArrays(dictionary(int32(), value_type), indices, values)); + CheckRoundtrip(array); + } +} + class RecursionLimits : public ::testing::Test, public io::MemoryMapFixture { public: void SetUp() { @@ -1948,6 +2288,72 @@ TEST(TestDictionaryBatch, DictionaryReplacement) { ASSERT_BATCHES_EQUAL(*in_batch, *out_batch); } +// GH-49704: delta and replacement dictionaries whose value type is an +// extension type +std::shared_ptr UuidDictionaryValues(const std::string& json) { + return ExtensionType::WrapArray(extension::uuid(), + ArrayFromJSON(fixed_size_binary(16), json)); +} + +TEST(TestDictionaryBatch, DictionaryOfExtensionDelta) { + auto type = dictionary(int32(), extension::uuid()); + auto schema = ::arrow::schema({field("f0", type)}); + auto dict1 = UuidDictionaryValues(R"(["0123456789abcdef", "abcdef0123456789"])"); + auto dict2 = UuidDictionaryValues(R"(["ffffffffffffffff"])"); + ASSERT_OK_AND_ASSIGN(auto full_dict, Concatenate({dict1, dict2})); + ASSERT_OK_AND_ASSIGN( + auto array, DictionaryArray::FromArrays( + type, ArrayFromJSON(int32(), "[2, 0, null, 1, 2]"), full_dict)); + auto in_batch = RecordBatch::Make(schema, array->length(), {array}); + + DictionaryBatchHelper helper(*schema); + ASSERT_OK(helper.Start()); + ASSERT_OK(helper.WriteDictionary(0L, dict1, /*is_delta=*/false)); + ASSERT_OK(helper.WriteDictionary(0L, dict2, /*is_delta=*/true)); + ASSERT_OK(helper.WriteBatchPayload(*in_batch)); + ASSERT_OK(helper.Close()); + + std::shared_ptr out_batch; + ASSERT_OK(helper.ReadBatch(&out_batch)); + ASSERT_OK(out_batch->ValidateFull()); + AssertSchemaEqual(*schema, *out_batch->schema()); + ASSERT_BATCHES_EQUAL(*in_batch, *out_batch); +} + +TEST(TestDictionaryBatch, DictionaryOfExtensionReplacement) { + auto type = dictionary(int32(), extension::uuid()); + auto schema = ::arrow::schema({field("f0", type)}); + auto make_batch = [&](const std::shared_ptr& dict) { + auto array = *DictionaryArray::FromArrays( + type, ArrayFromJSON(int32(), "[1, 0, null, 1]"), dict); + return RecordBatch::Make(schema, array->length(), {array}); + }; + auto dict1 = UuidDictionaryValues(R"(["0123456789abcdef", "abcdef0123456789"])"); + auto dict2 = UuidDictionaryValues(R"(["ffffffffffffffff", "0000000000000000"])"); + auto batch1 = make_batch(dict1); + auto batch2 = make_batch(dict2); + + // A stream with a dictionary replacement between two record batches + ASSERT_OK_AND_ASSIGN(auto sink, io::BufferOutputStream::Create()); + ASSERT_OK_AND_ASSIGN(auto writer, MakeStreamWriter(sink, schema)); + ASSERT_OK(writer->WriteRecordBatch(*batch1)); + ASSERT_OK(writer->WriteRecordBatch(*batch2)); + ASSERT_OK(writer->Close()); + ASSERT_EQ(writer->stats().num_replaced_dictionaries, 1); + ASSERT_OK_AND_ASSIGN(auto buffer, sink->Finish()); + + ASSERT_OK_AND_ASSIGN(auto reader, RecordBatchStreamReader::Open( + std::make_shared(buffer))); + AssertSchemaEqual(*schema, *reader->schema()); + for (const auto& expected : {batch1, batch2}) { + std::shared_ptr actual; + ASSERT_OK(reader->ReadNext(&actual)); + ASSERT_NE(actual, nullptr); + ASSERT_OK(actual->ValidateFull()); + ASSERT_BATCHES_EQUAL(*expected, *actual); + } +} + TEST_P(TestFileFormat, RoundTrip) { TestRoundTripWithOptions(*GetParam()); } TEST_P(TestFileFormatGenerator, RoundTrip) { TestRoundTripWithOptions(*GetParam()); } diff --git a/cpp/src/generated/Schema_generated.h b/cpp/src/generated/Schema_generated.h index 1090a120beaf..3b8fdbd4e28d 100644 --- a/cpp/src/generated/Schema_generated.h +++ b/cpp/src/generated/Schema_generated.h @@ -1960,7 +1960,8 @@ struct DictionaryEncoding FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table VT_ID = 4, VT_INDEXTYPE = 6, VT_ISORDERED = 8, - VT_DICTIONARYKIND = 10 + VT_DICTIONARYKIND = 10, + VT_CUSTOM_METADATA = 12 }; /// The known dictionary id in the application where this data is used. In /// the file or streaming formats, the dictionary ids are found in the @@ -1986,6 +1987,10 @@ struct DictionaryEncoding FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table org::apache::arrow::flatbuf::DictionaryKind dictionaryKind() const { return static_cast(GetField(VT_DICTIONARYKIND, 0)); } + /// Metadata for the dictionary value type, e.g. the extension type of the values + const ::flatbuffers::Vector<::flatbuffers::Offset> *custom_metadata() const { + return GetPointer> *>(VT_CUSTOM_METADATA); + } bool Verify(::flatbuffers::Verifier &verifier) const { return VerifyTableStart(verifier) && VerifyField(verifier, VT_ID, 8) && @@ -1993,6 +1998,9 @@ struct DictionaryEncoding FLATBUFFERS_FINAL_CLASS : private ::flatbuffers::Table verifier.VerifyTable(indexType()) && VerifyField(verifier, VT_ISORDERED, 1) && VerifyField(verifier, VT_DICTIONARYKIND, 2) && + VerifyOffset(verifier, VT_CUSTOM_METADATA) && + verifier.VerifyVector(custom_metadata()) && + verifier.VerifyVectorOfTables(custom_metadata()) && verifier.EndTable(); } }; @@ -2013,6 +2021,9 @@ struct DictionaryEncodingBuilder { void add_dictionaryKind(org::apache::arrow::flatbuf::DictionaryKind dictionaryKind) { fbb_.AddElement(DictionaryEncoding::VT_DICTIONARYKIND, static_cast(dictionaryKind), 0); } + void add_custom_metadata(::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset>> custom_metadata) { + fbb_.AddOffset(DictionaryEncoding::VT_CUSTOM_METADATA, custom_metadata); + } explicit DictionaryEncodingBuilder(::flatbuffers::FlatBufferBuilder &_fbb) : fbb_(_fbb) { start_ = fbb_.StartTable(); @@ -2029,15 +2040,34 @@ inline ::flatbuffers::Offset CreateDictionaryEncoding( int64_t id = 0, ::flatbuffers::Offset indexType = 0, bool isOrdered = false, - org::apache::arrow::flatbuf::DictionaryKind dictionaryKind = org::apache::arrow::flatbuf::DictionaryKind_DenseArray) { + org::apache::arrow::flatbuf::DictionaryKind dictionaryKind = org::apache::arrow::flatbuf::DictionaryKind_DenseArray, + ::flatbuffers::Offset<::flatbuffers::Vector<::flatbuffers::Offset>> custom_metadata = 0) { DictionaryEncodingBuilder builder_(_fbb); builder_.add_id(id); + builder_.add_custom_metadata(custom_metadata); builder_.add_indexType(indexType); builder_.add_dictionaryKind(dictionaryKind); builder_.add_isOrdered(isOrdered); return builder_.Finish(); } +inline ::flatbuffers::Offset CreateDictionaryEncodingDirect( + ::flatbuffers::FlatBufferBuilder &_fbb, + int64_t id = 0, + ::flatbuffers::Offset indexType = 0, + bool isOrdered = false, + org::apache::arrow::flatbuf::DictionaryKind dictionaryKind = org::apache::arrow::flatbuf::DictionaryKind_DenseArray, + const std::vector<::flatbuffers::Offset> *custom_metadata = nullptr) { + auto custom_metadata__ = custom_metadata ? _fbb.CreateVector<::flatbuffers::Offset>(*custom_metadata) : 0; + return org::apache::arrow::flatbuf::CreateDictionaryEncoding( + _fbb, + id, + indexType, + isOrdered, + dictionaryKind, + custom_metadata__); +} + /// ---------------------------------------------------------------------- /// A field represents a named column in a record / row batch or child of a /// nested type. diff --git a/dev/archery/archery/integration/datagen.py b/dev/archery/archery/integration/datagen.py index adffb2a1b21d..934aeabc68c2 100644 --- a/dev/archery/archery/integration/datagen.py +++ b/dev/archery/archery/integration/datagen.py @@ -1311,11 +1311,21 @@ def _get_children(self): return self.dictionary.field._get_children() def _get_dictionary(self): - return OrderedDict([ + entries = [ ('id', self.dictionary.id_), ('indexType', self.index_field._get_type()), ('isOrdered', self.dictionary.ordered) - ]) + ] + # The extension metadata of a dictionary value type is given in the + # "metadata" member of the "dictionary" object (GH-49704) + value_field = self.dictionary.field + if isinstance(value_field, ExtensionField): + extension_type = value_field.extension_type + entries.append(('metadata', metadata_key_values([ + ('ARROW:extension:name', extension_type.extension_name), + ('ARROW:extension:metadata', extension_type.serialized), + ]))) + return OrderedDict(entries) def generate_column(self, size, name=None): if name is None: @@ -1938,6 +1948,33 @@ def generate_extension_case(): dictionaries=[dict0]) +def generate_dictionary_extension_case(): + # Dictionaries whose value type is an extension type (GH-49704) + uuid_type = ExtensionType('arrow.uuid', '', + FixedSizeBinaryField('', 16)) + opaque_type = ExtensionType( + 'arrow.opaque', + '{"type_name":"geometry","vendor_name":"archery"}', + StringField('')) + + dict0 = Dictionary(0, ExtensionField('uuids', uuid_type), + size=5, name='DICT0') + dict1 = Dictionary(1, ExtensionField('opaques', opaque_type), + size=10, name='DICT1') + + fields = [ + DictionaryField('uuid_dict', get_field('', 'int8'), dict0), + DictionaryField('opaque_dict', get_field('', 'int32'), dict1, + metadata=[('some_key', 'some_value')]), + ListField('list_of_uuid_dict', + DictionaryField('item', get_field('', 'int16'), dict0)), + ] + + batch_sizes = [0, 13] + return _generate_file("dictionary_extension", fields, batch_sizes, + dictionaries=[dict0, dict1]) + + def generate_extension_wrapped_union_case(): # Unions wrapped in an extension type, exercising the extension/union # interaction across implementations (see the C++ fix in GH-50623). @@ -2072,6 +2109,24 @@ def get_generated_json_files(tempdir=None): .skip_tester('Ruby'), generate_extension_wrapped_union_case(), + + generate_dictionary_extension_case() + # TODO(https://github.com/apache/arrow/issues/49704): implement + # DictionaryEncoding.custom_metadata + .skip_tester('.NET') + .skip_tester('Java') + .skip_tester('JS') + .skip_tester('nanoarrow') + .skip_tester('Ruby') + # Go support is in the companion arrow-go pull request; enable it once + # that lands, since integration CI builds arrow-go from main + .skip_tester('Go') + # Rust cannot represent a dictionary of an extension type + # (https://github.com/apache/arrow-rs/issues/7982): it reads a + # dictionary of the storage type and does not write the value type + # metadata. It could consume this case but cannot produce it, and + # archery has no producer-only skip. + .skip_tester('Rust'), ] generated_paths = [] diff --git a/docs/source/format/Columnar.rst b/docs/source/format/Columnar.rst index f95d7e80769d..2578f973f6dc 100644 --- a/docs/source/format/Columnar.rst +++ b/docs/source/format/Columnar.rst @@ -21,7 +21,7 @@ Arrow Columnar Format ********************* -*Version: 1.5* +*Version: 1.6* .. seealso:: :ref:`Additions to the Arrow columnar format since version 1.0.0 ` @@ -1289,7 +1289,10 @@ array. This includes: * A ``dictionary`` property indicating whether the field is dictionary-encoded or not. If it is, a dictionary "id" is assigned to allow matching a subsequent dictionary IPC message with the - appropriate field. + appropriate field. In this case, the field's data type is the + dictionary value type, and the ``dictionary`` property may carry the + ``custom_metadata`` of the dictionary value type (see + :ref:`format_metadata_extension_types_dictionary`). We additionally provide both schema-level and field-level ``custom_metadata`` attributes allowing for systems to insert their @@ -1677,7 +1680,10 @@ Custom Application Metadata We provide a ``custom_metadata`` field at three levels to provide a mechanism for developers to pass application-specific metadata in Arrow protocol messages. This includes ``Field``, ``Schema``, and -``Message``. +``Message``. The ``DictionaryEncoding`` of a dictionary-encoded ``Field`` +additionally has a ``custom_metadata`` field, with the same semantics as +the ``Field``'s ``custom_metadata`` but applied to the dictionary value +type (see :ref:`format_metadata_extension_types_dictionary`). The colon symbol ``:`` is to be used as a namespace separator. It can be used multiple times in a key. @@ -1732,6 +1738,55 @@ extension types: metadata indicating the market trading calendar the data corresponds to +.. _format_metadata_extension_types_dictionary: + +Dictionary-encoded extension types +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +A dictionary-encoded field is serialized as a single ``Field`` whose data +type is the dictionary value type and whose ``dictionary`` property is a +``DictionaryEncoding`` giving the index type. The ``Field``'s +``custom_metadata`` describes the field's own type, so an extension type +whose storage type is a dictionary type stores its extension keys in the +``Field``'s ``custom_metadata``, as any other extension type does. + +The ``custom_metadata`` of the ``DictionaryEncoding`` is the +``custom_metadata`` of the dictionary value type. It has the same semantics +as a ``Field``'s ``custom_metadata``, including the reserved +``'ARROW:extension:*'`` keys, but applies to the dictionary values rather +than to the field. It is the IPC counterpart of the metadata of the +dictionary value ``ArrowSchema`` in the :ref:`C Data Interface +` (``ArrowSchema.dictionary->metadata``), while the +``Field``'s ``custom_metadata`` corresponds to the metadata of the parent +``ArrowSchema`` (``ArrowSchema.metadata``). + +If the dictionary *value* type is an extension type, for example a +dictionary of ``arrow.uuid`` values, writers should store the +``'ARROW:extension:name'`` and ``'ARROW:extension:metadata'`` keys of the +value type in the ``custom_metadata`` of the ``DictionaryEncoding`` rather +than in the ``Field``'s ``custom_metadata``. Both may be present, such as +for an extension type whose storage type is a dictionary of extension +values. Writers should omit the ``DictionaryEncoding``'s ``custom_metadata`` +when the dictionary value type has no metadata. + +When reading a dictionary-encoded field, the dictionary value type is +reconstructed from the ``Field``'s data type and the ``DictionaryEncoding``'s +``custom_metadata``. The result is wrapped in the dictionary type given by +the ``DictionaryEncoding``, and then the ``Field``'s ``custom_metadata`` is +applied to that dictionary type. As for other extension types, a reader +that does not recognize the extension name in the ``DictionaryEncoding``'s +``custom_metadata`` uses the dictionary value's storage type. + +Readers that predate this ``custom_metadata`` field (added in format version +1.6) ignore it and read a dictionary of the storage type. The field is +additive: the ``MetadataVersion`` remains ``V5``, and the IPC file footer +and Arrow Flight, which embed the same ``Schema``, need no other change. + +.. note:: + Before format version 1.6, a dictionary whose value type is an extension + type could not be distinguished from an extension type whose storage type + is a dictionary, as only the ``Field``'s ``custom_metadata`` was available. + .. seealso:: :ref:`format_canonical_extensions` diff --git a/docs/source/format/Integration.rst b/docs/source/format/Integration.rst index 5038132241c3..3129c95a1942 100644 --- a/docs/source/format/Integration.rst +++ b/docs/source/format/Integration.rst @@ -198,14 +198,19 @@ present if there are dictionary type fields in the schema. "dictionary": { "id": /* integer */, "indexType": /* Type */, - "isOrdered": /* boolean */ + "isOrdered": /* boolean */, + "metadata" : /* Metadata */ }, "metadata" : /* Metadata */ } The ``dictionary`` attribute is present if and only if the ``Field`` corresponds to a dictionary type, and its ``id`` maps onto a column in the ``DictionaryBatch``. In this -case the ``type`` attribute describes the value type of the dictionary. +case the ``type`` attribute describes the value type of the dictionary, and the +optional ``metadata`` attribute of ``dictionary`` holds the custom metadata of the +dictionary value type, as the ``custom_metadata`` of ``DictionaryEncoding`` does in +the IPC format (see :ref:`format_metadata_extension_types_dictionary`). It may be +omitted when the dictionary value type has no metadata. For primitive types, ``children`` is an empty array. @@ -424,6 +429,34 @@ would be represented:: ] } +A dictionary whose value type is an extension type carries the extension +metadata in the ``metadata`` attribute of its ``dictionary`` object instead of +the field's ``metadata``. For example, a dictionary of ``arrow.uuid`` values +with ``int8`` indices would be represented as:: + + { + "name" : "name_of_the_field", + "nullable" : /* boolean */, + "type" : { + "name" : "fixedsizebinary", + "byteWidth" : 16 + }, + "children" : [], + "dictionary": { + "id": 0, + "indexType": { + "name": "int", + "bitWidth": 8, + "isSigned": true + }, + "isOrdered": false, + "metadata" : [ + {"key": "ARROW:extension:name", "value": "arrow.uuid"}, + {"key": "ARROW:extension:metadata", "value": ""} + ] + } + } + **RecordBatch**:: { diff --git a/docs/source/format/Versioning.rst b/docs/source/format/Versioning.rst index 52cf207b959b..6e8e0e04d21c 100644 --- a/docs/source/format/Versioning.rst +++ b/docs/source/format/Versioning.rst @@ -76,7 +76,7 @@ that new clients can read serialized data produced by library version Post-1.0.0 Format Versions ========================== -Since version 1.0.0, there have been five new minor versions and zero new +Since version 1.0.0, there have been six new minor versions and zero new major versions of the Arrow format. Each new minor version added new features. When these new features are not used, the new minor format versions are compatible with format version 1.0.0. The new features added in each minor @@ -110,3 +110,10 @@ Version 1.5 ----------- * Expanded Decimal type bit widths to allow 32-bit and 64-bit types. + +Version 1.6 +----------- + +* Added ``custom_metadata`` to ``DictionaryEncoding``, carrying the custom + metadata of the dictionary value type, such as an extension type (see + :ref:`format_metadata_extension_types_dictionary`). diff --git a/format/Schema.fbs b/format/Schema.fbs index 67b771912828..46586f65bd11 100644 --- a/format/Schema.fbs +++ b/format/Schema.fbs @@ -25,6 +25,7 @@ /// Version 1.4 - Add BinaryView, Utf8View, variadicBufferCounts, ListView, and /// LargeListView. /// Version 1.5 - Add 32-bit and 64-bit as allowed bit widths for Decimal +/// Version 1.6 - Add custom_metadata to DictionaryEncoding. namespace org.apache.arrow.flatbuf; @@ -511,6 +512,9 @@ table DictionaryEncoding { isOrdered: bool; dictionaryKind: DictionaryKind; + + /// Metadata for the dictionary value type, e.g. the extension type of the values + custom_metadata: [ KeyValue ]; } /// ---------------------------------------------------------------------- diff --git a/python/pyarrow/tests/test_extension_type.py b/python/pyarrow/tests/test_extension_type.py index a9d15cfec10a..04140c00c4b4 100644 --- a/python/pyarrow/tests/test_extension_type.py +++ b/python/pyarrow/tests/test_extension_type.py @@ -1431,6 +1431,22 @@ def test_uuid_extension(): assert isinstance(array[0], pa.UuidScalar) +def test_uuid_dictionary_ipc(): + # GH-49704: a dictionary whose value type is an extension type + storage = pa.array([b"0123456789abcdef", b"zyxwvutsrqponmlk"], + pa.binary(16)) + dictionary = pa.ExtensionArray.from_storage(pa.uuid(), storage) + indices = pa.array([0, 1, None, 1, 0], pa.int32()) + array = pa.DictionaryArray.from_arrays(indices, dictionary) + assert array.type == pa.dictionary(pa.int32(), pa.uuid()) + + buf = ipc_write_batch(pa.RecordBatch.from_arrays([array], ["uuid_dict"])) + + batch = ipc_read_batch(buf) + assert batch.schema.field("uuid_dict").type == array.type + assert batch.column(0) == array + + @pytest.mark.numpy def test_uuid_to_numpy_uses_storage(): value = uuid4()