feat(udf): add initial support for JavaScript UDF

Cargo.toml

@@ -132,7 +132,8 @@ arrow-flight = "49"
 arrow-select = "49"
 arrow-ord = "49"
 arrow-row = "49"
-arrow-udf-wasm = { git = "https://github.com/risingwavelabs/arrow-udf.git", rev = "f9a9e0d" }
+arrow-udf-js = { git = "https://github.com/risingwavelabs/arrow-udf.git", rev = "51d941d" }
+arrow-udf-wasm = { git = "https://github.com/risingwavelabs/arrow-udf.git", rev = "51d941d" }
 arrow-array-deltalake = { package = "arrow-array", version = "48.0.1" }
 arrow-buffer-deltalake = { package = "arrow-buffer", version = "48.0.1" }
 arrow-cast-deltalake = { package = "arrow-cast", version = "48.0.1" }

e2e_test/udf/js_udf.slt

@@ -0,0 +1,52 @@
+statement ok
+create function int_42() returns int language javascript as $$
+    export function int_42() {
+        return 42;
+    }
+$$;
+
+statement ok
+create function gcd(int, int) returns int language javascript as $$
+    export function gcd(a, b) {
+        if(a == null || b == null) {
+            return null;
+        }
+        while (b != 0) {
+            let t = b;
+            b = a % b;
+            a = t;
+        }
+        return a;
+    }
+$$;
+
+statement ok
+create function to_string(boolean, smallint, int, bigint, real, float, varchar) returns varchar language javascript as $$
+    export function to_string(a, b, c, d, e, f, g) {
+        return a.toString() + b.toString() + c.toString() + d.toString() + e.toString() + f.toString() + g.toString();
+    }
+$$;
+
+query I
+select int_42();
+----
+42
+
+query I
+select gcd(25, 15);
+----
+5
+
+query T
+select to_string(false, 1::smallint, 2, 3, 4.5, 6.7, 'abc');
+----
+false1234.56.7abc
+
+statement ok
+drop function int_42;
+
+statement ok
+drop function gcd;
+
+statement ok
+drop function to_string;

e2e_test/udf/wasm/Cargo.toml

@@ -8,7 +8,7 @@ edition = "2021"
 crate-type = ["cdylib"]
 
 [dependencies]
-arrow-udf = { git = "https://github.com/risingwavelabs/arrow-udf.git", rev = "f9a9e0d" }
+arrow-udf = { git = "https://github.com/risingwavelabs/arrow-udf.git", rev = "51d941d" }
 genawaiter = "0.99"
 rust_decimal = "1"
 serde_json = "1"

proto/catalog.proto

@@ -216,8 +216,8 @@
   repeated data.DataType arg_types = 5;
   data.DataType return_type = 6;
   string language = 7;
-  string link = 8;
-  string identifier = 10;
+  optional string link = 8;
+  optional string identifier = 10;
   optional string body = 14;
 
   oneof kind {

proto/expr.proto

@@ -472,17 +472,21 @@
   string language = 4;
   // For external UDF: the link to the external function service.
   // For WASM UDF: the link to the wasm binary file.
-  string link = 5;
+  optional string link = 5;
   // An unique identifier for the function.
   // For external UDF, it's the name of the function in the external function service.
   // For WASM UDF, it's the name of the function in the wasm binary file.
-  string identifier = 6;
+  // For JavaScript UDF, it's the name of the function.
+  optional string identifier = 6;
+  // For JavaScript UDF, it's the body of the function.
+  optional string body = 7;
 }
 
 // Additional information for user defined table functions.
 message UserDefinedTableFunction {
   repeated data.DataType arg_types = 3;
   string language = 4;
-  string link = 5;
-  string identifier = 6;
+  optional string link = 5;
+  optional string identifier = 6;
+  optional string body = 7;
 }

src/expr/core/Cargo.toml

@@ -19,6 +19,7 @@ normal = ["workspace-hack", "ctor"]
 anyhow = "1"
 arrow-array = { workspace = true }
 arrow-schema = { workspace = true }
+arrow-udf-js = { workspace = true }
 arrow-udf-wasm = { workspace = true }
 async-trait = "0.1"
 auto_impl = "1"

src/expr/core/src/expr/expr_udf.rs

@@ -20,6 +20,7 @@ use std::time::Duration;
 
 use anyhow::Context;
 use arrow_schema::{Field, Fields, Schema};
+use arrow_udf_js::{CallMode, Runtime as JsRuntime};
 use arrow_udf_wasm::Runtime as WasmRuntime;
 use await_tree::InstrumentAwait;
 use cfg_or_panic::cfg_or_panic;
@@ -61,6 +62,7 @@ const INITIAL_RETRY_COUNT: u8 = 16;
 enum UdfImpl {
     External(Arc<ArrowFlightUdfClient>),
     Wasm(Arc<WasmRuntime>),
+    JavaScript(JsRuntime),
 }
 
 #[async_trait::async_trait]
