Repository: dask/dask

Issue:
DataFrame has incorrect dtype after .compute() when group by categorical column
If you join Dask DataFrame on a categorical column, then the outputted Dask DataFrame column is still `category` dtype. However, the moment you `.compute()` the outputted Dask DataFrame, then the column is the wrong dtype, not categorical.

Tested on Dask 2.14.0 and Pandas 1.0.3
This example where the category type looks like a float, so after .compute(), the dtype is float.
```
import dask.dataframe as dd
import numpy as np
import pandas as pd

df = pd.DataFrame({
    "GROUPBY_COL" : [0, 0, 1, 1],
    "TO_AGG_COL": [0, 0, 10, 10],
})

ddf = dd.from_pandas(df, npartitions=2)
ddf["GROUPBY_COL"] = ddf["GROUPBY_COL"].astype("category")
ddf2 = ddf.groupby("GROUPBY_COL").agg("sum").rename(columns={"TO_AGG_COL": "AGGED_COL"})

print(ddf.join(ddf2, on="GROUPBY_COL", how="left").dtypes) # prints category
print(ddf.join(ddf2, on="GROUPBY_COL", how="left").compute().dtypes) # prints int64
```
<br>

If the categorical column looks like an float, then Dask DataFrame says "category" type but upon `.compute()`, the pandas DataFrame says `float` type.  

```
import dask.dataframe as dd
import numpy as np
import pandas as pd

df = pd.DataFrame({
    "GROUPBY_COL" : [0.0, 0.0, 1.0, 1.0],
    "TO_AGG_COL": [0, 0, 10, 10],
})

ddf = dd.from_pandas(df, npartitions=2)
ddf["GROUPBY_COL"] = ddf["GROUPBY_COL"].astype("category")
ddf2 = ddf.groupby("GROUPBY_COL").agg("sum").rename(columns={"TO_AGG_COL": "AGGED_COL"})

print(ddf.join(ddf2, on="GROUPBY_COL", how="left").dtypes)
print(ddf.join(ddf2, on="GROUPBY_COL", how="left").compute().dtypes)
```
<br>


Similarly, if the categorical column is a string that does not look like an integer, then Dask DataFrame says "category" type but upon `.compute()`, the pandas DataFrame says `object` type.
```
import dask.dataframe as dd
import numpy as np
import pandas as pd

df = pd.DataFrame({
    "GROUPBY_COL" : ["a", "a", "b", "b"],
    "TO_AGG_COL": [0, 0, 10, 10],
})

ddf = dd.from_pandas(df, npartitions=2)
ddf["GROUPBY_COL"] = ddf["GROUPBY_COL"].astype("category")
ddf2 = ddf.groupby("GROUPBY_COL").agg("sum").rename(columns={"TO_AGG_COL": "AGGED_COL"})

print(ddf.join(ddf2, on="GROUPBY_COL", how="left").dtypes) # prints category
print(ddf.join(ddf2, on="GROUPBY_COL", how="left").compute().dtypes) # prints object
```
DataFrame.merge raises with AttributeError
The following raises with `AttributeError: 'numpy.ndarray' object has no attribute 'categories'`


```python
import pandas as pd
import dask.dataframe as dd
a = pd.DataFrame({"A": [0, 1, 2, 3], "B": [4, 5, 6, 7]})
b = pd.DataFrame({"A": [0, 1, 2, 4], "C": [4, 5, 7, 7]})


df1 = dd.from_pandas(a, 2)
df1['A'] = df1.A.astype('category')

df2 = dd.from_pandas(b, 2)
df2['A'] = df2.A.astype('category')

r = df1.merge(df2, on="A")
r.compute()
```

