# Buggy code
def df_apply(
    df,
    func,
    axis=0,
    raw=False,
    result_type=None,
    args=(),
    dtypes=None,
    output_type=None,
    index=None,
    elementwise=None,
    **kwds,
):
    if isinstance(func, (list, dict)):
        return df.aggregate(func)

    if isinstance(output_type, str):
        output_type = getattr(OutputType, output_type.lower())

    # calling member function
    if isinstance(func, str):
        func = getattr(df, func)
        sig = inspect.getfullargspec(func)
        if "axis" in sig.args:
            kwds["axis"] = axis
        return func(*args, **kwds)

    op = ApplyOperand(
        func=func,
        axis=axis,
        raw=raw,
        result_type=result_type,
        args=args,
        kwds=kwds,
        output_type=output_type,
        elementwise=elementwise,
    )
    return op(df, 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 df_apply(
    df,
    func,
    axis=0,
    raw=False,
    result_type=None,
    args=(),
    dtypes=None,
    output_type=None,
    index=None,
    elementwise=None,
    **kwds,
):
    if isinstance(func, (list, dict)):
        return df.aggregate(func)

    output_types = kwds.pop("output_types", None)
    object_type = kwds.pop("object_type", None)
    output_types = validate_output_types(
        output_type=output_type, output_types=output_types, object_type=object_type
    )
    output_type = output_types[0] if output_types else None

    # calling member function
    if isinstance(func, str):
        func = getattr(df, func)
        sig = inspect.getfullargspec(func)
        if "axis" in sig.args:
            kwds["axis"] = axis
        return func(*args, **kwds)

    op = ApplyOperand(
        func=func,
        axis=axis,
        raw=raw,
        result_type=result_type,
        args=args,
        kwds=kwds,
        output_types=output_type,
        elementwise=elementwise,
    )
    return op(df, dtypes=dtypes, index=index)