Create indicator for obs from left, right, or both

Author: Nick Eubank

pandas/core/frame.py

@@ -114,6 +114,13 @@
     side, respectively
 copy : boolean, default True
     If False, do not copy data unnecessarily
+indicator : boolean, default False
+    If True, adds a column to output DataFrame called "_merge" with 
+    information on the source of each row. "_merge" is Categorical-type 
+    and takes on a value of "left_only" for observations whose merge key 
+    only appears in 'left' DataFrame, "right_only" for observations whose 
+    merge key only appears in 'right' DataFrame, and "both" if the 
+    observation's merge key is found in both. 
 
 Examples
 --------

pandas/tools/merge.py

@@ -30,11 +30,11 @@
 @Appender(_merge_doc, indents=0)
 def merge(left, right, how='inner', on=None, left_on=None, right_on=None,
           left_index=False, right_index=False, sort=False,
-          suffixes=('_x', '_y'), copy=True):
+          suffixes=('_x', '_y'), copy=True, indicator=False):
     op = _MergeOperation(left, right, how=how, on=on, left_on=left_on,
                          right_on=right_on, left_index=left_index,
                          right_index=right_index, sort=sort, suffixes=suffixes,
-                         copy=copy)
+                         copy=copy, indicator=indicator)
     return op.get_result()
 if __debug__:
     merge.__doc__ = _merge_doc % '\nleft : DataFrame'
@@ -160,7 +160,7 @@ class _MergeOperation(object):
     def __init__(self, left, right, how='inner', on=None,
                  left_on=None, right_on=None, axis=1,
                  left_index=False, right_index=False, sort=True,
-                 suffixes=('_x', '_y'), copy=True):
+                 suffixes=('_x', '_y'), copy=True, indicator=False):
         self.left = self.orig_left = left
         self.right = self.orig_right = right
         self.how = how
@@ -177,12 +177,18 @@ def __init__(self, left, right, how='inner', on=None,
         self.left_index = left_index
         self.right_index = right_index
 
+        self.indicator = indicator
+
         # note this function has side effects
         (self.left_join_keys,
          self.right_join_keys,
          self.join_names) = self._get_merge_keys()
 
     def get_result(self):
+        if self.indicator:
+            # Has side-effects. Most cleaned up in `self._indicator_post_merge`          
+            self._indicator_pre_merge(self.left, self.right)
+
         join_index, left_indexer, right_indexer = self._get_join_info()
 
         ldata, rdata = self.left._data, self.right._data
@@ -202,10 +208,41 @@ def get_result(self):
         typ = self.left._constructor
         result = typ(result_data).__finalize__(self, method='merge')
 
+        if self.indicator:
+            # Has side-effects
+            self._indicator_post_merge(result, self.left, self.right)
+
         self._maybe_add_join_keys(result, left_indexer, right_indexer)
 
         return result
 
+    def _indicator_pre_merge(self, left, right):
+        
+        columns = left.columns.values.tolist() + right.columns.values.tolist()  
+
+        for i in ['_left_indicator', '_right_indicator', '_merge']:
+            if i in columns:
+                raise ValueError("Cannot use `indicator=True` option when data contains a column named {}".format(i))
+
+        left['_left_indicator'] = 1  
+        left['_left_indicator'] = left['_left_indicator'].astype('int8')  
+        
+        right['_right_indicator'] = 2     
+        right['_right_indicator'] = right['_right_indicator'].astype('int8')  
+
+
+    def _indicator_post_merge(self, result, left, right):
+        result['_left_indicator'].fillna(0, inplace=True)        
+        result['_right_indicator'].fillna(0, inplace=True)
+        
+        result['_merge'] = Categorical((result['_left_indicator'] + result['_right_indicator']), categories=[1,2,3])
+        result['_merge'].cat.rename_categories(['left_only', 'right_only', 'both'], inplace=True)        
+ 
+       # Cleanup
+        result.drop(labels=['_left_indicator', '_right_indicator'], axis=1, inplace=True)
+        left.drop(labels=['_left_indicator'], axis=1, inplace=True)
+        right.drop(labels=['_right_indicator'], axis=1, inplace=True)
+
     def _maybe_add_join_keys(self, result, left_indexer, right_indexer):
         # insert group keys
 

pandas/tools/tests/test_merge.py

@@ -856,6 +856,47 @@ def test_overlapping_columns_error_message(self):
 
         self.assertRaises(ValueError, merge, df, df2)
 
+    def test_indicator(self):
+        
+        df1 = pd.DataFrame({'col1':[0,1], 'col_left':['a','b']})
+        df1_copy = df1.copy()
+
+        df2 = pd.DataFrame({'col1':[1,2,3,4,5],'col_right':[2,2,2,2,2]})
+        df2_copy = df2.copy()
+        
+        df_result = pd.DataFrame({'col1':[0,1,2,3,4,5], 
+                'col_left':['a','b', np.nan,np.nan,np.nan,np.nan], 
+                'col_right':[np.nan, 2,2,2,2,2]}, 
+                dtype='float64')
+        df_result['_merge'] = pd.Categorical(['left_only','both','right_only',
+            'right_only','right_only','right_only']
+            , categories=['left_only', 'right_only', 'both'])
+
+        test = pd.merge(df1, df2, on='col1', how='outer', indicator=True)
+        assert_frame_equal(test, df_result)
+
+        # No side effects
+        assert_frame_equal(df1, df1_copy)
+        assert_frame_equal(df2, df2_copy)
+
+
+        test2 = pd.merge(df1, df2, on='col1', how='left', indicator=True)
+        assert (test2._merge != 'right_only').all()
+
+        test3 = pd.merge(df1, df2, on='col1', how='right', indicator=True)
+        assert (test3._merge != 'left_only').all()
+
+        test4 = pd.merge(df1, df2, on='col1', how='inner', indicator=True)
+        assert (test4._merge == 'both').all()
+
+        # Check if working name in df
+        for i in ['_right_indicator', '_left_indicator', '_merge']:
+            df_badcolumn = pd.DataFrame({'col1':[1,2], i:[2,2]})
+        
+            with tm.assertRaises(ValueError):
+                pd.merge(df1, df_badcolumn, on='col1', how='outer', indicator=True)
+
+
 def _check_merge(x, y):
     for how in ['inner', 'left', 'outer']:
         result = x.join(y, how=how)