Repository: JuliaML/TableTransforms.jl (julia)

Issue:
Add option to rename columns during selection
A nice feature of the DataFrames.jl select function is that it allows in-place renaming:

```julia
select(df, :x => :y, :a => :b)
```

Currently, we have to write a `Select(:x, :a)` and then use `Rename(:x => :y, :a => :b)` on the result. Maybe we could introduce a `SelectRename` transform or generalize the `Select` to accept renaming.

Pull request:
This PR is a major refactoring of the `Select` implementation to add support for selection with renaming.
I also took the liberty of renaming the variables to code style the rest of the TableTransforms.jl code.
I also made small optimizations in the apply function.

closes #106

Patch:
diff --git a/src/transforms/select.jl b/src/transforms/select.jl
index b0f4067b..996513fc 100644
--- a/src/transforms/select.jl
+++ b/src/transforms/select.jl
@@ -6,21 +6,24 @@ struct TableSelection{T,C}
   table::T
   cols::C
   names::Vector{Symbol}
-  function TableSelection(table::T, names::Vector{Symbol}) where {T}
+  onames::Vector{Symbol}
+  mapnames::Dict{Symbol,Symbol}
+  function TableSelection(table::T, names, onames) where {T}
     cols = Tables.columns(table)
-    @assert names ⊆ Tables.columnnames(cols)
-    new{T,typeof(cols)}(table, cols, names)
+    @assert onames ⊆ Tables.columnnames(cols)
+    new{T,typeof(cols)}(table, cols, names, onames, Dict(zip(names, onames)))
   end
 end
 
 function Base.:(==)(a::TableSelection, b::TableSelection)
-  a.names != b.names && return false
+  a.names  != b.names  && return false
+  a.onames != b.onames && return false
   all(Tables.getcolumn(a, nm) == Tables.getcolumn(b, nm) for nm in a.names)
 end
 
 function Base.show(io::IO, t::TableSelection)
   println(io, "TableSelection")
-  pretty_table(io, t, vcrop_mode=:middle)
+  pretty_table(io, t, vcrop_mode=:middle, newline_at_end=false)
 end
 
 # Tables.jl interface
@@ -32,14 +35,14 @@ Tables.getcolumn(t::TableSelection, i::Int) =
   Tables.getcolumn(t.cols, t.names[i])
 function Tables.getcolumn(t::TableSelection, nm::Symbol)
   nm ∉ t.names && error("Table has no column $nm.")
-  Tables.getcolumn(t.cols, nm)
+  Tables.getcolumn(t.cols, t.mapnames[nm])
 end
 
 function Tables.schema(t::TableSelection)
   schema = Tables.schema(t.table)
   names = schema.names
   types = schema.types
-  inds = indexin(t.names, collect(names))
+  inds = indexin(t.onames, collect(names))
   Tables.Schema(t.names, types[inds])
 end
 
@@ -52,6 +55,11 @@ Tables.materializer(t::TableSelection) =
     Select((col₁, col₂, ..., colₙ))
     
 The transform that selects columns `col₁`, `col₂`, ..., `colₙ`.
+
+    Select(col₁ => newcol₁, col₂ => newcol₂, ..., colₙ => newcolₙ)
+
+Selects the columns `col₁`, `col₂`, ..., `colₙ`
+and rename them to `newcol₁`, `newcol₂`, ..., `newcolₙ`.
     
     Select(regex)
 
@@ -63,22 +71,34 @@ Selects the columns that match with `regex`.
 Select(1, 3, 5)
 Select([:a, :c, :e])
 Select(("a", "c", "e"))
