Repository: modin-project/modin

Issue:
groupby does not always set column and index name correctly

### Describe the problem
<!-- Describe the problem clearly here. -->
groupby does not maintain the correct columns when the columns have changed. The index name also needs to be set correctly.

### Source code / logs
<!-- Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached. Try to provide a reproducible test case that is the bare minimum necessary to generate the problem. -->
`age_groups  = age_data.groupby('YEARS_BINNED').mean()`


Discussion:

Resolving patch:
diff --git a/modin/pandas/dataframe.py b/modin/pandas/dataframe.py
--- a/modin/pandas/dataframe.py
+++ b/modin/pandas/dataframe.py
@@ -356,9 +356,11 @@ def groupby(
             A new DataFrame resulting from the groupby.
         """
         axis = pandas.DataFrame()._get_axis_number(axis)
+        idx_name = ""
         if callable(by):
             by = by(self.index)
         elif isinstance(by, string_types):
+            idx_name = by
             by = self.__getitem__(by).values.tolist()
         elif is_list_like(by):
             if isinstance(by, pandas.Series):
@@ -378,7 +380,16 @@ def groupby(
         from .groupby import DataFrameGroupBy
 
         return DataFrameGroupBy(
-            self, by, axis, level, as_index, sort, group_keys, squeeze, **kwargs
+            self,
+            by,
+            axis,
+            level,
+            as_index,
+            sort,
+            group_keys,
+            squeeze,
+            idx_name,
+            **kwargs
         )
 
     def sum(
diff --git a/modin/pandas/groupby.py b/modin/pandas/groupby.py
--- a/modin/pandas/groupby.py
+++ b/modin/pandas/groupby.py
@@ -6,6 +6,7 @@
 import pandas.core.groupby
 from pandas.core.dtypes.common import is_list_like
 import pandas.core.common as com
+import numpy as np
 
 from modin.error_message import ErrorMessage
 from .utils import _inherit_docstrings
@@ -20,7 +21,17 @@
 )
 class DataFrameGroupBy(object):
     def __init__(
-        self, df, by, axis, level, as_index, sort, group_keys, squeeze, **kwargs
+        self,
+        df,
+        by,
+        axis,
+        level,
+        as_index,
+        sort,
+        group_keys,
+        squeeze,
+        idx_name,
+        **kwargs
     ):
 
         self._axis = axis
@@ -32,6 +43,7 @@ def __init__(
         # This tells us whether or not there are multiple columns/rows in the groupby
         self._is_multi_by = all(obj in self._df for obj in self._by)
         self._level = level
+        self._idx_name = idx_name
         self._kwargs = {
             "sort": sort,
             "as_index": as_index,
@@ -141,7 +153,9 @@ def sem(self, ddof=1):
         return self._default_to_pandas(lambda df: df.sem(ddof=ddof))
 
     def mean(self, *args, **kwargs):
-        return self._apply_agg_function(lambda df: df.mean(*args, **kwargs))
+        return self._apply_agg_function(
+            lambda df: df.mean(*args, **kwargs), numeric=True
+        )
 
     def any(self):
         return self._apply_agg_function(lambda df: df.any())
@@ -281,7 +295,7 @@ def size(self):
         return self._apply_agg_function(lambda df: df.size())
 
     def sum(self, **kwargs):
-        return self._apply_agg_function(lambda df: df.sum(**kwargs))
+        return self._apply_agg_function(lambda df: df.sum(**kwargs), numeric=True)
 
     def __unicode__(self):
         return self._default_to_pandas(lambda df: df.__unicode__())
@@ -332,7 +346,7 @@ def resample(self, rule, *args, **kwargs):
         return self._default_to_pandas(lambda df: df.resample(rule, *args, **kwargs))
 
     def median(self, **kwargs):
-        return self._apply_agg_function(lambda df: df.median(**kwargs))
+        return self._apply_agg_function(lambda df: df.median(**kwargs), numeric=True)
 
     def head(self, n=5):
         return self._default_to_pandas(lambda df: df.head(n))
@@ -393,7 +407,7 @@ def diff(self):
     def take(self, **kwargs):
         return self._default_to_pandas(lambda df: df.take(**kwargs))
 
-    def _apply_agg_function(self, f, **kwargs):
+    def _apply_agg_function(self, f, numeric=False, **kwargs):
         """Perform aggregation and combine stages based on a given function.
 
         Args:
@@ -411,6 +425,22 @@ def _apply_agg_function(self, f, **kwargs):
         new_manager = self._query_compiler.groupby_agg(
             self._by, self._axis, f, self._kwargs, kwargs
         )
+        if self._idx_name != "":
+            new_manager.index.name = self._idx_name
+            new_columns = self._df.columns.drop(self._idx_name)
+        # We preserve columns only if the grouped axis is the index
+        elif self._axis == 0:
+            new_columns = self._df.columns
+        # We just keep everything the same if it is column groups
+        else:
+            new_columns = new_manager.columns
+        if numeric and self._axis == 0:
+            new_columns = [
+                col
+                for col in new_columns
+                if np.issubdtype(self._df.dtypes[col], np.number)
+            ]
+        new_manager.columns = new_columns
         return DataFrame(query_compiler=new_manager)
 
     def _default_to_pandas(self, f, **kwargs):


Test patch: