Skip to content

Commit 7bab215

Browse files
authored
arrow: add concat structs benchmark (#7520)
1 parent 1a5999a commit 7bab215

File tree

1 file changed

+45
-2
lines changed

1 file changed

+45
-2
lines changed

arrow/benches/concatenate_kernel.rs

Lines changed: 45 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,13 @@
1515
// specific language governing permissions and limitations
1616
// under the License.
1717

18+
extern crate arrow;
1819
#[macro_use]
1920
extern crate criterion;
2021
use std::sync::Arc;
2122

2223
use criterion::Criterion;
2324

24-
extern crate arrow;
25-
2625
use arrow::array::*;
2726
use arrow::compute::concat;
2827
use arrow::datatypes::*;
@@ -178,6 +177,50 @@ fn add_benchmark(c: &mut Criterion) {
178177
c.bench_function("concat fixed size lists", |b| {
179178
b.iter(|| bench_concat(&v1, &v2))
180179
});
180+
181+
{
182+
let batch_size = 1024;
183+
let batch_count = 2;
184+
let struct_arrays = (0..batch_count)
185+
.map(|_| {
186+
let ints = create_primitive_array::<Int32Type>(batch_size, 0.0);
187+
let string_dict = create_sparse_dict_from_values::<Int32Type>(
188+
batch_size,
189+
0.0,
190+
&create_string_array_with_len::<i32>(20, 0.0, 10),
191+
0..10,
192+
);
193+
let int_dict = create_sparse_dict_from_values::<UInt16Type>(
194+
batch_size,
195+
0.0,
196+
&create_primitive_array::<Int64Type>(20, 0.0),
197+
0..10,
198+
);
199+
let fields = vec![
200+
Field::new("int_field", ints.data_type().clone(), false),
201+
Field::new("strings_dict_field", string_dict.data_type().clone(), false),
202+
Field::new("int_dict_field", int_dict.data_type().clone(), false),
203+
];
204+
205+
StructArray::try_new(
206+
fields.clone().into(),
207+
vec![Arc::new(ints), Arc::new(string_dict), Arc::new(int_dict)],
208+
None,
209+
)
210+
.unwrap()
211+
})
212+
.collect::<Vec<_>>();
213+
214+
let array_refs = struct_arrays
215+
.iter()
216+
.map(|a| a as &dyn Array)
217+
.collect::<Vec<_>>();
218+
219+
c.bench_function(
220+
&format!("concat struct with int32 and dicts size={batch_size} count={batch_count}"),
221+
|b| b.iter(|| bench_concat_arrays(&array_refs)),
222+
);
223+
}
181224
}
182225

183226
criterion_group!(benches, add_benchmark);

0 commit comments

Comments
 (0)