# Buggy code
def groupby_apply(
    groupby, func, *args, dtypes=None, index=None, output_types=None, **kwargs
):
    # todo this can be done with sort_index implemented
    if not groupby.op.groupby_params.get("as_index", True):
        raise NotImplementedError("apply when set_index == False is not supported")
    op = GroupByApply(func=func, args=args, kwds=kwargs, output_types=output_types)
    return op(groupby, dtypes=dtypes, index=index)


# Traceback (AssertionError)
Traceback (most recent call last):
File "/Users/wenjun/miniconda3/lib/python3.8/unittest/case.py", line 60, in testPartExecutor
yield
File "/Users/wenjun/miniconda3/lib/python3.8/unittest/case.py", line 676, in run
self._callTestMethod(testMethod)
File "/Users/wenjun/miniconda3/lib/python3.8/unittest/case.py", line 633, in _callTestMethod
method()
File "/Users/wenjun/Code/mars/mars/dataframe/groupby/tests/test_groupby_execution.py", line 401, in testGroupByApply
pd.testing.assert_frame_equal(self.executor.execute_dataframe(applied, concat=True)[0],
File "/Users/wenjun/Code/mars/mars/tests/core.py", line 686, in execute_tileable
result = super().execute_tileable(tileable, *args, **kwargs)
File "/Users/wenjun/Code/mars/mars/utils.py", line 439, in _inner
return func(*args, **kwargs)
File "/Users/wenjun/Code/mars/mars/executor.py", line 720, in execute_tileable
ret = self.execute_graph(chunk_graph, result_keys, n_parallel=n_parallel or n_thread,
File "/Users/wenjun/Code/mars/mars/tests/core.py", line 673, in execute_graph
return super().execute_graph(graph, keys, **kw)
File "/Users/wenjun/Code/mars/mars/executor.py", line 693, in execute_graph
res = graph_execution.execute(retval)
File "/Users/wenjun/Code/mars/mars/executor.py", line 574, in execute
future.result()
File "/Users/wenjun/miniconda3/lib/python3.8/concurrent/futures/_base.py", line 439, in result
return self.__get_result()
File "/Users/wenjun/miniconda3/lib/python3.8/concurrent/futures/_base.py", line 388, in __get_result
raise self._exception
File "/Users/wenjun/miniconda3/lib/python3.8/concurrent/futures/thread.py", line 57, in run
result = self.fn(*self.args, **self.kwargs)
File "/Users/wenjun/Code/mars/mars/tests/core.py", line 591, in _execute_operand
super()._execute_operand(op)
File "/Users/wenjun/Code/mars/mars/utils.py", line 439, in _inner
return func(*args, **kwargs)
File "/Users/wenjun/Code/mars/mars/executor.py", line 446, in _execute_operand
self.handle_op(first_op, results, self._mock)
File "/Users/wenjun/Code/mars/mars/executor.py", line 378, in handle_op
return Executor.handle(*args, **kw)
File "/Users/wenjun/Code/mars/mars/executor.py", line 644, in handle
return runner(results, op)
File "/Users/wenjun/Code/mars/mars/dataframe/groupby/apply.py", line 63, in execute
assert len(applied.index) == 1
AssertionError

# Fixed code
def groupby_apply(
    groupby, func, *args, dtypes=None, index=None, output_type=None, **kwargs
):
    # todo this can be done with sort_index implemented
    if not groupby.op.groupby_params.get("as_index", True):
        raise NotImplementedError("apply when set_index == False is not supported")

    output_types = kwargs.pop("output_types", None)
    object_type = kwargs.pop("object_type", None)
    output_types = validate_output_types(
        output_types=output_types, output_type=output_type, object_type=object_type
    )

    op = GroupByApply(func=func, args=args, kwds=kwargs, output_types=output_types)
    return op(groupby, dtypes=dtypes, index=index)