Merge pull request #528 from JuliaParallel/jps/darray-0dim

DArray: Improve 0-dim support

Author: jpsamaroo

Diff for: src/array/darray.jl

@@ -20,6 +20,7 @@ ArrayDomain(xs::NTuple{N,Base.OneTo}) where N =
     ArrayDomain{N,NTuple{N,UnitRange{Int}}}(ntuple(i->UnitRange(xs[i]), N))
 ArrayDomain(xs::NTuple{N,Int}) where N =
     ArrayDomain{N,NTuple{N,UnitRange{Int}}}(ntuple(i->xs[i]:xs[i], N))
+ArrayDomain(::Tuple{}) = ArrayDomain{0,Tuple{}}(())
 ArrayDomain(xs...) = ArrayDomain((xs...,))
 ArrayDomain(xs::Array) = ArrayDomain((xs...,))
 
@@ -31,6 +32,7 @@ chunks(a::ArrayDomain{N}) where {N} = DomainBlocks(
 
 (==)(a::ArrayDomain, b::ArrayDomain) = indexes(a) == indexes(b)
 Base.getindex(arr::AbstractArray, d::ArrayDomain) = arr[indexes(d)...]
+Base.getindex(arr::AbstractArray{T,0} where T, d::ArrayDomain{0}) = arr
 
 function intersect(a::ArrayDomain, b::ArrayDomain)
     if a === b
@@ -145,7 +147,6 @@ const WrappedDVector{T} = WrappedDArray{T,1}
 const DMatrix{T} = DArray{T,2}
 const DVector{T} = DArray{T,1}
 
-
 # mainly for backwards-compatibility
 DArray{T, N}(domain, subdomains, chunks, partitioning, concat=cat) where {T,N} =
     DArray(T, domain, subdomains, chunks, partitioning, concat)
@@ -178,6 +179,10 @@ function Base.collect(d::DArray; tree=false)
         return Array{eltype(d)}(undef, size(d)...)
     end
 
+    if ndims(d) == 0
+        return fetch(a.chunks[1])
+    end
+
     dimcatfuncs = [(x...) -> d.concat(x..., dims=i) for i in 1:ndims(d)]
     if tree
         collect(fetch(treereduce_nd(map(x -> ((args...,) -> Dagger.@spawn x(args...)) , dimcatfuncs), a.chunks)))
@@ -214,6 +219,7 @@ else
     Base.alignment(io::IO, x::ColorElement) =
         Base.alignment(io, something(x.value, "..."))
 #end
+Base.show(io::IO, x::ColorElement) = show(io, MIME("text/plain"), x)
 struct ColorArray{T,N} <: DenseArray{T,N}
     A::DArray{T,N}
     color_map::Vector{Symbol}
@@ -261,9 +267,21 @@ function Base.getindex(A::ColorArray{T,N}, idxs::Dims{S}) where {T,N,S}
     end
 end
 function Base.show(io::IO, ::MIME"text/plain", A::DArray{T,N}) where {T,N}
-    write(io, string(DArray{T,N}))
-    write(io, string(size(A)))
-    write(io, " with $(join(size(A.chunks), 'x')) partitions of size $(join(A.partitioning.blocksize, 'x')):")
+    if N == 1
+        write(io, "$(length(A))-element ")
+        write(io, string(DVector{T}))
+    elseif N == 2
+        write(io, string(DMatrix{T}))
+    elseif N == 0
+        write(io, "0-dimensional ")
+        write(io, "DArray{$T, $N}")
+    else
+        write(io, "$(join(size(A), 'x')) ")
+        write(io, "DArray{$T, $N}")
+    end
+    nparts = N > 0 ? size(A.chunks) : 1
+    partsize = N > 0 ? A.partitioning.blocksize : 1
+    write(io, " with $(join(nparts, 'x')) partitions of size $(join(partsize, 'x')):")
     pct_complete = 100 * (sum(c->c isa Chunk ? true : isready(c), A.chunks) / length(A.chunks))
     if pct_complete < 100
         println(io)
@@ -472,7 +490,7 @@ struct AutoBlocks end
 function auto_blocks(dims::Dims{N}) where N
     # TODO: Allow other partitioning schemes
     np = num_processors()
-    p = cld(dims[end], np)
+    p = N > 0 ? cld(dims[end], np) : 1
     return Blocks(ntuple(i->i == N ? p : dims[i], N))
 end
 auto_blocks(A::AbstractArray{T,N}) where {T,N} = auto_blocks(size(A))

Diff for: test/array/allocation.jl

@@ -1,9 +1,16 @@
 @testset "DVector/DMatrix/DArray constructor" begin
     for T in [Float32, Float64, Int32, Int64]
+        F = fill(one(T))
         V = rand(T, 64)
         M = rand(T, 64, 64)
         A = rand(T, 64, 64, 64)
 
+        # DArray ctor (empty)
+        DF = DArray(F, Blocks(()))
+        @test DF isa DArray{T,0}
+        @test collect(DF) == F
+        @test size(DF) == size(F)
+
         # DVector ctor
         DV = DVector(V, Blocks(8))
         @test DV isa DVector{T}
@@ -26,7 +33,8 @@ end
 
 @testset "random" begin
     for T in [Float32, Float64, Int32, Int64]
-        for dims in [(100,),
+        for dims in [(),
+                     (100,),
                      (100, 100),
                      (100, 100, 100)]
             dist = Blocks(ntuple(i->10, length(dims))...)
@@ -53,10 +61,9 @@ end
                 @test AXn isa Array{T,length(dims)}
                 @test AXn == collect(Xn)
                 @test AXn != collect(randn(dist, T, dims...))
-                @test !all(AXn .> 0)
             end
 
-            if length(dims) <= 2
+            if 1 <= length(dims) <= 2
                 # sprand
                 Xsp = sprand(dist, T, dims..., 0.1)
                 @test Xsp isa DArray{T,length(dims)}
@@ -76,7 +83,8 @@ end
 @testset "ones/zeros" begin
     for T in [Float32, Float64, Int32, Int64]
         for (fn, value) in [(ones, one(T)), (zeros, zero(T))]
-            for dims in [(100,),
+            for dims in [(),
+                         (100,),
                          (100, 100),
                          (100, 100, 100)]
                 dist = Blocks(ntuple(i->10, length(dims))...)
@@ -119,11 +127,16 @@ end
         for i in 1:(length(dims)-1)
             @test part_size[i] == 100
         end
-        @test part_size[end] == cld(100, np)
+        if length(dims) > 0
+            @test part_size[end] == cld(100, np)
+        else
+            @test part_size == ()
+        end
         @test size(DA) == ntuple(i->100, length(dims))
     end
 
-    for dims in [(100,),
+    for dims in [(),
+                 (100,),
                  (100, 100),
                  (100, 100, 100)]
         fn = if length(dims) == 1
@@ -133,15 +146,23 @@ end
         else
             DArray
         end
-        DA = fn(rand(dims...), AutoBlocks())
+        if length(dims) > 0
+            DA = fn(rand(dims...), AutoBlocks())
+        else
+            DA = fn(fill(rand()), AutoBlocks())
+        end
         test_auto_blocks(DA, dims)
 
-        DA = distribute(rand(dims...), AutoBlocks())
+        if length(dims) > 0
+            DA = distribute(rand(dims...), AutoBlocks())
+        else
+            DA = distribute(fill(rand()), AutoBlocks())
+        end
         test_auto_blocks(DA, dims)
 
         for fn in [rand, randn, sprand, ones, zeros]
             if fn === sprand
-                if length(dims) > 2
+                if length(dims) > 2 || length(dims) == 0
                     continue
                 end
                 DA = fn(AutoBlocks(), dims..., 0.1)
@@ -155,7 +176,8 @@ end
 
 @testset "Constructor variants" begin
     for fn in [ones, zeros, rand, randn, sprand]
-        for dims in [(100,),
+        for dims in [(),
+                     (100,),
                      (100, 100),
                      (100, 100, 100)]
             for dist in [Blocks(ntuple(i->10, length(dims))...),