Repository: modin-project/modin

Issue:
groupby() with concat() result columns can't be added
<!--
General questions should be asked on the mailing list modin-dev@googlegroups.com.

Before submitting an issue, please fill out the following form.
-->

### System information
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Ubuntu 18.04.2
- **Modin installed from (source or binary)**: pip
- **Modin version**: 0.6.3
- **Python version**: 3.7.3
- **Exact command to reproduce**: Snippet below

<!--
You can obtain the Modin version with

python -c "import modin; print(modin.__version__)"
-->

### Describe the problem
Cannot add a combination of results to a DF composed of grouped results.

### Source code / logs
```python
import modin.pandas as pd #import pandas as pd
df = pd.DataFrame(data={"i":[1,1 ], "a":[0, 0], "b":[0, 0]})
g = df.groupby(by="i")
cat = pd.concat([g["a"].first(), g["b"].first()], keys=["x", "y"], axis=1)
cat["z"] = cat.x + cat.y
```
doesn't work because the subtraction results in two columns being returned with modin vs one in pandas, probably something is mixed up with a multi-index on the columns.


Discussion:
Thanks @ecoughlan, you are right that it is setting a multi-index on the columns. When `keys` is set for `concat`, Modin seems to be adding a level to the columns instead of replacing the label names with the values. It should be a straightforward fix.

As a side note, `cat = pd.concat([g["a"].first(), g["b"].first()], keys=["x", "y"], axis=1)` could be rewritten to `cat = g.first().set_axis(["x", "y"], axis=1, inplace=False)`. I realize that isn't the point of the issue, but I want to add this as it is generally less efficient to do the way written.

Thanks again for posting!

Resolving patch:
diff --git a/modin/pandas/groupby.py b/modin/pandas/groupby.py
--- a/modin/pandas/groupby.py
+++ b/modin/pandas/groupby.py
@@ -230,8 +230,7 @@ def backfill(self, limit=None):
         return self.bfill(limit)
 
     def __getitem__(self, key):
-        # This operation requires a SeriesGroupBy Object
-        return self._default_to_pandas(lambda df: df.__getitem__(key))
+        return SeriesGroupBy(self._default_to_pandas(lambda df: df.__getitem__(key)))
 
     def cummin(self, axis=0, **kwargs):
         return self._apply_agg_function(lambda df: df.cummin(axis=axis, **kwargs))
@@ -491,3 +490,37 @@ def groupby_on_multiple_columns(df):
             return f(df.groupby(by=by, axis=self._axis, **self._kwargs), **kwargs)
 
         return self._df._default_to_pandas(groupby_on_multiple_columns)
+
+
+class SeriesGroupBy(object):  # pragma: no cover
+    def __init__(self, pandas_groupby_obj):
+        self._pandas_groupby_obj = pandas_groupby_obj
+
+    def __getattribute__(self, item):
+        if item in ["_pandas_groupby_obj"]:
+            return object.__getattribute__(self, item)
+        return_val = self._pandas_groupby_obj.__getattribute__(item)
+        from .series import Series
+        from .dataframe import DataFrame
+
+        if isinstance(return_val, pandas.Series):
+            return Series(return_val)
+        elif isinstance(return_val, pandas.DataFrame):
+            return DataFrame(return_val)
+        elif callable(return_val):
+
+            # This wraps the pandas callable and intercepts the return value before it
+            # is given back to the user to re-distribute it.
+            def wrapper(*args, **kwargs):
+                pandas_return_val = return_val(*args, **kwargs)
+                if isinstance(pandas_return_val, pandas.Series):
+                    return Series(pandas_return_val)
+                elif isinstance(pandas_return_val, pandas.DataFrame):
+                    return DataFrame(pandas_return_val)
+                else:
+                    return pandas_return_val
+
+            return wrapper
+
+        else:
+            return return_val
diff --git a/modin/pandas/series.py b/modin/pandas/series.py
--- a/modin/pandas/series.py
+++ b/modin/pandas/series.py
@@ -572,17 +572,21 @@ def groupby(
         observed=False,
         **kwargs
     ):
-        return self._default_to_pandas(
-            pandas.Series.groupby,
-            by=by,
-            axis=axis,
-            level=level,
-            as_index=as_index,
-            sort=sort,
-            group_keys=group_keys,
-            squeeze=squeeze,
-            observed=observed,
-            **kwargs
+        from .groupby import SeriesGroupBy
+
+        return SeriesGroupBy(
+            self._default_to_pandas(
+                pandas.Series.groupby,
+                by=by,
+                axis=axis,
+                level=level,
+                as_index=as_index,
+                sort=sort,
+                group_keys=group_keys,
+                squeeze=squeeze,
+                observed=observed,
+                **kwargs
+            )
         )
 
     def gt(self, other, level=None, fill_value=None, axis=0):


Test patch:
diff --git a/modin/pandas/test/test_series.py b/modin/pandas/test/test_series.py
--- a/modin/pandas/test/test_series.py
+++ b/modin/pandas/test/test_series.py
@@ -1431,12 +1431,14 @@ def test_get_values(data):
         modin_series.get_values()
 
 
-@pytest.mark.skip(reason="Using pandas Series.")
-def test_groupby():
-    modin_series = create_test_series()
+@pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)
+def test_groupby(data):
+    modin_series, _ = create_test_series(data)
 
-    with pytest.raises(NotImplementedError):
-        modin_series.groupby(None, None, None, None, None, None, None)
+    with pytest.warns(UserWarning):
+        result = modin_series.groupby(modin_series).count()
+
+    assert isinstance(result, pd.Series)
 
 
 @pytest.mark.parametrize("data", test_data_values, ids=test_data_keys)