test groupby.indices for multiple groupby and mix of types

Creates tests for GH26859

Author: alexifm

pandas/tests/groupby/test_groupby.py

@@ -7,6 +7,8 @@
 
 from pandas.errors import PerformanceWarning
 
+from pandas.core.dtypes.common import is_categorical_dtype, is_datetime64_any_dtype
+
 import pandas as pd
 from pandas import (
     DataFrame,
@@ -361,6 +363,119 @@ def f3(x):
         df2.groupby("a").apply(f3)
 
 
+def test_groupby_indices_error():
+    # GH 26860
+    # Test if DataFrame Groupby builds gb.indices
+    dt = pd.to_datetime(["2018-01-01", "2018-02-01", "2018-03-01"])
+    df = DataFrame(
+        {
+            "a": Series(list("abc")),
+            "b": Series(dt, dtype="category"),
+            "c": pd.Categorical.from_codes([-1, 0, 1], categories=[0, 1]),
+        }
+    )
+
+    df.groupby(["a", "b"]).indices
+
+
+@pytest.mark.parametrize(
+    "gb_cols",
+    [
+        "int_series",
+        "int_series_cat",
+        "float_series",
+        "float_series_cat",
+        "dt_series",
+        "dt_series_cat",
+        "period_series",
+        "period_series_cat",
+        [
+            "int_series",
+            "int_series_cat",
+            "float_series",
+            "float_series_cat",
+            "dt_series",
+            "dt_series_cat",
+            "period_series",
+            "period_series_cat",
+        ],
+    ],
+)
+def test_groupby_indices_output(gb_cols):
+    # GH 26860
+    # Test if DataFrame Groupby builds gb.indices correctly.
+    if isinstance(gb_cols, str):
+        gb_cols = [gb_cols]
+
+    cols = [
+        "int_series",
+        "int_series_cat",
+        "float_series",
+        "float_series_cat",
+        "dt_series",
+        "dt_series_cat",
+        "dttz_series",
+        "dttz_series_cat",
+        "period_series",
+        "period_series_cat",
+    ]
+
+    int_series = Series([1, 2, 3])
+    dt_series = pd.to_datetime(["2018Q1", "2018Q2", "2018Q3"])
+    dttz_series = dt_series.tz_localize("Europe/Berlin")
+    df = DataFrame(
+        data={
+            "int_series": int_series,
+            "int_series_cat": int_series.astype("category"),
+            "float_series": int_series.astype("float"),
+            "float_series_cat": int_series.astype("float").astype("category"),
+            "dt_series": dt_series,
+            "dt_series_cat": dt_series.astype("category"),
+            "dttz_series": dttz_series,
+            "dttz_series_cat": dttz_series.astype("category"),
+            "period_series": dt_series.to_period("Q"),
+            "period_series_cat": dt_series.to_period("Q").astype("category"),
+        },
+        columns=cols,
+    )
+
+    def dt_to_ts(elems):
+        return [Timestamp(el) for el in elems]
+
+    def ts_to_dt(elems):
+        return [el.to_datetime64() for el in elems]
+
+    num_gb_cols = len(gb_cols)
+
+    if num_gb_cols == 1:
+        s = df[gb_cols[0]]
+        col_vals = list(s.unique())
+
+        if is_datetime64_any_dtype(s):
+            col_vals = dt_to_ts(col_vals)
+
+        target = {key: np.array([i]) for i, key in enumerate(col_vals)}
+    else:
+        col_vals = {col: list(df[col].unique()) for col in gb_cols}
+
+        for col in gb_cols:
+            is_dt = is_datetime64_any_dtype(df[col])
+            is_cat_dt = is_categorical_dtype(df[col]) and is_datetime64_any_dtype(
+                df[col].cat.categories
+            )
+            if is_dt or is_cat_dt:
+                col_vals[col] = dt_to_ts(col_vals[col])
+
+        it = zip(*(col_vals[col] for col in gb_cols))
+        target = {key: np.array([i]) for i, key in enumerate(it)}
+
+    indices = df.groupby(gb_cols).indices
+
+    assert set(target.keys()) == set(indices.keys())
+    for key in target.keys():
+        assert pd.core.dtypes.missing.array_equivalent(target[key], indices[key])
+
+
 def test_attr_wrapper(ts):
     grouped = ts.groupby(lambda x: x.weekday())