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

vortex_array/scalar_fn/fns/
is_not_null.rs
