feat(udf): support inlined Rust UDF (#14903)

Signed-off-by: Runji Wang <[email protected]>

Cargo.toml

@@ -133,8 +133,8 @@ arrow-flight = "50"
 arrow-select = "50"
 arrow-ord = "50"
 arrow-row = "50"
-arrow-udf-js = { git = "https://github.com/risingwavelabs/arrow-udf.git", rev = "7ba1c22" }
-arrow-udf-wasm = "0.1"
+arrow-udf-js = "0.1"
+arrow-udf-wasm = { version = "0.1.2", features = ["build"] }
 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" }

ci/Dockerfile

@@ -41,6 +41,7 @@ RUN curl -sSL https://install.python-poetry.org | python3 -
 
 # add required rustup components
 RUN rustup component add rustfmt llvm-tools-preview clippy
+RUN rustup target add wasm32-wasi
 
 ENV CARGO_REGISTRIES_CRATES_IO_PROTOCOL=sparse
 

ci/build-ci-image.sh

@@ -10,7 +10,7 @@ cat ../rust-toolchain
 # shellcheck disable=SC2155
 
 # REMEMBER TO ALSO UPDATE ci/docker-compose.yml
-export BUILD_ENV_VERSION=v20240124_1
+export BUILD_ENV_VERSION=v20240204
 
 export BUILD_TAG="public.ecr.aws/x5u3w5h6/rw-build-env:${BUILD_ENV_VERSION}"
 

ci/docker-compose.yml

@@ -71,7 +71,7 @@ services:
       retries: 5
 
   source-test-env:
-    image: public.ecr.aws/x5u3w5h6/rw-build-env:v20240124_1
+    image: public.ecr.aws/x5u3w5h6/rw-build-env:v20240204
     depends_on:
       - mysql
       - db
@@ -81,7 +81,7 @@ services:
       - ..:/risingwave
 
   sink-test-env:
-    image: public.ecr.aws/x5u3w5h6/rw-build-env:v20240124_1
+    image: public.ecr.aws/x5u3w5h6/rw-build-env:v20240204
     depends_on:
       - mysql
       - db
@@ -93,12 +93,12 @@ services:
       - ..:/risingwave
 
   rw-build-env:
-    image: public.ecr.aws/x5u3w5h6/rw-build-env:v20240124_1
+    image: public.ecr.aws/x5u3w5h6/rw-build-env:v20240204
     volumes:
       - ..:/risingwave
 
   ci-flamegraph-env:
-    image: public.ecr.aws/x5u3w5h6/rw-build-env:v20240124_1
+    image: public.ecr.aws/x5u3w5h6/rw-build-env:v20240204
     # NOTE(kwannoel): This is used in order to permit
     # syscalls for `nperf` (perf_event_open),
     # so it can do CPU profiling.
@@ -109,7 +109,7 @@ services:
       - ..:/risingwave
 
   regress-test-env:
-    image: public.ecr.aws/x5u3w5h6/rw-build-env:v20240124_1
+    image: public.ecr.aws/x5u3w5h6/rw-build-env:v20240204
     depends_on:
       db:
         condition: service_healthy

ci/scripts/build-other.sh

@@ -8,7 +8,6 @@ source ci/scripts/common.sh
 
 echo "--- Build Rust UDF"
 cd e2e_test/udf/wasm
-rustup target add wasm32-wasi
 cargo build --release
 cd ../../..
 

docker/Dockerfile

@@ -33,7 +33,8 @@ COPY rust-toolchain rust-toolchain
 RUN rustup self update \
   && rustup set profile minimal \
   && rustup show \
-  && rustup component add rustfmt
+  && rustup component add rustfmt \
+  && rustup target add wasm32-wasi
 
 RUN cargo install flamegraph
 # TODO: cargo-chef doesn't work well now, because we update Cargo.lock very often.

docker/Dockerfile.hdfs

@@ -34,7 +34,8 @@ RUN unzip dashboard-artifact.zip && mv risingwave-dashboard-artifact /risingwave
 RUN rustup self update \
   && rustup set profile minimal \
   && rustup show \
-  && rustup component add rustfmt
+  && rustup component add rustfmt \
+  && rustup target add wasm32-wasi
 
 RUN cargo fetch
 

e2e_test/udf/wasm_udf.slt

@@ -91,3 +91,24 @@ drop function jsonb_access;
 
 statement ok
 drop function series;
+
+# inlined rust function
+statement ok
+create function gcd(int, int) returns int language rust as $$
+    fn gcd(mut a: i32, mut b: i32) -> i32 {
+        while b != 0 {
+            let t = b;
+            b = a % b;
+            a = t;
+        }
+        a
+    }
+$$;
+
+query I
+select gcd(25, 15);
+----
+5
+
+statement ok
+drop function gcd;

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

@@ -194,7 +194,7 @@ impl Build for UserDefinedFunction {
 
         let identifier = udf.get_identifier()?;
         let imp = match udf.language.as_str() {
-            "wasm" => {
+            "wasm" | "rust" => {
                 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.

src/frontend/src/handler/create_function.rs

@@ -24,9 +24,7 @@ use risingwave_expr::expr::get_or_create_wasm_runtime;
 use risingwave_object_store::object::{build_remote_object_store, ObjectStoreConfig};
 use risingwave_pb::catalog::function::{Kind, ScalarFunction, TableFunction};
 use risingwave_pb::catalog::Function;
-use risingwave_sqlparser::ast::{
-    CreateFunctionBody, FunctionDefinition, ObjectName, OperateFunctionArg,
-};
+use risingwave_sqlparser::ast::{CreateFunctionBody, ObjectName, OperateFunctionArg};
 use risingwave_storage::monitor::ObjectStoreMetrics;
 use risingwave_udf::ArrowFlightUdfClient;
 
@@ -54,7 +52,7 @@ pub async fn handle_create_function(
         Some(lang) => {
             let lang = lang.real_value().to_lowercase();
             match &*lang {
-                "python" | "java" | "wasm" | "javascript" => lang,
+                "python" | "java" | "wasm" | "rust" | "javascript" => lang,
                 _ => {
                     return Err(ErrorCode::InvalidParameterValue(format!(
                         "language {} is not supported",
@@ -134,12 +132,12 @@ pub async fn handle_create_function(
                 )
                 .into());
             };
-            let Some(FunctionDefinition::SingleQuotedDef(id)) = params.as_ else {
+            let Some(as_) = params.as_ else {
                 return Err(
                     ErrorCode::InvalidParameterValue("AS must be specified".to_string()).into(),
                 );
             };
-            identifier = id;
+            identifier = as_.into_string();
 
             // check UDF server
             {
@@ -173,16 +171,56 @@ pub async fn handle_create_function(
         }
         "javascript" => {
             identifier = function_name.to_string();
-            body = Some(match params.as_ {
-                Some(FunctionDefinition::SingleQuotedDef(s)) => s,
-                Some(FunctionDefinition::DoubleDollarDef(s)) => s,
-                _ => {
-                    return Err(ErrorCode::InvalidParameterValue(
-                        "AS must be specified".to_string(),
-                    )
-                    .into())
-                }
-            });
+            body = Some(
+                params
+                    .as_
+                    .ok_or_else(|| ErrorCode::InvalidParameterValue("AS must be specified".into()))?
+                    .into_string(),
+            );
+        }
+        "rust" => {
+            identifier = wasm_identifier(
+                &function_name,
+                &arg_types,
+                &return_type,
+                matches!(kind, Kind::Table(_)),
+            );
+            if params.using.is_some() {
+                return Err(ErrorCode::InvalidParameterValue(
+                    "USING is not supported for rust function".to_string(),
+                )
+                .into());
+            }
+            let function_body = params
+                .as_
+                .ok_or_else(|| ErrorCode::InvalidParameterValue("AS must be specified".into()))?
+                .into_string();
+            let script = format!("#[arrow_udf::function(\"{identifier}\")]\n{function_body}");
+            body = Some(function_body.clone());
+
+            let wasm_binary =
+                tokio::task::spawn_blocking(move || arrow_udf_wasm::build::build("", &script))
+                    .await?
+                    .context("failed to build rust function")?;
+
+            // below is the same as `wasm` language
+            let runtime = arrow_udf_wasm::Runtime::new(&wasm_binary)?;
+            check_wasm_function(&runtime, &identifier)?;
+
+            let system_params = session.env().meta_client().get_system_params().await?;
+            let object_name = format!("{:?}.wasm", md5::compute(&wasm_binary));
+            upload_wasm_binary(
+                system_params.wasm_storage_url(),
+                &object_name,
+                wasm_binary.into(),
+            )
+            .await?;
+
+            link = Some(format!(
+                "{}/{}",
+                system_params.wasm_storage_url(),
+                object_name
+            ));
         }
         "wasm" => {
             identifier = wasm_identifier(

src/sqlparser/src/ast/mod.rs

@@ -2642,6 +2642,24 @@ impl fmt::Display for FunctionDefinition {
     }
 }
 
+impl FunctionDefinition {
+    /// Returns the function definition as a string slice.
+    pub fn as_str(&self) -> &str {
+        match self {
+            FunctionDefinition::SingleQuotedDef(s) => s,
+            FunctionDefinition::DoubleDollarDef(s) => s,
+        }
+    }
+
+    /// Returns the function definition as a string.
+    pub fn into_string(self) -> String {
+        match self {
+            FunctionDefinition::SingleQuotedDef(s) => s,
+            FunctionDefinition::DoubleDollarDef(s) => s,
+        }
+    }
+}
+
 /// Return types of a function.
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]