GH-43807: [C++][Python] Add UUID extension type conversion support to/from Parquet

cpp/src/arrow/dataset/file_parquet.cc

@@ -132,6 +132,8 @@ parquet::ArrowReaderProperties MakeArrowReaderProperties(
   arrow_properties.set_io_context(
       parquet_scan_options.arrow_reader_properties->io_context());
   arrow_properties.set_use_threads(options.use_threads);
+  arrow_properties.set_arrow_extensions_enabled(
+      parquet_scan_options.arrow_reader_properties->get_arrow_extensions_enabled());
   return arrow_properties;
 }
 

cpp/src/arrow/extension/uuid.cc

@@ -40,7 +40,7 @@ Result<std::shared_ptr<DataType>> UuidType::Deserialize(
   if (!serialized.empty()) {
     return Status::Invalid("Unexpected serialized metadata: '", serialized, "'");
   }
-  if (!storage_type->Equals(*fixed_size_binary(16))) {
+  if (!IsSupportedStorageType(storage_type)) {
     return Status::Invalid("Invalid storage type for UuidType: ",
                            storage_type->ToString());
   }
@@ -55,4 +55,8 @@ std::string UuidType::ToString(bool show_metadata) const {
 
 std::shared_ptr<DataType> uuid() { return std::make_shared<UuidType>(); }
 
+bool UuidType::IsSupportedStorageType(const std::shared_ptr<DataType>& storage_type) {
+  return storage_type->Equals(*fixed_size_binary(16));
+}
+
 }  // namespace arrow::extension

cpp/src/arrow/extension/uuid.h

@@ -53,6 +53,8 @@ class ARROW_EXPORT UuidType : public ExtensionType {
 
   /// \brief Create a UuidType instance
   static Result<std::shared_ptr<DataType>> Make() { return std::make_shared<UuidType>(); }
+
+  static bool IsSupportedStorageType(const std::shared_ptr<DataType>& storage_type);
 };
 
 /// \brief Return a UuidType instance.

cpp/src/parquet/arrow/arrow_schema_test.cc

@@ -33,6 +33,7 @@
 
 #include "arrow/array.h"
 #include "arrow/extension/json.h"
+#include "arrow/extension/uuid.h"
 #include "arrow/ipc/writer.h"
 #include "arrow/testing/gtest_util.h"
 #include "arrow/type.h"
@@ -945,7 +946,7 @@ TEST_F(TestConvertParquetSchema, ParquetVariant) {
   }
 }
 
-TEST_F(TestConvertParquetSchema, ParquetSchemaArrowExtensions) {
+TEST_F(TestConvertParquetSchema, ParquetSchemaArrowJsonExtension) {
   std::vector<NodePtr> parquet_fields;
   parquet_fields.push_back(PrimitiveNode::Make(
       "json_1", Repetition::OPTIONAL, ParquetType::BYTE_ARRAY, ConvertedType::JSON));
@@ -1027,6 +1028,68 @@ TEST_F(TestConvertParquetSchema, ParquetSchemaArrowExtensions) {
   }
 }
 
+TEST_F(TestConvertParquetSchema, ParquetSchemaArrowUuidExtension) {
+  std::vector<NodePtr> parquet_fields;
+  parquet_fields.push_back(PrimitiveNode::Make("uuid", Repetition::OPTIONAL,
+                                               LogicalType::UUID(),
+                                               ParquetType::FIXED_LEN_BYTE_ARRAY, 16));
+
+  {
+    // Parquet file does not contain Arrow schema.
+    // By default, field should be treated as fixed_size_binary(16) in Arrow.
+    auto arrow_schema =
+        ::arrow::schema({::arrow::field("uuid", ::arrow::fixed_size_binary(16))});
+    std::shared_ptr<KeyValueMetadata> metadata{};
+    ASSERT_OK(ConvertSchema(parquet_fields, metadata));
+    CheckFlatSchema(arrow_schema);
+  }
+
+  {
+    // Parquet file does not contain Arrow schema.
+    // If Arrow extensions are enabled, field will be interpreted as uuid()
+    // extension field.
+    ArrowReaderProperties props;
+    props.set_arrow_extensions_enabled(true);
+    auto arrow_schema =
+        ::arrow::schema({::arrow::field("uuid", ::arrow::extension::uuid())});
+    std::shared_ptr<KeyValueMetadata> metadata{};
+    ASSERT_OK(ConvertSchema(parquet_fields, metadata, props));
+    CheckFlatSchema(arrow_schema);
+  }
+
+  {
+    // Parquet file contains Arrow schema.
+    // uuid will be interpreted as uuid() field even though extensions are not enabled.
+    ArrowReaderProperties props;
+    props.set_arrow_extensions_enabled(false);
+    std::shared_ptr<KeyValueMetadata> field_metadata =
+        ::arrow::key_value_metadata({"foo", "bar"}, {"biz", "baz"});
+    auto arrow_schema = ::arrow::schema({::arrow::field(
+        "uuid", ::arrow::extension::uuid(), /*nullable=*/true, field_metadata)});
+
+    std::shared_ptr<KeyValueMetadata> metadata;
+    ASSERT_OK(ArrowSchemaToParquetMetadata(arrow_schema, metadata));
+    ASSERT_OK(ConvertSchema(parquet_fields, metadata, props));
+    CheckFlatSchema(arrow_schema, true /* check_metadata */);
+  }
+
+  {
+    // Parquet file contains Arrow schema.
+    // uuid will be interpreted as uuid() field also when extensions *are* enabled
+    ArrowReaderProperties props;
+    props.set_arrow_extensions_enabled(true);
+    std::shared_ptr<KeyValueMetadata> field_metadata =
+        ::arrow::key_value_metadata({"foo", "bar"}, {"biz", "baz"});
+    auto arrow_schema = ::arrow::schema({::arrow::field(
+        "uuid", ::arrow::extension::uuid(), /*nullable=*/true, field_metadata)});
+
+    std::shared_ptr<KeyValueMetadata> metadata;
+    ASSERT_OK(ArrowSchemaToParquetMetadata(arrow_schema, metadata));
+    ASSERT_OK(ConvertSchema(parquet_fields, metadata, props));
+    CheckFlatSchema(arrow_schema, true /* check_metadata */);
+  }
+}
+
 class TestConvertArrowSchema : public ::testing::Test {
  public:
   virtual void SetUp() {}

cpp/src/parquet/arrow/schema.cc

@@ -22,6 +22,7 @@
 #include <vector>
 
 #include "arrow/extension/json.h"
+#include "arrow/extension/uuid.h"
 #include "arrow/extension_type.h"
 #include "arrow/io/memory.h"
 #include "arrow/ipc/api.h"
@@ -454,12 +455,18 @@ Status FieldToNode(const std::string& name, const std::shared_ptr<Field>& field,
         type = ParquetType::BYTE_ARRAY;
         logical_type = LogicalType::JSON();
         break;
+      } else if (ext_type->extension_name() == std::string("arrow.uuid")) {
+        type = ParquetType::FIXED_LEN_BYTE_ARRAY;
+        logical_type = LogicalType::UUID();
+        length = 16;
+        break;
       } else if (ext_type->extension_name() == std::string("parquet.variant")) {
         auto variant_type = std::static_pointer_cast<VariantExtensionType>(field->type());
 
         return VariantToNode(variant_type, name, field->nullable(), field_id, properties,
                              arrow_properties, out);
       }
+
       std::shared_ptr<::arrow::Field> storage_field = ::arrow::field(
           name, ext_type->storage_type(), field->nullable(), field->metadata());
       return FieldToNode(name, storage_field, properties, arrow_properties, out);
@@ -1052,60 +1059,62 @@ Result<bool> ApplyOriginalMetadata(const Field& origin_field, SchemaField* infer
   bool modified = false;
 
   auto& origin_type = origin_field.type();
-  const auto& inferred_type = inferred->field->type();
 
+  // The origin was an extension type. This occurs when the ARROW:extension:name field
+  // was present when the schema was written and that extension is registered when
+  // the schema is read.
   if (origin_type->id() == ::arrow::Type::EXTENSION) {
-    const auto& ex_type = checked_cast<const ::arrow::ExtensionType&>(*origin_type);
-    if (inferred_type->id() != ::arrow::Type::EXTENSION &&
-        ex_type.extension_name() == std::string("arrow.json") &&
-        ::arrow::extension::JsonExtensionType::IsSupportedStorageType(
-            inferred_type->id())) {
-      // Schema mismatch.
-      //
-      // Arrow extensions are DISABLED in Parquet.
-      // origin_type is ::arrow::extension::json()
-      // inferred_type is ::arrow::utf8()
-      //
-      // Origin type is restored as Arrow should be considered the source of truth.
-      inferred->field = inferred->field->WithType(origin_type);
-      RETURN_NOT_OK(ApplyOriginalStorageMetadata(origin_field, inferred));
-    } else if (inferred_type->id() == ::arrow::Type::EXTENSION &&
-               ex_type.extension_name() == std::string("arrow.json")) {
-      // Potential schema mismatch.
-      //
-      // Arrow extensions are ENABLED in Parquet.
-      // origin_type is arrow::extension::json(...)
-      // inferred_type is arrow::extension::json(arrow::utf8())
-      auto origin_storage_field = origin_field.WithType(ex_type.storage_type());
-
-      // Apply metadata recursively to storage type
-      RETURN_NOT_OK(ApplyOriginalStorageMetadata(*origin_storage_field, inferred));
-      inferred->field = inferred->field->WithType(origin_type);
-    } else if (inferred_type->id() == ::arrow::Type::EXTENSION &&
-               ex_type.extension_name() == std::string("parquet.variant")) {
-      // Potential schema mismatch.
-      //
-      // Arrow extensions are ENABLED in Parquet.
-      // origin_type is parquet::arrow::variant(...)
-      // inferred_type is
-      // parquet::arrow::variant(struct(arrow::binary(),arrow::binary()))
-      auto origin_storage_field = origin_field.WithType(ex_type.storage_type());
-
-      // Apply metadata recursively to storage type
-      RETURN_NOT_OK(ApplyOriginalStorageMetadata(*origin_storage_field, inferred));
-      inferred->field = inferred->field->WithType(origin_type);
+    const auto& origin_extension_type =
+        checked_cast<const ::arrow::ExtensionType&>(*origin_type);
+
+    // (Recursively) Apply the original storage metadata from the original storage field
+    // This applies extension types to child elements, if any.
+    auto origin_storage_field =
+        origin_field.WithType(origin_extension_type.storage_type());
+    RETURN_NOT_OK(ApplyOriginalStorageMetadata(*origin_storage_field, inferred));
+
+    // Use the inferred type after child updates for below checks to see if
+    // we can restore an extension type on the output.
+    const auto& inferred_type = inferred->field->type();
+
+    // Whether or not the inferred type is also an extension type. This can occur when
+    // arrow_extensions_enabled is true in the ArrowReaderProperties. Extension types
+    // are not currently inferred for any other reason.
+    bool arrow_extension_inferred = inferred_type->id() == ::arrow::Type::EXTENSION;
+
+    // Check if the inferred storage type is compatible with the extension type
+    // we're hoping to apply. We assume that if an extension type was inferred
+    // that it was constructed with a valid storage type. Otherwise, we check with
+    // extension types that we know about for valid storage, falling back to
+    // storage type equality for extension types that we don't know about.
+    std::string origin_extension_name = origin_extension_type.extension_name();
+    bool extension_supports_inferred_storage;
+
+    if (origin_extension_name == "arrow.json") {
+      extension_supports_inferred_storage =
+          arrow_extension_inferred ||
+          ::arrow::extension::JsonExtensionType::IsSupportedStorageType(
+              inferred_type->id());
+    } else if (origin_extension_name == "arrow.uuid") {
+      extension_supports_inferred_storage =
+          arrow_extension_inferred ||
+          ::arrow::extension::UuidType::IsSupportedStorageType(inferred_type);
+    } else if (origin_extension_name == "parquet.variant") {
+      extension_supports_inferred_storage =
+          arrow_extension_inferred ||
+          VariantExtensionType::IsSupportedStorageType(inferred_type);
     } else {
-      auto origin_storage_field = origin_field.WithType(ex_type.storage_type());
-
-      // Apply metadata recursively to storage type
-      RETURN_NOT_OK(ApplyOriginalStorageMetadata(*origin_storage_field, inferred));
+      extension_supports_inferred_storage =
+          origin_extension_type.storage_type()->Equals(*inferred_type);
+    }
 
-      // Restore extension type, if the storage type is the same as inferred
-      // from the Parquet type
-      if (ex_type.storage_type()->Equals(*inferred->field->type())) {
-        inferred->field = inferred->field->WithType(origin_type);
-      }
+    // If the origin extension of the metadata we are about to apply supports
+    // the Arrow storage type we would otherwise return, we restore the extension
+    // type to the output.
+    if (extension_supports_inferred_storage) {
+      inferred->field = inferred->field->WithType(origin_type);
     }
+
     modified = true;
   } else {
     ARROW_ASSIGN_OR_RAISE(modified, ApplyOriginalStorageMetadata(origin_field, inferred));

cpp/src/parquet/arrow/schema_internal.cc

@@ -18,6 +18,7 @@
 #include "parquet/arrow/schema_internal.h"
 
 #include "arrow/extension/json.h"
+#include "arrow/extension/uuid.h"
 #include "arrow/type.h"
 
 #include "parquet/properties.h"
@@ -134,20 +135,25 @@ Result<std::shared_ptr<ArrowType>> FromByteArray(
   }
 }
 
-Result<std::shared_ptr<ArrowType>> FromFLBA(const LogicalType& logical_type,
-                                            int32_t physical_length) {
+Result<std::shared_ptr<ArrowType>> FromFLBA(
+    const LogicalType& logical_type, int32_t physical_length,
+    const ArrowReaderProperties& reader_properties) {
   switch (logical_type.type()) {
     case LogicalType::Type::DECIMAL:
       return MakeArrowDecimal(logical_type);
     case LogicalType::Type::FLOAT16:
       return ::arrow::float16();
     case LogicalType::Type::NONE:
     case LogicalType::Type::INTERVAL:
+      return ::arrow::fixed_size_binary(physical_length);
     case LogicalType::Type::UUID:
+      if (physical_length == 16 && reader_properties.get_arrow_extensions_enabled()) {
+        return ::arrow::extension::uuid();
+      }
+
       return ::arrow::fixed_size_binary(physical_length);
     default:
-      return Status::NotImplemented("Unhandled logical logical_type ",
-                                    logical_type.ToString(),
+      return Status::NotImplemented("Unhandled logical_type ", logical_type.ToString(),
                                     " for fixed-length binary array");
   }
 }
@@ -216,7 +222,7 @@ Result<std::shared_ptr<ArrowType>> GetArrowType(
     case ParquetType::BYTE_ARRAY:
       return FromByteArray(logical_type, reader_properties);
     case ParquetType::FIXED_LEN_BYTE_ARRAY:
-      return FromFLBA(logical_type, type_length);
+      return FromFLBA(logical_type, type_length, reader_properties);
     default: {
       // PARQUET-1565: This can occur if the file is corrupt
       return Status::IOError("Invalid physical column type: ",

python/pyarrow/_dataset_parquet.pyx

@@ -703,7 +703,7 @@ cdef class ParquetFragmentScanOptions(FragmentScanOptions):
     cache_options : pyarrow.CacheOptions, default None
         Cache options used when pre_buffer is enabled. The default values should
         be good for most use cases. You may want to adjust these for example if
-        you have exceptionally high latency to the file system. 
+        you have exceptionally high latency to the file system.
     thrift_string_size_limit : int, default None
         If not None, override the maximum total string size allocated
         when decoding Thrift structures. The default limit should be
@@ -720,6 +720,10 @@ cdef class ParquetFragmentScanOptions(FragmentScanOptions):
         Parquet file.
     page_checksum_verification : bool, default False
         If True, verify the page checksum for each page read from the file.
+    arrow_extensions_enabled : bool, default False
+        If True, read Parquet logical types as Arrow extension types where possible,
+        (e.g., read JSON as the canonical `arrow.json` extension type or UUID as
+        the canonical `arrow.uuid` extension type).
     """
 
     # Avoid mistakingly creating attributes
@@ -733,7 +737,8 @@ cdef class ParquetFragmentScanOptions(FragmentScanOptions):
                  thrift_container_size_limit=None,
                  decryption_config=None,
                  decryption_properties=None,
-                 bint page_checksum_verification=False):
+                 bint page_checksum_verification=False,
+                 bint arrow_extensions_enabled=False):
         self.init(shared_ptr[CFragmentScanOptions](
             new CParquetFragmentScanOptions()))
         self.use_buffered_stream = use_buffered_stream
@@ -752,6 +757,7 @@ cdef class ParquetFragmentScanOptions(FragmentScanOptions):
         if decryption_properties is not None:
             self.decryption_properties = decryption_properties
         self.page_checksum_verification = page_checksum_verification
+        self.arrow_extensions_enabled = arrow_extensions_enabled
 
     cdef void init(self, const shared_ptr[CFragmentScanOptions]& sp):
         FragmentScanOptions.init(self, sp)
@@ -868,6 +874,14 @@ cdef class ParquetFragmentScanOptions(FragmentScanOptions):
     def page_checksum_verification(self, bint page_checksum_verification):
         self.reader_properties().set_page_checksum_verification(page_checksum_verification)
 
+    @property
+    def arrow_extensions_enabled(self):
+        return self.arrow_reader_properties().get_arrow_extensions_enabled()
+
+    @arrow_extensions_enabled.setter
+    def arrow_extensions_enabled(self, bint arrow_extensions_enabled):
+        self.arrow_reader_properties().set_arrow_extensions_enabled(arrow_extensions_enabled)
+
     def equals(self, ParquetFragmentScanOptions other):
         """
         Parameters
@@ -881,11 +895,12 @@ cdef class ParquetFragmentScanOptions(FragmentScanOptions):
         attrs = (
             self.use_buffered_stream, self.buffer_size, self.pre_buffer, self.cache_options,
             self.thrift_string_size_limit, self.thrift_container_size_limit,
-            self.page_checksum_verification)
+            self.page_checksum_verification, self.arrow_extensions_enabled)
         other_attrs = (
             other.use_buffered_stream, other.buffer_size, other.pre_buffer, other.cache_options,
             other.thrift_string_size_limit,
-            other.thrift_container_size_limit, other.page_checksum_verification)
+            other.thrift_container_size_limit, other.page_checksum_verification,
+            other.arrow_extensions_enabled)
         return attrs == other_attrs
 
     @staticmethod
@@ -902,7 +917,8 @@ cdef class ParquetFragmentScanOptions(FragmentScanOptions):
             cache_options=self.cache_options,
             thrift_string_size_limit=self.thrift_string_size_limit,
             thrift_container_size_limit=self.thrift_container_size_limit,
-            page_checksum_verification=self.page_checksum_verification
+            page_checksum_verification=self.page_checksum_verification,
+            arrow_extensions_enabled=self.arrow_extensions_enabled
         )
         return ParquetFragmentScanOptions._reconstruct, (kwargs,)
 

python/pyarrow/_parquet.pxd

@@ -404,6 +404,8 @@ cdef extern from "parquet/api/reader.h" namespace "parquet" nogil:
         CCacheOptions cache_options() const
         void set_coerce_int96_timestamp_unit(TimeUnit unit)
         TimeUnit coerce_int96_timestamp_unit() const
+        void set_arrow_extensions_enabled(c_bool extensions_enabled)
+        c_bool get_arrow_extensions_enabled() const
 
     ArrowReaderProperties default_arrow_reader_properties()