Use the DiskArrayEngineBackend type defined in DiskArrays

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#313 is merged and released
+[sources]
+DiskArrays = {url = "https://github.com/JuliaIO/DiskArrays.jl", rev = "as/static_backends"}
+
 [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)
@@ -100,7 +100,8 @@ wrap_reduction(a::OnlineStats.OnlineStat) = OnlineStats.value(a)
 for func in (:maximum, :minimum, :sum, :extrema, :median, :mean)
     fname = Symbol("diskarrays_$(func)_impl")
     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=:auto) where {A<:AbstractDiskArray}
             if dims === Colon()
                 dimspec = ntuple(i->i=>nothing, ndims(a))
                 res = aggregate_diskarray(a, $func, dimspec; skipmissing, preproc=ff, strategy)

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
 
@@ -101,16 +103,19 @@ function test_mapreduce(data; chunksize=ntuple(i -> max(1, size(data)[i] ÷ 2),
         @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)
+            # 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