Skip to content
59 changes: 53 additions & 6 deletions rust/arrow/benches/sort_kernel.rs
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@ fn create_array(size: usize, with_nulls: bool) -> ArrayRef {
Arc::new(array)
}

fn bench_sort(arr_a: &ArrayRef, array_b: &ArrayRef) {
fn bench_sort(arr_a: &ArrayRef, array_b: &ArrayRef, limit: Option<usize>) {
let columns = vec![
SortColumn {
values: arr_a.clone(),
Expand All @@ -45,29 +45,76 @@ fn bench_sort(arr_a: &ArrayRef, array_b: &ArrayRef) {
},
];

criterion::black_box(lexsort(&columns).unwrap());
criterion::black_box(lexsort(&columns, limit).unwrap());
}

fn add_benchmark(c: &mut Criterion) {
let arr_a = create_array(2u64.pow(10) as usize, false);
let arr_b = create_array(2u64.pow(10) as usize, false);

c.bench_function("sort 2^10", |b| b.iter(|| bench_sort(&arr_a, &arr_b)));
c.bench_function("sort 2^10", |b| b.iter(|| bench_sort(&arr_a, &arr_b, None)));

let arr_a = create_array(2u64.pow(12) as usize, false);
let arr_b = create_array(2u64.pow(12) as usize, false);

c.bench_function("sort 2^12", |b| b.iter(|| bench_sort(&arr_a, &arr_b)));
c.bench_function("sort 2^12", |b| b.iter(|| bench_sort(&arr_a, &arr_b, None)));

let arr_a = create_array(2u64.pow(10) as usize, true);
let arr_b = create_array(2u64.pow(10) as usize, true);

c.bench_function("sort nulls 2^10", |b| b.iter(|| bench_sort(&arr_a, &arr_b)));
c.bench_function("sort nulls 2^10", |b| {
b.iter(|| bench_sort(&arr_a, &arr_b, None))
});

let arr_a = create_array(2u64.pow(12) as usize, true);
let arr_b = create_array(2u64.pow(12) as usize, true);

c.bench_function("sort nulls 2^12", |b| b.iter(|| bench_sort(&arr_a, &arr_b)));
c.bench_function("sort nulls 2^12", |b| {
b.iter(|| bench_sort(&arr_a, &arr_b, None))
});

// with limit
{
let arr_a = create_array(2u64.pow(12) as usize, false);
let arr_b = create_array(2u64.pow(12) as usize, false);
c.bench_function("sort 2^12 limit 10", |b| {
b.iter(|| bench_sort(&arr_a, &arr_b, Some(10)))
});

let arr_a = create_array(2u64.pow(12) as usize, false);
let arr_b = create_array(2u64.pow(12) as usize, false);
c.bench_function("sort 2^12 limit 100", |b| {
b.iter(|| bench_sort(&arr_a, &arr_b, Some(100)))
});

let arr_a = create_array(2u64.pow(12) as usize, false);
let arr_b = create_array(2u64.pow(12) as usize, false);
c.bench_function("sort 2^12 limit 1000", |b| {
b.iter(|| bench_sort(&arr_a, &arr_b, Some(1000)))
});

let arr_a = create_array(2u64.pow(12) as usize, false);
let arr_b = create_array(2u64.pow(12) as usize, false);
c.bench_function("sort 2^12 limit 2^12", |b| {
b.iter(|| bench_sort(&arr_a, &arr_b, Some(2u64.pow(12) as usize)))
});

let arr_a = create_array(2u64.pow(12) as usize, true);
let arr_b = create_array(2u64.pow(12) as usize, true);

c.bench_function("sort nulls 2^12 limit 10", |b| {
b.iter(|| bench_sort(&arr_a, &arr_b, Some(10)))
});
c.bench_function("sort nulls 2^12 limit 100", |b| {
b.iter(|| bench_sort(&arr_a, &arr_b, Some(100)))
});
c.bench_function("sort nulls 2^12 limit 1000", |b| {
b.iter(|| bench_sort(&arr_a, &arr_b, Some(1000)))
});
c.bench_function("sort nulls 2^12 limit 2^12", |b| {
b.iter(|| bench_sort(&arr_a, &arr_b, Some(2u64.pow(12) as usize)))
});
}
}

criterion_group!(benches, add_benchmark);
Expand Down
Loading