Repository: dask/dask
Pull request: Auto-detect categorical columns in ArrowEngine-based read_parquet
Closes #5677 

Explicit use of the `categories` keyword is currently required to preserve the `"categorical"` dtype during the read phase of a parquet round-trip.  Since the categorical dtype is typically defined within the pandas metadata, there is no reason not to auto-detect the categorical columns.  This PR adds some simple logic to do this for the pyarrow-based reader (it is already done for fastparquet).

- [ ] Tests added / passed
- [ ] Passes `black dask` / `flake8 dask`


Problem statement:
Reading parquet with category dtype
Hi,
We have encountered a problem reading `category` dtype format from `parquet` file.  
versions:
 dask: 2.8.1
 pandas: 0.25.3
 pyarrow: 0.15.1 
 fastparquet:0.3.2

```python
import dask.dataframe as dd
import dask.datasets as ds
import pandas as pd

df = ds.timeseries()
df.name = df.name.astype('category')
df.dtypes
```  
prints out

> id                    int64
> name               category
> x                   float64
> y                   float64
  


```
df.to_parquet('test_parq', engine='pyarrow')
dff = dd.read_parquet('test_parq', engine='pyarrow')
dff.dtypes  
```

prints out 

> id                    int64
> name         datetime64[ns]
> x                   float64
> y                   float64
  

  
It is worth to note that:  
1. when running with `dask engine=fastparquet` the categorical column is preserved.  
2. when reading the parquet file directly with `pandas engine=pyarrow` the categorical column is preserved.

Hints from issue discussion:
Thanks for the report. Can you reproduce this issue when writing / reading using pandas or PyArrow directly? Or is it just with Dask?

>     2. when reading the parquet file directly with `pandas engine=pyarrow` the categorical column is preserved.

When using `pandas` there is no problem.

Also we used `PyArrow` directly (`read_table/write_table') and the dtype was successfully preserved. Thus it seams only an issue with `dask` implementation. 


Thanks for confirming. cc @rjzamora if you have a chance to take a look next week.
Thanks for raising @sephib - You are correct.  We don't seem to be using the dtype information stored in the pandas metadata.  Therefore, the type isn't corrected to  categorical until you call `compute()`, etc.  This is definitely something we should be able to fix.

In the mean time, you can still read categorical data *correctly* if you explicitly specify the categorical columns in the `read_parquet` command.   This is likely the reason we haven't bothered auto-detect the categorical columns so far. For example:

```python
df.to_parquet('test_parq', engine='pyarrow')
dff = dd.read_parquet('test_parq', engine='pyarrow', categories=["name"])
dff.dtypes  
```
**Out**:
```
id         int64
name    category
x        float64
y        float64
dtype: object
```


Thx for the work around.

We'll keep this open till #5690 is merged.

Commit message: only auto-populate categories if originally set to None

Patch to review:
diff --git a/dask/dataframe/io/parquet/arrow.py b/dask/dataframe/io/parquet/arrow.py
index 938c1dfa20e..708a3ab7bdd 100644
--- a/dask/dataframe/io/parquet/arrow.py
+++ b/dask/dataframe/io/parquet/arrow.py
@@ -168,6 +168,13 @@ def read_metadata(
                 storage_name_mapping,
                 column_index_names,
             ) = _parse_pandas_metadata(pandas_metadata)
+            if categories is None:
+                categories = []
+                for col in pandas_metadata["columns"]:
+                    if (col["pandas_type"] == "categorical") and (
+                        col["name"] not in categories
+                    ):
+                        categories.append(col["name"])
         else:
             index_names = []
             column_names = schema.names
diff --git a/dask/dataframe/io/tests/test_parquet.py b/dask/dataframe/io/tests/test_parquet.py
index 2782b18acc7..e4f141c0fc7 100644
--- a/dask/dataframe/io/tests/test_parquet.py
+++ b/dask/dataframe/io/tests/test_parquet.py
@@ -866,6 +866,10 @@ def test_categories(tmpdir, engine):
     ddf.to_parquet(fn, engine=engine)
     ddf2 = dd.read_parquet(fn, categories=["y"], engine=engine)
 
+    # Shouldn't need to specify categories explicitly
+    ddf3 = dd.read_parquet(fn, engine=engine)
+    assert_eq(ddf3, ddf2)
+
     with pytest.raises(NotImplementedError):
         ddf2.y.cat.categories
     assert set(ddf2.y.compute().cat.categories) == {"a", "b", "c"}
@@ -878,7 +882,7 @@ def test_categories(tmpdir, engine):
             # attempt to load as category that which is not so encoded
             ddf2 = dd.read_parquet(fn, categories=["x"], engine=engine).compute()
 
-    with pytest.raises(ValueError):
+    with pytest.raises((ValueError, FutureWarning)):
         # attempt to load as category unknown column
         ddf2 = dd.read_parquet(fn, categories=["foo"], engine=engine)
 
@@ -1794,13 +1798,13 @@ def test_append_cat_fp(tmpdir, engine):
         pytest.param(
             pd.DataFrame({"x": pd.Categorical(["a", "b", "a"])}),
             marks=pytest.mark.xfail(
-                reason="https://issues.apache.org/jira/browse/ARROW-3652"
+                reason="https://issues.apache.org/jira/browse/ARROW-3652", strict=False
             ),
         ),
         pytest.param(
             pd.DataFrame({"x": pd.Categorical([1, 2, 1])}),
             marks=pytest.mark.xfail(
-                reason="https://issues.apache.org/jira/browse/ARROW-3652"
+                reason="https://issues.apache.org/jira/browse/ARROW-3652", strict=False
             ),
         ),
         pd.DataFrame({"x": list(map(pd.Timestamp, [3000000, 2000000, 1000000]))}),  # ms


Code review comments:
diff_hunk: @@ -1794,13 +1798,13 @@ def test_append_cat_fp(tmpdir, engine):
         pytest.param(
             pd.DataFrame({"x": pd.Categorical(["a", "b", "a"])}),
             marks=pytest.mark.xfail(
-                reason="https://issues.apache.org/jira/browse/ARROW-3652"
+                reason="https://issues.apache.org/jira/browse/ARROW-3652", strict=False
path: dask/dataframe/io/tests/test_parquet.py
text: @user1:
Is this working for pyarrow 0.15+? If so, can you change these to be

```python
PYARROW_015 = LooseVersion(pa.__version__ >= "0.15.0"
strict=not PYARROW_0150
```

Just to ensure we don't regress on things.

@author:
Good call - Thanks!