vgi/
aggregate.rs

1// Copyright 2025, 2026 Query Farm LLC - https://query.farm
2
3//! Aggregate function model (update / combine / finalize).
4//!
5//! DuckDB drives aggregates with per-group state. The worker keeps state in
6//! the cross-process [`FunctionStorage`](crate::storage::FunctionStorage) (KV
7//! mode), keyed by `(execution_id, group_id)`, because update / combine /
8//! finalize can run in different pooled worker processes. States are opaque to
9//! the C++ extension — the worker chooses any encoding.
10
11use std::collections::HashMap;
12
13use arrow_array::{ArrayRef, Int64Array, RecordBatch};
14use arrow_schema::SchemaRef;
15use vgi_rpc::{Result, RpcError};
16
17use crate::arguments::Arguments;
18use crate::function::{ArgSpec, BindResponse, FunctionMetadata};
19use crate::settings::Settings;
20
21/// The reserved group-id column prepended to UPDATE input batches.
22pub const GROUP_COLUMN_NAME: &str = "__vgi_group_id";
23
24/// Parameters for `aggregate_bind`.
25pub struct AggregateBindParams {
26    pub arguments: Arguments,
27    pub input_schema: Option<SchemaRef>,
28    pub settings: Settings,
29}
30
31/// An aggregate VGI function.
32pub trait AggregateFunction: Send + Sync {
33    fn name(&self) -> &str;
34    fn metadata(&self) -> FunctionMetadata;
35    fn argument_specs(&self) -> Vec<ArgSpec>;
36    /// Resolve the (single-column `result`) output schema.
37    fn on_bind(&self, params: &AggregateBindParams) -> Result<BindResponse>;
38    /// The serialized initial state for a fresh group.
39    fn initial_state(&self) -> Vec<u8>;
40    /// Fold the batch's rows into the per-group `states` map. `states` is
41    /// pre-loaded (initial state for new groups) for every group id present in
42    /// `group_ids`. `columns` are the input columns with the group-id column
43    /// already stripped.
44    fn update(
45        &self,
46        states: &mut HashMap<i64, Vec<u8>>,
47        group_ids: &Int64Array,
48        columns: &[ArrayRef],
49    ) -> Result<()>;
50    /// Merge `source` state into `target` state, returning the new target.
51    fn combine(&self, target: Vec<u8>, source: Vec<u8>) -> Result<Vec<u8>>;
52    /// Build the single-column output batch: one row per `group_ids` entry.
53    /// `states[i]` is the loaded state for `group_ids[i]` (`None` if unseen).
54    fn finalize(
55        &self,
56        output_schema: &SchemaRef,
57        group_ids: &Int64Array,
58        states: &[Option<Vec<u8>>],
59    ) -> Result<RecordBatch>;
60    /// Evaluate the windowed aggregate for each output row. `frames[i]` is the
61    /// list of `(begin, end)` sub-frames for output row `i` over the
62    /// `partition`'s input columns; returns the output column (one element per
63    /// row), matching `output_schema`. Only `supports_window` functions
64    /// override this; the default errors.
65    fn window(
66        &self,
67        _partition: &RecordBatch,
68        _output_schema: &SchemaRef,
69        _frames: &[Vec<(i64, i64)>],
70        _filter_mask: Option<&[bool]>,
71    ) -> Result<arrow_array::ArrayRef> {
72        Err(RpcError::runtime_error(
73            "window() not supported by this aggregate",
74        ))
75    }
76
77    /// Process one chunk of a streaming-partitioned session. The chunk's
78    /// columns are `[partition_key_cols.., order_key_cols.., value_cols..]`.
79    /// `states` is the cross-chunk per-partition state map (partition-key bytes
80    /// → opaque state bytes), loaded and persisted by the framework. Returns a
81    /// same-length output column. Only `streaming_partitioned` functions
82    /// override this; the default errors.
83    fn streaming_chunk(
84        &self,
85        _chunk: &RecordBatch,
86        _partition_key_count: usize,
87        _order_key_count: usize,
88        _states: &mut HashMap<Vec<u8>, Vec<u8>>,
89    ) -> Result<ArrayRef> {
90        Err(RpcError::runtime_error(
91            "streaming_chunk() not supported by this aggregate",
92        ))
93    }
94
95    /// Like [`Self::finalize`], but with access to the bind-time arguments (stashed
96    /// at `aggregate_bind`, reloaded here). Override for `ConstParam(phase=
97    /// "finalize")` aggregates like `vgi_percentile`. The default ignores them.
98    fn finalize_with_args(
99        &self,
100        output_schema: &SchemaRef,
101        group_ids: &Int64Array,
102        states: &[Option<Vec<u8>>],
103        _args: &crate::arguments::Arguments,
104    ) -> Result<RecordBatch> {
105        self.finalize(output_schema, group_ids, states)
106    }
107}
vgi/aggregate.rs

vgi/
aggregate.rs
