Repository: JuliaML/TableTransforms.jl (julia)

Issue:
Sample
Add a Sample transform that samples rows of tables a la `sample` in StatsBase.jl. In particular, we should be able to specify the number of samples and specify if the sampling allows replacement or not.

Pull request:
This PR adds Sample transform.

closes #73

Patch:
diff --git a/Project.toml b/Project.toml
index c976cf5c..e85df9ab 100644
--- a/Project.toml
+++ b/Project.toml
@@ -8,8 +8,10 @@ CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ScientificTypes = "321657f4-b219-11e9-178b-2701a2544e81"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 Transducers = "28d57a85-8fef-5791-bfe6-a80928e7c999"
 
@@ -18,6 +20,7 @@ CategoricalArrays = "0.10"
 Distributions = "0.25"
 PrettyTables = "1.3"
 ScientificTypes = "2.3, 3.0"
+StatsBase = "0.33"
 Tables = "1.6"
 Transducers = "0.4"
 julia = "1.6"
diff --git a/docs/src/transforms/builtin.md b/docs/src/transforms/builtin.md
index 1e1596ea..f7eac48c 100644
--- a/docs/src/transforms/builtin.md
+++ b/docs/src/transforms/builtin.md
@@ -32,6 +32,12 @@ StdNames
 Sort
 ```
 
+## Sample
+
+```@docs
+Sample
+```
+
 ## Filter
 
 ```@docs
diff --git a/src/TableTransforms.jl b/src/TableTransforms.jl
index 46d6a295..6e154728 100644
--- a/src/TableTransforms.jl
+++ b/src/TableTransforms.jl
@@ -8,10 +8,12 @@ using Tables
 using ScientificTypes
 using Distributions: Normal
 using Transducers: tcollect
+using StatsBase: AbstractWeights, sample
 using LinearAlgebra
 using Statistics
 using PrettyTables
 using CategoricalArrays
+using Random: AbstractRNG, GLOBAL_RNG
 
 import Distributions: ContinuousUnivariateDistribution
 import Distributions: quantile, cdf
@@ -32,6 +34,7 @@ export
   Rename,
   StdNames,
   Sort,
+  Sample,
   Filter,
   DropMissing,
   Replace,
diff --git a/src/transforms.jl b/src/transforms.jl
index 2183face..ff294392 100644
--- a/src/transforms.jl
+++ b/src/transforms.jl
@@ -215,6 +215,7 @@ include("transforms/select.jl")
 include("transforms/rename.jl")
 include("transforms/stdnames.jl")
 include("transforms/sort.jl")
+include("transforms/sample.jl")
 include("transforms/filter.jl")
 include("transforms/replace.jl")
 include("transforms/coalesce.jl")
diff --git a/src/transforms/sample.jl b/src/transforms/sample.jl
new file mode 100644
index 00000000..b42bf6d7
--- /dev/null
+++ b/src/transforms/sample.jl
@@ -0,0 +1,73 @@
+# ------------------------------------------------------------------
+# Licensed under the MIT License. See LICENSE in the project root.
+# ------------------------------------------------------------------
+
+"""
+    Sample([rng], [wv::AbstractWeights], n; replace=true, ordered=false)
+
+Sample table rows by forwarding arguments to the `sample` function from StatsBase.jl.
+Optionally, a random number generator `rng` and a weight vector `wv` can be used.
+
+# Examples
+
+```julia
+Sample(1_000)
+Sample(1_000, replace=false)
+Sample(1_000, replace=false, ordered=true)
+
+# with rng
+using Random
+rng = MersenneTwister(2)
+Sample(rng, 1_000)
+
+# with weights
+using StatsBase
+wv = pweights([0.2, 0.1, 0.3])
+Sample(wv, 1_000)
+Sample(rng, wv, 1_000)
+```
+"""
+struct Sample{R<:AbstractRNG,W} <: Stateless
+  rng::R
+  wv::W
+  n::Int
+  replace::Bool
+  ordered::Bool
+end
+
+Sample(rng::AbstractRNG, wv::AbstractWeights, n::Int; replace=true, ordered=false) =
+  Sample(rng, wv, n, replace, ordered)
+
+Sample(rng::AbstractRNG, n::Int; replace=true, ordered=false) =
+  Sample(rng, nothing, n, replace, ordered)
+
+Sample(wv::AbstractWeights, n::Int; replace=true, ordered=false) =
+  Sample(GLOBAL_RNG, wv, n, replace, ordered)
+
+Sample(n::Int; replace=true, ordered=false) =
+  Sample(GLOBAL_RNG, nothing, n, replace, ordered)
+
+isrevertible(::Type{<:Sample}) = false
+
+function apply(transform::Sample, table)
+  rows = Tables.rowtable(table)
+
+  rng     = transform.rng
+  wv      = transform.wv
+  n       = transform.n
+  replace = transform.replace
+  ordered = transform.ordered
+
+  if isnothing(wv)
+    newrows = sample(rng, rows, n; replace, ordered)
+  else
+    newrows = sample(rng, rows, wv, n; replace, ordered)
+  end
+
+  newtable = newrows |> Tables.materializer(table)
+
+  newtable, nothing
+end
+
+revert(::Sample, newtable, cache) = 
+  throw(AssertionError("Transform is not revertible."))


