Repository: modin-project/modin

Issue:
`groupby(by).size()` implementation

<!--
You can obtain the Modin version with

python -c "import modin; print(modin.__version__)"
-->

### Describe the problem
<!-- Describe the problem clearly here. -->
`df.groupby(col).size()` is implemented using pandas, but it is not warning the user. Also, this operation returns a `pandas.Series`, which should never happen in Modin. This can be easily implemented based on the existing implementation.

There are also no tests for `DataFrameGroupBy.size`, which is related to this issue.

### Source code / logs

https://github.com/modin-project/modin/blob/e9bae1967ed3eddabc66241e73cc3cb6695fcd8d/modin/pandas/groupby.py#L321-L322
<!-- 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
@@ -1227,12 +1227,7 @@ def groupby_reduce(
             qc = self.getitem_column_array(self._modin_frame._numeric_columns(True))
         else:
             qc = self
-        first_column = qc.columns[0]
         as_index = groupby_args.get("as_index", True)
-        # When drop is False and as_index is False, we do not want to insert the `by`
-        # data as a new column in the dataframe. We will drop it.
-        drop_by = not drop
-
         # For simplicity we allow only one function to be passed in if both are the
         # same.
         if reduce_func is None:
@@ -1240,56 +1235,80 @@ def groupby_reduce(
             reduce_args = map_args
 
         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):
-                df = pandas.concat(
-                    [df] + [other[[o for o in other if o not in df]]], axis=1
+            def compute_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):
+                    df = pandas.concat(
+                        [df] + [other[[o for o in other if o not in df]]], axis=1
+                    )
+                    other = list(other.columns)
+                result = map_func(
+                    df.groupby(by=other, axis=axis, **groupby_args), **map_args
                 )
-                other = list(other.columns)
-            result = map_func(
-                df.groupby(by=other, axis=axis, **groupby_args), **map_args
-            )
-            if (
-                not isinstance(result.index, pandas.MultiIndex)
-                and result.index.name is not None
-                and result.index.name in result.columns
-            ):
-                result.index.name = "{}{}".format("_modin_groupby_", result.index.name)
-            return result
+                # The _modin_groupby_ prefix indicates that this is the first partition,
+                # and since we may need to insert the grouping data in the reduce phase
+                if (
+                    not isinstance(result.index, pandas.MultiIndex)
+                    and result.index.name is not None
+                    and result.index.name in result.columns
+                ):
+                    result.index.name = "{}{}".format(
+                        "_modin_groupby_", result.index.name
+                    )
+                return result
+
+            try:
+                return compute_map(df, other)
+            # This will happen with Arrow buffer read-only errors. We don't want to copy
+            # all the time, so this will try to fast-path the code first.
+            except (ValueError):
+                return compute_map(df.copy(), other.copy())
 
         def _reduce(df):
-            other_len = len(df.index.names)
-            df = df.reset_index(drop=False)
-            # See note above about setting `as_index`
-            groupby_args["as_index"] = as_index
-            if other_len > 1:
-                by_part = list(df.columns[0:other_len])
-            else:
-                by_part = df.columns[0]
-            result = reduce_func(
-                df.groupby(by=by_part, axis=axis, **groupby_args), **reduce_args
-            )
-            if (
-                not isinstance(result.index, pandas.MultiIndex)
-                and result.index.name is not None
-                and "_modin_groupby_" in result.index.name
-            ):
-                result.index.name = result.index.name[len("_modin_groupby_") :]
-            # Avoid inserting data after the first partition or if the data did not come
-            # from this query compiler.
-            if not as_index and (first_column not in df.columns or drop_by):
-                return result.drop(columns=by_part)
-            return result
+            def compute_reduce(df):
+                other_len = len(df.index.names)
+                df = df.reset_index(drop=False)
+                # See note above about setting `as_index`
+                groupby_args["as_index"] = as_index
+                if other_len > 1:
+                    by_part = list(df.columns[0:other_len])
+                else:
+                    by_part = df.columns[0]
+                result = reduce_func(
+                    df.groupby(by=by_part, axis=axis, **groupby_args), **reduce_args
+                )
+                if (
+                    not isinstance(result.index, pandas.MultiIndex)
+                    and result.index.name is not None
+                    and "_modin_groupby_" in result.index.name
+                ):
+                    result.index.name = result.index.name[len("_modin_groupby_") :]
+                if isinstance(by_part, str) and by_part in result.columns:
+                    if "_modin_groupby_" in by_part and drop:
+                        col_name = by_part[len("_modin_groupby_") :]
+                        new_result = result.drop(columns=col_name)
+                        new_result.columns = [
+                            col_name if "_modin_groupby_" in c else c
+                            for c in new_result.columns
+                        ]
+                        return new_result
+                    else:
+                        return result.drop(columns=by_part)
+                return result
+
+            try:
+                return compute_reduce(df)
+            # This will happen with Arrow buffer read-only errors. We don't want to copy
+            # all the time, so this will try to fast-path the code first.
+            except (ValueError):
+                return compute_reduce(df.copy())
 
         if axis == 0:
-            if not as_index and drop:
-                new_columns = by.columns.append(qc.columns)
-            else:
-                new_columns = qc.columns
+            new_columns = qc.columns
             new_index = None
         else:
             new_index = self.index
diff --git a/modin/pandas/groupby.py b/modin/pandas/groupby.py
--- a/modin/pandas/groupby.py
+++ b/modin/pandas/groupby.py
@@ -18,6 +18,7 @@
 
 from modin.error_message import ErrorMessage
 from .utils import _inherit_docstrings
+from .series import Series
 
 
 @_inherit_docstrings(
@@ -249,7 +250,9 @@ def cummax(self, axis=0, **kwargs):
 
     def apply(self, func, *args, **kwargs):
         return self._apply_agg_function(
-            lambda df: df.apply(func, *args, **kwargs), drop=self._as_index
+            # Grouping column in never dropped in groupby.apply, so drop=False
+            lambda df: df.apply(func, *args, **kwargs),
+            drop=False,
         )
 
     @property
@@ -270,7 +273,7 @@ def __getitem__(self, key):
         # special in which case SeriesGroupBy has to be returned. Such circumstances are when key equals to a single
         # column name and is not a list of column names or list of one column name.
         make_dataframe = True
-        if self._drop:
+        if self._drop and self._as_index:
             if not isinstance(key, list):
                 key = [key]
                 kwargs["squeeze"] = True
@@ -278,7 +281,17 @@ def __getitem__(self, key):
         # When `as_index` is False, pandas will always convert to a `DataFrame`, we
         # convert to a list here so that the result will be a `DataFrame`.
         elif not self._as_index and not isinstance(key, list):
-            key = [key]
+            # Sometimes `__getitem__` doesn't only get the item, it also gets the `by`
+            # column. This logic is here to ensure that we also get the `by` data so
+            # that it is there for `as_index=False`.
+            if (
+                isinstance(self._by, type(self._query_compiler))
+                and all(c in self._columns for c in self._by.columns)
+                and self._drop
+            ):
+                key = [key] + list(self._by.columns)
+            else:
+                key = [key]
         if isinstance(key, list) and (make_dataframe or not self._as_index):
             return DataFrameGroupBy(
                 self._df[key],
@@ -368,7 +381,42 @@ def all(self, **kwargs):
         )
 
     def size(self):
-        return pandas.Series({k: len(v) for k, v in self._index_grouped.items()})
+        if self._axis == 0:
+            if self._as_index:
+                work_object = self[self._df.columns[0]]
+            else:
+                # Size always works in as_index=True mode so it is necessary to make a copy
+                # of _kwargs and change as_index in it
+                kwargs = self._kwargs.copy()
+                kwargs["as_index"] = True
+                kwargs["squeeze"] = True
+                work_object = SeriesGroupBy(
+                    self._df[self._df.columns[0]],
+                    self._by,
+                    self._axis,
+                    idx_name=self._idx_name,
+                    drop=False,
+                    **kwargs,
+                )
+
+            result = work_object._groupby_reduce(
+                lambda df: pandas.DataFrame(df.size()),
+                lambda df: df.sum(),
+                numeric_only=False,
+            )
+            series_result = Series(query_compiler=result._query_compiler)
+            # Pandas does not name size() output
+            series_result.name = None
+            return series_result
+        else:
+            return DataFrameGroupBy(
+                self._df.T,
+                self._by,
+                0,
+                drop=self._drop,
+                idx_name=self._idx_name,
+                **self._kwargs,
+            ).size()
 
     def sum(self, **kwargs):
         return self._groupby_reduce(lambda df: df.sum(**kwargs), None)
@@ -515,8 +563,11 @@ def _groupby_reduce(
         # For aggregations, pandas behavior does this for the result.
         # For other operations it does not, so we wait until there is an aggregation to
         # actually perform this operation.
-        if self._idx_name is not None and drop and self._drop:
-            groupby_qc = self._query_compiler.drop(columns=[self._idx_name])
+        if drop and self._drop:
+            if self._as_index:
+                groupby_qc = self._query_compiler.drop(columns=self._by.columns)
+            else:
+                groupby_qc = self._query_compiler
         else:
             groupby_qc = self._query_compiler
 
@@ -605,8 +656,6 @@ def ndim(self):
 
     @property
     def _iter(self):
-        from .dataframe import Series
-
         group_ids = self._index_grouped.keys()
         if self._axis == 0:
             return (
@@ -632,3 +681,14 @@ def _iter(self):
                 )
                 for k in (sorted(group_ids) if self._sort else group_ids)
             )
+
+    def size(self):
+        result = self._groupby_reduce(
+            lambda df: pandas.DataFrame(df.size()),
+            lambda df: df.sum(),
+            numeric_only=False,
+        )
+        series_result = Series(query_compiler=result._query_compiler)
+        # Pandas does not name size() output
+        series_result.name = None
+        return series_result


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
@@ -16,7 +16,7 @@
 import numpy as np
 import modin.pandas as pd
 from modin.pandas.utils import from_pandas, to_pandas
-from .utils import df_equals
+from .utils import df_equals, check_df_columns_have_nans
 
 pd.DEFAULT_NPARTITIONS = 4
 
@@ -114,6 +114,8 @@ def test_mixed_dtypes_groupby(as_index):
 
         # TODO Add more apply functions
         apply_functions = [lambda df: df.sum(), min]
+        # Workaround for Pandas bug #34656. Recreate groupby object for Pandas
+        pandas_groupby = pandas_df.groupby(by=by[-1], as_index=as_index)
         for func in apply_functions:
             eval_apply(modin_groupby, pandas_groupby, func)
 
@@ -198,15 +200,46 @@ def test_mixed_dtypes_groupby(as_index):
 
 
 @pytest.mark.parametrize(
-    "by", [[1, 2, 1, 2], lambda x: x % 3, "col1", ["col1"], ["col1", "col2"]]
+    "by",
+    [
+        [1, 2, 1, 2],
+        lambda x: x % 3,
+        "col1",
+        ["col1"],
+        # col2 contains NaN, is it necessary to test functions like size()
+        "col2",
+        ["col2"],
+        pytest.param(
+            ["col1", "col2"],
+            marks=pytest.mark.xfail(reason="Excluded because of bug #1554"),
+        ),
+        pytest.param(
+            ["col2", "col4"],
+            marks=pytest.mark.xfail(reason="Excluded because of bug #1554"),
+        ),
+        pytest.param(
+            ["col4", "col2"],
+            marks=pytest.mark.xfail(reason="Excluded because of bug #1554"),
+        ),
+        pytest.param(
+            ["col3", "col4", "col2"],
+            marks=pytest.mark.xfail(reason="Excluded because of bug #1554"),
+        ),
+        # but cum* functions produce undefined results with NaNs so we need to test the same combinations without NaN too
+        ["col5"],
+        ["col1", "col5"],
+        ["col5", "col4"],
+        ["col4", "col5"],
+        ["col5", "col4", "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],
-            "col2": [4, 5, 6, 7],
-            "col3": [3, 8, 12, 10],
+            "col2": [4, 5, np.NaN, 7],
+            "col3": [np.NaN, np.NaN, 12, 10],
             "col4": [17, 13, 16, 15],
             "col5": [-4, -5, -6, -7],
         }
@@ -232,7 +265,13 @@ def test_simple_row_groupby(by, as_index):
     eval_min(modin_groupby, pandas_groupby)
     eval_general(modin_groupby, pandas_groupby, lambda df: df.idxmax(), is_default=True)
     eval_ndim(modin_groupby, pandas_groupby)
-    eval_cumsum(modin_groupby, pandas_groupby)
+    if not check_df_columns_have_nans(modin_df, by):
+        # cum* functions produce undefined results for columns with NaNs so we run them only when "by" columns contain no NaNs
+        eval_cumsum(modin_groupby, pandas_groupby)
+        eval_cummax(modin_groupby, pandas_groupby)
+        eval_cummin(modin_groupby, pandas_groupby)
+        eval_cumprod(modin_groupby, pandas_groupby)
+
     eval_general(
         modin_groupby,
         pandas_groupby,
@@ -240,22 +279,18 @@ def test_simple_row_groupby(by, as_index):
         modin_df_almost_equals_pandas,
         is_default=True,
     )
-    eval_cummax(modin_groupby, pandas_groupby)
 
-    # pandas is inconsistent between test environment and here, more investigation is
-    # required to understand why this is a mismatch because we default to pandas for
-    # this particular case.
-    if by != ["col1", "col2"]:
-        apply_functions = [lambda df: df.sum(), min]
-        for func in apply_functions:
-            eval_apply(modin_groupby, pandas_groupby, func)
+    # Workaround for Pandas bug #34656. Recreate groupby object for Pandas
+    pandas_groupby = pandas_df.groupby(by=by, as_index=as_index)
+    apply_functions = [lambda df: df.sum(), min]
+    for func in apply_functions:
+        eval_apply(modin_groupby, pandas_groupby, func)
 
     eval_dtypes(modin_groupby, pandas_groupby)
     eval_general(modin_groupby, pandas_groupby, lambda df: df.first(), is_default=True)
     eval_general(
         modin_groupby, pandas_groupby, lambda df: df.backfill(), is_default=True
     )
-    eval_cummin(modin_groupby, pandas_groupby)
     eval_general(modin_groupby, pandas_groupby, lambda df: df.bfill(), is_default=True)
     eval_general(modin_groupby, pandas_groupby, lambda df: df.idxmin(), is_default=True)
     eval_prod(modin_groupby, pandas_groupby)
@@ -285,7 +320,6 @@ def test_simple_row_groupby(by, as_index):
     eval_nunique(modin_groupby, pandas_groupby)
     eval_median(modin_groupby, pandas_groupby)
     eval_general(modin_groupby, pandas_groupby, lambda df: df.head(n), is_default=True)
-    eval_cumprod(modin_groupby, pandas_groupby)
     eval_general(
         modin_groupby,
         pandas_groupby,
@@ -294,9 +328,11 @@ def test_simple_row_groupby(by, as_index):
         is_default=True,
     )
 
-    transform_functions = [lambda df: df + 4, lambda df: -df - 10]
-    for func in transform_functions:
-        eval_transform(modin_groupby, pandas_groupby, func)
+    if not check_df_columns_have_nans(modin_df, by):
+        # Pandas groupby.transform does not work correctly with NaN values in grouping columns. See Pandas bug 17093.
+        transform_functions = [lambda df: df + 4, lambda df: -df - 10]
+        for func in transform_functions:
+            eval_transform(modin_groupby, pandas_groupby, func)
 
     pipe_functions = [lambda dfgb: dfgb.sum()]
     for func in pipe_functions:
@@ -311,6 +347,7 @@ def test_simple_row_groupby(by, as_index):
     )
     eval_fillna(modin_groupby, pandas_groupby)
     eval_count(modin_groupby, pandas_groupby)
+    eval_size(modin_groupby, pandas_groupby)
     eval_general(modin_groupby, pandas_groupby, lambda df: df.tail(n), is_default=True)
     eval_quantile(modin_groupby, pandas_groupby)
     eval_general(modin_groupby, pandas_groupby, lambda df: df.take(), is_default=True)
@@ -426,6 +463,7 @@ def test_single_group_row_groupby():
     )
     eval_fillna(modin_groupby, pandas_groupby)
     eval_count(modin_groupby, pandas_groupby)
+    eval_size(modin_groupby, pandas_groupby)
     eval_general(modin_groupby, pandas_groupby, lambda df: df.tail(n), is_default=True)
     eval_quantile(modin_groupby, pandas_groupby)
     eval_general(modin_groupby, pandas_groupby, lambda df: df.take(), is_default=True)
@@ -535,6 +573,7 @@ def test_large_row_groupby():
     )
     eval_fillna(modin_groupby, pandas_groupby)
     eval_count(modin_groupby, pandas_groupby)
+    eval_size(modin_groupby, pandas_groupby)
     eval_general(modin_groupby, pandas_groupby, lambda df: df.tail(n), is_default=True)
     eval_quantile(modin_groupby, pandas_groupby)
     eval_general(modin_groupby, pandas_groupby, lambda df: df.take(), is_default=True)
@@ -645,6 +684,7 @@ def test_simple_col_groupby():
     )
     eval_fillna(modin_groupby, pandas_groupby)
     eval_count(modin_groupby, pandas_groupby)
