[Python] support pandas' nullable Integer type in from_pandas

[asfimport]

From #4168. We should add support for pandas' nullable Integer extension dtypes, as those could map nicely to arrows integer types.

Ideally this happens in a generic way though, and not specific for this extension type, which is discussed in ARROW-5271

Reporter: Joris Van den Bossche / @jorisvandenbossche

Related issues:

• spark.createDataFrame fails with pandas DataFrame including pandas.NA (blocks)
• [Python] Add API to map Arrow types (including extension types) to pandas ExtensionArray instances for to_pandas conversions (is blocked by)
• [Python] Support protocols to extract Arrow objects from third-party classes (is blocked by)
• [Python] Interface for converting pandas ExtensionArray / other custom array objects to pyarrow Array (is blocked by)
• Support pandas nullable integer type Int64 (duplicates)

_{Note: This issue was originally created as ARROW-5379. Please see the migration documentation for further details.}

[asfimport]

Thomas Buhrmann / @buhrmann:
For the particular case of pd.Int64Dtype, the following may be a workaround for now, in case that's useful to anybody. In short, cast pandas Int64 columns to 'object' before converting to Arrow. When converting back to pandas, import with integer_object_nulls=True and cast back to Int64. Seems to work correctly for the below cases of pandas integer columns with or without NaNs, and different integer sizes:

 

import pandas as pd
import pyarrow as pa


def from_pandas(df):
    """Cast Int64 to object before 'serializing'"""
    for col in df:
        if isinstance(df[col].dtype, pd.Int64Dtype):
            df[col] = df[col].astype('object')
    return pa.Table.from_pandas(df)


def to_pandas(tbl):
    """After 'deserializing', recover the correct int type"""
    df = tbl.to_pandas(integer_object_nulls=True)

    for col in df:
        if (pa.types.is_integer(tbl.schema.field_by_name(col).type) and
            pd.api.types.is_object_dtype(df[col].dtype)):
                df[col] = df[col].astype('Int64')
    
    return df


df = pd.Series([0, 1, None, 2, 822215679726100500], dtype='Int64', name='x').to_frame()
# df = pd.Series([0, 1, 3, 2, 822215679726100500], dtype='Int64', name='x').to_frame()
# df = pd.Series([0, 1, 3, 2, 15], dtype='Int64', name='x').to_frame()
# df = pd.Series([0, 1, 3, 2, 15], dtype='int16', name='x').to_frame()

df2 = to_pandas(from_pandas(df))    
df2.dtypes

 

[asfimport]

Daniel Nugent / @nugend:
Is this actually something that would be appropriate to implement with extension types? It just requires that the mask parameter of pa.Array actually be used with Pandas integer columns, right?

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
It's indeed not needed to have an extension type for this.

With pandas master, the conversion of pandas -> arrow already works (using the \_\_arrow_array\_\_ protocol added in ARROW-3829, indeed using the mask of the pandas integer array).

The conversion the other way around (arrow -> pandas) is not yet working, but work in progress: see ARROW-2428

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
So the pandas -> arrow/feather conversion already works with pandas master and the latest Arrow release (0.15).

If you want to use this feature without relying on pandas master, you can use this monkeypatch (it's basically what is added in the development version of pandas master):

import pandas as pd
import pyarrow

pd.arrays.IntegerArray.__arrow_array__ = lambda self, type: pyarrow.array(self._data, mask=self._mask, type=type)

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
With the latest releases of pandas and pyarrow, full roundtrip for DataFrames with nullable integer dtype works (due to \_\_arrow_array\_\_ and \_\_from_arrow\_\_ being implemented on the pandas side) , so we can close this issue.

[asfimport]

Adam Lippai / @alippai:
What's the correct way to use this?

arr3 = pa.array([1, 2, 3, None], type=pa.int64())
pa.table([arr3], ['first']).to_pandas().dtypes

This results in:

first    float64
dtype: object```

Setting a `types_mapper` to the pandas Int64DType throw.s

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:

Setting a types_mapper to the pandas Int64DType throw.s

Can you show what you tried and the error? As that should normally work

[asfimport]

Adam Lippai / @alippai:
I was testing with pandas 1.1.3 and pyarrow 1.0.1. I don't know how I managed to get the exception :/

Now it works as expected without any problems and I can't reproduce the exception I got.

I'm sorry for the false (and incomplete) bug report, next time I'll be more careful. 

import pyarrow as pa
arr = pa.array([1, 2, 3, None], type=pa.int64())
pa.table([arr], ['first']).to_pandas(types_mapper={pa.int64(): pd.Int64Dtype()}.get).dtypes

first    Int64
dtype: object

[asfimport]

Joris Van den Bossche / @jorisvandenbossche:
Good to hear it is working!