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::arrays::ScalarFnArray;
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::ReduceNode;
23use crate::scalar_fn::ReduceNodeValidity;
24use crate::scalar_fn::ScalarFnId;
25use crate::scalar_fn::ScalarFnVTable;
26use crate::scalar_fn::ScalarFnVTableExt;
27use crate::scalar_fn::fns::binary::Binary;
28use crate::scalar_fn::fns::is_null::IsNull;
29use crate::scalar_fn::fns::not::Not;
30use crate::scalar_fn::fns::operators::Operator;
31use crate::scalar_fn::is_not_null_node;
32
33/// In array context, reduce
34///
35/// IsNull(x) -> if !x.nullable lit(false) else not(x.validity())
36/// IsNotNull(x) -> if !x.nullable lit(true) or x.validity()
37///
38/// In expression and array contexts for x, y where x is nullable but y is not,
39/// reduce
40///
41/// 1. IsNull(and(x, y)) -> and(IsNull(x), y)
42/// 2. IsNull(or(x, y)) -> and(IsNull(x), not(y))
43/// 3. IsNotNull(and(x, y)) -> or(IsNotNull(x), not(y))
44/// 4. IsNotNull(or(x, y)) -> or(IsNotNull(x), y)
45///
46/// Latter optimizations make sense because calculating IsNull(x)/IsNotNull(x)
47/// is at most expensive as calculating x, but usually much cheaper. Although
48/// in two cases you exchange 4 computations to 4 computations, the latter
49/// four are cheaper.
50pub(crate) fn reduce_null<T: ReduceNode>(is_null: bool, node: &T) -> VortexResult<Option<T>> {
51    let child = node.child(0);
52    if !child.node_dtype()?.is_nullable() {
53        return Ok(Some(node.new_constant((!is_null).into())));
54    }
55
56    if let ReduceNodeValidity::Reduced(validity) = child.validity()? {
57        return Ok(Some(if is_null {
58            validity.new_node(Not.bind(EmptyOptions), std::slice::from_ref(&validity))?
59        } else {
60            validity
61        }));
62    }
63
64    let Some(child_fn) = child.scalar_fn() else {
65        return Ok(None);
66    };
67    let Some(operator) = child_fn.as_opt::<Binary>() else {
68        return Ok(None);
69    };
70    let is_and = match operator {
71        Operator::And => true,
72        Operator::Or => false,
73        _ => return Ok(None),
74    };
75
76    let left = child.child(0);
77    let right = child.child(1);
78    let (nullable, non_nullable) = match (
79        left.node_dtype()?.is_nullable(),
80        right.node_dtype()?.is_nullable(),
81    ) {
82        (true, false) => (left, right),
83        (false, true) => (right, left),
84        // (false, false) already rewritten by constant folding
85        _ => return Ok(None),
86    };
87
88    // is_null(nullable) (rule 1, 2)
89    // is_not_null(nullable) (3, 4)
90    let left = if is_null {
91        nullable.new_node(IsNull.bind(EmptyOptions), std::slice::from_ref(&nullable))?
92    } else {
93        is_not_null_node(&nullable)?
94    };
95    // non_nullable (1, 4)
96    // not(non_nullable) (2, 3)
97    let right = if is_null == is_and {
98        non_nullable
99    } else {
100        non_nullable.new_node(Not.bind(EmptyOptions), std::slice::from_ref(&non_nullable))?
101    };
102    let combine = if is_null { Operator::And } else { Operator::Or };
103    Ok(Some(node.new_node(Binary.bind(combine), &[left, right])?))
104}
105
106/// Expression that checks for non-null values.
107#[derive(Clone)]
108pub struct IsNotNull;
109
110impl IsNotNull {
111    /// Creates a lazy non-null check over `input`.
112    #[expect(clippy::new_ret_no_self, reason = "constructs the lazy result array")]
113    pub fn new(input: ArrayRef) -> ScalarFnArray {
114        ScalarFnArray::try_new(IsNotNull.bind(EmptyOptions), vec![input])
115            .vortex_expect("IsNotNull has one child and an infallible return dtype")
116    }
117}
118
119impl ScalarFnVTable for IsNotNull {
120    type Options = EmptyOptions;
121
122    fn id(&self) -> ScalarFnId {
123        static ID: CachedId = CachedId::new("vortex.is_not_null");
124        *ID
125    }
126
127    fn serialize(&self, _instance: &Self::Options) -> VortexResult<Option<Vec<u8>>> {
128        Ok(Some(vec![]))
129    }
130
131    fn deserialize(
132        &self,
133        _metadata: &[u8],
134        _session: &VortexSession,
135    ) -> VortexResult<Self::Options> {
136        Ok(EmptyOptions)
137    }
138
139    fn arity(&self, _options: &Self::Options) -> Arity {
140        Arity::Exact(1)
141    }
142
143    fn child_name(&self, _instance: &Self::Options, child_idx: usize) -> ChildName {
144        match child_idx {
145            0 => ChildName::from("input"),
146            _ => unreachable!("Invalid child index {} for IsNotNull expression", child_idx),
147        }
148    }
149
150    fn fmt_sql(
151        &self,
152        _options: &Self::Options,
153        expr: &dyn ExprDisplay,
154        f: &mut Formatter<'_>,
155    ) -> std::fmt::Result {
156        write!(f, "is_not_null(")?;
157        Display::fmt(expr.display_child(0), f)?;
158        write!(f, ")")
159    }
160
161    fn return_dtype(&self, _options: &Self::Options, _arg_dtypes: &[DType]) -> VortexResult<DType> {
162        Ok(DType::Bool(Nullability::NonNullable))
163    }
164
165    fn execute(
166        &self,
167        _data: &Self::Options,
168        args: &dyn ExecutionArgs,
169        _ctx: &mut ExecutionCtx,
170    ) -> VortexResult<ArrayRef> {
171        Ok(args.get(0)?.validity()?.to_array(args.row_count()))
172    }
173
174    fn reduce<T: ReduceNode>(&self, _options: &Self::Options, node: &T) -> VortexResult<Option<T>> {
175        reduce_null(false, node)
176    }
177
178    fn is_strict(&self, _instance: &Self::Options) -> bool {
179        // Null input produces the non-null boolean value `false`.
180        false
181    }
182
183    fn is_infallible(&self, _instance: &Self::Options) -> bool {
184        true
185    }
186}
187
188#[cfg(test)]
189mod tests {
190    use std::sync::LazyLock;
191
192    use vortex_buffer::buffer;
193    use vortex_error::VortexExpect as _;
194    use vortex_error::VortexResult;
195    use vortex_session::VortexSession;
196
197    use crate::IntoArray;
198    use crate::VortexSessionExecute;
199    use crate::array_session;
200    use crate::arrays::PrimitiveArray;
201    use crate::arrays::StructArray;
202    use crate::dtype::DType;
203    use crate::dtype::Nullability;
204    use crate::dtype::StructFields;
205    use crate::expr::BoundExpression;
206    use crate::expr::Expression;
207    use crate::expr::and;
208    use crate::expr::col;
209    use crate::expr::eq;
210    use crate::expr::get_item;
211    use crate::expr::is_not_null;
212    use crate::expr::is_null;
213    use crate::expr::lit;
214    use crate::expr::not;
215    use crate::expr::or;
216    use crate::expr::root;
217    use crate::expr::test_harness;
218    use crate::scalar::Scalar;
219    use crate::scalar_fn::EmptyOptions;
220    use crate::scalar_fn::ScalarFnVTableExt;
221    use crate::scalar_fn::internal::row_count::RowCount;
222    use crate::stats::StatsSession;
223    use crate::stats::all_null;
224    use crate::stats::null_count;
225
226    static STATS_SESSION: LazyLock<VortexSession> =
227        LazyLock::new(|| VortexSession::empty().with::<StatsSession>());
228
229    #[test]
230    fn dtype() {
231        let dtype = test_harness::struct_dtype();
232        assert_eq!(
233            is_not_null(root()).return_dtype(&dtype).unwrap(),
234            DType::Bool(Nullability::NonNullable)
235        );
236    }
237
238    fn bool_dtype() -> DType {
239        let names = ["a", "b"].into();
240        let fields = vec![
241            DType::Bool(Nullability::Nullable),
242            DType::Bool(Nullability::NonNullable),
243        ];
244        DType::Struct(StructFields::new(names, fields), Nullability::NonNullable)
245    }
246
247    fn optimized(expr: Expression, scope: &DType) -> VortexResult<BoundExpression> {
248        expr.bind(scope)?.optimize()
249    }
250
251    #[test]
252    fn reduce_to_constant() -> VortexResult<()> {
253        let dtype = bool_dtype();
254        assert_eq!(
255            optimized(is_not_null(col("b")), &dtype)?,
256            lit(true).bind(&dtype)?
257        );
258        assert_eq!(
259            optimized(is_null(col("b")), &dtype)?,
260            lit(false).bind(&dtype)?
261        );
262        Ok(())
263    }
264
265    #[test]
266    fn reduce_kleene() -> VortexResult<()> {
267        let dtype = bool_dtype();
268        assert_eq!(
269            optimized(is_not_null(and(col("a"), col("b"))), &dtype)?,
270            or(is_not_null(col("a")), not(col("b"))).bind(&dtype)?
271        );
272        assert_eq!(
273            optimized(is_not_null(or(col("a"), col("b"))), &dtype)?,
274            or(is_not_null(col("a")), col("b")).bind(&dtype)?
275        );
276        assert_eq!(
277            optimized(is_null(and(col("a"), col("b"))), &dtype)?,
278            and(is_null(col("a")), col("b")).bind(&dtype)?
279        );
280        assert_eq!(
281            optimized(is_null(or(col("a"), col("b"))), &dtype)?,
282            and(is_null(col("a")), not(col("b"))).bind(&dtype)?
283        );
284        Ok(())
285    }
286
287    #[test]
288    fn replace_children() {
289        let expr = is_not_null(root());
290        expr.with_children([root()])
291            .vortex_expect("operation should succeed in test");
292    }
293
294    #[test]
295    fn evaluate_mask() {
296        let test_array =
297            PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)])
298                .into_array();
299        let expected = [true, false, true, false, true];
300
301        let result = test_array.clone().apply(&is_not_null(root())).unwrap();
302
303        assert_eq!(result.len(), test_array.len());
304        assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable));
305
306        for (i, expected_value) in expected.iter().enumerate() {
307            assert_eq!(
308                result
309                    .execute_scalar(i, &mut array_session().create_execution_ctx())
310                    .unwrap(),
311                Scalar::bool(*expected_value, Nullability::NonNullable)
312            );
313        }
314    }
315
316    #[test]
317    fn evaluate_all_true() {
318        let test_array = buffer![1, 2, 3, 4, 5].into_array();
319
320        let result = test_array.clone().apply(&is_not_null(root())).unwrap();
321
322        assert_eq!(result.len(), test_array.len());
323        for i in 0..result.len() {
324            assert_eq!(
325                result
326                    .execute_scalar(i, &mut array_session().create_execution_ctx())
327                    .unwrap(),
328                Scalar::bool(true, Nullability::NonNullable)
329            );
330        }
331    }
332
333    #[test]
334    fn evaluate_all_false() {
335        let test_array =
336            PrimitiveArray::from_option_iter(vec![None::<i32>, None, None, None, None])
337                .into_array();
338
339        let result = test_array.clone().apply(&is_not_null(root())).unwrap();
340
341        assert_eq!(result.len(), test_array.len());
342        for i in 0..result.len() {
343            assert_eq!(
344                result
345                    .execute_scalar(i, &mut array_session().create_execution_ctx())
346                    .unwrap(),
347                Scalar::bool(false, Nullability::NonNullable)
348            );
349        }
350    }
351
352    #[test]
353    fn evaluate_struct() {
354        let test_array = StructArray::from_fields(&[(
355            "a",
356            PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)])
357                .into_array(),
358        )])
359        .unwrap()
360        .into_array();
361        let expected = [true, false, true, false, true];
362
363        let result = test_array
364            .clone()
365            .apply(&is_not_null(get_item("a", root())))
366            .unwrap();
367
368        assert_eq!(result.len(), test_array.len());
369        assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable));
370
371        for (i, expected_value) in expected.iter().enumerate() {
372            assert_eq!(
373                result
374                    .execute_scalar(i, &mut array_session().create_execution_ctx())
375                    .unwrap(),
376                Scalar::bool(*expected_value, Nullability::NonNullable)
377            );
378        }
379    }
380
381    #[test]
382    fn test_display() {
383        let expr = is_not_null(get_item("name", root()));
384        assert_eq!(expr.to_string(), "is_not_null($.name)");
385
386        let expr2 = is_not_null(root());
387        assert_eq!(expr2.to_string(), "is_not_null($)");
388    }
389
390    #[test]
391    fn test_is_not_null_is_not_strict() {
392        assert!(
393            !is_not_null(col("a"))
394                .as_scalar()
395                .is_some_and(|f| f.signature().is_strict())
396        );
397    }
398
399    #[test]
400    fn test_is_not_null_falsification() -> VortexResult<()> {
401        let expr = is_not_null(col("a"));
402        let dtype = test_harness::struct_dtype();
403
404        assert_eq!(
405            expr.bind(&dtype)?.falsify(&STATS_SESSION)?,
406            Some(
407                or(
408                    eq(null_count(col("a")), RowCount.new_expr(EmptyOptions, []),),
409                    all_null(col("a")),
410                )
411                .bind(&dtype)?
412            )
413        );
414        Ok(())
415    }
416}
vortex_array/scalar_fn/fns/is_not_null.rs

vortex_array/scalar_fn/fns/
is_not_null.rs
