better_duck_core/udf/scalar/
mod.rs

1//! DuckDB scalar functions: row-wise functions used in a `SELECT` list or
2//! `WHERE` clause, e.g. `SELECT my_func(x) FROM t`.
3
4mod function;
5
6use std::ffi::CString;
7
8use crate::{
9    connection::Connection,
10    error::Result,
11    ffi::{duckdb_bind_info, duckdb_data_chunk, duckdb_function_info, duckdb_vector},
12};
13
14use self::function::{ScalarFunction, ScalarFunctionInfo, ScalarFunctionSet};
15use super::{callback::contain_callback, data_chunk::DataChunkHandle, vector::VectorMut};
16use crate::types::LogicalType;
17
18pub use self::function::ScalarBindInfo;
19
20/// A DuckDB scalar function: computes one value per row.
21///
22/// See the callback containment contract in [`crate::udf`].
23pub trait VScalar: Sized {
24    /// State set at registration time, shared across every invocation and every
25    /// worker thread. Persists for the lifetime of the catalog entry, so it must
26    /// be `'static`; any interior mutation must be synchronized.
27    type State: Send + Sync + 'static;
28
29    /// Per-query data produced by [`VScalar::bind`] and shared, read-only, by every
30    /// [`VScalar::invoke`] call for that query. Use `()` if the function needs none.
31    type BindData: Send + Sync + 'static;
32
33    /// The possible signatures of this function. Each becomes a DuckDB overload;
34    /// [`VScalar::invoke`] must be able to handle every one of them.
35    ///
36    /// # Errors
37    ///
38    /// Returns an error if a signature's logical type cannot be built.
39    fn signatures() -> Result<Vec<ScalarSignature>>;
40
41    /// Runs once per query that references this function, before any [`invoke`](VScalar::invoke).
42    /// Inspects the call's argument expressions (via [`ScalarBindInfo`] — e.g. to
43    /// fold a constant argument or reject an unsupported one) and produces the
44    /// per-query [`BindData`](VScalar::BindData). Functions needing no bind data
45    /// return `Ok(())` (with `type BindData = ()`).
46    ///
47    /// # Errors
48    ///
49    /// Returns an error to reject the query at bind time with that message.
50    fn bind(bind: &ScalarBindInfo) -> super::UdfResult<Self::BindData>;
51
52    /// Computes `output[row]` for every `row` in `0..input.len()`.
53    ///
54    /// DuckDB guarantees `input` and `output` stay live for the duration of this
55    /// call, and that `output`'s capacity is at least `input.len()`. `bind_data` is
56    /// the value [`bind`](VScalar::bind) produced for this query.
57    ///
58    /// # Errors
59    ///
60    /// Returns an error to fail the query with that message.
61    fn invoke(
62        state: &Self::State,
63        bind_data: &Self::BindData,
64        input: &DataChunkHandle,
65        output: &mut VectorMut<'_>,
66    ) -> super::UdfResult<()>;
67
68    /// Whether this function is volatile — re-evaluated for every row even with
69    /// no parameters, rather than optimized to a constant. Needed for functions
70    /// like random-number or UUID generators.
71    fn volatile() -> bool {
72        false
73    }
74
75    /// Whether this function should be invoked for rows containing `NULL`
76    /// parameters. By default DuckDB substitutes `NULL` as the result for any
77    /// row with a `NULL` argument without calling [`VScalar::invoke`] at all;
78    /// returning `true` here disables that shortcut.
79    fn special_handling() -> bool {
80        false
81    }
82}
83
84/// The parameter shape of one [`ScalarSignature`].
85enum ScalarParams {
86    /// A fixed list of parameter types.
87    Exact(Vec<LogicalType>),
88    /// Any number of arguments of a single type.
89    Variadic(LogicalType),
90}
91
92/// One overload of a scalar function: a parameter shape and a return type.
93pub struct ScalarSignature {
94    parameters: Option<ScalarParams>,
95    return_type: LogicalType,
96}
97
98impl ScalarSignature {
99    /// A signature with a fixed list of parameter types.
100    pub fn exact(
101        parameters: Vec<LogicalType>,
102        return_type: LogicalType,
103    ) -> Self {
104        Self { parameters: Some(ScalarParams::Exact(parameters)), return_type }
105    }
106
107    /// A signature accepting any number of arguments of `parameter`'s type.
108    pub fn variadic(
109        parameter: LogicalType,
110        return_type: LogicalType,
111    ) -> Self {
112        Self { parameters: Some(ScalarParams::Variadic(parameter)), return_type }
113    }
114
115    fn apply(
116        &self,
117        f: &ScalarFunction,
118    ) {
119        f.set_return_type(&self.return_type);
120        match &self.parameters {
121            Some(ScalarParams::Exact(params)) => {
122                for p in params {
123                    f.add_parameter(p);
124                }
125            },
126            Some(ScalarParams::Variadic(p)) => f.set_varargs(p),
127            None => {},
128        }
129    }
130}
131
132/// The C trampoline installed via `duckdb_scalar_function_set_function`.
133///
134/// Nothing above the `contain_callback` call may panic or carry a `Drop` impl:
135/// since Rust 1.81 a panic escaping a plain `extern "C"` frame aborts the
136/// process regardless of panic strategy, so containment must be the outermost
137/// thing here. See [`crate::udf`] for the full `panic = "abort"` caveat.
138unsafe extern "C" fn scalar_trampoline<S: VScalar>(
139    info: duckdb_function_info,
140    input: duckdb_data_chunk,
141    output: duckdb_vector,
142) {
143    let sink = ScalarFunctionInfo::from(info);
144    contain_callback(&sink, || {
145        // SAFETY: DuckDB owns `input` and guarantees it stays live and unmutated
146        // by other code for the duration of this call.
147        let chunk = unsafe { DataChunkHandle::borrowed(input) };
148        // SAFETY: DuckDB owns `output`, guarantees it stays live for the
149        // duration of this call, and that no other code writes through it
150        // concurrently.
151        let mut out = unsafe { VectorMut::new(output) };
152        // SAFETY: `register_scalar_function`/`register_scalar_function_with_state`
153        // always call `ScalarFunction::set_extra_info::<S::State>`, so the state
154        // stored for this catalog entry always has type `S::State`.
155        let state = unsafe { sink.state::<S::State>() };
156        // SAFETY: `scalar_bind_trampoline::<S>` runs before any invoke and stores a
157        // `Box<S::BindData>` via `set_bind_data`, so the bind data has type
158        // `S::BindData` and is live for every invocation of this query.
159        let bind_data = unsafe { sink.bind_data::<S::BindData>() };
160        S::invoke(state, bind_data, &chunk, &mut out)
161    });
162}
163
164/// The C trampoline installed via `duckdb_scalar_function_set_bind`. Runs `S::bind`
165/// and stores its result as the query's bind data. See [`scalar_trampoline`] for why
166/// containment is the outermost thing here.
167unsafe extern "C" fn scalar_bind_trampoline<S: VScalar>(info: duckdb_bind_info) {
168    let bind = ScalarBindInfo::from(info);
169    contain_callback(&bind, || {
170        let data = S::bind(&bind)?;
171        bind.set_bind_data(data);
172        Ok(())
173    });
174}
175
176impl Connection {
177    /// Registers `S` as a scalar function named `name`, using `S::State`'s
178    /// default value as the shared state for every overload.
179    ///
180    /// # Errors
181    ///
182    /// Returns an error if `name` contains a NUL byte, a signature's logical
183    /// type cannot be built, or DuckDB rejects the registration (e.g. a name
184    /// conflict).
185    pub fn register_scalar_function<S: VScalar>(
186        &mut self,
187        name: &str,
188    ) -> Result<()>
189    where
190        S::State: Default,
191    {
192        // A fresh default per overload, rather than one shared value cloned —
193        // avoids requiring `S::State: Clone` in addition to `Default`.
194        register_scalar_function_impl::<S>(self, name, S::State::default)
195    }
196
197    /// Registers `S` as a scalar function named `name`, with explicit shared
198    /// state. The state is cloned once per overload.
199    ///
200    /// # Errors
201    ///
202    /// Returns an error if `name` contains a NUL byte, a signature's logical
203    /// type cannot be built, or DuckDB rejects the registration (e.g. a name
204    /// conflict).
205    pub fn register_scalar_function_with_state<S: VScalar>(
206        &mut self,
207        name: &str,
208        state: S::State,
209    ) -> Result<()>
210    where
211        S::State: Clone,
212    {
213        register_scalar_function_impl::<S>(self, name, move || state.clone())
214    }
215}
216
217fn register_scalar_function_impl<S: VScalar>(
218    conn: &mut Connection,
219    name: &str,
220    mut make_state: impl FnMut() -> S::State,
221) -> Result<()> {
222    let c_name = CString::new(name)?;
223    let set = ScalarFunctionSet::new(&c_name);
224    for signature in S::signatures()? {
225        let f = ScalarFunction::new(&c_name);
226        signature.apply(&f);
227        f.set_function(Some(scalar_trampoline::<S>));
228        f.set_bind(Some(scalar_bind_trampoline::<S>));
229        if S::volatile() {
230            f.set_volatile();
231        }
232        if S::special_handling() {
233            f.set_special_handling();
234        }
235        f.set_extra_info(make_state());
236        set.add_function(&f)?;
237    }
238    set.register(conn.raw_con(), name)
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244    use crate::connection::Connection;
245
246    /// A hand-written `VScalar`, exercising the trait directly (not through the
247    /// `#[duckdb_scalar]` macro, which lands in a later phase).
248    struct AddOne;
249
250    impl VScalar for AddOne {
251        type State = ();
252        type BindData = ();
253
254        fn signatures() -> Result<Vec<ScalarSignature>> {
255            Ok(vec![ScalarSignature::exact(
256                vec![LogicalType::of::<i32>()?],
257                LogicalType::of::<i32>()?,
258            )])
259        }
260
261        fn bind(_bind: &ScalarBindInfo) -> super::super::UdfResult<()> {
262            Ok(())
263        }
264
265        fn invoke(
266            _state: &(),
267            _bind_data: &(),
268            input: &DataChunkHandle,
269            output: &mut VectorMut<'_>,
270        ) -> super::super::UdfResult<()> {
271            let col = input.vector(0)?;
272            for row in 0..input.len() {
273                let v: i32 = col.get(row)?;
274                output.set(row, v + 1)?;
275            }
276            Ok(())
277        }
278    }
279
280    #[test]
281    fn register_and_call_scalar_function() {
282        let mut conn = Connection::open_in_memory().unwrap();
283        conn.register_scalar_function::<AddOne>("add_one").unwrap();
284        conn.execute_batch("CREATE TABLE t (v INTEGER)").unwrap();
285        conn.execute_batch("INSERT INTO t VALUES (1), (2), (41)").unwrap();
286        let result = conn.execute("SELECT add_one(v) AS r FROM t ORDER BY v").unwrap();
287        let rows: Vec<_> = result.collect::<Result<_>>().unwrap();
288        assert_eq!(rows.len(), 3);
289        assert_eq!(rows[2].get("r"), Some(&crate::types::value::DuckValue::Int(42)));
290    }
291
292    /// A `VScalar` whose `invoke` always errors, verifying the error surfaces as
293    /// a normal query error and the connection stays usable afterward.
294    struct AlwaysFails;
295
296    impl VScalar for AlwaysFails {
297        type State = ();
298        type BindData = ();
299
300        fn signatures() -> Result<Vec<ScalarSignature>> {
301            Ok(vec![ScalarSignature::exact(
302                vec![LogicalType::of::<i32>()?],
303                LogicalType::of::<i32>()?,
304            )])
305        }
306
307        fn bind(_bind: &ScalarBindInfo) -> super::super::UdfResult<()> {
308            Ok(())
309        }
310
311        fn invoke(
312            _state: &(),
313            _bind_data: &(),
314            _input: &DataChunkHandle,
315            _output: &mut VectorMut<'_>,
316        ) -> super::super::UdfResult<()> {
317            Err("deliberate failure".into())
318        }
319    }
320
321    #[test]
322    fn invoke_error_surfaces_as_query_error_and_connection_stays_usable() {
323        let mut conn = Connection::open_in_memory().unwrap();
324        conn.register_scalar_function::<AlwaysFails>("always_fails").unwrap();
325        conn.execute_batch("CREATE TABLE t (v INTEGER)").unwrap();
326        conn.execute_batch("INSERT INTO t VALUES (1)").unwrap();
327        let err = match conn.execute("SELECT always_fails(v) FROM t") {
328            Ok(_) => panic!("expected an error"),
329            Err(e) => e,
330        };
331        assert!(err.to_string().contains("deliberate failure"), "{err}");
332        // The connection must still be usable after a UDF error.
333        conn.execute_batch("INSERT INTO t VALUES (2)").unwrap();
334    }
335
336    /// A panicking `VScalar`, verifying panic containment.
337    struct AlwaysPanics;
338
339    impl VScalar for AlwaysPanics {
340        type State = ();
341        type BindData = ();
342
343        fn signatures() -> Result<Vec<ScalarSignature>> {
344            Ok(vec![ScalarSignature::exact(
345                vec![LogicalType::of::<i32>()?],
346                LogicalType::of::<i32>()?,
347            )])
348        }
349
350        fn bind(_bind: &ScalarBindInfo) -> super::super::UdfResult<()> {
351            Ok(())
352        }
353
354        fn invoke(
355            _state: &(),
356            _bind_data: &(),
357            _input: &DataChunkHandle,
358            _output: &mut VectorMut<'_>,
359        ) -> super::super::UdfResult<()> {
360            panic!("deliberate panic")
361        }
362    }
363
364    #[test]
365    #[cfg(panic = "unwind")]
366    fn invoke_panic_is_contained_and_connection_stays_usable() {
367        let mut conn = Connection::open_in_memory().unwrap();
368        conn.register_scalar_function::<AlwaysPanics>("always_panics").unwrap();
369        conn.execute_batch("CREATE TABLE t (v INTEGER)").unwrap();
370        conn.execute_batch("INSERT INTO t VALUES (1)").unwrap();
371        let err = match conn.execute("SELECT always_panics(v) FROM t") {
372            Ok(_) => panic!("expected an error"),
373            Err(e) => e,
374        };
375        assert!(err.to_string().contains("deliberate panic"), "{err}");
376        conn.execute_batch("INSERT INTO t VALUES (2)").unwrap();
377    }
378
379    /// `add_const(x, c)`: the second argument must be a *constant*. `bind` inspects
380    /// the argument expressions, folds the constant to an `i64`, and stores it as
381    /// bind data; `invoke` adds it to every row. This exercises the whole bind/fold
382    /// path: set_bind, argument_count/argument, Expression::is_foldable/fold (via the
383    /// bind client context), set_bind_data, and get_bind_data.
384    struct AddConst;
385
386    impl VScalar for AddConst {
387        type State = ();
388        type BindData = i64;
389
390        fn signatures() -> Result<Vec<ScalarSignature>> {
391            Ok(vec![ScalarSignature::exact(
392                vec![LogicalType::of::<i32>()?, LogicalType::of::<i32>()?],
393                LogicalType::of::<i32>()?,
394            )])
395        }
396
397        fn bind(bind: &ScalarBindInfo) -> super::super::UdfResult<i64> {
398            // Registration state is readable at bind time too (here it is `()`).
399            // SAFETY: this function is registered with `State = ()`, so the stored
400            // extra-info has type `()`.
401            let _state: &() = unsafe { bind.extra_info::<()>() };
402            assert_eq!(bind.argument_count(), 2, "add_const has two arguments");
403            let arg = bind.argument(1).ok_or("add_const needs a second argument")?;
404            if !arg.is_foldable() {
405                return Err("add_const's second argument must be a constant".into());
406            }
407            let ctx = bind.client_context().ok_or("no client context for folding")?;
408            let folded = arg.fold(&ctx).map_err(|e| format!("fold failed: {e:?}"))?;
409            match folded {
410                crate::types::value::DuckValue::Int(n) => Ok(i64::from(n)),
411                crate::types::value::DuckValue::BigInt(n) => Ok(n),
412                other => Err(format!("expected an integer constant, got {other:?}").into()),
413            }
414        }
415
416        fn invoke(
417            _state: &(),
418            bind_data: &i64,
419            input: &DataChunkHandle,
420            output: &mut VectorMut<'_>,
421        ) -> super::super::UdfResult<()> {
422            let col = input.vector(0)?;
423            let c = i32::try_from(*bind_data).map_err(|_| "constant out of range")?;
424            for row in 0..input.len() {
425                let v: i32 = col.get(row)?;
426                output.set(row, v + c)?;
427            }
428            Ok(())
429        }
430    }
431
432    #[test]
433    fn bind_folds_a_constant_argument_and_invoke_uses_it() {
434        let mut conn = Connection::open_in_memory().unwrap();
435        conn.register_scalar_function::<AddConst>("add_const").unwrap();
436        conn.execute_batch("CREATE TABLE t (v INTEGER)").unwrap();
437        conn.execute_batch("INSERT INTO t VALUES (1), (2), (40)").unwrap();
438        // The second argument (10) is folded at bind time and added to every row.
439        let rows: Vec<_> = conn
440            .execute("SELECT add_const(v, 2 + 8) AS r FROM t ORDER BY v")
441            .unwrap()
442            .collect::<Result<_>>()
443            .unwrap();
444        assert_eq!(rows.len(), 3);
445        assert_eq!(rows[0].get("r"), Some(&crate::types::value::DuckValue::Int(11)));
446        assert_eq!(rows[2].get("r"), Some(&crate::types::value::DuckValue::Int(50)));
447    }
448
449    #[test]
450    fn bind_rejects_a_non_constant_argument() {
451        let mut conn = Connection::open_in_memory().unwrap();
452        conn.register_scalar_function::<AddConst>("add_const").unwrap();
453        conn.execute_batch("CREATE TABLE t (v INTEGER)").unwrap();
454        conn.execute_batch("INSERT INTO t VALUES (1)").unwrap();
455        // The second argument references a column, so it is not foldable — bind
456        // rejects the query with our error message, and the connection stays usable.
457        let err = match conn.execute("SELECT add_const(v, v) FROM t") {
458            Ok(_) => panic!("expected a bind error"),
459            Err(e) => e,
460        };
461        assert!(err.to_string().contains("must be a constant"), "{err}");
462        conn.execute_batch("INSERT INTO t VALUES (2)").unwrap();
463    }
464}
better_duck_core/udf/scalar/mod.rs

better_duck_core/udf/scalar/
mod.rs
