add test case for proto udf decode fallback

Author: leoyvens

datafusion/proto/tests/cases/roundtrip_logical_plan.rs

@@ -113,15 +113,21 @@ fn roundtrip_json_test(proto: &protobuf::LogicalExprNode) {
 #[cfg(not(feature = "json"))]
 fn roundtrip_json_test(_proto: &protobuf::LogicalExprNode) {}
 
-// Given a DataFusion logical Expr, convert it to protobuf and back, using debug formatting to test
-// equality.
 fn roundtrip_expr_test(initial_struct: Expr, ctx: SessionContext) {
     let extension_codec = DefaultLogicalExtensionCodec {};
-    let proto: protobuf::LogicalExprNode =
-        serialize_expr(&initial_struct, &extension_codec)
-            .unwrap_or_else(|e| panic!("Error serializing expression: {:?}", e));
-    let round_trip: Expr =
-        from_proto::parse_expr(&proto, &ctx, &extension_codec).unwrap();
+    roundtrip_expr_test_with_codec(initial_struct, ctx, &extension_codec);
+}
+
+// Given a DataFusion logical Expr, convert it to protobuf and back, using debug formatting to test
+// equality.
+fn roundtrip_expr_test_with_codec(
+    initial_struct: Expr,
+    ctx: SessionContext,
+    codec: &dyn LogicalExtensionCodec,
+) {
+    let proto: protobuf::LogicalExprNode = serialize_expr(&initial_struct, codec)
+        .unwrap_or_else(|e| panic!("Error serializing expression: {:?}", e));
+    let round_trip: Expr = from_proto::parse_expr(&proto, &ctx, codec).unwrap();
 
     assert_eq!(format!("{:?}", &initial_struct), format!("{round_trip:?}"));
 
@@ -2185,22 +2191,26 @@ fn roundtrip_aggregate_udf() {
     roundtrip_expr_test(test_expr, ctx);
 }
 
-#[test]
-fn roundtrip_scalar_udf() {
+fn dummy_udf() -> ScalarUDF {
     let scalar_fn = Arc::new(|args: &[ColumnarValue]| {
         let ColumnarValue::Array(array) = &args[0] else {
             panic!("should be array")
         };
         Ok(ColumnarValue::from(Arc::new(array.clone()) as ArrayRef))
     });
 
-    let udf = create_udf(
+    create_udf(
         "dummy",
         vec![DataType::Utf8],
         DataType::Utf8,
         Volatility::Immutable,
         scalar_fn,
-    );
+    )
+}
+
+#[test]
+fn roundtrip_scalar_udf() {
+    let udf = dummy_udf();
 
     let test_expr = Expr::ScalarFunction(ScalarFunction::new_udf(
         Arc::new(udf.clone()),
@@ -2210,7 +2220,57 @@ fn roundtrip_scalar_udf() {
     let ctx = SessionContext::new();
     ctx.register_udf(udf);
 
-    roundtrip_expr_test(test_expr, ctx);
+    roundtrip_expr_test(test_expr.clone(), ctx);
+
+    // Now test loading the UDF without registering it in the context, but rather creating it in the
+    // extension codec.
+    #[derive(Debug)]
+    struct DummyUDFExtensionCodec;
+
+    impl LogicalExtensionCodec for DummyUDFExtensionCodec {
+        fn try_decode(
+            &self,
+            _buf: &[u8],
+            _inputs: &[LogicalPlan],
+            _ctx: &SessionContext,
+        ) -> Result<Extension> {
+            not_impl_err!("LogicalExtensionCodec is not provided")
+        }
+
+        fn try_encode(&self, _node: &Extension, _buf: &mut Vec<u8>) -> Result<()> {
+            not_impl_err!("LogicalExtensionCodec is not provided")
+        }
+
+        fn try_decode_table_provider(
+            &self,
+            _buf: &[u8],
+            _table_ref: &TableReference,
+            _schema: SchemaRef,
+            _ctx: &SessionContext,
+        ) -> Result<Arc<dyn TableProvider>> {
+            not_impl_err!("LogicalExtensionCodec is not provided")
+        }
+
+        fn try_encode_table_provider(
+            &self,
+            _table_ref: &TableReference,
+            _node: Arc<dyn TableProvider>,
+            _buf: &mut Vec<u8>,
+        ) -> Result<()> {
+            not_impl_err!("LogicalExtensionCodec is not provided")
+        }
+
+        fn try_decode_udf(&self, name: &str, _buf: &[u8]) -> Result<Arc<ScalarUDF>> {
+            if name == "dummy" {
+                Ok(Arc::new(dummy_udf()))
+            } else {
+                return Err(DataFusionError::Internal(format!("UDF {name} not found")));
+            }
+        }
+    }
+
+    let ctx = SessionContext::new();
+    roundtrip_expr_test_with_codec(test_expr, ctx, &DummyUDFExtensionCodec)
 }
 
 #[test]