Allow constructing ConcatDiskArray from mixed type arrays (#210)

* Improve concat diskarray

* Test mixed type concatenation

* Make cat of unchunked chunked

Author: meggart

src/cat.jl

@@ -4,12 +4,16 @@
 
 Joins multiple AbstractArrays or AbstractDiskArrays in lazy concatination.
 """
-struct ConcatDiskArray{T,N,P} <: AbstractDiskArray{T,N}
+struct ConcatDiskArray{T,N,P,C,HC} <: AbstractDiskArray{T,N}
     parents::P
     startinds::NTuple{N,Vector{Int}}
     size::NTuple{N,Int}
+    chunks::C
+    haschunks::HC
 end
-function ConcatDiskArray(arrays::AbstractArray{<:AbstractArray{T,N},M}) where {T,N,M}
+function ConcatDiskArray(arrays::AbstractArray{<:AbstractArray{<:Any,N},M}) where {N,M}
+    T = mapreduce(eltype,promote_type, init = eltype(first(arrays)),arrays)
+        
     function othersize(x, id)
         return (x[1:(id - 1)]..., x[(id + 1):end]...)
     end
@@ -51,12 +55,13 @@ function ConcatDiskArray(arrays::AbstractArray{<:AbstractArray{T,N},M}) where {T
     startinds = map(first, si)
     sizes = map(last, si)
 
-    return ConcatDiskArray{T,D,typeof(arrays1)}(arrays1, startinds, sizes)
+    chunks = concat_chunksize(D, 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 -> eltype(a) == eltype(first(arrays)), arrays) ||
-        error("Arrays don't have the same element type")
     all(a -> ndims(a) == ndims(first(arrays)), arrays) ||
         error("Arrays don't have the same dimensions")
     return error("Should not be reached")
@@ -98,11 +103,10 @@ function _concat_diskarray_block_io(f, a::ConcatDiskArray, inds...)
     end
 end
 
-haschunks(::ConcatDiskArray) = Chunked()
+haschunks(c::ConcatDiskArray) = c.haschunks
 
-function eachchunk(aconc::ConcatDiskArray{T,N}) where {T,N}
-    s = size(aconc)
-    oldchunks = map(eachchunk, aconc.parents)
+function concat_chunksize(N, parents)
+    oldchunks = map(eachchunk, parents)
     newchunks = ntuple(N) do i
         sliceinds = Base.setindex(ntuple(_ -> 1, N), :, i)
         v = map(c -> c.chunks[i], oldchunks[sliceinds...])
@@ -113,6 +117,10 @@ function eachchunk(aconc::ConcatDiskArray{T,N}) where {T,N}
     return GridChunks(newchunks...)
 end
 
+function eachchunk(aconc::ConcatDiskArray{T,N}) where {T,N}
+    aconc.chunks
+end
+
 function mergechunks(a::RegularChunks, b::RegularChunks)
     if a.s == 0 || (a.cs == b.cs && length(last(a)) == a.cs)
         RegularChunks(a.cs, a.offset, a.s + b.s)

test/runtests.jl

@@ -428,6 +428,15 @@ end
         @test sum(ca) == sum(0:23)
     end
 
+    @testset "Concatenation of unchunked arrays" begin
+        a = UnchunkedDiskArray(rand(10,20))
+        b = UnchunkedDiskArray(rand(15,20))
+        c = UnchunkedDiskArray(rand(18,20))
+        d = cat(a, b, c; dims=1)
+        @test DiskArrays.eachchunk(d) == [(1:10, 1:20); (11:25, 1:20); (26:43, 1:20);;]
+        @test DiskArrays.haschunks(d) isa DiskArrays.Chunked
+    end
+
     @testset "cat mixed chunk size" begin
         a = AccessCountDiskArray(collect(1:10); chunksize=(3,))
         b = AccessCountDiskArray(collect(1:9); chunksize=(4,))
@@ -450,6 +459,19 @@ end
         @test d == 1:26
         @test c == 20:26
     end
+
+    @testset "ConcatDiskArray works for mixed element types" begin
+        da1 = AccessCountDiskArray(collect(Float64,reshape(1:24, 4, 6, 1)))
+        da2 = AccessCountDiskArray(collect(Float32,reshape(1:24, 4, 6, 1)))
+        @test eltype(da1) <: Float64
+        @test eltype(da2) <: Float32
+        c = DiskArrays.ConcatDiskArray([da1, da2])
+
+        @test eltype(c) <: Float64
+        slic = c[:,1,1]
+        @test slic isa Vector{Float64}
+        @test slic == Float64[1, 2, 3, 4, 1, 2, 3, 4]
+    end
 end
 
 @testset "Broadcast with length 1 and 0 final dim" begin