apache / apache/arrow

[Python] Support Python enums in pyarrow

Open
#33,729 2 comments 1 reaction 0 assignees View on GitHub
Component: Python Type: bug
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.