Merge pull request numpy#3107 from MrBago/improve_searchsorted

charris · charris · commit aeb13affe23a · 2013-08-17T10:34:16.000-07:00
Improve searchsorted
diff --git a/numpy/core/src/multiarray/item_selection.c b/numpy/core/src/multiarray/item_selection.c
@@ -1866,10 +1866,9 @@ PyArray_LexSort(PyObject *sort_keys, int axis)
  *
  * For each key use bisection to find the first index i s.t. key <= arr[i].
  * When there is no such index i, set i = len(arr). Return the results in ret.
- * All arrays are assumed contiguous on entry and both arr and key must be of
- * the same comparable type.
+ * Both arr and key must be of the same comparable type.
  *
- * @param arr contiguous sorted array to be searched.
+ * @param arr 1d, strided, sorted array to be searched.
  * @param key contiguous array of keys.
  * @param ret contiguous array of intp for returned indices.
  * @return void
@@ -1883,15 +1882,16 @@ local_search_left(PyArrayObject *arr, PyArrayObject *key, PyArrayObject *ret)
     char *parr = PyArray_DATA(arr);
     char *pkey = PyArray_DATA(key);
     npy_intp *pret = (npy_intp *)PyArray_DATA(ret);
-    int elsize = PyArray_DESCR(arr)->elsize;
+    int elsize = PyArray_DESCR(key)->elsize;
+    npy_intp arrstride = *PyArray_STRIDES(arr);
     npy_intp i;
 
     for (i = 0; i < nkeys; ++i) {
         npy_intp imin = 0;
         npy_intp imax = nelts;
         while (imin < imax) {
             npy_intp imid = imin + ((imax - imin) >> 1);
-            if (compare(parr + elsize*imid, pkey, key) < 0) {
+            if (compare(parr + arrstride*imid, pkey, key) < 0) {
                 imin = imid + 1;
             }
             else {
@@ -1909,10 +1909,9 @@ local_search_left(PyArrayObject *arr, PyArrayObject *key, PyArrayObject *ret)
  *
  * For each key use bisection to find the first index i s.t. key < arr[i].
  * When there is no such index i, set i = len(arr). Return the results in ret.
- * All arrays are assumed contiguous on entry and both arr and key must be of
- * the same comparable type.
+ * Both arr and key must be of the same comparable type.
  *
- * @param arr contiguous sorted array to be searched.
+ * @param arr 1d, strided, sorted array to be searched.
  * @param key contiguous array of keys.
  * @param ret contiguous array of intp for returned indices.
  * @return void
@@ -1926,15 +1925,16 @@ local_search_right(PyArrayObject *arr, PyArrayObject *key, PyArrayObject *ret)
     char *parr = PyArray_DATA(arr);
     char *pkey = PyArray_DATA(key);
     npy_intp *pret = (npy_intp *)PyArray_DATA(ret);
-    int elsize = PyArray_DESCR(arr)->elsize;
+    int elsize = PyArray_DESCR(key)->elsize;
+    npy_intp arrstride = *PyArray_STRIDES(arr);
     npy_intp i;
 
     for(i = 0; i < nkeys; ++i) {
         npy_intp imin = 0;
         npy_intp imax = nelts;
         while (imin < imax) {
             npy_intp imid = imin + ((imax - imin) >> 1);
-            if (compare(parr + elsize*imid, pkey, key) <= 0) {
+            if (compare(parr + arrstride*imid, pkey, key) <= 0) {
                 imin = imid + 1;
             }
             else {
@@ -1951,11 +1951,11 @@ local_search_right(PyArrayObject *arr, PyArrayObject *key, PyArrayObject *ret)
  *
  * For each key use bisection to find the first index i s.t. key <= arr[i].
  * When there is no such index i, set i = len(arr). Return the results in ret.
- * All arrays are assumed contiguous on entry and both arr and key must be of
- * the same comparable type.
+ * Both arr and key must be of the same comparable type.
  *
- * @param arr contiguous sorted array to be searched.
+ * @param arr 1d, strided array to be searched.
  * @param key contiguous array of keys.
+ * @param sorter 1d, strided array of intp that sorts arr.
  * @param ret contiguous array of intp for returned indices.
  * @return int
  */
@@ -1968,22 +1968,24 @@ local_argsearch_left(PyArrayObject *arr, PyArrayObject *key,
     npy_intp nkeys = PyArray_SIZE(key);
     char *parr = PyArray_DATA(arr);
     char *pkey = PyArray_DATA(key);
-    npy_intp *psorter = (npy_intp *)PyArray_DATA(sorter);
+    char *psorter = PyArray_DATA(sorter);
     npy_intp *pret = (npy_intp *)PyArray_DATA(ret);
-    int elsize = PyArray_DESCR(arr)->elsize;
+    int elsize = PyArray_DESCR(key)->elsize;
+    npy_intp arrstride = *PyArray_STRIDES(arr);
+    npy_intp sorterstride = *PyArray_STRIDES(sorter);
     npy_intp i;
 
     for (i = 0; i < nkeys; ++i) {
         npy_intp imin = 0;
         npy_intp imax = nelts;
         while (imin < imax) {
             npy_intp imid = imin + ((imax - imin) >> 1);
-            npy_intp indx = psorter[imid];
+            npy_intp indx = *(npy_intp *)(psorter + sorterstride * imid);
 
             if (indx < 0 || indx >= nelts) {
                 return -1;
             }
-            if (compare(parr + elsize*indx, pkey, key) < 0) {
+            if (compare(parr + arrstride*indx, pkey, key) < 0) {
                 imin = imid + 1;
             }
             else {
@@ -2002,11 +2004,11 @@ local_argsearch_left(PyArrayObject *arr, PyArrayObject *key,
  *
  * For each key use bisection to find the first index i s.t. key < arr[i].
  * When there is no such index i, set i = len(arr). Return the results in ret.
- * All arrays are assumed contiguous on entry and both arr and key must be of
- * the same comparable type.
+ * Both arr and key must be of the same comparable type.
  *
- * @param arr contiguous sorted array to be searched.
+ * @param arr 1d, strided array to be searched.
  * @param key contiguous array of keys.
+ * @param sorter 1d, strided array of intp that sorts arr.
  * @param ret contiguous array of intp for returned indices.
  * @return int
  */
@@ -2019,22 +2021,24 @@ local_argsearch_right(PyArrayObject *arr, PyArrayObject *key,
     npy_intp nkeys = PyArray_SIZE(key);
     char *parr = PyArray_DATA(arr);
     char *pkey = PyArray_DATA(key);
-    npy_intp *psorter = (npy_intp *)PyArray_DATA(sorter);
+    char *psorter = PyArray_DATA(sorter);
     npy_intp *pret = (npy_intp *)PyArray_DATA(ret);
-    int elsize = PyArray_DESCR(arr)->elsize;
+    int elsize = PyArray_DESCR(key)->elsize;
+    npy_intp arrstride = *PyArray_STRIDES(arr);
+    npy_intp sorterstride = *PyArray_STRIDES(sorter);
     npy_intp i;
 
     for(i = 0; i < nkeys; ++i) {
         npy_intp imin = 0;
         npy_intp imax = nelts;
         while (imin < imax) {
             npy_intp imid = imin + ((imax - imin) >> 1);
-            npy_intp indx = psorter[imid];
+            npy_intp indx = *(npy_intp *)(psorter + sorterstride * imid);
 
             if (indx < 0 || indx >= nelts) {
                 return -1;
             }
-            if (compare(parr + elsize*indx, pkey, key) <= 0) {
+            if (compare(parr + arrstride*indx, pkey, key) <= 0) {
                 imin = imid + 1;
             }
             else {
@@ -2087,6 +2091,7 @@ PyArray_SearchSorted(PyArrayObject *op1, PyObject *op2,
     PyArrayObject *sorter = NULL;
     PyArrayObject *ret = NULL;
     PyArray_Descr *dtype;
+    int ap1_flags = NPY_ARRAY_NOTSWAPPED | NPY_ARRAY_ALIGNED;
     NPY_BEGIN_THREADS_DEF;
 
     /* Find common type */
@@ -2095,23 +2100,28 @@ PyArray_SearchSorted(PyArrayObject *op1, PyObject *op2,
         return NULL;
     }
 
-    /* need ap1 as contiguous array and of right type */
+    /* need ap2 as contiguous array and of right type */
     Py_INCREF(dtype);
-    ap1 = (PyArrayObject *)PyArray_CheckFromAny((PyObject *)op1, dtype,
-                                1, 1,
-                                NPY_ARRAY_DEFAULT | NPY_ARRAY_NOTSWAPPED,
+    ap2 = (PyArrayObject *)PyArray_CheckFromAny(op2, dtype,
+                                0, 0,
+                                NPY_ARRAY_CARRAY_RO | NPY_ARRAY_NOTSWAPPED,
                                 NULL);
-    if (ap1 == NULL) {
+    if (ap2 == NULL) {
         Py_DECREF(dtype);
         return NULL;
     }
 
-    /* need ap2 as contiguous array and of right type */
-    ap2 = (PyArrayObject *)PyArray_CheckFromAny(op2, dtype,
-                                0, 0,
-                                NPY_ARRAY_DEFAULT | NPY_ARRAY_NOTSWAPPED,
-                                NULL);
-    if (ap2 == NULL) {
+    /*
+     * If the needle (ap2) is larger than the haystack (op1) we copy the
+     * haystack to a continuous array for improved cache utilization.
+     */
+    if (PyArray_SIZE(ap2) > PyArray_SIZE(op1)) {
+        ap1_flags |= NPY_ARRAY_CARRAY_RO;
+    }
+
+    ap1 = (PyArrayObject *)PyArray_CheckFromAny((PyObject *)op1, dtype,
+                                1, 1, ap1_flags, NULL);
+    if (ap1 == NULL) {
         goto fail;
     }
     /* check that comparison function exists */
@@ -2125,7 +2135,7 @@ PyArray_SearchSorted(PyArrayObject *op1, PyObject *op2,
         /* need ap3 as contiguous array and of right type */
         ap3 = (PyArrayObject *)PyArray_CheckFromAny(perm, NULL,
                                     1, 1,
-                                    NPY_ARRAY_DEFAULT | NPY_ARRAY_NOTSWAPPED,
+                                    NPY_ARRAY_ALIGNED | NPY_ARRAY_NOTSWAPPED,
                                     NULL);
         if (ap3 == NULL) {
             PyErr_SetString(PyExc_TypeError,
@@ -2140,7 +2150,7 @@ PyArray_SearchSorted(PyArrayObject *op1, PyObject *op2,
         /* convert to known integer size */
         sorter = (PyArrayObject *)PyArray_FromArray(ap3,
                                     PyArray_DescrFromType(NPY_INTP),
-                                    NPY_ARRAY_DEFAULT | NPY_ARRAY_NOTSWAPPED);
+                                    NPY_ARRAY_ALIGNED | NPY_ARRAY_NOTSWAPPED);
         if (sorter == NULL) {
             PyErr_SetString(PyExc_ValueError,
                         "could not parse sorter argument");
diff --git a/numpy/core/tests/test_multiarray.py b/numpy/core/tests/test_multiarray.py
@@ -870,7 +870,6 @@ def test_searchsorted(self):
         # order.
 
         # check double
-        a = np.array([np.nan, 1, 0])
         a = np.array([0, 1, np.nan])
         msg = "Test real searchsorted with nans, side='l'"
         b = a.searchsorted(a, side='l')
@@ -897,6 +896,41 @@ def test_searchsorted(self):
         b = a.searchsorted(np.array(128,dtype='>i4'))
         assert_equal(b, 1, msg)
 
+        # Check 0 elements
+        a = np.ones(0)
+        b = a.searchsorted([0, 1, 2], 'l')
+        assert_equal(b, [0, 0, 0])
+        b = a.searchsorted([0, 1, 2], 'r')
+        assert_equal(b, [0, 0, 0])
+        a = np.ones(1)
+        # Check 1 element
+        b = a.searchsorted([0, 1, 2], 'l')
+        assert_equal(b, [0, 0, 1])
+        b = a.searchsorted([0, 1, 2], 'r')
+        assert_equal(b, [0, 1, 1])
+        # Check all elements equal
+        a = np.ones(2)
+        b = a.searchsorted([0, 1, 2], 'l')
+        assert_equal(b, [0, 0, 2])
+        b = a.searchsorted([0, 1, 2], 'r')
+        assert_equal(b, [0, 2, 2])
+
+        # Test searching unaligned array
+        a = np.arange(10)
+        aligned = np.empty(a.itemsize * a.size + 1, 'uint8')
+        unaligned = aligned[1:].view(a.dtype)
+        unaligned[:] = a
+        # Test searching unaligned array
+        b = unaligned.searchsorted(a, 'l')
+        assert_equal(b, a)
+        b = unaligned.searchsorted(a, 'r')
+        assert_equal(b, a + 1)
+        # Test searching for unaligned keys
+        b = a.searchsorted(unaligned, 'l')
+        assert_equal(b, a)
+        b = a.searchsorted(unaligned, 'r')
+        assert_equal(b, a + 1)
+
     def test_searchsorted_unicode(self):
         # Test searchsorted on unicode strings.
 
@@ -935,8 +969,9 @@ def test_searchsorted_with_sorter(self):
         # bounds check
         assert_raises(ValueError, np.searchsorted, a, 4, sorter=[0,1,2,3,5])
         assert_raises(ValueError, np.searchsorted, a, 0, sorter=[-1,0,1,2,3])
+        assert_raises(ValueError, np.searchsorted, a, 0, sorter=[4,0,-1,2,3])
 
-        a = np.random.rand(100)
+        a = np.random.rand(300)
         s = a.argsort()
         b = np.sort(a)
         k = np.linspace(0, 1, 20)
@@ -945,8 +980,30 @@ def test_searchsorted_with_sorter(self):
         a = np.array([0, 1, 2, 3, 5]*20)
         s = a.argsort()
         k = [0, 1, 2, 3, 5]
-        assert_equal(a.searchsorted(k, side='l', sorter=s), [0, 20, 40, 60, 80])
-        assert_equal(a.searchsorted(k, side='r', sorter=s), [20, 40, 60, 80, 100])
+        expected = [0, 20, 40, 60, 80]
+        assert_equal(a.searchsorted(k, side='l', sorter=s), expected)
+        expected = [20, 40, 60, 80, 100]
+        assert_equal(a.searchsorted(k, side='r', sorter=s), expected)
+
+        # Test searching unaligned array
+        keys = np.arange(10)
+        a = keys.copy()
+        np.random.shuffle(s)
+        s = a.argsort()
+        aligned = np.empty(a.itemsize * a.size + 1, 'uint8')
+        unaligned = aligned[1:].view(a.dtype)
+        # Test searching unaligned array
+        unaligned[:] = a
+        b = unaligned.searchsorted(keys, 'l', s)
+        assert_equal(b, keys)
+        b = unaligned.searchsorted(keys, 'r', s)
+        assert_equal(b, keys + 1)
+        # Test searching for unaligned keys
+        unaligned[:] = keys
+        b = a.searchsorted(unaligned, 'l', s)
+        assert_equal(b, keys)
+        b = a.searchsorted(unaligned, 'r', s)
+        assert_equal(b, keys + 1)
 
 
     def test_partition(self):
