Fix some issues with samplers

.gitignore

@@ -25,3 +25,7 @@ Manifest.toml
 
 benchmark/Manifest.toml
 benchmark/random_data.arrow
+
+# Clickstream dataset used by comparison/benchmark_clickstream.jl (download from https://dumps.wikimedia.org/other/clickstream/2024-12/)
+comparison/clickstream-*.tsv
+comparison/clickstream-*.tsv.gz

comparison/benchmark_data_results.csv

@@ -0,0 +1,13 @@
+alg,N,m,repetitions,t_add,t_get
+WRSWR-BIN,34240686,3424069,100,28.234017344161856,3.63
+WRSWR-BIN,34240686,342407,100,17.949567220703457,58.800000000000004
+WRSWR-BIN,34240686,34241,100,17.01158814195487,2.56
+WRSWR-BIN,34240686,3425,100,16.926156700248352,4.390000000000001
+WRAExp-J,34240686,3424069,100,63.313972090979725,6.268304051400001e8
+WRAExp-J,34240686,342407,100,7.436388725973536,4.524304526e7
+WRAExp-J,34240686,34241,100,3.5189679575929054,3.76732991e6
+WRAExp-J,34240686,3425,100,3.1451885654393714,349277.97000000003
+WRSWR-SKIP,34240686,3424069,100,8.621851440417988,9.450000000000001
+WRSWR-SKIP,34240686,342407,100,2.5124268161566627,9.22
+WRSWR-SKIP,34240686,34241,100,1.74814293673906,8.180000000000001
+WRSWR-SKIP,34240686,3425,100,1.6209248611432612,8.120000000000001

comparison/benchmark_results.csv

@@ -0,0 +1,37 @@
+alg,N,m,d,repetitions,t_add,t_get
+WRSWR-BIN,10000000,1000,d_decr,100,15.351767258000008,3.6700000000000004
+WRSWR-BIN,10000000,10000,d_decr,100,15.476184351000002,3.94
+WRSWR-BIN,10000000,100000,d_decr,100,15.719699089999997,3.66
+WRSWR-BIN,10000000,1000000,d_decr,100,19.063183873999996,4.680000000000001
+WRAExp-J,10000000,1000,d_decr,100,3.0741342810000005,81035.38
+WRAExp-J,10000000,10000,d_decr,100,3.3204783489999996,914356.77
+WRAExp-J,10000000,100000,d_decr,100,5.671898599000001,1.042682492e7
+WRAExp-J,10000000,1000000,d_decr,100,30.680991763000012,1.2113268227000001e8
+WRSWR-SKIP,10000000,1000,d_decr,100,2.5568252619999994,10.120000000000001
+WRSWR-SKIP,10000000,10000,d_decr,100,2.6231644160000007,7.570000000000001
+WRSWR-SKIP,10000000,100000,d_decr,100,3.0551440570000015,9.71
+WRSWR-SKIP,10000000,1000000,d_decr,100,6.853534331000001,9.91
+WRSWR-BIN,10000000,1000,d_const,100,15.514981697000003,4.670000000000001
+WRSWR-BIN,10000000,10000,d_const,100,15.661928006000004,4.65
+WRSWR-BIN,10000000,100000,d_const,100,18.322097050000004,4.7700000000000005
+WRSWR-BIN,10000000,1000000,d_const,100,43.01236574400001,4.700000000000001
+WRAExp-J,10000000,1000,d_const,100,3.1586977939999996,103805.02
+WRAExp-J,10000000,10000,d_const,100,4.0101058410000014,1.06259209e6
+WRAExp-J,10000000,100000,d_const,100,11.213429000999996,1.255766748e7
+WRAExp-J,10000000,1000000,d_const,100,65.902981587,1.5659078689e8
+WRSWR-SKIP,10000000,1000,d_const,100,2.6010059800000014,6.240000000000001
+WRSWR-SKIP,10000000,10000,d_const,100,2.889858204,8.270000000000001
+WRSWR-SKIP,10000000,100000,d_const,100,5.284384767000001,9.72
+WRSWR-SKIP,10000000,1000000,d_const,100,20.140484113000003,9.660000000000002
+WRSWR-BIN,10000000,1000,d_incr,100,15.512825132000005,4.6000000000000005
+WRSWR-BIN,10000000,10000,d_incr,100,15.966000717999995,5.12
+WRSWR-BIN,10000000,100000,d_incr,100,21.065067476000003,6.3
+WRSWR-BIN,10000000,1000000,d_incr,100,66.010967182,3.8900000000000006
+WRAExp-J,10000000,1000,d_incr,100,3.4165506430000008,96650.21
+WRAExp-J,10000000,10000,d_incr,100,4.7180384370000015,1.0320041799999999e6
+WRAExp-J,10000000,100000,d_incr,100,17.500391315000005,1.2637909870000001e7
+WRAExp-J,10000000,1000000,d_incr,100,103.122172538,1.5031790519000003e8
+WRSWR-SKIP,10000000,1000,d_incr,100,2.6298687370000007,7.170000000000001
+WRSWR-SKIP,10000000,10000,d_incr,100,3.216991788000002,6.71
+WRSWR-SKIP,10000000,100000,d_incr,100,7.403542585000004,8.180000000000001
+WRSWR-SKIP,10000000,1000000,d_incr,100,30.919183713000013,9.800000000000002

