328 lines
9.7 KiB
Rust
328 lines
9.7 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.
|
|
|
|
use criterion::measurement::WallTime;
|
|
use criterion::{BenchmarkGroup, BenchmarkId, Criterion, criterion_group, criterion_main};
|
|
use rand::distr::{Distribution, StandardUniform};
|
|
use rand::prelude::StdRng;
|
|
use rand::{Rng, SeedableRng};
|
|
use std::hint;
|
|
use std::ops::Range;
|
|
use std::sync::Arc;
|
|
|
|
use arrow::array::*;
|
|
use arrow::datatypes::*;
|
|
use arrow::util::bench_util::*;
|
|
use arrow_select::zip::zip;
|
|
|
|
trait InputGenerator {
|
|
fn name(&self) -> &str;
|
|
|
|
/// Return an ArrayRef containing a single null value
|
|
fn generate_scalar_with_null_value(&self) -> ArrayRef;
|
|
|
|
/// Generate a `number_of_scalars` unique scalars
|
|
fn generate_non_null_scalars(&self, seed: u64, number_of_scalars: usize) -> Vec<ArrayRef>;
|
|
|
|
/// Generate array with specified length and null percentage
|
|
fn generate_array(&self, seed: u64, array_length: usize, null_percentage: f32) -> ArrayRef;
|
|
}
|
|
|
|
struct GeneratePrimitive<T: ArrowPrimitiveType> {
|
|
description: String,
|
|
_marker: std::marker::PhantomData<T>,
|
|
}
|
|
|
|
impl<T> InputGenerator for GeneratePrimitive<T>
|
|
where
|
|
T: ArrowPrimitiveType,
|
|
StandardUniform: Distribution<T::Native>,
|
|
{
|
|
fn name(&self) -> &str {
|
|
self.description.as_str()
|
|
}
|
|
|
|
fn generate_scalar_with_null_value(&self) -> ArrayRef {
|
|
new_null_array(&T::DATA_TYPE, 1)
|
|
}
|
|
|
|
fn generate_non_null_scalars(&self, seed: u64, number_of_scalars: usize) -> Vec<ArrayRef> {
|
|
let rng = StdRng::seed_from_u64(seed);
|
|
|
|
rng.sample_iter::<T::Native, _>(StandardUniform)
|
|
.take(number_of_scalars)
|
|
.map(|v: T::Native| {
|
|
Arc::new(PrimitiveArray::<T>::new_scalar(v).into_inner()) as ArrayRef
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn generate_array(&self, seed: u64, array_length: usize, null_percentage: f32) -> ArrayRef {
|
|
Arc::new(create_primitive_array_with_seed::<T>(
|
|
array_length,
|
|
null_percentage,
|
|
seed,
|
|
))
|
|
}
|
|
}
|
|
|
|
struct GenerateBytes<Byte: ByteArrayType> {
|
|
range_length: std::ops::Range<usize>,
|
|
description: String,
|
|
|
|
_marker: std::marker::PhantomData<Byte>,
|
|
}
|
|
|
|
impl<Byte> InputGenerator for GenerateBytes<Byte>
|
|
where
|
|
Byte: ByteArrayType,
|
|
{
|
|
fn name(&self) -> &str {
|
|
self.description.as_str()
|
|
}
|
|
|
|
fn generate_scalar_with_null_value(&self) -> ArrayRef {
|
|
new_null_array(&Byte::DATA_TYPE, 1)
|
|
}
|
|
|
|
fn generate_non_null_scalars(&self, seed: u64, number_of_scalars: usize) -> Vec<ArrayRef> {
|
|
let array = self.generate_array(seed, number_of_scalars, 0.0);
|
|
|
|
(0..number_of_scalars).map(|i| array.slice(i, 1)).collect()
|
|
}
|
|
|
|
fn generate_array(&self, seed: u64, array_length: usize, null_percentage: f32) -> ArrayRef {
|
|
let is_binary =
|
|
Byte::DATA_TYPE == DataType::Binary || Byte::DATA_TYPE == DataType::LargeBinary;
|
|
if is_binary {
|
|
Arc::new(create_binary_array_with_len_range_and_prefix_and_seed::<
|
|
Byte::Offset,
|
|
>(
|
|
array_length,
|
|
null_percentage,
|
|
self.range_length.start,
|
|
self.range_length.end - 1,
|
|
&[],
|
|
seed,
|
|
))
|
|
} else {
|
|
Arc::new(create_string_array_with_len_range_and_prefix_and_seed::<
|
|
Byte::Offset,
|
|
>(
|
|
array_length,
|
|
null_percentage,
|
|
self.range_length.start,
|
|
self.range_length.end - 1,
|
|
"",
|
|
seed,
|
|
))
|
|
}
|
|
}
|
|
}
|
|
|
|
struct GenerateStringView {
|
|
range: Range<usize>,
|
|
description: String,
|
|
_marker: std::marker::PhantomData<StringViewType>,
|
|
}
|
|
|
|
impl InputGenerator for GenerateStringView {
|
|
fn name(&self) -> &str {
|
|
self.description.as_str()
|
|
}
|
|
fn generate_scalar_with_null_value(&self) -> ArrayRef {
|
|
new_null_array(&DataType::Utf8View, 1)
|
|
}
|
|
|
|
fn generate_non_null_scalars(&self, seed: u64, number_of_scalars: usize) -> Vec<ArrayRef> {
|
|
let array = self.generate_array(seed, number_of_scalars, 0.0);
|
|
(0..number_of_scalars).map(|i| array.slice(i, 1)).collect()
|
|
}
|
|
|
|
fn generate_array(&self, seed: u64, array_length: usize, null_percentage: f32) -> ArrayRef {
|
|
Arc::new(create_string_view_array_with_len_range_and_seed(
|
|
array_length,
|
|
null_percentage,
|
|
self.range.clone(),
|
|
seed,
|
|
))
|
|
}
|
|
}
|
|
|
|
fn mask_cases(len: usize) -> Vec<(&'static str, BooleanArray)> {
|
|
vec![
|
|
("all_true", create_boolean_array(len, 0.0, 1.0)),
|
|
("99pct_true", create_boolean_array(len, 0.0, 0.99)),
|
|
("90pct_true", create_boolean_array(len, 0.0, 0.9)),
|
|
("50pct_true", create_boolean_array(len, 0.0, 0.5)),
|
|
("10pct_true", create_boolean_array(len, 0.0, 0.1)),
|
|
("1pct_true", create_boolean_array(len, 0.0, 0.01)),
|
|
("all_false", create_boolean_array(len, 0.0, 0.0)),
|
|
("50pct_nulls", create_boolean_array(len, 0.5, 0.5)),
|
|
]
|
|
}
|
|
|
|
fn bench_zip_on_input_generator(c: &mut Criterion, input_generator: &impl InputGenerator) {
|
|
const ARRAY_LEN: usize = 8192;
|
|
|
|
let mut group =
|
|
c.benchmark_group(format!("zip_{ARRAY_LEN}_from_{}", input_generator.name()).as_str());
|
|
|
|
let null_scalar = input_generator.generate_scalar_with_null_value();
|
|
let [non_null_scalar_1, non_null_scalar_2]: [_; 2] = input_generator
|
|
.generate_non_null_scalars(42, 2)
|
|
.try_into()
|
|
.unwrap();
|
|
|
|
let array_1_10pct_nulls = input_generator.generate_array(42, ARRAY_LEN, 0.1);
|
|
let array_2_10pct_nulls = input_generator.generate_array(18, ARRAY_LEN, 0.1);
|
|
|
|
let masks = mask_cases(ARRAY_LEN);
|
|
|
|
// Benchmarks for different scalar combinations
|
|
for (description, truthy, falsy) in &[
|
|
("null_vs_non_null_scalar", &null_scalar, &non_null_scalar_1),
|
|
(
|
|
"non_null_scalar_vs_null_scalar",
|
|
&non_null_scalar_1,
|
|
&null_scalar,
|
|
),
|
|
("non_nulls_scalars", &non_null_scalar_1, &non_null_scalar_2),
|
|
] {
|
|
bench_zip_input_on_all_masks(
|
|
description,
|
|
&mut group,
|
|
&masks,
|
|
&Scalar::new(truthy),
|
|
&Scalar::new(falsy),
|
|
);
|
|
}
|
|
|
|
bench_zip_input_on_all_masks(
|
|
"array_vs_non_null_scalar",
|
|
&mut group,
|
|
&masks,
|
|
&array_1_10pct_nulls,
|
|
&non_null_scalar_1,
|
|
);
|
|
|
|
bench_zip_input_on_all_masks(
|
|
"non_null_scalar_vs_array",
|
|
&mut group,
|
|
&masks,
|
|
&non_null_scalar_1,
|
|
&array_1_10pct_nulls,
|
|
);
|
|
|
|
bench_zip_input_on_all_masks(
|
|
"array_vs_array",
|
|
&mut group,
|
|
&masks,
|
|
&array_1_10pct_nulls,
|
|
&array_2_10pct_nulls,
|
|
);
|
|
|
|
group.finish();
|
|
}
|
|
|
|
fn bench_zip_input_on_all_masks(
|
|
description: &str,
|
|
group: &mut BenchmarkGroup<WallTime>,
|
|
masks: &[(&str, BooleanArray)],
|
|
truthy: &impl Datum,
|
|
falsy: &impl Datum,
|
|
) {
|
|
for (mask_description, mask) in masks {
|
|
let id = BenchmarkId::new(description, mask_description);
|
|
group.bench_with_input(id, mask, |b, mask| {
|
|
b.iter(|| hint::black_box(zip(mask, truthy, falsy)))
|
|
});
|
|
}
|
|
}
|
|
|
|
fn add_benchmark(c: &mut Criterion) {
|
|
// Primitive
|
|
bench_zip_on_input_generator(
|
|
c,
|
|
&GeneratePrimitive::<Int32Type> {
|
|
description: "i32".to_string(),
|
|
_marker: std::marker::PhantomData,
|
|
},
|
|
);
|
|
|
|
// Short strings
|
|
bench_zip_on_input_generator(
|
|
c,
|
|
&GenerateBytes::<GenericStringType<i32>> {
|
|
description: "short strings (3..10)".to_string(),
|
|
range_length: 3..10,
|
|
_marker: std::marker::PhantomData,
|
|
},
|
|
);
|
|
|
|
// Long strings
|
|
bench_zip_on_input_generator(
|
|
c,
|
|
&GenerateBytes::<GenericStringType<i32>> {
|
|
description: "long strings (100..400)".to_string(),
|
|
range_length: 100..400,
|
|
_marker: std::marker::PhantomData,
|
|
},
|
|
);
|
|
|
|
// Short Bytes
|
|
bench_zip_on_input_generator(
|
|
c,
|
|
&GenerateBytes::<GenericBinaryType<i32>> {
|
|
description: "short bytes (3..10)".to_string(),
|
|
range_length: 3..10,
|
|
_marker: std::marker::PhantomData,
|
|
},
|
|
);
|
|
|
|
// Long Bytes
|
|
bench_zip_on_input_generator(
|
|
c,
|
|
&GenerateBytes::<GenericBinaryType<i32>> {
|
|
description: "long bytes (100..400)".to_string(),
|
|
range_length: 100..400,
|
|
_marker: std::marker::PhantomData,
|
|
},
|
|
);
|
|
|
|
bench_zip_on_input_generator(
|
|
c,
|
|
&GenerateStringView {
|
|
description: "string_views size (3..10)".to_string(),
|
|
range: 3..10,
|
|
_marker: std::marker::PhantomData,
|
|
},
|
|
);
|
|
|
|
bench_zip_on_input_generator(
|
|
c,
|
|
&GenerateStringView {
|
|
description: "string_views size (10..100)".to_string(),
|
|
range: 10..100,
|
|
_marker: std::marker::PhantomData,
|
|
},
|
|
);
|
|
}
|
|
|
|
criterion_group!(benches, add_benchmark);
|
|
criterion_main!(benches);
|