Repository: dask/dask
Pull request: [WIP] [Parquet] Handle new (dtype) pandas metadata with pyarrow
Closes #6085

Uses `ignore_metadata=False` during arrow -> pandas conversion in `ArrowEngine` parquet reader.  Before these changes, `ignore_metadata=True` was used to simplify the setting of the index.  However, the pandas metadata is needed to get certain dtype conversions correct.

- [ ] Tests added / passed
- [ ] Passes `black dask` / `flake8 dask`

@martindurant @msebbah


Problem statement:
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!

Hints from issue 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)`

Commit message: use issubset for clarity

Patch to review:
diff --git a/dask/dataframe/io/parquet/arrow.py b/dask/dataframe/io/parquet/arrow.py
index b75a764d000..2fae52a1e9e 100644
--- 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)
diff --git a/dask/dataframe/io/tests/test_parquet.py b/dask/dataframe/io/tests/test_parquet.py
index aace56121d4..77fcf54a893 100644
--- a/dask/dataframe/io/tests/test_parquet.py
+++ b/dask/dataframe/io/tests/test_parquet.py
@@ -2358,3 +2358,23 @@ 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
+
+
+@pytest.mark.parametrize("dtype", ["Int64", "str"])
+def test_pandas_metadata_nullable_pyarrow(tmpdir, dtype):
+
+    check_pyarrow()
+    if pa.__version__ < LooseVersion("0.16.0"):
+        pytest.skip("PyArrow 0.16 Required.")
+    if pd.__version__ < LooseVersion("1.0.0"):
+        pytest.skip("Pandas 1.0.0 Required.")
+    tmpdir = str(tmpdir)
+
+    ddf1 = dd.from_pandas(
+        pd.DataFrame({"A": ["dog" if dtype == "str" else 1, None]}, dtype=dtype),
+        npartitions=1,
+    )
+    ddf1.to_parquet(tmpdir, engine="pyarrow")
+    ddf2 = dd.read_parquet(tmpdir, engine="pyarrow")
+
+    assert_eq(ddf1, ddf2, check_index=False)


Code review comments:
diff_hunk: @@ -2358,3 +2358,23 @@ 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
+
+
+@pytest.mark.parametrize("dtype", ["Int64", "str"])
path: dask/dataframe/io/tests/test_parquet.py
text: @user1:
Can avoid the parametrization by doing both in the same dataframe

```python
df = pd.DataFrame({"A": pd.array([1, 2]), "B": pd.array(["a", "b"])})
...
```

@author:
Makes sense -  I was trying to isolate the two different cases, but I agree your suggestion is much cleaner :)