+    eval_size(modin_groupby, pandas_groupby)
     eval_general(modin_groupby, pandas_groupby, lambda df: df.take(), is_default=True)
     eval_groups(modin_groupby, pandas_groupby)
 
@@ -924,6 +964,10 @@ def eval_count(modin_groupby, pandas_groupby):
     df_equals(modin_groupby.count(), pandas_groupby.count())
 
 
+def eval_size(modin_groupby, pandas_groupby):
+    df_equals(modin_groupby.size(), pandas_groupby.size())
+
+
 def eval_pipe(modin_groupby, pandas_groupby, func):
     df_equals(modin_groupby.pipe(func), pandas_groupby.pipe(func))
 
diff --git a/modin/pandas/test/utils.py b/modin/pandas/test/utils.py
--- a/modin/pandas/test/utils.py
+++ b/modin/pandas/test/utils.py
@@ -528,3 +528,20 @@ def name_contains(test_name, vals):
         True if a substring in vals is in test_name, else False.
     """
     return any(val in test_name for val in vals)
+
+
+def check_df_columns_have_nans(df, cols):
+    """Checks if there are NaN values in specified columns of a dataframe.
+
+    :param df: Dataframe to check.
+    :param cols: One column name or list of column names.
+    :return:
+        True if specified columns of dataframe contains NaNs.
+    """
+    return (
+        pandas.api.types.is_list_like(cols)
+        and any(x in df.columns and df[x].hasnans for x in cols)
+        or not pandas.api.types.is_list_like(cols)
+        and cols in df.columns
+        and df[cols].hasnans
+    )