Backport PR #47057 on branch 1.4.x (BUG: groupby.transform/agg with engine='numba' and a MultiIndex) (#47062)

BUG: groupby.transform/agg with engine='numba' and a MultiIndex (#47057)

Co-authored-by: Jeff Reback <[email protected]>
(cherry picked from commit c4027ad)

Author: mroeschke

doc/source/whatsnew/v1.4.3.rst

@@ -16,7 +16,8 @@ Fixed regressions
 ~~~~~~~~~~~~~~~~~
 - Fixed regression in :meth:`DataFrame.nsmallest` led to wrong results when ``np.nan`` in the sorting column (:issue:`46589`)
 - Fixed regression in :func:`read_fwf` raising ``ValueError`` when ``widths`` was specified with ``usecols`` (:issue:`46580`)
--
+- Fixed regression in :meth:`.Groupby.transform` and :meth:`.Groupby.agg` failing with ``engine="numba"`` when the index was a :class:`MultiIndex` (:issue:`46867`)
+- Fixed regression is :meth:`.Styler.to_latex` and :meth:`.Styler.to_html` where ``buf`` failed in combination with ``encoding`` (:issue:`47053`)
 
 .. ---------------------------------------------------------------------------
 

pandas/core/groupby/groupby.py

@@ -1257,7 +1257,16 @@ def _numba_prep(self, func, data):
         sorted_ids = algorithms.take_nd(ids, sorted_index, allow_fill=False)
 
         sorted_data = data.take(sorted_index, axis=self.axis).to_numpy()
-        sorted_index_data = data.index.take(sorted_index).to_numpy()
+        if len(self.grouper.groupings) > 1:
+            raise NotImplementedError(
+                "More than 1 grouping labels are not supported with engine='numba'"
+            )
+        # GH 46867
+        index_data = data.index
+        if isinstance(index_data, MultiIndex):
+            group_key = self.grouper.groupings[0].name
+            index_data = index_data.get_level_values(group_key)
+        sorted_index_data = index_data.take(sorted_index).to_numpy()
 
         starts, ends = lib.generate_slices(sorted_ids, ngroups)
         return (

pandas/tests/groupby/aggregate/test_numba.py

@@ -187,3 +187,30 @@ def f(values, index):
         [-1.5, -3.0], columns=["v"], index=Index(["A", "B"], name="group")
     )
     tm.assert_frame_equal(result, expected)
+
+
+@td.skip_if_no("numba")
+def test_multiindex_one_key(nogil, parallel, nopython):
+    def numba_func(values, index):
+        return 1
+
+    df = DataFrame([{"A": 1, "B": 2, "C": 3}]).set_index(["A", "B"])
+    engine_kwargs = {"nopython": nopython, "nogil": nogil, "parallel": parallel}
+    result = df.groupby("A").agg(
+        numba_func, engine="numba", engine_kwargs=engine_kwargs
+    )
+    expected = DataFrame([1.0], index=Index([1], name="A"), columns=["C"])
+    tm.assert_frame_equal(result, expected)
+
+
+@td.skip_if_no("numba")
+def test_multiindex_multi_key_not_supported(nogil, parallel, nopython):
+    def numba_func(values, index):
+        return 1
+
+    df = DataFrame([{"A": 1, "B": 2, "C": 3}]).set_index(["A", "B"])
+    engine_kwargs = {"nopython": nopython, "nogil": nogil, "parallel": parallel}
+    with pytest.raises(NotImplementedError, match="More than 1 grouping labels"):
+        df.groupby(["A", "B"]).agg(
+            numba_func, engine="numba", engine_kwargs=engine_kwargs
+        )

pandas/tests/groupby/transform/test_numba.py

@@ -176,3 +176,30 @@ def f(values, index):
     result = df.groupby("group").transform(f, engine="numba")
     expected = DataFrame([-4.0, -3.0, -2.0], columns=["v"], index=[-1, -2, -3])
     tm.assert_frame_equal(result, expected)
+
+
+@td.skip_if_no("numba")
+def test_multiindex_one_key(nogil, parallel, nopython):
+    def numba_func(values, index):
+        return 1
+
+    df = DataFrame([{"A": 1, "B": 2, "C": 3}]).set_index(["A", "B"])
+    engine_kwargs = {"nopython": nopython, "nogil": nogil, "parallel": parallel}
+    result = df.groupby("A").transform(
+        numba_func, engine="numba", engine_kwargs=engine_kwargs
+    )
+    expected = DataFrame([{"A": 1, "B": 2, "C": 1.0}]).set_index(["A", "B"])
+    tm.assert_frame_equal(result, expected)
+
+
+@td.skip_if_no("numba")
+def test_multiindex_multi_key_not_supported(nogil, parallel, nopython):
+    def numba_func(values, index):
+        return 1
+
+    df = DataFrame([{"A": 1, "B": 2, "C": 3}]).set_index(["A", "B"])
+    engine_kwargs = {"nopython": nopython, "nogil": nogil, "parallel": parallel}
+    with pytest.raises(NotImplementedError, match="More than 1 grouping labels"):
+        df.groupby(["A", "B"]).transform(
+            numba_func, engine="numba", engine_kwargs=engine_kwargs
+        )