1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
use super::cover::K;
use alloc::vec::Vec;
use core::f64::consts::E;
use fastrand;
use std::{io, vec};
/// Creates a representative sample of `input` of `size` bytes.
pub fn create_sample<R: io::Read>(input: &mut R, size: usize) -> Vec<u8> {
let reservoir = Reservoir::new(size);
reservoir.fill(input)
}
/// A reservoir is created from an input stream.
///
/// Once filled, it will contain a best effort sample of a dataset,
/// where each input value has an equivalent probability of being included.
struct Reservoir {
/// Where the sampled data is stored.
///
/// Once the lake is filled, then this should contain a representative sample
/// of the larger dataset.
lake: Vec<u8>,
/// K is the size of each sample.
///
/// The original Zstd dictionary implementation states that values
/// between 16 and 2048+ are reasonable.
k: u16,
}
impl Reservoir {
/// Initialize a new empty reservoir, creating an allocation of `size`.
pub fn new(size: usize) -> Self {
assert!(size >= 16, "Reservoirs cannot be below 16 bytes in size");
let lake: Vec<u8> = vec![0; size];
let k = K as u16;
Self { lake, k }
}
/// Filling the reservoir is performed using Algorithm L.
///
/// The return value is the populated reservoir.
pub fn fill<R: io::Read>(mut self, source: &mut R) -> Vec<u8> {
// https://en.wikipedia.org/wiki/Reservoir_sampling#:~:text=end%0A%20%20end%0Aend-,Optimal%3A%20Algorithm,-L%5Bedit
// https://richardstartin.github.io/posts/reservoir-sampling#algorithm-l:~:text=%3B%0A%20%20%20%20%7D%0A%7D-,Algorithm%20L,-Algorithm%20L%20was
// First fill the reservoir with the start of the input stream
let mut total_bytes_read: usize = 0;
while let Ok(num_bytes) = source.read(self.lake.as_mut_slice()) {
total_bytes_read += num_bytes;
// Stop when we've completely filled the buffer
if total_bytes_read == self.lake.len() {
break;
}
// If we haven't filled the lake all the way, resize it
if num_bytes == 0 {
self.lake.resize(total_bytes_read, 0);
}
}
let mut threshold = E.powf(fastrand::f64().ln() / f64::from(self.k));
// An index into the stream of the next sample to take
let mut next = self.lake.len();
// Because we're sampling k-mers of size K into the lake,
// split the lake into chunks of k size for simplicity
let mut lake_chunks = self
.lake
.chunks_mut(self.k as usize)
.collect::<Vec<&mut [u8]>>();
// Used when discarding chunks
let end_of_lake = lake_chunks.len();
let mut counter = end_of_lake / self.k as usize;
// Algorithm L is considered better than algorithm R because it
// determines how many inputs can be skipped, rather than
// processing every input.
// This is done by abusing the statistics in ways
// I do not understand.
// Items with a weight smaller than the threshold enter the lake,
// replacing the item in the lake with the largest threshold
let mut dumpster = Vec::with_capacity(self.k as usize);
loop {
// `num_bytes_read` is kept track of to watch for EOD.
let num_bytes_read: u64;
if counter == next {
num_bytes_read = source
.read(lake_chunks[fastrand::usize(0..end_of_lake)])
.unwrap() as u64;
// Advance at least to the next sample, skipping forward a few samples
next += ((fastrand::f64().ln() / f64::ln(1.0 - threshold)).floor() as usize + 1)
* self.k as usize;
// Update the threshold to reflect changes
threshold *= E.powf(fastrand::f64().ln() / f64::from(end_of_lake as u32))
} else {
// Drop the next chunk
num_bytes_read = source.read(&mut dumpster).unwrap() as u64;
//source.seek_relative(self.k.into()).unwrap();
}
if num_bytes_read == 0 {
break;
}
counter += self.k as usize;
}
self.lake.shrink_to_fit();
self.lake
}
}
#[cfg(test)]
mod tests;