Skip to content

refactor: Use Signature::coercible for isnan/iszero #19604

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jefffrey merged 9 commits into from
Jan 7, 2026
Merged

refactor: Use Signature::coercible for isnan/iszero #19604

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
6b1b0e4
refactor: Use TypeSignatureClass::Float for isnan/iszero/nanvl signat…
kumarUjjawal Jan 2, 2026
f4eb9cb
change nanvl for consistency
kumarUjjawal Jan 2, 2026
7b482aa
fix sql test error
kumarUjjawal Jan 2, 2026
f2b196e
fix panic
kumarUjjawal Jan 2, 2026
bd86ee5
fix issues with nanvl
kumarUjjawal Jan 2, 2026
c6a65fe
switch to new_implicit
kumarUjjawal Jan 2, 2026
8dd92eb
add float16 support
kumarUjjawal Jan 4, 2026
1b1dfa6
suggestion from jefffrey
kumarUjjawal Jan 5, 2026
28a242b
reverted coercible
kumarUjjawal Jan 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 22 additions & 10 deletions datafusion/functions/src/math/iszero.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,12 +18,13 @@
use std::any::Any;
use std::sync::Arc;

use arrow::array::{ArrayRef, AsArray, BooleanArray};
use arrow::datatypes::DataType::{Boolean, Float32, Float64};
use arrow::datatypes::{DataType, Float32Type, Float64Type};
use arrow::array::{ArrayRef, ArrowNativeTypeOp, AsArray, BooleanArray};
use arrow::datatypes::DataType::{Boolean, Float16, Float32, Float64};
use arrow::datatypes::{DataType, Float16Type, Float32Type, Float64Type};

use datafusion_common::{Result, exec_err};
use datafusion_expr::TypeSignature::Exact;
use datafusion_common::types::NativeType;
use datafusion_common::{Result, ScalarValue, exec_err};
use datafusion_expr::{Coercion, TypeSignatureClass};
use datafusion_expr::{
ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
Volatility,
Expand Down Expand Up @@ -59,12 +60,14 @@ impl Default for IsZeroFunc {

impl IsZeroFunc {
pub fn new() -> Self {
use DataType::*;
// Accept any numeric type and coerce to float
let float = Coercion::new_implicit(
TypeSignatureClass::Float,
Copy link
Member

@martin-g martin-g Jan 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This also includes Float16 but fn iszero() below currently handles only Float32 and Float64.

vec![TypeSignatureClass::Numeric],
NativeType::Float64,
);
Self {
signature: Signature::one_of(
vec![Exact(vec![Float32]), Exact(vec![Float64])],
Volatility::Immutable,
),
signature: Signature::coercible(vec![float], Volatility::Immutable),
}
}
}
Expand All @@ -87,6 +90,10 @@ impl ScalarUDFImpl for IsZeroFunc {
}

fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
// Handle NULL input
if args.args[0].data_type().is_null() {
return Ok(ColumnarValue::Scalar(ScalarValue::Boolean(None)));
}
make_scalar_function(iszero, vec![])(&args.args)
}

Expand All @@ -108,6 +115,11 @@ fn iszero(args: &[ArrayRef]) -> Result<ArrayRef> {
|x| x == 0.0,
)) as ArrayRef),

Float16 => Ok(Arc::new(BooleanArray::from_unary(
args[0].as_primitive::<Float16Type>(),
|x| x.is_zero(),
)) as ArrayRef),

other => exec_err!("Unsupported data type {other:?} for function iszero"),
}
}
Expand Down
32 changes: 21 additions & 11 deletions datafusion/functions/src/math/nans.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,10 @@

//! Math function: `isnan()`.

use arrow::datatypes::{DataType, Float32Type, Float64Type};
use datafusion_common::{Result, exec_err};
use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, TypeSignature};
use arrow::datatypes::{DataType, Float16Type, Float32Type, Float64Type};
use datafusion_common::types::NativeType;
use datafusion_common::{Result, ScalarValue, exec_err};
use datafusion_expr::{Coercion, ColumnarValue, ScalarFunctionArgs, TypeSignatureClass};

