Update colwise (#28)

Return a vector when a single function is passed, and a matrix when several functions are passed. Rely on comprehensions to infer the return type, and return a NullableArray when applicable. Allow passing a tuple of functions, not only a vector.

Author: cjprybol

src/groupeddatatable/grouping.jl

@@ -232,21 +232,31 @@ dt = DataTable(a = repeat([1, 2, 3, 4], outer=[2]),
                b = repeat([2, 1], outer=[4]),
                c = randn(8))
 colwise(sum, dt)
+colwise([sum, lenth], dt)
+colwise((minimum, maximum), dt)
 colwise(sum, groupby(dt, :a))
 ```
 
 """
-colwise(f::Function, d::AbstractDataTable) = Any[vcat(f(d[idx])) for idx in 1:size(d, 2)]
-colwise(f::Function, gd::GroupedDataTable) = map(colwise(f), gd)
-colwise(f::Function) = x -> colwise(f, x)
-colwise(f) = x -> colwise(f, x)
+function colwise(f, d::AbstractDataTable)
+    x = [f(d[i]) for i in 1:ncol(d)]
+    if eltype(x) <: Nullable
+        return NullableArray(x)
+    else
+        return x
+    end
+end
 # apply several functions to each column in a DataTable
-colwise{T<:Function}(fns::Vector{T}, d::AbstractDataTable) =
-    reshape(Any[vcat(f(d[idx])) for f in fns, idx in 1:size(d, 2)],
-            length(fns)*size(d, 2))
-colwise{T<:Function}(fns::Vector{T}, gd::GroupedDataTable) = map(colwise(fns), gd)
-colwise{T<:Function}(fns::Vector{T}) = x -> colwise(fns, x)
-
+function colwise(fns::Union{AbstractVector, Tuple}, d::AbstractDataTable)
+    x = [f(d[i]) for f in fns, i in 1:ncol(d)]
+    if eltype(x) <: Nullable
+        return NullableArray(x)
+    else
+        return x
+    end
+end
+colwise(f, gd::GroupedDataTable) = [colwise(f, g) for g in gd]
+colwise(f) = x -> colwise(f, x)
 
 """
 Split-apply-combine in one step; apply `f` to each grouping in `d`
@@ -371,12 +381,11 @@ end
 
 function _makeheaders{T<:Function}(fs::Vector{T}, cn::Vector{Symbol})
     fnames = _fnames(fs) # see other/utils.jl
-    reshape([Symbol(colname,'_',fname) for fname in fnames, colname in cn],
-            length(fnames)*length(cn))
+    [Symbol(colname,'_',fname) for fname in fnames for colname in cn]
 end
 
 function _aggregate{T<:Function}(d::AbstractDataTable, fs::Vector{T}, headers::Vector{Symbol}, sort::Bool=false)
-    res = DataTable(colwise(fs, d), headers)
+    res = DataTable(Any[vcat(f(d[i])) for f in fs for i in 1:size(d, 2)], headers)
     sort && sort!(res, cols=headers)
     res
 end

test/grouping.jl

@@ -2,14 +2,104 @@ module TestGrouping
     using Base.Test
     using DataTables
 
+    srand(1)
     dt = DataTable(a = repeat([1, 2, 3, 4], outer=[2]),
                    b = repeat([2, 1], outer=[4]),
                    c = randn(8))
     #dt[6, :a] = Nullable()
     #dt[7, :b] = Nullable()
 
-    cols = [:a, :b]
+    nullfree = DataTable(Any[collect(1:10)], [:x1])
+    @testset "colwise" begin
+        @testset "::Function, ::AbstractDataTable" begin
+            cw = colwise(sum, dt)
+            answer = NullableArray([20, 12, -0.4283098098931877])
+            @test isa(cw, NullableArray{Any, 1})
+            @test size(cw) == (ncol(dt),)
+            @test isequal(cw, answer)
+
+            cw = colwise(sum, nullfree)
+            answer = [55]
+            @test isa(cw, Array{Int, 1})
+            @test size(cw) == (ncol(nullfree),)
+            @test cw == answer
+        end
+
+        @testset "::Function, ::GroupedDataTable" begin
+            gd = groupby(DataTable(A = [:A, :A, :B, :B], B = 1:4), :A)
+            @test colwise(length, gd) == [[2,2], [2,2]]
+        end
+
+        @testset "::Vector, ::AbstractDataTable" begin
+            cw = colwise([sum], dt)
+            answer = NullableArray([20 12 -0.4283098098931877])
+            @test isa(cw, NullableArray{Any, 2})
+            @test size(cw) == (length([sum]),ncol(dt))
+            @test isequal(cw, answer)
+
+            cw = colwise([sum, minimum], nullfree)
+            answer = reshape([55, 1], (2,1))
+            @test isa(cw, Array{Int, 2})
+            @test size(cw) == (length([sum, minimum]), ncol(nullfree))
+            @test cw == answer
+
+            cw = colwise([NullableArray], nullfree)
+            answer = reshape([NullableArray(1:10)], (1,1))
+            @test isa(cw, Array{NullableArray{Int,1},2})
+            @test size(cw) == (length([NullableArray]), ncol(nullfree))
+            @test isequal(cw, answer)
+
+            @test_throws MethodError colwise(["Bob", :Susie], DataTable(A = 1:10, B = 11:20))
+        end
+
+        @testset "::Vector, ::GroupedDataTable" begin
+            gd = groupby(DataTable(A = [:A, :A, :B, :B], B = 1:4), :A)
+            @test colwise([length], gd) == [[2 2], [2 2]]
+        end
+
+        @testset "::Tuple, ::AbstractDataTable" begin
+            cw = colwise((sum, length), dt)
+            answer = Any[Nullable(20) Nullable(12) Nullable(-0.4283098098931877); 8 8 8]
+            @test isa(cw, Array{Any, 2})
+            @test size(cw) == (length((sum, length)), ncol(dt))
+            @test isequal(cw, answer)
+
+            cw = colwise((sum, length), nullfree)
+            answer = reshape([55, 10], (2,1))
+            @test isa(cw, Array{Int, 2})
+            @test size(cw) == (length((sum, length)), ncol(nullfree))
+            @test cw == answer
+
+            cw = colwise((CategoricalArray, NullableArray), nullfree)
+            answer = reshape([CategoricalArray(1:10), NullableArray(1:10)],
+                             (length((CategoricalArray, NullableArray)), ncol(nullfree)))
+            @test typeof(cw) == Array{AbstractVector,2}
+            @test size(cw) == (length((CategoricalArray, NullableArray)), ncol(nullfree))
+            @test isequal(cw, answer)
+
+            @test_throws MethodError colwise(("Bob", :Susie), DataTable(A = 1:10, B = 11:20))
+        end
+
+        @testset "::Tuple, ::GroupedDataTable" begin
+            gd = groupby(DataTable(A = [:A, :A, :B, :B], B = 1:4), :A)
+            @test colwise((length), gd) == [[2,2],[2,2]]
+        end
+
+        @testset "::Function" begin
+            cw = map(colwise(sum), (nullfree, dt))
+            answer = ([55], NullableArray(Any[20, 12, -0.4283098098931877]))
+            @test isequal(cw, answer)
+
+            cw = map(colwise((sum, length)), (nullfree, dt))
+            answer = (reshape([55, 10], (2,1)), Any[Nullable(20) Nullable(12) Nullable(-0.4283098098931877); 8 8 8])
+            @test isequal(cw, answer)
+
+            cw = map(colwise([sum, length]), (nullfree, dt))
+            @test isequal(cw, answer)
+        end
+    end
 
+    cols = [:a, :b]
     f(dt) = DataTable(cmax = maximum(dt[:c]))
 
     sdt = unique(dt[cols])
@@ -24,9 +114,6 @@ module TestGrouping
 
     byf = by(dt, :a, dt -> DataTable(bsum = sum(dt[:b])))
 
-    @test all(T -> T <: AbstractVector, map(typeof, colwise([sum], dt)))
-    @test all(T -> T <: AbstractVector, map(typeof, colwise(sum, dt)))
-
     # groupby() without groups sorting
     gd = groupby(dt, cols)
     ga = map(f, gd)