BUG/API: .merge() and .join() on category dtype columns will now preserve the category dtype when possible

closes #10409

Author: jreback

asv_bench/benchmarks/join_merge.py

@@ -6,7 +6,7 @@
     from pandas import ordered_merge as merge_ordered
 
 
-#----------------------------------------------------------------------
+# ----------------------------------------------------------------------
 # Append
 
 class Append(object):
@@ -35,7 +35,7 @@ def time_append_mixed(self):
         self.mdf1.append(self.mdf2)
 
 
-#----------------------------------------------------------------------
+# ----------------------------------------------------------------------
 # Concat
 
 class Concat(object):
@@ -120,7 +120,7 @@ def time_f_ordered_axis1(self):
         concat(self.frames_f, axis=1, ignore_index=True)
 
 
-#----------------------------------------------------------------------
+# ----------------------------------------------------------------------
 # Joins
 
 class Join(object):
@@ -202,7 +202,7 @@ def time_join_non_unique_equal(self):
         (self.fracofday * self.temp[self.fracofday.index])
 
 
-#----------------------------------------------------------------------
+# ----------------------------------------------------------------------
 # Merges
 
 class Merge(object):
@@ -257,7 +257,31 @@ def time_i8merge(self):
         merge(self.left, self.right, how='outer')
 
 
-#----------------------------------------------------------------------
+class MergeCategoricals(object):
+    goal_time = 0.2
+
+    def setup(self):
+        self.left_object = pd.DataFrame(
+            {'X': np.random.choice(range(0, 10), size=(10000,)),
+             'Y': np.random.choice(['one', 'two', 'three'], size=(10000,))})
+
+        self.right_object = pd.DataFrame(
+            {'X': np.random.choice(range(0, 10), size=(10000,)),
+             'Z': np.random.choice(['jjj', 'kkk', 'sss'], size=(10000,))})
+
+        self.left_cat = self.left_object.assign(
+            Y=self.left_object['Y'].astype('category'))
+        self.right_cat = self.right_object.assign(
+            Z=self.right_object['Z'].astype('category'))
+
+    def time_merge_object(self):
+        merge(self.left_object, self.right_object, on='X')
+
+    def time_merge_cat(self):
+        merge(self.left_cat, self.right_cat, on='X')
+
+
+# ----------------------------------------------------------------------
 # Ordered merge
 
 class MergeOrdered(object):
@@ -332,7 +356,7 @@ def time_multiby(self):
         merge_asof(self.df1e, self.df2e, on='time', by=['key', 'key2'])
 
 
-#----------------------------------------------------------------------
+# ----------------------------------------------------------------------
 # data alignment
 
 class Align(object):

doc/source/whatsnew/v0.20.0.txt

@@ -519,6 +519,7 @@ Other API Changes
 - The :func:`pd.read_gbq` method now stores ``INTEGER`` columns as ``dtype=object`` if they contain ``NULL`` values. Otherwise they are stored as ``int64``. This prevents precision lost for integers greather than 2**53. Furthermore ``FLOAT`` columns with values above 10**4 are no longer casted to ``int64`` which also caused precision loss (:issue:`14064`, :issue:`14305`).
 - Reorganization of timeseries development tests (:issue:`14854`)
 - Specific support for ``copy.copy()`` and ``copy.deepcopy()`` functions on NDFrame objects (:issue:`15444`)
+- ``.merge()`` and ``.join()`` on ``category`` dtype columns will now preserve the category dtype when possible (:issue:`10409`)
 
 .. _whatsnew_0200.deprecations:
 
@@ -567,7 +568,7 @@ Performance Improvements
 - Improved performance and reduced memory when indexing with a ``MultiIndex`` (:issue:`15245`)
 - When reading buffer object in ``read_sas()`` method without specified format, filepath string is inferred rather than buffer object. (:issue:`14947`)
 - Improved performance of `rank()` for categorical data (:issue:`15498`)
-
+- Improved performance of merge/join on ``category`` columns (:issue:`10409`)
 
 
 .. _whatsnew_0200.bug_fixes:

pandas/core/internals.py

@@ -5228,6 +5228,8 @@ def get_reindexed_values(self, empty_dtype, upcasted_na):
                 # External code requested filling/upcasting, bool values must
                 # be upcasted to object to avoid being upcasted to numeric.
                 values = self.block.astype(np.object_).values
+            elif self.block.is_categorical:
+                values = self.block.values
             else:
                 # No dtype upcasting is done here, it will be performed during
                 # concatenation itself.

pandas/tests/test_categorical.py

@@ -4097,9 +4097,12 @@ def test_merge(self):
         expected = df.copy()
 
         # object-cat
+        # note that we propogate the category
+        # because we don't have any matching rows
         cright = right.copy()
         cright['d'] = cright['d'].astype('category')
         result = pd.merge(left, cright, how='left', left_on='b', right_on='c')
+        expected['d'] = expected['d'].astype('category', categories=['null'])
         tm.assert_frame_equal(result, expected)
 
         # cat-object

pandas/tests/tools/test_merge.py

