Fix the DiskArrays backend implementation so CI passes

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/enginearrays.jl

@@ -3,8 +3,8 @@ using OnlineStats: OnlineStats
 using Statistics
 export engine, compute, compute!
 
-import DiskArrays: ComputeBackend
-struct DiskArrayEngineBackend <: ComputeBackend end
+# 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)
@@ -34,14 +34,26 @@ function collect_bcdims(A)
     last.(oc), first.(oc)
 end
 
-DiskArrays.diskarrays_mapreduce_impl(f, op, a, dims, init, ::DiskArrayEngineBackend) =
-    mapreduce_engine(f, op, a; dims, init)
+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[], xin...)
+(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)
@@ -99,13 +111,16 @@ wrap_reduction(a::OnlineStats.OnlineStat) = OnlineStats.value(a)
 
 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
-        function DiskArrays.$(fname)(ff, a, ::DiskArrayEngineBackend; dims=:, skipmissing=false, strategy=:auto)
+        # `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)
 
-                compute(res)
+                only(compute(res))
             else
                 dimspec = ntuple(i->i=>nothing, dims)
                 aggregate_diskarray(a, $func, dimspec; skipmissing, preproc=ff, strategy)

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/test_compute_backend.jl

@@ -1,5 +1,6 @@
 using DiskArrayEngine
-using DiskArrays: AbstractDiskArray, AccessCountDiskArray
+using DiskArrays: AbstractDiskArray, DiskArrayEngineBackend, withbackend
+using DiskArrays.TestTypes: AccessCountDiskArray
 using Statistics
 using Test
 
@@ -9,7 +10,8 @@ import DiskArrayEngine as DAE
 
 function make_arrays(data; chunksize=size(data))
     mat = data
-    da = AccessCountDiskArray(data; chunksize=chunksize)
+    # 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
 
@@ -38,6 +40,9 @@ function test_scalar_reductions(data; chunksize=ntuple(i -> max(1, size(data)[i]
             @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
@@ -90,7 +95,8 @@ function test_mapreduce(data; chunksize=ntuple(i -> max(1, size(data)[i] ÷ 2),
 
         @testset "mapreduce (no dims, no init)" begin
             @test mapreduce(x -> 2x, +, da) ≈ mapreduce(x -> 2x, +, mat)
-            @test mapreduce(*, -, da) ≈ mapreduce(*, -, 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
@@ -99,18 +105,21 @@ function test_mapreduce(data; chunksize=ntuple(i -> max(1, size(data)[i] ÷ 2),
         end
 
         @testset "mapreduce (init)" begin
-            @test Array(mapreduce(identity, +, da; init=0)) ≈ mapreduce(identity, +, mat; init=0)
-            @test Array(mapreduce(identity, *, da; init=1)) ≈ mapreduce(identity, *, mat; init=1)
-            @test Array(mapreduce(identity, +, da; dims=1, init=0)) ≈ mapreduce(identity, +, mat; dims=1, init=0)
+            @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
-            R = zeros(size(da, 1), size(da, 2), 1)
-            mapreducedim!(x -> 2x, +, R, da)
-            @test R ≈ mapreducedim!(x -> 2x, +, similar(R, size(da, 1), size(da, 2), 1), mat)
-            R2 = zeros(1, size(da, 2), size(da, 3))
-            mapreducedim!(x -> x^2, +, R2, da)
-            @test R2 ≈ mapreducedim!(x -> x^2, +, similar(R2), mat)
+            # 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/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[:,:,:]