last fixes

Author: meggart

src/cat.jl

@@ -18,30 +18,54 @@ struct ConcatDiskArray{T,N,P,C,HC} <: AbstractDiskArray{T,N}
     chunks::C
     haschunks::HC
 end
-function ConcatDiskArray(arrays::AbstractArray{<:AbstractArray{<:Any,N},M}) where {N,M}
-    T = mapreduce(eltype, promote_type, init=eltype(first(arrays)), arrays)
 
+function ConcatDiskArray(arrays::AbstractArray{Union{<:AbstractArray,Missing}})
+    et = Base.nonmissingtype(eltype(arrays))
+    T = Union{Missing,eltype(et)}
+    N = ndims(arrays)
+    M = ndims(et)
+    _ConcatDiskArray(arrays, T, Val(N), Val(M))
+end
+function ConcatDiskArray(arrays::AbstractArray{<:AbstractArray})
+    T = eltype(eltype(arrays))
+    N = ndims(arrays)
+    M = ndims(eltype(arrays))
+    _ConcatDiskArray(arrays, T, Val(N), Val(M))
+end
+function ConcatDiskArray(arrays::AbstractArray)
+    N = ndims(arrays)
+    M, T = foldl(arrays, init=(-1, Union{})) do (M, T), a
+        if ismissing(a)
+            (M, promote_type(Missing, T))
+        else
+            M == -1 || ndims(a) == M || throw(ArgumentError("All arrays to concatenate must have equal ndims"))
+            (ndims(a), promote_type(eltype(a), T))
+        end
+    end
+    _ConcatDiskArray(arrays, T, Val(N), Val(M))
+end
+
+
+function _ConcatDiskArray(arrays, T, ::Val{N}, ::Val{M}) where {N,M}
     if N > M
-        newshape = extenddims(size(arrays), size(first(arrays)), 1)
+        newshape = extenddims(size(arrays), ntuple(_ -> 1, N), 1)
         arrays1 = reshape(arrays, newshape)
         D = N
     else
         arrays1 = arrays
         D = M
     end
+    _ConcatDiskArray(arrays1::AbstractArray, T, Val(D))
+end
+function _ConcatDiskArray(arrays1::AbstractArray, T, ::Val{D}) where {D}
     startinds, sizes = arraysize_and_startinds(arrays1)
 
     chunks = concat_chunksize(arrays1)
     hc = Chunked(batchstrategy(chunks))
 
     return ConcatDiskArray{T,D,typeof(arrays1),typeof(chunks),typeof(hc)}(arrays1, startinds, sizes, chunks, hc)
 end
-function ConcatDiskArray(arrays::AbstractArray)
-    # Validate array eltype and dimensionality
-    all(a -> ndims(a) == ndims(first(arrays)), arrays) ||
-        error("Arrays don't have the same dimensions")
-    return error("Should not be reached")
-end
+
 extenddims(a::Tuple{Vararg{Any,N}}, b::Tuple{Vararg{Any,M}}, fillval) where {N,M} = extenddims((a..., fillval), b, fillval)
 extenddims(a::Tuple{Vararg{Any,N}}, _::Tuple{Vararg{Any,N}}, _) where {N} = a
 
@@ -51,6 +75,7 @@ function arraysize_and_startinds(arrays1)
     sizes = map(i -> zeros(Int, i), size(arrays1))
     for i in CartesianIndices(arrays1)
         ai = arrays1[i]
+        ismissing(ai) && continue
         sizecur = extenddims(size(ai), size(arrays1), 1)
         foreach(sizecur, i.I, sizes) do si, ind, sizeall
             if sizeall[ind] == 0
@@ -62,6 +87,9 @@ function arraysize_and_startinds(arrays1)
         end
     end
     r = map(sizes) do sizeall
+        #Replace missing sizes with size 1
+        replace!(sizeall, 0 => 1)
+        #Add starting 1
         pushfirst!(sizeall, 1)
         for i in 2:length(sizeall)
             sizeall[i] = sizeall[i-1] + sizeall[i]
@@ -80,13 +108,24 @@ function readblock!(a::ConcatDiskArray, aout, inds::AbstractUnitRange...)
     # Find affected blocks and indices in blocks
     _concat_diskarray_block_io(a, inds...) do outer_range, array_range, I
         vout = view(aout, outer_range...)
-        readblock!(a.parents[I], vout, array_range...)
+        if ismissing(I)
+            vout .= missing
+        else
+            readblock!(a.parents[I], vout, array_range...)
+        end
     end
 end
 function writeblock!(a::ConcatDiskArray, aout, inds::AbstractUnitRange...)
     _concat_diskarray_block_io(a, inds...) do outer_range, array_range, I
         data = view(aout, outer_range...)
-        writeblock!(a.parents[I], data, array_range...)
+        if ismissing(I)
+            if !all(ismissing, data)
+                @warn "Trying to write data to missing array tile, skipping write"
+            end
+            return
+        else
+            writeblock!(a.parents[I], data, array_range...)
+        end
     end
 end
 
@@ -101,17 +140,24 @@ function _concat_diskarray_block_io(f, a::ConcatDiskArray, inds...)
     end
     map(CartesianIndices(blockinds)) do cI
         myar = a.parents[cI]
-        mysize = extenddims(size(myar), cI.I, 1)
+        size_inferred = map(a.startinds, size(a), cI.I) do si, sa, ii
+            ii == length(si) ? sa - si[ii] + 1 : si[ii+1] - si[ii]
+        end
+        mysize = extenddims(size_inferred, cI.I, 1)
         array_range = map(cI.I, a.startinds, mysize, inds) do ii, si, ms, indstoread
             max(first(indstoread) - si[ii] + 1, 1):min(last(indstoread) - si[ii] + 1, ms)
         end
         outer_range = map(cI.I, a.startinds, array_range, inds) do ii, si, ar, indstoread
             (first(ar)+si[ii]-first(indstoread)):(last(ar)+si[ii]-first(indstoread))
         end
         #Shorten array range to shape of actual array
