REF/PERF: deduplicate kth_smallest (#40559)

Author: mzeitlin11

pandas/_libs/algos.pxd

@@ -1,21 +1,4 @@
 from pandas._libs.util cimport numeric
 
 
-cdef inline Py_ssize_t swap(numeric *a, numeric *b) nogil:
-    cdef:
-        numeric t
-
-    # cython doesn't allow pointer dereference so use array syntax
-    t = a[0]
-    a[0] = b[0]
-    b[0] = t
-    return 0
-
-
-cdef enum TiebreakEnumType:
-    TIEBREAK_AVERAGE
-    TIEBREAK_MIN,
-    TIEBREAK_MAX
-    TIEBREAK_FIRST
-    TIEBREAK_FIRST_DESCENDING
-    TIEBREAK_DENSE
+cdef numeric kth_smallest_c(numeric* arr, Py_ssize_t k, Py_ssize_t n) nogil

pandas/_libs/algos.pyx

@@ -64,6 +64,14 @@ cdef:
     float64_t NaN = <float64_t>np.NaN
     int64_t NPY_NAT = get_nat()
 
+cdef enum TiebreakEnumType:
+    TIEBREAK_AVERAGE
+    TIEBREAK_MIN,
+    TIEBREAK_MAX
+    TIEBREAK_FIRST
+    TIEBREAK_FIRST_DESCENDING
+    TIEBREAK_DENSE
+
 tiebreakers = {
     "average": TIEBREAK_AVERAGE,
     "min": TIEBREAK_MIN,
@@ -237,34 +245,75 @@ def groupsort_indexer(const int64_t[:] index, Py_ssize_t ngroups):
     return indexer, counts
 
 
+cdef inline Py_ssize_t swap(numeric *a, numeric *b) nogil:
+    cdef:
+        numeric t
+
+    # cython doesn't allow pointer dereference so use array syntax
+    t = a[0]
+    a[0] = b[0]
+    b[0] = t
+    return 0
+
+
+cdef inline numeric kth_smallest_c(numeric* arr, Py_ssize_t k, Py_ssize_t n) nogil:
+    """
+    See kth_smallest.__doc__. The additional parameter n specifies the maximum
+    number of elements considered in arr, needed for compatibility with usage
+    in groupby.pyx
+    """
+    cdef:
+        Py_ssize_t i, j, l, m
+        numeric x
+
+    l = 0
+    m = n - 1
+
+    while l < m:
+        x = arr[k]
+        i = l
+        j = m
+
+        while 1:
+            while arr[i] < x: i += 1
+            while x < arr[j]: j -= 1
+            if i <= j:
+                swap(&arr[i], &arr[j])
+                i += 1; j -= 1
+
+            if i > j: break
+
+        if j < k: l = i
+        if k < i: m = j
+    return arr[k]
+
+
 @cython.boundscheck(False)
 @cython.wraparound(False)
-def kth_smallest(numeric[:] a, Py_ssize_t k) -> numeric:
+def kth_smallest(numeric[::1] arr, Py_ssize_t k) -> numeric:
+    """
+    Compute the kth smallest value in arr. Note that the input
+    array will be modified.
+
+    Parameters
+    ----------
+    arr : numeric[::1]
+        Array to compute the kth smallest value for, must be
+        contiguous
+    k : Py_ssize_t
+
+    Returns
+    -------
+    numeric
+        The kth smallest value in arr
+    """
     cdef:
-        Py_ssize_t i, j, l, m, n = a.shape[0]
-        numeric x
+        numeric result
 
     with nogil:
-        l = 0
-        m = n - 1
-
-        while l < m:
-            x = a[k]
-            i = l
-            j = m
-
-            while 1:
-                while a[i] < x: i += 1
-                while x < a[j]: j -= 1
-                if i <= j:
-                    swap(&a[i], &a[j])
-                    i += 1; j -= 1
-
-                if i > j: break
-
-            if j < k: l = i
-            if k < i: m = j
-    return a[k]
+        result = kth_smallest_c(&arr[0], k, arr.shape[0])
+
+    return result
 
 
 # ----------------------------------------------------------------------

pandas/_libs/groupby.pyx

@@ -30,7 +30,7 @@ from numpy.math cimport NAN
 
 cnp.import_array()
 
-from pandas._libs.algos cimport swap
+from pandas._libs.algos cimport kth_smallest_c
 from pandas._libs.util cimport (
     get_nat,
     numeric,
@@ -88,7 +88,7 @@ cdef inline float64_t median_linear(float64_t* a, int n) nogil:
         n -= na_count
 
     if n % 2:
-        result = kth_smallest_c( a, n // 2, n)
+        result = kth_smallest_c(a, n // 2, n)
     else:
         result = (kth_smallest_c(a, n // 2, n) +
                   kth_smallest_c(a, n // 2 - 1, n)) / 2
@@ -99,35 +99,6 @@ cdef inline float64_t median_linear(float64_t* a, int n) nogil:
     return result
 
 
-# TODO: Is this redundant with algos.kth_smallest
-cdef inline float64_t kth_smallest_c(float64_t* a,
-                                     Py_ssize_t k,
-                                     Py_ssize_t n) nogil:
-    cdef:
-        Py_ssize_t i, j, l, m
-        float64_t x, t
-
-    l = 0
-    m = n - 1
-    while l < m:
-        x = a[k]
-        i = l
-        j = m
-
-        while 1:
-            while a[i] < x: i += 1
-            while x < a[j]: j -= 1
-            if i <= j:
-                swap(&a[i], &a[j])
-                i += 1; j -= 1
-
-            if i > j: break
-
-        if j < k: l = i
-        if k < i: m = j
-    return a[k]
-
-
 @cython.boundscheck(False)
 @cython.wraparound(False)
 def group_median_float64(ndarray[float64_t, ndim=2] out,

pandas/core/algorithms.py

@@ -1307,7 +1307,9 @@ def compute(self, method: str) -> Series:
         narr = len(arr)
         n = min(n, narr)
 
-        kth_val = algos.kth_smallest(arr.copy(), n - 1)
+        # arr passed into kth_smallest must be contiguous. We copy
+        # here because kth_smallest will modify its input
+        kth_val = algos.kth_smallest(arr.copy(order="C"), n - 1)
         (ns,) = np.nonzero(arr <= kth_val)
         inds = ns[arr[ns].argsort(kind="mergesort")]