Improve quantile performance v2

src/Statistics.jl

@@ -937,8 +937,7 @@ function quantile!(q::AbstractArray, v::AbstractVector, p::AbstractArray;
     end
     isempty(q) && return q
 
-    minp, maxp = extrema(p)
-    _quantilesort!(v, sorted, minp, maxp)
+    _quantilesort!(v, sorted, p)
 
     for (i, j) in zip(eachindex(p), eachindex(q))
         @inbounds q[j] = _quantile(v,p[i], alpha=alpha, beta=beta)
@@ -949,35 +948,34 @@ end
 function quantile!(v::AbstractVector, p::Union{AbstractArray, Tuple{Vararg{Real}}};
                    sorted::Bool=false, alpha::Real=1., beta::Real=alpha)
     if !isempty(p)
-        minp, maxp = extrema(p)
-        _quantilesort!(v, sorted, minp, maxp)
+        _quantilesort!(v, sorted, p isa Tuple ? collect(p) : p)
     end
     return map(x->_quantile(v, x, alpha=alpha, beta=beta), p)
 end
 
 quantile!(v::AbstractVector, p::Real; sorted::Bool=false, alpha::Real=1., beta::Real=alpha) =
-    _quantile(_quantilesort!(v, sorted, p, p), p, alpha=alpha, beta=beta)
+    _quantile(_quantilesort!(v, sorted, [p]), p, alpha=alpha, beta=beta)
 
-# Function to perform partial sort of v for quantiles in given range
-function _quantilesort!(v::AbstractArray, sorted::Bool, minp::Real, maxp::Real)
+function _quantilesort!(v::AbstractArray, sorted::Bool, p::AbstractArray)
     isempty(v) && throw(ArgumentError("empty data vector"))
     require_one_based_indexing(v)
 
     if !sorted
-        lv = length(v)
-        lo = floor(Int,minp*(lv))
-        hi = ceil(Int,1+maxp*(lv))
-
-        # only need to perform partial sort
-        sort!(v, 1, lv, Base.Sort.PartialQuickSort(lo:hi), Base.Sort.Forward)
+        start = 1
+        for pv in sort(p)
+            lv = length(v)
+            lo = floor(Int,pv*(lv))
+            hi = ceil(Int,1+pv*(lv))
+            sort!(v, start, lv, Base.Sort.PartialQuickSort(lo:hi), Base.Sort.Forward)
+            start = hi + 1
+        end
     end
     if (sorted && (ismissing(v[end]) || (v[end] isa Number && isnan(v[end])))) ||
         any(x -> ismissing(x) || (x isa Number && isnan(x)), v)
         throw(ArgumentError("quantiles are undefined in presence of NaNs or missing values"))
     end
     return v
 end
-
 # Core quantile lookup function: assumes `v` sorted
 @inline function _quantile(v::AbstractVector, p::Real; alpha::Real=1.0, beta::Real=alpha)
     0 <= p <= 1 || throw(ArgumentError("input probability out of [0,1] range"))

test/runtests.jl

@@ -601,6 +601,14 @@ end
     @test quantile(skipmissing([1, missing, 2]), 0.5) === 1.5
     @test quantile([1], 0.5) === 1.0
 
+    # randomized partialsort correctness test
+    Random.seed!(1234)
+    for i in 1:200, j in 1:20
+        x = rand(2000)
+        p = rand(j)
+        @test quantile(x, p) == [quantile(x, v) for v in p]
+    end
+
     # make sure that type inference works correctly in normal cases
     for T in [Int, BigInt, Float64, Float16, BigFloat, Rational{Int}, Rational{BigInt}]
         for S in [Float64, Float16, BigFloat, Rational{Int}, Rational{BigInt}]