# Buggy code
def _get_grouped(cls, op: "DataFrameGroupByAgg", df, ctx, copy=False, grouper=None):
    if copy:
        df = df.copy()

    params = op.groupby_params.copy()
    params.pop("as_index", None)
    selection = params.pop("selection", None)

    if grouper is not None:
        params["by"] = grouper
        params.pop("level", None)
    elif isinstance(params.get("by"), list):
        new_by = []
        for v in params["by"]:
            if isinstance(v, Base):
                new_by.append(ctx[v.key])
            else:
                new_by.append(v)
        params["by"] = new_by

    if op.stage == OperandStage.agg:
        grouped = df.groupby(**params)
    else:
        # for the intermediate phases, do not sort
        params["sort"] = False
        grouped = df.groupby(**params)

    if selection:
        grouped = grouped[selection]
    return grouped


# Traceback (KeyError)
Attempt 1: Unexpected error KeyError occurred in executing operand bc8dccd428eb0f6261420866f7206b73 in 0.0.0.0:23252
Traceback (most recent call last):
File "/Users/wenjun.swj/Code/mars/mars/promise.py", line 378, in _wrapped
return func(*args, **kwargs)
File "/Users/wenjun.swj/Code/mars/mars/utils.py", line 383, in _wrapped
return func(*args, **kwargs)
File "/Users/wenjun.swj/Code/mars/mars/worker/execution.py", line 501, in execute_graph
graph_record = self._graph_records[(session_id, graph_key)] = GraphExecutionRecord(
File "/Users/wenjun.swj/Code/mars/mars/worker/execution.py", line 57, in __init__
graph = self.graph = deserialize_graph(graph_serialized)
File "/Users/wenjun.swj/Code/mars/mars/utils.py", line 317, in deserialize_graph
return graph_cls.from_json(json_obj)
File "mars/graph.pyx", line 440, in mars.graph.DirectedGraph.from_json
return cls.deserialize(SerializableGraph.from_json(json_obj))
File "mars/serialize/core.pyx", line 718, in mars.serialize.core.Serializable.from_json
return cls.deserialize(provider, obj)
File "mars/serialize/core.pyx", line 689, in mars.serialize.core.Serializable.deserialize
[cb(key_to_instance) for cb in callbacks]
File "mars/serialize/jsonserializer.pyx", line 787, in mars.serialize.jsonserializer.JsonSerializeProvider.deserialize_field.cb
o = subs[val.key, val.id]
KeyError: ('b6888c78a929d77f42d7f3953fc813d9', '140581826655232')

# Fixed code
def _get_grouped(cls, op: "DataFrameGroupByAgg", df, ctx, copy=False, grouper=None):
    if copy:
        df = df.copy()

    params = op.groupby_params.copy()
    params.pop("as_index", None)
    selection = params.pop("selection", None)

    if grouper is not None:
        params["by"] = grouper
        params.pop("level", None)
    elif isinstance(params.get("by"), list):
        new_by = []
        for v in params["by"]:
            if isinstance(v, Base):
                new_by.append(ctx[v.key])
            else:
                new_by.append(v)
        params["by"] = new_by

    if op.stage == OperandStage.agg:
        grouped = df.groupby(**params)
    else:
        # for the intermediate phases, do not sort
        params["sort"] = False
        grouped = df.groupby(**params)

    if selection is not None:
        grouped = grouped[selection]
    return grouped