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 9483db47f0 perf: vectorize `spark_unscaled_value` (9x faster) (#4972)
9483db47f0 is described below
commit 9483db47f04293b3639a9df14c044194e19ed4c5
Author: kid <[email protected]>
AuthorDate: Mon Jul 20 22:34:39 2026 +0800
perf: vectorize `spark_unscaled_value` (9x faster) (#4972)
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(-)
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::<Decimal128Array>()
+ .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<Columnar
},
ColumnarValue::Array(a) => {
let arr = a.as_primitive::<Decimal128Type>();
- 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::<Decimal128Array>()
+ .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::<arrow::datatypes::Int64Type>();
+ 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(())
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
