use arrow::{
datatypes::DataType,
};
use datafusion_common::{internal_err, Result, DataFusionError};
use datafusion_expr::ColumnarValue;
use datafusion_expr::TypeSignature::*;
use datafusion_expr::{ScalarUDFImpl, Signature, Volatility};
use std::any::Any;
use std::sync::Arc;
use arrow::array::{ArrayRef, BooleanArray, Float32Array, Float64Array};
#[derive(Debug)]
pub(super) struct IsNanFunc {
signature: Signature,
}
impl IsNanFunc {
pub fn new() -> Self {
use DataType::*;
Self {
signature:
Signature::one_of(
vec![Exact(vec![Float32]), Exact(vec![Float64])],
Volatility::Immutable,
)
}
}
}
impl ScalarUDFImpl for IsNanFunc {
fn as_any(&self) -> &dyn Any {
self
}
fn name(&self) -> &str {
"isnan"
}
fn signature(&self) -> &Signature {
&self.signature
}
fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
Ok(DataType::Boolean)
}
fn invoke(&self, args: &[ColumnarValue]) -> Result<ColumnarValue> {
let args = ColumnarValue::values_to_arrays(args)?;
let arr: ArrayRef = match args[0].data_type() {
DataType::Float64 => {
Arc::new(make_function_scalar_inputs_return_type!(
&args[0],
"x",
Float64Array,
BooleanArray,
{ f64::is_nan }
))
},
DataType::Float32 => {
Arc::new(make_function_scalar_inputs_return_type!(
&args[0],
"x",
Float32Array,
BooleanArray,
{ f32::is_nan }
))
},
other => return internal_err!("Unsupported data type {other:?} for function isnan"),
};
Ok(ColumnarValue::Array(arr))
}
}