Skip to main content

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, ScalarValue};
26use datafusion_expr::ColumnarValue;
27use std::hash::Hash;
28use std::sync::Arc;
29
30/// IS NOT NULL expression
31#[derive(Debug, Eq)]
32pub struct IsNotNullExpr {
33    /// The input expression
34    arg: Arc<dyn PhysicalExpr>,
35}
36
37// Manually derive PartialEq and Hash to work around https://github.com/rust-lang/rust/issues/78808
38impl PartialEq for IsNotNullExpr {
39    fn eq(&self, other: &Self) -> bool {
40        self.arg.eq(&other.arg)
41    }
42}
43
44impl Hash for IsNotNullExpr {
45    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
46        self.arg.hash(state);
47    }
48}
49
50impl IsNotNullExpr {
51    /// Create new not expression
52    pub fn new(arg: Arc<dyn PhysicalExpr>) -> Self {
53        Self { arg }
54    }
55
56    /// Get the input expression
57    pub fn arg(&self) -> &Arc<dyn PhysicalExpr> {
58        &self.arg
59    }
60}
61
62impl std::fmt::Display for IsNotNullExpr {
63    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
64        write!(f, "{} IS NOT NULL", self.arg)
65    }
66}
67
68impl PhysicalExpr for IsNotNullExpr {
69    fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
70        Ok(DataType::Boolean)
71    }
72
73    fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
74        Ok(false)
75    }
76
77    fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
78        let arg = self.arg.evaluate(batch)?;
79        match arg {
80            ColumnarValue::Array(array) => {
81                let is_not_null = arrow::compute::is_not_null(&array)?;
82                Ok(ColumnarValue::Array(Arc::new(is_not_null)))
83            }
84            ColumnarValue::Scalar(scalar) => Ok(ColumnarValue::Scalar(
85                ScalarValue::Boolean(Some(!scalar.is_null())),
86            )),
87        }
88    }
89
90    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
91        vec![&self.arg]
92    }
93
94    fn with_new_children(
95        self: Arc<Self>,
96        children: Vec<Arc<dyn PhysicalExpr>>,
97    ) -> Result<Arc<dyn PhysicalExpr>> {
98        Ok(Arc::new(IsNotNullExpr::new(Arc::clone(&children[0]))))
99    }
100
101    fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
102        self.arg.fmt_sql(f)?;
103        write!(f, " IS NOT NULL")
104    }
105
106    #[cfg(feature = "proto")]
107    fn try_to_proto(
108        &self,
109        ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
110    ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> {
111        use datafusion_proto_models::protobuf;
112
113        Ok(Some(protobuf::PhysicalExprNode {
114            expr_id: None,
115            expr_type: Some(protobuf::physical_expr_node::ExprType::IsNotNullExpr(
116                Box::new(protobuf::PhysicalIsNotNull {
117                    expr: Some(Box::new(ctx.encode_child(&self.arg)?)),
118                }),
119            )),
120        }))
121    }
122}
123
124#[cfg(feature = "proto")]
125impl IsNotNullExpr {
126    /// Reconstruct an [`IsNotNullExpr`] from its protobuf representation.
127    pub fn try_from_proto(
128        node: &datafusion_proto_models::protobuf::PhysicalExprNode,
129        ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>,
130    ) -> Result<Arc<dyn PhysicalExpr>> {
131        use datafusion_physical_expr_common::expect_expr_variant;
132        use datafusion_proto_models::protobuf;
133
134        let node = expect_expr_variant!(
135            node,
136            protobuf::physical_expr_node::ExprType::IsNotNullExpr,
137            "IsNotNullExpr",
138        );
139        let expr = ctx.decode_required_expression(
140            node.expr.as_deref(),
141            "IsNotNullExpr",
142            "expr",
143        )?;
144
145        Ok(Arc::new(IsNotNullExpr::new(expr)))
146    }
147}
148
149/// Create an IS NOT NULL expression
150pub fn is_not_null(arg: Arc<dyn PhysicalExpr>) -> Result<Arc<dyn PhysicalExpr>> {
151    Ok(Arc::new(IsNotNullExpr::new(arg)))
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157    use crate::expressions::col;
158    use arrow::array::{
159        Array, BooleanArray, Float64Array, Int32Array, StringArray, UnionArray,
160    };
161    use arrow::buffer::ScalarBuffer;
162    use arrow::datatypes::*;
163    use datafusion_common::cast::as_boolean_array;
164    use datafusion_physical_expr_common::physical_expr::fmt_sql;
165
166    #[test]
167    fn is_not_null_op() -> Result<()> {
168        let schema = Schema::new(vec![Field::new("a", DataType::Utf8, true)]);
169        let a = StringArray::from(vec![Some("foo"), None]);
170        let expr = is_not_null(col("a", &schema)?).unwrap();
171        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)])?;
172
173        // expression: "a is not null"
174        let result = expr
175            .evaluate(&batch)?
176            .into_array(batch.num_rows())
177            .expect("Failed to convert to array");
178        let result =
179            as_boolean_array(&result).expect("failed to downcast to BooleanArray");
180
181        let expected = &BooleanArray::from(vec![true, false]);
182
183        assert_eq!(expected, result);
184
185        Ok(())
186    }
187
188    #[test]
189    fn union_is_not_null_op() {
190        // union of [{A=1}, {A=}, {B=1.1}, {B=1.2}, {B=}]
191        let int_array = Int32Array::from(vec![Some(1), None, None, None, None]);
192        let float_array =
193            Float64Array::from(vec![None, None, Some(1.1), Some(1.2), None]);
194        let type_ids = [0, 0, 1, 1, 1].into_iter().collect::<ScalarBuffer<i8>>();
195
196        let children = vec![Arc::new(int_array) as Arc<dyn Array>, Arc::new(float_array)];
197
198        let union_fields: UnionFields = [
199            (0, Arc::new(Field::new("A", DataType::Int32, true))),
200            (1, Arc::new(Field::new("B", DataType::Float64, true))),
201        ]
202        .into_iter()
203        .collect();
204
205        let array =
206            UnionArray::try_new(union_fields.clone(), type_ids, None, children).unwrap();
207
208        let field = Field::new(
209            "my_union",
210            DataType::Union(union_fields, UnionMode::Sparse),
211            true,
212        );
213
214        let schema = Schema::new(vec![field]);
215        let expr = is_not_null(col("my_union", &schema).unwrap()).unwrap();
216        let batch =
217            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array)]).unwrap();
218
219        // expression: "a is not null"
220        let actual = expr
221            .evaluate(&batch)
222            .unwrap()
223            .into_array(batch.num_rows())
224            .expect("Failed to convert to array");
225        let actual = as_boolean_array(&actual).unwrap();
226
227        let expected = &BooleanArray::from(vec![true, false, true, true, false]);
228
229        assert_eq!(expected, actual);
230    }
231
232    #[test]
233    fn test_fmt_sql() -> Result<()> {
234        let union_fields: UnionFields = [
235            (0, Arc::new(Field::new("A", DataType::Int32, true))),
236            (1, Arc::new(Field::new("B", DataType::Float64, true))),
237        ]
238        .into_iter()
239        .collect();
240
241        let field = Field::new(
242            "my_union",
243            DataType::Union(union_fields, UnionMode::Sparse),
244            true,
245        );
246
247        let schema = Schema::new(vec![field]);
248        let expr = is_not_null(col("my_union", &schema).unwrap()).unwrap();
249        let display_string = expr.to_string();
250        assert_eq!(display_string, "my_union@0 IS NOT NULL");
251        let sql_string = fmt_sql(expr.as_ref()).to_string();
252        assert_eq!(sql_string, "my_union IS NOT NULL");
253
254        Ok(())
255    }
256}
257
258#[cfg(all(test, feature = "proto"))]
259mod proto_tests {
260    use super::*;
261    use crate::expressions::{Column, col};
262    use crate::proto_test_util::{
263        StubDecoder, StubEncoder, UnreachableDecoder, column_node,
264    };
265    use arrow::datatypes::Field;
266    use datafusion_common::DataFusionError;
267    use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx;
268    use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx;
269    use datafusion_proto_models::protobuf::{
270        PhysicalExprNode, PhysicalIsNotNull, physical_expr_node,
271    };
272
273    fn is_not_null_node(expr: Option<Box<PhysicalExprNode>>) -> PhysicalExprNode {
274        PhysicalExprNode {
275            expr_id: None,
276            expr_type: Some(physical_expr_node::ExprType::IsNotNullExpr(Box::new(
277                PhysicalIsNotNull { expr },
278            ))),
279        }
280    }
281
282    fn is_not_null_fixture() -> IsNotNullExpr {
283        let schema = Schema::new(vec![Field::new("a", DataType::Int64, true)]);
284        IsNotNullExpr::new(col("a", &schema).unwrap())
285    }
286
287    #[test]
288    fn try_to_proto_encodes_is_not_null_expr() {
289        let is_not_null = is_not_null_fixture();
290        let encoder = StubEncoder::ok();
291        let ctx = PhysicalExprEncodeCtx::new(&encoder);
292
293        let node = is_not_null
294            .try_to_proto(&ctx)
295            .unwrap()
296            .expect("IsNotNullExpr should encode to Some(node)");
297
298        assert!(node.expr_id.is_none());
299        let is_not_null_node = match node.expr_type {
300            Some(physical_expr_node::ExprType::IsNotNullExpr(boxed)) => *boxed,
301            other => panic!("expected an IsNotNullExpr node, got {other:?}"),
302        };
303        assert!(is_not_null_node.expr.is_some());
304    }
305
306    #[test]
307    fn try_to_proto_propagates_expr_encode_error() {
308        let is_not_null = is_not_null_fixture();
309        let encoder = StubEncoder::failing_on(1);
310        let ctx = PhysicalExprEncodeCtx::new(&encoder);
311        let err = is_not_null.try_to_proto(&ctx).unwrap_err();
312        assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1")));
313    }
314
315    #[test]
316    fn try_from_proto_decodes_is_not_null_expr() {
317        let node = is_not_null_node(Some(Box::new(column_node("a"))));
318        let schema = Schema::empty();
319        let decoder = StubDecoder::ok();
320        let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
321
322        let decoded = IsNotNullExpr::try_from_proto(&node, &ctx).unwrap();
323        let is_not_null = decoded
324            .downcast_ref::<IsNotNullExpr>()
325            .expect("decoded expr should be an IsNotNullExpr");
326        assert!(is_not_null.arg().downcast_ref::<Column>().is_some());
327    }
328
329    #[test]
330    fn try_from_proto_rejects_non_is_not_null_node() {
331        let node = column_node("a");
332        let schema = Schema::empty();
333        let decoder = UnreachableDecoder;
334        let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
335        let err = IsNotNullExpr::try_from_proto(&node, &ctx).unwrap_err();
336        assert!(
337            matches!(err, DataFusionError::Internal(msg) if msg.contains("PhysicalExprNode is not a IsNotNullExpr"))
338        );
339    }
340
341    #[test]
342    fn try_from_proto_rejects_missing_expr() {
343        let node = is_not_null_node(None);
344        let schema = Schema::empty();
345        let decoder = UnreachableDecoder;
346        let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
347        let err = IsNotNullExpr::try_from_proto(&node, &ctx).unwrap_err();
348        assert!(
349            matches!(err, DataFusionError::Internal(msg) if msg.contains("IsNotNullExpr is missing required field 'expr'"))
350        );
351    }
352
353    #[test]
354    fn try_from_proto_propagates_expr_decode_error() {
355        let node = is_not_null_node(Some(Box::new(column_node("a"))));
356        let schema = Schema::empty();
357        let decoder = StubDecoder::failing_on(1);
358        let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
359        let err = IsNotNullExpr::try_from_proto(&node, &ctx).unwrap_err();
360        assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1")));
361    }
362}