Sorting in-place several Vector's simultaneously

@eldee I just wish to thank you for your help. : ) In the last 3 days, I’ve learnt something new (@generated) and I think it’s super cool that Julia allows us to hack into its functionality so easily (in my case, using all of the sorting capabilities on my custom struct so long as I define size, length, getindex, setindex) without losing any of the performance!!!

For anyone else interested in this, I am posting the finalized code and its benchmarks:

struct Zip{N, T<:Tuple} <: AbstractVector{Tuple}
   vec::T
   Zip{N,T}(vecs::T) where {N,T} = length(vecs)==N ? new(vecs) : error("lengths don't match!") 
end
Zip(vecs) = Zip{length(vecs), typeof(vecs)}(vecs)
Base.size(z::Zip) = (length(z),)  # for AbstractVector compliance
@generated function Base.length(z::Zip{N,T}) where {N,T}
   Expr(:call, :min, (:(length(z.vec[$j])) for j=1:N)...) end
@generated function Base.getindex(z::Zip{N,T}, i::Int) where {N,T}
   Expr(:tuple, (:(@inbounds(z.vec[$j][i])) for j=1:N)...) end
@generated function Base.setindex!(z::Zip{N,T}, vals::Tuple, i::Int)  where {N,T}
   Expr(:block, (:(@inbounds(z.vec[$j][i] = vals[$j])) for j=1:N)...) end;

function tmp!(x,y)  
   n = min(length.((x,y))...)
   p = sortperm(x)[1:n]
   view(x,1:n) .= x[p]
   view(y,1:n) .= y[p] 
end

julia> begin n=10^6;   
   @btime sort!(x, alg=QuickSort)                                       setup=(x=rand(1:$n,$n));                    
   @btime tmp!(x,y)                                                     setup=(x=rand(1:$n,$n); y=rand(1.0:$n,$n+10));  
   @btime sort!(Zip((x,y)), by=first, alg=QuickSort)                    setup=(x=rand(1:$n,$n); y=rand(1.0:$n,$n+10)); 
   @btime sort!(StructArray((x,view(y,1:$n))), by=first, alg=QuickSort) setup=(x=rand(1:$n,$n); y=rand(1.0:$n,$n+10)); end; 
  71.717 ms (0 allocations: 0 bytes)
  159.054 ms (15 allocations: 38.15 MiB)
  75.413 ms (0 allocations: 0 bytes)
  81.581 ms (0 allocations: 0 bytes)
5 Likes

I don’t think you need @generated for this sort of thing. The ntuple function is your friend here. For example:

Base.length(z::Zip{N}) where {N} = minimum(ntuple(j -> length(z.vec[j]), Val{N}()))

should do the trick, I think?

Note that this is this container type has abstractly typed elements (it indicates that every element could be a different type of Tuple), and is probably not what you want:

julia> AbstractVector{Tuple{Int,Int}} <: AbstractVector{Tuple}
false

(Implementing this kind of thing properly is tricky, but should be taken care of by StructArrays.jl here.)

1 Like