Repository: modin-project/modin

Issue:
`groupby` with `as_index=False` still drops the column from the dataset

### Describe the problem
<!-- Describe the problem clearly here. -->
When a `groupby` is done with `as_index=False`, Modin still drops the column. See below.

### Source code / logs
Example from taxi data:

```python
df.groupby('passenger_count', as_index=False).mean()[['passenger_count','total_amount']]
```

<!-- 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. -->



Discussion:

Resolving patch:
diff --git a/modin/backends/pandas/query_compiler.py b/modin/backends/pandas/query_compiler.py
--- a/modin/backends/pandas/query_compiler.py
+++ b/modin/backends/pandas/query_compiler.py
@@ -1176,14 +1176,19 @@ def groupby_reduce(
         reduce_func=None,
         reduce_args=None,
         numeric_only=True,
+        drop=False,
     ):
         assert isinstance(
             by, type(self)
         ), "Can only use groupby reduce with another Query Compiler"
 
         other_len = len(by.columns)
+        as_index = groupby_args.get("as_index", True)
 
         def _map(df, other):
+            # Set `as_index` to True to track the metadata of the grouping object
+            # It is used to make sure that between phases we are constructing the
+            # right index and placing columns in the correct order.
             groupby_args["as_index"] = True
             other = other.squeeze(axis=axis ^ 1)
             if isinstance(other, pandas.DataFrame):
@@ -1205,6 +1210,8 @@ def _map(df, other):
         if reduce_func is not None:
 
             def _reduce(df):
+                # See note above about setting `as_index`
+                groupby_args["as_index"] = True
                 if other_len > 1:
                     by = list(df.columns[0:other_len])
                 else:
@@ -1223,6 +1230,8 @@ def _reduce(df):
         else:
 
             def _reduce(df):
+                # See note above about setting `as_index`
+                groupby_args["as_index"] = True
                 if other_len > 1:
                     by = list(df.columns[0:other_len])
                 else:
@@ -1260,10 +1269,25 @@ def _reduce(df):
             new_columns=new_columns,
             new_index=new_index,
         )
-        return self.__constructor__(new_modin_frame)
+        result = self.__constructor__(new_modin_frame)
+        # Reset `as_index` because it was edited inplace.
+        groupby_args["as_index"] = as_index
+        if as_index:
+            return result
+        else:
+            if result.index.name is None or result.index.name in result.columns:
+                drop = False
+            return result.reset_index(drop=not drop)
+
+    def groupby_agg(self, by, axis, agg_func, groupby_args, agg_args, drop=False):
+        as_index = groupby_args.get("as_index", True)
 
-    def groupby_agg(self, by, axis, agg_func, groupby_args, agg_args):
         def groupby_agg_builder(df):
+            # Set `as_index` to True to track the metadata of the grouping object
+            # It is used to make sure that between phases we are constructing the
+            # right index and placing columns in the correct order.
+            groupby_args["as_index"] = True
+
             def compute_groupby(df):
                 grouped_df = df.groupby(by=by, axis=axis, **groupby_args)
                 try:
@@ -1285,7 +1309,15 @@ def compute_groupby(df):
         new_modin_frame = self._modin_frame._apply_full_axis(
             axis, lambda df: groupby_agg_builder(df)
         )
-        return self.__constructor__(new_modin_frame)
+        result = self.__constructor__(new_modin_frame)
+        # Reset `as_index` because it was edited inplace.
+        groupby_args["as_index"] = as_index
+        if as_index:
+            return result
+        else:
+            if result.index.name is None or result.index.name in result.columns:
+                drop = False
+            return result.reset_index(drop=not drop)
 
     # END Manual Partitioning methods
 
diff --git a/modin/pandas/groupby.py b/modin/pandas/groupby.py
--- a/modin/pandas/groupby.py
+++ b/modin/pandas/groupby.py
@@ -229,13 +229,15 @@ def cummax(self, axis=0, **kwargs):
         return result
 
     def apply(self, func, *args, **kwargs):
-        return self._apply_agg_function(lambda df: df.apply(func, *args, **kwargs))
+        return self._apply_agg_function(
+            lambda df: df.apply(func, *args, **kwargs), drop=self._as_index
+        )
 
     @property
     def dtypes(self):
         if self._axis == 1:
             raise ValueError("Cannot call dtypes on groupby with axis=1")
