data.rs

1use crate::rng::Rng;2use crate::Json;34/// A dataset source: a named, seeded generator of JSON rows.5///6/// Pouring obeys the determinism law: the same seed pours the same rows in7/// the same order, and a shorter pour is a prefix of a longer one. A well8/// may tag a row `{"split": "eval"}` to claim it for evaluation; untagged9/// rows are training rows.10///11/// ```12/// use mrlycore::data::Well;13/// use mrlycore::{json, Json};14///15/// struct Squares;16///17/// impl Well for Squares {18///     fn name(&self) -> &str {19///         "squares"20///     }21///     fn about(&self) -> &str {22///         "The squares of the naturals."23///     }24///     fn pour(&self, _seed: u64, count: usize) -> Vec<Json> {25///         (0..count).map(|n| json!({ "n": n, "square": n * n })).collect()26///     }27/// }28///29/// let well: Box<dyn Well> = Box::new(Squares);30/// assert_eq!(well.pour(7, 2)[..], well.pour(7, 3)[..2]);31/// ```32pub trait Well {33    /// Returns the dataset's name.34    fn name(&self) -> &str;35    /// Returns the dataset's one-line description.36    fn about(&self) -> &str;37    /// Pours count seeded rows, identical for identical seeds, earlier rows first.38    fn pour(&self, seed: u64, count: usize) -> Vec<Json>;39}4041/// Shuffles rows into the seed's order, so a short pour samples an enumeration fairly.42pub fn shuffle(rows: &mut [Json], seed: u64) {43    let mut rng = Rng::new(seed);44    for i in (1..rows.len()).rev() {45        let j = rng.below(i + 1);46        rows.swap(i, j);47    }48}4950#[cfg(test)]51mod tests {52    use super::*;53    use crate::json;5455    #[test]56    fn shuffle_replays_and_reorders() {57        let fresh = || (0..32).map(|n| json!({ "n": n })).collect::<Vec<Json>>();58        let (mut a, mut b, straight) = (fresh(), fresh(), fresh());59        shuffle(&mut a, 7);60        shuffle(&mut b, 7);61        assert_eq!(a, b);62        assert_ne!(a, straight);63        let mut sorted = a.clone();64        sorted.sort_by_key(|row| row["n"].as_i64());65        assert_eq!(sorted, straight);66    }6768    #[test]69    fn different_seeds_deal_different_orders() {70        let fresh = || (0..32).map(|n| json!({ "n": n })).collect::<Vec<Json>>();71        let (mut a, mut b) = (fresh(), fresh());72        shuffle(&mut a, 1);73        shuffle(&mut b, 2);74        assert_ne!(a, b);75    }76}