Merge pull request #168 from influxdata/crepererum/lazy-component

feat: allow lazy compilation of WASM component during query

Author: crepererum

query/src/lib.rs

@@ -2,6 +2,7 @@
 #![allow(unused_crate_dependencies)]
 
 use std::collections::HashMap;
+use std::pin::Pin;
 
 use datafusion::execution::TaskContext;
 use datafusion_common::{DataFusionError, Result as DataFusionResult};
@@ -17,12 +18,74 @@ use crate::format::UdfCodeFormatter;
 /// Module for UDF code formatting implementations
 pub mod format;
 
+/// Inner type of [`ComponentFn`].
+///
+/// This is deliberately NOT exposed directly to the user because:
+///
+/// - it's a rather complex type
+/// - embedding it into a struct allows us to implement some convenience methods
+type ComponentFnInner<'a> = Box<
+    dyn Fn() -> Pin<Box<dyn Future<Output = &'a WasmComponentPrecompiled> + Send + 'a>> + Sync + 'a,
+>;
+
+/// A type-erased async function that returns a [`WasmComponentPrecompiled`].
+pub struct ComponentFn<'a>(ComponentFnInner<'a>);
+
+impl<'a> ComponentFn<'a> {
+    /// Create function from pre-compiled component.
+    pub fn eager(component: &'a WasmComponentPrecompiled) -> Self {
+        Self(Box::new(move || Box::pin(async move { component })))
+    }
+
+    /// Create function lazily.
+    ///
+    /// # Example
+    /// You can either use an async closure, in which case `'a` can be a local lifetime:
+    ///
+    /// ```
+    /// # use datafusion_udf_wasm_query::ComponentFn;
+    /// let component = ComponentFn::lazy(async || {
+    ///     todo!("create Arc<WasmComponentPrecompiled>")
+    /// });
+    /// ```
+    ///
+    /// or some existing async function, in which `'a` must be `'static`:
+    ///
+    /// ```
+    /// # use datafusion_udf_wasm_host::WasmComponentPrecompiled;
+    /// # use datafusion_udf_wasm_query::ComponentFn;
+    /// async fn get_component() -> &'static WasmComponentPrecompiled {
+    ///     todo!("create Arc<WasmComponentPrecompiled>")
+    /// }
+    ///
+    /// let component = ComponentFn::lazy(get_component);
+    /// ```
+    pub fn lazy<F, Fut>(f: F) -> Self
+    where
+        F: Fn() -> Fut + Sync + 'a,
+        Fut: Future<Output = &'a WasmComponentPrecompiled> + Send + 'a,
+    {
+        Self(Box::new(move || Box::pin(f())))
+    }
+
+    /// Get underlying component.
+    pub async fn get(&self) -> &'a WasmComponentPrecompiled {
+        (self.0)().await
+    }
+}
+
+impl<'a> std::fmt::Debug for ComponentFn<'a> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ComponentFn").finish_non_exhaustive()
+    }
+}
+
 /// Represents a supported UDF language with its associated WASM component
 /// and code formatter.
 #[derive(Debug)]
 pub struct Lang<'a> {
     /// Pre-compiled WASM component for the language
-    pub component: &'a WasmComponentPrecompiled,
+    pub component: ComponentFn<'a>,
     /// Code formatter for the language
     pub formatter: Box<dyn UdfCodeFormatter>,
 }
@@ -81,7 +144,13 @@ impl<'a> UdfQueryParser<'a> {
             for code in blocks {
                 let code = lang.formatter.format(code);
                 udfs.extend(
-                    WasmScalarUdf::new(lang.component, permissions, io_rt.clone(), code).await?,
+                    WasmScalarUdf::new(
+                        lang.component.get().await,
+                        permissions,
+                        io_rt.clone(),
+                        code,
+                    )
+                    .await?,
                 );
             }
         }

query/tests/integration.rs

