Repository: mars-project/mars

Issue:
Add `inplace` parameter in `reset_index`
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Is your feature request related to a problem? Please describe.**
Add support for `inplace` parameter in `df.reset_index` and `series.reset_index`.



Discussion:

Resolving patch:
diff --git a/mars/base_app.py b/mars/base_app.py
--- a/mars/base_app.py
+++ b/mars/base_app.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import argparse
+import faulthandler
 import json
 import logging
 import os
@@ -85,6 +86,9 @@ def _main(self, argv=None):
         self.config_args(parser)
         args = self.args = self.parse_args(parser, argv)
 
+        if int(os.environ.get('MARS_ENABLE_FAULTHANDLER', '1')):  # pragma: no branch
+            faulthandler.enable()
+
         endpoint = args.endpoint
         if endpoint is not None:
             ep_parts = endpoint.split(':')
diff --git a/mars/dataframe/indexing/reset_index.py b/mars/dataframe/indexing/reset_index.py
--- a/mars/dataframe/indexing/reset_index.py
+++ b/mars/dataframe/indexing/reset_index.py
@@ -205,12 +205,277 @@ def __call__(self, a):
             return self._call_series(a)
 
 
-def df_reset_index(df, level=None, drop=False, col_level=None, col_fill=None):
+def df_reset_index(df, level=None, drop=False, inplace=False, col_level=0, col_fill=''):
+    """
+    Reset the index, or a level of it.
+
+    Reset the index of the DataFrame, and use the default one instead.
+    If the DataFrame has a MultiIndex, this method can remove one or more
+    levels.
+
+    Parameters
+    ----------
+    level : int, str, tuple, or list, default None
+        Only remove the given levels from the index. Removes all levels by
+        default.
+    drop : bool, default False
+        Do not try to insert index into dataframe columns. This resets
+        the index to the default integer index.
+    inplace : bool, default False
+        Modify the DataFrame in place (do not create a new object).
+    col_level : int or str, default 0
+        If the columns have multiple levels, determines which level the
+        labels are inserted into. By default it is inserted into the first
+        level.
+    col_fill : object, default ''
+        If the columns have multiple levels, determines how the other
+        levels are named. If None then the index name is repeated.
+
+    Returns
+    -------
+    DataFrame or None
+        DataFrame with the new index or None if ``inplace=True``.
+
+    See Also
+    --------
+    DataFrame.set_index : Opposite of reset_index.
+    DataFrame.reindex : Change to new indices or expand indices.
+    DataFrame.reindex_like : Change to same indices as other DataFrame.
+
+    Examples
+    --------
+    >>> import mars.tensor as mt
+    >>> import mars.dataframe as md
+    >>> df = md.DataFrame([('bird', 389.0),
+    ...                    ('bird', 24.0),
+    ...                    ('mammal', 80.5),
+    ...                    ('mammal', mt.nan)],
+    ...                   index=['falcon', 'parrot', 'lion', 'monkey'],
+    ...                   columns=('class', 'max_speed'))
+    >>> df.execute()
+             class  max_speed
+    falcon    bird      389.0
+    parrot    bird       24.0
+    lion    mammal       80.5
+    monkey  mammal        NaN
+
+    When we reset the index, the old index is added as a column, and a
+    new sequential index is used:
+
+    >>> df.reset_index().execute()
+        index   class  max_speed
+    0  falcon    bird      389.0
+    1  parrot    bird       24.0
+    2    lion  mammal       80.5
+    3  monkey  mammal        NaN
+
+    We can use the `drop` parameter to avoid the old index being added as
+    a column:
+
+    >>> df.reset_index(drop=True).execute()
+        class  max_speed
+    0    bird      389.0
+    1    bird       24.0
+    2  mammal       80.5
+    3  mammal        NaN
+
+    You can also use `reset_index` with `MultiIndex`.
+
+    >>> import pandas as pd
+    >>> index = pd.MultiIndex.from_tuples([('bird', 'falcon'),
+    ...                                    ('bird', 'parrot'),
+    ...                                    ('mammal', 'lion'),
+    ...                                    ('mammal', 'monkey')],
+    ...                                   names=['class', 'name'])
+    >>> columns = pd.MultiIndex.from_tuples([('speed', 'max'),
+    ...                                      ('species', 'type')])
+    >>> df = md.DataFrame([(389.0, 'fly'),
+    ...                    ( 24.0, 'fly'),
+    ...                    ( 80.5, 'run'),
+    ...                    (mt.nan, 'jump')],
+    ...                   index=index,
+    ...                   columns=columns)
+    >>> df.execute()
+                   speed species
+                     max    type
+    class  name
+    bird   falcon  389.0     fly
+           parrot   24.0     fly
+    mammal lion     80.5     run
+           monkey    NaN    jump
+
+    If the index has multiple levels, we can reset a subset of them:
+
+    >>> df.reset_index(level='class').execute()
+             class  speed species
+                      max    type
+    name
+    falcon    bird  389.0     fly
+    parrot    bird   24.0     fly
+    lion    mammal   80.5     run
+    monkey  mammal    NaN    jump
+
+    If we are not dropping the index, by default, it is placed in the top
+    level. We can place it in another level:
+
+    >>> df.reset_index(level='class', col_level=1).execute()
+                    speed species
+             class    max    type
+    name
+    falcon    bird  389.0     fly
+    parrot    bird   24.0     fly
+    lion    mammal   80.5     run
+    monkey  mammal    NaN    jump
+
+    When the index is inserted under another level, we can specify under
+    which one with the parameter `col_fill`:
+
+    >>> df.reset_index(level='class', col_level=1, col_fill='species').execute()
+                  species  speed species
+                    class    max    type
+    name
+    falcon           bird  389.0     fly
+    parrot           bird   24.0     fly
+    lion           mammal   80.5     run
+    monkey         mammal    NaN    jump
+
+    If we specify a nonexistent level for `col_fill`, it is created:
+
+    >>> df.reset_index(level='class', col_level=1, col_fill='genus').execute()
+                    genus  speed species
+                    class    max    type
+    name
+    falcon           bird  389.0     fly
+    parrot           bird   24.0     fly
+    lion           mammal   80.5     run
+    monkey         mammal    NaN    jump
+    """
     op = DataFrameResetIndex(level=level, drop=drop, col_level=col_level,
                              col_fill=col_fill, output_types=[OutputType.dataframe])
