# Buggy code
def _execute_dataframe_map(cls, ctx, op):
    a, pivots = [ctx[c.key] for c in op.inputs]
    out = op.outputs[0]

    if isinstance(a, pd.DataFrame):
        # use numpy.searchsorted to find split positions.
        by = op.by

        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:
            by += [distinct_col]

        records = a[by].to_records(index=False)
        p_records = pivots.to_records(index=False)
        if op.ascending:
            poses = records.searchsorted(p_records, side="right")
        else:
            poses = len(records) - records[::-1].searchsorted(p_records, side="right")
        del records, p_records

        poses = (None,) + tuple(poses) + (None,)
        for i in range(op.n_partition):
            values = a.iloc[poses[i] : poses[i + 1]]
            ctx[(out.key, str(i))] = values
    else:  # pragma: no cover
        # for cudf, find split positions in loops.
        if op.ascending:
            pivots.append(a.iloc[-1][op.by])
            for i in range(op.n_partition):
                selected = a
                for label in op.by:
                    selected = selected.loc[a[label] <= pivots.iloc[i][label]]
                ctx[(out.key, str(i))] = selected
        else:
            pivots.append(a.iloc[-1][op.by])
            for i in range(op.n_partition):
                selected = a
                for label in op.by:
                    selected = selected.loc[a[label] >= pivots.iloc[i][label]]
                ctx[(out.key, str(i))] = selected


# 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_dataframe_map(cls, ctx, op):
    a, pivots = [ctx[c.key] for c in op.inputs]
    out = op.outputs[0]

    if isinstance(a, pd.DataFrame):
        # use numpy.searchsorted to find split positions.
        by = op.by

        distinct_col = (
            _PSRS_DISTINCT_COL
            if a.columns.nlevels == 1
            else (_PSRS_DISTINCT_COL,) + ("",) * (a.columns.nlevels - 1)
        )
        if distinct_col in a.columns:
            by += [distinct_col]

        records = a[by].to_records(index=False)
        p_records = pivots.to_records(index=False)
        if op.ascending:
            poses = records.searchsorted(p_records, side="right")
        else:
            poses = len(records) - records[::-1].searchsorted(p_records, side="right")
        del records, p_records

        poses = (None,) + tuple(poses) + (None,)
        for i in range(op.n_partition):
            values = a.iloc[poses[i] : poses[i + 1]]
            ctx[(out.key, str(i))] = values
    else:  # pragma: no cover
        # for cudf, find split positions in loops.
        if op.ascending:
            pivots.append(a.iloc[-1][op.by])
            for i in range(op.n_partition):
                selected = a
                for label in op.by:
                    selected = selected.loc[a[label] <= pivots.iloc[i][label]]
                ctx[(out.key, str(i))] = selected
        else:
            pivots.append(a.iloc[-1][op.by])
            for i in range(op.n_partition):
                selected = a
                for label in op.by:
                    selected = selected.loc[a[label] >= pivots.iloc[i][label]]
                ctx[(out.key, str(i))] = selected