datafusion_physical_expr/expressions/
is_not_null.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! IS NOT NULL expression
19
20use crate::PhysicalExpr;
21use arrow::{
22    datatypes::{DataType, Schema},
23    record_batch::RecordBatch,
24};
25use datafusion_common::Result;
26use datafusion_common::ScalarValue;
27use datafusion_expr::ColumnarValue;
28use std::hash::Hash;
29use std::sync::Arc;
30
31/// IS NOT NULL expression
32#[derive(Debug, Eq)]
33pub struct IsNotNullExpr {
34    /// The input expression
35    arg: Arc<dyn PhysicalExpr>,
36}
37
38// Manually derive PartialEq and Hash to work around https://github.com/rust-lang/rust/issues/78808
39impl PartialEq for IsNotNullExpr {
40    fn eq(&self, other: &Self) -> bool {
41        self.arg.eq(&other.arg)
42    }
43}
44
45impl Hash for IsNotNullExpr {
46    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
47        self.arg.hash(state);
48    }
49}
50
51impl IsNotNullExpr {
52    /// Create new not expression
53    pub fn new(arg: Arc<dyn PhysicalExpr>) -> Self {
54        Self { arg }
55    }
56
57    /// Get the input expression
58    pub fn arg(&self) -> &Arc<dyn PhysicalExpr> {
59        &self.arg
60    }
61}
62
63impl std::fmt::Display for IsNotNullExpr {
64    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
65        write!(f, "{} IS NOT NULL", self.arg)
66    }
67}
68
69impl PhysicalExpr for IsNotNullExpr {
70    fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
71        Ok(DataType::Boolean)
72    }
73
74    fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
75        Ok(false)
76    }
77
78    fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
79        let arg = self.arg.evaluate(batch)?;
80        match arg {
81            ColumnarValue::Array(array) => {
82                let is_not_null = arrow::compute::is_not_null(&array)?;
83                Ok(ColumnarValue::Array(Arc::new(is_not_null)))
84            }
85            ColumnarValue::Scalar(scalar) => Ok(ColumnarValue::Scalar(
86                ScalarValue::Boolean(Some(!scalar.is_null())),
87            )),
88        }
89    }
90
91    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
92        vec![&self.arg]
93    }
94
95    fn with_new_children(
96        self: Arc<Self>,
97        children: Vec<Arc<dyn PhysicalExpr>>,
98    ) -> Result<Arc<dyn PhysicalExpr>> {
99        Ok(Arc::new(IsNotNullExpr::new(Arc::clone(&children[0]))))
100    }
101
102    fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
103        self.arg.fmt_sql(f)?;
104        write!(f, " IS NOT NULL")
105    }
106}
107
108/// Create an IS NOT NULL expression
109pub fn is_not_null(arg: Arc<dyn PhysicalExpr>) -> Result<Arc<dyn PhysicalExpr>> {
110    Ok(Arc::new(IsNotNullExpr::new(arg)))
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116    use crate::expressions::col;
117    use arrow::array::{
118        Array, BooleanArray, Float64Array, Int32Array, StringArray, UnionArray,
119    };
120    use arrow::buffer::ScalarBuffer;
121    use arrow::datatypes::*;
122    use datafusion_common::cast::as_boolean_array;
123    use datafusion_physical_expr_common::physical_expr::fmt_sql;
124
125    #[test]
126    fn is_not_null_op() -> Result<()> {
127        let schema = Schema::new(vec![Field::new("a", DataType::Utf8, true)]);
128        let a = StringArray::from(vec![Some("foo"), None]);
129        let expr = is_not_null(col("a", &schema)?).unwrap();
130        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)])?;
131
132        // expression: "a is not null"
133        let result = expr
134            .evaluate(&batch)?
135            .into_array(batch.num_rows())
136            .expect("Failed to convert to array");
137        let result =
138            as_boolean_array(&result).expect("failed to downcast to BooleanArray");
139
140        let expected = &BooleanArray::from(vec![true, false]);
141
142        assert_eq!(expected, result);
143
144        Ok(())
145    }
146
147    #[test]
148    fn union_is_not_null_op() {
149        // union of [{A=1}, {A=}, {B=1.1}, {B=1.2}, {B=}]
150        let int_array = Int32Array::from(vec![Some(1), None, None, None, None]);
151        let float_array =
152            Float64Array::from(vec![None, None, Some(1.1), Some(1.2), None]);
153        let type_ids = [0, 0, 1, 1, 1].into_iter().collect::<ScalarBuffer<i8>>();
154
155        let children = vec![Arc::new(int_array) as Arc<dyn Array>, Arc::new(float_array)];
156
157        let union_fields: UnionFields = [
158            (0, Arc::new(Field::new("A", DataType::Int32, true))),
159            (1, Arc::new(Field::new("B", DataType::Float64, true))),
160        ]
161        .into_iter()
162        .collect();
163
164        let array =
165            UnionArray::try_new(union_fields.clone(), type_ids, None, children).unwrap();
166
167        let field = Field::new(
168            "my_union",
169            DataType::Union(union_fields, UnionMode::Sparse),
170            true,
171        );
172
173        let schema = Schema::new(vec![field]);
174        let expr = is_not_null(col("my_union", &schema).unwrap()).unwrap();
175        let batch =
176            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array)]).unwrap();
177
178        // expression: "a is not null"
179        let actual = expr
180            .evaluate(&batch)
181            .unwrap()
182            .into_array(batch.num_rows())
183            .expect("Failed to convert to array");
184        let actual = as_boolean_array(&actual).unwrap();
185
186        let expected = &BooleanArray::from(vec![true, false, true, true, false]);
187
188        assert_eq!(expected, actual);
189    }
190
191    #[test]
192    fn test_fmt_sql() -> Result<()> {
193        let union_fields: UnionFields = [
194            (0, Arc::new(Field::new("A", DataType::Int32, true))),
195            (1, Arc::new(Field::new("B", DataType::Float64, true))),
196        ]
197        .into_iter()
198        .collect();
199
200        let field = Field::new(
201            "my_union",
202            DataType::Union(union_fields, UnionMode::Sparse),
203            true,
204        );
205
206        let schema = Schema::new(vec![field]);
207        let expr = is_not_null(col("my_union", &schema).unwrap()).unwrap();
208        let display_string = expr.to_string();
209        assert_eq!(display_string, "my_union@0 IS NOT NULL");
210        let sql_string = fmt_sql(expr.as_ref()).to_string();
211        assert_eq!(sql_string, "my_union IS NOT NULL");
212
213        Ok(())
214    }
215}
datafusion_physical_expr/expressions/is_not_null.rs

datafusion_physical_expr/expressions/
is_not_null.rs
