From 63c88a5df27d8fda960580966dfc48b93ff7898c Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 10 Jul 2026 12:26:04 -0600 Subject: [PATCH 1/3] perf: optimize nanvl in datafusion-functions --- datafusion/functions/src/math/nanvl.rs | 150 ++++++++++++++++++------- 1 file changed, 111 insertions(+), 39 deletions(-) diff --git a/datafusion/functions/src/math/nanvl.rs b/datafusion/functions/src/math/nanvl.rs index b1f69032efae6..2f6e32a84d027 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,81 @@ fn scalar_is_nan(scalar: &ScalarValue) -> bool { /// - otherwise -> output is x (which may itself be NULL) fn nanvl(args: &[ArrayRef]) -> Result { match args[0].data_type() { - Float64 => { - let x = args[0].as_primitive::(); - let y = args[1].as_primitive::(); - 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::(); - let y = args[1].as_primitive::(); - let result: Float32Array = x + Float64 => Ok(Arc::new(nanvl_impl::( + args[0].as_primitive(), + args[1].as_primitive(), + ))), + Float32 => Ok(Arc::new(nanvl_impl::( + args[0].as_primitive(), + args[1].as_primitive(), + ))), + Float16 => Ok(Arc::new(nanvl_impl::( + 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`] exactly as the +/// `FromIterator` implementation does, so the result is byte-for-byte the same. +fn nanvl_impl(x: &PrimitiveArray, y: &PrimitiveArray) -> PrimitiveArray +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 = 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::::new(values.into(), None) } - Float16 => { - let x = args[0].as_primitive::(); - let y = args[1].as_primitive::(); - 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::::new(values.into(), nulls.finish()) } - other => exec_err!("Unsupported data type {other:?} for function nanvl"), } } @@ -197,7 +236,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 +274,37 @@ 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 = 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); + } } From c99084ef2a183d5752197613a1b6c56eb8d87689 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 10 Jul 2026 14:29:58 -0600 Subject: [PATCH 2/3] test: add partially-null nanvl benchmarks and unit tests Adds benchmarks for the three null-configuration match arms (only-x-null, only-y-null, both-null) and unit tests covering the single-null-input cases. These benchmarks were used to evaluate splitting the null-aware path into per-configuration arms; measurements showed that split regresses performance by 8-16%, so only the coverage additions are kept. --- datafusion/functions/benches/nanvl.rs | 74 ++++++++++++++++++++++++++ datafusion/functions/src/math/nanvl.rs | 54 +++++++++++++++++++ 2 files changed, 128 insertions(+) diff --git a/datafusion/functions/benches/nanvl.rs b/datafusion/functions/benches/nanvl.rs index 206eebd81eb81..d3d2c7ebff998 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::>(), + )); + // `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::>(), + )); + // `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::>(), + )); + 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 2f6e32a84d027..d675bd598f145 100644 --- a/datafusion/functions/src/math/nanvl.rs +++ b/datafusion/functions/src/math/nanvl.rs @@ -307,4 +307,58 @@ mod test { 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 = 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 = 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)); + } } From d261479c571c2e649b87cc98ab36d0b2b345bb24 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20Heres?= Date: Sat, 11 Jul 2026 08:36:13 +0200 Subject: [PATCH 3/3] Apply suggestion from @Dandandan --- datafusion/functions/src/math/nanvl.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/datafusion/functions/src/math/nanvl.rs b/datafusion/functions/src/math/nanvl.rs index d675bd598f145..cc146983067be 100644 --- a/datafusion/functions/src/math/nanvl.rs +++ b/datafusion/functions/src/math/nanvl.rs @@ -174,8 +174,7 @@ fn nanvl(args: &[ArrayRef]) -> Result { /// 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`] exactly as the -/// `FromIterator` implementation does, so the result is byte-for-byte the same. +/// path builds its null buffer lazily via [`NullBufferBuilder`]. fn nanvl_impl(x: &PrimitiveArray, y: &PrimitiveArray) -> PrimitiveArray where T: ArrowPrimitiveType,