From 675d3eb796559a035aafdc38434a4b8713fcfdfe Mon Sep 17 00:00:00 2001 From: kid Date: Sat, 18 Jul 2026 00:15:09 +0800 Subject: [PATCH] perf: vectorize `spark_unscaled_value` (9x faster) Replace the per-element Int64Builder loop in the array branch with the vectorized Arrow unary kernel: the Decimal128 values buffer is mapped to Int64 in a single pass and the null buffer is carried over untouched (zero-copy). The scalar branch is unchanged. Benchmark (8192 rows, Decimal128(20,2), benches/unscaled_value.rs): - no nulls: 29.5 us -> 3.2 us (~9.3x) - sparse nulls (1 in 10): 35.5 us -> 3.0 us (~11.7x) - dense nulls (1 in 2): 30.5 us -> 2.9 us (~10.4x) Part of #4942 --- native/spark-expr/Cargo.toml | 4 ++ native/spark-expr/benches/unscaled_value.rs | 62 ++++++++++++++++++ .../src/math_funcs/internal/unscaled_value.rs | 63 ++++++++++++++++--- 3 files changed, 120 insertions(+), 9 deletions(-) create mode 100644 native/spark-expr/benches/unscaled_value.rs diff --git a/native/spark-expr/Cargo.toml b/native/spark-expr/Cargo.toml index 9f98b4fcb6..0c34bb7810 100644 --- a/native/spark-expr/Cargo.toml +++ b/native/spark-expr/Cargo.toml @@ -167,3 +167,7 @@ harness = false [[bench]] name = "check_overflow" harness = false + +[[bench]] +name = "unscaled_value" +harness = false diff --git a/native/spark-expr/benches/unscaled_value.rs b/native/spark-expr/benches/unscaled_value.rs new file mode 100644 index 0000000000..e4348add40 --- /dev/null +++ b/native/spark-expr/benches/unscaled_value.rs @@ -0,0 +1,62 @@ +// 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, Decimal128Array}; +use criterion::{criterion_group, criterion_main, Criterion}; +use datafusion::physical_plan::ColumnarValue; +use datafusion_comet_spark_expr::spark_unscaled_value; +use std::hint::black_box; +use std::sync::Arc; + +/// Build a Decimal128(20, 2) column of `rows` rows, with every `null_every`-th +/// row null (`null_every == 0` means no nulls). +fn create_decimal_array(rows: usize, null_every: usize) -> ArrayRef { + let arr: Decimal128Array = (0..rows) + .map(|i| { + if null_every != 0 && i % null_every == 0 { + None + } else { + Some((i as i128 % 100_000) * 100) + } + }) + .collect::() + .with_precision_and_scale(20, 2) + .unwrap(); + Arc::new(arr) +} + +fn criterion_benchmark(c: &mut Criterion) { + let rows = 8192; + + let mut bench = |name: &str, arr: &ArrayRef| { + let args = vec![ColumnarValue::Array(Arc::clone(arr))]; + c.bench_function(name, |b| { + b.iter(|| black_box(spark_unscaled_value(black_box(&args)).unwrap())) + }); + }; + + let no_nulls = create_decimal_array(rows, 0); + let sparse_nulls = create_decimal_array(rows, 10); + let dense_nulls = create_decimal_array(rows, 2); + + bench("spark_unscaled_value: no nulls", &no_nulls); + bench("spark_unscaled_value: sparse nulls", &sparse_nulls); + bench("spark_unscaled_value: dense nulls", &dense_nulls); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/native/spark-expr/src/math_funcs/internal/unscaled_value.rs b/native/spark-expr/src/math_funcs/internal/unscaled_value.rs index a8f2e56a9f..ca25b6d920 100644 --- a/native/spark-expr/src/math_funcs/internal/unscaled_value.rs +++ b/native/spark-expr/src/math_funcs/internal/unscaled_value.rs @@ -15,10 +15,8 @@ // specific language governing permissions and limitations // under the License. -use arrow::{ - array::{AsArray, Int64Builder}, - datatypes::Decimal128Type, -}; +use arrow::array::AsArray; +use arrow::datatypes::{Decimal128Type, Int64Type}; use datafusion::common::{internal_err, Result as DataFusionResult, ScalarValue}; use datafusion::physical_plan::ColumnarValue; use std::sync::Arc; @@ -34,11 +32,58 @@ pub fn spark_unscaled_value(args: &[ColumnarValue]) -> DataFusionResult { let arr = a.as_primitive::(); - let mut result = Int64Builder::new(); - for v in arr.into_iter() { - result.append_option(v.map(|v| v as i64)); - } - Ok(ColumnarValue::Array(Arc::new(result.finish()))) + Ok(ColumnarValue::Array(Arc::new( + arr.unary::<_, Int64Type>(|v| v as i64), + ))) } } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Array, Decimal128Array, Int64Array}; + + #[test] + fn test_unscaled_value_array_with_nulls() -> DataFusionResult<()> { + let arr: Decimal128Array = + vec![Some(12345), None, Some(-678), None, Some(i64::MAX as i128)] + .into_iter() + .collect::() + .with_precision_and_scale(20, 2) + .unwrap(); + let args = vec![ColumnarValue::Array(Arc::new(arr))]; + let result = spark_unscaled_value(&args)?; + + let ColumnarValue::Array(result) = result else { + panic!("Expected array result"); + }; + let result = result.as_primitive::(); + let expected = Int64Array::from(vec![Some(12345), None, Some(-678), None, Some(i64::MAX)]); + assert_eq!(result, &expected); + assert_eq!(result.null_count(), 2); + Ok(()) + } + + #[test] + fn test_unscaled_value_scalar() -> DataFusionResult<()> { + let args = vec![ColumnarValue::Scalar(ScalarValue::Decimal128( + Some(12345), + 20, + 2, + ))]; + let result = spark_unscaled_value(&args)?; + assert!(matches!( + result, + ColumnarValue::Scalar(ScalarValue::Int64(Some(12345))) + )); + + let args = vec![ColumnarValue::Scalar(ScalarValue::Decimal128(None, 20, 2))]; + let result = spark_unscaled_value(&args)?; + assert!(matches!( + result, + ColumnarValue::Scalar(ScalarValue::Int64(None)) + )); + Ok(()) + } +}