Add a sparse keyword to reductions along a dimension and reduce column-range views through the sparse kernels

docs/src/index.md

@@ -186,6 +186,12 @@ In many cases it may be better to convert the sparse matrix into `(I,J,V)` forma
 manipulate the values or the structure in the dense vectors `(I,J,V)`, and then reconstruct
 the sparse matrix.
 
+Reductions along a dimension, such as `sum(S; dims = 2)`, return a dense `Matrix`, as for dense
+input. To keep the result sparse, pass `sparse = true`: `sum(S; dims = 2, sparse = true)` stores an
+entry only for the rows of `S` that store one, at a cost proportional to the number of stored
+entries rather than to the number of rows. `prod`, `maximum`, `minimum`, `count`, `any`, `all` and
+`mapreduce` accept the keyword in the same way.
+
 ## Correspondence of dense and sparse methods
 
 The following table gives a correspondence between built-in methods on sparse matrices and their

src/sparsematrix.jl

@@ -199,6 +199,11 @@ getrowval(S::SparseMatrixCSCColumnSubset) = rowvals(parent(S))
 getnzval( S::AbstractSparseMatrixCSC) = nonzeros(S)
 getnzval( S::SparseMatrixCSCColumnSubset) = nonzeros(parent(S))
 nzvalview(S::AbstractSparseMatrixCSC) = view(nonzeros(S), 1:nnz(S))
