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 std:: sync:: Arc ;
19+ use arrow:: array:: { BooleanBuilder , RecordBatch } ;
20+ use arrow:: datatypes:: { DataType , Field , Schema } ;
21+ use criterion:: { criterion_group, criterion_main, Criterion } ;
22+ use datafusion:: physical_expr:: expressions:: Column ;
23+ use datafusion:: physical_expr:: PhysicalExpr ;
24+ use datafusion_comet_spark_expr:: { Cast , EvalMode , SparkCastOptions } ;
25+
26+ fn criterion_benchmark ( c : & mut Criterion ) {
27+ let expr = Arc :: new ( Column :: new ( "a" , 0 ) ) ;
28+ let boolean_batch = create_boolean_batch ( ) ;
29+ let spark_cast_options = SparkCastOptions :: new ( EvalMode :: Legacy , "UTC" , false ) ;
30+ Cast :: new ( expr. clone ( ) , DataType :: Int8 , spark_cast_options. clone ( ) ) ;
31+ let cast_to_i8 = Cast :: new ( expr. clone ( ) , DataType :: Boolean , spark_cast_options. clone ( ) ) ;
32+ let cast_to_i16 = Cast :: new ( expr. clone ( ) , DataType :: Boolean , spark_cast_options. clone ( ) ) ;
33+ let cast_to_i32 = Cast :: new ( expr. clone ( ) , DataType :: Boolean , spark_cast_options. clone ( ) ) ;
34+ let cast_to_i64 = Cast :: new ( expr. clone ( ) , DataType :: Boolean , spark_cast_options) ;
35+
36+ let mut group = c. benchmark_group ( format ! ( "cast_bool_to_int" ) ) ;
37+ group. bench_function ( "i8" , |b| {
38+ b. iter ( || cast_to_i8. evaluate ( & boolean_batch) . unwrap ( ) ) ;
39+ } ) ;
40+ group. bench_function ( "i16" , |b| {
41+ b. iter ( || cast_to_i16. evaluate ( & boolean_batch) . unwrap ( ) ) ;
42+ } ) ;
43+ group. bench_function ( "i32" , |b| {
44+ b. iter ( || cast_to_i32. evaluate ( & boolean_batch) . unwrap ( ) ) ;
45+ } ) ;
46+ group. bench_function ( "i64" , |b| {
47+ b. iter ( || cast_to_i64. evaluate ( & boolean_batch) . unwrap ( ) ) ;
48+ } ) ;
49+ }
50+
51+ fn create_boolean_batch ( ) -> RecordBatch {
52+ let schema = Arc :: new ( Schema :: new ( vec ! [ Field :: new( "a" , DataType :: Boolean , true ) ] ) ) ;
53+ let mut b = BooleanBuilder :: with_capacity ( 1000 ) ;
54+ for i in 0 ..1000 {
55+ if i % 10 == 0 {
56+ b. append_null ( ) ;
57+ } else {
58+ b. append_value ( rand:: random :: < bool > ( ) ) ;
59+ }
60+ }
61+ let array = b. finish ( ) ;
62+ RecordBatch :: try_new ( schema, vec ! [ Arc :: new( array) ] ) . unwrap ( )
63+ }
64+
65+ fn config ( ) -> Criterion {
66+ Criterion :: default ( )
67+ }
68+
69+ criterion_group ! {
70+ name = benches;
71+ config = config( ) ;
72+ targets = criterion_benchmark
73+ }
74+ criterion_main ! ( benches) ;
0 commit comments