Repository: modin-project/modin

Issue:
BUG: groupby with timestamp in `by` raises KeyError
### Modin version checks

- [X] I have checked that this issue has not already been reported.

- [X] I have confirmed this bug exists on the latest released version of Modin.

- [X] I have confirmed this bug exists on the main branch of Modin. (In order to do this you can follow [this guide](https://modin.readthedocs.io/en/stable/getting_started/installation.html#installing-from-the-github-master-branch).)


### Reproducible Example

```python
import modin.pandas as pd

df = pd.DataFrame({'timestamp': [pd.to_datetime(1490195805, unit='s')], 'numeric': [0]})
print(df.groupby('timestamp').mean())
```


### Issue Description

I get KeyError in Modin, but the groupby works in pandas.

### Expected Behavior

pandas gives output like

```
                     numeric
timestamp
2017-03-22 15:16:45      0.0
```

### Error Logs

<details>

```python-traceback

UserWarning: Distributing <class 'dict'> object. This may take some time.
pd.---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
File ~/opt/anaconda3/envs/modin-dev/lib/python3.10/site-packages/pandas/core/indexes/base.py:3800, in Index.get_loc(self, key, method, tolerance)
   3799 try:
-> 3800     return self._engine.get_loc(casted_key)
   3801 except KeyError as err:

File ~/opt/anaconda3/envs/modin-dev/lib/python3.10/site-packages/pandas/_libs/index.pyx:138, in pandas._libs.index.IndexEngine.get_loc()

File ~/opt/anaconda3/envs/modin-dev/lib/python3.10/site-packages/pandas/_libs/index.pyx:165, in pandas._libs.index.IndexEngine.get_loc()

File pandas/_libs/hashtable_class_helper.pxi:5745, in pandas._libs.hashtable.PyObjectHashTable.get_item()

File pandas/_libs/hashtable_class_helper.pxi:5753, in pandas._libs.hashtable.PyObjectHashTable.get_item()

KeyError: 'timestamp'

The above exception was the direct cause of the following exception:

KeyError                                  Traceback (most recent call last)
Input In [5], in <cell line: 4>()
      1 import modin.pandas as pd
      3 df = pd.DataFrame({'timestamp': [pd.to_datetime(1490195805, unit='s')], 'numeric': [0]})
----> 4 print(df.groupby('timestamp').mean())

File ~/software_sources/modin/modin/logging/logger_decorator.py:128, in enable_logging.<locals>.decorator.<locals>.run_and_log(*args, **kwargs)
    113 """
    114 Compute function with logging if Modin logging is enabled.
    115
   (...)
    125 Any
    126 """
    127 if LogMode.get() == "disable":
--> 128     return obj(*args, **kwargs)
    130 logger = get_logger()
    131 logger_level = getattr(logger, log_level)

File ~/software_sources/modin/modin/pandas/groupby.py:138, in DataFrameGroupBy.mean(self, numeric_only)
    136 def mean(self, numeric_only=None):
    137     return self._check_index(
--> 138         self._wrap_aggregation(
    139             type(self._query_compiler).groupby_mean,
    140             numeric_only=numeric_only,
    141             agg_kwargs=dict(numeric_only=numeric_only),
    142         )
    143     )

File ~/software_sources/modin/modin/logging/logger_decorator.py:128, in enable_logging.<locals>.decorator.<locals>.run_and_log(*args, **kwargs)
    113 """
    114 Compute function with logging if Modin logging is enabled.
    115
   (...)
    125 Any
    126 """
    127 if LogMode.get() == "disable":
--> 128     return obj(*args, **kwargs)
    130 logger = get_logger()
    131 logger_level = getattr(logger, log_level)

File ~/software_sources/modin/modin/pandas/groupby.py:1095, in DataFrameGroupBy._wrap_aggregation(self, qc_method, numeric_only, agg_args, agg_kwargs, **kwargs)
   1091 else:
   1092     groupby_qc = self._query_compiler
   1094 result = type(self._df)(
-> 1095     query_compiler=qc_method(
   1096         groupby_qc,
   1097         by=self._by,
   1098         axis=self._axis,
   1099         groupby_kwargs=self._kwargs,
   1100         agg_args=agg_args,
   1101         agg_kwargs=agg_kwargs,
   1102         drop=self._drop,
   1103         **kwargs,
   1104     )
   1105 )
   1106 if self._squeeze:
   1107     return result.squeeze()

File ~/software_sources/modin/modin/logging/logger_decorator.py:128, in enable_logging.<locals>.decorator.<locals>.run_and_log(*args, **kwargs)
    113 """
    114 Compute function with logging if Modin logging is enabled.
    115
   (...)
    125 Any
    126 """
    127 if LogMode.get() == "disable":
--> 128     return obj(*args, **kwargs)
    130 logger = get_logger()
    131 logger_level = getattr(logger, log_level)

File ~/software_sources/modin/modin/core/storage_formats/pandas/query_compiler.py:2635, in PandasQueryCompiler.groupby_mean(self, by, axis, groupby_kwargs, agg_args, agg_kwargs, drop)
   2616 result = GroupByReduce.register(
   2617     lambda dfgb, **kwargs: pandas.concat(
   2618         [dfgb.sum(**kwargs), dfgb.count()],
   (...)
   2631     drop=drop,
   2632 )
   2634 if len(datetime_cols) > 0:
-> 2635     result = result.astype({col: dtype for col, dtype in datetime_cols.items()})
   2636 return result

File ~/software_sources/modin/modin/logging/logger_decorator.py:128, in enable_logging.<locals>.decorator.<locals>.run_and_log(*args, **kwargs)
    113 """
    114 Compute function with logging if Modin logging is enabled.
    115
   (...)
    125 Any
    126 """
    127 if LogMode.get() == "disable":
--> 128     return obj(*args, **kwargs)
    130 logger = get_logger()
    131 logger_level = getattr(logger, log_level)

File ~/software_sources/modin/modin/core/storage_formats/pandas/query_compiler.py:1537, in PandasQueryCompiler.astype(self, col_dtypes, **kwargs)
   1536 def astype(self, col_dtypes, **kwargs):
-> 1537     return self.__constructor__(self._modin_frame.astype(col_dtypes))

File ~/software_sources/modin/modin/logging/logger_decorator.py:128, in enable_logging.<locals>.decorator.<locals>.run_and_log(*args, **kwargs)
    113 """
    114 Compute function with logging if Modin logging is enabled.
    115
   (...)
    125 Any
    126 """
    127 if LogMode.get() == "disable":
--> 128     return obj(*args, **kwargs)
    130 logger = get_logger()
    131 logger_level = getattr(logger, log_level)

File ~/software_sources/modin/modin/core/dataframe/pandas/dataframe/dataframe.py:124, in lazy_metadata_decorator.<locals>.decorator.<locals>.run_f_on_minimally_updated_metadata(self, *args, **kwargs)
    122     elif apply_axis == "rows":
    123         obj._propagate_index_objs(axis=0)
--> 124 result = f(self, *args, **kwargs)
    125 if apply_axis is None and not transpose:
    126     result._deferred_index = self._deferred_index

File ~/software_sources/modin/modin/core/dataframe/pandas/dataframe/dataframe.py:1146, in PandasDataframe.astype(self, col_dtypes)
   1143 for i, column in enumerate(columns):
   1144     dtype = col_dtypes[column]
   1145     if (
-> 1146         not isinstance(dtype, type(self.dtypes[column]))
   1147         or dtype != self.dtypes[column]
   1148     ):
   1149         # Update the new dtype series to the proper pandas dtype
   1150         try:
   1151             new_dtype = np.dtype(dtype)

File ~/opt/anaconda3/envs/modin-dev/lib/python3.10/site-packages/pandas/core/series.py:982, in Series.__getitem__(self, key)
    979     return self._values[key]
    981 elif key_is_scalar:
--> 982     return self._get_value(key)
    984 if is_hashable(key):
    985     # Otherwise index.get_value will raise InvalidIndexError
    986     try:
    987         # For labels that don't resolve as scalars like tuples and frozensets

File ~/opt/anaconda3/envs/modin-dev/lib/python3.10/site-packages/pandas/core/series.py:1092, in Series._get_value(self, label, takeable)
   1089     return self._values[label]
   1091 # Similar to Index.get_value, but we do not fall back to positional
-> 1092 loc = self.index.get_loc(label)
   1093 return self.index._get_values_for_loc(self, loc, label)

File ~/opt/anaconda3/envs/modin-dev/lib/python3.10/site-packages/pandas/core/indexes/base.py:3802, in Index.get_loc(self, key, method, tolerance)
   3800     return self._engine.get_loc(casted_key)
   3801 except KeyError as err:
-> 3802     raise KeyError(key) from err
   3803 except TypeError:
   3804     # If we have a listlike key, _check_indexing_error will raise
   3805     #  InvalidIndexError. Otherwise we fall through and re-raise
   3806     #  the TypeError.
   3807     self._check_indexing_error(key)

KeyError: 'timestamp'

```

</details>


### Installed Versions

<details>
INSTALLED VERSIONS
------------------
commit           : 621bc10e16c21ca4270a4c2a25aa05f814f2686b
python           : 3.10.4.final.0
python-bits      : 64
OS               : Darwin
OS-release       : 21.5.0
Version          : Darwin Kernel Version 21.5.0: Tue Apr 26 21:08:22 PDT 2022; root:xnu-8020.121.3~4/RELEASE_X86_64
machine          : x86_64
processor        : i386
byteorder        : little
LC_ALL           : None
LANG             : en_US.UTF-8
LOCALE           : en_US.UTF-8

Modin dependencies
------------------
modin            : 0.16.0
ray              : 2.0.0
dask             : 2022.7.1
distributed      : 2022.7.1
hdk              : None

pandas dependencies
-------------------
pandas           : 1.5.0
numpy            : 1.23.2
pytz             : 2022.2.1
dateutil         : 2.8.2
setuptools       : 61.2.0
pip              : 22.2.2
Cython           : None
pytest           : 7.1.2
hypothesis       : None
sphinx           : 4.5.0
blosc            : None
feather          : 0.4.1
xlsxwriter       : None
lxml.etree       : 4.9.1
html5lib         : None
pymysql          : None
psycopg2         : None
jinja2           : 3.1.2
IPython          : 8.4.0
pandas_datareader: None
bs4              : 4.11.1
bottleneck       : None
brotli           : 1.0.9
fastparquet      : 0.8.1
fsspec           : 2022.7.1
gcsfs            : None
matplotlib       : 3.5.2
numba            : None
numexpr          : 2.8.3
odfpy            : None
openpyxl         : 3.0.10
pandas_gbq       : 0.17.7
pyarrow          : 8.0.0
pyreadstat       : None
pyxlsb           : None
s3fs             : 2022.7.1
scipy            : 1.9.0
snappy           : None
sqlalchemy       : 1.4.39
tables           : 3.7.0
tabulate         : None
xarray           : 2022.6.0
xlrd             : 2.0.1
xlwt             : None
zstandard        : None
tzdata           : None

</details>



Discussion:
The problem is that we convert the `by` timestamp column to int64 [here](https://github.com/modin-project/modin/blob/621bc10e16c21ca4270a4c2a25aa05f814f2686b/modin/core/storage_formats/pandas/query_compiler.py#L2589) because we think we have to mean it, but we don't have to mean it because it's a key and not a value for the groupby. The key column is not present in the groupby result, so we get a KeyError when we try to convert the type back to a datetime type [here](https://github.com/modin-project/modin/blob/621bc10e16c21ca4270a4c2a25aa05f814f2686b/modin/core/storage_formats/pandas/query_compiler.py#L2634).

I have confirmed that this is a regression due to a76e2a17a1c0168cffa97f08d164e48025817f05.

Resolving patch:
diff --git a/modin/core/storage_formats/pandas/query_compiler.py b/modin/core/storage_formats/pandas/query_compiler.py
--- a/modin/core/storage_formats/pandas/query_compiler.py
+++ b/modin/core/storage_formats/pandas/query_compiler.py
@@ -2524,6 +2524,40 @@ def _callable_func(self, func, axis, *args, **kwargs):
     # nature. They require certain data to exist on the same partition, and
     # after the shuffle, there should be only a local map required.
 
+    def _groupby_internal_columns(self, by, drop):
+        """
+        Extract internal columns from by argument of groupby.
+
+        Parameters
+        ----------
+        by : BaseQueryCompiler, column or index label, Grouper or list
+        drop : bool
+            Indicates whether or not by data came from self frame.
+            True, by data came from self. False, external by data.
+
+        Returns
+        -------
+        by : list of BaseQueryCompiler, column or index label, or Grouper
+        internal_by : list of str
+            List of internal column name to be dropped during groupby.
+        """
+        if isinstance(by, type(self)):
+            if drop:
+                internal_by = by.columns
+                by = [by]
+            else:
+                internal_by = []
+                by = [by]
+        else:
+            if not isinstance(by, list):
+                by = [by] if by is not None else []
+            internal_by = [o for o in by if hashable(o) and o in self.columns]
+            internal_qc = (
+                [self.getitem_column_array(internal_by)] if len(internal_by) else []
+            )
+            by = internal_qc + by[len(internal_by) :]
+        return by, internal_by
+
     groupby_all = GroupByReduce.register("all")
     groupby_any = GroupByReduce.register("any")
     groupby_count = GroupByReduce.register("count")
@@ -2533,12 +2567,14 @@ def _callable_func(self, func, axis, *args, **kwargs):
     groupby_sum = GroupByReduce.register("sum")
 
     def groupby_mean(self, by, axis, groupby_kwargs, agg_args, agg_kwargs, drop=False):
+        _, internal_by = self._groupby_internal_columns(by, drop)
+
         numeric_only = agg_kwargs.get("numeric_only", False)
         datetime_cols = (
             {
                 col: dtype
                 for col, dtype in zip(self.dtypes.index, self.dtypes)
-                if is_datetime64_any_dtype(dtype)
+                if is_datetime64_any_dtype(dtype) and col not in internal_by
             }
             if not numeric_only
             else dict()
@@ -2776,26 +2812,7 @@ def groupby_agg(
         groupby_kwargs = groupby_kwargs.copy()
 
         as_index = groupby_kwargs.get("as_index", True)
-        if isinstance(by, type(self)):
-            # `drop` parameter indicates whether or not 'by' data came
-            # from the `self` frame:
-            # True: 'by' data came from the `self`
-            # False: external 'by' data
-            if drop:
-                internal_by = by.columns
-                by = [by]
-            else:
-                internal_by = []
-                by = [by]
-        else:
-            if not isinstance(by, list):
-                by = [by] if by is not None else []
-            internal_by = [o for o in by if hashable(o) and o in self.columns]
-            internal_qc = (
-                [self.getitem_column_array(internal_by)] if len(internal_by) else []
-            )
-
-            by = internal_qc + by[len(internal_by) :]
+        by, internal_by = self._groupby_internal_columns(by, drop)
 
         broadcastable_by = [o._modin_frame for o in by if isinstance(o, type(self))]
         not_broadcastable_by = [o for o in by if not isinstance(o, type(self))]


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
@@ -2165,3 +2165,22 @@ def test_groupby_with_frozenlist():
     pandas_df = pandas_df.set_index(["a", "b"])
     modin_df = from_pandas(pandas_df)
     eval_general(modin_df, pandas_df, lambda df: df.groupby(df.index.names).count())
+
+
+@pytest.mark.parametrize(
+    "by_func",
+    [
+        lambda df: "timestamp0",
+        lambda df: ["timestamp0", "timestamp1"],
+        lambda df: ["timestamp0", df["timestamp1"]],
+    ],
+)
+def test_mean_with_datetime(by_func):
+    data = {
+        "timestamp0": [pd.to_datetime(1490195805, unit="s")],
+        "timestamp1": [pd.to_datetime(1490195805, unit="s")],
+        "numeric": [0],
+    }
+
+    modin_df, pandas_df = create_test_dfs(data)
+    eval_general(modin_df, pandas_df, lambda df: df.groupby(by=by_func(df)).mean())