diff --git a/src/reductions.jl b/src/reductions.jl index 3930d243..b29b420d 100644 --- a/src/reductions.jl +++ b/src/reductions.jl @@ -590,37 +590,35 @@ function _mapreducecols!(f, op::typeof(+), R::AbstractArray, A::SparseMatrixCSCO end # any(pred, A, dims = 1) => mapreduce(pred, |, A, dims = 1) -# final argument `post` is to allow post-mapping each columnar mapreduce +# `R` holds the initial value of each column, which the column's reduction is folded into. +# With `negate`, `pred` is the negated predicate of an `&` reduction, see below. function _mapreducerows!(pred::P, ::typeof(|), R::AbstractMatrix{Bool}, A::SparseMatrixCSCOrColumnSubset{Tv}, - post::F = identity) where {P, F, Tv} + negate::Bool = false) where {P, Tv} nzval = getnzval(A) m, n = size(A) @inbounds for ii in axes(A,2) + # an initial value that already decides the column + R[1, ii] == negate || continue rng = getnzrange(A, ii) len = length(rng) - # An empty column is trivial - if len == 0 - R[1, ii] = post(pred(zero(Tv))) - continue - end - # If predicate on zero is true, then sparse column can be short-circuited - if pred(zero(Tv)) && len < m - R[1, ii] = post(true) - continue - end - # Otherwise reduce over the stored values - r = false - for jj in rng - r = pred(nzval[jj]) - r && break + if len < m && pred(zero(Tv)) + # an unstored entry satisfies the predicate + r = true + else + # reduce over the stored values + r = false + for jj in rng + r = pred(nzval[jj]) + r && break + end end - R[1, ii] = post(r) + R[1, ii] = r != negate end return R end # all(pred, A, dims = 1) => mapreduce(pred, &, A, dims = 1) == .!mapreduce(!pred, |, A, dims = 1) _mapreducerows!(pred::P, ::typeof(&), R::AbstractMatrix{Bool}, - A::SparseMatrixCSCOrColumnSubset) where {P} = _mapreducerows!(!pred, |, R, A, !) + A::SparseMatrixCSCOrColumnSubset) where {P} = _mapreducerows!(!pred, |, R, A, true) # findmax/min and argmax/min methods # find the first entry of the sparse matrix that is not stored - return its cartesian index diff --git a/test/reductions.jl b/test/reductions.jl index 24f310da..24a1ec37 100644 --- a/test/reductions.jl +++ b/test/reductions.jl @@ -372,6 +372,10 @@ end @test all(pred, Bs, dims = 1) == all(pred, Bd, dims = 1) @test all(pred, Cs, dims = 1) == all(pred, Cd, dims = 1) end + # `|` and `&` along the rows fold into `init` + D = [1.0 2.0; 0.0 3.0] + @test mapreduce(>(5), |, sparse(D); dims=1, init=true) == mapreduce(>(5), |, D; dims=1, init=true) + @test mapreduce(!iszero, &, sparse(D); dims=1, init=false) == mapreduce(!iszero, &, D; dims=1, init=false) end end