Repository: JuliaML/TableTransforms.jl (julia)

Issue:
Add `Unit` transform to convert units
The transform should behave like `uconvert` with column selections.

Pull request:
closes #282

Patch:
diff --git a/docs/src/transforms.md b/docs/src/transforms.md
index ba467a73..f1d6c19b 100644
--- a/docs/src/transforms.md
+++ b/docs/src/transforms.md
@@ -116,6 +116,12 @@ AbsoluteUnits
 Unitify
 ```
 
+## Unit
+
+```@docs
+Unit
+```
+
 ## Map
 
 ```@docs
diff --git a/src/TableTransforms.jl b/src/TableTransforms.jl
index f4d9142a..bc7cb165 100644
--- a/src/TableTransforms.jl
+++ b/src/TableTransforms.jl
@@ -67,6 +67,7 @@ export
   DropUnits,
   AbsoluteUnits,
   Unitify,
+  Unit,
   Map,
   Replace,
   Coalesce,
diff --git a/src/transforms.jl b/src/transforms.jl
index 7abc755b..48594542 100644
--- a/src/transforms.jl
+++ b/src/transforms.jl
@@ -281,6 +281,7 @@ include("transforms/dropunits.jl")
 include("transforms/dropconstant.jl")
 include("transforms/absoluteunits.jl")
 include("transforms/unitify.jl")
+include("transforms/unit.jl")
 include("transforms/map.jl")
 include("transforms/replace.jl")
 include("transforms/coalesce.jl")
diff --git a/src/transforms/unit.jl b/src/transforms/unit.jl
new file mode 100644
index 00000000..cb98f203
--- /dev/null
+++ b/src/transforms/unit.jl
@@ -0,0 +1,88 @@
+# ------------------------------------------------------------------
+# Licensed under the MIT License. See LICENSE in the project root.
+# ------------------------------------------------------------------
+
+"""
+    Unit(unit)
+
+Converts the units of all columns in the table to `unit`.
+
+    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).
+
+# Examples
+
+```julia
+Unit(u"m")
+Unit(1 => u"km", :b => u"K", "c" => u"s")
+Unit([2, 3] => u"cm")
+Unit([:a, :c] => u"cm")
+Unit(["a", "c"] => u"cm")
+Unit(r"[abc]" => u"km")
+```
+"""
+struct Unit <: StatelessFeatureTransform
+  selectors::Vector{ColumnSelector}
+  units::Vector{Units}
+end
+
+Unit() = throw(ArgumentError("cannot create Unit transform without arguments"))
+
+Unit(unit::Units) = Unit([AllSelector()], [unit])
+
+Unit(pairs::Pair...) = Unit(collect(selector.(first.(pairs))), collect(last.(pairs)))
+
+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)
+
+  selectors = transform.selectors
+  units = transform.units
+  pairs = mapreduce(vcat, selectors, units) do selector, u
+    snames = selector(names)
+    snames .=> u
+  end
+  unitdict = Dict(pairs) 
+
+  tuples = map(names) do name
+    x = Tables.getcolumn(cols, name)
+    if haskey(unitdict, name)
+      u = unitdict[name]
+      _uconvert(u, x)
+    else
+      (x, nothing)
+    end
+  end
+
+  columns = first.(tuples)
+  ounits = last.(tuples)
+
+  𝒯 = (; zip(names, columns)...)
+  newfeat = 𝒯 |> Tables.materializer(feat)
+  newfeat, ounits
+end
+
+function revertfeat(::Unit, newfeat, fcache)
+  cols = Tables.columns(newfeat)
+  names = Tables.columnnames(cols)
+
+  ounits = fcache
+  columns = map(names, ounits) do name, u
+    x = Tables.getcolumn(cols, name)
+    isnothing(u) ? x : map(v -> uconvert(u, v), x)
+  end
+
+  𝒯 = (; zip(names, columns)...)
+  𝒯 |> Tables.materializer(newfeat)
+end


Test patch:
diff --git a/test/shows.jl b/test/shows.jl
index 987150fe..cdd810e1 100644
--- a/test/shows.jl
+++ b/test/shows.jl
@@ -257,6 +257,21 @@
     @test iostr == "Unitify transform"
   end
 
+  @testset "Unit" begin
+    T = Unit(:a => u"m", [:b, :c] => u"s")
+
+    # compact mode
+    iostr = sprint(show, T)
+    @test iostr == "Unit(selectors: ColumnSelector[:a, [:b, :c]], units: Units[m, s])"
+
+    # full mode
+    iostr = sprint(show, MIME("text/plain"), T)
+    @test iostr == """
+    Unit transform
+    ├─ selectors: ColumnSelectors.ColumnSelector[:a, [:b, :c]]
+    └─ units: Unitful.Units[m, s]"""
+  end
+
   @testset "Map" begin
     fun = (a, b) -> 2a + b
     T = Map(:a => sin, [:a, :b] => fun => :c)
