Improve sorted sampling (#103)

Author: Tortar

.github/workflows/ci.yml

@@ -15,7 +15,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.8'
+          - '1.10'
           - '1'
           - 'nightly'
         os:

src/SamplingInterface.jl

@@ -157,8 +157,8 @@ Base.@constprop :aggressive function itsample(rng::AbstractRNG, iter, n::Int, me
         s = ReservoirSample(rng, iter_type, n, method, ImmutSample(), ordered ? Ord() : Unord())
         return update_all!(s, iter, ordered)
     else
-        replace = method isa AlgL || method isa AlgR ? false : true
-        sortedindices_sample(rng, iter, n; iter_type, replace, ordered)
+        replace = method isa AlgL || method isa AlgR ? NoReplace() : Replace()
+        sortedindices_sample(rng, iter, n, replace; iter_type, ordered)
     end
 end
 function itsample(rng::AbstractRNG, iter, wv::Function, method = AlgWRSWRSKIP(); iter_type = infer_eltype(iter))

src/SamplingUtils.jl

@@ -13,17 +13,26 @@ function infer_eltype(itr)
     ifelse(T2 !== Union{} && T2 <: T1, T2, T1)
 end
 
-function sortedrandrange(rng, range, n)
-    s = Vector{Int}(undef, n)
-    curmax = -log(Float64(range.stop))
-    for i in n:-1:1
-        curmax += randexp(rng)/i
-        @inbounds s[i] = ceil(Int, exp(-curmax))
-    end
-    return s
+get_sorted_indices(rng, n, N, ::Replace) = SortedRandRangeIter(rng, 1:N, n)
+get_sorted_indices(rng, n, N, ::NoReplace) = sort!(sample(rng, 1:N, n; replace=false))
+
+struct SortedRandRangeIter{R}
+    rng::R
+    range::UnitRange{Int}
+    n::Int
 end
 
-function get_sorted_indices(rng, n, N, replace)
-    replace == true && return sortedrandrange(rng, 1:N, n)
-    return sort!(sample(rng, 1:N, n; replace=replace))
+@inline function Base.iterate(s::SortedRandRangeIter)
+    curmax = -log(Float64(s.range.stop)) + randexp(s.rng)/s.n
+    return (s.range.stop - ceil(Int, exp(-curmax)) + 1, (s.n-1, curmax))
+end
+@inline function Base.iterate(s::SortedRandRangeIter, state)
+    state[1] == 0 && return nothing
+    curmax = state[2] + randexp(s.rng)/state[1]
+    return (s.range.stop - ceil(Int, exp(-curmax)) + 1, (state[1]-1, curmax))
 end
+
+Base.IteratorEltype(::SortedRandRangeIter) = Base.HasEltype()
+Base.eltype(::SortedRandRangeIter) = Int
+Base.IteratorSize(::SortedRandRangeIter) = Base.HasLength()
+Base.length(s::SortedRandRangeIter) = s.n

src/SortedSamplingMulti.jl

@@ -7,39 +7,29 @@ Algorithm which generates sorted random indices used to retrieve the sample
 from the iterable. The number of elements in the iterable needs to be known 
 before starting the sampling.
 """
-function sortedindices_sample(rng, iter, n::Int; 
-        iter_type = infer_eltype(iter), replace = false, ordered = false)
+function sortedindices_sample(rng, iter, n::Int, replace; 
+        iter_type = infer_eltype(iter), ordered = false)
     N = length(iter)
     if N <= n
         reservoir = collect(iter)
-        replace && return sample(rng, reservoir, n, ordered=ordered)
+        replace isa Replace && return sample(rng, reservoir, n, ordered=ordered)
         return ordered ? reservoir : shuffle!(rng, reservoir)
     end
     reservoir = Vector{iter_type}(undef, n)
     indices = get_sorted_indices(rng, n, N, replace)
-    first_idx = indices[1]
-    el, state = iterate(iter)::Tuple
-    if first_idx != 1
-        el, state = skip_ahead_no_end(iter, state, first_idx - 2)
+    curr_idx, state_idx = iterate(indices)
+    el, state_el = iterate(iter)::Tuple
+    for _ in 1:curr_idx-1
+        el, state_el = iterate(iter, state_el)::Tuple
     end
     reservoir[1] = el
-    i = 2
-    @inbounds while i <= n
-        skip_k = indices[i] - indices[i-1] - 1
-        if skip_k >= 0
-            el, state = skip_ahead_no_end(iter, state, skip_k)
+    @inbounds for i in 2:n
+        next_idx, state_idx = iterate(indices, state_idx)
+        for _ in 1:next_idx-curr_idx
+            el, state_el = iterate(iter, state_el)::Tuple
         end
         reservoir[i] = el
-        i += 1
+        curr_idx = next_idx
     end
     return ordered ? reservoir : shuffle!(rng, reservoir)
 end
-
-function skip_ahead_no_end(iter, state, n)
-    for _ in 1:n
-        it = iterate(iter, state)::Tuple
-        state = it[2]
-    end
-    it = iterate(iter, state)::Tuple
-    return it
-end

src/StreamSampling.jl

@@ -25,6 +25,9 @@ struct MutSample end
 struct Ord end
 struct Unord end
 
+struct Replace end
+struct NoReplace end
+
 abstract type AbstractReservoirSample <: OnlineStat{Any} end
 
 abstract type ReservoirAlgorithm end