17
17
18
18
use std:: sync:: Arc ;
19
19
20
- use arrow:: util:: bench_util:: create_primitive_array;
20
+ use arrow:: util:: bench_util:: { create_primitive_array, create_string_array } ;
21
21
use criterion:: { black_box, criterion_group, criterion_main, Criterion } ;
22
22
use datafusion:: {
23
23
arrow:: {
@@ -27,14 +27,28 @@ use datafusion::{
27
27
} ,
28
28
logical_expr:: Accumulator ,
29
29
} ;
30
- use datafusion_functions_extra:: common:: mode:: PrimitiveModeAccumulator ;
30
+ use datafusion_functions_extra:: common:: mode:: { BytesModeAccumulator , PrimitiveModeAccumulator } ;
31
31
32
- fn prepare_mode_accumulator ( ) -> Box < dyn Accumulator > {
32
+ fn prepare_primitive_mode_accumulator ( ) -> Box < dyn Accumulator > {
33
33
Box :: new ( PrimitiveModeAccumulator :: < Int32Type > :: new ( & DataType :: Int32 ) )
34
34
}
35
35
36
- fn mode_bench ( c : & mut Criterion , name : & str , values : ArrayRef ) {
37
- let mut accumulator = prepare_mode_accumulator ( ) ;
36
+ fn prepare_bytes_mode_accumulator ( ) -> Box < dyn Accumulator > {
37
+ Box :: new ( BytesModeAccumulator :: new ( & DataType :: Utf8 ) )
38
+ }
39
+
40
+ fn mode_bench_primitive ( c : & mut Criterion , name : & str , values : ArrayRef ) {
41
+ let mut accumulator = prepare_primitive_mode_accumulator ( ) ;
42
+ c. bench_function ( name, |b| {
43
+ b. iter ( || {
44
+ accumulator. update_batch ( & [ values. clone ( ) ] ) . unwrap ( ) ;
45
+ black_box ( accumulator. evaluate ( ) . unwrap ( ) ) ;
46
+ } ) ;
47
+ } ) ;
48
+ }
49
+
50
+ fn mode_bench_bytes ( c : & mut Criterion , name : & str , values : ArrayRef ) {
51
+ let mut accumulator = prepare_bytes_mode_accumulator ( ) ;
38
52
c. bench_function ( name, |b| {
39
53
b. iter ( || {
40
54
accumulator. update_batch ( & [ values. clone ( ) ] ) . unwrap ( ) ;
@@ -44,17 +58,32 @@ fn mode_bench(c: &mut Criterion, name: &str, values: ArrayRef) {
44
58
}
45
59
46
60
fn mode_benchmark ( c : & mut Criterion ) {
47
- // Case: No nulls
48
- let values = Arc :: new ( create_primitive_array :: < Int32Type > ( 8192 , 0.0 ) ) as ArrayRef ;
49
- mode_bench ( c, "mode benchmark no nulls" , values) ;
61
+ let sizes = [ 100_000 , 1_000_000 ] ;
62
+ let null_percentages = [ 0.0 , 0.3 , 0.7 ] ;
50
63
51
- // Case: 30% nulls
52
- let values = Arc :: new ( create_primitive_array :: < Int32Type > ( 8192 , 0.3 ) ) as ArrayRef ;
53
- mode_bench ( c, "mode benchmark 30% nulls" , values) ;
64
+ for & size in & sizes {
65
+ for & null_percentage in & null_percentages {
66
+ let values = Arc :: new ( create_primitive_array :: < Int32Type > ( size, null_percentage) ) as ArrayRef ;
67
+ let name = format ! (
68
+ "PrimitiveModeAccumulator: {} elements, {}% nulls" ,
69
+ size,
70
+ null_percentage * 100.0
71
+ ) ;
72
+ mode_bench_primitive ( c, & name, values) ;
73
+ }
74
+ }
54
75
55
- // Case: 70% nulls
56
- let values = Arc :: new ( create_primitive_array :: < Int32Type > ( 8192 , 0.7 ) ) as ArrayRef ;
57
- mode_bench ( c, "mode benchmark 70% nulls" , values) ;
76
+ for & size in & sizes {
77
+ for & null_percentage in & null_percentages {
78
+ let values = Arc :: new ( create_string_array :: < i32 > ( size, null_percentage) ) as ArrayRef ;
79
+ let name = format ! (
80
+ "BytesModeAccumulator: {} elements, {}% nulls" ,
81
+ size,
82
+ null_percentage * 100.0
83
+ ) ;
84
+ mode_bench_bytes ( c, & name, values) ;
85
+ }
86
+ }
58
87
}
59
88
60
89
criterion_group ! ( benches, mode_benchmark) ;
0 commit comments