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), [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::VortexResult;
194    use vortex_session::VortexSession;
195
196    use crate::IntoArray;
197    use crate::VortexSessionExecute;
198    use crate::array_session;
199    use crate::arrays::PrimitiveArray;
200    use crate::arrays::StructArray;
201    use crate::dtype::DType;
202    use crate::dtype::Nullability;
203    use crate::dtype::StructFields;
204    use crate::expr::BoundExpression;
205    use crate::expr::Expression;
206    use crate::expr::and;
207    use crate::expr::col;
208    use crate::expr::eq;
209    use crate::expr::get_item;
210    use crate::expr::is_not_null;
211    use crate::expr::is_null;
212    use crate::expr::lit;
213    use crate::expr::not;
214    use crate::expr::or;
215    use crate::expr::root;
216    use crate::expr::test_harness;
217    use crate::scalar::Scalar;
218    use crate::scalar_fn::EmptyOptions;
219    use crate::scalar_fn::ScalarFnVTableExt;
220    use crate::scalar_fn::internal::row_count::RowCount;
221    use crate::stats::StatsSession;
222    use crate::stats::all_null;
223    use crate::stats::null_count;
224
225    static STATS_SESSION: LazyLock<VortexSession> =
226        LazyLock::new(|| VortexSession::empty().with::<StatsSession>());
227
228    #[test]
229    fn dtype() {
230        let dtype = test_harness::struct_dtype();
231        assert_eq!(
232            is_not_null(root()).return_dtype(&dtype).unwrap(),
233            DType::Bool(Nullability::NonNullable)
234        );
235    }
236
237    fn bool_dtype() -> DType {
238        let names = ["a", "b"].into();
239        let fields = vec![
240            DType::Bool(Nullability::Nullable),
241            DType::Bool(Nullability::NonNullable),
242        ];
243        DType::Struct(StructFields::new(names, fields), Nullability::NonNullable)
244    }
245
246    fn optimized(expr: Expression, scope: &DType) -> VortexResult<BoundExpression> {
247        expr.bind(scope)?.optimize()
248    }
249
250    #[test]
251    fn reduce_to_constant() -> VortexResult<()> {
252        let dtype = bool_dtype();
253        assert_eq!(
254            optimized(is_not_null(col("b")), &dtype)?,
255            lit(true).bind(&dtype)?
256        );
257        assert_eq!(
258            optimized(is_null(col("b")), &dtype)?,
259            lit(false).bind(&dtype)?
260        );
261        Ok(())
262    }
263
264    #[test]
265    fn reduce_kleene() -> VortexResult<()> {
266        let dtype = bool_dtype();
267        assert_eq!(
268            optimized(is_not_null(and(col("a"), col("b"))), &dtype)?,
269            or(is_not_null(col("a")), not(col("b"))).bind(&dtype)?
270        );
271        assert_eq!(
272            optimized(is_not_null(or(col("a"), col("b"))), &dtype)?,
273            or(is_not_null(col("a")), col("b")).bind(&dtype)?
274        );
275        assert_eq!(
276            optimized(is_null(and(col("a"), col("b"))), &dtype)?,
277            and(is_null(col("a")), col("b")).bind(&dtype)?
278        );
279        assert_eq!(
280            optimized(is_null(or(col("a"), col("b"))), &dtype)?,
281            and(is_null(col("a")), not(col("b"))).bind(&dtype)?
282        );
283        Ok(())
284    }
285
286    #[test]
287    fn evaluate_mask() {
288        let test_array =
289            PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)])
290                .into_array();
291        let expected = [true, false, true, false, true];
292
293        let result = test_array.clone().apply(&is_not_null(root())).unwrap();
294
295        assert_eq!(result.len(), test_array.len());
296        assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable));
297
298        for (i, expected_value) in expected.iter().enumerate() {
299            assert_eq!(
300                result
301                    .execute_scalar(i, &mut array_session().create_execution_ctx())
302                    .unwrap(),
303                Scalar::bool(*expected_value, Nullability::NonNullable)
304            );
305        }
306    }
307
308    #[test]
309    fn evaluate_all_true() {
310        let test_array = buffer![1, 2, 3, 4, 5].into_array();
311
312        let result = test_array.clone().apply(&is_not_null(root())).unwrap();
313
314        assert_eq!(result.len(), test_array.len());
315        for i in 0..result.len() {
316            assert_eq!(
317                result
318                    .execute_scalar(i, &mut array_session().create_execution_ctx())
319                    .unwrap(),
320                Scalar::bool(true, Nullability::NonNullable)
321            );
322        }
323    }
324
325    #[test]
326    fn evaluate_all_false() {
327        let test_array =
328            PrimitiveArray::from_option_iter(vec![None::<i32>, None, None, None, None])
329                .into_array();
330
331        let result = test_array.clone().apply(&is_not_null(root())).unwrap();
332
333        assert_eq!(result.len(), test_array.len());
334        for i in 0..result.len() {
335            assert_eq!(
336                result
337                    .execute_scalar(i, &mut array_session().create_execution_ctx())
338                    .unwrap(),
339                Scalar::bool(false, Nullability::NonNullable)
340            );
341        }
342    }
343
344    #[test]
345    fn evaluate_struct() {
346        let test_array = StructArray::from_fields(&[(
347            "a",
348            PrimitiveArray::from_option_iter(vec![Some(1), None, Some(2), None, Some(3)])
349                .into_array(),
350        )])
351        .unwrap()
352        .into_array();
353        let expected = [true, false, true, false, true];
354
355        let result = test_array
356            .clone()
357            .apply(&is_not_null(get_item("a", root())))
358            .unwrap();
359
360        assert_eq!(result.len(), test_array.len());
361        assert_eq!(result.dtype(), &DType::Bool(Nullability::NonNullable));
362
363        for (i, expected_value) in expected.iter().enumerate() {
364            assert_eq!(
365                result
366                    .execute_scalar(i, &mut array_session().create_execution_ctx())
367                    .unwrap(),
368                Scalar::bool(*expected_value, Nullability::NonNullable)
369            );
370        }
371    }
372
373    #[test]
374    fn test_display() {
375        let expr = is_not_null(get_item("name", root()));
376        assert_eq!(expr.to_string(), "is_not_null($.name)");
377
378        let expr2 = is_not_null(root());
379        assert_eq!(expr2.to_string(), "is_not_null($)");
380    }
381
382    #[test]
383    fn test_is_not_null_is_not_strict() {
384        assert!(
385            !is_not_null(col("a"))
386                .as_scalar()
387                .is_some_and(|f| f.signature().is_strict())
388        );
389    }
390
391    #[test]
392    fn test_is_not_null_falsification() -> VortexResult<()> {
393        let expr = is_not_null(col("a"));
394        let dtype = test_harness::struct_dtype();
395
396        assert_eq!(
397            expr.bind(&dtype)?.falsify(&STATS_SESSION)?,
398            Some(
399                or(
400                    eq(null_count(col("a")), RowCount.new_expr(EmptyOptions, []),),
401                    all_null(col("a")),
402                )
403                .bind(&dtype)?
404            )
405        );
406        Ok(())
407    }
408}
vortex_array/scalar_fn/fns/is_not_null.rs

vortex_array/scalar_fn/fns/
is_not_null.rs
