This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 479a0aa96d perf: optimize nanvl in datafusion-functions (#23458)
479a0aa96d is described below
commit 479a0aa96d73e73e6ea182488f8fb4b5e9485663
Author: Andy Grove <[email protected]>
AuthorDate: Sat Jul 11 01:00:19 2026 -0600
perf: optimize nanvl in datafusion-functions (#23458)
## Which issue does this PR close?
N/A
## Rationale for this change
Improve performance of existing function.
## What changes are included in this PR?
Added a null-free fast path to the nanvl array kernel that iterates the
raw value slices (generic across `Float16`/`Float32`/`Float64`) instead
of per-element `Option` iteration plus collect, eliminating
null-bookkeeping overhead on the common no-null input.
## Are these changes tested?
Existing tests + one new unit test
Benchmark (criterion):
Wins:
- nanvl_array_f64_1024: 92.229% faster (base 3250ns -> cand 252ns)
- nanvl_array_f32_1024: 95.031% faster (base 3189ns -> cand 158ns)
- nanvl_array_f64_4096: 93.826% faster (base 12679ns -> cand 782ns)
- nanvl_array_f32_4096: 96.952% faster (base 12549ns -> cand 382ns)
- nanvl_array_f64_8192: 93.375% faster (base 25792ns -> cand 1708ns)
- nanvl_array_f32_8192: 96.997% faster (base 24943ns -> cand 749ns)
Within noise:
- nanvl_scalar_f64: -0.406% faster (base 36ns -> cand 36ns)
- nanvl_scalar_f32: -1.367% faster (base 36ns -> cand 37ns)
Partially-null inputs (base = `main`, cand = this PR). Even on inputs
containing nulls the rewritten kernel is faster than `main`:
- nanvl_array_f64_x_nulls_1024: 45.9% faster (base 3496ns -> cand
1893ns)
- nanvl_array_f64_y_nulls_1024: 39.2% faster (base 3300ns -> cand
2007ns)
- nanvl_array_f64_both_nulls_1024: 47.5% faster (base 3713ns -> cand
1949ns)
- nanvl_array_f64_x_nulls_4096: 41.5% faster (base 13360ns -> cand
7818ns)
- nanvl_array_f64_y_nulls_4096: 34.5% faster (base 12887ns -> cand
8439ns)
- nanvl_array_f64_both_nulls_4096: 46.0% faster (base 14506ns -> cand
7835ns)
- nanvl_array_f64_x_nulls_8192: 47.0% faster (base 27516ns -> cand
14582ns)
- nanvl_array_f64_y_nulls_8192: 42.1% faster (base 26585ns -> cand
15398ns)
- nanvl_array_f64_both_nulls_8192: 49.0% faster (base 29647ns -> cand
15119ns)
Also added benchmarks and unit tests for partially-null inputs
(only-x-null, only-y-null, both-null). Splitting the null-aware path
into per-configuration match arms was evaluated with these benchmarks
but regressed performance by 8-16% (the larger function body penalizes
even the unchanged both-null arm), so only the coverage additions were
kept.
## Are there any user-facing changes?
No
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please add the `api
change` label.
-->
---------
Co-authored-by: Daniël Heres <[email protected]>
---
datafusion/functions/benches/nanvl.rs | 74 ++++++++++++
datafusion/functions/src/math/nanvl.rs | 203 ++++++++++++++++++++++++++-------
2 files changed, 238 insertions(+), 39 deletions(-)
diff --git a/datafusion/functions/benches/nanvl.rs
b/datafusion/functions/benches/nanvl.rs
index 206eebd81e..d3d2c7ebff 100644
--- a/datafusion/functions/benches/nanvl.rs
+++ b/datafusion/functions/benches/nanvl.rs
@@ -108,6 +108,80 @@ fn criterion_benchmark(c: &mut Criterion) {
bench.iter(||
black_box(nanvl_fn.invoke_with_args(args.clone()).unwrap()))
});
}
+
+ // Partially-null array benchmarks exercise the null-aware match arms that
+ // the fully-populated benchmarks above never reach: only-x-null,
+ // only-y-null, and both-null.
+ let bench_pair =
+ |c: &mut Criterion, name: &str, x: ArrayRef, y: ArrayRef, size: usize|
{
+ c.bench_function(name, |bench| {
+ let args = ScalarFunctionArgs {
+ args: vec![
+ ColumnarValue::Array(Arc::clone(&x)),
+ ColumnarValue::Array(Arc::clone(&y)),
+ ],
+ arg_fields: vec![
+ Field::new("a", DataType::Float64, true).into(),
+ Field::new("b", DataType::Float64, true).into(),
+ ],
+ number_rows: size,
+ return_field: Field::new("f", DataType::Float64,
true).into(),
+ config_options: Arc::clone(&config_options),
+ };
+ bench.iter(||
black_box(nanvl_fn.invoke_with_args(args.clone()).unwrap()))
+ });
+ };
+
+ for size in [1024, 4096, 8192] {
+ // `x` mixes non-NaN, NaN, and null so every code path is taken.
+ let x_nulls: ArrayRef = Arc::new(Float64Array::from(
+ (0..size)
+ .map(|i| match i % 3 {
+ 0 => Some(1.0),
+ 1 => Some(f64::NAN),
+ _ => None,
+ })
+ .collect::<Vec<_>>(),
+ ));
+ // `x` without nulls, alternating non-NaN and NaN.
+ let x_full: ArrayRef = Arc::new(Float64Array::from(
+ (0..size)
+ .map(|i| if i % 2 == 0 { 1.0 } else { f64::NAN })
+ .collect::<Vec<_>>(),
+ ));
+ // `y` with roughly a quarter nulls.
+ let y_nulls: ArrayRef = Arc::new(Float64Array::from(
+ (0..size)
+ .map(|i| if i % 4 == 3 { None } else { Some(2.0) })
+ .collect::<Vec<_>>(),
+ ));
+ let y_full: ArrayRef = Arc::new(Float64Array::from(vec![2.0; size]));
+
+ // (Some, None): only `x` has nulls.
+ bench_pair(
+ c,
+ &format!("nanvl/array_f64_x_nulls/{size}"),
+ Arc::clone(&x_nulls),
+ Arc::clone(&y_full),
+ size,
+ );
+ // (None, Some): only `y` has nulls.
+ bench_pair(
+ c,
+ &format!("nanvl/array_f64_y_nulls/{size}"),
+ Arc::clone(&x_full),
+ Arc::clone(&y_nulls),
+ size,
+ );
+ // (Some, Some): both inputs have nulls.
+ bench_pair(
+ c,
+ &format!("nanvl/array_f64_both_nulls/{size}"),
+ Arc::clone(&x_nulls),
+ Arc::clone(&y_nulls),
+ size,
+ );
+ }
}
criterion_group!(benches, criterion_benchmark);
diff --git a/datafusion/functions/src/math/nanvl.rs
b/datafusion/functions/src/math/nanvl.rs
index b1f69032ef..cc14698306 100644
--- a/datafusion/functions/src/math/nanvl.rs
+++ b/datafusion/functions/src/math/nanvl.rs
@@ -17,9 +17,12 @@
use std::sync::Arc;
-use arrow::array::{ArrayRef, AsArray, Float16Array, Float32Array,
Float64Array};
+use arrow::array::builder::NullBufferBuilder;
+use arrow::array::{Array, ArrayRef, AsArray, PrimitiveArray};
use arrow::datatypes::DataType::{Float16, Float32, Float64};
-use arrow::datatypes::{DataType, Float16Type, Float32Type, Float64Type};
+use arrow::datatypes::{
+ ArrowPrimitiveType, DataType, Float16Type, Float32Type, Float64Type,
+};
use datafusion_common::types::{NativeType, logical_float64};
use datafusion_common::{Result, ScalarValue, exec_err,
utils::take_function_args};
use datafusion_expr::{
@@ -27,6 +30,7 @@ use datafusion_expr::{
TypeSignature, TypeSignatureClass, Volatility,
};
use datafusion_macros::user_doc;
+use num_traits::Float;
#[user_doc(
doc_section(label = "Math Functions"),
@@ -148,46 +152,80 @@ fn scalar_is_nan(scalar: &ScalarValue) -> bool {
/// - otherwise -> output is x (which may itself be NULL)
fn nanvl(args: &[ArrayRef]) -> Result<ArrayRef> {
match args[0].data_type() {
- Float64 => {
- let x = args[0].as_primitive::<Float64Type>();
- let y = args[1].as_primitive::<Float64Type>();
- let result: Float64Array = x
- .iter()
- .zip(y.iter())
- .map(|(x_value, y_value)| match x_value {
- Some(x_value) if x_value.is_nan() => y_value,
- _ => x_value,
- })
- .collect();
- Ok(Arc::new(result) as ArrayRef)
- }
- Float32 => {
- let x = args[0].as_primitive::<Float32Type>();
- let y = args[1].as_primitive::<Float32Type>();
- let result: Float32Array = x
+ Float64 => Ok(Arc::new(nanvl_impl::<Float64Type>(
+ args[0].as_primitive(),
+ args[1].as_primitive(),
+ ))),
+ Float32 => Ok(Arc::new(nanvl_impl::<Float32Type>(
+ args[0].as_primitive(),
+ args[1].as_primitive(),
+ ))),
+ Float16 => Ok(Arc::new(nanvl_impl::<Float16Type>(
+ args[0].as_primitive(),
+ args[1].as_primitive(),
+ ))),
+ other => exec_err!("Unsupported data type {other:?} for function
nanvl"),
+ }
+}
+
+/// Element-wise `nanvl`: selects `y[i]` where `x[i]` is `NaN`, otherwise
`x[i]`
+/// (a null `x` selects `x`, i.e. propagates null).
+///
+/// This produces output identical to collecting an iterator of `Option`s but
+/// splits out a null-free fast path that iterates the raw value slices,
+/// skipping per-element validity checks and `Option` handling. The null-aware
+/// path builds its null buffer lazily via [`NullBufferBuilder`].
+fn nanvl_impl<T>(x: &PrimitiveArray<T>, y: &PrimitiveArray<T>) ->
PrimitiveArray<T>
+where
+ T: ArrowPrimitiveType,
+ T::Native: Float,
+{
+ let xv = x.values();
+ let yv = y.values();
+
+ match (x.nulls(), y.nulls()) {
+ // No nulls in either input means no nulls in the output, so we can
+ // iterate values directly and avoid the null bookkeeping entirely.
+ (None, None) => {
+ let values: Vec<T::Native> = xv
.iter()
- .zip(y.iter())
- .map(|(x_value, y_value)| match x_value {
- Some(x_value) if x_value.is_nan() => y_value,
- _ => x_value,
- })
+ .zip(yv.iter())
+ .map(
+ |(&x_value, &y_value)| {
+ if x_value.is_nan() { y_value } else { x_value }
+ },
+ )
.collect();
- Ok(Arc::new(result) as ArrayRef)
+ PrimitiveArray::<T>::new(values.into(), None)
}
- Float16 => {
- let x = args[0].as_primitive::<Float16Type>();
- let y = args[1].as_primitive::<Float16Type>();
- let result: Float16Array = x
- .iter()
- .zip(y.iter())
- .map(|(x_value, y_value)| match x_value {
- Some(x_value) if x_value.is_nan() => y_value,
- _ => x_value,
- })
- .collect();
- Ok(Arc::new(result) as ArrayRef)
+ _ => {
+ let len = x.len();
+ let mut nulls = NullBufferBuilder::new(len);
+ let mut values = Vec::with_capacity(len);
+ for i in 0..len {
+ // `y` is only consulted when `x` is a (non-null) NaN, matching
+ // the original short-circuiting match.
+ if x.is_valid(i) {
+ let x_value = xv[i];
+ if x_value.is_nan() {
+ if y.is_valid(i) {
+ values.push(yv[i]);
+ nulls.append_non_null();
+ } else {
+ values.push(T::Native::default());
+ nulls.append_null();
+ }
+ } else {
+ values.push(x_value);
+ nulls.append_non_null();
+ }
+ } else {
+ values.push(T::Native::default());
+ nulls.append_null();
+ }
+ }
+ PrimitiveArray::<T>::new(values.into(), nulls.finish())
}
- other => exec_err!("Unsupported data type {other:?} for function
nanvl"),
}
}
@@ -197,7 +235,7 @@ mod test {
use crate::math::nanvl::nanvl;
- use arrow::array::{ArrayRef, Float32Array, Float64Array};
+ use arrow::array::{Array, ArrayRef, Float32Array, Float64Array};
use datafusion_common::cast::{as_float32_array, as_float64_array};
#[test]
@@ -235,4 +273,91 @@ mod test {
assert_eq!(floats.value(2), 3.0);
assert!(floats.value(3).is_nan());
}
+
+ #[test]
+ fn test_nanvl_f64_with_nulls() {
+ // Covers the null-aware path and null propagation:
+ // - x null -> null (regardless of y)
+ // - x NaN, y non-null -> y
+ // - x NaN, y null -> null
+ // - x non-NaN -> x
+ let args: Vec<ArrayRef> = vec![
+ Arc::new(Float64Array::from(vec![
+ None,
+ Some(f64::NAN),
+ Some(f64::NAN),
+ Some(2.5),
+ ])), // x
+ Arc::new(Float64Array::from(vec![
+ Some(9.0),
+ Some(6.0),
+ None,
+ Some(7.0),
+ ])), // y
+ ];
+
+ let result = nanvl(&args).expect("failed to initialize function
nanvl");
+ let floats =
+ as_float64_array(&result).expect("failed to initialize function
nanvl");
+
+ assert_eq!(floats.len(), 4);
+ assert!(floats.is_null(0));
+ assert_eq!(floats.value(1), 6.0);
+ assert!(floats.is_null(2));
+ assert_eq!(floats.value(3), 2.5);
+ }
+
+ #[test]
+ fn test_nanvl_f64_only_y_nulls() {
+ // `x` has no nulls, `y` does:
+ // - x non-NaN -> x
+ // - x NaN, y non-null -> y
+ // - x NaN, y null -> null (propagated from y)
+ let args: Vec<ArrayRef> = vec![
+ Arc::new(Float64Array::from(vec![1.0, f64::NAN, f64::NAN, 4.0])),
// x
+ Arc::new(Float64Array::from(vec![
+ Some(5.0),
+ Some(6.0),
+ None,
+ Some(8.0),
+ ])), // y
+ ];
+
+ let result = nanvl(&args).expect("failed to initialize function
nanvl");
+ let floats =
+ as_float64_array(&result).expect("failed to initialize function
nanvl");
+
+ assert_eq!(floats.len(), 4);
+ assert_eq!(floats.value(0), 1.0);
+ assert_eq!(floats.value(1), 6.0);
+ assert!(floats.is_null(2));
+ assert_eq!(floats.value(3), 4.0);
+ }
+
+ #[test]
+ fn test_nanvl_f64_only_x_nulls() {
+ // `x` has nulls, `y` does not:
+ // - x null -> null (propagated from x)
+ // - x NaN -> y
+ // - x non-NaN -> x
+ let args: Vec<ArrayRef> = vec![
+ Arc::new(Float64Array::from(vec![
+ None,
+ Some(f64::NAN),
+ Some(3.0),
+ None,
+ ])), // x
+ Arc::new(Float64Array::from(vec![5.0, 6.0, 7.0, 8.0])), // y
+ ];
+
+ let result = nanvl(&args).expect("failed to initialize function
nanvl");
+ let floats =
+ as_float64_array(&result).expect("failed to initialize function
nanvl");
+
+ assert_eq!(floats.len(), 4);
+ assert!(floats.is_null(0));
+ assert_eq!(floats.value(1), 6.0);
+ assert_eq!(floats.value(2), 3.0);
+ assert!(floats.is_null(3));
+ }
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
