This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23458-fa40708411f59bed7a0ce9dbf34306a63b660d9a in repository https://gitbox.apache.org/repos/asf/datafusion.git
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]
