BUG: fix compatibility issue of df.sort_index() and df.groupby(sort=True)

python/xorbits/_mars/dataframe/groupby/sort.py

@@ -122,13 +122,21 @@ def _get_out_df(p_index, in_df):
             if p_index == 0:
                 out_df = in_df.loc[: pivots[p_index]]
             elif p_index == op.n_partition - 1:
-                out_df = in_df.loc[pivots[p_index - 1] :].drop(
-                    index=pivots[p_index - 1], errors="ignore"
-                )
+                # in_df may not have the index in pivots, and throw ZeroDivisionError
+                # in this case, just return an empty dataframe
+                try:
+                    out_df = in_df.loc[pivots[p_index - 1] :].drop(
+                        index=pivots[p_index - 1], errors="ignore"
+                    )
+                except ZeroDivisionError:
+                    out_df = pd.DataFrame(columns=in_df.columns)
             else:
-                out_df = in_df.loc[pivots[p_index - 1] : pivots[p_index]].drop(
-                    index=pivots[p_index - 1], errors="ignore"
-                )
+                try:
+                    out_df = in_df.loc[pivots[p_index - 1] : pivots[p_index]].drop(
+                        index=pivots[p_index - 1], errors="ignore"
+                    )
+                except ZeroDivisionError:
+                    out_df = pd.DataFrame(columns=in_df.columns)
             return out_df
 
         for i in range(op.n_partition):

python/xorbits/_mars/dataframe/sort/sort_index.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import pandas as pd
+from pandas.core.indexes.api import default_index
 
 from ... import opcodes as OperandDef
 from ...core import OutputType, recursive_tile
@@ -122,16 +123,12 @@ def execute(cls, ctx, op: "DataFrameSortIndex"):
         ctx[op.outputs[0].key] = result
 
     def _call_dataframe(self, df):
-        if self.ignore_index:
-            index_value = parse_index(pd.RangeIndex(df.shape[0]))
-        else:
-            index_value = df.index_value
         if self.axis == 0:
             return self.new_dataframe(
                 [df],
                 shape=df.shape,
                 dtypes=df.dtypes,
-                index_value=index_value,
+                index_value=df.index_value,
                 columns_value=df.columns_value,
             )
         else:
@@ -141,7 +138,7 @@ def _call_dataframe(self, df):
                 [df],
                 shape=df.shape,
                 dtypes=dtypes,
-                index_value=index_value,
+                index_value=df.index_value,
                 columns_value=columns_value,
             )
 
@@ -246,6 +243,11 @@ def sort_index(
         gpu=a.op.is_gpu(),
     )
     sorted_a = op(a)
+    if ignore_index:
+        if axis == 1:
+            sorted_a.columns = default_index(len(sorted_a.columns))
+        else:
+            sorted_a.index = default_index(len(sorted_a.index))
     if inplace:
         a.data = sorted_a.data
     else:

python/xorbits/_mars/dataframe/sort/tests/test_sort_execution.py

@@ -289,6 +289,8 @@ def test_sort_values_execution(setup, distinct_opt):
 
 
 def test_sort_index_execution(setup):
+    from ...utils import is_pandas_2
+
     raw = pd.DataFrame(np.random.rand(100, 20), index=np.random.rand(100))
 
     mdf = DataFrame(raw)
@@ -341,13 +343,11 @@ def test_sort_index_execution(setup):
 
     mdf = DataFrame(raw, chunk_size=4)
 
-    result = mdf.sort_index(axis=1, ignore_index=True).execute().fetch()
-    try:  # for python3.5
+    # behavior of sort_index(axis=1, ignore_index=True) change after 2.0
+    if is_pandas_2():
+        result = mdf.sort_index(axis=1, ignore_index=True).execute().fetch()
         expected = raw.sort_index(axis=1, ignore_index=True)
-    except TypeError:
-        expected = raw.sort_index(axis=1)
-        expected.index = pd.RangeIndex(len(expected))
-    pd.testing.assert_frame_equal(result, expected)
+        pd.testing.assert_frame_equal(result, expected)
 
     # test series
     raw = pd.Series(np.random.rand(10), index=np.random.rand(10))