Function main

benchmarks/src/bin/gen_wide_data.rs:252–351

()

}

fn main() -> Result<()> {
    let cli = Cli::parse();
    if cli.width_factor == 0 {
        return Err(exec_datafusion_err!("--width-factor must be >= 1"));
    }
    if cli.num_files == 0 {
        return Err(exec_datafusion_err!("--num-files must be >= 1"));
    }
    if cli.rows_per_file == 0 {
        return Err(exec_datafusion_err!("--rows-per-file must be >= 1"));
    }
    if cli.batch_size == 0 {
        return Err(exec_datafusion_err!("--batch-size must be >= 1"));
    }

    let raw_schema = base_schema();
    let wide_schema = widen_schema(&raw_schema, cli.width_factor);
    let cap_rows = cli.rows_per_file.saturating_mul(cli.num_files);

    println!(
        "Synthesizing: {} base cols × {} factor = {} cols × {} files of {} rows each (total {} rows, ZSTD(1), 1 row group/file).",
        raw_schema.fields().len(),
        cli.width_factor,
        wide_schema.fields().len(),
        cli.num_files,
        cli.rows_per_file,
        cap_rows,
    );

    std::fs::create_dir_all(&cli.dst_dir)
        .map_err(|e| exec_datafusion_err!("creating {}: {e}", cli.dst_dir.display()))?;

    let mut file_idx: usize = 0;
    let mut writer: Option<ArrowWriter<File>> = None;
    let mut rows_in_current: usize = 0;
    let mut total_written: usize = 0;
    let mut row_cursor: usize = 0;

    'outer: while total_written < cap_rows {
        let batch_n = cli.batch_size.min(cap_rows - total_written);
        let base_batch = synthesize_batch(row_cursor, batch_n, &raw_schema)?;
        let wide_batch = widen_batch(&base_batch, &wide_schema, cli.width_factor)?;
        row_cursor += batch_n;

        let mut remaining = wide_batch;
        while remaining.num_rows() > 0 {
            if writer.is_none() {
                if file_idx >= cli.num_files {
                    break 'outer;
                }
                file_idx += 1;
                let path = if cli.num_files == 1 {
                    cli.dst_dir.join("events.parquet")
                } else {
                    cli.dst_dir.join(format!("events_{file_idx:04}.parquet"))
                };
                writer = Some(open_writer(
                    &path,