Merge pull request #401 from JuliaParallel/jps/darray-correct-eltype

DArray: Figure out correct eltype

Author: jpsamaroo

.buildkite/pipeline.yml

@@ -4,7 +4,6 @@
     queue: "juliaecosystem"
     sandbox_capable: "true"
     os: linux
-    arch: x86_64
   command: "julia --project -e 'using Pkg; Pkg.develop(;path=\"lib/TimespanLogging\")'"
 .bench: &bench
   if: build.message =~ /\[run benchmarks\]/

src/array/darray.jl

@@ -116,8 +116,8 @@ mutable struct DArray{T,N,F} <: ArrayOp{T, N}
     subdomains::AbstractArray{ArrayDomain{N}, N}
     chunks::AbstractArray{Any, N}
     concat::F
-    function DArray{T,N,F}(domain, subdomains, chunks, concat::Function) where {T, N,F}
-        new(domain, subdomains, chunks, concat)
+    function DArray{T,N,F}(domain, subdomains, chunks, concat::Function) where {T,N,F}
+        new{T,N,F}(domain, subdomains, chunks, concat)
     end
 end
 
@@ -227,14 +227,14 @@ end
 
 If a `DArray` tree has a `Thunk` in it, make the whole thing a big thunk.
 """
-function Base.fetch(c::DArray)
+function Base.fetch(c::DArray{T}) where T
     if any(istask, chunks(c))
         thunks = chunks(c)
         sz = size(thunks)
         dmn = domain(c)
         dmnchunks = domainchunks(c)
         fetch(Dagger.spawn(Options(meta=true), thunks...) do results...
-            t = eltype(results[1])
+            t = eltype(fetch(results[1]))
             DArray(t, dmn, dmnchunks, reshape(Any[results...], sz))
         end)
     else

src/array/map-reduce.jl

@@ -19,9 +19,10 @@ function stage(ctx::Context, node::Map)
     f = node.f
     for i=eachindex(domains)
         inps = map(x->chunks(x)[i], inputs)
-        thunks[i] = Dagger.@spawn map(f, inps...) 
+        thunks[i] = Dagger.@spawn map(f, inps...)
     end
-    DArray(Any, domain(primary), domainchunks(primary), thunks)
+    RT = Base.promote_op(node.f, map(eltype, node.inputs)...)
+    return DArray(RT, domain(primary), domainchunks(primary), thunks)
 end
 
 map(f, x::ArrayOp, xs::ArrayOp...) = Map(f, (x, xs...))

test/array.jl

@@ -13,6 +13,7 @@ end
 @testset "DArray constructor" begin
     x = fetch(rand(Blocks(2,2), 3,3))
     @test collect(x) == DArray{Float64, 2}(x.domain, x.subdomains, x.chunks) |> collect
+    @test x isa DArray{Float64, 2}
 end
 
 @testset "rand" begin
@@ -39,6 +40,13 @@ end
     @test r[1:10] != r[11:20]
 end
 
+@testset "map" begin
+    X1 = fetch(ones(Blocks(10, 10), 100, 100))
+    X2 = fetch(map(x->x+1, X1))
+    @test typeof(X1) === typeof(X2)
+    @test collect(X1) .+ 1 == collect(X2)
+end
+
 @testset "sum" begin
     X = ones(Blocks(10, 10), 100, 100)
     @test sum(X) == 10000