implement efficient sortperm (JuliaArrays#31)

Author: Pietro Vertechi

src/StructArrays.jl

@@ -12,10 +12,10 @@ include("sort.jl")
 function __init__()
     Requires.@require Tables="bd369af6-aec1-5ad0-b16a-f7cc5008161c" include("tables.jl")
     Requires.@require PooledArrays="2dfb63ee-cc39-5dd5-95bd-886bf059d720" begin
-        isdiscrete(::PooledArrays.PooledArray) = true
+        ispooledarray(::PooledArrays.PooledArray) = true
     end
     Requires.@require WeakRefStrings="ea10d353-3f73-51f8-a26c-33c1cb351aa5" begin
-        isdiscrete(::WeakRefStrings.StringArray) = true
+        isstringarray(::WeakRefStrings.StringArray) = true
     end
 end
 

src/sort.jl

@@ -1,12 +1,137 @@
-isdiscrete(v) = false
+using Base.Sort, Base.Order
+
+isstringarray(::AbstractArray) = false
+ispooledarray(::AbstractArray) = false
 
 function Base.permute!(c::StructVector, p::AbstractVector)
     foreachfield(c) do v
-        if isdiscrete(v) || v isa StructVector
+        if  v isa StructVector || isstringarray(v) || ispooledarray(v)
             permute!(v, p)
         else
             copyto!(v, v[p])
         end
     end
     return c
 end
+
+struct TiedIndices{T <: AbstractVector}
+    vec::T
+    perm::Vector{Int}
+    within::Tuple{Int, Int}
+end
+
+TiedIndices(vec::AbstractVector, perm=sortperm(vec)) =
+    TiedIndices(vec, perm, extrema(axes(vec, 1)))
+
+Base.IteratorSize(::Type{<:TiedIndices}) = Base.SizeUnknown()
+
+Base.eltype(::Type{<:TiedIndices{T}}) where {T} =
+    Pair{eltype(T), UnitRange{Int}}
+
+function Base.iterate(n::TiedIndices, i = n.within[1])
+    vec, perm = n.vec, n.perm
+    l = n.within[2]
+    i > l && return nothing
+    row = vec[perm[i]]
+    i1 = i
+    @inbounds while i1 <= l && isequal(row, vec[perm[i1]])
+        i1 += 1
+    end
+    return (row => i:(i1-1), i1)
+end
+
+tiedindices(args...) = TiedIndices(args...)
+
+function groupindices(args...)
+    t = tiedindices(args...)
+    p = t.perm
+    ((row => t.perm[idxs]) for (row, idxs) in t)
+end
+
+function Base.sortperm(c::StructVector{T};
+    alg = DEFAULT_UNSTABLE) where {T<:Union{Tuple, NamedTuple}}
+
+    cols = fieldarrays(c)
+    x = cols[1]
+    p = sortperm(x; alg = alg)
+    if length(cols) > 1
+        y = cols[2]
+        refine_perm!(p, cols, 1, x, y, 1, length(x))
+    end
+    return p
+end
+
+Base.sort!(c::StructArray{<:Union{Tuple, NamedTuple}}) = permute!(c, sortperm(c))
+Base.sort(c::StructArray{<:Union{Tuple, NamedTuple}}) = c[sortperm(c)]
+
+# Methods from IndexedTables to refine sorting:
+# # assuming x[p] is sorted, sort by remaining columns where x[p] is constant
+function refine_perm!(p, cols, c, x, y, lo, hi)
+    temp = similar(p, 0)
+    order = Base.Order.By(j->(@inbounds k=y[j]; k))
+    nc = length(cols)
+    for (_, idxs) in TiedIndices(x, p, (lo, hi))
+        i, i1 = extrema(idxs)
+        if i1 > i
+            sort_sub_by!(p, i, i1, y, order, temp)
+            if c < nc-1
+                z = cols[c+2]
+                refine_perm!(p, cols, c+1, y, z, i, i1)
+            end
+        end
+    end
+end
+
+# sort the values in v[i0:i1] in place, by array `by`
+Base.@noinline function sort_sub_by!(v, i0, i1, by, order, temp)
+    empty!(temp)
+    sort!(v, i0, i1, MergeSort, order, temp)
+end
+
+Base.@noinline function sort_sub_by!(v, i0, i1, by::Vector{T}, order, temp) where T<:Integer
+    min = max = by[v[i0]]
+    @inbounds for i = i0+1:i1
+        val = by[v[i]]
+        if val < min
+            min = val
+        elseif val > max
+            max = val
+        end
+    end
+    rangelen = max-min+1
+    n = i1-i0+1
+    if rangelen <= n
+        sort_int_range_sub_by!(v, i0-1, n, by, rangelen, min, temp)
+    else
+        empty!(temp)
+        sort!(v, i0, i1, MergeSort, order, temp)
+    end
+    v
+end
+
+# in-place counting sort of x[ioffs+1:ioffs+n] by values in `by`
+function sort_int_range_sub_by!(x, ioffs, n, by, rangelen, minval, temp)
+    offs = 1 - minval
+
+    where = fill(0, rangelen+1)
+    where[1] = 1
+    @inbounds for i = 1:n
+        where[by[x[i+ioffs]] + offs + 1] += 1
+    end
+    cumsum!(where, where)
+
+    length(temp) < n && resize!(temp, n)
+    @inbounds for i = 1:n
+        xi = x[i+ioffs]
+        label = by[xi] + offs
+        wl = where[label]
+        temp[wl] = xi
+        where[label] = wl+1
+    end
+
+    @inbounds for i = 1:n
+        x[i+ioffs] = temp[i]
+    end
+    x
+end
+

test/runtests.jl

@@ -22,9 +22,12 @@ end
     a = WeakRefStrings.StringVector(["a", "b", "c"])
     b = PooledArrays.PooledArray([1, 2, 3])
     c = [:a, :b, :c]
-    @test StructArrays.isdiscrete(a)
-    @test StructArrays.isdiscrete(b)
-    @test !StructArrays.isdiscrete(c)
+    @test !StructArrays.ispooledarray(a)
+    @test StructArrays.isstringarray(a)
+    @test StructArrays.ispooledarray(b)
+    @test !StructArrays.isstringarray(b)
+    @test !StructArrays.ispooledarray(c)
+    @test !StructArrays.isstringarray(c)
     s = StructArray(a=a, b=b, c=c)
     permute!(s, [2, 3, 1])
     @test s.a == ["b", "c", "a"]
@@ -36,6 +39,44 @@ end
     @test s == t
 end
 
+@testset "sortperm" begin
+    c = StructArray(a=[1,1,2,2], b=[1,2,3,3], c=["a","b","c","d"])
+    d = StructArray(a=[1,1,2,2], b=[1,2,3,3], c=["a","b","c","d"])
+    @test issorted(c)
+    @test sortperm(c) == [1,2,3,4]
+    permute!(c, [2,3,4,1])
+    @test c == StructArray(a=[1,2,2,1], b=[2,3,3,1], c=["b","c","d","a"])
+    @test sortperm(c) == [4,1,2,3]
+    @test !issorted(c)
+    @test sort(c) == d
+    sort!(c)
+    @test c == d
+
+    c = StructArray(a=[1,1,2,2], b=[1,2,3,3], c=PooledArrays.PooledArray(["a","b","c","d"]))
+    d = StructArray(a=[1,1,2,2], b=[1,2,3,3], c=PooledArrays.PooledArray(["a","b","c","d"]))
+    @test issorted(c)
+    @test sortperm(c) == [1,2,3,4]
+    permute!(c, [2,3,4,1])
+    @test c == StructArray(a=[1,2,2,1], b=[2,3,3,1], c=PooledArrays.PooledArray(["b","c","d","a"]))
+    @test sortperm(c) == [4,1,2,3]
+    @test !issorted(c)
+    @test sort(c) == d
+    sort!(c)
+    @test c == d
+end
+
+@testset "iterators" begin
+    c = [1, 2, 3, 1, 1]
+    d = StructArrays.tiedindices(c)
+    @test eltype(d) == Pair{Int, UnitRange{Int}}
+    s = collect(d)
+    @test first.(s) == [1, 2, 3]
+    @test last.(s) == [1:3, 4:4, 5:5]
+    t = collect(StructArrays.groupindices(c))
+    @test first.(t) == [1, 2, 3]
+    @test last.(t) == [[1, 4, 5], [2], [3]]
+end
+
 @testset "similar" begin
     t = StructArray(a = rand(10), b = rand(Bool, 10))
     s = similar(t)