Skip to main content

vortex_array/scalar_fn/fns/
is_not_null.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::fmt::Formatter;
6
7use vortex_error::VortexExpect as _;
8use vortex_error::VortexResult;
9use vortex_session::VortexSession;
10use vortex_session::registry::CachedId;
11
12use crate::ArrayRef;
13use crate::ExecutionCtx;
14use crate::IntoArray;
15use crate::arrays::ConstantArray;
16use crate::arrays::ScalarFnArray;
17use crate::dtype::DType;
18use crate::dtype::Nullability;
19use crate::expr::display::ExprDisplay;
20use crate::scalar_fn::Arity;
21use crate::scalar_fn::ChildName;
22use crate::scalar_fn::EmptyOptions;
23use crate::scalar_fn::ExecutionArgs;
24use crate::scalar_fn::ScalarFnId;
25use crate::scalar_fn::ScalarFnVTable;
26use crate::scalar_fn::ScalarFnVTableExt;
27use crate::validity::Validity;
28
29/// Expression that checks for non-null values.
30#[derive(Clone)]
31pub struct IsNotNull;
32
33impl IsNotNull {
34    /// Creates a lazy non-null check over `input`.
35    #[expect(clippy::new_ret_no_self, reason = "constructs the lazy result array")]
36    pub fn new(input: ArrayRef) -> ScalarFnArray {
37        ScalarFnArray::try_new(IsNotNull.bind(EmptyOptions), vec![input])
38            .vortex_expect("IsNotNull has one child and an infallible return dtype")
39    }
40}
41
42impl ScalarFnVTable for IsNotNull {
43    type Options = EmptyOptions;
44
45    fn id(&self) -> ScalarFnId {
46        static ID: CachedId = CachedId::new("vortex.is_not_null");
47        *ID
48    }
49
50    fn serialize(&self, _instance: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
51        Ok(Some(vec![]))
52    }
53
54    fn deserialize(
55        &self,
56        _metadata: &[u8],
57        _session: &VortexSession,
58    ) -> VortexResult<Self::Options> {
59        Ok(EmptyOptions)
60    }
61
62    fn arity(&self, _options: &Self::Options) -> Arity {
63        Arity::Exact(1)
64    }
65
66    fn child_name(&self, _instance: &Self::Options, child_idx: usize) -> ChildName {
67        match child_idx {
68            0 => ChildName::from("input"),
69            _ => unreachable!("Invalid child index {} for IsNotNull expression", child_idx),
70        }
71    }
72
73    fn fmt_sql(
74        &self,
75        _options: &Self::Options,
76        expr: &dyn ExprDisplay,
77        f: &mut Formatter<'_>,
78    ) -> std::fmt::Result {
79        write!(f, "is_not_null(")?;
80        Display::fmt(expr.display_child(0), f)?;
81        write!(f, ")")
82    }
83
84    fn return_dtype(&self, _options: &Self::Options, _arg_dtypes: &[DType]) -> VortexResult<DType> {
85        Ok(DType::Bool(Nullability::NonNullable))
86    }
87
88    fn execute(
89        &self,
90        _data: &Self::Options,
91        args: &dyn ExecutionArgs,
92        _ctx: &mut ExecutionCtx,
93    ) -> VortexResult<ArrayRef> {
94        let child = args.get(0)?;
95        match child.validity()? {
96            Validity::NonNullable | Validity::AllValid => {
97                Ok(ConstantArray::new(true, args.row_count()).into_array())
98            }
99            Validity::AllInvalid => Ok(ConstantArray::new(false, args.row_count()).into_array()),
100            Validity::Array(a) => Ok(a),
101        }
102    }
103
104    fn is_strict(&self, _instance: &Self::Options) -> bool {
105        // Null input produces the non-null boolean value `false`.
106        false
107    }
108
109    fn is_infallible(&self, _instance: &Self::Options) -> bool {
110        true
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use std::sync::LazyLock;
117
118    use vortex_buffer::buffer;
119    use vortex_error::VortexExpect as _;
120    use vortex_error::VortexResult;
121    use vortex_session::VortexSession;
122
123    use crate::IntoArray;
124    use crate::VortexSessionExecute;
125    use crate::array_session;
126    use crate::arrays::PrimitiveArray;
127    use crate::arrays::StructArray;
128    use crate::dtype::DType;
129    use crate::dtype::Nullability;
130    use crate::expr::col;
131    use crate::expr::eq;
132    use crate::expr::get_item;
133    use crate::expr::is_not_null;
134    use crate::expr::or;
135    use crate::expr::root;
136    use crate::expr::test_harness;
137    use crate::scalar::Scalar;
138    use crate::scalar_fn::EmptyOptions;
139    use crate::scalar_fn::ScalarFnVTableExt;
140    use crate::scalar_fn::internal::row_count::RowCount;
141    use crate::stats::StatsSession;
142    use crate::stats::all_null;
143    use crate::stats::null_count;
144
145    static STATS_SESSION: LazyLock<VortexSession> =
146        LazyLock::new(|| VortexSession::empty().with::<StatsSession>());
147
148    #[test]
149    fn dtype() {
150        let dtype = test_harness::struct_dtype();
151        assert_eq!(
152            is_not_null(root()).return_dtype(&dtype).unwrap(),
153            DType::Bool(Nullability::NonNullable)
154        );
155    }
156
157    #[test]
158    fn replace_children() {
159        let expr = is_not_null(root());
160        expr.with_children([root()])
161            .vortex_expect("operation should succeed in test");
162    }
163
164    #[test]
165    fn evaluate_mask() {
166        let test_array =
167            PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)])
168                .into_array();
169        let expected = [true, false, true, false, true];
170
171        let result = test_array.clone().apply(&is_not_null(root())).unwrap();
172
173        assert_eq!(result.len(), test_array.len());
174        assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable));
175
176        for (i, expected_value) in expected.iter().enumerate() {
177            assert_eq!(
178                result
179                    .execute_scalar(i, &mut array_session().create_execution_ctx())
180                    .unwrap(),
181                Scalar::bool(*expected_value, Nullability::NonNullable)
182            );
183        }
184    }
185
186    #[test]
187    fn evaluate_all_true() {
188        let test_array = buffer![1, 2, 3, 4, 5].into_array();
189
190        let result = test_array.clone().apply(&is_not_null(root())).unwrap();
191
192        assert_eq!(result.len(), test_array.len());
193        for i in 0..result.len() {
194            assert_eq!(
195                result
196                    .execute_scalar(i, &mut array_session().create_execution_ctx())
197                    .unwrap(),
198                Scalar::bool(true, Nullability::NonNullable)
199            );
200        }
201    }
202
203    #[test]
204    fn evaluate_all_false() {
205        let test_array =
206            PrimitiveArray::from_option_iter(vec![None::<i32>, None, None, None, None])
207                .into_array();
208
209        let result = test_array.clone().apply(&is_not_null(root())).unwrap();
210
211        assert_eq!(result.len(), test_array.len());
212        for i in 0..result.len() {
213            assert_eq!(
214                result
215                    .execute_scalar(i, &mut array_session().create_execution_ctx())
216                    .unwrap(),
217                Scalar::bool(false, Nullability::NonNullable)
218            );
219        }
220    }
221
222    #[test]
223    fn evaluate_struct() {
224        let test_array = StructArray::from_fields(&[(
225            "a",
226            PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)])
227                .into_array(),
228        )])
229        .unwrap()
230        .into_array();
231        let expected = [true, false, true, false, true];
232
233        let result = test_array
234            .clone()
235            .apply(&is_not_null(get_item("a", root())))
236            .unwrap();
237
238        assert_eq!(result.len(), test_array.len());
239        assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable));
240
241        for (i, expected_value) in expected.iter().enumerate() {
242            assert_eq!(
243                result
244                    .execute_scalar(i, &mut array_session().create_execution_ctx())
245                    .unwrap(),
246                Scalar::bool(*expected_value, Nullability::NonNullable)
247            );
248        }
249    }
250
251    #[test]
252    fn test_display() {
253        let expr = is_not_null(get_item("name", root()));
254        assert_eq!(expr.to_string(), "is_not_null($.name)");
255
256        let expr2 = is_not_null(root());
257        assert_eq!(expr2.to_string(), "is_not_null($)");
258    }
259
260    #[test]
261    fn test_is_not_null_is_not_strict() {
262        assert!(
263            !is_not_null(col("a"))
264                .as_scalar()
265                .is_some_and(|f| f.signature().is_strict())
266        );
267    }
268
269    #[test]
270    fn test_is_not_null_falsification() -> VortexResult<()> {
271        let expr = is_not_null(col("a"));
272        let dtype = test_harness::struct_dtype();
273
274        assert_eq!(
275            expr.bind(&dtype)?.falsify(&STATS_SESSION)?,
276            Some(
277                or(
278                    eq(null_count(col("a")), RowCount.new_expr(EmptyOptions, []),),
279                    all_null(col("a")),
280                )
281                .bind(&dtype)?
282            )
283        );
284        Ok(())
285    }
286}