Repository: pandas-dev/pandas
Pull request: BUG/TST: fix arrow roundtrip / parquet tests for recent pyarrow
Closes https://github.com/pandas-dev/pandas/issues/29976

Problem statement:
test_additional_extension_arrays fails with pd.NA
cc @jorisvandenbossche 

im using pyarrow 0.15.1 locally

```
_____________________________________________ TestParquetPyArrow.test_additional_extension_arrays _____________________________________________

self = <pandas.tests.io.test_parquet.TestParquetPyArrow object at 0x129bc4050>, pa = 'pyarrow'

    @td.skip_if_no("pyarrow", min_version="0.15.0")
    def test_additional_extension_arrays(self, pa):
        # test additional ExtensionArrays that are supported through the
        # __arrow_array__ protocol
        df = pd.DataFrame(
            {
                "a": pd.Series([1, 2, 3], dtype="Int64"),
                "b": pd.Series(["a", None, "c"], dtype="string"),
            }
        )
        if LooseVersion(pyarrow.__version__) >= LooseVersion("0.15.1.dev"):
            expected = df
        else:
            # de-serialized as plain int / object
            expected = df.assign(a=df.a.astype("int64"), b=df.b.astype("object"))
>       check_round_trip(df, pa, expected=expected)

pandas/tests/io/test_parquet.py:522: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
pandas/tests/io/test_parquet.py:175: in check_round_trip
    compare(repeat)
pandas/tests/io/test_parquet.py:167: in compare
    df.to_parquet(path, **write_kwargs)
pandas/core/frame.py:2061: in to_parquet
    **kwargs,
pandas/io/parquet.py:255: in to_parquet
    **kwargs,
pandas/io/parquet.py:101: in write
    table = self.api.Table.from_pandas(df, **from_pandas_kwargs)
pyarrow/table.pxi:1057: in pyarrow.lib.Table.from_pandas
    ???
/usr/local/lib/python3.7/site-packages/pyarrow/pandas_compat.py:555: in dataframe_to_arrays
    for c, f in zip(columns_to_convert, convert_fields)]
/usr/local/lib/python3.7/site-packages/pyarrow/pandas_compat.py:555: in <listcomp>
    for c, f in zip(columns_to_convert, convert_fields)]
/usr/local/lib/python3.7/site-packages/pyarrow/pandas_compat.py:546: in convert_column
    raise e
/usr/local/lib/python3.7/site-packages/pyarrow/pandas_compat.py:540: in convert_column
    result = pa.array(col, type=type_, from_pandas=True, safe=safe)
pyarrow/array.pxi:196: in pyarrow.lib.array
    ???
pyarrow/array.pxi:92: in pyarrow.lib._handle_arrow_array_protocol
    ???
pandas/core/arrays/string_.py:200: in __arrow_array__
    return pa.array(self._ndarray, type=type, from_pandas=True)
pyarrow/array.pxi:207: in pyarrow.lib.array
    ???
pyarrow/array.pxi:78: in pyarrow.lib._ndarray_to_array
    ???
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

>   ???
E   pyarrow.lib.ArrowTypeError: ("Expected a string or bytes object, got a 'NAType' object", 'Conversion failed for column b with type string')
```

Hints from issue discussion:
Hmm, thanks for catching that! We should really be testing this on CI .. 
Will look into it.

Commit message: BUG/TST: fix arrow roundtrip / parquet tests for recent pyarrow

Patch to review:
diff --git a/pandas/core/arrays/boolean.py b/pandas/core/arrays/boolean.py
index aec3397bddd16..05f68356163dc 100644
--- a/pandas/core/arrays/boolean.py
+++ b/pandas/core/arrays/boolean.py
@@ -101,6 +101,24 @@ def __repr__(self) -> str:
     def _is_boolean(self) -> bool:
         return True
 