+nzvalview(S::SparseMatrixCSCView) = view(nonzeros(S), _storedrange(S))
+# the stored entries of a column-range view are a contiguous range of the parent's
+_storedrange(S::AbstractSparseMatrixCSC) = 1:nnz(S)
+_storedrange(S::SparseMatrixCSCView) = (colptr = getcolptr(S); Int(colptr[1]):Int(colptr[end]) - 1)
+widelength(S::SparseMatrixCSCView) = prod(Int64.(size(S)))
 
 """
     nnz(A)
@@ -223,9 +228,10 @@ nnz(S::AdjOrTrans{<:Any,<:AbstractSparseMatrixCSC}) = nnz(parent(S))
 nnz(S::UpperTriangular{<:Any,<:AbstractSparseMatrixCSC}) = nnz1(S)
 nnz(S::LowerTriangular{<:Any,<:AbstractSparseMatrixCSC}) = nnz1(S)
 nnz(S::SparseMatrixCSCColumnSubset) = nnz1(S)
+nnz(S::SparseMatrixCSCView) = length(_storedrange(S))
 nnz1(S) = @inbounds sum(length.(nzrange.(Ref(S), axes(S, 2))))
 
-function Base._simple_count(pred, S::AbstractSparseMatrixCSC, init::T) where T
+function Base._simple_count(pred, S::SparseMatrixCSCUnion, init::T) where T
     init + T(count(pred, nzvalview(S)) + pred(zero(eltype(S)))*(prod(size(S)) - nnz(S)))
 end
 
@@ -2495,10 +2501,11 @@ Base.isequal(A::Transpose{<:Any,<:AbstractSparseMatrixCSCInclAdjointAndTranspose
 
 ## Reductions
 
-# In general, output of sparse matrix reductions will not be sparse,
-# and computing reductions along columns into SparseMatrixCSC is
-# non-trivial, so use Arrays for output. Array element type is given by `R`.
-function Base.reducedim_initarray(A::AbstractSparseMatrixCSC, region, v0, ::Type{R}) where {R}
+# Reductions along a dimension return a dense `Array`, as for dense input. A sparse result
+# (issue #43) is opt-in with the keyword `sparse = true`, see `_mapreduce_dim_sparse` below.
+# Covers column-range views so they reduce like their copy (#377), where Base's `similar`
+# would otherwise give a sparse result.
+function Base.reducedim_initarray(A::SparseMatrixCSCUnion, region, v0, ::Type{R}) where {R}
     fill!(Array{R}(undef, Base.to_shape(Base.reduced_indices(A, region))), v0)
 end
 
@@ -2522,7 +2529,7 @@ function _mapreducezeros(f::F, op::G, ::Type{T}, nzeros::Integer, v0) where {F,G
     v
 end
 
-function Base._mapreduce(f::F, op::G, ::Base.IndexCartesian, A::AbstractSparseMatrixCSC{T}) where {F,G,T}
+function Base._mapreduce(f::F, op::G, ::Base.IndexCartesian, A::SparseMatrixCSCUnion{T}) where {F,G,T}
     z = nnz(A)
     n = widelength(A)
     if z == 0
@@ -2547,12 +2554,12 @@ _mapreducezeros(f::Base.ExtremaMap, op::typeof(Base._extrema_rf), ::Type{T}, nze
     nzeros == 0 ? v0 : op(v0, f(zero(T)))
 
 # Specialized mapreduce for any and all
-Base._any(f, A::AbstractSparseMatrixCSC, ::Colon) =
+Base._any(f, A::SparseMatrixCSCUnion, ::Colon) =
     iszero(widelength(A)) ? false : Base._mapreduce(f, |, IndexCartesian(), A)
-Base._all(f, A::AbstractSparseMatrixCSC, ::Colon) =
+Base._all(f, A::SparseMatrixCSCUnion, ::Colon) =
     iszero(widelength(A)) ? true  : Base._mapreduce(f, &, IndexCartesian(), A)
 
-function Base._mapreduce(f::F, op::Union{typeof(Base.mul_prod),typeof(*)}, ::Base.IndexCartesian, A::AbstractSparseMatrixCSC{T}) where {F,T}
+function Base._mapreduce(f::F, op::Union{typeof(Base.mul_prod),typeof(*)}, ::Base.IndexCartesian, A::SparseMatrixCSCUnion{T}) where {F,T}
     nnzA = nnz(A)
     nzeros = widelength(A) - nnzA
     if nzeros == 0
@@ -2565,8 +2572,166 @@ function Base._mapreduce(f::F, op::Union{typeof(Base.mul_prod),typeof(*)}, ::Bas
     end
 end
 
+# The `sparse` keyword, the opt-in for a sparse result (issue #43): `sum(A; dims = 2, sparse = true)`
+# and the other reductions along a dimension return a `SparseMatrixCSC` that stores an entry
+# only for the rows or columns of `A` that store one (all of them when a slice that stores
+# nothing reduces to something nonzero, as for `f(0) != 0`), in time proportional to
+# nnz(A) + length(result) rather than to length(A). Base's `sum`, `prod`, `maximum` and
+# `minimum` forward unknown keywords to `mapreduce`, so this method serves them all; `any`,
+# `all` and `count` do not and get their own methods below.
+function Base.mapreduce(f, op, A::SparseMatrixCSCUnion; dims=:, init=Base._InitialValue(), sparse::Bool=false)
+    sparse || return Base._mapreduce_dim(f, op, init, A, dims)
+    dims === (:) && throw(ArgumentError("a sparse result needs a reduction along a dimension, pass `dims`"))
+    return _mapreduce_dim_sparse(f, op, init, A, dims)
+end
+for (fname, _fname, op) in ((:any, :_any, :(Base.or_any)), (:all, :_all, :(Base.and_all)))
+    @eval begin
+        Base.$fname(A::SparseMatrixCSCUnion; dims=:, sparse::Bool=false) = Base.$fname(identity, A; dims, sparse)
+        Base.$fname(f, A::SparseMatrixCSCUnion; dims=:, sparse::Bool=false) =
+            sparse ? mapreduce(f, $op, A; dims, sparse) : Base.$_fname(f, A, dims)
+    end
+end
+Base.count(A::SparseMatrixCSCUnion; dims=:, init=0, sparse::Bool=false) = count(identity, A; dims, init, sparse)
+Base.count(f, A::SparseMatrixCSCUnion; dims=:, init=0, sparse::Bool=false) =
+    sparse ? mapreduce(Base._bool(f), Base.add_sum, A; dims, init, sparse) : Base._count(f, A, dims, init)
+
+# The slices are reduced the way Base's dense result is: seeded with `init` when given and
+# otherwise with `mapreduce_first` of their first element, with the entries a slice does not
+# store folded in through `_mapreducezeros`.
+_seed(f, op, ::Base._InitialValue, x) = Base.mapreduce_first(f, op, x)
+_seed(f, op, init, x) = op(init, f(x))
+# element type of the dense result, taken from Base's initialization of a 1 x 1 stand-in,
+# which evaluates f(0) as Base does
+_reduced_eltype(f, op, ::Base._InitialValue, ::Type{T}) where T =
+    eltype(Base.reducedim_init(f, op, fill(zero(T), 1, 1), 1))
+_reduced_eltype(f, op, init, ::Type{T}) where T = typeof(init)
+# the reduction of a slice with no entries at all, as Base initializes the dense result
+_reduced_empty(f, op, ::Base._InitialValue, ::Type{T}) where T =
+    Base.reducedim_init(f, op, Matrix{T}(undef, 0, 1), 1)[1]
+_reduced_empty(f, op, init, ::Type{T}) where T = init
+# the reduction of a slice of length `len` that stores nothing
+_reduce_unstored(f, op, init, ::Type{T}, len) where T =
+    len == 0 ? _reduced_empty(f, op, init, T) : _mapreducezeros(f, op, T, len - 1, _seed(f, op, init, zero(T)))
+
+function _mapreduce_dim_sparse(f, op, init, A::SparseMatrixCSCUnion{T,Ti}, dims) where {T,Ti}
+    m, n = size(A)
+    rm, rn = map(length, Base.reduced_indices(A, dims))   # also validates `dims`
+    Tr = _reduced_eltype(f, op, init, T)
+    applicable(zero, Tr) || throw(ArgumentError("cannot store a sparse result of element type $Tr, " *
+        "which has no zero (as for `extrema`); reduce without `sparse = true`"))
+    if rm == rn == 1
+        R = spzeros(Tr, Ti, 1, 1)
+        v = isempty(A) ? _reduced_empty(f, op, init, T) : _seed(identity, op, init, mapreduce(f, op, A))
+        if nnz(A) > 0 || !isequal(v, zero(Tr))
+            push!(rowvals(R), 1)
+            push!(nonzeros(R), v)
+            getcolptr(R)[2] = 2
+        end
+        return R
+    elseif rm == 1
+        return _mapreducerows_sparse!(f, op, init, spzeros(Tr, Ti, 1, n), A)
+    elseif rn == 1
+        return _mapreducecols_sparse!(f, op, init, spzeros(Tr, Ti, m, 1), A)
+    else
+        # a dimension beyond 2: every entry is a slice of its own
+        return convert(SparseMatrixCSC{Tr,Ti}, map(x -> _seed(f, op, init, x), A))
+    end
+end
+
+# `R` is a structurally empty `1 x n` sparse matrix: its columns are built in order
+function _mapreducerows_sparse!(f, op, init, R::SparseMatrixCSC, A::SparseMatrixCSCUnion{T}) where T
+    colptr = getcolptr(A)
+    nzval = nonzeros(A)
+    m, n = size(A)
+    z = zero(eltype(R))
+    # the reduction of a column that stores nothing; when every column is full it is not
+    # needed and f(0) is not evaluated, since it might throw
+    zunstored = nnz(A) == m*n && m > 0 ? z : _reduce_unstored(f, op, init, T, m)
+    store_unstored = !isequal(zunstored, z)
+    Rcolptr, Rrowval, Rnzval = getcolptr(R), rowvals(R), nonzeros(R)
+    nstored = store_unstored ? n : count(col -> colptr[col+1] > colptr[col], 1:n)
+    resize!(Rrowval, nstored)
+    fill!(Rrowval, 1)
+    resize!(Rnzval, nstored)
+    k = 0
+    @inbounds for col in 1:n
+        rng = colptr[col]:colptr[col+1]-1
+        if isempty(rng)
+            store_unstored || (Rcolptr[col+1] = k + 1; continue)
+            v = zunstored
+        else
+            r = _seed(f, op, init, nzval[first(rng)])
+            @simd for j in first(rng)+1:last(rng)
+                r = op(r, f(nzval[j]))
+            end
+            v = _mapreducezeros(f, op, T, m - length(rng), r)
+        end
+        k += 1
+        Rnzval[k] = v
+        Rcolptr[col+1] = k + 1
+    end
+    return R
+end
+
+# `R` is a structurally empty `m x 1` sparse matrix. With enough stored entries the rows are
+# reduced in a dense workspace that is then compressed into `R`; a hypersparse `A` instead has
+# its stored entries sorted by row so that only the rows storing something are ever visited.
+function _mapreducecols_sparse!(f, op, init, R::SparseMatrixCSC, A::SparseMatrixCSCUnion{T}) where T
+    m, n = size(A)
+    Tr = eltype(R)
+    z = zero(Tr)
+    rows = view(rowvals(A), _storedrange(A))
+    vals = view(nonzeros(A), _storedrange(A))
+    nz = length(rows)
+    zunstored = nz == m*n && n > 0 ? z : _reduce_unstored(f, op, init, T, n)
+    store_unstored = !isequal(zunstored, z)
+    Rcolptr, Rrowval, Rnzval = getcolptr(R), rowvals(R), nonzeros(R)
+    if store_unstored || 8 * nz >= m
+        W = Vector{Tr}(undef, m)
+        cnt = zeros(Int, m)   # stored entries seen in each row
+        @inbounds for j in eachindex(rows, vals)
+            i = rows[j]
+            W[i] = cnt[i] == 0 ? _seed(f, op, init, vals[j]) : op(W[i], f(vals[j]))
+            cnt[i] += 1
+        end
+        resize!(Rrowval, m)
+        resize!(Rnzval, m)
+        k = 0
+        @inbounds for i in 1:m
+            if cnt[i] > 0
+                k += 1
+                Rrowval[k] = i
+                Rnzval[k] = _mapreducezeros(f, op, T, n - cnt[i], W[i])
+            elseif store_unstored
+                k += 1
+                Rrowval[k] = i
+                Rnzval[k] = zunstored
+            end
+        end
+        resize!(Rrowval, k)
+        resize!(Rnzval, k)
+    else
+        perm = sortperm(rows; alg=Base.Sort.DEFAULT_STABLE)   # keeps each row's entries in column order
+        s = 1
+        @inbounds while s <= nz
+            row = rows[perm[s]]
+            r = _seed(f, op, init, vals[perm[s]])
+            t = s + 1
+            while t <= nz && rows[perm[t]] == row
+                r = op(r, f(vals[perm[t]]))
+                t += 1
+            end
+            push!(Rrowval, row)
+            push!(Rnzval, _mapreducezeros(f, op, T, n - (t - s), r))
+            s = t
+        end
+    end
+    Rcolptr[2] = length(Rnzval) + 1
+    return R
+end
+
 # General mapreducedim
-function _mapreducerows!(f, op, R::AbstractArray, A::AbstractSparseMatrixCSC{T}) where T
+function _mapreducerows!(f, op, R::AbstractArray, A::SparseMatrixCSCUnion{T}) where T
     require_one_based_indexing(A, R)
     colptr = getcolptr(A)
     rowval = rowvals(A)
@@ -2582,7 +2747,7 @@ function _mapreducerows!(f, op, R::AbstractArray, A::AbstractSparseMatrixCSC{T})
     R
 end
 
-function _mapreducecols!(f, op, R::AbstractArray, A::AbstractSparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
+function _mapreducecols!(f, op, R::AbstractArray, A::SparseMatrixCSCUnion{Tv,Ti}) where {Tv,Ti}
     require_one_based_indexing(A, R)
     colptr = getcolptr(A)
     rowval = rowvals(A)
@@ -2602,7 +2767,7 @@ function _mapreducecols!(f, op, R::AbstractArray, A::AbstractSparseMatrixCSC{Tv,
     R
 end
 
-function Base._mapreducedim!(f::F, op::G, R::AbstractArray, A::AbstractSparseMatrixCSC{T}) where {F,G,T}
+function Base._mapreducedim!(f::F, op::G, R::AbstractArray, A::SparseMatrixCSCUnion{T}) where {F,G,T}
     require_one_based_indexing(A, R)
     lsiz = Base.check_reducedims(R,A)
     isempty(A) && return R
@@ -2620,17 +2785,17 @@ function Base._mapreducedim!(f::F, op::G, R::AbstractArray, A::AbstractSparseMat
         # Reduction along a dimension > 2
         # Compute op(R, f(A))
         m, n = size(A)
+        colptr = getcolptr(A)
+        rowval = rowvals(A)
         nzval = nonzeros(A)
-        if length(nzval) == m*n
+        if nnz(A) == m*n
             # No zeros, so don't compute f(0) since it might throw
-            for col in axes(A,2)
-                @simd for row in axes(A,1)
-                    @inbounds R[row, col] = op(R[row, col], f(nzval[(col-1)*m+row]))
+            @inbounds for col in axes(A,2)
+                @simd for j = colptr[col]:colptr[col+1]-1
+                    R[rowval[j], col] = op(R[rowval[j], col], f(nzval[j]))
                 end
             end
         else
-            colptr = getcolptr(A)
-            rowval = rowvals(A)
             zeroval = f(zero(T))
             @inbounds for col in axes(A,2)
                 lastrow = 0
@@ -2653,20 +2818,20 @@ end
 
 # Specialized mapreducedim for + cols to avoid allocating a
 # temporary array when f(0) == 0
-function _mapreducecols!(f, op::typeof(+), R::AbstractArray, A::AbstractSparseMatrixCSC{Tv,Ti}) where {Tv,Ti}
+function _mapreducecols!(f, op::typeof(+), R::AbstractArray, A::SparseMatrixCSCUnion{Tv,Ti}) where {Tv,Ti}
     require_one_based_indexing(A, R)
+    colptr = getcolptr(A)
+    rowval = rowvals(A)
     nzval = nonzeros(A)
     m, n = size(A)
-    if length(nzval) == m*n
+    if nnz(A) == m*n
         # No zeros, so don't compute f(0) since it might throw
-        for col in axes(A,2)
-            @simd for row in axes(A,1)
-                @inbounds R[row, 1] = op(R[row, 1], f(nzval[(col-1)*m+row]))
+        @inbounds for col in axes(A,2)
+            @simd for j = colptr[col]:colptr[col+1]-1
+                R[rowval[j], 1] = op(R[rowval[j], 1], f(nzval[j]))
             end
         end
     else
-        colptr = getcolptr(A)
-        rowval = rowvals(A)
         zeroval = f(zero(Tv))
         if isequal(zeroval, zero(Tv))
             # Case where f(0) == 0
@@ -2695,7 +2860,7 @@ end
 
 # any(pred, A, dims = 1) => mapreduce(pred, |, A, dims = 1)
 # final argument `post` is to allow post-mapping each columnar mapreduce
-function _mapreducerows!(pred::P, ::typeof(|), R::AbstractMatrix{Bool}, A::AbstractSparseMatrixCSC{Tv},
+function _mapreducerows!(pred::P, ::typeof(|), R::AbstractMatrix{Bool}, A::SparseMatrixCSCUnion{Tv},
                          post::F = identity) where {P, F, Tv}
     nzval = nonzeros(A)
     colptr = getcolptr(A)
@@ -2725,7 +2890,7 @@ function _mapreducerows!(pred::P, ::typeof(|), R::AbstractMatrix{Bool}, A::Abstr
 end
 # all(pred, A, dims = 1) => mapreduce(pred, &, A, dims = 1) == .!mapreduce(!pred, |, A, dims = 1)
 _mapreducerows!(pred::P, ::typeof(&), R::AbstractMatrix{Bool},
-                A::AbstractSparseMatrixCSC) where {P} = _mapreducerows!(!pred, |, R, A, !)
+                A::SparseMatrixCSCUnion) where {P} = _mapreducerows!(!pred, |, R, A, !)
 
 # findmax/min and argmax/min methods
 # find first zero value in sparse matrix - return linear index in full matrix