GH-50312: [Python] Fix UUID extension type round-trip to pandas returning bytes#50325
GH-50312: [Python] Fix UUID extension type round-trip to pandas returning bytes#50325parker-cassar wants to merge 2 commits into
Conversation
|
|
parker-cassar
force-pushed
the
gh-50312-uuid-pandas-roundtrip
branch
from
5b15a46 to
386b1b6
Compare
|
Hey @parker-cassar, thanks so much for jumping on this so quickly. Just out of curiosity—since this delegates to to_pylist(), do you think this pattern might end up being useful for other Arrow extension types down the way??? |
rok
left a comment
rok
left a comment
There was a problem hiding this comment.
Thanks for working on this @parker-cassar!
A couple of questions since we're aiming to support Pandas 3.x and Pandas 2.x too for a while.
| assert result_df.loc[2, "id"] == u2 | ||
| assert pd.isna(result_df.loc[1, "id"]) | ||
|
|
||
|
|
There was a problem hiding this comment.
Would this test pass with the current proposal?
| @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) |
|
|
||
| class _UuidPandasDtype: | ||
| def __from_arrow__(self, array): | ||
| return np.asarray(array.to_pylist(), dtype=object).reshape(1, -1) |
There was a problem hiding this comment.
Wouldn't this make UuidType.__from_arrow__ return a 2-D array to satisfy old pandas DataFrame block construction, but the same hook is also used by Array.to_pandas(), which expects 1-D data for pd.Series?
Perhaps a better fix is to keep __from_arrow__ returning normal 1-D object data, and handle the old pandas block-manager reshape only in the DataFrame reconstruction path where 2-D block shape is
actually needed?
Something like:
diff --git a/python/pyarrow/types.pxi b/python/pyarrow/types.pxi
--- a/python/pyarrow/types.pxi
+++ b/python/pyarrow/types.pxi
@@ -1971,7 +1971,7 @@ cdef class JsonType(BaseExtensionType):
class _UuidPandasDtype:
def __from_arrow__(self, array):
- return np.asarray(array.to_pylist(), dtype=object).reshape(1, -1)
+ # 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):
diff --git a/python/pyarrow/pandas_compat.py b/python/pyarrow/pandas_compat.py
--- a/python/pyarrow/pandas_compat.py
+++ b/python/pyarrow/pandas_compat.py
@@ -779,6 +779,13 @@ def _reconstruct_block(item, columns=None, extension_columns=None,
if not hasattr(pandas_dtype, '__from_arrow__'):
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_arr3 participants
Rationale for this change
Converting a Table with an
arrow.uuidextension column to pandas currently produces a column ofbytesinstead ofuuid.UUIDobjects. This happens becauseUuidTypedoes not implementto_pandas_dtype(), soTable.to_pandas()falls back to the storage type (fixed_size_binary(16)) and producesbytes. The bug occurs even without a Parquet roundtrip.Note: the original issue suggested this might be specific to Python 3.14 but I tested on Python versions 3.10 - 3.14 and still had the issue since
UuidTypehas never implementedto_pandas_dtype().What changes are included in this PR?
Added
UuidType.to_pandas_dtype(): returns a dtype wrapper implementing__from_arrow__, which delegates toto_pylist()sinceUuidScalar.as_py()already producesuuid.UUIDobjects.Are these changes tested?
Yes. Added
test_uuid_roundtripwhich covers pandas DataFrame with a UUID column -> pyarrow Table -> Parquet on disk -> pyarrow Table -> pandas DataFrame. The final conversion is what this PR fixes.Are there any user-facing changes?
Yes.
Table.to_pandas()now returnsuuid.UUIDforarrow.uuidcolumns instead ofbytes.