fix: UDF extraction and EXPLAIN queries

Author: crepererum

host/src/udf_query.rs

@@ -89,11 +89,7 @@ impl<'a> UdfQueryParser<'a> {
         let mut sql = String::new();
         let mut udf_blocks: HashMap<String, Vec<String>> = HashMap::new();
         for s in statements {
-            let Statement::Statement(stmt) = s else {
-                continue;
-            };
-
-            match parse_udf(*stmt)? {
+            match parse_udf(s)? {
                 Parsed::Udf { code, language } => {
                     if let Some(existing) = udf_blocks.get_mut(&language) {
                         existing.push(code);
@@ -130,31 +126,34 @@ enum Parsed {
 }
 
 /// Parse a single SQL statement to extract a UDF
-fn parse_udf(stmt: SqlStatement) -> DataFusionResult<Parsed> {
+fn parse_udf(stmt: Statement) -> DataFusionResult<Parsed> {
     match stmt {
-        SqlStatement::CreateFunction(cf) => {
-            let function_body = cf.function_body.as_ref();
-
-            let language = if let Some(lang) = cf.language.as_ref() {
-                lang.to_string()
-            } else {
-                return Err(DataFusionError::Plan(
-                    "function language is required for UDFs".to_string(),
-                ));
-            };
-
-            let code = match function_body {
-                Some(body) => extract_function_body(body),
-                None => Err(DataFusionError::Plan(
-                    "function body is required for UDFs".to_string(),
-                )),
-            }?;
-
-            Ok(Parsed::Udf {
-                code: code.to_string(),
-                language,
-            })
-        }
+        Statement::Statement(stmt) => match *stmt {
+            SqlStatement::CreateFunction(cf) => {
+                let function_body = cf.function_body.as_ref();
+
+                let language = if let Some(lang) = cf.language.as_ref() {
+                    lang.to_string()
+                } else {
+                    return Err(DataFusionError::Plan(
+                        "function language is required for UDFs".to_string(),
+                    ));
+                };
+
+                let code = match function_body {
+                    Some(body) => extract_function_body(body),
+                    None => Err(DataFusionError::Plan(
+                        "function body is required for UDFs".to_string(),
+                    )),
+                }?;
+
+                Ok(Parsed::Udf {
+                    code: code.to_string(),
+                    language,
+                })
+            }
+            _ => Ok(Parsed::Other(stmt.to_string())),
+        },
         _ => Ok(Parsed::Other(stmt.to_string())),
     }
 }

host/tests/integration_tests/udf_query.rs

@@ -4,7 +4,7 @@ use datafusion::{
     assert_batches_eq,
     prelude::{DataFrame, SessionContext},
 };
-use datafusion_common::Result as DataFusionResult;
+use datafusion_common::{Result as DataFusionResult, test_util::batches_to_string};
 use datafusion_udf_wasm_host::{
     WasmPermissions,
     udf_query::{ParsedQuery, UdfQueryParser},
@@ -77,7 +77,7 @@ def add_one(x: int) -> int:
 ';
 
 CREATE FUNCTION multiply_two()
-LANGUAGE python  
+LANGUAGE python
 AS '
 def multiply_two(x: int) -> int:
     return x * 2
@@ -175,3 +175,43 @@ SELECT add_one(1)
     let err = r.err().unwrap();
     assert!(err.message().contains("Invalid function 'add_one'"));
 }
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_explain() {
+    let query = r#"
+CREATE FUNCTION add_one()
+LANGUAGE python
+AS '
+def add_one(x: int) -> int:
+    return x + 1
+';
+
+EXPLAIN SELECT add_one(1);
+"#;
+
+    let ctx = SessionContext::new();
+    let component = python_component().await;
+
+    let parser = UdfQueryParser::new(HashMap::from_iter([("python".to_string(), component)]));
+    let parsed_query = parser
+        .parse(query, &WasmPermissions::new(), ctx.task_ctx().as_ref())
+        .await
+        .unwrap();
+
+    let df = UdfQueryInvocator::invoke(&ctx, parsed_query).await.unwrap();
+    let batch = df.collect().await.unwrap();
+
+    insta::assert_snapshot!(
+        batches_to_string(&batch),
+        @r"
+    +---------------+--------------------------------------------------------+
+    | plan_type     | plan                                                   |
+    +---------------+--------------------------------------------------------+
+    | logical_plan  | Projection: add_one(Int64(1))                          |
+    |               |   EmptyRelation                                        |
+    | physical_plan | ProjectionExec: expr=[add_one(1) as add_one(Int64(1))] |
+    |               |   PlaceholderRowExec                                   |
+    |               |                                                        |
+    +---------------+--------------------------------------------------------+
+    ");
+}