sparse findnext findprev hash performance improved (#31354)

* sparse findnext findprev hash performance improved

* added tests and minor changes

Author: KlausC

stdlib/SparseArrays/src/sparsematrix.jl

@@ -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])
 

stdlib/SparseArrays/test/sparse.jl

@@ -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))
@@ -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