-        return self._apply_agg_function(lambda df: df.dtypes)
+        return self._apply_agg_function(lambda df: df.dtypes, drop=self._as_index)
 
     def first(self, **kwargs):
         return self._default_to_pandas(lambda df: df.first(**kwargs))
@@ -274,7 +276,9 @@ def aggregate(self, arg, *args, **kwargs):
             return self._default_to_pandas(
                 lambda df: df.aggregate(arg, *args, **kwargs)
             )
-        return self._apply_agg_function(lambda df: df.aggregate(arg, *args, **kwargs))
+        return self._apply_agg_function(
+            lambda df: df.aggregate(arg, *args, **kwargs), drop=self._as_index
+        )
 
     def last(self, **kwargs):
         return self._default_to_pandas(lambda df: df.last(**kwargs))
@@ -393,7 +397,9 @@ def corr(self, **kwargs):
         return self._default_to_pandas(lambda df: df.corr(**kwargs))
 
     def fillna(self, **kwargs):
-        result = self._apply_agg_function(lambda df: df.fillna(**kwargs))
+        result = self._apply_agg_function(
+            lambda df: df.fillna(**kwargs), drop=self._as_index
+        )
         # pandas does not name the index on fillna
         result.index.name = None
         return result
@@ -472,6 +478,7 @@ def _groupby_reduce(
                 reduce_func=reduce_func,
                 reduce_args=kwargs,
                 numeric_only=numeric_only,
+                drop=self._drop,
             )
         )
 
@@ -501,7 +508,9 @@ def _apply_agg_function(self, f, drop=True, **kwargs):
             groupby_qc = self._query_compiler.drop(columns=[self._idx_name])
         else:
             groupby_qc = self._query_compiler
-        new_manager = groupby_qc.groupby_agg(by, self._axis, f, self._kwargs, kwargs)
+        new_manager = groupby_qc.groupby_agg(
+            by, self._axis, f, self._kwargs, kwargs, drop=self._drop
+        )
         if self._idx_name is not None and self._as_index:
             new_manager.index.name = self._idx_name
         return DataFrame(query_compiler=new_manager)


Test patch:
diff --git a/modin/pandas/test/test_groupby.py b/modin/pandas/test/test_groupby.py
--- a/modin/pandas/test/test_groupby.py
+++ b/modin/pandas/test/test_groupby.py
@@ -25,7 +25,8 @@ def ray_groupby_equals_pandas(ray_groupby, pandas_groupby):
         df_equals(g1[1], g2[1])
 
 
-def test_mixed_dtypes_groupby():
+@pytest.mark.parametrize("as_index", [True, False])
+def test_mixed_dtypes_groupby(as_index):
     frame_data = np.random.randint(97, 198, size=(2 ** 6, 2 ** 4))
     pandas_df = pandas.DataFrame(frame_data).add_prefix("col")
     # Convert every other column to string
@@ -40,12 +41,11 @@ def test_mixed_dtypes_groupby():
     by_values = [("col1",), (lambda x: x % 2,), (ray_df["col0"], pandas_df["col0"])]
 
     for by in by_values:
-        ray_groupby = ray_df.groupby(by=by[0])
-        pandas_groupby = pandas_df.groupby(by=by[-1])
+        ray_groupby = ray_df.groupby(by=by[0], as_index=as_index)
+        pandas_groupby = pandas_df.groupby(by=by[-1], as_index=as_index)
 
         ray_groupby_equals_pandas(ray_groupby, pandas_groupby)
         eval_ngroups(ray_groupby, pandas_groupby)
-        eval_skew(ray_groupby, pandas_groupby)
         eval_ffill(ray_groupby, pandas_groupby)
         eval_sem(ray_groupby, pandas_groupby)
         eval_mean(ray_groupby, pandas_groupby)
@@ -69,7 +69,10 @@ def test_mixed_dtypes_groupby():
         eval_bfill(ray_groupby, pandas_groupby)
         eval_idxmin(ray_groupby, pandas_groupby)
         eval_prod(ray_groupby, pandas_groupby)
-        eval_std(ray_groupby, pandas_groupby)
+        if as_index:
+            eval_std(ray_groupby, pandas_groupby)
+            eval_var(ray_groupby, pandas_groupby)
+            eval_skew(ray_groupby, pandas_groupby)
 
         agg_functions = ["min", "max"]
         for func in agg_functions:
