REF: hide ArrayManager implementation details from GroupBy (pandas-dev#41086)

Author: jbrockmendel

pandas/core/groupby/generic.py

@@ -97,7 +97,6 @@
     all_indexes_same,
 )
 import pandas.core.indexes.base as ibase
-from pandas.core.internals import ArrayManager
 from pandas.core.series import Series
 from pandas.core.util.numba_ import maybe_use_numba
 
@@ -1103,17 +1102,11 @@ def _cython_agg_general(
         if numeric_only:
             data = data.get_numeric_data(copy=False)
 
-        using_array_manager = isinstance(data, ArrayManager)
-
         def cast_agg_result(result: ArrayLike, values: ArrayLike) -> ArrayLike:
             # see if we can cast the values to the desired dtype
             # this may not be the original dtype
 
-            if (
-                not using_array_manager
-                and isinstance(result.dtype, np.dtype)
-                and result.ndim == 1
-            ):
+            if isinstance(result.dtype, np.dtype) and result.ndim == 1:
                 # We went through a SeriesGroupByPath and need to reshape
                 # GH#32223 includes case with IntegerArray values
                 # We only get here with values.dtype == object
@@ -1794,8 +1787,6 @@ def count(self) -> DataFrame:
         ids, _, ngroups = self.grouper.group_info
         mask = ids != -1
 
-        using_array_manager = isinstance(data, ArrayManager)
-
         def hfunc(bvalues: ArrayLike) -> ArrayLike:
             # TODO(2DEA): reshape would not be necessary with 2D EAs
             if bvalues.ndim == 1:
@@ -1805,10 +1796,6 @@ def hfunc(bvalues: ArrayLike) -> ArrayLike:
                 masked = mask & ~isna(bvalues)
 
             counted = lib.count_level_2d(masked, labels=ids, max_bin=ngroups, axis=1)
-            if using_array_manager:
-                # count_level_2d return (1, N) array for single column
-                # -> extract 1D array
-                counted = counted[0, :]
             return counted
 
         new_mgr = data.grouped_reduce(hfunc)

pandas/core/internals/array_manager.py

@@ -932,12 +932,20 @@ def grouped_reduce(self: T, func: Callable, ignore_failures: bool = False) -> T:
         result_indices: list[int] = []
 
         for i, arr in enumerate(self.arrays):
+            # grouped_reduce functions all expect 2D arrays
+            arr = ensure_block_shape(arr, ndim=2)
             try:
                 res = func(arr)
             except (TypeError, NotImplementedError):
                 if not ignore_failures:
                     raise
                 continue
+
+            if res.ndim == 2:
+                # reverse of ensure_block_shape
+                assert res.shape[0] == 1
+                res = res[0]
+
             result_arrays.append(res)
             result_indices.append(i)