From 3b6976b8a698b5ebd5791232939dfa64cdc4adb9 Mon Sep 17 00:00:00 2001 From: "Viral B. Shah" Date: Tue, 8 Sep 2026 12:21:55 +0000 Subject: [PATCH 1/2] Add fast `sort!`/`sort` for `SparseMatrixCSC` (#302) The generic `Base.sort!(A; dims)` fallback sorts a sparse matrix through `getindex`/`setindex!` on slices, which densifies the structure and is orders of magnitude slower than sorting the equivalent dense matrix. Sorting a column only permutes its stored entries: the values sorting before `zero(eltype(A))` end up at the top of the column, the rest at the bottom, and the structural zeros stay contiguous in between. So each column can be sorted in place by sorting its stored values and rewriting its row indices, exactly as `sort!` already does for sparse vectors (#303). Rows are sorted the same way, on a transposed copy that is transposed back afterwards. `nnz` is now preserved instead of growing, and `sort` returns a `SparseMatrixCSC` for `dims = 1` rather than a `ReshapedArray` wrapping a `SparseVector`. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01PddHRZV2UXjpTPRkoQoM8j --- src/sparsematrix.jl | 58 ++++++++++++++++++++++++++++++++++++++++ src/sparsevector.jl | 6 ++--- test/sparsematrix_ops.jl | 55 +++++++++++++++++++++++++++++++++++++ 3 files changed, 115 insertions(+), 4 deletions(-) diff --git a/src/sparsematrix.jl b/src/sparsematrix.jl index f2f8673b..f5ac4ad5 100644 --- a/src/sparsematrix.jl +++ b/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 diff --git a/src/sparsevector.jl b/src/sparsevector.jl index 7bb56f6e..84b72054 100644 --- a/src/sparsevector.jl +++ b/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...) diff --git a/test/sparsematrix_ops.jl b/test/sparsematrix_ops.jl index a0d941a3..ecc8416f 100644 --- a/test/sparsematrix_ops.jl +++ b/test/sparsematrix_ops.jl @@ -759,4 +759,59 @@ 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 "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 From d2200ce4c7fb4fb64d28b866a152b66477d981e9 Mon Sep 17 00:00:00 2001 From: "Viral B. Shah" Date: Tue, 8 Sep 2026 14:07:14 +0000 Subject: [PATCH 2/2] Test empty and zero-size sparse matrices explicitly `Base.sort` on a zero-size dense matrix throws `ArgumentError: step cannot be zero`, so the size-parametrised testset compares those cases against the input rather than against a dense reference. Pin down what the sparse methods actually do there instead of leaving it implicit: 0xn, mx0 and 0x0 round-trip unchanged with the right size, element/index types and `nnz == 0`, and structurally empty matrices of non-zero size still match dense and leave the column pointers alone. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01PddHRZV2UXjpTPRkoQoM8j --- test/sparsematrix_ops.jl | 38 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/test/sparsematrix_ops.jl b/test/sparsematrix_ops.jl index ecc8416f..c2639e6d 100644 --- a/test/sparsematrix_ops.jl +++ b/test/sparsematrix_ops.jl @@ -804,6 +804,44 @@ end @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])