Repository: mars-project/mars

Issue:
[BUG] sort_values cannot find column __PSRS_TMP_DISTINCT_COL
**Describe the bug**
When executing code below in distributed environment:

```python
rs = np.random.RandomState(0)
raw2 = pd.DataFrame({'a': rs.rand(10),
                    'b': [f's{rs.randint(1000)}' for _ in range(10)]
                    })
raw2['b'] = raw2['b'].astype(md.ArrowStringDtype())
mdf = md.DataFrame(raw2, chunk_size=4)
filtered = mdf[mdf['a'] > 0.5]
df2 = filtered.sort_values(by='b')
result = df2.execute(session=sess, timeout=self.timeout)
```

Error raised

```
2020-10-19 19:46:44,463 Unexpected exception occurred in BaseCalcActor._calc_results. graph_key=cfd4b1a2cc914a2b30aa228eda1e7ea8
Traceback (most recent call last):
  File "/Users/wenjun.swj/Code/mars/mars/utils.py", line 365, in _wrapped
    return func(*args, **kwargs)
  File "/Users/wenjun.swj/Code/mars/mars/worker/calc.py", line 201, in _calc_results
    self._execution_pool.submit(executor.execute_graph, graph,
  File "src/gevent/event.py", line 383, in gevent._gevent_cevent.AsyncResult.result
  File "src/gevent/event.py", line 305, in gevent._gevent_cevent.AsyncResult.get
  File "src/gevent/event.py", line 335, in gevent._gevent_cevent.AsyncResult.get
  File "src/gevent/event.py", line 323, in gevent._gevent_cevent.AsyncResult.get
  File "src/gevent/event.py", line 303, in gevent._gevent_cevent.AsyncResult._raise_exception
  File "/Users/wenjun.swj/miniconda3/lib/python3.8/site-packages/gevent/_compat.py", line 65, in reraise
    raise value.with_traceback(tb)
  File "/Users/wenjun.swj/miniconda3/lib/python3.8/site-packages/gevent/threadpool.py", line 142, in __run_task
    thread_result.set(func(*args, **kwargs))
  File "mars/actors/pool/gevent_pool.pyx", line 127, in mars.actors.pool.gevent_pool.GeventThreadPool._wrap_watch.inner
  File "/Users/wenjun.swj/Code/mars/mars/executor.py", line 693, in execute_graph
    res = graph_execution.execute(retval)
  File "/Users/wenjun.swj/Code/mars/mars/executor.py", line 574, in execute
    future.result()
  File "/Users/wenjun.swj/miniconda3/lib/python3.8/concurrent/futures/_base.py", line 439, in result
    return self.__get_result()
  File "/Users/wenjun.swj/miniconda3/lib/python3.8/concurrent/futures/_base.py", line 388, in __get_result
    raise self._exception
  File "/Users/wenjun.swj/miniconda3/lib/python3.8/concurrent/futures/thread.py", line 57, in run
    result = self.fn(*self.args, **self.kwargs)
  File "/Users/wenjun.swj/Code/mars/mars/utils.py", line 439, in _inner
    return func(*args, **kwargs)
  File "/Users/wenjun.swj/Code/mars/mars/executor.py", line 446, in _execute_operand
    self.handle_op(first_op, results, self._mock)
  File "/Users/wenjun.swj/Code/mars/mars/executor.py", line 378, in handle_op
    return Executor.handle(*args, **kw)
  File "/Users/wenjun.swj/Code/mars/mars/executor.py", line 644, in handle
    return runner(results, op)
  File "/Users/wenjun.swj/Code/mars/mars/dataframe/sort/psrs.py", line 357, in execute
    ctx[op.outputs[-1].key] = res[by].iloc[slc]
  File "/Users/wenjun.swj/miniconda3/lib/python3.8/site-packages/pandas/core/frame.py", line 2908, in __getitem__
    indexer = self.loc._get_listlike_indexer(key, axis=1, raise_missing=True)[1]
  File "/Users/wenjun.swj/miniconda3/lib/python3.8/site-packages/pandas/core/indexing.py", line 1254, in _get_listlike_indexer
    self._validate_read_indexer(keyarr, indexer, axis, raise_missing=raise_missing)
  File "/Users/wenjun.swj/miniconda3/lib/python3.8/site-packages/pandas/core/indexing.py", line 1304, in _validate_read_indexer
    raise KeyError(f"{not_found} not in index")
KeyError: "['__PSRS_TMP_DISTINCT_COL'] not in index"
```


