Handle unitless columns in Unit transform (#297)

* Handle unitless columns in Unit transform

* Add string column to new test

Author: juliohm

src/transforms/unit.jl

@@ -5,15 +5,14 @@
 """
     Unit(unit)
 
-Converts the units of all columns in the table to `unit`.
+Converts the units of all columns in the table to `unit` from Unitful.jl.
 
     Unit(cols₁ => unit₁, cols₂ => unit₂, ..., colsₙ => unitₙ)
 
 Converts the units of selected columns `cols₁`, `cols₂`, ..., `colsₙ`
 to `unit₁`, `unit₂`, ... `unitₙ`.
 
-The column selection can be a single column identifier (index or name),
-a collection of identifiers or a regular expression (regex).
+Unitless columns become unitful if they are explicitly selected.
 
 # Examples
 
@@ -39,10 +38,6 @@ Unit(pairs::Pair...) = Unit(collect(selector.(first.(pairs))), collect(last.(pai
 
 isrevertible(::Type{<:Unit}) = true
 
-_uconvert(u, x) = _uconvert(nonmissingtype(eltype(x)), u, x)
-_uconvert(::Type, _, x) = (x, nothing)
-_uconvert(::Type{Q}, u, x) where {Q<:AbstractQuantity} = (map(v -> uconvert(u, v), x), unit(Q))
-
 function applyfeat(transform::Unit, feat, prep)
   cols = Tables.columns(feat)
   names = Tables.columnnames(cols)
@@ -59,7 +54,7 @@ function applyfeat(transform::Unit, feat, prep)
     x = Tables.getcolumn(cols, name)
     if haskey(unitdict, name)
       u = unitdict[name]
-      _uconvert(u, x)
+      _withunit(u, x)
     else
       (x, nothing)
     end
@@ -80,9 +75,22 @@ function revertfeat(::Unit, newfeat, fcache)
   ounits = fcache
   columns = map(names, ounits) do name, u
     x = Tables.getcolumn(cols, name)
-    isnothing(u) ? x : map(v -> uconvert(u, v), x)
+    _withoutunit(u, x)
   end
 
   𝒯 = (; zip(names, columns)...)
   𝒯 |> Tables.materializer(newfeat)
 end
+
+_withunit(u, x) = _withunit(nonmissingtype(eltype(x)), u, x)
+_withunit(::Type{Q}, u, x) where {Q<:AbstractQuantity} = (map(v -> uconvert(u, v), x), unit(Q))
+_withunit(::Type{Q}, u, x) where {Q<:Number} = (x * u, NoUnits)
+
+function _withoutunit(u, x)
+  if u === NoUnits
+    map(ustrip, x)
+  else
+    map(xᵢ -> uconvert(u, xᵢ), x)
+  end
+end
+_withoutunit(::Nothing, x) = x

test/transforms/unit.jl

@@ -5,25 +5,23 @@
   b = [300, 500, missing, 800, missing, 400] * u"cm"
   c = [8, 2, 5, 7, 9, 4] * u"km"
   d = [0.3, 0.1, 0.9, 0.2, 0.7, 0.4]
-  e = ["no", "no", "yes", "yes", "no", "yes"]
-  t = Table(; a, b, c, d, e)
+  t = Table(; a, b, c, d)
 
   T = Unit(u"m")
   n, c = apply(T, t)
   @test unit(eltype(n.a)) == u"m"
   @test unit(eltype(n.b)) == u"m"
   @test unit(eltype(n.c)) == u"m"
-  @test eltype(n.d) <: Float64
-  @test eltype(n.e) <: String
+  @test unit(eltype(n.d)) == u"m"
   tₒ = revert(T, n, c)
   @test unit(eltype(tₒ.a)) == u"m"
   @test unit(eltype(tₒ.b)) == u"cm"
   @test unit(eltype(tₒ.c)) == u"km"
+  @test unit(eltype(tₒ.d)) == NoUnits
   @test all(isapprox.(tₒ.a, t.a))
   @test all(isapprox.(skipmissing(tₒ.b), skipmissing(t.b)))
   @test all(isapprox.(tₒ.c, t.c))
   @test tₒ.d == t.d
-  @test tₒ.e == t.e
 
   a = [2.7, 2.9, 2.2, 1.4, 1.8, 3.3] * u"m"
   b = [300, 500, missing, 800, missing, 400] * u"cm"
@@ -160,6 +158,24 @@
   @test tₒ.e == t.e
   @test tₒ.f == t.f
 
+  # table with unitless columns
+  a = [2.0, 2.0, 3.0] * u"m"
+  b = [4.0, 5.0, 6.0]
+  c = [7.0, 8.0, 9.0]
+  d = ["no", "yes", "no"]
+  t = Table(; a, b, c, d)
+  T = Unit("b" => u"kg")
+  n, c = apply(T, t)
+  @test unit(eltype(n.a)) == u"m"
+  @test unit(eltype(n.b)) == u"kg"
+  @test unit(eltype(n.c)) == NoUnits
+  @test eltype(n.d) <: String
+  tₒ = revert(T, n, c)
+  @test unit(eltype(tₒ.a)) == u"m"
+  @test unit(eltype(tₒ.b)) == NoUnits
+  @test unit(eltype(tₒ.c)) == NoUnits
+  @test eltype(tₒ.d) <: String
+
   # error: cannot create Unit transform without arguments
   @test_throws ArgumentError Unit()
 end