GH-50312: [Python] Fix UUID extension type round-trip to pandas returning bytes
Project: apache/arrow | Language: C++ | Concepts: api-design, backwards-compatibility, testing-strategy | Difficulty: mid
Link: https://github.com/apache/arrow/pull/50325
The problem
When converting an Arrow Table or Array with the arrow.uuid extension type to a pandas DataFrame or Series using .to_pandas(), the resulting pandas column contained raw bytes objects instead of Python uuid.UUID objects. This was inconsistent with the behavior of .to_pylist() and did not match user or test expectations, especially for round-tripping data through pandas and Parquet.
Why it's non-trivial
A naive fix might simply change the Python conversion logic, but:
- Arrow extension types are handled generically, often falling back to their storage type (in this case, 16-byte fixed-size binary), so the conversion path must be customized.
- The change needs to be implemented in the C++ conversion layer for performance and to avoid reliance on Python-level hooks, which are less efficient and can be inconsistent.
- The solution must preserve backwards compatibility, especially for NumPy conversion, and avoid breaking the extension array protocol for other types.
The approach
1. Detecting the UUID Extension Type
A helper function is added in C++ to identify the Arrow UUID extension type:
+bool IsUuidExtension(const DataType& type) {
+ if (type.id() != Type::EXTENSION) {
+ return false;
+ }
+ const auto& extension_type = checked_cast<const ExtensionType&>(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<const FixedSizeBinaryType&>(storage_type).byte_width() == 16;
+}
This function checks if a given Arrow type is the UUID extension with the correct storage layout.
2. Converting Storage Bytes to uuid.UUID Python Objects
A new C++ helper constructs a uuid.UUID Python object from 16 bytes:
+Result<PyObject*> UuidFromBytes(std::string_view bytes, PyObject* kwargs) {
+ PyObject* uuid_class = GetUuidClass();
+ if (!uuid_class) {
+ return Status::Invalid("Could not import uuid.UUID");
+ }
+ OwnedRef py_bytes(
+ PyBytes_FromStringAndSize(bytes.data(), static_cast<Py_ssize_t>(bytes.size())));
+ RETURN_IF_PYERROR();
+ if (PyDict_SetItemString(kwargs, "bytes", py_bytes.obj()) < 0) {
+ RETURN_IF_PYERROR();
+ }
+ PyObject* empty_args = Py_GetConstantBorrowed(Py_CONSTANT_EMPTY_TUPLE);
+ PyObject* result = PyObject_Call(uuid_class, empty_args, kwargs);
+ RETURN_IF_PYERROR();
+ return result;
+}
This function creates a Python uuid.UUID instance from raw bytes, reusing empty argument/kwargs objects for efficiency.
3. Integrating the Conversion in the Pandas Export Path
The main conversion logic in arrow_to_pandas.cc is updated to use the above for UUID extension arrays:
+ 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<const ExtensionArray&>(*data.chunk(c));
+ storage_arrays.push_back(extension_array.storage());
+ }
+ ChunkedArray storage(std::move(storage_arrays), type.storage_type());
+ OwnedRef kwargs(PyDict_New());
+ RETURN_IF_PYERROR();
+ auto WrapUuid = [&](const std::string_view& view, PyObject** out) {
+ ARROW_ASSIGN_OR_RAISE(*out, internal::UuidFromBytes(view, kwargs.obj()));
+ return Status::OK();
+ };
+ return ConvertAsPyObjects<FixedSizeBinaryType>(options, storage, WrapUuid,
+ out_values);
+ }
This visitor method handles UUID extension arrays by converting each 16-byte value to a uuid.UUID object, chunk by chunk.
4. Ensuring Only UUID Uses This Path
Where the conversion path is selected, UUID gets special treatment, other extensions remain unchanged:
- *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;
This ensures only UUID extension columns are converted to Python objects (uuid.UUID) in pandas, others use the default extension handling.
5. Tests
New and extended tests verify:
- Pandas round-trip produces
uuid.UUIDobjects, not bytes. - NumPy conversion still returns raw bytes.
- Deduplication and pandas options behave as expected.
Example test:
+@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"])
From the review thread
GiTaDi-CrEaTe: 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???
parker-cassar: Maybe! It works here because the target is object dtype anyway, so building the Python objects is unavoidable. Other extension types could use the same from_arrow + reshape structure, but anything performance-sensitive would want a zero-copy path instead of going through Python objects.
rok: @parker-cassar I think we better implement this in C++ kind of like proposed here. This will likely improve performance and give us a cleaner API/make compat layer with pandas neater. Sorry for not suggesting this sooner. Do you think you can continue in this vein?
parker-cassar: Yes I agree and am happy you suggested this (: the Python compat layer approach doesn't really make sense here, especially since a UUID column in pandas isn't exactly a rare case. Doing it natively is cleaner. I'll give it a go and I'll try to have a PR in a few days or less based on your proposal.
Takeaways
- Specialized conversion logic is sometimes necessary for extension types to ensure correct, idiomatic results in downstream libraries (e.g., pandas).
- Implementing type conversions in the core (C++) layer can improve both performance and maintainability, compared to patching at the Python level.
- Testing should cover round-trips and edge cases (such as nulls and deduplication) to verify both correctness and compatibility.