Repository: mars-project/mars

Issue:
[BUG] Tensor sort failed with `IndexError`
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**
Tensor sort failed with `IndexError`.

**To Reproduce**
``` Python
In [59]: a = mt.random.randint(0, 2, size=(1000,), chunk_size=300)                                  

In [60]: mt.sort(mt.sort(a)).execute()                                                              
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-60-6c95dab88665> in <module>
----> 1 mt.sort(mt.sort(a)).execute()

~/Documents/mars_dev/mars/mars/core.py in execute(self, session, **kw)
    644 
    645         if wait:
--> 646             return run()
    647         else:
    648             thread_executor = ThreadPoolExecutor(1)

~/Documents/mars_dev/mars/mars/core.py in run()
    640 
    641         def run():
--> 642             self.data.execute(session, **kw)
    643             return self
    644 

~/Documents/mars_dev/mars/mars/core.py in execute(self, session, **kw)
    377 
    378         if wait:
--> 379             return run()
    380         else:
    381             # leverage ThreadPoolExecutor to submit task,

~/Documents/mars_dev/mars/mars/core.py in run()
    372         def run():
    373             # no more fetch, thus just fire run
--> 374             session.run(self, **kw)
    375             # return Tileable or ExecutableTuple itself
    376             return self

~/Documents/mars_dev/mars/mars/session.py in run(self, *tileables, **kw)
    503         tileables = tuple(mt.tensor(t) if not isinstance(t, (Entity, Base)) else t
    504                           for t in tileables)
--> 505         result = self._sess.run(*tileables, **kw)
    506 
    507         for t in tileables:

~/Documents/mars_dev/mars/mars/session.py in run(self, *tileables, **kw)
    109             # set number of running cores
    110             self.context.set_ncores(kw['n_parallel'])
--> 111             res = self._executor.execute_tileables(tileables, **kw)
    112             return res
    113 

~/Documents/mars_dev/mars/mars/utils.py in _inner(*args, **kwargs)
    456         def _inner(*args, **kwargs):
    457             with self:
--> 458                 return func(*args, **kwargs)
    459 
    460         return _inner

~/Documents/mars_dev/mars/mars/executor.py in execute_tileables(self, tileables, fetch, n_parallel, n_thread, print_progress, mock, compose, name)
    884                                    n_parallel=n_parallel or n_thread,
    885                                    print_progress=print_progress, mock=mock,
--> 886                                    chunk_result=chunk_result)
    887 
    888                 # update shape of tileable and its chunks whatever it's successful or not

~/Documents/mars_dev/mars/mars/executor.py in execute_graph(self, graph, keys, n_parallel, print_progress, mock, no_intermediate, compose, retval, chunk_result)
    696             print_progress=print_progress, mock=mock, mock_max_memory=self._mock_max_memory,
    697             fetch_keys=fetch_keys, no_intermediate=no_intermediate)
--> 698         res = graph_execution.execute(retval)
    699         self._mock_max_memory = max(self._mock_max_memory, graph_execution._mock_max_memory)
    700         if mock:

~/Documents/mars_dev/mars/mars/executor.py in execute(self, retval)
    577         # wait until all the futures completed
    578         for future in executed_futures:
--> 579             future.result()
    580 
    581         if retval:

~/miniconda3/envs/py3.7.2/lib/python3.7/concurrent/futures/_base.py in result(self, timeout)
    423                 raise CancelledError()
    424             elif self._state == FINISHED:
--> 425                 return self.__get_result()
    426 
    427             self._condition.wait(timeout)

~/miniconda3/envs/py3.7.2/lib/python3.7/concurrent/futures/_base.py in __get_result(self)
    382     def __get_result(self):
    383         if self._exception:
--> 384             raise self._exception
    385         else:
    386             return self._result

~/miniconda3/envs/py3.7.2/lib/python3.7/concurrent/futures/thread.py in run(self)
     55 
     56         try:
---> 57             result = self.fn(*self.args, **self.kwargs)
     58         except BaseException as exc:
     59             self.future.set_exception(exc)

~/Documents/mars_dev/mars/mars/utils.py in _inner(*args, **kwargs)
    456         def _inner(*args, **kwargs):
    457             with self:
--> 458                 return func(*args, **kwargs)
    459 
    460         return _inner

~/Documents/mars_dev/mars/mars/executor.py in _execute_operand(self, op)
    444             # so we pass the first operand's first output to Executor.handle
    445             first_op = ops[0]
--> 446             self.handle_op(first_op, results, self._mock)
    447 
    448             # update maximal memory usage during execution

~/Documents/mars_dev/mars/mars/executor.py in handle_op(self, *args, **kw)
    376 
    377     def handle_op(self, *args, **kw):
--> 378         return Executor.handle(*args, **kw)
    379 
    380     def _order_starts(self):

~/Documents/mars_dev/mars/mars/executor.py in handle(cls, op, results, mock)
    647                 # The `UFuncTypeError` was introduced by numpy#12593 since v1.17.0.
    648                 try:
--> 649                     return runner(results, op)
    650                 except UFuncTypeError as e:
    651                     raise TypeError(str(e)).with_traceback(sys.exc_info()[2]) from None

~/Documents/mars_dev/mars/mars/tensor/base/psrs.py in execute(cls, ctx, op)
    426                               num=n, endpoint=False).astype(int)
    427             slc = (slice(None),) * op.axis + (slc,)
--> 428             ctx[op.outputs[-1].key] = res[slc]
    429 
    430 

IndexError: index 0 is out of bounds for axis 0 with size 0
```



Discussion:

Resolving patch:
diff --git a/mars/tensor/base/psrs.py b/mars/tensor/base/psrs.py
--- a/mars/tensor/base/psrs.py
+++ b/mars/tensor/base/psrs.py
@@ -19,8 +19,8 @@
 
 from ... import opcodes as OperandDef
 from ...tiles import TilesError
-from ...serialize import ValueType, Int32Field, Int64Field, \
-    ListField, StringField, BoolField
+from ...serialize import ValueType, Int32Field, \
+    ListField, StringField, BoolField, AnyField
 from ...operands import OperandStage
 from ...utils import get_shuffle_input_keys_idxes, flatten, stack_back
 from ..core import TensorOrder
@@ -351,7 +351,7 @@ class PSRSSortRegularSample(TensorOperand, TensorOperandMixin):
     _kind = StringField('kind')
     _return_indices = BoolField('return_indices')
     _n_partition = Int32Field('n_partition')
-    _axis_offset = Int64Field('axis_offset')
+    _axis_offset = AnyField('axis_offset')
 
     def __init__(self, axis=None, order=None, kind=None, return_indices=None,
                  n_partition=None, axis_offset=None, dtype=None, gpu=None, **kw):
@@ -393,6 +393,11 @@ def execute(cls, ctx, op):
         (a,), device_id, xp = as_same_device(
             [ctx[c.key] for c in op.inputs], device=op.device, ret_extra=True)
 
+        if len(a) == 0:
+            # when chunk is empty, return the empty chunk itself
+            ctx[op.outputs[0].key] = ctx[op.outputs[-1].key] = a
+            return
+
         with device(device_id):
             n = op.n_partition
             w = a.shape[op.axis] * 1.0 / (n + 1)
@@ -453,12 +458,12 @@ def kind(self):
     @classmethod
     def execute(cls, ctx, op):
         inputs, device_id, xp = as_same_device(
-            [ctx[c.key] for c in op.inputs], device=op.device, ret_extra=True)
+            [ctx[c.key] for c in op.inputs if len(ctx[c.key]) > 0], device=op.device, ret_extra=True)
 
         with device(device_id):
             a = xp.concatenate(inputs, axis=op.axis)
             p = len(inputs)
-            assert a.shape[op.axis] == p ** 2
+            assert a.shape[op.axis] == p * len(op.inputs)
 
             if op.kind is not None:
                 # sort
@@ -470,10 +475,9 @@ def execute(cls, ctx, op):
                 a.partition(kth, axis=op.axis)
 
             select = xp.linspace(p - 1, a.shape[op.axis] - 1,
-                                 num=p - 1, endpoint=False).astype(int)
+                                 num=len(op.inputs) - 1, endpoint=False).astype(int)
             slc = (slice(None),) * op.axis + (select,)
-            ctx[op.outputs[0].key] = result = a[slc]
-            assert result.shape[op.axis] == p - 1
+            ctx[op.outputs[0].key] = a[slc]
 
 
 class PSRSShuffle(TensorMapReduceOperand, TensorOperandMixin):
diff --git a/mars/web/session.py b/mars/web/session.py
--- a/mars/web/session.py
+++ b/mars/web/session.py
@@ -149,7 +149,8 @@ def _check_response_finished(self, graph_url, timeout=None):
             resp = self._req_session.get(graph_url, params={'wait_timeout': timeout})
         except requests.ConnectionError as ex:
             err_msg = str(ex)
-            if 'ConnectionResetError' in err_msg or 'Connection refused' in err_msg:
+            if 'ConnectionResetError' in err_msg or 'Connection refused' in err_msg or \
+                    'Connection aborted' in err_msg:
                 return False
             raise
 


Test patch:
diff --git a/mars/tensor/base/tests/test_base_execute.py b/mars/tensor/base/tests/test_base_execute.py
--- a/mars/tensor/base/tests/test_base_execute.py
+++ b/mars/tensor/base/tests/test_base_execute.py
@@ -1196,6 +1196,16 @@ def testSortExecution(self):
         res = self.executor.execute_tensor(a, concat=True)[0]
         np.testing.assert_array_equal(res, np.sort(np.sort(raw, axis=1), axis=0))
 
+        # test with empty chunk
+        raw = np.random.rand(20, 10)
+        raw[:, :8] = 1
+        a = tensor(raw, chunk_size=5)
+        filtered = a[a < 1]
+        filtered.sort()
+
+        res = self.executor.execute_tensor(filtered, concat=True)[0]
+        np.testing.assert_array_equal(res, np.sort(raw[raw < 1]))
+
     def testSortIndicesExecution(self):
         # only 1 chunk when axis = -1
         raw = np.random.rand(100, 10)