better_duck_core/udf/table/
mod.rs

1//! DuckDB table functions: functions used in a `FROM` clause that produce rows
2//! and columns, e.g. `SELECT * FROM my_func(1, 2)`.
3
4mod function;
5mod row;
6
7use std::ffi::CString;
8
9use crate::{
10    connection::Connection,
11    error::Result,
12    ffi::{duckdb_bind_info, duckdb_data_chunk, duckdb_function_info, duckdb_init_info},
13};
14
15use self::function::TableFunction;
16pub use self::function::{BindInfo, InitInfo, TableFunctionInfo};
17pub use self::row::{run_table_func, TableInitData, TableRow};
18use super::{callback::contain_callback, data_chunk::DataChunkHandle};
19use crate::types::LogicalType;
20
21/// A DuckDB table function: produces rows and columns for use in a `FROM`
22/// clause, e.g. `SELECT * FROM my_func(1, 2)`.
23///
24/// See the callback containment contract in [`crate::udf`].
25pub trait VTab: Sized {
26    /// Data produced once by [`VTab::bind`] and shared, read-only, by every
27    /// later call to [`VTab::init`] and [`VTab::func`] for this query.
28    type BindData: Send + Sync;
29
30    /// Data produced once by [`VTab::init`], shared across every worker thread
31    /// executing this query. Any interior mutation must be synchronized.
32    type InitData: Send + Sync;
33
34    /// The positional parameter types this function accepts, in order. Empty by
35    /// default (no positional parameters).
36    ///
37    /// # Errors
38    ///
39    /// Returns an error if a parameter's logical type cannot be built.
40    fn parameters() -> Result<Vec<LogicalType>> {
41        Ok(Vec::new())
42    }
43
44    /// The named (keyword) parameters this function accepts, as
45    /// `(name, type)` pairs. Empty by default (no named parameters). Read back
46    /// during [`VTab::bind`] via [`BindInfo::get_named_parameter`].
47    ///
48    /// # Errors
49    ///
50    /// Returns an error if a parameter's logical type cannot be built.
51    fn named_parameters() -> Result<Vec<(String, LogicalType)>> {
52        Ok(Vec::new())
53    }
54
55    /// Whether this function can consume a projected column list — i.e. honors
56    /// [`InitInfo::column_indices`] in [`VTab::init`]/[`VTab::func`] to skip
57    /// producing columns the query doesn't need. `false` by default (every
58    /// column is always written).
59    fn supports_projection_pushdown() -> bool {
60        false
61    }
62
63    /// Determines the function's output schema (via
64    /// [`BindInfo::add_result_column`]) and produces the data shared by every
65    /// later call for this query.
66    ///
67    /// # Errors
68    ///
69    /// Returns an error to fail the query with that message.
70    fn bind(bind: &BindInfo) -> super::UdfResult<Self::BindData>;
71
72    /// Produces data shared across every worker thread executing this query,
73    /// e.g. the initial position of a cursor.
74    ///
75    /// # Errors
76    ///
77    /// Returns an error to fail the query with that message.
78    fn init(init: &InitInfo<Self>) -> super::UdfResult<Self::InitData>;
79
80    /// Writes up to `output.capacity()` rows into `output`, then calls
81    /// `output.set_len(k)` with the number actually written. Called repeatedly
82    /// until `set_len(0)` (or `output.len() == 0` on entry, if untouched) signals
83    /// the scan is complete.
84    ///
85    /// # Errors
86    ///
87    /// Returns an error to fail the query with that message.
88    fn func(
89        func: &TableFunctionInfo<Self>,
90        output: &mut DataChunkHandle,
91    ) -> super::UdfResult<()>;
92}
93
94/// A [`VTab`] that also produces per-worker-thread ("local") init data, in
95/// addition to the query-wide init data every `VTab` already produces.
96///
97/// Registered via [`Connection::register_table_function_ext`] instead of
98/// [`Connection::register_table_function`]. Useful for state that must not be
99/// shared across threads (e.g. a non-`Sync` cursor or scratch buffer) when
100/// [`InitInfo::set_max_threads`] allows more than one worker.
101pub trait VTabLocalInit: VTab {
102    /// Data produced once per worker thread, retrievable on that same thread
103    /// via [`TableFunctionInfo::local_init_data`].
104    type LocalInitData: Send + Sync;
105
106    /// Produces this thread's local init data.
107    ///
108    /// # Errors
109    ///
110    /// Returns an error to fail the query with that message.
111    fn local_init(init: &InitInfo<Self>) -> super::UdfResult<Self::LocalInitData>;
112}
113
114/// The C trampoline installed via `duckdb_table_function_set_bind`.
115///
116/// See [`scalar_trampoline`](super::scalar) for why containment must be the
117/// outermost thing here.
118unsafe extern "C" fn table_bind_trampoline<T: VTab>(info: duckdb_bind_info) {
119    let bind = BindInfo::from(info);
120    contain_callback(&bind, || {
121        let data = T::bind(&bind)?;
122        bind.set_bind_data(data);
123        Ok(())
124    });
125}
126
127/// The C trampoline installed via `duckdb_table_function_set_init`.
128unsafe extern "C" fn table_init_trampoline<T: VTab>(info: duckdb_init_info) {
129    let init = InitInfo::<T>::from(info);
130    contain_callback(&init, || {
131        let data = T::init(&init)?;
132        init.set_init_data(data);
133        Ok(())
134    });
135}
136
137/// The C trampoline installed via `duckdb_table_function_set_local_init`.
138///
139/// DuckDB calls this once per worker thread that executes the query, using
140/// the same `duckdb_init_info` shape as the global init callback above — the
141/// data it sets is stored as *this thread's* local init data instead, purely
142/// by virtue of which callback DuckDB is currently invoking.
143unsafe extern "C" fn table_local_init_trampoline<T: VTabLocalInit>(info: duckdb_init_info) {
144    let init = InitInfo::<T>::from(info);
145    contain_callback(&init, || {
146        let data = T::local_init(&init)?;
147        init.set_local_init_data(data);
148        Ok(())
149    });
150}
151
152/// The C trampoline installed via `duckdb_table_function_set_function`.
153unsafe extern "C" fn table_func_trampoline<T: VTab>(
154    info: duckdb_function_info,
155    output: duckdb_data_chunk,
156) {
157    let func = TableFunctionInfo::<T>::from(info);
158    contain_callback(&func, || {
159        // SAFETY: DuckDB owns `output` and guarantees it stays live and
160        // unmutated by other code for the duration of this call.
161        let mut chunk = unsafe { DataChunkHandle::borrowed(output) };
162        T::func(&func, &mut chunk)
163    });
164}
165
166impl Connection {
167    /// Builds and configures a [`TableFunction`] for `T`: positional and named
168    /// parameters, the projection-pushdown flag, and the bind/init/func
169    /// trampolines. Shared by every `register_table_function*` entry point
170    /// below; callers add anything extra (local init, extra info) before
171    /// calling [`TableFunction::register`].
172    fn build_table_function<T: VTab>(name: &str) -> Result<TableFunction> {
173        let c_name = CString::new(name)?;
174        let f = TableFunction::new(&c_name);
175        for param in T::parameters()? {
176            f.add_parameter(&param);
177        }
178        for (param_name, ty) in T::named_parameters()? {
179            f.add_named_parameter(&param_name, &ty)?;
180        }
181        f.set_supports_projection_pushdown(T::supports_projection_pushdown());
182        f.set_bind(Some(table_bind_trampoline::<T>));
183        f.set_init(Some(table_init_trampoline::<T>));
184        f.set_function(Some(table_func_trampoline::<T>));
185        Ok(f)
186    }
187
188    /// Registers `T` as a table function named `name`.
189    ///
190    /// # Errors
191    ///
192    /// Returns an error if `name` contains a NUL byte, a parameter's logical
193    /// type cannot be built, or DuckDB rejects the registration (e.g. a name
194    /// conflict).
195    pub fn register_table_function<T: VTab>(
196        &mut self,
197        name: &str,
198    ) -> Result<()> {
199        let f = Self::build_table_function::<T>(name)?;
200        f.register(self.raw_con(), name)
201    }
202
203    /// Registers `T` as a table function named `name`, additionally wiring its
204    /// per-worker-thread [`VTabLocalInit::local_init`] callback.
205    ///
206    /// # Errors
207    ///
208    /// Same as [`register_table_function`](Self::register_table_function).
209    pub fn register_table_function_ext<T: VTabLocalInit>(
210        &mut self,
211        name: &str,
212    ) -> Result<()> {
213        let f = Self::build_table_function::<T>(name)?;
214        f.set_local_init(Some(table_local_init_trampoline::<T>));
215        f.register(self.raw_con(), name)
216    }
217
218    /// Registers `T` as a table function named `name`, sharing `extra_info`
219    /// read-only across every `bind`/`init`/`func` call for this function via
220    /// [`BindInfo::extra_info`]/[`InitInfo::extra_info`]/
221    /// [`TableFunctionInfo::extra_info`].
222    ///
223    /// # Errors
224    ///
225    /// Same as [`register_table_function`](Self::register_table_function).
226    pub fn register_table_function_with_extra_info<T: VTab, E: Send + Sync + 'static>(
227        &mut self,
228        name: &str,
229        extra_info: E,
230    ) -> Result<()> {
231        let f = Self::build_table_function::<T>(name)?;
232        f.set_extra_info(extra_info);
233        f.register(self.raw_con(), name)
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240    use crate::connection::Connection;
241    use crate::types::value::DuckValue;
242
243    /// `series(start, stop)`: the integers in `[start, stop)`, one per row,
244    /// exercising `bind` (result schema + cardinality), `init` (cursor state
245    /// shared across the scan), and chunked `func` output.
246    struct Series;
247
248    struct SeriesBind {
249        start: i64,
250        stop: i64,
251    }
252
253    struct SeriesInit {
254        next: std::sync::atomic::AtomicI64,
255    }
256
257    impl VTab for Series {
258        type BindData = SeriesBind;
259        type InitData = SeriesInit;
260
261        fn parameters() -> Result<Vec<LogicalType>> {
262            Ok(vec![LogicalType::of::<i64>()?, LogicalType::of::<i64>()?])
263        }
264
265        fn bind(bind: &BindInfo) -> super::super::UdfResult<Self::BindData> {
266            bind.add_result_column("n", &LogicalType::of::<i64>()?)?;
267            let start: i64 = bind.get_parameter(0)?;
268            let stop: i64 = bind.get_parameter(1)?;
269            bind.set_cardinality((stop - start).max(0) as u64, true);
270            Ok(SeriesBind { start, stop })
271        }
272
273        fn init(init: &InitInfo<Self>) -> super::super::UdfResult<Self::InitData> {
274            Ok(SeriesInit { next: std::sync::atomic::AtomicI64::new(init.bind_data().start) })
275        }
276
277        fn func(
278            func: &TableFunctionInfo<Self>,
279            output: &mut DataChunkHandle,
280        ) -> super::super::UdfResult<()> {
281            let stop = func.bind_data().stop;
282            let init = func.init_data();
283            let cap = output.capacity();
284            let mut written = 0usize;
285            {
286                let mut col = output.vector_mut(0)?;
287                while written < cap {
288                    let n = init.next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
289                    if n >= stop {
290                        init.next.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
291                        break;
292                    }
293                    col.set(written, n)?;
294                    written += 1;
295                }
296            }
297            output.set_len(written)?;
298            Ok(())
299        }
300    }
301
302    #[test]
303    fn register_and_scan_table_function() {
304        let mut conn = Connection::open_in_memory().unwrap();
305        conn.register_table_function::<Series>("series").unwrap();
306        let result = conn.execute("SELECT n FROM series(1, 6) ORDER BY n").unwrap();
307        let rows: Vec<_> = result.collect::<Result<_>>().unwrap();
308        let got: Vec<i64> = rows
309            .iter()
310            .map(|r| match r.get("n").unwrap() {
311                DuckValue::BigInt(n) => *n,
312                other => panic!("expected BigInt, got {other:?}"),
313            })
314            .collect();
315        assert_eq!(got, vec![1, 2, 3, 4, 5]);
316    }
317
318    #[test]
319    fn aggregate_over_table_function_matches_expected_sum() {
320        let mut conn = Connection::open_in_memory().unwrap();
321        conn.register_table_function::<Series>("series").unwrap();
322        let mut result = conn.execute("SELECT sum(n) AS total FROM series(1, 101)").unwrap();
323        let row = result.next().unwrap().unwrap();
324        assert_eq!(row.get("total"), Some(&DuckValue::HugeInt(5050)));
325    }
326
327    /// A large enough range to force multiple chunks through `func`, proving the
328    /// cursor state in `InitData` survives across repeated calls.
329    #[test]
330    fn scan_spanning_multiple_chunks() {
331        let mut conn = Connection::open_in_memory().unwrap();
332        conn.register_table_function::<Series>("series").unwrap();
333        let result = conn.execute("SELECT count(*) AS n FROM series(0, 10000)").unwrap();
334        let rows: Vec<_> = result.collect::<Result<_>>().unwrap();
335        assert_eq!(rows[0].get("n"), Some(&DuckValue::BigInt(10000)));
336    }
337
338    /// A `VTab` whose `bind` always errors, verifying it surfaces as a normal
339    /// query error and the connection stays usable afterward.
340    struct AlwaysFailsBind;
341    impl VTab for AlwaysFailsBind {
342        type BindData = ();
343        type InitData = ();
344
345        fn bind(_bind: &BindInfo) -> super::super::UdfResult<Self::BindData> {
346            Err("deliberate bind failure".into())
347        }
348
349        fn init(_init: &InitInfo<Self>) -> super::super::UdfResult<Self::InitData> {
350            Ok(())
351        }
352
353        fn func(
354            _func: &TableFunctionInfo<Self>,
355            output: &mut DataChunkHandle,
356        ) -> super::super::UdfResult<()> {
357            output.set_len(0)?;
358            Ok(())
359        }
360    }
361
362    #[test]
363    fn bind_error_surfaces_as_query_error_and_connection_stays_usable() {
364        let mut conn = Connection::open_in_memory().unwrap();
365        conn.register_table_function::<AlwaysFailsBind>("always_fails_bind").unwrap();
366        let err = match conn.execute("SELECT * FROM always_fails_bind()") {
367            Ok(_) => panic!("expected an error"),
368            Err(e) => e,
369        };
370        assert!(err.to_string().contains("deliberate bind failure"), "{err}");
371        conn.execute_batch("CREATE TABLE t (v INTEGER)").unwrap();
372    }
373
374    /// `doubler(stop)`: the even numbers `0, 2, 4, …` below `2 * stop`,
375    /// exercising `VTabLocalInit` — each worker thread's local init data (a
376    /// fixed multiplier here) is read back in `func` via
377    /// `TableFunctionInfo::local_init_data`.
378    struct Doubler;
379
380    struct DoublerInit {
381        next: std::sync::atomic::AtomicI64,
382        stop: i64,
383    }
384
385    impl VTab for Doubler {
386        type BindData = i64;
387        type InitData = DoublerInit;
388
389        fn parameters() -> Result<Vec<LogicalType>> {
390            Ok(vec![LogicalType::of::<i64>()?])
391        }
392
393        fn bind(bind: &BindInfo) -> super::super::UdfResult<Self::BindData> {
394            bind.add_result_column("n", &LogicalType::of::<i64>()?)?;
395            let stop: i64 = bind.get_parameter(0)?;
396            Ok(stop)
397        }
398
399        fn init(init: &InitInfo<Self>) -> super::super::UdfResult<Self::InitData> {
400            let stop = *init.bind_data();
401            Ok(DoublerInit { next: std::sync::atomic::AtomicI64::new(0), stop })
402        }
403
404        fn func(
405            func: &TableFunctionInfo<Self>,
406            output: &mut DataChunkHandle,
407        ) -> super::super::UdfResult<()> {
408            let init = func.init_data();
409            let multiplier = func.local_init_data::<i64>().copied().expect("local init ran");
410            let cap = output.capacity();
411            let mut written = 0usize;
412            {
413                let mut col = output.vector_mut(0)?;
414                while written < cap {
415                    let n = init.next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
416                    if n >= init.stop {
417                        init.next.fetch_sub(1, std::sync::atomic::Ordering::Relaxed);
418                        break;
419                    }
420                    col.set(written, n * multiplier)?;
421                    written += 1;
422                }
423            }
424            output.set_len(written)?;
425            Ok(())
426        }
427    }
428
429    impl VTabLocalInit for Doubler {
430        type LocalInitData = i64;
431
432        fn local_init(_init: &InitInfo<Self>) -> super::super::UdfResult<Self::LocalInitData> {
433            Ok(2)
434        }
435    }
436
437    #[test]
438    fn local_init_data_is_readable_from_func() {
439        let mut conn = Connection::open_in_memory().unwrap();
440        conn.register_table_function_ext::<Doubler>("doubler").unwrap();
441        let result = conn.execute("SELECT n FROM doubler(5) ORDER BY n").unwrap();
442        let rows: Vec<_> = result.collect::<Result<_>>().unwrap();
443        let got: Vec<i64> = rows
444            .iter()
445            .map(|r| match r.get("n").unwrap() {
446                DuckValue::BigInt(n) => *n,
447                other => panic!("expected BigInt, got {other:?}"),
448            })
449            .collect();
450        assert_eq!(got, vec![0, 2, 4, 6, 8]);
451    }
452
453    /// `with_context()`: a single row equal to the `i64` extra info shared
454    /// across `bind`/`func` via `register_table_function_with_extra_info`.
455    struct WithContext;
456
457    impl VTab for WithContext {
458        type BindData = ();
459        type InitData = std::sync::atomic::AtomicBool;
460
461        fn bind(bind: &BindInfo) -> super::super::UdfResult<Self::BindData> {
462            bind.add_result_column("ctx", &LogicalType::of::<i64>()?)?;
463            // Prove `extra_info` is reachable during `bind` too, not just `func`.
464            assert!(bind.extra_info::<i64>().is_some());
465            Ok(())
466        }
467
468        fn init(_init: &InitInfo<Self>) -> super::super::UdfResult<Self::InitData> {
469            Ok(std::sync::atomic::AtomicBool::new(false))
470        }
471
472        fn func(
473            func: &TableFunctionInfo<Self>,
474            output: &mut DataChunkHandle,
475        ) -> super::super::UdfResult<()> {
476            let already_emitted = func.init_data().swap(true, std::sync::atomic::Ordering::Relaxed);
477            if already_emitted {
478                output.set_len(0)?;
479                return Ok(());
480            }
481            let ctx = *func.extra_info::<i64>().expect("extra info was registered");
482            output.vector_mut(0)?.set(0, ctx)?;
483            output.set_len(1)?;
484            Ok(())
485        }
486    }
487
488    #[test]
489    fn extra_info_is_shared_and_readable_from_bind_and_func() {
490        let mut conn = Connection::open_in_memory().unwrap();
491        conn.register_table_function_with_extra_info::<WithContext, i64>("with_context", 777)
492            .unwrap();
493        let mut result = conn.execute("SELECT ctx FROM with_context()").unwrap();
494        let row = result.next().unwrap().unwrap();
495        assert_eq!(row.get("ctx"), Some(&DuckValue::BigInt(777)));
496    }
497}
better_duck_core/udf/table/mod.rs

better_duck_core/udf/table/
mod.rs