-    return op(df)
+    ret = op(df)
+    if not inplace:
+        return ret
+    else:
+        df.data = ret.data
+
+
+def series_reset_index(series, level=None, drop=False, name=None, inplace=False):
+    """
+    Generate a new DataFrame or Series with the index reset.
+
+    This is useful when the index needs to be treated as a column, or
+    when the index is meaningless and needs to be reset to the default
+    before another operation.
+
+    Parameters
+    ----------
+    level : int, str, tuple, or list, default optional
+        For a Series with a MultiIndex, only remove the specified levels
+        from the index. Removes all levels by default.
+    drop : bool, default False
+        Just reset the index, without inserting it as a column in
+        the new DataFrame.
+    name : object, optional
+        The name to use for the column containing the original Series
+        values. Uses ``self.name`` by default. This argument is ignored
+        when `drop` is True.
+    inplace : bool, default False
+        Modify the Series in place (do not create a new object).
+
+    Returns
+    -------
+    Series or DataFrame
+        When `drop` is False (the default), a DataFrame is returned.
+        The newly created columns will come first in the DataFrame,
+        followed by the original Series values.
+        When `drop` is True, a `Series` is returned.
+        In either case, if ``inplace=True``, no value is returned.
+
+    See Also
+    --------
+    DataFrame.reset_index: Analogous function for DataFrame.
+
+    Examples
+    --------
+    >>> import mars.tensor as mt
+    >>> import mars.dataframe as md
+    >>> s = md.Series([1, 2, 3, 4], name='foo',
+    ...               index=md.Index(['a', 'b', 'c', 'd'], name='idx'))
+
+    Generate a DataFrame with default index.
+
+    >>> s.reset_index().execute()
+      idx  foo
+    0   a    1
+    1   b    2
+    2   c    3
+    3   d    4
+
+    To specify the name of the new column use `name`.
+
+    >>> s.reset_index(name='values').execute()
+      idx  values
+    0   a       1
+    1   b       2
+    2   c       3
+    3   d       4
+
+    To generate a new Series with the default set `drop` to True.
+
+    >>> s.reset_index(drop=True).execute()
+    0    1
+    1    2
+    2    3
+    3    4
+    Name: foo, dtype: int64
+
+    To update the Series in place, without generating a new one
+    set `inplace` to True. Note that it also requires ``drop=True``.
+
+    >>> s.reset_index(inplace=True, drop=True)
+    >>> s.execute()
+    0    1
+    1    2
+    2    3
+    3    4
+    Name: foo, dtype: int64
+
+    The `level` parameter is interesting for Series with a multi-level
+    index.
+
+    >>> import numpy as np
+    >>> import pandas as pd
+    >>> arrays = [np.array(['bar', 'bar', 'baz', 'baz']),
+    ...           np.array(['one', 'two', 'one', 'two'])]
+    >>> s2 = md.Series(
+    ...     range(4), name='foo',
+    ...     index=pd.MultiIndex.from_arrays(arrays,
+    ...                                     names=['a', 'b']))
+
+    To remove a specific level from the Index, use `level`.
+
+    >>> s2.reset_index(level='a').execute()
+           a  foo
+    b
+    one  bar    0
+    two  bar    1
+    one  baz    2
+    two  baz    3
 
+    If `level` is not set, all levels are removed from the Index.
 
-def series_reset_index(series, level=None, drop=False, name=None):
+    >>> s2.reset_index().execute()
+         a    b  foo
+    0  bar  one    0
+    1  bar  two    1
+    2  baz  one    2
+    3  baz  two    3
+    """
     op = DataFrameResetIndex(level=level, drop=drop, name=name, output_types=[OutputType.series])
