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    /// Statically pre-resolved secrets, delivered on `AggregateBindRequest.secrets`
30    /// when the aggregate advertises a required secret (see
31    /// [`FunctionMetadata::required_secrets`](crate::function::FunctionMetadata::required_secrets)).
32    /// Aggregates read a secret *value* here at bind time only — two-phase
33    /// `.get()`-style secret resolution is not supported for aggregates. Empty
34    /// when nothing was requested / resolved.
35    pub secrets: crate::secrets::Secrets,
36}
37
38/// An aggregate VGI function.
39pub trait AggregateFunction: Send + Sync {
40    fn name(&self) -> &str;
41    fn metadata(&self) -> FunctionMetadata;
42    fn argument_specs(&self) -> Vec<ArgSpec>;
43    /// Resolve the (single-column `result`) output schema.
44    fn on_bind(&self, params: &AggregateBindParams) -> Result<BindResponse>;
45    /// The serialized initial state for a fresh group.
46    fn initial_state(&self) -> Vec<u8>;
47    /// Fold the batch's rows into the per-group `states` map. `states` is
48    /// pre-loaded (initial state for new groups) for every group id present in
49    /// `group_ids`. `columns` are the input columns with the group-id column
50    /// already stripped.
51    fn update(
52        &self,
53        states: &mut HashMap<i64, Vec<u8>>,
54        group_ids: &Int64Array,
55        columns: &[ArrayRef],
56    ) -> Result<()>;
57    /// Merge `source` state into `target` state, returning the new target.
58    fn combine(&self, target: Vec<u8>, source: Vec<u8>) -> Result<Vec<u8>>;
59    /// Build the single-column output batch: one row per `group_ids` entry.
60    /// `states[i]` is the loaded state for `group_ids[i]` (`None` if unseen).
61    fn finalize(
62        &self,
63        output_schema: &SchemaRef,
64        group_ids: &Int64Array,
65        states: &[Option<Vec<u8>>],
66    ) -> Result<RecordBatch>;
67    /// Evaluate the windowed aggregate for each output row. `frames[i]` is the
68    /// list of `(begin, end)` sub-frames for output row `i` over the
69    /// `partition`'s input columns; returns the output column (one element per
70    /// row), matching `output_schema`. Only `supports_window` functions
71    /// override this; the default errors.
72    fn window(
73        &self,
74        _partition: &RecordBatch,
75        _output_schema: &SchemaRef,
76        _frames: &[Vec<(i64, i64)>],
77        _filter_mask: Option<&[bool]>,
78    ) -> Result<arrow_array::ArrayRef> {
79        Err(RpcError::runtime_error(
80            "window() not supported by this aggregate",
81        ))
82    }
83
84    /// Process one chunk of a streaming-partitioned session. The chunk's
85    /// columns are `[partition_key_cols.., order_key_cols.., value_cols..]`.
86    /// `states` is the cross-chunk per-partition state map (partition-key bytes
87    /// → opaque state bytes), loaded and persisted by the framework. Returns a
88    /// same-length output column. Only `streaming_partitioned` functions
89    /// override this; the default errors.
90    fn streaming_chunk(
91        &self,
92        _chunk: &RecordBatch,
93        _partition_key_count: usize,
94        _order_key_count: usize,
95        _states: &mut HashMap<Vec<u8>, Vec<u8>>,
96    ) -> Result<ArrayRef> {
97        Err(RpcError::runtime_error(
98            "streaming_chunk() not supported by this aggregate",
99        ))
100    }
101
102    /// Like [`Self::finalize`], but with access to the bind-time arguments (stashed
103    /// at `aggregate_bind`, reloaded here). Override for `ConstParam(phase=
104    /// "finalize")` aggregates like `vgi_percentile`. The default ignores them.
105    fn finalize_with_args(
106        &self,
107        output_schema: &SchemaRef,
108        group_ids: &Int64Array,
109        states: &[Option<Vec<u8>>],
110        _args: &crate::arguments::Arguments,
111    ) -> Result<RecordBatch> {
112        self.finalize(output_schema, group_ids, states)
113    }
114}
vgi/aggregate.rs

vgi/
aggregate.rs
