Make common impl. with Index.searchsorted

Author: topper-123

doc/source/whatsnew/v0.24.0.txt

@@ -496,6 +496,7 @@ Performance Improvements
  both when indexing by label (using .loc) and position(.iloc).
  Likewise, slicing a ``CategoricalIndex`` itself (i.e. ``ci[100:200]``) shows similar speed improvements (:issue:`21659`)
 - Improved performance of :func:`Series.searchsorted` (:issue:`22034`)
+- Improved performance of :func:`Index.searchsorted` when dtype is uint64, float64 or object (:issue:`22034`)
 - Improved performance of :func:`Series.describe` in case of numeric dtypes (:issue:`21274`)
 - Improved performance of :func:`pandas.core.groupby.GroupBy.rank` when dealing with tied rankings (:issue:`21237`)
 - Improved performance of :func:`DataFrame.set_index` with columns consisting of :class:`Period` objects (:issue:`21582`,:issue:`21606`)

pandas/core/base.py

@@ -15,7 +15,8 @@
  is_list_like,
  is_scalar,
  is_extension_type,
- is_extension_array_dtype)
+ is_extension_array_dtype,
+ ensure_platform_int)
 
 from pandas.util._validators import validate_bool_kwarg
 from pandas.errors import AbstractMethodError
@@ -1230,7 +1231,9 @@ def factorize(self, sort=False, na_sentinel=-1):
  @Appender(_shared_docs['searchsorted'])
  @deprecate_kwarg(old_arg_name='key', new_arg_name='value')
  def searchsorted(self, value, side='left', sorter=None):
- # needs coercion on the key (DatetimeIndex does already)
+ if sorter is not None:
+ sorter = ensure_platform_int(sorter)
+ value = com.maybe_convert_numeric_dtype(value, dtype=self.dtype)
  return self.values.searchsorted(value, side=side, sorter=sorter)
 
  def drop_duplicates(self, keep='first', inplace=False):

pandas/core/common.py

@@ -15,7 +15,9 @@
 from pandas import compat
 from pandas.compat import iteritems, PY36, OrderedDict
 from pandas.core.dtypes.generic import ABCSeries, ABCIndex, ABCIndexClass
-from pandas.core.dtypes.common import is_integer
+from pandas.core.dtypes.common import (is_integer, is_integer_dtype, is_float,
+ is_float_dtype, is_object_dtype,
+ is_scalar)
 from pandas.core.dtypes.inference import _iterable_not_string
 from pandas.core.dtypes.missing import isna, isnull, notnull # noqa
 from pandas.core.dtypes.cast import construct_1d_object_array_from_listlike
@@ -430,3 +432,47 @@ def _pipe(obj, func, *args, **kwargs):
  return func(*args, **kwargs)
  else:
  return func(obj, *args, **kwargs)
+
+
+def maybe_convert_numeric_dtype(value, dtype):
+ """
+ Convert value to have dtype 'dtype' if 'dtype' is int or float.
+
+ :func:`numpy.searchsorted` is only fast if value is of same dtype
+ as the searched array. Below we ensure that value has the right
+ dtype for giving fast results for arr.searchsorted.
+
+ Notes
+ -----
+ We do not recast the value if it is a float and the array is dtype int.
+ This is because loc of float 2.1 should be behind loc of int 2,
+ *not* before it.
+
+ Parameters
+ ----------
+ value : scalar or list-like
+ dtype : a dtype
+
+ Returns
+ -------
+ return_value : value as array with appropriate dtype for fast calling
+ of numpy.searchsorted
+ """
+ if is_float_dtype(dtype):
+ value = np.asarray(value, dtype=dtype)
+ elif is_integer_dtype(dtype):
+ # check bounds
+ iinfo = np.iinfo(dtype)
+ iinfo_val = np.array([value]) if is_scalar(value) else value
+ if (iinfo_val < iinfo.min).any() or (iinfo_val > iinfo.max).any():
+ msg = "Value {} out of bound for dtype {}".format(value, dtype)
+ raise ValueError(msg)
+
+ # convert 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
+ # but float 2.2 should *not* be converted to int 2.0
+ value = np.asarray(value, dtype=dtype)
+ return value

pandas/core/series.py

@@ -19,6 +19,7 @@
  is_categorical_dtype,
  is_bool,
  is_integer, is_integer_dtype,
+ is_numeric_dtype,
  is_float_dtype,
  is_extension_type,
  is_extension_array_dtype,
@@ -2082,14 +2083,19 @@ def __rmatmul__(self, other):
  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
-
- return self._values.searchsorted(value, side=side, sorter=sorter)
+
+ if is_numeric_dtype(self):
+ value = com.maybe_convert_numeric_dtype(value, dtype=self.dtype)
+ elif not (is_object_dtype(self) or is_categorical_dtype(self)):
+ value = Series(value)._values
+
+ result = self._values.searchsorted(value, side=side, sorter=sorter)
+
+ if is_scalar(result):
+ # ensure that a 1-dim array is returned
+ result = np.array([result])
+
+ return result
 
  # -------------------------------------------------------------------
  # Combination