diff --git a/pandas/core/array_algos/take.py b/pandas/core/array_algos/take.py index 2179085d9ad9b..054497089c5ab 100644 --- a/pandas/core/array_algos/take.py +++ b/pandas/core/array_algos/take.py @@ -121,6 +121,43 @@ def _take_nd_ndarray( return out +def take_1d( + arr: ArrayLike, + indexer: np.ndarray, + fill_value=None, + allow_fill: bool = True, +) -> ArrayLike: + """ + Specialized version for 1D arrays. Differences compared to take_nd: + + - Assumes input (arr, indexer) has already been converted to numpy array / EA + - Only works for 1D arrays + + To ensure the lowest possible overhead. + + TODO(ArrayManager): mainly useful for ArrayManager, otherwise can potentially + be removed again if we don't end up with ArrayManager. + """ + if not isinstance(arr, np.ndarray): + # ExtensionArray -> dispatch to their method + return arr.take(indexer, fill_value=fill_value, allow_fill=allow_fill) + + indexer, dtype, fill_value, mask_info = _take_preprocess_indexer_and_fill_value( + arr, indexer, 0, None, fill_value, allow_fill + ) + + # at this point, it's guaranteed that dtype can hold both the arr values + # and the fill_value + out = np.empty(indexer.shape, dtype=dtype) + + func = _get_take_nd_function( + arr.ndim, arr.dtype, out.dtype, axis=0, mask_info=mask_info + ) + func(arr, indexer, out, fill_value) + + return out + + def take_2d_multi( arr: np.ndarray, indexer: np.ndarray, fill_value=np.nan ) -> np.ndarray: diff --git a/pandas/core/internals/array_manager.py b/pandas/core/internals/array_manager.py index e53f39e0a3a53..905fa448ff033 100644 --- a/pandas/core/internals/array_manager.py +++ b/pandas/core/internals/array_manager.py @@ -58,7 +58,7 @@ import pandas.core.algorithms as algos from pandas.core.array_algos.quantile import quantile_compat -from pandas.core.array_algos.take import take_nd +from pandas.core.array_algos.take import take_1d from pandas.core.arrays import ( DatetimeArray, ExtensionArray, @@ -991,7 +991,7 @@ def _reindex_indexer( else: validate_indices(indexer, len(self._axes[0])) new_arrays = [ - take_nd( + take_1d( arr, indexer, allow_fill=True, @@ -1073,7 +1073,7 @@ def unstack(self, unstacker, fill_value) -> ArrayManager: new_arrays = [] for arr in self.arrays: for i in range(unstacker.full_shape[1]): - new_arr = take_nd( + new_arr = take_1d( arr, new_indexer2D[:, i], allow_fill=True, fill_value=fill_value ) new_arrays.append(new_arr)