Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions python/pyarrow/pandas_compat.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
33 changes: 33 additions & 0 deletions python/pyarrow/tests/parquet/test_data_types.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would this test pass with the current proposal?

Suggested change
@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)

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No, this did not pass with the 2-D return. Applied your approach, from_arrow returns 1-D now and the reshape moved to _reconstruct_block. test passes locally.

Note: I squashed my 4 iterative commits into one during a rebase (3rd commit was just resyncing main). Then amended that squashed commit with a small pep 8 fix. The changes since your review are your two suggestions. Sorry for the noise, will avoid the mid-review force-pushes going forward.


def test_undefined_logical_type(parquet_test_datadir):
test_file = f"{parquet_test_datadir}/unknown-logical-type.parquet"

Expand Down
11 changes: 11 additions & 0 deletions python/pyarrow/types.pxi
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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):
"""
Expand Down
Loading