From 8ce9de59130f8e52a5bf03ff4a56cbed0ff8052c Mon Sep 17 00:00:00 2001 From: Fabian Gans Date: Fri, 11 Sep 2026 17:47:19 +0200 Subject: [PATCH 1/6] First implementation of DiskArrays backend --- src/diskarrayresults.jl | 2 +- src/enginearrays.jl | 69 +++++++++++++++-------------------------- 2 files changed, 26 insertions(+), 45 deletions(-) diff --git a/src/diskarrayresults.jl b/src/diskarrayresults.jl index 91bb083..87eb9c1 100644 --- a/src/diskarrayresults.jl +++ b/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 diff --git a/src/enginearrays.jl b/src/enginearrays.jl index 09fd6b0..0b4be49 100644 --- a/src/enginearrays.jl +++ b/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[] + +import DiskArrays: ComputeBackend +struct DiskArrayEngineBackend <: ComputeBackend end bcdims(p) = ntuple(identity,ndims(p)) function collect_bcdims(A) @@ -55,7 +34,10 @@ function collect_bcdims(A) last.(oc), first.(oc) end -function Base.mapreduce(f, op, A::AbstractEngineArray...; dims=:, init = nothing, fin = identity) +DiskArrays.diskarrays_mapreduce_impl(f, op, a, dims, init, ::DiskArrayEngineBackend) = + mapreduce_engine(f, op, a; dims, init) + +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 +50,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) @@ -83,7 +65,7 @@ function Base.mapreduce(f, op, A::AbstractEngineArray...; dims=:, init = nothing 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 +90,24 @@ 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, red, i) in ( + (:(maximum), :max, :init_max), + (:(minimum), :min, :init_min), + (:(extrema), :extrred, :init_ex), + (:(sum), :(+), :init_sum), ) + fname = Symbol("diskarrays_$(func)_impl") 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) + function DiskArrays.$(fname)(ff, a, ::DiskArrayEngineBackend; dims=:, skipmissing=false) + red = skipmissing ? missred($red) : $red + init = skipmissing ? missing : $i + mapreduce(ff, red, a; dims, init=($(i)(a))) + 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 +154,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)}() From 4575933d1a80f1def1ed07c062a8b4ea668155bd Mon Sep 17 00:00:00 2001 From: Fabian Gans Date: Mon, 14 Sep 2026 15:30:34 +0200 Subject: [PATCH 2/6] remove engine array and implement DiskArrays backend --- src/enginearrays.jl | 30 +++++++++++++++++++----------- src/util/aggregate.jl | 10 +++++----- 2 files changed, 24 insertions(+), 16 deletions(-) diff --git a/src/enginearrays.jl b/src/enginearrays.jl index 0b4be49..131549b 100644 --- a/src/enginearrays.jl +++ b/src/enginearrays.jl @@ -37,6 +37,13 @@ end DiskArrays.diskarrays_mapreduce_impl(f, op, a, dims, init, ::DiskArrayEngineBackend) = mapreduce_engine(f, op, a; dims, init) +struct _DAEReducer{F,R} <: Function + f::F + r::R +end +(d::_DAEReducer{F,R})(xout, xin...) where {F,R} = xout[] = d.r(xout[], xin...) + + function mapreduce_engine(f, op, A...; dims=:, init=nothing, fin=identity) s, bcd = collect_bcdims(A) nd = maximum(bcd) @@ -59,7 +66,7 @@ function mapreduce_engine(f, op, A...; dims=:, init=nothing, fin=identity) # 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 @@ -90,18 +97,19 @@ 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 (func, red, i) in ( - (:(maximum), :max, :init_max), - (:(minimum), :min, :init_min), - (:(extrema), :extrred, :init_ex), - (:(sum), :(+), :init_sum), - ) +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) - red = skipmissing ? missred($red) : $red - init = skipmissing ? missing : $i - mapreduce(ff, red, a; dims, init=($(i)(a))) + function DiskArrays.$(fname)(ff, a, ::DiskArrayEngineBackend; dims=:, skipmissing=false, strategy=:auto) + if dims === Colon() + dimspec = ntuple(i->i=>nothing, ndims(a)) + res = aggregate_diskarray(a, $func, dimspec; skipmissing, preproc=ff, strategy) + + compute(res) + else + dimspec = ntuple(i->i=>nothing, dims) + aggregate_diskarray(a, $func, dimspec; skipmissing, preproc=ff, strategy) + end end end) end diff --git a/src/util/aggregate.jl b/src/util/aggregate.jl index b69e588..0df33c3 100644 --- a/src/util/aggregate.jl +++ b/src/util/aggregate.jl @@ -83,11 +83,11 @@ 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 @@ -95,7 +95,7 @@ function aggregate_diskarray(a, f, dimspec; skipmissing=false, strategy=:auto,ou 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 @@ -103,8 +103,8 @@ function aggregate_diskarray(a, f, dimspec; skipmissing=false, strategy=:auto,ou 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) From 62b15bcb372742a700e83ffecc53d8306794b363 Mon Sep 17 00:00:00 2001 From: Fabian Gans Date: Tue, 15 Sep 2026 11:12:14 +0200 Subject: [PATCH 3/6] Restructure tests --- test/runtests.jl | 3 +- test/test_compute_backend.jl | 145 +++++++++++++++++++++++++++++++++++ 2 files changed, 147 insertions(+), 1 deletion(-) create mode 100644 test/test_compute_backend.jl diff --git a/test/runtests.jl b/test/runtests.jl index 8556bfc..480f023 100644 --- a/test/runtests.jl +++ b/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") \ No newline at end of file +include("test_interpolation.jl") +include("test_compute_backend.jl") \ No newline at end of file diff --git a/test/test_compute_backend.jl b/test/test_compute_backend.jl new file mode 100644 index 0000000..3253806 --- /dev/null +++ b/test/test_compute_backend.jl @@ -0,0 +1,145 @@ +using DiskArrayEngine +using DiskArrays: AbstractDiskArray, AccessCountDiskArray +using Statistics +using Test + +import DiskArrayEngine as DAE + +# ── Helpers ───────────────────────────────────────────────────────────────── + +function make_arrays(data; chunksize=size(data)) + mat = data + da = AccessCountDiskArray(data; chunksize=chunksize) + 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) + 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) + @test mapreduce(*, -, da) ≈ mapreduce(*, -, 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 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) + 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) + 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 From 7bdabb6c830e7035437eb129120d1c29926a23fa Mon Sep 17 00:00:00 2001 From: Anshul Singhvi Date: Fri, 18 Sep 2026 12:15:52 -0400 Subject: [PATCH 4/6] Use the `DiskArrayEngineBackend` type defined in DiskArrays DiskArrays now owns the backend type (JuliaIO/DiskArrays.jl#313), so that it can be a member of its `DynamicBackend` sum type. We only add methods. - Type `a` in the reduction impls to resolve an ambiguity with the `::ComputeBackend` fallbacks in DiskArrays. - Tests select the backend per array with `withbackend` instead of relying on the backend preference, and import `AccessCountDiskArray` from `DiskArrays.TestTypes`. - Temporary `[sources]` entry pointing to the DiskArrays branch. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_018ra7EAJeFE1dtekwf4Tzxh --- Project.toml | 4 ++++ src/enginearrays.jl | 7 ++++--- test/test_compute_backend.jl | 6 ++++-- 3 files changed, 12 insertions(+), 5 deletions(-) diff --git a/Project.toml b/Project.toml index bbec3e0..cbbdb00 100644 --- a/Project.toml +++ b/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" diff --git a/src/enginearrays.jl b/src/enginearrays.jl index 131549b..6e38703 100644 --- a/src/enginearrays.jl +++ b/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) diff --git a/test/test_compute_backend.jl b/test/test_compute_backend.jl index 3253806..e233cb1 100644 --- a/test/test_compute_backend.jl +++ b/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 From d18414076c98df9e548fe79cd1fef1370777becf Mon Sep 17 00:00:00 2001 From: Anshul Singhvi Date: Sun, 20 Sep 2026 13:07:54 -0400 Subject: [PATCH 5/6] Fix mistakes in the compute backend tests - Qualify `Base.mapreducedim!`, size the output from `ndims` instead of assuming 3 dimensions, and start the reference from `zero(R)` instead of uninitialized memory. - `init=0.0` for float data, an Int `init` with `dims` throws in Base too. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_018ra7EAJeFE1dtekwf4Tzxh --- test/test_compute_backend.jl | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/test/test_compute_backend.jl b/test/test_compute_backend.jl index e233cb1..a1ffed3 100644 --- a/test/test_compute_backend.jl +++ b/test/test_compute_backend.jl @@ -103,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 From 1fed2f10fe7f977ab723a09cc4e726839c5006d6 Mon Sep 17 00:00:00 2001 From: Anshul Singhvi Date: Sun, 20 Sep 2026 13:45:38 -0400 Subject: [PATCH 6/6] Fix the DiskArrays backend implementation so the test suite passes - Reductions over all dimensions return a scalar like in Base, instead of the 1x...x1 array from `compute`. Same for `mapreduce` without `dims`. - `mapreduce` applied only `op` and never `f`. - Without `init`, `mapreduce` got DiskArrays' `_InitialValue` sentinel to fill the buffers with. Use a neutral element of `op` instead. - `value(::OnlineStats.Extrema)` is a NamedTuple `(min, max, nmin, nmax)`, so `maximum` with `strategy=:reduce` returned `nmax`, a count, and `extrema` failed to convert. Use the accessors of the stat. This fixes the `@test_broken` "gives all ones for some reason" in the Aggregate tests. - `extrema` defaults to `strategy=:reduce`, the direct aggregator broadcasts its result into the output, which fails for a Tuple. - Tests: scalar results for `mapreduce` without `dims`, an associative `op`, and regression tests for the OnlineStats path. - Point `[sources]` to the DiskArrays branch with the `count` fallback. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_018ra7EAJeFE1dtekwf4Tzxh --- Project.toml | 4 ++-- src/enginearrays.jl | 24 +++++++++++++++++++----- src/util/onlinestats_helpers.jl | 7 ++++--- test/test_compute_backend.jl | 10 +++++++--- test/test_interpolation.jl | 3 +-- 5 files changed, 33 insertions(+), 15 deletions(-) diff --git a/Project.toml b/Project.toml index cbbdb00..a31dbe7 100644 --- a/Project.toml +++ b/Project.toml @@ -26,9 +26,9 @@ Zarr = "0a941bbe-ad1d-11e8-39d9-ab76183a1d99" [weakdeps] Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" -# TODO remove once JuliaIO/DiskArrays.jl#313 is merged and released +# TODO remove once JuliaIO/DiskArrays.jl#315 is merged and released [sources] -DiskArrays = {url = "https://github.com/JuliaIO/DiskArrays.jl", rev = "as/static_backends"} +DiskArrays = {url = "https://github.com/JuliaIO/DiskArrays.jl", rev = "as/backend_fallbacks"} [extensions] DaggerExt = "Dagger" diff --git a/src/enginearrays.jl b/src/enginearrays.jl index 6e38703..506edb8 100644 --- a/src/enginearrays.jl +++ b/src/enginearrays.jl @@ -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,14 +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 # `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} + 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) diff --git a/src/util/onlinestats_helpers.jl b/src/util/onlinestats_helpers.jl index 0172e34..1bc04a6 100644 --- a/src/util/onlinestats_helpers.jl +++ b/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}) ]) diff --git a/test/test_compute_backend.jl b/test/test_compute_backend.jl index a1ffed3..4d3b97d 100644 --- a/test/test_compute_backend.jl +++ b/test/test_compute_backend.jl @@ -40,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 @@ -92,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 @@ -101,8 +105,8 @@ 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 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 diff --git a/test/test_interpolation.jl b/test/test_interpolation.jl index e6f60c2..f2521f4 100644 --- a/test/test_interpolation.jl +++ b/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[:,:,:]