Repository: JuliaML/TableTransforms.jl (julia)

Issue:
Consider restricting Levels to columns with categorical scitype
Currently the Levels transform contributed by @Frank-III and refactored by @eliascarv can operate directly on any type of column, including columns with the wrong scitype. I think we could consider a more consistent approach where we ask users to be explicit about the scitypes using the Coerce transform first. Even though this is a breaking change, it will provide more safety in real-world applications of pipelines.

Pull request:
closes #95

Patch:
diff --git a/src/transforms/levels.jl b/src/transforms/levels.jl
index 9a5624b4..7d6b9ebc 100644
--- a/src/transforms/levels.jl
+++ b/src/transforms/levels.jl
@@ -29,30 +29,30 @@ Levels(; kwargs...) = throw(ArgumentError("Cannot create a Levels object without
 
 isrevertible(transform::Levels) = true
 
-_categorical(x::AbstractVector, l, o) =
-  categorical(x, levels=l, ordered=o), y -> unwrap.(y)
-
-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)
   snames = choose(transform.colspec, names)
   ordered = choose(transform.ordered, snames)
-  levels = transform.levels
+  tlevels = transform.levels
   
   results = map(names) do nm
     x = Tables.getcolumn(cols, nm)
+    
     if nm ∈ snames
+      assert_categorical(x)
+      
       o = nm ∈ ordered
-      l = levels[findfirst(==(nm), snames)]
-      return _categorical(x, l, o)
+      l = tlevels[findfirst(==(nm), snames)]
+      y = categorical(x, levels=l, ordered=o)
+      
+      xl, xo = levels(x), isordered(x)
+      revfunc = y -> categorical(y, levels=xl, ordered=xo)
+    else
+      y, revfunc = x, identity
     end
-    x, identity
+
+    y, revfunc
   end
 
   columns, cache = first.(results), last.(results)


Test patch:
diff --git a/test/transforms.jl b/test/transforms.jl
index 4d607327..5fbf66ea 100644
--- a/test/transforms.jl
+++ b/test/transforms.jl
@@ -1061,9 +1061,9 @@
   end
 
   @testset "Levels" begin
-    a = rand([true, false], 50)
-    b = rand(["y", "n"], 50)
-    c = rand(1:3, 50)
+    a = categorical(rand([true, false], 50))
+    b = categorical(rand(["y", "n"], 50))
+    c = categorical(rand(1:3, 50))
     t = Table(; a, b, c)
 
     T = Levels(2 => ["n", "y", "m"])
@@ -1091,11 +1091,10 @@
     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]) 
-    d = [1, 23, 5, 7, 7]
+    d = categorical([1, 23, 5, 7, 7])
     e = categorical([2, 3, 1, 4, 1])
     t = Table(; a, b, c, d, e)
 
@@ -1148,12 +1147,11 @@
     @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
 
     a = rand([true, false], 50)
-    b = rand(["y", "n"], 50)
-    c = rand(1:3, 50)
+    b = categorical(rand(["y", "n"], 50))
+    c = categorical(rand(1:3, 50))
     t = Table(; a, b, c)
 
     # throws: Levels without arguments
@@ -1165,6 +1163,10 @@
     T = Levels("x" => ["n", "y", "m"], "y" => 1:4)
     @test_throws AssertionError apply(T, t)
 
+    # throws: non categorical collumn
+    T = Levels(:a => [true, false], ordered=[:a])
+    @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)