Preserve dtype in merge keys when possible

Author: Mike Kelly

pandas/tools/merge.py

@@ -280,19 +280,28 @@ def _indicator_post_merge(self, result):
         return result
 
     def _maybe_add_join_keys(self, result, left_indexer, right_indexer):
-        # insert group keys
+
+        consolidate = False
+
+        left_has_missing = None
+        right_has_missing = None
 
         keys = zip(self.join_names, self.left_on, self.right_on)
         for i, (name, lname, rname) in enumerate(keys):
             if not _should_fill(lname, rname):
                 continue
 
+            take_left, take_right = None, None
+
             if name in result:
+<<<<<<< HEAD
                 key_indexer = result.columns.get_loc(name)
+=======
+>>>>>>> e79b978... Preserve dtype in merge keys when possible
 
                 if left_indexer is not None and right_indexer is not None:
-
                     if name in self.left:
+<<<<<<< HEAD
                         if len(self.left) == 0:
                             continue
 
@@ -316,19 +325,71 @@ def _maybe_add_join_keys(self, result, left_indexer, right_indexer):
                         result.iloc[na_indexer, key_indexer] = (
                             algos.take_1d(self.left_join_keys[i],
                                           left_na_indexer))
+=======
+
+                        if left_has_missing is None:
+                            left_has_missing = any(left_indexer == -1)
+
+                        if left_has_missing:
+                            take_right = self.right_join_keys[i]
+
+                            if result[name].dtype != self.left[name].dtype:
+                                take_left = self.left[name].values
+
+                    elif name in self.right:
+
+                        if right_has_missing is None:
+                            right_has_missing = any(right_indexer == -1)
+
+                        if right_has_missing:
+                            take_left = self.left_join_keys[i]
+
+                            if result[name].dtype != self.right[name].dtype:
+                                take_right = self.right[name].values
+
+>>>>>>> e79b978... Preserve dtype in merge keys when possible
             elif left_indexer is not None \
                     and isinstance(self.left_join_keys[i], np.ndarray):
 
-                if name is None:
-                    name = 'key_%d' % i
+                take_left = self.left_join_keys[i]
+                take_right = self.right_join_keys[i]
+
+            if take_left is not None or take_right is not None:
+
+                if take_left is None:
+                    lvals = result[name].values
+                else:
+                    lfill = take_left.dtype.type()
+                    lvals = com.take_1d(take_left, left_indexer, fill_value=lfill)
+
+                if take_right is None:
+                    rvals = result[name].values
+                else:
+                    rfill = take_right.dtype.type()
+                    rvals = com.take_1d(take_right, right_indexer, fill_value=rfill)
+
+                key_col = np.where(left_indexer != -1, lvals, rvals)
+
+                if name in result:
+                    if result[name].dtype != key_col.dtype:
+                        consolidate = True
+                    result[name] = key_col
+                else:
+                    result.insert(i, name or 'key_%d' % i, key_col)
+                    consolidate = True
 
+<<<<<<< HEAD
                 # a faster way?
                 key_col = algos.take_1d(self.left_join_keys[i], left_indexer)
                 na_indexer = (left_indexer == -1).nonzero()[0]
                 right_na_indexer = right_indexer.take(na_indexer)
                 key_col.put(na_indexer, algos.take_1d(self.right_join_keys[i],
                                                       right_na_indexer))
                 result.insert(i, name, key_col)
+=======
+        if consolidate:
+            result.consolidate(inplace=True)
+>>>>>>> e79b978... Preserve dtype in merge keys when possible
 
     def _get_join_info(self):
         left_ax = self.left._data.axes[self.axis]

pandas/tools/tests/test_merge.py

@@ -507,8 +507,8 @@ def test_join_many_non_unique_index(self):
 
         result = result.reset_index()
 
-        result['a'] = result['a'].astype(np.float64)
-        result['b'] = result['b'].astype(np.float64)
+        # result['a'] = result['a'].astype(np.float64)
+        # result['b'] = result['b'].astype(np.float64)
 
         assert_frame_equal(result, expected.ix[:, result.columns])
 
@@ -1033,6 +1033,7 @@ def test_overlapping_columns_error_message(self):
         df2.columns = ['key1', 'foo', 'foo']
         self.assertRaises(ValueError, merge, df, df2)
 
+<<<<<<< HEAD
     def test_merge_on_datetime64tz(self):
 
         # GH11405
@@ -1426,6 +1427,27 @@ def test_indicator(self):
         test5 = df3.merge(df4, on=['col1', 'col2'],
                           how='outer', indicator=True)
         assert_frame_equal(test5, hand_coded_result)
+=======
+    def test_merge_join_key_dtype_cast(self):
+        # #8596
+
+        df1 = DataFrame({'key': [1], 'v1': [10]})
+        df2 = DataFrame({'key': [2], 'v1': [20]})
+        df = merge(df1, df2, how='outer')
+        self.assertEqual(df['key'].dtype, 'int64')
+
+        df1 = DataFrame({'key': [True], 'v1': [1]})
+        df2 = DataFrame({'key': [False],'v1': [0]})
+        df = merge(df1, df2, how='outer')
+        self.assertEqual(df['key'].dtype, 'bool')
+
+        df1 = DataFrame({'val': [1]})
+        df2 = DataFrame({'val': [2]})
+        lkey = np.array([1])
+        rkey = np.array([2])
+        df = merge(df1, df2, left_on=lkey, right_on=rkey, how='outer')
+        self.assertEqual(df['key_0'].dtype, 'int64')
+>>>>>>> e79b978... Preserve dtype in merge keys when possible
 
 
 def _check_merge(x, y):