diff --git a/test/transforms.jl b/test/transforms.jl
index 46a2c4ab..6fd34f26 100644
--- a/test/transforms.jl
+++ b/test/transforms.jl
@@ -15,6 +15,7 @@ transformfiles = [
   "dropconstant.jl",
   "absoluteunits.jl",
   "unitify.jl",
+  "unit.jl",
   "map.jl",
   "replace.jl",
   "coalesce.jl",
diff --git a/test/transforms/unit.jl b/test/transforms/unit.jl
new file mode 100644
index 00000000..06ed9efc
--- /dev/null
+++ b/test/transforms/unit.jl
@@ -0,0 +1,165 @@
+@testset "Unit" begin
+  @test isrevertible(Unit(u"m"))
+
+  a = [2.7, 2.9, 2.2, 1.4, 1.8, 3.3] * u"m"
+  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 = 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
+  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 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"
+  c = [8, 2, 5, 7, 9, 4] * u"km"
+  d = [29.1, missing, 29.2, missing, 28.4, 26.4] * u"°C"
+  e = [0.9, 0.4, 0.5, 0.1, 0.3, 0.6] * u"kg"
+  f = 0.5u"ppm" * e
+  t = Table(; a, b, c, d, e, f)
+
+  T = Unit(4 => u"K")
+  n, c = apply(T, t)
+  @test unit(eltype(n.a)) == u"m"
+  @test unit(eltype(n.b)) == u"cm"
+  @test unit(eltype(n.c)) == u"km"
+  @test unit(eltype(n.d)) == u"K"
+  @test unit(eltype(n.e)) == u"kg"
+  @test unit(eltype(n.f)) == u"kg * ppm"
+  tₒ = revert(T, n, c)
+  @test unit(eltype(tₒ.d)) == u"°C"
+  @test tₒ.a == t.a
+  @test isequal(tₒ.b, t.b)
+  @test tₒ.c == t.c
+  @test all(isapprox.(skipmissing(tₒ.d), skipmissing(t.d)))
+  @test tₒ.e == t.e
+  @test tₒ.f == t.f
+
+  T = Unit(:e => u"g")
+  n, c = apply(T, t)
+  @test unit(eltype(n.a)) == u"m"
+  @test unit(eltype(n.b)) == u"cm"
+  @test unit(eltype(n.c)) == u"km"
+  @test unit(eltype(n.d)) == u"°C"
+  @test unit(eltype(n.e)) == u"g"
+  @test unit(eltype(n.f)) == u"kg * ppm"
+  tₒ = revert(T, n, c)
+  @test unit(eltype(tₒ.e)) == u"kg"
+  @test tₒ.a == t.a
+  @test isequal(tₒ.b, t.b)
+  @test tₒ.c == t.c
+  @test isequal(tₒ.d, t.d)
+  @test all(isapprox.(tₒ.e, t.e))
+  @test tₒ.f == t.f
+
+  T = Unit("f" => u"kg")
+  n, c = apply(T, t)
+  @test unit(eltype(n.a)) == u"m"
+  @test unit(eltype(n.b)) == u"cm"
+  @test unit(eltype(n.c)) == u"km"
+  @test unit(eltype(n.d)) == u"°C"
+  @test unit(eltype(n.e)) == u"kg"
+  @test unit(eltype(n.f)) == u"kg"
+  tₒ = revert(T, n, c)
+  @test unit(eltype(tₒ.f)) == u"kg * ppm"
+  @test tₒ.a == t.a
+  @test isequal(tₒ.b, t.b)
+  @test tₒ.c == t.c
+  @test isequal(tₒ.d, t.d)
+  @test tₒ.e == t.e
+  @test all(isapprox.(tₒ.f, t.f))
+
+  T = Unit([1, 2, 3] => 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 unit(eltype(n.d)) == u"°C"
+  @test unit(eltype(n.e)) == u"kg"
+  @test unit(eltype(n.f)) == u"kg * ppm"
+  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 all(isapprox.(tₒ.a, t.a))
+  @test all(isapprox.(skipmissing(tₒ.b), skipmissing(t.b)))
+  @test all(isapprox.(tₒ.c, t.c))
+  @test isequal(tₒ.d, t.d)
+  @test tₒ.e == t.e
+  @test tₒ.f == t.f
+
+  T = Unit([:a, :b, :c] => u"cm")
+  n, c = apply(T, t)
+  @test unit(eltype(n.a)) == u"cm"
+  @test unit(eltype(n.b)) == u"cm"
+  @test unit(eltype(n.c)) == u"cm"
+  @test unit(eltype(n.d)) == u"°C"
+  @test unit(eltype(n.e)) == u"kg"
+  @test unit(eltype(n.f)) == u"kg * ppm"
+  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 all(isapprox.(tₒ.a, t.a))
+  @test all(isapprox.(skipmissing(tₒ.b), skipmissing(t.b)))
+  @test all(isapprox.(tₒ.c, t.c))
+  @test isequal(tₒ.d, t.d)
+  @test tₒ.e == t.e
+  @test tₒ.f == t.f
+
+  T = Unit(["a", "b", "c"] => u"km")
+  n, c = apply(T, t)
+  @test unit(eltype(n.a)) == u"km"
+  @test unit(eltype(n.b)) == u"km"
+  @test unit(eltype(n.c)) == u"km"
+  @test unit(eltype(n.d)) == u"°C"
+  @test unit(eltype(n.e)) == u"kg"
+  @test unit(eltype(n.f)) == u"kg * ppm"
+  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 all(isapprox.(tₒ.a, t.a))
+  @test all(isapprox.(skipmissing(tₒ.b), skipmissing(t.b)))
+  @test all(isapprox.(tₒ.c, t.c))
+  @test isequal(tₒ.d, t.d)
+  @test tₒ.e == t.e
+  @test tₒ.f == t.f
+
+  T = Unit(r"[abc]" => 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 unit(eltype(n.d)) == u"°C"
+  @test unit(eltype(n.e)) == u"kg"
+  @test unit(eltype(n.f)) == u"kg * ppm"
+  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 all(isapprox.(tₒ.a, t.a))
+  @test all(isapprox.(skipmissing(tₒ.b), skipmissing(t.b)))
+  @test all(isapprox.(tₒ.c, t.c))
+  @test isequal(tₒ.d, t.d)
+  @test tₒ.e == t.e
+  @test tₒ.f == t.f
+
+  # error: cannot create Unit transform without arguments
+  @test_throws ArgumentError Unit()
+end