ENH: be a bit more clever to avoid consolidation with join keys, GH #733

Author: wesm

pandas/core/internals.py

@@ -802,12 +802,8 @@ def merge(self, other, lsuffix=None, rsuffix=None):
 
         return BlockManager(consolidated, new_axes)
 
-    def _maybe_rename_join(self, other, lsuffix, rsuffix, exclude=None,
-                           copydata=True):
+    def _maybe_rename_join(self, other, lsuffix, rsuffix, copydata=True):
         to_rename = self.items.intersection(other.items)
-        if exclude is not None and len(exclude) > 0:
-            to_rename = to_rename - exclude
-
         if len(to_rename) > 0:
             if not lsuffix and not rsuffix:
                 raise Exception('columns overlap: %s' % to_rename)
@@ -822,7 +818,6 @@ def rrenamer(x):
                     return '%s%s' % (x, rsuffix)
                 return x
 
-            # XXX: COPIES DATA!
             this = self.rename_items(lrenamer, copydata=copydata)
             other = other.rename_items(rrenamer, copydata=copydata)
         else:

pandas/tools/merge.py

@@ -97,13 +97,27 @@ def _maybe_add_join_keys(self, result, left_indexer, right_indexer):
 
         # insert group keys
         for i, name in enumerate(self.join_names):
-            # a faster way?
-            key_col = com.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, com.take_1d(self.right_join_keys[i],
-                                                right_na_indexer))
-            result.insert(i, name, key_col)
+            if name in result:
+                key_col = result[name]
+
+                if name in self.left:
+                    na_indexer = (left_indexer == -1).nonzero()[0]
+                    right_na_indexer = right_indexer.take(na_indexer)
+                    key_col.put(na_indexer, com.take_1d(self.right_join_keys[i],
+                                                        right_na_indexer))
+                else:
+                    na_indexer = (right_indexer == -1).nonzero()[0]
+                    left_na_indexer = right_indexer.take(na_indexer)
+                    key_col.put(na_indexer, com.take_1d(self.left_join_keys[i],
+                                                        left_na_indexer))
+            else:
+                # a faster way?
+                key_col = com.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, com.take_1d(self.right_join_keys[i],
+                                                    right_na_indexer))
+                result.insert(i, name, key_col)
 
     def _get_join_info(self):
         left_ax = self.left._data.axes[self.axis]
@@ -144,17 +158,8 @@ def _get_merge_data(self):
         """
         ldata, rdata = self.left._data, self.right._data
         lsuf, rsuf = self.suffixes
-        exclude_names = set(x for x in self.join_names if x is not None)
-        if self.left_on is not None:
-            exclude_names -= set(c.name if hasattr(c, 'name') else c
-                                 for c in self.left_on)
-        if self.right_on is not None:
-            exclude_names -= set(c.name if hasattr(c, 'name') else c
-                                 for c in self.right_on)
         ldata, rdata = ldata._maybe_rename_join(rdata, lsuf, rsuf,
-                                                exclude=exclude_names,
                                                 copydata=False)
-
         return ldata, rdata
 
     def _get_merge_keys(self):
@@ -172,8 +177,6 @@ def _get_merge_keys(self):
         left_keys, right_keys
         """
         # Hm, any way to make this logic less complicated??
-        join_names = []
-
         if (self.on is None and self.left_on is None
             and self.right_on is None):
 
@@ -198,18 +201,56 @@ def _get_merge_keys(self):
             self.left_on = self.right_on = self.on
             self.drop_keys = True
 
