[Python] Support Python enums in pyarrow
- Dominant language
- C++
- Stars
- 17.1k
- Forks
- 4.3k
- Avg merge
- 3d 18h
- Merged PRs (30d)
- 91
Description
### Describe the bug, including details regarding any error messages, version, and platform.
Hello!
Filing this as a bug, though it could be feature request or even usage question.
Code:
```
import pyarrow
from enum import Enum
import pandas as pd
class Unit(Enum):
A = "A"
B = "B"
df = pd.DataFrame({'x': [Unit.A, Unit.B]})
print(pyarrow.Table.from_pandas(df))
```
Expected:
smth like
```
pyarrow.Table
x: dictionary
----
x: [ -- dictionary:
["A","B"] -- indices:
[0,1]]
```
Got:
```
Traceback (most recent call last):
File "x.py", line 12, in
print(pyarrow.Table.from_pandas(df))
File "pyarrow/table.pxi", line 3475, in pyarrow.lib.Table.from_pandas
File "/Users/a/venv/lib/python3.8/site-packages/pyarrow/pandas_compat.py", line 611, in dataframe_to_arrays
arrays = [convert_column(c, f)
File "/Users/avenv/lib/python3.8/site-packages/pyarrow/pandas_compat.py", line 611, in
arrays = [convert_column(c, f)
File "/Users/a/venv/lib/python3.8/site-packages/pyarrow/pandas_compat.py", line 598, in convert_column
raise e
File "/Users/a/venv/lib/python3.8/site-packages/pyarrow/pandas_compat.py", line 592, in convert_column
result = pa.array(col, type=type_, from_pandas=True, safe=safe)
File "pyarrow/array.pxi", line 316, in pyarrow.lib.array
File "pyarrow/array.pxi", line 83, in pyarrow.lib._ndarray_to_array
File "pyarrow/error.pxi", line 100, in pyarrow.lib.check_status
pyarrow.lib.ArrowInvalid: ("Could not convert with type Unit: did not recognize Python value type when inferring an Arrow data type", 'Conversion failed for column x with type object')
```
Extracting `.name` from enum values and converting to `category` works as expected.
### Component(s)
Python
Contributor guide
Research direction
Start with pyarrow/pandas_compat.py and the Table.from_pandas and pa.array conversion paths named in the traceback. Reproduce the example with the Python Enum values, then verify that conversion produces the expected dictionary-backed string column without breaking existing pandas conversions.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- pandas, python
- Domain
- data
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100