```pytb
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-2-d6a6fee6936c> in <module>
     12
     13 r = df1.merge(df2, on="A")
---> 14 r.compute()

~/sandbox/dask/dask/base.py in compute(self, **kwargs)
    164         dask.base.compute
    165         """
--> 166         (result,) = compute(self, traverse=False, **kwargs)
    167         return result
    168

~/sandbox/dask/dask/base.py in compute(*args, **kwargs)
    436     postcomputes = [x.__dask_postcompute__() for x in collections]
    437     results = schedule(dsk, keys, **kwargs)
--> 438     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    439
    440

~/sandbox/dask/dask/base.py in <listcomp>(.0)
    436     postcomputes = [x.__dask_postcompute__() for x in collections]
    437     results = schedule(dsk, keys, **kwargs)
--> 438     return repack([f(r, *a) for r, (f, a) in zip(results, postcomputes)])
    439
    440

~/sandbox/dask/dask/dataframe/core.py in finalize(results)
    102
    103 def finalize(results):
--> 104     return _concat(results)
    105
    106

~/sandbox/dask/dask/dataframe/core.py in _concat(args, ignore_index)
     97         args[0]
     98         if not args2
---> 99         else methods.concat(args2, uniform=True, ignore_index=ignore_index)
    100     )
    101

~/sandbox/dask/dask/dataframe/methods.py in concat(dfs, axis, join, uniform, filter_warning, ignore_index)
    381             uniform=uniform,
    382             filter_warning=filter_warning,
--> 383             ignore_index=ignore_index,
    384         )
    385

~/sandbox/dask/dask/dataframe/methods.py in concat_pandas(dfs, axis, join, uniform, filter_warning, ignore_index)
    490                         )
    491                         parts.append(data)
--> 492                 out[col] = union_categoricals(parts)
    493                 # Pandas resets index type on assignment if frame is empty
    494                 # https://github.com/pandas-dev/pandas/issues/17101

~/Envs/dask-dev/lib/python3.7/site-packages/pandas/core/dtypes/concat.py in union_categoricals(to_union, sort_categories, ignore_order)
    304     if not all(
    305         is_dtype_equal(other.categories.dtype, first.categories.dtype)
--> 306         for other in to_union[1:]
    307     ):
    308         raise TypeError("dtype of categories must be the same")

~/Envs/dask-dev/lib/python3.7/site-packages/pandas/core/dtypes/concat.py in <genexpr>(.0)
    304     if not all(
    305         is_dtype_equal(other.categories.dtype, first.categories.dtype)
--> 306         for other in to_union[1:]
    307     ):
    308         raise TypeError("dtype of categories must be the same")

AttributeError: 'numpy.ndarray' object has no attribute 'categories'
```

This only happens with unknown categories. With known categories there's no exception.


Discussion:
@eugeneh101 thanks for the report. It looks like the issue is somewhere in https://github.com/dask/dask/blob/340103bb56e4979bd42ff669835b3476594d93bb/dask/dataframe/multi.py#L231-L233.

With unknown categories (like you have) the pandas dataframes that get there might have different concrete CategoricalDtypes since the observed values differ in each partition. When pandas has multiple categorical dtypes in concat, it upcasts to object dtype. Ideally we would to union the dtypes for all the categoricals that are involved in the merge (and maybe the other columns as well?)

Are you interested in working on this? Pandas has a `pd.api.types.union_categorical` that might be helpful. And the function calling `merge_chunk` can say exactly which columns (and index) are categorical and should be unioned.
Hello Tom!
Thank you for the quick reply. Thank you for explaining the logic on what goes behind the scenes. I would like to respectfully decline as I don't think I know enough to make it work. Nonetheless, thank you for your fantastic work on Dask. I use it every single day!
The fix for https://github.com/dask/dask/issues/6134 might actually fix this too. If we preserve categorical in `merge_chunk` I think things will be OK later on.

Resolving patch:
diff --git a/dask/dataframe/multi.py b/dask/dataframe/multi.py
--- a/dask/dataframe/multi.py
+++ b/dask/dataframe/multi.py
@@ -59,7 +59,7 @@
 from tlz import merge_sorted, unique, first
 import numpy as np
 import pandas as pd
-from pandas.api.types import is_dtype_equal, is_categorical_dtype
+from pandas.api.types import is_dtype_equal, is_categorical_dtype, union_categoricals
 
 from ..base import tokenize, is_dask_collection
 from ..highlevelgraph import HighLevelGraph