@@ -79,7 +82,6 @@ def test_mixed_dtypes_groupby():
         eval_last(ray_groupby, pandas_groupby)
         eval_mad(ray_groupby, pandas_groupby)
         eval_max(ray_groupby, pandas_groupby)
-        eval_var(ray_groupby, pandas_groupby)
         eval_len(ray_groupby, pandas_groupby)
         eval_sum(ray_groupby, pandas_groupby)
         eval_ngroup(ray_groupby, pandas_groupby)
@@ -89,7 +91,7 @@ def test_mixed_dtypes_groupby():
         eval_cumprod(ray_groupby, pandas_groupby)
         eval_cov(ray_groupby, pandas_groupby)
 
-        transform_functions = [lambda df: df + 4, lambda df: -df - 10]
+        transform_functions = [lambda df: df, lambda df: df + df]
         for func in transform_functions:
             eval_transform(ray_groupby, pandas_groupby, func)
 
@@ -107,7 +109,9 @@ def test_mixed_dtypes_groupby():
         eval_groups(ray_groupby, pandas_groupby)
 
 
-def test_simple_row_groupby():
+@pytest.mark.parametrize("by", [[1, 2, 1, 2], lambda x: x % 3, "col1"])
+@pytest.mark.parametrize("as_index", [True, False])
+def test_simple_row_groupby(by, as_index):
     pandas_df = pandas.DataFrame(
         {
             "col1": [0, 1, 2, 3],
@@ -119,75 +123,73 @@ def test_simple_row_groupby():
     )
 
     ray_df = from_pandas(pandas_df)
-
-    by_values = [[1, 2, 1, 2], lambda x: x % 3, "col1"]
     n = 1
-    for by in by_values:
-        ray_groupby = ray_df.groupby(by=by)
-        pandas_groupby = pandas_df.groupby(by=by)
+    ray_groupby = ray_df.groupby(by=by, as_index=as_index)
+    pandas_groupby = pandas_df.groupby(by=by, as_index=as_index)
 
-        ray_groupby_equals_pandas(ray_groupby, pandas_groupby)
-        eval_ngroups(ray_groupby, pandas_groupby)
-        eval_skew(ray_groupby, pandas_groupby)
-        eval_ffill(ray_groupby, pandas_groupby)
-        eval_sem(ray_groupby, pandas_groupby)
-        eval_mean(ray_groupby, pandas_groupby)
-        eval_any(ray_groupby, pandas_groupby)
-        eval_min(ray_groupby, pandas_groupby)
-        eval_idxmax(ray_groupby, pandas_groupby)
-        eval_ndim(ray_groupby, pandas_groupby)
-        eval_cumsum(ray_groupby, pandas_groupby)
-        eval_pct_change(ray_groupby, pandas_groupby)
-        eval_cummax(ray_groupby, pandas_groupby)
+    ray_groupby_equals_pandas(ray_groupby, pandas_groupby)
+    eval_ngroups(ray_groupby, pandas_groupby)
+    eval_ffill(ray_groupby, pandas_groupby)
+    eval_sem(ray_groupby, pandas_groupby)
+    eval_mean(ray_groupby, pandas_groupby)
+    eval_any(ray_groupby, pandas_groupby)
+    eval_min(ray_groupby, pandas_groupby)
+    eval_idxmax(ray_groupby, pandas_groupby)
+    eval_ndim(ray_groupby, pandas_groupby)
+    eval_cumsum(ray_groupby, pandas_groupby)
+    eval_pct_change(ray_groupby, pandas_groupby)
+    eval_cummax(ray_groupby, pandas_groupby)
 
-        apply_functions = [lambda df: df.sum(), min]
-        for func in apply_functions:
-            eval_apply(ray_groupby, pandas_groupby, func)
+    apply_functions = [lambda df: df.sum(), min]
+    for func in apply_functions:
+        eval_apply(ray_groupby, pandas_groupby, func)
 
-        eval_dtypes(ray_groupby, pandas_groupby)
-        eval_first(ray_groupby, pandas_groupby)
-        eval_backfill(ray_groupby, pandas_groupby)
-        eval_cummin(ray_groupby, pandas_groupby)
-        eval_bfill(ray_groupby, pandas_groupby)
-        eval_idxmin(ray_groupby, pandas_groupby)
-        eval_prod(ray_groupby, pandas_groupby)
+    eval_dtypes(ray_groupby, pandas_groupby)
+    eval_first(ray_groupby, pandas_groupby)
+    eval_backfill(ray_groupby, pandas_groupby)
+    eval_cummin(ray_groupby, pandas_groupby)
+    eval_bfill(ray_groupby, pandas_groupby)
+    eval_idxmin(ray_groupby, pandas_groupby)
+    eval_prod(ray_groupby, pandas_groupby)
+    if as_index:
         eval_std(ray_groupby, pandas_groupby)
+        eval_var(ray_groupby, pandas_groupby)
+        eval_skew(ray_groupby, pandas_groupby)
 
-        agg_functions = ["min", "max"]
-        for func in agg_functions:
-            eval_agg(ray_groupby, pandas_groupby, func)
-            eval_aggregate(ray_groupby, pandas_groupby, func)
+    agg_functions = ["min", "max"]
+    for func in agg_functions:
+        eval_agg(ray_groupby, pandas_groupby, func)
+        eval_aggregate(ray_groupby, pandas_groupby, func)
 
-        eval_last(ray_groupby, pandas_groupby)
-        eval_mad(ray_groupby, pandas_groupby)
-        eval_rank(ray_groupby, pandas_groupby)
-        eval_max(ray_groupby, pandas_groupby)
-        eval_var(ray_groupby, pandas_groupby)
-        eval_len(ray_groupby, pandas_groupby)
-        eval_sum(ray_groupby, pandas_groupby)
-        eval_ngroup(ray_groupby, pandas_groupby)
-        eval_nunique(ray_groupby, pandas_groupby)
-        eval_median(ray_groupby, pandas_groupby)
-        eval_head(ray_groupby, pandas_groupby, n)
-        eval_cumprod(ray_groupby, pandas_groupby)
-        eval_cov(ray_groupby, pandas_groupby)
+    eval_last(ray_groupby, pandas_groupby)
+    eval_mad(ray_groupby, pandas_groupby)
+    eval_rank(ray_groupby, pandas_groupby)
+    eval_max(ray_groupby, pandas_groupby)
+    eval_len(ray_groupby, pandas_groupby)
+    eval_sum(ray_groupby, pandas_groupby)
+    eval_ngroup(ray_groupby, pandas_groupby)
+    eval_nunique(ray_groupby, pandas_groupby)
+    eval_median(ray_groupby, pandas_groupby)
+    eval_head(ray_groupby, pandas_groupby, n)
+    eval_cumprod(ray_groupby, pandas_groupby)
+    eval_cov(ray_groupby, pandas_groupby)
 
-        transform_functions = [lambda df: df + 4, lambda df: -df - 10]
-        for func in transform_functions:
-            eval_transform(ray_groupby, pandas_groupby, func)
+    transform_functions = [lambda df: df + 4, lambda df: -df - 10]
+    for func in transform_functions:
+        eval_transform(ray_groupby, pandas_groupby, func)
 
-        pipe_functions = [lambda dfgb: dfgb.sum()]
-        for func in pipe_functions:
-            eval_pipe(ray_groupby, pandas_groupby, func)
+    pipe_functions = [lambda dfgb: dfgb.sum()]
+    for func in pipe_functions:
+        eval_pipe(ray_groupby, pandas_groupby, func)
 
-        eval_corr(ray_groupby, pandas_groupby)
-        eval_fillna(ray_groupby, pandas_groupby)
-        eval_count(ray_groupby, pandas_groupby)
-        eval_tail(ray_groupby, pandas_groupby, n)
-        eval_quantile(ray_groupby, pandas_groupby)
-        eval_take(ray_groupby, pandas_groupby)
-        eval___getattr__(ray_groupby, pandas_groupby)
-        eval_groups(ray_groupby, pandas_groupby)
+    eval_corr(ray_groupby, pandas_groupby)
+    eval_fillna(ray_groupby, pandas_groupby)
+    eval_count(ray_groupby, pandas_groupby)
+    eval_tail(ray_groupby, pandas_groupby, n)
+    eval_quantile(ray_groupby, pandas_groupby)
+    eval_take(ray_groupby, pandas_groupby)
+    eval___getattr__(ray_groupby, pandas_groupby)
+    eval_groups(ray_groupby, pandas_groupby)
 
 
 def test_single_group_row_groupby():