Add fast sort!/sort for SparseMatrixCSC

src/sparsematrix.jl

@@ -1777,6 +1777,64 @@ function permute(A::AbstractSparseMatrixCSC{Tv,Ti}, p::AbstractVector{<:Integer}
     unchecked_noalias_permute!(X, A, p, q, C)
 end
 
+## Sorting
+
+#sorting TODO: integrate with `Base.Sort.IEEEFloatOptimization`'s partitioning by zero
+searchsortedfirst_discard_keywords(v::AbstractVector, x; lt=isless, by=identity,
+    rev::Union{Bool,Nothing}=nothing, order::Base.Order.Ordering=Forward, kws...) =
+        searchsortedfirst(v, x, Base.Order.ord(lt,by,rev,order))
+
+"""
+Sort the stored entries of each column of `A` in place, rewriting the row indices so that
+the values sorting before `zero(eltype(A))` end up at the top of their column and the
+remaining values at the bottom, with the structural zeros in between. `nnz(A)` and the
+column pointers are left untouched.
+"""
+function _sortcolumns!(A::AbstractSparseMatrixCSC; kws...)
+    require_one_based_indexing(A)
+    rows = rowvals(A)
+    vals = nonzeros(A)
+    m = size(A, 1)
+    z = zero(eltype(A))
+    for j in axes(A, 2)
+        r = nzrange(A, j)
+        isempty(r) && continue
+        col = view(vals, r)
+        sort!(col; kws...)
+        # `i-1` stored values sort before the structural zeros and `length(r)-i+1` after
+        i = searchsortedfirst_discard_keywords(col, z; kws...)
+        k = first(r)
+        @inbounds for t in 1:i-1
+            rows[k] = t
+            k += 1
+        end
+        @inbounds for t in (m - length(r) + i):m
+            rows[k] = t
+            k += 1
+        end
+    end
+    return A
+end
+
+function Base.sort!(A::AbstractSparseMatrixCSC; dims::Integer, kws...)
+    if dims == 1
+        _sortcolumns!(A; kws...)
+    elseif dims == 2
+        # the rows of `A` are the columns of `transpose(A)`, which is cheap to form and
+        # cheap to transpose back once its columns are sorted
+        At = ftranspose(A, identity)
+        _sortcolumns!(At; kws...)
+        transpose!(A, At)
+    else
+        throw(ArgumentError(lazy"dimension out of range, got dims = $dims, expected 1 or 2"))
+    end
+    return A
+end
+
+# the generic `Base.sort` for matrices goes through `permutedims`/`reshape` and does not
+# return a `SparseMatrixCSC` for `dims = 1`
+Base.sort(A::AbstractSparseMatrixCSC; kws...) = sort!(copy(A); kws...)
+
 ## fkeep! and children tril!, triu!, droptol!, dropzeros[!]
 
 function _fkeep!(f::F, A::AbstractSparseMatrixCSC) where F<:Function

src/sparsevector.jl

@@ -2272,10 +2272,8 @@ function _densifystarttolastnz!(x::SparseVector)
     x
 end
 
-#sorting TODO: integrate with `Base.Sort.IEEEFloatOptimization`'s partitioning by zero
-searchsortedfirst_discard_keywords(v::AbstractVector, x; lt=isless, by=identity,
-    rev::Union{Bool,Nothing}=nothing, order::Base.Order.Ordering=Forward, kws...) =
-        searchsortedfirst(v,x,Base.Order.ord(lt,by,rev,order))
+# `searchsortedfirst_discard_keywords` is defined alongside the sparse matrix sorting
+# methods in sparsematrix.jl
 function sort!(x::AbstractCompressedVector; kws...)
     nz = nonzeros(x)
     sort!(nz; kws...)

test/sparsematrix_ops.jl

@@ -759,4 +759,97 @@ end
     @test isdiag(S)
 end
 
+@testset "sort/sort! of a sparse matrix" begin
+    # `sort` of a dense 0-dimension-along-`dims` matrix errors in Base, so those sizes are
+    # compared against the input itself rather than against a dense reference
+    @testset "size = ($m, $n), density = $d" for (m, n) in ((6, 5), (1, 1), (0, 3), (3, 0),
+                                                            (1, 9), (9, 1), (20, 13)),
+                                                 d in (0.0, 0.05, 0.3, 1.0)
+        A = sprand(m, n, d)
+        M = Matrix(A)
+        for dims in (1, 2), kws in ((;), (; rev=true), (; by=abs), (; by=x -> -x),
+                                    (; lt=(x, y) -> isless(y, x)),
+                                    (; alg=Base.DEFAULT_STABLE))
+            expected = (m == 0 || n == 0) ? M : sort(M; dims, kws...)
+            B = copy(A)
+            @test sort!(B; dims, kws...) === B
+            @test B isa SparseMatrixCSC
+            @test Matrix(B) == expected
+            # sorting only moves the stored entries around
+            @test nnz(B) == nnz(A)
+            S = sort(A; dims, kws...)
+            @test S isa SparseMatrixCSC
+            @test Matrix(S) == expected
+            @test A == sparse(M) # `sort` leaves its argument alone
+        end
+    end
+
+    @testset "index type $Ti" for Ti in (Int32, Int64)
+        A = SparseMatrixCSC{Float64,Ti}(sprand(11, 7, 0.4))
+        for dims in (1, 2)
+            @test sort(A; dims) isa SparseMatrixCSC{Float64,Ti}
+            @test Matrix(sort(A; dims)) == sort(Matrix(A); dims)
+        end
+    end
+
+    @testset "keyword arguments" begin
+        A = sprand(50, 50, 0.1)
+        # `scratch` is forwarded to the underlying `sort!` and ignored by the search for
+        # where the structural zeros belong (see #335)
+        @test Matrix(sort!(copy(A); dims=1, scratch=Vector{Float64}(undef, 50))) ==
+            sort(Matrix(A); dims=1)
+        @test_throws MethodError sort!(copy(A); dims=1, banana=:blue)
+        @test_throws ArgumentError sort!(copy(A); dims=3)
+        @test_throws ArgumentError sort!(copy(A); dims=0)
+        @test_throws UndefKeywordError sort!(copy(A))
+    end
+
+    @testset "empty and zero-size matrices" begin
+        # `Base.sort` on a zero-size *dense* matrix throws `ArgumentError: step cannot be
+        # zero`, so there is no dense reference to compare against here; the sparse methods
+        # just return the (empty) matrix unchanged
+        @testset "size = ($m, $n)" for (m, n) in ((0, 3), (3, 0), (0, 0))
+            A = spzeros(m, n)
+            for dims in (1, 2)
+                B = copy(A)
+                @test sort!(B; dims) === B
+                @test size(B) == (m, n)
+                @test nnz(B) == 0
+                @test B == A
+                S = sort(A; dims)
+                @test S isa SparseMatrixCSC{Float64,Int}
+                @test size(S) == (m, n)
+                @test nnz(S) == 0
+            end
+        end
+
+        # structurally empty, but not zero-size: here dense does give a reference
+        @testset "all structural zeros, size = ($m, $n)" for (m, n) in ((1, 1), (5, 4))
+            A = spzeros(m, n)
+            for dims in (1, 2)
+                B = sort!(copy(A); dims)
+                @test Matrix(B) == sort(Matrix(A); dims)
+                @test nnz(B) == 0
+                @test getcolptr(B) == getcolptr(A)
+            end
+        end
+
+        # a single column/row that is entirely structural next to a populated one
+        A = SparseMatrixCSC(4, 3, [1, 1, 5, 5], [1, 2, 3, 4], [1.0, -2.0, 0.0, 3.0])
+        for dims in (1, 2)
+            @test Matrix(sort(A; dims)) == sort(Matrix(A); dims)
+            @test nnz(sort(A; dims)) == nnz(A)
+        end
+    end
+
+    @testset "stored zeros" begin
+        # column 1 stores an explicit zero next to structural zeros
+        A = SparseMatrixCSC(4, 2, [1, 3, 4], [1, 3, 2], [0.0, -1.0, 2.0])
+        for dims in (1, 2)
+            @test Matrix(sort(A; dims)) == sort(Matrix(A); dims)
+            @test nnz(sort(A; dims)) == nnz(A)
+        end
+    end
+end
+
 end # module