Avoid allocation of temp array for dense data reads and writes (#161)

* introduce diskindex type

* Add Aliasing types, breaks inference

* revert alignment stuff

* don't allocate when output array can be used

* Add more getindex/setindex tests

* more tests

* more tests

* More tests (again)

* Remove unused methods

* Remove assume_effects

* Fix signature

* fix dispatch for breaking out stepranges

* bump julia compat to 1.9

* test on 1.9

* make 1.9 tests pass

* Add type for ChunkStrategy

Author: meggart

.github/workflows/ci.yml

@@ -15,6 +15,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
+          - '1.9'
           - '1' # Leave this line unchanged. '1' will automatically expand to the latest stable 1.x release of Julia.
           - 'nightly'
         os:

Project.toml

@@ -10,7 +10,7 @@ OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 [compat]
 Aqua = "0.8"
 LRUCache = "1"
-julia = "1.6"
+julia = "1.9"
 OffsetArrays = "1"
 Statistics = "1.9"
 Test = "1.9"

src/batchgetindex.jl

@@ -86,8 +86,8 @@ function is_sparse_index(ids; density_threshold = 0.5)
 end
 
 function process_index(i, cs, strategy::Union{ChunkRead,SubRanges})
-    outsize, tempsize, outinds,tempinds,datainds,cs = process_index(i,cs, NoBatch(strategy))
-    outsize, tempsize, ([outinds],), ([tempinds],), ([datainds],), cs
+    ii,cs = process_index(i,cs, NoBatch(strategy))
+    DiskIndex(ii.output_size, ii.temparray_size, ([ii.output_indices],), ([ii.temparray_indices],), ([ii.data_indices],)), cs
 end
 
 
@@ -111,7 +111,7 @@ function process_index(i::AbstractArray{<:Integer,N}, cs, ::ChunkRead) where N
         push!(tempinds, (tempind,))
         maxtempind = max(maxtempind,maximum(tempind))
     end
-    size(i), ((maxtempind),), (outinds,), (tempinds,), (datainds,), Base.tail(cs)
+    DiskIndex(size(i), ((maxtempind),), (outinds,), (tempinds,), (datainds,)), Base.tail(cs)
 end
 
 function find_subranges_sorted(inds,allow_steprange=false)
@@ -150,7 +150,11 @@ function find_subranges_sorted(inds,allow_steprange=false)
             end
         end
     end
-    push!(rangelist,inds[current_base]:last(inds))
+    if current_step == 1 || current_step == 0
+        push!(rangelist,inds[current_base]:last(inds))
+    else
+        push!(rangelist,inds[current_base]:current_step:last(inds))
+    end
     push!(outputinds,current_base:length(inds))
     rangelist, outputinds
 end
@@ -174,8 +178,8 @@ function process_index(i::AbstractArray{<:Integer,N}, cs, s::SubRanges) where N
             (r,)
         end
         outinds = tuple.(outputinds)
-        tempsize = maximum(length(rangelist))
-        (length(i),), (tempsize,), (outinds,), (tempinds,), (datainds,), Base.tail(cs)
+        tempsize = maximum(length,rangelist)
+        DiskIndex((length(i),), (tempsize,), (outinds,), (tempinds,), (datainds,)), Base.tail(cs)
     else
         p = mysortperm(i)
         i_sorted = view(i,p)
@@ -190,7 +194,7 @@ function process_index(i::AbstractArray{<:Integer,N}, cs, s::SubRanges) where N
             (view(p,oi),)
         end
         tempsize = maximum(length(rangelist))
-        size(i), (tempsize,), (outinds,), (tempinds,), (datainds,), Base.tail(cs)
+        DiskIndex(size(i), (tempsize,), (outinds,), (tempinds,), (datainds,)), Base.tail(cs)
     end
 end
 function process_index(i::AbstractArray{Bool,N}, cs, cr::ChunkRead) where N
@@ -199,8 +203,14 @@ end
 function process_index(i::AbstractArray{Bool,N}, cs, cr::SubRanges) where N
     process_index(findall(i),cs,cr)
 end
-function process_index(i::StepRange{<:Integer}, cs, ::ChunkStrategy{CanStepRange})
-    (length(i),), (length(i),), (Colon(),), (Colon(),), (i,), Base.tail(cs)
+function process_index(i::StepRange{<:Integer}, cs, ::ChunkRead{CanStepRange})
+    DiskIndex((length(i),), (length(i),), ([(Colon(),)],), ([(Colon(),)],), ([(i,)],)), Base.tail(cs)
+end
+function process_index(i::StepRange{<:Integer}, cs, ::SubRanges{CanStepRange})
+    DiskIndex((length(i),), (length(i),), ([(Colon(),)],), ([(Colon(),)],), ([(i,)],)), Base.tail(cs)
+end
+function process_index(i::StepRange{<:Integer}, cs, ::NoBatch{CanStepRange})
+    DiskIndex((length(i),), (length(i),), (Colon(),), (Colon(),), (i,)), Base.tail(cs)
 end
 function process_index(i::AbstractArray{<:CartesianIndex{N},M}, cs, ::Union{ChunkRead,SubRanges}) where {N,M}
     csnow, csrem = splitcs(i,cs)
@@ -217,14 +227,16 @@ function process_index(i::AbstractArray{<:CartesianIndex{N},M}, cs, ::Union{Chun
     for (cI,a) in chunksdict
         datamin,datamax = extrema(first,a)
         aa = first.(a)
-        tempind = aa .- datamin .+ oneunit(CartesianIndex{N})
+        tempind = map(aa) do ind
+            ind - datamin + oneunit(CartesianIndex{N})
+        end
         push!(outinds, tuple(map(last,a)))
         push!(datainds, range.(datamin.I,datamax.I))
         push!(tempinds, tuple(tempind))
         s = datamax.I .- datamin.I .+ 1
         tempsize = max.(s,tempsize)
     end
-    size(i), tempsize, (outinds,), (tempinds,), (datainds,), csrem
+    DiskIndex(size(i), tempsize, (outinds,), (tempinds,), (datainds,)), csrem
 end
 
 

src/diskarray.jl

@@ -40,12 +40,10 @@ Determines a list of tuples used to perform the read or write operations. The re
 - `temp_indices` indices for reading from temp array
 - `data_indices` indices for reading from data array
 """
-Base.@assume_effects :removable resolve_indices(a,i) = resolve_indices(a,i,batchstrategy(a))
-Base.@assume_effects :removable resolve_indices(a, i, batch_strategy) = _resolve_indices(eachchunk(a).chunks, i, (), (), (), (), (), batch_strategy)
-Base.@assume_effects :removable resolve_indices(a::AbstractVector, i::Tuple{AbstractVector{<:Integer}}, batch_strategy::NoBatch) = _resolve_indices(eachchunk(a).chunks, i, (), (), (), (), (), batch_strategy)
-Base.@assume_effects :removable resolve_indices(a::AbstractVector, i::Tuple{AbstractVector{<:Integer}}, batch_strategy::ChunkRead) = _resolve_indices(eachchunk(a).chunks, i, (), (), (), (), (), batch_strategy)
-Base.@assume_effects :removable resolve_indices(a::AbstractVector, i::Tuple{AbstractVector{<:Integer}}, batch_strategy::SubRanges) = _resolve_indices(eachchunk(a).chunks, i, (), (), (), (), (), batch_strategy)
-resolve_indices(a, ::Tuple{Colon}, _) = (length(a),), size(a), (Colon(),), (Colon(),), map(s->1:s,size(a))
+resolve_indices(a,i) = resolve_indices(a,i,batchstrategy(a))
+resolve_indices(a, i, batch_strategy) = _resolve_indices(eachchunk(a).chunks, i, DiskIndex((),(),(),(),()), batch_strategy)
+resolve_indices(a::AbstractVector, i::Tuple{AbstractVector{<:Integer}}, batch_strategy) = _resolve_indices(eachchunk(a).chunks, i, DiskIndex((), (), (), (), ()), batch_strategy)
+resolve_indices(a, ::Tuple{Colon}, _) = DiskIndex((length(a),), size(a), (Colon(),), (Colon(),), map(s->1:s,size(a)))
 resolve_indices(a, i::Tuple{<:CartesianIndex}, batch_strategy=NoBatch()) =
     resolve_indices(a, only(i).I, batch_strategy)
 function resolve_indices(a, i::Tuple{<:AbstractVector{<:Integer}}, batchstrategy)
@@ -73,49 +71,65 @@ function need_batch_index(i, cs, batchstrat)
     nb = (allow_multi || has_chunk_gap(approx_chunksize.(csnow), i)) && is_sparse_index(i; density_threshold=density_thresh)
     nb, csrem
 end
-function _resolve_indices(cs, i, output_size, temp_sizes, output_indices, temp_indices, data_indices, nb)
+
+
+struct DiskIndex{N,M,A<:Tuple,B<:Tuple,C<:Tuple}
+    output_size::NTuple{N,Int}
+    temparray_size::NTuple{M,Int}
+    output_indices::A
+    temparray_indices::B
+    data_indices::C
+end
+@inline function merge_index(a::DiskIndex,b::DiskIndex)
+    DiskIndex(
+        (a.output_size...,b.output_size...),
+        (a.temparray_size...,b.temparray_size...),
+        (a.output_indices...,b.output_indices...),
+        (a.temparray_indices...,b.temparray_indices...),
+        (a.data_indices...,b.data_indices...),
+    )
+end
+
+function _resolve_indices(cs, i, indices_pre::DiskIndex, nb::ChunkStrategy)
     inow = first(i)
-    outsize, tempsize, outinds, tempinds, datainds, cs = process_index(inow, cs, nb)
-    output_size = (output_size..., outsize...)
-    output_indices = (output_indices..., outinds...)
-    temp_sizes = (temp_sizes..., tempsize...)
-    temp_indices = (temp_indices..., tempinds...)
-    data_indices = (data_indices..., datainds...)
-    _resolve_indices(cs, Base.tail(i), output_size, temp_sizes, output_indices, temp_indices, data_indices, nb)
-end
-_resolve_indices(::Tuple{}, ::Tuple{}, output_size, temp_sizes, output_indices, temp_indices, data_indices, nb) = output_size, temp_sizes, output_indices, temp_indices, data_indices
+    indices_new, cs_rem = process_index(inow, cs, nb)
+    _resolve_indices(cs_rem, Base.tail(i), merge_index(indices_pre,indices_new), nb)
+end
+_resolve_indices(::Tuple{}, ::Tuple{}, indices::DiskIndex, nb::ChunkStrategy) = indices
 #No dimension left in array, only singular indices allowed
-function _resolve_indices(::Tuple{}, i, output_size, temp_sizes, output_indices, temp_indices, data_indices, nb)
+function _resolve_indices(::Tuple{}, i, indices_pre::DiskIndex, nb::ChunkStrategy)
     inow = first(i)
     (length(inow) == 1 && only(inow) == 1) || throw(ArgumentError("Trailing indices must be 1"))
-    output_size = (output_size..., size(inow)...)
-    output_indices = (output_indices..., size(inow)...)
-    _resolve_indices((), Base.tail(i), output_size, temp_sizes, output_indices, temp_indices, data_indices, nb)
+    indices_new = DiskIndex(size(inow),(),size(inow),(),())
+    indices = merge_index(indices_pre,indices_new)
+    _resolve_indices((), Base.tail(i), indices, nb)
 end
 #Still dimensions left, but no indices available
-function _resolve_indices(cs, ::Tuple{}, output_size, temp_sizes, output_indices, temp_indices, data_indices, nb)
+function _resolve_indices(cs, ::Tuple{}, indices_pre::DiskIndex, nb::ChunkStrategy) 
     csnow = first(cs)
     arraysize_from_chunksize(csnow) == 1 || throw(ArgumentError("Indices can only be omitted for trailing singleton dimensions"))
-    data_indices = (data_indices..., 1:1)
-    temp_sizes = (temp_sizes..., 1)
-    temp_indices = (temp_indices..., 1)
-    _resolve_indices(Base.tail(cs), (), output_size, temp_sizes, output_indices, temp_indices, data_indices, nb)
+    indices_new = add_dimension_index(nb)
+    indices = merge_index(indices_pre,indices_new)
+    _resolve_indices(Base.tail(cs), (), indices, nb)
 end
 
+add_dimension_index(::NoBatch) = DiskIndex((),(1,),(),(1,),(1:1,))
+add_dimension_index(::Union{ChunkRead,SubRanges}) = DiskIndex((),(1,),([()],),([(1,)],),([(1:1,)],))
+
 
 #outsize, tempsize, outinds,tempinds,datainds,cs
 process_index(i, cs, ::NoBatch) = process_index(i, cs)
-process_index(inow::Integer, cs) = ((), 1, (), (1,), (inow:inow,), Base.tail(cs))
+process_index(inow::Integer, cs) = DiskIndex((), (1,), (), (1,), (inow:inow,)), Base.tail(cs)
 function process_index(::Colon, cs)
     s = arraysize_from_chunksize(first(cs))
-    (s,), (s,), (Colon(),), (Colon(),), (1:s,), Base.tail(cs)
+    DiskIndex((s,), (s,), (Colon(),), (Colon(),), (1:s,),), Base.tail(cs)
 end
-function process_index(i::AbstractUnitRange, cs)
-    (length(i),), (length(i),), (Colon(),), (Colon(),), (i,), Base.tail(cs)
+function process_index(i::AbstractUnitRange{<:Integer}, cs, ::NoBatch)
+    DiskIndex((length(i),), (length(i),), (Colon(),), (Colon(),), (i,)), Base.tail(cs)
 end
 function process_index(i::AbstractArray{<:Integer}, cs, ::NoBatch)
     indmin, indmax = extrema(i)
-    size(i), ((indmax - indmin + 1),), map(_->Colon(),size(i)), ((i .- (indmin - 1)),), (indmin:indmax,), Base.tail(cs)
+    DiskIndex(size(i), ((indmax - indmin + 1),), map(_->Colon(),size(i)), ((i .- (indmin - 1)),), (indmin:indmax,)), Base.tail(cs)
 end
 function process_index(i::AbstractArray{Bool,N}, cs, ::NoBatch) where {N}
     csnow, csrem = splitcs(i, cs)
@@ -124,7 +138,7 @@ function process_index(i::AbstractArray{Bool,N}, cs, ::NoBatch) where {N}
     indmin, indmax = cindmin.I, cindmax.I
     tempsize = indmax .- indmin .+ 1
     tempinds = view(i, range.(indmin, indmax)...)
-    (sum(i),), tempsize, (Colon(),), (tempinds,), range.(indmin, indmax), csrem
+    DiskIndex((sum(i),), tempsize, (Colon(),), (tempinds,), range.(indmin, indmax)), csrem
 end
 function process_index(i::AbstractArray{<:CartesianIndex{N}}, cs, ::NoBatch) where {N}
     csnow, csrem = splitcs(i, cs)
@@ -133,14 +147,14 @@ function process_index(i::AbstractArray{<:CartesianIndex{N}}, cs, ::NoBatch) whe
     indmin, indmax = cindmin.I, cindmax.I
     tempsize = indmax .- indmin .+ 1
     tempoffset = cindmin - oneunit(cindmin)
-    tempinds = i .- tempoffset
+    tempinds = i .- (CartesianIndex(tempoffset),)
     outinds = map(_->Colon(),size(i))
-    size(i), tempsize, outinds, (tempinds,), range.(indmin, indmax), csrem
+    DiskIndex(size(i), tempsize, outinds, (tempinds,), range.(indmin, indmax)), csrem
 end
 function process_index(i::CartesianIndices{N}, cs, ::NoBatch) where {N}
     _, csrem = splitcs(i, cs)
     cols = map(_ -> Colon(), i.indices)
-    length.(i.indices), length.(i.indices), cols, cols, i.indices, csrem
+    DiskIndex(length.(i.indices), length.(i.indices), cols, cols, i.indices), csrem
 end
 splitcs(i::AbstractArray{<:CartesianIndex}, cs) = splitcs(first(i).I, (), cs)
 splitcs(i::AbstractArray{Bool}, cs) = splitcs(size(i), (), cs)
@@ -149,6 +163,20 @@ splitcs(_, cs) = (first(cs),), Base.tail(cs)
 splitcs(si, csnow, csrem) = splitcs(Base.tail(si), (csnow..., first(csrem)), Base.tail(csrem))
 splitcs(::Tuple{}, csnow, csrem) = (csnow, csrem)
 
+#Determine wether output and temp array can a) be identical b) share memory through reshape or 
+# c) need to be allocated individually
+function output_aliasing(di::DiskIndex)
+    if all(i->isa(i,Union{Int,AbstractUnitRange,Colon}),di.temparray_indices) && 
+        all(i->isa(i,Union{Int,AbstractUnitRange,Colon}),di.output_indices)
+        if di.output_size == di.temparray_size
+            return :identical
+        else 
+            return :reshapeoutput
+        end
+    else
+        return :noalign
+    end
+end
 
 
 function getindex_disk(a, i::Union{Integer,CartesianIndex}...)
@@ -162,6 +190,14 @@ function getindex_disk(a, i::Union{Integer,CartesianIndex}...)
     readblock!(a, outputarray, j...)
     only(outputarray)
 end
+function getindex_disk(a, i::Integer)
+    checkscalar(i)
+    checkbounds(a,i)
+    outputarray = Array{eltype(a)}(undef, map(_ -> 1, size(a))...)
+    j = map(k->k:k,CartesianIndices(a)[i].I)
+    readblock!(a, outputarray, j...)
+    only(outputarray)
+end
 
 function create_outputarray(out, a, output_size)
     size(out) == output_size || throw(ArgumentError("Expected output array size of $output_size"))
@@ -171,17 +207,13 @@ create_outputarray(::Nothing, a, output_size) = Array{eltype(a)}(undef, output_s
 
 getindex_disk(a, i...) = getindex_disk!(nothing, a, i...)
 
-function _getindex_do_rest(out,a,output_size, temparray_size, output_indices, temparray_indices, data_indices)
-    
-end
-
 function getindex_disk_batch!(out,a,i)
-    output_size, temparray_size, output_indices, temparray_indices, data_indices = resolve_indices(a, i)
-    moutput_indices = MRArray(output_indices)
-    mtemparray_indices = MRArray(temparray_indices)
-    mdata_indicess = MRArray(data_indices)
-    outputarray = create_outputarray(out, a, output_size)
-    temparray = Array{eltype(a)}(undef, temparray_size...)
+    indices = resolve_indices(a, i)
+    moutput_indices = MRArray(indices.output_indices)
+    mtemparray_indices = MRArray(indices.temparray_indices)
+    mdata_indicess = MRArray(indices.data_indices)
+    outputarray = create_outputarray(out, a, indices.output_size)
+    temparray = Array{eltype(a)}(undef, indices.temparray_size...)
     for ii in eachindex(moutput_indices)
         data_indices = mdata_indicess[ii]
         output_indices = moutput_indices[ii]
@@ -194,12 +226,20 @@ function getindex_disk_batch!(out,a,i)
 end
 
 function getindex_disk_nobatch!(out,a,i)
-    output_size, temparray_size, output_indices, temparray_indices, data_indices = resolve_indices(a, i, NoBatch(allow_steprange(a), 1.0))
+    indices = resolve_indices(a, i, NoBatch(allow_steprange(a), 1.0))
     #@debug output_size, temparray_size, output_indices, temparray_indices, data_indices
-    outputarray = create_outputarray(out, a, output_size)
-    temparray = Array{eltype(a)}(undef, temparray_size...)
-    readblock!(a, temparray, data_indices...)
-    transfer_results!(outputarray, temparray, output_indices, temparray_indices)
+    outputarray = create_outputarray(out, a, indices.output_size)
+    outalias = output_aliasing(indices)
+    if outalias === :identical
+        readblock!(a, outputarray, indices.data_indices...)
+    elseif outalias === :reshapeoutput
+        temparray = reshape(outputarray,indices.temparray_size)
+        readblock!(a, temparray, indices.data_indices...)
+    else
+        temparray = Array{eltype(a)}(undef, indices.temparray_size...)
+        readblock!(a, temparray, indices.data_indices...)
+        transfer_results!(outputarray, temparray, indices.output_indices, indices.temparray_indices)
+    end
     outputarray
 end
 
@@ -243,11 +283,11 @@ end
 
 function setindex_disk_batch!(a,v,i)
     batch_strategy = batchstrategy(a)
-    output_size, temparray_size, output_indices, temparray_indices, data_indices = resolve_indices(a, i, batch_strategy)
-    moutput_indices = MRArray(output_indices)
-    mtemparray_indices = MRArray(temparray_indices)
-    mdata_indicess = MRArray(data_indices)
-    temparray = Array{eltype(a)}(undef, temparray_size...)
+    indices = resolve_indices(a, i, batch_strategy)
+    moutput_indices = MRArray(indices.output_indices)
+    mtemparray_indices = MRArray(indices.temparray_indices)
+    mdata_indicess = MRArray(indices.data_indices)
+    temparray = Array{eltype(a)}(undef, indices.temparray_size...)
     for (output_indices, temparray_indices, data_indices) in zip(moutput_indices, mtemparray_indices, mdata_indicess)
         transfer_results_write!(v, temparray, output_indices, temparray_indices)
         vtemparray = maybeshrink(temparray, a, data_indices)
@@ -256,10 +296,18 @@ function setindex_disk_batch!(a,v,i)
 end
 
 function setindex_disk_nobatch!(a,v,i)
-    output_size, temparray_size, output_indices, temparray_indices, data_indices = resolve_indices(a, i, NoBatch())
-    temparray = Array{eltype(a)}(undef, temparray_size...)
-    transfer_results_write!(v, temparray, output_indices, temparray_indices)
-    writeblock!(a, temparray, data_indices...)
+    indices = resolve_indices(a, i, NoBatch())
+    outalias = output_aliasing(indices)
+    if outalias === :identical
+        writeblock!(a, v, indices.data_indices...)
+    elseif outalias === :reshapeoutput
+        temparray = reshape(v,indices.temparray_size)
+        writeblock!(a, temparray, indices.data_indices...)
+    else
+        temparray = Array{eltype(a)}(undef, indices.temparray_size...)
+        transfer_results_write!(v, temparray, indices.output_indices, indices.temparray_indices)
+        writeblock!(a, temparray, indices.data_indices...)
+    end
 end
 
 function setindex_disk!(a::AbstractArray, v::AbstractArray, i...)
@@ -268,6 +316,7 @@ function setindex_disk!(a::AbstractArray, v::AbstractArray, i...)
     else
         setindex_disk_nobatch!(a,v,i)
     end
+    v
 end
 
 "Find the indices of elements containing integers in a Tuple"

src/util/testtypes.jl

@@ -34,7 +34,7 @@ DiskArrays.haschunks(::AccessCountDiskArray) = DiskArrays.Chunked()
 DiskArrays.eachchunk(a::AccessCountDiskArray) = DiskArrays.GridChunks(a, a.chunksize)
 function DiskArrays.readblock!(a::AccessCountDiskArray, aout, i::OrdinalRange...)
     ndims(a) == length(i) || error("Number of indices is not correct")
-    foreach(i) do r  
+    foreach(i) do r
         isa(r,AbstractUnitRange) || DiskArrays.allow_steprange(a) || error("StepRange passed although trait is false")
     end
     # println("reading from indices ", join(string.(i)," "))