306 lines
12 KiB
Rust
306 lines
12 KiB
Rust
// Licensed to the Apache Software Foundation (ASF) under one
|
|
// or more contributor license agreements. See the NOTICE file
|
|
// distributed with this work for additional information
|
|
// regarding copyright ownership. The ASF licenses this file
|
|
// to you under the Apache License, Version 2.0 (the
|
|
// "License"); you may not use this file except in compliance
|
|
// with the License. You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing,
|
|
// software distributed under the License is distributed on an
|
|
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
// KIND, either express or implied. See the License for the
|
|
// specific language governing permissions and limitations
|
|
// under the License.
|
|
|
|
extern crate arrow;
|
|
#[macro_use]
|
|
extern crate criterion;
|
|
|
|
use criterion::{Criterion, Throughput};
|
|
use std::hint;
|
|
|
|
use arrow::array::*;
|
|
use arrow::util::bench_util::*;
|
|
use arrow_array::types::{Int8Type, Int16Type, Int32Type, Int64Type};
|
|
|
|
const BATCH_SIZE: usize = 64 * 1024;
|
|
|
|
/// Run [`ArrayIter::fold`] while using black_box on each item and the result of the cb to prevent compiler optimizations.
|
|
fn fold_black_box_item_and_cb_res<ArrayAcc, F, B>(array: ArrayAcc, init: B, mut f: F)
|
|
where
|
|
ArrayAcc: ArrayAccessor,
|
|
F: FnMut(B, Option<ArrayAcc::Item>) -> B,
|
|
{
|
|
let result = ArrayIter::new(array).fold(hint::black_box(init), |acc, item| {
|
|
let res = f(acc, hint::black_box(item));
|
|
hint::black_box(res)
|
|
});
|
|
|
|
hint::black_box(result);
|
|
}
|
|
/// Run [`ArrayIter::fold`] while using black_box on each item to prevent compiler optimizations.
|
|
fn fold_black_box_item<ArrayAcc, F, B>(array: ArrayAcc, init: B, mut f: F)
|
|
where
|
|
ArrayAcc: ArrayAccessor,
|
|
F: FnMut(B, Option<ArrayAcc::Item>) -> B,
|
|
{
|
|
let result = ArrayIter::new(array).fold(hint::black_box(init), |acc, item| {
|
|
f(acc, hint::black_box(item))
|
|
});
|
|
|
|
hint::black_box(result);
|
|
}
|
|
|
|
/// Run [`ArrayIter::fold`] without using black_box on each item, but only on the result
|
|
/// to see if the compiler can do more optimizations.
|
|
fn fold_black_box_result<ArrayAcc, F, B>(array: ArrayAcc, init: B, f: F)
|
|
where
|
|
ArrayAcc: ArrayAccessor,
|
|
F: FnMut(B, Option<ArrayAcc::Item>) -> B,
|
|
{
|
|
let result = ArrayIter::new(array).fold(hint::black_box(init), f);
|
|
|
|
hint::black_box(result);
|
|
}
|
|
|
|
/// Run [`ArrayIter::any`] while using black_box on each item and the predicate return value to prevent compiler optimizations.
|
|
fn any_black_box_item_and_predicate<ArrayAcc>(
|
|
array: ArrayAcc,
|
|
mut any_predicate: impl FnMut(Option<ArrayAcc::Item>) -> bool,
|
|
) where
|
|
ArrayAcc: ArrayAccessor,
|
|
{
|
|
let any_res = ArrayIter::new(array).any(|item| {
|
|
let item = hint::black_box(item);
|
|
let res = any_predicate(item);
|
|
hint::black_box(res)
|
|
});
|
|
|
|
hint::black_box(any_res);
|
|
}
|
|
|
|
/// Run [`ArrayIter::any`] without using black_box in the loop, but only on the result
|
|
/// to see if the compiler can do more optimizations.
|
|
fn any_black_box_result<ArrayAcc>(
|
|
array: ArrayAcc,
|
|
any_predicate: impl FnMut(Option<ArrayAcc::Item>) -> bool,
|
|
) where
|
|
ArrayAcc: ArrayAccessor,
|
|
{
|
|
let any_res = ArrayIter::new(array).any(any_predicate);
|
|
|
|
hint::black_box(any_res);
|
|
}
|
|
|
|
/// Benchmark [`ArrayIter`] functions,
|
|
///
|
|
/// The passed `predicate_that_will_always_evaluate_to_false` function should be a predicate
|
|
/// that always returns `false` to ensure that the full array is always iterated over.
|
|
///
|
|
/// The predicate function should:
|
|
/// 1. always return false
|
|
/// 2. be impossible for the compiler to optimize away
|
|
/// 3. not use `hint::black_box` internally (unless impossible) to allow for more compiler optimizations
|
|
///
|
|
/// the way to achieve this is to make the predicate check for a value that is not presented in the array.
|
|
///
|
|
/// The reason for these requirements is that we want to iterate over the entire array while
|
|
/// letting the compiler have room for optimizations so it will be more representative of real world usage.
|
|
fn benchmark_array_iter<ArrayAcc, FoldFn, FoldInit>(
|
|
c: &mut Criterion,
|
|
name: &str,
|
|
nonnull_array: ArrayAcc,
|
|
nullable_array: ArrayAcc,
|
|
fold_init: FoldInit,
|
|
fold_fn: FoldFn,
|
|
predicate_that_will_always_evaluate_to_false: impl Fn(Option<ArrayAcc::Item>) -> bool,
|
|
) where
|
|
ArrayAcc: ArrayAccessor + Copy,
|
|
FoldInit: Copy,
|
|
FoldFn: Fn(FoldInit, Option<ArrayAcc::Item>) -> FoldInit,
|
|
{
|
|
let predicate_that_will_always_evaluate_to_false =
|
|
&predicate_that_will_always_evaluate_to_false;
|
|
let fold_fn = &fold_fn;
|
|
|
|
// Assert always false return false
|
|
{
|
|
let found = ArrayIter::new(nonnull_array).any(predicate_that_will_always_evaluate_to_false);
|
|
assert!(!found, "The predicate must always evaluate to false");
|
|
}
|
|
{
|
|
let found =
|
|
ArrayIter::new(nullable_array).any(predicate_that_will_always_evaluate_to_false);
|
|
assert!(!found, "The predicate must always evaluate to false");
|
|
}
|
|
|
|
c.benchmark_group(name)
|
|
.throughput(Throughput::Elements(BATCH_SIZE as u64))
|
|
// Most of the Rust default iterator functions are implemented on top of 2 functions:
|
|
// `fold` and `try_fold`
|
|
// so we are benchmarking `fold` first
|
|
.bench_function("nonnull fold black box item and fold result", |b| {
|
|
b.iter(|| fold_black_box_item_and_cb_res(nonnull_array, fold_init, fold_fn))
|
|
})
|
|
.bench_function("nonnull fold black box item", |b| {
|
|
b.iter(|| fold_black_box_item(nonnull_array, fold_init, fold_fn))
|
|
})
|
|
.bench_function("nonnull fold black box only result", |b| {
|
|
b.iter(|| fold_black_box_result(nonnull_array, fold_init, fold_fn))
|
|
})
|
|
.bench_function("null fold black box item and fold result", |b| {
|
|
b.iter(|| fold_black_box_item_and_cb_res(nullable_array, fold_init, fold_fn))
|
|
})
|
|
.bench_function("null fold black box item", |b| {
|
|
b.iter(|| fold_black_box_item(nullable_array, fold_init, fold_fn))
|
|
})
|
|
.bench_function("null fold black box only result", |b| {
|
|
b.iter(|| fold_black_box_result(nullable_array, fold_init, fold_fn))
|
|
})
|
|
// Due to `try_fold` not being available in stable Rust,
|
|
// we are benchmarking `any` instead which the default Rust implementation
|
|
// uses `try_fold` under the hood.
|
|
.bench_function("nonnull any black box item and predicate", |b| {
|
|
b.iter(|| {
|
|
any_black_box_item_and_predicate(
|
|
nonnull_array,
|
|
predicate_that_will_always_evaluate_to_false,
|
|
)
|
|
})
|
|
})
|
|
.bench_function("nonnull any black box only result", |b| {
|
|
b.iter(|| {
|
|
any_black_box_result(nonnull_array, predicate_that_will_always_evaluate_to_false)
|
|
})
|
|
})
|
|
.bench_function("null any black box item and predicate", |b| {
|
|
b.iter(|| {
|
|
any_black_box_item_and_predicate(
|
|
nullable_array,
|
|
predicate_that_will_always_evaluate_to_false,
|
|
)
|
|
})
|
|
})
|
|
.bench_function("null any black box only result", |b| {
|
|
b.iter(|| {
|
|
any_black_box_result(nullable_array, predicate_that_will_always_evaluate_to_false)
|
|
})
|
|
});
|
|
}
|
|
|
|
/// Replace all occurrences of `item_to_replace` with `replace_with` in the given `PrimitiveArray`.
|
|
/// will make it so we can filter by missing value
|
|
fn replace_primitive_value<T>(
|
|
array: PrimitiveArray<T>,
|
|
item_to_replace: T::Native,
|
|
replace_with: T::Native,
|
|
) -> PrimitiveArray<T>
|
|
where
|
|
T: ArrowPrimitiveType,
|
|
<T as ArrowPrimitiveType>::Native: Eq,
|
|
{
|
|
array.unary(|item| {
|
|
if item == item_to_replace {
|
|
replace_with
|
|
} else {
|
|
item
|
|
}
|
|
})
|
|
}
|
|
|
|
fn add_benchmark(c: &mut Criterion) {
|
|
benchmark_array_iter(
|
|
c,
|
|
"int8",
|
|
&replace_primitive_value(create_primitive_array::<Int8Type>(BATCH_SIZE, 0.0), 42, 1),
|
|
&replace_primitive_value(create_primitive_array::<Int8Type>(BATCH_SIZE, 0.5), 42, 1),
|
|
// fold init
|
|
0i8,
|
|
// fold function
|
|
|acc, item| acc.wrapping_add(item.unwrap_or_default()),
|
|
// predicate that will always evaluate to false while allowing us to avoid using hint::black_box and let the compiler optimize more
|
|
|item| item == Some(42),
|
|
);
|
|
benchmark_array_iter(
|
|
c,
|
|
"int16",
|
|
&replace_primitive_value(create_primitive_array::<Int16Type>(BATCH_SIZE, 0.0), 42, 1),
|
|
&replace_primitive_value(create_primitive_array::<Int16Type>(BATCH_SIZE, 0.5), 42, 1),
|
|
// fold init
|
|
0i16,
|
|
// fold function
|
|
|acc, item| acc.wrapping_add(item.unwrap_or_default()),
|
|
// predicate that will always evaluate to false while allowing us to avoid using hint::black_box and let the compiler optimize more
|
|
|item| item == Some(42),
|
|
);
|
|
benchmark_array_iter(
|
|
c,
|
|
"int32",
|
|
&replace_primitive_value(create_primitive_array::<Int32Type>(BATCH_SIZE, 0.0), 42, 1),
|
|
&replace_primitive_value(create_primitive_array::<Int32Type>(BATCH_SIZE, 0.5), 42, 1),
|
|
// fold init
|
|
0i32,
|
|
// fold function
|
|
|acc, item| acc.wrapping_add(item.unwrap_or_default()),
|
|
// predicate that will always evaluate to false while allowing us to avoid using hint::black_box and let the compiler optimize more
|
|
|item| item == Some(42),
|
|
);
|
|
benchmark_array_iter(
|
|
c,
|
|
"int64",
|
|
&replace_primitive_value(create_primitive_array::<Int64Type>(BATCH_SIZE, 0.0), 42, 1),
|
|
&replace_primitive_value(create_primitive_array::<Int64Type>(BATCH_SIZE, 0.5), 42, 1),
|
|
// fold init
|
|
0i64,
|
|
// fold function
|
|
|acc, item| acc.wrapping_add(item.unwrap_or_default()),
|
|
// predicate that will always evaluate to false while allowing us to avoid using hint::black_box and let the compiler optimize more
|
|
|item| item == Some(42),
|
|
);
|
|
|
|
benchmark_array_iter(
|
|
c,
|
|
"string with len 16",
|
|
&create_string_array_with_len::<i32>(BATCH_SIZE, 0.0, 16),
|
|
&create_string_array_with_len::<i32>(BATCH_SIZE, 0.5, 16),
|
|
// fold init
|
|
0_usize,
|
|
// fold function
|
|
|acc, item| acc.wrapping_add(item.map(|item| item.len()).unwrap_or_default()),
|
|
// predicate that will always evaluate to false while allowing us to avoid using hint::black_box and let the compiler optimize more
|
|
|item| item.is_some_and(|item| item.is_empty()),
|
|
);
|
|
|
|
benchmark_array_iter(
|
|
c,
|
|
"string view with len 16",
|
|
&create_string_view_array_with_len(BATCH_SIZE, 0.0, 16, false),
|
|
&create_string_view_array_with_len(BATCH_SIZE, 0.5, 16, false),
|
|
// fold init
|
|
0_usize,
|
|
// fold function
|
|
|acc, item| acc.wrapping_add(item.map(|item| item.len()).unwrap_or_default()),
|
|
// predicate that will always evaluate to false while allowing us to avoid using hint::black_box and let the compiler optimize more
|
|
|item| item.is_some_and(|item| item.is_empty()),
|
|
);
|
|
|
|
benchmark_array_iter(
|
|
c,
|
|
"boolean mixed true and false",
|
|
&create_boolean_array(BATCH_SIZE, 0.0, 0.5),
|
|
&create_boolean_array(BATCH_SIZE, 0.5, 0.5),
|
|
// fold init
|
|
0_usize,
|
|
// fold function
|
|
|acc, item| acc.wrapping_add(item.unwrap_or_default() as usize),
|
|
// Must use black_box here as this can be optimized away
|
|
|_item| hint::black_box(false),
|
|
);
|
|
}
|
|
|
|
criterion_group!(benches, add_benchmark);
|
|
criterion_main!(benches);
|