GH-49704: [C++] Read and write DictionaryEncoding.custom_metadata for dictionary value types

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<SchemaExporter>();
-      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

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<Field>& input) -> Result<std::shared_ptr<Field>> {
+    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<const DictionaryType&>(
+      *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();

cpp/src/arrow/compare.cc

@@ -875,7 +875,7 @@ class TypeEqualsVisitor {
   Status Visit(const DictionaryType& left) {
     const auto& right = checked_cast<const DictionaryType&>(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();
   }

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>& 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<RecordBatch> 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<std::pair<std::string, std::string>>;
+  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<RecordBatch>& input)
+      -> Result<std::shared_ptr<RecordBatch>> {
+    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<Buffer>(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<const DictionaryType&>(
+      *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<RecordBatch> batch;
   ASSERT_OK((*GetParam())(&batch));  // NOLINT clang-tidy gtest issue