|
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