diff --git a/cpp/src/arrow/c/bridge.cc b/cpp/src/arrow/c/bridge.cc index 184be3ab8eb9..6e8bdaa02316 100644 --- a/cpp/src/arrow/c/bridge.cc +++ b/cpp/src/arrow/c/bridge.cc @@ -269,7 +269,8 @@ struct SchemaExporter { // child dictionary struct describes value type. RETURN_NOT_OK(VisitTypeInline(*dict_type.index_type(), this)); dict_exporter_ = std::make_unique(); - RETURN_NOT_OK(dict_exporter_->ExportType(*dict_type.value_type())); + RETURN_NOT_OK(dict_exporter_->ExportField( + Field("", dict_type.value_type(), true, dict_type.value_metadata()))); } else { RETURN_NOT_OK(VisitTypeInline(type, this)); } @@ -1064,7 +1065,8 @@ struct SchemaImporter { SchemaImporter dict_importer; RETURN_NOT_OK(dict_importer.ImportDict(this, c_struct_->dictionary)); bool ordered = (c_struct_->flags & ARROW_FLAG_DICTIONARY_ORDERED) != 0; - type_ = dictionary(type_, dict_importer.type_, ordered); + type_ = dictionary(type_, dict_importer.type_, ordered, + std::move(dict_importer.metadata_.metadata)); } // Import metadata diff --git a/cpp/src/arrow/c/bridge_test.cc b/cpp/src/arrow/c/bridge_test.cc index 9a4d104d0004..60b3eb1494ad 100644 --- a/cpp/src/arrow/c/bridge_test.cc +++ b/cpp/src/arrow/c/bridge_test.cc @@ -2280,6 +2280,38 @@ TEST_F(TestSchemaImport, RegisteredExtension) { } } +TEST(TestDictionaryValueMetadata, CDataRoundTrip) { + auto metadata = key_value_metadata({{"extra", std::string("a\0b", 3)}}); + auto parent_metadata = key_value_metadata({{"extra", "parent"}}); + auto type = dictionary(int32(), uuid(), true, metadata); + auto expected = + field("d", struct_({field("values", list(type))}), true, parent_metadata); + auto roundtrip = + [](const std::shared_ptr& input) -> Result> { + ArrowSchema c_schema{}; + RETURN_NOT_OK(ExportField(*input, &c_schema)); + SchemaExportGuard guard(&c_schema); + return ImportField(&c_schema); + }; + + ASSERT_EQ(GetExtensionType("uuid"), nullptr); + ASSERT_OK_AND_ASSIGN(auto unknown, roundtrip(expected)); + const auto& dictionary_type = checked_cast( + *unknown->type()->field(0)->type()->field(0)->type()); + AssertTypeEqual(*fixed_size_binary(16), *dictionary_type.value_type()); + auto unknown_metadata = metadata->Copy(); + unknown_metadata->Append(kExtensionTypeKeyName, "uuid"); + unknown_metadata->Append(kExtensionMetadataKeyName, "uuid-serialized"); + ASSERT_TRUE(unknown_metadata->Equals(*dictionary_type.value_metadata())); + ASSERT_TRUE(parent_metadata->Equals(*unknown->metadata())); + + ASSERT_OK_AND_ASSIGN(auto rewritten, roundtrip(unknown)); + ASSERT_TRUE(unknown->Equals(rewritten, true)); + ExtensionTypeGuard extension(uuid()); + ASSERT_OK_AND_ASSIGN(auto restored, roundtrip(rewritten)); + ASSERT_TRUE(expected->Equals(restored, true)); +} + TEST_F(TestSchemaImport, FormatStringError) { FillPrimitive(""); CheckImportError(); diff --git a/cpp/src/arrow/compare.cc b/cpp/src/arrow/compare.cc index 06f1d14e73a3..20b36fe43025 100644 --- a/cpp/src/arrow/compare.cc +++ b/cpp/src/arrow/compare.cc @@ -875,7 +875,7 @@ class TypeEqualsVisitor { Status Visit(const DictionaryType& left) { const auto& right = checked_cast(right_); result_ = left.index_type()->Equals(right.index_type()) && - left.value_type()->Equals(right.value_type()) && + left.value_type()->Equals(right.value_type(), check_metadata_) && (left.ordered() == right.ordered()); return Status::OK(); } diff --git a/cpp/src/arrow/integration/json_integration_test.cc b/cpp/src/arrow/integration/json_integration_test.cc index 0f7e6c85dc55..f7235d03d66c 100644 --- a/cpp/src/arrow/integration/json_integration_test.cc +++ b/cpp/src/arrow/integration/json_integration_test.cc @@ -31,6 +31,8 @@ #include "arrow/array.h" #include "arrow/array/builder_binary.h" #include "arrow/array/builder_primitive.h" +#include "arrow/extension/opaque.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 +50,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 +728,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 +1289,162 @@ 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 + auto unknown_type = + dictionary(int8(), fixed_size_binary(16), false, + key_value_metadata({{kExtensionTypeKeyName, "unknown.extension"}, + {kExtensionMetadataKeyName, ""}})); + Schema ex_schema({field("uuid_dict", dictionary(int8(), uuid_type), /*nullable=*/true, + key_value_metadata({{"some_key", "some_value"}})), + field("unknown_dict", unknown_type)}); + + 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( + unknown_type, 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()), + field("f3", dictionary(int32(), + extension::opaque( + struct_({field("inner", dictionary(int8(), uuid_type))}), + "nested", "test"))), + }); + 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(TestJsonFileReadWrite, DictionaryValueMetadataRoundTrip) { + auto metadata = key_value_metadata({{"extra", std::string("a\0b", 3)}}); + auto type = dictionary(int32(), uuid(), true, metadata); + auto values = ExtensionType::WrapArray( + uuid(), ArrayFromJSON(fixed_size_binary(16), R"(["0123456789abcdef", null])")); + ASSERT_OK_AND_ASSIGN( + auto dict, + DictionaryArray::FromArrays(type, ArrayFromJSON(int32(), "[0, null, 1]"), values)); + ASSERT_OK_AND_ASSIGN(auto nested, StructArray::Make({dict}, {field("d", type)})); + auto batch = + RecordBatch::Make(schema({field("nested", nested->type(), true, + key_value_metadata({{"extra", "parent"}}))}), + nested->length(), {nested}); + auto roundtrip = [](const std::shared_ptr& input) + -> Result> { + ARROW_ASSIGN_OR_RAISE(auto writer, IntegrationJsonWriter::Open(input->schema())); + RETURN_NOT_OK(writer->WriteRecordBatch(*input)); + ARROW_ASSIGN_OR_RAISE(auto json_data, writer->Finish()); + ARROW_ASSIGN_OR_RAISE( + auto reader, IntegrationJsonReader::Open(std::make_shared(json_data))); + return reader->ReadRecordBatch(0); + }; + ASSERT_EQ(GetExtensionType("uuid"), nullptr); + ASSERT_OK_AND_ASSIGN(auto unknown, roundtrip(batch)); + const auto& unknown_type = checked_cast( + *unknown->schema()->field(0)->type()->field(0)->type()); + AssertTypeEqual(*fixed_size_binary(16), *unknown_type.value_type()); + auto expected_metadata = metadata->Copy(); + expected_metadata->Append(kExtensionTypeKeyName, "uuid"); + expected_metadata->Append(kExtensionMetadataKeyName, "uuid-serialized"); + ASSERT_TRUE(expected_metadata->Equals(*unknown_type.value_metadata())); + ASSERT_OK_AND_ASSIGN(auto rewritten, roundtrip(unknown)); + AssertSchemaEqual(*unknown->schema(), *rewritten->schema(), true); + AssertBatchesEqual(*unknown, *rewritten); + ExtensionTypeGuard guard(uuid()); + ASSERT_OK_AND_ASSIGN(auto restored, roundtrip(rewritten)); + ASSERT_OK(restored->ValidateFull()); + AssertSchemaEqual(*batch->schema(), *restored->schema(), true); + AssertBatchesEqual(*batch, *restored); +} + 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..87aa37176af4 100644 --- a/cpp/src/arrow/integration/json_internal.cc +++ b/cpp/src/arrow/integration/json_internal.cc @@ -164,10 +164,10 @@ class SchemaWriter { writer_->StartObject(); writer_->Key("key"); - writer_->String(key.c_str()); + writer_->String(key); writer_->Key("value"); - writer_->String(value.c_str()); + writer_->String(value); writer_->EndObject(); } @@ -187,6 +187,17 @@ 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(type.value_metadata(), + {{kExtensionTypeKeyName, ext_type.extension_name()}, + {kExtensionMetadataKeyName, ext_type.Serialize()}}); + } else { + WriteKeyValueMetadata(type.value_metadata()); + } writer_->EndObject(); return Status::OK(); @@ -220,7 +231,11 @@ class SchemaWriter { const auto& dict_type = checked_cast(*type); // Ensure we visit child fields first so that, in the case of nested // dictionaries, inner dictionaries get a smaller id than outer dictionaries. - RETURN_NOT_OK(WriteChildren(dict_type.value_type()->fields(), field_pos)); + const DataType* value_type = dict_type.value_type().get(); + while (value_type->id() == Type::EXTENSION) { + value_type = checked_cast(*value_type).storage_type().get(); + } + RETURN_NOT_OK(WriteChildren(value_type->fields(), field_pos)); ARROW_ASSIGN_OR_RAISE(const int64_t dictionary_id, mapper_.GetFieldId(field_pos.path())); RETURN_NOT_OK(WriteDictionaryMetadata(dictionary_id, dict_type)); @@ -451,7 +466,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,8 +1323,27 @@ 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); + type = ::arrow::dictionary(index_type, type, is_ordered, std::move(value_metadata)); } } diff --git a/cpp/src/arrow/ipc/dictionary.cc b/cpp/src/arrow/ipc/dictionary.cc index bb5539f2b4bb..9d3e7e614e7a 100644 --- a/cpp/src/arrow/ipc/dictionary.cc +++ b/cpp/src/arrow/ipc/dictionary.cc @@ -101,14 +101,17 @@ struct DictionaryFieldMapper::Impl { void ImportField(const FieldPosition& pos, const Field& field) { const DataType* type = field.type().get(); - if (type->id() == Type::EXTENSION) { + while (type->id() == Type::EXTENSION) { type = checked_cast(*type).storage_type().get(); } if (type->id() == Type::DICTIONARY) { InsertPath(pos); // Import nested dictionaries - ImportFields(pos, - checked_cast(*type).value_type()->fields()); + type = checked_cast(*type).value_type().get(); + while (type->id() == Type::EXTENSION) { + type = checked_cast(*type).storage_type().get(); + } + ImportFields(pos, type->fields()); } else { ImportFields(pos, type->fields()); } @@ -304,7 +307,7 @@ struct DictionaryCollector { const Array* array) { const DataType* type = array->type().get(); - if (type->id() == Type::EXTENSION) { + while (type->id() == Type::EXTENSION) { type = checked_cast(*type).storage_type().get(); array = checked_cast(*array).storage().get(); } @@ -314,8 +317,11 @@ struct DictionaryCollector { // Traverse the dictionary to first gather any nested dictionaries // (so that they appear in the output before their parent) - const auto& dict_type = checked_cast(*type); - RETURN_NOT_OK(WalkChildren(position, *dict_type.value_type(), *dictionary)); + const Array* storage = dictionary.get(); + while (storage->type_id() == Type::EXTENSION) { + storage = checked_cast(*storage).storage().get(); + } + RETURN_NOT_OK(WalkChildren(position, *storage->type(), *storage)); // Then record the dictionary itself ARROW_ASSIGN_OR_RAISE(int64_t id, mapper_.GetFieldId(position.path())); diff --git a/cpp/src/arrow/ipc/metadata_internal.cc b/cpp/src/arrow/ipc/metadata_internal.cc index 02134aa805ab..41cf53641161 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,25 @@ 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 (type.value_metadata()) { + const auto& metadata = *type.value_metadata(); + for (int64_t i = 0; i < metadata.size(); ++i) { + dictionary_value_metadata_.emplace_back(metadata.key(i), metadata.value(i)); + } + } + 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 +842,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 +889,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,35 +959,20 @@ 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) + 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())); + DictionaryType::Make(index_type, type, encoding->isOrdered(), + std::move(value_metadata))); dictionary_id = encoding->id(); } // 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..c2d2b9c3620a 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,372 @@ 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)); + auto uuid_metadata = key_value_metadata( + {{kExtensionTypeKeyName, "uuid"}, {kExtensionMetadataKeyName, "uuid-serialized"}}); + AssertSchemaEqual(Schema({field("f0", dictionary(int32(), fixed_size_binary(16), false, + uuid_metadata))}), + *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"}})); + auto unknown_metadata = + key_value_metadata({{kExtensionTypeKeyName, "unknown.extension"}, + {kExtensionMetadataKeyName, "unknown-serialized"}}); + AssertSchemaEqual(Schema({field("f0", dictionary(int32(), fixed_size_binary(16), false, + unknown_metadata))}), + *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, DictionaryEncodingMetadataPreserved) { + auto extra_metadata = key_value_metadata({{"some_key", "some_value"}}); + auto expected_storage = Schema( + {field("f0", dictionary(int32(), fixed_size_binary(16), false, extra_metadata))}); + ASSERT_OK_AND_ASSIGN( + auto schema, ReadSingleDictionaryFieldSchema(flatbuf::Type::Type_FixedSizeBinary, + 32, {}, {{"some_key", "some_value"}})); + AssertSchemaEqual(expected_storage, *schema, /*check_metadata=*/true); + CheckSchemaRoundtrip(*schema); + + ASSERT_OK_AND_ASSIGN(schema, ReadSingleDictionaryFieldSchema( + flatbuf::Type::Type_FixedSizeBinary, 32, {}, + {{"ARROW:extension:name", "unknown.extension"}, + {"ARROW:extension:metadata", ""}, + {"some_key", "some_value"}})); + auto unknown_metadata = + key_value_metadata({{kExtensionTypeKeyName, "unknown.extension"}, + {kExtensionMetadataKeyName, ""}, + {"some_key", "some_value"}}); + AssertSchemaEqual(Schema({field("f0", dictionary(int32(), fixed_size_binary(16), false, + unknown_metadata))}), + *schema, /*check_metadata=*/true); + CheckSchemaRoundtrip(*schema); + + 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(), false, + extra_metadata))}), + *schema, /*check_metadata=*/true); + CheckSchemaRoundtrip(*schema); +} + +TEST(TestDictionaryValueMetadata, UnknownExtensionRoundTrip) { + auto value_type = uuid(); // Deliberately unregistered during the first two reads. + auto extra = key_value_metadata({{"extra", std::string("a\0b", 3)}}); + auto field_metadata = key_value_metadata({{"extra", "field value"}}); + auto type = dictionary(int32(), value_type, true, extra); + auto storage = ArrayFromJSON(fixed_size_binary(16), + R"(["0123456789abcdef", null, "abcdef0123456789"])"); + auto values = ExtensionType::WrapArray(value_type, storage); + ASSERT_OK_AND_ASSIGN(auto array, + DictionaryArray::FromArrays( + type, ArrayFromJSON(int32(), "[2, null, 0, 1]"), values)); + // An extension value can itself contain a dictionary. Both dictionary + // annotations must survive while neither extension is registered. + ASSERT_OK_AND_ASSIGN(auto nested_values, + StructArray::Make({array}, {field("inner", type)})); + auto nested_extension = + std::make_shared(nested_values->type()); + auto nested_type = dictionary(int16(), nested_extension, false, extra); + ASSERT_OK_AND_ASSIGN(auto nested, + DictionaryArray::FromArrays( + nested_type, ArrayFromJSON(int16(), "[0, 2, null, 3]"), + ExtensionType::WrapArray(nested_extension, nested_values))); + auto batch = RecordBatch::Make( + schema({field("f0", type, true, field_metadata), field("nested", nested_type)}), + array->length(), {array, nested}); + + for (bool file : {false, true}) { + ARROW_SCOPED_TRACE("file = ", file); + auto roundtrip = [file](const std::shared_ptr& input) + -> Result> { + ARROW_ASSIGN_OR_RAISE(auto sink, io::BufferOutputStream::Create()); + ARROW_ASSIGN_OR_RAISE(auto writer, file ? MakeFileWriter(sink, input->schema()) + : MakeStreamWriter(sink, input->schema())); + RETURN_NOT_OK(writer->WriteRecordBatch(*input)); + RETURN_NOT_OK(writer->Close()); + ARROW_ASSIGN_OR_RAISE(auto buffer, sink->Finish()); + auto source = std::make_shared(buffer); + if (file) { + ARROW_ASSIGN_OR_RAISE(auto reader, RecordBatchFileReader::Open(source)); + return reader->ReadRecordBatch(0); + } + ARROW_ASSIGN_OR_RAISE(auto reader, RecordBatchStreamReader::Open(source)); + return reader->Next(); + }; + + ASSERT_EQ(GetExtensionType("uuid"), nullptr); + ASSERT_OK_AND_ASSIGN(auto unknown, roundtrip(batch)); + ASSERT_OK(unknown->ValidateFull()); + const auto& unknown_type = + checked_cast(*unknown->schema()->field(0)->type()); + AssertTypeEqual(*fixed_size_binary(16), *unknown_type.value_type()); + auto expected_metadata = extra->Copy(); + expected_metadata->Append(kExtensionTypeKeyName, "uuid"); + expected_metadata->Append(kExtensionMetadataKeyName, "uuid-serialized"); + ASSERT_TRUE(expected_metadata->Equals(*unknown_type.value_metadata())); + ASSERT_TRUE(field_metadata->Equals(*unknown->schema()->field(0)->metadata())); + AssertArraysEqual( + *storage, + *checked_cast(*unknown->column(0)).dictionary()); + + ASSERT_OK_AND_ASSIGN(auto rewritten, roundtrip(unknown)); + AssertSchemaEqual(*unknown->schema(), *rewritten->schema(), true); + AssertBatchesEqual(*unknown, *rewritten); + + ExtensionTypeGuard guard({value_type, nested_extension}); + ASSERT_OK_AND_ASSIGN(auto restored, roundtrip(rewritten)); + ASSERT_OK(restored->ValidateFull()); + AssertSchemaEqual(*batch->schema(), *restored->schema(), true); + AssertBatchesEqual(*batch, *restored); + } +} + +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 +1415,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 +2375,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/arrow/type.cc b/cpp/src/arrow/type.cc index e841f0918097..ee7c4674f43d 100644 --- a/cpp/src/arrow/type.cc +++ b/cpp/src/arrow/type.cc @@ -34,6 +34,7 @@ #include "arrow/array.h" #include "arrow/compare.h" +#include "arrow/extension_type.h" #include "arrow/record_batch.h" #include "arrow/result.h" #include "arrow/status.h" @@ -629,7 +630,8 @@ Result> MergeDictionaryTypes( MergeTypes(left.value_type(), right.value_type(), options)); auto ordered = left.ordered() && right.ordered(); if (indices && values) { - return dictionary(indices, values, ordered); + // Like Field::MergeWith, retain the metadata from the first occurrence. + return dictionary(indices, values, ordered, left.value_metadata()); } else if (values) { return Status::TypeError("Could not merge dictionary index types"); } @@ -1562,17 +1564,45 @@ int DictionaryType::bit_width() const { Result> DictionaryType::Make( const std::shared_ptr& index_type, const std::shared_ptr& value_type, bool ordered) { + return Make(index_type, value_type, ordered, nullptr); +} + +Result> DictionaryType::Make( + const std::shared_ptr& index_type, + const std::shared_ptr& value_type, bool ordered, + std::shared_ptr value_metadata) { RETURN_NOT_OK(ValidateParameters(*index_type, *value_type)); - return std::make_shared(index_type, value_type, ordered); + return std::make_shared(index_type, value_type, ordered, + std::move(value_metadata)); } DictionaryType::DictionaryType(const std::shared_ptr& index_type, const std::shared_ptr& value_type, bool ordered) + : DictionaryType(index_type, value_type, ordered, nullptr) {} + +DictionaryType::DictionaryType(const std::shared_ptr& index_type, + const std::shared_ptr& value_type, bool ordered, + std::shared_ptr value_metadata) : FixedWidthType(Type::DICTIONARY), index_type_(index_type), value_type_(value_type), - ordered_(ordered) { + ordered_(ordered), + value_metadata_(std::move(value_metadata)) { ARROW_CHECK_OK(ValidateParameters(*index_type_, *value_type_)); + if (value_type_->id() == Type::EXTENSION && value_metadata_ && + (value_metadata_->Contains(kExtensionTypeKeyName) || + value_metadata_->Contains(kExtensionMetadataKeyName))) { + // The ExtensionType is authoritative. Do not serialize conflicting or + // duplicate extension keys, or mutate the caller's metadata. + auto remaining = std::make_shared(); + for (int64_t i = 0; i < value_metadata_->size(); ++i) { + const auto& key = value_metadata_->key(i); + if (key != kExtensionTypeKeyName && key != kExtensionMetadataKeyName) { + remaining->Append(key, value_metadata_->value(i)); + } + } + value_metadata_ = std::move(remaining); + } } DataTypeLayout DictionaryType::layout() const { @@ -1586,6 +1616,9 @@ std::string DictionaryType::ToString(bool show_metadata) const { ss << this->name() << "ToString(show_metadata) << ", indices=" << index_type_->ToString(show_metadata) << ", ordered=" << ordered_ << ">"; + if (show_metadata && value_metadata_ && value_metadata_->size() > 0) { + ss << value_metadata_->ToString(); + } return ss.str(); } @@ -2931,7 +2964,21 @@ std::string DictionaryType::ComputeFingerprint() const { return TypeIdFingerprint(*this) + index_fingerprint + value_fingerprint + ordered_fingerprint; } - return ordered_fingerprint; + // Extension types generally have no fingerprint. Fall back to comparing + // the actual types instead of treating every unordered dictionary as equal. + return ""; +} + +std::string DictionaryType::ComputeMetadataFingerprint() const { + std::stringstream ss; + if (value_metadata_) { + AppendMetadataFingerprint(*value_metadata_, &ss); + } + const auto& value_fingerprint = value_type_->metadata_fingerprint(); + if (!value_fingerprint.empty()) { + ss << "+{" << value_fingerprint << "}"; + } + return ss.str(); } std::string ListType::ComputeFingerprint() const { @@ -3338,6 +3385,14 @@ std::shared_ptr dictionary(const std::shared_ptr& index_type return std::make_shared(index_type, dict_type, ordered); } +std::shared_ptr dictionary( + const std::shared_ptr& index_type, + const std::shared_ptr& dict_type, bool ordered, + std::shared_ptr value_metadata) { + return std::make_shared(index_type, dict_type, ordered, + std::move(value_metadata)); +} + std::shared_ptr field(std::string name, std::shared_ptr type, bool nullable, std::shared_ptr metadata) { diff --git a/cpp/src/arrow/type.h b/cpp/src/arrow/type.h index 14a05c369e3b..dd84a67ae6b0 100644 --- a/cpp/src/arrow/type.h +++ b/cpp/src/arrow/type.h @@ -1933,11 +1933,20 @@ class ARROW_EXPORT DictionaryType : public FixedWidthType { DictionaryType(const std::shared_ptr& index_type, const std::shared_ptr& value_type, bool ordered = false); + DictionaryType(const std::shared_ptr& index_type, + const std::shared_ptr& value_type, bool ordered, + std::shared_ptr value_metadata); + // A constructor variant that validates its input parameters static Result> Make( const std::shared_ptr& index_type, const std::shared_ptr& value_type, bool ordered = false); + static Result> Make( + const std::shared_ptr& index_type, + const std::shared_ptr& value_type, bool ordered, + std::shared_ptr value_metadata); + std::string ToString(bool show_metadata = false) const override; std::string name() const override { return "dictionary"; } @@ -1948,6 +1957,14 @@ class ARROW_EXPORT DictionaryType : public FixedWidthType { const std::shared_ptr& index_type() const { return index_type_; } const std::shared_ptr& value_type() const { return value_type_; } + /// Metadata of the dictionary values, separate from the enclosing Field metadata. + /// This includes unregistered extension annotations. When value_type() is an + /// ExtensionType, its extension keys are represented by that type instead and + /// are removed from the supplied metadata without modifying it. + const std::shared_ptr& value_metadata() const { + return value_metadata_; + } + bool ordered() const { return ordered_; } protected: @@ -1955,11 +1972,13 @@ class ARROW_EXPORT DictionaryType : public FixedWidthType { const DataType& value_type); std::string ComputeFingerprint() const override; + std::string ComputeMetadataFingerprint() const override; // Must be an integer type (not currently checked) std::shared_ptr index_type_; std::shared_ptr value_type_; bool ordered_; + std::shared_ptr value_metadata_; }; // ---------------------------------------------------------------------- diff --git a/cpp/src/arrow/type_fwd.h b/cpp/src/arrow/type_fwd.h index a1c1ac20f694..8d33408c3452 100644 --- a/cpp/src/arrow/type_fwd.h +++ b/cpp/src/arrow/type_fwd.h @@ -678,6 +678,17 @@ std::shared_ptr dictionary(const std::shared_ptr& index_type const std::shared_ptr& dict_type, bool ordered = false); +/// \brief Create a DictionaryType with metadata for its values +/// +/// Unregistered extension annotations and other value metadata are preserved +/// independently of the enclosing Field's metadata. If dict_type is an extension +/// type, its extension annotations take precedence over those in value_metadata. +ARROW_EXPORT +std::shared_ptr dictionary( + const std::shared_ptr& index_type, + const std::shared_ptr& dict_type, bool ordered, + std::shared_ptr value_metadata); + /// @} /// \defgroup schema-factories Factory functions for fields and schemas diff --git a/cpp/src/arrow/type_test.cc b/cpp/src/arrow/type_test.cc index 133afbfa0c58..257ff70809ce 100644 --- a/cpp/src/arrow/type_test.cc +++ b/cpp/src/arrow/type_test.cc @@ -32,6 +32,7 @@ #include "arrow/array.h" #include "arrow/memory_pool.h" #include "arrow/table.h" +#include "arrow/testing/extension_type.h" #include "arrow/testing/gtest_util.h" #include "arrow/testing/matchers.h" #include "arrow/testing/random.h" @@ -2298,6 +2299,56 @@ TEST(TestDictionaryType, Equals) { AssertTypeNotEqual(*t5, *t6); } +TEST(TestDictionaryType, ValueMetadata) { + auto metadata = key_value_metadata({{"key", "value"}}); + auto plain = dictionary(int8(), int32()); + auto annotated = dictionary(int8(), int32(), false, metadata); + auto same = dictionary(int8(), int32(), false, metadata->Copy()); + auto different = + dictionary(int8(), int32(), false, key_value_metadata({{"key", "different"}})); + ASSERT_TRUE(plain->Equals(annotated)); + ASSERT_FALSE(plain->Equals(annotated, true)); + ASSERT_TRUE(annotated->Equals(same, true)); + ASSERT_FALSE(annotated->Equals(different, true)); + ASSERT_TRUE(plain->Equals( + dictionary(int8(), int32(), false, std::make_shared()), true)); + ASSERT_EQ(plain->ToString(), annotated->ToString()); + ASSERT_NE(plain->ToString(true), annotated->ToString(true)); + + // Equality and schema fingerprints must include metadata nested inside values. + auto nested = dictionary(int8(), struct_({field("x", annotated)})); + auto nested_different = dictionary(int8(), struct_({field("x", different)})); + ASSERT_TRUE(nested->Equals(nested_different)); + ASSERT_FALSE(nested->Equals(nested_different, true)); + ASSERT_FALSE(schema({field("d", nested)}) + ->Equals(*schema({field("d", nested_different)}), true)); + + auto options = Field::MergeOptions::Permissive(); + ASSERT_OK_AND_ASSIGN( + auto merged, field("d", annotated) + ->MergeWith(field("d", dictionary(int32(), int32())), options)); + AssertTypeEqual(*dictionary(int32(), int32(), false, metadata), *merged->type(), true); +} + +TEST(TestDictionaryType, ExtensionValueMetadata) { + auto metadata = key_value_metadata({{kExtensionTypeKeyName, "conflicting"}, + {kExtensionMetadataKeyName, "stale"}, + {"extra", "preserved"}}); + metadata->Append(kExtensionTypeKeyName, "duplicate"); + auto type = dictionary(int32(), uuid(), false, metadata); + const auto& dict = checked_cast(*type); + ASSERT_TRUE( + dict.value_metadata()->Equals(*key_value_metadata({{"extra", "preserved"}}))); + ASSERT_EQ(metadata->size(), 4); // Construction must not mutate its input. + ASSERT_FALSE(type->Equals(dictionary(int16(), uuid()))); + ASSERT_FALSE(type->Equals(dictionary(int32(), complex128()))); + ASSERT_TRUE(type->Equals(dictionary(int32(), uuid()))); + ASSERT_FALSE(type->Equals(dictionary(int32(), uuid()), true)); + ASSERT_TRUE(type->Equals( + dictionary(int32(), uuid(), false, key_value_metadata({{"extra", "preserved"}})), + true)); +} + TEST(TypesTest, SmallestDecimal) { for (int32_t i = 1; i < 76; ++i) { auto t = smallest_decimal(i, 4); 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/CDataInterface.rst b/docs/source/format/CDataInterface.rst index 0c71b0cd8b0b..82b9eefd0c48 100644 --- a/docs/source/format/CDataInterface.rst +++ b/docs/source/format/CDataInterface.rst @@ -537,6 +537,14 @@ For dictionary-encoded arrays, the :c:member:`ArrowSchema.format` string encodes the *index* type. The dictionary *value* type can be read from the :c:member:`ArrowSchema.dictionary` structure. +The :c:member:`ArrowSchema.metadata` of the :c:member:`ArrowSchema.dictionary` +structure holds the metadata of the dictionary value type, such as the +extension type of the dictionary values, while the parent structure's +:c:member:`ArrowSchema.metadata` applies to the dictionary-encoded type as a +whole. In the IPC format, these correspond to the ``custom_metadata`` of the +``DictionaryEncoding`` and of the ``Field``, respectively (see +:ref:`format_metadata_extension_types_dictionary`). + The same holds for :c:member:`ArrowArray` structure: while the parent structure points to the index data, the :c:member:`ArrowArray.dictionary` points to the dictionary values array. 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()