PERF: optimize algos.take for repeated calls

pandas/core/algorithms.py

@@ -4,6 +4,7 @@
 """
 from __future__ import annotations
 
+import functools
 import operator
 from textwrap import dedent
 from typing import TYPE_CHECKING, Dict, Optional, Tuple, Union, cast
@@ -1534,40 +1535,79 @@ def _take_nd_object(arr, indexer, out, axis: int, fill_value, mask_info):
 }
 
 
+@functools.lru_cache(maxsize=128)
+def __get_take_nd_function_cached(ndim, arr_dtype, out_dtype, axis):
+    """
+    Part of _get_take_nd_function below that doesn't need `mask_info` and thus
+    can be cached (mask_info potentially contains a numpy ndarray which is not
+    hashable and thus cannot be used as argument for cached function).
+    """
+    tup = (arr_dtype.name, out_dtype.name)
+    if ndim == 1:
+        func = _take_1d_dict.get(tup, None)
+    elif ndim == 2:
+        if axis == 0:
+            func = _take_2d_axis0_dict.get(tup, None)
+        else:
+            func = _take_2d_axis1_dict.get(tup, None)
+    if func is not None:
+        return func
+
+    tup = (out_dtype.name, out_dtype.name)
+    if ndim == 1:
+        func = _take_1d_dict.get(tup, None)
+    elif ndim == 2:
+        if axis == 0:
+            func = _take_2d_axis0_dict.get(tup, None)
+        else:
+            func = _take_2d_axis1_dict.get(tup, None)
+    if func is not None:
+        func = _convert_wrapper(func, out_dtype)
+        return func
+
+    return None
+
+
 def _get_take_nd_function(
     ndim: int, arr_dtype, out_dtype, axis: int = 0, mask_info=None
 ):
+    """
+    Get the appropriate "take" implementation for the given dimension, axis
+    and dtypes.
+    """
+    func = None
     if ndim <= 2:
-        tup = (arr_dtype.name, out_dtype.name)
-        if ndim == 1:
-            func = _take_1d_dict.get(tup, None)
-        elif ndim == 2:
-            if axis == 0:
-                func = _take_2d_axis0_dict.get(tup, None)
-            else:
-                func = _take_2d_axis1_dict.get(tup, None)
-        if func is not None:
-            return func
-
-        tup = (out_dtype.name, out_dtype.name)
-        if ndim == 1:
-            func = _take_1d_dict.get(tup, None)
-        elif ndim == 2:
-            if axis == 0:
-                func = _take_2d_axis0_dict.get(tup, None)
-            else:
-                func = _take_2d_axis1_dict.get(tup, None)
-        if func is not None:
-            func = _convert_wrapper(func, out_dtype)
-            return func
+        # for this part we don't need `mask_info` -> use the cached algo lookup
+        func = __get_take_nd_function_cached(ndim, arr_dtype, out_dtype, axis)
+
+    if func is None:
+
+        def func(arr, indexer, out, fill_value=np.nan):
+            indexer = ensure_int64(indexer)
+            _take_nd_object(
+                arr, indexer, out, axis=axis, fill_value=fill_value, mask_info=mask_info
+            )
+
+    return func
 
-    def func2(arr, indexer, out, fill_value=np.nan):
-        indexer = ensure_int64(indexer)
-        _take_nd_object(
-            arr, indexer, out, axis=axis, fill_value=fill_value, mask_info=mask_info
-        )
 
-    return func2
+@functools.lru_cache(maxsize=128)
+def _maybe_promote_cached(dtype, fill_value, fill_value_type):
+    # also use fill_value_type as (unused) argument to use this in the cache
+    # lookup -> differentiate 1 and True
+    return maybe_promote(dtype, fill_value)
+
+
+def _maybe_promote(dtype, fill_value):
+    try:
+        # error: Argument 3 to "__call__" of "_lru_cache_wrapper" has incompatible type
+        # "Type[Any]"; expected "Hashable"  [arg-type]
+        return _maybe_promote_cached(
+            dtype, fill_value, type(fill_value)
+        )  # type: ignore[arg-type]
+    except TypeError:
+        # if fill_value is not hashable (required for caching)
+        return maybe_promote(dtype, fill_value)
 
 
 def take(arr, indices, axis: int = 0, allow_fill: bool = False, fill_value=None):
@@ -1677,7 +1717,7 @@ def _take_preprocess_indexer_and_fill_value(
         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)
+            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
@@ -1786,6 +1826,44 @@ def take_nd(
 take_1d = take_nd
 
 
+def take_1d_array(
+    arr: ArrayLike,
+    indexer: np.ndarray,
+    out=None,
+    fill_value=lib.no_default,
+    allow_fill: bool = True,
+):
+    """
+    Specialized version for 1D arrays. Differences compared to take_nd/take_1d:
+
+    - Assumes input (arr, indexer) has already been converted to numpy array / EA
+    - Only works for 1D arrays
+
+    To ensure the lowest possible overhead.
+    """
+    if fill_value is lib.no_default:
+        fill_value = na_value_for_dtype(arr.dtype, compat=False)
+
+    if isinstance(arr, ABCExtensionArray):
+        # Check for EA to catch DatetimeArray, TimedeltaArray
+        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, out, 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, indexer, fill_value=np.nan):
     """
     Specialized Cython take which sets NaN values in one pass.

pandas/core/internals/array_manager.py

@@ -862,7 +862,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 = algos.take(
+                new_arr = algos.take_1d_array(
                     arr, new_indexer2D[:, i], allow_fill=True, fill_value=fill_value
                 )
                 new_arrays.append(new_arr)

pandas/tests/test_take.py

@@ -417,6 +417,18 @@ def test_take_axis_1(self):
         with pytest.raises(IndexError, match="indices are out-of-bounds"):
             algos.take(arr, [0, 3], axis=1, allow_fill=True, fill_value=0)
 
+    def test_take_non_hashable_fill_value(self):
+        arr = np.array([1, 2, 3])
+        indexer = np.array([1, -1])
+        with pytest.raises(ValueError, match="fill_value must be a scalar"):
+            algos.take(arr, indexer, allow_fill=True, fill_value=[1])
+
+        # with object dtype it is allowed
+        arr = np.array([1, 2, 3], dtype=object)
+        result = algos.take(arr, indexer, allow_fill=True, fill_value=[1])
+        expected = np.array([2, [1]], dtype=object)
+        tm.assert_numpy_array_equal(result, expected)
+
 
 class TestExtensionTake:
     # The take method found in pd.api.extensions