improve topk/perm's worst case scenario (#79)

* improve topk/perm's worst case scenario

* using binary search for larger k

Author: sl-solution

src/stat/hp_stat.jl

@@ -87,6 +87,46 @@ Base.@propagate_inbounds function hp_topk_vals(x::AbstractVector{T}, k::Int, lt_
     end
     topk_vals(res_out, k, lt_fun, by)
 end
+Base.@propagate_inbounds function hp_topk_vals(x::Union{Vector{T}, SubArray{T, N, Vector{T}, Tuple{I}, L}}, k::Int, lt_fun::F, by) where {T<:Union{Missing, FLOATS, INTEGERS}} where {F} where N where I <: UnitRange{Int} where L
+    k < 1 && throw(ArgumentError("k must be greater than 1"))
+    all(ismissing, x) && return Union{Missing,T}[missing]
+    nt = Threads.nthreads()
+    res = Vector{T}(undef, k * nt)
+    res_out = allowmissing(res)
+    fill!(res_out, missing)
+    cz = div(length(x), nt)
+    Threads.@threads for i in 1:nt
+        lo = (i - 1) * cz + 1
+        i == nt ? hi = length(x) : hi = i * cz
+        th_res = view(res, (i-1)*k+1:i*k)
+        th_x = view(x, lo:hi)
+        th_res_out = view(res_out, (i-1)*k+1:i*k)
+        idx, cnt = initiate_topk_res!(th_res, th_x, by)
+        topk_sort!(th_res, 1, cnt, lt_fun)
+        if k < 21
+            for i in idx+1:length(th_x)
+                if !ismissing(by(th_x[i]))
+                    insert_fixed_sorted!(th_res, th_x[i], lt_fun)
+                    cnt += 1
+                end
+            end
+        else
+            for i in idx+1:length(th_x)
+                if !ismissing(by(th_x[i]))
+                    insert_fixed_sorted_binary!(th_res, th_x[i], lt_fun)
+                    cnt += 1
+                end
+            end
+        end
+        if cnt < k
+            view(th_res_out, 1:cnt) .= view(th_res, 1:cnt)
+        else
+            th_res_out .= th_res
+        end
+    end
+    topk_vals(res_out, k, lt_fun, by)
+end
+
 
 Base.@propagate_inbounds function hp_topk_perm(x::AbstractVector{T}, k::Int, lt_fun::F, by) where {T} where {F}
     k < 1 && throw(ArgumentError("k must be greater than 1"))
@@ -124,4 +164,51 @@ Base.@propagate_inbounds function hp_topk_perm(x::AbstractVector{T}, k::Int, lt_
         end
     end
     perm_out[topk_perm(res_out, k, lt_fun, by)]
+end
+
+Base.@propagate_inbounds function hp_topk_perm(x::Union{Vector{T}, SubArray{T, N, Vector{T}, Tuple{I}, L}}, k::Int, lt_fun::F, by) where {T<:Union{Missing, FLOATS, INTEGERS}} where {F} where N where I <: UnitRange{Int} where L
+    k < 1 && throw(ArgumentError("k must be greater than 1"))
+    all(ismissing, x) && return Union{Missing,Int}[missing]
+    nt = Threads.nthreads()
+    res = Vector{T}(undef, k * nt)
+    res_out = allowmissing(res)
+    fill!(res_out, missing)
+    perm = zeros(Int, k * nt)
+    perm_out = allowmissing(perm)
+    fill!(perm_out, missing)
+    cz = div(length(x), nt)
+    Threads.@threads for i in 1:nt
+        lo = (i - 1) * cz + 1
+        i == nt ? hi = length(x) : hi = i * cz
+        th_res = view(res, (i-1)*k+1:i*k)
+        th_perm = view(perm, (i-1)*k+1:i*k)
+        th_x = view(x, lo:hi)
+        th_res_out = view(res_out, (i-1)*k+1:i*k)
+        th_perm_out = view(perm_out, (i-1)*k+1:i*k)
+        idx, cnt = initiate_topk_res_perm!(th_perm, th_res, th_x, by, offset = lo - 1)
+        topk_sort_permute!(th_res, th_perm, 1, cnt, lt_fun)
+        if k < 16
+            for i in idx+1:length(th_x)
+                if !ismissing(by(th_x[i]))
+                    insert_fixed_sorted_perm!(th_perm, th_res, i + lo - 1, th_x[i], lt_fun)
+                    cnt += 1
+                end
+            end
+        else
+            for i in idx+1:length(th_x)
+                if !ismissing(by(th_x[i]))
+                    insert_fixed_sorted_perm_binary!(th_perm, th_res, i + lo - 1, th_x[i], lt_fun)
+                    cnt += 1
+                end
+            end
+        end
+        if cnt < k
+            view(th_res_out, 1:cnt) .= view(th_res, 1:cnt)
+            view(th_perm_out, 1:cnt) .= view(th_perm, 1:cnt)
+        else
+            th_res_out .= th_res
+            th_perm_out .= th_perm
+        end
+    end
+    perm_out[topk_perm(res_out, k, lt_fun, by)]
 end

src/stat/non_hp_stat.jl

@@ -426,7 +426,22 @@ function initiate_topk_res_perm!(perm, res, x, by; offset=0)
     end
     idx, cnt - 1
 end
-
+# it is unsafe because x must be continuous in memory
+function unsafe_shift_insert!(x::AbstractVector{T}, i, item) where T<:Union{Missing, FLOATS, INTEGERS}
+    n = length(x)
+    ccall(:memmove, Ptr{Cvoid}, (Ptr{Cvoid}, Ptr{Cvoid}, Csize_t), pointer(x, i + 1), pointer(x, i), (n - i) * Base.aligned_sizeof(T))
+    x[i] = item
+    x
+end
+Base.@propagate_inbounds function insert_fixed_sorted_binary!(x, item, lt_fun)
+    if !lt_fun(item, x[end])
+        return
+    end
+    idx = searchsortedlast(x, item, lt=lt_fun)
+    unsafe_shift_insert!(x, idx + 1, item)
+    nothing
+    
+end
 Base.@propagate_inbounds function insert_fixed_sorted!(x, item, lt_fun)
     if !lt_fun(item, x[end])
         return
@@ -445,6 +460,16 @@ Base.@propagate_inbounds function insert_fixed_sorted!(x, item, lt_fun)
     nothing
 end
 # TODO we do not need x, this is just easier to implement, later we may fix this
+Base.@propagate_inbounds function insert_fixed_sorted_perm_binary!(perm, x, idx, item, lt_fun)
+    if !lt_fun(item, x[end])
+        return
+    end
+    i = searchsortedlast(x, item, lt=lt_fun)
+    unsafe_shift_insert!(x, i + 1, item)
+    unsafe_shift_insert!(perm, i + 1, idx)
+    nothing
+end
+
 Base.@propagate_inbounds function insert_fixed_sorted_perm!(perm, x, idx, item, lt_fun)
     if !lt_fun(item, x[end])
         return
@@ -468,6 +493,7 @@ end
 Base.@propagate_inbounds function topk_vals(x::AbstractVector{T}, k::Int, lt_fun::F, by) where {T} where {F}
     k < 1 && throw(ArgumentError("k must be greater than 1"))
     all(ismissing, x) && return Union{Missing,T}[missing]
+    # TODO should we use similar() here?
     res = Vector{T}(undef, k)
     idx, cnt = initiate_topk_res!(res, x, by)
     topk_sort!(res, 1, cnt, lt_fun)
@@ -484,6 +510,35 @@ Base.@propagate_inbounds function topk_vals(x::AbstractVector{T}, k::Int, lt_fun
     end
 end
 
+#if k is greater than 20 (15 in topkperm) we switch to binary search - 21 and 16 are selected based on simulation study
+Base.@propagate_inbounds function topk_vals(x::Union{Vector{T}, SubArray{T, N, Vector{T}, Tuple{I}, L}}, k::Int, lt_fun::F, by) where {T<:Union{Missing, FLOATS, INTEGERS}} where {F} where N where I <: UnitRange{Int} where L
+    k < 1 && throw(ArgumentError("k must be greater than 1"))
+    all(ismissing, x) && return Union{Missing,T}[missing]
+    res = Vector{T}(undef, k)
+    idx, cnt = initiate_topk_res!(res, x, by)
+    topk_sort!(res, 1, cnt, lt_fun)
+    if k < 21
+        for i in idx+1:length(x)
+            if !ismissing(by(x[i]))
+                insert_fixed_sorted!(res, x[i], lt_fun)
+                cnt += 1
+            end
+        end
+    else
+        for i in idx+1:length(x)
+            if !ismissing(by(x[i]))
+                insert_fixed_sorted_binary!(res, x[i], lt_fun)
+                cnt += 1
+            end
+        end
+    end
+    if cnt < k
+        allowmissing(resize!(res, cnt))
+    else
+        allowmissing(res)
+    end
+end
+
 # ktop permutation
 #TODO should we return [missing] or Int[] when all elements are missings?
 Base.@propagate_inbounds function topk_perm(x::AbstractVector{T}, k::Int, lt_fun::F, by) where {T} where {F}
@@ -505,6 +560,34 @@ Base.@propagate_inbounds function topk_perm(x::AbstractVector{T}, k::Int, lt_fun
         allowmissing(perm)
     end
 end
+Base.@propagate_inbounds function topk_perm(x::Union{Vector{T}, SubArray{T, N, Vector{T}, Tuple{I}, L}}, k::Int, lt_fun::F, by) where {T<:Union{Missing, FLOATS, INTEGERS}} where {F} where N where I <: UnitRange{Int} where L
+    k < 1 && throw(ArgumentError("k must be greater than 1"))
+    all(ismissing, x) && return Union{Missing,Int}[missing]
+    res = Vector{T}(undef, k)
+    perm = zeros(Int, k)
+    idx, cnt = initiate_topk_res_perm!(perm, res, x, by)
+    topk_sort_permute!(res, perm, 1, cnt, lt_fun)
+    if k < 16
+        for i in idx+1:length(x)
+            if !ismissing(by(x[i]))
+                insert_fixed_sorted_perm!(perm, res, i, x[i], lt_fun)
+                cnt += 1
+            end
+        end
+    else
+        for i in idx+1:length(x)
+            if !ismissing(by(x[i]))
+                insert_fixed_sorted_perm_binary!(perm, res, i, x[i], lt_fun)
+                cnt += 1
+            end
+        end
+    end
+    if cnt < k
+        allowmissing(resize!(perm, cnt))
+    else
+        allowmissing(perm)
+    end
+end
 
 """
     topk(x, k; rev = false, lt = <, by = identity, threads = false)

test/stats.jl

@@ -1,4 +1,4 @@
-using Random
+using Random,PooledArrays,CategoricalArrays
 @testset "topk" begin
     # general usage
     for i in 1:100
@@ -24,7 +24,7 @@ using Random
             @test partialsortperm(x, 1:j, rev=true) == topkperm(x, j)
         end
         x = rand(Int8, 10000)
-        for j in 1:15
+        for j in 1:30
             @test partialsort(x, 1:j) == topk(x, j, rev=true) == topk(x, j, rev=true, threads=true)
             @test partialsort(x, 1:j, rev=true) == topk(x, j) == topk(x, j, threads=true)
             @test partialsortperm(x, 1:j) == topkperm(x, j, rev=true) == topkperm(x, j, rev=true, threads=true)
@@ -49,12 +49,27 @@ using Random
             @test partialsortperm(x, 1:min(11, j), rev=true) == topkperm(x, j)
         end
         x = [randstring() for _ in 1:101]
-        for j in 1:15
+        for j in 1:30
             @test partialsort(x, 1:j) == topk(x, j, rev=true) == topk(x, j, rev=true, threads=true)
             @test partialsort(x, 1:j, rev=true) == topk(x, j) == topk(x, j, threads=true)
             @test partialsortperm(x, 1:j) == topkperm(x, j, rev=true) == topkperm(x, j, rev=true, threads=true)
             @test partialsortperm(x, 1:j, rev=true) == topkperm(x, j) == topkperm(x, j, threads = true)
         end
+        x = PooledArray(rand(1:100, 100))
+        for j in 1:50
+            @test partialsort(x, 1:j) == topk(x, j, rev=true) == topk(x, j, rev=true, threads=true)
+            @test partialsort(x, 1:j, rev=true) == topk(x, j) == topk(x, j, threads=true)
+            @test partialsortperm(x, 1:j) == topkperm(x, j, rev=true) == topkperm(x, j, rev=true, threads=true)
+            @test partialsortperm(x, 1:j, rev=true) == topkperm(x, j) == topkperm(x, j, threads = true)
+        end
+        x = CategoricalArray(rand(100))
+        for j in 1:50
+            @test partialsort(x, 1:j) == topk(x, j, rev=true, lt = isless) 
+            @test partialsort(x, 1:j, rev=true) == topk(x, j, lt = isless) 
+            @test partialsortperm(x, 1:j) == topkperm(x, j, rev=true, lt = isless)
+            @test partialsortperm(x, 1:j, rev=true) == topkperm(x, j, lt = isless)
+        end
+
     end
     x = [1, 10, missing, 100, -1000, 32, 54, 0, missing, missing, -1]
     @test topk(x, 2) == [100, 54] == topk(x, 2, threads = true)
@@ -93,4 +108,20 @@ using Random
     @test topkperm(x, 3, by=ff678) == [8, 1]
     @test topk(x, 3, by=ff678, rev=true) == [-1,-100]
     @test topkperm(x, 3, by=ff678, rev=true) == [1,8]
+
+    x=[missing for _ in 1:1000]
+    @test isequal(topk(x, 10), topk(x,10,threads=true))
+    @test isequal(topk(x, 10), [missing])
+    @test isequal(topk(x, 100), topk(x,100,threads=true))
+    @test isequal(topk(x, 100), [missing])
+    @test isequal(topkperm(x, 100), topkperm(x,100,threads=true))
+    @test isequal(topkperm(x, 100), [missing])
+    @test isequal(topkperm(x, 10), topkperm(x,10,threads=true))
+    @test isequal(topkperm(x, 10), [missing])
+    @test isequal(topkperm(x, 10,rev=true), topkperm(x,10,threads=true,rev=true))
+    @test isequal(topkperm(x, 10,rev=true), [missing])
+
+    x=CategoricalArray(rand(1000))
+    # TODO categorical array is not thread safe - fortunately, it throws Errors - however, in future we may need to fix it
+    @test_throws UndefRefError topk(x,10,lt=isless,threads=true)
 end