Julia: sortslices(a; dims=1) is slow for numerical arrays

[PetrKryslUCSD]

Julia sorting functions are slow compared to Matlab.

---

Outdated begin

Matlab R2012b sorting of integer arrays:

>> a=randi(30000000,30000000,2);
>> tic; b=sort(a,2); toc
Elapsed time is 0.882043 seconds.
>> tic; b=sort(a,2); toc
Elapsed time is 0.903528 seconds.
>> tic; b=sort(a,2); toc
Elapsed time is 0.900367 seconds.
``

Julia sorting of integer arrays:
julia> a=rand(Int64,30000000,2);

julia> @time sort(a,2);
elapsed time: 46.625749977 seconds (17279978752 bytes allocated, 19.15% gc time)

The Julia sort() is more than 40 times slower than Matlab.

Outdated end

---

The sortslices(a; dims=1) has a similar problem. Still working with the same 30,000,000 x 2 array size:

>>  a=rand(Int64,30000000,2);
>> tic; c=sortrows(a); toc
Elapsed time is 12.071053 seconds.

Compare with

julia> @time sortslices(a; dims=1);
100.747272 seconds (31.13 M allocations: 2.737 GiB, 1.78% gc time)

[jakebolewski]

It won't explain all of the current gap but does matlab do parallel sort under the hood? Do the timings change if you set maxNumCompThreads(1)?

[mbaz]

maxNumCompThreads(1) doesn't produce any change for me on Matlab R2013b.

[PetrKryslUCSD]

Good guess: for sort() using R2012b:

>> maxNumCompThreads(1)
Warning: maxNumCompThreads will be removed in a future release. Please
remove any instances of this function from your
code.
> In maxNumCompThreads at 27

ans =

     2

>> tic; b=sort(a,2); toc
Elapsed time is 7.099417 seconds.

>> maxNumCompThreads(8)
Warning: maxNumCompThreads will be removed in a future release. Please
remove any instances of this function from your
code.
> In maxNumCompThreads at 27

ans =

     1

>> tic; b=sort(a,2); toc
Elapsed time is 1.343386 seconds.

[ViralBShah]

Cc: @kmsquire

[kmsquire]

So, the problem is that sort along a particular dimension is implemented using mapslices, which has the same problem map does, in that it is taking an anonymous function and has no type information about that function.

I'm willing to take a crack at fixing this, but I'd love some suggestions. In particular, at a high level, it would seem that staged functions and/or cartesian indexing might be helpful, but I'm not familiar with either of these enough to know for sure. Any suggestions on which I should try to learn more about, if either?

Cc: @timholy

[ViralBShah]

Also Cc: @jutho

[timholy]

Ah, an interesting challenge.

The easy way would be to do the slicing yourself, and call sort on the resulting vector:

indexes = [Colon() for i = 1:N]
slices = [(indexes[dim] = i; slice(A, indexes...)) for i = 1:size(A,dim)]
p = sortperm(slices; kws...)
indexes[dim] = p
A[indexes...]

If profiling reveals that creating slices is a bottleneck, there may be ways to do better.

[Jutho]

As an alternative, this would be a (mutating) stagedfunction definition

using Base.Cartesian
stagedfunction sort!{T,N}(a::AbstractArray{T,N}, dim::Int; kwargs...)
    return quote
        @nif $N d->(dim==d) d->begin
            @nexprs $(N-1) n->(d_{n}=size(a,n<d ? n : n+1))
            @nloops $(N-1) i n->(1:d_{n}) begin
                s=@ncall $N slice a m->(m<d ? i_{m} : (m==d ? Colon() : i_{m-1}))
                sort!(s; kwargs...)
            end
        end
    end
end

Not really tested, so I hope this works.

[kmsquire]

I had planned to get to this right after it was submitted, but haven't been able to make the time. I might be able to look at it tomorrow (or I might not), but if someone wants to take @timholy's or @jutho's ideas and run with them, please feel free.

Cc: @tcovert

[tcovert]

For what its worth, whatever is going on in the sort routine for DataFrames.jl seems to be much faster than sortrows().

[ViralBShah]

I was just experimenting with slicing instead of using subarrays, and it is not any faster for sortrows. I also tried reinterpreting to a Vector of tuples for sortrows, but that blows up the memory usage and is significantly slower.

For sortrows on a few columns, I think the simplest thing may be to do something like:

function mysortrows{T}(A::AbstractMatrix{T}; kws...)
    for i in size(A,2):-1:1
        p = sortperm(A[:,i])
        A = A[p,:]
    end
    A
end

[ViralBShah]

The fastest thing to do for sortrows on any number of columns is probably something like (not tested at all):

function mysortrows{T}(A::AbstractMatrix{T}; kws...)
   X = copy(A)
   for i in 1:size(A,2)
      x = X[:,i]
      p = sortperm(x)
      X = X[p,:]
      x = x[p]
      f = find(x[1:end-1] == x[2:end]) # Find rows that need sorting on secondary keys
      if length(f)
          X = sub(X, f, 1:size(A,2))
      else
          break
       end
       X
   end

This way, unless you need to sort on secondary keys, sortrows can be completed for the cost of just one sort.

[lindahua]

I think this needs to be done using a specialized algorithm instead of something like mapslice.

Here is the pseudo code:

