allow transformation destination to be a function (#2897)

Author: bkamins

NEWS.md

@@ -72,6 +72,11 @@
   ([#2859](https://github.com/JuliaData/DataFrames.jl/pull/2859))
 * `Cols` now accepts a predicate accepting column names as strings.
   ([#2881](https://github.com/JuliaData/DataFrames.jl/pull/2881))
+* In `source => transformation => destination` transformation specification
+  minilanguage now `destination` can be also a `Function` generating
+  target column names and taking column names specified by `source`
+  as an argument.
+  ([#2897](https://github.com/JuliaData/DataFrames.jl/pull/2897))
 
 ## Bug fixes
 

docs/src/man/split_apply_combine.md

@@ -51,7 +51,11 @@ each subset of the `DataFrame`. This specification can be of the following forms
    `function` returns a single value or a vector; the generated name is created by
    concatenating source column name and `function` name by default (see examples below).
 3. a `cols => function => target_cols` form additionally explicitly specifying
-   the target column or columns.
+   the target column or columns, which must be a single name (as a `Symbol` or a string),
+   a vector of names or `AsTable`. Additionally it can be a `Function` which
+   takes a string or a vector of strings as an argument containing names of columns
+   selected by `cols`, and returns the target columns names (all accepted types
+   except `AsTable` are allowed).
 4. a `col => target_cols` pair, which renames the column `col` to `target_cols`, which
    must be single name (as a `Symbol` or a string), a vector of names or `AsTable`.
 5. a `nrow` or `nrow => target_cols` form which efficiently computes the number of rows

src/abstractdataframe/selection.jl

@@ -45,7 +45,11 @@ const TRANSFORMATION_COMMON_RULES =
        `function` returns a single value or a vector; the generated name is created by
        concatenating source column name and `function` name by default (see examples below).
     3. a `cols => function => target_cols` form additionally explicitly specifying
-       the target column or columns.
+       the target column or columns, which must be a single name (as a `Symbol` or a string),
+       a vector of names or `AsTable`. Additionally it can be a `Function` which
+       takes a string or a vector of strings as an argument containing names of columns
+   selected by `cols`, and returns the target columns names (all accepted types
+   except `AsTable` are allowed).
     4. a `col => target_cols` pair, which renames the column `col` to `target_cols`, which
        must be single name (as a `Symbol` or a string), a vector of names or `AsTable`.
     5. a `nrow` or `nrow => target_cols` form which efficiently computes the number of rows
@@ -244,7 +248,8 @@ function normalize_selection(idx::AbstractIndex,
                                                      <:Pair{<:Base.Callable,
                                                             <:Union{Symbol, AbstractString, DataType,
                                                                     AbstractVector{Symbol},
-                                                                    AbstractVector{<:AbstractString}}}}),
+                                                                    AbstractVector{<:AbstractString},
+                                                                    Function}}}),
                              renamecols::Bool)
     lls = last(last(sel))
 
@@ -276,6 +281,20 @@ function normalize_selection(idx::AbstractIndex,
             end
     end
 
+    if lls isa Function
+        fun_colnames = _names(idx)[c]
+        # if AsTable was used as source we always treat it as multicolumn selector
+        if wanttable && fun_colnames isa Symbol
+            fun_colnames = [fun_colnames]
+        end
+        lls = lls(string.(fun_colnames))
+        if !(lls isa Union{Symbol, AbstractString, AbstractVector{Symbol},
+                           AbstractVector{<:AbstractString}})
+            throw(ArgumentError("function producing target column names must " *
+                                "return a Symbol, a string, a vector of Symbols " *
+                                "or a vector of strings"))
+        end
+    end
     if lls isa AbstractString
         combine_target_col = Symbol(lls)
     elseif lls isa AbstractVector{<:AbstractString}
@@ -798,6 +817,15 @@ julia> select(df, AsTable(:) => ByRow(mean), renamecols=false)
    2 │     3.5
    3 │     4.5
 
+julia> select(df, AsTable(:) => ByRow(mean) => x -> join(x, "_"))
+3×1 DataFrame
+ Row │ a_b
+     │ Float64
+─────┼─────────
+   1 │     2.5
+   2 │     3.5
+   3 │     4.5
+
 julia> select(first, df)
 3×2 DataFrame
  Row │ a      b
@@ -1125,6 +1153,15 @@ julia> combine(df, AsTable(:) => ByRow(mean), renamecols=false)
    2 │     3.5
    3 │     4.5
 
+julia> combine(df, AsTable(:) => ByRow(mean) => x -> join(x, "_"))
+3×1 DataFrame
+ Row │ a_b
+     │ Float64
+─────┼─────────
+   1 │     2.5
+   2 │     3.5
+   3 │     4.5
+
 julia> combine(first, df)
 1×2 DataFrame
  Row │ a      b

test/select.jl

@@ -1692,4 +1692,34 @@ end
     @test_throws ArgumentError select(sdf, [:x => length => :a, 1 => :b], copycols=false)
 end
 
+@testset "function as target column names specifier" begin
+    df_ref = DataFrame(x=[[1, 2], [3, 4]], id=1:2)
+    for v in (df_ref, groupby(df_ref, :id))
+        @test select(v, :id, :x => ByRow(first) => identity) == DataFrame(id=1:2, x=[1, 3])
+        @test select(v, :id, "x" => ByRow(first) => identity) == DataFrame(id=1:2, x=[1, 3])
+        @test select(v, :id, 1 => ByRow(first) => identity) == DataFrame(id=1:2, x=[1, 3])
+        @test select(v, :id, 1 => ByRow(first) => uppercase) == DataFrame(id=1:2, X=[1, 3])
+        @test select(v, :id, 1 => ByRow(first) => string) == DataFrame(id=1:2, x=[1, 3])
+        @test select(v, :id, 1 => ByRow(first) => x -> Symbol(x)) == DataFrame(id=1:2, x=[1, 3])
+        @test select(v, :id, 1 => identity => x -> ["p", "q"]) ==
+              DataFrame(id=1:2, p=[1, 3], q=[2, 4])
+        @test select(v, :id, 1 => identity => x -> [:p, :q]) ==
+              DataFrame(id=1:2, p=[1, 3], q=[2, 4])
+        @test_throws ArgumentError select(v, :id, 1 => identity => x -> [:p, "q"])
+        @test_throws ArgumentError select(v, :id, 1 => identity => x -> AsTable)
+        @test select(v, :id, AsTable(1) => first => string) ==
+              DataFrame("id" => 1:2, "[\"x\"]" => [[1, 2], [3, 4]])
+        @test select(v, :id, ["x", "x"] => ByRow((p,q) -> first(p)) => string) ==
+            DataFrame("id" => 1:2, "[\"x\", \"x\"]" => [1, 3])
+        @test select(v, :id, 1:2 => ((p, q) -> q) => x -> join(x, "_")) ==
+              DataFrame(id=1:2, x_id=1:2)
+        @test select(v, :id, AsTable(1:2) => last => x -> join(x, "_")) ==
+              DataFrame(id=1:2, x_id=1:2)
+        # we could make this work, but I skip it to keep the code simpler
+        # The problem is that Symbol and String are not Function
+        @test_throws ArgumentError select(v, :id, 1 => ByRow(first) => Symbol)
+        @test_throws ArgumentError select(v, :id, 1 => ByRow(first) => String)
+    end
+end
+
 end # module