+Select(1 => :x, 3 => :y)
+Select(:a => :x, :b => :y)
+Select("a" => "x", "b" => "y")
 Select(r"[ace]")
 ```
 """
 struct Select{S<:ColSpec} <: Stateless
   colspec::S
+  newnames::Union{Vector{Symbol},Nothing}
 end
 
-Select(spec) = Select(colspec(spec))
+Select(spec) = Select(colspec(spec), nothing)
+Select(cols::T...) where {T<:Col} = Select(cols)
+
+Select(pairs::Pair{T,Symbol}...) where {T<:Col} = 
+  Select(colspec(first.(pairs)), collect(last.(pairs)))
 
-Select(cols::T...) where {T<:Col} = 
-  Select(colspec(cols))
+Select(pairs::Pair{T,S}...) where {T<:Col,S<:AbstractString} = 
+  Select(colspec(first.(pairs)), collect(Symbol.(last.(pairs))))
 
 Select() = throw(ArgumentError("Cannot create a Select object without arguments."))
 
 isrevertible(::Type{<:Select}) = true
 
+# utils
+_newnames(::Nothing, select) = select
+_newnames(names::Vector{Symbol}, select) = names
+
 function apply(transform::Select, table)
   # original columns
   cols = Tables.columns(table)
@@ -86,6 +106,7 @@ function apply(transform::Select, table)
   # retrieve relevant column names
   allcols = collect(Tables.columnnames(cols))
   select  = choose(transform.colspec, allcols)
+  names   = _newnames(transform.newnames, select)
   reject  = setdiff(allcols, select)
 
   # keep track of indices to revert later
@@ -93,35 +114,33 @@ function apply(transform::Select, table)
   rinds = indexin(reject, allcols)
 
   # sort indices to facilitate reinsertion
-  sperm  = sortperm(sinds)
-  sorted = sortperm(rinds)
-  reject = reject[sorted]
-  rinds  = rinds[sorted]
+  sperm = sortperm(sinds)
 
   # rejected columns
-  rcols = [Tables.getcolumn(cols, name) for name in reject]
+  rcolumns = [Tables.getcolumn(cols, name) for name in reject]
 
-  TableSelection(table, select), (reject, rcols, sperm, rinds)
+  TableSelection(table, names, select), (select, sperm, reject, rcolumns, rinds)
 end
 
 function revert(::Select, newtable, cache)
   # selected columns
-  cols   = Tables.columns(newtable)
-  select = Tables.columnnames(cols)
+  cols  = Tables.columns(newtable)
+  names = Tables.columnnames(cols)
   # https://github.com/JuliaML/TableTransforms.jl/issues/76
-  scols  = Any[Tables.getcolumn(cols, name) for name in select]
+  columns = Any[Tables.getcolumn(cols, name) for name in names]
 
   # rejected columns
-  reject, rcols, sperm, rinds = cache
+  select, sperm, reject, rcolumns, rinds = cache
 
   # restore rejected columns
-  anames = collect(select[sperm])
-  acols  = collect(scols[sperm])
+  onames = select[sperm]
+  ocolumns = columns[sperm]
   for (i, rind) in enumerate(rinds)
-    insert!(anames, rind, reject[i])
-    insert!(acols, rind, rcols[i])
+    insert!(onames, rind, reject[i])
+    insert!(ocolumns, rind, rcolumns[i])
   end
-  𝒯 = (; zip(anames, acols)...)
+
+  𝒯 = (; zip(onames, ocolumns)...)
   𝒯 |> Tables.materializer(newtable)
 end
 
@@ -164,7 +183,7 @@ Reject(::AllSpec) = throw(ArgumentError("Cannot reject all columns."))
 isrevertible(::Type{<:Reject}) = true
 
 function apply(transform::Reject, table)
-  cols = Tables.columns(table)
+  cols    = Tables.columns(table)
   allcols = Tables.columnnames(cols)
   reject  = choose(transform.colspec, allcols)
   select  = setdiff(allcols, reject)


Test patch:
diff --git a/test/colspec.jl b/test/colspec.jl
index de8a90df..414caafe 100644
--- a/test/colspec.jl
+++ b/test/colspec.jl
@@ -115,7 +115,9 @@ end
   t = Table(; a, b, c, d, e, f)
 
   # Tables.jl interface
-  s = TableTransforms.TableSelection(t, [:a, :b, :e])
+  select = [:a, :b, :e]
+  newnames = select
+  s = TableTransforms.TableSelection(t, newnames, select)
   @test Tables.istable(s) == true
   @test Tables.columnaccess(s) == true
   @test Tables.rowaccess(s) == false
@@ -131,16 +133,30 @@ end
   @test Tables.getcolumn(s, 1) == Tables.getcolumn(cols, 1)
   @test Tables.getcolumn(s, 3) == Tables.getcolumn(cols, :e)
 
+  # selectin with renaming
+  select = [:c, :d, :f]
+  newnames = [:x, :y, :z]
+  s = TableTransforms.TableSelection(t, newnames, select)
+  @test Tables.columnnames(s) == [:x, :y, :z]
+  @test Tables.getcolumn(s, :x) == t.c
+  @test Tables.getcolumn(s, :y) == t.d
+  @test Tables.getcolumn(s, :z) == t.f
+
   # row table
+  select = [:a, :b, :e]
+  newnames = select
   rt = Tables.rowtable(t)
-  s = TableTransforms.TableSelection(rt, [:a, :b, :e])
+  s = TableTransforms.TableSelection(rt, newnames, select)
   cols = Tables.columns(rt)
   @test Tables.getcolumn(s, :a) == Tables.getcolumn(cols, :a)
   @test Tables.getcolumn(s, 1) == Tables.getcolumn(cols, 1)
   @test Tables.getcolumn(s, 3) == Tables.getcolumn(cols, :e)
 
   # throws
-  @test_throws AssertionError TableTransforms.TableSelection(t, [:a, :b, :z])
-  s = TableTransforms.TableSelection(t, [:a, :b, :e])
+  @test_throws AssertionError TableTransforms.TableSelection(t, [:a, :b, :z], [:a, :b, :z])
+  @test_throws AssertionError TableTransforms.TableSelection(t, [:x, :y, :z], [:c, :d, :k])
+  s = TableTransforms.TableSelection(t, [:a, :b, :e], [:a, :b, :e])
   @test_throws ErrorException Tables.getcolumn(s, :f)
+  s = TableTransforms.TableSelection(t, [:x, :y, :z], [:c, :d, :f])
+  @test_throws ErrorException Tables.getcolumn(s, :c)
 end
diff --git a/test/shows.jl b/test/shows.jl
index 9dd65d1f..89631215 100644
--- a/test/shows.jl
+++ b/test/shows.jl
@@ -4,13 +4,28 @@
 
     # compact mode
     iostr = sprint(show, T)
-    @test iostr == "Select([:a, :b, :c])"
+    @test iostr == "Select([:a, :b, :c], nothing)"
 
     # full mode
     iostr = sprint(show, MIME("text/plain"), T)
     @test iostr == """
     Select transform
-    └─ colspec = [:a, :b, :c]"""
+    ├─ colspec = [:a, :b, :c]
+    └─ newnames = nothing"""
+
+    # selection with renaming
+    T = Select(:a => :x, :b => :y)
+
+    # compact mode
+    iostr = sprint(show, T)
+    @test iostr == "Select([:a, :b], [:x, :y])"
+
+    # full mode
+    iostr = sprint(show, MIME("text/plain"), T)
+    @test iostr == """
+    Select transform
+    ├─ colspec = [:a, :b]
+    └─ newnames = [:x, :y]"""
   end
 
   @testset "Reject" begin
@@ -321,13 +336,13 @@
 
     # compact mode
     iostr = sprint(show, pipeline)
-    @test iostr == "Select([:x, :z]) → ZScore() → Scale(0, 1)"
+    @test iostr == "Select([:x, :z], nothing) → ZScore() → Scale(0, 1)"
 
     # full mode
     iostr = sprint(show, MIME("text/plain"), pipeline)
     @test iostr == """
     Sequential
-    ├─ Select([:x, :z])
+    ├─ Select([:x, :z], nothing)
     ├─ ZScore()
     └─ Scale(0, 1)"""
   end
diff --git a/test/transforms.jl b/test/transforms.jl
index ac8ff260..bae2abda 100644
--- a/test/transforms.jl
+++ b/test/transforms.jl
@@ -98,6 +98,89 @@
     n2 = reapply(T, t, c1)
     @test n1 == n2
 
+    # selection with renaming
+    a = rand(4000)
+    b = rand(4000)
+    c = rand(4000)
+    d = rand(4000)
+    t = Table(; a, b, c, d)
+
+    # integer => symbol
+    T = Select(1 => :x, 3 => :y)
+    n, c = apply(T, t)
+    @test Tables.columnnames(n) == [:x, :y]
+    @test Tables.getcolumn(n, :x) == t.a
+    @test Tables.getcolumn(n, :y) == t.c
+
+    # integer => string
+    T = Select(2 => "x", 4 => "y")
+    n, c = apply(T, t)
+    @test Tables.columnnames(n) == [:x, :y]
+    @test Tables.getcolumn(n, :x) == t.b
+    @test Tables.getcolumn(n, :y) == t.d
+
+    # symbol => symbol
+    T = Select(:a => :x, :c => :y)
+    n, c = apply(T, t)
+    @test Tables.columnnames(n) == [:x, :y]
+    @test Tables.getcolumn(n, :x) == t.a
+    @test Tables.getcolumn(n, :y) == t.c
+
+    # symbol => string
+    T = Select(:b => "x", :d => "y")
+    n, c = apply(T, t)
+    @test Tables.columnnames(n) == [:x, :y]
+    @test Tables.getcolumn(n, :x) == t.b
+    @test Tables.getcolumn(n, :y) == t.d
+
+    T = Select(:a => :x1, :b => :x2, :c => :x3, :d => :x4)
+    n, c = apply(T, t)
+    @test Tables.columnnames(n) == [:x1, :x2, :x3, :x4]
+    @test Tables.getcolumn(n, :x1) == t.a
+    @test Tables.getcolumn(n, :x2) == t.b
+    @test Tables.getcolumn(n, :x3) == t.c
+    @test Tables.getcolumn(n, :x4) == t.d
+
+    # string => symbol
+    T = Select("a" => :x, "c" => :y)
+    n, c = apply(T, t)
+    @test Tables.columnnames(n) == [:x, :y]
+    @test Tables.getcolumn(n, :x) == t.a
+    @test Tables.getcolumn(n, :y) == t.c
+
+    # string => string
+    T = Select("b" => "x", "d" => "y")
+    n, c = apply(T, t)
+    @test Tables.columnnames(n) == [:x, :y]
+    @test Tables.getcolumn(n, :x) == t.b
+    @test Tables.getcolumn(n, :y) == t.d
+
+    T = Select("a" => "x1", "b" => "x2", "c" => "x3", "d" => "x4")
+    n, c = apply(T, t)
+    @test Tables.columnnames(n) == [:x1, :x2, :x3, :x4]
+    @test Tables.getcolumn(n, :x1) == t.a
+    @test Tables.getcolumn(n, :x2) == t.b
+    @test Tables.getcolumn(n, :x3) == t.c
+    @test Tables.getcolumn(n, :x4) == t.d
+
+    # row table
+    rt = Tables.rowtable(t)
+    cols = Tables.columns(rt)
+
+    T = Select(:a => :x, :c => :y)
+    n, c = apply(T, rt)
+    @test Tables.columnnames(n) == [:x, :y]
+    @test Tables.getcolumn(n, :x) == Tables.getcolumn(cols, :a)
+    @test Tables.getcolumn(n, :y) == Tables.getcolumn(cols, :c)
+    rtₒ = revert(T, n, c)
+    @test rt == rtₒ
+
+    # reapply test
+    T = Select(:b => :x, :d => :y)
+    n1, c1 = apply(T, t)
+    n2 = reapply(T, t, c1)
+    @test n1 == n2
+
     # selection with Regex
     T = Select(r"[dcb]")
     n, c = apply(T, t)