Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(core): support to define param_names and param_types in the list file #925

Merged
merged 1 commit into from
Nov 19, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 34 additions & 6 deletions wren-core-py/src/remote_functions.rs
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,10 @@ pub struct PyRemoteFunction {
pub function_type: String,
pub name: String,
pub return_type: Option<String>,
pub param_names: Option<Vec<String>>,
pub param_types: Option<Vec<String>>,
/// It's a comma separated string of parameter names
pub param_names: Option<String>,
/// It's a comma separated string of parameter types
pub param_types: Option<String>,
pub description: Option<String>,
}

Expand All @@ -54,12 +56,26 @@ impl PyRemoteFunction {

impl From<wren_core::mdl::function::RemoteFunction> for PyRemoteFunction {
fn from(remote_function: wren_core::mdl::function::RemoteFunction) -> Self {
let param_names = remote_function.param_names.map(|names| {
names
.iter()
.map(|name| name.to_string())
.collect::<Vec<String>>()
.join(",")
});
let param_types = remote_function.param_types.map(|types| {
types
.iter()
.map(|t| t.to_string())
.collect::<Vec<String>>()
.join(",")
});
Self {
function_type: remote_function.function_type.to_string(),
name: remote_function.name,
return_type: Some(remote_function.return_type),
param_names: remote_function.param_names,
param_types: remote_function.param_types,
param_names,
param_types,
description: remote_function.description,
}
}
Expand All @@ -69,14 +85,26 @@ impl From<PyRemoteFunction> for wren_core::mdl::function::RemoteFunction {
fn from(
remote_function: PyRemoteFunction,
) -> wren_core::mdl::function::RemoteFunction {
let param_names = remote_function.param_names.map(|names| {
names
.split(",")
.map(|name| name.to_string())
.collect::<Vec<String>>()
});
let param_types = remote_function.param_types.map(|types| {
types
.split(",")
.map(|t| t.to_string())
.collect::<Vec<String>>()
});
wren_core::mdl::function::RemoteFunction {
function_type: FunctionType::from_str(&remote_function.function_type)
.unwrap(),
name: remote_function.name,
// TODO: Get the return type form DataFusion SessionState
return_type: remote_function.return_type.unwrap_or("string".to_string()),
param_names: remote_function.param_names,
param_types: remote_function.param_types,
param_names,
param_types,
description: remote_function.description,
}
}
Expand Down
6 changes: 3 additions & 3 deletions wren-core-py/tests/functions.csv
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
function_type,name,return_type,description
scalar,add_two,int,"Adds two numbers together."
window,max_if,int,"If the condition is true, returns the maximum value in the window."
function_type,name,return_type,param_names,param_types,description
scalar,add_two,int,"f1,f2","int,int","Adds two numbers together."
window,max_if,int,,,"If the condition is true, returns the maximum value in the window."
10 changes: 10 additions & 0 deletions wren-core-py/tests/test_modeling_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -63,3 +63,13 @@ def test_get_available_functions():
assert add_two["name"] == "add_two"
assert add_two["function_type"] == "scalar"
assert add_two["description"] == "Adds two numbers together."
assert add_two["return_type"] == "int"
assert add_two["param_names"] == "f1,f2"
assert add_two["param_types"] == "int,int"

max_if = next(filter(lambda x: x["name"] == "max_if", map(lambda x: x.to_dict(), functions)))
assert max_if["name"] == "max_if"
assert max_if["function_type"] == "window"
assert max_if["param_names"] is None
assert max_if["param_types"] is None