support json and decimal

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

Cargo.toml

@@ -131,7 +131,7 @@ arrow-flight = "49"
 arrow-select = "49"
 arrow-ord = "49"
 arrow-row = "49"
-arrow-udf-wasm = { git = "https://github.com/risingwavelabs/arrow-udf.git", rev = "34e1c9d" }
+arrow-udf-wasm = { git = "https://github.com/risingwavelabs/arrow-udf.git", rev = "3ac8371" }
 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/wasm/Cargo.toml

@@ -8,4 +8,6 @@ edition = "2021"
 crate-type = ["cdylib"]
 
 [dependencies]
-arrow-udf = { git = "https://github.com/risingwavelabs/arrow-udf-wasm.git", rev = "34e1c9d" }
+arrow-udf = { git = "https://github.com/risingwavelabs/arrow-udf.git", rev = "3ac8371" }
+rust_decimal = "1"
+serde_json = "1"

e2e_test/udf/wasm/src/lib.rs

@@ -1,4 +1,5 @@
 use arrow_udf::function;
+use rust_decimal::Decimal;
 
 #[function("int_42() -> int")]
 fn int_42() -> i32 {
@@ -62,3 +63,13 @@ fn extract_tcp_info(tcp_packet: &[u8]) -> (String, String, i16, i16) {
         dst_port as i16,
     )
 }
+
+#[function("decimal_add(decimal, decimal) -> decimal")]
+fn decimal_add(a: Decimal, b: Decimal) -> Decimal {
+    a + b
+}
+
+#[function("jsonb_access(json, int) -> json")]
+fn jsonb_access(json: serde_json::Value, index: i32) -> Option<serde_json::Value> {
+    json.get(index as usize).cloned()
+}

e2e_test/udf/wasm_udf.slt

@@ -17,6 +17,14 @@ statement ok
 create function extract_tcp_info(bytea) returns struct<src_addr varchar, dst_addr varchar, src_port smallint, dst_port smallint>
 language wasm using link 'fs://e2e_test/udf/wasm/target/wasm32-wasi/release/udf.wasm';
 
+statement ok
+create function decimal_add(decimal, decimal) returns decimal
+language wasm using link 'fs://e2e_test/udf/wasm/target/wasm32-wasi/release/udf.wasm';
+
+statement ok
+create function jsonb_access(jsonb, int) returns jsonb
+language wasm using link 'fs://e2e_test/udf/wasm/target/wasm32-wasi/release/udf.wasm';
+
 query I
 select int_42();
 ----
@@ -37,6 +45,19 @@ select extract_tcp_info(E'\\x45000034a8a8400040065b8ac0a8000ec0a80001035d20b6d97
 ----
 (192.168.0.14,192.168.0.1,861,8374)
 
+query R
+select decimal_add(1.11, 2.22);
+----
+3.33
+
+query T
+select jsonb_access(a::jsonb, 1) from
+(values ('["a", "b", "c"]'), (null), ('[0, false]')) t(a);
+----
+"b"
+NULL
+false
+
 statement ok
 drop function int_42;
 
@@ -48,3 +69,9 @@ drop function gcd(int,int,int);
 
 statement ok
 drop function extract_tcp_info;
+
+statement ok
+drop function decimal_add;
+
+statement ok
+drop function jsonb_access;

src/expr/udf/README.md

@@ -50,7 +50,7 @@ fn gcd(mut x: i32, mut y: i32) -> i32 {
 You can find more usages in the [documentation](https://docs.rs/arrow_udf/0.1.0/arrow_udf/attr.function.html) and more examples in the [tests](https://github.com/risingwavelabs/arrow-udf/blob/main/arrow-udf/tests/tests.rs).
 
 Currently we only support scalar functions with a limited set of data types.
-`decimal`, `timestamptz`, `jsonb` and complex array types are not supported yet.
+`timestamptz` and complex array types are not supported yet.
 
 ## 3. Build the project
 

src/frontend/src/handler/create_function.rs

@@ -284,7 +284,7 @@ fn datatype_name(ty: &DataType) -> String {
         DataType::Timestamptz => "timestamptz".to_string(),
         DataType::Interval => "interval".to_string(),
         DataType::Decimal => "decimal".to_string(),
-        DataType::Jsonb => "jsonb".to_string(),
+        DataType::Jsonb => "json".to_string(),
         DataType::Serial => "serial".to_string(),
         DataType::Int256 => "int256".to_string(),
         DataType::Bytea => "bytea".to_string(),