Speed up permsort by utilizing stability of the default sorting algorithm

base/ordering.jl

@@ -13,7 +13,7 @@ import .Base:
 
 export # not exported by Base
     Ordering, Forward, Reverse,
-    By, Lt, Perm,
+    By, Lt, Perm, PermFast,
     ReverseOrdering, ForwardOrdering,
     DirectOrdering,
     lt, ord, ordtype
@@ -106,8 +106,21 @@ struct Perm{O<:Ordering,V<:AbstractVector} <: Ordering
     data::V
 end
 
+"""
+    PermFast(order::Ordering, data::AbstractVector)
+`Ordering` on the indices of `data` where `i` is less than `j` if `data[i]` is
+less than `data[j]` according to `order`. In the case that `data[i]` and
+`data[j]` are equal, the ordering is undefined. Thus, it is designed to be
+faster than `Perm` when a stable sorting algorithm is used.
+"""
+struct PermFast{O<:Ordering,V<:AbstractVector} <: Ordering
+    order::O
+    data::V
+end
+
 ReverseOrdering(by::By) = By(by.by, ReverseOrdering(by.order))
 ReverseOrdering(perm::Perm) = Perm(ReverseOrdering(perm.order), perm.data)
+ReverseOrdering(perm::PermFast) = PermFast(ReverseOrdering(perm.order), perm.data)
 
 """
     lt(o::Ordering, a, b)
@@ -125,6 +138,12 @@ lt(o::Lt,                    a, b) = o.lt(a,b)
     (lt(p.order, da, db)::Bool) | (!(lt(p.order, db, da)::Bool) & (a < b))
 end
 
+@propagate_inbounds function lt(p::PermFast, a::Integer, b::Integer)
+    da = p.data[a]
+    db = p.data[b]
+    lt(p.order, da, db)::Bool
+end
+
 _ord(lt::typeof(isless), by::typeof(identity), order::Ordering) = order
 _ord(lt::typeof(isless), by,                   order::Ordering) = By(by, order)
 

base/sort.jl

@@ -560,15 +560,52 @@ elements that are not
 @inline send_to_end!(f::F, v::AbstractVector, ::ReverseOrdering, end_stable=false; lo, hi) where F <: Function =
     end_stable ? (send_to_end!(!f, v; lo, hi)+1, hi) : (hi-send_to_end!(f, view(v, hi:-1:lo))+1, hi)
 
+"""
+    send_to_end_stable!(f::Function, v::AbstractVector; [lo, hi])
+
+Send every element of `v` for which `f` returns `true` to the end of the vector and return
+the index of the last element which for which `f` returns `false`.
+
+`send_to_end_stable!(f, v, lo, hi)` is equivalent to `send_to_end_stable!(f, view(v, lo:hi))+lo-1`
+
+Preserves the order of the elements.
+"""
+function send_to_end_stable!(f::F, v::AbstractVector; lo=firstindex(v), hi=lastindex(v)) where F <: Function
+    tmp=copy(v)
+    offset = 0
+    @inbounds begin
+        while lo <= hi
+            x = tmp[lo]
+            fx = f(x)::Bool
+            v[(fx ? hi : lo) - offset] = x
+            offset += fx
+            lo += 1
+        end
+    end
+
+    # This is similar to the partition function
+    pivot_index = lo-offset-1
+    # t[<=pivot_index] <* f(x) = false
+    # t[>pivot_index]  >* f(x) = true
+
+    # Make the results stable
+    reverse!(v, pivot_index+1, hi)
+    return pivot_index
+end
+
+@inline send_to_end_stable!(f::F, v::AbstractVector, ::ForwardOrdering; lo, hi) where F <: Function = (lo, send_to_end_stable!(f, v; lo, hi))
+@inline send_to_end_stable!(f::F, v::AbstractVector, ::ReverseOrdering; lo, hi) where F <: Function = (hi-send_to_end_stable!(f, view(v, hi:-1:lo))+1, hi)
 
 function _sort!(v::AbstractVector, a::MissingOptimization, o::Ordering, kw)
     @getkw lo hi
     if nonmissingtype(eltype(v)) != eltype(v) && o isa DirectOrdering
         lo, hi = send_to_end!(ismissing, v, o; lo, hi)
         _sort!(WithoutMissingVector(v, unsafe=true), a.next, o, (;kw..., lo, hi))
-    elseif eltype(v) <: Integer && o isa Perm{DirectOrdering} && nonmissingtype(eltype(o.data)) != eltype(o.data)
+    elseif eltype(v) <: Integer && (o isa Perm{DirectOrdering} || o isa PermFast{DirectOrdering}) &&
+            nonmissingtype(eltype(o.data)) != eltype(o.data)
+        PermT = o isa Perm{DirectOrdering} ? Perm : PermFast
         lo, hi = send_to_end!(i -> ismissing(@inbounds o.data[i]), v, o)
-        _sort!(v, a.next, Perm(o.order, WithoutMissingVector(o.data, unsafe=true)), (;kw..., lo, hi))
+        _sort!(v, a.next, PermT(o.order, WithoutMissingVector(o.data, unsafe=true)), (;kw..., lo, hi))
     else
         _sort!(v, a.next, o, kw)
     end
@@ -603,7 +640,7 @@ function _sort!(v::AbstractVector, a::IEEEFloatOptimization, o::Ordering, kw)
         j = send_to_end!(x -> after_zero(o, x), v; lo, hi)
         scratch = _sort!(iv, a.next, Reverse, (;kw..., lo, hi=j))
         if scratch === nothing # Union split
-            _sort!(iv, a.next, Forward, (;kw..., lo=j+1, hi, scratch))
+            _sort!(iv, a.next, Forward, (;kw..., lo=j+1, hi, nothing))
         else
             _sort!(iv, a.next, Forward, (;kw..., lo=j+1, hi, scratch))
         end
@@ -613,10 +650,20 @@ function _sort!(v::AbstractVector, a::IEEEFloatOptimization, o::Ordering, kw)
         j = send_to_end!(i -> after_zero(o.order, @inbounds o.data[i]), v; lo, hi)
         scratch = _sort!(v, a.next, Perm(Reverse, ip), (;kw..., lo, hi=j))
         if scratch === nothing # Union split
-            _sort!(v, a.next, Perm(Forward, ip), (;kw..., lo=j+1, hi, scratch))
+            _sort!(v, a.next, Perm(Forward, ip), (;kw..., lo=j+1, hi, nothing))
         else
             _sort!(v, a.next, Perm(Forward, ip), (;kw..., lo=j+1, hi, scratch))
         end
+    elseif eltype(v) <: Integer && o isa PermFast && o.order isa DirectOrdering && is_concrete_IEEEFloat(eltype(o.data))
+        lo, hi = send_to_end_stable!(i -> isnan(@inbounds o.data[i]), v, o.order; lo, hi)
+        ip = reinterpret(UIntType(eltype(o.data)), o.data)
+        j = send_to_end_stable!(i -> after_zero(o.order, @inbounds o.data[i]), v; lo, hi)
+        scratch = _sort!(v, a.next, PermFast(Reverse, ip), (;kw..., lo, hi=j))
+        if scratch === nothing # Union split
+            _sort!(v, a.next, PermFast(Forward, ip), (;kw..., lo=j+1, hi, nothing))
+        else
+            _sort!(v, a.next, PermFast(Forward, ip), (;kw..., lo=j+1, hi, scratch))
+        end
     else
         _sort!(v, a.next, o, kw)
     end
@@ -1558,7 +1605,11 @@ function sortperm(A::AbstractArray;
         end
     end
     ix = copymutable(LinearIndices(A))
-    sort!(ix; alg, order = Perm(ordr, vec(A)), scratch, dims...)
+    if alg == DEFAULT_STABLE
+        sort!(ix; alg, order = PermFast(ordr, vec(A)), scratch, dims...)
+    else
+        sort!(ix; alg, order = Perm(ordr, vec(A)), scratch, dims...)
+    end
 end
 
 
@@ -1615,7 +1666,11 @@ function sortperm!(ix::AbstractArray{T}, A::AbstractArray;
     if !initialized
         ix .= LinearIndices(A)
     end
-    sort!(ix; alg, order = Perm(ord(lt, by, rev, order), vec(A)), scratch, dims...)
+    if alg == DEFAULT_STABLE
+        sort!(ix; alg, order = PermFast(ord(lt, by, rev, order), vec(A)), scratch, dims...)
+    else
+        sort!(ix; alg, order = Perm(ord(lt, by, rev, order), vec(A)), scratch, dims...)
+    end
 end
 
 # sortperm for vectors of few unique integers