Repository: databricks/koalas

Issue:
Duplicated key when groupby -> apply is used
Hello, I opened this thread in [SO](https://stackoverflow.com/questions/60790960/koalas-groupby-apply-returns-cannot-insert-key-already-exists), but viewing other issues reported here ([[1]](https://github.com/databricks/koalas/issues/1300), [[2]](https://github.com/databricks/koalas/issues/1007)), I'm suspecting that this could also be a bug.

I'm not sure if I'm missing something about the groupby - apply when moving from pandas to koalas, but the only requirement that we have is that the function to be applied cannot be type annotated since it will be used in different contexts, with different dataframes, so we don't know all the columns in advance. We're aware that koalas must run the function to infer the schema before applying the function, and we're ok with that overhead
Solves issue with key handling during groupby - apply
Solves #1361 

When the function to be applied after a `groupby` operation doesn't have type annotations, koalas has to infer the schema of the result. In this context koalas raises an exception due to the duplication between index and column names in `pandas.groupby.apply` [[1]](https://github.com/pandas-dev/pandas/issues/32384).

One option to deal with this issue is get rid of the duplicated columns, getting this results:

```python
>>> df = pd.DataFrame({'d': [1.0, 1.0, 1.0, 2.0, 2.0, 2.0], 'v': [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]})
>>> df
     d    v
0  1.0  1.0
1  1.0  2.0
2  1.0  3.0
3  2.0  4.0
4  2.0  5.0
5  2.0  6.0

>>> df.groupby('d').sum()
        v
d
1.0   6.0
2.0  15.0

>>> df.groupby('d').apply(sum)
       d     v
d
1.0  3.0   6.0
2.0  6.0  15.0

# Before PR
>>> ks.from_pandas(df).groupby('d').apply(sum)
...
ValueError: cannot insert d, already exists

# After PR
>>> ks.from_pandas(df).groupby('d').apply(sum)
       v
d
1.0  6.0
2.0  15.0
```

Another option that come to mind would be to rename the columns used as groupby keys (see below), but either way its still 'inconsistent' with pandas behavior.

```python
>>> ks.from_pandas(df).groupby('d').apply(sum)
       d_dup     v
d
1.0  3.0   6.0
2.0  6.0  15.0
```

how should koalas handle this?


Discussion:
Hi again, I encountered this [closed issue in pandas](https://github.com/pandas-dev/pandas/issues/32384), it seems that the grouping column is kept when one performs apply, and it is explicitly tested for (so it's a feature of pandas)

The problem arises then koala's apply has to infer the schema, runs the pandas groupby - apply and then initiates a koalas.DataFrame with the result (that after groupby-apply has index and a column with the same name):

```python
def apply(self, func):
    ...
    if should_infer_schema:
        # Here we execute with the first 1000 to get the return type.
        limit = get_option("compute.shortcut_limit")
        pdf = self._kdf.head(limit)._to_internal_pandas()
        pdf = pdf.groupby(input_groupnames).apply(func)  # <- pandas apply receive the key as column and index
        kdf = DataFrame(pdf) #  <- ValueError: cannot insert "key", already exists
        return_schema = kdf._sdf.drop(*HIDDEN_COLUMNS).schema
```

for now I can drop the key column inside `func`, I'm not sure if koalas should handle the logic for deduplicating index/column names before the initialization of `ks.DataFrame`
Sincerely appreciate for the report an issue & detailed investigation 👍 

I think we also should handle this logic and add the test explicitly like pandas.

@ueshin , @HyukjinKwon , WDYT about this?

Or, maybe are you interested in creating a PR to fix this, @syonekura ?
@ueshin , @HyukjinKwon 

For saving your time, the summary of this issue is like the below.

- in pandas.
```python
>>> df = pd.DataFrame({'d': [1.0, 1.0, 1.0, 2.0, 2.0, 2.0], 'v': [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]})
>>> df
     d    v
0  1.0  1.0
1  1.0  2.0
2  1.0  3.0
3  2.0  4.0
4  2.0  5.0
5  2.0  6.0

>>> df.groupby('d').sum()
        v
d
1.0   6.0
2.0  15.0

>>> df.groupby('d').apply(sum)
       d     v
d
1.0  3.0   6.0
2.0  6.0  15.0
```

- but in koalas
```python
>>> ks.from_pandas(df).groupby('d').sum()
        v
d
1.0   6.0
2.0  15.0

>>> ks.from_pandas(df).groupby('d').apply(sum)
...
ValueError: cannot insert d, already exists
```
Thanks, yeah. It would be great to fix to match at this moment.

In Spark, it supports the duplicated columns in some contexts. In some contexts, it doesn't.
So, as a future plan, we're thinking about disallowing duplicated column & index names in general [as described in the best practice](https://koalas.readthedocs.io/en/latest/user_guide/best_practices.html#do-not-use-duplicated-column-names) due to the limitation of Spark.
@itholic I would love to submit a PR for this, I'll give it a try during this weekend
@syonekura Sure, let's try :smile_cat: 
# [Codecov](https://codecov.io/gh/databricks/koalas/pull/1375?src=pr&el=h1) Report
> Merging [#1375](https://codecov.io/gh/databricks/koalas/pull/1375?src=pr&el=desc) into [master](https://codecov.io/gh/databricks/koalas/commit/710b753987aaae46d45b961d013897f6142aea04&el=desc) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/databricks/koalas/pull/1375/graphs/tree.svg?width=650&height=150&src=pr&token=EtnrpE8dVx)](https://codecov.io/gh/databricks/koalas/pull/1375?src=pr&el=tree)

```diff
@@           Coverage Diff           @@
##           master    #1375   +/-   ##
=======================================
  Coverage   95.26%   95.26%           
=======================================
  Files          34       34           
  Lines        7785     7789    +4     
=======================================
+ Hits         7416     7420    +4     
  Misses        369      369           
```


| [Impacted Files](https://codecov.io/gh/databricks/koalas/pull/1375?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [databricks/koalas/groupby.py](https://codecov.io/gh/databricks/koalas/pull/1375/diff?src=pr&el=tree#diff-ZGF0YWJyaWNrcy9rb2FsYXMvZ3JvdXBieS5weQ==) | `90.54% <100.00%> (+0.07%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/databricks/koalas/pull/1375?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/databricks/koalas/pull/1375?src=pr&el=footer). Last update [710b753...ac187de](https://codecov.io/gh/databricks/koalas/pull/1375?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Resolving patch:
diff --git a/databricks/koalas/frame.py b/databricks/koalas/frame.py
--- a/databricks/koalas/frame.py
+++ b/databricks/koalas/frame.py
@@ -3274,7 +3274,7 @@ def rename(index):
                 idx = []
                 for l in level:
                     try:
-                        i = self._internal.index_spark_column_names.index(l)
+                        i = self._internal.index_names.index((l,))
                         idx.append(i)
                     except ValueError:
                         if multi_index:
@@ -3282,7 +3282,7 @@ def rename(index):
                         else:
                             raise KeyError(
                                 "Level unknown must be same as name ({})".format(
-                                    self._internal.index_spark_column_names[0]
+                                    name_like_string(self._internal.index_names[0])
                                 )
                             )
             else:
diff --git a/databricks/koalas/groupby.py b/databricks/koalas/groupby.py
--- a/databricks/koalas/groupby.py
+++ b/databricks/koalas/groupby.py
@@ -49,7 +49,6 @@
     HIDDEN_COLUMNS,
     NATURAL_ORDER_COLUMN_NAME,
     SPARK_INDEX_NAME_FORMAT,
-    SPARK_DEFAULT_INDEX_NAME,
 )
 from databricks.koalas.missing.groupby import (
     _MissingPandasLikeDataFrameGroupBy,
@@ -1073,32 +1072,12 @@ def rename_output(pdf):
                 # https://github.com/databricks/koalas/issues/628.
 
                 # TODO: deduplicate this logic with _InternalFrame.from_pandas
-                columns = pdf.columns
-
-                index = pdf.index
-
-                index_map = []
-                if isinstance(index, pd.MultiIndex):
-                    if index.names is None:
-                        index_map = [
-                            (SPARK_INDEX_NAME_FORMAT(i), None) for i in range(len(index.levels))
-                        ]
-                    else:
-                        index_map = [
-                            (SPARK_INDEX_NAME_FORMAT(i) if name is None else name, name)
-                            for i, name in enumerate(index.names)
-                        ]
-                else:
-                    index_map = [
-                        (
-                            index.name if index.name is not None else SPARK_DEFAULT_INDEX_NAME,
-                            index.name,
-                        )
-                    ]
-
-                new_index_columns = [index_column for index_column, _ in index_map]
-                new_data_columns = [str(col) for col in columns]
+                new_index_columns = [
+                    SPARK_INDEX_NAME_FORMAT(i) for i in range(len(pdf.index.names))
+                ]
+                new_data_columns = [name_like_string(col) for col in pdf.columns]
 
+                pdf.index.names = new_index_columns
                 reset_index = pdf.reset_index()
                 reset_index.columns = new_index_columns + new_data_columns
                 for name, col in reset_index.iteritems():
diff --git a/databricks/koalas/internal.py b/databricks/koalas/internal.py
--- a/databricks/koalas/internal.py
+++ b/databricks/koalas/internal.py
@@ -446,16 +446,17 @@ def __init__(
         >>> internal = kdf[('a', 'y')]._internal
 
         >>> internal._sdf.show()  # doctest: +NORMALIZE_WHITESPACE +ELLIPSIS
-        +-----------+-----------+------+------+------+...
-        |row_index_a|row_index_b|(a, x)|(a, y)|(b, z)|...
-        +-----------+-----------+------+------+------+...
-        |        foo|        bar|     1|     2|     3|...
-        |        foo|        bar|     4|     5|     6|...
-        |        zoo|        bar|     7|     8|     9|...
-        +-----------+-----------+------+------+------+...
+        +-----------------+-----------------+------+------+------+...
+        |__index_level_0__|__index_level_1__|(a, x)|(a, y)|(b, z)|...
+        +-----------------+-----------------+------+------+------+...
+        |              foo|              bar|     1|     2|     3|...
+        |              foo|              bar|     4|     5|     6|...
+        |              zoo|              bar|     7|     8|     9|...
+        +-----------------+-----------------+------+------+------+...
 
         >>> internal._index_map  # doctest: +NORMALIZE_WHITESPACE
-        OrderedDict([('row_index_a', ('row_index_a',)), ('row_index_b', ('row_index_b',)),
+        OrderedDict([('__index_level_0__', ('row_index_a',)),
+         ('__index_level_1__', ('row_index_b',)),
          ('(a, x)', ('a', 'x'))])
 
         >>> internal._column_labels
@@ -991,43 +992,19 @@ def from_pandas(pdf: pd.DataFrame) -> "_InternalFrame":
             column_labels = None
         column_label_names = columns.names
 
-        index = pdf.index
-
-        index_map = []  # type: ignore
-        if isinstance(index, pd.MultiIndex):
-            if index.names is None:
-                index_map = [(SPARK_INDEX_NAME_FORMAT(i), None) for i in range(len(index.levels))]
-            else:
-                index_map = [
-                    (
-                        SPARK_INDEX_NAME_FORMAT(i)  # type: ignore
-                        if name is None
-                        else name_like_string(name),
-                        name if name is None or isinstance(name, tuple) else (name,),
-                    )
-                    for i, name in enumerate(index.names)
-                ]
-        else:
-            name = index.name
-            index_map = [
-                (
-                    name_like_string(name)  # type: ignore
-                    if name is not None
-                    else SPARK_DEFAULT_INDEX_NAME,
-                    name if name is None or isinstance(name, tuple) else (name,),
-                )
-            ]
-
-        index_columns = [index_column for index_column, _ in index_map]
+        index_names = [
+            name if name is None or isinstance(name, tuple) else (name,) for name in pdf.index.names
+        ]
+        index_columns = [SPARK_INDEX_NAME_FORMAT(i) for i in range(len(index_names))]
 
+        pdf = pdf.copy()
+        pdf.index.names = index_columns
         reset_index = pdf.reset_index()
         reset_index.columns = index_columns + data_columns
         schema = StructType(
             [
                 StructField(
-                    name_like_string(name),
-                    infer_pd_series_spark_type(col),
-                    nullable=bool(col.isnull().any()),
+                    name, infer_pd_series_spark_type(col), nullable=bool(col.isnull().any()),
                 )
                 for name, col in reset_index.iteritems()
             ]
@@ -1040,7 +1017,7 @@ def from_pandas(pdf: pd.DataFrame) -> "_InternalFrame":
         sdf = default_session().createDataFrame(reset_index, schema=schema)
         return _InternalFrame(
             spark_frame=sdf,
-            index_map=OrderedDict(index_map),
+            index_map=OrderedDict(zip(index_columns, index_names)),
             column_labels=column_labels,
             data_spark_columns=[scol_for(sdf, col) for col in data_columns],
             column_label_names=column_label_names,


Test patch:
diff --git a/databricks/koalas/tests/test_groupby.py b/databricks/koalas/tests/test_groupby.py
--- a/databricks/koalas/tests/test_groupby.py
+++ b/databricks/koalas/tests/test_groupby.py
@@ -1161,6 +1161,21 @@ def test_apply_with_new_dataframe_without_shortcut(self):
         with option_context("compute.shortcut_limit", 0):
             self.test_apply_with_new_dataframe()
 
+    def test_apply_key_handling(self):
+        pdf = pd.DataFrame(
+            {"d": [1.0, 1.0, 1.0, 2.0, 2.0, 2.0], "v": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]}
+        )
+        kdf = ks.from_pandas(pdf)
+
+        self.assert_eq(
+            kdf.groupby("d").apply(sum).sort_index(), pdf.groupby("d").apply(sum).sort_index()
+        )
+
+        with ks.option_context("compute.shortcut_limit", 1):
+            self.assert_eq(
+                kdf.groupby("d").apply(sum).sort_index(), pdf.groupby("d").apply(sum).sort_index()
+            )
+
     def test_transform(self):
         pdf = pd.DataFrame(
             {"a": [1, 2, 3, 4, 5, 6], "b": [1, 1, 2, 3, 5, 8], "c": [1, 4, 9, 16, 25, 36]},
diff --git a/databricks/koalas/tests/test_internal.py b/databricks/koalas/tests/test_internal.py
--- a/databricks/koalas/tests/test_internal.py
+++ b/databricks/koalas/tests/test_internal.py
@@ -17,7 +17,11 @@
 
 import pandas as pd
 
-from databricks.koalas.internal import _InternalFrame, SPARK_DEFAULT_INDEX_NAME
+from databricks.koalas.internal import (
+    _InternalFrame,
+    SPARK_DEFAULT_INDEX_NAME,
+    SPARK_INDEX_NAME_FORMAT,
+)
 from databricks.koalas.testing.utils import ReusedSQLTestCase, SQLTestUtils
 
 
@@ -43,7 +47,8 @@ def test_from_pandas(self):
         sdf = internal.spark_frame
 
         self.assert_eq(
-            internal.index_map, OrderedDict([(SPARK_DEFAULT_INDEX_NAME, None), ("a", ("a",))])
+            internal.index_map,
+            OrderedDict([(SPARK_INDEX_NAME_FORMAT(0), None), (SPARK_INDEX_NAME_FORMAT(1), ("a",))]),
         )
         self.assert_eq(internal.column_labels, [("b",)])
         self.assert_eq(internal.data_spark_column_names, ["b"])
@@ -58,7 +63,8 @@ def test_from_pandas(self):
         sdf = internal.spark_frame
 
         self.assert_eq(
-            internal.index_map, OrderedDict([(SPARK_DEFAULT_INDEX_NAME, None), ("a", ("a",))])
+            internal.index_map,
+            OrderedDict([(SPARK_INDEX_NAME_FORMAT(0), None), (SPARK_INDEX_NAME_FORMAT(1), ("a",))]),
         )
         self.assert_eq(internal.column_labels, [("x", "b")])
         self.assert_eq(internal.data_spark_column_names, ["(x, b)"])