Test patch:
diff --git a/test/Project.toml b/test/Project.toml
index b5359a20..1c3678f0 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -9,10 +9,11 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReferenceTests = "324d217c-45ce-50fc-942e-d289b448e8cf"
 ScientificTypes = "321657f4-b219-11e9-178b-2701a2544e81"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TypedTables = "9d95f2ec-7b3d-5a63-8d20-e2491e220bb9"
 
 [compat]
 GR = "=0.59.0"
-Plots = "=1.22.4"
\ No newline at end of file
+Plots = "=1.22.4"
diff --git a/test/runtests.jl b/test/runtests.jl
index 8473bf91..014a4add 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -8,6 +8,7 @@ using LinearAlgebra
 using Statistics
 using Test, Random, Plots
 using ReferenceTests, ImageIO
+using StatsBase
 
 # set default configurations for plots
 gr(ms=2, mc=:black, aspectratio=:equal,
diff --git a/test/transforms.jl b/test/transforms.jl
index afbe2960..4d607327 100644
--- a/test/transforms.jl
+++ b/test/transforms.jl
@@ -544,6 +544,57 @@
     @test_throws MethodError apply(Sort(:a, :b, test=1), t)
   end
 
+  @testset "Sample" begin
+    a = [3, 6, 2, 7, 8, 3]
+    b = [8, 5, 1, 2, 3, 4]
+    c = [1, 8, 5, 2, 9, 4]
+    t = Table(; a, b, c)
+    trows = Tables.rowtable(t) 
+
+    T = Sample(30)
+    n, c = apply(T, t)
+    @test length(n.a) == 30
+
+    T = Sample(6, replace=false)
+    n, c = apply(T, t)
+    @test n.a ⊆ t.a
+    @test n.b ⊆ t.b
+    @test n.c ⊆ t.c
+
+    T = Sample(30, ordered=true)
+    n, c = apply(T, t)
+    @test unique(Tables.rowtable(n)) == trows
+
+    T = Sample(6, replace=false, ordered=true)
+    n, c = apply(T, t)
+    @test n.a ⊆ t.a
+    @test n.b ⊆ t.b
+    @test n.c ⊆ t.c
+    @test Tables.rowtable(n) == trows
+
+    # with rng
+    T = Sample(MersenneTwister(2), 8)
+    n, c = apply(T, t)
+    @test n.a == [3, 7, 8, 2, 2, 6, 2, 6]
+    @test n.b == [8, 2, 3, 1, 1, 5, 1, 5]
+    @test n.c == [1, 2, 9, 5, 5, 8, 5, 8]
+
+    #with weights
+    wv = pweights([0.1, 0.25, 0.15, 0.25, 0.1, 0.15])
+    T = Sample(MersenneTwister(2), wv, 10_000)
+    n, c = apply(T, t)
+    nrows = Tables.rowtable(n)
+    @test isapprox(count(==(trows[1]), nrows) / 10_000, 0.10, atol=0.01)
+    @test isapprox(count(==(trows[2]), nrows) / 10_000, 0.25, atol=0.01)
+    @test isapprox(count(==(trows[3]), nrows) / 10_000, 0.15, atol=0.01)
+    @test isapprox(count(==(trows[4]), nrows) / 10_000, 0.25, atol=0.01)
+    @test isapprox(count(==(trows[5]), nrows) / 10_000, 0.10, atol=0.01)
+    @test isapprox(count(==(trows[6]), nrows) / 10_000, 0.15, atol=0.01)
+
+    # throws
+    @test_throws AssertionError revert(T, n, c)
+  end
+
   @testset "Filter" begin
     a = [3, 2, 1, 4, 5, 3]
     b = [2, 4, 4, 5, 8, 5]