Refactor Levels Tranform (#91)

* Refactor Levels transform

* Add nothing to ColSpec

* Update ColSpec tests

* Remove unused constructor

* Update docstring

* Update Levels tests

* Add more tests

* Add header in levels.jl

* Update _categorical function

* Update choose docstring

* Add and update old tests

* Apply suggestions

Co-authored-by: Júlio Hoffimann <[email protected]>

* Fix typo

Co-authored-by: Júlio Hoffimann <[email protected]>

Author: eliascarv

src/colspec.jl

@@ -40,7 +40,7 @@ function MyTransform(args::T...) where {T<:ColSelector}
 end
 ```
 """
-const ColSpec = Union{Vector{T},NTuple{N,T},Regex,Colon} where {N,T<:ColSelector}
+const ColSpec = Union{Vector{T},NTuple{N,T},Regex,Colon,Nothing} where {N,T<:ColSelector}
 
 """
     choose(colspec::ColSpec, names) -> Vector{Symbol}
@@ -69,6 +69,18 @@ julia> choose(r"[ace]", names)
  :a
  :c
  :e
+
+julia> choose(:, names)
+6-element Vector{Symbol}:
+  :a
+  :b
+  :c
+  :d
+  :e
+  :f
+
+julia> choose(nothing, names)
+Symbol[]
 ```
 """
 function choose(colspec::Vector{Symbol}, names)
@@ -100,3 +112,4 @@ choose(colspec::Regex, names::Tuple) =
 
 choose(::Colon, names::Vector) = names
 choose(::Colon, names::Tuple) = collect(names)
+choose(::Nothing, names) = Symbol[]

src/transforms/levels.jl

@@ -1,72 +1,76 @@
+# ------------------------------------------------------------------
+# Licensed under the MIT License. See LICENSE in the project root.
+# ------------------------------------------------------------------
+
 """
-    Levels(:a => ["yes", "no"])
+    Levels(col₁ => levels₁, col₂ => levels₂, ..., colₙ => levelsₙ; ordered=nothing)
 
-Return a copy of the table with specified levels and orders for categorical columns
-allowing only changing the order of the column.
+Convert columns `col₁`, `col₂`, ..., `colₙ` to categorical arrays with given levels `levels₁`, `levels₂`, ..., `levelsₙ`.
+Optionally, specify which columns are `ordered`.
 
 # Examples
 
 ```julia
-Levels(:a => ["yes, "no"], :c => [1, 2, 4], :d => ["a", "b", "c"])
-Levels("a" => ["yes", "no"], "c" => [1, 2, 4], ordered = ["a", "c"])
-Levels(:a => ["yes", "no"], :c => [1, 2, 4], :d => [1, 23, 5, 7], ordered = [:a, :b, :c])
+Levels(1 => 1:3, 2 => ["a", "b"], ordered=r"a")
+Levels(:a => 1:3, :b => ["a", "b"], ordered=[:a])
+Levels("a" => 1:3, "b" => ["a", "b"], ordered=["b"])
 ```
 """
-struct Levels{K} <: Stateless
-  levelspec::K
-  ordered::Vector{Symbol}
+struct Levels{S<:ColSpec,O<:ColSpec,L} <: Stateless
+  colspec::S
+  ordered::O
+  levels::L
 end
 
-Levels(pairs::Pair{Symbol}...; ordered=Symbol[]) =
-  Levels(NamedTuple(pairs), ordered)
+Levels(pairs::Pair{T}...; ordered::ColSpec=nothing) where {T<:ColSelector} =
+  Levels(first.(pairs), ordered, last.(pairs))
 
-Levels(pairs::Pair{K}...; ordered=K[]) where {K<:AbstractString} =
-  Levels(NamedTuple(Symbol(k) => v for (k,v) in pairs), Symbol.(ordered))
+Levels(; kwargs...) = throw(ArgumentError("Cannot create a Levels object without arguments."))
 
 isrevertible(transform::Levels) = true
 
-# when the col is already a categorical array and wanna change levels and order
-_categorify(l::AbstractVector, x::CategoricalVector, o) =
-  categorical(x, levels=l, ordered=o), levels(x)
-
-# when the col is normal array and want to change to categorical array
-_categorify(l::AbstractVector, x::AbstractVector, o) =
-  categorical(x, levels=l, ordered=o), unwrap 
+_categorical(x::AbstractVector, l, o) =
+  categorical(x, levels=l, ordered=o), y -> unwrap.(y)
 
-# when the col is not need for change or convert back to normal array
-_categorify(f::Function, x::AbstractVector, o) =
-  o ? (categorical(x, ordered=true), levels(x)) : (f.(x), f)
+function _categorical(x::CategoricalArray, l, o)
+  xl, xo = levels(x), isordered(x)
+  revfunc = y -> categorical(y, levels=xl, ordered=xo)
+  categorical(x, levels=l, ordered=o), revfunc
+end
 
 function apply(transform::Levels, table)
   cols = Tables.columns(table)
   names = Tables.columnnames(cols)
-
-  result = map(names) do nm
+  snames = choose(transform.colspec, names)
+  ordered = choose(transform.ordered, snames)
+  levels = transform.levels
+  
+  results = map(names) do nm
     x = Tables.getcolumn(cols, nm)
-    l = get(transform.levelspec, nm, identity)
-    o = nm ∈ transform.ordered
-    _categorify(l, x, o)
+    if nm ∈ snames
+      o = nm ∈ ordered
+      l = levels[findfirst(==(nm), snames)]
+      return _categorical(x, l, o)
+    end
+    x, identity
   end
-  
-  categ = first.(result)
-  cache = last.(result)
 
-  𝒯 = (; zip(names, categ)...)
+  columns, cache = first.(results), last.(results)
+
+  𝒯 = (; zip(names, columns)...)
   newtable = 𝒯 |> Tables.materializer(table)
-  
   newtable, cache
 end
 
-function revert(transform::Levels, newtable, cache)
+function revert(::Levels, newtable, cache)
   cols = Tables.columns(newtable)
   names = Tables.columnnames(cols)
 
-  ocols = map(zip(cache, names)) do (f, nm)
+  columns = map(names, cache) do nm, revfunc
     x = Tables.getcolumn(cols, nm)
-    c, _ = _categorify(f, x, false)
-    c
+    revfunc(x)
   end
 
-  𝒯 = (; zip(names, ocols)...)
+  𝒯 = (; zip(names, columns)...)
   𝒯 |> Tables.materializer(newtable)
-end
+end

test/colspec.jl

@@ -57,6 +57,12 @@
   snames = TableTransforms.choose(:, tupnames)
   @test snames == [:a, :b, :c, :d, :e, :f]
 
+  # nothing
+  snames = TableTransforms.choose(nothing, vecnames)
+  @test snames == Symbol[]
+  snames = TableTransforms.choose(nothing, tupnames)
+  @test snames == Symbol[]
+
   # throws
   @test_throws AssertionError TableTransforms.choose(r"x", vecnames)
   @test_throws AssertionError TableTransforms.choose(r"x", tupnames)

test/transforms.jl

@@ -1275,8 +1275,39 @@
     n, c = apply(T, t)
     @test_throws AssertionError revert(T, n, c)
   end
-  
+
   @testset "Levels" begin
+    a = rand([true, false], 50)
+    b = rand(["y", "n"], 50)
+    c = rand(1:3, 50)
+    t = Table(; a, b, c)
+
+    T = Levels(2 => ["n", "y", "m"])
+    n, c = apply(T, t)
+    @test levels(n.b) == ["n", "y", "m"]
+    @test isordered(n.b) == false
+    tₒ = revert(T, n, c)
+    @test tₒ == t
+
+    T = Levels(:b => ["n", "y", "m"], :c => 1:4, ordered=[:c])
+    n, c = apply(T, t)
+    @test levels(n.b) == ["n", "y", "m"]
+    @test isordered(n.b) == false
+    @test levels(n.c) == [1, 2, 3, 4]
+    @test isordered(n.c) == true
+    tₒ = revert(T, n, c)
+    @test tₒ == t
+
+    T = Levels("b" => ["n", "y", "m"], "c" => 1:4, ordered=["b"])
+    n, c = apply(T, t)
+    @test levels(n.b) == ["n", "y", "m"]
+    @test isordered(n.b) == true
+    @test levels(n.c) == [1, 2, 3, 4]
+    @test isordered(n.c) == false
+    tₒ = revert(T, n, c)
+    @test tₒ == t
+
+    # categorical columns
     a = categorical(["yes", "no", "no", "no", "yes"]) 
     b = categorical([1, 2, 4, 2, 8], ordered=false) 
     c = categorical([1, 2, 1, 2, 1]) 
@@ -1288,6 +1319,7 @@
     n, c = apply(T, t)
     @test levels(n.a) == ["yes", "no"]
     @test levels(n.c) == [1, 2, 4]
+    @test levels(n.d) == [1, 23, 5, 7]
     @test levels(n.e) == [1, 2, 3, 4, 5]
     tₒ = revert(T, n, c)
     @test levels(tₒ.a) == ["no", "yes"]
@@ -1302,35 +1334,60 @@
     @test levels(tₒ.a) == ["no", "yes"]
     @test levels(tₒ.c) == [1, 2]
 
-    T = Levels(:a => ["yes", "no"], :c => [1, 2, 4], :d => [1, 23, 5, 7], ordered=[:b])
+    T = Levels(:a => ["yes", "no"], :c => [1, 2, 4], :d => [1, 23, 5, 7])
     n, c = apply(T, t)
     @test levels(n.a) == ["yes", "no"]
-    @test isordered(n.b) == true
     @test levels(n.c) == [1, 2, 4]
+    @test levels(n.d) == [1, 23, 5, 7]
     tₒ = revert(T, n, c)
     @test levels(tₒ.a) == ["no", "yes"]
     @test levels(tₒ.c) == [1, 2]
-    @test isordered(tₒ.b) == false
 
-    T = Levels("a" => ["yes", "no"], "c" => [1, 2, 4], "e" => 5:-1:1, ordered=["b", "e"])
+    T = Levels("a" => ["yes", "no"], "c" => [1, 2, 4], "e" => 5:-1:1, ordered=["e"])
     n, c = apply(T, t)
-    @test isordered(n.b) == true
-    @test levels(n.e) == [5, 4, 3, 2, 1]
+    @test levels(n.a) == ["yes", "no"]
     @test levels(n.c) == [1, 2, 4]
+    @test levels(n.e) == [5, 4, 3, 2, 1]
+    @test isordered(n.a) == false
+    @test isordered(n.c) == false
+    @test isordered(n.e) == true
     tₒ = revert(T, n, c)
-    @test isordered(tₒ.b) == false
     @test levels(tₒ.e) == [1, 2, 3, 4]
+    @test isordered(tₒ.e) == false
 
-    T = Levels(:a => ["yes", "no"], :c => [1, 2, 4], :d => [1, 23, 5, 7], ordered=[:a, :b, :d])
+    T = Levels(:a => ["yes", "no"], :c => [1, 2, 4], :d => [1, 23, 5, 7], ordered=[:a, :d])
     n, c = apply(T, t)
     @test levels(n.a) == ["yes", "no"]
+    @test levels(n.c) == [1, 2, 4]
+    @test levels(n.d) == [1, 23, 5, 7]
     @test isordered(n.a) == true
-    @test isordered(n.b) == true
+    @test isordered(n.c) == false
     @test isordered(n.d) == true
     tₒ = revert(T, n, c)
     @test typeof(tₒ.d) == Vector{Int64}
     @test isordered(tₒ.a) == false
-    @test isordered(tₒ.b) == false
+
+    a = rand([true, false], 50)
+    b = rand(["y", "n"], 50)
+    c = rand(1:3, 50)
+    t = Table(; a, b, c)
+
+    # throws: Levels without arguments
+    @test_throws ArgumentError Levels()
+
+    # throws: columns that do not exist in the original table
+    T = Levels(:x => ["n", "y", "m"], :y => 1:4)
+    @test_throws AssertionError apply(T, t)
+    T = Levels("x" => ["n", "y", "m"], "y" => 1:4)
+    @test_throws AssertionError apply(T, t)
+
+    # throws: invalid ordered column selection
+    T = Levels(:b => ["n", "y", "m"], :c => 1:4, ordered=[:a])
+    @test_throws AssertionError apply(T, t)
+    T = Levels("b" => ["n", "y", "m"], "c" => 1:4, ordered=["a"])
+    @test_throws AssertionError apply(T, t)
+    T = Levels("b" => ["n", "y", "m"], "c" => 1:4, ordered=r"xy")
+    @test_throws AssertionError apply(T, t)
   end
 
   @testset "Sort" begin