diff --git a/stdlib/SparseArrays/src/higherorderfns.jl b/stdlib/SparseArrays/src/higherorderfns.jl index 66ab8b3f60e52..029f4abd1fcf7 100644 --- a/stdlib/SparseArrays/src/higherorderfns.jl +++ b/stdlib/SparseArrays/src/higherorderfns.jl @@ -2,6 +2,11 @@ module HigherOrderFns +_goodbuffers(n, colptr, rowval, nzval) = length(colptr) == n + 1 && colptr[end] - 1 == length(rowval) == length(nzval) +_goodbuffers(S::SparseMatrixCSC) = _goodbuffers(S.n, S.colptr, S.rowval, S.nzval) +_checkbuffers(S::SparseMatrixCSC) = (@assert _goodbuffers(S); S) +_checkbuffers(S::SparseVector) = (@assert length(S.nzval) == length(S.nzind); S) + # This module provides higher order functions specialized for sparse arrays, # particularly map[!]/broadcast[!] for SparseVectors and SparseMatrixCSCs at present. import Base: map, map!, broadcast, copy, copyto! @@ -9,7 +14,8 @@ import Base: map, map!, broadcast, copy, copyto! using Base: front, tail, to_shape using ..SparseArrays: SparseVector, SparseMatrixCSC, AbstractSparseVector, AbstractSparseMatrix, AbstractSparseArray, indtype, nnz, nzrange, - SparseVectorUnion, AdjOrTransSparseVectorUnion, nonzeroinds, nonzeros + SparseVectorUnion, AdjOrTransSparseVectorUnion, nonzeroinds, nonzeros, + spzeros using Base.Broadcast: BroadcastStyle, Broadcasted, flatten using LinearAlgebra @@ -126,17 +132,27 @@ const SpBroadcasted2{Style<:SPVM,Axes,F,Args<:Tuple{SparseVecOrMat,SparseVecOrMa @inline storedvals(A::SparseVecOrMat) = A.nzval @inline setcolptr!(A::SparseVector, j, val) = val @inline setcolptr!(A::SparseMatrixCSC, j, val) = A.colptr[j] = val -function trimstorage!(A::SparseVecOrMat, maxstored) - resize!(storedinds(A), maxstored) - resize!(storedvals(A), maxstored) - return maxstored -end + function expandstorage!(A::SparseVecOrMat, maxstored) - length(storedinds(A)) < maxstored && resize!(storedinds(A), maxstored) - length(storedvals(A)) < maxstored && resize!(storedvals(A), maxstored) + if _capacity(A) < maxstored + sizehint!(storedinds(A), maxstored) + sizehint!(storedvals(A), maxstored) + end return maxstored end +function _reset!(C::SparseVector) + empty!(C.nzind) + empty!(C.nzval) +end + +function _reset!(C::SparseMatrixCSC) + fill!(C.colptr, 1) + empty!(C.rowval) + empty!(C.nzval) +end + +_capacity(A::SparseVecOrMat) = capacity(storedinds(A)) # (2) map[!] entry points map(f::Tf, A::SparseVector) where {Tf} = _noshapecheck_map(f, A) @@ -169,18 +185,17 @@ end copy(bc::SpBroadcasted1) = _noshapecheck_map(bc.f, bc.args[1]) @inline function copyto!(C::SparseVecOrMat, bc::Broadcasted0{Nothing}) - isempty(C) && return _finishempty!(C) + isempty(C) && return _checkbuffers(C) f = bc.f fofnoargs = f() if _iszero(fofnoargs) # f() is zero, so empty C - trimstorage!(C, 0) - _finishempty!(C) + _reset!(C) else # f() is nonzero, so densify C and fill with independent calls to f() _densestructure!(C) storedvals(C)[1] = fofnoargs broadcast!(f, view(storedvals(C), 2:length(storedvals(C)))) end - return C + return _checkbuffers(C) end function _diffshape_broadcast(f::Tf, A::SparseVecOrMat, Bs::Vararg{SparseVecOrMat,N}) where {Tf,N} @@ -223,22 +238,13 @@ _maxnnzfrom(shape::NTuple{2}, A::SparseMatrixCSC) = nnz(A) * div(shape[1], A.m) @inline _unchecked_maxnnzbcres(shape, As...) = _unchecked_maxnnzbcres(shape, As) @inline _checked_maxnnzbcres(shape::NTuple{1}, As...) = shape[1] != 0 ? _unchecked_maxnnzbcres(shape, As) : 0 @inline _checked_maxnnzbcres(shape::NTuple{2}, As...) = shape[1] != 0 && shape[2] != 0 ? _unchecked_maxnnzbcres(shape, As) : 0 -@inline function _allocres(shape::NTuple{1}, indextype, entrytype, maxnnz) - storedinds = Vector{indextype}(undef, maxnnz) - storedvals = Vector{entrytype}(undef, maxnnz) - return SparseVector(shape..., storedinds, storedvals) -end -@inline function _allocres(shape::NTuple{2}, indextype, entrytype, maxnnz) - pointers = Vector{indextype}(undef, shape[2] + 1) - storedinds = Vector{indextype}(undef, maxnnz) - storedvals = Vector{entrytype}(undef, maxnnz) - return SparseMatrixCSC(shape..., pointers, storedinds, storedvals) -end +@inline _allocres(shape::NTuple, indextype, entrytype, maxnnz) = sizehint!(spzeros(entrytype, indextype, shape...), maxnnz) # (4) _map_zeropres!/_map_notzeropres! specialized for a single sparse vector/matrix "Stores only the nonzero entries of `map(f, Array(A))` in `C`." function _map_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat) where Tf - spaceC::Int = min(length(storedinds(C)), length(storedvals(C))) + _reset!(C) + spaceC::Int = _capacity(C) Ck = 1 @inbounds for j in columns(C) setcolptr!(C, j, Ck) @@ -246,15 +252,14 @@ function _map_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat) where Tf Cx = f(storedvals(A)[Ak]) if !_iszero(Cx) Ck > spaceC && (spaceC = expandstorage!(C, Ck + nnz(A) - (Ak - 1))) - storedinds(C)[Ck] = storedinds(A)[Ak] - storedvals(C)[Ck] = Cx + push!(storedinds(C), storedinds(A)[Ak]) + push!(storedvals(C), Cx) Ck += 1 end end end @inbounds setcolptr!(C, numcols(C) + 1, Ck) - trimstorage!(C, Ck - 1) - return C + return _checkbuffers(C) end """ Densifies `C`, storing `fillvalue` in place of each unstored entry in `A` and @@ -273,19 +278,21 @@ function _map_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseVecOrMa end # NOTE: Combining the fill! above into the loop above to avoid multiple sweeps over / # nonsequential access of storedvals(C) does not appear to improve performance. - return C + return _checkbuffers(C) end # helper functions for these methods and some of those below @inline _densecoloffsets(A::SparseVector) = 0 @inline _densecoloffsets(A::SparseMatrixCSC) = 0:A.m:(A.m*(A.n - 1)) function _densestructure!(A::SparseVector) - expandstorage!(A, A.n) + resize!(A.nzval, A.n) + resize!(A.nzind, A.n) copyto!(A.nzind, 1:A.n) return A end function _densestructure!(A::SparseMatrixCSC) nnzA = A.m * A.n - expandstorage!(A, nnzA) + resize!(A.rowval, nnzA) + resize!(A.nzval, nnzA) copyto!(A.colptr, 1:A.m:(nnzA + 1)) for k in _densecoloffsets(A) copyto!(A.rowval, k + 1, 1:A.m) @@ -296,7 +303,8 @@ end # (5) _map_zeropres!/_map_notzeropres! specialized for a pair of sparse vectors/matrices function _map_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat, B::SparseVecOrMat) where Tf - spaceC::Int = min(length(storedinds(C)), length(storedvals(C))) + _reset!(C) + spaceC::Int = _capacity(C) rowsentinelA = convert(indtype(A), numrows(C) + 1) rowsentinelB = convert(indtype(B), numrows(C) + 1) Ck = 1 @@ -327,15 +335,14 @@ function _map_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat, B::SparseVe # corresponding broadcast code (below). if !_iszero(Cx) Ck > spaceC && (spaceC = expandstorage!(C, Ck + (nnz(A) - (Ak - 1)) + (nnz(B) - (Bk - 1)))) - storedinds(C)[Ck] = Ci - storedvals(C)[Ck] = Cx + push!(storedinds(C), Ci) + push!(storedvals(C), Cx) Ck += 1 end end end @inbounds setcolptr!(C, numcols(C) + 1, Ck) - trimstorage!(C, Ck - 1) - return C + return _checkbuffers(C) end function _map_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseVecOrMat, B::SparseVecOrMat) where Tf # Build dense matrix structure in C, expanding storage if necessary @@ -367,13 +374,14 @@ function _map_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseVecOrMa Cx != fillvalue && (storedvals(C)[jo + Ci] = Cx) end end - return C + return _checkbuffers(C) end # (6) _map_zeropres!/_map_notzeropres! for more than two sparse matrices / vectors function _map_zeropres!(f::Tf, C::SparseVecOrMat, As::Vararg{SparseVecOrMat,N}) where {Tf,N} - spaceC::Int = min(length(storedinds(C)), length(storedvals(C))) + _reset!(C) + spaceC::Int = _capacity(C) rowsentinel = numrows(C) + 1 Ck = 1 stopks = _colstartind_all(1, As) @@ -388,16 +396,15 @@ function _map_zeropres!(f::Tf, C::SparseVecOrMat, As::Vararg{SparseVecOrMat,N}) Cx = f(vals...) if !_iszero(Cx) Ck > spaceC && (spaceC = expandstorage!(C, Ck + min(length(C), _sumnnzs(As...)) - (sum(ks) - N))) - storedinds(C)[Ck] = activerow - storedvals(C)[Ck] = Cx + push!(storedinds(C), activerow) + push!(storedvals(C), Cx) Ck += 1 end activerow = min(rows...) end end @inbounds setcolptr!(C, numcols(C) + 1, Ck) - trimstorage!(C, Ck - 1) - return C + return _checkbuffers(C) end function _map_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, As::Vararg{SparseVecOrMat,N}) where {Tf,N} # Build dense matrix structure in C, expanding storage if necessary @@ -420,7 +427,7 @@ function _map_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, As::Vararg{Spars activerow = min(rows...) end end - return C + return _checkbuffers(C) end # helper methods for map/map! methods just above @@ -460,8 +467,9 @@ end # (7) _broadcast_zeropres!/_broadcast_notzeropres! specialized for a single (input) sparse vector/matrix function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat) where Tf - isempty(C) && return _finishempty!(C) - spaceC::Int = min(length(storedinds(C)), length(storedvals(C))) + isempty(C) && return _checkbuffers(C) + _reset!(C) + spaceC::Int = _capacity(C) # C and A cannot have the same shape, as we directed that case to map in broadcast's # entry point; here we need efficiently handle only heterogeneous C-A combinations where # one or both of C and A has at least one singleton dimension. @@ -479,8 +487,8 @@ function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat) where Cx = f(storedvals(A)[Ak]) if !_iszero(Cx) Ck > spaceC && (spaceC = expandstorage!(C, _unchecked_maxnnzbcres(size(C), A))) - storedinds(C)[Ck] = storedinds(A)[Ak] - storedvals(C)[Ck] = Cx + push!(storedinds(C), storedinds(A)[Ak]) + push!(storedvals(C), Cx) Ck += 1 end end @@ -499,16 +507,15 @@ function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat) where if !_iszero(fofAx) for Ci::indtype(C) in 1:numrows(C) Ck > spaceC && (spaceC = expandstorage!(C, _unchecked_maxnnzbcres(size(C), A))) - storedinds(C)[Ck] = Ci - storedvals(C)[Ck] = fofAx + push!(storedinds(C), Ci) + push!(storedvals(C), fofAx) Ck += 1 end end end end @inbounds setcolptr!(C, numcols(C) + 1, Ck) - trimstorage!(C, Ck - 1) - return C + return _checkbuffers(C) end function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseVecOrMat) where Tf # For information on this code, see comments in similar code in _broadcast_zeropres! above @@ -539,14 +546,15 @@ function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseV end end end - return C + return _checkbuffers(C) end # (8) _broadcast_zeropres!/_broadcast_notzeropres! specialized for a pair of (input) sparse vectors/matrices function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat, B::SparseVecOrMat) where Tf - isempty(C) && return _finishempty!(C) - spaceC::Int = min(length(storedinds(C)), length(storedvals(C))) + isempty(C) && return _checkbuffers(C) + _reset!(C) + spaceC::Int = _capacity(C) rowsentinelA = convert(indtype(A), numrows(C) + 1) rowsentinelB = convert(indtype(B), numrows(C) + 1) # C, A, and B cannot all have the same shape, as we directed that case to map in broadcast's @@ -601,8 +609,8 @@ function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat, B::Sp # chain above differs from that in the corresponding map code. if !_iszero(Cx) Ck > spaceC && (spaceC = expandstorage!(C, _unchecked_maxnnzbcres(size(C), A, B))) - storedinds(C)[Ck] = Ci - storedvals(C)[Ck] = Cx + push!(storedinds(C), Ci) + push!(storedvals(C), Cx) Ck += 1 end end @@ -619,8 +627,8 @@ function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat, B::Sp if !_iszero(Cx) for Ci::indtype(C) in 1:numrows(C) Ck > spaceC && (spaceC = expandstorage!(C, _unchecked_maxnnzbcres(size(C), A, B))) - storedinds(C)[Ck] = Ci - storedvals(C)[Ck] = Cx + push!(storedinds(C), Ci) + push!(storedvals(C), Cx) Ck += 1 end end @@ -640,8 +648,8 @@ function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat, B::Sp Cx = f(Ax, storedvals(B)[Bk]) if !_iszero(Cx) Ck > spaceC && (spaceC = expandstorage!(C, _unchecked_maxnnzbcres(size(C), A, B))) - storedinds(C)[Ck] = storedinds(B)[Bk] - storedvals(C)[Ck] = Cx + push!(storedinds(C), storedinds(B)[Bk]) + push!(storedvals(C), Cx) Ck += 1 end Bk += oneunit(Bk) @@ -659,8 +667,8 @@ function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat, B::Sp end if !_iszero(Cx) Ck > spaceC && (spaceC = expandstorage!(C, _unchecked_maxnnzbcres(size(C), A, B))) - storedinds(C)[Ck] = Ci - storedvals(C)[Ck] = Cx + push!(storedinds(C), Ci) + push!(storedvals(C), Cx) Ck += 1 end end @@ -681,8 +689,8 @@ function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat, B::Sp Cx = f(storedvals(A)[Ak], Bx) if !_iszero(Cx) Ck > spaceC && (spaceC = expandstorage!(C, _unchecked_maxnnzbcres(size(C), A, B))) - storedinds(C)[Ck] = storedinds(A)[Ak] - storedvals(C)[Ck] = Cx + push!(storedinds(C), storedinds(A)[Ak]) + push!(storedvals(C), Cx) Ck += 1 end Ak += oneunit(Ak) @@ -700,8 +708,8 @@ function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat, B::Sp end if !_iszero(Cx) Ck > spaceC && (spaceC = expandstorage!(C, _unchecked_maxnnzbcres(size(C), A, B))) - storedinds(C)[Ck] = Ci - storedvals(C)[Ck] = Cx + push!(storedinds(C), Ci) + push!(storedvals(C), Cx) Ck += 1 end end @@ -709,8 +717,7 @@ function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat, B::Sp end end @inbounds setcolptr!(C, numcols(C) + 1, Ck) - trimstorage!(C, Ck - 1) - return C + return _checkbuffers(C) end function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseVecOrMat, B::SparseVecOrMat) where Tf # For information on this code, see comments in similar code in _broadcast_zeropres! above @@ -809,10 +816,8 @@ function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseV end end end - return C + return _checkbuffers(C) end -_finishempty!(C::SparseVector) = C -_finishempty!(C::SparseMatrixCSC) = (fill!(C.colptr, 1); C) # special case - vector outer product _copy(f::typeof(*), x::SparseVectorUnion, y::AdjOrTransSparseVectorUnion) = _outer(x, y) @@ -853,14 +858,16 @@ function _outer(trans::Tf, x, y) where Tf end end cumsum!(colptrC, colptrC) - + resize!(rowvalC, idx) + resize!(nzvalsC, idx) return SparseMatrixCSC(nx, ny, colptrC, rowvalC, nzvalsC) end # (9) _broadcast_zeropres!/_broadcast_notzeropres! for more than two (input) sparse vectors/matrices function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, As::Vararg{SparseVecOrMat,N}) where {Tf,N} - isempty(C) && return _finishempty!(C) - spaceC::Int = min(length(storedinds(C)), length(storedvals(C))) + isempty(C) && return _checkbuffers(C) + _reset!(C) + spaceC::Int = _capacity(C) expandsverts = _expandsvert_all(C, As) expandshorzs = _expandshorz_all(C, As) rowsentinel = numrows(C) + 1 @@ -882,8 +889,8 @@ function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, As::Vararg{SparseVecOrMa Cx = f(args...) if !_iszero(Cx) Ck > spaceC && (spaceC = expandstorage!(C, _unchecked_maxnnzbcres(size(C), As))) - storedinds(C)[Ck] = activerow - storedvals(C)[Ck] = Cx + push!(storedinds(C), activerow) + push!(storedvals(C), Cx) Ck += 1 end activerow = min(rows...) @@ -899,19 +906,18 @@ function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, As::Vararg{SparseVecOrMa end if !_iszero(Cx) Ck > spaceC && (spaceC = expandstorage!(C, _unchecked_maxnnzbcres(size(C), As))) - storedinds(C)[Ck] = Ci - storedvals(C)[Ck] = Cx + push!(storedinds(C), Ci) + push!(storedvals(C), Cx) Ck += 1 end end end end @inbounds setcolptr!(C, numcols(C) + 1, Ck) - trimstorage!(C, Ck - 1) - return C + return _checkbuffers(C) end function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, As::Vararg{SparseVecOrMat,N}) where {Tf,N} - isempty(C) && return _finishempty!(C) + isempty(C) && return _checkbuffers(C) # Build dense matrix structure in C, expanding storage if necessary _densestructure!(C) # Populate values @@ -949,7 +955,7 @@ function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, As::Vararg end end end - return C + return _checkbuffers(C) end # helper method for broadcast/broadcast! methods just above diff --git a/stdlib/SparseArrays/src/sparsematrix.jl b/stdlib/SparseArrays/src/sparsematrix.jl index 6c8037bd6bfe2..9e6f816c79333 100644 --- a/stdlib/SparseArrays/src/sparsematrix.jl +++ b/stdlib/SparseArrays/src/sparsematrix.jl @@ -26,6 +26,7 @@ struct SparseMatrixCSC{Tv,Ti<:Integer} <: AbstractSparseMatrix{Tv,Ti} throw(ArgumentError("number of $str ($lbl) must be ≥ 0, got $k")) m < 0 && throwsz("rows", 'm', m) n < 0 && throwsz("columns", 'n', n) + _goodbuffers(Int(n), colptr, rowval, nzval) || throw(ArgumentError("Illegal buffers for SparseMatrixCSC construction $n $colptr $rowval $nzval")) new(Int(m), Int(n), colptr, rowval, nzval) end end @@ -35,6 +36,13 @@ function SparseMatrixCSC(m::Integer, n::Integer, colptr::Vector, rowval::Vector, SparseMatrixCSC{Tv,Ti}(m, n, colptr, rowval, nzval) end + + +_goodbuffers(n, colptr, rowval, nzval) = length(colptr) == n + 1 && colptr[end] - 1 == length(rowval) == length(nzval) +_goodbuffers(S::SparseMatrixCSC) = _goodbuffers(S.n, S.colptr, S.rowval, S.nzval) +_checkbuffers(S::SparseMatrixCSC) = (@assert _goodbuffers(S); S) + + size(S::SparseMatrixCSC) = (S.m, S.n) # Define an alias for views of a SparseMatrixCSC which include all rows and a unit range of the columns. @@ -152,6 +160,7 @@ nzrange(S::SparseMatrixCSC, col::Integer) = S.colptr[col]:(S.colptr[col+1]-1) nzrange(S::SparseMatrixCSCView, col::Integer) = nzrange(S.parent, S.indices[2][col]) function Base.show(io::IO, ::MIME"text/plain", S::SparseMatrixCSC) + _checkbuffers(S) xnnz = nnz(S) print(io, S.m, "×", S.n, " ", typeof(S), " with ", xnnz, " stored ", xnnz == 1 ? "entry" : "entries") @@ -163,6 +172,7 @@ end Base.show(io::IO, S::SparseMatrixCSC) = Base.show(convert(IOContext, io), S::SparseMatrixCSC) function Base.show(io::IOContext, S::SparseMatrixCSC) + _checkbuffers(S) if nnz(S) == 0 return show(io, MIME("text/plain"), S) end @@ -216,7 +226,7 @@ end ## Reshape -function sparse_compute_reshaped_colptr_and_rowval(colptrS::Vector{Ti}, rowvalS::Vector{Ti}, +function sparse_compute_reshaped_colptr_and_rowval!(colptrS::Vector{Ti}, rowvalS::Vector{Ti}, mS::Int, nS::Int, colptrA::Vector{Ti}, rowvalA::Vector{Ti}, mA::Int, nA::Int) where Ti lrowvalA = length(rowvalA) @@ -259,7 +269,7 @@ function copy(ra::ReshapedArray{<:Any,2,<:SparseMatrixCSC}) rowval = similar(a.rowval) nzval = copy(a.nzval) - sparse_compute_reshaped_colptr_and_rowval(colptr, rowval, mS, nS, a.colptr, a.rowval, mA, nA) + sparse_compute_reshaped_colptr_and_rowval!(colptr, rowval, mS, nS, a.colptr, a.rowval, mA, nA) return SparseMatrixCSC(mS, nS, colptr, rowval, nzval) end @@ -286,7 +296,7 @@ function copyto!(A::SparseMatrixCSC, B::SparseMatrixCSC) copyto!(A.rowval, B.rowval) else # This is like a "reshape B into A". - sparse_compute_reshaped_colptr_and_rowval(A.colptr, A.rowval, A.m, A.n, B.colptr, B.rowval, B.m, B.n) + sparse_compute_reshaped_colptr_and_rowval!(A.colptr, A.rowval, A.m, A.n, B.colptr, B.rowval, B.m, B.n) end else length(A) >= length(B) || throw(BoundsError()) @@ -316,10 +326,10 @@ function copyto!(A::SparseMatrixCSC, B::SparseMatrixCSC) @inbounds for i in 2:length(A.colptr) A.colptr[i] += nnzB - lastmodptrA end - sparse_compute_reshaped_colptr_and_rowval(A.colptr, A.rowval, A.m, lastmodcolA-1, B.colptr, B.rowval, B.m, B.n) + sparse_compute_reshaped_colptr_and_rowval!(A.colptr, A.rowval, A.m, lastmodcolA-1, B.colptr, B.rowval, B.m, B.n) end copyto!(A.nzval, B.nzval) - return A + return _checkbuffers(A) end ## similar @@ -431,6 +441,7 @@ SparseMatrixCSC{Tv,Ti}(M::Transpose{<:Any,SparseMatrixCSC}) where {Tv,Ti} = Spar # converting from SparseMatrixCSC to other matrix types function Matrix(S::SparseMatrixCSC{Tv}) where Tv # Handle cases where zero(Tv) is not defined but the array is dense. + _checkbuffers(S) A = length(S) == nnz(S) ? Matrix{Tv}(undef, S.m, S.n) : zeros(Tv, S.m, S.n) for Sj in 1:S.n for Sk in nzrange(S, Sj) @@ -800,6 +811,8 @@ respectively. Simultaneously fixes the one-position-forward shift in `X.colptr`. """ @noinline function _distributevals_halfperm!(X::SparseMatrixCSC{Tv,Ti}, A::SparseMatrixCSC{Tv,Ti}, q::AbstractVector{<:Integer}, f::Function) where {Tv,Ti} + resize!(X.rowval, nnz(A)) + resize!(X.nzval, nnz(A)) @inbounds for Xi in 1:A.n Aj = q[Xi] for Ak in nzrange(A, Aj) @@ -836,11 +849,7 @@ transpose!(X::SparseMatrixCSC{Tv,Ti}, A::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti} = adjoint!(X::SparseMatrixCSC{Tv,Ti}, A::SparseMatrixCSC{Tv,Ti}) where {Tv,Ti} = ftranspose!(X, A, conj) function ftranspose(A::SparseMatrixCSC{Tv,Ti}, f::Function) where {Tv,Ti} - X = SparseMatrixCSC(A.n, A.m, - Vector{Ti}(undef, A.m+1), - Vector{Ti}(undef, nnz(A)), - Vector{Tv}(undef, nnz(A))) - halfperm!(X, A, 1:A.n, f) + halfperm!(sizehint!(spzeros(Tv, Ti, A.n, A.m), nnz(A)), A, 1:A.n, f) end adjoint(A::SparseMatrixCSC) = Adjoint(A) transpose(A::SparseMatrixCSC) = Transpose(A) @@ -1055,10 +1064,7 @@ function permute!(X::SparseMatrixCSC{Tv,Ti}, A::SparseMatrixCSC{Tv,Ti}, p::AbstractVector{<:Integer}, q::AbstractVector{<:Integer}) where {Tv,Ti} _checkargs_sourcecompatdest_permute!(A, X) _checkargs_sourcecompatperms_permute!(A, p, q) - C = SparseMatrixCSC(A.n, A.m, - Vector{Ti}(undef, A.m + 1), - Vector{Ti}(undef, nnz(A)), - Vector{Tv}(undef, nnz(A))) + C = sizehint!(spzeros(Tv,Ti,A.n,A.m), nnz(A)) _checkargs_permutationsvalid_permute!(p, C.colptr, q, X.colptr) unchecked_noalias_permute!(X, A, p, q, C) end @@ -1074,10 +1080,7 @@ end function permute!(A::SparseMatrixCSC{Tv,Ti}, p::AbstractVector{<:Integer}, q::AbstractVector{<:Integer}) where {Tv,Ti} _checkargs_sourcecompatperms_permute!(A, p, q) - C = SparseMatrixCSC(A.n, A.m, - Vector{Ti}(undef, A.m + 1), - Vector{Ti}(undef, nnz(A)), - Vector{Tv}(undef, nnz(A))) + C = sizehint!(spzeros(Tv,Ti, A.n, A.m), nnz(A)) workcolptr = Vector{Ti}(undef, A.n + 1) _checkargs_permutationsvalid_permute!(p, C.colptr, q, workcolptr) unchecked_aliasing_permute!(A, p, q, C, workcolptr) @@ -1145,14 +1148,8 @@ julia> permute(A, [1, 2, 3, 4], [4, 3, 2, 1]) function permute(A::SparseMatrixCSC{Tv,Ti}, p::AbstractVector{<:Integer}, q::AbstractVector{<:Integer}) where {Tv,Ti} _checkargs_sourcecompatperms_permute!(A, p, q) - X = SparseMatrixCSC(A.m, A.n, - Vector{Ti}(undef, A.n + 1), - Vector{Ti}(undef, nnz(A)), - Vector{Tv}(undef, nnz(A))) - C = SparseMatrixCSC(A.n, A.m, - Vector{Ti}(undef, A.m + 1), - Vector{Ti}(undef, nnz(A)), - Vector{Tv}(undef, nnz(A))) + X = sizehint!(spzeros(Tv, Ti, A.m, A.n), nnz(A)) + C = sizehint!(spzeros(Tv, Ti, A.n, A.m), nnz(A)) _checkargs_permutationsvalid_permute!(p, C.colptr, q, X.colptr) unchecked_noalias_permute!(X, A, p, q, C) end @@ -1246,6 +1243,7 @@ For an out-of-place version, see [`dropzeros`](@ref). For algorithmic information, see `fkeep!`. """ dropzeros!(A::SparseMatrixCSC; trim::Bool = true) = fkeep!(A, (i, j, x) -> x != 0, trim) + """ dropzeros(A::SparseMatrixCSC; trim::Bool = true) @@ -2214,7 +2212,7 @@ function permute_rows!(S::SparseMatrixCSC{Tv,Ti}, pI::Vector{Int}) where {Tv,Ti} k += 1 end end - S + _checkbuffers(S) end function getindex_general(A::SparseMatrixCSC, I::AbstractVector, J::AbstractVector) @@ -2351,6 +2349,7 @@ function Base.fill!(V::SubArray{Tv, <:Any, <:SparseMatrixCSC, Tuple{Vararg{Union else _spsetnz_setindex!(A, convert(Tv, x), I, J) end + _checkbuffers(A) end """ Helper method for immediately preceding setindex! method. For all (i,j) such that i in I and @@ -2630,7 +2629,7 @@ function setindex!(A::SparseMatrixCSC{Tv,Ti}, V::AbstractVecOrMat, Ix::Union{Int deleteat!(rowvalA, colptrA[end]:length(rowvalA)) deleteat!(nzvalA, colptrA[end]:length(nzvalA)) - return A + return _checkbuffers(A) end # Logical setindex! @@ -2739,7 +2738,7 @@ function setindex!(A::SparseMatrixCSC, x::AbstractArray, I::AbstractMatrix{Bool} deleteat!(rowvalB, bidx:n) end end - A + _checkbuffers(A) end function setindex!(A::SparseMatrixCSC, x::AbstractArray, Ix::AbstractVector{<:Integer}) @@ -2850,7 +2849,7 @@ function setindex!(A::SparseMatrixCSC, x::AbstractArray, Ix::AbstractVector{<:In deleteat!(rowvalB, bidx:n) end end - A + _checkbuffers(A) end ## dropstored! methods @@ -2885,7 +2884,7 @@ function dropstored!(A::SparseMatrixCSC, i::Integer, j::Integer) @inbounds A.colptr[m] -= 1 end end - return A + return _checkbuffers(A) end """ dropstored!(A::SparseMatrixCSC, I::AbstractVector{<:Integer}, J::AbstractVector{<:Integer}) @@ -2971,7 +2970,7 @@ function dropstored!(A::SparseMatrixCSC, deleteat!(rowvalA, rowidx:nnzA) deleteat!(nzvalA, rowidx:nnzA) end - return A + return _checkbuffers(A) end dropstored!(A::SparseMatrixCSC, i::Integer, J::AbstractVector{<:Integer}) = dropstored!(A, [i], J) dropstored!(A::SparseMatrixCSC, I::AbstractVector{<:Integer}, j::Integer) = dropstored!(A, I, [j]) @@ -3485,7 +3484,7 @@ function circshift!(O::SparseMatrixCSC, X::SparseMatrixCSC, (r,c)::Base.DimsInte @inbounds for i=1:O.n subvector_shifter!(O.rowval, O.nzval, O.colptr[i], O.colptr[i+1]-1, O.m, r) end - return O + return _checkbuffers(O) end circshift!(O::SparseMatrixCSC, X::SparseMatrixCSC, (r,)::Base.DimsInteger{1}) = circshift!(O, X, (r,0)) diff --git a/stdlib/SparseArrays/src/sparsevector.jl b/stdlib/SparseArrays/src/sparsevector.jl index 8487cd5a4aab5..f5b1627430599 100644 --- a/stdlib/SparseArrays/src/sparsevector.jl +++ b/stdlib/SparseArrays/src/sparsevector.jl @@ -65,6 +65,13 @@ function nnz(x::SparseColumnView) return length(nzrange(parent(x), colidx)) end + +function Base.sizehint!(v::SparseVector, newlen::Integer) + sizehint!(v.nzind, newlen) + sizehint!(v.nzval, newlen) + return v +end + ## similar # # parent method for similar that preserves stored-entry structure (for when new and old dims match) diff --git a/stdlib/SparseArrays/test/sparse.jl b/stdlib/SparseArrays/test/sparse.jl index df28e1aa0d6df..11e2f8b788630 100644 --- a/stdlib/SparseArrays/test/sparse.jl +++ b/stdlib/SparseArrays/test/sparse.jl @@ -25,7 +25,7 @@ end S = sparse(I, 3, 3) fill!(S, 0) @test iszero(S) # test success with stored zeros via fill! - @test iszero(SparseMatrixCSC(2, 2, [1,2,3], [1,2], [0,0,1])) # test success with nonzeros beyond data range + @test_throws ArgumentError iszero(SparseMatrixCSC(2, 2, [1,2,3], [1,2], [0,0,1])) # test failure with nonzeros beyond data range end @testset "isone specialization for SparseMatrixCSC" begin @test isone(sparse(I, 3, 3)) # test success @@ -2117,19 +2117,6 @@ end @test repr("text/plain", sparse([true true])) == "1×2 SparseArrays.SparseMatrixCSC{Bool,$Int} with 2 stored entries:\n [1, 1] = 1\n [1, 2] = 1" end -@testset "check buffers" for n in 1:3 - local A - rowval = [1,2,3] - nzval1 = Int[] - nzval2 = [1,1,1] - A = SparseMatrixCSC(n, n, [1:n+1;], rowval, nzval1) - @test nnz(A) == n - @test_throws BoundsError A[n,n] - A = SparseMatrixCSC(n, n, [1:n+1;], rowval, nzval2) - @test nnz(A) == n - @test A == Matrix(I, n, n) -end - @testset "reverse search direction if step < 0 #21986" begin local A, B A = guardseed(1234) do @@ -2231,8 +2218,6 @@ end # count should throw for sparse arrays for which zero(eltype) does not exist @test_throws MethodError count(SparseMatrixCSC(2, 2, Int[1, 2, 3], Int[1, 2], Any[true, true])) @test_throws MethodError count(SparseVector(2, Int[1], Any[true])) - # count should run only over S.nzval[1:nnz(S)], not S.nzval in full - @test count(SparseMatrixCSC(2, 2, Int[1, 2, 3], Int[1, 2], Bool[true, true, true])) == 2 end @testset "sparse findprev/findnext operations" begin @@ -2280,15 +2265,6 @@ end @test sum(s, dims=2) == reshape([1, 2, 3], 3, 1) end -@testset "mapreduce of sparse matrices with trailing elements in nzval #26534" begin - B = SparseMatrixCSC{Int,Int}(2, 3, - [1, 3, 4, 5], - [1, 2, 1, 2, 999, 999, 999, 999], - [1, 2, 3, 6, 999, 999, 999, 999] - ) - @test maximum(B) == 6 -end - _length_or_count_or_five(::Colon) = 5 _length_or_count_or_five(x::AbstractVector{Bool}) = count(x) _length_or_count_or_five(x) = length(x)