#[macro_use]
extern crate criterion;
use criterion::Criterion;
use std::sync::{Arc, Mutex};
use tokio::runtime::Runtime;
extern crate arrow;
extern crate datafusion;
use arrow::{
array::{Float32Array, Float64Array},
datatypes::{DataType, Field, Schema},
record_batch::RecordBatch,
};
use datafusion::error::Result;
use datafusion::datasource::MemTable;
use datafusion::execution::context::ExecutionContext;
fn query(ctx: Arc<Mutex<ExecutionContext>>, sql: &str) {
let mut rt = Runtime::new().unwrap();
let df = ctx.lock().unwrap().sql(&sql).unwrap();
rt.block_on(df.collect()).unwrap();
}
fn create_context(
array_len: usize,
batch_size: usize,
) -> Result<Arc<Mutex<ExecutionContext>>> {
let schema = Arc::new(Schema::new(vec![
Field::new("f32", DataType::Float32, false),
Field::new("f64", DataType::Float64, false),
]));
let batches = (0..array_len / batch_size)
.map(|i| {
RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Float32Array::from(vec![i as f32; batch_size])),
Arc::new(Float64Array::from(vec![i as f64; batch_size])),
],
)
.unwrap()
})
.collect::<Vec<_>>();
let mut ctx = ExecutionContext::new();
let provider = MemTable::try_new(schema, vec![batches])?;
ctx.register_table("t", Box::new(provider));
Ok(Arc::new(Mutex::new(ctx)))
}
fn criterion_benchmark(c: &mut Criterion) {
let array_len = 1048576; let batch_size = 512; let ctx = create_context(array_len, batch_size).unwrap();
c.bench_function("sqrt_20_9", |b| {
b.iter(|| query(ctx.clone(), "SELECT sqrt(f32) FROM t"))
});
let array_len = 1048576; let batch_size = 4096; let ctx = create_context(array_len, batch_size).unwrap();
c.bench_function("sqrt_20_12", |b| {
b.iter(|| query(ctx.clone(), "SELECT sqrt(f32) FROM t"))
});
let array_len = 4194304; let batch_size = 4096; let ctx = create_context(array_len, batch_size).unwrap();
c.bench_function("sqrt_22_12", |b| {
b.iter(|| query(ctx.clone(), "SELECT sqrt(f32) FROM t"))
});
let array_len = 4194304; let batch_size = 16384; let ctx = create_context(array_len, batch_size).unwrap();
c.bench_function("sqrt_22_14", |b| {
b.iter(|| query(ctx.clone(), "SELECT sqrt(f32) FROM t"))
});
}
criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);