GH-50312: [Python] Fix UUID extension type round-trip to pandas returning bytes

python/pyarrow/src/arrow/python/arrow_to_pandas.cc

@@ -116,6 +116,17 @@ void BufferCapsule_Destructor(PyObject* capsule) {
 using internal::arrow_traits;
 using internal::npy_traits;
 
+bool IsUuidExtension(const DataType& type) {
+  if (type.id() != Type::EXTENSION) {
+    return false;
+  }
+  const auto& extension_type = checked_cast<const ExtensionType&>(type);
+  const auto& storage_type = *extension_type.storage_type();
+  return extension_type.extension_name() == "arrow.uuid" &&
+         storage_type.id() == Type::FIXED_SIZE_BINARY &&
+         checked_cast<const FixedSizeBinaryType&>(storage_type).byte_width() == 16;
+}
+
 template <typename T>
 struct WrapBytes {};
 
@@ -1378,6 +1389,26 @@ struct ObjectWriterVisitor {
     return ConvertStruct(options, data, out_values);
   }
 
+  Status Visit(const ExtensionType& type) {
+    if (!IsUuidExtension(type)) {
+      return Status::NotImplemented("No implemented conversion to object dtype: ",
+                                    type.ToString());
+    }
+    ArrayVector storage_arrays;
+    storage_arrays.reserve(data.num_chunks());
+    for (int c = 0; c < data.num_chunks(); ++c) {
+      const auto& extension_array = checked_cast<const ExtensionArray&>(*data.chunk(c));
+      storage_arrays.push_back(extension_array.storage());
+    }
+    ChunkedArray storage(std::move(storage_arrays), type.storage_type());
+    auto WrapUuid = [&](const std::string_view& view, PyObject** out) {
+      ARROW_ASSIGN_OR_RAISE(*out, internal::UuidFromBytes(view));
+      return Status::OK();
+    };
+    return ConvertAsPyObjects<FixedSizeBinaryType>(options, storage, WrapUuid,
+                                                   out_values);
+  }
+
   template <typename Type>
   enable_if_t<is_floating_type<Type>::value ||
                   std::is_same<DictionaryType, Type>::value ||
@@ -2253,7 +2284,10 @@ static Status GetPandasWriterType(const ChunkedArray& data, const PandasOptions&
       *output_type = PandasWriter::CATEGORICAL;
       break;
     case Type::EXTENSION:
-      *output_type = PandasWriter::EXTENSION;
+      // UUID has a native object conversion to uuid.UUID. Other extension
+      // types continue through the pandas ExtensionArray protocol.
+      *output_type =
+          IsUuidExtension(*data.type()) ? PandasWriter::OBJECT : PandasWriter::EXTENSION;
       break;
     default:
       return Status::NotImplemented(
@@ -2358,8 +2392,10 @@ class ConsolidatedBlockCreator : public PandasBlockCreator {
       *out = PandasWriter::EXTENSION;
       return Status::OK();
     } else {
-      // In case of an extension array default to the storage type
-      if (arrays_[column_index]->type()->id() == Type::EXTENSION) {
+      // In case of an extension array default to the storage type, except for
+      // UUID which has a native Python object conversion.
+      if (arrays_[column_index]->type()->id() == Type::EXTENSION &&
+          !IsUuidExtension(*arrays_[column_index]->type())) {
         arrays_[column_index] = GetStorageChunkedArray(arrays_[column_index]);
       }
       // In case of a RunEndEncodedArray default to the values type
@@ -2599,8 +2635,9 @@ Status ConvertChunkedArrayToPandas(const PandasOptions& options,
   // Table->DataFrame
   modified_options.allow_zero_copy_blocks = true;
 
-  // In case of an extension array default to the storage type
-  if (arr->type()->id() == Type::EXTENSION) {
+  // In case of an extension array default to the storage type, except for UUID
+  // which has a native Python object conversion.
+  if (arr->type()->id() == Type::EXTENSION && !IsUuidExtension(*arr->type())) {
     arr = GetStorageChunkedArray(arr);
   }
   // In case of a RunEndEncodedArray decode the array

python/pyarrow/src/arrow/python/helpers.cc

@@ -338,24 +338,48 @@ struct ModuleOnceRunner {
 static PyObject* uuid_UUID = nullptr;
 static ModuleOnceRunner uuid_runner("uuid");
 
-}  // namespace
-
-bool IsPyUuid(PyObject* obj) {
+PyObject* GetUuidClass() {
   uuid_runner.RunOnce([](OwnedRef& module) {
     OwnedRef ref;
     if (ImportFromModule(module.obj(), "UUID", &ref).ok()) {
       uuid_UUID = ref.obj();
     }
   });
-  if (!uuid_UUID) return false;
-  int result = PyObject_IsInstance(obj, uuid_UUID);
+  return uuid_UUID;
+}
+
+}  // namespace
+
+bool IsPyUuid(PyObject* obj) {
+  PyObject* uuid_class = GetUuidClass();
+  if (!uuid_class) return false;
+  int result = PyObject_IsInstance(obj, uuid_class);
   if (result < 0) {
     PyErr_Clear();
     return false;
   }
   return result != 0;
 }
 
+Result<PyObject*> UuidFromBytes(std::string_view bytes) {
+  PyObject* uuid_class = GetUuidClass();
+  if (!uuid_class) {
+    return Status::Invalid("Could not import uuid.UUID");
+  }
+  OwnedRef py_bytes(
+      PyBytes_FromStringAndSize(bytes.data(), static_cast<Py_ssize_t>(bytes.size())));
+  RETURN_IF_PYERROR();
+  OwnedRef kwargs(PyDict_New());
+  RETURN_IF_PYERROR();
+  if (PyDict_SetItemString(kwargs.obj(), "bytes", py_bytes.obj()) < 0) {
+    RETURN_IF_PYERROR();
+  }
+  OwnedRef args(PyTuple_New(0));
+  PyObject* result = PyObject_Call(uuid_class, args.obj(), kwargs.obj());
+  RETURN_IF_PYERROR();
+  return result;
+}
+
 namespace {
 
 // Once initialized, these variables hold borrowed references to Pandas static data.

python/pyarrow/src/arrow/python/helpers.h

@@ -22,6 +22,7 @@
 #include <limits>
 #include <memory>
 #include <string>
+#include <string_view>
 #include <utility>
 
 #include "arrow/python/numpy_interop.h"
@@ -96,6 +97,10 @@ bool PyFloat_IsNaN(PyObject* obj);
 ARROW_PYTHON_EXPORT
 bool IsPyUuid(PyObject* obj);
 
+// \brief Construct a uuid.UUID from 16 raw bytes
+ARROW_PYTHON_EXPORT
+Result<PyObject*> UuidFromBytes(std::string_view bytes);
+
 inline bool IsPyBinary(PyObject* obj) {
   return PyBytes_Check(obj) || PyByteArray_Check(obj) || PyMemoryView_Check(obj);
 }

python/pyarrow/tests/parquet/test_data_types.py

@@ -604,6 +604,27 @@ def test_uuid_extension_type():
         store_schema=False)
 
 
+@pytest.mark.pandas
+def test_uuid_roundtrip(tempdir):
+    import uuid
+    u1, u2 = uuid.uuid4(), uuid.uuid4()
+    df = pd.DataFrame({"id": [u1, None, u2]})
+    table = pa.Table.from_pandas(df)
+    assert table.column("id").type == pa.uuid()
+
+    path = tempdir / "uuid_pandas_roundtrip.parquet"
+    pq.write_table(table, path)
+    read_table = pq.read_table(path)
+    assert read_table.column("id").type == pa.uuid()
+
+    result_df = read_table.to_pandas()
+    assert isinstance(result_df.loc[0, "id"], uuid.UUID)
+    assert isinstance(result_df.loc[2, "id"], uuid.UUID)
+    assert result_df.loc[0, "id"] == u1
+    assert result_df.loc[2, "id"] == u2
+    assert pd.isna(result_df.loc[1, "id"])
+
+
 def test_undefined_logical_type(parquet_test_datadir):
     test_file = f"{parquet_test_datadir}/unknown-logical-type.parquet"
 

python/pyarrow/tests/test_extension_type.py

@@ -1400,6 +1400,40 @@ def test_uuid_extension():
     assert isinstance(array[0], pa.UuidScalar)
 
 
+@pytest.mark.pandas
+def test_uuid_to_pandas():
+    import pandas as pd
+    import pandas.testing as tm
+    values = [uuid4(), None, uuid4()]
+    array = pa.array(values, type=pa.uuid())
+    chunked_array = pa.chunked_array([array.slice(0, 1), array.slice(1)])
+    expected = pd.Series(values, dtype=object)
+    tm.assert_series_equal(array.to_pandas(), expected)
+    tm.assert_series_equal(chunked_array.to_pandas(), expected)
+    tm.assert_frame_equal(
+        pa.table({"uuid": chunked_array}).to_pandas(),
+        expected.to_frame(name="uuid"),
+    )
+
+
+@pytest.mark.pandas
+def test_uuid_to_pandas_options():
+    values = [uuid4(), uuid4()] * 2
+    array = pa.array(values, type=pa.uuid())
+    chunked_array = pa.chunked_array([array.slice(0, 2), array.slice(2)])
+    for obj in [array, chunked_array, pa.table({"uuid": chunked_array})]:
+        with pytest.raises(pa.ArrowInvalid):
+            obj.to_pandas(zero_copy_only=True)
+        result = obj.to_pandas()
+        if result.ndim == 2:
+            result = result["uuid"]
+        assert len({id(value) for value in result}) == 2
+        result = obj.to_pandas(deduplicate_objects=False)
+        if result.ndim == 2:
+            result = result["uuid"]
+        assert len({id(value) for value in result}) == 4
+
+
 def test_uuid_scalar_from_python():
     # Test with explicit type
     py_uuid = uuid4()