Port IsNotNullExpr proto serialization hooks

Author: chakkk309

datafusion/physical-expr/src/expressions/is_not_null.rs

@@ -22,8 +22,7 @@ use arrow::{
     datatypes::{DataType, Schema},
     record_batch::RecordBatch,
 };
-use datafusion_common::Result;
-use datafusion_common::ScalarValue;
+use datafusion_common::{Result, ScalarValue};
 use datafusion_expr::ColumnarValue;
 use std::hash::Hash;
 use std::sync::Arc;
@@ -103,6 +102,48 @@ impl PhysicalExpr for IsNotNullExpr {
         self.arg.fmt_sql(f)?;
         write!(f, " IS NOT NULL")
     }
+
+    #[cfg(feature = "proto")]
+    fn try_to_proto(
+        &self,
+        ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>,
+    ) -> Result<Option<datafusion_proto_models::protobuf::PhysicalExprNode>> {
+        use datafusion_proto_models::protobuf;
+
+        Ok(Some(protobuf::PhysicalExprNode {
+            expr_id: None,
+            expr_type: Some(protobuf::physical_expr_node::ExprType::IsNotNullExpr(
+                Box::new(protobuf::PhysicalIsNotNull {
+                    expr: Some(Box::new(ctx.encode_child(&self.arg)?)),
+                }),
+            )),
+        }))
+    }
+}
+
+#[cfg(feature = "proto")]
+impl IsNotNullExpr {
+    /// Reconstruct an [`IsNotNullExpr`] from its protobuf representation.
+    pub fn try_from_proto(
+        node: &datafusion_proto_models::protobuf::PhysicalExprNode,
+        ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>,
+    ) -> Result<Arc<dyn PhysicalExpr>> {
+        use datafusion_physical_expr_common::expect_expr_variant;
+        use datafusion_proto_models::protobuf;
+
+        let node = expect_expr_variant!(
+            node,
+            protobuf::physical_expr_node::ExprType::IsNotNullExpr,
+            "IsNotNullExpr",
+        );
+        let expr = ctx.decode_required_expression(
+            node.expr.as_deref(),
+            "IsNotNullExpr",
+            "expr",
+        )?;
+
+        Ok(Arc::new(IsNotNullExpr::new(expr)))
+    }
 }
 
 /// Create an IS NOT NULL expression
@@ -213,3 +254,109 @@ mod tests {
         Ok(())
     }
 }
