Repository: JuliaML/TableTransforms.jl (julia)

Issue:
Lazy Select/Reject
Currently our `Select/Reject` are greedy and generate copies of the input tables. We could create a lazy type for selections similar to what TableOperations.jl does.

Pull request:
This PR implements a lazy type for selections/rejections.

Fix #16

Patch:
diff --git a/Project.toml b/Project.toml
index 29281ae4..2b7b45c0 100644
--- a/Project.toml
+++ b/Project.toml
@@ -6,6 +6,7 @@ version = "0.2.0"
 [deps]
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
 ScientificTypes = "321657f4-b219-11e9-178b-2701a2544e81"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
diff --git a/src/TableTransforms.jl b/src/TableTransforms.jl
index 9e24008d..557350f5 100644
--- a/src/TableTransforms.jl
+++ b/src/TableTransforms.jl
@@ -10,6 +10,7 @@ using Distributions: Normal
 using Transducers: tcollect
 using LinearAlgebra
 using Statistics
+using PrettyTables
 
 import Distributions: ContinuousUnivariateDistribution
 import Distributions: quantile, cdf
diff --git a/src/transforms/select.jl b/src/transforms/select.jl
index 89019c47..de5d9529 100644
--- a/src/transforms/select.jl
+++ b/src/transforms/select.jl
@@ -2,6 +2,46 @@
 # Licensed under the MIT License. See LICENSE in the project root.
 # ------------------------------------------------------------------
 
+struct TableSelection{T} 
+  table::T
+  cols::Vector{Symbol}
+  function TableSelection(table::T, cols::Vector{Symbol}) where {T}
+    @assert cols ⊆ Tables.columnnames(table)
+    new{T}(table, cols)
+  end
+end
+
+function Base.:(==)(a::TableSelection, b::TableSelection)
+  a.cols != b.cols && return false
+  all(Tables.getcolumn(a, col) == Tables.getcolumn(b, col) for col in a.cols)
+end
+
+function Base.show(io::IO, t::TableSelection)
+  println(io, "TableSelection")
+  pretty_table(io, t, vcrop_mode=:middle)
+end
+
+# Tables.jl interface
+Tables.istable(::Type{<:TableSelection}) = true
+Tables.columnaccess(::Type{<:TableSelection}) = true
+Tables.columns(t::TableSelection) = t
+Tables.columnnames(t::TableSelection) = t.cols
+Tables.getcolumn(t::TableSelection, col::Int) =
+  Tables.getcolumn(t.table, t.cols[col])
+Tables.getcolumn(t::TableSelection, col::Symbol) = 
+  Tables.getcolumn(t.table, col)
+
+function Tables.schema(t::TableSelection)
+  schema = Tables.schema(t.table)
+  names = schema.names
+  types = schema.types
+  inds = indexin(t.cols, collect(names))
+  Tables.Schema(t.cols, types[inds])
+end
+
+Tables.materializer(t::TableSelection) = 
+  Tables.materializer(t.table)
+
 const ColSpec = Union{Vector{Symbol}, Regex}
 
 """
@@ -19,6 +59,9 @@ struct Select{S<:ColSpec} <: Stateless
   cols::S
 end
 
+# to avoid StackOverflowError in Select() and Select(())
+Select(::Tuple{}) = throw(ArgumentError("Cannot create a Select object without arguments."))
+
 Select(cols::T...) where {T<:Union{AbstractString, Symbol}} = 
   Select(cols)
 
@@ -42,6 +85,10 @@ function apply(transform::Select, table)
   select  = _select(transform.cols, allcols)
   reject  = setdiff(allcols, select)
 
+  # validate selections
+  @assert !isempty(select) "Invalid selection"
+  @assert select ⊆ Tables.columnnames(table) "Invalid selection"
+
   # keep track of indices to revert later
   sinds = indexin(select, allcols)
   rinds = indexin(reject, allcols)
@@ -55,17 +102,10 @@ function apply(transform::Select, table)
   # original columns
   cols = Tables.columns(table)
 
-  # selected columns
-  scols = [Tables.getcolumn(cols, name) for name in select]
-
   # rejected columns
   rcols = [Tables.getcolumn(cols, name) for name in reject]
 
-  # table with selected columns
-  𝒯 = (; zip(select, scols)...)
-  stable = 𝒯 |> Tables.materializer(table)
-
-  stable, (reject, rcols, sperm, rinds)
+  TableSelection(table, select), (reject, rcols, sperm, rinds)
 end
 
 function revert(::Select, newtable, cache)
@@ -88,6 +128,9 @@ function revert(::Select, newtable, cache)
   𝒯 |> Tables.materializer(newtable)
 end
 
+# reverting a single TableSelection is trivial
+revert(::Select, newtable::TableSelection, cache) = newtable.table
+
 """
     Reject(col₁, col₂, ..., colₙ)
     Reject([col₁, col₂, ..., colₙ])