sort the index based on the first column --> p
for j = 2 : size(A, 2)
     find consecutive sections, 
         such that each section i0:i1, A[p[i0:i1], j-1] is the same value.
     sort each section.
end

The implementation should be very careful to avoid unnecessary temporary allocation.

[rened]

I can offer the following as a starting point:

sortrows(a; kargs...) = a[sortrows(a, collect(1:size(a,1)), 1, 1, size(a,1); kargs...),:]
function sortrows(a, perm, col, from, to; kargs...)
    (to - from == 0 || col > size(a,2)) && return perm

    ind = from:to
    perm[ind] = perm[ind][sortperm(a[perm[ind],col]; kargs...)]

    i = from;
    for m = from+1:to
        if a[perm[m], col] != a[perm[i], col]
            sortrows(a, perm, col+1, i, m-1;kargs...)
            i = m
        end
    end
    i < to && sortrows(a, perm, col+1, i, to;kargs...)
    perm
end
sortcols(a;kargs...) = a[:,sortrows(a', collect(1:size(a,2)), 1, 1, size(a,2);kargs...)]

a = rand(int(1e6),2)
@time Base.sortrows(a)
@time sortrows(a)
assert(isequal(Base.sortrows(a),sortrows(a)))
a = rand(1:10,int(1e6),10)
@time Base.sortrows(a)
@time sortrows(a)
assert(isequal(Base.sortrows(a),sortrows(a)))

elapsed time: 5.078052154 seconds (463967744 bytes allocated, 3.32% gc time)
elapsed time: 0.486893081 seconds (160000752 bytes allocated, 36.42% gc time)
elapsed time: 8.196978319 seconds (527967744 bytes allocated, 2.19% gc time)
elapsed time: 1.154467671 seconds (613507216 bytes allocated)

[ViralBShah]

This looks pretty good. How about a PR?

[rened]

The only problem: these timings are from 0.3.7 - using master I get:

elapsed time: 1.329124183 seconds (114 MB allocated, 8.37% gc time in 4 pauses with 1 full sweep)
elapsed time: 0.378710152 seconds (183 MB allocated, 18.08% gc time in 8 pauses with 1 full sweep)
elapsed time: 2.17216094 seconds (175 MB allocated, 4.17% gc time in 4 pauses with 1 full sweep)
elapsed time: 2.465555558 seconds (749 MB allocated, 4.53% gc time in 30 pauses with 1 full sweep)

So in the case where all the items in the first col are dissimilar it would still be 4x faster, but in the second case the built-in seems to be better. I can try to bring the times down some more and then submit a PR.

[abhishekmalali]

I made an attempt to incorporate @lindahua 's idea and write a basic code and see its performance. Code that I have written has a long way to go.

function mysortrows(B::AbstractMatrix,cols::Array; kws...)
       for i = 1:length(cols) 
        if i == 1
            p =sortperm(B[:,cols[i]]; kws...);
            B = B[p,:];
        else
            i0_old = 0;
            i1_old = 0;
            i0_new = 0;
            i1_new = 0;
            for j = 1:size(B,1)-1
                if B[j,cols[1:i-1]] == B[j+1,cols[1:i-1]] && i0_old == i0_new
                    i0_new = j;
                elseif B[j,cols[1:i-1]] != B[j+1,cols[1:i-1]] && i0_old != i0_new && i1_new == i1_old
                    i1_new = j;
                elseif i0_old != i0_new && j == size(B,1)-1
                    i1_new = j+1;
                end
                if i0_new != i0_old && i1_new != i1_old
                    p = sortperm(B[i0_new:i1_new,cols[i]]; kws...);
                    B[i0_new:i1_new,:] = B[i0_new:i1_new,:][p,:];
                    i0_old = i0_new;
                    i1_old = i1_new;
                end
            end
            end
    end
    return B
end

A = rand(1:10, 10^6, 8);
mysortrows(A,[1])
@time p_mysortrows = mysortrows(A,[1:8])
sortrows(A)
@time p_sortrows = sortrows(A,by = x->(x[1],x[2],x[3],x[4],x[5],x[6],x[7],x[8]))
print(p_sortrows == p_mysortrows)

The results for the above code were as following

elapsed time: 8.650744886 seconds (6510021088 bytes allocated, 28.04% gc time)
elapsed time: 81.553844845 seconds (7932920840 bytes allocated, 33.34% gc time)
true

Most test cases, I saw a considerable improvement considering this comes with a performance overhead the way it was written.

[ViralBShah]

Bump @rened. Would be nice to get even some of the speed improvements you have.

[rened]

@ViralBShah I'll try to improve on what I have in the coming days and submit a PR!

[ViralBShah]

Thank you!

[rened]

@ViralBShah I tried, by I can't consistently beat sortrows on master any more...

[ViralBShah]

:-( or should that be :-)

[kmsquire]

Unfortunately, the original problem is still there.

[KristofferC]

julia> @time sort(a,dims=2);
  1.146972 seconds (46 allocations: 915.529 MiB, 3.75% gc time)

Think this can be closed.

[kmsquire]

From a note at the bottom of the description:

[ViralBShah note: The sort performance discussed here is no longer an issue, only the sortrows]

@KristofferC, what doessortrows performance look like on your machine?

[KristofferC]

That was hard to see, I updated with the current syntax.