Repository: JuliaML/TableTransforms.jl (julia)

Issue:
OneHot should probably produce categorical arrays
Currently the resulting columns are just Vector{Bool}, maybe they should be categorical arrays.

Pull request description:
…form

fixes #158 

### Default Behavior (categorical is true by default)

![image](https://user-images.githubusercontent.com/52173002/208770312-0c20f49b-f0bc-4878-b365-552671a15d5d.png)

### Behavior When Categorical is set to false

![image](https://user-images.githubusercontent.com/52173002/208770417-9bfb8a19-a762-4b53-9045-7853f3f614a9.png)

Patch:
diff --git a/src/transforms/onehot.jl b/src/transforms/onehot.jl
index 8147321f..0235c548 100644
--- a/src/transforms/onehot.jl
+++ b/src/transforms/onehot.jl
@@ -3,10 +3,12 @@
 # ------------------------------------------------------------------
 
 """
-    OneHot(col)
-    
+    OneHot(col; categ=true)
+
 Transforms categorical column `col` into one-hot columns of levels
 returned by the `levels` function of CategoricalArrays.jl.
+The `categ` option can be used to convert resulting
+columns to categorical arrays as opposed to boolean vectors.
 
 # Examples
 
@@ -14,23 +16,27 @@ returned by the `levels` function of CategoricalArrays.jl.
 OneHot(1)
 OneHot(:a)
 OneHot("a")
+OneHot("a", categ=false)
 ```
 """
 struct OneHot{S<:ColSpec} <: StatelessFeatureTransform
   colspec::S
-  function OneHot(col::Col)
+  categ::Bool
+  function OneHot(col, categ)
     cs = colspec([col])
-    new{typeof(cs)}(cs)
+    new{typeof(cs)}(cs, categ)
   end
 end
 
+OneHot(col; categ=true) = OneHot(col, categ)
+
 isrevertible(::Type{<:OneHot}) = true
 
 function applyfeat(transform::OneHot, feat, prep)
   cols = Tables.columns(feat)
   names = Tables.columnnames(cols) |> collect
   columns = Any[Tables.getcolumn(cols, nm) for nm in names]
-  
+
   name = choose(transform.colspec, names)[1]
   ind = findfirst(==(name), names)
   x = columns[ind]
@@ -48,6 +54,9 @@ function applyfeat(transform::OneHot, feat, prep)
 
   newnms, newcols = first.(onehot), last.(onehot)
 
+  # convert to categorical arrays if necessary
+  newcols = transform.categ ? categorical.(newcols, levels=[false, true]) : newcols
+  
   splice!(names, ind, newnms)
   splice!(columns, ind, newcols)
   
@@ -65,7 +74,7 @@ function revertfeat(::OneHot, newfeat, fcache)
   
   oname, inds, levels, ordered = fcache
   x = map(zip(columns[inds]...)) do row
-    levels[findfirst(row)]
+    levels[findfirst(==(true), row)]
   end
 
   ocolumn = categorical(x; levels, ordered)


Test patch:
diff --git a/test/shows.jl b/test/shows.jl
index deb7f018..c0c3a9e3 100644
--- a/test/shows.jl
+++ b/test/shows.jl
@@ -200,13 +200,14 @@
 
     # compact mode
     iostr = sprint(show, T)
-    @test iostr == "OneHot([:a])"
+    @test iostr == "OneHot([:a], true)"
 
     # full mode
     iostr = sprint(show, MIME("text/plain"), T)
     @test iostr == """
     OneHot transform
-    └─ colspec = [:a]"""
+    ├─ colspec = [:a]
+    └─ categ = true"""
   end
 
   @testset "Identity" begin
diff --git a/test/transforms/onehot.jl b/test/transforms/onehot.jl
index f5a0ea3a..41ee2ef5 100644
--- a/test/transforms/onehot.jl
+++ b/test/transforms/onehot.jl
@@ -4,7 +4,39 @@
   c = categorical([3, 2, 2, 1, 1, 3])
   t = Table(; a, b, c)
 
-  T = OneHot(1)
+  T = OneHot(1; categ=true)
+  n, c = apply(T, t)
+  @test Tables.columnnames(n) == (:a_false, :a_true, :b, :c)
+  @test n.a_false == categorical(Bool[1, 0, 0, 1, 0, 0])
+  @test n.a_true  == categorical(Bool[0, 1, 1, 0, 1, 1])
+  @test n.a_false isa CategoricalVector{Bool}
+  @test n.a_true  isa CategoricalVector{Bool}
+  tₒ = revert(T, n, c)
+  @test t == tₒ
+
+  T = OneHot(:b; categ=true)
+  n, c = apply(T, t)
+  @test Tables.columnnames(n) == (:a, :b_f, :b_m, :c)
+  @test n.b_f == categorical(Bool[0, 1, 0, 0, 0, 1])
+  @test n.b_m == categorical(Bool[1, 0, 1, 1, 1, 0])
+  @test n.b_f isa CategoricalVector{Bool}
+  @test n.b_m isa CategoricalVector{Bool}
+  tₒ = revert(T, n, c)
+  @test t == tₒ
+
+  T = OneHot("c"; categ=true)
+  n, c = apply(T, t)
+  @test Tables.columnnames(n) == (:a, :b, :c_1, :c_2, :c_3)
+  @test n.c_1 == categorical(Bool[0, 0, 0, 1, 1, 0])
+  @test n.c_2 == categorical(Bool[0, 1, 1, 0, 0, 0])
+  @test n.c_3 == categorical(Bool[1, 0, 0, 0, 0, 1])
+  @test n.c_1 isa CategoricalVector{Bool}
+  @test n.c_2 isa CategoricalVector{Bool}
+  @test n.c_3 isa CategoricalVector{Bool}
+  tₒ = revert(T, n, c)
+  @test t == tₒ
+
+  T = OneHot(1; categ=false)
   n, c = apply(T, t)
   @test Tables.columnnames(n) == (:a_false, :a_true, :b, :c)
   @test n.a_false == Bool[1, 0, 0, 1, 0, 0]
@@ -12,7 +44,7 @@
   tₒ = revert(T, n, c)
   @test t == tₒ
 
-  T = OneHot(:b)
+  T = OneHot(:b; categ=false)
   n, c = apply(T, t)
   @test Tables.columnnames(n) == (:a, :b_f, :b_m, :c)
   @test n.b_f == Bool[0, 1, 0, 0, 0, 1]
@@ -20,7 +52,7 @@
   tₒ = revert(T, n, c)
   @test t == tₒ
 
-  T = OneHot("c")
+  T = OneHot("c"; categ=false)
   n, c = apply(T, t)
   @test Tables.columnnames(n) == (:a, :b, :c_1, :c_2, :c_3)
   @test n.c_1 == Bool[0, 0, 0, 1, 1, 0]
@@ -35,7 +67,7 @@
   b_m = rand(10)
   t   = Table(; b, b_f, b_m)
 
-  T = OneHot(:b)
+  T = OneHot(:b; categ=false)
   n, c = apply(T, t)
   @test Tables.columnnames(n) == (:b_f_, :b_m_, :b_f, :b_m)
   @test n.b_f_ == Bool[0, 1, 0, 0, 0, 1]
@@ -50,7 +82,7 @@
   b_m_ = rand(10)
   t    = Table(; b, b_f, b_m, b_f_, b_m_)
 
-  T = OneHot(:b)
+  T = OneHot(:b; categ=false)
   n, c = apply(T, t)
   @test Tables.columnnames(n) == (:b_f__, :b_m__, :b_f, :b_m, :b_f_, :b_m_)
   @test n.b_f__ == Bool[0, 1, 0, 0, 0, 1]
@@ -70,4 +102,4 @@
   # invalid column selection
   @test_throws AssertionError apply(OneHot(:c), t)
   @test_throws AssertionError apply(OneHot("c"), t)
-end
+end
\ No newline at end of file