diff --git a/pandas/core/algorithms.py b/pandas/core/algorithms.py index cdbef673643e8..98243431e42df 100644 --- a/pandas/core/algorithms.py +++ b/pandas/core/algorithms.py @@ -1661,6 +1661,40 @@ def take(arr, indices, axis: int = 0, allow_fill: bool = False, fill_value=None) return result +def _take_preprocess_indexer_and_fill_value( + arr, indexer, axis, out, fill_value, allow_fill +): + mask_info = None + + if indexer is None: + indexer = np.arange(arr.shape[axis], dtype=np.int64) + dtype, fill_value = arr.dtype, arr.dtype.type() + else: + indexer = ensure_int64(indexer, copy=False) + if not allow_fill: + dtype, fill_value = arr.dtype, arr.dtype.type() + mask_info = None, False + else: + # check for promotion based on types only (do this first because + # it's faster than computing a mask) + dtype, fill_value = maybe_promote(arr.dtype, fill_value) + if dtype != arr.dtype and (out is None or out.dtype != dtype): + # check if promotion is actually required based on indexer + mask = indexer == -1 + needs_masking = mask.any() + mask_info = mask, needs_masking + if needs_masking: + if out is not None and out.dtype != dtype: + raise TypeError("Incompatible type for fill_value") + else: + # if not, then depromote, set fill_value to dummy + # (it won't be used but we don't want the cython code + # to crash when trying to cast it to dtype) + dtype, fill_value = arr.dtype, arr.dtype.type() + + return indexer, dtype, fill_value, mask_info + + def take_nd( arr, indexer, @@ -1700,8 +1734,6 @@ def take_nd( subarray : array-like May be the same type as the input, or cast to an ndarray. """ - mask_info = None - if fill_value is lib.no_default: fill_value = na_value_for_dtype(arr.dtype, compat=False) @@ -1712,31 +1744,9 @@ def take_nd( arr = extract_array(arr) arr = np.asarray(arr) - if indexer is None: - indexer = np.arange(arr.shape[axis], dtype=np.int64) - dtype, fill_value = arr.dtype, arr.dtype.type() - else: - indexer = ensure_int64(indexer, copy=False) - if not allow_fill: - dtype, fill_value = arr.dtype, arr.dtype.type() - mask_info = None, False - else: - # check for promotion based on types only (do this first because - # it's faster than computing a mask) - dtype, fill_value = maybe_promote(arr.dtype, fill_value) - if dtype != arr.dtype and (out is None or out.dtype != dtype): - # check if promotion is actually required based on indexer - mask = indexer == -1 - needs_masking = mask.any() - mask_info = mask, needs_masking - if needs_masking: - if out is not None and out.dtype != dtype: - raise TypeError("Incompatible type for fill_value") - else: - # if not, then depromote, set fill_value to dummy - # (it won't be used but we don't want the cython code - # to crash when trying to cast it to dtype) - dtype, fill_value = arr.dtype, arr.dtype.type() + indexer, dtype, fill_value, mask_info = _take_preprocess_indexer_and_fill_value( + arr, indexer, axis, out, fill_value, allow_fill + ) flip_order = False if arr.ndim == 2 and arr.flags.f_contiguous: