diff --git a/src/mapreduce.jl b/src/mapreduce.jl index a28e431..f5af91e 100644 --- a/src/mapreduce.jl +++ b/src/mapreduce.jl @@ -51,9 +51,13 @@ end # Implementation for special cases and if fallback breaks in future julia versions -for fname in [:sum, :prod, :all, :any, :minimum, :maximum] - @eval Base.$fname(v::AbstractDiskArray) = Base.$fname(identity, v::AbstractDiskArray) - @eval function Base.$fname(f::Function, v::AbstractDiskArray) +for (fname, _fname) in ((:sum, :_sum), (:prod, :_prod), + (:all, :_all), (:any, :_any), + (:minimum, :_minimum), (:maximum, :_maximum)) + @eval Base.$fname(v::AbstractDiskArray; dims=:) = Base.$_fname(v, dims) + @eval Base.$fname(f::Function, v::AbstractDiskArray; dims=:) = Base.$_fname(f, v, dims) + @eval Base.$_fname(v::AbstractDiskArray, ::Colon) = Base.$_fname(identity, v, :) + @eval function Base.$_fname(f::Function, v::AbstractDiskArray, ::Colon) $fname(eachchunk(v)) do chunk $fname(f, v[chunk...]) end @@ -66,3 +70,24 @@ function Base.count(f, v::AbstractDiskArray) count(f, v[chunk...]) end end + +for (_fname, init, acum) in ((:_sum, :zero, :+), (:_prod, :one, :*), + (:_all, _->:true, :&), (:_any, _->:false, :|), + (:_maximum, :typemin, :max), (:_minimum, :typemax, :min)) + @eval Base.$_fname(a::AbstractDiskArray, dims) = Base.$_fname(identity, a, dims) + @eval function Base.$_fname(f::Function, a::AbstractDiskArray{T,N}, dims) where {T,N} + _dims = typeof(dims)<:Tuple ? [dims...] : typeof(dims)<:Number ? [dims] : dims + out_dims = ntuple(Val(N)) do i + i in _dims ? 1 : size(a)[i] + end + T1 = Base.promote_op(f, T) + out = fill($init(T1), out_dims...) + for c in eachchunk(a) + out_c = ntuple(Val(N)) do i + i in _dims ? (1:1) : c[i] + end + out[out_c...] .= $acum.(out[out_c...], Base.$_fname(f, a[c...], dims)) + end + return out + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 8ecefa2..e466ebb 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1059,3 +1059,28 @@ end end @test count(a) + count(!, a) == length(a) end + +@testset "dims kwarg" begin + a = ChunkedDiskArray(reshape(1:3*4*5, 3,4,5); chunksize=(2,3,4)) + for fname in (:sum, :prod, :minimum, :maximum) + @test @eval $fname($a) == $fname(Array($a)) + @test @eval $fname($a, dims=1) == $fname(Array($a), dims=1) + @test @eval $fname($a, dims=2) == $fname(Array($a), dims=2) + @test @eval $fname($a, dims=(1,2)) == $fname(Array($a), dims=(1,2)) + @test @eval $fname(x->x/2, $a) ≈ $fname(x->x/2, Array($a)) + @test @eval $fname(x->x/2, $a, dims=(1,2)) ≈ $fname(x->x/2, Array($a), dims=(1,2)) + @eval out = @capture_out @trace $fname($a) DiskArrays + @test occursin("DiskGenerator", out) == false + end + b = ChunkedDiskArray(reshape(1:3*4*5 .> 10, 3,4,5); chunksize=(2,3,4)) + for fname in (:all, :any) + @test @eval $fname($b) == $fname(Array($b)) + @test @eval $fname($b, dims=1) == $fname(Array($b), dims=1) + @test @eval $fname($b, dims=2) == $fname(Array($b), dims=2) + @test @eval $fname($b, dims=(1,2)) == $fname(Array($b), dims=(1,2)) + @test @eval $fname(!, $b) == $fname(!, Array($b)) + @test @eval $fname(!, $b, dims=(1,2)) == $fname(!, Array($b), dims=(1,2)) + @eval out = @capture_out @trace $fname($b) DiskArrays + @test occursin("DiskGenerator", out) == false + end +end