Fix "map does not respect sources index type" (#350) (#351)

dpinol · Dani Pinyol · web-flow · commit d938744fcd2d · 2023-02-24T23:33:03.000-05:00
Co-authored-by: Dani Pinyol &lt;dani@avatarcognition.com&gt;
diff --git a/src/sparsevector.jl b/src/sparsevector.jl
@@ -1380,6 +1380,7 @@ function _binarymap(f::Function,
                     mode::Int) where {Tx,Ty}
     0 <= mode <= 2 || throw(ArgumentError("Incorrect mode $mode."))
     R = Base.Broadcast.combine_eltypes(f, (x, y))
+    I = promote_type(eltype(nonzeroinds(x)), eltype(nonzeroinds(y)))
     n = length(x)
     length(y) == n || throw(DimensionMismatch())
 
@@ -1391,7 +1392,7 @@ function _binarymap(f::Function,
     my = length(ynzind)
     cap = (mode == 0 ? min(mx, my) : mx + my)::Int
 
-    rind = Vector{Int}(undef, cap)
+    rind = Vector{I}(undef, cap)
     rval = Vector{R}(undef, cap)
     ir = 0
     ir = (
diff --git a/test/sparsevector.jl b/test/sparsevector.jl
@@ -1601,6 +1601,12 @@ end
     @test length(nonzeros(simA)) == 0
 end
 
+@testset "map preserves index types" begin
+    v1 = spzeros(Float32, Int16, 10)
+    v2 = spzeros(Float32, Int32, 10)
+    @test eltype(typeof(SparseArrays.nonzeroinds(map(max, v1, v2)))) == Int32
+end
+
 @testset "Fast operations on full column views" begin
     n = 1000
     A = sprandn(n, n, 0.01)
