Skip to content

vgi.aggregate_function

Module overview

Framework for implementing aggregate functions.

AggregateFunction provides a batch-oriented API for DuckDB aggregate functions (e.g., SELECT my_agg(col) FROM t GROUP BY category). The C++ side manages trivial per-group state (just an int64 group_id), while Python holds the real accumulation state in FunctionStorage.

Three phases:

UPDATE: accumulate input rows into per-group state
COMBINE: merge states from parallel workers
FINALIZE: produce one result per group

class `AggregateBindParams`

Description

Parameters passed to AggregateFunction.on_bind().

Attributes

attribute `args`

Arguments | None

The bound function Arguments, or None if none.

attribute `input_schema`

pa.Schema | None

Arrow schema of the aggregate’s input columns, or None.

attribute `settings`

dict[str, Any]

DuckDB session settings relevant to the function.

attribute `secrets`

SecretsAccessor

Accessor for the resolved secrets the function declared.

attribute `auth_context`

The caller’s authentication context (anonymous by default).

class `AggregateFunction`

Bases: vgi.function.Function

Description

Base class for aggregate functions.

Aggregate functions accumulate input rows into per-group state during UPDATE, merge parallel worker states during COMBINE, and produce one result row per group during FINALIZE.

Input columns are declared via Param annotations on update(), and the output type via Returns annotation — the same pattern as ScalarFunction.compute().

Example:

class SumFunction(AggregateFunction[SumState]):
  class Meta:
      name = "vgi_sum"

  @classmethod
  def initial_state(cls, params):
      return SumState()

  @classmethod
  def update(
      cls,
      states: dict[int, SumState],
      group_ids: pa.Int64Array,
      value: Annotated[pa.Int64Array, Param(doc="Column to sum")],
  ) -> None:
      ...

  @classmethod
  def combine(cls, source, target, params):
      return SumState(total=source.total + target.total)

  @classmethod
  def finalize(
      cls,
      group_ids: pa.Int64Array,
      states: dict[int, SumState],
      params: ProcessParams,
  ) -> Annotated[pa.RecordBatch, Returns(pa.int64())]:
      ...

Attributes

attribute `state_class`

type[TState] | None

The resolved TState type used for per-group accumulation state, inferred from the generic parameter; None until resolved by __init_subclass__.

Methods

method `on_bind`

on_bind(
params: AggregateBindParams,
**kwargs: Any = {},
) -> BindResponse

Override to provide output schema and optional bind-time logic.

Must return a BindResponse with an output_schema containing exactly one field (the aggregate result column).

method `catalog_output_schema`

catalog_output_schema() -> pa.Schema

Return output schema for catalog introspection.

method `initial_state`

initial_state(params: ProcessParams[Any]) -> TState

Create the initial state for a new group.

Called when a group_id is first encountered during UPDATE. Must return a valid TState instance representing the identity element (e.g., 0 for SUM, empty list for LISTAGG).

method `update`

update(*args: Any = (), **kwargs: Any = {}) -> None

Accumulate input rows into per-group state.

Declare input columns as Param-annotated parameters:

@classmethod
def update(
  cls,
  states: dict[int, MyState],
  group_ids: pa.Int64Array,
  value: Annotated[pa.Int64Array, Param(doc="Column to sum")],
) -> None:
  ...

The states dict is pre-populated with initial_state() for all new group_ids. group_ids is parallel to each column array.

IMPORTANT — reassign to persist. Treat state as immutable: to record a change you MUST write it back with states[gid] = new_state. The framework only persists a group whose entry you assigned during this call (plus groups already saved from an earlier batch). Mutating the existing object in place — e.g. states[gid].items.append(x) — is NOT detected for a group first seen in this batch, so its data is silently dropped and finalize() sees only initial_state(). This bites single-group / single-batch aggregates hardest. Do:

s = states[gid]
states[gid] = MyState(items=s.items + new_items)  # reassign

not:

states[gid].items.extend(new_items)  # in-place: may be lost

method `combine`

combine(
source: TState,
target: TState,
params: ProcessParams[Any],
) -> TState

Merge two partial states from parallel workers.

Returns the merged TState. Framework replaces target and removes source.

method `finalize`

finalize(*args: Any = (), **kwargs: Any = {}) -> Any

Produce results for the requested group_ids.

Annotate the return type with Returns:

@classmethod
def finalize(
  cls,
  group_ids: pa.Int64Array,
  states: dict[int, MyState],
  params: ProcessParams,
) -> Annotated[pa.RecordBatch, Returns(pa.int64())]:
  ...

Must return a RecordBatch with one row per group_id.

method `ensure_state`

ensure_state(
states: dict[int, TState],
group_id: int,
params: ProcessParams[Any],
) -> TState

Get or create state for a group_id.

The framework pre-populates the states dict before calling update() and finalize(), so this helper should not normally be needed. Provided for defensive coding.

Parameters

states: Mapping of group_id to its accumulated state.
group_id: The group to fetch or create state for.
params: The current process parameters (used to build fresh state).

Returns

The state for the given group_id.

method `window_init`

window_init(
partition: WindowPartition,
params: ProcessParams[Any],
) -> Any

Derive optional per-partition state from the raw partition.

Called once per partition before any window() call. Return any StreamStateCodec — anything with serialize_to_bytes() and deserialize_from_bytes(), which ArrowSerializableDataclass provides — so it can round-trip through storage, or None if no derived state is required. The return value is passed back to window() as window_state.

Default implementation returns None.

method `window_prepare`

window_prepare(
partition: WindowPartition,
window_state: Any,
params: ProcessParams[Any],
) -> Any

Derive per-partition state for the window() loop (optional hook).

