Introduce PermUnstable to speed up sortperm

Author: petvana

base/ordering.jl

@@ -13,7 +13,7 @@ import .Base:
 
 export # not exported by Base
     Ordering, Forward, Reverse,
-    By, Lt, Perm,
+    By, Lt, Perm, PermUnstable,
     ReverseOrdering, ForwardOrdering,
     DirectOrdering,
     lt, ord, ordtype
@@ -106,8 +106,21 @@ struct Perm{O<:Ordering,V<:AbstractVector} <: Ordering
     data::V
 end
 
+"""
+    PermUnstable(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 PermUnstable{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::PermUnstable) = PermUnstable(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::PermUnstable, 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

@@ -573,9 +573,11 @@ function _sort!(v::AbstractVector, a::MissingOptimization, o::Ordering, kw)
     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 && nonmissingtype(eltype(o.data)) != eltype(o.data) &&
+            (o isa Perm{DirectOrdering} || o isa PermUnstable{DirectOrdering})
+        PermT = o isa Perm{DirectOrdering} ? Perm : PermUnstable
         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
@@ -614,15 +616,17 @@ function _sort!(v::AbstractVector, a::IEEEFloatOptimization, o::Ordering, kw)
         else
             _sort!(iv, a.next, Forward, (;kw..., lo=j+1, hi, scratch))
         end
-    elseif eltype(v) <: Integer && o isa Perm && o.order isa DirectOrdering && is_concrete_IEEEFloat(eltype(o.data))
+    elseif eltype(v) <: Integer && (o isa Perm || o isa PermUnstable) &&
+            o.order isa DirectOrdering && is_concrete_IEEEFloat(eltype(o.data))
         lo, hi = send_to_end!(i -> isnan(@inbounds o.data[i]), v, o.order, true; lo, hi)
         ip = reinterpret(UIntType(eltype(o.data)), o.data)
         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))
+        PermT = o isa Perm ? Perm : PermUnstable
+        scratch = _sort!(v, a.next, PermT(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, PermT(Forward, ip), (;kw..., lo=j+1, hi, scratch))
         else
-            _sort!(v, a.next, Perm(Forward, ip), (;kw..., lo=j+1, hi, scratch))
+            _sort!(v, a.next, PermT(Forward, ip), (;kw..., lo=j+1, hi, scratch))
         end
     else
         _sort!(v, a.next, o, kw)
@@ -1482,7 +1486,7 @@ function partialsortperm!(ix::AbstractVector{<:Integer}, v::AbstractVector,
     end
 
     # do partial quicksort
-    _sort!(ix, _PartialQuickSort(k), Perm(ord(lt, by, rev, order), v), (;))
+    _sort!(ix, _PartialQuickSort(k), PermUnstable(ord(lt, by, rev, order), v), (;))
 
     maybeview(ix, k)
 end
@@ -1554,7 +1558,8 @@ function sortperm(A::AbstractArray;
         end
     end
     ix = copymutable(LinearIndices(A))
-    sort!(ix; alg, order = Perm(ordr, vec(A)), scratch, dims...)
+    PermT = alg == DEFAULT_STABLE ? PermUnstable : Perm
+    sort!(ix; alg, order = PermT(ordr, vec(A)), scratch, dims...)
 end
 
 
@@ -1608,7 +1613,8 @@ 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...)
+    PermT = alg == DEFAULT_STABLE ? PermUnstable : Perm
+    sort!(ix; alg, order = PermT(ord(lt, by, rev, order), vec(A)), scratch, dims...)
 end
 
 # sortperm for vectors of few unique integers