[Python] Custom Python type/array subclasses for ExtensionTypes implemented in C++
- Dominant language
- C++
- Stars
- 17.1k
- Forks
- 4.3k
- Avg merge
- 3d 18h
- Merged PRs (30d)
- 91
Description
When wrapping a type (or array) in a pyarrow object, we need to define which Python class to use. Currently, for extension types, this logic lives here in `pyarrow_wrap_data_type`:
https://github.com/apache/arrow/blob/b413ac4f2b6911af5e8241803277caccc43aa3c4/python/pyarrow/public-api.pxi#L114-L120
So there are currently two options:
- The ExtensionType is implemented in Python, by subclassing `pyarrow.(Py)ExtensionType`, and which links to the C++ `arrow::py::PyExtensionType` (a subclass of `arrow::ExtensionType`). In this case, we store the python type instance on the C++ instance, and return this as python object in `pyarrow_wrap_data_type`.
- The ExtensionType is implemented in C++, and then we currently always fall back to wrap this in the `pyarrow.BaseExtenstionType` base class (there is currently a bug in this, but that is getting fixed in [GH-33802](https://github.com/apache/arrow/pull/33802)).
However, that means that for such extension types implemented in C++, there is currently no way to have a "richer" python Type object (or Array object, since that is determined by the Type, and for a BaseExtensionType, that will always use the base ExtensionArray). While for an extension type, you might want to add type-specific attributes or methods.
For canonical extension types that are implemented in Arrow C++ itself (for example, the currently discussed Tensor extension type in https://github.com/apache/arrow/pull/8510, or a previous effort to add complex type as extension type in https://github.com/apache/arrow/pull/10565), I think it will work today to create a custom subclass of `pyarrow.BaseExtensionType` for the specific canonical type, and then we could add a special case to `pyarrow_wrap_data_type` checking the name of the extension type, and if it is a canonical one we implement ourselves, use the python subclass we implemented ourselves.
But for extension types that are implemented in C++ externally (or for extension types that are implemented in Arrow C++, but for which we don't provide a custom python subclass), that doesn't work.
I am wondering to what extent we want to allow "registering" a python class that should be used when wrapping a specific C++ extension type (and to what extent this would be useful for
Contributor guide
Research direction
Start by reading pyarrow/public-api.pxi, especially pyarrow_wrap_data_type, and compare the current handling of Python- and C++-implemented extension types with GH-33802. Determine whether and how Python Type and Array subclasses should be registered for C++ extension types. Done means the desired registration behavior and its scope are defined clearly enough to support type-specific Python attributes or methods.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- api, backend
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100