kumarUjjawal commented on code in PR #19919:
URL: https://github.com/apache/datafusion/pull/19919#discussion_r2712286743
##########
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:
Test | Coverage Status
-- | --
test_iszero_f64 - array [1.0, 0.0, 3.0, -0.0] | Covered by scalar.slt:828
(column test with small_floats)
test_iszero_f32 - array [1.0, 0.0, 3.0, -0.0] | Covered by scalar.slt:828
(column test with small_floats)
test_iszero_scalar_f64_zero - 0.0 | iszero(0.0) in scalar.slt:816 and
math.slt:116
test_iszero_scalar_f64_neg_zero - -0.0 | iszero(-0.0) in scalar.slt:816 and
math.slt:116
test_iszero_scalar_f64_non_zero - 1.5 | iszero(10.0) returns false in
scalar.slt:816, iszero(1.0) in math.slt:116
test_iszero_scalar_null - null | iszero(NULL) in scalar.slt:822 and
math.slt:116
--
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]
