rsv_lib/io/
sample.rs

1use crate::args::Sample;
2use crate::utils::cli_result::CliResult;
3use crate::utils::filename::new_file;
4use crate::utils::priority_queue::PriorityQueue;
5use crate::utils::reader::IoReader;
6use crate::utils::table::Table;
7use crate::utils::writer::Writer;
8use rand::rng;
9use rand::rngs::StdRng;
10use rand::{Rng, SeedableRng};
11use std::borrow::Cow;
12use std::time::Instant;
13
14impl Sample {
15    pub fn io_run(&self) -> CliResult {
16        let time_limit = (self.time_limit - 0.7).clamp(0.0, f32::MAX);
17
18        // open files
19        let lines = IoReader::new().lines();
20
21        if lines.is_empty() {
22            return Ok(());
23        }
24
25        // header
26        let header = match self.no_header {
27            true => None,
28            false => Some(lines[0].to_owned()),
29        };
30
31        // seed
32        let mut rng = match self.seed {
33            Some(s) => StdRng::seed_from_u64(s as u64),
34            None => StdRng::from_rng(&mut rng()),
35        };
36
37        // read
38        let mut queue = PriorityQueue::with_capacity(self.n);
39        let time = Instant::now();
40        for (line_n, r) in lines
41            .into_iter()
42            .skip(1 - self.no_header as usize)
43            .enumerate()
44        {
45            let priority = rng.random::<f64>();
46            if queue.can_insert(priority) {
47                queue.push(line_n, priority, r);
48            }
49
50            if time_limit > 0.0 && line_n % 10000 == 0 && time.elapsed().as_secs_f32() >= time_limit
51            {
52                break;
53            }
54        }
55
56        match (self.export, self.show_number) {
57            (true, _) => write_to_file(header, queue)?,
58            (false, true) => print_to_stdout(header, queue)?,
59            (false, false) => print_to_stdout_no_number(header, queue)?,
60        };
61
62        Ok(())
63    }
64}
65
66fn write_to_file(header: Option<String>, queue: PriorityQueue<String>) -> CliResult {
67    // new file
68    let out = new_file("sampled.csv");
69    let mut wtr = Writer::new(&out)?;
70    if let Some(r) = header {
71        wtr.write_str_unchecked(r);
72    }
73    for r in queue.into_sorted_items() {
74        wtr.write_str_unchecked(&r.item);
75    }
76
77    println!("Saved to file: {}", out.display());
78
79    Ok(())
80}
81
82fn print_to_stdout(header: Option<String>, queue: PriorityQueue<String>) -> CliResult {
83    let mut table = Table::new();
84
85    if let Some(h) = header {
86        table.add_record([Cow::Borrowed("#"), Cow::Borrowed(""), Cow::from(h)]);
87    }
88
89    queue.into_sorted_items().into_iter().for_each(|i| {
90        table.add_record([
91            Cow::from(i.line_n_as_string()),
92            Cow::Borrowed("->"),
93            Cow::from(i.item),
94        ])
95    });
96
97    table.print_blank_unchecked();
98
99    Ok(())
100}
101
102fn print_to_stdout_no_number(header: Option<String>, queue: PriorityQueue<String>) -> CliResult {
103    let mut wtr = Writer::stdout().unwrap();
104
105    // header
106    if let Some(h) = header {
107        wtr.write_str_unchecked(h);
108    }
109
110    // samples
111    queue.into_sorted_items().into_iter().for_each(|i| {
112        wtr.write_str_unchecked(i.item);
113    });
114
115    Ok(())
116}
rsv_lib/io/sample.rs

rsv_lib/io/
sample.rs
