Refactor Replace transform (#198)

* [WIP] Refactor 'Replace' transform

* Add support for ColSpec in Replace & Update Map docstring

* Update variable name

* Add comments

* Fix type promotion

* Add tests

* Fix typo

* Update docstring

* Fix typo

* Update docstring

* Update docstring

Author: eliascarv

src/transforms/map.jl

@@ -3,16 +3,18 @@
 # ------------------------------------------------------------------
 
 """
-    Map(cols₁ => fun₁, cols₂ => fun₂ => target₂, ..., colsₙ => funₙ => targetₙ)
+    Map(cols₁ => fun₁ => target₁, cols₂ => fun₂, ..., colsₙ => funₙ => targetₙ)
 
 Applies the `funᵢ` function to the columns selected by `colsᵢ` using 
 the `map` function and saves the result in a new column named `targetᵢ`.
-The target column name is optional and when omitted a new name is generated
-by joining the selected column names with the function name.
+
+The column selection can be a single column identifier (index or name),
+a collection of identifiers or a regular expression (regex).
+
+Passing a target column name is optional and when omitted a new name
+is generated by joining the selected column names with the function name.
 If the target column already exists in the table, the original
-column will be replaced. The column selection can be a single
-column identifier (index or name), a collection of identifiers
-or a regular expression (regex).
+column will be replaced.
 
 # Examples
 
@@ -24,14 +26,18 @@ Map([:a, :c] => ((a, c) -> 2a * 3c) => :col1)
 Map(["c", "a"] => ((c, a) -> 3c / a) => :col1, "c" => tan)
 Map(r"[abc]" => ((a, b, c) -> a^2 - 2b + c) => "col1")
 ```
+
+## Notes
+
+* Anonymous functions must be passed with parentheses as in the examples above.
 """
 struct Map <: StatelessFeatureTransform
   colspecs::Vector{ColSpec}
   funs::Vector{Function}
   targets::Vector{Union{Nothing,Symbol}}
 end
 
-isrevertible(::Type{Map}) = true
+Map() = throw(ArgumentError("cannot create a Map transform without arguments"))
 
 # utility types
 const TargetName = Union{Symbol,AbstractString}
@@ -40,23 +46,19 @@ const PairWithoutTarget = Pair{<:Any,<:Function}
 const MapPair = Union{PairWithTarget,PairWithoutTarget}
 
 # utility functions
-_extract(p::PairWithTarget) = first(p), first(last(p)), last(last(p))
-_extract(p::PairWithoutTarget) = first(p), last(p), nothing
-
-_target(name) = name
-_target(name::AbstractString) = Symbol(name)
+_extract(p::PairWithTarget) = colspec(first(p)), first(last(p)), Symbol(last(last(p)))
+_extract(p::PairWithoutTarget) = colspec(first(p)), last(p), nothing
 
 function Map(pairs::MapPair...)
-  tuples = map(pairs) do p
-    spec, fun, name = _extract(p)
-    (colspec(spec), fun, _target(name))
-  end
-  colspecs = map(t -> t[1], tuples) |> collect
-  funs = map(t -> t[2], tuples) |> collect
-  targets = map(t -> t[3], tuples) |> collect
+  tuples = map(_extract, pairs)
+  colspecs = [t[1] for t in tuples]
+  funs = [t[2] for t in tuples]
+  targets = [t[3] for t in tuples]
   Map(colspecs, funs, targets)
 end
 
+isrevertible(::Type{Map}) = true
+
 _makename(snames, fun) = Symbol(join([snames; nameof(fun)], "_"))
 
 function preprocess(transform::Map, table)

src/transforms/replace.jl

@@ -3,42 +3,111 @@
 # ------------------------------------------------------------------
 
 """
-    Replace(old₁ => new₁, old₂ => new₂, ..., oldₙ => newₙ)
+    Replace(cols₁ => pred₁ => new₁, pred₂ => new₂, ..., colsₙ => predₙ => newₙ)
 
-Replaces `oldᵢ` value with `newᵢ` value in the table.
+Replaces all values where `predᵢ` predicate returns `true` with `newᵢ` value 
+in the the columns selected by `colsᵢ`.
+
+Passing a column selection is optional and when omitted all columns in the table 
+will be selected. The column selection can be a single column identifier (index or name), 
+a collection of identifiers, or a regular expression (regex).
+
+The predicate can be a function that accepts a single argument
+and returns a boolean, or a value. If the predicate is a value,
+it will be transformed into the following function: `x -> x === value`.
 
 # Examples
 
 ```julia
 Replace(1 => -1, 5 => -5)
-Replace(1 => 1.5, 5 => 5.5, 4 => true)
+Replace(2 => 0.0 => 1.5, 5.0 => 5.5)
+Replace(:b => 0.0 => 1.5, 5.0 => 5.5)
+Replace("b" => 0.0 => 1.5, 5.0 => 5.5)
+Replace([1, 3] => >(5) => 5)
+Replace([:a, :c] => isequal(2) => -2)
+Replace(["a", "c"] => (x -> 4 < x < 6) => 0)
+Replace(r"[abc]" => (x -> isodd(x) && x > 10) => 2)
 ```
+
+## Notes
+
+* Anonymous functions must be passed with parentheses as in the examples above.
+* Replacements are applied in the sequence in which they are defined, therefore,
+  if there is more than one replacement for the same column, the first valid one will be applied.
 """
-struct Replace{K,V} <: StatelessFeatureTransform
-  pairs::IdDict{K,V}
+struct Replace <: StatelessFeatureTransform
+  colspecs::Vector{ColSpec}
+  preds::Vector{Function}
+  news::Vector{Any}
 end
 
-Replace() = throw(ArgumentError("Cannot create a Replace object without arguments."))
+Replace() = throw(ArgumentError("cannot create a Replace transform without arguments"))
 
-Replace(pairs::Pair...) = Replace(IdDict(values(pairs)))
+# utility functions
+_extract(p::Pair) = AllSpec(), _pred(first(p)), last(p)
+_extract(p::Pair{<:Any,<:Pair}) = colspec(first(p)), _pred(first(last(p))), last(last(p))
+
+_pred(f::Function) = f
+_pred(v) = Base.Fix2(===, v)
+
+function Replace(pairs::Pair...)
+  tuples = map(_extract, pairs)
+  colspecs = [t[1] for t in tuples]
+  preds = [t[2] for t in tuples]
+  news = Any[t[3] for t in tuples]
+  Replace(colspecs, preds, news)
+end
 
 isrevertible(::Type{<:Replace}) = true
 
-function applyfeat(transform::Replace, feat, prep)
+function preprocess(transform::Replace, table)
+  cols = Tables.columns(table)
+  names = Tables.columnnames(cols)
+
+  colspecs = transform.colspecs
+  preds = transform.preds
+  news = transform.news
+
+  # column replacements
+  colreps = map(colspecs, preds, news) do colspec, pred, new
+    snames = choose(colspec, names)
+    snames => pred => new
+  end
+
+  # join replacements of each column
+  map(names) do name
+    pairs = filter(p -> name ∈ first(p), colreps)
+    reps = isempty(pairs) ? nothing : map(last, pairs)
+    name => reps
+  end
+end
+
+function applyfeat(::Replace, feat, prep)
   cols = Tables.columns(feat)
   names = Tables.columnnames(cols)
 
-  olds = keys(transform.pairs)
-  values = map(names) do nm
-    x = Tables.getcolumn(cols, nm)
-    y = [get(transform.pairs, xᵢ, xᵢ) for xᵢ in x]
-    inds = [findall(xᵢ -> xᵢ === old, x) .=> old for old in olds]
-    rev = Dict(reduce(vcat, inds))
-    y, rev
+  tuples = map(prep) do (name, reps)
+    x = Tables.getcolumn(cols, name)
+    if isnothing(reps)
+      x, nothing
+    else
+      # reversal dict
+      rev = Dict{Int,eltype(x)}()
+      y = map(enumerate(x)) do (i, v)
+        for (pred, new) in reps
+          if pred(v)
+            rev[i] = v
+            return new
+          end
+        end
+        v
+      end
+      y, rev
+    end
   end
 
-  columns = first.(values)
-  fcache = last.(values)
+  columns = first.(tuples)
+  fcache = last.(tuples)
 
   𝒯 = (; zip(names, columns)...)
   newfeat = 𝒯 |> Tables.materializer(feat)
@@ -49,9 +118,9 @@ function revertfeat(::Replace, newfeat, fcache)
   cols = Tables.columns(newfeat)
   names = Tables.columnnames(cols)
 
-  columns = map(names, fcache) do nm, rev
-    y = Tables.getcolumn(cols, nm)
-    [get(rev, i, y[i]) for i in eachindex(y)]
+  columns = map(names, fcache) do name, rev
+    y = Tables.getcolumn(cols, name)
+    isnothing(rev) ? y : [get(rev, i, y[i]) for i in 1:length(y)]
   end
 
   𝒯 = (; zip(names, columns)...)

test/shows.jl

@@ -152,18 +152,19 @@
   end
 
   @testset "Replace" begin
-    pairs = IdDict(1 => -1, 5 => -5)
-    T = Replace(pairs)
+    T = Replace(1 => -1, 5 => -5)
 
     # compact mode
     iostr = sprint(show, T)
-    @test iostr == "Replace($pairs)"
+    @test iostr == "Replace(TableTransforms.ColSpec[all, all], Function[Base.Fix2{typeof(===), Int64}(===, 1), Base.Fix2{typeof(===), Int64}(===, 5)], Any[-1, -5])"
 
     # full mode
     iostr = sprint(show, MIME("text/plain"), T)
     @test iostr == """
     Replace transform
-    └─ pairs = $pairs"""
+    ├─ colspecs = TableTransforms.ColSpec[all, all]
+    ├─ preds = Function[Base.Fix2{typeof(===), Int64}(===, 1), Base.Fix2{typeof(===), Int64}(===, 5)]
+    └─ news = Any[-1, -5]"""
   end
 
   @testset "Coalesce" begin

test/transforms/map.jl

@@ -68,4 +68,7 @@
   @test n.op1 == @. t.a^2 - 2 * t.b + t.c
   tₒ = revert(T, n, c)
   @test t == tₒ
+
+  # throws
+  @test_throws ArgumentError Map()
 end

test/transforms/replace.jl

@@ -20,6 +20,92 @@
   tₒ = revert(T, n, c)
   @test t == tₒ
 
+  # with colspec
+  T = Replace(2 => 4 => -4, :c => 1 => -1, "e" => 5 => -5)
+  n, c = apply(T, t)
+  @test n.a == [3, 2, 1, 4, 5, 3]
+  @test n.b == [2, -4, -4, 5, 8, 5]
+  @test n.c == [-1, -1, 6, 2, 4, -1]
+  @test n.d == [4, 3, 7, 5, 4, 1]
+  @test n.e == [-5, -5, 2, 6, -5, 2]
+  @test n.f == [4, 4, 3, 4, 5, 2]
+  @test isrevertible(T)
+  tₒ = revert(T, n, c)
+  @test t == tₒ
+
+  T = Replace([2, 5] => 5 => -5)
+  n, c = apply(T, t)
+  @test n.a == [3, 2, 1, 4, 5, 3]
+  @test n.b == [2, 4, 4, -5, 8, -5]
+  @test n.c == [1, 1, 6, 2, 4, 1]
+  @test n.d == [4, 3, 7, 5, 4, 1]
+  @test n.e == [-5, -5, 2, 6, -5, 2]
+  @test n.f == [4, 4, 3, 4, 5, 2]
+  @test isrevertible(T)
+  tₒ = revert(T, n, c)
+  @test t == tₒ
+
+  T = Replace([:b, :e] => 5 => -5)
+  n, c = apply(T, t)
+  @test n.a == [3, 2, 1, 4, 5, 3]
+  @test n.b == [2, 4, 4, -5, 8, -5]
+  @test n.c == [1, 1, 6, 2, 4, 1]
+  @test n.d == [4, 3, 7, 5, 4, 1]
+  @test n.e == [-5, -5, 2, 6, -5, 2]
+  @test n.f == [4, 4, 3, 4, 5, 2]
+  @test isrevertible(T)
+  tₒ = revert(T, n, c)
+  @test t == tₒ
+
+  T = Replace(["b", "e"] => 5 => -5)
+  n, c = apply(T, t)
+  @test n.a == [3, 2, 1, 4, 5, 3]
+  @test n.b == [2, 4, 4, -5, 8, -5]
+  @test n.c == [1, 1, 6, 2, 4, 1]
+  @test n.d == [4, 3, 7, 5, 4, 1]
+  @test n.e == [-5, -5, 2, 6, -5, 2]
+  @test n.f == [4, 4, 3, 4, 5, 2]
+  @test isrevertible(T)
+  tₒ = revert(T, n, c)
+  @test t == tₒ
+
+  T = Replace(r"[be]" => 5 => -5)
+  n, c = apply(T, t)
+  @test n.a == [3, 2, 1, 4, 5, 3]
+  @test n.b == [2, 4, 4, -5, 8, -5]
+  @test n.c == [1, 1, 6, 2, 4, 1]
+  @test n.d == [4, 3, 7, 5, 4, 1]
+  @test n.e == [-5, -5, 2, 6, -5, 2]
+  @test n.f == [4, 4, 3, 4, 5, 2]
+  @test isrevertible(T)
+  tₒ = revert(T, n, c)
+  @test t == tₒ
+
+  # with predicates
+  T = Replace([:b, :d] => >(4) => 0)
+  n, c = apply(T, t)
+  @test n.a == [3, 2, 1, 4, 5, 3]
+  @test n.b == [2, 4, 4, 0, 0, 0]
+  @test n.c == [1, 1, 6, 2, 4, 1]
+  @test n.d == [4, 3, 0, 0, 4, 1]
+  @test n.e == [5, 5, 2, 6, 5, 2]
+  @test n.f == [4, 4, 3, 4, 5, 2]
+  @test isrevertible(T)
+  tₒ = revert(T, n, c)
+  @test t == tₒ
+
+  T = Replace([:a, :f] => (x -> 1 < x < 5) => 0)
+  n, c = apply(T, t)
+  @test n.a == [0, 0, 1, 0, 5, 0]
+  @test n.b == [2, 4, 4, 5, 8, 5]
+  @test n.c == [1, 1, 6, 2, 4, 1]
+  @test n.d == [4, 3, 7, 5, 4, 1]
+  @test n.e == [5, 5, 2, 6, 5, 2]
+  @test n.f == [0, 0, 0, 0, 5, 0]
+  @test isrevertible(T)
+  tₒ = revert(T, n, c)
+  @test t == tₒ
+
   # table schema after apply and revert
   T = Replace(1 => -1, 5 => -5)
   n, c = apply(T, t)