Jefffrey commented on code in PR #23458:
URL: https://github.com/apache/datafusion/pull/23458#discussion_r3562803823
##########
datafusion/functions/src/math/nanvl.rs:
##########
@@ -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<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`] exactly as the
+/// `FromIterator` implementation does, so the result is byte-for-byte the
same.
Review Comment:
probably can remove reference to previous version (like how it mentions
result is byte-for-byte the same)
##########
datafusion/functions/src/math/nanvl.rs:
##########
@@ -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<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`] exactly as the
+/// `FromIterator` implementation does, so the result is byte-for-byte the
same.
+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.
Review Comment:
same here, no need to mention original short-circuiting match considering
the code will be gone
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
