vortex_array/scalar_fn/fns/
is_nan.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use vortex_error::VortexResult;
5use vortex_error::vortex_bail;
6use vortex_session::registry::CachedId;
7
8use crate::dtype::DType;
9use crate::match_each_float_ptype;
10use crate::scalar_fn::EmptyOptions;
11use crate::scalar_fn::ScalarFnId;
12use crate::scalar_fn::unstable::row::RowFn;
13use crate::scalar_fn::unstable::row::RowVisitor;
14
15/// Expression that checks for NaN values.
16///
17/// The function is strict: null inputs produce null outputs. Only primitive float inputs
18/// are supported.
19#[derive(Clone)]
20pub struct IsNan;
21
22impl RowFn for IsNan {
23    type Options = EmptyOptions;
24
25    const ARG_NAMES: &'static [&'static str] = &["input"];
26    const INFALLIBLE: bool = true;
27
28    fn id(&self) -> ScalarFnId {
29        static ID: CachedId = CachedId::new("vortex.is_nan");
30        *ID
31    }
32
33    fn dispatch<V: RowVisitor>(
34        &self,
35        _options: &Self::Options,
36        args: &[DType],
37        visitor: V,
38    ) -> VortexResult<V::VisitResult> {
39        let ptype = match args {
40            [DType::Primitive(ptype, _)] if ptype.is_float() => *ptype,
41            _ => vortex_bail!("is_nan expects a single float input, got {args:?}"),
42        };
43        match_each_float_ptype!(ptype, |T| {
44            visitor.visit_bool::<(T,), false>(|(value,)| value.is_nan())
45        })
46    }
47}
48
49#[cfg(test)]
50mod tests {
51    use vortex_buffer::buffer;
52
53    use crate::IntoArray;
54    use crate::VortexSessionExecute;
55    use crate::array_session;
56    use crate::arrays::PrimitiveArray;
57    use crate::dtype::DType;
58    use crate::dtype::Nullability;
59    use crate::dtype::PType;
60    use crate::expr::get_item;
61    use crate::expr::is_nan;
62    use crate::expr::lit;
63    use crate::expr::root;
64    use crate::scalar::Scalar;
65    use crate::validity::Validity;
66
67    #[test]
68    fn dtype() {
69        assert_eq!(
70            is_nan(root())
71                .return_dtype(&DType::Primitive(PType::F32, Nullability::Nullable))
72                .unwrap(),
73            DType::Bool(Nullability::Nullable)
74        );
75        assert_eq!(
76            is_nan(root())
77                .return_dtype(&DType::Primitive(PType::F64, Nullability::NonNullable))
78                .unwrap(),
79            DType::Bool(Nullability::NonNullable)
80        );
81    }
82
83    #[test]
84    fn dtype_rejects_non_float() {
85        assert!(
86            is_nan(root())
87                .return_dtype(&DType::Primitive(PType::I32, Nullability::NonNullable))
88                .is_err()
89        );
90        assert!(
91            is_nan(root())
92                .return_dtype(&DType::Bool(Nullability::NonNullable))
93                .is_err()
94        );
95    }
96
97    #[test]
98    fn evaluate_floats() {
99        let test_array = PrimitiveArray::from_option_iter([
100            Some(1.0f32),
101            Some(f32::NAN),
102            None,
103            Some(f32::NEG_INFINITY),
104            Some(-f32::NAN),
105        ])
106        .into_array();
107        let expected = [Some(false), Some(true), None, Some(false), Some(true)];
108
109        let result = test_array.clone().apply(&is_nan(root())).unwrap();
110
111        assert_eq!(result.len(), test_array.len());
112        assert_eq!(result.dtype(), &DType::Bool(Nullability::Nullable));
113        for (i, expected_value) in expected.iter().enumerate() {
114            let expected_scalar = match expected_value {
115                Some(value) => Scalar::bool(*value, Nullability::Nullable),
116                None => Scalar::null(DType::Bool(Nullability::Nullable)),
117            };
118            assert_eq!(
119                result
120                    .execute_scalar(i, &mut array_session().create_execution_ctx())
121                    .unwrap(),
122                expected_scalar
123            );
124        }
125    }
126
127    #[test]
128    fn evaluate_all_valid_floats() {
129        let test_array = PrimitiveArray::new(
130            buffer![1.0f64, f64::NAN, f64::INFINITY, -f64::NAN],
131            Validity::NonNullable,
132        )
133        .into_array();
134        let expected = [false, true, false, true];
135
136        let result = test_array.clone().apply(&is_nan(root())).unwrap();
137
138        assert_eq!(result.len(), test_array.len());
139        assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable));
140        for (i, expected_value) in expected.iter().enumerate() {
141            assert_eq!(
142                result
143                    .execute_scalar(i, &mut array_session().create_execution_ctx())
144                    .unwrap(),
145                Scalar::bool(*expected_value, Nullability::NonNullable)
146            );
147        }
148    }
149
150    #[test]
151    fn evaluate_all_null_floats() {
152        let test_array = PrimitiveArray::from_option_iter([None::<f32>, None, None]).into_array();
153
154        let result = test_array.clone().apply(&is_nan(root())).unwrap();
155
156        assert_eq!(result.len(), test_array.len());
157        for i in 0..result.len() {
158            assert_eq!(
159                result
160                    .execute_scalar(i, &mut array_session().create_execution_ctx())
161                    .unwrap(),
162                Scalar::null(DType::Bool(Nullability::Nullable))
163            );
164        }
165    }
166
167    #[test]
168    fn evaluate_constant() {
169        let test_array = buffer![1.0f32, 2.0, 3.0].into_array();
170        let cases = [
171            (lit(f32::NAN), Some(true)),
172            (lit(1.0f32), Some(false)),
173            (
174                lit(Scalar::null(DType::Primitive(
175                    PType::F32,
176                    Nullability::Nullable,
177                ))),
178                None,
179            ),
180        ];
181
182        for (expr_child, expected_value) in cases {
183            let result = test_array.clone().apply(&is_nan(expr_child)).unwrap();
184            for i in 0..result.len() {
185                let expected_scalar = match expected_value {
186                    Some(value) => Scalar::bool(value, Nullability::Nullable),
187                    None => Scalar::null(DType::Bool(Nullability::Nullable)),
188                };
189                assert_eq!(
190                    result
191                        .execute_scalar(i, &mut array_session().create_execution_ctx())
192                        .unwrap(),
193                    expected_scalar
194                );
195            }
196        }
197    }
198
199    #[test]
200    fn evaluate_rejects_non_float() {
201        let test_array = buffer![1i32, 2, 3].into_array();
202        assert!(test_array.apply(&is_nan(root())).is_err());
203    }
204
205    #[test]
206    fn evaluate_sliced() {
207        let test_array = buffer![1.0f32, f32::NAN, 2.0, f32::NAN, 3.0]
208            .into_array()
209            .slice(1..4)
210            .unwrap();
211        let expected = [true, false, true];
212
213        let result = test_array.clone().apply(&is_nan(root())).unwrap();
214
215        assert_eq!(result.len(), test_array.len());
216        for (i, expected_value) in expected.iter().enumerate() {
217            assert_eq!(
218                result
219                    .execute_scalar(i, &mut array_session().create_execution_ctx())
220                    .unwrap(),
221                Scalar::bool(*expected_value, Nullability::NonNullable)
222            );
223        }
224    }
225
226    #[test]
227    fn test_display() {
228        let expr = is_nan(get_item("name", root()));
229        assert_eq!(expr.to_string(), "vortex.is_nan($.name)");
230
231        let expr2 = is_nan(root());
232        assert_eq!(expr2.to_string(), "vortex.is_nan($)");
233    }
234
235    #[test]
236    fn test_is_nan_is_strict() {
237        assert!(
238            is_nan(root())
239                .as_scalar()
240                .is_some_and(|f| f.signature().is_strict())
241        );
242    }
243}
vortex_array/scalar_fn/fns/is_nan.rs

vortex_array/scalar_fn/fns/
is_nan.rs
