[REVIEW] Add support merging b/w categorical data

python/cudf/cudf/api/__init__.py

@@ -1 +1,3 @@
 # Copyright (c) 2020, NVIDIA CORPORATION.
+
+from cudf.api import types

python/cudf/cudf/api/types/__init__.py

@@ -0,0 +1,3 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.
+
+from cudf.api.types.categoricals import _union_categoricals

python/cudf/cudf/api/types/categoricals.py

@@ -0,0 +1,40 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.
+
+from __future__ import annotations
+
+from typing import List, Union
+
+import cudf
+
+ParentType = Union["cudf.Series", "cudf.Index"]
+
+
+def _union_categoricals(
+    to_union: List[Union[cudf.Series, cudf.Index]],
+    sort_categories: bool = False,
+    ignore_order: bool = False,
+):
+    """
+    This is an internal API which combines categorical data.
+    """
+    # TODO(s) in the order specified :
+    # 1. The return type needs to be changed
+    #    to cudf.Categorical once it is implemented.
+    # 2. Make this API public (i.e., to resemble
+    #    pd.api.types.union_categoricals)
+
+    if ignore_order:
+        raise TypeError("ignore_order is not yet implemented")
+
+    result_col = cudf.core.column.CategoricalColumn._concat(
+        [obj._column for obj in to_union]
+    )
+    if sort_categories:
+        sorted_categories = result_col.categories.sort_by_values(
+            ascending=True
+        )[0]
+        result_col = result_col.cat().reorder_categories(
+            new_categories=sorted_categories
+        )
+
+    return cudf.Index(result_col)

python/dask_cudf/dask_cudf/backends.py

@@ -244,6 +244,21 @@ def is_categorical_dtype_cudf(obj):
     return cudf.utils.dtypes.is_categorical_dtype(obj)
 
 
+try:
+    from dask.dataframe.dispatch import union_categoricals_dispatch
+
+    @union_categoricals_dispatch.register((cudf.Series, cudf.Index))
+    def union_categoricals_cudf(
+        to_union, sort_categories=False, ignore_order=False
+    ):
+        return cudf.api.types._union_categoricals(
+            to_union, sort_categories=False, ignore_order=False
+        )
+
+
+except ImportError:
+    pass
+
 try:
 
     from dask.dataframe.core import group_split_dispatch, hash_object_dispatch

python/dask_cudf/dask_cudf/tests/test_core.py

@@ -780,6 +780,46 @@ def test_index_map_partitions():
     dd.assert_eq(mins_pd, mins_gd)
 
 
+def test_merging_categorical_columns():
+    try:
+        from dask.dataframe.dispatch import (  # noqa: F401
+            union_categoricals_dispatch,
+        )
+    except ImportError:
+        pytest.skip(
+            "need a version of dask that has union_categoricals_dispatch"
+        )
+
+    df_1 = cudf.DataFrame(
+        {"id_1": [0, 1, 2, 3], "cat_col": ["a", "b", "f", "f"]}
+    )
+
+    ddf_1 = dgd.from_cudf(df_1, npartitions=2)
+
+    ddf_1 = dd.categorical.categorize(ddf_1, columns=["cat_col"])
+
+    df_2 = cudf.DataFrame(
+        {"id_2": [111, 112, 113], "cat_col": ["g", "h", "f"]}
+    )
+
+    ddf_2 = dgd.from_cudf(df_2, npartitions=2)
+
+    ddf_2 = dd.categorical.categorize(ddf_2, columns=["cat_col"])
+    expected = cudf.DataFrame(
+        {
+            "id_1": [2, 3],
+            "cat_col": cudf.Series(
+                ["f", "f"],
+                dtype=cudf.CategoricalDtype(
+                    categories=["a", "b", "f", "g", "h"], ordered=False
+                ),
+            ),
+            "id_2": [113, 113],
+        }
+    )
+    dd.assert_eq(ddf_1.merge(ddf_2), expected)
+
+
 def test_correct_meta():
     # Need these local imports in this specific order.
     # For context: https://github.com/rapidsai/cudf/issues/7946