Implement a DiskArrays compute backend

Project.toml

@@ -26,6 +26,10 @@ Zarr = "0a941bbe-ad1d-11e8-39d9-ab76183a1d99"
 [weakdeps]
 Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54"
 
+# TODO remove once JuliaIO/DiskArrays.jl#315 is merged and released
+[sources]
+DiskArrays = {url = "https://github.com/JuliaIO/DiskArrays.jl", rev = "as/backend_fallbacks"}
+
 [extensions]
 DaggerExt = "Dagger"
 

src/diskarrayresults.jl

@@ -2,7 +2,7 @@ export results_as_diskarrays
 using DiskArrays: AbstractDiskArray, RegularChunks
 using OffsetArrays: OffsetArray
 
-struct GMWOPResult{T,N,G<:GMDWop,CS,ISPEC} <: AbstractEngineArray{T,N}
+struct GMWOPResult{T,N,G<:GMDWop,CS,ISPEC} <: AbstractDiskArray{T,N}
   op::G
   ires::Val{ISPEC}
   chunksize::CS

src/enginearrays.jl

@@ -2,30 +2,9 @@ using DiskArrays
 using OnlineStats: OnlineStats
 using Statistics
 export engine, compute, compute!
-abstract type AbstractEngineArray{T,N} <: AbstractDiskArray{T,N} end
-struct EngineArray{T,N,P} <: AbstractEngineArray{T,N}
-    parent::P
-    bcdims::Base.RefValue{NTuple{N,Int}}
-end
-EngineArray(p) = EngineArray{eltype(p),ndims(p),typeof(p)}(p,Ref(ntuple(identity,ndims(p))))
-unengine(a::EngineArray) = a.parent
-unengine(a) = a
-engine(p;bcdims=ntuple(identity,ndims(p))) = EngineArray{eltype(p),ndims(p),typeof(p)}(p,Ref(bcdims))
-engine(p::EngineArray;bcdims=ntuple(identity,ndims(p))) = EngineArray{eltype(p),ndims(p),typeof(p)}(p.parent,bcdims)
-function engine(p::AbstractEngineArray;bcdims=nothing)
-    if bcdims === nothing
-        p
-    else
-        bcdims=ntuple(identity,ndims(p))
-        EngineArray{eltype(p),ndims(p),typeof(p)}(p,Ref(bcdims))
-    end
-end
-Base.parent(p::EngineArray) = p.parent
-DiskArrays.readblock!(a::EngineArray,xout,r::OrdinalRange...) = DiskArrays.readblock!(a.parent,xout,r...)
-DiskArrays.writeblock!(a::EngineArray,xout,r::OrdinalRange...) = DiskArrays.writeblock!(a.parent,xout,r...)
-Base.size(a::EngineArray) = size(a.parent)
-DiskArrays.eachchunk(a::EngineArray) = DiskArrays.eachchunk(a.parent)
-bcdims(a::EngineArray) = a.bcdims[]
+
+# The backend type is owned by DiskArrays, we only add methods for it
+import DiskArrays: DiskArrayEngineBackend
 bcdims(p) = ntuple(identity,ndims(p))
 
 function collect_bcdims(A)
@@ -55,7 +34,29 @@ function collect_bcdims(A)
     last.(oc), first.(oc)
 end
 
-function Base.mapreduce(f, op, A::AbstractEngineArray...; dims=:, init = nothing, fin = identity)
+function DiskArrays.diskarrays_mapreduce_impl(f, op, a, dims, init, ::DiskArrayEngineBackend)
+    res = mapreduce_engine(f, op, a; dims, init=neutral_element(f, op, a, init))
+    # Like in Base, a reduction over all dimensions returns a scalar
+    dims === Colon() ? only(compute(res)) : res
+end
+
+# The output buffers are initialized with `init`, so without one we need a neutral element of `op`
+neutral_element(f, op, a, init) = init
+function neutral_element(f, op, a, ::Base._InitialValue)
+    T = Base.promote_op(f, eltype(a))
+    op === max && return typemin(T)
+    op === min && return typemax(T)
+    Base.reduce_empty(op, T)
+end
+
+struct _DAEReducer{F,R} <: Function
+    f::F
+    r::R
+end
+(d::_DAEReducer{F,R})(xout, xin...) where {F,R} = xout[] = d.r(xout[], d.f(xin...))
+
+
+function mapreduce_engine(f, op, A...; dims=:, init=nothing, fin=identity)
     s, bcd = collect_bcdims(A)
     nd = maximum(bcd)
     ia = map(A) do ar
