diff --git a/pandas/_libs/algos.pyi b/pandas/_libs/algos.pyi index bff7d117c97da..6dd1c7c9fb209 100644 --- a/pandas/_libs/algos.pyi +++ b/pandas/_libs/algos.pyi @@ -1,4 +1,5 @@ -# Note: this covers algos.pyx and algos_common_helper but NOT algos_take_helper +from __future__ import annotations + from typing import Any import numpy as np @@ -82,7 +83,7 @@ def pad( old: np.ndarray, # ndarray[algos_t] new: np.ndarray, # ndarray[algos_t] limit=..., -) -> np.ndarray: ... # np.ndarray[np.intp, ndim=1] +) -> npt.NDArray[np.intp]: ... # np.ndarray[np.intp, ndim=1] def pad_inplace( values: np.ndarray, # algos_t[:] mask: np.ndarray, # uint8_t[:] @@ -97,7 +98,7 @@ def backfill( old: np.ndarray, # ndarray[algos_t] new: np.ndarray, # ndarray[algos_t] limit=..., -) -> np.ndarray: ... # np.ndarray[np.intp, ndim=1] +) -> npt.NDArray[np.intp]: ... # np.ndarray[np.intp, ndim=1] def backfill_inplace( values: np.ndarray, # algos_t[:] mask: np.ndarray, # uint8_t[:] @@ -163,230 +164,287 @@ def ensure_uint16(arr: object, copy=...) -> npt.NDArray[np.uint16]: ... def ensure_uint32(arr: object, copy=...) -> npt.NDArray[np.uint32]: ... def ensure_uint64(arr: object, copy=...) -> npt.NDArray[np.uint64]: ... def take_1d_int8_int8( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_1d_int8_int32( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_1d_int8_int64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_1d_int8_float64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_1d_int16_int16( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_1d_int16_int32( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_1d_int16_int64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_1d_int16_float64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_1d_int32_int32( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_1d_int32_int64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_1d_int32_float64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_1d_int64_int64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_1d_int64_float64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_1d_float32_float32( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_1d_float32_float64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_1d_float64_float64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_1d_object_object( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_1d_bool_bool( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_1d_bool_object( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis0_int8_int8( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis0_int8_int32( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis0_int8_int64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis0_int8_float64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis0_int16_int16( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis0_int16_int32( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis0_int16_int64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis0_int16_float64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis0_int32_int32( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis0_int32_int64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis0_int32_float64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis0_int64_int64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis0_int64_float64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis0_float32_float32( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis0_float32_float64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis0_float64_float64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis0_object_object( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis0_bool_bool( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis0_bool_object( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis1_int8_int8( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis1_int8_int32( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis1_int8_int64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis1_int8_float64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis1_int16_int16( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis1_int16_int32( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis1_int16_int64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis1_int16_float64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis1_int32_int32( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis1_int32_int64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis1_int32_float64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis1_int64_int64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis1_int64_float64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis1_float32_float32( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis1_float32_float64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis1_float64_float64( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis1_object_object( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis1_bool_bool( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_axis1_bool_object( - values: np.ndarray, indexer: np.ndarray, out: np.ndarray, fill_value=... + values: np.ndarray, indexer: npt.NDArray[np.intp], out: np.ndarray, fill_value=... ) -> None: ... def take_2d_multi_int8_int8( - values: np.ndarray, indexer, out: np.ndarray, fill_value=... + values: np.ndarray, + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], + out: np.ndarray, + fill_value=..., ) -> None: ... def take_2d_multi_int8_int32( - values: np.ndarray, indexer, out: np.ndarray, fill_value=... + values: np.ndarray, + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], + out: np.ndarray, + fill_value=..., ) -> None: ... def take_2d_multi_int8_int64( - values: np.ndarray, indexer, out: np.ndarray, fill_value=... + values: np.ndarray, + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], + out: np.ndarray, + fill_value=..., ) -> None: ... def take_2d_multi_int8_float64( - values: np.ndarray, indexer, out: np.ndarray, fill_value=... + values: np.ndarray, + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], + out: np.ndarray, + fill_value=..., ) -> None: ... def take_2d_multi_int16_int16( - values: np.ndarray, indexer, out: np.ndarray, fill_value=... + values: np.ndarray, + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], + out: np.ndarray, + fill_value=..., ) -> None: ... def take_2d_multi_int16_int32( - values: np.ndarray, indexer, out: np.ndarray, fill_value=... + values: np.ndarray, + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], + out: np.ndarray, + fill_value=..., ) -> None: ... def take_2d_multi_int16_int64( - values: np.ndarray, indexer, out: np.ndarray, fill_value=... + values: np.ndarray, + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], + out: np.ndarray, + fill_value=..., ) -> None: ... def take_2d_multi_int16_float64( - values: np.ndarray, indexer, out: np.ndarray, fill_value=... + values: np.ndarray, + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], + out: np.ndarray, + fill_value=..., ) -> None: ... def take_2d_multi_int32_int32( - values: np.ndarray, indexer, out: np.ndarray, fill_value=... + values: np.ndarray, + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], + out: np.ndarray, + fill_value=..., ) -> None: ... def take_2d_multi_int32_int64( - values: np.ndarray, indexer, out: np.ndarray, fill_value=... + values: np.ndarray, + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], + out: np.ndarray, + fill_value=..., ) -> None: ... def take_2d_multi_int32_float64( - values: np.ndarray, indexer, out: np.ndarray, fill_value=... + values: np.ndarray, + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], + out: np.ndarray, + fill_value=..., ) -> None: ... def take_2d_multi_int64_float64( - values: np.ndarray, indexer, out: np.ndarray, fill_value=... + values: np.ndarray, + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], + out: np.ndarray, + fill_value=..., ) -> None: ... def take_2d_multi_float32_float32( - values: np.ndarray, indexer, out: np.ndarray, fill_value=... + values: np.ndarray, + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], + out: np.ndarray, + fill_value=..., ) -> None: ... def take_2d_multi_float32_float64( - values: np.ndarray, indexer, out: np.ndarray, fill_value=... + values: np.ndarray, + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], + out: np.ndarray, + fill_value=..., ) -> None: ... def take_2d_multi_float64_float64( - values: np.ndarray, indexer, out: np.ndarray, fill_value=... + values: np.ndarray, + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], + out: np.ndarray, + fill_value=..., ) -> None: ... def take_2d_multi_object_object( - values: np.ndarray, indexer, out: np.ndarray, fill_value=... + values: np.ndarray, + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], + out: np.ndarray, + fill_value=..., ) -> None: ... def take_2d_multi_bool_bool( - values: np.ndarray, indexer, out: np.ndarray, fill_value=... + values: np.ndarray, + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], + out: np.ndarray, + fill_value=..., ) -> None: ... def take_2d_multi_bool_object( - values: np.ndarray, indexer, out: np.ndarray, fill_value=... + values: np.ndarray, + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], + out: np.ndarray, + fill_value=..., ) -> None: ... def take_2d_multi_int64_int64( - values: np.ndarray, indexer, out: np.ndarray, fill_value=... + values: np.ndarray, + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], + out: np.ndarray, + fill_value=..., ) -> None: ... diff --git a/pandas/_libs/hashing.pyx b/pandas/_libs/hashing.pyx index 2dd2f1feadd70..75eee4d432637 100644 --- a/pandas/_libs/hashing.pyx +++ b/pandas/_libs/hashing.pyx @@ -120,13 +120,6 @@ cdef inline uint64_t _rotl(uint64_t x, uint64_t b) nogil: return (x << b) | (x >> (64 - b)) -cdef inline void u32to8_le(uint8_t* p, uint32_t v) nogil: - p[0] = (v) - p[1] = (v >> 8) - p[2] = (v >> 16) - p[3] = (v >> 24) - - cdef inline uint64_t u8to64_le(uint8_t* p) nogil: return (p[0] | p[1] << 8 | diff --git a/pandas/_libs/hashtable_class_helper.pxi.in b/pandas/_libs/hashtable_class_helper.pxi.in index 6d51ea7d5de7b..77d3f954a9a5d 100644 --- a/pandas/_libs/hashtable_class_helper.pxi.in +++ b/pandas/_libs/hashtable_class_helper.pxi.in @@ -1046,7 +1046,7 @@ cdef class StringHashTable(HashTable): not None, then _additionally_ any value "val" satisfying val == na_value is considered missing. mask : ndarray[bool], optional - Not yet implementd for StringHashTable. + Not yet implemented for StringHashTable. Returns ------- diff --git a/pandas/_libs/index.pyi b/pandas/_libs/index.pyi index 18ee216d5c4a1..446a980487cde 100644 --- a/pandas/_libs/index.pyi +++ b/pandas/_libs/index.pyi @@ -62,4 +62,4 @@ class BaseMultiIndexCodesEngine: values: np.ndarray, # np.ndarray[object] of tuples method: str, limit: int | None, - ) -> np.ndarray: ... # np.ndarray[np.int64] + ) -> npt.NDArray[np.intp]: ... diff --git a/pandas/core/array_algos/take.py b/pandas/core/array_algos/take.py index 6d350cfa2c1d6..7da1b571fcd27 100644 --- a/pandas/core/array_algos/take.py +++ b/pandas/core/array_algos/take.py @@ -13,7 +13,10 @@ algos as libalgos, lib, ) -from pandas._typing import ArrayLike +from pandas._typing import ( + ArrayLike, + npt, +) from pandas.core.dtypes.cast import maybe_promote from pandas.core.dtypes.common import ( @@ -110,7 +113,7 @@ def take_nd( def _take_nd_ndarray( arr: np.ndarray, - indexer, + indexer: npt.NDArray[np.intp] | None, axis: int, fill_value, allow_fill: bool, @@ -122,7 +125,7 @@ def _take_nd_ndarray( else: indexer = ensure_platform_int(indexer) - indexer, dtype, fill_value, mask_info = _take_preprocess_indexer_and_fill_value( + dtype, fill_value, mask_info = _take_preprocess_indexer_and_fill_value( arr, indexer, fill_value, allow_fill ) @@ -160,7 +163,7 @@ def _take_nd_ndarray( def take_1d( arr: ArrayLike, - indexer: np.ndarray, + indexer: npt.NDArray[np.intp], fill_value=None, allow_fill: bool = True, ) -> ArrayLike: @@ -168,7 +171,7 @@ def take_1d( Specialized version for 1D arrays. Differences compared to `take_nd`: - Assumes input array has already been converted to numpy array / EA - - Assumes indexer is already guaranteed to be int64 dtype ndarray + - Assumes indexer is already guaranteed to be intp dtype ndarray - Only works for 1D arrays To ensure the lowest possible overhead. @@ -176,6 +179,8 @@ def take_1d( Note: similarly to `take_nd`, this function assumes that the indexer is a valid(ated) indexer with no out of bound indices. """ + indexer = ensure_platform_int(indexer) + if not isinstance(arr, np.ndarray): # ExtensionArray -> dispatch to their method return arr.take(indexer, fill_value=fill_value, allow_fill=allow_fill) @@ -183,7 +188,7 @@ def take_1d( if not allow_fill: return arr.take(indexer) - indexer, dtype, fill_value, mask_info = _take_preprocess_indexer_and_fill_value( + dtype, fill_value, mask_info = _take_preprocess_indexer_and_fill_value( arr, indexer, fill_value, True ) @@ -200,7 +205,9 @@ def take_1d( def take_2d_multi( - arr: np.ndarray, indexer: tuple[np.ndarray, np.ndarray], fill_value=np.nan + arr: np.ndarray, + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], + fill_value=np.nan, ) -> np.ndarray: """ Specialized Cython take which sets NaN values in one pass. @@ -249,6 +256,7 @@ def take_2d_multi( if func is not None: func(arr, indexer, out=out, fill_value=fill_value) else: + # test_reindex_multi _take_2d_multi_object( arr, indexer, out, fill_value=fill_value, mask_info=mask_info ) @@ -457,7 +465,7 @@ def wrapper( def _take_nd_object( arr: np.ndarray, - indexer: np.ndarray, # np.ndarray[np.intp] + indexer: npt.NDArray[np.intp], out: np.ndarray, axis: int, fill_value, @@ -480,7 +488,7 @@ def _take_nd_object( def _take_2d_multi_object( arr: np.ndarray, - indexer: tuple[np.ndarray, np.ndarray], + indexer: tuple[npt.NDArray[np.intp], npt.NDArray[np.intp]], out: np.ndarray, fill_value, mask_info, @@ -509,7 +517,7 @@ def _take_2d_multi_object( def _take_preprocess_indexer_and_fill_value( arr: np.ndarray, - indexer: np.ndarray, + indexer: npt.NDArray[np.intp], fill_value, allow_fill: bool, ): @@ -533,5 +541,4 @@ def _take_preprocess_indexer_and_fill_value( # to crash when trying to cast it to dtype) dtype, fill_value = arr.dtype, arr.dtype.type() - indexer = ensure_platform_int(indexer) - return indexer, dtype, fill_value, mask_info + return dtype, fill_value, mask_info diff --git a/pandas/core/frame.py b/pandas/core/frame.py index 76a00071c8adc..0aea55613a78a 100644 --- a/pandas/core/frame.py +++ b/pandas/core/frame.py @@ -4672,17 +4672,23 @@ def _reindex_columns( allow_dups=False, ) - def _reindex_multi(self, axes, copy: bool, fill_value) -> DataFrame: + def _reindex_multi( + self, axes: dict[str, Index], copy: bool, fill_value + ) -> DataFrame: """ We are guaranteed non-Nones in the axes. """ + new_index, row_indexer = self.index.reindex(axes["index"]) new_columns, col_indexer = self.columns.reindex(axes["columns"]) if row_indexer is not None and col_indexer is not None: + # Fastpath. By doing two 'take's at once we avoid making an + # unnecessary copy. + # We only get here with `not self._is_mixed_type`, which (almost) + # ensures that self.values is cheap. It may be worth making this + # condition more specific. indexer = row_indexer, col_indexer - # error: Argument 2 to "take_2d_multi" has incompatible type "Tuple[Any, - # Any]"; expected "ndarray" new_values = take_2d_multi(self.values, indexer, fill_value=fill_value) return self._constructor(new_values, index=new_index, columns=new_columns) else: diff --git a/pandas/core/indexes/extension.py b/pandas/core/indexes/extension.py index 47ae88a52d919..034ca091c6fbf 100644 --- a/pandas/core/indexes/extension.py +++ b/pandas/core/indexes/extension.py @@ -41,7 +41,7 @@ _T = TypeVar("_T", bound="NDArrayBackedExtensionIndex") -def inherit_from_data(name: str, delegate, cache: bool = False, wrap: bool = False): +def _inherit_from_data(name: str, delegate, cache: bool = False, wrap: bool = False): """ Make an alias for a method of the underlying ExtensionArray. @@ -130,7 +130,7 @@ def inherit_names(names: list[str], delegate, cache: bool = False, wrap: bool = def wrapper(cls): for name in names: - meth = inherit_from_data(name, delegate, cache=cache, wrap=wrap) + meth = _inherit_from_data(name, delegate, cache=cache, wrap=wrap) setattr(cls, name, meth) return cls diff --git a/pandas/tests/apply/test_series_apply.py b/pandas/tests/apply/test_series_apply.py index 18f96f9c61ab8..48f984c21623b 100644 --- a/pandas/tests/apply/test_series_apply.py +++ b/pandas/tests/apply/test_series_apply.py @@ -774,7 +774,7 @@ def test_apply_series_on_date_time_index_aware_series(dti, exp, aware): tm.assert_frame_equal(result, exp) -def test_apply_scaler_on_date_time_index_aware_series(): +def test_apply_scalar_on_date_time_index_aware_series(): # GH 25959 # Calling apply on a localized time series should not cause an error series = tm.makeTimeSeries(nper=30).tz_localize("UTC")