# Buggy code
def execute(cls, ctx, op):
    a = ctx[op.inputs[0].key]

    if op.sort_type == "sort_values":
        ctx[op.outputs[0].key] = res = execute_sort_values(a, op)
    else:
        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"
    ):
        # 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))
        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)
    w = int(res.shape[op.axis] // n)

    slc = (slice(None),) * op.axis + (slice(0, n * w, w),)
    if op.sort_type == "sort_values":
        # do regular sample
        if op.by is not None:
            ctx[op.outputs[-1].key] = res[by].iloc[slc]
        else:
            ctx[op.outputs[-1].key] = res.iloc[slc]
    else:
        # do regular sample
        ctx[op.outputs[-1].key] = res.iloc[slc]


# Traceback (KeyError)
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"

# Fixed code
def execute(cls, ctx, op):
    a = ctx[op.inputs[0].key]

    if op.sort_type == "sort_values":
        ctx[op.outputs[0].key] = res = execute_sort_values(a, op)
    else:
        ctx[op.outputs[0].key] = res = execute_sort_index(a, op)

    by = op.by
    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), 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([res] * num), op, by=by)
    w = int(res.shape[op.axis] // n)

    slc = (slice(None),) * op.axis + (slice(0, n * w, w),)
    if op.sort_type == "sort_values":
        # do regular sample
        if op.by is not None:
            ctx[op.outputs[-1].key] = res[by].iloc[slc]
        else:
            ctx[op.outputs[-1].key] = res.iloc[slc]
    else:
        # do regular sample
        ctx[op.outputs[-1].key] = res.iloc[slc]