@@ -68,7 +69,7 @@ function Base.mapreduce(f, op, A::AbstractEngineArray...; dims=:, init = nothing
                 fill(1,sloop)
             end
         end
-        InputArray(unengine,windows = lw, dimsmap = bcdims(ar))
+        InputArray(ar, windows=lw, dimsmap=bcdims(ar))
     end
     # tf = Base.promote_op(f,Base.nonmissingtype.(eltype.(A))...)
     # top = Base.promote_op(op,tf,tf)
@@ -77,13 +78,13 @@ function Base.mapreduce(f, op, A::AbstractEngineArray...; dims=:, init = nothing
     # end
     tf = Base.promote_op(f,eltype.(A)...)
     top = Base.promote_op(op,tf,tf)
-    func = create_userfunction(f,top,red=op,init=init,buftype=top,finalize=fin)
+    func = create_userfunction(_DAEReducer(f, op), top, red=op, init=init, buftype=top, finalize=fin, is_mutating=true, allow_threads=false)
     if dims === Colon()
         dims = ntuple(identity,nd)
     end
     windows = map(ntuple(identity,nd)) do idim
         if idim in dims
-            Repeated(s[idim])
+            Repeated(1, s[idim])
         else
             1:s[idim]
         end
@@ -108,25 +109,28 @@ extrred(x,y) = min(first(x),first(y)),max(last(x),last(y))
 wrap_reduction(a) = a
 wrap_reduction(a::OnlineStats.OnlineStat) = OnlineStats.value(a)
 
-for (f,red,i) in (
-    (:(Base.maximum),:max,:init_max),
-    (:(Base.minimum),:min,:init_min),
-    (:(Base.extrema),:extrred,:init_ex),
-    (:(Base.sum),:(+),:init_sum),
-    )
+for func in (:maximum, :minimum, :sum, :extrema, :median, :mean)
+    fname = Symbol("diskarrays_$(func)_impl")
+    # The direct aggregator broadcasts its result into the output, which does not work for the Tuple from `extrema`
+    default_strategy = QuoteNode(func === :extrema ? :reduce : :auto)
     eval(quote
-    $f(a::AbstractEngineArray;dims=:,skipmissing=false) = $f(identity,a;dims,skipmissing)
-    function $f(ff::Base.Callable,a::AbstractEngineArray;dims=:,skipmissing=false)
-        red = skipmissing ? missred($red) : $red
-        init = skipmissing ? missing : $i
-        mapreduce(ff,red,a;dims,init=$(i)(a))
-    end 
-
-
-end)
+        # `a::AbstractDiskArray` avoids an ambiguity with the `::ComputeBackend` fallback in DiskArrays
+        function DiskArrays.$(fname)(ff, a::A, ::DiskArrayEngineBackend; dims=:, skipmissing=false, strategy=$default_strategy) where {A<:AbstractDiskArray}
+            if dims === Colon()
+                dimspec = ntuple(i->i=>nothing, ndims(a))
+                res = aggregate_diskarray(a, $func, dimspec; skipmissing, preproc=ff, strategy)
+
+                only(compute(res))
+            else
+                dimspec = ntuple(i->i=>nothing, dims)
+                aggregate_diskarray(a, $func, dimspec; skipmissing, preproc=ff, strategy)
+            end
+        end
+    end)
 end
 
-Base.mapslices(f,A::AbstractEngineArray...;dims,outchunks=nothing) = mapslices_engine(f,A;dims,outchunks)
+diskarrays_mapslices_impl(f, A...; dims, outchunks=nothing) =
+    mapslices_engine(Base.Fix1(Statistics.median, f), A; dims, outchunks)
 
 function mapslices_engine(f,A...;dims,outchunks=nothing)
     s, bcd = collect_bcdims(A)
@@ -173,14 +177,14 @@ end
 
 get_infered_types(f,inputtypes) = Base.promote_op(f,inputtypes...)
 
-Statistics.median(a::AbstractEngineArray;dims=()) = mapslices_engine(median,a;dims)
+diskarrays_median_impl(f, a; dims=()) = mapslices_engine(median, a; dims)
 
 
 struct EngineStyle{N} <: Base.Broadcast.AbstractArrayStyle{N} end
 using DiskArrays: ChunkStyle
 using Base.Broadcast: DefaultArrayStyle
 
-Base.BroadcastStyle(::Type{<:AbstractEngineArray{<:Any,N}}) where N = EngineStyle{N}()
+diskarrays_broadcaststyle(::Type{T}, ::DiskArrayEngineBackend) where T = EngineStyle{ndims(T)}()
 Base.BroadcastStyle(::EngineStyle{N}, ::EngineStyle{M}) where {N,M} = EngineStyle{max(N, M)}()
 function Base.BroadcastStyle(::EngineStyle{N}, ::DefaultArrayStyle{M}) where {N,M}
     return EngineStyle{max(N, M)}()

src/util/aggregate.jl

@@ -83,28 +83,28 @@ This will never happen if you use a do-block.
 ```julia
     aggregate_diskarray(a,mean,(2=>nothing,3=>4,4=>[1,1,1,2,2,2,3,3,3,4,4,4,4,4,5,5,5]))
 """    
-function aggregate_diskarray(a, f, dimspec; skipmissing=false, strategy=:auto,outchunks=nothing)
+function aggregate_diskarray(a, f, dimspec; skipmissing=false, strategy=:auto, outchunks=nothing, preproc=identity)
 
     hasmissings = Missing <: eltype(a)
     if strategy == :reduce || ((isa(f,DataType) || isa(f,UnionAll)) && f <: OnlineStat)
-        agg = ReduceAggregator(disk_onlinestat(f))
+        agg = ReduceAggregator(disk_onlinestat(f, preproc))
         op = gmwop_for_aggregator(agg,dimspec,a;outchunks)
         results_as_diskarrays(op)[1]
     elseif strategy == :direct
         rett = Base.promote_op(f,Vector{Base.nonmissingtype(eltype(a))})
         if hasmissings
             rett = Union{rett,Missing}
         end
-        agg = DirectAggregator(create_userfunction(f,rett))
+        agg = DirectAggregator(create_userfunction(Base.Fix1(f, preproc), rett))
         op = gmwop_for_aggregator(agg,dimspec,a;outchunks)
         results_as_diskarrays(op)[1]
     elseif strategy == :auto
         rett = Base.promote_op(f,Vector{Base.nonmissingtype(eltype(a))})
         if hasmissings
             rett = Union{rett,Missing}
         end
-        agg1 = DirectAggregator(create_userfunction(f,rett))
-        agg2 = ReduceAggregator(disk_onlinestat(f))
+        agg1 = DirectAggregator(create_userfunction(Base.Fix1(f, preproc), rett))
+        agg2 = ReduceAggregator(disk_onlinestat(f, preproc))
 
         op1 = gmwop_for_aggregator(agg1,dimspec,a;outchunks)
         p1 = optimize_loopranges(op1,5e8)

src/util/onlinestats_helpers.jl

@@ -44,8 +44,9 @@ has_onlineversion(f) = f in keys(func_to_online)
 const func_to_online = Dict([
     mean => (OnlineStats.Mean,Union{Float64,Missing}),
     sum => (OnlineStats.Sum,Union{Float64,Missing}),
-    extrema => (OnlineStats.Extrema, Union{Tuple{Float64,Float64},Missing}),
-    maximum => (DerivedOnlineStat{OnlineStats.Extrema,v -> last(OnlineStats.value(v)),OnlineStats.fit!,()}, Union{Float64,Missing}),
-    minimum => (DerivedOnlineStat{OnlineStats.Extrema,v -> first(OnlineStats.value(v)),OnlineStats.fit!,()}, Union{Float64,Missing}),
+    # `value(::Extrema)` is a NamedTuple `(min, max, nmin, nmax)`, so use the accessors instead of `first`/`last`
+    extrema => (DerivedOnlineStat{OnlineStats.Extrema,extrema,OnlineStats.fit!,()}, Union{Tuple{Float64,Float64},Missing}),
+    maximum => (DerivedOnlineStat{OnlineStats.Extrema,maximum,OnlineStats.fit!,()}, Union{Float64,Missing}),
+    minimum => (DerivedOnlineStat{OnlineStats.Extrema,minimum,OnlineStats.fit!,()}, Union{Float64,Missing}),
     median => (DerivedOnlineStat{OnlineStats.ExpandingHist,OnlineStats.median,OnlineStats.fit!,(200,)}, Union{Float64,Missing})
 ])

test/runtests.jl

@@ -9,4 +9,5 @@ include("test_distribute.jl")
 include("test_restart.jl")
 include("test_graph.jl")
 include("test_mergeops.jl")
-include("test_interpolation.jl")
+include("test_interpolation.jl")
+include("test_compute_backend.jl")

test/test_compute_backend.jl

@@ -0,0 +1,154 @@
+using DiskArrayEngine
+using DiskArrays: AbstractDiskArray, DiskArrayEngineBackend, withbackend
+using DiskArrays.TestTypes: AccessCountDiskArray
+using Statistics
+using Test
+
+import DiskArrayEngine as DAE
+
+# ── Helpers ─────────────────────────────────────────────────────────────────
+
+function make_arrays(data; chunksize=size(data))
+    mat = data
+    # Select the backend per array, so the tests do not depend on the backend preference
+    da = withbackend(AccessCountDiskArray(data; chunksize=chunksize), DiskArrayEngineBackend())
+    return (materialized=mat, disk=da)
+end
+
+# ── Scalar reductions ───────────────────────────────────────────────────────
+
+function test_scalar_reductions(data; chunksize=ntuple(i -> max(1, size(data)[i] ÷ 2), ndims(data)))
+    @testset "DiskArrayEngine: scalar reductions" begin
+        mat, da = make_arrays(data; chunksize)
+
+        @testset "sum" begin
+            @test sum(da) ≈ sum(mat)
+            @test sum(identity, da) ≈ sum(mat)
+            @test sum(x -> 2x, da) ≈ sum(x -> 2x, mat)
+        end
+
+        @testset "prod" begin
+            @test prod(da) ≈ prod(mat)
+            @test prod(identity, da) ≈ prod(mat)
+            @test prod(x -> 2x, da) ≈ prod(x -> 2x, mat)
+        end
+
+        @testset "minimum / maximum" begin
+            @test minimum(da) ≈ minimum(mat)
+            @test minimum(identity, da) ≈ minimum(mat)
+            @test minimum(x -> abs(x), da) ≈ minimum(x -> abs(x), mat)
+            @test maximum(da) ≈ maximum(mat)
+            @test maximum(identity, da) ≈ maximum(mat)
+            @test maximum(x -> abs(x), da) ≈ maximum(x -> abs(x), mat)
+            # The OnlineStats path, `:auto` mostly picks the direct one for these small arrays
+            @test minimum(da; strategy=:reduce) ≈ minimum(mat)
+            @test maximum(da; strategy=:reduce) ≈ maximum(mat)
+        end
+
+        @testset "extrema" begin
+            @test extrema(da) == extrema(mat)
+            @test extrema(identity, da) == extrema(mat)
+            @test extrema(x -> abs(x), da) == extrema(x -> abs(x), mat)
+        end
+
+        @testset "count" begin
+            @test count(x -> x > 0, da) == count(x -> x > 0, mat)
+        end
+
+        @testset "mean" begin
+            @test mean(da) ≈ mean(mat)
+            @test mean(identity, da) ≈ mean(mat)
+            @test mean(x -> 2x, da) ≈ mean(x -> 2x, mat)
+        end
+
+        @testset "median" begin
+            @test median(da) ≈ median(mat)
+        end
+    end
+end
+
+# ── Lazy result verification ────────────────────────────────────────────────
+
+function test_lazy_results(data; chunksize=ntuple(i -> max(1, size(data)[i] ÷ 2), ndims(data)))
+    @testset "DiskArrayEngine: lazy results" begin
+        da = make_arrays(data; chunksize).disk
+
+        # sum/mean with dims≠: should return lazy GMWOPResult (DAE's impl
+        # calls aggregate_diskarray which returns a lazy array when dims is set).
+        # The DefaultBackend fallback would return a materialized Matrix, so
+        # checking the type proves the DAE impl path is exercised.
+        s = sum(da, dims=1)
+        @test s isa DAE.GMWOPResult
+        @test Array(s) ≈ sum(data, dims=1)
+
+        m = mean(da, dims=1)
+        @test m isa DAE.GMWOPResult
+        @test Array(m) ≈ mean(data, dims=1)
+    end
+end
+
+# ── mapreduce ───────────────────────────────────────────────────────────────
+
+function test_mapreduce(data; chunksize=ntuple(i -> max(1, size(data)[i] ÷ 2), ndims(data)))
+    @testset "DiskArrayEngine: mapreduce" begin
+        mat, da = make_arrays(data; chunksize)
+
+        @testset "mapreduce (no dims, no init)" begin
+            @test mapreduce(x -> 2x, +, da) ≈ mapreduce(x -> 2x, +, mat)
+            # `op` must be associative, chunks are reduced in an arbitrary order
+            @test mapreduce(abs, max, da) ≈ mapreduce(abs, max, mat)
+        end
+
+        @testset "mapreduce (dims=)" begin
+            @test Array(mapreduce(x -> 2x, +, da; dims=1)) ≈ mapreduce(x -> 2x, +, mat; dims=1)
+            @test Array(mapreduce(x -> 2x, +, da; dims=(1, 2))) ≈ mapreduce(x -> 2x, +, mat; dims=(1, 2))
+        end
+
+        @testset "mapreduce (init)" begin
+            @test mapreduce(identity, +, da; init=0) ≈ mapreduce(identity, +, mat; init=0)
+            @test mapreduce(identity, *, da; init=1) ≈ mapreduce(identity, *, mat; init=1)
+            # An Int `init` with float data and `dims` is an InexactError in Base as well
+            @test Array(mapreduce(identity, +, da; dims=1, init=0.0)) ≈ mapreduce(identity, +, mat; dims=1, init=0.0)
+        end
+
+        @testset "mapreducedim!" begin
+            # Reduce over the last and over the first dimension
+            redsize(d) = ntuple(i -> i == d ? 1 : size(mat, i), ndims(mat))
+            R = zeros(redsize(ndims(mat)))
+            Base.mapreducedim!(x -> 2x, +, R, da)
+            @test R ≈ Base.mapreducedim!(x -> 2x, +, zero(R), mat)
+            R2 = zeros(redsize(1))
+            Base.mapreducedim!(x -> x^2, +, R2, da)
+            @test R2 ≈ Base.mapreducedim!(x -> x^2, +, zero(R2), mat)
+        end
+
+        @testset "mapfoldl (no init)" begin
+            @test mapfoldl(x -> 2x, +, da) ≈ mapfoldl(x -> 2x, +, mat)
+        end
+
+        @testset "mapfoldl (with init)" begin
+            @test mapfoldl(identity, +, da; init=0) ≈ mapfoldl(identity, +, mat; init=0)
+            @test mapfoldl(identity, *, da; init=1) ≈ mapfoldl(identity, *, mat; init=1)
+        end
+    end
+end
+
+# ── Run tests over multiple array shapes ────────────────────────────────────
+
+@testset "DiskArrayEngine: scalar reductions" begin
+    for data in (randn(5, 4, 2), randn(10), randn(10, 20), randn(3, 3, 3))
+        test_scalar_reductions(data)
+    end
+end
+
+@testset "DiskArrayEngine: lazy results" begin
+    for data in (randn(5, 4, 2), randn(10), randn(10, 20), randn(3, 3, 3))
+        test_lazy_results(data)
+    end
+end
+
+@testset "DiskArrayEngine: mapreduce" begin
+    for data in (randn(5, 4, 2), randn(10), randn(10, 20), randn(3, 3, 3))
+        test_mapreduce(data)
+    end
+end

test/test_interpolation.jl

@@ -26,8 +26,7 @@ end
     @test size(agg_mean) == (1,5,6)
     @test agg_mean[:,:,:] == mean(a, dims=1)
     agg_max = aggregate_diskarray(a, maximum, (2=>nothing,), strategy=:reduce)
-    # This gives all ones for some reason
-    @test_broken agg_max[:,:,:] == maximum(a, dims=2)
+    @test agg_max[:,:,:] == maximum(a, dims=2)
     agg_sec = aggregate_diskarray(a, mean, (2=>2,))
     # This should work but currently throws a bounds error
     @test_throws BoundsError agg_sec[:,:,:]