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 c546bc1532ac..a585871ebe5d 100644 --- a/python/pyarrow/tests/parquet/test_data_types.py +++ b/python/pyarrow/tests/parquet/test_data_types.py @@ -604,6 +604,27 @@ 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"]) + + 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()