@@ -231,7 +231,48 @@ def require(divisions, parts, required=None):
 
 def merge_chunk(lhs, *args, **kwargs):
     empty_index_dtype = kwargs.pop("empty_index_dtype", None)
-    out = lhs.merge(*args, **kwargs)
+    categorical_columns = kwargs.pop("categorical_columns", None)
+
+    rhs, *args = args
+    left_index = kwargs.get("left_index", False)
+    right_index = kwargs.get("right_index", False)
+
+    if categorical_columns is not None and PANDAS_GT_100:
+        for col in categorical_columns:
+            left = None
+            right = None
+
+            if col in lhs:
+                left = lhs[col]
+            elif col == kwargs.get("right_on", None) and left_index:
+                if is_categorical_dtype(lhs.index):
+                    left = lhs.index
+
+            if col in rhs:
+                right = rhs[col]
+            elif col == kwargs.get("left_on", None) and right_index:
+                if is_categorical_dtype(rhs.index):
+                    right = rhs.index
+
+            dtype = "category"
+            if left is not None and right is not None:
+                dtype = union_categoricals(
+                    [left.astype("category").values, right.astype("category").values]
+                ).dtype
+
+            if left is not None:
+                if isinstance(left, pd.Index):
+                    lhs.index = left.astype(dtype)
+                else:
+                    lhs[col] = left.astype(dtype)
+            if right is not None:
+                if isinstance(right, pd.Index):
+                    rhs.index = right.astype(dtype)
+                else:
+                    rhs[col] = right.astype(dtype)
+
+    out = lhs.merge(rhs, *args, **kwargs)
+
     # Workaround pandas bug where if the output result of a merge operation is
     # an empty dataframe, the output index is `int64` in all cases, regardless
     # of input dtypes.