@@ -123,6 +125,7 @@ impl UserDefinedFunction {
 
         let output: arrow_array::RecordBatch = match &self.imp {
             UdfImpl::Wasm(runtime) => runtime.call(&self.identifier, &input)?,
+            UdfImpl::JavaScript(runtime) => runtime.call(&self.identifier, &input)?,
             UdfImpl::External(client) => {
                 let disable_retry_count = self.disable_retry_count.load(Ordering::Relaxed);
                 let result = if disable_retry_count != 0 {
@@ -189,16 +192,30 @@ impl Build for UserDefinedFunction {
         let return_type = DataType::from(prost.get_return_type().unwrap());
         let udf = prost.get_rex_node().unwrap().as_udf().unwrap();
 
+        let identifier = udf.get_identifier()?;
         let imp = match udf.language.as_str() {
             "wasm" => {
+                let link = udf.get_link()?;
                 // Use `block_in_place` as an escape hatch to run async code here in sync context.
                 // Calling `block_on` directly will panic.
                 UdfImpl::Wasm(tokio::task::block_in_place(|| {
-                    tokio::runtime::Handle::current()
-                        .block_on(get_or_create_wasm_runtime(&udf.link))
+                    tokio::runtime::Handle::current().block_on(get_or_create_wasm_runtime(link))
                 })?)
             }
-            _ => UdfImpl::External(get_or_create_flight_client(&udf.link)?),
+            "javascript" => {
+                let mut rt = JsRuntime::new()?;
+                rt.add_function(
+                    identifier,
+                    arrow_schema::DataType::try_from(&return_type)?,
+                    CallMode::CalledOnNullInput,
+                    udf.get_body()?,
+                )?;
+                UdfImpl::JavaScript(rt)
+            }
+            _ => {
+                let link = udf.get_link()?;
+                UdfImpl::External(get_or_create_flight_client(link)?)
+            }
         };
 
         let arg_schema = Arc::new(Schema::new(
@@ -222,8 +239,8 @@ impl Build for UserDefinedFunction {
             return_type,
             arg_schema,
             imp,
-            identifier: udf.identifier.clone(),
-            span: format!("udf_call({})", udf.identifier).into(),
+            identifier: identifier.clone(),
+            span: format!("udf_call({})", identifier).into(),
             disable_retry_count: AtomicU8::new(0),
         })
     }

src/expr/core/src/table_function/user_defined.rs

@@ -177,28 +177,26 @@ pub fn new_user_defined(prost: &PbTableFunction, chunk_size: usize) -> Result<Bo
             .try_collect::<_, Fields, _>()?,
     ));
 
+    let link = udtf.get_link()?;
     let client = match udtf.language.as_str() {
         "wasm" => {
             // Use `block_in_place` as an escape hatch to run async code here in sync context.
             // Calling `block_on` directly will panic.
             UdfImpl::Wasm(tokio::task::block_in_place(|| {
-                tokio::runtime::Handle::current().block_on(
-                    crate::expr::expr_udf::get_or_create_wasm_runtime(&udtf.link),
-                )
+                tokio::runtime::Handle::current()
+                    .block_on(crate::expr::expr_udf::get_or_create_wasm_runtime(link))
             })?)
         }
         // connect to UDF service
-        _ => UdfImpl::External(crate::expr::expr_udf::get_or_create_flight_client(
-            &udtf.link,
-        )?),
+        _ => UdfImpl::External(crate::expr::expr_udf::get_or_create_flight_client(link)?),
     };
 
     Ok(UserDefinedTableFunction {
         children: prost.args.iter().map(expr_build_from_prost).try_collect()?,
         return_type: prost.return_type.as_ref().expect("no return type").into(),
         arg_schema,
         client,
-        identifier: udtf.identifier.clone(),
+        identifier: udtf.get_identifier()?.clone(),
         chunk_size,
     }
     .boxed())

src/frontend/src/catalog/function_catalog.rs

@@ -29,9 +29,9 @@ pub struct FunctionCatalog {
     pub arg_types: Vec<DataType>,
     pub return_type: DataType,
     pub language: String,
-    pub identifier: String,
+    pub identifier: Option<String>,
     pub body: Option<String>,
-    pub link: String,
+    pub link: Option<String>,
 }
 
 #[derive(Clone, Display, PartialEq, Eq, Hash, Debug)]

src/frontend/src/catalog/system_catalog/rw_catalog/rw_functions.rs

@@ -70,7 +70,7 @@ impl SysCatalogReaderImpl {
                         ))),
                         Some(ScalarImpl::Int32(function.return_type.to_oid())),
                         Some(ScalarImpl::Utf8(function.language.clone().into())),
-                        Some(ScalarImpl::Utf8(function.link.clone().into())),
+                        function.link.clone().map(|s| ScalarImpl::Utf8(s.into())),
                         Some(ScalarImpl::Utf8(
                             get_acl_items(
                                 &Object::FunctionId(function.id.function_id()),

src/frontend/src/expr/table_function.rs

@@ -77,6 +77,7 @@ impl TableFunction {
                     language: c.language.clone(),
                     link: c.link.clone(),
                     identifier: c.identifier.clone(),
+                    body: c.body.clone(),
                 }),
         }
     }

src/frontend/src/expr/user_defined_function.rs

@@ -54,10 +54,9 @@ impl UserDefinedFunction {
             arg_types,
             return_type,
             language: udf.get_language().clone(),
-            identifier: udf.get_identifier().clone(),
-            // TODO: Ensure if we need `body` here
-            body: None,
-            link: udf.get_link().clone(),
+            identifier: udf.identifier.clone(),
+            body: udf.body.clone(),
+            link: udf.link.clone(),
         };
 
         Ok(Self {
@@ -90,6 +89,7 @@ impl Expr for UserDefinedFunction {
                 language: self.catalog.language.clone(),
                 identifier: self.catalog.identifier.clone(),
                 link: self.catalog.link.clone(),
+                body: self.catalog.body.clone(),
             })),
         }
     }