diff --git a/Cargo.toml b/Cargo.toml index 7b4d201..81581f2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,6 +21,8 @@ log = "0.4" datafusion-execution = "39" [dev-dependencies] +codspeed-criterion-compat = "2.3" +criterion = "0.5.1" datafusion = "39" tokio = { version = "1.37", features = ["full"] } @@ -33,3 +35,7 @@ print_stdout = "deny" pedantic = { level = "deny", priority = -1 } missing_errors_doc = "allow" cast_possible_truncation = "allow" + +[[bench]] +name = "main" +harness = false diff --git a/benches/main.rs b/benches/main.rs new file mode 100644 index 0000000..92ec395 --- /dev/null +++ b/benches/main.rs @@ -0,0 +1,38 @@ +use codspeed_criterion_compat::{criterion_group, criterion_main, Bencher, Criterion}; + +use datafusion_common::ScalarValue; +use datafusion_expr::ColumnarValue; +use datafusion_functions_json::udfs::{json_contains_udf, json_get_str_udf}; + +fn bench_json_contains(b: &mut Bencher) { + let json_contains = json_contains_udf(); + let args = &[ + ColumnarValue::Scalar(ScalarValue::Utf8(Some( + r#"{"a": {"aa": "x", "ab: "y"}, "b": []}"#.to_string(), + ))), + ColumnarValue::Scalar(ScalarValue::Utf8(Some("a".to_string()))), + ColumnarValue::Scalar(ScalarValue::Utf8(Some("aa".to_string()))), + ]; + + b.iter(|| json_contains.invoke(args).unwrap()); +} + +fn bench_json_get_str(b: &mut Bencher) { + let json_get_str = json_get_str_udf(); + let args = &[ + ColumnarValue::Scalar(ScalarValue::Utf8(Some( + r#"{"a": {"aa": "x", "ab: "y"}, "b": []}"#.to_string(), + ))), + ColumnarValue::Scalar(ScalarValue::Utf8(Some("a".to_string()))), + ColumnarValue::Scalar(ScalarValue::Utf8(Some("aa".to_string()))), + ]; + + b.iter(|| json_get_str.invoke(args).unwrap()); +} +fn criterion_benchmark(c: &mut Criterion) { + c.bench_function("json_contains", bench_json_contains); + c.bench_function("json_get_str", bench_json_get_str); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/src/lib.rs b/src/lib.rs index e5dd590..ebd9c8c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -28,6 +28,17 @@ pub mod functions { pub use crate::json_length::json_length; } +pub mod udfs { + pub use crate::json_contains::json_contains_udf; + pub use crate::json_get::json_get_udf; + pub use crate::json_get_bool::json_get_bool_udf; + pub use crate::json_get_float::json_get_float_udf; + pub use crate::json_get_int::json_get_int_udf; + pub use crate::json_get_json::json_get_json_udf; + pub use crate::json_get_str::json_get_str_udf; + pub use crate::json_length::json_length_udf; +} + /// Register all JSON UDFs pub fn register_all(registry: &mut dyn FunctionRegistry) -> Result<()> { let functions: Vec> = vec![