You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by ag...@apache.org on 2026/07/15 16:43:02 UTC
(datafusion-comet) branch main updated: perf: optimize try_arithmetic_kernel in datafusion-comet-spark-expr (#4910)
This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 078fed3741 perf: optimize try_arithmetic_kernel in datafusion-comet-spark-expr (#4910)
078fed3741 is described below
commit 078fed374179e39818e16188acc345bfc7263474
Author: Andy Grove <ag...@apache.org>
AuthorDate: Wed Jul 15 10:42:56 2026 -0600
perf: optimize try_arithmetic_kernel in datafusion-comet-spark-expr (#4910)
---
native/spark-expr/Cargo.toml | 4 +
native/spark-expr/benches/checked_arithmetic.rs | 153 ++++++++++++
native/spark-expr/src/lib.rs | 9 +-
.../src/math_funcs/checked_arithmetic.rs | 259 ++++++++++++++-------
native/spark-expr/src/math_funcs/mod.rs | 1 +
5 files changed, 332 insertions(+), 94 deletions(-)
diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml
index 6f90181ee0..1c0fd37f62 100644
--- a/native/spark-expr/Cargo.toml
+++ b/native/spark-expr/Cargo.toml
@@ -136,6 +136,10 @@ harness = false
name = "parse_url"
harness = false
+[[bench]]
+name = "checked_arithmetic"
+harness = false
+
[[bench]]
name = "to_json"
harness = false
diff --git a/native/spark-expr/benches/checked_arithmetic.rs b/native/spark-expr/benches/checked_arithmetic.rs
new file mode 100644
index 0000000000..280852171c
--- /dev/null
+++ b/native/spark-expr/benches/checked_arithmetic.rs
@@ -0,0 +1,153 @@
+// 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 arrow::array::{ArrayRef, Float64Array, Int32Array, Int64Array};
+use arrow::datatypes::DataType;
+use criterion::{criterion_group, criterion_main, Criterion};
+use datafusion::physical_plan::ColumnarValue;
+use datafusion_comet_spark_expr::{checked_add, checked_div, checked_mul, checked_sub, EvalMode};
+use std::hint::black_box;
+use std::sync::Arc;
+
+const NUM_ROWS: usize = 8192;
+
+/// Deterministic pseudo-random generator so the benchmark inputs are stable across runs.
+fn next(state: &mut u64) -> u64 {
+ *state ^= *state << 13;
+ *state ^= *state >> 7;
+ *state ^= *state << 17;
+ *state
+}
+
+fn i32_array(seed: u64, nulls: bool) -> ArrayRef {
+ let mut state = seed;
+ let values: Vec<Option<i32>> = (0..NUM_ROWS)
+ .map(|_| {
+ let v = (next(&mut state) % 100_000) as i32 - 50_000;
+ if nulls && next(&mut state).is_multiple_of(10) {
+ None
+ } else {
+ Some(v)
+ }
+ })
+ .collect();
+ Arc::new(Int32Array::from(values))
+}
+
+fn i64_array(seed: u64, nulls: bool) -> ArrayRef {
+ let mut state = seed;
+ let values: Vec<Option<i64>> = (0..NUM_ROWS)
+ .map(|_| {
+ let v = (next(&mut state) % 1_000_000) as i64 - 500_000;
+ if nulls && next(&mut state).is_multiple_of(10) {
+ None
+ } else {
+ Some(v)
+ }
+ })
+ .collect();
+ Arc::new(Int64Array::from(values))
+}
+
+fn f64_array(seed: u64, nulls: bool) -> ArrayRef {
+ let mut state = seed;
+ let values: Vec<Option<f64>> = (0..NUM_ROWS)
+ .map(|_| {
+ let v = (next(&mut state) % 1_000_000) as f64 / 7.0 + 1.0;
+ if nulls && next(&mut state).is_multiple_of(10) {
+ None
+ } else {
+ Some(v)
+ }
+ })
+ .collect();
+ Arc::new(Float64Array::from(values))
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+ let mut group = c.benchmark_group("checked_arithmetic");
+
+ for (label, nulls) in [("no_nulls", false), ("with_nulls", true)] {
+ let i32_args = [
+ ColumnarValue::Array(i32_array(0x1234_5678, nulls)),
+ ColumnarValue::Array(i32_array(0x9e37_79b9, nulls)),
+ ];
+ let i64_args = [
+ ColumnarValue::Array(i64_array(0x1234_5678, nulls)),
+ ColumnarValue::Array(i64_array(0x9e37_79b9, nulls)),
+ ];
+ let f64_args = [
+ ColumnarValue::Array(f64_array(0x1234_5678, nulls)),
+ ColumnarValue::Array(f64_array(0x9e37_79b9, nulls)),
+ ];
+
+ group.bench_function(format!("checked_add_i32_ansi_{label}"), |b| {
+ b.iter(|| {
+ black_box(checked_add(
+ black_box(&i32_args),
+ &DataType::Int32,
+ EvalMode::Ansi,
+ ))
+ })
+ });
+
+ group.bench_function(format!("checked_add_i32_try_{label}"), |b| {
+ b.iter(|| {
+ black_box(checked_add(
+ black_box(&i32_args),
+ &DataType::Int32,
+ EvalMode::Try,
+ ))
+ })
+ });
+
+ group.bench_function(format!("checked_sub_i64_ansi_{label}"), |b| {
+ b.iter(|| {
+ black_box(checked_sub(
+ black_box(&i64_args),
+ &DataType::Int64,
+ EvalMode::Ansi,
+ ))
+ })
+ });
+
+ group.bench_function(format!("checked_mul_i64_ansi_{label}"), |b| {
+ b.iter(|| {
+ black_box(checked_mul(
+ black_box(&i64_args),
+ &DataType::Int64,
+ EvalMode::Ansi,
+ ))
+ })
+ });
+
+ group.bench_function(format!("checked_div_f64_ansi_{label}"), |b| {
+ b.iter(|| {
+ black_box(checked_div(
+ black_box(&f64_args),
+ &DataType::Float64,
+ EvalMode::Ansi,
+ ))
+ })
+ });
+ }
+
+ 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 174a4ada9a..40422cde07 100644
--- a/native/spark-expr/src/lib.rs
+++ b/native/spark-expr/src/lib.rs
@@ -84,10 +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::{
- create_modulo_expr, create_negate_expr, spark_ceil, spark_decimal_div,
- spark_decimal_integral_div, spark_floor, spark_log, spark_make_decimal, spark_round,
- spark_unhex, spark_unscaled_value, CheckOverflow, DecimalRescaleCheckOverflow, NegativeExpr,
- NormalizeNaNAndZero, WideDecimalBinaryExpr, WideDecimalOp,
+ 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_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/checked_arithmetic.rs b/native/spark-expr/src/math_funcs/checked_arithmetic.rs
index bb4118f868..a76574e16d 100644
--- a/native/spark-expr/src/math_funcs/checked_arithmetic.rs
+++ b/native/spark-expr/src/math_funcs/checked_arithmetic.rs
@@ -15,14 +15,16 @@
// specific language governing permissions and limitations
// under the License.
-use arrow::array::{Array, ArrowNativeTypeOp, PrimitiveArray, PrimitiveBuilder};
+use arrow::array::{Array, ArrowNativeTypeOp, BooleanBufferBuilder, PrimitiveArray};
use arrow::array::{ArrayRef, AsArray};
use crate::{divide_by_zero_error, EvalMode, SparkError};
+use arrow::buffer::NullBuffer;
use arrow::datatypes::{
ArrowPrimitiveType, DataType, Float16Type, Float32Type, Float64Type, Int16Type, Int32Type,
Int64Type, Int8Type,
};
+use arrow::error::ArrowError;
use datafusion::common::DataFusionError;
use datafusion::physical_plan::ColumnarValue;
use std::sync::Arc;
@@ -36,108 +38,94 @@ pub fn try_arithmetic_kernel<T>(
where
T: ArrowPrimitiveType,
{
- let len = left.len();
- let mut builder = PrimitiveBuilder::<T>::with_capacity(len);
match op {
- "checked_add" => {
- for i in 0..len {
- if left.is_null(i) || right.is_null(i) {
- builder.append_null();
- } else {
- match left.value(i).add_checked(right.value(i)) {
- Ok(v) => builder.append_value(v),
- Err(_e) => {
- if is_ansi_mode {
- return Err(SparkError::ArithmeticOverflow {
- from_type: String::from("integer"),
- }
- .into());
- } else {
- builder.append_null();
- }
+ "checked_add" => checked_binary(left, right, is_ansi_mode, false, |l, r| l.add_checked(r)),
+ "checked_sub" => checked_binary(left, right, is_ansi_mode, false, |l, r| l.sub_checked(r)),
+ "checked_mul" => checked_binary(left, right, is_ansi_mode, false, |l, r| l.mul_checked(r)),
+ "checked_div" => checked_binary(left, right, is_ansi_mode, true, |l, r| l.div_checked(r)),
+ _ => Err(DataFusionError::Internal(format!(
+ "Unsupported operation: {:?}",
+ op
+ ))),
+ }
+}
+
+fn checked_binary<T, F>(
+ left: &PrimitiveArray<T>,
+ right: &PrimitiveArray<T>,
+ is_ansi_mode: bool,
+ is_div: bool,
+ op: F,
+) -> Result<ArrayRef, DataFusionError>
+where
+ T: ArrowPrimitiveType,
+ F: Fn(T::Native, T::Native) -> Result<T::Native, ArrowError>,
+{
+ if is_ansi_mode {
+ return arrow::compute::kernels::arity::try_binary::<_, _, _, T>(left, right, op)
+ .map(|array| Arc::new(array) as ArrayRef)
+ .map_err(|e| match e {
+ ArrowError::DivideByZero => divide_by_zero_error().into(),
+ _ => DataFusionError::from(SparkError::ArithmeticOverflow {
+ from_type: String::from("integer"),
+ }),
+ });
+ }
+
+ let len = left.len();
+ let lhs = &left.values()[..len];
+ let rhs = &right.values()[..len];
+ let nulls = NullBuffer::union(left.nulls(), right.nulls());
+
+ let mut values = vec![T::Native::default(); len];
+ let mut overflowed: Vec<usize> = Vec::new();
+
+ for (i, (out, (&l, &r))) in values.iter_mut().zip(lhs.iter().zip(rhs)).enumerate() {
+ match op(l, r) {
+ Ok(v) => *out = v,
+ Err(_) => {
+ if !is_ansi_mode {
+ overflowed.push(i);
+ } else if nulls.as_ref().is_none_or(|n| n.is_valid(i)) {
+ return if is_div && r.is_zero() {
+ Err(divide_by_zero_error().into())
+ } else {
+ Err(SparkError::ArithmeticOverflow {
+ from_type: String::from("integer"),
}
- }
+ .into())
+ };
}
}
}
- "checked_sub" => {
- for i in 0..len {
- if left.is_null(i) || right.is_null(i) {
- builder.append_null();
- } else {
- match left.value(i).sub_checked(right.value(i)) {
- Ok(v) => builder.append_value(v),
- Err(_e) => {
- if is_ansi_mode {
- return Err(SparkError::ArithmeticOverflow {
- from_type: String::from("integer"),
- }
- .into());
- } else {
- builder.append_null();
- }
- }
- }
- }
- }
+ }
+
+ let nulls = if overflowed.is_empty() {
+ nulls
+ } else {
+ let mut validity = BooleanBufferBuilder::new(len);
+ match &nulls {
+ Some(n) => validity.append_buffer(n.inner()),
+ None => validity.append_n(len, true),
}
- "checked_mul" => {
- for i in 0..len {
- if left.is_null(i) || right.is_null(i) {
- builder.append_null();
- } else {
- match left.value(i).mul_checked(right.value(i)) {
- Ok(v) => builder.append_value(v),
- Err(_e) => {
- if is_ansi_mode {
- return Err(SparkError::ArithmeticOverflow {
- from_type: String::from("integer"),
- }
- .into());
- } else {
- builder.append_null();
- }
- }
- }
- }
- }
+ for i in overflowed {
+ validity.set_bit(i, false);
}
- "checked_div" => {
- for i in 0..len {
- if left.is_null(i) || right.is_null(i) {
- builder.append_null();
- } else {
- match left.value(i).div_checked(right.value(i)) {
- Ok(v) => builder.append_value(v),
- Err(_e) => {
- if is_ansi_mode {
- return if right.value(i).is_zero() {
- Err(divide_by_zero_error().into())
- } else {
- return Err(SparkError::ArithmeticOverflow {
- from_type: String::from("integer"),
- }
- .into());
- };
- } else {
- builder.append_null();
- }
- }
- }
+ Some(NullBuffer::new(validity.finish()))
+ };
+
+ if let Some(n) = &nulls {
+ if n.null_count() > 0 {
+ for (out, valid) in values.iter_mut().zip(n.iter()) {
+ if !valid {
+ *out = T::Native::default();
}
}
}
- _ => {
- return Err(DataFusionError::Internal(format!(
- "Unsupported operation: {:?}",
- op
- )))
- }
}
- Ok(Arc::new(builder.finish()) as ArrayRef)
+ Ok(Arc::new(PrimitiveArray::<T>::new(values.into(), nulls)) as ArrayRef)
}
-
pub fn checked_add(
args: &[ColumnarValue],
data_type: &DataType,
@@ -254,3 +242,94 @@ fn checked_arithmetic_internal(
Ok(ColumnarValue::Array(result_array?))
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use arrow::array::{Float64Array, Int32Array};
+
+ fn int32_args(left: Vec<Option<i32>>, right: Vec<Option<i32>>) -> Vec<ColumnarValue> {
+ vec![
+ ColumnarValue::Array(Arc::new(Int32Array::from(left))),
+ ColumnarValue::Array(Arc::new(Int32Array::from(right))),
+ ]
+ }
+
+ fn as_int32(value: ColumnarValue) -> Int32Array {
+ let ColumnarValue::Array(array) = value else {
+ unreachable!()
+ };
+ array.as_primitive::<Int32Type>().clone()
+ }
+
+ #[test]
+ fn test_checked_add_propagates_nulls() {
+ let args = int32_args(
+ vec![Some(1), None, Some(3), None],
+ vec![Some(10), Some(20), None, None],
+ );
+ let result = as_int32(checked_add(&args, &DataType::Int32, EvalMode::Ansi).unwrap());
+ assert_eq!(result, Int32Array::from(vec![Some(11), None, None, None]));
+ }
+
+ #[test]
+ fn test_checked_add_overflow_is_null_in_try_mode() {
+ let args = int32_args(vec![Some(i32::MAX), Some(1)], vec![Some(1), Some(1)]);
+ let result = as_int32(checked_add(&args, &DataType::Int32, EvalMode::Try).unwrap());
+ assert_eq!(result, Int32Array::from(vec![None, Some(2)]));
+ }
+
+ #[test]
+ fn test_checked_add_overflow_errors_in_ansi_mode() {
+ let args = int32_args(vec![Some(i32::MAX)], vec![Some(1)]);
+ assert!(checked_add(&args, &DataType::Int32, EvalMode::Ansi).is_err());
+ }
+
+ #[test]
+ fn test_checked_sub_and_mul_overflow() {
+ let args = int32_args(vec![Some(i32::MIN), Some(5)], vec![Some(1), Some(3)]);
+ let result = as_int32(checked_sub(&args, &DataType::Int32, EvalMode::Try).unwrap());
+ assert_eq!(result, Int32Array::from(vec![None, Some(2)]));
+
+ let args = int32_args(vec![Some(i32::MAX), Some(5)], vec![Some(2), Some(3)]);
+ let result = as_int32(checked_mul(&args, &DataType::Int32, EvalMode::Try).unwrap());
+ assert_eq!(result, Int32Array::from(vec![None, Some(15)]));
+ }
+
+ #[test]
+ fn test_checked_div_by_zero() {
+ let args = vec![
+ ColumnarValue::Array(Arc::new(Float64Array::from(vec![Some(1.0), Some(8.0)]))),
+ ColumnarValue::Array(Arc::new(Float64Array::from(vec![Some(0.0), Some(2.0)]))),
+ ];
+ let ColumnarValue::Array(array) =
+ checked_div(&args, &DataType::Float64, EvalMode::Try).unwrap()
+ else {
+ unreachable!()
+ };
+ assert_eq!(
+ array.as_primitive::<Float64Type>(),
+ &Float64Array::from(vec![None, Some(4.0)])
+ );
+
+ assert!(checked_div(&args, &DataType::Float64, EvalMode::Ansi).is_err());
+ }
+
+ /// A null slot may hold any value in its underlying buffer (e.g. after a filter or a slice).
+ /// Such a row must not overflow-error in ANSI mode, and its result slot must read back as a
+ /// null holding the default value.
+ #[test]
+ fn test_null_row_with_garbage_value_does_not_error_in_ansi_mode() {
+ let nulls = NullBuffer::from(vec![false, true]);
+ let left = Int32Array::new(vec![i32::MAX, 1].into(), Some(nulls));
+ let right = Int32Array::from(vec![Some(1), Some(1)]);
+ let args = vec![
+ ColumnarValue::Array(Arc::new(left)),
+ ColumnarValue::Array(Arc::new(right)),
+ ];
+
+ let result = as_int32(checked_add(&args, &DataType::Int32, EvalMode::Ansi).unwrap());
+ assert_eq!(result, Int32Array::from(vec![None, Some(2)]));
+ assert_eq!(result.values()[0], 0);
+ }
+}
diff --git a/native/spark-expr/src/math_funcs/mod.rs b/native/spark-expr/src/math_funcs/mod.rs
index f66c584e2f..51bd016a2a 100644
--- a/native/spark-expr/src/math_funcs/mod.rs
+++ b/native/spark-expr/src/math_funcs/mod.rs
@@ -30,6 +30,7 @@ mod utils;
mod wide_decimal_binary_expr;
pub use ceil::spark_ceil;
+pub use checked_arithmetic::{checked_add, checked_div, checked_mul, checked_sub};
pub use div::spark_decimal_div;
pub use div::spark_decimal_integral_div;
pub use floor::spark_floor;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org