Repository: pyjanitor-devs/pyjanitor

Issue:
[ENH] Modify dataframe list with union of their categoricals
The problem: vertical `pd.concat([df1, df2, ...])` converts categoricals to strings if the categories do not perfectly overlap between all dataframes. 

The solution: use pandas' `union_categoricals` to modify the original dataframes to add all categories to all dataframes prior to the concatenation.

How to use `union_categoricals`:

From https://stackoverflow.com/questions/45639350/retaining-categorical-dtype-upon-dataframe-concatenation :

```python
from pandas.api.types import union_categoricals

uc = union_categoricals([df1.x,df2.x])
df1.x = pd.Categorical( df1.x, categories=uc.categories )
df2.x = pd.Categorical( df2.x, categories=uc.categories )
```

The final function could iterate over all columns whose names match and are categorical. It can alternatively take an explicit list of particular columns.


Discussion:

Resolving patch:
diff --git a/janitor/functions.py b/janitor/functions.py
--- a/janitor/functions.py
+++ b/janitor/functions.py
@@ -11,6 +11,7 @@
 import numpy as np
 import pandas as pd
 import pandas_flavor as pf
+from pandas.api.types import union_categoricals
 from scipy.stats import mode
 from sklearn.preprocessing import LabelEncoder
 
@@ -27,6 +28,96 @@
 )
 
 
