/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 | | } |