|
30 | 30 | import numpy as np |
31 | 31 |
|
32 | 32 | from pandas._libs import lib |
33 | | -from pandas._typing import FrameOrSeries, FrameOrSeriesUnion |
| 33 | +from pandas._typing import ArrayLike, FrameOrSeries, FrameOrSeriesUnion |
34 | 34 | from pandas.util._decorators import Appender, Substitution, doc |
35 | 35 |
|
36 | 36 | from pandas.core.dtypes.cast import ( |
|
59 | 59 | validate_func_kwargs, |
60 | 60 | ) |
61 | 61 | import pandas.core.algorithms as algorithms |
| 62 | +from pandas.core.arrays import ExtensionArray |
62 | 63 | from pandas.core.base import DataError, SpecificationError |
63 | 64 | import pandas.core.common as com |
64 | 65 | from pandas.core.construction import create_series_with_explicit_dtype |
@@ -1033,32 +1034,31 @@ def _cython_agg_blocks( |
1033 | 1034 |
|
1034 | 1035 | no_result = object() |
1035 | 1036 |
|
1036 | | - def cast_result_block(result, block: "Block", how: str) -> "Block": |
1037 | | - # see if we can cast the block to the desired dtype |
| 1037 | + def cast_agg_result(result, values: ArrayLike, how: str) -> ArrayLike: |
| 1038 | + # see if we can cast the values to the desired dtype |
1038 | 1039 | # this may not be the original dtype |
1039 | 1040 | assert not isinstance(result, DataFrame) |
1040 | 1041 | assert result is not no_result |
1041 | 1042 |
|
1042 | | - dtype = maybe_cast_result_dtype(block.dtype, how) |
| 1043 | + dtype = maybe_cast_result_dtype(values.dtype, how) |
1043 | 1044 | result = maybe_downcast_numeric(result, dtype) |
1044 | 1045 |
|
1045 | | - if block.is_extension and isinstance(result, np.ndarray): |
1046 | | - # e.g. block.values was an IntegerArray |
1047 | | - # (1, N) case can occur if block.values was Categorical |
| 1046 | + if isinstance(values, ExtensionArray) and isinstance(result, np.ndarray): |
| 1047 | + # e.g. values was an IntegerArray |
| 1048 | + # (1, N) case can occur if values was Categorical |
1048 | 1049 | # and result is ndarray[object] |
1049 | 1050 | # TODO(EA2D): special casing not needed with 2D EAs |
1050 | 1051 | assert result.ndim == 1 or result.shape[0] == 1 |
1051 | 1052 | try: |
1052 | 1053 | # Cast back if feasible |
1053 | | - result = type(block.values)._from_sequence( |
1054 | | - result.ravel(), dtype=block.values.dtype |
| 1054 | + result = type(values)._from_sequence( |
| 1055 | + result.ravel(), dtype=values.dtype |
1055 | 1056 | ) |
1056 | 1057 | except (ValueError, TypeError): |
1057 | 1058 | # reshape to be valid for non-Extension Block |
1058 | 1059 | result = result.reshape(1, -1) |
1059 | 1060 |
|
1060 | | - agg_block: "Block" = block.make_block(result) |
1061 | | - return agg_block |
| 1061 | + return result |
1062 | 1062 |
|
1063 | 1063 | def blk_func(block: "Block") -> List["Block"]: |
1064 | 1064 | new_blocks: List["Block"] = [] |
@@ -1092,28 +1092,25 @@ def blk_func(block: "Block") -> List["Block"]: |
1092 | 1092 | # Categoricals. This will done by later self._reindex_output() |
1093 | 1093 | # Doing it here creates an error. See GH#34951 |
1094 | 1094 | sgb = get_groupby(obj, self.grouper, observed=True) |
1095 | | - try: |
1096 | | - result = sgb.aggregate(lambda x: alt(x, axis=self.axis)) |
1097 | | - except TypeError: |
1098 | | - # we may have an exception in trying to aggregate |
1099 | | - # continue and exclude the block |
1100 | | - raise |
1101 | | - else: |
1102 | | - assert isinstance(result, (Series, DataFrame)) # for mypy |
1103 | | - # In the case of object dtype block, it may have been split |
1104 | | - # in the operation. We un-split here. |
1105 | | - result = result._consolidate() |
1106 | | - assert isinstance(result, (Series, DataFrame)) # for mypy |
1107 | | - assert len(result._mgr.blocks) == 1 |
1108 | | - |
1109 | | - # unwrap DataFrame to get array |
1110 | | - result = result._mgr.blocks[0].values |
1111 | | - if isinstance(result, np.ndarray) and result.ndim == 1: |
1112 | | - result = result.reshape(1, -1) |
1113 | | - agg_block = cast_result_block(result, block, how) |
1114 | | - new_blocks = [agg_block] |
| 1095 | + result = sgb.aggregate(lambda x: alt(x, axis=self.axis)) |
| 1096 | + |
| 1097 | + assert isinstance(result, (Series, DataFrame)) # for mypy |
| 1098 | + # In the case of object dtype block, it may have been split |
| 1099 | + # in the operation. We un-split here. |
| 1100 | + result = result._consolidate() |
| 1101 | + assert isinstance(result, (Series, DataFrame)) # for mypy |
| 1102 | + assert len(result._mgr.blocks) == 1 |
| 1103 | + |
| 1104 | + # unwrap DataFrame to get array |
| 1105 | + result = result._mgr.blocks[0].values |
| 1106 | + if isinstance(result, np.ndarray) and result.ndim == 1: |
| 1107 | + result = result.reshape(1, -1) |
| 1108 | + res_values = cast_agg_result(result, block.values, how) |
| 1109 | + agg_block = block.make_block(res_values) |
| 1110 | + new_blocks = [agg_block] |
1115 | 1111 | else: |
1116 | | - agg_block = cast_result_block(result, block, how) |
| 1112 | + res_values = cast_agg_result(result, block.values, how) |
| 1113 | + agg_block = block.make_block(res_values) |
1117 | 1114 | new_blocks = [agg_block] |
1118 | 1115 | return new_blocks |
1119 | 1116 |
|
|
0 commit comments