Fix sparse array setindex(::Int, ::Vector) (#43678)

JuliaLang · Jan 7, 2022 · dc61f29 · dc61f29
1 parent cbf6c1c
commit dc61f29
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 17 deletions.
diff --git a/stdlib/SparseArrays/src/sparsematrix.jl b/stdlib/SparseArrays/src/sparsematrix.jl
@@ -2878,6 +2878,7 @@ end
 
 # Nonscalar A[I,J] = B: Convert B to a SparseMatrixCSC of the appropriate shape first
 _to_same_csc(::AbstractSparseMatrixCSC{Tv, Ti}, V::AbstractMatrix, I...) where {Tv,Ti} = convert(SparseMatrixCSC{Tv,Ti}, V)
+_to_same_csc(::AbstractSparseMatrixCSC{Tv, Ti}, V::AbstractMatrix, i::Integer, J) where {Tv,Ti} = convert(SparseMatrixCSC{Tv,Ti}, reshape(V, (1, length(J))))
 _to_same_csc(::AbstractSparseMatrixCSC{Tv, Ti}, V::AbstractVector, I...) where {Tv,Ti} = convert(SparseMatrixCSC{Tv,Ti}, reshape(V, map(length, I)))
 
 setindex!(A::AbstractSparseMatrixCSC{Tv}, B::AbstractVecOrMat, I::Integer, J::Integer) where {Tv} = _setindex_scalar!(A, B, I, J)
@@ -2886,12 +2887,20 @@ function setindex!(A::AbstractSparseMatrixCSC{Tv,Ti}, V::AbstractVecOrMat, Ix::U
     require_one_based_indexing(A, V, Ix, Jx)
     (I, J) = Base.ensure_indexable(to_indices(A, (Ix, Jx)))
     checkbounds(A, I, J)
-    Base.setindex_shape_check(V, length(I), length(J))
+    nJ = length(J)
+    Base.setindex_shape_check(V, length(I), nJ)
     B = _to_same_csc(A, V, I, J)
 
+    m, n = size(A)
+    if (!isempty(I) && (I[1] < 1 || I[end] > m)) || (!isempty(J) && (J[1] < 1 || J[end] > n))
+        throw(BoundsError(A, (I, J)))
+    end
+    if isempty(I) || isempty(J)
+        return A
+    end
+
     issortedI = issorted(I)
     issortedJ = issorted(J)
-
     if !issortedI && !issortedJ
         pI = sortperm(I); @inbounds I = I[pI]
         pJ = sortperm(J); @inbounds J = J[pJ]
@@ -2904,20 +2913,6 @@ function setindex!(A::AbstractSparseMatrixCSC{Tv,Ti}, V::AbstractVecOrMat, Ix::U
         B = B[:, pJ]
     end
 
-    m, n = size(A)
-    mB, nB = size(B)
-
-    if (!isempty(I) && (I[1] < 1 || I[end] > m)) || (!isempty(J) && (J[1] < 1 || J[end] > n))
-        throw(BoundsError(A, (I, J)))
-    end
-
-    if isempty(I) || isempty(J)
-        return A
-    end
-
-    nI = length(I)
-    nJ = length(J)
-
     colptrA = getcolptr(A); rowvalA = rowvals(A); nzvalA = nonzeros(A)
     colptrB = getcolptr(B); rowvalB = rowvals(B); nzvalB = nonzeros(B)
 
@@ -2931,7 +2926,6 @@ function setindex!(A::AbstractSparseMatrixCSC{Tv,Ti}, V::AbstractVecOrMat, Ix::U
     resize!(nzvalA, nnzS)
 
     colB = 1
-    asgn_col = J[colB]
 
     I_asgn = falses(m)
     fill!(view(I_asgn, I), true)

diff --git a/test/arrayops.jl b/test/arrayops.jl
@@ -2964,3 +2964,18 @@ end
         @test c + zero(c) == c
     end
 end
+
+@testset "Allow assignment of singleton array to sparse array #43644" begin
+    K = spzeros(3,3)
+    b = zeros(3,3)
+    b[3,:] = [1,2,3]
+    K[3,1:3] += [1.0 2.0 3.0]'
+    @test K == b
+    K[3:3,1:3] += zeros(1, 3)
+    @test K == b
+    K[3,1:3] += zeros(3)
+    @test K == b
+    K[3,:] += zeros(3,1)
+    @test K == b
+    @test_throws DimensionMismatch K[3,1:2] += [1.0 2.0 3.0]'
+end