+
+#[cfg(all(test, feature = "proto"))]
+mod proto_tests {
+    use super::*;
+    use crate::expressions::{Column, col};
+    use crate::proto_test_util::{
+        StubDecoder, StubEncoder, UnreachableDecoder, column_node,
+    };
+    use arrow::datatypes::Field;
+    use datafusion_common::DataFusionError;
+    use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx;
+    use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx;
+    use datafusion_proto_models::protobuf::{
+        PhysicalExprNode, PhysicalIsNotNull, physical_expr_node,
+    };
+
+    fn is_not_null_node(expr: Option<Box<PhysicalExprNode>>) -> PhysicalExprNode {
+        PhysicalExprNode {
+            expr_id: None,
+            expr_type: Some(physical_expr_node::ExprType::IsNotNullExpr(Box::new(
+                PhysicalIsNotNull { expr },
+            ))),
+        }
+    }
+
+    fn is_not_null_fixture() -> IsNotNullExpr {
+        let schema = Schema::new(vec![Field::new("a", DataType::Int64, true)]);
+        IsNotNullExpr::new(col("a", &schema).unwrap())
+    }
+
+    #[test]
+    fn try_to_proto_encodes_is_not_null_expr() {
+        let is_not_null = is_not_null_fixture();
+        let encoder = StubEncoder::ok();
+        let ctx = PhysicalExprEncodeCtx::new(&encoder);
+
+        let node = is_not_null
+            .try_to_proto(&ctx)
+            .unwrap()
+            .expect("IsNotNullExpr should encode to Some(node)");
+
+        assert!(node.expr_id.is_none());
+        let is_not_null_node = match node.expr_type {
+            Some(physical_expr_node::ExprType::IsNotNullExpr(boxed)) => *boxed,
+            other => panic!("expected an IsNotNullExpr node, got {other:?}"),
+        };
+        assert!(is_not_null_node.expr.is_some());
+    }
+
+    #[test]
+    fn try_to_proto_propagates_expr_encode_error() {
+        let is_not_null = is_not_null_fixture();
+        let encoder = StubEncoder::failing_on(1);
+        let ctx = PhysicalExprEncodeCtx::new(&encoder);
+        let err = is_not_null.try_to_proto(&ctx).unwrap_err();
+        assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1")));
+    }
+
+    #[test]
+    fn try_from_proto_decodes_is_not_null_expr() {
+        let node = is_not_null_node(Some(Box::new(column_node("a"))));
+        let schema = Schema::empty();
+        let decoder = StubDecoder::ok();
+        let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
+
+        let decoded = IsNotNullExpr::try_from_proto(&node, &ctx).unwrap();
+        let is_not_null = decoded
+            .downcast_ref::<IsNotNullExpr>()
+            .expect("decoded expr should be an IsNotNullExpr");
+        assert!(is_not_null.arg().downcast_ref::<Column>().is_some());
+    }
+
+    #[test]
+    fn try_from_proto_rejects_non_is_not_null_node() {
+        let node = column_node("a");
+        let schema = Schema::empty();
+        let decoder = UnreachableDecoder;
+        let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
+        let err = IsNotNullExpr::try_from_proto(&node, &ctx).unwrap_err();
+        assert!(
+            matches!(err, DataFusionError::Internal(msg) if msg.contains("PhysicalExprNode is not a IsNotNullExpr"))
+        );
+    }
+
+    #[test]
+    fn try_from_proto_rejects_missing_expr() {
+        let node = is_not_null_node(None);
+        let schema = Schema::empty();
+        let decoder = UnreachableDecoder;
+        let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
+        let err = IsNotNullExpr::try_from_proto(&node, &ctx).unwrap_err();
+        assert!(
+            matches!(err, DataFusionError::Internal(msg) if msg.contains("IsNotNullExpr is missing required field 'expr'"))
+        );
+    }
+
+    #[test]
+    fn try_from_proto_propagates_expr_decode_error() {
+        let node = is_not_null_node(Some(Box::new(column_node("a"))));
+        let schema = Schema::empty();
+        let decoder = StubDecoder::failing_on(1);
+        let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder);
+        let err = IsNotNullExpr::try_from_proto(&node, &ctx).unwrap_err();
+        assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1")));
+    }
+}

datafusion/proto/src/physical_plan/from_proto.rs

@@ -307,15 +307,7 @@ pub fn parse_physical_expr_with_converter(
                 proto_converter,
             )?))
         }
-        ExprType::IsNotNullExpr(e) => {
-            Arc::new(IsNotNullExpr::new(parse_required_physical_expr(
-                e.expr.as_deref(),
-                ctx,
-                "expr",
-                input_schema,
-                proto_converter,
-            )?))
-        }
+        ExprType::IsNotNullExpr(_) => IsNotNullExpr::try_from_proto(proto, &decode_ctx)?,
         ExprType::NotExpr(_) => NotExpr::try_from_proto(proto, &decode_ctx)?,
         ExprType::Negative(_) => NegativeExpr::try_from_proto(proto, &decode_ctx)?,
         ExprType::InList(_) => InListExpr::try_from_proto(proto, &decode_ctx)?,

datafusion/proto/src/physical_plan/to_proto.rs

@@ -36,7 +36,7 @@ use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr;
 use datafusion_physical_expr::window::{SlidingAggregateWindowExpr, StandardWindowExpr};
 use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
 use datafusion_physical_plan::expressions::{
-    CaseExpr, DynamicFilterPhysicalExpr, IsNotNullExpr, IsNullExpr, Literal, TryCastExpr,
+    CaseExpr, DynamicFilterPhysicalExpr, IsNullExpr, Literal, TryCastExpr,
 };
 use datafusion_physical_plan::udaf::AggregateFunctionExpr;
 use datafusion_physical_plan::windows::{PlainAggregateWindowExpr, WindowUDFExpr};
@@ -356,17 +356,6 @@ pub fn serialize_physical_expr_with_converter(
                 }),
             )),
         })
-    } else if let Some(expr) = expr.downcast_ref::<IsNotNullExpr>() {
-        Ok(protobuf::PhysicalExprNode {
-            expr_id,
-            expr_type: Some(protobuf::physical_expr_node::ExprType::IsNotNullExpr(
-                Box::new(protobuf::PhysicalIsNotNull {
-                    expr: Some(Box::new(
-                        proto_converter.physical_expr_to_proto(expr.arg(), codec)?,
-                    )),
-                }),
-            )),
-        })
     } else if let Some(lit) = expr.downcast_ref::<Literal>() {
         Ok(protobuf::PhysicalExprNode {
             expr_id,