Make common impl. with Index.searchsorted

Author: topper-123

pandas/core/base.py

@@ -1497,8 +1497,12 @@ def factorize(self, sort=False, na_sentinel=-1):
     @Substitution(klass='IndexOpsMixin')
     @Appender(_shared_docs['searchsorted'])
     def searchsorted(self, value, side='left', sorter=None):
-        # needs coercion on the key (DatetimeIndex does already)
-        return self._values.searchsorted(value, side=side, sorter=sorter)
+        result = com.searchsorted(self._values, value,
+                                  side=side, sorter=sorter)
+
+        if is_scalar(value):
+            return result if is_scalar(result) else result[0]
+        return result
 
     def drop_duplicates(self, keep='first', inplace=False):
         inplace = validate_bool_kwarg(inplace, 'inplace')

pandas/core/common.py

@@ -12,13 +12,16 @@
 import numpy as np
 
 from pandas._libs import lib, tslibs
-import pandas.compat as compat
-from pandas.compat import PY36, OrderedDict, iteritems
-
 from pandas.core.dtypes.cast import construct_1d_object_array_from_listlike
-from pandas.core.dtypes.common import (
-    is_array_like, is_bool_dtype, is_extension_array_dtype, is_integer)
-from pandas.core.dtypes.generic import ABCIndex, ABCIndexClass, ABCSeries
+from pandas import compat
+from pandas.compat import iteritems, PY2, PY36, OrderedDict
+from pandas.core.dtypes.generic import ABCSeries, ABCIndex, ABCIndexClass
+from pandas.core.dtypes.common import (is_integer, is_integer_dtype,
+                                       is_bool_dtype, is_extension_array_dtype,
+                                       is_array_like,
+                                       is_float_dtype, is_object_dtype,
+                                       is_categorical_dtype, is_numeric_dtype,
+                                       is_scalar, ensure_platform_int)
 from pandas.core.dtypes.inference import _iterable_not_string
 from pandas.core.dtypes.missing import isna, isnull, notnull  # noqa
 
@@ -481,3 +484,83 @@ def f(x):
         f = mapper
 
     return f
+
+
+def ensure_integer_dtype(arr, value):
+    """
+    Ensure optimal dtype for :func:`searchsorted_integer` is returned.
+
+    Parameters
+    ----------
+    arr : a numpy integer array
+    value : a number or array of numbers
+
+    Returns
+    -------
+    dtype : an numpy integer dtype
+
+    Raises
+    ------
+    TypeError : if value is not a number
+    """
+    value_arr = np.array([value]) if is_scalar(value) else np.array(value)
+
+    if PY2 and not is_numeric_dtype(value_arr):
+        # python 2 allows "a" < 1, avoid such nonsense
+        msg = "value must be numeric, was type {}"
+        raise TypeError(msg.format(value))
+
+    iinfo = np.iinfo(arr.dtype)
+    if not ((value_arr < iinfo.min).any() or (value_arr > iinfo.max).any()):
+        return arr.dtype
+    else:
+        return value_arr.dtype
+
+
+def searchsorted_integer(arr, value, side="left", sorter=None):
+    """
+    searchsorted implementation, but only for integer arrays.
+
+    We get a speedup if the dtype of arr and value is the same.
+
+    See :func:`searchsorted` for a more general searchsorted implementation.
+    """
+    if sorter is not None:
+        sorter = ensure_platform_int(sorter)
+
+    dtype = ensure_integer_dtype(arr, value)
+
+    if is_integer(value) or is_integer_dtype(value):
+        value = np.asarray(value, dtype=dtype)
+    elif hasattr(value, 'is_integer') and value.is_integer():
+        # float 2.0 can be converted to int 2 for better speed,
+        # but float 2.2 should *not* be converted to int 2
+        value = np.asarray(value, dtype=dtype)
+
+    return np.searchsorted(arr, value, side=side, sorter=sorter)
+
+
+def searchsorted(arr, value, side="left", sorter=None):
+    """
+    Find indices where elements should be inserted to maintain order.
+
+    Find the indices into a sorted array-like `arr` such that, if the
+    corresponding elements in `value` were inserted before the indices,
+    the order of `arr` would be preserved.
+
+    See :class:`IndexOpsMixin.searchsorted` for more details and examples.
+    """
+    if sorter is not None:
+        sorter = ensure_platform_int(sorter)
+
+    if is_integer_dtype(arr):
+        return searchsorted_integer(arr, value, side=side, sorter=sorter)
+    elif (is_object_dtype(arr) or is_float_dtype(arr) or
+          is_categorical_dtype(arr)):
+        return arr.searchsorted(value, side=side, sorter=sorter)
+    else:
+        # fallback solution. E.g. arr is an array with dtype='datetime64[ns]'
+        # and value is a pd.Timestamp, need to convert value
+        from pandas.core.series import Series
+        value = Series(value)._values
+        return arr.searchsorted(value, side=side, sorter=sorter)

pandas/core/series.py

@@ -2331,17 +2331,12 @@ def __rmatmul__(self, other):
     @Substitution(klass='Series')
     @Appender(base._shared_docs['searchsorted'])
     def searchsorted(self, value, side='left', sorter=None):
-        if sorter is not None:
-            sorter = ensure_platform_int(sorter)
-        if not is_extension_type(self._values):
-            # numpy searchsorted is only fast if value is of same dtype as the
-            # searched array. Below we ensure that value has the right dtype,
-            # and is not 0-dimensional.
-            value = np.asarray(value, dtype=self._values.dtype)
-            value = value[..., np.newaxis] if value.ndim == 0 else value
-
-        result = self._values.searchsorted(value, side=side, sorter=sorter)
-        return result[0] if is_scalar(value) else result
+        result = com.searchsorted(self._values, value,
+                                  side=side, sorter=sorter)
+
+        if is_scalar(value):
+            return result if is_scalar(result) else result[0]
+        return result
 
     # -------------------------------------------------------------------
     # Combination