datafusion_physical_expr/expressions/
is_not_null.rs1use 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#[derive(Debug, Eq)]
32pub struct IsNotNullExpr {
33 arg: Arc<dyn PhysicalExpr>,
35}
36
37impl 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 pub fn new(arg: Arc<dyn PhysicalExpr>) -> Self {
53 Self { arg }
54 }
55
56 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 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
149pub 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 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 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 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}