+def unionize_dataframe_categories(
+    *dataframes, column_names: Union[str, Iterable[str], Any] = None
+) -> List[pd.DataFrame]:
+    """
+    Given a group of dataframes which contain some categorical columns, for
+    each categorical column present, find all the possible categories across
+    all the dataframes which have that column.
+    Update each dataframes' corresponding column with a new categorical object
+    that contains the original data
+    but has labels for all the possible categories from all dataframes.
+    This is useful when concatenating a list of dataframes which all have the
+    same categorical columns into one dataframe.
+
+    If, for a given categorical column, all input dataframes do not have at
+    least one instance of all the possible categories,
+    Pandas will change the output dtype of that column from ``category`` to
+    ``object``, losing out on dramatic speed gains you get from the former
+    format.
+
+    Usage example for concatentation of categorical column-containing
+    dataframes:
+
+    Instead of:
+
+    .. code-block:: python
+
+        concatenated_df = pd.concat([df1, df2, df3], ignore_index=True)
+
+    which in your case has resulted in ``category`` -> ``object`` conversion,
+    use:
+
+    .. code-block:: python
+
+        unionized_dataframes = unionize_dataframe_categories(df1, df2, df2)
+        concatenated_df = pd.concat(unionized_dataframes, ignore_index=True)
+
+    :param dataframes: The dataframes you wish to unionize the categorical
+        objects for.
+    :param column_names: If supplied, only unionize this subset of columns.
+    :returns: A list of the category-unioned dataframes in the same order they
+        were provided.
+    """
+
+    if any(not isinstance(df, pd.DataFrame) for df in dataframes):
+        raise TypeError("Inputs must all be dataframes.")
+
+    if column_names is None:
+        # Find all columns across all dataframes that are categorical
+
+        column_names = set()
+
+        for df in dataframes:
+            column_names = column_names.union(
+                [
+                    column_name
+                    for column_name in df.columns
+                    if isinstance(df[column_name].dtype, pd.CategoricalDtype)
+                ]
+            )
+    elif isinstance(column_names, str):
+        column_names = [column_names]
+
+    # For each categorical column, find all possible values across the DFs
+
+    category_unions = {
+        column_name: union_categoricals(
+            [df[column_name] for df in dataframes if column_name in df.columns]
+        )
+        for column_name in column_names
+    }
+
+    # Make a shallow copy of all DFs and modify the categorical columns
+    # such that they can encode the union of all possible categories for each.
+
+    refactored_dfs = []
+
+    for df in dataframes:
+        df = df.copy(deep=False)
+
+        for column_name, categorical in category_unions.items():
+            if column_name in df.columns:
+                df[column_name] = pd.Categorical(
+                    df[column_name], categories=categorical.categories
+                )
+
+        refactored_dfs.append(df)
+
+    return refactored_dfs
+
+
 @pf.register_dataframe_method
 def move(
     df: pd.DataFrame,


Test patch:
diff --git a/tests/functions/test_unionize_dataframe_categories.py b/tests/functions/test_unionize_dataframe_categories.py
new file mode 100644
--- /dev/null
+++ b/tests/functions/test_unionize_dataframe_categories.py
@@ -0,0 +1,145 @@
+import pandas as pd
+import pytest
+
+import janitor
+
+
+@pytest.fixture
+def uniontest_df1():
+    return pd.DataFrame(
+        dict(
+            jerbs=pd.Categorical(["fireman", "programmer", "astronaut"]),
+            fruits=pd.Categorical(["apple", "banana", "orange"]),
+            data=[1, 2, 3],
+            df1_exclusive="somedata",
+        )
+    )
+
+
+@pytest.fixture
+def uniontest_df2():
+    return pd.DataFrame(
+        dict(
+            jerbs=pd.Categorical(["fireman", "actor", "astronaut"]),
+            fruits=pd.Categorical(["grape", "strawberry", "cherry"]),
+            data=[4, 5, 6],
+            df2_exclusive="otherdata",
+            animals=pd.Categorical(["bear", "tiger", "sloth"]),
+        )
+    )
+
+
+@pytest.fixture
+def uniontest_df3():
+    return pd.DataFrame(
+        dict(
+            jerbs=pd.Categorical(["salesman", "actor", "programmer"]),
+            fruits=pd.Categorical(["grape", "banana", "cherry"]),
+            data=[7, 8, 9],
+            df3_exclusive="evenmoredata",
+            animals=pd.Categorical(["bear", "capybara", "sloth"]),
+        )
+    )
+
+
+def test_unionize_dataframe_categories_type(uniontest_df1):
+    with pytest.raises(TypeError):
+        janitor.unionize_dataframe_categories(
+            uniontest_df1, uniontest_df1["jerbs"]
+        )
+
+
+def test_unionize_dataframe_categories(
+    uniontest_df1, uniontest_df2, uniontest_df3
+):
+    udf1, udf2, udf3 = janitor.unionize_dataframe_categories(
+        uniontest_df1, uniontest_df2, uniontest_df3
+    )
+
+    # test categories were unioned properly
+
+    assert set(udf1["jerbs"].dtype.categories) == set(
+        udf2["jerbs"].dtype.categories
+    )
+
+    assert set(udf1["jerbs"].dtype.categories) == set(
+        udf3["jerbs"].dtype.categories
+    )
+
+    assert set(udf1["fruits"].dtype.categories) == set(
+        udf3["fruits"].dtype.categories
+    )
+
+    assert set(udf1["fruits"].dtype.categories) == set(
+        udf3["fruits"].dtype.categories
+    )
+
+    assert set(udf2["animals"].dtype.categories) == set(
+        udf3["animals"].dtype.categories
+    )
+
+    # test columns did not bleed in
+
+    assert "df2_exclusive" not in udf1.columns
+    assert "df2_exclusive" not in udf3.columns
+    assert "df1_exclusive" not in udf2.columns
+    assert "df1_exclusive" not in udf3.columns
+    assert "df3_exclusive" not in udf1.columns
+    assert "df3_exclusive" not in udf2.columns
+
+    assert "animals" not in udf1.columns
+
+    # test that pd.concat now does not destroy categoricals
+
+    # NOTE: 'animals' column will not be preserved as categorical after concat
+    # because it is not present in df1. Instead, this column will be set to
+    # object with NaN filled in for df1 missing values.
+
+    udf = pd.concat([udf1, udf2, udf3], ignore_index=True)
+
+    assert isinstance(udf["jerbs"].dtype, pd.CategoricalDtype)
+    assert isinstance(udf["fruits"].dtype, pd.CategoricalDtype)
+
+    # test that the data is the same
+
+    pd.testing.assert_frame_equal(udf1, uniontest_df1, check_categorical=False)
+    pd.testing.assert_frame_equal(udf2, uniontest_df2, check_categorical=False)
+    pd.testing.assert_frame_equal(udf3, uniontest_df3, check_categorical=False)
+
+
+def test_unionize_dataframe_categories_original_preservation(
+    uniontest_df1, uniontest_df2
+):
+    udf1, udf2 = janitor.unionize_dataframe_categories(
+        uniontest_df1, uniontest_df2
+    )
+
+    assert not (
+        set(uniontest_df1["fruits"].dtype.categories)
+        == set(udf1["fruits"].dtype.categories)
+    )
+
+
+def test_unionize_dataframe_categories_single(
+    uniontest_df1, uniontest_df2, uniontest_df3
+):
+    udf1, udf2, udf3 = janitor.unionize_dataframe_categories(
+        uniontest_df1, uniontest_df2, uniontest_df3, column_names="fruits"
+    )
+
+    # check that fruits did get unionized
+
+    assert set(udf1["fruits"].dtype.categories) == set(
+        udf2["fruits"].dtype.categories
+    )
+
+    assert set(udf1["fruits"].dtype.categories) == set(
+        udf3["fruits"].dtype.categories
+    )
+
+    # check that jerbs did not when we didn't want it to
+
+    assert not (
+        set(udf1["jerbs"].dtype.categories)
+        == set(udf2["jerbs"].dtype.categories)
+    )