#![allow(clippy::unwrap_used)]
use std::iter::Iterator;
use arrow_buffer::{ArrowNativeType, MutableBuffer, ScalarBuffer};
use divan::Bencher;
use num_traits::PrimInt;
use vortex_buffer::{Buffer, BufferMut};
use vortex_error::{VortexExpect, vortex_err};
fn main() {
divan::main();
}
pub struct Arrow<T>(T);
impl<T: ArrowNativeType> FromIterator<T> for Arrow<ScalarBuffer<T>> {
fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
Self(ScalarBuffer::from_iter(iter))
}
}
impl<T: ArrowNativeType> FromIterator<T> for Arrow<MutableBuffer> {
fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
Self(MutableBuffer::from_iter(iter))
}
}
const INPUT_SIZE: &[i32] = &[128, 1024, 2048, 16_384, 65_536];
const INPUT_SIZE_USIZE: &[usize] = &[128, 1024, 2048, 16_384, 65_536];
#[divan::bench(
types = [Arrow<ScalarBuffer<i32>>,Buffer<i32>],
args = INPUT_SIZE,
)]
fn from_iter<B: FromIterator<i32>>(n: i32) {
B::from_iter((0..n).map(|i| i % i32::MAX));
}
trait MapEach<T, R> {
type Output;
fn map_each<F>(self, f: F) -> Self::Output
where
F: FnMut(T) -> R;
}
impl<T: ArrowNativeType, R: ArrowNativeType> MapEach<T, R> for Arrow<ScalarBuffer<T>> {
type Output = Arrow<ScalarBuffer<R>>;
fn map_each<F>(self, f: F) -> Self::Output
where
F: FnMut(T) -> R,
{
Arrow(ScalarBuffer::from(
self.0
.into_inner()
.into_vec::<T>()
.map_err(|_| vortex_err!("Failed to convert Arrow buffer into a mut vec"))
.vortex_expect("Failed to convert Arrow buffer into a mut vec")
.into_iter()
.map(f)
.collect::<Vec<R>>(),
))
}
}
impl<T: Copy, R> MapEach<T, R> for BufferMut<T> {
type Output = BufferMut<R>;
fn map_each<F>(self, f: F) -> Self::Output
where
F: FnMut(T) -> R,
{
BufferMut::<T>::map_each(self, f)
}
}
#[divan::bench(
types = [Arrow<ScalarBuffer<i32>>, BufferMut<i32>],
args = INPUT_SIZE,
)]
fn map_each<B: MapEach<i32, u32> + FromIterator<i32>>(bencher: Bencher, n: i32) {
bencher
.with_inputs(|| B::from_iter((0..n).map(|i| i % i32::MAX)))
.bench_values(|buffer| B::map_each(buffer, |i| (i as u32) + 1));
}
#[divan::bench(args = INPUT_SIZE)]
fn push_vortex_buffer(bencher: Bencher, length: i32) {
bencher
.with_inputs(|| BufferMut::<i32>::with_capacity(length as usize))
.bench_refs(|buffer| {
for idx in 0..length {
buffer.push(divan::black_box(idx));
}
});
}
#[divan::bench(args = INPUT_SIZE)]
fn push_arrow_buffer(bencher: Bencher, length: i32) {
bencher
.with_inputs(|| {
Arrow(MutableBuffer::with_capacity(
length as usize * size_of::<i32>(),
))
})
.bench_refs(|buffer| {
for idx in 0..length {
buffer.0.push(divan::black_box(idx));
}
});
}
#[divan::bench(types = [u8, u16, u32, u64], args = INPUT_SIZE_USIZE)]
fn push_n_vortex_buffer<T: PrimInt>(bencher: Bencher, length: usize) {
bencher
.with_inputs(|| BufferMut::<T>::with_capacity(length))
.bench_refs(|buffer| {
for _ in 0..100 {
unsafe {
buffer.push_n_unchecked(
divan::black_box(T::one()),
divan::black_box(length / 100),
)
};
}
});
}