@@ -11,6 +11,8 @@
 from pandas.tools.concat import concat
 from pandas.tools.merge import merge, MergeError
 from pandas.util.testing import assert_frame_equal, assert_series_equal
+from pandas.types.dtypes import CategoricalDtype
+from pandas.types.common import is_categorical_dtype, is_object_dtype
 from pandas import DataFrame, Index, MultiIndex, Series, Categorical
 import pandas.util.testing as tm
 
@@ -1024,38 +1026,6 @@ def test_left_join_index_multi_match(self):
         expected.index = np.arange(len(expected))
         tm.assert_frame_equal(result, expected)
 
-    def test_join_multi_dtypes(self):
-
-        # test with multi dtypes in the join index
-        def _test(dtype1, dtype2):
-            left = DataFrame({'k1': np.array([0, 1, 2] * 8, dtype=dtype1),
-                              'k2': ['foo', 'bar'] * 12,
-                              'v': np.array(np.arange(24), dtype=np.int64)})
-
-            index = MultiIndex.from_tuples([(2, 'bar'), (1, 'foo')])
-            right = DataFrame(
-                {'v2': np.array([5, 7], dtype=dtype2)}, index=index)
-
-            result = left.join(right, on=['k1', 'k2'])
-
-            expected = left.copy()
-
-            if dtype2.kind == 'i':
-                dtype2 = np.dtype('float64')
-            expected['v2'] = np.array(np.nan, dtype=dtype2)
-            expected.loc[(expected.k1 == 2) & (expected.k2 == 'bar'), 'v2'] = 5
-            expected.loc[(expected.k1 == 1) & (expected.k2 == 'foo'), 'v2'] = 7
-
-            tm.assert_frame_equal(result, expected)
-
-            result = left.join(right, on=['k1', 'k2'], sort=True)
-            expected.sort_values(['k1', 'k2'], kind='mergesort', inplace=True)
-            tm.assert_frame_equal(result, expected)
-
-        for d1 in [np.int64, np.int32, np.int16, np.int8, np.uint8]:
-            for d2 in [np.int64, np.float64, np.float32, np.float16]:
-                _test(np.dtype(d1), np.dtype(d2))
-
     def test_left_merge_na_buglet(self):
         left = DataFrame({'id': list('abcde'), 'v1': randn(5),
                           'v2': randn(5), 'dummy': list('abcde'),
@@ -1242,3 +1212,140 @@ def f():
         def f():
             household.join(log_return, how='outer')
         self.assertRaises(NotImplementedError, f)
+
+
+class TestMergeDtypes(tm.TestCase):
+
+    def setUp(self):
+
+        self.df = DataFrame(
+            {'A': ['foo', 'bar'],
+             'B': Series(['foo', 'bar']).astype('category'),
+             'C': [1, 2],
+             'D': [1.0, 2.0],
+             'E': Series([1, 2], dtype='uint64'),
+             'F': Series([1, 2], dtype='int32')})
+
+    def test_different(self):
+
+        # we expect differences by kind
+        # to be ok, while other differences should return object
+
+        left = self.df
+        for col in self.df.columns:
+            right = DataFrame({'A': self.df[col]})
+            result = pd.merge(left, right, on='A')
+            assert is_object_dtype(result.A.dtype)
+
+    def test_join_multi_dtypes(self):
+
+        # test with multi dtypes in the join index
+        def _test(dtype1, dtype2):
+            left = DataFrame({'k1': np.array([0, 1, 2] * 8, dtype=dtype1),
+                              'k2': ['foo', 'bar'] * 12,
+                              'v': np.array(np.arange(24), dtype=np.int64)})
+
+            index = MultiIndex.from_tuples([(2, 'bar'), (1, 'foo')])
+            right = DataFrame(
+                {'v2': np.array([5, 7], dtype=dtype2)}, index=index)
+
+            result = left.join(right, on=['k1', 'k2'])
+
+            expected = left.copy()
+
+            if dtype2.kind == 'i':
+                dtype2 = np.dtype('float64')
+            expected['v2'] = np.array(np.nan, dtype=dtype2)
+            expected.loc[(expected.k1 == 2) & (expected.k2 == 'bar'), 'v2'] = 5
+            expected.loc[(expected.k1 == 1) & (expected.k2 == 'foo'), 'v2'] = 7
+
+            tm.assert_frame_equal(result, expected)
+
+            result = left.join(right, on=['k1', 'k2'], sort=True)
+            expected.sort_values(['k1', 'k2'], kind='mergesort', inplace=True)
+            tm.assert_frame_equal(result, expected)
+
+        for d1 in [np.int64, np.int32, np.int16, np.int8, np.uint8]:
+            for d2 in [np.int64, np.float64, np.float32, np.float16]:
+                _test(np.dtype(d1), np.dtype(d2))
+
+
+class TestMergeCategorical(tm.TestCase):
+    _multiprocess_can_split_ = True
+
+    def setUp(self):
+        np.random.seed(1234)
+        self.left = DataFrame(
+            {'X': Series(np.random.choice(
+                ['foo', 'bar'],
+                size=(10,))).astype('category', categories=['foo', 'bar']),
+             'Y': np.random.choice(['one', 'two', 'three'], size=(10,))})
+        self.right = pd.DataFrame(
+            {'X': Series(['foo', 'bar']).astype('category',
+                                                categories=['foo', 'bar']),
+             'Z': [1, 2]})
+
+    def test_identical(self):
+        # merging on the same, should preserve dtypes
+        merged = pd.merge(self.left, self.left, on='X')
+        result = merged.dtypes.sort_index()
+        expected = Series([CategoricalDtype(),
+                           np.dtype('O'),
+                           np.dtype('O')],
+                          index=['X', 'Y_x', 'Y_y'])
+        assert_series_equal(result, expected)
+
+    def test_basic(self):
+        # we have matching Categorical dtypes in X
+        # so should preserve the merged column
+        merged = pd.merge(self.left, self.right, on='X')
+        result = merged.dtypes.sort_index()
+        expected = Series([CategoricalDtype(),
+                           np.dtype('O'),
+                           np.dtype('int64')],
+                          index=['X', 'Y', 'Z'])
+        assert_series_equal(result, expected)
+
+    def test_other_columns(self):
+        # non-merge columns should preserve if possible
+        left = self.left
+        right = self.right.assign(Z=self.right.Z.astype('category'))
+
+        merged = pd.merge(left, right, on='X')
+        result = merged.dtypes.sort_index()
+        expected = Series([CategoricalDtype(),
+                           np.dtype('O'),
+                           CategoricalDtype()],
+                          index=['X', 'Y', 'Z'])
+        assert_series_equal(result, expected)
+
+        # categories are preserved
+        assert left.X.values.is_dtype_equal(merged.X.values)
+        assert right.Z.values.is_dtype_equal(merged.Z.values)
+
+    def test_dtype_on_merged_different(self):
+        # our merging columns, X now has 2 different dtypes
+        # so we must be object as a result
+        left = self.left
+
+        for change in [lambda x: x,
+                       lambda x: x.astype('category',
+                                          categories=['bar', 'foo']),
+                       lambda x: x.astype('category',
+                                          categories=['foo', 'bar', 'bah']),
+                       lambda x: x.astype('category', ordered=True)]:
+            for how in ['inner', 'outer', 'left', 'right']:
+
+                X = change(self.right.X.astype('object'))
+                right = self.right.assign(X=X)
+                assert is_categorical_dtype(left.X.values)
+                assert not left.X.values.is_dtype_equal(right.X.values)
+
+                merged = pd.merge(left, right, on='X', how=how)
+
+                result = merged.dtypes.sort_index()
+                expected = Series([np.dtype('O'),
+                                   np.dtype('O'),
+                                   np.dtype('int64')],
+                                  index=['X', 'Y', 'Z'])
+                assert_series_equal(result, expected)

pandas/tests/tools/test_merge_asof.py

@@ -147,6 +147,7 @@ def test_basic_categorical(self):
         trades.ticker = trades.ticker.astype('category')
         quotes = self.quotes.copy()
         quotes.ticker = quotes.ticker.astype('category')
+        expected.ticker = expected.ticker.astype('category')
 
         result = merge_asof(trades, quotes,
                             on='time',

pandas/tests/types/test_common.py

@@ -39,16 +39,33 @@ def test_period_dtype(self):
 
 
 def test_dtype_equal():
-    assert is_dtype_equal(np.int64, np.int64)
-    assert not is_dtype_equal(np.int64, np.float64)
 
-    p1 = PeriodDtype('D')
-    p2 = PeriodDtype('D')
-    assert is_dtype_equal(p1, p2)
-    assert not is_dtype_equal(np.int64, p1)
+    dtypes = dict(dt_tz=pandas_dtype('datetime64[ns, US/Eastern]'),
+                  dt=pandas_dtype('datetime64[ns]'),
+                  td=pandas_dtype('timedelta64[ns]'),
+                  p=PeriodDtype('D'),
+                  i=np.int64,
+                  f=np.float64,
+                  o=np.object)
 
-    p3 = PeriodDtype('2D')
-    assert not is_dtype_equal(p1, p3)
+    # match equal to self, but not equal to other
+    for name, dtype in dtypes.items():
+        assert is_dtype_equal(dtype, dtype)
 
-    assert not DatetimeTZDtype.is_dtype(np.int64)
-    assert not PeriodDtype.is_dtype(np.int64)
+        for name2, dtype2 in dtypes.items():
+            if name != name2:
+                assert not is_dtype_equal(dtype, dtype2)
+
+    # we are strict on kind equality
+    for dtype in [np.int8, np.int16, np.int32]:
+        assert not is_dtype_equal(dtypes['i'], dtype)
+
+    for dtype in [np.float32]:
+        assert not is_dtype_equal(dtypes['f'], dtype)
+
+    # strict w.r.t. PeriodDtype
+    assert not is_dtype_equal(dtypes['p'], PeriodDtype('2D'))
+
+    # strict w.r.t. datetime64
+    assert not is_dtype_equal(dtypes['dt_tz'],
+                              pandas_dtype('datetime64[ns, CET]'))