Handle zero on arrays of unions of number types and missings (#53602)

(cherry picked from commit fc6c618)
JuliaLang · Mar 15, 2024 · 8c9015f · 8c9015f
1 parent fcadc06
commit 8c9015f
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/base/abstractarray.jl b/base/abstractarray.jl
@@ -1218,6 +1218,7 @@ end
 copymutable(itr) = collect(itr)
 
 zero(x::AbstractArray{T}) where {T<:Number} = fill!(similar(x, typeof(zero(T))), zero(T))
+zero(x::AbstractArray{S}) where {S<:Union{Missing, Number}} = fill!(similar(x, typeof(zero(S))), zero(S))
 zero(x::AbstractArray) = map(zero, x)
 
 ## iteration support for arrays by iterating over `eachindex` in the array ##

diff --git a/test/abstractarray.jl b/test/abstractarray.jl
@@ -1972,6 +1972,16 @@ end
 
     @test zero([[2,2], [3,3,3]]) isa Vector{Vector{Int}}
     @test zero([[2,2], [3,3,3]]) == [[0,0], [0, 0, 0]]
+
+
+    @test zero(Union{Float64, Missing}[missing]) == [0.0]
+    struct CustomNumber <: Number
+        val::Float64
+    end
+    Base.zero(::Type{CustomNumber}) = CustomNumber(0.0)
+    @test zero([CustomNumber(5.0)]) == [CustomNumber(0.0)]
+    @test zero(Union{CustomNumber, Missing}[missing]) == [CustomNumber(0.0)]
+    @test zero(Vector{Union{CustomNumber, Missing}}(undef, 1)) == [CustomNumber(0.0)]
 end
 
 @testset "`_prechecked_iterate` optimization" begin