11use codspeed_criterion_compat:: { criterion_group, criterion_main, Bencher , Criterion } ;
22
3- use datafusion:: common :: ScalarValue ;
3+ use datafusion:: arrow :: datatypes :: DataType ;
44use datafusion:: logical_expr:: ColumnarValue ;
5+ use datafusion:: { common:: ScalarValue , logical_expr:: ScalarFunctionArgs } ;
56use datafusion_functions_json:: udfs:: { json_contains_udf, json_get_str_udf} ;
67
78fn bench_json_contains ( b : & mut Bencher ) {
89 let json_contains = json_contains_udf ( ) ;
9- let args = & [
10+ let args = vec ! [
1011 ColumnarValue :: Scalar ( ScalarValue :: Utf8 ( Some (
1112 r#"{"a": {"aa": "x", "ab: "y"}, "b": []}"# . to_string( ) ,
1213 ) ) ) ,
1314 ColumnarValue :: Scalar ( ScalarValue :: Utf8 ( Some ( "a" . to_string( ) ) ) ) ,
1415 ColumnarValue :: Scalar ( ScalarValue :: Utf8 ( Some ( "aa" . to_string( ) ) ) ) ,
1516 ] ;
1617
17- b. iter ( || json_contains. invoke_batch ( args, 1 ) . unwrap ( ) ) ;
18+ b. iter ( || {
19+ json_contains
20+ . invoke_with_args ( ScalarFunctionArgs {
21+ args : args. clone ( ) ,
22+ number_rows : 1 ,
23+ return_type : & DataType :: Boolean ,
24+ } )
25+ . unwrap ( )
26+ } ) ;
1827}
1928
2029fn bench_json_get_str ( b : & mut Bencher ) {
@@ -27,7 +36,15 @@ fn bench_json_get_str(b: &mut Bencher) {
2736 ColumnarValue :: Scalar ( ScalarValue :: Utf8 ( Some ( "aa" . to_string ( ) ) ) ) ,
2837 ] ;
2938
30- b. iter ( || json_get_str. invoke_batch ( args, 1 ) . unwrap ( ) ) ;
39+ b. iter ( || {
40+ json_get_str
41+ . invoke_with_args ( ScalarFunctionArgs {
42+ args : args. to_vec ( ) ,
43+ number_rows : 1 ,
44+ return_type : & DataType :: Utf8 ,
45+ } )
46+ . unwrap ( )
47+ } ) ;
3148}
3249
3350fn criterion_benchmark ( c : & mut Criterion ) {
0 commit comments