NVIDIA / NVIDIA/cudf

[BUG] `pyarrow.table` does not accept `cudf.pandas.DataFrame`

Open
#14,521 5 comments 0 reactions 0 assignees View on GitHub
bug cudf.pandas Python
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.