Repository: numba/numba

Issue:
Poor np.median performance
For certain arrays the `np.median` function slows down dramatically. The slowdown occurs when all (or several) values in the array are identical, which hints at some sorting issue. 

See also some discussion on the mailing list, Antoine Pitrou has some great comments on the cause.
https://groups.google.com/a/continuum.io/forum/#!topic/numba-users/ENbEof1CmVU

Consider the following example which calculates the median repeatedly. A pure Numpy version and an equivalent Numba njit function:

```python
import numba
import numpy as np

def numpy_median(x):    
    for i in range(x.size):
        np.median(x)
        
@numba.njit
def numba_median(x):    
    for i in range(x.size):
        np.median(x)
```

Two arrays, one with all identical values, and another with random data:
```
n = 3000
x_equal = np.zeros(n) 
x_random = np.random.randn(n)
assert x_equal.dtype == x_random.dtype # both are f64
```

The timings on my machine (Win 10, Py3.6 and Numba 0.34):

`np_eq = %timeit -o numpy_median(x_equal)`
298 ms ± 37.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

`np_ra = %timeit -o numpy_median(x_random)`
288 ms ± 24.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

`nb_eq = %timeit -o numba_median(x_equal)`
32.4 s ± 787 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

`nb_ra = %timeit -o numba_median(x_random)`
93.3 ms ± 8.03 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

So for random data Numba is a few times faster, but for identical data, Numba is over 300x times slower. I get similar results with Numba 0.33, 0.32 and 0.30.

Here is a notebook containing the code above, and a custom median implementation using `np.sort`, which doesn't suffer from the same performance hit. 
https://gist.github.com/RutgerK/57297c37cd6a62349e789bc23c8f04ec

edit: 
Btw i encountered this while working on an implementation of "repeated medians regression". I don't mind adding this to the examples in this repo if its considered a nice addition. I'm not sure if it really adds something in the sense of explaining how Numba works.
https://gist.github.com/RutgerK/66b3edc665c0897ea4e777e54ec6a985





Discussion:

Resolving patch:
diff --git a/numba/targets/arraymath.py b/numba/targets/arraymath.py
--- a/numba/targets/arraymath.py
+++ b/numba/targets/arraymath.py
@@ -459,18 +459,25 @@ def _partition(A, low, high):
         A[low], A[mid] = A[mid], A[low]
     if A[high] < A[mid]:
         A[high], A[mid] = A[mid], A[high]
-        if A[mid] < A[low]:
-            A[low], A[mid] = A[mid], A[low]
+    if A[mid] < A[low]:
+        A[low], A[mid] = A[mid], A[low]
     pivot = A[mid]
 
     A[high], A[mid] = A[mid], A[high]
-
     i = low
-    for j in range(low, high):
-        if A[j] <= pivot:
-            A[i], A[j] = A[j], A[i]
+    j = high - 1
+    while True:
+        while i < high and A[i] < pivot:
             i += 1
-
+        while j >= low and pivot < A[j]:
+            j -= 1
+        if i >= j:
+            break
+        A[i], A[j] = A[j], A[i]
+        i += 1
+        j -= 1
+    # Put the pivot back in its final place (all items before `i`
+    # are smaller than the pivot, all items at/after `i` are larger)
     A[i], A[high] = A[high], A[i]
     return i
 


Test patch: