From 457ff20652169bf73e4be1adf323fdd351bb5f74 Mon Sep 17 00:00:00 2001 From: parker-cassar Date: Tue, 7 Jul 2026 20:11:34 -0700 Subject: [PATCH] GH-50312: [Python] Fix UUID extension type round-trip to pandas returning bytes --- python/pyarrow/pandas_compat.py | 7 ++++ .../pyarrow/tests/parquet/test_data_types.py | 33 +++++++++++++++++++ python/pyarrow/types.pxi | 11 +++++++ 3 files changed, 51 insertions(+) diff --git a/python/pyarrow/pandas_compat.py b/python/pyarrow/pandas_compat.py index d8fd383d31d1..999c92e75c1f 100644 --- a/python/pyarrow/pandas_compat.py +++ b/python/pyarrow/pandas_compat.py @@ -780,6 +780,13 @@ def _reconstruct_block(item, columns=None, extension_columns=None, return_block= raise ValueError("This column does not support to be converted " "to a pandas ExtensionArray") arr = pandas_dtype.__from_arrow__(arr) + if isinstance(arr, np.ndarray) and arr.ndim == 1: + # A plain (non-ExtensionArray) 1-D result — e.g. from the UUID + # extension type — must be reshaped to the (1, n) single-column + # block layout that make_block and create_dataframe_from_blocks + # expect. pandas ExtensionArrays are 1-D and handled directly, so + # they are intentionally left untouched here. + arr = arr.reshape(1, -1) else: arr = block_arr diff --git a/python/pyarrow/tests/parquet/test_data_types.py b/python/pyarrow/tests/parquet/test_data_types.py index c546bc1532ac..7bd8238a3491 100644 --- a/python/pyarrow/tests/parquet/test_data_types.py +++ b/python/pyarrow/tests/parquet/test_data_types.py @@ -604,6 +604,39 @@ 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"]) + + +@pytest.mark.pandas +def test_uuid_array_to_pandas(): + from uuid import uuid4 + import pandas as pd + import pandas.testing as tm + values = [uuid4(), None, uuid4()] + arr = pa.array(values, type=pa.uuid()) + result = arr.to_pandas() + expected = pd.Series(values, dtype=object) + tm.assert_series_equal(result, expected) + + def test_undefined_logical_type(parquet_test_datadir): test_file = f"{parquet_test_datadir}/unknown-logical-type.parquet" diff --git a/python/pyarrow/types.pxi b/python/pyarrow/types.pxi index ec1a5a2ba9a3..7ce186972bcf 100644 --- a/python/pyarrow/types.pxi +++ b/python/pyarrow/types.pxi @@ -1969,6 +1969,14 @@ cdef class JsonType(BaseExtensionType): return JsonScalar +class _UuidPandasDtype: + def __from_arrow__(self, array): + # Return a 1-D object array of uuid.UUID values (nulls become None). + # Per pandas' __from_arrow__ contract this is 1-D; the table-to-blocks + # path reshapes it to the single-column block layout as needed. + return np.asarray(array.to_pylist(), dtype=object) + + cdef class UuidType(BaseExtensionType): """ Concrete class for UUID extension type. @@ -1987,6 +1995,9 @@ cdef class UuidType(BaseExtensionType): def __arrow_ext_scalar_class__(self): return UuidScalar + def to_pandas_dtype(self): + return _UuidPandasDtype() + cdef class FixedShapeTensorType(BaseExtensionType): """