Called once per partition, after window_init (or after the state is rehydrated from storage on a cold reload), before any window() call. The return value is passed as window_state to every window() call against this partition, replacing the opaque _WindowStatePlaceholder user code would otherwise receive.

Use this hook for one-shot per-partition work that window() would otherwise have to redo on every call: deserialise the _WindowStatePlaceholder, reshape NumPy buffers from window_init’s state, build symbol→index lookups, etc. Anything you would otherwise be tempted to memoise via a module-level dict.

The result lives in the framework’s per-partition cache and is dropped automatically when the partition is evicted from the LRU or its destructor fires.

Default implementation returns window_state unchanged — for aggregates that don’t define this hook, window() receives the placeholder (or None) exactly as it did before. Backward compatible.

method `window`

window(
rid: int,
subframes: list[tuple[int, int]],
partition: WindowPartition,
window_state: Any,
params: ProcessParams[Any],
) -> Any

Compute the aggregate value for one output row.

Parameters

rid: Partition-local row index being filled.
subframes: Frame ranges [(begin, end), …] — 1 for the default frame, 3 when EXCLUDE produces multiple subframes.
partition: The cached partition data.
window_state: window_prepare()’s return value if the function defines that hook; otherwise the value returned by window_init() (may be None), wrapped in a _WindowStatePlaceholder on cold reload.
params: Shared ProcessParams.

Returns

A Python scalar or Arrow-compatible value; the worker wraps it into an IPC batch matching the function’s output schema.

method `window_batch`

window_batch(
row_ids: list[int],
subframes: list[list[tuple[int, int]]],
partition: WindowPartition,
window_state: Any,
params: ProcessParams[Any],
) -> pa.Array[Any] | list[Any]

Compute the aggregate value for count consecutive output rows.

Default implementation calls :meth:window once per row. Override when per-row Python object construction dominates the call cost and you want to build the output as an Arrow array directly, bypassing the framework’s default pa.array(results, ...) conversion.

Parameters

row_ids: Partition-local row indices being filled. Length is the batch size.
subframes: subframes[i] is the frame ranges for output row row_ids[i]. Same shape as :meth:window’s subframes argument, one per row.
partition: The cached partition data.
window_state: As :meth:window.
params: As :meth:window.

Returns

Either a :class:pa.Array of length len(row_ids) matching the function’s output type — shipped directly as the response with no further conversion — or a list[Any] of the same length, fed through pa.array(results, type=output_type) (equivalent to the default per-row path).

method `streaming_open`

streaming_open(params: ProcessParams[Any]) -> Any

Build cross-partition global state for a streaming session.

Called once when aggregate_streaming_open arrives, before any chunk is processed. Return any object (it lives in an in-process cache keyed by execution_id for the duration of the session).

Typical contents: a dict of per-partition aggregate states (populated lazily as new partition keys appear in input chunks), plus any cross-partition resources to share — symbol intern tables, allocator pools, prepared output buffers.

Default implementation returns None (no shared state); the function still works if streaming_chunk keeps everything in local variables, but per-partition state would have to live somewhere caller-supplied.

method `streaming_chunk`

streaming_chunk(
chunk: pa.RecordBatch,
streaming_state: Any,
partition_key_count: int,
order_key_count: int,
params: ProcessParams[Any],
) -> pa.Array[Any] | list[Any]

Process one chunk of streaming input.

Parameters

chunk: Input rows for this batch. Schema layout is [partition_key_cols…, order_key_cols…, value_cols…] — the first partition_key_count columns are partition keys (used to dispatch to the right per-partition state), the next order_key_count are order keys (informational; may be used to verify monotonicity), the rest are the function’s value arguments in declaration order.
streaming_state: Whatever streaming_open returned. The framework passes the same object on every chunk; mutate in place to accumulate state across chunks.
partition_key_count: Number of leading columns that form the partition key.
order_key_count: Number of columns following the partition key that form the order key.
params: Shared ProcessParams.

Returns

Either a :class:pa.Array of length chunk.num_rows matching the function’s output type, or a list of the same length (which the framework converts via pa.array). Each output value is the cumulative aggregate snapshot at that input row’s position in its partition’s order.

method `streaming_close`

streaming_close(
streaming_state: Any,
params: ProcessParams[Any],
) -> None

Tear down streaming session state.

Called once when aggregate_streaming_close arrives, after the last chunk. Use to release any external resources held by streaming_state. The framework drops its reference after this call, so anything not held elsewhere is GCed naturally.

Default implementation is a no-op.

Inherited members (4)

get_metadata method · from MetadataMixin — Get the resolved metadata for this function class.
describe method · from MetadataMixin — Get metadata as a dictionary (for JSON serialization).
logger attribute · from Function
storage attribute · from Function

class `WindowPartition`

Description

Full partition data passed to a windowed aggregate callback.

Constructed by the worker from the aggregate_window_init RPC payload and re-hydrated on every aggregate_window call via storage.

Attributes

attribute `inputs`

pa.RecordBatch

The partition’s input RecordBatch (all input columns, all rows).

attribute `row_count`

int

Total number of rows in the partition.

attribute `filter_mask`

pa.BooleanArray

Boolean mask from an optional FILTER (WHERE ...) clause. Length equals row_count.

attribute `frame_stats`

tuple[tuple[int, int], tuple[int, int]]

((begin_delta, end_delta), (begin_delta, end_delta)) — DuckDB’s per-partition frame statistics for planning.

attribute `all_valid`

list[bool]

Per-input-column validity flag (True if no nulls in column).

Methods

method `filter`

filter(start: int, end: int) -> pa.RecordBatch

Slice the partition inputs for rows [start, end).