Handle sparse matrices of 2^63 or more entries in reductions, sprand and linear indexing

src/SparseArrays.jl

@@ -380,7 +380,11 @@ julia> collect(SparseArrays.iternz(x))
 """
 function iternz end
 
-widelength(x::AbstractSparseArray) = prod(Int64.(size(x)))
+widelength(x::AbstractSparseArray) = prod(widen, size(x))
+# Linear indexing compares an index with `length(A)`, which wraps for 2^63 entries or more
+_checklinearindexing(A::AbstractSparseMatrixCSC) = widelength(A) <= typemax(Int) ||
+    throw(OverflowError(string("a ", size(A, 1), "×", size(A, 2), " sparse matrix has more ",
+        "entries than an `Int` can index linearly; index it by row and column")))
 
 
 const _restore_scalar_indexing = Expr[]

src/constructors.jl

@@ -714,20 +714,43 @@ julia> sprand(Float64, 3, 0.75)
 """
 function sprand(r::AbstractRNG, m::Integer, n::Integer, density::AbstractFloat, rfn::Function, ::Type{T}) where T
     m, n = Int(m), Int(n)
+    if _sprandoverflows(m, n, density)
+        colptr, rowval, V = _sprand_bycolumn(r, m, n, density, rfn)
+        return SparseMatrixCSC(m, n, colptr, rowval, convert(Vector{T}, V))
+    end
     I, V = _sprandvalues(r, m, n, density, rfn)
     return sparse_sortedlinearindices!(I, convert(Vector{T}, V), m, n)
 end
 function sprand(r::AbstractRNG, m::Integer, n::Integer, density::AbstractFloat, rfn::Function)
     m, n = Int(m), Int(n)
+    if _sprandoverflows(m, n, density)
+        colptr, rowval, V = _sprand_bycolumn(r, m, n, density, rfn)
+        return SparseMatrixCSC(m, n, colptr, rowval, convert(Vector{eltype(V)}, V))
+    end
     I, V = _sprandvalues(r, m, n, density, rfn)
     return sparse_sortedlinearindices!(I, convert(Vector{eltype(V)}, V), m, n)
 end
-function _sprandvalues(r::AbstractRNG, m::Int, n::Int, density::AbstractFloat, rfn::F) where {F}
+# the arguments are checked first, since negative dimensions can overflow too
+function _sprandoverflows(m::Int, n::Int, density::AbstractFloat)
     (m < 0 || n < 0) && throw(ArgumentError("invalid Array dimensions"))
     0 <= density <= 1 || throw(ArgumentError("$density not in [0,1]"))
+    return Base.mul_with_overflow(m, n)[2]
+end
+function _sprandvalues(r::AbstractRNG, m::Int, n::Int, density::AbstractFloat, rfn::F) where {F}
     I = randsubseq(r, 1:(m*n), density)
     return I, rfn(r, length(I))
 end
+# an `Int` cannot index the `m*n` entries, so the stored rows are drawn column by column
+function _sprand_bycolumn(r::AbstractRNG, m::Int, n::Int, density::AbstractFloat, rfn::F) where {F}
+    colptr = Vector{Int}(undef, n + 1)
+    colptr[1] = 1
+    rowval, rows = Int[], Int[]
+    for col in 1:n
+        append!(rowval, randsubseq!(r, rows, 1:m, density))
+        colptr[col + 1] = length(rowval) + 1
+    end
+    return colptr, rowval, rfn(r, length(rowval))
+end
 
 sprand(m::Integer, n::Integer, density::AbstractFloat, rfn::Function) =
     sprand(default_rng(), m, n, density, (r, i) -> rfn(i))

src/indexing.jl

@@ -35,8 +35,10 @@ _integer_indices(I::AbstractArray{CartesianIndex{1}}) = map(first ∘ Tuple, I)
 
 # Index types no method below matches (`CartesianIndex`, custom indices) are lowered and
 # dispatched again; indices that are already lowered fall back to Base.
