Repository: dask/dask

Issue:
Categorical column dtype is not preserved during parquet roundtrip with pandas 2.0
**Describe the issue**:

With `pandas==2.0`, when reading parquet with partitioning on an integer column, the column is read back as a `CategoryDtype`, but its `dtype` is not consistent with `pyarrow`.

**Minimal Complete Verifiable Example**:

See better example in [this comment](https://github.com/dask/dask/issues/9966#issuecomment-1432485259).

```python
import pandas as pd
import dask.dataframe as dd
import pyarrow as pa
import pyarrow.parquet as pq
from pyarrow.dataset import HivePartitioning, write_dataset

df = pd.DataFrame({"col1": [1, 2], "col2": ["a", "b"]})

pa_path = "foo.parquet"
table = pa.Table.from_pandas(df)

# specify that partitioning column has int32 dtype
partitioning = HivePartitioning(pa.schema([("col1", pa.int32())]))

# write partitioned dataset with pyarrow
write_dataset(
    data=table,
    base_dir=pa_path,
    basename_template="part.{i}.parquet",
    format="parquet",
    partitioning=partitioning,
)

# check that reading pyarrow-written data is the same for pyarrow and dask
df_dask = dd.read_parquet(pa_path, engine="pyarrow").compute()
table2 = pq.read_table(pa_path)
df_arrow = table2.to_pandas()
print(f"{df_dask['col1'].dtype.categories=}")
print(f"{df_arrow['col1'].dtype.categories=}")
print(f"{table2=}")
```

This prints:

```
df_dask['col1'].dtype.categories=Index([1, 2], dtype='int64')
df_arrow['col1'].dtype.categories=Index([1, 2], dtype='int32')
table2=pyarrow.Table
col2: string
col1: dictionary<values=int32, indices=int32, ordered=0>
----
col2: [["a"],["b"]]
col1: [  -- dictionary:
[1,2]  -- indices:
[0],  -- dictionary:
[1,2]  -- indices:
[1]]
```

**Anything else we need to know?**:

This happens with pandas 2.0. With 1.5, `categories` were of dtype `int64`, whether they were read with `pyarrow` or `dask`.

The problem remains when specifying `engine="fastparquet"` with `dd.read_parquet`.

* Xref https://github.com/dask/dask/issues/9736.

This is responsible for a failing test with upstream pandas 2.0:

```
dask/dataframe/io/tests/test_parquet.py::test_roundtrip_partitioned_pyarrow_dataset[fastparquet]: AssertionError: Attributes of DataFrame.iloc[:, 0] (column name="col1") are different

Attribute "dtype" are different
[left]:  CategoricalDtype(categories=[1, 2], ordered=False)
[right]: CategoricalDtype(categories=[1, 2], ordered=False)
dask/dataframe/io/tests/test_parquet.py::test_roundtrip_partitioned_pyarrow_dataset[pyarrow]: AssertionError: Attributes of DataFrame.iloc[:, 0] (column name="col1") are different

Attribute "dtype" are different
[left]:  CategoricalDtype(categories=[1, 2], ordered=False)
[right]: CategoricalDtype(categories=[1, 2], ordered=False)
dask/dataframe/tests/test_arithmetics_reduction.py::test_datetime_std_across_axis1_null_results[False]: TypeError: float() argument must be a string or a real number, not 'Timestamp'
dask/dataframe/tests/test_arithmetics_reduction.py::test_datetime_std_across_axis1_null_results[True]: TypeError: float() argument must be a string or a real number, not 'Timestamp'
dask/dataframe/tests/test_groupby.py::test_groupby_unaligned_index[disk]: AssertionError: assert FrozenList(['a', None]) == FrozenList([None])
  At index 0 diff: 'a' != None
  Left contains one more item: None
  Full diff:
  - FrozenList([None])
  + FrozenList(['a', None])
  ?             +++++
```

**Environment**:

- Dask version: 2023.2.0+11.g0890b96b
- Python version: 3.11
- Operating System: macOS
- Install method (conda, pip, source): conda

cc @jrbourbeau @rjzamora 


Discussion:
Thanks for writing this up @j-bennet. I think we can remove the `HivePartitioning` part of this example and still produce the same issue with creating a parquet dataset with `pandas` and then reading it in with `pyarrow`, `pandas`, and `dask`:

```python
import pandas as pd
import pyarrow.parquet as pq

import dask.dataframe as dd

# Create parquet dataset with pandas
df = pd.DataFrame({"col1": [1, 2], "col2": ["a", "b"]})
outdir = "foo.parquet"
df.to_parquet(outdir, engine="pyarrow", partition_cols=["col1"])

# Load the dataset using pyarrow, pandas, and dask
table = pq.read_table(outdir)
df = pd.read_parquet(outdir, engine="pyarrow")
ddf = dd.read_parquet(outdir, engine="pyarrow")
print(f"{table.schema.field('col1') = }")
print(f"{df['col1'].dtype.categories = }")
print(f"{ddf['col1'].dtype.categories = }")
```

When running with `pandas=1.5.3` and `pyarrow=11.0.0`, the above snippet outputs

```
table.schema.field('col1') = pyarrow.Field<col1: dictionary<values=int32, indices=int32, ordered=0>>
df['col1'].dtype.categories = Int64Index([1, 2], dtype='int64')
ddf['col1'].dtype.categories = Int64Index([1, 2], dtype='int64')
```

note that `pyarrow` has `int32` while `pandas` and `dask` give `int64`.

When running with `pandas=2.0.0.dev0+1525.g9764f3dbf0` (a recent nightly release) and `pyarrow=11.0.0`, the above snippet outputs

```
table.schema.field('col1') = pyarrow.Field<col1: dictionary<values=int32, indices=int32, ordered=0>>
df['col1'].dtype.categories = Index([1, 2], dtype='int32')
ddf['col1'].dtype.categories = Index([1, 2], dtype='int64')
```

note that the `pandas` result has changed to `int32` while `dask` continues to give `int64`.


Hi, quick question:

Is the problem that the pyarrow table has dtype ``int32`` or that the categories in pandas have dtype ``int32``?

The pandas dtype is now consistent with the arrow dtype, because an Index can support arbitrary NumPy dtypes starting from pandas 2.0.
In my original example, the parquet was written with `int32` specified in schema, so this is how we should read it back.

In simplified example from @jrbourbeau, there was no explicit schema, but I would assume `dask` should preserve dtypes from `pyarrow` and `pandas` and read it as `int32`.

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
@@ -179,7 +179,7 @@ class PartitionObj:
 
     def __init__(self, name, keys):
         self.name = name
-        self.keys = sorted(keys)
+        self.keys = pd.Index(keys.sort_values(), copy=False)
 
 
 def _frag_subset(old_frag, row_groups):
@@ -1739,7 +1739,7 @@ def _read_table(
                     # columns are NOT in our table yet.
                     cat = keys_dict.get(partition.name, None)
                     cat_ind = np.full(
-                        len(arrow_table), partition.keys.index(cat), dtype="i4"
+                        len(arrow_table), partition.keys.get_loc(cat), dtype="i4"
                     )
                     arr = pa.DictionaryArray.from_arrays(
                         cat_ind, pa.array(partition.keys)


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
@@ -3553,6 +3553,12 @@ def test_partitioned_preserve_index(tmpdir, write_engine, read_engine):
     df1.to_parquet(tmp, partition_on="B", engine=write_engine)
 
     expect = data[data["B"] == 1]
+    if PANDAS_GT_200 and read_engine == "pyarrow":
+        # fastparquet does not preserve dtype of cats
+        expect = expect.copy()  # SettingWithCopyWarning
+        expect["B"] = expect["B"].astype(
+            pd.CategoricalDtype(expect["B"].dtype.categories.astype("int32"))
+        )
     got = dd.read_parquet(tmp, engine=read_engine, filters=[("B", "==", 1)])
     assert_eq(expect, got)
 
@@ -4461,10 +4467,13 @@ def test_custom_filename_with_partition(tmpdir, engine):
 
 
 @PYARROW_MARK
-@pytest.mark.xfail(PANDAS_GT_200, reason="https://github.com/dask/dask/issues/9966")
 def test_roundtrip_partitioned_pyarrow_dataset(tmpdir, engine):
     # See: https://github.com/dask/dask/issues/8650
 
+    if engine == "fastparquet" and PANDAS_GT_200:
+        # https://github.com/dask/dask/issues/9966
+        pytest.xfail("fastparquet reads as int64 while pyarrow does as int32")
+
     import pyarrow.parquet as pq
     from pyarrow.dataset import HivePartitioning, write_dataset
 
@@ -4832,3 +4841,21 @@ def test_read_parquet_convert_string_fastparquet_warns(tmp_path):
             UserWarning, match="`dataframe.convert_string` is not supported"
         ):
             dd.read_parquet(outfile, engine="fastparquet")
+
+
+@PYARROW_MARK
+@pytest.mark.skipif(
+    not PANDAS_GT_200, reason="pd.Index does not support int32 before 2.0"
+)
+def test_read_parquet_preserve_categorical_column_dtype(tmp_path):
+    df = pd.DataFrame({"a": [1, 2], "b": ["x", "y"]})
+
+    outdir = tmp_path / "out.parquet"
+    df.to_parquet(outdir, engine="pyarrow", partition_cols=["a"])
+    ddf = dd.read_parquet(outdir, engine="pyarrow")
+
+    expected = pd.DataFrame(
+        {"b": ["x", "y"], "a": pd.Categorical(pd.Index([1, 2], dtype="int32"))},
+        index=[0, 0],
+    )
+    assert_eq(ddf, expected)