[WIP] better performance in collect_similar

src/generator.jl

@@ -47,6 +47,41 @@ function Base.collect(itr::DiskGenerator{<:AbstractArray{<:Any,N}}) where {N}
     return dest
 end
 
+# Warning: this is not public API!
+function Base.collect_similar(A::AbstractArray, itr::DiskGenerator{<:AbstractArray{<:Any,N}}) where {N}
+    input = itr.iter # this is known to be an array
+    y = iterate(itr)
+    shp = axes(input)
+    if y === nothing
+        et = Base.@default_eltype(itr)
+        return similar(A, et, shp)
+    end
+    v1, st = y
+    dest = similar(A, typeof(v1), shp)# TODO: should this be `Base.return_type(itr.f, Tuple{eltype(input)})`?
+    i = y
+    # If the array is chunked, read each chunk and apply the function
+    # via broadcasting.
+    if DiskArrays.haschunks(input) isa DiskArrays.Chunked
+        chunks = eachchunk(input)
+        for chunk_inds in chunks
+            dest[chunk_inds...] .= itr.f.(input[chunk_inds...])
+        end
+    else # iterate as normal array
+        for I in eachindex(itr.iter)
+            if i isa Nothing # Mainly to keep JET clean 
+                error(
+                    "Should not be reached: iterator is shorter than its `eachindex` iterator"
+                )
+            else
+                dest[I] = first(i)
+                i = iterate(itr, last(i))
+            end
+        end
+    end
+    return dest
+
+end
+
 macro implement_generator(t)
     t = esc(t)
     quote

test/runtests.jl

@@ -953,3 +953,17 @@ end
     @test getindex_count(A) == 0 
 end
 
+@testset "Map over indices correctly" begin
+    # This is a regression test for issue #144
+    # `map` should always work over the correct indices,
+    # especially since we overload generators to `DiskArrayGenerator`.
+
+    data = [i+j for i in 1:200, j in 1:100]
+    da = AccessCountDiskArray(data, chunksize=(10,10))
+    @test map(identity, da) == data
+    @test all(map(identity, da) .== data)
+
+    # Make sure that type inference works
+    @inferred Matrix{Int} map(identity, da)
+    @inferred Matrix{Float64} map(x -> x * 5.0, da)
+end