-@propagate_inbounds getindex(A::AbstractSparseMatrixCSC{Tv,Ti}, I...) where {Tv,Ti} =
-    _getindex_lowered(A, I, _lower_indices(A, I...))
+@propagate_inbounds function getindex(A::AbstractSparseMatrixCSC{Tv,Ti}, I...) where {Tv,Ti}
+    I isa Tuple{Integer} && _checklinearindexing(A)
+    return _getindex_lowered(A, I, _lower_indices(A, I...))
+end
 @propagate_inbounds getindex(x::AbstractSparseVector{Tv,Ti}, I...) where {Tv,Ti} =
     _getindex_lowered(x, I, _lower_indices(x, I...))
 @propagate_inbounds _getindex_lowered(A, I, L) = A[L...]
@@ -503,6 +505,7 @@ end
 
 function getindex(A::AbstractSparseMatrixCSC{Tv,Ti}, I::AbstractArray) where {Tv,Ti}
     require_one_based_indexing(A, I)
+    eltype(I) <: Integer && _checklinearindexing(A)
     ndims(I) == 2 || return _getindex_dense(A, I)
     szA = size(A)
     nA = szA[1]*szA[2]
@@ -596,6 +599,13 @@ _checkscalarindex(::Integer) = nothing
 
 # dispatch helper for #29034
 @RCI setindex!(A::AbstractSparseMatrixCSC, _v, _i::Integer, _j::Integer) = _setindex_scalar!(A, _v, _i, _j)
+@RCI function setindex!(A::AbstractSparseMatrixCSC, _v, _i::Integer)
+    _checkscalarindex(_i)
+    _checklinearindexing(A)
+    checkbounds(A, _i)
+    col, row = divrem(Int(_i) - 1, size(A, 1))
+    return _setindex_scalar!(A, _v, row + 1, col + 1)
+end
 
 function _setindex_scalar!(A::AbstractSparseMatrixCSC{Tv,Ti}, _v, _i::Integer, _j::Integer) where {Tv,Ti<:Integer}
     _checkscalarindex(_i); _checkscalarindex(_j)
@@ -1124,6 +1134,7 @@ end
 
 function setindex!(A::AbstractSparseMatrixCSC, x::AbstractArray, Ix::AbstractVector{<:Integer})
     require_one_based_indexing(A, x, Ix)
+    _checklinearindexing(A)
     # to_indices does not compare a Bool mask's length with A's
     Ix isa AbstractVector{Bool} && checkbounds(A, Ix)
     (I,) = Base.ensure_indexable(to_indices(A, (Ix,)))
@@ -1564,6 +1575,7 @@ getindex(A::AbstractSparseMatrixCSC, ::Colon) = A[1:end]
 
 function getindex(A::AbstractSparseMatrixCSC{Tv}, I::AbstractUnitRange) where Tv
     require_one_based_indexing(A, I)
+    _checklinearindexing(A)
     checkbounds(A, I)
     szA = size(A)
     nA = szA[1]*szA[2]
@@ -1606,6 +1618,7 @@ end
 function getindex(A::AbstractSparseMatrixCSC{Tv,Ti}, I::AbstractVector) where {Tv,Ti}
     I isa Base.LogicalIndex && return _logical_index(A, I.mask)
     require_one_based_indexing(A, I)
+    eltype(I) <: Integer && _checklinearindexing(A)
     @boundscheck checkbounds(A, I)
     szA = size(A)
     nA = szA[1]*szA[2]

src/reductions.jl

@@ -95,13 +95,13 @@ function _mapreduce_ordered(f::F, op::G, A::SparseMatrixCSCOrColumnSubset{T}) wh
     v, dropout = _foldstored(f, op, z, nzvalview(A))
     dropout && return v
     v = Base._InitialValue()
-    unstored = Int64(0)   # spans columns, so it can exceed typemax(Int32)
+    unstored = widen(0)   # spans columns, so it can exceed typemax(Int)
     @inbounds for col in axes(A, 2)
         prev = 0
         for k in getnzrange(A, col)
             row = Int(rowval[k])
             v = rf(v, _afterunstored(op, z, idem, unstored + (row - prev - 1), f(nzval[k])))
