Merge pull request #44 from JuliaComputing/s/pick

make `pick` more useful

.travis.yml

@@ -5,7 +5,7 @@ os:
   - osx
 julia:
   - 0.5
-  - nightly
+  - 0.6
 notifications:
   email: false
 # uncomment the following lines to override the default test script

REQUIRE

@@ -1,4 +1,4 @@
-julia 0.5
+julia 0.5 0.7-
 Compat 0.19
 NamedTuples 2.1.0
 PooledArrays

src/columns.jl

@@ -122,8 +122,8 @@ end
 sort!(c::Columns) = permute!(c, sortperm(c))
 sort(c::Columns) = c[sortperm(c)]
 
+map(p::ProjFn, c::Columns) = Columns(p(c.columns))
 map(p::Proj, c::Columns) = p(c.columns)
-(p::Proj)(c::Columns) = p(c.columns)
 
 vcat{D<:Tup,C<:Tuple}(c::Columns{D,C}, cs::Columns{D,C}...) = Columns{D,C}((map(vcat, map(x->x.columns, (c,cs...))...)...,))
 vcat{D<:Tup,C<:NamedTuple}(c::Columns{D,C}, cs::Columns{D,C}...) = Columns{D,C}(C(map(vcat, map(x->x.columns, (c,cs...))...)...,))

src/utils.jl

@@ -2,6 +2,8 @@ using Base.Test
 import Base: tuple_type_cons, tuple_type_head, tuple_type_tail, in, ==, isless, convert,
              length, eltype, start, next, done, show
 
+export @pick, pick
+
 eltypes(::Type{Tuple{}}) = Tuple{}
 eltypes{T<:Tuple}(::Type{T}) =
     tuple_type_cons(eltype(tuple_type_head(T)), eltypes(tuple_type_tail(T)))
@@ -56,9 +58,55 @@ end
 
 # family of projection functions
 
-immutable Proj{field}; end
+immutable ProjFn{F}
+    f::F
+end
+
+(p::ProjFn)(x::Tup) = p.f(x)
 
-(::Proj{field}){field}(x) = getfield(x, field)
+immutable Proj{f} end
+
+function (p::Proj{f}){f}(x::Tup)
+    getfield(x, f)
+end
+
+"""
+    @pick(fields...)
+
+Returns a callable object `f` such that `f(x::Tuple)` returns a `Tuple` with only
+elements of index specified by `fields`, `f(x::NamedTuple)` return a `Tuple` if
+`fields` are integers, or a `NamedTuple` if `fields` are symbols with only the specified
+fields in the output.
+
+The callable is specialized to work efficiently on `Columns` by calling it once
+on `.columns` field to get the equivalent result.
+
+Calling `map` on an `IndexedTable` with a `@pick` callable will run the callable on
+the data columns.
+
+# Examples
+    c = Columns(x=[1], y=[2.0])
+    @pick(2)(c) == Columns([2.0])
+    @pick(y)(c) == Columns(y=[2.0])
+    @pick(2,1)(c) == Columns([2.0], [1])
+    @pick(y,x)(c) == Columns(y=[2.0], x=[1])
+
+    t = IndexedTable([1], c)
+    map(@pick(y, x), t) == IndexedTables([1], Columns(y=[2.0], x=[1]))
+"""
+macro pick(ex...)
+    tup = if all([isa(x, Symbol) for x in ex])
+        # Named tuple
+        args = [:(getfield(x, $(Expr(:quote, f)))) for f in ex]
+        T = Expr(:macrocall,
+                 :(NamedTuples.$(Symbol("@NT"))),
+                   map((x) -> :($(esc(x))), ex)...)
+        :($T($(args...)))
+    else
+        :(($([:(getfield(x, $f)) for f in ex]...),))
+    end
+    :(IndexedTables.ProjFn(x -> $tup))
+end
 
 pick(fld) = Proj{fld}()
 

test/test_core.jl

@@ -1,6 +1,7 @@
 using Base.Test
 using IndexedTables
 using PooledArrays
+using NamedTuples
 
 let a = Columns([1,2,1],["foo","bar","baz"]),
     b = Columns([2,1,1],["bar","baz","foo"]),
@@ -23,13 +24,28 @@ let c = Columns([1,1,1,2,2], [1,2,4,3,5]),
     f = Columns([1,1,1], sort([rand(),0.5,rand()]))
     @test merge(IndexedTable(c,ones(5)),IndexedTable(d,ones(5))).index == Columns([1,1,1,1,2,2,2,2],[1,2,3,4,1,3,4,5])
     @test eltype(merge(IndexedTable(c,Columns(ones(Int, 5))),IndexedTable(d,Columns(ones(Float64, 5)))).data) == Tuple{Float64}
-    @test eltype(merge(IndexedTable(c,Columns(x=ones(Int, 5))),IndexedTable(d,Columns(x=ones(Float64, 5)))).data) == NamedTuples.@NT(x){Float64}
+    @test eltype(merge(IndexedTable(c,Columns(x=ones(Int, 5))),IndexedTable(d,Columns(x=ones(Float64, 5)))).data) == @NT(x){Float64}
     @test length(merge(IndexedTable(e,ones(3)),IndexedTable(f,ones(3)))) == 5
     @test vcat(Columns(x=[1]), Columns(x=[1.0])) == Columns(x=[1,1.0])
     @test vcat(Columns(x=PooledArray(["x"])), Columns(x=["y"])) == Columns(x=["x", "y"])
+
     @test summary(c) == "Columns{Tuple{Int64,Int64}}"
 end
 
+let
+    x = Columns([1], [2.0])
+    @test map(pick(2), x) == [2.0]
+    @test map(@pick(2), x) == Columns([2.0])
+    @test map(@pick(2,1), x) == Columns([2.0], [1])
+
+    y = Columns(x=[1], y=[2.0])
+    @test map(pick(2), y) == [2.0]
+    @test map(@pick(2), y) == Columns([2.0])
+    @test map(@pick(y), y) == Columns(y=[2.0])
+    @test map(@pick(2,1), y) == Columns([2.0], [1])
+    @test map(@pick(y,x), y) == Columns(y=[2.0], x=[1])
+end
+
 let c = Columns([1,1,1,2,2], [1,2,4,3,5]),
     d = Columns([1,1,2,2,2], [1,3,1,4,5]),
     e = Columns([1,1,1], sort([rand(),0.5,rand()])),