Repository: dask/dask

Issue:
DataFrames with nullable integer or new string dtype don't support the full roundtrip to parquet
Hi guys,

Regarding the nullable integer dtype, pandas v1.0.0 offers the writing and reading back in with to_parquet() / read_parquet() starting with pyarrow >= 0.16.
It seems Dask Dataframes do not comply yet
```python
import dask.dataframe as dd

df = dd.from_pandas(pd.DataFrame({'A': [1, None]}, dtype='Int64'), npartitions=1)
df.to_parquet('foo', engine='pyarrow')
dd.read_parquet('foo', engine='pyarrow').compute().dtypes
```
Out: A    float64
dtype: object

However the _meta attribute looks good:
```python
dd.read_parquet('foo', engine='pyarrow')._meta.dtypes
```
Out: A    Int64
dtype: object

Same issue with the new string dtype

I am happy to help, by the way!


Discussion:
Thanks for the report. The fact that the metadata is correct is helpful.

The issue seems to be the `ignore_metadata` at https://github.com/dask/dask/blob/263187bbc263a395dc871e5874ee4fda21abb190/dask/dataframe/io/parquet/arrow.py#L549

```
517  ->	        df = df.to_pandas(categories=categories, use_threads=False, ignore_metadata=True)[
518  	            list(columns_and_parts)
519  	        ]
520
521  	        if index:
522  	            df = df.set_index(index)
523  	        return df
(Pdb) df
pyarrow.Table
A: int64
metadata
--------
{b'pandas': b'{"index_columns": [{"kind": "range", "name": null, "start": 0, "'
            b'stop": 2, "step": 1}], "column_indexes": [{"name": null, "field_'
            b'name": null, "pandas_type": "unicode", "numpy_type": "object", "'
            b'metadata": {"encoding": "UTF-8"}}], "columns": [{"name": "A", "f'
            b'ield_name": "A", "pandas_type": "int64", "numpy_type": "Int64", '
            b'"metadata": null}], "creator": {"library": "pyarrow", "version":'
            b' "0.16.0"}, "pandas_version": "1.0.3"}'}
(Pdb) df.to_pandas(ignore_metadata=True).dtypes
A    float64
dtype: object
(Pdb) df.to_pandas(ignore_metadata=False).dtypes
A    Int64
dtype: object
```

@msebbah I would start by removing that `ignore_metadata` and seeing what breaks. Are you interested in working on that? Or we could ping @rjzamora and ask if he recalls why it's there :)
`ignore_metadata` was probably included specifically for setting the index - in case the user wants a different index from the default one.
> ignore_metadata was probably included specifically for setting the index - in case the user wants a different index from the default one.

Right - The `ignore_metadata=True` is used to prevent pyarrow from using the metadata to set the index.  The user may want a different column to be used as the index (or just a default RangeIndex).   We do not **need** to do this if there is other metadata information that we really **do not** want to ignore...  I will try to think of a resonable way to clean up the logic
Is setting the index a cheap operation? You could always  
```
df.to_pandas(categories=categories, use_threads=False, ignore_metadata=False).reset_index()[list(columns_and_parts)]
```
> Is setting the index a cheap operation? You could always

I did confirm that it **does** work to do that.  However, it seems to take a few seconds on a large (30mil rows) `dask.datasets.timeseries` partition.
@TomAugspurger , is there a reason that re/setting the index should take any time. I can only think of the case that it's a range that would get materialised into a data array.
Sorry @martindurant - explicitly using `inplace=True` after the read + pyarrow->pandas conversion seems to be much faster.
Trying to set up an environment with pyarrow 16.0 to add a test.  However, the following changes seem to work fine on exisiting tests:

```python
df = piece.read(
    columns=columns,
    partitions=partitions,
    use_pandas_metadata=True,
    use_threads=False,
    **kwargs.get("read", {}),
).to_pandas(
    categories=categories, use_threads=False, ignore_metadata=False
)
df.reset_index(drop=False, inplace=True)
df = df[list(columns_and_parts)]

if index:
    df = df.set_index(index)
return df
```
I wonder if it's worth looking out for a range-index (i.e., one that has no data column, only metadata) 
> I wonder if it's worth looking out for a range-index (i.e., one that has no data column, only metadata)

Good call - There should be no reason to reset the index if `isinstance(df.index, pd.RangeIndex)`

Resolving patch:
diff --git a/dask/dataframe/io/parquet/arrow.py b/dask/dataframe/io/parquet/arrow.py
--- a/dask/dataframe/io/parquet/arrow.py
+++ b/dask/dataframe/io/parquet/arrow.py
@@ -546,9 +546,34 @@ def read_partition(
             use_pandas_metadata=True,
             use_threads=False,
             **kwargs.get("read", {}),
-        ).to_pandas(categories=categories, use_threads=False, ignore_metadata=True)[
-            list(columns_and_parts)
-        ]
+        ).to_pandas(categories=categories, use_threads=False, ignore_metadata=False)
+
+        # Note that `to_pandas(ignore_metadata=False)` means
+        # pyarrow will use the pandas metadata to set the index.
+        index_in_columns_and_parts = set(df.index.names).issubset(
+            set(columns_and_parts)
+        )
+        if not index:
+            if index_in_columns_and_parts:
+                # User does not want to set index and a desired
+                # column/partition has been set to the index
+                df.reset_index(drop=False, inplace=True)
+            else:
+                # User does not want to set index and an
+                # "unwanted" column has been set to the index
+                df.reset_index(drop=True, inplace=True)
+        else:
+            if set(df.index.names) != set(index) and index_in_columns_and_parts:
+                # The wrong index has been set and it contains
+                # one or more desired columns/partitions
+                df.reset_index(drop=False, inplace=True)
+            elif index_in_columns_and_parts:
+                # The correct index has already been set
+                index = False
+                columns_and_parts = list(
+                    set(columns_and_parts).difference(set(df.index.names))
+                )
+        df = df[list(columns_and_parts)]
 
         if index:
             df = df.set_index(index)


Test patch:
diff --git a/dask/dataframe/io/tests/test_parquet.py b/dask/dataframe/io/tests/test_parquet.py
--- a/dask/dataframe/io/tests/test_parquet.py
+++ b/dask/dataframe/io/tests/test_parquet.py
@@ -2365,3 +2365,27 @@ def test_filter_nonpartition_columns(
     df_read = ddf_read.compute()
     assert len(df_read) == len(df_read[df_read["time"] < 5])
     assert df_read["time"].max() < 5
+
+
+def test_pandas_metadata_nullable_pyarrow(tmpdir):
+
+    check_pyarrow()
+    if pa.__version__ < LooseVersion("0.16.0") or pd.__version__ < LooseVersion(
+        "1.0.0"
+    ):
+        pytest.skip("PyArrow>=0.16 and Pandas>=1.0.0 Required.")
+    tmpdir = str(tmpdir)
+
+    ddf1 = dd.from_pandas(
+        pd.DataFrame(
+            {
+                "A": pd.array([1, None, 2], dtype="Int64"),
+                "B": pd.array(["dog", "cat", None], dtype="str"),
+            }
+        ),
+        npartitions=1,
+    )
+    ddf1.to_parquet(tmpdir, engine="pyarrow")
+    ddf2 = dd.read_parquet(tmpdir, engine="pyarrow")
+
+    assert_eq(ddf1, ddf2, check_index=False)