ARROW-2428: [Python] Support pandas ExtensionArray in Table.to_pandas conversion

[jorisvandenbossche]

Prototype for https://issues.apache.org/jira/browse/ARROW-2428

What does this PR do?

• Based on the pandas_metadata (stored when creating a Table from a pandas DataFrame), we infer which columns originally had a pandas extension dtype, and support a custom conversion (based on a __from_arrow__ method defined on the pandas extension dtype)
• The user can also specify explicitly with the extension_column keyword which columns should be converted to an extension dtype

This only covers use case 1 discussed in the issue: automatic roundtrip for pandas DataFrames that have extension dtypes.
So it eg does not yet provide a way to do this if the arrow.Table has no pandas metadata (did not originate from a pandas DataFrame)

[github-actions]

https://issues.apache.org/jira/browse/ARROW-2428

[wesm]

The implementation is somewhat gnarly but I don't see an easy way to achieve the desired outcome otherwise.

[wesm]

Incomplete?

[wesm]

Aside: is there a monkey-patch context manager somewhere?

[jorisvandenbossche]

pytest has a monkeypatch fixture that we can use. This would look like

def test_convert_to_extension_array(monkeypatch):
    ...
    monkeypatch.setattr(pd.Int64Dtype, "__from_arrow__", _Int64Dtype__from_arrow__, raising=False)
    ....

without the need for the try/except to ensure clean-up

[wesm]

I also don't know whether exposing the extension_columns parameter publicly makes sense.

[jorisvandenbossche]

I also don't know whether exposing the extension_columns parameter publicly makes sense.

Yes, I agree, I am also not yet fully clear about that part of the API. For now, it was mainly for being able to test this a bit more.

Currently, I provide an "automatic detection" based on the stored pandas_metadata (which enables automatic roundtrip).

But, I think we need other mechanisms as well:

• arrow extension types need to be able to specify a custom pandas dtype to convert to as well (one option is that they can point to a pandas ExtensionDtype object, and then use the same ExtensionDtype.__from_arrow__ logic).
• It would be nice to have some way to override all this logic. For example, consider the use case where you have a parquet file (or arrow memory in general) that does not originate from pandas (and thus has no pandas metadata), and you want to ensure that all integer columns are converted to the pandas nullable integer type instead of the default numpy int64 dtype. One option could be to specify (somehow somewhere) a type map like {pa.int64(): pd.IntegerDtype} (but how would that work with parametrized types?)

[wesm]

Seems like a little more work is needed here, at least some docstrings. Let me know when you want me to take another look

[jorisvandenbossche]

I think #5162 could already be merged (based on your comment there #5162 (comment))? Then I will update this PR after that.

[jorisvandenbossche]

OK, this should be ready for review now.

I added now support for 2 of the cases of https://issues.apache.org/jira/browse/ARROW-2428:

• If there is pandas_metadata, we will check which columns originated from pandas extension dtypes
• If you defined a pyarrow.ExtensionType, we check if the ExtensionType.to_pandas_dtype method returns a pandas extension dtype

In both cases, the actual conversion of the pyarrow array to the pandas ExtensionArray is dispatched to the pandas ExtensionDtype.__from_arrow__ method, which is a method on the pandas extension dtype that knows how to convert a arrow array into an extension array of its dtype.

[jorisvandenbossche]

I did some performance profiling to see the impact of this change, as checking dtypes can be quite expensive, and given that for dataframes with basic numpy types this will be only overhead without any benefit.

For a relatively normal sized dataframe (100_000 rows x 100 columns) of only floats, this gives a small overhead (from 26ms to 29ms for Table.to_pandas()).
But since this check scales linearly with the number of columns, for a "wide" dataframe (100 rows x 1000 columns) there is a big slowdown: 10ms to 22ms.
See https://gist.github.com/jorisvandenbossche/d36bceb82fd2dda38ee419ba51dff5ed

Now, the most expensive part (~90% of the overhead) of the code I added, is the pandas_dtype(..) function (a function from pandas that converts a string to either an (registered) ExtensionDtype or either a numpy dtype). Since we only need this in case it would return an ExtensionDtype, we could avoid calling this function if we know in advance we have a string representation of a numpy dtype. And since numpy supports a limited set of numpy dtypes, we could hardcode those "known" strings to avoid calling pandas_dtype. For the case above of simple floats dataframe, this would remove most of the overhead.
-> since it is an easy way to avoid overhead, will implement this.

Another (additional) option could also be to add a keyword to to_pandas to disable the extension type checks, so a user that wants the best performance and knows it only has basic types can disable this.

[wesm]

+1. This looks OK here, thanks @jorisvandenbossche!