Refactor TableSelection

src/tableselection.jl

@@ -2,56 +2,46 @@
 # Licensed under the MIT License. See LICENSE in the project root.
 # ------------------------------------------------------------------
 
-struct TableSelection{T,C}
-  table::T
-  cols::C
-  ncols::Int
-  names::Vector{Symbol}
-  onames::Vector{Symbol}
-  mapnames::Dict{Symbol,Symbol}
-
-  function TableSelection(table::T, names, onames) where {T}
-    cols = Tables.columns(table)
-    _assert(onames ⊆ Tables.columnnames(cols), "all selected columns must exist in the table")
-    ncols = length(names)
-    mapnames = Dict(zip(names, onames))
-    new{T,typeof(cols)}(table, cols, ncols, names, onames, mapnames)
-  end
-end
+"""
+    TableSelection(table, names)
 
-function Base.:(==)(a::TableSelection, b::TableSelection)
-  a.names != b.names && return false
-  a.onames != b.onames && return false
-  all(nm -> Tables.getcolumn(a, nm) == Tables.getcolumn(b, nm), a.names)
+Stores a sub-`table` with given column `names`.
+"""
+struct TableSelection{T,N}
+  table::T
+  names::NTuple{N,Symbol}
+  oinds::NTuple{N,Int}
 end
 
-function Base.show(io::IO, t::TableSelection)
-  println(io, "TableSelection")
-  pretty_table(io, t, vcrop_mode=:middle, newline_at_end=false)
+function TableSelection(table, names)
+  cols = Tables.columns(table)
+  onames = Tables.columnnames(cols)
+  _assert(names ⊆ onames, "invalid columns for table selection")
+  oinds = indexin(collect(names), collect(onames))
+  TableSelection(table, Tuple(names), Tuple(oinds))
 end
 
-# Tables.jl interface
 Tables.istable(::Type{<:TableSelection}) = true
+
 Tables.columnaccess(::Type{<:TableSelection}) = true
+
 Tables.columns(t::TableSelection) = t
+
 Tables.columnnames(t::TableSelection) = t.names
 
-function Tables.getcolumn(t::TableSelection, i::Int)
-  1 ≤ i ≤ t.ncols || error("Table has no column with index $i.")
-  Tables.getcolumn(t.cols, t.mapnames[t.names[i]])
-end
+Tables.getcolumn(t::TableSelection, i::Int) = Tables.getcolumn(Tables.columns(t.table), t.oinds[i])
 
-function Tables.getcolumn(t::TableSelection, nm::Symbol)
-  nm ∉ t.names && error("Table has no column $nm.")
-  Tables.getcolumn(t.cols, t.mapnames[nm])
-end
+Tables.getcolumn(t::TableSelection, nm::Symbol) = Tables.getcolumn(Tables.columns(t.table), nm)
 
 Tables.materializer(t::TableSelection) = Tables.materializer(t.table)
-
 function Tables.schema(t::TableSelection)
-  schema = Tables.schema(t.cols)
-  names = schema.names
-  types = schema.types
-  inds = indexin(t.onames, collect(names))
-  Tables.Schema(t.names, types[inds])
+  schema = Tables.schema(t.table)
+  names = [schema.names[i] for i in t.oinds]
+  types = [schema.types[i] for i in t.oinds]
+  Tables.Schema(names, types)
+end
+
+function Base.show(io::IO, t::TableSelection)
+  println(io, "TableSelection")
+  pretty_table(io, t, vcrop_mode=:bottom, newline_at_end=false)
 end

src/transforms/select.jl

@@ -45,16 +45,19 @@ Select(pairs::Pair{C,S}...) where {C<:Column,S<:AbstractString} =
 
 Select() = throw(ArgumentError("cannot create Select transform without arguments"))
 
-# utils
-_newnames(::Nothing, select) = select
-_newnames(names::Vector{Symbol}, select) = names
-
 function applyfeat(transform::Select, feat, prep)
   cols = Tables.columns(feat)
   names = collect(Tables.columnnames(cols))
-  select = transform.selector(names)
-  newnames = _newnames(transform.newnames, select)
-  newfeat = TableSelection(feat, newnames, select)
+
+  # lazy selection of columns
+  snames = transform.selector(names)
+  stable = TableSelection(feat, snames)
+
+  # rename if necessary
+  nnames = transform.newnames
+  rename = isnothing(nnames) ? Identity() : Rename(nnames)
+  newfeat = stable |> rename
+
   newfeat, nothing
 end
 
@@ -92,9 +95,8 @@ Reject(::AllSelector) = throw(ArgumentError("cannot reject all columns"))
 function applyfeat(transform::Reject, feat, prep)
   cols = Tables.columns(feat)
   names = Tables.columnnames(cols)
-  reject = transform.selector(names)
-  select = setdiff(names, reject)
-  strans = Select(select)
-  newfeat, _ = applyfeat(strans, feat, prep)
+  snames = transform.selector(names)
+  select = Select(setdiff(names, snames))
+  newfeat, _ = applyfeat(select, feat, prep)
   newfeat, nothing
 end

test/tableselection.jl

@@ -8,14 +8,13 @@
   t = Table(; a, b, c, d, e, f)
 
   # Tables.jl interface
-  select = [:a, :b, :e]
-  newnames = select
-  s = TT.TableSelection(t, newnames, select)
+  names = [:a, :b, :e]
+  s = TT.TableSelection(t, names)
   @test Tables.istable(s) == true
   @test Tables.columnaccess(s) == true
   @test Tables.rowaccess(s) == false
   @test Tables.columns(s) === s
-  @test Tables.columnnames(s) == [:a, :b, :e]
+  @test Tables.columnnames(s) == (:a, :b, :e)
   @test Tables.schema(s).names == (:a, :b, :e)
   @test Tables.schema(s).types == (Float64, Float64, Float64)
   @test Tables.materializer(s) == Tables.materializer(t)
@@ -26,36 +25,16 @@
   @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 = TT.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
-  @test Tables.getcolumn(s, 1) == t.c
-  @test Tables.getcolumn(s, 2) == t.d
-  @test Tables.getcolumn(s, 3) == t.f
-
   # row table
-  select = [:a, :b, :e]
-  newnames = select
+  names = [:a, :b, :e]
   rt = Tables.rowtable(t)
-  s = TT.TableSelection(rt, newnames, select)
+  s = TT.TableSelection(rt, names)
   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 TT.TableSelection(t, [:a, :b, :z], [:a, :b, :z])
-  @test_throws AssertionError TT.TableSelection(t, [:x, :y, :z], [:c, :d, :k])
-  s = TT.TableSelection(t, [:a, :b, :e], [:a, :b, :e])
-  @test_throws ErrorException Tables.getcolumn(s, :f)
-  @test_throws ErrorException Tables.getcolumn(s, 4)
-  s = TT.TableSelection(t, [:x, :y, :z], [:c, :d, :f])
-  @test_throws ErrorException Tables.getcolumn(s, :c)
-  @test_throws ErrorException Tables.getcolumn(s, 4)
-  @test_throws ErrorException Tables.getcolumn(s, -2)
+  @test_throws AssertionError TT.TableSelection(t, [:a, :b, :z])
+  @test_throws AssertionError TT.TableSelection(t, [:x, :y, :z])
 end