Implement revert for Sample (#119)

Author: juliohm

src/transforms/sample.jl

@@ -46,7 +46,7 @@ Sample(size::Int, weights::AbstractWeights;
 Sample(size::Int, weights; kwargs...) =
   Sample(size, Weights(collect(weights)); kwargs...)
 
-isrevertible(::Type{<:Sample}) = false
+isrevertible(::Type{<:Sample}) = true
 
 function preprocess(transform::Sample, table)
   # retrieve valid indices
@@ -60,24 +60,47 @@ function preprocess(transform::Sample, table)
   rng     = transform.rng
 
   # sample a subset of indices
-  if isnothing(weights)
+  sinds = if isnothing(weights)
     sample(rng, inds, size; replace, ordered)
   else
     sample(rng, inds, weights, size; replace, ordered)
   end
+  rinds = setdiff(inds, sinds)
+
+  sinds, rinds
 end
 
 function applyfeat(::Sample, table, prep)
   # collect all rows
   rows = Tables.rowtable(table)
 
   # preprocessed indices
-  sinds = prep
+  sinds, rinds = prep
 
   # select rows
   srows = view(rows, sinds)
+  rrows = view(rows, rinds)
+
+  stable = srows |> Tables.materializer(table)
+
+  stable, (sinds, rinds, rrows)
+end
+
+function revertfeat(::Sample, newtable, fcache)
+  # collect all rows
+  rows = Tables.rowtable(newtable)
+
+  sinds, rinds, rrows = fcache
 
-  newtable = srows |> Tables.materializer(table)
+  uinds = sort(unique(sinds))
+  urows = map(uinds) do i
+    j = findfirst(==(i), sinds)
+    rows[j]
+  end
+
+  for (i, row) in zip(rinds, rrows)
+    insert!(urows, i, row)
+  end
 
-  newtable, nothing
+  urows |> Tables.materializer(newtable)
 end

test/transforms.jl

@@ -678,20 +678,27 @@
     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, replace=true)
     n, c = apply(T, t)
     @test length(n.a) == 30
 
+    @test isrevertible(T)
+    r = revert(T, n, c)
+    @test r == t
+
     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
 
+    r = revert(T, n, c)
+    @test r == t
+
     T = Sample(30, replace=true, ordered=true)
     n, c = apply(T, t)
+    trows = Tables.rowtable(t) 
     @test unique(Tables.rowtable(n)) == trows
 
     T = Sample(6, replace=false, ordered=true)