Discussion:

Resolving patch:
diff --git a/mars/dataframe/sort/psrs.py b/mars/dataframe/sort/psrs.py
--- a/mars/dataframe/sort/psrs.py
+++ b/mars/dataframe/sort/psrs.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+
 import numpy as np
 import pandas as pd
 
@@ -189,7 +191,7 @@ def _tile_psrs(cls, op, in_data):
                                       dtype=out.dtype, name=out.name)
 
 
-def execute_sort_values(data, op, inplace=None):
+def execute_sort_values(data, op, inplace=None, by=None):
     if inplace is None:
         inplace = op.inplace
     # ignore_index is new in Pandas version 1.0.0.
@@ -198,7 +200,7 @@ def execute_sort_values(data, op, inplace=None):
         kwargs = dict(axis=op.axis, ascending=op.ascending, ignore_index=ignore_index,
                       na_position=op.na_position, kind=op.kind)
         if isinstance(data, pd.DataFrame):
-            kwargs['by'] = op.by
+            kwargs['by'] = by if by is not None else op.by
         if inplace:
             kwargs['inplace'] = True
             try:
@@ -334,20 +336,21 @@ def execute(cls, ctx, op):
             ctx[op.outputs[0].key] = res = execute_sort_index(a, op)
 
         by = op.by
-        if getattr(ctx, 'running_mode', None) == RunningMode.distributed \
-                and isinstance(a, pd.DataFrame) and op.sort_type == 'sort_values':
+        add_distinct_col = bool(int(os.environ.get('PSRS_DISTINCT_COL', '0'))) \
+            or getattr(ctx, 'running_mode', None) == RunningMode.distributed
+        if add_distinct_col and isinstance(a, pd.DataFrame) and op.sort_type == 'sort_values':
             # when running under distributed mode, we introduce an extra column
             # to make sure pivots are distinct
             chunk_idx = op.inputs[0].index[0]
             distinct_col = _PSRS_DISTINCT_COL if a.columns.nlevels == 1 \
                 else (_PSRS_DISTINCT_COL,) + ('',) * (a.columns.nlevels - 1)
-            res[distinct_col] = np.arange(chunk_idx << 32, (chunk_idx << 32) + len(a))
+            res[distinct_col] = np.arange(chunk_idx << 32, (chunk_idx << 32) + len(a), dtype=np.int64)
             by = list(by) + [distinct_col]
 
         n = op.n_partition
         if a.shape[op.axis] < n:
             num = n // a.shape[op.axis] + 1
