Skip to main content

datafusion_functions/math/
iszero.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::sync::Arc;
19
20use arrow::array::{ArrowNativeTypeOp, AsArray, BooleanArray};
21use arrow::datatypes::DataType::{
22    Boolean, Decimal32, Decimal64, Decimal128, Decimal256, Float16, Float32, Float64,
23    Int8, Int16, Int32, Int64, Null, UInt8, UInt16, UInt32, UInt64,
24};
25use arrow::datatypes::{
26    DataType, Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, Float16Type,
27    Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, UInt8Type,
28    UInt16Type, UInt32Type, UInt64Type,
29};
30
31use datafusion_common::utils::take_function_args;
32use datafusion_common::{Result, ScalarValue, internal_err};
33use datafusion_expr::{Coercion, TypeSignatureClass};
34use datafusion_expr::{
35    ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature,
36    Volatility,
37};
38use datafusion_macros::user_doc;
39
40#[user_doc(
41    doc_section(label = "Math Functions"),
42    description = "Returns true if a given number is +0.0 or -0.0 otherwise returns false.",
43    syntax_example = "iszero(numeric_expression)",
44    sql_example = r#"```sql
45> SELECT iszero(0);
46+------------+
47| iszero(0)  |
48+------------+
49| true       |
50+------------+
51```"#,
52    standard_argument(name = "numeric_expression", prefix = "Numeric")
53)]
54#[derive(Debug, PartialEq, Eq, Hash)]
55pub struct IsZeroFunc {
56    signature: Signature,
57}
58
59impl Default for IsZeroFunc {
60    fn default() -> Self {
61        IsZeroFunc::new()
62    }
63}
64
65impl IsZeroFunc {
66    pub fn new() -> Self {
67        // Accept any numeric type (ints, uints, floats, decimals) without implicit casts.
68        let numeric = Coercion::new_exact(TypeSignatureClass::Numeric);
69        Self {
70            signature: Signature::coercible(vec![numeric], Volatility::Immutable),
71        }
72    }
73}
74
75impl ScalarUDFImpl for IsZeroFunc {
76    fn name(&self) -> &str {
77        "iszero"
78    }
79
80    fn signature(&self) -> &Signature {
81        &self.signature
82    }
83
84    fn return_type(&self, _arg_types: &[DataType]) -> Result<DataType> {
85        Ok(Boolean)
86    }
87
88    fn is_strict(&self) -> bool {
89        true
90    }
91
92    fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result<ColumnarValue> {
93        let [arg] = take_function_args(self.name(), args.args)?;
94
95        match arg {
96            ColumnarValue::Scalar(scalar) => {
97                if scalar.is_null() {
98                    return Ok(ColumnarValue::Scalar(ScalarValue::Boolean(None)));
99                }
100
101                match scalar {
102                    ScalarValue::Float64(Some(v)) => {
103                        Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(v == 0.0))))
104                    }
105                    ScalarValue::Float32(Some(v)) => {
106                        Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(v == 0.0))))
107                    }
108                    ScalarValue::Float16(Some(v)) => Ok(ColumnarValue::Scalar(
109                        ScalarValue::Boolean(Some(v.is_zero())),
110                    )),
111
112                    ScalarValue::Int8(Some(v)) => {
113                        Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(v == 0))))
114                    }
115                    ScalarValue::Int16(Some(v)) => {
116                        Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(v == 0))))
117                    }
118                    ScalarValue::Int32(Some(v)) => {
119                        Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(v == 0))))
120                    }
121                    ScalarValue::Int64(Some(v)) => {
122                        Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(v == 0))))
123                    }
124                    ScalarValue::UInt8(Some(v)) => {
125                        Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(v == 0))))
126                    }
127                    ScalarValue::UInt16(Some(v)) => {
128                        Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(v == 0))))
129                    }
130                    ScalarValue::UInt32(Some(v)) => {
131                        Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(v == 0))))
132                    }
133                    ScalarValue::UInt64(Some(v)) => {
134                        Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(v == 0))))
135                    }
136
137                    ScalarValue::Decimal32(Some(v), ..) => {
138                        Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(v == 0))))
139                    }
140                    ScalarValue::Decimal64(Some(v), ..) => {
141                        Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(v == 0))))
142                    }
143                    ScalarValue::Decimal128(Some(v), ..) => {
144                        Ok(ColumnarValue::Scalar(ScalarValue::Boolean(Some(v == 0))))
145                    }
146                    ScalarValue::Decimal256(Some(v), ..) => Ok(ColumnarValue::Scalar(
147                        ScalarValue::Boolean(Some(v.is_zero())),
148                    )),
149
150                    _ => {
151                        internal_err!(
152                            "Unexpected scalar type for iszero: {:?}",
153                            scalar.data_type()
154                        )
155                    }
156                }
157            }
158            ColumnarValue::Array(array) => match array.data_type() {
159                Null => Ok(ColumnarValue::Array(Arc::new(BooleanArray::new_null(
160                    array.len(),
161                )))),
162
163                Float64 => Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
164                    array.as_primitive::<Float64Type>(),
165                    |x| x == 0.0,
166                )))),
167                Float32 => Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
168                    array.as_primitive::<Float32Type>(),
169                    |x| x == 0.0,
170                )))),
171                Float16 => Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
172                    array.as_primitive::<Float16Type>(),
173                    |x| x.is_zero(),
174                )))),
175
176                Int8 => Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
177                    array.as_primitive::<Int8Type>(),
178                    |x| x == 0,
179                )))),
180                Int16 => Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
181                    array.as_primitive::<Int16Type>(),
182                    |x| x == 0,
183                )))),
184                Int32 => Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
185                    array.as_primitive::<Int32Type>(),
186                    |x| x == 0,
187                )))),
188                Int64 => Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
189                    array.as_primitive::<Int64Type>(),
190                    |x| x == 0,
191                )))),
192                UInt8 => Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
193                    array.as_primitive::<UInt8Type>(),
194                    |x| x == 0,
195                )))),
196                UInt16 => Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
197                    array.as_primitive::<UInt16Type>(),
198                    |x| x == 0,
199                )))),
200                UInt32 => Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
201                    array.as_primitive::<UInt32Type>(),
202                    |x| x == 0,
203                )))),
204                UInt64 => Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
205                    array.as_primitive::<UInt64Type>(),
206                    |x| x == 0,
207                )))),
208
209                Decimal32(_, _) => {
210                    Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
211                        array.as_primitive::<Decimal32Type>(),
212                        |x| x == 0,
213                    ))))
214                }
215                Decimal64(_, _) => {
216                    Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
217                        array.as_primitive::<Decimal64Type>(),
218                        |x| x == 0,
219                    ))))
220                }
221                Decimal128(_, _) => {
222                    Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
223                        array.as_primitive::<Decimal128Type>(),
224                        |x| x == 0,
225                    ))))
226                }
227                Decimal256(_, _) => {
228                    Ok(ColumnarValue::Array(Arc::new(BooleanArray::from_unary(
229                        array.as_primitive::<Decimal256Type>(),
230                        |x| x.is_zero(),
231                    ))))
232                }
233
234                other => {
235                    internal_err!("Unexpected data type {other:?} for function iszero")
236                }
237            },
238        }
239    }
240
241    fn documentation(&self) -> Option<&Documentation> {
242        self.doc()
243    }
244}