@@ -103,6 +146,9 @@ struct Reject{S<:ColSpec} <: Stateless
   cols::S
 end
 
+# to avoid StackOverflowError in Reject() and Reject(())
+Reject(::Tuple{}) = throw(ArgumentError("Cannot create a Reject object with no arguments."))
+
 Reject(cols::T...) where {T<:Union{AbstractString, Symbol}} = 
   Reject(cols)
 


Test patch:
diff --git a/test/transforms.jl b/test/transforms.jl
index 63c58b0f..aaed9c9d 100644
--- a/test/transforms.jl
+++ b/test/transforms.jl
@@ -12,60 +12,60 @@
 
     T = Select(:f, :d)
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:f, :d)
+    @test Tables.columnnames(n) == [:f, :d]
     tₒ = revert(T, n, c)
     @test t == tₒ
 
     T = Select(:f, :d, :b)
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:f, :d, :b)
+    @test Tables.columnnames(n) == [:f, :d, :b]
     tₒ = revert(T, n, c)
     @test t == tₒ
 
     T = Select(:d, :c, :b)
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:d, :c, :b)
+    @test Tables.columnnames(n) == [:d, :c, :b]
     tₒ = revert(T, n, c)
     @test t == tₒ
 
     T = Select(:e, :c, :b, :a)
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:e, :c, :b, :a)
+    @test Tables.columnnames(n) == [:e, :c, :b, :a]
     tₒ = revert(T, n, c)
     @test t == tₒ
 
     # selection with tuples
     T = Select((:e, :c, :b, :a))
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:e, :c, :b, :a)
+    @test Tables.columnnames(n) == [:e, :c, :b, :a]
     tₒ = revert(T, n, c)
     @test t == tₒ
 
     # selection with vectors
     T = Select([:e, :c, :b, :a])
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:e, :c, :b, :a)
+    @test Tables.columnnames(n) == [:e, :c, :b, :a]
     tₒ = revert(T, n, c)
     @test t == tₒ
 
     # selection with strings
     T = Select("d", "c", "b")
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:d, :c, :b)
+    @test Tables.columnnames(n) == [:d, :c, :b]
     tₒ = revert(T, n, c)
     @test t == tₒ
 
     # selection with tuple of strings
     T = Select(("d", "c", "b"))
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:d, :c, :b)
+    @test Tables.columnnames(n) == [:d, :c, :b]
     tₒ = revert(T, n, c)
     @test t == tₒ
 
     # selection with vector of strings
     T = Select(["d", "c", "b"])
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:d, :c, :b)
+    @test Tables.columnnames(n) == [:d, :c, :b]
     tₒ = revert(T, n, c)
     @test t == tₒ
 
@@ -83,7 +83,7 @@
     # selection with Regex
     T = Select(r"[dcb]")
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:b, :c, :d) # the order of columns is preserved
+    @test Tables.columnnames(n) == [:b, :c, :d] # the order of columns is preserved
     tₒ = revert(T, n, c)
     @test t == tₒ
 
@@ -96,16 +96,33 @@
     # select columns whose names contain the character x
     T = Select(r"x")
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:x1, :x2)
+    @test Tables.columnnames(n) == [:x1, :x2]
     tₒ = revert(T, n, c)
     @test t == tₒ
 
     # select columns whose names contain the character y
     T = Select(r"y")
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:y1, :y2)
+    @test Tables.columnnames(n) == [:y1, :y2]
     tₒ = revert(T, n, c)
     @test t == tₒ
+
+    # throws: Select without arguments
+    @test_throws ArgumentError Select()
+    @test_throws ArgumentError Select(())
+
+    # throws: empty selection
+    @test_throws AssertionError apply(Select(r"a"), t)
+    @test_throws AssertionError apply(Select(Symbol[]), t)
+    @test_throws AssertionError apply(Select(String[]), t)
+
+    # throws: columns that do not exist in the original table
+    @test_throws AssertionError apply(Select(:x3, :y3), t)
+    @test_throws AssertionError apply(Select([:x3, :y3]), t)
+    @test_throws AssertionError apply(Select((:x3, :y3)), t)
+    @test_throws AssertionError apply(Select("x3", "y3"), t)
+    @test_throws AssertionError apply(Select(["x3", "y3"]), t)
+    @test_throws AssertionError apply(Select(("x3", "y3")), t)
   end
 
   @testset "Reject" begin
