Add support for multiple column selection in DropExtrema (#255)

* Add support for multiple column selection in 'DropExtrema'

* Update tests

Author: eliascarv

src/transforms/dropextrema.jl

@@ -3,46 +3,68 @@
 # ------------------------------------------------------------------
 
 """
-    DropExtrema(col; low=0.25, high=0.75)
+    DropExtrema(; low=0.25, high=0.75)
 
-Drops the rows where the values in the column `col` are outside the interval
-`[quantile(col, low), quantile(col, high)]`.
+Drops rows where any of the values in all columns
+are outside the interval (`[quantile(col, low), quantile(col, high)]`).
+
+    DropExtrema(col₁, col₂, ..., colₙ; low=0.25, high=0.75)
+    DropExtrema([col₁, col₂, ..., colₙ]; low=0.25, high=0.75)
+    DropExtrema((col₁, col₂, ..., colₙ); low=0.25, high=0.75)
+
+Drops rows where any of the values in columns `col₁`, `col₂`, ..., `colₙ`
+are outside the interval.
+
+    DropExtrema(regex; low=0.25, high=0.75)
+
+Drops rows where any of the values in columns that match with `regex`
+are outside the interval.
 
 # Examples
 
 ```julia
-DropExtrema(1)
+DropExtrema(low=0.3, high=0.7)
+DropExtrema(1, low=0.3, high=0.7)
 DropExtrema(:a, low=0.2, high=0.8)
 DropExtrema("a", low=0.3, high=0.7)
+DropExtrema(1, 3, 5, low=0, high=1)
+DropExtrema([:a, :c, :e], low=0.3, high=0.7)
+DropExtrema(("a", "c", "e"), low=0.25, high=0.75)
+DropExtrema(r"[ace]", low=0.3, high=0.7)
 ```
 """
-struct DropExtrema{S<:SingleColumnSelector,T} <: StatelessFeatureTransform
+struct DropExtrema{S<:ColumnSelector,T} <: StatelessFeatureTransform
   selector::S
   low::T
   high::T
 
-  function DropExtrema(selector::S, low::T, high::T) where {S<:SingleColumnSelector,T}
+  function DropExtrema(selector::S, low::T, high::T) where {S<:ColumnSelector,T}
     _assert(0 ≤ low ≤ high ≤ 1, "invalid quantiles")
     new{S,T}(selector, low, high)
   end
 end
 
-DropExtrema(selector::SingleColumnSelector, low, high) = DropExtrema(selector, promote(low, high)...)
-DropExtrema(col::Column; low=0.25, high=0.75) = DropExtrema(selector(col), low, high)
+DropExtrema(selector::ColumnSelector, low, high) = DropExtrema(selector, promote(low, high)...)
+
+DropExtrema(; low=0.25, high=0.75) = DropExtrema(AllSelector(), low, high)
+DropExtrema(cols; low=0.25, high=0.75) = DropExtrema(selector(cols), low, high)
+DropExtrema(cols::C...; low=0.25, high=0.75) where {C<:Column} = DropExtrema(selector(cols), low, high)
 
 isrevertible(::Type{<:DropExtrema}) = true
 
 function preprocess(transform::DropExtrema, feat)
   cols = Tables.columns(feat)
   names = Tables.columnnames(cols)
-  sname = selectsingle(transform.selector, names)
+  snames = transform.selector(names)
 
-  x = Tables.getcolumn(cols, sname)
-  low = convert(eltype(x), transform.low)
-  high = convert(eltype(x), transform.high)
-  xl, xh = quantile(x, (low, high))
+  limits = map(snames) do name
+    x = Tables.getcolumn(cols, name)
+    low = convert(eltype(x), transform.low)
+    high = convert(eltype(x), transform.high)
+    name => quantile(x, (low, high))
+  end
 
-  ftrans = Filter(row -> xl ≤ row[sname] ≤ xh)
+  ftrans = Filter(row -> all(xl ≤ row[nm] ≤ xh for (nm, (xl, xh)) in limits))
   fprep = preprocess(ftrans, feat)
   ftrans, fprep
 end

test/shows.jl

@@ -180,13 +180,13 @@
 
     # compact mode
     iostr = sprint(show, T)
-    @test iostr == "DropExtrema(:a, 0.25, 0.75)"
+    @test iostr == "DropExtrema([:a], 0.25, 0.75)"
 
     # full mode
     iostr = sprint(show, MIME("text/plain"), T)
     @test iostr == """
     DropExtrema transform
-    ├─ selector = :a
+    ├─ selector = [:a]
     ├─ low = 0.25
     └─ high = 0.75"""
   end

test/transforms/dropextrema.jl

@@ -20,6 +20,17 @@
   tₒ = revert(T, n, c)
   @test t == tₒ
 
+  T = DropExtrema(1, 2)
+  n, c = apply(T, t)
+  @test n.a == [5.1, 5.4]
+  @test n.b == [4.4, 1.9]
+  @test n.c == [2.8, 8.9]
+  @test n.d == [8.1, 4.1]
+  @test n.e == [5.9, 2.4]
+  @test n.f == [1.3, 5.1]
+  tₒ = revert(T, n, c)
+  @test t == tₒ
+
   T = DropExtrema(:c, low=0.3, high=0.7)
   n, c = apply(T, t)
   @test n.a == [6.9, 7.8, 4.8, 8.0]
@@ -31,6 +42,17 @@
   tₒ = revert(T, n, c)
   @test t == tₒ
 
+  T = DropExtrema([:c, :d], low=0.3, high=0.7)
+  n, c = apply(T, t)
+  @test n.a == [7.8, 4.8]
+  @test n.b == [6.3, 0.5]
+  @test n.c == [5.7, 6.7]
+  @test n.d == [6.2, 6.2]
+  @test n.e == [4.1, 1.3]
+  @test n.f == [9.0, 8.9]
+  tₒ = revert(T, n, c)
+  @test t == tₒ
+
   T = DropExtrema("e", low=0.2, high=0.8)
   n, c = apply(T, t)
   @test n.a == [7.8, 0.0, 5.1, 1.1, 8.0, 5.4]
@@ -42,6 +64,17 @@
   tₒ = revert(T, n, c)
   @test t == tₒ
 
-  # throws
+  T = DropExtrema(("e", "f"), low=0.2, high=0.8)
+  n, c = apply(T, t)
+  @test n.a == [0.0, 1.1, 8.0, 5.4]
+  @test n.b == [1.4, 3.0, 6.1, 1.9]
+  @test n.c == [2.8, 8.4, 5.0, 8.9]
+  @test n.d == [1.9, 4.0, 6.9, 4.1]
+  @test n.e == [1.6, 4.9, 3.5, 2.4]
+  @test n.f == [6.2, 6.2, 3.8, 5.1]
+  tₒ = revert(T, n, c)
+  @test t == tₒ
+
+  # error: invalid quantiles
   @test_throws AssertionError DropExtrema(:a, low=0, high=1.4)
 end