Skip to content

Commit

Permalink
fix(udf): make field name case-insensitive in Rust UDF (#16096)
Browse files Browse the repository at this point in the history
Signed-off-by: Runji Wang <[email protected]>
  • Loading branch information
wangrunji0408 authored Apr 3, 2024
1 parent 02ad72e commit 3c24444
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 1 deletion.
27 changes: 27 additions & 0 deletions e2e_test/udf/rust_udf.slt
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,33 @@ statement ok
drop function key_value;


statement ok
create function key_values(varchar) returns table (Key varchar, Value varchar) language rust as $$
#[derive(StructType)]
struct KeyValue<'a> {
// note that field names are case-insensitive
key: &'a str,
value: &'a str,
}
#[function("key_values(varchar) -> setof struct KeyValue")]
fn key_values(kv: &str) -> impl Iterator<Item = KeyValue<'_>> {
kv.split(',').filter_map(|kv| {
kv.split_once('=')
.map(|(key, value)| KeyValue { key, value })
})
}
$$;

query T
select * from key_values('a=1,b=2');
----
a 1
b 2

statement ok
drop function key_values;


statement ok
create function series(n int) returns table (x int) language rust as $$
fn series(n: i32) -> impl Iterator<Item = i32> {
Expand Down
28 changes: 27 additions & 1 deletion src/frontend/src/handler/create_function.rs
Original file line number Diff line number Diff line change
Expand Up @@ -325,8 +325,34 @@ fn find_wasm_identifier_v2(
runtime: &arrow_udf_wasm::Runtime,
inlined_signature: &str,
) -> Result<String> {
// Inline types in function signature.
//
// # Example
//
// ```text
// types = { "KeyValue": "key:varchar,value:varchar" }
// input = "keyvalue(varchar, varchar) -> struct KeyValue"
// output = "keyvalue(varchar, varchar) -> struct<key:varchar,value:varchar>"
// ```
let inline_types = |s: &str| -> String {
let mut inlined = s.to_string();
// iteratively replace `struct Xxx` with `struct<...>` until no replacement is made.
loop {
let replaced = inlined.clone();
for (k, v) in runtime.types() {
inlined = inlined.replace(&format!("struct {k}"), &format!("struct<{v}>"));
}
if replaced == inlined {
return inlined;
}
}
};
// Function signature in arrow-udf is case sensitive.
// However, SQL identifiers are usually case insensitive and stored in lowercase.
// So we should convert the signature to lowercase before comparison.
let identifier = runtime
.find_function_by_inlined_signature(inlined_signature)
.functions()
.find(|f| inline_types(f).to_lowercase() == inlined_signature)
.ok_or_else(|| {
ErrorCode::InvalidParameterValue(format!(
"function not found in wasm binary: \"{}\"\nHINT: available functions:\n {}",
Expand Down

0 comments on commit 3c24444

Please sign in to comment.