-    return op(series)
+    ret = op(series)
+    if not inplace:
+        return ret
+    elif ret.ndim == 2:
+        raise TypeError('Cannot reset_index inplace on a Series to create a DataFrame')
+    else:
+        series.data = ret.data


Test patch:
diff --git a/mars/dataframe/indexing/tests/test_indexing.py b/mars/dataframe/indexing/tests/test_indexing.py
--- a/mars/dataframe/indexing/tests/test_indexing.py
+++ b/mars/dataframe/indexing/tests/test_indexing.py
@@ -705,6 +705,9 @@ def testResetIndex(self):
         self.assertEqual(s2.chunks[1].shape, (2, 2))
         pd.testing.assert_index_equal(s2.chunks[1].index_value.to_pandas(), pd.RangeIndex(2, 4))
 
+        with self.assertRaises(TypeError):
+            md.Series(series_data, chunk_size=2).reset_index(inplace=True)
+
     def testHeadTailOptimize(self):
         raw = pd.DataFrame(np.random.rand(4, 3))
 
diff --git a/mars/dataframe/indexing/tests/test_indexing_execution.py b/mars/dataframe/indexing/tests/test_indexing_execution.py
--- a/mars/dataframe/indexing/tests/test_indexing_execution.py
+++ b/mars/dataframe/indexing/tests/test_indexing_execution.py
@@ -665,6 +665,12 @@ def testResetIndexExecution(self):
         expected = data.reset_index(level='class', col_level=1, col_fill='species')
         pd.testing.assert_frame_equal(result, expected)
 
+        df = md.DataFrame(data, chunk_size=3)
+        df.reset_index(level='class', col_level=1, col_fill='species', inplace=True)
+        result = self.executor.execute_dataframe(df, concat=True)[0]
+        expected = data.reset_index(level='class', col_level=1, col_fill='species')
+        pd.testing.assert_frame_equal(result, expected)
+
         # Test Series
 
         s = pd.Series([1, 2, 3, 4], name='foo',
@@ -708,6 +714,11 @@ def testResetIndexExecution(self):
         expected = (data1 + data2).reset_index()
         np.testing.assert_array_equal(result.to_numpy(), expected.to_numpy())
 
+        series1 = md.Series(data1, chunk_size=3)
+        series1.reset_index(inplace=True, drop=True)
+        result = self.executor.execute_dataframe(series1, concat=True)[0]
+        pd.testing.assert_index_equal(result.index, pd.RangeIndex(10))
+
         # case from https://github.com/mars-project/mars/issues/1286
         data = pd.DataFrame(np.random.rand(10, 3), columns=list('abc'))
         df = md.DataFrame(data, chunk_size=3)