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

vortex_array/scalar_fn/fns/
is_not_null.rs