@@ -119,60 +136,60 @@
 
     T = Reject(:f, :d)
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:a, :b, :c, :e)
+    @test Tables.columnnames(n) == [:a, :b, :c, :e]
     tₒ = revert(T, n, c)
     @test t == tₒ
 
     T = Reject(:f, :d, :b)
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:a, :c, :e)
+    @test Tables.columnnames(n) == [:a, :c, :e]
     tₒ = revert(T, n, c)
     @test t == tₒ
 
     T = Reject(:d, :c, :b)
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:a, :e, :f)
+    @test Tables.columnnames(n) == [:a, :e, :f]
     tₒ = revert(T, n, c)
     @test t == tₒ
 
     T = Reject(:e, :c, :b, :a)
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:d, :f)
+    @test Tables.columnnames(n) == [:d, :f]
     tₒ = revert(T, n, c)
     @test t == tₒ
 
     # rejection with tuples
     T = Reject((:e, :c, :b, :a))
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:d, :f)
+    @test Tables.columnnames(n) == [:d, :f]
     tₒ = revert(T, n, c)
     @test t == tₒ
 
     # rejection with vectors
     T = Reject([:e, :c, :b, :a])
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:d, :f)
+    @test Tables.columnnames(n) == [:d, :f]
     tₒ = revert(T, n, c)
     @test t == tₒ
 
     # rejection with strings
     T = Reject("d", "c", "b")
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:a, :e, :f)
+    @test Tables.columnnames(n) == [:a, :e, :f]
     tₒ = revert(T, n, c)
     @test t == tₒ
 
     # rejection with tuple of strings
     T = Reject(("d", "c", "b"))
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:a, :e, :f)
+    @test Tables.columnnames(n) == [:a, :e, :f]
     tₒ = revert(T, n, c)
     @test t == tₒ
 
     # rejection with vector of strings
     T = Reject(["d", "c", "b"])
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:a, :e, :f)
+    @test Tables.columnnames(n) == [:a, :e, :f]
     tₒ = revert(T, n, c)
     @test t == tₒ
 
@@ -190,7 +207,7 @@
     # rejection with Regex
     T = Reject(r"[dcb]")
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:a, :e, :f) # the order of columns is preserved
+    @test Tables.columnnames(n) == [:a, :e, :f] # the order of columns is preserved
     tₒ = revert(T, n, c)
     @test t == tₒ
 
@@ -203,16 +220,59 @@
     # reject columns whose names contain the character x
     T = Reject(r"x")
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:y1, :y2)
+    @test Tables.columnnames(n) == [:y1, :y2]
     tₒ = revert(T, n, c)
     @test t == tₒ
 
     # reject columns whose names contain the character y
     T = Reject(r"y")
     n, c = apply(T, t)
-    @test Tables.columnnames(n) == (:x1, :x2)
+    @test Tables.columnnames(n) == [:x1, :x2]
     tₒ = revert(T, n, c)
     @test t == tₒ
+
+    # throws: Reject without arguments
+    @test_throws ArgumentError Reject()
+    @test_throws ArgumentError Reject(())
+
+    # throws: reject all columns
+    @test_throws AssertionError apply(Reject(r"[xy]"), t)
+    @test_throws AssertionError apply(Reject(:x1, :x2, :y1, :y2), t)
+    @test_throws AssertionError apply(Reject([:x1, :x2, :y1, :y2]), t)
+    @test_throws AssertionError apply(Reject((:x1, :x2, :y1, :y2)), t)
+    @test_throws AssertionError apply(Reject("x1", "x2", "y1", "y2"), t)
+    @test_throws AssertionError apply(Reject(["x1", "x2", "y1", "y2"]), t)
+    @test_throws AssertionError apply(Reject(["x1", "x2", "y1", "y2"]), t)
+  end
+
+  @testset "TableSelection" begin
+    a = rand(4000)
+    b = rand(4000)
+    c = rand(4000)
+    d = rand(4000)
+    e = rand(4000)
+    f = rand(4000)
+    t = Table(; a, b, c, d, e, f)
+
+    # Tables.jl interface
+    s = TableTransforms.TableSelection(t, [:a, :b, :e])
+    @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.schema(s).names == (:a, :b, :e)
+    @test Tables.schema(s).types == (Float64, Float64, Float64)
+    @test Tables.materializer(s) == Tables.materializer(t)
+
+    # getcolumn
+    cols = Tables.columns(t)
+    @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])
   end
 
   @testset "Rename" begin