@@ -464,7 +464,7 @@ def is_view(self) -> bool:
464
464
def is_single_block (self ) -> bool :
465
465
return len (self .arrays ) == 1
466
466
467
- def _get_data_subset (self : T , predicate : Callable ) -> T :
467
+ def _get_data_subset (self : T , predicate : Callable ) -> tuple [ T , npt . NDArray [ np . intp ]] :
468
468
indices = [i for i , arr in enumerate (self .arrays ) if predicate (arr )]
469
469
arrays = [self .arrays [i ] for i in indices ]
470
470
# TODO copy?
@@ -473,9 +473,9 @@ def _get_data_subset(self: T, predicate: Callable) -> T:
473
473
taker = np .array (indices , dtype = "intp" )
474
474
new_cols = self ._axes [1 ].take (taker )
475
475
new_axes = [self ._axes [0 ], new_cols ]
476
- return type (self )(arrays , new_axes , verify_integrity = False )
476
+ return type (self )(arrays , new_axes , verify_integrity = False ), taker
477
477
478
- def get_bool_data (self : T , copy : bool = False ) -> T :
478
+ def get_bool_data (self : T , copy : bool = False ) -> tuple [ T , npt . NDArray [ np . intp ]] :
479
479
"""
480
480
Select columns that are bool-dtype and object-dtype columns that are all-bool.
481
481
@@ -485,9 +485,8 @@ def get_bool_data(self: T, copy: bool = False) -> T:
485
485
Whether to copy the blocks
486
486
"""
487
487
return self ._get_data_subset (is_inferred_bool_dtype )
488
- # FIXME: return indexer
489
488
490
- def get_numeric_data (self : T , copy : bool = False ) -> T :
489
+ def get_numeric_data (self : T , copy : bool = False ) -> tuple [ T , npt . NDArray [ np . intp ]] :
491
490
"""
492
491
Select columns that have a numeric dtype.
493
492
@@ -935,7 +934,7 @@ def idelete(self, indexer) -> ArrayManager:
935
934
# --------------------------------------------------------------------
936
935
# Array-wise Operation
937
936
938
- def grouped_reduce (self : T , func : Callable , ignore_failures : bool = False ) -> T :
937
+ def grouped_reduce (self : T , func : Callable , ignore_failures : bool = False ) -> tuple [ T , npt . NDArray [ np . intp ]] :
939
938
"""
940
939
Apply grouped reduction function columnwise, returning a new ArrayManager.
941
940
@@ -948,6 +947,7 @@ def grouped_reduce(self: T, func: Callable, ignore_failures: bool = False) -> T:
948
947
Returns
949
948
-------
950
949
ArrayManager
950
+ np.ndarray[intp]
951
951
"""
952
952
result_arrays : list [np .ndarray ] = []
953
953
result_indices : list [int ] = []
@@ -975,14 +975,16 @@ def grouped_reduce(self: T, func: Callable, ignore_failures: bool = False) -> T:
975
975
else :
976
976
index = Index (range (result_arrays [0 ].shape [0 ]))
977
977
978
+ taker = None
978
979
if ignore_failures :
979
- columns = self .items [np .array (result_indices , dtype = "int64" )]
980
+ taker = np .array (result_indices , dtype = np .intp )
981
+ columns = self .items [taker ]
980
982
else :
981
983
columns = self .items
982
984
983
985
# error: Argument 1 to "ArrayManager" has incompatible type "List[ndarray]";
984
986
# expected "List[Union[ndarray, ExtensionArray]]"
985
- return type (self )(result_arrays , [index , columns ]) # type: ignore[arg-type]
987
+ return type (self )(result_arrays , [index , columns ]), taker # type: ignore[arg-type]
986
988
987
989
def reduce (
988
990
self : T , func : Callable , ignore_failures : bool = False
0 commit comments