Skip to content

Commit b7f87d1

Browse files
committed
GH-50312: [Python] Fix UUID extension type round-trip to pandas returning bytes
1 parent 5f6ee02 commit b7f87d1

3 files changed

Lines changed: 49 additions & 0 deletions

File tree

python/pyarrow/pandas_compat.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -780,6 +780,13 @@ def _reconstruct_block(item, columns=None, extension_columns=None, return_block=
780780
raise ValueError("This column does not support to be converted "
781781
"to a pandas ExtensionArray")
782782
arr = pandas_dtype.__from_arrow__(arr)
783+
if isinstance(arr, np.ndarray) and arr.ndim == 1:
784+
# A plain (non-ExtensionArray) 1-D result — e.g. from the UUID
785+
# extension type — must be reshaped to the (1, n) single-column
786+
# block layout that make_block and create_dataframe_from_blocks
787+
# expect. pandas ExtensionArrays are 1-D and handled directly, so
788+
# they are intentionally left untouched here.
789+
arr = arr.reshape(1, -1)
783790
else:
784791
arr = block_arr
785792

python/pyarrow/tests/parquet/test_data_types.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -604,6 +604,37 @@ def test_uuid_extension_type():
604604
store_schema=False)
605605

606606

607+
@pytest.mark.pandas
608+
def test_uuid_roundtrip(tempdir):
609+
import uuid
610+
u1, u2 = uuid.uuid4(), uuid.uuid4()
611+
df = pd.DataFrame({"id": [u1, None, u2]})
612+
table = pa.Table.from_pandas(df)
613+
assert table.column("id").type == pa.uuid()
614+
615+
path = tempdir / "uuid_pandas_roundtrip.parquet"
616+
pq.write_table(table, path)
617+
read_table = pq.read_table(path)
618+
assert read_table.column("id").type == pa.uuid()
619+
620+
result_df = read_table.to_pandas()
621+
assert isinstance(result_df.loc[0, "id"], uuid.UUID)
622+
assert isinstance(result_df.loc[2, "id"], uuid.UUID)
623+
assert result_df.loc[0, "id"] == u1
624+
assert result_df.loc[2, "id"] == u2
625+
assert pd.isna(result_df.loc[1, "id"])
626+
627+
@pytest.mark.pandas
628+
def test_uuid_array_to_pandas():
629+
from uuid import uuid4
630+
import pandas as pd
631+
import pandas.testing as tm
632+
values = [uuid4(), None, uuid4()]
633+
arr = pa.array(values, type=pa.uuid())
634+
result = arr.to_pandas()
635+
expected = pd.Series(values, dtype=object)
636+
tm.assert_series_equal(result, expected)
637+
607638
def test_undefined_logical_type(parquet_test_datadir):
608639
test_file = f"{parquet_test_datadir}/unknown-logical-type.parquet"
609640

python/pyarrow/types.pxi

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1969,6 +1969,14 @@ cdef class JsonType(BaseExtensionType):
19691969
return JsonScalar
19701970

19711971

1972+
class _UuidPandasDtype:
1973+
def __from_arrow__(self, array):
1974+
# Return a 1-D object array of uuid.UUID values (nulls become None).
1975+
# Per pandas' __from_arrow__ contract this is 1-D; the table-to-blocks
1976+
# path reshapes it to the single-column block layout as needed.
1977+
return np.asarray(array.to_pylist(), dtype=object)
1978+
1979+
19721980
cdef class UuidType(BaseExtensionType):
19731981
"""
19741982
Concrete class for UUID extension type.
@@ -1987,6 +1995,9 @@ cdef class UuidType(BaseExtensionType):
19871995
def __arrow_ext_scalar_class__(self):
19881996
return UuidScalar
19891997

1998+
def to_pandas_dtype(self):
1999+
return _UuidPandasDtype()
2000+
19902001

19912002
cdef class FixedShapeTensorType(BaseExtensionType):
19922003
"""

0 commit comments

Comments
 (0)