-        # this is a touch kludgy, but accomplishes the goal
-        left_keys = None
-        if self.left_on is not None:
-            self.left, left_keys, left_names = \
-                _get_keys(self.left, self.left_on, drop=self.drop_keys)
-            join_names = left_names
-
-        right_keys = None
-        if self.right_on is not None:
-            self.right, right_keys, right_names = \
-                _get_keys(self.right, self.right_on, drop=self.drop_keys)
-            join_names = right_names
+        left_keys = []
+        right_keys = []
+        join_names = []
+        left_drop, right_drop = [], []
+        left, right = self.left, self.right
+
+        is_lkey = lambda x: isinstance(x, np.ndarray) and len(x) == len(left)
+        is_rkey = lambda x: isinstance(x, np.ndarray) and len(x) == len(right)
+
+        # ugh, spaghetti re #733
+        if _any(self.left_on) and _any(self.right_on):
+            for lk, rk in zip(self.left_on, self.right_on):
+                if is_lkey(lk):
+                    left_keys.append(lk)
+                    if is_rkey(rk):
+                        right_keys.append(rk)
+                        join_names.append(None)  # what to do?
+                    else:
+                        right_keys.append(right[rk].values)
+                        join_names.append(rk)
+                else:
+                    if not is_rkey(rk):
+                        right_keys.append(right[rk].values)
+                        if lk == rk:
+                            right_drop.append(rk)
+                    else:
+                        right_keys.append(rk)
+                    left_keys.append(left[lk].values)
+                    join_names.append(lk)
+        elif _any(self.left_on):
+            for k in self.left_on:
+                if is_lkey(k):
+                    left_keys.append(k)
+                    join_names.append(None)
+                else:
+                    left_keys.append(left[k].values)
+                    join_names.append(k)
+        elif _any(self.right_on):
+            for k in self.right_on:
+                if is_rkey(k):
+                    right_keys.append(k)
+                    join_names.append(None)
+                else:
+                    right_keys.append(right[k].values)
+                    join_names.append(k)
+
+        if right_drop:
+            self.right = self.right.drop(right_drop, axis=1)
+        if left_drop:
+            self.left = self.left.drop(left_drop, axis=1)
 
         return left_keys, right_keys, join_names
 
@@ -271,7 +312,8 @@ def _get_group_keys(self):
                              sort=self.sort)
         return left_group_key, right_group_key, max_groups
 
-def _get_keys(frame, on, drop=False):
+
+def _get_join_keys(left, right, left_on, right_on, drop=False):
     to_drop = []
     keys = []
     names = []
@@ -993,3 +1035,5 @@ def _all_indexes_same(indexes):
             return False
     return True
 
+def _any(x):
+    return x is not None and len(x) > 0

pandas/tools/tests/test_merge.py

@@ -180,25 +180,9 @@ def test_handle_overlap_arbitrary_key(self):
         joined = merge(self.df, self.df2,
                        left_on='key2', right_on='key1',
                        suffixes=['.foo', '.bar'])
-
         self.assert_('key1.foo' in joined)
         self.assert_('key2.bar' in joined)
 
-        # result = merge(self.df, self.df2, on='key1')
-        # left_on = self.df['key2'].copy()
-        # left_on.name = 'baz'
-        # right_on = self.df2['key1'].copy()
-        # right_on.name = 'baz'
-
-        # grouped = self.df2.groupby('key1').mean()
-        # self.assert_('key2' in grouped)
-
-        # joined = merge(self.df, grouped, left_on='key1',
-        #                right_index=True, suffixes=['.foo', '.bar'])
-        # foo
-        # self.assert_('key2.foo' in joined)
-        # self.assert_('key2.bar' in joined)
-
     def test_merge_common(self):
         joined = merge(self.df, self.df2)
         exp = merge(self.df, self.df2, on=['key1', 'key2'])
@@ -524,6 +508,28 @@ def test_join_sort(self):
                              index=[1, 2, 0, 3])
         assert_frame_equal(joined, expected)
 
+        # smoke test
+        joined = left.join(right, on='key', sort=False)
+        self.assert_(np.array_equal(joined.index, range(4)))
+
+    def test_intelligently_handle_join_key(self):
+        # #733, be a bit more 1337 about not returning unconsolidated DataFrame
+
+        left = DataFrame({'key' : [1, 1, 2, 2, 3],
+                          'value' : range(5)}, columns=['value', 'key'])
+        right = DataFrame({'key' : [1, 1, 2, 3, 4, 5],
+                           'rvalue' : range(6)})
+
+        joined = merge(left, right, on='key', how='outer')
+        expected = DataFrame({'key' : [1, 1, 1, 1, 2, 2, 3, 4, 5.],
+                              'value' : np.array([0, 0, 1, 1, 2, 3, 4,
+                                                  np.nan, np.nan]),
+                              'rvalue' : np.array([0, 1, 0, 1, 2, 2, 3, 4, 5])},
+                             columns=['value', 'key', 'rvalue'])
+        assert_frame_equal(joined, expected)
+
+        self.assert_(joined._data.is_consolidated())
+
 class TestMergeMulti(unittest.TestCase):
 
     def setUp(self):