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