Sorting wrt multiple arrays

[matthieugomez]

What's the best way to sort lexicographically wrt 2 arrays a and b ?
The code below is quite slow, so I'm wondering whether there is a better solution

n = 1_000_000
a = rand(1:n, n)
b = rand(1:10, n)

@time sortperm([(a[i], b[i]) for i in 1:length(a)])
#elapsed time: 9.676994989 seconds   (287943584 bytes allocated, 16.39% gc time)
@time sortperm([(sub(a,i)[1], sub(b,i)[1]) for i in 1:length(a)], order= Base.Order.Lexicographic)
#elapsed time: 9.446759537 seconds (560247620 bytes allocated, 17.83% gc time)

[matthieugomez]

The following works well actually.

@time sortperm(collect(zip(a, b)))
#elapsed time: 1.230827695 seconds (63991856 bytes allocated, 9.36% gc time)

Is this the correct way to do it then? I'm not sure I understand the difference between these syntaxes.

[stevengj]

Don't time code in global scope (see the Julia performance tips). If you do:

s0(a,b) = sortperm([(a[i], b[i]) for i in 1:length(a)])
s1(a,b) = sortperm(collect(zip(a, b)))
@time s0(a,b);
@time s1(a,b);

then you will get almost identical timings.

[stevengj]

It should in principle be possible to do this without making a copy of the arrays, by defining a new AbstractVector type that wraps a and b. e.g. I tried:

type ZipVector{T,S} <: AbstractVector{Tuple{T,S}}
    a::Vector{T}
    b::Vector{S}
end
Base.getindex(z::ZipVector, i::Int) = (a[i],b[i])
Base.size(z::ZipVector) = size(a)

n = 1_000_000
a = rand(1:n, n)
b = rand(1:10, n)
@time sortperm(ZipVector(a,b));

but it is insanely slow (in Julia 0.4): 20.672 seconds (187 M allocations: 3570 MB, 2.12% gc time).

It seems like getindex is allocating (a[i],b[i]) on the heap each time it is called, hence the huge number of allocations. @JeffBezanson, I thought that returning tuples was supposed to be cheap, especially after JuliaLang/julia#10380?

[stevengj]

It definitely seems to be getindex that is causing the heap-allocation. I tried:

function foo(X)
    for (a,b) in X
        a + b
    end
end
@time foo(ZipVector(a,b));

and got 2.156 seconds (19153 k allocations: 451 MB, 2.92% gc time).

[stevengj]

Okay, it is a type-inference problem:

Z = ZipVector(a,b)::ZipVector{Int64,Int64}
Base.Test.@inferred getindex(Z, 1)

gives

ERROR: return type Tuple{Int64,Int64} does not match inferred return type Tuple{Any,Any}
 in error at ./error.jl:21

[simonster]

I think you mean:

Base.getindex(z::ZipVector, i::Int) = (z.a[i],z.b[i])
Base.size(z::ZipVector) = size(z.a)

[stevengj]

Oh, it is just a stupid typo. I was returning (a[i],b[i]) when it should have been (z.a[i],z.b[i]). Correcting the getindex method to

Base.getindex(z::ZipVector, i::Int) = (z.a[i],z.b[i])

gives

julia> @time sortperm(ZipVector(a,b));
 647.231 milliseconds (11 allocations: 7813 KB)

which is about 30% slower than the version above that makes a copy of the array. As @JeffBezanson remarked to me, making a copy is O(n) while the sort performs O(n log n) comparisons, so it is typically going to be preferable for cache locality to make a copy (unless maybe the arrays are already in-cache).

But it may be nice for other reasons to do it without making a copy, and it is nice that this is (a) possible and (b) not too much slower.

[stevengj]

@simonster, yes, you spotted it as I was writing up my last post.

[matthieugomez]

Thanks a lot for getting back to me. The Zipvector trick is very good to know. I don't get the conclusion though. I get half a second for sortperm(Zipvector), and more than 1s for s0 and s1. What is the faster deep copy version you mention?

[stevengj]

For me, s0 and s1 (which both make copies of the data) were slightly faster. Probably it is because I'm using Julia 0.4, and arrays of tuples are much more efficient now than in Julia 0.3.x (which is what I guess you are using?).