Allow to pass extra info to table functions (#637)

Via `register_table_function_with_options`

Author: mlafeldt

crates/duckdb/src/vtab/function.rs

@@ -1,4 +1,5 @@
 use super::{
+    drop_boxed,
     ffi::{
         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,
@@ -36,6 +37,7 @@ impl BindInfo {
             duckdb_bind_add_result_column(self.ptr, c_str.as_ptr() as *const c_char, column_type.ptr);
         }
     }
+
     /// Report that an error has occurred while calling bind.
     ///
     /// # Arguments
@@ -49,18 +51,20 @@ impl BindInfo {
     /// Sets the user-provided bind data in the bind object. This object can be retrieved again during execution.
     ///
     /// # Arguments
-    ///  * `extra_data`: The bind data object.
-    ///  * `destroy`: The callback that will be called to destroy the bind data (if any)
+    ///  * `data`: The bind data object.
+    ///  * `free_function`: The callback that will be called to destroy the bind data (if any)
     ///
     /// # Safety
-    ///
+    /// `data` must be a valid pointer, and `free_function` must properly free it.
     pub unsafe fn set_bind_data(&self, data: *mut c_void, free_function: Option<unsafe extern "C" fn(*mut c_void)>) {
         duckdb_bind_set_bind_data(self.ptr, data, free_function);
     }
+
     /// Retrieves the number of regular (non-named) parameters to the function.
     pub fn get_parameter_count(&self) -> u64 {
         unsafe { duckdb_bind_get_parameter_count(self.ptr) }
     }
+
     /// Retrieves the parameter at the given index.
     ///
     /// # Arguments
@@ -107,10 +111,7 @@ impl BindInfo {
     pub fn set_cardinality(&self, cardinality: idx_t, is_exact: bool) {
         unsafe { duckdb_bind_set_cardinality(self.ptr, cardinality, is_exact) }
     }
-    /// Retrieves the extra info of the function as set in [`TableFunction::set_extra_info`]
-    ///
-    /// # Arguments
-    /// * `returns`: The extra info
+    /// Retrieves the extra info of the function as set in [`TableFunction::with_extra_info`].
     pub fn get_extra_info<T>(&self) -> *const T {
         unsafe { duckdb_bind_get_extra_info(self.ptr).cast() }
     }
@@ -162,13 +163,11 @@ impl InitInfo {
         indices
     }
 
-    /// Retrieves the extra info of the function as set in [`TableFunction::set_extra_info`]
-    ///
-    /// # Arguments
-    /// * `returns`: The extra info
+    /// Retrieves the extra info of the function as set in [`TableFunction::with_extra_info`].
     pub fn get_extra_info<T>(&self) -> *const T {
         unsafe { duckdb_init_get_extra_info(self.0).cast() }
     }
+
     /// Gets the bind data set by [`BindInfo::set_bind_data`] during the bind.
     ///
     /// Note that the bind data should be considered as read-only.
@@ -179,13 +178,15 @@ impl InitInfo {
     pub fn get_bind_data<T>(&self) -> *const T {
         unsafe { duckdb_init_get_bind_data(self.0).cast() }
     }
+
     /// Sets how many threads can process this table function in parallel (default: 1)
     ///
     /// # Arguments
     /// * `max_threads`: The maximum amount of threads that can process this table function
     pub fn set_max_threads(&self, max_threads: idx_t) {
         unsafe { duckdb_init_set_max_threads(self.0, max_threads) }
     }
+
     /// Report that an error has occurred while calling init.
     ///
     /// # Arguments
@@ -307,15 +308,35 @@ impl TableFunction {
 
     /// Assigns extra information to the table function that can be fetched during binding, etc.
     ///
+    /// For most use cases, prefer [`with_extra_info`](Self::with_extra_info) which handles memory management automatically.
+    ///
     /// # Arguments
     /// * `extra_info`: The extra information
     /// * `destroy`: The callback that will be called to destroy the bind data (if any)
     ///
     /// # Safety
+    /// The caller must ensure that `extra_info` is a valid pointer and that `destroy`
+    /// properly cleans up the data when called.
     pub unsafe fn set_extra_info(&self, extra_info: *mut c_void, destroy: duckdb_delete_callback_t) {
+        duckdb_table_function_set_extra_info(self.ptr, extra_info, destroy);
+    }
+
+    /// Assigns extra information to the table function that can be fetched during binding, init, and execution.
+    ///
+    /// This is a safe wrapper around [`set_extra_info`](Self::set_extra_info) that handles memory management automatically.
+    ///
+    /// # Arguments
+    /// * `info`: The extra information to store
+    pub fn with_extra_info<T>(&self, info: T) -> &Self
+    where
+        T: Send + Sync + 'static,
+    {
         unsafe {
-            duckdb_table_function_set_extra_info(self.ptr, extra_info, destroy);
+            let boxed = Box::new(info);
+            let ptr = Box::into_raw(boxed) as *mut c_void;
+            self.set_extra_info(ptr, Some(drop_boxed::<T>));
         }
+        self
     }
 
     /// Sets the thread-local init function of the table function
@@ -383,13 +404,11 @@ impl<V: VTab> TableFunctionInfo<V> {
         }
     }
 
-    /// Retrieves the extra info of the function as set in [`TableFunction::set_extra_info`]
-    ///
-    /// # Arguments
-    /// * `returns`: The extra info
+    /// Retrieves the extra info of the function as set in [`TableFunction::with_extra_info`].
     pub fn get_extra_info<T>(&self) -> *mut T {
         unsafe { duckdb_function_get_extra_info(self.ptr).cast() }
     }
+
     /// Gets the thread-local init data set by [`InitInfo::set_init_data`] during the local_init.
     ///
     /// # Arguments

crates/duckdb/src/vtab/mod.rs

@@ -23,6 +23,20 @@ mod excel;
 pub use function::{BindInfo, InitInfo, TableFunction, TableFunctionInfo};
 pub use value::Value;
 
+/// Options for registering a table function.
+pub struct TableFunctionOptions<E> {
+    /// Extra info passed to the function at runtime.
+    /// Accessible via `BindInfo::get_extra_info`, `InitInfo::get_extra_info`,
+    /// or `TableFunctionInfo::get_extra_info`.
+    pub extra_info: Option<E>,
+}
+
+impl<E> Default for TableFunctionOptions<E> {
+    fn default() -> Self {
+        Self { extra_info: None }
+    }
+}
+
 use crate::core::{DataChunkHandle, LogicalTypeHandle};
 use ffi::{duckdb_bind_info, duckdb_data_chunk, duckdb_function_info, duckdb_init_info};
 
@@ -150,6 +164,35 @@ impl Connection {
         }
         self.db.borrow_mut().register_table_function(table_function)
     }
+
+    /// Register the given TableFunction with options.
+    #[inline]
+    pub fn register_table_function_with_options<T: VTab, E>(
+        &self,
+        name: &str,
+        options: TableFunctionOptions<E>,
+    ) -> Result<()>
+    where
+        E: Send + Sync + 'static,
+    {
+        let table_function = TableFunction::default();
+        table_function
+            .set_name(name)
+            .supports_pushdown(T::supports_pushdown())
+            .set_bind(Some(bind::<T>))
+            .set_init(Some(init::<T>))
+            .set_function(Some(func::<T>));
+        if let Some(extra_info) = options.extra_info {
+            table_function.with_extra_info(extra_info);
+        }
+        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)
+    }
 }
 
 impl InnerConnection {
@@ -287,6 +330,62 @@ mod test {
         Ok(())
     }
 
+    // Test table function with extra info
+    struct PrefixVTab;
+
+    impl VTab for PrefixVTab {
+        type InitData = HelloInitData;
+        type BindData = HelloBindData;
+
+        fn bind(bind: &BindInfo) -> Result<Self::BindData, Box<dyn Error>> {
+            bind.add_result_column("column0", LogicalTypeHandle::from(LogicalTypeId::Varchar));
+            let name = bind.get_parameter(0).to_string();
+            Ok(HelloBindData { name })
+        }
+
+        fn init(_: &InitInfo) -> Result<Self::InitData, Box<dyn Error>> {
+            Ok(HelloInitData {
+                done: AtomicBool::new(false),
+            })
+        }
+
+        fn func(func: &TableFunctionInfo<Self>, output: &mut DataChunkHandle) -> Result<(), Box<dyn Error>> {
+            let init_data = func.get_init_data();
+            let bind_data = func.get_bind_data();
+            let prefix = unsafe { &*func.get_extra_info::<String>() };
+
+            if init_data.done.swap(true, Ordering::Relaxed) {
+                output.set_len(0);
+            } else {
+                let vector = output.flat_vector(0);
+                let result = CString::new(format!("{prefix} {}", bind_data.name))?;
+                vector.insert(0, result);
+                output.set_len(1);
+            }
+            Ok(())
+        }
+
+        fn parameters() -> Option<Vec<LogicalTypeHandle>> {
+            Some(vec![LogicalTypeHandle::from(LogicalTypeId::Varchar)])
+        }
+    }
+
+    #[test]
+    fn test_table_function_with_options() -> Result<(), Box<dyn Error>> {
+        let conn = Connection::open_in_memory()?;
+        conn.register_table_function_with_options::<PrefixVTab, _>(
+            "greet",
+            TableFunctionOptions {
+                extra_info: Some("Howdy".to_string()),
+            },
+        )?;
+
+        let val = conn.query_row("select * from greet('partner')", [], |row| <(String,)>::try_from(row))?;
+        assert_eq!(val, ("Howdy partner".to_string(),));
+
+        Ok(())
+    }
+
     #[cfg(feature = "vtab-loadable")]
     use duckdb_loadable_macros::duckdb_entrypoint;