Repository: dask/dask

Issue:
Nullable types saved incorrectly in partitioned parquet files, leading to ValueError when appending
**What happened**:

Same as #6085, but for partitioned parquet files. The dtypes change as such:

| original | reloaded via `read_parquet()` |
| --- | --- |
| Int64 | float64 |
| string | object |
| boolean | object |

This happens only when saving the parquet files with partitions; without partitions, the metadata is saved and reloaded correctly to match the original dtypes.

This has a side effect when appending to partitioned parquet files involving any of these dtypes, resulting in `ValueError: Appended dtypes differ.`.

**What you expected to happen**:

dtypes are saved and reloaded to match the original. This is the case when **not** partitioning the data:

```python
import dask.dataframe as dd
import pandas as pd
path = '/some/path/no_partition'

df = dd.from_pandas(pd.DataFrame({
        'part': ["X", "Y", "Z"],
        'pd_int': pd.Series([1, 2, None], dtype="Int64"),
        'pd_str': pd.Series(["a", "b", None], dtype="string"),
        'pd_bool': pd.Series([True, False, None], dtype="boolean"),
        'int': pd.Series([1, 2, 3], dtype="int"),
        'str': pd.Series(["a", "b", None], dtype="str"),
        'bool': pd.Series([True, False, True], dtype="bool"),
    }), npartitions=1)
df.to_parquet(path, engine='pyarrow')
```

The following outputs match each other:

```python
>>> df.dtypes    # Original
part        object
pd_int       Int64
pd_str      string
pd_bool    boolean
int          int64
str         object
bool          bool
dtype: object
>>> dd.read_parquet(path, engine='pyarrow').compute().dtypes    # Reload data
part        object
pd_int       Int64
pd_str      string
pd_bool    boolean
int          int64
str         object
bool          bool
dtype: object
>>> dd.read_parquet(path, engine='pyarrow')._meta.dtypes    # Meta
part        object
pd_int       Int64
pd_str      string
pd_bool    boolean
int          int64
str         object
bool          bool
dtype: object
```

**Minimal Complete Verifiable Example**:

```python
import dask.dataframe as dd
import pandas as pd
path2 = '/some/path/partitioned'

df = dd.from_pandas(pd.DataFrame({
        'part': ["X", "Y", "Z"],
        'pd_int': pd.Series([1, 2, None], dtype="Int64"),
        'pd_str': pd.Series(["a", "b", None], dtype="string"),
        'pd_bool': pd.Series([True, False, None], dtype="boolean"),
        'int': pd.Series([1, 2, 3], dtype="int"),
        'str': pd.Series(["a", "b", None], dtype="str"),
        'bool': pd.Series([True, False, True], dtype="bool"),
    }), npartitions=1)
df.to_parquet(path2, engine='pyarrow', partition_on=["part"])
```

In the following outputs, the "primitive" dtypes are correct, but extension types do not match the original:

```python
>>> dd.read_parquet(path2, engine='pyarrow').compute().dtypes
pd_int      float64
pd_str       object
pd_bool      object
int           int64
str          object
bool           bool
part       category
dtype: object
>>> dd.read_parquet(path2, engine='pyarrow')._meta.dtypes
pd_int      float64
pd_str       object
pd_bool      object
int           int64
str          object
bool           bool
part       category
dtype: object
```

Side effect: Cannot append to partitioned parquet files.

