Repository: mars-project/mars

Issue:
[BUG] `KeyError` was raised when execute dataframe getitem
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**To Reproduce**
Append an empty dataframe leads to a `KeyError` when execute getitem.

``` Python
In [9]: d = md.DataFrame(pd.DataFrame([[0, 1], [2, 3]], columns=['col1', 'col2']))                  

In [10]: d2 = d.append(md.DataFrame(pd.DataFrame(columns=['col1', 'col2'])))                        

In [11]: def f(df): 
    ...:     return pd.Series([1, 2], index=['c', 'd']) 
    ...:                                                                                            

In [12]: d2.apply(f, axis=1)['c'].execute()                                                         
  0%|                                                                       | 0/100 [00:00<?, ?it/s]Failed to run subtask r8o56UzDolvyhNrfLbGHmWOA on band numa-0
Traceback (most recent call last):
  File "/Users/hekaisheng/miniconda3/lib/python3.7/site-packages/pandas/core/indexes/base.py", line 3080, in get_loc
    return self._engine.get_loc(casted_key)
  File "pandas/_libs/index.pyx", line 70, in pandas._libs.index.IndexEngine.get_loc
  File "pandas/_libs/index.pyx", line 101, in pandas._libs.index.IndexEngine.get_loc
  File "pandas/_libs/hashtable_class_helper.pxi", line 4554, in pandas._libs.hashtable.PyObjectHashTable.get_item
  File "pandas/_libs/hashtable_class_helper.pxi", line 4562, in pandas._libs.hashtable.PyObjectHashTable.get_item
KeyError: 'c'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/scheduling/worker/execution.py", line 266, in internal_run_subtask
    subtask, band_name, subtask_api, batch_quota_req)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/scheduling/worker/execution.py", line 361, in _retry_run_subtask
    return await _retry_run(subtask, subtask_info, _run_subtask_once)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/scheduling/worker/execution.py", line 84, in _retry_run
    raise ex
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/scheduling/worker/execution.py", line 68, in _retry_run
    return await target_async_func(*args)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/scheduling/worker/execution.py", line 319, in _run_subtask_once
    return await asyncio.shield(aiotask)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/subtask/api.py", line 59, in run_subtask_in_slot
    return await ref.run_subtask(subtask)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/backends/context.py", line 154, in send
    return self._process_result_message(result)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/backends/context.py", line 59, in _process_result_message
    raise message.error.with_traceback(message.traceback)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/backends/pool.py", line 496, in send
    result = await future
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/api.py", line 118, in __on_receive__
    return await super().__on_receive__(message)
  File "mars/oscar/core.pyx", line 351, in __on_receive__
    raise ex
  File "mars/oscar/core.pyx", line 345, in mars.oscar.core._BaseActor.__on_receive__
    return await self._handle_actor_result(result)
  File "mars/oscar/core.pyx", line 250, in _handle_actor_result
    result = list(dones)[0].result()
  File "mars/oscar/core.pyx", line 273, in _run_actor_async_generator
    with debug_async_timeout('actor_lock_timeout',
  File "mars/oscar/core.pyx", line 275, in mars.oscar.core._BaseActor._run_actor_async_generator
    async with self._lock:
  File "mars/oscar/core.pyx", line 279, in mars.oscar.core._BaseActor._run_actor_async_generator
    res = await gen.athrow(*res)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/subtask/worker/runner.py", line 102, in run_subtask
    result = yield self._running_processor.run(subtask)
  File "mars/oscar/core.pyx", line 284, in mars.oscar.core._BaseActor._run_actor_async_generator
    res = await self._handle_actor_result(res)
  File "mars/oscar/core.pyx", line 219, in _handle_actor_result
    result = await result
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/backends/context.py", line 154, in send
    return self._process_result_message(result)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/backends/context.py", line 59, in _process_result_message
    raise message.error.with_traceback(message.traceback)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/backends/pool.py", line 496, in send
    result = await future
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/api.py", line 118, in __on_receive__
    return await super().__on_receive__(message)
  File "mars/oscar/core.pyx", line 351, in __on_receive__
    raise ex
  File "mars/oscar/core.pyx", line 345, in mars.oscar.core._BaseActor.__on_receive__
    return await self._handle_actor_result(result)
  File "mars/oscar/core.pyx", line 250, in _handle_actor_result
    result = list(dones)[0].result()
  File "mars/oscar/core.pyx", line 273, in _run_actor_async_generator
    with debug_async_timeout('actor_lock_timeout',
  File "mars/oscar/core.pyx", line 275, in mars.oscar.core._BaseActor._run_actor_async_generator
    async with self._lock:
  File "mars/oscar/core.pyx", line 279, in mars.oscar.core._BaseActor._run_actor_async_generator
    res = await gen.athrow(*res)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/subtask/worker/processor.py", line 470, in run
    result = yield self._running_aio_task
  File "mars/oscar/core.pyx", line 284, in mars.oscar.core._BaseActor._run_actor_async_generator
    res = await self._handle_actor_result(res)
  File "mars/oscar/core.pyx", line 219, in _handle_actor_result
    result = await result
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/subtask/worker/processor.py", line 359, in run
    await self._execute_graph(chunk_graph)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/subtask/worker/processor.py", line 185, in _execute_graph
    await to_wait
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/lib/aio/_threads.py", line 36, in to_thread
    return await loop.run_in_executor(None, func_call)
  File "/Users/hekaisheng/miniconda3/lib/python3.7/concurrent/futures/thread.py", line 57, in run
    result = self.fn(*self.args, **self.kwargs)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/core/mode.py", line 79, in _inner
    return func(*args, **kwargs)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/subtask/worker/processor.py", line 158, in _execute_operand
    return execute(ctx, op)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/core/operand/core.py", line 473, in execute
    result = executor(results, op)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/dataframe/indexing/getitem.py", line 391, in execute
    ctx[op.outputs[0].key] = df[op.col_names]
  File "/Users/hekaisheng/miniconda3/lib/python3.7/site-packages/pandas/core/frame.py", line 3024, in __getitem__
    indexer = self.columns.get_loc(key)
  File "/Users/hekaisheng/miniconda3/lib/python3.7/site-packages/pandas/core/indexes/base.py", line 3082, in get_loc
    raise KeyError(key) from err
KeyError: 'c'
Subtask r8o56UzDolvyhNrfLbGHmWOA errored
Traceback (most recent call last):
  File "/Users/hekaisheng/miniconda3/lib/python3.7/site-packages/pandas/core/indexes/base.py", line 3080, in get_loc
    return self._engine.get_loc(casted_key)
  File "pandas/_libs/index.pyx", line 70, in pandas._libs.index.IndexEngine.get_loc
  File "pandas/_libs/index.pyx", line 101, in pandas._libs.index.IndexEngine.get_loc
  File "pandas/_libs/hashtable_class_helper.pxi", line 4554, in pandas._libs.hashtable.PyObjectHashTable.get_item
  File "pandas/_libs/hashtable_class_helper.pxi", line 4562, in pandas._libs.hashtable.PyObjectHashTable.get_item
KeyError: 'c'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/scheduling/worker/execution.py", line 266, in internal_run_subtask
    subtask, band_name, subtask_api, batch_quota_req)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/scheduling/worker/execution.py", line 361, in _retry_run_subtask
    return await _retry_run(subtask, subtask_info, _run_subtask_once)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/scheduling/worker/execution.py", line 84, in _retry_run
    raise ex
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/scheduling/worker/execution.py", line 68, in _retry_run
    return await target_async_func(*args)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/scheduling/worker/execution.py", line 319, in _run_subtask_once
    return await asyncio.shield(aiotask)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/subtask/api.py", line 59, in run_subtask_in_slot
    return await ref.run_subtask(subtask)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/backends/context.py", line 154, in send
    return self._process_result_message(result)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/backends/context.py", line 59, in _process_result_message
    raise message.error.with_traceback(message.traceback)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/backends/pool.py", line 496, in send
    result = await future
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/api.py", line 118, in __on_receive__
    return await super().__on_receive__(message)
  File "mars/oscar/core.pyx", line 351, in __on_receive__
    raise ex
  File "mars/oscar/core.pyx", line 345, in mars.oscar.core._BaseActor.__on_receive__
    return await self._handle_actor_result(result)
  File "mars/oscar/core.pyx", line 250, in _handle_actor_result
    result = list(dones)[0].result()
  File "mars/oscar/core.pyx", line 273, in _run_actor_async_generator
    with debug_async_timeout('actor_lock_timeout',
  File "mars/oscar/core.pyx", line 275, in mars.oscar.core._BaseActor._run_actor_async_generator
    async with self._lock:
  File "mars/oscar/core.pyx", line 279, in mars.oscar.core._BaseActor._run_actor_async_generator
    res = await gen.athrow(*res)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/subtask/worker/runner.py", line 102, in run_subtask
    result = yield self._running_processor.run(subtask)
  File "mars/oscar/core.pyx", line 284, in mars.oscar.core._BaseActor._run_actor_async_generator
    res = await self._handle_actor_result(res)
  File "mars/oscar/core.pyx", line 219, in _handle_actor_result
    result = await result
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/backends/context.py", line 154, in send
    return self._process_result_message(result)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/backends/context.py", line 59, in _process_result_message
    raise message.error.with_traceback(message.traceback)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/backends/pool.py", line 496, in send
    result = await future
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/oscar/api.py", line 118, in __on_receive__
    return await super().__on_receive__(message)
  File "mars/oscar/core.pyx", line 351, in __on_receive__
    raise ex
  File "mars/oscar/core.pyx", line 345, in mars.oscar.core._BaseActor.__on_receive__
    return await self._handle_actor_result(result)
  File "mars/oscar/core.pyx", line 250, in _handle_actor_result
    result = list(dones)[0].result()
  File "mars/oscar/core.pyx", line 273, in _run_actor_async_generator
    with debug_async_timeout('actor_lock_timeout',
  File "mars/oscar/core.pyx", line 275, in mars.oscar.core._BaseActor._run_actor_async_generator
    async with self._lock:
  File "mars/oscar/core.pyx", line 279, in mars.oscar.core._BaseActor._run_actor_async_generator
    res = await gen.athrow(*res)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/subtask/worker/processor.py", line 470, in run
    result = yield self._running_aio_task
  File "mars/oscar/core.pyx", line 284, in mars.oscar.core._BaseActor._run_actor_async_generator
    res = await self._handle_actor_result(res)
  File "mars/oscar/core.pyx", line 219, in _handle_actor_result
    result = await result
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/subtask/worker/processor.py", line 359, in run
    await self._execute_graph(chunk_graph)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/subtask/worker/processor.py", line 185, in _execute_graph
    await to_wait
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/lib/aio/_threads.py", line 36, in to_thread
    return await loop.run_in_executor(None, func_call)
  File "/Users/hekaisheng/miniconda3/lib/python3.7/concurrent/futures/thread.py", line 57, in run
    result = self.fn(*self.args, **self.kwargs)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/core/mode.py", line 79, in _inner
    return func(*args, **kwargs)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/services/subtask/worker/processor.py", line 158, in _execute_operand
    return execute(ctx, op)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/core/operand/core.py", line 473, in execute
    result = executor(results, op)
  File "/Users/hekaisheng/Documents/mars_dev/mars/mars/dataframe/indexing/getitem.py", line 391, in execute
    ctx[op.outputs[0].key] = df[op.col_names]
  File "/Users/hekaisheng/miniconda3/lib/python3.7/site-packages/pandas/core/frame.py", line 3024, in __getitem__
    indexer = self.columns.get_loc(key)
  File "/Users/hekaisheng/miniconda3/lib/python3.7/site-packages/pandas/core/indexes/base.py", line 3082, in get_loc
    raise KeyError(key) from err
KeyError: 'c'
```