-            unstored, prev = Int64(0), row
+            unstored, prev = widen(0), row
         end
         unstored += m - prev
     end
@@ -187,14 +187,30 @@ end
 _mapreducezeros(f::F, op::Union{typeof(Base.add_sum),typeof(+)}, ::Type{T}, nzeros::Integer, v0) where {F,T} =
     nzeros == 0 ? Base.reduce_first(op, v0) : op(_addzeros(_unstored(op, f, T, v0), nzeros), v0)
 _mapreducezeros(f::F, op::Union{typeof(Base.mul_prod),typeof(*)},::Type{T}, nzeros::Integer, v0) where {F,T} =
-    nzeros == 0 ? Base.reduce_first(op, v0) : op(_unstored(op, f, T, v0)^nzeros, v0)
+    nzeros == 0 ? Base.reduce_first(op, v0) : op(_powzeros(_unstored(op, f, T, v0), nzeros), v0)
+# `z` to the power of `nzeros`. An eltype may define `^` for an `Int` power only, so a count
+# from `widelength` is split into `Int` powers, as `_addzeros` splits it into factors.
+_powzeros(z, nzeros) = z^nzeros
+function _powzeros(z, nzeros::Int128)
+    nzeros <= typemax(Int) && return z^Int(nzeros)
+    q, r = divrem(nzeros, typemax(Int))
+    return (z^typemax(Int))^Int(q) * z^Int(r)
+end
 _unstored(op, f, ::Type{T}, v0) where T = _widen(Base.reduce_first(op, f(zero(T))), v0)
 _widen(z, v0) = _widen(z, typeof(v0))
 _widen(z, ::Type{S}) where S = convert(promote_type(typeof(z), S), z)
 # `nzeros` copies of `z` added up, in the type `+` keeps for them: a small integer wraps
-# rather than widening to the count's `Int`, as when dense input is summed with `+`
+# rather than widening to the count's `Int`, as when dense input is summed with `+`. A count
+# from `widelength` is an `Int128`, which would widen `Complex{Int}` and the like, so it is
+# split into `Int` factors.
 _addzeros(z, nzeros) = z * nzeros
+function _addzeros(z, nzeros::Int128)
+    nzeros <= typemax(Int) && return z * Int(nzeros)
+    q, r = divrem(nzeros - 1, typemax(Int))
+    return z * Int(q) * typemax(Int) + z * Int(r + 1)
+end
 _addzeros(z::Base.BitInteger, nzeros) = z * (nzeros % typeof(z))
+_addzeros(z::Base.BitInteger, nzeros::Int128) = z * (nzeros % typeof(z))
 _mapreducezeros(f::F, op::Union{typeof(min),typeof(max)}, ::Type{T}, nzeros::Integer, v0) where {F,T} =
     nzeros == 0 ? v0 : op(v0, f(zero(T)))
 _mapreducezeros(f::Base.ExtremaMap, op::typeof(Base._extrema_rf), ::Type{T}, nzeros::Integer, v0) where {T} =
