better_duck_core/udf/aggregate/
mod.rs

1//! DuckDB aggregate functions: functions that fold many rows into one value per
2//! group, e.g. `SELECT my_sum(x) FROM t GROUP BY k`.
3//!
4//! An aggregate is defined by a per-group [`State`](VAggregate::State) and four
5//! callbacks — init, update, combine, finalize — plus a destructor. DuckDB allocates
6//! the raw state bytes uninitialised, so the state is wrapped in a `RawState` that
7//! pairs the value with an `initialised` flag (a `MaybeUninit` guard), and every
8//! callback that runs user code is wrapped in `contain_callback` so a panic can
9//! never cross the C boundary.
10
11mod function;
12
13use std::ffi::CString;
14use std::mem::MaybeUninit;
15use std::panic::{catch_unwind, AssertUnwindSafe};
16
17use crate::{
18    connection::Connection,
19    error::Result,
20    ffi::{duckdb_aggregate_state, duckdb_data_chunk, duckdb_function_info, duckdb_vector, idx_t},
21};
22
23use self::function::{AggregateFunction, AggregateFunctionInfo, AggregateFunctionSet};
24use super::{callback::contain_callback, data_chunk::DataChunkHandle, vector::VectorMut};
25use crate::types::LogicalType;
26
27/// A DuckDB aggregate function: folds a set of input rows into one result value.
28///
29/// See the callback-containment contract in [`crate::udf`].
30pub trait VAggregate: Sized {
31    /// The per-group accumulator. Allocated by DuckDB (uninitialised), set up by
32    /// [`init`](VAggregate::init), and dropped by the generated destructor.
33    type State: Send + Sync + 'static;
34
35    /// Registration-time shared state, available read-only to every callback via
36    /// DuckDB's extra-info slot. Use `()` if the aggregate needs none.
37    type Shared: Send + Sync + 'static;
38
39    /// The input parameter types (one overload).
40    ///
41    /// # Errors
42    ///
43    /// Returns an error if a parameter's logical type cannot be built.
44    fn parameters() -> Result<Vec<LogicalType>>;
45
46    /// The result type produced by [`finalize`](VAggregate::finalize).
47    ///
48    /// # Errors
49    ///
50    /// Returns an error if the return type cannot be built.
51    fn return_type() -> Result<LogicalType>;
52
53    /// The initial accumulator value for a fresh group.
54    fn init() -> Self::State;
55
56    /// Folds `input`'s row `row` into `state`. `shared` is the registration state.
57    ///
58    /// # Errors
59    ///
60    /// Returns an error to fail the query with that message.
61    fn update(
62        shared: &Self::Shared,
63        state: &mut Self::State,
64        input: &DataChunkHandle,
65        row: usize,
66    ) -> super::UdfResult<()>;
67
68    /// Merges `source` into `target` (parallel/partial aggregation).
69    fn combine(
70        shared: &Self::Shared,
71        source: &Self::State,
72        target: &mut Self::State,
73    );
74
75    /// Writes the finalised value of `state` into `output` at `row`.
76    ///
77    /// # Errors
78    ///
79    /// Returns an error to fail the query with that message.
80    fn finalize(
81        shared: &Self::Shared,
82        state: &Self::State,
83        output: &mut VectorMut<'_>,
84        row: usize,
85    ) -> super::UdfResult<()>;
86
87    /// Whether the aggregate should still be invoked for `NULL` inputs. `false` by
88    /// default (DuckDB skips rows whose argument is `NULL`).
89    fn special_handling() -> bool {
90        false
91    }
92}
93
94/// DuckDB-allocated aggregate state: the accumulator plus an `initialised` flag, so
95/// the (uninitialised) bytes DuckDB hands out are only read after `init` runs and
96/// dropped at most once.
97struct RawState<S> {
98    initialised: bool,
99    data: MaybeUninit<S>,
100}
101
102impl<S> RawState<S> {
103    fn ready(value: S) -> Self {
104        Self { initialised: true, data: MaybeUninit::new(value) }
105    }
106
107    fn poisoned() -> Self {
108        Self { initialised: false, data: MaybeUninit::uninit() }
109    }
110
111    /// # Safety
112    /// The state must have been initialised by `init`.
113    unsafe fn get(&self) -> &S {
114        debug_assert!(self.initialised, "aggregate state read before init");
115        // SAFETY: `initialised` implies `data` holds a valid `S` (caller contract).
116        unsafe { &*self.data.as_ptr() }
117    }
118
119    /// # Safety
120    /// The state must have been initialised by `init`.
121    unsafe fn get_mut(&mut self) -> &mut S {
122        debug_assert!(self.initialised, "aggregate state written before init");
123        // SAFETY: `initialised` implies `data` holds a valid `S` (caller contract).
124        unsafe { &mut *self.data.as_mut_ptr() }
125    }
126
127    /// Drops the contained value if present, leaving the slot poisoned. Idempotent.
128    unsafe fn drop_value(&mut self) {
129        if self.initialised {
130            self.initialised = false;
131            // SAFETY: `initialised` was true, so `data` held a valid `S`.
132            unsafe { self.data.assume_init_drop() };
133        }
134    }
135}
136
137/// `state_size` callback: the byte size of one wrapped state.
138unsafe extern "C" fn agg_state_size<A: VAggregate>(_info: duckdb_function_info) -> idx_t {
139    std::mem::size_of::<RawState<A::State>>() as idx_t
140}
141
142/// `init` callback: write the initial wrapped state into DuckDB's raw bytes.
143unsafe extern "C" fn agg_init<A: VAggregate>(
144    _info: duckdb_function_info,
145    state: duckdb_aggregate_state,
146) {
147    let slot = state as *mut RawState<A::State>;
148    // A user `init` panic must not cross the C boundary; on panic leave the slot
149    // poisoned so the destructor is still safe.
150    let value = catch_unwind(AssertUnwindSafe(A::init));
151    // SAFETY: `slot` points at `size_of::<RawState<A::State>>()` uninitialised bytes
152    // DuckDB just allocated; writing a fresh `RawState` initialises them exactly once.
153    unsafe {
154        match value {
155            Ok(v) => slot.write(RawState::ready(v)),
156            Err(_) => slot.write(RawState::poisoned()),
157        }
158    }
159}
160
161/// `update` callback: fold each input row into its group's state.
162unsafe extern "C" fn agg_update<A: VAggregate>(
163    info: duckdb_function_info,
164    input: duckdb_data_chunk,
165    states: *mut duckdb_aggregate_state,
166) {
167    let sink = AggregateFunctionInfo::from(info);
168    contain_callback(&sink, || {
169        // SAFETY: the extra-info stored at registration has type `A::Shared`.
170        let shared = unsafe { sink.extra_info::<A::Shared>() };
171        // SAFETY: DuckDB owns `input` for the call and guarantees it stays live.
172        let chunk = unsafe { DataChunkHandle::borrowed(input) };
173        for row in 0..chunk.len() {
174            // SAFETY: `states` has one entry per input row; `states[row]` points at a
175            // live, init'd `RawState<A::State>` for that row's group.
176            let slot = unsafe { &mut *((*states.add(row)) as *mut RawState<A::State>) };
177            // SAFETY: the slot was initialised by `agg_init` before any update.
178            let state = unsafe { slot.get_mut() };
179            A::update(shared, state, &chunk, row)?;
180        }
181        Ok(())
182    });
183}
184
185/// `combine` callback: merge each source state into the matching target state.
186unsafe extern "C" fn agg_combine<A: VAggregate>(
187    info: duckdb_function_info,
188    source: *mut duckdb_aggregate_state,
189    target: *mut duckdb_aggregate_state,
190    count: idx_t,
191) {
192    let sink = AggregateFunctionInfo::from(info);
193    contain_callback(&sink, || {
194        // SAFETY: the extra-info stored at registration has type `A::Shared`.
195        let shared = unsafe { sink.extra_info::<A::Shared>() };
196        for i in 0..count as usize {
197            // SAFETY: `source[i]` points at a live, init'd `RawState<A::State>`.
198            let src = unsafe { &*((*source.add(i)) as *const RawState<A::State>) };
199            // SAFETY: `target[i]` points at a live, init'd `RawState<A::State>`.
200            let tgt = unsafe { &mut *((*target.add(i)) as *mut RawState<A::State>) };
201            // SAFETY: both slots were initialised by `agg_init` before any combine.
202            let (s, t) = unsafe { (src.get(), tgt.get_mut()) };
203            A::combine(shared, s, t);
204        }
205        Ok(())
206    });
207}
208
209/// `finalize` callback: write each state's finalised value into the result vector.
210unsafe extern "C" fn agg_finalize<A: VAggregate>(
211    info: duckdb_function_info,
212    source: *mut duckdb_aggregate_state,
213    result: duckdb_vector,
214    count: idx_t,
215    offset: idx_t,
216) {
217    let sink = AggregateFunctionInfo::from(info);
218    contain_callback(&sink, || {
219        // SAFETY: the extra-info stored at registration has type `A::Shared`.
220        let shared = unsafe { sink.extra_info::<A::Shared>() };
221        // SAFETY: DuckDB owns `result` for the call and guarantees it stays live.
222        let mut out = unsafe { VectorMut::new(result) };
223        for i in 0..count as usize {
224            // SAFETY: `source[i]` points at a live, init'd `RawState<A::State>`.
225            let slot = unsafe { &*((*source.add(i)) as *const RawState<A::State>) };
226            // SAFETY: the slot was initialised by `agg_init` before finalize.
227            let state = unsafe { slot.get() };
228            A::finalize(shared, state, &mut out, offset as usize + i)?;
229        }
230        Ok(())
231    });
232}
233
234/// `destroy` callback: drop each state's value exactly once.
235unsafe extern "C" fn agg_destroy<A: VAggregate>(
236    states: *mut duckdb_aggregate_state,
237    count: idx_t,
238) {
239    // A `Drop` panic must not cross the C boundary.
240    let _ = catch_unwind(AssertUnwindSafe(|| {
241        for i in 0..count as usize {
242            // SAFETY: `states[i]` points at a live `RawState<A::State>` DuckDB is
243            // destroying; `drop_value` drops the contained value at most once.
244            let slot = unsafe { &mut *((*states.add(i)) as *mut RawState<A::State>) };
245            // SAFETY: as above.
246            unsafe { slot.drop_value() };
247        }
248    }));
249}
250
251impl Connection {
252    /// Registers `A` as an aggregate function named `name`, with a default shared
253    /// state.
254    ///
255    /// # Errors
256    ///
257    /// Returns an error if `name` contains a NUL byte, a logical type cannot be
258    /// built, or DuckDB rejects the registration (e.g. a name conflict).
259    pub fn register_aggregate_function<A: VAggregate>(
260        &mut self,
261        name: &str,
262    ) -> Result<()>
263    where
264        A::Shared: Default,
265    {
266        self.register_aggregate_function_with_state::<A>(name, A::Shared::default())
267    }
268
269    /// Registers `A` as an aggregate function named `name`, with an explicit shared
270    /// registration state (available read-only to every callback).
271    ///
272    /// # Errors
273    ///
274    /// As [`register_aggregate_function`](Connection::register_aggregate_function).
275    pub fn register_aggregate_function_with_state<A: VAggregate>(
276        &mut self,
277        name: &str,
278        shared: A::Shared,
279    ) -> Result<()> {
280        let c_name = CString::new(name)?;
281        let f = build_aggregate::<A>(&c_name, shared)?;
282        f.register(self.raw_con(), name)
283    }
284
285    /// Registers several `VAggregate` overloads under one SQL name as an aggregate
286    /// function set (DuckDB dispatches on argument types). Each `(name-independent)`
287    /// overload `A_i` is built with its own default shared state.
288    ///
289    /// This takes a closure that adds overloads to a builder, so overloads of
290    /// *different* Rust types can share one SQL name.
291    ///
292    /// # Errors
293    ///
294    /// As [`register_aggregate_function`](Connection::register_aggregate_function),
295    /// plus a conflict between two overloads in the set.
296    pub fn register_aggregate_function_set(
297        &mut self,
298        name: &str,
299        build: impl FnOnce(&mut AggregateSetBuilder) -> Result<()>,
300    ) -> Result<()> {
301        let c_name = CString::new(name)?;
302        let set = AggregateFunctionSet::new(&c_name);
303        let mut builder = AggregateSetBuilder { set: &set, name };
304        build(&mut builder)?;
305        set.register(self.raw_con(), name)
306    }
307}
308
309/// Builds a configured [`AggregateFunction`] for `A` with shared state `shared`.
310fn build_aggregate<A: VAggregate>(
311    c_name: &std::ffi::CStr,
312    shared: A::Shared,
313) -> Result<AggregateFunction> {
314    let f = AggregateFunction::new(c_name);
315    for p in A::parameters()? {
316        f.add_parameter(&p);
317    }
318    f.set_return_type(&A::return_type()?);
319    f.set_extra_info(shared);
320    f.set_functions(
321        Some(agg_state_size::<A>),
322        Some(agg_init::<A>),
323        Some(agg_update::<A>),
324        Some(agg_combine::<A>),
325        Some(agg_finalize::<A>),
326    );
327    f.set_destructor(Some(agg_destroy::<A>));
328    if A::special_handling() {
329        f.set_special_handling();
330    }
331    Ok(f)
332}
333
334/// Collects aggregate overloads into a function set (see
335/// [`Connection::register_aggregate_function_set`]).
336pub struct AggregateSetBuilder<'a> {
337    set: &'a AggregateFunctionSet,
338    name: &'a str,
339}
340
341impl AggregateSetBuilder<'_> {
342    /// Adds overload `A` (with a default shared state) to the set.
343    ///
344    /// # Errors
345    ///
346    /// Returns an error if a logical type cannot be built or the overload conflicts
347    /// with one already in the set.
348    pub fn add<A: VAggregate>(&mut self) -> Result<&mut Self>
349    where
350        A::Shared: Default,
351    {
352        self.add_with_state::<A>(A::Shared::default())
353    }
354
355    /// Adds overload `A` with an explicit shared state to the set.
356    ///
357    /// # Errors
358    ///
359    /// As [`add`](AggregateSetBuilder::add).
360    pub fn add_with_state<A: VAggregate>(
361        &mut self,
362        shared: A::Shared,
363    ) -> Result<&mut Self> {
364        let c_name = CString::new(self.name)?;
365        let f = build_aggregate::<A>(&c_name, shared)?;
366        self.set.add_function(&f)?;
367        Ok(self)
368    }
369}
370
371#[cfg(test)]
372mod tests {
373    use super::*;
374    use crate::types::value::DuckValue;
375
376    /// `int_sum(x)`: a BIGINT running sum over INTEGER inputs.
377    struct IntSum;
378
379    impl VAggregate for IntSum {
380        type State = i64;
381        type Shared = ();
382
383        fn parameters() -> Result<Vec<LogicalType>> {
384            Ok(vec![LogicalType::of::<i32>()?])
385        }
386
387        fn return_type() -> Result<LogicalType> {
388            LogicalType::of::<i64>()
389        }
390
391        fn init() -> i64 {
392            0
393        }
394
395        fn update(
396            _shared: &(),
397            state: &mut i64,
398            input: &DataChunkHandle,
399            row: usize,
400        ) -> super::super::UdfResult<()> {
401            let v: i32 = input.vector(0)?.get(row)?;
402            *state += i64::from(v);
403            Ok(())
404        }
405
406        fn combine(
407            _shared: &(),
408            source: &i64,
409            target: &mut i64,
410        ) {
411            *target += *source;
412        }
413
414        fn finalize(
415            _shared: &(),
416            state: &i64,
417            output: &mut VectorMut<'_>,
418            row: usize,
419        ) -> super::super::UdfResult<()> {
420            output.set(row, *state)?;
421            Ok(())
422        }
423    }
424
425    #[test]
426    fn aggregate_sums_and_groups() {
427        let mut conn = Connection::open_in_memory().unwrap();
428        conn.register_aggregate_function::<IntSum>("int_sum").unwrap();
429        conn.execute_batch("CREATE TABLE t (k INTEGER, v INTEGER)").unwrap();
430        conn.execute_batch("INSERT INTO t VALUES (1,10),(1,20),(2,5),(2,7),(2,100)").unwrap();
431
432        // Grand total.
433        let total = conn.execute("SELECT int_sum(v) AS s FROM t").unwrap();
434        let rows: Vec<_> = total.collect::<Result<_>>().unwrap();
435        assert_eq!(rows[0].get("s"), Some(&DuckValue::BigInt(142)));
436
437        // Grouped — exercises combine across partial states too.
438        let mut grouped =
439            conn.execute("SELECT k, int_sum(v) AS s FROM t GROUP BY k ORDER BY k").unwrap();
440        let g0 = grouped.next().unwrap().unwrap();
441        assert_eq!(g0.get("s"), Some(&DuckValue::BigInt(30)));
442        let g1 = grouped.next().unwrap().unwrap();
443        assert_eq!(g1.get("s"), Some(&DuckValue::BigInt(112)));
444    }
445
446    /// A second overload over BIGINT inputs, so a set can dispatch on argument type.
447    struct BigIntSum;
448
449    impl VAggregate for BigIntSum {
450        type State = i64;
451        type Shared = ();
452
453        fn parameters() -> Result<Vec<LogicalType>> {
454            Ok(vec![LogicalType::of::<i64>()?])
455        }
456
457        fn return_type() -> Result<LogicalType> {
458            LogicalType::of::<i64>()
459        }
460
461        fn init() -> i64 {
462            0
463        }
464
465        fn update(
466            _shared: &(),
467            state: &mut i64,
468            input: &DataChunkHandle,
469            row: usize,
470        ) -> super::super::UdfResult<()> {
471            let v: i64 = input.vector(0)?.get(row)?;
472            *state += v;
473            Ok(())
474        }
475
476        fn combine(
477            _shared: &(),
478            source: &i64,
479            target: &mut i64,
480        ) {
481            *target += *source;
482        }
483
484        fn finalize(
485            _shared: &(),
486            state: &i64,
487            output: &mut VectorMut<'_>,
488            row: usize,
489        ) -> super::super::UdfResult<()> {
490            output.set(row, *state)?;
491            Ok(())
492        }
493    }
494
495    #[test]
496    fn aggregate_set_dispatches_on_argument_type() {
497        let mut conn = Connection::open_in_memory().unwrap();
498        // One SQL name `my_sum` with INTEGER and BIGINT overloads.
499        conn.register_aggregate_function_set("my_sum", |b| {
500            b.add::<IntSum>()?;
501            b.add::<BigIntSum>()?;
502            Ok(())
503        })
504        .unwrap();
505
506        // INTEGER overload.
507        let mut r_int = conn
508            .execute("SELECT my_sum(CAST(v AS INTEGER)) AS s FROM (VALUES (1),(2),(3)) t(v)")
509            .unwrap();
510        assert_eq!(r_int.next().unwrap().unwrap().get("s"), Some(&DuckValue::BigInt(6)));
511
512        // BIGINT overload.
513        let mut r_big = conn
514            .execute("SELECT my_sum(CAST(v AS BIGINT)) AS s FROM (VALUES (10),(20)) t(v)")
515            .unwrap();
516        assert_eq!(r_big.next().unwrap().unwrap().get("s"), Some(&DuckValue::BigInt(30)));
517    }
518}
better_duck_core/udf/aggregate/mod.rs

better_duck_core/udf/aggregate/
mod.rs
