ENH: enable joining with overlap and lsuffix/rsuffix options

Author: wesm

pandas/core/frame.py

@@ -1410,8 +1410,8 @@ def add_prefix(self, prefix):
         -------
         LongPanel
         """
-        f = (('%s' % prefix) + '%s').__mod__
-        return self.rename(columns=f)
+        new_data = self._data.add_prefix(prefix)
+        return self._constructor(new_data)
 
     def add_suffix(self, suffix):
         """
@@ -1425,8 +1425,8 @@ def add_suffix(self, suffix):
         -------
         with_suffix : DataFrame
         """
-        f = ('%s' + ('%s' % suffix)).__mod__
-        return self.rename(columns=f)
+        new_data = self._data.add_suffix(suffix)
+        return self._constructor(new_data)
 
     #----------------------------------------------------------------------
     # Arithmetic / combination related
@@ -2065,7 +2065,7 @@ def append(self, other):
         return self._constructor(data=new_data, index=new_index,
                                  columns=new_columns)
 
-    def join(self, other, on=None, how=None):
+    def join(self, other, on=None, how=None, lsuffix='', rsuffix=''):
         """
         Join columns with other DataFrame either on index or on a key
         column
@@ -2083,6 +2083,10 @@ def join(self, other, on=None, how=None):
             * right: use input frame's index
             * outer: form union of indexes
             * inner: use intersection of indexes
+        lsuffix : string
+            Suffix to use from left frame's overlapping columns
+        rsuffix : string
+            Suffix to use from right frame's overlapping columns
 
         Returns
         -------
@@ -2092,29 +2096,32 @@ def join(self, other, on=None, how=None):
             if how is not None:
                 raise Exception('how parameter is not valid when '
                                 '*on* specified')
-            return self._join_on(other, on)
+            return self._join_on(other, on, lsuffix, rsuffix)
         else:
             if how is None:
                 how = 'left'
-            return self._join_index(other, how)
+            return self._join_index(other, how, lsuffix, rsuffix)
 
-    def _join_on(self, other, on):
+    def _join_on(self, other, on, lsuffix, rsuffix):
         if len(other.index) == 0:
             return self
 
         if on not in self:
             raise Exception('%s column not contained in this frame!' % on)
 
+        this, other = self._maybe_rename_join(other, lsuffix, rsuffix)
         new_data = self._data.join_on(other._data, self[on], axis=1)
         return self._constructor(new_data)
 
-    def _join_index(self, other, how):
+    def _join_index(self, other, how, lsuffix, rsuffix):
         join_index = self._get_join_index(other, how)
-        this_data = self.reindex(join_index)._data
-        other_data = other.reindex(join_index)._data
+
+        this = self.reindex(join_index)
+        other = other.reindex(join_index)
+        this, other = this._maybe_rename_join(other, lsuffix, rsuffix)
 
         # merge blocks
-        merged_data = this_data.merge(other_data)
+        merged_data = this._data.merge(other._data)
         return self._constructor(merged_data)
 
     def _get_join_index(self, other, how):
@@ -2131,6 +2138,30 @@ def _get_join_index(self, other, how):
 
         return join_index
 
+    def _maybe_rename_join(self, other, lsuffix, rsuffix):
+        intersection = self.columns.intersection(other.columns)
+
+        if len(intersection) > 0:
+            if not lsuffix and not rsuffix:
+                raise Exception('columns overlap: %s' % intersection)
+
+            def lrenamer(x):
+                if x in intersection:
+                    return '%s%s' % (x, lsuffix)
+                return x
+
+            def rrenamer(x):
+                if x in intersection:
+                    return '%s%s' % (x, rsuffix)
+                return x
+
+            this = self.rename(columns=lrenamer)
+            other = other.rename(columns=rrenamer)
+        else:
+            this = self
+
+        return this, other
+
     #----------------------------------------------------------------------
     # Statistical methods, etc.
 

pandas/core/internals.py

@@ -673,6 +673,14 @@ def rename_items(self, mapper):
         new_axes[0] = new_items
         return BlockManager(new_blocks, new_axes)
 
+    def add_prefix(self, prefix):
+        f = (('%s' % prefix) + '%s').__mod__
+        return self.rename_items(f)
+
+    def add_suffix(self, suffix):
+        f = ('%s' + ('%s' % suffix)).__mod__
+        return self.rename_items(f)
+
     def fillna(self, value):
         """
 

pandas/core/sparse.py

@@ -1058,16 +1058,16 @@ def _join_on(self, other, on):
         # need to implement?
         raise NotImplementedError
 
-    def _join_index(self, other, how):
+    def _join_index(self, other, how, lsuffix, rsuffix):
         join_index = self._get_join_index(other, how)
 
-        result_series = self.reindex(join_index)._series
-        other_series = other.reindex(join_index)._series
+        this = self.reindex(join_index)
+        other = other.reindex(join_index)
 
-        for col in other_series:
-            if col in result_series:
-                raise Exception('Overlapping columns!')
+        this, other = this._maybe_rename_join(other, lsuffix, rsuffix)
 
+        result_series = this._series
+        other_series = other._series
         result_series.update(other_series)
 
         return self._constructor(result_series, index=join_index)

pandas/tests/test_frame.py

@@ -2299,6 +2299,16 @@ def test_join(self):
         self.assertRaises(Exception, target.join, source, on='C',
                           how='left')
 
+    def test_join_overlap(self):
+        df1 = self.frame.copy()
+        df2 = self.frame.copy()
+
+        joined = df1.join(df2, lsuffix='_df1', rsuffix='_df2')
+        df1_suf = df1.add_suffix('_df1')
+        df2_suf = df2.add_suffix('_df2')
+        expected = df1_suf.join(df2_suf)
+        assert_frame_equal(joined, expected)
+
     def test_clip(self):
         median = self.frame.median().median()