From 457ff20652169bf73e4be1adf323fdd351bb5f74 Mon Sep 17 00:00:00 2001 From: parker-cassar Date: Tue, 7 Jul 2026 20:11:34 -0700 Subject: [PATCH 1/2] 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): """ From d0d60cb033fa3a4098d8aea08fdf95bee9de94cc Mon Sep 17 00:00:00 2001 From: Rok Mihevc Date: Tue, 21 Jul 2026 12:26:27 +0200 Subject: [PATCH 2/2] Convert UUIDs natively for pandas --- python/pyarrow/pandas_compat.py | 7 -- .../src/arrow/python/arrow_to_pandas.cc | 64 +++++++++++++++++-- .../pyarrow/tests/parquet/test_data_types.py | 12 ---- python/pyarrow/tests/test_extension_type.py | 39 +++++++++++ python/pyarrow/types.pxi | 11 ---- 5 files changed, 98 insertions(+), 35 deletions(-) diff --git a/python/pyarrow/pandas_compat.py b/python/pyarrow/pandas_compat.py index 999c92e75c1f..d8fd383d31d1 100644 --- a/python/pyarrow/pandas_compat.py +++ b/python/pyarrow/pandas_compat.py @@ -780,13 +780,6 @@ 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/src/arrow/python/arrow_to_pandas.cc b/python/pyarrow/src/arrow/python/arrow_to_pandas.cc index f163266f3b87..4a11195efc04 100644 --- a/python/pyarrow/src/arrow/python/arrow_to_pandas.cc +++ b/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(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(storage_type).byte_width() == 16; +} + template struct WrapBytes {}; @@ -1378,6 +1389,43 @@ 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(*data.chunk(c)); + storage_arrays.push_back(extension_array.storage()); + } + ChunkedArray storage(std::move(storage_arrays), type.storage_type()); + + OwnedRef uuid_module; + OwnedRef uuid_constructor; + RETURN_NOT_OK(internal::ImportModule("uuid", &uuid_module)); + RETURN_NOT_OK( + internal::ImportFromModule(uuid_module.obj(), "UUID", &uuid_constructor)); + + auto WrapUuid = [&](const std::string_view& view, PyObject** out) { + OwnedRef args(PyTuple_New(0)); + OwnedRef kwargs(PyDict_New()); + OwnedRef bytes( + PyBytes_FromStringAndSize(view.data(), static_cast(view.size()))); + RETURN_IF_PYERROR(); + if (PyDict_SetItemString(kwargs.obj(), "bytes", bytes.obj()) < 0) { + RETURN_IF_PYERROR(); + } + *out = PyObject_Call(uuid_constructor.obj(), args.obj(), kwargs.obj()); + RETURN_IF_PYERROR(); + return Status::OK(); + }; + return ConvertAsPyObjects(options, storage, WrapUuid, + out_values); + } + template enable_if_t::value || std::is_same::value || @@ -2253,7 +2301,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 +2409,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 +2652,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 diff --git a/python/pyarrow/tests/parquet/test_data_types.py b/python/pyarrow/tests/parquet/test_data_types.py index 7bd8238a3491..a585871ebe5d 100644 --- a/python/pyarrow/tests/parquet/test_data_types.py +++ b/python/pyarrow/tests/parquet/test_data_types.py @@ -625,18 +625,6 @@ def test_uuid_roundtrip(tempdir): 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/tests/test_extension_type.py b/python/pyarrow/tests/test_extension_type.py index 35b801eca87e..48855ca524cc 100644 --- a/python/pyarrow/tests/test_extension_type.py +++ b/python/pyarrow/tests/test_extension_type.py @@ -1400,6 +1400,45 @@ 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() diff --git a/python/pyarrow/types.pxi b/python/pyarrow/types.pxi index 7ce186972bcf..ec1a5a2ba9a3 100644 --- a/python/pyarrow/types.pxi +++ b/python/pyarrow/types.pxi @@ -1969,14 +1969,6 @@ 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. @@ -1995,9 +1987,6 @@ cdef class UuidType(BaseExtensionType): def __arrow_ext_scalar_class__(self): return UuidScalar - def to_pandas_dtype(self): - return _UuidPandasDtype() - cdef class FixedShapeTensorType(BaseExtensionType): """