Skip to content

Commit

Permalink
feat: add named parameter support to virtual tables (#250)
Browse files Browse the repository at this point in the history
* feat: add named parameter support

* safely handle missing parameters
  • Loading branch information
Mause authored Feb 8, 2024
1 parent 9af377a commit d48b967
Show file tree
Hide file tree
Showing 2 changed files with 100 additions and 8 deletions.
58 changes: 50 additions & 8 deletions src/vtab/function.rs
Original file line number Diff line number Diff line change
@@ -1,12 +1,13 @@
use super::{
ffi::{
duckdb_bind_add_result_column, duckdb_bind_get_extra_info, duckdb_bind_get_parameter,
duckdb_bind_get_parameter_count, duckdb_bind_info, duckdb_bind_set_bind_data, duckdb_bind_set_cardinality,
duckdb_bind_set_error, duckdb_create_table_function, duckdb_data_chunk, duckdb_delete_callback_t,
duckdb_destroy_table_function, duckdb_table_function, duckdb_table_function_add_parameter,
duckdb_table_function_init_t, duckdb_table_function_set_bind, duckdb_table_function_set_extra_info,
duckdb_table_function_set_function, duckdb_table_function_set_init, duckdb_table_function_set_local_init,
duckdb_table_function_set_name, duckdb_table_function_supports_projection_pushdown, idx_t,
duckdb_bind_add_result_column, duckdb_bind_get_extra_info, duckdb_bind_get_named_parameter,
duckdb_bind_get_parameter, duckdb_bind_get_parameter_count, duckdb_bind_info, duckdb_bind_set_bind_data,
duckdb_bind_set_cardinality, duckdb_bind_set_error, duckdb_create_table_function, duckdb_data_chunk,
duckdb_delete_callback_t, duckdb_destroy_table_function, duckdb_table_function,
duckdb_table_function_add_named_parameter, duckdb_table_function_add_parameter, duckdb_table_function_init_t,
duckdb_table_function_set_bind, duckdb_table_function_set_extra_info, duckdb_table_function_set_function,
duckdb_table_function_set_init, duckdb_table_function_set_local_init, duckdb_table_function_set_name,
duckdb_table_function_supports_projection_pushdown, idx_t,
},
LogicalType, Value,
};
Expand Down Expand Up @@ -64,8 +65,36 @@ impl BindInfo {
/// * `index`: The index of the parameter to get
///
/// returns: The value of the parameter
///
/// # Panics
/// If requested parameter is out of range for function definition
pub fn get_parameter(&self, param_index: u64) -> Value {
unsafe { Value::from(duckdb_bind_get_parameter(self.ptr, param_index)) }
unsafe {
let ptr = duckdb_bind_get_parameter(self.ptr, param_index);
if ptr.is_null() {
panic!("{} is out of range for function definition", param_index);
} else {
Value::from(ptr)
}
}
}

/// Retrieves the named parameter with the given name.
///
/// # Arguments
/// * `name`: The name of the parameter to get
///
/// returns: The value of the parameter
pub fn get_named_parameter(&self, name: &str) -> Option<Value> {
unsafe {
let name = &CString::new(name).unwrap();
let ptr = duckdb_bind_get_named_parameter(self.ptr, name.as_ptr());
if ptr.is_null() {
None
} else {
Some(Value::from(ptr))
}
}
}

/// Sets the cardinality estimate for the table function, used for optimization.
Expand Down Expand Up @@ -204,6 +233,19 @@ impl TableFunction {
self
}

/// Adds a named parameter to the table function.
///
/// # Arguments
/// * `name`: The name of the parameter to add.
/// * `logical_type`: The type of the parameter to add.
pub fn add_named_parameter(&self, name: &str, logical_type: &LogicalType) -> &Self {
unsafe {
let string = CString::new(name).unwrap();
duckdb_table_function_add_named_parameter(self.ptr, string.as_ptr(), logical_type.ptr);
}
self
}

/// Sets the main function of the table function
///
/// # Arguments
Expand Down
50 changes: 50 additions & 0 deletions src/vtab/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,11 @@ pub trait VTab: Sized {
fn parameters() -> Option<Vec<LogicalType>> {
None
}
/// The named parameters of the table function
/// default is None
fn named_parameters() -> Option<Vec<(String, LogicalType)>> {
None
}
}

unsafe extern "C" fn func<T>(info: duckdb_function_info, output: duckdb_data_chunk)
Expand Down Expand Up @@ -135,6 +140,9 @@ impl Connection {
for ty in T::parameters().unwrap_or_default() {
table_function.add_parameter(&ty);
}
for (name, ty) in T::named_parameters().unwrap_or_default() {
table_function.add_named_parameter(&name, &ty);
}
self.db.borrow_mut().register_table_function(table_function)
}
}
Expand Down Expand Up @@ -232,13 +240,55 @@ mod test {
}
}

struct HelloWithNamedVTab {}
impl VTab for HelloWithNamedVTab {
type InitData = HelloInitData;
type BindData = HelloBindData;

fn bind(bind: &BindInfo, data: *mut HelloBindData) -> Result<(), Box<dyn Error>> {
bind.add_result_column("column0", LogicalType::new(LogicalTypeId::Varchar));
let param = bind.get_named_parameter("name").unwrap().to_string();
assert!(bind.get_named_parameter("unknown_name").is_none());
unsafe {
(*data).name = CString::new(param).unwrap().into_raw();
}
Ok(())
}

fn init(init_info: &InitInfo, data: *mut HelloInitData) -> Result<(), Box<dyn Error>> {
HelloVTab::init(init_info, data)
}

fn func(func: &FunctionInfo, output: &mut DataChunk) -> Result<(), Box<dyn Error>> {
HelloVTab::func(func, output)
}

fn named_parameters() -> Option<Vec<(String, LogicalType)>> {
Some(vec![("name".to_string(), LogicalType::new(LogicalTypeId::Varchar))])
}
}

#[test]
fn test_table_function() -> Result<(), Box<dyn Error>> {
let conn = Connection::open_in_memory()?;
conn.register_table_function::<HelloVTab>("hello")?;

let val = conn.query_row("select * from hello('duckdb')", [], |row| <(String,)>::try_from(row))?;
assert_eq!(val, ("Hello duckdb".to_string(),));

Ok(())
}

#[test]
fn test_named_table_function() -> Result<(), Box<dyn Error>> {
let conn = Connection::open_in_memory()?;
conn.register_table_function::<HelloWithNamedVTab>("hello_named")?;

let val = conn.query_row("select * from hello_named(name = 'duckdb')", [], |row| {
<(String,)>::try_from(row)
})?;
assert_eq!(val, ("Hello duckdb".to_string(),));

Ok(())
}

Expand Down

0 comments on commit d48b967

Please sign in to comment.