# Buggy code
def loads(buf):
    mv = memoryview(buf)
    header = read_file_header(mv)
    compress = header.compress

    if compress == CompressType.NONE:
        data = buf[HEADER_LENGTH:]
    else:
        data = decompressors[compress](mv[HEADER_LENGTH:])

    if header.type == SerialType.ARROW:
        try:
            return pyarrow.deserialize(memoryview(data), mars_serialize_context())
        except pyarrow.lib.ArrowInvalid:  # pragma: no cover
            # reconstruct value from buffers of arrow components
            data_view = memoryview(data)
            meta_block_size = np.frombuffer(data_view[0:4], dtype="int32").item()
            meta = pickle.loads(data_view[4 : 4 + meta_block_size])  # nosec
            buffer_sizes = meta.pop("buffer_sizes")
            bounds = np.cumsum([4 + meta_block_size] + buffer_sizes)
            meta["data"] = [
                pyarrow.py_buffer(data_view[bounds[idx] : bounds[idx + 1]])
                for idx in range(len(buffer_sizes))
            ]
            return pyarrow.deserialize_components(meta, mars_serialize_context())
    else:
        return pickle.loads(data)


# 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 loads(buf):
    mv = memoryview(buf)
    header = read_file_header(mv)
    compress = header.compress

    if compress == CompressType.NONE:
        data = buf[HEADER_LENGTH:]
    else:
        data = decompressors[compress](mv[HEADER_LENGTH:])

    if header.type == SerialType.ARROW:
        try:
            return deserialize(memoryview(data))
        except pyarrow.lib.ArrowInvalid:  # pragma: no cover
            # reconstruct value from buffers of arrow components
            data_view = memoryview(data)
            meta_block_size = np.frombuffer(data_view[0:4], dtype="int32").item()
            meta = pickle.loads(data_view[4 : 4 + meta_block_size])  # nosec
            buffer_sizes = meta.pop("buffer_sizes")
            bounds = np.cumsum([4 + meta_block_size] + buffer_sizes)
            meta["data"] = [
                pyarrow.py_buffer(data_view[bounds[idx] : bounds[idx + 1]])
                for idx in range(len(buffer_sizes))
            ]
            return pyarrow.deserialize_components(meta, mars_serialize_context())
    else:
        return pickle.loads(data)