Infer outputtype by return_types

Author: felixcremer

src/DAT/broadcast.jl

@@ -17,8 +17,11 @@ function Base.materialize(bc::Broadcast.Broadcasted{XStyle})
     args2 = map(arg -> arg isa Broadcast.Broadcasted ? Base.materialize(arg) : arg, bc.args)
     args2 = map(to_yax, args2)
     # determine output type by calling `eltype` on a dummy function call
-    dummy_args = map(a -> first(a.data), args2)
-    outtype = typeof(bc.f(dummy_args...))
+    intypes = (eltype.(args2)...,)
+    @debug intypes
+    outtypes = Base.return_types(bc.f, intypes)
+    outtype = Union{outtypes...}
+    @debug outtype
     return xmap(XFunction(bc.f; inplace=false), args2..., output=XOutput(; outtype))
 end
 function Base.materialize!(bc::Broadcast.Broadcasted{XStyle})

test/DAT/broadcast.jl

@@ -60,4 +60,12 @@ a, b, c = sample_arrays()
     xscalar = a .* 3 .+ 1
     @test all(xscalar[:] .== 4.0)
     @test isa(a .+ b, YAXArray)
+end
+
+@testset "missing handling" begin
+    am = YAXArray([missing 1 ; 1 2])
+    aeq = am .== am
+    @test eltype(aeq) == Union{Missing, Bool}
+    @test ismissing(aeq[1,1])
+    @test aeq[1,2]
 end