Jefffrey commented on code in PR #19919:
URL: https://github.com/apache/datafusion/pull/19919#discussion_r2712266555
##########
datafusion/functions/src/math/iszero.rs:
##########
@@ -90,79 +89,225 @@ 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)));
+ let [arg] = take_function_args(self.name(), args.args)?;
+
+ match arg {
+ ColumnarValue::Scalar(scalar) => {
+ if scalar.is_null() {
+ return
Ok(ColumnarValue::Scalar(ScalarValue::Boolean(None)));
+ }
+
+ match scalar {
+ ScalarValue::Float64(Some(v)) => {
+ Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(v
== 0.0))))
+ }
+ ScalarValue::Float32(Some(v)) => {
+ Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(v
== 0.0))))
+ }
+ ScalarValue::Float16(Some(v)) => Ok(ColumnarValue::Scalar(
+ ScalarValue::Boolean(Some(v.is_zero())),
+ )),
+ _ => {
+ internal_err!(
+ "Unexpected scalar type for iszero: {:?}",
+ scalar.data_type()
+ )
+ }
+ }
+ }
+ ColumnarValue::Array(array) => match array.data_type() {
+ Float64 =>
Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
+ array.as_primitive::<Float64Type>(),
+ |x| x == 0.0,
+ )))),
+ Float32 =>
Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
+ array.as_primitive::<Float32Type>(),
+ |x| x == 0.0,
+ )))),
+ Float16 =>
Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
+ array.as_primitive::<Float16Type>(),
+ |x| x.is_zero(),
+ )))),
+ other => {
+ internal_err!("Unexpected data type {other:?} for function
iszero")
+ }
+ },
}
- make_scalar_function(iszero, vec![])(&args.args)
}
fn documentation(&self) -> Option<&Documentation> {
self.doc()
}
}
-/// Iszero SQL function
-fn iszero(args: &[ArrayRef]) -> Result<ArrayRef> {
- match args[0].data_type() {
- Float64 => Ok(Arc::new(BooleanArray::from_unary(
- args[0].as_primitive::<Float64Type>(),
- |x| x == 0.0,
- )) as ArrayRef),
-
- Float32 => Ok(Arc::new(BooleanArray::from_unary(
- args[0].as_primitive::<Float32Type>(),
- |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"),
- }
-}
-
#[cfg(test)]
mod test {
use std::sync::Arc;
use arrow::array::{ArrayRef, Float32Array, Float64Array};
-
+ use arrow::datatypes::{DataType, Field};
+ use datafusion_common::ScalarValue;
use datafusion_common::cast::as_boolean_array;
+ use datafusion_common::config::ConfigOptions;
+ use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl};
- use crate::math::iszero::iszero;
+ use crate::math::iszero::IsZeroFunc;
#[test]
fn test_iszero_f64() {
- let args: Vec<ArrayRef> =
- vec![Arc::new(Float64Array::from(vec![1.0, 0.0, 3.0, -0.0]))];
-
- let result = iszero(&args).expect("failed to initialize function
iszero");
- let booleans =
- as_boolean_array(&result).expect("failed to initialize function
iszero");
-
- assert_eq!(booleans.len(), 4);
- assert!(!booleans.value(0));
- assert!(booleans.value(1));
- assert!(!booleans.value(2));
- assert!(booleans.value(3));
+ let array = Arc::new(Float64Array::from(vec![1.0, 0.0, 3.0, -0.0]));
+ let arg_fields = vec![Field::new("a", DataType::Float64,
false).into()];
+ let args = ScalarFunctionArgs {
+ args: vec![ColumnarValue::Array(Arc::clone(&array) as ArrayRef)],
+ arg_fields,
+ number_rows: array.len(),
+ return_field: Field::new("f", DataType::Boolean, true).into(),
+ config_options: Arc::new(ConfigOptions::default()),
+ };
+ let result = IsZeroFunc::new()
+ .invoke_with_args(args)
+ .expect("failed to initialize function iszero");
+
+ match result {
+ ColumnarValue::Array(arr) => {
+ let booleans =
+ as_boolean_array(&arr).expect("failed to convert to
BooleanArray");
+ assert_eq!(booleans.len(), 4);
+ assert!(!booleans.value(0));
+ assert!(booleans.value(1));
+ assert!(!booleans.value(2));
+ assert!(booleans.value(3));
+ }
+ ColumnarValue::Scalar(_) => panic!("Expected an array value"),
+ }
}
#[test]
fn test_iszero_f32() {
- let args: Vec<ArrayRef> =
- vec![Arc::new(Float32Array::from(vec![1.0, 0.0, 3.0, -0.0]))];
-
- let result = iszero(&args).expect("failed to initialize function
iszero");
- let booleans =
- as_boolean_array(&result).expect("failed to initialize function
iszero");
-
- assert_eq!(booleans.len(), 4);
- assert!(!booleans.value(0));
- assert!(booleans.value(1));
- assert!(!booleans.value(2));
- assert!(booleans.value(3));
+ let array = Arc::new(Float32Array::from(vec![1.0, 0.0, 3.0, -0.0]));
Review Comment:
It would be good if we can do it for all the tests there
--
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]