Discussion:

![image](https://user-images.githubusercontent.com/12445254/132281425-54d72227-9c60-4155-ad72-d926748ff8b2.png)

`to_pandas` will not work either.
If concat empty df and another df with value, `concat_df.execute()` may update wrong dtypes.

```
In [23]: d = md.DataFrame(pd.DataFrame(columns=['col1', 'col2']))               

In [24]: d2 = md.DataFrame(pd.DataFrame([[0, 1], [2, 3]], columns=['col1', 'col2
    ...: ']))                                                                   

In [25]: d3 = d.append(d2)                                                      

In [26]: d4 = d3.apply(f, axis=1)                                               

In [27]: d4.execute()  # wrong                                                       
100%|██████████████████████████████████████| 100.0/100 [00:00<00:00, 796.97it/s]
Out[27]: 
  col1 col2    c    d
0  NaN  NaN  1.0  2.0
1  NaN  NaN  1.0  2.0


In [28]: d4.dtypes  # wrong                                                            
Out[28]: 
col1    object
col2    object
dtype: object
```
And I think `apply`, `groupby().apply()` as well as `map_chunk` are all effected.

Resolving patch:
diff --git a/mars/dataframe/base/apply.py b/mars/dataframe/base/apply.py
--- a/mars/dataframe/base/apply.py
+++ b/mars/dataframe/base/apply.py
@@ -26,7 +26,7 @@
 from ...utils import enter_current_session, quiet_stdio
 from ..operands import DataFrameOperandMixin, DataFrameOperand
 from ..utils import build_df, build_series, parse_index, validate_axis, \
-    validate_output_types, make_dtypes, make_dtype
+    validate_output_types, make_dtypes, make_dtype, build_empty_df, build_empty_series
 
 
 class ApplyOperand(DataFrameOperand, DataFrameOperandMixin):
@@ -86,13 +86,22 @@ def kwds(self):
     @enter_current_session
     def execute(cls, ctx, op):
         input_data = ctx[op.inputs[0].key]
+        out = op.outputs[0]
+        if len(input_data) == 0:
+            if op.output_types[0] == OutputType.dataframe:
+                ctx[out.key] = build_empty_df(out.dtypes)
+            else:
+                ctx[out.key] = build_empty_series(
+                    out.dtype, name=out.name)
+            return
+
         if isinstance(input_data, pd.DataFrame):
             result = input_data.apply(op.func, axis=op.axis, raw=op.raw, result_type=op.result_type,
                                       args=op.args, **op.kwds)
         else:
             result = input_data.apply(op.func, convert_dtype=op.convert_dtype, args=op.args,
                                       **op.kwds)
-        ctx[op.outputs[0].key] = result
+        ctx[out.key] = result
 
     @classmethod
     def _tile_df(cls, op):
@@ -285,7 +294,6 @@ def _call_series(self, series, dtypes=None, dtype=None, name=None, index=None):
                     else:
                         raise TypeError('Cannot determine dtypes, '
                                         'please specify `dtypes` as argument')
-
                 columns_value = parse_index(dtypes.index, store_data=True)
 
                 return self.new_dataframe([series], shape=(series.shape[0], len(dtypes)),
diff --git a/mars/dataframe/base/map_chunk.py b/mars/dataframe/base/map_chunk.py
--- a/mars/dataframe/base/map_chunk.py
+++ b/mars/dataframe/base/map_chunk.py
@@ -23,7 +23,7 @@
 from ...utils import enter_current_session, has_unknown_shape, quiet_stdio
 from ..operands import DataFrameOperand, DataFrameOperandMixin, OutputType
 from ..utils import build_df, build_empty_df, build_series, parse_index, \
-    validate_output_types
+    validate_output_types, build_empty_series
 
 
 class DataFrameMapChunk(DataFrameOperand, DataFrameOperandMixin):
@@ -175,12 +175,21 @@ def tile(cls, op: "DataFrameMapChunk"):
     @enter_current_session
     def execute(cls, ctx, op: "DataFrameMapChunk"):
         inp = ctx[op.input.key]
-        out_chunk = op.outputs[0]
+        out = op.outputs[0]
+        if len(inp) == 0:
+            if op.output_types[0] == OutputType.dataframe:
+                ctx[out.key] = build_empty_df(out.dtypes)
+            elif op.output_types[0] == OutputType.series:
+                ctx[out.key] = build_empty_series(
+                    out.dtype, name=out.name)
+            else:
+                raise ValueError(f'Chunk can not be empty except for dataframe/series.')
+            return
 
         kwargs = op.kwargs or dict()
         if op.with_chunk_index:
-            kwargs['chunk_index'] = out_chunk.index
-        ctx[out_chunk.key] = op.func(inp, *op.args, **kwargs)
+            kwargs['chunk_index'] = out.index
+        ctx[out.key] = op.func(inp, *op.args, **kwargs)
 
 
 def map_chunk(df_or_series, func, args=(), **kwargs):


Test patch:
diff --git a/mars/dataframe/tests/test_core.py b/mars/dataframe/tests/test_core.py
--- a/mars/dataframe/tests/test_core.py
+++ b/mars/dataframe/tests/test_core.py
@@ -208,6 +208,36 @@ def test_to_frame_or_series(setup):
     pd.testing.assert_series_equal(raw.to_series(name='new_name'), result)
 
 
+def test_to_frame_or_series_apply(setup):
+    df1 = DataFrame(pd.DataFrame([[0, 1], [2, 3]], columns=['col1', 'col2']))
+    df2 = df1.append(DataFrame(pd.DataFrame(columns=['col1', 'col2'])))
+    pd_df2 = df2.apply(lambda row: pd.Series([1, 2], index=['c', 'd']), axis=1).to_pandas()
+    assert pd_df2.columns.tolist() == ['c', 'd']
+
+    def f(df):
+        df['col3'] = df['col2']
+        return df
+
+    pd_df3 = df2.groupby(['col1']).apply(f).to_pandas()
+    assert pd_df3.columns.tolist() == ['col1', 'col2', 'col3']
+
+    pd_df4 = df2.map_chunk(lambda chunk_df: chunk_df.apply(
+        lambda row: pd.Series([1, 2], index=['c', 'd']), axis=1)).to_pandas()
+    assert pd_df4.columns.tolist() == ['c', 'd']
+
+    ser1 = Series(pd.Series(data={'a': 1, 'b': 2, 'c': 3}, index=['a', 'b', 'c']))
+    ser2 = ser1.append(Series(pd.Series(dtype=np.int64)))
+    pd_ser2 = ser2.apply(lambda v: str(v)).execute()
+    assert pd_ser2.dtype == object
+
+    ser3 = ser2.map_chunk(lambda chunk_series: chunk_series.apply(lambda x: float(x))).execute()
+
+    def check_dtype(s):
+        assert s.dtypes == np.float64
+        return s
+    ser3.map_chunk(check_dtype).execute()
+
+
 def test_assign(setup):
     rs = np.random.RandomState(0)
     raw = pd.DataFrame({"A": rs.rand(10), "B": rs.rand(10)})