Skip to content

Commit f131b54

Browse files
authored
bench: create zip kernel benchmarks (#8654)
# Which issue does this PR close? N/A # Rationale for this change I have a PR to improve zip perf for scalar but I don't see any benchmarks for it: - #8653 # What changes are included in this PR? created zip benchmarks for scalar and non scalar with different masks # Are these changes tested? N/A # Are there any user-facing changes? Nope
1 parent d49f017 commit f131b54

File tree

3 files changed

+364
-3
lines changed

3 files changed

+364
-3
lines changed

arrow/Cargo.toml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -177,6 +177,11 @@ name = "interleave_kernels"
177177
harness = false
178178
required-features = ["test_utils"]
179179

180+
[[bench]]
181+
name = "zip_kernels"
182+
harness = false
183+
required-features = ["test_utils"]
184+
180185
[[bench]]
181186
name = "length_kernel"
182187
harness = false

arrow/benches/zip_kernels.rs

Lines changed: 279 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,279 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
use criterion::measurement::WallTime;
19+
use criterion::{BenchmarkGroup, BenchmarkId, Criterion, criterion_group, criterion_main};
20+
use rand::distr::{Distribution, StandardUniform};
21+
use rand::prelude::StdRng;
22+
use rand::{Rng, SeedableRng};
23+
use std::hint;
24+
use std::sync::Arc;
25+
26+
use arrow::array::*;
27+
use arrow::datatypes::*;
28+
use arrow::util::bench_util::*;
29+
use arrow_select::zip::zip;
30+
31+
trait InputGenerator {
32+
fn name(&self) -> &str;
33+
34+
/// Return an ArrayRef containing a single null value
35+
fn generate_scalar_with_null_value(&self) -> ArrayRef;
36+
37+
/// Generate a `number_of_scalars` unique scalars
38+
fn generate_non_null_scalars(&self, seed: u64, number_of_scalars: usize) -> Vec<ArrayRef>;
39+
40+
/// Generate array with specified length and null percentage
41+
fn generate_array(&self, seed: u64, array_length: usize, null_percentage: f32) -> ArrayRef;
42+
}
43+
44+
struct GeneratePrimitive<T: ArrowPrimitiveType> {
45+
description: String,
46+
_marker: std::marker::PhantomData<T>,
47+
}
48+
49+
impl<T> InputGenerator for GeneratePrimitive<T>
50+
where
51+
T: ArrowPrimitiveType,
52+
StandardUniform: Distribution<T::Native>,
53+
{
54+
fn name(&self) -> &str {
55+
self.description.as_str()
56+
}
57+
58+
fn generate_scalar_with_null_value(&self) -> ArrayRef {
59+
new_null_array(&T::DATA_TYPE, 1)
60+
}
61+
62+
fn generate_non_null_scalars(&self, seed: u64, number_of_scalars: usize) -> Vec<ArrayRef> {
63+
let rng = StdRng::seed_from_u64(seed);
64+
65+
rng.sample_iter::<T::Native, _>(StandardUniform)
66+
.take(number_of_scalars)
67+
.map(|v: T::Native| {
68+
Arc::new(PrimitiveArray::<T>::new_scalar(v).into_inner()) as ArrayRef
69+
})
70+
.collect()
71+
}
72+
73+
fn generate_array(&self, seed: u64, array_length: usize, null_percentage: f32) -> ArrayRef {
74+
Arc::new(create_primitive_array_with_seed::<T>(
75+
array_length,
76+
null_percentage,
77+
seed,
78+
))
79+
}
80+
}
81+
82+
struct GenerateBytes<Byte: ByteArrayType> {
83+
range_length: std::ops::Range<usize>,
84+
description: String,
85+
86+
_marker: std::marker::PhantomData<Byte>,
87+
}
88+
89+
impl<Byte> InputGenerator for GenerateBytes<Byte>
90+
where
91+
Byte: ByteArrayType,
92+
{
93+
fn name(&self) -> &str {
94+
self.description.as_str()
95+
}
96+
97+
fn generate_scalar_with_null_value(&self) -> ArrayRef {
98+
new_null_array(&Byte::DATA_TYPE, 1)
99+
}
100+
101+
fn generate_non_null_scalars(&self, seed: u64, number_of_scalars: usize) -> Vec<ArrayRef> {
102+
let array = self.generate_array(seed, number_of_scalars, 0.0);
103+
104+
(0..number_of_scalars).map(|i| array.slice(i, 1)).collect()
105+
}
106+
107+
fn generate_array(&self, seed: u64, array_length: usize, null_percentage: f32) -> ArrayRef {
108+
let is_binary =
109+
Byte::DATA_TYPE == DataType::Binary || Byte::DATA_TYPE == DataType::LargeBinary;
110+
if is_binary {
111+
Arc::new(create_binary_array_with_len_range_and_prefix_and_seed::<
112+
Byte::Offset,
113+
>(
114+
array_length,
115+
null_percentage,
116+
self.range_length.start,
117+
self.range_length.end - 1,
118+
&[],
119+
seed,
120+
))
121+
} else {
122+
Arc::new(create_string_array_with_len_range_and_prefix_and_seed::<
123+
Byte::Offset,
124+
>(
125+
array_length,
126+
null_percentage,
127+
self.range_length.start,
128+
self.range_length.end - 1,
129+
"",
130+
seed,
131+
))
132+
}
133+
}
134+
}
135+
136+
fn mask_cases(len: usize) -> Vec<(&'static str, BooleanArray)> {
137+
vec![
138+
("all_true", create_boolean_array(len, 0.0, 1.0)),
139+
("99pct_true", create_boolean_array(len, 0.0, 0.99)),
140+
("90pct_true", create_boolean_array(len, 0.0, 0.9)),
141+
("50pct_true", create_boolean_array(len, 0.0, 0.5)),
142+
("10pct_true", create_boolean_array(len, 0.0, 0.1)),
143+
("1pct_true", create_boolean_array(len, 0.0, 0.01)),
144+
("all_false", create_boolean_array(len, 0.0, 0.0)),
145+
("50pct_nulls", create_boolean_array(len, 0.5, 0.5)),
146+
]
147+
}
148+
149+
fn bench_zip_on_input_generator(c: &mut Criterion, input_generator: &impl InputGenerator) {
150+
const ARRAY_LEN: usize = 8192;
151+
152+
let mut group =
153+
c.benchmark_group(format!("zip_{ARRAY_LEN}_from_{}", input_generator.name()).as_str());
154+
155+
let null_scalar = input_generator.generate_scalar_with_null_value();
156+
let [non_null_scalar_1, non_null_scalar_2]: [_; 2] = input_generator
157+
.generate_non_null_scalars(42, 2)
158+
.try_into()
159+
.unwrap();
160+
161+
let array_1_10pct_nulls = input_generator.generate_array(42, ARRAY_LEN, 0.1);
162+
let array_2_10pct_nulls = input_generator.generate_array(18, ARRAY_LEN, 0.1);
163+
164+
let masks = mask_cases(ARRAY_LEN);
165+
166+
// Benchmarks for different scalar combinations
167+
for (description, truthy, falsy) in &[
168+
("null_vs_non_null_scalar", &null_scalar, &non_null_scalar_1),
169+
(
170+
"non_null_scalar_vs_null_scalar",
171+
&non_null_scalar_1,
172+
&null_scalar,
173+
),
174+
("non_nulls_scalars", &non_null_scalar_1, &non_null_scalar_2),
175+
] {
176+
bench_zip_input_on_all_masks(
177+
description,
178+
&mut group,
179+
&masks,
180+
&Scalar::new(truthy),
181+
&Scalar::new(falsy),
182+
);
183+
}
184+
185+
bench_zip_input_on_all_masks(
186+
"array_vs_non_null_scalar",
187+
&mut group,
188+
&masks,
189+
&array_1_10pct_nulls,
190+
&non_null_scalar_1,
191+
);
192+
193+
bench_zip_input_on_all_masks(
194+
"non_null_scalar_vs_array",
195+
&mut group,
196+
&masks,
197+
&array_1_10pct_nulls,
198+
&non_null_scalar_1,
199+
);
200+
201+
bench_zip_input_on_all_masks(
202+
"array_vs_array",
203+
&mut group,
204+
&masks,
205+
&array_1_10pct_nulls,
206+
&array_2_10pct_nulls,
207+
);
208+
209+
group.finish();
210+
}
211+
212+
fn bench_zip_input_on_all_masks(
213+
description: &str,
214+
group: &mut BenchmarkGroup<WallTime>,
215+
masks: &[(&str, BooleanArray)],
216+
truthy: &impl Datum,
217+
falsy: &impl Datum,
218+
) {
219+
for (mask_description, mask) in masks {
220+
let id = BenchmarkId::new(description, mask_description);
221+
group.bench_with_input(id, mask, |b, mask| {
222+
b.iter(|| hint::black_box(zip(mask, truthy, falsy)))
223+
});
224+
}
225+
}
226+
227+
fn add_benchmark(c: &mut Criterion) {
228+
// Primitive
229+
bench_zip_on_input_generator(
230+
c,
231+
&GeneratePrimitive::<Int32Type> {
232+
description: "i32".to_string(),
233+
_marker: std::marker::PhantomData,
234+
},
235+
);
236+
237+
// Short strings
238+
bench_zip_on_input_generator(
239+
c,
240+
&GenerateBytes::<GenericStringType<i32>> {
241+
description: "short strings (3..10)".to_string(),
242+
range_length: 3..10,
243+
_marker: std::marker::PhantomData,
244+
},
245+
);
246+
247+
// Long strings
248+
bench_zip_on_input_generator(
249+
c,
250+
&GenerateBytes::<GenericStringType<i32>> {
251+
description: "long strings (100..400)".to_string(),
252+
range_length: 100..400,
253+
_marker: std::marker::PhantomData,
254+
},
255+
);
256+
257+
// Short Bytes
258+
bench_zip_on_input_generator(
259+
c,
260+
&GenerateBytes::<GenericBinaryType<i32>> {
261+
description: "short bytes (3..10)".to_string(),
262+
range_length: 3..10,
263+
_marker: std::marker::PhantomData,
264+
},
265+
);
266+
267+
// Long Bytes
268+
bench_zip_on_input_generator(
269+
c,
270+
&GenerateBytes::<GenericBinaryType<i32>> {
271+
description: "long bytes (100..400)".to_string(),
272+
range_length: 100..400,
273+
_marker: std::marker::PhantomData,
274+
},
275+
);
276+
}
277+
278+
criterion_group!(benches, add_benchmark);
279+
criterion_main!(benches);

0 commit comments

Comments
 (0)