Skip to content

Commit

Permalink
sparse findnext findprev hash performance improved (#31354)
Browse files Browse the repository at this point in the history
* sparse findnext findprev hash performance improved

* added tests and minor changes
  • Loading branch information
KlausC authored and ViralBShah committed Apr 3, 2019
1 parent aee211b commit e0bef65
Show file tree
Hide file tree
Showing 2 changed files with 90 additions and 4 deletions.
71 changes: 71 additions & 0 deletions stdlib/SparseArrays/src/sparsematrix.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1364,6 +1364,77 @@ function sparse_sortedlinearindices!(I::Vector{Ti}, V::Vector, m::Int, n::Int) w
return SparseMatrixCSC(m, n, colptr, I, V)
end

# findfirst/next/prev/last
function _idxfirstnz(A::SparseMatrixCSC, ij::CartesianIndex{2})
nzr = nzrange(A, ij[2])
searchk = searchsortedfirst(A.rowval, ij[1], first(nzr), last(nzr), Forward)
return _idxnextnz(A, searchk)
end

function _idxlastnz(A::SparseMatrixCSC, ij::CartesianIndex{2})
nzr = nzrange(A, ij[2])
searchk = searchsortedlast(A.rowval, ij[1], first(nzr), last(nzr), Forward)
return _idxprevnz(A, searchk)
end

function _idxnextnz(A::SparseMatrixCSC, idx::Integer)
nnza = nnz(A)
nzval = nonzeros(A)
z = zero(eltype(A))
while idx <= nnza
nzv = nzval[idx]
!isequal(nzv, z) && return idx, nzv
idx += 1
end
return zero(idx), z
end

function _idxprevnz(A::SparseMatrixCSC, idx::Integer)
nzval = nonzeros(A)
z = zero(eltype(A))
while idx > 0
nzv = nzval[idx]
!isequal(nzv, z) && return idx, nzv
idx -= 1
end
return zero(idx), z
end

function _idx_to_cartesian(A::SparseMatrixCSC, idx::Integer)
rowval = rowvals(A)
i = rowval[idx]
j = searchsortedlast(A.colptr, idx, 1, size(A, 2), Base.Order.Forward)
return CartesianIndex(i, j)
end

function Base.findnext(pred::Function, A::SparseMatrixCSC, ij::CartesianIndex{2})
if nnz(A) == length(A) || pred(zero(eltype(A)))
return invoke(findnext, Tuple{Function,Any,Any}, pred, A, ij)
end
idx, nzv = _idxfirstnz(A, ij)
while idx > 0
if pred(nzv)
return _idx_to_cartesian(A, idx)
end
idx, nzv = _idxnextnz(A, idx + 1)
end
return nothing
end

function Base.findprev(pred::Function, A::SparseMatrixCSC, ij::CartesianIndex{2})
if nnz(A) == length(A) || pred(zero(eltype(A)))
return invoke(findprev, Tuple{Function,Any,Any}, pred, A, ij)
end
idx, nzv = _idxlastnz(A, ij)
while idx > 0
if pred(nzv)
return _idx_to_cartesian(A, idx)
end
idx, nzv = _idxprevnz(A, idx - 1)
end
return nothing
end

"""
sprand([rng],[type],m,[n],p::AbstractFloat,[rfn])
Expand Down
23 changes: 19 additions & 4 deletions stdlib/SparseArrays/test/sparse.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2245,16 +2245,20 @@ end
@test findprev(!iszero, x,i) == findprev(!iszero, x_sp,i)
end

y = [0 0 0 0 0;
y = [7 0 0 0 0;
1 0 1 0 0;
1 0 0 0 1;
1 7 0 7 1;
0 0 1 0 0;
1 0 1 1 0]
y_sp = sparse(y)
1 0 1 1 0.0]
y_sp = [x == 7 ? -0.0 : x for x in sparse(y)]
y = Array(y_sp)
@test isequal(y_sp[1,1], -0.0)

for i in keys(y)
@test findnext(!iszero, y,i) == findnext(!iszero, y_sp,i)
@test findprev(!iszero, y,i) == findprev(!iszero, y_sp,i)
@test findnext(iszero, y,i) == findnext(iszero, y_sp,i)
@test findprev(iszero, y,i) == findprev(iszero, y_sp,i)
end

z_sp = sparsevec(Dict(1=>1, 5=>1, 8=>0, 10=>1))
Expand All @@ -2264,6 +2268,17 @@ end
@test findnext(!iszero, z,i) == findnext(!iszero, z_sp,i)
@test findprev(!iszero, z,i) == findprev(!iszero, z_sp,i)
end

w = [ "a" ""; "" "b"]
w_sp = sparse(w)

for i in keys(w)
@test findnext(!isequal(""), w,i) == findnext(!isequal(""), w_sp,i)
@test findprev(!isequal(""), w,i) == findprev(!isequal(""), w_sp,i)
@test findnext(isequal(""), w,i) == findnext(isequal(""), w_sp,i)
@test findprev(isequal(""), w,i) == findprev(isequal(""), w_sp,i)
end

end

# #20711
Expand Down

0 comments on commit e0bef65

Please sign in to comment.