@@ -220,7 +236,7 @@ function Base._mapreduce(f::F, op::Union{typeof(Base.mul_prod),typeof(*)}, ::Bas
     elseif !_commutes(f, op, T)
         _mapreduce_ordered(f, op, A)
     else
-        v = Base.reduce_first(op, f(zero(T)))^nzeros
+        v = _powzeros(Base.reduce_first(op, f(zero(T))), nzeros)
         (nnzA == 0 || (_iszero(v) && Base.promote_op(f, T) <: Integer)) && return v
         r = v*Base._mapreduce(f, op, nzvalview(A))
         # A zero ends the product unless a stored `Inf` or `NaN` makes it `NaN`. The product
@@ -266,7 +282,7 @@ end
 # A reduction along a dimension, an empty array and a shape mismatch are left to Base's method.
 for T in (:SparseVectorOrView, :SparseMatrixCSCOrColumnSubset)
     @eval function Base.mapreduce(f::F, op::G, A::$T, B::$T; dims=:, init=Base._InitialValue(), kw...) where {F,G}
-        (dims === (:) && isempty(kw) && size(A) == size(B) && !isempty(A)) ||
+        (dims === (:) && isempty(kw) && size(A) == size(B) && !iszero(widelength(A))) ||
             return reduce(op, map(f, A, B); dims, init, kw...)
         return _commutes(f, op, eltype(A), eltype(B)) ? _mapreduce_unordered(f, op, init, A, B) :
                                                          _mapreduce_ordered(f, op, init, A, B)
@@ -280,7 +296,7 @@ function _mapreduce_ordered(f::F, op::G, init, A::AbstractArray{Ta}, B::Abstract
     Ai, Av, Bi, Bv = getrowval(A), getnzval(A), getrowval(B), getnzval(B)
     m = size(A, 1)
     v = init
-    unstored = Int64(0)   # positions that neither array stores, since the last one folded
+    unstored = widen(0)   # positions that neither array stores, since the last one folded
     @inbounds for col in axes(A, 2)
         ra, rb = getnzrange(A, col), getnzrange(B, col)
         ia, ib = first(ra), first(rb)
@@ -290,18 +306,18 @@ function _mapreduce_ordered(f::F, op::G, init, A::AbstractArray{Ta}, B::Abstract
             row = Int(min(Ai[ia], Bi[ib]))
             x, ia, ib = _merge_step(f, Ai, Av, ia, Bi, Bv, ib)
             v = rf(_foldunstored(f, rf, v, Ta, Tb, unstored + (row - prev - 1)), x)
-            unstored, prev = Int64(0), row
+            unstored, prev = widen(0), row
         end
         while ia <= ea
             row = Int(Ai[ia])
             v = rf(_foldunstored(f, rf, v, Ta, Tb, unstored + (row - prev - 1)), f(Av[ia], zero(Tb)))
-            unstored, prev = Int64(0), row
+            unstored, prev = widen(0), row
             ia += 1
         end
         while ib <= eb
             row = Int(Bi[ib])
             v = rf(_foldunstored(f, rf, v, Ta, Tb, unstored + (row - prev - 1)), f(zero(Ta), Bv[ib]))
-            unstored, prev = Int64(0), row
+            unstored, prev = widen(0), row
             ib += 1
         end
         unstored += m - prev
@@ -365,7 +381,7 @@ function _mapreduce_unordered(f::F, op::G, init, A::AbstractArray{Ta}, B::Abstra
         ib == eb && (v = rf(v, f(zero(Ta), Bv[ib])))
     end
     w isa Base._InitialValue || (v = v isa Base._InitialValue ? w : op(v, w))
-    nzeros = prod(Int64.(size(A))) - stored
+    nzeros = prod(widen, size(A)) - stored
     if v isa Base._InitialValue   # nothing stored and no `init`: seed with one unstored position
         v = rf(v, f(zero(Ta), zero(Tb)))
         nzeros -= 1
@@ -627,7 +643,7 @@ end
 # The kernels above fold each slice's stored entries first and its unstored ones last. That
 # matches element order only when `op` commutes or every entry is stored; otherwise the
 # kernels below fold the unstored entries where they lie.
-_inorder(f, op, A::AbstractArray{T}) where {T} = _commutes(f, op, T) || nnz(A) == length(A)
+_inorder(f, op, A::AbstractArray{T}) where {T} = _commutes(f, op, T) || nnz(A) == widelength(A)
 function _mapreducerows_ordered!(f::F, op::G, R::AbstractArray, A::SparseMatrixCSCOrColumnSubset{T}) where {F,G,T}
     require_one_based_indexing(A, R)
     rowval, nzval = getrowval(A), getnzval(A)

src/sparsematrix.jl

@@ -319,7 +319,8 @@ function _storedinds(S::SparseMatrixCSCColumnSubset)
     end
     return inds
 end
-widelength(S::SparseMatrixCSCColumnSubset) = prod(Int64.(size(S)))
+widelength(S::SparseMatrixCSCColumnSubset) = prod(widen, size(S))
+widelength(x::Union{SparseColumnView,SparseVectorView}) = length(x)
 # `A` itself where it exposes compressed storage through `nzrange`, `rowvals` and `nonzeros`,
 # otherwise its sparse copy
 _compressed(A) = A
@@ -353,7 +354,7 @@ nnz(S::SparseMatrixCSCView) = length(_storedinds(S))
 nnz1(S) = @inbounds sum(length.(getnzrange.(Ref(S), axes(S, 2))))
 
 function Base._simple_count(pred, S::SparseMatrixCSCOrColumnSubset, init::T) where T
-    init + T(count(pred, nzvalview(S)) + pred(zero(eltype(S)))*(prod(size(S)) - nnz(S)))
+    init + T(count(pred, nzvalview(S)) + pred(zero(eltype(S)))*(widelength(S) - nnz(S)))
 end
 # the entries of a lazy adjoint or transpose are those of its parent, each under the wrapper
 Base._simple_count(pred, S::AdjOrTrans{<:Any,<:SparseMatrixCSCOrColumnSubset}, init) =

test/issues.jl

@@ -968,6 +968,20 @@ end
     @inferred hcat(fixture(Float64, 2, 1), Matrix(1.0I, 2, 2))
 end
 
+@testset "a matrix of 2^64 entries does not overflow its length" begin
+    B = spzeros(2^62, 4)
+    @test any(iszero, B) && !all(isone, B)
+    @test sum(x -> 1.0, B) == 2.0^64 && prod(x -> 2.0, B) == Inf
+    @test count(iszero, B; init = big(0)) == big(2)^64
+    @test findmax(B) == (0.0, CartesianIndex(1, 1))
+    @test mapreduce(*, +, B, B) == 0.0
+    @test_throws OverflowError B[[1]]
+    @test_throws OverflowError B[[1]] = [1.0]
+    @test_throws OverflowError B[1] = 1.0
+    @test nnz(B) == 0
+    @test nnz(sprand(Xoshiro(1), 2^62, 4, 1e-17)) > 0
+end
+
 if isdefined(Docs, :undocumented_names) # new in Julia 1.11
     @testset "docstrings (issue julia#52725)" begin
         @test isempty(Docs.undocumented_names(SparseArrays))

test/trim/src/SparseTrim.jl

@@ -91,6 +91,11 @@ function indexing()
     B[:, 2] .= 1.0
     bad += check(B[:, 2] == [1.0, 1.0, 1.0])
     bad += check(Array(view(A, :, 3)) == [0.0, 0.0, 5.0])
+    bad += check(A[3] == 4.0 && A[[3, 5]] == sparsevec([1, 2], [4.0, 3.0], 2) && A[2:3] == sparsevec([2], [4.0], 2))
+    C = copy(A)
+    C[2] = 1.0
+    C[[4]] = [7.0]
+    bad += check(C[2, 1] == 1.0 && C[1, 2] == 7.0 && nnz(C) == 6)
     return bad
 end
 
@@ -113,6 +118,8 @@ function algebra()
     bad += check(A - B == sparse([1, 2, 3, 3, 1], [1, 2, 1, 3, 3], [2.0, 3.0, 5.0, 5.0, -1.0], 3, 3))
     bad += check(A * B == sparse([1, 3, 3], [3, 1, 3], [2.0, -5.0, 4.0], 3, 3))
     bad += check(Matrix(A * D) == [4.0 0.0 0.0; 0.0 9.0 0.0; 28.0 0.0 25.0])
+    bad += check(mapreduce(*, +, A, A) == 54.0 && prod(x -> x + 1.0, A) == 360.0)
+    bad += check(mapreduce(identity, (a, b) -> a + b, A) == 14.0)
     bad += check(sum(A) == 14.0 && sum(A; dims = 1) == [6.0 3.0 5.0])
     bad += check(sum(A; dims = 2) == reshape([2.0, 3.0, 9.0], 3, 1))
     bad += check(maximum(A) == 5.0 && minimum(B) == -1.0 && prod(A) == 0.0)