-        array_range = map((i, j) -> j, size(myar), array_range)
+        array_range = map((i, j) -> j, size_inferred, array_range)
         outer_range = fix_outerrangeshape(outer_range, array_range)
-        f(outer_range, array_range, cI)
+        if ismissing(myar)
+            f(outer_range, array_range, missing)
+        else
+            f(outer_range, array_range, cI)
+        end
     end
 end
 fix_outerrangeshape(outer_range, array_range) = fix_outerrangeshape((), outer_range, array_range)
@@ -126,6 +172,7 @@ function concat_chunksize(parents)
     newchunks = map(s -> Vector{Union{RegularChunks,IrregularChunks}}(undef, s), size(parents))
     for i in CartesianIndices(parents)
         array = parents[i]
+        ismissing(array) && continue
         chunks = eachchunk(array)
         foreach(chunks.chunks, i.I, newchunks) do c, ind, newc
             if !isassigned(newc, ind)

src/diskarray.jl

@@ -28,7 +28,7 @@ function readblock!(a::AbstractArray, aout, r...)
     if isdisk(a)
         @warn "Using fallback readblock! for array $(typeof(a)). This should not happen but there should be a custom implementation."
     end
-    copyto!(aout, CartesianIndices(aout), a, CartesianIndices(r))
+    aout .= view(a, CartesianIndices(r))
 end
 
 """

test/runtests.jl

@@ -491,6 +491,33 @@ end
         @test slic isa Vector{Float64}
         @test slic == Float64[1, 2, 3, 4, 1, 2, 3, 4]
     end
+
+    @testset "Concat DiskArray with missing tiles" begin
+        a = zeros(Int, 3, 4)
+        b = ones(Int, 2, 4)
+        c = fill(2, 3, 5)
+        d = fill(missing, 2, 5)
+        aconc = DiskArrays.ConcatDiskArray(reshape([a, b, c, missing], 2, 2))
+        abase = [a c; b d]
+        @test all(isequal.(aconc[:, :], abase))
+        @test all(isequal.(aconc[3:4, 4:6], abase[3:4, 4:6]))
+        ch = DiskArrays.eachchunk(aconc)
+        @test ch.chunks[1] == [1:3, 4:5]
+        @test ch.chunks[2] == [1:4, 5:9]
+
+        a = ones(100, 50)
+        b = [rem(i.I[3], 5) == 0 ? missing : a for i in CartesianIndices((1, 1, 100))]
+        b[1] = missing
+        a_conc = DiskArrays.ConcatDiskArray(b)
+        ch = eachchunk(a_conc)
+        @test ch.chunks[1] == [1:100]
+        @test ch.chunks[2] == [1:50]
+        @test ch.chunks[3] === DiskArrays.RegularChunks(1, 0, 100)
+
+        @test all(isequal.(a_conc[2, 2, 1:5], [missing, 1.0, 1.0, 1.0, missing]))
+        @test all(isequal.(a_conc[end, end, 95:100], [missing, 1.0, 1.0, 1.0, 1.0, missing]))
+
+    end
 end
 
 @testset "Broadcast with length 1 and 0 final dim" begin
@@ -979,7 +1006,7 @@ end
 @testset "Padded disk arrays" begin
     M = (1:100) * (1:120)'
     A = cat(M, 2M, 3M, 4M; dims=3)
-    ch = ChunkedDiskArray(A, (128, 128, 2)) 
+    ch = ChunkedDiskArray(A, (128, 128, 2))
     pa = DiskArrays.pad(ch, ((10, 20), (30, 40), (1, 2)); fill=999)
     @test size(pa) == (130, 190, 7)
     # All outside
@@ -1009,9 +1036,9 @@ end
         @test DiskArrays._pad_offset(c1, (10, 10)) == DiskArrays.RegularChunks(10, 0, 120)
         @test DiskArrays._pad_offset(c1, (0, 0)) == c1
 
-        c2 = DiskArrays.IrregularChunks(chunksizes = [10, 10, 20, 30, 40])
+        c2 = DiskArrays.IrregularChunks(chunksizes=[10, 10, 20, 30, 40])
         #The following test would assume padding ends up in a separate chunk:
-        @test DiskArrays._pad_offset(c2, (5, 5)) == DiskArrays.IrregularChunks(chunksizes = [5, 10, 10, 20, 30, 40, 5])
+        @test DiskArrays._pad_offset(c2, (5, 5)) == DiskArrays.IrregularChunks(chunksizes=[5, 10, 10, 20, 30, 40, 5])
         @test DiskArrays._pad_offset(c2, (0, 0)) == c2
     end
 end
@@ -1054,10 +1081,10 @@ end
 end
 
 @testset "identity function" begin
-    a = ChunkedDiskArray(1:10 .> 3; chunksize=(3, ))
-	for fname in [:sum, :prod, :all, :any, :minimum, :maximum, :count]
-		@eval out = @capture_out @trace $fname($a) DiskArrays
-		@test occursin("DiskGenerator", out) == false
+    a = ChunkedDiskArray(1:10 .> 3; chunksize=(3,))
+    for fname in [:sum, :prod, :all, :any, :minimum, :maximum, :count]
+        @eval out = @capture_out @trace $fname($a) DiskArrays
+        @test occursin("DiskGenerator", out) == false
     end
     @test count(a) + count(!, a) == length(a)
 end