@@ -14,7 +14,7 @@ use datafusion::{
 use datafusion_common::{Result as DataFusionResult, test_util::batches_to_string};
 use datafusion_udf_wasm_host::WasmPermissions;
 use datafusion_udf_wasm_query::{
-    Lang, ParsedQuery, UdfQueryParser,
+    ComponentFn, Lang, ParsedQuery, UdfQueryParser,
     format::{NoOpFormatter, StripIndentationFormatter},
 };
 use tokio::runtime::Handle;
@@ -53,13 +53,12 @@ SELECT add_one(1);
 "#;
 
     let ctx = SessionContext::new();
-    let component = python_component().await;
     let formatter = Box::new(NoOpFormatter);
 
     let parser = UdfQueryParser::new(HashMap::from_iter([(
         "python".to_string(),
         Lang {
-            component,
+            component: ComponentFn::lazy(python_component),
             formatter,
         },
     )]));
@@ -109,13 +108,12 @@ SELECT add_one(1), multiply_two(3);
 "#;
 
     let ctx = SessionContext::new();
-    let component = python_component().await;
     let formatter = Box::new(NoOpFormatter);
 
     let parser = UdfQueryParser::new(HashMap::from_iter([(
         "python".to_string(),
         Lang {
-            component,
+            component: ComponentFn::lazy(python_component),
             formatter,
         },
     )]));
@@ -161,13 +159,12 @@ SELECT add_one(1), multiply_two(3);
 "#;
 
     let ctx = SessionContext::new();
-    let component = python_component().await;
     let formatter = Box::new(NoOpFormatter);
 
     let parser = UdfQueryParser::new(HashMap::from_iter([(
         "python".to_string(),
         Lang {
-            component,
+            component: ComponentFn::lazy(python_component),
             formatter,
         },
     )]));
@@ -207,13 +204,12 @@ SELECT add_one(1)
 "#;
 
     let ctx = SessionContext::new();
-    let component = python_component().await;
     let formatter = Box::new(NoOpFormatter);
 
     let parser = UdfQueryParser::new(HashMap::from_iter([(
         "python".to_string(),
         Lang {
-            component,
+            component: ComponentFn::lazy(python_component),
             formatter,
         },
     )]));
@@ -248,13 +244,12 @@ EXPLAIN SELECT add_one(1);
 "#;
 
     let ctx = SessionContext::new();
-    let component = python_component().await;
     let formatter = Box::new(NoOpFormatter);
 
     let parser = UdfQueryParser::new(HashMap::from_iter([(
         "python".to_string(),
         Lang {
-            component,
+            component: ComponentFn::lazy(python_component),
             formatter,
         },
     )]));
@@ -304,13 +299,12 @@ async fn test_strip_indentation_everything_indented() {
     let query = query_lines.join("\n");
 
     let ctx = SessionContext::new();
-    let component = python_component().await;
     let formatter = Box::new(StripIndentationFormatter);
 
     let parser = UdfQueryParser::new(HashMap::from_iter([(
         "python".to_string(),
         Lang {
-            component,
+            component: ComponentFn::lazy(python_component),
             formatter,
         },
     )]));
@@ -355,13 +349,12 @@ async fn test_strip_indentation_empty_lines_not_indented() {
     let query = query_lines.join("\n");
 
     let ctx = SessionContext::new();
-    let component = python_component().await;
     let formatter = Box::new(StripIndentationFormatter);
 
     let parser = UdfQueryParser::new(HashMap::from_iter([(
         "python".to_string(),
         Lang {
-            component,
+            component: ComponentFn::lazy(python_component),
             formatter,
         },
     )]));
@@ -405,13 +398,12 @@ async fn test_strip_indentation_python_further_indented() {
     let query = query_lines.join("\n");
 
     let ctx = SessionContext::new();
-    let component = python_component().await;
     let formatter = Box::new(StripIndentationFormatter);
 
     let parser = UdfQueryParser::new(HashMap::from_iter([(
         "python".to_string(),
         Lang {
-            component,
+            component: ComponentFn::lazy(python_component),
             formatter,
         },
     )]));