From 49b0ab0227a56d9332d0a589eb66db1748087be1 Mon Sep 17 00:00:00 2001 From: Joris Van den Bossche Date: Wed, 24 Feb 2021 18:06:01 +0100 Subject: [PATCH] REF: share common code in ArrayManager.get_numeric_data/get_bool_data --- pandas/core/internals/array_manager.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/pandas/core/internals/array_manager.py b/pandas/core/internals/array_manager.py index 1d954d52147fb..cd8d3e547abbd 100644 --- a/pandas/core/internals/array_manager.py +++ b/pandas/core/internals/array_manager.py @@ -480,31 +480,34 @@ def is_view(self) -> bool: def is_single_block(self) -> bool: return False + def _get_data_subset(self, predicate: Callable) -> ArrayManager: + indices = [i for i, arr in enumerate(self.arrays) if predicate(arr)] + arrays = [self.arrays[i] for i in indices] + # TODO copy? + new_axes = [self._axes[0], self._axes[1][np.array(indices, dtype="int64")]] + return type(self)(arrays, new_axes, verify_integrity=False) + def get_bool_data(self, copy: bool = False) -> ArrayManager: """ + Select columns that are bool-dtype. + Parameters ---------- copy : bool, default False Whether to copy the blocks """ - mask = np.array([is_bool_dtype(t) for t in self.get_dtypes()], dtype="object") - arrays = [self.arrays[i] for i in np.nonzero(mask)[0]] - # TODO copy? - new_axes = [self._axes[0], self._axes[1][mask]] - return type(self)(arrays, new_axes) + return self._get_data_subset(lambda arr: is_bool_dtype(arr.dtype)) def get_numeric_data(self, copy: bool = False) -> ArrayManager: """ + Select columns that have a numeric dtype. + Parameters ---------- copy : bool, default False Whether to copy the blocks """ - mask = np.array([is_numeric_dtype(t) for t in self.get_dtypes()]) - arrays = [self.arrays[i] for i in np.nonzero(mask)[0]] - # TODO copy? - new_axes = [self._axes[0], self._axes[1][mask]] - return type(self)(arrays, new_axes) + return self._get_data_subset(lambda arr: is_numeric_dtype(arr.dtype)) def copy(self: T, deep=True) -> T: """