Function generate_data

benchmarks/src/cancellation.rs:271–312

(
    data_dir: impl AsRef<Path>,
    num_files: usize,
    num_rows_per_file: usize,
)

}

async fn generate_data(
    data_dir: impl AsRef<Path>,
    num_files: usize,
    num_rows_per_file: usize,
) -> Result<(), Box<dyn std::error::Error + Send + Sync + 'static>> {
    let absolute = std::env::current_dir().unwrap().join(data_dir);
    let store = Arc::new(object_store::local::LocalFileSystem::new_with_prefix(
        absolute,
    )?);

    let columns = [
        ("A", DataType::Float64),
        ("B", DataType::Float64),
        ("C", DataType::Float64),
        ("D", DataType::Boolean),
        ("E", DataType::Utf8),
        ("F", DataType::Utf8),
        ("G", DataType::Utf8),
        ("H", DataType::Utf8),
        ("I", DataType::Utf8),
        ("J", DataType::Utf8),
        ("K", DataType::Utf8),
    ];

    for file_num in 1..=num_files {
        println!("Generating file {file_num} of {num_files}");
        let data = columns.iter().map(|(column_name, column_type)| {
            let column = random_data(column_type, num_rows_per_file);
            (column_name, column)
        });
        let to_write = RecordBatch::try_from_iter(data).unwrap();
        let path = object_store::path::Path::from(format!("{file_num}.parquet").as_str());
        let object_store_writer = ParquetObjectWriter::new(Arc::clone(&store) as _, path);

        let mut writer =
            AsyncArrowWriter::try_new(object_store_writer, to_write.schema(), None)?;
        writer.write(&to_write).await?;
        writer.close().await?;
    }

    Ok(())
}

fn random_data(column_type: &DataType, rows: usize) -> Arc<dyn Array> {
    let mut rng = rand::rng();