static USAGE: &str = r#"
Randomly samples CSV data.
It supports eight sampling methods:
* RESERVOIR: the default sampling method when NO INDEX is present and no sampling method
is specified. Visits every CSV record exactly once, using MEMORY PROPORTIONAL to the
sample size (k) - O(k).
https://en.wikipedia.org/wiki/Reservoir_sampling
* INDEXED: the default sampling method when an INDEX is present and no sampling method
is specified. Uses random I/O to sample efficiently, as it only visits records selected
by random indexing, using MEMORY PROPORTIONAL to the sample size (k) - O(k).
https://en.wikipedia.org/wiki/Random_access
* BERNOULLI: the sampling method when the --bernoulli option is specified.
Each record has an independent probability p of being selected, where p is
specified by the <sample-size> argument. For example, if p=0.1, then each record
has a 10% chance of being selected, regardless of the other records. The final
sample size is random and follows a binomial distribution. Uses CONSTANT MEMORY - O(1).
When sampling from a remote URL, processes the file in chunks without downloading it
entirely, making it especially efficient for sampling large remote files.
https://en.wikipedia.org/wiki/Bernoulli_sampling
* SYSTEMATIC: the sampling method when the --systematic option is specified.
Selects every nth record from the input, where n is the integer part of <sample-size>
and the fraction part is the percentage of the population to sample.
For example, if <sample-size> is 10.5, it will select every 10th record and 50% of the
population. If <sample-size> is a whole number (no fractional part), it will select
every nth record for the whole population. Uses CONSTANT memory - O(1). The starting
point can be specified as "random" or "first". Useful for time series data or when you
want evenly spaced samples.
https://en.wikipedia.org/wiki/Systematic_sampling
* STRATIFIED: the sampling method when the --stratified option is specified.
Stratifies the population by the specified column and then samples from each stratum.
Particularly useful when a population has distinct subgroups (strata) that are
heterogeneous within but homogeneous between in terms of the variable of interest.
For example, if you want to sample 1,000 records from a population of 100,000 across the US,
you can stratify the population by US state and then sample 20 records from each stratum.
This will ensure that you have a representative sample from each of the 50 states.
The sample size must be a whole number. Uses MEMORY PROPORTIONAL to the
number of strata (s) and samples per stratum (k) as specified by <sample-size> - O(s*k).
https://en.wikipedia.org/wiki/Stratified_sampling
* WEIGHTED: the sampling method when the --weighted option is specified.
Samples records with probabilities proportional to values in a specified weight column.
Records with higher weights are more likely to be selected. For example, if you have
sales data and want to sample transactions weighted by revenue, high-value transactions
will have a higher chance of being included. Non-numeric weights are treated as zero.
The weights are automatically normalized using the maximum weight in the dataset.
Specify the desired sample size with <sample-size>. Uses MEMORY PROPORTIONAL to the
sample size (k) - O(k).
"Weighted random sampling with a reservoir" https://doi.org/10.1016/j.ipl.2005.11.003
* CLUSTER: the sampling method when the --cluster option is specified.
Samples entire groups of records together based on a cluster identifier column.
The number of clusters is specified by the <sample-size> argument.
Useful when records are naturally grouped (e.g., by household, neighborhood, etc.).
For example, if you have records grouped by neighborhood and specify a sample size of 10,
it will randomly select 10 neighborhoods and include ALL records from those neighborhoods
in the output. This ensures that natural groupings in the data are preserved.
Uses MEMORY PROPORTIONAL to the number of clusters (c) - O(c).
https://en.wikipedia.org/wiki/Cluster_sampling
* TIMESERIES: the sampling method when the --timeseries option is specified.
Samples records based on time intervals from a time-series dataset. Groups records by
time windows (e.g., hourly, daily, weekly) and selects one record per interval.
Supports adaptive sampling (e.g., prefer business hours or weekends) and aggregation
(e.g., mean, sum, min, max) within each interval. The starting point can be "first"
(earliest), "last" (most recent), or "random". Particularly useful for time-series data
where simple row-based sampling would always return the same records due to sorting.
Uses MEMORY PROPORTIONAL to the number of records - O(n).
Supports sampling from CSVs on remote URLs. Note that the entire file is downloaded first
to a temporary file before sampling begins for all sampling methods except Bernoulli, which
streams the file as it samples it, stopping when the desired sample size is reached or the
end of the file is reached.
Sampling from stdin is also supported for all sampling methods, copying stdin to a in-memory
buffer first before sampling begins.
If a stats cache is available, it will be used to do extra checks on systematic,
weighted and cluster sampling, and to speed up sampling in general.
This command is intended to provide a means to sample from a CSV data set that
is too big to fit into memory (for example, for use with commands like
'qsv stats' with the '--everything' option).
Examples:
# Take a sample of 1000 records from data.csv using RESERVOIR or INDEXED sampling
# depending on whether an INDEX is present.
qsv sample 1000 data.csv
# Take a sample of approximately 10% of the records from data.csv using RESERVOIR
# or INDEXED sampling depending on whether an INDEX is present.
qsv sample 0.1 data.csv
# Take a sample using BERNOULLI sampling where each record has a 5% chance of being selected
qsv sample --bernoulli 0.05 data.csv
# Take a sample using SYSTEMATIC sampling where every 10th record is selected
# and approximately 50% of the population is sampled, starting from a random point.
qsv sample --systematic random 10.5 data.csv
# Take a sample using STRATIFIED sampling where 20 records are sampled from each
# stratum defined by the 'State' column.
qsv sample --stratified State 20 data.csv
# Take a sample using WEIGHTED sampling where records are sampled with probabilities
# proportional to the 'Revenue' column, for a total sample size of 1000 records.
qsv sample --weighted Revenue 1000 data.csv
# Take a sample using CLUSTER sampling where 10 clusters defined by the
# 'Neighborhood' column are randomly selected and all records from those clusters
# are included in the sample.
qsv sample --cluster Neighborhood 10 data.csv
For more examples, see https://github.com/dathere/qsv/blob/master/tests/test_sample.rs.
Usage:
qsv sample [options] <sample-size> [<input>]
qsv sample --help
sample arguments:
<input> The CSV file to sample. This can be a local file,
stdin, or a URL (http and https schemes supported).
<sample-size> When using INDEXED, RESERVOIR or WEIGHTED sampling, the sample size.
Can either be a whole number or a value between value between 0 and 1.
If a fraction, specifies the sample size as a percentage of the population.
(e.g. 0.15 - 15 percent of the CSV)
When using BERNOULLI sampling, the probability of selecting each record
(between 0 and 1).
When using SYSTEMATIC sampling, the integer part is the interval between
records to sample & the fractional part is the percentage of the
population to sample. When there is no fractional part, it will
select every nth record for the entire population.
When using STRATIFIED sampling, the stratum sample size.
When using CLUSTER sampling, the number of clusters.
When using TIMESERIES sampling, the interval number (treated as hours
by default, e.g., 1 = 1 hour). Use --ts-interval for custom intervals
like "1d" (daily), "1w" (weekly), "1m" (monthly), "1y" (yearly), etc.
sample options:
--seed <number> Random Number Generator (RNG) seed.
--rng <kind> The Random Number Generator (RNG) algorithm to use.
Three RNGs are supported:
* standard: Use the standard RNG.
1.5 GB/s throughput.
* faster: Use faster RNG using the Xoshiro256Plus algorithm.
8 GB/s throughput.
* cryptosecure: Use cryptographically secure HC128 algorithm.
Recommended by eSTREAM (https://www.ecrypt.eu.org/stream/).
2.1 GB/s throughput though slow initialization.
[default: standard]
SAMPLING METHODS:
--bernoulli Use Bernoulli sampling instead of indexed or reservoir sampling.
When this flag is set, <sample-size> must be between
0 and 1 and represents the probability of selecting each record.
--systematic <arg> Use systematic sampling (every nth record as specified by <sample-size>).
If <arg> is "random", the starting point is randomly chosen between 0 & n.
If <arg> is "first", the starting point is the first record.
The sample size must be a whole number. Uses CONSTANT memory - O(1).
--stratified <col> Use stratified sampling. The strata column is specified by <col>.
Can be either a column name or 0-based column index.
The sample size must be a whole number. Uses MEMORY PROPORTIONAL to the
number of strata (s) and samples per stratum (k) - O(s*k).
--weighted <col> Use weighted sampling. The weight column is specified by <col>.
Can be either a column name or 0-based column index.
The column will be parsed as a number. Records with non-number weights
will be skipped.
Uses MEMORY PROPORTIONAL to the sample size (k) - O(k).
--cluster <col> Use cluster sampling. The cluster column is specified by <col>.
Can be either a column name or 0-based column index.
Uses MEMORY PROPORTIONAL to the number of clusters (c) - O(c).
--timeseries <col> Use time-series sampling. The time column is specified by <col>.
Can be either a column name or 0-based column index.
Sorts records by the specified time column and then groups by time intervals
and selects one record per interval.
Supports various date formats (19 formats recognized by qsv-dateparser).
Uses MEMORY PROPORTIONAL to the number of records - O(n).
TIME-SERIES SAMPLING OPTIONS:
--ts-interval <intvl> Time interval for grouping records. Format: <number><unit>
where unit is h (hour), d (day), w (week), m (month), y (year).
Examples: "1h", "1d", "1w", "2d" (every 2 days).
If not specified, <sample-size> is treated as hours.
--ts-start <mode> Starting point for time-series sampling.
Options: "first" (earliest timestamp, default), "last" (most recent timestamp),
"random" (random starting point).
[default: first]
--ts-adaptive <mode> Adaptive sampling mode for time-series data.
Options: "business-hours" (prefer 9am-5pm Mon-Fri),
"weekends" (prefer weekends), "business-days" (prefer weekdays),
"both" (combine business-hours and weekends).
--ts-aggregate <func> Aggregation function to apply within each time interval.
Options: "first", "last", "mean", "sum", "count", "min", "max", "median".
When specified, aggregates all records in each interval instead of selecting a single record.
--ts-input-tz <tz> Timezone for parsing input timestamps. Can be an IANA timezone name or "local" for the local timezone.
[default: UTC]
--ts-prefer-dmy Prefer to parse dates in dmy format. Otherwise, use mdy format.
REMOTE FILE OPTIONS:
--user-agent <agent> Specify custom user agent to use when the input is a URL.
It supports the following variables -
$QSV_VERSION, $QSV_TARGET, $QSV_BIN_NAME, $QSV_KIND and $QSV_COMMAND.
Try to follow the syntax here -
https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/User-Agent
--timeout <secs> Timeout for downloading URLs in seconds. If 0, no timeout is used.
[default: 30]
--max-size <mb> Maximum size of the file to download in MB before sampling.
Will download the entire file if not specified.
If the CSV is partially downloaded, the sample will be taken
only from the downloaded portion.
--force Do not use stats cache, even if its available.
Common options:
-h, --help Display this message
-o, --output <file> Write output to <file> instead of stdout.
-n, --no-headers When set, the first row will be considered as part of
the population to sample from. (When not set, the
first row is the header row and will always appear
in the output.)
-d, --delimiter <arg> The field delimiter for reading/writing CSV data.
Must be a single character. (default: ,)
"#;
use std::{io, str::FromStr};
use chrono::{DateTime, Datelike, Duration, TimeZone, Timelike, Utc, Weekday};
use chrono_tz::Tz;
use foldhash::{HashMap, HashMapExt, HashSet, HashSetExt};
use futures_util::StreamExt;
use qsv_dateparser::parse_with_preference_and_timezone;
use rand::{
Rng, RngExt, SeedableRng,
distr::{Bernoulli, Distribution},
prelude::IndexedRandom,
rngs::StdRng,
};
use rand_hc::Hc128Rng;
use rand_xoshiro::Xoshiro256Plus;
use rayon::prelude::ParallelSliceMut;
use serde::Deserialize;
use strum_macros::EnumString;
use tempfile::NamedTempFile;
use url::Url;
use crate::{
CliResult,
config::{Config, Delimiter},
select::SelectColumns,
util,
util::{SchemaArgs, StatsMode, get_stats_records},
};
#[derive(Deserialize)]
struct Args {
arg_input: Option<String>,
arg_sample_size: f64,
flag_output: Option<String>,
flag_no_headers: bool,
flag_delimiter: Option<Delimiter>,
flag_seed: Option<u64>,
flag_rng: String,
flag_user_agent: Option<String>,
flag_timeout: Option<u16>,
flag_max_size: Option<u64>,
flag_bernoulli: bool,
flag_systematic: Option<String>,
flag_stratified: Option<String>,
flag_weighted: Option<String>,
flag_cluster: Option<String>,
flag_timeseries: Option<String>,
flag_ts_interval: Option<String>,
flag_ts_start: Option<String>,
flag_ts_adaptive: Option<String>,
flag_ts_aggregate: Option<String>,
flag_ts_input_tz: Option<String>,
flag_ts_prefer_dmy: bool,
flag_force: bool,
}
impl Args {
fn get_column_index(
header: &csv::ByteRecord,
column_spec: &str,
purpose: &str,
) -> CliResult<usize> {
if let Ok(idx) = column_spec.parse::<usize>() {
if idx < header.len() {
return Ok(idx);
}
return fail_incorrectusage_clierror!(
"{} column index {} is out of bounds (max: {})",
purpose,
idx,
header.len() - 1
);
}
for (i, field) in header.iter().enumerate() {
if column_spec == String::from_utf8_lossy(field) {
return Ok(i);
}
}
fail_incorrectusage_clierror!("Could not find {} column named '{}'", purpose, column_spec)
}
fn get_strata_column(&self, header: &csv::ByteRecord) -> CliResult<usize> {
match &self.flag_stratified {
Some(col) => Self::get_column_index(header, col, "strata"),
None => {
fail_incorrectusage_clierror!(
"--stratified <col> is required for stratified sampling"
)
},
}
}
fn get_weight_column(&self, header: &csv::ByteRecord) -> CliResult<usize> {
match &self.flag_weighted {
Some(col) => Self::get_column_index(header, col, "weight"),
None => {
fail_incorrectusage_clierror!("--weighted <col> is required for weighted sampling")
},
}
}
fn get_cluster_column(&self, header: &csv::ByteRecord) -> CliResult<usize> {
match &self.flag_cluster {
Some(col) => Self::get_column_index(header, col, "cluster"),
None => {
fail_incorrectusage_clierror!("--cluster <col> is required for cluster sampling")
},
}
}
fn get_timeseries_column(&self, header: &csv::ByteRecord) -> CliResult<usize> {
match &self.flag_timeseries {
Some(col) => Self::get_column_index(header, col, "timeseries"),
None => {
fail_incorrectusage_clierror!(
"--timeseries <col> is required for timeseries sampling"
)
},
}
}
}
#[derive(Debug, EnumString, PartialEq)]
#[strum(ascii_case_insensitive)]
enum RngKind {
Standard,
Faster,
Cryptosecure,
}
#[derive(PartialEq)]
enum SamplingMethod {
Bernoulli,
Systematic,
Stratified,
Weighted,
Cluster,
Timeseries,
Default,
}
trait RngProvider: Sized {
type RngType: Rng + SeedableRng;
fn get_name() -> &'static str;
fn create(seed: Option<u64>) -> Self::RngType {
if let Some(seed) = seed {
Self::RngType::seed_from_u64(seed) } else {
rand::make_rng::<Self::RngType>()
}
}
}
struct StandardRng;
impl RngProvider for StandardRng {
type RngType = StdRng;
fn get_name() -> &'static str {
"standard"
}
}
struct FasterRng;
impl RngProvider for FasterRng {
type RngType = Xoshiro256Plus;
fn get_name() -> &'static str {
"faster"
}
}
struct CryptoRng;
impl RngProvider for CryptoRng {
type RngType = Hc128Rng;
fn get_name() -> &'static str {
"cryptosecure"
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum TSStartMode {
First,
Last,
Random,
}
impl std::str::FromStr for TSStartMode {
type Err = &'static str;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.to_lowercase().as_str() {
"first" => Ok(TSStartMode::First),
"last" => Ok(TSStartMode::Last),
"random" => Ok(TSStartMode::Random),
_ => Err("Time-series start mode must be 'first', 'last' or 'random'"),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
enum AggregationFunction {
First,
Last,
Mean,
Sum,
Count,
Min,
Max,
Median,
}
impl FromStr for AggregationFunction {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.to_lowercase().as_str() {
"first" => Ok(AggregationFunction::First),
"last" => Ok(AggregationFunction::Last),
"mean" => Ok(AggregationFunction::Mean),
"sum" => Ok(AggregationFunction::Sum),
"count" => Ok(AggregationFunction::Count),
"min" => Ok(AggregationFunction::Min),
"max" => Ok(AggregationFunction::Max),
"median" => Ok(AggregationFunction::Median),
_ => Err(format!(
"Invalid aggregation function: {s}. Supported: first, last, mean, sum, count, \
min, max, median"
)),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq)]
enum AdaptiveMode {
BusinessHours,
Weekends,
BusinessDays,
Both,
}
impl FromStr for AdaptiveMode {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.to_lowercase().as_str() {
"business-hours" | "businesshours" => Ok(AdaptiveMode::BusinessHours),
"weekends" => Ok(AdaptiveMode::Weekends),
"business-days" | "businessdays" => Ok(AdaptiveMode::BusinessDays),
"both" => Ok(AdaptiveMode::Both),
_ => Err(format!(
"Invalid adaptive mode: {s}. Supported: business-hours, weekends, business-days, \
both"
)),
}
}
}
fn parse_time_interval(interval_str: &str) -> CliResult<Duration> {
let s = interval_str.trim().to_lowercase();
if s.len() < 2 {
return fail_incorrectusage_clierror!(
"Invalid time interval format: {interval_str}. Expected format: <number><unit> (e.g., \
1h, 1d, 1w, 1m, 1y)"
);
}
let (num_str, unit) = s.split_at(s.len() - 1);
let num: i64 = num_str.parse().map_err(|_| {
format!(
"Invalid time interval number: {num_str}. Expected format: <number><unit> (e.g., 1h, \
1d, 1w, 1m, 1y)"
)
})?;
if num <= 0 {
return fail_incorrectusage_clierror!("Time interval must be positive");
}
let duration = match unit {
"h" => Duration::hours(num),
"d" => Duration::days(num),
"w" => Duration::weeks(num),
"m" => Duration::days(num * 30), "y" => Duration::days(num * 365), _ => {
return fail_incorrectusage_clierror!(
"Invalid time interval unit: {unit}. Supported units: h (hour), d (day), w \
(week), m (month), y (year)"
);
},
};
Ok(duration)
}
fn parse_timestamp(
value: &[u8],
prefer_dmy: bool,
input_tz: Option<&str>,
) -> CliResult<DateTime<Utc>> {
let Ok(value_str) = simdutf8::basic::from_utf8(value) else {
return fail_incorrectusage_clierror!("Time column value is not valid UTF-8");
};
if let Ok(ts_val) = atoi_simd::parse::<i64>(value) {
if let Some(dt) = Utc.timestamp_opt(ts_val, 0).single() {
return Ok(dt);
}
if let Some(dt) = Utc.timestamp_millis_opt(ts_val).single() {
return Ok(dt);
}
}
let tz: Tz = if let Some(tz_str) = input_tz {
if tz_str.eq_ignore_ascii_case("local") {
if let Ok(tz_name) = iana_time_zone::get_timezone() {
tz_name.parse::<Tz>().unwrap_or(chrono_tz::UTC)
} else {
chrono_tz::UTC
}
} else {
tz_str.parse::<Tz>().unwrap_or(chrono_tz::UTC)
}
} else {
chrono_tz::UTC
};
parse_with_preference_and_timezone(value_str, prefer_dmy, &tz)
.map_err(|e| format!("Failed to parse timestamp '{value_str}': {e}").into())
}
fn is_business_hours(dt: &DateTime<Utc>) -> bool {
let hour = dt.hour();
(9..=17).contains(&hour)
}
fn is_weekend(dt: &DateTime<Utc>) -> bool {
matches!(dt.weekday(), Weekday::Sat | Weekday::Sun)
}
fn is_business_day(dt: &DateTime<Utc>) -> bool {
!is_weekend(dt)
}
fn check_stats_cache(
args: &Args,
method: &SamplingMethod,
) -> CliResult<(Option<f64>, Option<u64>)> {
if args.flag_force {
return Ok((None, None));
}
let schema_args = SchemaArgs {
arg_input: args.arg_input.clone(),
flag_no_headers: args.flag_no_headers,
flag_delimiter: args.flag_delimiter,
flag_jobs: None,
flag_polars: false,
flag_memcheck: false,
flag_force: args.flag_force,
flag_prefer_dmy: false,
flag_dates_whitelist: String::new(),
flag_enum_threshold: 0,
flag_ignore_case: false,
flag_strict_dates: false,
flag_strict_formats: false,
flag_pattern_columns: SelectColumns::parse("")?,
flag_stdout: false,
flag_output: None,
};
match get_stats_records(&schema_args, StatsMode::Frequency) {
Ok((csv_fields, stats)) => {
let mut max_weight = None;
let mut cardinality = None;
match method {
SamplingMethod::Weighted => {
if let Some(weight_col) = &args.flag_weighted {
let idx = if weight_col.chars().all(char::is_numeric) {
weight_col.parse::<usize>().ok()
} else {
csv_fields
.iter()
.position(|field| field == weight_col.as_bytes())
};
if let Some(idx) = idx
&& let Some(col_stats) = stats.get(idx)
{
let min_weight = col_stats
.min
.clone()
.unwrap_or_default()
.parse::<f64>()
.unwrap_or_default();
if min_weight < 0.0 {
return fail_incorrectusage_clierror!(
"Weights must be non-negative. Lowest weight: {min_weight}"
);
}
max_weight = col_stats.max.clone().unwrap().parse::<f64>().ok();
}
}
},
SamplingMethod::Cluster => {
if let Some(cluster_col) = &args.flag_cluster {
let idx = if cluster_col.chars().all(char::is_numeric) {
cluster_col.parse::<usize>().ok()
} else {
csv_fields
.iter()
.position(|field| field == cluster_col.as_bytes())
};
if let Some(idx) = idx {
cardinality = stats.get(idx).map(|col_stats| col_stats.cardinality);
}
}
},
_ => {},
}
Ok((max_weight, cardinality))
},
_ => Ok((None, None)),
}
}
#[allow(clippy::future_not_send)]
async fn stream_bernoulli_sampling(uri: &str, args: &Args, rng_kind: &RngKind) -> CliResult<()> {
let default_delim = match std::env::var("QSV_DEFAULT_DELIMITER") {
Ok(delim) => Delimiter::decode_delimiter(&delim).unwrap().as_byte(),
_ => b',',
};
let mut wtr = Config::new(args.flag_output.as_ref())
.delimiter(args.flag_delimiter)
.writer()?;
let client = util::create_reqwest_async_client(
args.flag_user_agent.clone(),
util::timeout_secs(args.flag_timeout.unwrap_or(30)).map(|t| t as u16)?,
Some(uri.to_string()),
)?;
let response = client.get(uri).send().await?;
let mut stream = response.bytes_stream();
if !args.flag_no_headers {
let mut header_bytes = Vec::new();
while let Some(chunk) = stream.next().await {
let chunk = chunk?;
header_bytes.extend_from_slice(&chunk);
let mut rdr = csv::ReaderBuilder::new()
.has_headers(true)
.delimiter(default_delim)
.from_reader(&header_bytes[..]);
if let Ok(headers) = rdr.headers() {
wtr.write_record(headers)?;
break;
}
}
}
let mut std_rng = StandardRng::create(args.flag_seed);
let mut faster_rng = FasterRng::create(args.flag_seed);
let mut crypto_rng = CryptoRng::create(args.flag_seed);
let mut record = csv::ByteRecord::new();
let mut buffer = Vec::new();
while let Some(chunk) = stream.next().await {
let chunk = chunk?;
buffer.extend_from_slice(&chunk);
if let Some(pos) = buffer.iter().rposition(|&b| b == b'\n') {
let mut csv_reader = csv::ReaderBuilder::new()
.has_headers(args.flag_no_headers)
.delimiter(default_delim)
.from_reader(&buffer[..=pos]);
while matches!(csv_reader.read_byte_record(&mut record), Ok(true)) {
match rng_kind {
RngKind::Standard => {
if std_rng.random_bool(args.arg_sample_size) {
wtr.write_byte_record(&record)?;
}
},
RngKind::Faster => {
if faster_rng.random_bool(args.arg_sample_size) {
wtr.write_byte_record(&record)?;
}
},
RngKind::Cryptosecure => {
if crypto_rng.random_bool(args.arg_sample_size) {
wtr.write_byte_record(&record)?;
}
},
}
}
buffer.drain(..=pos);
}
}
if !buffer.is_empty() {
let mut csv_reader = csv::ReaderBuilder::new()
.has_headers(args.flag_no_headers)
.delimiter(default_delim)
.from_reader(&buffer[..]);
while matches!(csv_reader.read_byte_record(&mut record), Ok(true)) {
match rng_kind {
RngKind::Standard => {
if std_rng.random_bool(args.arg_sample_size) {
wtr.write_byte_record(&record)?;
}
},
RngKind::Faster => {
if faster_rng.random_bool(args.arg_sample_size) {
wtr.write_byte_record(&record)?;
}
},
RngKind::Cryptosecure => {
if crypto_rng.random_bool(args.arg_sample_size) {
wtr.write_byte_record(&record)?;
}
},
}
}
}
Ok(wtr.flush()?)
}
pub fn run(argv: &[&str]) -> CliResult<()> {
let mut args: Args = util::get_args(USAGE, argv)?;
if args.arg_sample_size.is_sign_negative() {
return fail_incorrectusage_clierror!("Sample size cannot be negative.");
}
let methods = [
args.flag_bernoulli,
args.flag_systematic.is_some(),
args.flag_stratified.is_some(),
args.flag_weighted.is_some(),
args.flag_cluster.is_some(),
args.flag_timeseries.is_some(),
];
if methods.iter().filter(|&&x| x).count() > 1 {
return fail_incorrectusage_clierror!("Only one sampling method can be specified");
}
let Ok(rng_kind) = RngKind::from_str(&args.flag_rng) else {
return fail_incorrectusage_clierror!(
"Invalid RNG algorithm `{}`. Supported RNGs are: standard, faster, cryptosecure.",
args.flag_rng
);
};
let sampling_method = match (
args.flag_bernoulli,
args.flag_systematic.is_some(),
args.flag_stratified.is_some(),
args.flag_weighted.is_some(),
args.flag_cluster.is_some(),
args.flag_timeseries.is_some(),
) {
(true, _, _, _, _, _) => SamplingMethod::Bernoulli,
(_, true, _, _, _, _) => SamplingMethod::Systematic,
(_, _, true, _, _, _) => SamplingMethod::Stratified,
(_, _, _, true, _, _) => SamplingMethod::Weighted,
(_, _, _, _, true, _) => SamplingMethod::Cluster,
(_, _, _, _, _, true) => SamplingMethod::Timeseries,
(false, false, false, false, false, false) => SamplingMethod::Default,
};
let temp_download = NamedTempFile::new()?;
args.arg_input = match args.arg_input {
Some(ref uri) if Url::parse(uri).is_ok() && uri.starts_with("http") => {
if sampling_method == SamplingMethod::Bernoulli {
log::info!("Streaming Bernoulli sampling remote file");
let rt = tokio::runtime::Runtime::new()?;
rt.block_on(stream_bernoulli_sampling(uri, &args, &rng_kind))?;
return Ok(());
}
let max_size_bytes = args.flag_max_size.map(|mb| mb * 1024 * 1024);
let future = util::download_file(
uri,
temp_download.path().to_path_buf(),
false,
Some(util::set_user_agent(args.flag_user_agent.clone())?),
args.flag_timeout,
max_size_bytes,
);
tokio::runtime::Runtime::new()?.block_on(future)?;
Some(temp_download.path().to_str().unwrap().to_string())
},
Some(uri) => Some(uri), None => None,
};
let rconfig = Config::new(args.arg_input.as_ref())
.delimiter(args.flag_delimiter)
.no_headers_flag(args.flag_no_headers)
.flexible(true)
.skip_format_check(true);
let mut rdr = rconfig.reader()?;
let mut wtr = Config::new(args.flag_output.as_ref())
.delimiter(args.flag_delimiter)
.writer()?;
rconfig.write_headers(&mut rdr, &mut wtr)?;
let mut sample_size = args.arg_sample_size;
match sampling_method {
SamplingMethod::Bernoulli => {
if args.arg_sample_size >= 1.0 || args.arg_sample_size <= 0.0 {
return fail_incorrectusage_clierror!(
"Bernoulli sampling requires a probability between 0 and 1"
);
}
sample_bernoulli(
&mut rdr,
&mut wtr,
args.arg_sample_size,
args.flag_seed,
&rng_kind,
)?;
},
SamplingMethod::Systematic => {
let starting_point = match args.flag_systematic.as_deref().map(str::to_lowercase) {
Some(arg) if arg == "random" || arg == "first" => arg,
Some(_) => {
return fail_incorrectusage_clierror!(
"Systematic sampling starting point must be either 'random' or 'first'"
);
},
None => String::from("random"),
};
let row_count: u64 = if let Ok(rc) = util::count_rows(&rconfig) {
rc
} else {
return fail!("Systematic sampling requires rowcount.");
};
sample_systematic(
&mut rdr,
&mut wtr,
args.arg_sample_size,
row_count,
&starting_point,
args.flag_seed,
&rng_kind,
)?;
},
SamplingMethod::Stratified => {
let strata_column = args.get_strata_column(&rdr.byte_headers()?.clone())?;
sample_stratified(
&mut rdr,
&mut wtr,
strata_column,
args.arg_sample_size as usize,
args.flag_seed,
&rng_kind,
)?;
},
SamplingMethod::Weighted => {
let weight_column = args.get_weight_column(&rdr.byte_headers()?.clone())?;
let (max_weight, _) = check_stats_cache(&args, &SamplingMethod::Weighted)?;
#[allow(clippy::cast_precision_loss)]
let sample_size = if args.arg_sample_size < 1.0 {
let row_count: u64 = if let Ok(rc) = util::count_rows(&rconfig) {
rc
} else {
return fail!("Weighted fractional sampling requires rowcount.");
};
(row_count as f64 * args.arg_sample_size).round() as usize
} else {
args.arg_sample_size as usize
};
sample_weighted(
&rconfig,
&mut rdr,
&mut wtr,
weight_column,
max_weight,
sample_size,
args.flag_seed,
&rng_kind,
)?;
},
SamplingMethod::Cluster => {
let cluster_column = args.get_cluster_column(&rdr.byte_headers()?.clone())?;
let (_, cardinality) = check_stats_cache(&args, &SamplingMethod::Cluster)?;
sample_cluster(
&rconfig,
&mut rdr,
&mut wtr,
cluster_column,
cardinality,
args.arg_sample_size as usize,
args.flag_seed,
&rng_kind,
)?;
},
SamplingMethod::Timeseries => {
let time_column = args.get_timeseries_column(&rdr.byte_headers()?.clone())?;
let interval_str = if let Some(interval) = &args.flag_ts_interval {
interval.clone()
} else if args.arg_sample_size.fract() == 0.0 && args.arg_sample_size > 0.0 {
format!("{}h", args.arg_sample_size as i64)
} else {
return fail_incorrectusage_clierror!(
"Time-series sampling requires either --ts-interval (e.g., '1h', '1d', '1w', \
'1m', '1y') or a positive whole number for <sample-size> (treated as hours)"
);
};
let start_mode = match args
.flag_ts_start
.as_deref()
.unwrap_or("first")
.parse::<TSStartMode>()
{
Ok(mode) => mode,
Err(msg) => return fail_incorrectusage_clierror!("{msg}"),
};
let adaptive_mode = if let Some(adaptive_str) = &args.flag_ts_adaptive {
Some(
AdaptiveMode::from_str(adaptive_str)
.map_err(|e| format!("Invalid adaptive mode: {e}"))?,
)
} else {
None
};
let aggregate_func = if let Some(agg_str) = &args.flag_ts_aggregate {
Some(
AggregationFunction::from_str(agg_str)
.map_err(|e| format!("Invalid aggregation function: {e}"))?,
)
} else {
None
};
let prefer_dmy = args.flag_ts_prefer_dmy || rconfig.get_dmy_preference();
let input_tz = match args.flag_ts_input_tz.as_deref() {
Some(tz_str) => {
if tz_str.eq_ignore_ascii_case("local") {
if let Ok(tz_name) = iana_time_zone::get_timezone() {
if tz_name.parse::<chrono_tz::Tz>().is_ok() {
Some(tz_str)
} else {
wwarn!(
"Invalid local timezone from iana_time_zone, falling back to \
UTC."
);
None
}
} else {
wwarn!("Could not determine local timezone, falling back to UTC.");
None
}
} else if tz_str.parse::<chrono_tz::Tz>().is_ok() {
Some(tz_str)
} else {
wwarn!("Invalid timezone '{tz_str}', falling back to UTC.");
None
}
},
None => None,
};
sample_timeseries(
&rconfig,
&mut rdr,
&mut wtr,
time_column,
&interval_str,
start_mode,
adaptive_mode,
aggregate_func,
prefer_dmy,
input_tz,
args.flag_seed,
&rng_kind,
)?;
},
SamplingMethod::Default => {
if let Some(mut idx) = rconfig.indexed()? {
#[allow(clippy::cast_precision_loss)]
if sample_size < 1.0 {
sample_size *= idx.count() as f64;
}
let sample_count = sample_size as usize;
let total_count = idx.count().try_into().unwrap();
match rng_kind {
RngKind::Standard => {
log::info!("doing standard INDEXED sampling...");
let mut rng = StandardRng::create(args.flag_seed);
sample_indices(&mut rng, total_count, sample_count, |i| {
idx.seek(i as u64)?;
Ok(wtr.write_byte_record(&idx.byte_records().next().unwrap()?)?)
})?;
},
RngKind::Faster => {
log::info!("doing --faster INDEXED sampling...");
let mut rng = FasterRng::create(args.flag_seed);
sample_indices(&mut rng, total_count, sample_count, |i| {
idx.seek(i as u64)?;
Ok(wtr.write_byte_record(&idx.byte_records().next().unwrap()?)?)
})?;
},
RngKind::Cryptosecure => {
log::info!("doing --cryptosecure INDEXED sampling...");
let mut rng = CryptoRng::create(args.flag_seed);
sample_indices(&mut rng, total_count, sample_count, |i| {
idx.seek(i as u64)?;
Ok(wtr.write_byte_record(&idx.byte_records().next().unwrap()?)?)
})?;
},
}
} else {
#[allow(clippy::cast_precision_loss)]
let sample_size = if args.arg_sample_size < 1.0 {
let row_count: u64 = if let Ok(rc) = util::count_rows(&rconfig) {
rc
} else {
return fail!("Fractional sampling requires rowcount.");
};
(row_count as f64 * args.arg_sample_size).round() as u64
} else {
args.arg_sample_size as u64
};
sample_reservoir(&mut rdr, &mut wtr, sample_size, args.flag_seed, &rng_kind)?;
}
},
}
Ok(wtr.flush()?)
}
fn sample_reservoir<R: io::Read, W: io::Write>(
rdr: &mut csv::Reader<R>,
wtr: &mut csv::Writer<W>,
sample_size: u64,
seed: Option<u64>,
rng_kind: &RngKind,
) -> CliResult<()> {
let mut reservoir = Vec::with_capacity(sample_size as usize);
let mut records = rdr.byte_records().enumerate();
for (_, row) in records.by_ref().take(sample_size as usize) {
reservoir.push(row?);
}
match rng_kind {
RngKind::Standard => {
do_reservoir_sampling::<StandardRng>(&mut records, &mut reservoir, sample_size, seed)
},
RngKind::Faster => {
do_reservoir_sampling::<FasterRng>(&mut records, &mut reservoir, sample_size, seed)
},
RngKind::Cryptosecure => {
do_reservoir_sampling::<CryptoRng>(&mut records, &mut reservoir, sample_size, seed)
},
}?;
for record in reservoir {
wtr.write_byte_record(&record)?;
}
Ok(())
}
fn do_reservoir_sampling<T: RngProvider>(
records: &mut impl Iterator<Item = (usize, Result<csv::ByteRecord, csv::Error>)>,
reservoir: &mut [csv::ByteRecord],
sample_size: u64,
seed: Option<u64>,
) -> CliResult<()> {
log::info!("doing {} RESERVOIR sampling...", T::get_name());
let mut rng = T::create(seed);
let mut random_idx: usize;
for (i, row) in records {
random_idx = rng.random_range(0..=i);
if random_idx < sample_size as usize {
unsafe {
*reservoir.get_unchecked_mut(random_idx) = row?;
}
}
}
Ok(())
}
fn sample_bernoulli<R: io::Read, W: io::Write>(
rdr: &mut csv::Reader<R>,
wtr: &mut csv::Writer<W>,
probability: f64,
seed: Option<u64>,
rng_kind: &RngKind,
) -> CliResult<()> {
let mut records = rdr.byte_records();
match rng_kind {
RngKind::Standard => {
do_bernoulli_sampling::<StandardRng>(&mut records, wtr, probability, seed)
},
RngKind::Faster => do_bernoulli_sampling::<FasterRng>(&mut records, wtr, probability, seed),
RngKind::Cryptosecure => {
do_bernoulli_sampling::<CryptoRng>(&mut records, wtr, probability, seed)
},
}
}
fn do_bernoulli_sampling<T: RngProvider>(
records: &mut impl Iterator<Item = Result<csv::ByteRecord, csv::Error>>,
wtr: &mut csv::Writer<impl io::Write>,
probability: f64,
seed: Option<u64>,
) -> CliResult<()> {
log::info!("doing {} BERNOULLI sampling...", T::get_name());
let mut rng = T::create(seed);
let dist =
Bernoulli::new(probability).map_err(|_| "probability must be between 0.0 and 1.0")?;
for row in records {
if dist.sample(&mut rng) {
wtr.write_byte_record(&row?)?;
}
}
Ok(())
}
fn sample_indices<F>(
rng: &mut impl Rng,
total_count: usize,
sample_count: usize,
mut process_index: F,
) -> CliResult<()>
where
F: FnMut(usize) -> CliResult<()>,
{
if sample_count > total_count {
return fail!("Sample size cannot be larger than population size");
}
let mut selected = Vec::with_capacity(sample_count);
for i in 0..sample_count {
selected.push(i);
}
for i in sample_count..total_count {
let j = rng.random_range(0..=i);
if j < sample_count {
unsafe { *selected.get_unchecked_mut(j) = i };
}
}
selected.par_sort_unstable();
for idx in selected {
process_index(idx)?;
}
Ok(())
}
fn sample_systematic<R: io::Read, W: io::Write>(
rdr: &mut csv::Reader<R>,
wtr: &mut csv::Writer<W>,
sample_size: f64,
row_count: u64,
starting_point: &str,
seed: Option<u64>,
rng_kind: &RngKind,
) -> CliResult<()> {
if sample_size <= 0.0 {
return fail_incorrectusage_clierror!("Sample size must be positive");
}
let interval = sample_size.trunc() as usize;
let percentage = sample_size.fract();
if interval == 0 {
return fail_incorrectusage_clierror!("Interval must be at least 1");
}
#[allow(clippy::cast_precision_loss)]
let target_count = if percentage > 0.0 {
((row_count as f64) * percentage).round() as u64
} else {
row_count
};
let start = if starting_point == "random" {
match rng_kind {
RngKind::Standard => {
let mut rng = StandardRng::create(seed);
rng.random_range(0..interval)
},
RngKind::Faster => {
let mut rng = FasterRng::create(seed);
rng.random_range(0..interval)
},
RngKind::Cryptosecure => {
let mut rng = CryptoRng::create(seed);
rng.random_range(0..interval)
},
}
} else {
0 };
let mut selected_count = 0;
for (i, record) in rdr.byte_records().enumerate().skip(start) {
if i.is_multiple_of(interval) && selected_count < target_count {
wtr.write_byte_record(&record?)?;
selected_count += 1;
}
}
Ok(())
}
fn sample_stratified<R: io::Read, W: io::Write>(
rdr: &mut csv::Reader<R>,
wtr: &mut csv::Writer<W>,
strata_column: usize,
samples_per_stratum: usize,
seed: Option<u64>,
rng_kind: &RngKind,
) -> CliResult<()> {
const ESTIMATED_STRATA_COUNT: usize = 100;
let mut strata_counts: HashMap<Vec<u8>, usize> = HashMap::with_capacity(ESTIMATED_STRATA_COUNT);
let mut records = Vec::with_capacity(ESTIMATED_STRATA_COUNT * samples_per_stratum);
let mut curr_record;
for record in rdr.byte_records() {
curr_record = record?;
let stratum = curr_record
.get(strata_column)
.ok_or_else(|| format!("Strata column index {strata_column} out of bounds"))?
.to_vec();
*strata_counts.entry(stratum.clone()).or_default() += 1;
records.push(curr_record);
}
let strata_count = strata_counts.len();
if strata_count == 0 {
return fail_incorrectusage_clierror!("No valid strata found in the data");
}
let mut reservoirs: HashMap<Vec<u8>, Vec<csv::ByteRecord>> =
HashMap::with_capacity(strata_count);
for stratum in strata_counts.keys() {
reservoirs.insert(stratum.clone(), Vec::with_capacity(samples_per_stratum));
}
match rng_kind {
RngKind::Standard => {
let mut rng = StandardRng::create(seed);
do_stratified_sampling(
records.into_iter(),
&mut reservoirs,
strata_column,
samples_per_stratum,
&mut rng,
)?;
},
RngKind::Faster => {
let mut rng = FasterRng::create(seed);
do_stratified_sampling(
records.into_iter(),
&mut reservoirs,
strata_column,
samples_per_stratum,
&mut rng,
)?;
},
RngKind::Cryptosecure => {
let mut rng = CryptoRng::create(seed);
do_stratified_sampling(
records.into_iter(),
&mut reservoirs,
strata_column,
samples_per_stratum,
&mut rng,
)?;
},
}
let mut strata: Vec<_> = reservoirs.keys().collect();
strata.par_sort_unstable();
for stratum in strata {
if let Some(records) = reservoirs.get(stratum) {
for record in records {
wtr.write_byte_record(record)?;
}
}
}
Ok(())
}
fn do_stratified_sampling<T: Rng + ?Sized>(
records: impl Iterator<Item = csv::ByteRecord>,
reservoirs: &mut HashMap<Vec<u8>, Vec<csv::ByteRecord>>,
strata_column: usize,
samples_per_stratum: usize,
rng: &mut T,
) -> CliResult<()> {
let mut records_seen: HashMap<Vec<u8>, usize> = HashMap::with_capacity(reservoirs.len());
for record in records {
let stratum = record
.get(strata_column)
.ok_or_else(|| format!("Strata column index {strata_column} out of bounds"))?
.to_vec();
let seen = records_seen.entry(stratum.clone()).or_default();
if let Some(reservoir) = reservoirs.get_mut(&stratum) {
if reservoir.len() < samples_per_stratum {
reservoir.push(record);
} else {
let j = rng.random_range(0..=*seen);
if j < samples_per_stratum {
unsafe { *reservoir.get_unchecked_mut(j) = record };
}
}
*seen += 1;
}
}
Ok(())
}
fn sample_weighted<R: io::Read, W: io::Write>(
rconfig: &Config,
rdr: &mut csv::Reader<R>,
wtr: &mut csv::Writer<W>,
weight_column: usize,
max_weight_stats: Option<f64>,
sample_size: usize,
seed: Option<u64>,
rng_kind: &RngKind,
) -> CliResult<()> {
let max_weight = if let Some(wt) = max_weight_stats {
wt
} else {
let mut max_weight_scan = 0.0f64;
let mut curr_record;
for record in rdr.byte_records() {
curr_record = record?;
let weight: f64 = fast_float2::parse(
curr_record
.get(weight_column)
.ok_or_else(|| format!("Weight column index {weight_column} out of bounds"))?,
)
.unwrap_or(0.0);
if weight < 0.0 {
return fail_incorrectusage_clierror!("Weights must be non-negative: ({weight})");
}
max_weight_scan = max_weight_scan.max(weight);
}
max_weight_scan
};
if max_weight == 0.0 {
return fail_incorrectusage_clierror!("All weights are zero");
}
let mut rdr2 = rconfig.reader()?;
match rng_kind {
RngKind::Standard => {
log::info!("doing standard WEIGHTED sampling...");
let mut rng = StandardRng::create(seed);
do_weighted_sampling(
&mut rdr2.byte_records(),
wtr,
weight_column,
sample_size,
max_weight,
&mut rng,
)?;
},
RngKind::Faster => {
log::info!("doing --faster WEIGHTED sampling...");
let mut rng = FasterRng::create(seed);
do_weighted_sampling(
&mut rdr2.byte_records(),
wtr,
weight_column,
sample_size,
max_weight,
&mut rng,
)?;
},
RngKind::Cryptosecure => {
log::info!("doing --cryptosecure WEIGHTED sampling...");
let mut rng = CryptoRng::create(seed);
do_weighted_sampling(
&mut rdr2.byte_records(),
wtr,
weight_column,
sample_size,
max_weight,
&mut rng,
)?;
},
}
Ok(())
}
fn do_weighted_sampling<T: Rng + ?Sized>(
records: &mut impl Iterator<Item = Result<csv::ByteRecord, csv::Error>>,
wtr: &mut csv::Writer<impl io::Write>,
weight_column: usize,
sample_size: usize,
max_weight: f64,
rng: &mut T,
) -> CliResult<()> {
let mut selected = HashSet::with_capacity(sample_size);
let mut attempts = 0;
let max_attempts = sample_size * 100; let mut curr_record;
let mut selected_len = 0;
let mut records_exhausted = false;
while selected_len < sample_size && attempts < max_attempts && !records_exhausted {
let mut any_records = false;
for (i, record) in records.enumerate() {
any_records = true;
if selected_len >= sample_size {
break;
}
curr_record = record?;
let weight: f64 = fast_float2::parse(
curr_record
.get(weight_column)
.ok_or_else(|| format!("Weight column index {weight_column} out of bounds"))?,
)
.unwrap_or(0.0);
if weight < 0.0 {
return fail_incorrectusage_clierror!("Weights must be non-negative: ({weight})");
}
let include_flag = if weight == 0.0 {
false
} else {
rng.random::<f64>() <= (weight / max_weight)
};
if include_flag && !selected.contains(&i) {
selected.insert(i);
selected_len += 1;
wtr.write_byte_record(&curr_record)?;
}
attempts += 1;
if attempts >= max_attempts {
break;
}
}
records_exhausted = !any_records;
}
if selected_len < sample_size {
wwarn!("Could only sample {selected_len} records out of requested {sample_size}");
}
Ok(())
}
fn aggregate_numeric_values(values: &[f64], func: AggregationFunction) -> f64 {
if values.is_empty() {
return 0.0;
}
#[allow(clippy::cast_precision_loss)]
match func {
AggregationFunction::First => *values.first().unwrap_or(&0.0),
AggregationFunction::Last => *values.last().unwrap_or(&0.0),
AggregationFunction::Mean => {
let sum: f64 = values.iter().sum();
sum / values.len() as f64
},
AggregationFunction::Sum => values.iter().sum(),
AggregationFunction::Count => values.len() as f64,
AggregationFunction::Min => values.iter().copied().fold(f64::INFINITY, f64::min),
AggregationFunction::Max => values.iter().copied().fold(f64::NEG_INFINITY, f64::max),
AggregationFunction::Median => {
let mut sorted = values.to_vec();
sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
let mid = sorted.len() / 2;
if sorted.len().is_multiple_of(2) {
f64::midpoint(sorted[mid - 1], sorted[mid])
} else {
sorted[mid]
}
},
}
}
fn aggregate_records(
records: &[csv::ByteRecord],
headers: &csv::ByteRecord,
func: AggregationFunction,
) -> CliResult<csv::ByteRecord> {
if records.is_empty() {
return fail_incorrectusage_clierror!("Cannot aggregate empty record set");
}
let mut result_fields = Vec::with_capacity(headers.len());
for col_idx in 0..headers.len() {
let mut numeric_values = Vec::new();
let mut all_numeric = true;
for record in records {
if let Some(field) = record.get(col_idx) {
if let Ok(num) = fast_float2::parse::<f64, &[u8]>(field) {
numeric_values.push(num);
} else {
all_numeric = false;
break;
}
} else {
all_numeric = false;
break;
}
}
if all_numeric && !numeric_values.is_empty() {
let aggregated = aggregate_numeric_values(&numeric_values, func);
result_fields.push(aggregated.to_string().into_bytes());
} else {
let value = match func {
AggregationFunction::First
| AggregationFunction::Min
| AggregationFunction::Mean
| AggregationFunction::Sum
| AggregationFunction::Count => records[0].get(col_idx).unwrap_or(b"").to_vec(),
AggregationFunction::Last
| AggregationFunction::Max
| AggregationFunction::Median => records[records.len() - 1]
.get(col_idx)
.unwrap_or(b"")
.to_vec(),
};
result_fields.push(value);
}
}
Ok(csv::ByteRecord::from(result_fields))
}
fn sample_timeseries<R: io::Read, W: io::Write>(
_rconfig: &Config,
rdr: &mut csv::Reader<R>,
wtr: &mut csv::Writer<W>,
time_column: usize,
interval_str: &str,
start_mode: TSStartMode,
adaptive_mode: Option<AdaptiveMode>,
aggregate_func: Option<AggregationFunction>,
prefer_dmy: bool,
input_tz: Option<&str>,
seed: Option<u64>,
rng_kind: &RngKind,
) -> CliResult<()> {
let interval = parse_time_interval(interval_str)?;
let headers = rdr.byte_headers()?.clone();
let mut records_with_times: Vec<(DateTime<Utc>, csv::ByteRecord)> = Vec::new();
for record_result in rdr.byte_records() {
let record = record_result?;
if let Some(time_field) = record.get(time_column) {
match parse_timestamp(time_field, prefer_dmy, input_tz) {
Ok(dt) => {
records_with_times.push((dt, record));
},
Err(e) => {
log::warn!("Skipping record with invalid timestamp: {e}");
},
}
} else {
log::warn!("Skipping record with missing time column");
}
}
if records_with_times.is_empty() {
return fail_incorrectusage_clierror!("No valid timestamps found in time column");
}
records_with_times.par_sort_unstable_by(|a, b| a.0.cmp(&b.0));
let start_time = match start_mode {
TSStartMode::Last => {
records_with_times.last().unwrap().0
},
TSStartMode::Random => {
let earliest = records_with_times.first().unwrap().0;
let latest = records_with_times.last().unwrap().0;
let range_secs = (latest - earliest).num_seconds();
if range_secs > 0 {
let random_offset = match rng_kind {
RngKind::Standard => {
let mut rng = StandardRng::create(seed);
rng.random_range(0..=range_secs)
},
RngKind::Faster => {
let mut rng = FasterRng::create(seed);
rng.random_range(0..=range_secs)
},
RngKind::Cryptosecure => {
let mut rng = CryptoRng::create(seed);
rng.random_range(0..=range_secs)
},
};
earliest + Duration::seconds(random_offset)
} else {
earliest
}
},
TSStartMode::First => {
records_with_times.first().unwrap().0
},
};
let mut interval_groups: HashMap<i64, Vec<(DateTime<Utc>, csv::ByteRecord)>> = HashMap::new();
for (dt, record) in records_with_times {
let elapsed = dt - start_time;
let interval_num = if elapsed.num_seconds() >= 0 {
elapsed.num_seconds() / interval.num_seconds()
} else {
(elapsed.num_seconds() - interval.num_seconds() + 1) / interval.num_seconds()
};
interval_groups
.entry(interval_num)
.or_default()
.push((dt, record));
}
let mut interval_keys: Vec<i64> = interval_groups.keys().copied().collect();
interval_keys.sort_unstable();
for interval_key in interval_keys {
let group = interval_groups.get(&interval_key).unwrap();
if let Some(agg_func) = aggregate_func {
let records_only: Vec<csv::ByteRecord> = group.iter().map(|(_, r)| r.clone()).collect();
let aggregated = aggregate_records(&records_only, &headers, agg_func)?;
wtr.write_byte_record(&aggregated)?;
} else {
let selected = match adaptive_mode {
Some(AdaptiveMode::BusinessHours) => {
group
.iter()
.find(|(dt, _)| is_business_hours(dt) && is_business_day(dt))
.or_else(|| group.first())
},
Some(AdaptiveMode::Weekends) => {
group
.iter()
.find(|(dt, _)| is_weekend(dt))
.or_else(|| group.first())
},
Some(AdaptiveMode::BusinessDays) => {
group
.iter()
.find(|(dt, _)| is_business_day(dt))
.or_else(|| group.first())
},
Some(AdaptiveMode::Both) => {
group
.iter()
.find(|(dt, _)| {
(is_business_hours(dt) && is_business_day(dt)) || is_weekend(dt)
})
.or_else(|| group.first())
},
None => {
group.first()
},
};
if let Some((_, record)) = selected {
wtr.write_byte_record(record)?;
}
}
}
Ok(())
}
fn sample_cluster<R: io::Read, W: io::Write>(
rconfig: &Config,
rdr: &mut csv::Reader<R>,
wtr: &mut csv::Writer<W>,
cluster_column: usize,
cluster_cardinality: Option<u64>,
requested_clusters: usize,
seed: Option<u64>,
rng_kind: &RngKind,
) -> CliResult<()> {
const ESTIMATED_CLUSTER_COUNT: usize = 100;
let cluster_count = if let Some(cardinality) = cluster_cardinality {
if requested_clusters > cardinality as usize {
return fail_incorrectusage_clierror!(
"Requested sample size ({requested_clusters}) exceeds number of clusters \
({cardinality})",
);
}
requested_clusters
} else {
ESTIMATED_CLUSTER_COUNT
};
let mut unique_clusters: HashSet<Vec<u8>> = HashSet::with_capacity(cluster_count);
let mut all_clusters: Vec<Vec<u8>> = Vec::with_capacity(cluster_count);
let mut curr_record;
for record in rdr.byte_records() {
curr_record = record?;
let cluster = curr_record
.get(cluster_column)
.ok_or_else(|| format!("Cluster column index {cluster_column} out of bounds"))?
.to_vec();
if unique_clusters.insert(cluster.clone()) {
all_clusters.push(cluster);
}
}
if unique_clusters.is_empty() {
return fail_incorrectusage_clierror!("No valid clusters found in the data");
}
let selected_clusters: HashSet<Vec<u8>> = match rng_kind {
RngKind::Standard => {
let mut rng = StandardRng::create(seed);
all_clusters
.sample(&mut rng, requested_clusters.min(all_clusters.len()))
.cloned()
.collect()
},
RngKind::Faster => {
let mut rng = FasterRng::create(seed);
all_clusters
.sample(&mut rng, requested_clusters.min(all_clusters.len()))
.cloned()
.collect()
},
RngKind::Cryptosecure => {
let mut rng = CryptoRng::create(seed);
all_clusters
.sample(&mut rng, requested_clusters.min(all_clusters.len()))
.cloned()
.collect()
},
};
let mut rdr2 = rconfig.reader()?;
let mut curr_record;
for record in rdr2.byte_records() {
curr_record = record?;
let cluster = curr_record
.get(cluster_column)
.ok_or_else(|| format!("Cluster column index {cluster_column} out of bounds"))?
.to_vec();
if selected_clusters.contains(&cluster) {
wtr.write_byte_record(&curr_record)?;
}
}
Ok(())
}