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    use vortex_error::VortexExpect as _;
53
54    use crate::IntoArray;
55    use crate::VortexSessionExecute;
56    use crate::array_session;
57    use crate::arrays::PrimitiveArray;
58    use crate::dtype::DType;
59    use crate::dtype::Nullability;
60    use crate::dtype::PType;
61    use crate::expr::get_item;
62    use crate::expr::is_nan;
63    use crate::expr::lit;
64    use crate::expr::root;
65    use crate::scalar::Scalar;
66    use crate::validity::Validity;
67
68    #[test]
69    fn dtype() {
70        assert_eq!(
71            is_nan(root())
72                .return_dtype(&DType::Primitive(PType::F32, Nullability::Nullable))
73                .unwrap(),
74            DType::Bool(Nullability::Nullable)
75        );
76        assert_eq!(
77            is_nan(root())
78                .return_dtype(&DType::Primitive(PType::F64, Nullability::NonNullable))
79                .unwrap(),
80            DType::Bool(Nullability::NonNullable)
81        );
82    }
83
84    #[test]
85    fn dtype_rejects_non_float() {
86        assert!(
87            is_nan(root())
88                .return_dtype(&DType::Primitive(PType::I32, Nullability::NonNullable))
89                .is_err()
90        );
91        assert!(
92            is_nan(root())
93                .return_dtype(&DType::Bool(Nullability::NonNullable))
94                .is_err()
95        );
96    }
97
98    #[test]
99    fn replace_children() {
100        let expr = is_nan(root());
101        expr.with_children([root()])
102            .vortex_expect("operation should succeed in test");
103    }
104
105    #[test]
106    fn evaluate_floats() {
107        let test_array = PrimitiveArray::from_option_iter([
108            Some(1.0f32),
109            Some(f32::NAN),
110            None,
111            Some(f32::NEG_INFINITY),
112            Some(-f32::NAN),
113        ])
114        .into_array();
115        let expected = [Some(false), Some(true), None, Some(false), Some(true)];
116
117        let result = test_array.clone().apply(&is_nan(root())).unwrap();
118
119        assert_eq!(result.len(), test_array.len());
120        assert_eq!(result.dtype(), &DType::Bool(Nullability::Nullable));
121        for (i, expected_value) in expected.iter().enumerate() {
122            let expected_scalar = match expected_value {
123                Some(value) => Scalar::bool(*value, Nullability::Nullable),
124                None => Scalar::null(DType::Bool(Nullability::Nullable)),
125            };
126            assert_eq!(
127                result
128                    .execute_scalar(i, &mut array_session().create_execution_ctx())
129                    .unwrap(),
130                expected_scalar
131            );
132        }
133    }
134
135    #[test]
136    fn evaluate_all_valid_floats() {
137        let test_array = PrimitiveArray::new(
138            buffer![1.0f64, f64::NAN, f64::INFINITY, -f64::NAN],
139            Validity::NonNullable,
140        )
141        .into_array();
142        let expected = [false, true, false, true];
143
144        let result = test_array.clone().apply(&is_nan(root())).unwrap();
145
146        assert_eq!(result.len(), test_array.len());
147        assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable));
148        for (i, expected_value) in expected.iter().enumerate() {
149            assert_eq!(
150                result
151                    .execute_scalar(i, &mut array_session().create_execution_ctx())
152                    .unwrap(),
153                Scalar::bool(*expected_value, Nullability::NonNullable)
154            );
155        }
156    }
157
158    #[test]
159    fn evaluate_all_null_floats() {
160        let test_array = PrimitiveArray::from_option_iter([None::<f32>, None, None]).into_array();
161
162        let result = test_array.clone().apply(&is_nan(root())).unwrap();
163
164        assert_eq!(result.len(), test_array.len());
165        for i in 0..result.len() {
166            assert_eq!(
167                result
168                    .execute_scalar(i, &mut array_session().create_execution_ctx())
169                    .unwrap(),
170                Scalar::null(DType::Bool(Nullability::Nullable))
171            );
172        }
173    }
174
175    #[test]
176    fn evaluate_constant() {
177        let test_array = buffer![1.0f32, 2.0, 3.0].into_array();
178        let cases = [
179            (lit(f32::NAN), Some(true)),
180            (lit(1.0f32), Some(false)),
181            (
182                lit(Scalar::null(DType::Primitive(
183                    PType::F32,
184                    Nullability::Nullable,
185                ))),
186                None,
187            ),
188        ];
189
190        for (expr_child, expected_value) in cases {
191            let result = test_array.clone().apply(&is_nan(expr_child)).unwrap();
192            for i in 0..result.len() {
193                let expected_scalar = match expected_value {
194                    Some(value) => Scalar::bool(value, Nullability::Nullable),
195                    None => Scalar::null(DType::Bool(Nullability::Nullable)),
196                };
197                assert_eq!(
198                    result
199                        .execute_scalar(i, &mut array_session().create_execution_ctx())
200                        .unwrap(),
201                    expected_scalar
202                );
203            }
204        }
205    }
206
207    #[test]
208    fn evaluate_rejects_non_float() {
209        let test_array = buffer![1i32, 2, 3].into_array();
210        assert!(test_array.apply(&is_nan(root())).is_err());
211    }
212
213    #[test]
214    fn evaluate_sliced() {
215        let test_array = buffer![1.0f32, f32::NAN, 2.0, f32::NAN, 3.0]
216            .into_array()
217            .slice(1..4)
218            .unwrap();
219        let expected = [true, false, true];
220
221        let result = test_array.clone().apply(&is_nan(root())).unwrap();
222
223        assert_eq!(result.len(), test_array.len());
224        for (i, expected_value) in expected.iter().enumerate() {
225            assert_eq!(
226                result
227                    .execute_scalar(i, &mut array_session().create_execution_ctx())
228                    .unwrap(),
229                Scalar::bool(*expected_value, Nullability::NonNullable)
230            );
231        }
232    }
233
234    #[test]
235    fn test_display() {
236        let expr = is_nan(get_item("name", root()));
237        assert_eq!(expr.to_string(), "vortex.is_nan($.name)");
238
239        let expr2 = is_nan(root());
240        assert_eq!(expr2.to_string(), "vortex.is_nan($)");
241    }
242
243    #[test]
244    fn test_is_nan_is_strict() {
245        assert!(
246            is_nan(root())
247                .as_scalar()
248                .is_some_and(|f| f.signature().is_strict())
249        );
250    }
251}
vortex_array/scalar_fn/fns/is_nan.rs

vortex_array/scalar_fn/fns/
is_nan.rs