-            res = execute_sort_values(pd.concat([a] * num), op)
+            res = execute_sort_values(pd.concat([res] * num), op, by=by)
         w = int(res.shape[op.axis] // n)
 
         slc = (slice(None),) * op.axis + (slice(0, n * w, w),)
@@ -471,7 +474,7 @@ def _execute_dataframe_map(cls, ctx, op):
 
             distinct_col = _PSRS_DISTINCT_COL if a.columns.nlevels == 1 \
                 else (_PSRS_DISTINCT_COL,) + ('',) * (a.columns.nlevels - 1)
-            if _PSRS_DISTINCT_COL in a.columns:
+            if distinct_col in a.columns:
                 by += [distinct_col]
 
             records = a[by].to_records(index=False)
diff --git a/mars/serialize/dataserializer.py b/mars/serialize/dataserializer.py
--- a/mars/serialize/dataserializer.py
+++ b/mars/serialize/dataserializer.py
@@ -21,6 +21,7 @@
 from enum import Enum
 from io import BytesIO
 import pickle  # nosec
+import warnings
 
 import numpy as np
 import pandas as pd
@@ -244,7 +245,7 @@ def loads(buf):
 
     if header.type == SerialType.ARROW:
         try:
-            return pyarrow.deserialize(memoryview(data), mars_serialize_context())
+            return deserialize(memoryview(data))
         except pyarrow.lib.ArrowInvalid:  # pragma: no cover
             # reconstruct value from buffers of arrow components
             data_view = memoryview(data)
@@ -290,6 +291,19 @@ def dumps(obj, *, serial_type=None, compress=None, pickle_protocol=None):
     return sio.getvalue()
 
 
+def _wrap_deprecates(fun):
+    @functools.wraps(fun)
+    def _wrapped(*args, **kwargs):
+        with warnings.catch_warnings():
+            warnings.filterwarnings('ignore', category=DeprecationWarning)
+            return fun(*args, **kwargs)
+
+    if pyarrow is not None and LooseVersion(pyarrow.__version__) >= '2.0':
+        return _wrapped
+    return fun
+
+
+@_wrap_deprecates
 def serialize(data):
     if isinstance(data, pyarrow.SerializedPyObject):
         return data
@@ -306,6 +320,7 @@ def serialize(data):
         raise SerializationFailed(obj=data)
 
 
+@_wrap_deprecates
 def deserialize(data):
     return pyarrow.deserialize(data, mars_serialize_context())
 


Test patch:
diff --git a/mars/dataframe/sort/tests/test_sort_execute.py b/mars/dataframe/sort/tests/test_sort_execute.py
--- a/mars/dataframe/sort/tests/test_sort_execute.py
+++ b/mars/dataframe/sort/tests/test_sort_execute.py
@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
+import os
+import sys
 import unittest
 
 import numpy as np
@@ -28,156 +29,163 @@ def setUp(self) -> None:
         super().setUp()
         self.executor = ExecutorForTest('numpy')
 
+    def tearDown(self) -> None:
+        os.environ.pop('PSRS_DISTINCT_COL', None)
+        super().tearDown()
+
     def testSortValuesExecution(self):
-        df = pd.DataFrame(np.random.rand(100, 10), columns=['a' + str(i) for i in range(10)])
+        distinct_opts = ['0'] if sys.platform.lower().startswith('win') else ['0', '1']
+        for add_distinct in distinct_opts:
+            os.environ['PSRS_DISTINCT_COL'] = add_distinct
+            df = pd.DataFrame(np.random.rand(100, 10), columns=['a' + str(i) for i in range(10)])
 
-        # test one chunk
-        mdf = DataFrame(df)
-        result = self.executor.execute_dataframe(mdf.sort_values('a0'), concat=True)[0]
-        expected = df.sort_values('a0')
+            # test one chunk
+            mdf = DataFrame(df)
+            result = self.executor.execute_dataframe(mdf.sort_values('a0'), concat=True)[0]
+            expected = df.sort_values('a0')
 
-        pd.testing.assert_frame_equal(result, expected)
+            pd.testing.assert_frame_equal(result, expected)
 
-        result = self.executor.execute_dataframe(mdf.sort_values(['a6', 'a7'], ascending=False), concat=True)[0]
-        expected = df.sort_values(['a6', 'a7'], ascending=False)
+            result = self.executor.execute_dataframe(mdf.sort_values(['a6', 'a7'], ascending=False), concat=True)[0]
+            expected = df.sort_values(['a6', 'a7'], ascending=False)
 
-        pd.testing.assert_frame_equal(result, expected)
+            pd.testing.assert_frame_equal(result, expected)
 
-        # test psrs
-        mdf = DataFrame(df, chunk_size=10)
-        result = self.executor.execute_dataframe(mdf.sort_values('a0'), concat=True)[0]
-        expected = df.sort_values('a0')
+            # test psrs
+            mdf = DataFrame(df, chunk_size=10)
+            result = self.executor.execute_dataframe(mdf.sort_values('a0'), concat=True)[0]
+            expected = df.sort_values('a0')
 
-        pd.testing.assert_frame_equal(result, expected)
+            pd.testing.assert_frame_equal(result, expected)
 
-        result = self.executor.execute_dataframe(mdf.sort_values(['a3', 'a4']), concat=True)[0]
-        expected = df.sort_values(['a3', 'a4'])
+            result = self.executor.execute_dataframe(mdf.sort_values(['a3', 'a4']), concat=True)[0]
+            expected = df.sort_values(['a3', 'a4'])
 
-        pd.testing.assert_frame_equal(result, expected)
+            pd.testing.assert_frame_equal(result, expected)
 
-        # test ascending=False
-        result = self.executor.execute_dataframe(mdf.sort_values(['a0', 'a1'], ascending=False), concat=True)[0]
-        expected = df.sort_values(['a0', 'a1'], ascending=False)
+            # test ascending=False
+            result = self.executor.execute_dataframe(mdf.sort_values(['a0', 'a1'], ascending=False), concat=True)[0]
+            expected = df.sort_values(['a0', 'a1'], ascending=False)
 
-        pd.testing.assert_frame_equal(result, expected)
+            pd.testing.assert_frame_equal(result, expected)
 
-        result = self.executor.execute_dataframe(mdf.sort_values(['a7'], ascending=False), concat=True)[0]
-        expected = df.sort_values(['a7'], ascending=False)
+            result = self.executor.execute_dataframe(mdf.sort_values(['a7'], ascending=False), concat=True)[0]
+            expected = df.sort_values(['a7'], ascending=False)
 
-        pd.testing.assert_frame_equal(result, expected)
+            pd.testing.assert_frame_equal(result, expected)
 
-        # test multiindex
-        df2 = df.copy(deep=True)
-        df2.columns = pd.MultiIndex.from_product([list('AB'), list('CDEFG')])
-        mdf = DataFrame(df2, chunk_size=10)
+            # test multiindex
+            df2 = df.copy(deep=True)
+            df2.columns = pd.MultiIndex.from_product([list('AB'), list('CDEFG')])
+            mdf = DataFrame(df2, chunk_size=10)
 
-        result = self.executor.execute_dataframe(mdf.sort_values([('A', 'C')]), concat=True)[0]
-        expected = df2.sort_values([('A', 'C')])
+            result = self.executor.execute_dataframe(mdf.sort_values([('A', 'C')]), concat=True)[0]
+            expected = df2.sort_values([('A', 'C')])
 
-        pd.testing.assert_frame_equal(result, expected)
+            pd.testing.assert_frame_equal(result, expected)
 
-        # test rechunk
-        mdf = DataFrame(df, chunk_size=3)
-        result = self.executor.execute_dataframe(mdf.sort_values('a0'), concat=True)[0]
-        expected = df.sort_values('a0')
+            # test rechunk
+            mdf = DataFrame(df, chunk_size=3)
+            result = self.executor.execute_dataframe(mdf.sort_values('a0'), concat=True)[0]
+            expected = df.sort_values('a0')
 
-        pd.testing.assert_frame_equal(result, expected)
+            pd.testing.assert_frame_equal(result, expected)
 
-        result = self.executor.execute_dataframe(mdf.sort_values(['a3', 'a4']), concat=True)[0]
-        expected = df.sort_values(['a3', 'a4'])
+            result = self.executor.execute_dataframe(mdf.sort_values(['a3', 'a4']), concat=True)[0]
+            expected = df.sort_values(['a3', 'a4'])
 
-        pd.testing.assert_frame_equal(result, expected)
+            pd.testing.assert_frame_equal(result, expected)
 
-        # test other types
-        raw = pd.DataFrame({'a': np.random.rand(10),
-                            'b': np.random.randint(1000, size=10),
-                            'c': np.random.rand(10),
-                            'd': [np.random.bytes(10) for _ in range(10)],
-                            'e': [pd.Timestamp(f'201{i}') for i in range(10)],
-                            'f': [pd.Timedelta(f'{i} days') for i in range(10)]
-                            },)
-        mdf = DataFrame(raw, chunk_size=3)
+            # test other types
+            raw = pd.DataFrame({'a': np.random.rand(10),
+                                'b': np.random.randint(1000, size=10),
+                                'c': np.random.rand(10),
+                                'd': [np.random.bytes(10) for _ in range(10)],
+                                'e': [pd.Timestamp(f'201{i}') for i in range(10)],
+                                'f': [pd.Timedelta(f'{i} days') for i in range(10)]
+                                },)
+            mdf = DataFrame(raw, chunk_size=3)
 
-        for label in raw.columns:
-            result = self.executor.execute_dataframe(mdf.sort_values(label), concat=True)[0]
-            expected = raw.sort_values(label)
-            pd.testing.assert_frame_equal(result, expected)
+            for label in raw.columns:
+                result = self.executor.execute_dataframe(mdf.sort_values(label), concat=True)[0]
+                expected = raw.sort_values(label)
+                pd.testing.assert_frame_equal(result, expected)
 
-        result = self.executor.execute_dataframe(mdf.sort_values(['a', 'b', 'e'], ascending=False), concat=True)[0]
-        expected = raw.sort_values(['a', 'b', 'e'], ascending=False)
+            result = self.executor.execute_dataframe(mdf.sort_values(['a', 'b', 'e'], ascending=False), concat=True)[0]
+            expected = raw.sort_values(['a', 'b', 'e'], ascending=False)
 
-        pd.testing.assert_frame_equal(result, expected)
+            pd.testing.assert_frame_equal(result, expected)
 
-        # test nan
-        df = pd.DataFrame({
-            'col1': ['A', 'A', 'B', 'B', 'D', 'C'],
-            'col2': [2, 1, 9, np.nan, 7, 4],
-            'col3': [0, 1, 9, 4, 2, 3],
-        })
-        mdf = DataFrame(df)
-        result = self.executor.execute_dataframe(mdf.sort_values(['col2']), concat=True)[0]
-        expected = df.sort_values(['col2'])
+            # test nan
+            df = pd.DataFrame({
+                'col1': ['A', 'A', 'B', 'B', 'D', 'C'],
+                'col2': [2, 1, 9, np.nan, 7, 4],
+                'col3': [0, 1, 9, 4, 2, 3],
+            })
+            mdf = DataFrame(df)
+            result = self.executor.execute_dataframe(mdf.sort_values(['col2']), concat=True)[0]
+            expected = df.sort_values(['col2'])
 
-        pd.testing.assert_frame_equal(result, expected)
+            pd.testing.assert_frame_equal(result, expected)
 
-        mdf = DataFrame(df, chunk_size=3)
-        result = self.executor.execute_dataframe(mdf.sort_values(['col2']), concat=True)[0]
-        expected = df.sort_values(['col2'])
+            mdf = DataFrame(df, chunk_size=3)
+            result = self.executor.execute_dataframe(mdf.sort_values(['col2']), concat=True)[0]
+            expected = df.sort_values(['col2'])
 
-        pd.testing.assert_frame_equal(result, expected)
+            pd.testing.assert_frame_equal(result, expected)
 
-        # test ignore_index
-        executor = ExecutorForTest(storage=new_session().context)
+            # test ignore_index
+            executor = ExecutorForTest(storage=new_session().context)
 
-        df = pd.DataFrame(np.random.rand(10, 3), columns=['a' + str(i) for i in range(3)])
+            df = pd.DataFrame(np.random.rand(10, 3), columns=['a' + str(i) for i in range(3)])
 
-        mdf = DataFrame(df, chunk_size=3)
-        result = executor.execute_dataframe(
-            mdf.sort_values(['a0', 'a1'], ignore_index=True), concat=True)[0]
-        try:  # for python3.5
-            expected = df.sort_values(['a0', 'a1'], ignore_index=True)
-        except TypeError:
-            expected = df.sort_values(['a0', 'a1'])
-            expected.index = pd.RangeIndex(len(expected))
+            mdf = DataFrame(df, chunk_size=3)
+            result = executor.execute_dataframe(
+                mdf.sort_values(['a0', 'a1'], ignore_index=True), concat=True)[0]
+            try:  # for python3.5
+                expected = df.sort_values(['a0', 'a1'], ignore_index=True)
+            except TypeError:
+                expected = df.sort_values(['a0', 'a1'])
+                expected.index = pd.RangeIndex(len(expected))
 
-        pd.testing.assert_frame_equal(result, expected)
+            pd.testing.assert_frame_equal(result, expected)
 
-        # test inplace
-        mdf = DataFrame(df)
-        mdf.sort_values('a0', inplace=True)
-        result = self.executor.execute_dataframe(mdf, concat=True)[0]
-        df.sort_values('a0', inplace=True)
+            # test inplace
+            mdf = DataFrame(df)
+            mdf.sort_values('a0', inplace=True)
+            result = self.executor.execute_dataframe(mdf, concat=True)[0]
+            df.sort_values('a0', inplace=True)
 
-        pd.testing.assert_frame_equal(result, df)
+            pd.testing.assert_frame_equal(result, df)
 
-        # test unknown shape
-        df = pd.DataFrame({'a': list(range(10)),
-                           'b': np.random.random(10)})
-        mdf = DataFrame(df, chunk_size=4)
-        filtered = mdf[mdf['a'] > 2]
-        result = self.executor.execute_dataframe(filtered.sort_values(by='b'), concat=True)[0]
+            # test unknown shape
+            df = pd.DataFrame({'a': list(range(10)),
+                               'b': np.random.random(10)})
+            mdf = DataFrame(df, chunk_size=4)
+            filtered = mdf[mdf['a'] > 2]
+            result = self.executor.execute_dataframe(filtered.sort_values(by='b'), concat=True)[0]
 
-        pd.testing.assert_frame_equal(result, df[df['a'] > 2].sort_values(by='b'))
+            pd.testing.assert_frame_equal(result, df[df['a'] > 2].sort_values(by='b'))
 
-        # test Sereis.sort_values
-        raw = pd.Series(np.random.rand(10))
-        series = Series(raw)
-        result = self.executor.execute_dataframe(series.sort_values(), concat=True)[0]
-        expected = raw.sort_values()
+            # test Series.sort_values
+            raw = pd.Series(np.random.rand(10))
+            series = Series(raw)
+            result = self.executor.execute_dataframe(series.sort_values(), concat=True)[0]
+            expected = raw.sort_values()
 
-        pd.testing.assert_series_equal(result, expected)
+            pd.testing.assert_series_equal(result, expected)
 
-        series = Series(raw, chunk_size=3)
-        result = self.executor.execute_dataframe(series.sort_values(), concat=True)[0]
-        expected = raw.sort_values()
+            series = Series(raw, chunk_size=3)
+            result = self.executor.execute_dataframe(series.sort_values(), concat=True)[0]
+            expected = raw.sort_values()
 
-        pd.testing.assert_series_equal(result, expected)
+            pd.testing.assert_series_equal(result, expected)
 
-        series = Series(raw, chunk_size=2)
-        result = self.executor.execute_dataframe(series.sort_values(ascending=False), concat=True)[0]
-        expected = raw.sort_values(ascending=False)
+            series = Series(raw, chunk_size=2)
+            result = self.executor.execute_dataframe(series.sort_values(ascending=False), concat=True)[0]
+            expected = raw.sort_values(ascending=False)
 
-        pd.testing.assert_series_equal(result, expected)
+            pd.testing.assert_series_equal(result, expected)
 
     def testSortIndexExecution(self):
         raw = pd.DataFrame(np.random.rand(100, 20), index=np.random.rand(100))
diff --git a/mars/scheduler/tests/integrated/test_normal_execution.py b/mars/scheduler/tests/integrated/test_normal_execution.py
--- a/mars/scheduler/tests/integrated/test_normal_execution.py
+++ b/mars/scheduler/tests/integrated/test_normal_execution.py
@@ -151,10 +151,11 @@ def testMainDataFrameWithoutEtcd(self):
                             'b': [f's{rs.randint(1000)}' for _ in range(10)]
                             })
         raw2['b'] = raw2['b'].astype(md.ArrowStringDtype())
-        mdf = md.DataFrame(raw2, chunk_size=3)
-        df2 = mdf.sort_values(by='b')
+        mdf = md.DataFrame(raw2, chunk_size=4)
+        filtered = mdf[mdf['a'] > 0.5]
+        df2 = filtered.sort_values(by='b')
         result = df2.execute(session=sess, timeout=self.timeout).fetch(session=sess)
-        expected = raw2.sort_values(by='b')
+        expected = raw2[raw2['a'] > 0.5].sort_values(by='b')
         pd.testing.assert_frame_equal(result, expected)
 
         s1 = pd.Series(np.random.rand(10), index=[11, 1, 2, 5, 7, 6, 8, 9, 10, 3])