Add fast reductions

src/abstractdataframe/selection.jl

@@ -376,12 +376,60 @@ _transformation_helper(df::AbstractDataFrame, col_idx::Int, (fun,)::Ref{Any}) =
 
 _empty_astable_helper(fun, len) = [fun(NamedTuple()) for _ in 1:len]
 
+@noinline _table_transformation(df_sel::AbstractDataFrame, fun) = fun(Tables.columntable(df_sel))
+
+function _sum_fast(cols::Vector{<:AbstractVector})
+    isempty(cols) && throw(ArgumentError("No columns selected for reduction"))
+    return sum(cols)
+end
+
+function _sum_skipmissing_fast(cols::Vector{<:AbstractVector})
+    isempty(cols) && throw(ArgumentError("No columns selected for reduction"))
+    sumz = nothing
+    for col in cols
+        try
+            zi = zero(eltype(col))
+            if isnothing(sumz)
+                sumz = zi
+            elseif !ismissing(zi) # zi is missing if eltype is Missing
+                sumz += zi
+            end
+        catch e
+            if e isa MethodError && e.f === Base.zero
+                sumz = nothing
+                break
+            else
+                throw(e)
+            end
+        end
+    end
+    # this will happen if eltype of some columns do not support zero
+    # or all columns have eltype Missing
+    isnothing(sumz) && return nothing
+    init = fill(sumz, length(cols[1]))
+
+    return foldl(cols, init=init) do l, r
+        l .= ifelse.(ismissing.(r), l, l .+ r)
+    end
+end
+
 function _transformation_helper(df::AbstractDataFrame, col_idx::AsTable, (fun,)::Ref{Any})
-    tbl = Tables.columntable(select(df, col_idx.cols, copycols=false))
-    if isempty(tbl) && fun isa ByRow
+    df_sel = select(df, col_idx.cols, copycols=false)
+    if ncol(df_sel) == 0 && fun isa ByRow
         return _empty_astable_helper(fun.fun, nrow(df))
+    elseif typeof(fun) === typeof(Base.sum) || typeof(fun) === ByRow{typeof(Base.sum)}
+        return _sum_fast(map(identity, eachcol(df_sel)))
+    elseif typeof(fun) === ByRow{typeof(Base.:∘(Base.sum, Base.skipmissing))}
+        fastsum = _sum_skipmissing_fast(map(identity, eachcol(df_sel)))
+        if isnothing(fastsum)
+            slowsum = _table_transformation(df_sel, fun)
+            T = mapreduce(typeof, promote_type, slowsum)
+            return convert(Vector{T}, slowsum)
+        else
+            return fastsum
+        end
     else
-        return fun(tbl)
+        return _table_transformation(df_sel, fun)
     end
 end
 
@@ -392,6 +440,9 @@ function _transformation_helper(df::AbstractDataFrame, col_idx::AbstractVector{I
         return _empty_selector_helper(fun.fun, nrow(df))
     else
         cdf = eachcol(df)
+        if typeof(fun) === typeof(Base.:+) || typeof(fun) === ByRow{typeof(Base.:+)}
+            return _sum_fast(map(c -> cdf[c], col_idx))
+        end
         return fun(map(c -> cdf[c], col_idx)...)
     end
 end

test/select.jl

@@ -1678,4 +1678,26 @@ end
     @test select(df, :a => ByRow(f) => :a) == DataFrame(a=1:3)
 end
 
+@testset "wide reductions" begin
+    Random.seed!(1234)
+    df = DataFrame(rand(Int, 2, 20_000), :auto)
+    df.x100 = [1, 2]
+    df.x1000 = Union{Int, Missing}[1, 2]
+    df.x10000 = Union{Float64, Missing}[1, missing]
+
+    @test @elapsed(select(df, All() => (+) => :res)) < 10.0
+    @test @elapsed(select(df, Between(:x2,:x19999) => ByRow(+) => :res)) < 10.0
+    @test @elapsed(select(df, AsTable(All()) => sum => :res)) < 10.0
+    @test @elapsed(select(df, AsTable(Not(:x99)) => ByRow(sum) => :res)) < 10.0
+
+    for sel in (All(), Between(:x2,:x19999), Not(:x99))
+        @show sel
+        @test select(df, sel => (+) => :res) ≅
+              select(df, sel => ByRow(+) => :res) ≅
+              select(df, AsTable(sel) => sum => :res) ≅
+              select(df, AsTable(sel) => ByRow(sum) => :res) ≅
+              DataFrame(res=sum(collect(eachcol(select(df, sel)))))
+    end
+end
+
 end # module