Reset level

pandas/core/frame.py

@@ -2258,7 +2258,7 @@ def reindex_like(self, other, method=None, copy=True, limit=None):
 
     truncate = generic.truncate
 
-    def set_index(self, keys, drop=True, inplace=False,
+    def set_index(self, keys, drop=True, append=False, inplace=False,
                   verify_integrity=False):
         """
         Set the DataFrame index (row labels) using one or more existing
@@ -2269,6 +2269,8 @@ def set_index(self, keys, drop=True, inplace=False,
         keys : column label or list of column labels / arrays
         drop : boolean, default True
             Delete columns to be used as the new index
+        append : boolean, default False
+            Whether to append columns to existing index
         inplace : boolean, default False
             Modify the DataFrame in place (do not create a new object)
         verify_integrity : boolean, default False
@@ -2295,7 +2297,18 @@ def set_index(self, keys, drop=True, inplace=False,
         else:
             frame = self.copy()
 
+        names = keys
+
         arrays = []
+        if append:
+            names = [x for x in self.index.names]
+            if isinstance(self.index, MultiIndex):
+                for i in range(self.index.nlevels):
+                    arrays.append(self.index.get_level_values(i))
+            else:
+                arrays.append(np.asarray(self.index))
+            names.extend(keys)
+
         for col in keys:
             if isinstance(col, (list, Series, np.ndarray)):
                 level = col
@@ -2305,7 +2318,7 @@ def set_index(self, keys, drop=True, inplace=False,
                     del frame[col]
             arrays.append(level)
 
-        index = MultiIndex.from_arrays(arrays, names=keys)
+        index = MultiIndex.from_arrays(arrays, names=names)
 
         if verify_integrity and not index.is_unique:
             duplicates = index.get_duplicates()
@@ -2317,7 +2330,7 @@ def set_index(self, keys, drop=True, inplace=False,
         frame.index = index
         return frame
 
-    def reset_index(self, drop=False):
+    def reset_index(self, level=None, drop=False):
         """
         For DataFrame with multi-level index, return new DataFrame with
         labeling information in the columns under the index names, defaulting
@@ -2327,6 +2340,9 @@ def reset_index(self, drop=False):
 
         Parameters
         ----------
+        level : int, str, tuple, or list, default None
+            Only remove the given levels from the index. Removes all levels by
+            default
         drop : boolean, default False
             Do not try to insert index into dataframe columns
 
@@ -2341,24 +2357,36 @@ def _maybe_cast(values):
                 values = lib.maybe_convert_objects(values)
             return values
 
+        new_index = np.arange(len(new_obj))
         if not drop:
             if isinstance(self.index, MultiIndex):
                 names = self.index.names
                 zipped = zip(self.index.levels, self.index.labels)
+
+                if level is not None:
+                    if not isinstance(level, (tuple, list)):
+                        level = [level]
+
+                    level = [self.index._get_level_number(lev) for lev in level]
+
                 for i, (lev, lab) in reversed(list(enumerate(zipped))):
                     col_name = names[i]
                     if col_name is None:
                         col_name = 'level_%d' % i
 
                     # to ndarray and maybe infer different dtype
                     level_values = _maybe_cast(lev.values)
-                    new_obj.insert(0, col_name, level_values.take(lab))
+                    if level is None or i in level:
+                        new_obj.insert(0, col_name, level_values.take(lab))
+
+                if level is not None and len(level) < len(self.index.levels):
+                    new_index = self.index.droplevel(level)
             else:
                 name = self.index.name
                 if name is None or name == 'index':
                     name = 'index' if 'index' not in self else 'level_0'
                 new_obj.insert(0, name, _maybe_cast(self.index.values))
-        new_obj.index = np.arange(len(new_obj))
+        new_obj.index = new_index
         return new_obj
 
     delevel = deprecate('delevel', reset_index)

pandas/tests/test_frame.py

@@ -1386,6 +1386,12 @@ def test_set_index2(self):
         # corner case
         self.assertRaises(Exception, df.set_index, 'A', verify_integrity=True)
 
+        # append
+        result = df.set_index(['A', 'B'], append=True)
+        xp = df.reset_index().set_index(['index', 'A', 'B'])
+        xp.index.names = [None, 'A', 'B']
+        assert_frame_equal(result, xp)
+
     def test_set_index_pass_arrays(self):
         df = DataFrame({'A' : ['foo', 'bar', 'foo', 'bar',
                                'foo', 'bar', 'foo', 'foo'],
@@ -6160,6 +6166,21 @@ def test_reset_index(self):
         resetted = self.frame.reset_index()
         self.assertEqual(resetted.columns.name, 'columns')
 
+        # only remove certain columns
+        frame = self.frame.reset_index().set_index(['index', 'A', 'B'])
+        rs = frame.reset_index(['A', 'B'])
+        assert_frame_equal(rs, self.frame)
+
+        rs = frame.reset_index(['index', 'A', 'B'])
+        assert_frame_equal(rs, self.frame.reset_index())
+
+        rs = frame.reset_index(['index', 'A', 'B'])
+        assert_frame_equal(rs, self.frame.reset_index())
+
+        rs = frame.reset_index('A')
+        xp = self.frame.reset_index().set_index(['index', 'B'])
+        assert_frame_equal(rs, xp)
+
     def test_reset_index_right_dtype(self):
         time = np.arange(0.0, 10, np.sqrt(2)/2)
         s1 = Series((9.81 * time ** 2) /2,