+    def __from_arrow__(self, array):
+        """Construct BooleanArray from passed pyarrow Array/ChunkedArray"""
+        import pyarrow
+
+        if isinstance(array, pyarrow.Array):
+            chunks = [array]
+        else:
+            # pyarrow.ChunkedArray
+            chunks = array.chunks
+
+        results = []
+        for arr in chunks:
+            # TODO should optimize this without going through object array
+            bool_arr = BooleanArray._from_sequence(np.array(arr))
+            results.append(bool_arr)
+
+        return BooleanArray._concat_same_type(results)
+
 
 def coerce_to_array(values, mask=None, copy: bool = False):
     """
diff --git a/pandas/core/arrays/string_.py b/pandas/core/arrays/string_.py
index f6af05ab4d9e7..c2f797cb4b6c2 100644
--- a/pandas/core/arrays/string_.py
+++ b/pandas/core/arrays/string_.py
@@ -86,7 +86,7 @@ def __from_arrow__(self, array):
 
         results = []
         for arr in chunks:
-            # using _from_sequence to ensure None is convered to np.nan
+            # using _from_sequence to ensure None is convered to NA
             str_arr = StringArray._from_sequence(np.array(arr))
             results.append(str_arr)
 
@@ -197,7 +197,10 @@ def __arrow_array__(self, type=None):
 
         if type is None:
             type = pa.string()
-        return pa.array(self._ndarray, type=type, from_pandas=True)
+
+        values = self._ndarray.copy()
+        values[self.isna()] = None
+        return pa.array(values, type=type, from_pandas=True)
 
     def _values_for_factorize(self):
         arr = self._ndarray.copy()
diff --git a/pandas/tests/arrays/string_/test_string.py b/pandas/tests/arrays/string_/test_string.py
index 0dfd75a2042b0..aaa6d1cacd0a7 100644
--- a/pandas/tests/arrays/string_/test_string.py
+++ b/pandas/tests/arrays/string_/test_string.py
@@ -204,5 +204,5 @@ def test_arrow_roundtrip():
     result = table.to_pandas()
     assert isinstance(result["a"].dtype, pd.StringDtype)
     tm.assert_frame_equal(result, df)
-    # ensure the missing value is represented by NaN and not None
-    assert np.isnan(result.loc[2, "a"])
+    # ensure the missing value is represented by NA and not np.nan or None
+    assert result.loc[2, "a"] is pd.NA
diff --git a/pandas/tests/arrays/test_boolean.py b/pandas/tests/arrays/test_boolean.py
index a13bb8edc8e48..1e331429975a5 100644
--- a/pandas/tests/arrays/test_boolean.py
+++ b/pandas/tests/arrays/test_boolean.py
@@ -546,12 +546,29 @@ def test_reductions_return_types(dropna, data, all_numeric_reductions):
 #         result = arr[mask]
 
 
-@pytest.mark.skip(reason="broken test")
 @td.skip_if_no("pyarrow", min_version="0.15.0")
 def test_arrow_array(data):
     # protocol added in 0.15.0
     import pyarrow as pa
 
     arr = pa.array(data)
-    expected = pa.array(np.array(data, dtype=object), type=pa.bool_(), from_pandas=True)
+
+    # TODO use to_numpy(na_value=None) here
+    data_object = np.array(data, dtype=object)
+    data_object[data.isna()] = None
+    expected = pa.array(data_object, type=pa.bool_(), from_pandas=True)
     assert arr.equals(expected)
+
+
+@td.skip_if_no("pyarrow", min_version="0.15.1.dev")
+def test_arrow_roundtrip():
+    # roundtrip possible from arrow 1.0.0
+    import pyarrow as pa
+
+    data = pd.array([True, False, None], dtype="boolean")
+    df = pd.DataFrame({"a": data})
+    table = pa.table(df)
+    assert table.field("a").type == "bool"
+    result = table.to_pandas()
+    assert isinstance(result["a"].dtype, pd.BooleanDtype)
+    tm.assert_frame_equal(result, df)
diff --git a/pandas/tests/io/test_parquet.py b/pandas/tests/io/test_parquet.py
index a98c93c250070..3e687d185df84 100644
--- a/pandas/tests/io/test_parquet.py
+++ b/pandas/tests/io/test_parquet.py
@@ -504,7 +504,6 @@ def test_empty_dataframe(self, pa):
         df = pd.DataFrame()
         check_round_trip(df, pa)
 
-    @pytest.mark.skip(reason="broken test")
     @td.skip_if_no("pyarrow", min_version="0.15.0")
     def test_additional_extension_arrays(self, pa):
         # test additional ExtensionArrays that are supported through the


Code review comments:
diff_hunk: @@ -546,12 +546,29 @@ def test_reductions_return_types(dropna, data, all_numeric_reductions):
 #         result = arr[mask]
 
 
-@pytest.mark.skip(reason="broken test")
 @td.skip_if_no("pyarrow", min_version="0.15.0")
 def test_arrow_array(data):
     # protocol added in 0.15.0
     import pyarrow as pa
 
     arr = pa.array(data)
-    expected = pa.array(np.array(data, dtype=object), type=pa.bool_(), from_pandas=True)
+
+    # TODO use to_numpy(na_value=None) here
path: pandas/tests/arrays/test_boolean.py
text: @user1:
is this todo actionable now or do we have to wait until [condition]?

@author:
It's not yet implemented, open issue for this is https://github.com/pandas-dev/pandas/issues/30038