[BUG] `pyarrow.table` does not accept `cudf.pandas.DataFrame`
- Dominant language
- C++
- Stars
- 9.8k
- Forks
- 1.1k
- Avg merge
- 3d 6m
- Merged PRs (30d)
- 278
Description
### Background
Most of pyarrow is implemented in Cython. They have a lazy-loaded pandas-shim which they use to interoperate with pandas. This is implemented as the `_PandasAPIShim` `cdef` class. There is a singleton shim object that is accessible as `pyarrow.lib._pandas_api` from python (and as both `pyarrow.lib._pandas_api` and `pyarrow.lib.pandas_api` from cython):
```
In [1]: import pyarrow
In [2]: pyarrow.lib.pandas_api
---------------------------------------------------------------------------
AttributeError Traceback (most recent call last)
Cell In[2], line 1
----> 1 pyarrow.lib.pandas_api
AttributeError: module 'pyarrow.lib' has no attribute 'pandas_api'
In [3]: pyarrow.lib._pandas_api
Out[3]:
```
So at import time we make this API shim, and then it lazily initialises itself on first use.
This object saves a number of things:
1. The observed pandas module it exported (by doing `import pandas; self.pandas = pandas`)
2. Various type constructors (e.g. `self.data_frame = pandas.DataFrame`)
The first is relatively unproblematic. What we would like is for that module to be our intercepted wrapped module, which we can arrange with a little bit of rejigging of imports in cudf. It is the memoisation of the type constructors that is the problematic thing.
### cudf.pandas wrapping scheme
Recall that the way our wrapping scheme works is that we deliver wrapped _modules_ and decide at `__getattr__` time whether any attribute lookups deliver real or wrapped attributes. So:
```
%load_ext cudf.pandas
import pandas as pd # pd is _always_ a wrapped module
DataFrame = pd.DataFrame # this is context-dependant either a real or wrapped constructor
```
This works well, as long as someone doesn't memoise an attribute lookup. If they do, we only get to make the decision about what type of attribute to deliver once:
```
In [1]: %load_ext cudf.pandas
In [3]: from cudf.pandas.module_finder import disable_transparent_mode_if_enabled
In [4]: with disable_transparent_mode_if_enabled():
...: from pandas import DataFrame
...:
In [5]: DataFrame
Out[5]: pandas.core.frame.DataFrame
In [6]: from pandas import DataFrame
In [7]: DataFrame
Out[7]: cudf.pandas._wrappers.pandas.DataFrame
```
Unfortunately, pyarrow's pandas shim does exactly this. And we can't make the right decision, because sometimes (when used inside cudf) we need to deliver real objects, other times (when the user is using pyarrow) we need to deliver wrapped ones.
I said it would be a miracle if @shwina's approach worked, and it _kind of_ does, but unfortunately it's not quite miraculous enough. Here's what's going on:
`pa.table` uses `pa.lib._pandas_api.is_data_frame` to determine if the passed object is a pandas dataframe. The leading underscore here is crucial! This is a python object that we can replace. Similarly, `pa.Table.from_pandas` calls out to some Python code that uses `pa.lib._pandas_api` (which we can control).
However, the `to_pandas` method on the resulting object calls `pyarrow.lib.pandas_api.data_frame` note _no_ leading underscore. This is a Cython level module attribute that we _can't_ replace from Python.
So, we have this:
```
%load_ext cudf.pandas
import pyarrow as pa
pa.lib._pandas_api = pa.lib._PandasAPIShim()
pa.lib.pandas_api = pa.lib._pandas_api # this doesn't do anything at the Cython level
import pandas as pd
df = pd.DataFrame({"a": [1, 2, 3]})
tab = pa.table(df) # works
df2 = tab.to_pandas()
print(type(df))
#
print(type(df2))
#
```
### What if we initialise the Cython level object with wrapped attributes?
This seems like it might work, we have to be a bit careful about how we're importing pyarrow in cudf, but we can make this work so that `pa.lib.pandas_api` is a shim wrapper that sees our wrapped attributes.
_But_ the memoisation breaks things, because inside cudf we use `to_arrow().to_pandas()` in various places to produce an honest-to-goodness pandas object, but this will now produce a wrapped object (and wrapping things in a `disable_transparent_mode_if_enabled()` context manager won't help because we _already_ took the decision about what constructor to deliver).
### Options
1. Convince the arrow folks that the memoisation of attribute lookup in their pandas shim is not really a performance win, and that it would be convenient if they just used `self.pandas.DataFrame`. This would allow us to (context-dependently) provide a real or wrapped object as appropriate. We would likely do this after releasing cudf-PAM since then the motivation is clear.
2. Rather than letting module attribute lookup be context dependent, always deliver wrapped types such that the constructor context-dependently decides whether or not to deliver a real or wrapped type.
3. Something else?
**Steps/Code to reproduce bug**
```python
In [1]: %load_ext cudf.pandas
In [2]: import pandas as pd
In [3]: df = pd.DataFrame({'a': [1, 2, 3]})
In [4]: import pyarrow as pa
In [5]: pa.table(df)
---------------------------------------------------------------------------
TypeError Traceback (most recent call last)
Cell In[5], line 1
----> 1 pa.table(df)
File /nvme/0/pgali/envs/cudfdev/lib/python3.10/site-packages/pyarrow/table.pxi:5165, in pyarrow.lib.table()
TypeError: Expected pandas DataFrame, python dictionary or list of arrays
```
Contributor guide
Assessment
This issue has not been assessed yet.