comparison/samplers.jl

@@ -78,7 +78,7 @@ macro quantile_fast(k)
     append!(block.args, firstv.args)
     for i in 2:k
         nextv = quote
-            $(esc(:s)) *= ($(esc(:n)) - $i) * $(esc(:p))
+            $(esc(:s)) *= ($(esc(:n)) - $(i - 1)) * $(esc(:p))
             $(esc(:q)) *= 1. - $(esc(:p))
             $(esc(:x)) += $(esc(:s)) / ($(esc(:q)) * $(factorial(i)))
             $(esc(:x)) > $(esc(:nt)) && return $i
@@ -98,7 +98,7 @@ end
 
 function get(s::AlgWRSWRSKIPSampler)
 	if s.seen_k < s.n
-		return sample(s.rng, s.value[1:s.seen_k], Weights(s.weigths[1:s.seen_k]), s.n)
+		return sample(s.rng, s.value[1:s.seen_k], Weights(diff([0.0; s.weights[1:s.seen_k]])), s.n)
 	else 
 		return s.value
 	end
@@ -253,17 +253,17 @@ function get(s::AlgWRAExpJSampler{<:BinaryHeap{Pair{T, Tuple{Float64,Float64}}}}
     m = 2
     @inbounds for j in 2:s.n
         i = rand(s.rng, sampler)
-        if i <= j-1
+        if i < m
             out[j] = kvs[i][1]
         else
             out[j] = kvs[m][1]
+            if m < s.n
+                sampler[m] = kvs[m][2][2]
+                Wnew -= kvs[m][2][2]
+                sampler[m+1] = Wnew
+            end
             m += 1
         end
-        if j < s.n
-            sampler[j] = kvs[j][2][2]
-            Wnew -= kvs[j][2][2]
-            sampler[j+1] = Wnew
-        end
     end
     return out
 end

src/SamplingInterface.jl

@@ -223,10 +223,10 @@ function itsample(iter, method = AlgRSWRSKIP(); iter_type = infer_eltype(iter))
     return itsample(Random.default_rng(), iter, method; iter_type)
 end
 function itsample(iter, n::Int, method = AlgL(); iter_type = infer_eltype(iter), ordered = false)
-    return itsample(Random.default_rng(), iter, n, method; ordered)
+    return itsample(Random.default_rng(), iter, n, method; iter_type, ordered)
 end
 function itsample(iter, wv::Function, method = AlgWRSWRSKIP(); iter_type = infer_eltype(iter))
-    return itsample(Random.default_rng(), iter, wv, method)
+    return itsample(Random.default_rng(), iter, wv, method; iter_type)
 end
 function itsample(iter, wv::Function, n::Int, method = AlgAExpJ(); iter_type = infer_eltype(iter), 
         ordered = false)
@@ -247,7 +247,7 @@ Base.@constprop :aggressive function itsample(rng::AbstractRNG, iter, n::Int, me
         s = ReservoirSampler{iter_type,Float64}(rng, n, method, ImmutSampler(), ordered ? Ord() : Unord())
         return update_all!(s, iter, ordered)
     else
-        m = method isa AlgL || method isa AlgR || method isa AlgD ? AlgD() : AlgORDSWR()
+        m = method isa Union{AlgL, AlgR, AlgD} ? AlgD() : method isa AlgHiddenShuffle ? method : AlgORDSWR()
         s = collect(SequentialSampler{iter_type}(rng, iter, n, length(iter), m))
         return ordered ? s : fshuffle!(rng, s)
     end

src/SamplingReduction.jl

@@ -3,19 +3,19 @@ const SMWR = Union{MultiAlgRSWRSKIPSampler, MultiAlgWRSWRSKIPSampler}
 const SMWOWR = Union{MultiAlgAResSampler, MultiAlgAExpJSampler}
 
 reduce_samples(t) = error()
-function reduce_samples(t::Union{TypeS,TypeUnion}, ss::BinaryHeap...)
-    nt = length(ss)
-    n = minimum(length.(ss))
-    lkeys = sort(reduce(vcat, [s.valtree for s in ss]), by=(x->x[end]), rev=true)[1:n]
-    return lkeys
+function reduce_samples(t::Union{TypeS,TypeUnion}, n::Integer, ss::BinaryHeap...)
+    lkeys = sort(reduce(vcat, [s.valtree for s in ss]), by=(x->x[end]), rev=true)
+    return lkeys[1:min(n, length(lkeys))]
 end
 function reduce_samples(ps::AbstractArray, rngs, t::Union{TypeS,TypeUnion}, ss::AbstractArray...)
+    return reduce_samples(ps, rngs, t, minimum(length.(ss)), ss...)
+end
+function reduce_samples(ps::AbstractArray, rngs, t::Union{TypeS,TypeUnion}, n::Integer, ss::AbstractArray...)
     nt = length(ss)
     T = get_type_rs(t, ss...)
     v = Vector{Vector{T}}(undef, nt)
-    n = minimum(length.(ss))
     ns = rand(extract_rng(rngs, 1), Multinomial(n, ps))
-    Threads.@threads for i in 1:nt
+    for i in 1:nt
         s = ss[i]
         vi = Vector{T}(undef, ns[i])
         @inbounds for (q, j) in enumerate(SequentialSampler(extract_rng(rngs, 1), 
@@ -27,6 +27,29 @@ function reduce_samples(ps::AbstractArray, rngs, t::Union{TypeS,TypeUnion}, ss::
     return reduce(vcat, v)
 end
 
+# reservoirs which together saw fewer than n elements are all still filling up, so
+# they hold their elements as seen and merging them amounts to concatenating these
+function append_unfilled!(value, k, ss::MultiAlgRSWRSKIPSampler...)
+    for s in ss
+        @inbounds for i in 1:s.seen_k
+            value[k+i] = s.value[i]
+        end
+        k += s.seen_k
+    end
+    return value
+end
+function append_unfilled!(value, weights, k, w, ss::MultiAlgWRSWRSKIPSampler...)
+    for s in ss
+        @inbounds for i in 1:s.seen_k
+            value[k+i] = s.value[i]
+            weights[k+i] = w + s.weights[i]
+        end
+        k += s.seen_k
+        w += s.state
+    end
+    return value
+end
+
 extract_rng(v::AbstractArray, i) = v[i]
 extract_rng(v::AbstractRNG, i) = v
 

src/SamplingUtils.jl

@@ -43,6 +43,7 @@ struct SeqSampleIter{R}
 end
 
 @inline function Base.iterate(it::SeqSampleIter)
+    it.n == 0 && return nothing
     i = 0
     q1 = it.N - it.n + 1
     q2 = q1 / it.N

src/SortedSamplingMulti.jl

@@ -71,7 +71,9 @@ end
 
 @inline function Base.iterate(s::MultiAlgORDSampler)
     indices, iter = s.inds, s.it
-    curr_idx, state_idx = iterate(indices)::Tuple
+    it_indices = iterate(indices)
+    it_indices === nothing && return nothing
+    curr_idx, state_idx = it_indices
     el, state_el = iterate(iter)::Tuple
     for _ in 1:curr_idx-1
         el, state_el = iterate(iter, state_el)::Tuple

src/UnweightedSamplingMulti.jl

@@ -173,7 +173,7 @@ macro quantile_fast(k)
     append!(block.args, firstv.args)
     for i in 2:k
         nextv = quote
-            $(esc(:s)) *= ($(esc(:n)) - $i) * $(esc(:p))
+            $(esc(:s)) *= ($(esc(:n)) - $(i - 1)) * $(esc(:p))
             $(esc(:q)) *= 1. - $(esc(:p))
             $(esc(:x)) += $(esc(:s)) / ($(esc(:q)) * $(factorial(i)))
             $(esc(:x)) > $(esc(:nt)) && return $i
@@ -216,11 +216,16 @@ function Base.merge(ss::MultiAlgLSampler...)
     error("To Be Implemented")
 end
 function Base.merge(ss::MultiAlgRSWRSKIPSampler...)
-    newvalue = reduce_samples(get_ps(ss...), [s.rng for s in ss], TypeUnion(), value.(ss)...)
-    skip_k = sum(getfield(s, :skip_k) for s in ss)
-    seen_k = sum(getfield(s, :seen_k) for s in ss)
     n = minimum(s.n for s in ss)
-    return MultiAlgRSWRSKIPSampler_Mut(n, skip_k, seen_k, ss[1].rng, newvalue, nothing)
+    seen_k = sum(getfield(s, :seen_k) for s in ss)
+    if seen_k < n
+        newvalue = Vector{get_type_rs(TypeUnion(), (s.value for s in ss)...)}(undef, n)
+        append_unfilled!(newvalue, 0, ss...)
+        return MultiAlgRSWRSKIPSampler_Mut(n, 0, seen_k, ss[1].rng, newvalue, nothing)
+    end
+    newvalue = reduce_samples(get_ps(ss...), [s.rng for s in ss], TypeUnion(), n, value.(ss)...)
+    s = MultiAlgRSWRSKIPSampler_Mut(n, 0, seen_k, ss[1].rng, newvalue, nothing)
+    return recompute_skip!(s, n)
 end
 
 function Base.merge!(ss::MultiAlgRSampler...)
@@ -231,12 +236,17 @@ function Base.merge!(ss::MultiAlgLSampler...)
 end
 function Base.merge!(s1::MultiAlgRSWRSKIPSampler{<:Nothing}, ss::MultiAlgRSWRSKIPSampler...)
     s1.n > minimum(s.n for s in ss) && error("The size of the mutated reservoir should be the minimum size between all merged reservoir")
-    newvalue = reduce_samples(get_ps(s1, ss...), [s1.rng, [s.rng for s in ss]...], TypeS(), value(s1), value.(ss)...)
-    for i in 1:length(newvalue)
-        @inbounds s1.value[i] = newvalue[i]
+    seen_k = s1.seen_k + sum(getfield(s, :seen_k) for s in ss)
+    if seen_k < s1.n
+        append_unfilled!(s1.value, s1.seen_k, ss...)
+    else
+        newvalue = reduce_samples(get_ps(s1, ss...), [s1.rng, [s.rng for s in ss]...], TypeS(), s1.n, value(s1), value.(ss)...)
+        for i in 1:length(newvalue)
+            @inbounds s1.value[i] = newvalue[i]
+        end
     end
-    s1.skip_k += sum(getfield(s, :skip_k) for s in ss)
-    s1.seen_k += sum(getfield(s, :seen_k) for s in ss)
+    s1.seen_k = seen_k
+    seen_k >= s1.n && recompute_skip!(s1, s1.n)
     return s1
 end
 

src/UnweightedSamplingSingle.jl

@@ -46,7 +46,7 @@ function Base.merge(ss::SingleAlgRSWRSKIPSampler...)
     ps = cumsum(ns ./ n_tot)
     r = rand(ss[1].rng)
     value = ss[findfirst(p -> r < p, ps)].rvalue
-    return typeof(ss[1])(n_tot, sum(s.skip_k for s in ss), ss[1].rng, value)
+    return typeof(ss[1])(n_tot, ceil(Int, n_tot/rand(ss[1].rng)), ss[1].rng, value)
 end
 
 function Base.merge!(s1::SingleAlgRSWRSKIPSampler_Mut, ss::SingleAlgRSWRSKIPSampler_Mut...)
@@ -59,6 +59,6 @@ function Base.merge!(s1::SingleAlgRSWRSKIPSampler_Mut, ss::SingleAlgRSWRSKIPSamp
         s1.rvalue = RefVal_Immut(ss[i-1].rvalue.value)
     end
     s1.seen_k += sum(s.seen_k for s in ss)
-    s1.skip_k += sum(s.skip_k for s in ss)
+    s1.skip_k = ceil(Int, s1.seen_k/rand(s1.rng))
     return s1
 end

src/WeightedSamplingMulti.jl

@@ -137,7 +137,7 @@ end
                 while s.weights[j] < curx * s.state
                     j += 1
                 end
-                newvalues[i] = s.value[j]
+                newvalues[n-i+1] = s.value[j]
             end
             s.value .= newvalues
             s = @inline recompute_skip!(s, n)
@@ -189,66 +189,75 @@ end
 extract_T(::DataStructures.BinaryHeap{T}) where T = T
 
 function Base.merge(ss::MultiAlgAResSampler...)
-    newvalue = reduce_samples(TypeUnion(), [s.value for s in ss]...)
+    n = minimum(s.n for s in ss)
+    newvalue = reduce_samples(TypeUnion(), n, [s.value for s in ss]...)
     newheap = BinaryHeap(Base.By(last, DataStructures.FasterForward()), newvalue)
     seen_k = sum(getfield(s, :seen_k) for s in ss)
-    n = minimum(s.n for s in ss)
     s = MultiAlgAResSampler_Mut(seen_k, n, ss[1].rng, newheap)
     return s
 end
 function Base.merge(ss::MultiAlgAExpJSampler...)
-    newvalue = reduce_samples(TypeUnion(), [s.value for s in ss]...)
+    n = minimum(s.n for s in ss)
+    newvalue = reduce_samples(TypeUnion(), n, [s.value for s in ss]...)
     newheap = BinaryHeap(Base.By(last, DataStructures.FasterForward()), newvalue)
     seen_k = sum(getfield(s, :seen_k) for s in ss)
-    state = sum(getfield(s, :state) for s in ss)
-    min_priority = minimum(getfield(s, :min_priority) for s in ss)
-    n = minimum(s.n for s in ss)
-    s = MultiAlgAExpJSampler_Mut(state, min_priority, seen_k, n, ss[1].rng, newheap)
+    z = zero(ss[1].state)
+    s = MultiAlgAExpJSampler_Mut(z, z, seen_k, n, ss[1].rng, newheap)
+    seen_k >= n && recompute_skip!(s)
     return s
 end
 function Base.merge(ss::MultiAlgWRSWRSKIPSampler...)
-    newvalue = reduce_samples(get_ps(ss...), [s.rng for s in ss], TypeUnion(), value.(ss)...)
-    skip_w = sum(getfield(s, :skip_w) for s in ss)
+    n = minimum(s.n for s in ss)
     state = sum(getfield(s, :state) for s in ss)
     seen_k = sum(getfield(s, :seen_k) for s in ss)
-    n = minimum(s.n for s in ss)
-    s = MultiAlgWRSWRSKIPSampler_Mut(n, state, skip_w, seen_k, ss[1].rng, Memory{Float64}(undef,0), newvalue, nothing)
-    return s
+    weights = Memory{typeof(state)}(undef, n)
+    if seen_k < n
+        newvalue = Vector{get_type_rs(TypeUnion(), (s.value for s in ss)...)}(undef, n)
+        append_unfilled!(newvalue, weights, 0, zero(state), ss...)
+        return MultiAlgWRSWRSKIPSampler_Mut(n, state, zero(state), seen_k, ss[1].rng, weights, newvalue, nothing)
+    end
+    newvalue = reduce_samples(get_ps(ss...), [s.rng for s in ss], TypeUnion(), n, value.(ss)...)
+    s = MultiAlgWRSWRSKIPSampler_Mut(n, state, zero(state), seen_k, ss[1].rng, weights, newvalue, nothing)
+    return recompute_skip!(s, n)
 end
 
 function Base.merge!(s1::MultiAlgAResSampler, ss::MultiAlgAResSampler...)
-    length(typeof(s1.value.valtree).parameters) == 3 && error("Merging ordered reservoirs is not possible")
+    eltype(s1.value.valtree) <: Tuple && error("Merging ordered reservoirs is not possible")
     s1.n > minimum(s.n for s in ss) && error("The size of the mutated reservoir should be the minimum size between all merged reservoir")
+    newvalue = reduce_samples(TypeS(), s1.n, s1.value, [s.value for s in ss]...)
     empty!(s1.value.valtree)
-    newvalue = reduce_samples(TypeS(), s1.value, [s.value for s in ss]...)
     for e in newvalue
         push!(s1.value, e[1] => e[2])
     end
     s1.seen_k += sum(getfield(s, :seen_k) for s in ss)
     return s1
 end
 function Base.merge!(s1::MultiAlgAExpJSampler, ss::MultiAlgAExpJSampler...)
-    length(typeof(s1.value.valtree).parameters) == 3 && error("Merging ordered reservoirs is not possible")
+    eltype(s1.value.valtree) <: Tuple && error("Merging ordered reservoirs is not possible")
     s1.n > minimum(s.n for s in ss) && error("The size of the mutated reservoir should be the minimum size between all merged reservoir")
+    newvalue = reduce_samples(TypeS(), s1.n, s1.value, [s.value for s in ss]...)
     empty!(s1.value.valtree)
-    newvalue = reduce_samples(TypeS(), s1.value, [s.value for s in ss]...)
     for e in newvalue
         push!(s1.value, e[1] => e[2])
     end
     s1.seen_k += sum(getfield(s, :seen_k) for s in ss)
-    s1.state += sum(getfield(s, :state) for s in ss)
-    s1.min_priority = min(s1.min_priority, minimum(getfield(s, :min_priority) for s in ss))
+    s1.seen_k >= s1.n && recompute_skip!(s1)
     return s1
 end
 function Base.merge!(s1::MultiAlgWRSWRSKIPSampler{<:Nothing}, ss::MultiAlgWRSWRSKIPSampler...)
     s1.n > minimum(s.n for s in ss) && error("The size of the mutated reservoir should be the minimum size between all merged reservoir")
-    newvalue = reduce_samples(get_ps(s1, ss...), [s1.rng, [s.rng for s in ss]...], TypeS(), value(s1), value.(ss)...)
-    for i in 1:length(newvalue)
-        @inbounds s1.value[i] = newvalue[i]
+    seen_k = s1.seen_k + sum(getfield(s, :seen_k) for s in ss)
+    if seen_k < s1.n
+        append_unfilled!(s1.value, s1.weights, s1.seen_k, s1.state, ss...)
+    else
+        newvalue = reduce_samples(get_ps(s1, ss...), [s1.rng, [s.rng for s in ss]...], TypeS(), s1.n, value(s1), value.(ss)...)
+        for i in 1:length(newvalue)
+            @inbounds s1.value[i] = newvalue[i]
+        end
     end
-    s1.skip_w += sum(getfield(s, :skip_w) for s in ss)
     s1.state += sum(getfield(s, :state) for s in ss)
-    s1.seen_k += sum(getfield(s, :seen_k) for s in ss)
+    s1.seen_k = seen_k
+    seen_k >= s1.n && recompute_skip!(s1, s1.n)
     return s1
 end
 

src/WeightedSamplingSingle.jl

@@ -51,8 +51,7 @@ function Base.merge(ss::SingleAlgWRSWRSKIPSampler...)
     ps = cumsum(ns ./ n_tot)
     r = rand(ss[1].rng)
     value = ss[findfirst(p -> r < p, ps)].rvalue
-    return typeof(ss[1])(sum(s.seen_k for s in ss), sum(s.total_w for s in ss), sum(s.skip_w for s in ss), 
-                      ss[1].rng, value)
+    return typeof(ss[1])(sum(s.seen_k for s in ss), n_tot, n_tot/rand(ss[1].rng), ss[1].rng, value)
 end
 
 function Base.merge!(s1::SingleAlgWRSWRSKIPSampler_Mut, ss::SingleAlgWRSWRSKIPSampler_Mut...)
@@ -65,8 +64,8 @@ function Base.merge!(s1::SingleAlgWRSWRSKIPSampler_Mut, ss::SingleAlgWRSWRSKIPSa
         s1.rvalue = RefVal_Immut(ss[i-1].rvalue.value)
     end
     s1.seen_k += sum(s.seen_k for s in ss)
-    s1.skip_w += sum(s.skip_w for s in ss)
     s1.total_w += sum(s.total_w for s in ss)
+    s1.skip_w = s1.total_w/rand(s1.rng)
     return s1
 end