diff --git a/pandas/core/groupby/generic.py b/pandas/core/groupby/generic.py index 80351a832ec7e..ab03cce0d6476 100644 --- a/pandas/core/groupby/generic.py +++ b/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) diff --git a/pandas/core/internals/array_manager.py b/pandas/core/internals/array_manager.py index ff76228646a02..71e6d14e6a716 100644 --- a/pandas/core/internals/array_manager.py +++ b/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)