@@ -253,6 +294,7 @@ def merge_indexed_dataframes(lhs, rhs, left_index=True, right_index=True, **kwar
 
     meta = lhs._meta_nonempty.merge(rhs._meta_nonempty, **kwargs)
     kwargs["empty_index_dtype"] = meta.index.dtype
+    kwargs["categorical_columns"] = meta.select_dtypes(include="category").columns
 
     dsk = dict()
     for i, (a, b) in enumerate(parts):
@@ -323,6 +365,8 @@ def hash_join(
     name = "hash-join-" + token
 
     kwargs["empty_index_dtype"] = meta.index.dtype
+    kwargs["categorical_columns"] = meta.select_dtypes(include="category").columns
+
     dsk = {
         (name, i): (apply, merge_chunk, [(lhs2._name, i), (rhs2._name, i)], kwargs)
         for i in range(npartitions)
@@ -334,11 +378,13 @@ def hash_join(
 
 
 def single_partition_join(left, right, **kwargs):
-    # if the merge is perfomed on_index, divisions can be kept, otherwise the
-    # new index will not necessarily correspond the current divisions
+    # if the merge is performed on_index, divisions can be kept, otherwise the
+    # new index will not necessarily correspond with the current divisions
 
     meta = left._meta_nonempty.merge(right._meta_nonempty, **kwargs)
     kwargs["empty_index_dtype"] = meta.index.dtype
+    kwargs["categorical_columns"] = meta.select_dtypes(include="category").columns
+
     name = "merge-" + tokenize(left, right, **kwargs)
     if left.npartitions == 1 and kwargs["how"] in allowed_right:
         left_key = first(left.__dask_keys__())
@@ -531,6 +577,8 @@ def merge(
             suffixes=suffixes,
             indicator=indicator,
         )
+        categorical_columns = meta.select_dtypes(include="category").columns
+
         if merge_indexed_left and left.known_divisions:
             right = rearrange_by_divisions(
                 right, right_on, left.divisions, max_branch, shuffle=shuffle
@@ -555,6 +603,7 @@ def merge(
             suffixes=suffixes,
             indicator=indicator,
             empty_index_dtype=meta.index.dtype,
+            categorical_columns=categorical_columns,
         )
     # Catch all hash join
     else:


Test patch:
diff --git a/dask/dataframe/tests/test_multi.py b/dask/dataframe/tests/test_multi.py
--- a/dask/dataframe/tests/test_multi.py
+++ b/dask/dataframe/tests/test_multi.py
@@ -2141,3 +2141,86 @@ def test_groupby_concat_cudf(engine):
         res_dd = dd.concat([grouped_dd1, grouped_dd2], axis=1)
 
     assert_eq(res_dd.compute().sort_index(), res.sort_index())
+
+
+def test_categorical_join():
+    # https://github.com/dask/dask/issues/6134
+    df = pd.DataFrame(
+        {
+            "join_col": ["a", "a", "b", "b"],
+            "a": [0, 0, 10, 10],
+        }
+    )
+    df2 = pd.DataFrame({"b": [1, 2, 1, 2]}, index=["a", "a", "b", "b"])
+
+    ddf = dd.from_pandas(df, npartitions=2)
+    ddf2 = dd.from_pandas(df2, npartitions=1)
+    ddf["join_col"] = ddf["join_col"].astype("category")
+    ddf2.index = ddf2.index.astype("category")
+
+    expected = ddf.compute().join(ddf2.compute(), on="join_col", how="left")
+
+    actual_dask = ddf.join(ddf2, on="join_col", how="left")
+    assert actual_dask.join_col.dtype == "category"
+
+    actual = actual_dask.compute()
+    if dd._compat.PANDAS_GT_100:
+        assert actual.join_col.dtype == "category"
+        assert assert_eq(expected, actual)
+    else:
+        assert actual.join_col.dtype == "object"
+        assert (expected.values == actual.values).all()
+
+
+def test_categorical_merge_with_columns_missing_from_left():
+    df1 = pd.DataFrame({"A": [0, 1], "B": pd.Categorical(["a", "b"])})
+    df2 = pd.DataFrame({"C": pd.Categorical(["a", "b"])})
+
+    expected = pd.merge(df2, df1, left_index=True, right_on="A")
+
+    ddf1 = dd.from_pandas(df1, npartitions=2)
+    ddf2 = dd.from_pandas(df2, npartitions=2)
+
+    actual = dd.merge(ddf2, ddf1, left_index=True, right_on="A").compute()
+    assert actual.C.dtype == "category"
+    assert actual.B.dtype == "category"
+    assert actual.A.dtype == "int64"
+    assert actual.index.dtype == "int64"
+    assert assert_eq(expected, actual)
+
+
+@pytest.mark.skipif(not dd._compat.PANDAS_GT_0250, reason="Changes in categoricals")
+def test_categorical_merge_with_merge_column_cat_in_one_and_not_other_upcasts():
+    df1 = pd.DataFrame({"A": pd.Categorical([0, 1]), "B": pd.Categorical(["a", "b"])})
+    df2 = pd.DataFrame({"C": pd.Categorical(["a", "b"])})
+
+    expected = pd.merge(df2, df1, left_index=True, right_on="A")
+
+    ddf1 = dd.from_pandas(df1, npartitions=2)
+    ddf2 = dd.from_pandas(df2, npartitions=2)
+
+    actual = dd.merge(ddf2, ddf1, left_index=True, right_on="A").compute()
+    assert actual.C.dtype == "category"
+    assert actual.B.dtype == "category"
+    assert actual.A.dtype == "int64"
+    assert actual.index.dtype == "int64"
+    assert assert_eq(expected, actual)
+
+
+def test_categorical_merge_retains_category_dtype():
+    # https://github.com/dask/dask/issues/6142
+    a = pd.DataFrame({"A": [0, 1, 2, 3], "B": [4, 5, 6, 7]})
+    b = pd.DataFrame({"A": [0, 1, 2, 4], "C": [4, 5, 7, 7]})
+
+    df1 = dd.from_pandas(a, 2)
+    df1["A"] = df1.A.astype("category")
+
+    df2 = dd.from_pandas(b, 2)
+    df2["A"] = df2.A.astype("category")
+
+    actual_dask = df1.merge(df2, on="A")
+    assert actual_dask.A.dtype == "category"
+
+    if dd._compat.PANDAS_GT_100:
+        actual = actual_dask.compute()
+        assert actual.A.dtype == "category"