use arrow::array::{ArrayRef, AsArray, BooleanArray};
use datafusion_expr::{Documentation, ScalarUDFImpl, Signature, Volatility};
Expand Down Expand Up @@ -54,15 +55,14 @@ impl Default for IsNanFunc {

impl IsNanFunc {
pub fn new() -> Self {
use DataType::*;
// Accept any numeric type and coerce to float
let float = Coercion::new_implicit(
TypeSignatureClass::Float,
vec![TypeSignatureClass::Numeric],
NativeType::Float64,
);
Self {
signature: Signature::one_of(
vec![
TypeSignature::Exact(vec![Float32]),
TypeSignature::Exact(vec![Float64]),
],
Volatility::Immutable,
),
signature: Signature::coercible(vec![float], Volatility::Immutable),
}
}
}
Expand All @@ -84,6 +84,11 @@ impl ScalarUDFImpl for IsNanFunc {
}

fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
// Handle NULL input
if args.args[0].data_type().is_null() {
return Ok(ColumnarValue::Scalar(ScalarValue::Boolean(None)));
}

let args = ColumnarValue::values_to_arrays(&args.args)?;

let arr: ArrayRef = match args[0].data_type() {
Expand All @@ -96,6 +101,11 @@ impl ScalarUDFImpl for IsNanFunc {
args[0].as_primitive::<Float32Type>(),
f32::is_nan,
)) as ArrayRef,

DataType::Float16 => Arc::new(BooleanArray::from_unary(
args[0].as_primitive::<Float16Type>(),
|x| x.is_nan(),
)) as ArrayRef,
other => {
return exec_err!(
"Unsupported data type {other:?} for function {}",
Expand Down
27 changes: 22 additions & 5 deletions datafusion/functions/src/math/nanvl.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,9 @@ use std::sync::Arc;

use crate::utils::make_scalar_function;

use arrow::array::{ArrayRef, AsArray, Float32Array, Float64Array};
use arrow::datatypes::DataType::{Float32, Float64};
use arrow::datatypes::{DataType, Float32Type, Float64Type};
use arrow::array::{ArrayRef, AsArray, Float16Array, Float32Array, Float64Array};
use arrow::datatypes::DataType::{Float16, Float32, Float64};
use arrow::datatypes::{DataType, Float16Type, Float32Type, Float64Type};
use datafusion_common::{DataFusionError, Result, exec_err};
use datafusion_expr::TypeSignature::Exact;
use datafusion_expr::{
Expand Down Expand Up @@ -66,10 +66,13 @@ impl Default for NanvlFunc {

impl NanvlFunc {
pub fn new() -> Self {
use DataType::*;
Self {
signature: Signature::one_of(
vec![Exact(vec![Float32, Float32]), Exact(vec![Float64, Float64])],
vec![
Exact(vec![Float16, Float16]),
Exact(vec![Float32, Float32]),
Exact(vec![Float64, Float64]),
],
Volatility::Immutable,
),
}
Expand All @@ -91,6 +94,7 @@ impl ScalarUDFImpl for NanvlFunc {

fn return_type(&self, arg_types: &[DataType]) -> Result<DataType> {
match &arg_types[0] {
Float16 => Ok(Float16),
Float32 => Ok(Float32),
_ => Ok(Float64),
}
Expand Down Expand Up @@ -130,6 +134,19 @@ fn nanvl(args: &[ArrayRef]) -> Result<ArrayRef> {
.map(|res| Arc::new(res) as _)
.map_err(DataFusionError::from)
}
Float16 => {
let compute_nanvl =
|x: <Float16Type as arrow::datatypes::ArrowPrimitiveType>::Native,
y: <Float16Type as arrow::datatypes::ArrowPrimitiveType>::Native| {
if x.is_nan() { y } else { x }
};

let x = args[0].as_primitive() as &Float16Array;
let y = args[1].as_primitive() as &Float16Array;
arrow::compute::binary::<_, _, _, Float16Type>(x, y, compute_nanvl)
.map(|res| Arc::new(res) as _)
.map_err(DataFusionError::from)
}
other => exec_err!("Unsupported data type {other:?} for function nanvl"),
}
}
Expand Down