```python
>>> df.to_parquet(path2, engine='pyarrow', partition_on=["part"], append=True)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/var/folders/4t/cc5xr5r51y944twv96s5lcsh0000gn/T/ipykernel_63576/1882213619.py in <module>
----> 1 df.to_parquet(path2, engine='pyarrow', partition_on=["part"], append=True)

~/mambaforge/envs/swarm/lib/python3.8/site-packages/dask/dataframe/core.py in to_parquet(self, path, *args, **kwargs)
   4575         from .io import to_parquet
   4576 
-> 4577         return to_parquet(self, path, *args, **kwargs)
   4578 
   4579     def to_orc(self, path, *args, **kwargs):

~/mambaforge/envs/swarm/lib/python3.8/site-packages/dask/dataframe/io/parquet/core.py in to_parquet(df, path, engine, compression, write_index, append, overwrite, ignore_divisions, partition_on, storage_options, custom_metadata, write_metadata_file, compute, compute_kwargs, schema, name_function, **kwargs)
    681     # Engine-specific initialization steps to write the dataset.
    682     # Possibly create parquet metadata, and load existing stuff if appending
--> 683     meta, schema, i_offset = engine.initialize_write(
    684         df,
    685         fs,

~/mambaforge/envs/swarm/lib/python3.8/site-packages/dask/dataframe/io/parquet/arrow.py in initialize_write(cls, df, fs, path, append, partition_on, ignore_divisions, division_info, schema, index_cols, **kwargs)
    597             elif (pd.Series(dtypes).loc[names] != df[names].dtypes).any():
    598                 # TODO Coerce values for compatible but different dtypes
--> 599                 raise ValueError(
    600                     "Appended dtypes differ.\n{}".format(
    601                         set(dtypes.items()) ^ set(df.dtypes.iteritems())

ValueError: Appended dtypes differ.
{('pd_int', Int64Dtype()), ('pd_str', 'object'), ('bool', 'bool'), ('part', dtype('O')), ('str', 'object'), ('pd_bool', 'object'), ('__null_dask_index__', 'int64'), ('pd_int', 'float64'), ('bool', dtype('bool')), ('int', dtype('int64')), ('int', 'int64'), ('str', dtype('O')), ('pd_bool', BooleanDtype), ('pd_str', string[python]), ('__null_dask_index__', dtype('int64'))}
```

**Anything else we need to know?**:

**Environment**:

- Dask version: 2021.11.1
- Python version: 3.8.12
- Operating System: macOS 11.6
- Install method (conda, pip, source): pip



Discussion:
Thank you for submitting this issue! The examples you provided were great, I was able to easily reproduce the problem. Since this is a parquet-related problem, @rjzamora would you be able to recommend the best way to proceed?

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
@@ -63,7 +63,15 @@ def _append_row_groups(metadata, md):
 
 
 def _write_partitioned(
-    table, root_path, filename, partition_cols, fs, index_cols=(), **kwargs
+    table,
+    df,
+    root_path,
+    filename,
+    partition_cols,
+    fs,
+    preserve_index,
+    index_cols=(),
+    **kwargs,
 ):
     """Write table to a partitioned dataset with pyarrow.
 
@@ -78,7 +86,10 @@ def _write_partitioned(
     """
     fs.mkdirs(root_path, exist_ok=True)
 
-    df = table.to_pandas(ignore_metadata=True)
+    if preserve_index:
+        df.reset_index(inplace=True)
+    df = df[table.schema.names]
+
     index_cols = list(index_cols) if index_cols else []
     preserve_index = False
     if index_cols:
@@ -673,10 +684,12 @@ def write_partition(
         if partition_on:
             md_list = _write_partitioned(
                 t,
+                df,
                 path,
                 filename,
                 partition_on,
                 fs,
+                preserve_index,
                 index_cols=index_cols,
                 compression=compression,
                 **kwargs,


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
@@ -685,6 +685,13 @@ def test_append_with_partition(tmpdir, engine):
         }
     )
     df1.index.name = "index"
+
+    # Check that nullable dtypes work
+    # (see: https://github.com/dask/dask/issues/8373)
+    df0["lat"] = df0["lat"].astype("Int64")
+    df1["lat"].iloc[0] = np.nan
+    df1["lat"] = df1["lat"].astype("Int64")
+
     dd_df0 = dd.from_pandas(df0, npartitions=1)
     dd_df1 = dd.from_pandas(df1, npartitions=1)
     dd.to_parquet(dd_df0, tmp, partition_on=["lon"], engine=engine)