diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 2b078268559..b8c03945585 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -225,4 +225,16 @@ harness = false [[bench]] name = "cast_int_to_decimal" +harness = false + +[[bench]] +name = "log" +harness = false + +[[bench]] +name = "pow" +harness = false + +[[bench]] +name = "abs" harness = false \ No newline at end of file diff --git a/native/spark-expr/benches/abs.rs b/native/spark-expr/benches/abs.rs new file mode 100644 index 00000000000..59dc4ca812c --- /dev/null +++ b/native/spark-expr/benches/abs.rs @@ -0,0 +1,51 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; +use datafusion::physical_plan::ColumnarValue; +use datafusion_comet_spark_expr::abs; +use std::hint::black_box; + +#[path = "common/mod.rs"] +mod common; +use common::{i64_array, NULL_RATIOS, ROW_COUNTS}; + +fn criterion_benchmark(c: &mut Criterion) { + let mut group = c.benchmark_group("abs"); + for rows in ROW_COUNTS { + for (null_ratio, tag) in NULL_RATIOS { + let arr = i64_array(rows, null_ratio, |i| { + let v = (i as i64) % 1000; + if i % 2 == 0 { + -v + } else { + v + } + }); + let args = vec![ColumnarValue::Array(arr)]; + group.bench_with_input( + BenchmarkId::from_parameter(format!("{rows}/{tag}")), + &args, + |b, args| b.iter(|| black_box(abs(black_box(args)).unwrap())), + ); + } + } + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/benches/common/mod.rs b/native/spark-expr/benches/common/mod.rs index 7b833022fe1..92b734b1a52 100644 --- a/native/spark-expr/benches/common/mod.rs +++ b/native/spark-expr/benches/common/mod.rs @@ -18,11 +18,53 @@ //! Helpers shared by the cast-from-string benchmarks, pulled in with //! `#[path = "common/mod.rs"] mod common;`. This lives in a subdirectory so that Cargo's bench //! auto-discovery, which only looks at `benches/*.rs`, does not treat it as a bench target. +#![allow(dead_code)] -use arrow::array::{builder::StringBuilder, RecordBatch}; +use arrow::array::{builder::StringBuilder, ArrayRef, Float64Array, Int64Array, RecordBatch}; use arrow::datatypes::{DataType, Field, Schema}; use std::sync::Arc; +pub const ROW_COUNTS: [usize; 3] = [8_192, 65_536, 524_288]; + +pub const NULL_RATIOS: [(f64, &str); 3] = [(0.0, "no_nulls"), (0.1, "sparse"), (1.0, "all_null")]; + +fn is_null(i: usize, null_ratio: f64) -> bool { + if null_ratio <= 0.0 { + false + } else if null_ratio >= 1.0 { + true + } else { + let stride = (1.0 / null_ratio).round() as usize; + stride != 0 && i.is_multiple_of(stride) + } +} + +pub fn f64_array(rows: usize, null_ratio: f64, value: impl Fn(usize) -> f64) -> ArrayRef { + let arr: Float64Array = (0..rows) + .map(|i| { + if is_null(i, null_ratio) { + None + } else { + Some(value(i)) + } + }) + .collect(); + Arc::new(arr) +} + +pub fn i64_array(rows: usize, null_ratio: f64, value: impl Fn(usize) -> i64) -> ArrayRef { + let arr: Int64Array = (0..rows) + .map(|i| { + if is_null(i, null_ratio) { + None + } else { + Some(value(i)) + } + }) + .collect(); + Arc::new(arr) +} + /// A single-column `Utf8` batch of `rows` rows, where row `i` holds `value(i)` unless /// `i % null_modulus == 0`, in which case it is null. pub fn string_batch( diff --git a/native/spark-expr/benches/log.rs b/native/spark-expr/benches/log.rs new file mode 100644 index 00000000000..c70c85ce0f1 --- /dev/null +++ b/native/spark-expr/benches/log.rs @@ -0,0 +1,49 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; +use datafusion::physical_plan::ColumnarValue; +use datafusion_comet_spark_expr::spark_log; +use std::hint::black_box; +use std::sync::Arc; + +#[path = "common/mod.rs"] +mod common; +use common::{f64_array, NULL_RATIOS, ROW_COUNTS}; + +fn criterion_benchmark(c: &mut Criterion) { + let mut group = c.benchmark_group("spark_log"); + for rows in ROW_COUNTS { + let base = f64_array(rows, 0.0, |_| 10.0); + for (null_ratio, tag) in NULL_RATIOS { + let value = f64_array(rows, null_ratio, |i| (i % 1000) as f64 + 1.0); + let args = vec![ + ColumnarValue::Array(Arc::clone(&base)), + ColumnarValue::Array(value), + ]; + group.bench_with_input( + BenchmarkId::from_parameter(format!("{rows}/{tag}")), + &args, + |b, args| b.iter(|| black_box(spark_log(black_box(args)).unwrap())), + ); + } + } + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/benches/pow.rs b/native/spark-expr/benches/pow.rs new file mode 100644 index 00000000000..38d32be8ed4 --- /dev/null +++ b/native/spark-expr/benches/pow.rs @@ -0,0 +1,49 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; +use datafusion::physical_plan::ColumnarValue; +use datafusion_comet_spark_expr::spark_pow; +use std::hint::black_box; +use std::sync::Arc; + +#[path = "common/mod.rs"] +mod common; +use common::{f64_array, NULL_RATIOS, ROW_COUNTS}; + +fn criterion_benchmark(c: &mut Criterion) { + let mut group = c.benchmark_group("spark_pow"); + for rows in ROW_COUNTS { + let base = f64_array(rows, 0.0, |_| 2.0); + for (null_ratio, tag) in NULL_RATIOS { + let exp = f64_array(rows, null_ratio, |i| (i % 10) as f64 + 1.0); + let args = vec![ + ColumnarValue::Array(Arc::clone(&base)), + ColumnarValue::Array(exp), + ]; + group.bench_with_input( + BenchmarkId::from_parameter(format!("{rows}/{tag}")), + &args, + |b, args| b.iter(|| black_box(spark_pow(black_box(args)).unwrap())), + ); + } + } + group.finish(); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/src/lib.rs b/native/spark-expr/src/lib.rs index c396604f7f7..1a1ce83bee0 100644 --- a/native/spark-expr/src/lib.rs +++ b/native/spark-expr/src/lib.rs @@ -84,11 +84,11 @@ pub use error::{decimal_overflow_error, SparkError, SparkErrorWithContext, Spark pub use hash_funcs::*; pub use json_funcs::{FromJson, ToJson}; pub use math_funcs::{ - checked_add, checked_div, checked_mul, checked_sub, create_modulo_expr, create_negate_expr, - spark_ceil, spark_decimal_div, spark_decimal_integral_div, spark_floor, spark_log, - spark_make_decimal, spark_pow, spark_round, spark_unhex, spark_unscaled_value, CheckOverflow, - DecimalRescaleCheckOverflow, NegativeExpr, NormalizeNaNAndZero, WideDecimalBinaryExpr, - WideDecimalOp, + abs, checked_add, checked_div, checked_mul, checked_sub, create_modulo_expr, + create_negate_expr, spark_ceil, spark_decimal_div, spark_decimal_integral_div, spark_floor, + spark_log, spark_make_decimal, spark_pow, spark_round, spark_unhex, spark_unscaled_value, + CheckOverflow, DecimalRescaleCheckOverflow, NegativeExpr, NormalizeNaNAndZero, + WideDecimalBinaryExpr, WideDecimalOp, }; pub use query_context::{create_query_context_map, QueryContext, QueryContextMap}; pub use string_funcs::*; diff --git a/native/spark-expr/src/math_funcs/mod.rs b/native/spark-expr/src/math_funcs/mod.rs index 56809a01fbc..55e78e41e6b 100644 --- a/native/spark-expr/src/math_funcs/mod.rs +++ b/native/spark-expr/src/math_funcs/mod.rs @@ -30,6 +30,7 @@ pub(crate) mod unhex; mod utils; mod wide_decimal_binary_expr; +pub use abs::abs; pub use ceil::spark_ceil; pub use checked_arithmetic::{checked_add, checked_div, checked_mul, checked_sub}; pub use div::spark_decimal_div;