Unitufy: Add support for non-numeric columns (#281)

JuliaML · Jun 17, 2024 · 1211dd5 · 1211dd5
1 parent 27cad23
commit 1211dd5
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/src/transforms/unitify.jl b/src/transforms/unitify.jl
@@ -45,7 +45,7 @@ function applyfeat(::Unitify, feat, prep)
   pairs = map(names) do name
     x = Tables.getcolumn(cols, name)
     newname, unit = _unitify(name)
-    newname => x * unit
+    newname => _addunit(x, unit)
   end
 
   newfeat = (; pairs...) |> Tables.materializer(feat)
@@ -59,7 +59,7 @@ function revertfeat(::Unitify, newfeat, fcache)
   onames = fcache
   ocolumns = map(names) do name
     x = Tables.getcolumn(cols, name)
-    ustrip.(x)
+    first(_dropunit(x))
   end
 
   𝒯 = (; zip(onames, ocolumns)...)

diff --git a/test/transforms/unitify.jl b/test/transforms/unitify.jl
@@ -42,4 +42,15 @@
   @test unit(eltype(n.a)) === NoUnits
   tₒ = revert(T, n, c)
   @test tₒ == t
+
+  # non-numeric columns
+  t = Table(; anm => a, :b => rand('a':'z', 10), :c => categorical(rand(["yes", "no"], 10)))
+  T = Unitify()
+  n, c = apply(T, t)
+  @test Tables.schema(n).names == (:a, :b, :c)
+  @test unit(eltype(n.a)) === u"m"
+  @test eltype(n.b) <: Char
+  @test eltype(n.c) <: CategoricalValue
+  tₒ = revert(T, n, c)
+  @test tₒ == t
 end