Solve the overflow in mean() on integers by promoting accumulator (#25)

Author: kagalenko-m-b

src/Statistics.jl

@@ -61,17 +61,17 @@ julia> mean([√1, √2, √3])
 function mean(f, itr)
     y = iterate(itr)
     if y === nothing
-        return Base.mapreduce_empty_iter(f, Base.add_sum, itr,
+        return Base.mapreduce_empty_iter(f, +, itr,
                                          Base.IteratorEltype(itr)) / 0
     end
     count = 1
     value, state = y
-    f_value = f(value)
-    total = Base.reduce_first(Base.add_sum, f_value)
+    f_value = f(value)/1
+    total = Base.reduce_first(+, f_value)
     y = iterate(itr, state)
     while y !== nothing
         value, state = y
-        total += f(value)
+        total += _mean_promote(total, f(value))
         count += 1
         y = iterate(itr, state)
     end
@@ -103,9 +103,6 @@ julia> mean(√, [1 2 3; 4 5 6], dims=2)
 """
 mean(f, A::AbstractArray; dims=:) = _mean(f, A, dims)
 
-_mean(f, A::AbstractArray, ::Colon) = sum(f, A) / length(A)
-_mean(f, A::AbstractArray, dims) = sum(f, A, dims=dims) / mapreduce(i -> size(A, i), *, unique(dims); init=1)
-
 """
     mean!(r, v)
 
@@ -164,10 +161,25 @@ julia> mean(A, dims=2)
  3.5
 ```
 """
-mean(A::AbstractArray; dims=:) = _mean(A, dims)
+mean(A::AbstractArray; dims=:) = _mean(identity, A, dims)
+
+_mean_promote(x::T, y::S) where {T,S} = convert(promote_type(T, S), y)
 
-_mean(A::AbstractArray{T}, region) where {T} = mean!(Base.reducedim_init(t -> t/2, +, A, region), A)
-_mean(A::AbstractArray, ::Colon) = sum(A) / length(A)
+function _mean(f, A::AbstractArray, dims=:)
+    isempty(A) && return sum(f, A, dims=dims)/0
+    if dims === (:)
+        n = length(A)
+    else
+        n = mapreduce(i -> size(A, i), *, unique(dims); init=1)
+    end
+    x1 = f(first(A)) / 1
+    result = sum(x -> _mean_promote(x1, f(x)), A, dims=dims)
+    if dims === (:)
+        return result / n
+    else
+        return result ./= n
+    end
+end
 
 function mean(r::AbstractRange{<:Real})
     isempty(r) && return oftype((first(r) + last(r)) / 2, NaN)

test/runtests.jl

@@ -130,6 +130,23 @@ end
         @test mean(identity, x) == mean(identity, g) == typemax(T)
         @test mean(x, dims=2) == [typemax(T)]'
     end
+    # Check that mean avoids integer overflow (#22)
+    let x = fill(typemax(Int), 10), a = tuple(x...)
+        @test (mean(x) == mean(x, dims=1)[] == mean(float, x)
+               == mean(a) == mean(v for v in x)  == mean(v for v in a)
+               ≈ float(typemax(Int)))
+    end
+    let x = rand(10000)  # mean should use sum's accurate pairwise algorithm
+        @test mean(x) == sum(x) / length(x)
+    end
+    @test mean(Number[1, 1.5, 2+3im]) === 1.5+1im # mixed-type array
+    @test mean(v for v in Number[1, 1.5, 2+3im]) === 1.5+1im
+    @test (@inferred mean(Int[])) === 0/0
+    @test (@inferred mean(Float32[])) === 0.f0/0    
+    @test (@inferred mean(Float64[])) === 0/0
+    @test (@inferred mean(Iterators.filter(x -> true, Int[]))) === 0/0
+    @test (@inferred mean(Iterators.filter(x -> true, Float32[]))) === 0.f0/0
+    @test (@inferred mean(Iterators.filter(x -> true, Float64[]))) === 0/0
 end
 
 @testset "mean/median for ranges" begin
@@ -710,7 +727,7 @@ end
     x = Any[1, 2, 4, 10]
     y = Any[1, 2, 4, 10//1]
     @test var(x) === 16.25
-    @test var(y) === 65//4
+    @test var(y) === 16.25
     @test std(x) === sqrt(16.25)
     @test quantile(x, 0.5)  === 3.0
     @test quantile(x, 1//2) === 3//1