# Buggy code
def execute_sort_values(data, op, inplace=None):
    if inplace is None:
        inplace = op.inplace
    # ignore_index is new in Pandas version 1.0.0.
    ignore_index = getattr(op, "ignore_index", False)
    if isinstance(data, (pd.DataFrame, pd.Series)):
        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
        if inplace:
            kwargs["inplace"] = True
            try:
                data.sort_values(**kwargs)
            except TypeError:  # pragma: no cover
                kwargs.pop("ignore_index", None)
                data.sort_values(**kwargs)
            return data
        else:
            try:
                return data.sort_values(**kwargs)
            except TypeError:  # pragma: no cover
                kwargs.pop("ignore_index", None)
                return data.sort_values(**kwargs)

    else:  # pragma: no cover
        # cudf doesn't support axis and kind
        if isinstance(data, cudf.DataFrame):
            return data.sort_values(
                op.by, ascending=op.ascending, na_position=op.na_position
            )
        else:
            return data.sort_values(ascending=op.ascending, na_position=op.na_position)


# 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_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.
    ignore_index = getattr(op, "ignore_index", False)
    if isinstance(data, (pd.DataFrame, pd.Series)):
        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"] = by if by is not None else op.by
        if inplace:
            kwargs["inplace"] = True
            try:
                data.sort_values(**kwargs)
            except TypeError:  # pragma: no cover
                kwargs.pop("ignore_index", None)
                data.sort_values(**kwargs)
            return data
        else:
            try:
                return data.sort_values(**kwargs)
            except TypeError:  # pragma: no cover
                kwargs.pop("ignore_index", None)
                return data.sort_values(**kwargs)

    else:  # pragma: no cover
        # cudf doesn't support axis and kind
        if isinstance(data, cudf.DataFrame):
            return data.sort_values(
                op.by, ascending=op.ascending, na_position=op.na_position
            )
        else:
            return data.sort_values(ascending=op.ascending, na_position=op.na_position)