Repository: mars-project/mars

Issue:
Add `incremental_index` argument for `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.**

For now, `reset_index` will force the index incremental, however, if the index does not matter for users, an argument `incremental_index` can be provided.


Discussion:

Resolving patch:
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
@@ -33,10 +33,13 @@ class DataFrameResetIndex(DataFrameOperand, DataFrameOperandMixin):
     _name = StringField('name')
     _col_level = AnyField('col_level')
     _col_fill = AnyField('col_fill')
+    _incremental_index = BoolField('incremental_index')
 
-    def __init__(self, level=None, drop=None, name=None, col_level=None, col_fill=None, output_types=None, **kwargs):
+    def __init__(self, level=None, drop=None, name=None, col_level=None, col_fill=None,
+                 incremental_index=None, output_types=None, **kwargs):
         super().__init__(_level=level, _drop=drop, _name=name, _col_level=col_level,
-                         _col_fill=col_fill, _output_types=output_types, **kwargs)
+                         _col_fill=col_fill, _incremental_index=incremental_index,
+                         _output_types=output_types, **kwargs)
 
     @property
     def level(self):
@@ -58,8 +61,12 @@ def col_level(self):
     def col_fill(self):
         return self._col_fill
 
+    @property
+    def incremental_index(self):
+        return self._incremental_index
+
     @classmethod
-    def _tile_series(cls, op):
+    def _tile_series(cls, op: "DataFrameResetIndex"):
         out_chunks = []
         out = op.outputs[0]
         is_range_index = out.index_value.has_value()
@@ -78,7 +85,8 @@ def _tile_series(cls, op):
                 out_chunk = chunk_op.new_chunk([c], shape=shape, index=c.index + (0,), dtypes=out.dtypes,
                                                index_value=index_value, columns_value=out.columns_value)
             out_chunks.append(out_chunk)
-        if not is_range_index and isinstance(out.index_value._index_value, IndexValue.RangeIndex):
+        if not is_range_index and isinstance(out.index_value.value, IndexValue.RangeIndex) and \
+                op.incremental_index:
             out_chunks = standardize_range_index(out_chunks)
         new_op = op.copy()
         if op.drop:
@@ -92,7 +100,7 @@ def _tile_series(cls, op):
                                          dtypes=out.dtypes)
 
     @classmethod
-    def _tile_dataframe(cls, op):
+    def _tile_dataframe(cls, op: "DataFrameResetIndex"):
         in_df = op.inputs[0]
         out_df = op.outputs[0]
         added_columns_num = len(out_df.dtypes) - len(in_df.dtypes)
@@ -122,7 +130,7 @@ def _tile_dataframe(cls, op):
                                                index=c.index, columns_value=c.columns_value, dtypes=c.dtypes)
             out_chunks.append(new_chunk)
         if not index_has_value or chunk_has_nan:
-            if isinstance(out_df.index_value._index_value, IndexValue.RangeIndex):
+            if isinstance(out_df.index_value.value, IndexValue.RangeIndex) and op.incremental_index:
                 out_chunks = standardize_range_index(out_chunks)
         new_op = op.copy()
         columns_splits = list(in_df.nsplits[1])
@@ -205,7 +213,8 @@ def __call__(self, a):
             return self._call_series(a)
 
 
-def df_reset_index(df, level=None, drop=False, inplace=False, col_level=0, col_fill=''):
+def df_reset_index(df, level=None, drop=False, inplace=False, col_level=0,
+                   col_fill='', incremental_index=False):
     """
     Reset the index, or a level of it.
 
@@ -230,6 +239,12 @@ def df_reset_index(df, level=None, drop=False, inplace=False, col_level=0, col_f
     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.
+    incremental_index: bool, default False
+        Ensure RangeIndex incremental, when output DataFrame has multiple chunks,
+        ensuring index incremental costs more computation,
+        so by default, each chunk will have index which starts from 0,
+        setting incremental_index=True，reset_index will garantee that
+        output DataFrame's index is from 0 to n - 1.
 
     Returns
     -------
@@ -351,7 +366,8 @@ class    max    type
     monkey         mammal    NaN    jump
     """
     op = DataFrameResetIndex(level=level, drop=drop, col_level=col_level,
-                             col_fill=col_fill, output_types=[OutputType.dataframe])
+                             col_fill=col_fill, incremental_index=incremental_index,
+                             output_types=[OutputType.dataframe])
     ret = op(df)
     if not inplace:
         return ret
@@ -359,7 +375,8 @@ class    max    type
         df.data = ret.data
 
 
-def series_reset_index(series, level=None, drop=False, name=None, inplace=False):
+def series_reset_index(series, level=None, drop=False, name=None,
+                       inplace=False, incremental_index=False):
     """
     Generate a new DataFrame or Series with the index reset.
 
@@ -381,6 +398,12 @@ def series_reset_index(series, level=None, drop=False, name=None, inplace=False)
         when `drop` is True.
     inplace : bool, default False
         Modify the Series in place (do not create a new object).
+    incremental_index: bool, default False
+        Ensure RangeIndex incremental, when output Series has multiple chunks,
+        ensuring index incremental costs more computation,
+        so by default, each chunk will have index which starts from 0,
+        setting incremental_index=True，reset_index will garantee that
+        output Series's index is from 0 to n - 1.
 
     Returns
     -------
@@ -471,7 +494,9 @@ def series_reset_index(series, level=None, drop=False, name=None, inplace=False)
     2  baz  one    2
     3  baz  two    3
     """
-    op = DataFrameResetIndex(level=level, drop=drop, name=name, output_types=[OutputType.series])
+    op = DataFrameResetIndex(level=level, drop=drop, name=name,
+                             incremental_index=incremental_index,
+                             output_types=[OutputType.series])
     ret = op(series)
     if not inplace:
         return ret
diff --git a/mars/dataframe/merge/concat.py b/mars/dataframe/merge/concat.py
--- a/mars/dataframe/merge/concat.py
+++ b/mars/dataframe/merge/concat.py
@@ -91,11 +91,16 @@ def _tile_dataframe(cls, op):
 
         out_df = op.outputs[0]
         inputs = op.inputs
+        axis = op.axis
 
-        check_chunks_unknown_shape(inputs, TilesError)
+        if not all(inputs[i].nsplits[1 - axis] == inputs[i + 1].nsplits[1 - axis]
+                   for i in range(len(inputs) - 1)):
+            # need rechunk
+            check_chunks_unknown_shape(inputs, TilesError)
+            normalized_nsplits = {1 - axis: inputs[0].nsplits[1 - axis]}
+            inputs = [inp.rechunk(normalized_nsplits)._inplace_tile()
+                      for inp in inputs]
 
-        normalized_nsplits = {1: inputs[0].nsplits[1]} if op.axis == 0 else {0: inputs[0].nsplits[0]}
-        inputs = [item.rechunk(normalized_nsplits)._inplace_tile() for item in inputs]
         out_chunks = []
         nsplits = []
         cum_index = 0
diff --git a/mars/tensor/datasource/from_dense.py b/mars/tensor/datasource/from_dense.py
--- a/mars/tensor/datasource/from_dense.py
+++ b/mars/tensor/datasource/from_dense.py
@@ -14,7 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from collections import Iterable
+from collections.abc import Iterable
 from functools import reduce
 from operator import and_
 


Test patch:
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
@@ -708,7 +708,7 @@ def testResetIndexExecution(self):
         df1 = md.DataFrame(data1, chunk_size=5)
         data2 = pd.DataFrame(np.random.rand(10, 3), index=[11, 1, 2, 5, 7, 6, 8, 9, 10, 3])
         df2 = md.DataFrame(data2, chunk_size=6)
-        df = (df1 + df2).reset_index()
+        df = (df1 + df2).reset_index(incremental_index=True)
         result = sess.run(df)
         pd.testing.assert_index_equal(result.index, pd.RangeIndex(12))
         # Inconsistent with Pandas when input dataframe's shape is unknown.
@@ -720,7 +720,7 @@ def testResetIndexExecution(self):
         series1 = md.Series(data1, chunk_size=3)
         data2 = pd.Series(np.random.rand(10,), index=[11, 1, 2, 5, 7, 6, 8, 9, 10, 3])
         series2 = md.Series(data2, chunk_size=3)
-        df = (series1 + series2).reset_index()
+        df = (series1 + series2).reset_index(incremental_index=True)
         result = sess.run(df)
         pd.testing.assert_index_equal(result.index, pd.RangeIndex(12))
         # Inconsistent with Pandas when input dataframe's shape is unknown.
@@ -737,7 +737,7 @@ def testResetIndexExecution(self):
         data = pd.DataFrame(np.random.rand(10, 3), columns=list('abc'))
         df = md.DataFrame(data, chunk_size=3)
 
-        r = df.sort_values('a').reset_index(drop=True)
+        r = df.sort_values('a').reset_index(drop=True, incremental_index=True)
         result = self.executor.execute_dataframe(r, concat=True)[0]
         expected = data.sort_values('a').reset_index(drop=True)
         pd.testing.assert_frame_equal(result, expected)