Coverage Report

Created: 2024-10-13 08:39

/Users/andrewlamb/Software/datafusion/datafusion/physical-expr/src/expressions/is_not_null.rs
Line
Count
Source (jump to first uncovered line)
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
20
use std::hash::{Hash, Hasher};
21
use std::{any::Any, sync::Arc};
22
23
use crate::physical_expr::down_cast_any_ref;
24
use crate::PhysicalExpr;
25
use arrow::{
26
    datatypes::{DataType, Schema},
27
    record_batch::RecordBatch,
28
};
29
use datafusion_common::Result;
30
use datafusion_common::ScalarValue;
31
use datafusion_expr::ColumnarValue;
32
33
/// IS NOT NULL expression
34
#[derive(Debug, Hash)]
35
pub struct IsNotNullExpr {
36
    /// The input expression
37
    arg: Arc<dyn PhysicalExpr>,
38
}
39
40
impl IsNotNullExpr {
41
    /// Create new not expression
42
0
    pub fn new(arg: Arc<dyn PhysicalExpr>) -> Self {
43
0
        Self { arg }
44
0
    }
45
46
    /// Get the input expression
47
0
    pub fn arg(&self) -> &Arc<dyn PhysicalExpr> {
48
0
        &self.arg
49
0
    }
50
}
51
52
impl std::fmt::Display for IsNotNullExpr {
53
0
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
54
0
        write!(f, "{} IS NOT NULL", self.arg)
55
0
    }
56
}
57
58
impl PhysicalExpr for IsNotNullExpr {
59
    /// Return a reference to Any that can be used for downcasting
60
0
    fn as_any(&self) -> &dyn Any {
61
0
        self
62
0
    }
63
64
0
    fn data_type(&self, _input_schema: &Schema) -> Result<DataType> {
65
0
        Ok(DataType::Boolean)
66
0
    }
67
68
0
    fn nullable(&self, _input_schema: &Schema) -> Result<bool> {
69
0
        Ok(false)
70
0
    }
71
72
0
    fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
73
0
        let arg = self.arg.evaluate(batch)?;
74
0
        match arg {
75
0
            ColumnarValue::Array(array) => {
76
0
                let is_not_null = arrow::compute::is_not_null(&array)?;
77
0
                Ok(ColumnarValue::Array(Arc::new(is_not_null)))
78
            }
79
0
            ColumnarValue::Scalar(scalar) => Ok(ColumnarValue::Scalar(
80
0
                ScalarValue::Boolean(Some(!scalar.is_null())),
81
0
            )),
82
        }
83
0
    }
84
85
0
    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
86
0
        vec![&self.arg]
87
0
    }
88
89
0
    fn with_new_children(
90
0
        self: Arc<Self>,
91
0
        children: Vec<Arc<dyn PhysicalExpr>>,
92
0
    ) -> Result<Arc<dyn PhysicalExpr>> {
93
0
        Ok(Arc::new(IsNotNullExpr::new(Arc::clone(&children[0]))))
94
0
    }
95
96
0
    fn dyn_hash(&self, state: &mut dyn Hasher) {
97
0
        let mut s = state;
98
0
        self.hash(&mut s);
99
0
    }
100
}
101
102
impl PartialEq<dyn Any> for IsNotNullExpr {
103
0
    fn eq(&self, other: &dyn Any) -> bool {
104
0
        down_cast_any_ref(other)
105
0
            .downcast_ref::<Self>()
106
0
            .map(|x| self.arg.eq(&x.arg))
107
0
            .unwrap_or(false)
108
0
    }
109
}
110
/// Create an IS NOT NULL expression
111
0
pub fn is_not_null(arg: Arc<dyn PhysicalExpr>) -> Result<Arc<dyn PhysicalExpr>> {
112
0
    Ok(Arc::new(IsNotNullExpr::new(arg)))
113
0
}
114
115
#[cfg(test)]
116
mod tests {
117
    use super::*;
118
    use crate::expressions::col;
119
    use arrow::{
120
        array::{BooleanArray, StringArray},
121
        datatypes::*,
122
    };
123
    use arrow_array::{Array, Float64Array, Int32Array, UnionArray};
124
    use arrow_buffer::ScalarBuffer;
125
    use datafusion_common::cast::as_boolean_array;
126
127
    #[test]
128
    fn is_not_null_op() -> Result<()> {
129
        let schema = Schema::new(vec![Field::new("a", DataType::Utf8, true)]);
130
        let a = StringArray::from(vec![Some("foo"), None]);
131
        let expr = is_not_null(col("a", &schema)?).unwrap();
132
        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)])?;
133
134
        // expression: "a is not null"
135
        let result = expr
136
            .evaluate(&batch)?
137
            .into_array(batch.num_rows())
138
            .expect("Failed to convert to array");
139
        let result =
140
            as_boolean_array(&result).expect("failed to downcast to BooleanArray");
141
142
        let expected = &BooleanArray::from(vec![true, false]);
143
144
        assert_eq!(expected, result);
145
146
        Ok(())
147
    }
148
149
    #[test]
150
    fn union_is_not_null_op() {
151
        // union of [{A=1}, {A=}, {B=1.1}, {B=1.2}, {B=}]
152
        let int_array = Int32Array::from(vec![Some(1), None, None, None, None]);
153
        let float_array =
154
            Float64Array::from(vec![None, None, Some(1.1), Some(1.2), None]);
155
        let type_ids = [0, 0, 1, 1, 1].into_iter().collect::<ScalarBuffer<i8>>();
156
157
        let children = vec![Arc::new(int_array) as Arc<dyn Array>, Arc::new(float_array)];
158
159
        let union_fields: UnionFields = [
160
            (0, Arc::new(Field::new("A", DataType::Int32, true))),
161
            (1, Arc::new(Field::new("B", DataType::Float64, true))),
162
        ]
163
        .into_iter()
164
        .collect();
165
166
        let array =
167
            UnionArray::try_new(union_fields.clone(), type_ids, None, children).unwrap();
168
169
        let field = Field::new(
170
            "my_union",
171
            DataType::Union(union_fields, UnionMode::Sparse),
172
            true,
173
        );
174
175
        let schema = Schema::new(vec![field]);
176
        let expr = is_not_null(col("my_union", &schema).unwrap()).unwrap();
177
        let batch =
178
            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(array)]).unwrap();
179
180
        // expression: "a is not null"
181
        let actual = expr
182
            .evaluate(&batch)
183
            .unwrap()
184
            .into_array(batch.num_rows())
185
            .expect("Failed to convert to array");
186
        let actual = as_boolean_array(&actual).unwrap();
187
188
        let expected = &BooleanArray::from(vec![true, false, true, true, false]);
189
190
        assert_eq!(expected, actual);
191
    }
192
}