Feature request: custom aggregate functions

[jakeboone02]

Problem

TanStack DB's aggregate function support is hardcoded in packages/db/src/query/compiler/group-by.ts (getAggregateFunction switch). Only sum, count, avg, min, max are supported. Any other aggregate name throws UnsupportedAggregateFunctionError.

This makes it impossible for users to define domain-specific aggregations (e.g. group_concat, array_agg, json_agg, bitwise OR, geometric mean, etc.) without forking the package.

Related

• SQL STRING_AGG function in groupBy #422 — requests built-in STRING_AGG / group_concat

Proposal

Expose a registration API for custom aggregate functions, leveraging the existing { preMap, reduce, postMap? } contract from @tanstack/db-ivm's groupByOperators.

API

import {
  registerAggregate,
  unregisterAggregate,
  getRegisteredAggregates,
  Aggregate,
  toExpression,
} from '@tanstack/db'
import type { ExpressionLike } from '@tanstack/db'

// Register at app startup — factory receives valueExtractor + additional compiled args
registerAggregate('group_concat', (valueExtractor, additionalArgs) => {
  const separator = (additionalArgs[0] as string | undefined) ?? ','
  return {
    preMap: (data) => {
      const value = valueExtractor(data)
      return value != null ? String(value) : ''
    },
    reduce: (values) => {
      const parts: string[] = []
      for (const [value, multiplicity] of values) {
        if (value !== '') {
          for (let i = 0; i < multiplicity; i++) {
            parts.push(value as string)
          }
        }
      }
      return parts.join(separator)
    },
  }
})

// Helper to create aggregate IR nodes — extra args become additional expressions
function groupConcat(
  arg: ExpressionLike,
  separator: string = ',',
): Aggregate<string> {
  return new Aggregate('group_concat', [
    toExpression(arg),
    toExpression(separator),
  ])
}

const result = useLiveQuery((q) =>
  q
    .from({ todo: todosCollection })
    .groupBy(({ todo }) => todo.listId)
    .select(({ todo }) => ({
      listId: todo.listId,
      allNames: groupConcat(todo.text, ' | '),
    })),
)

Types

// Shape of a custom aggregate implementation (mirrors BasicAggregateFunction from @tanstack/db-ivm)
type CustomAggregateImpl<V = unknown, R = unknown> = {
  preMap: (data: unknown) => V
  reduce: (values: Array<[V, number]>) => V
  postMap?: (result: V) => R
}

// Factory receives the compiled value extractor AND additional compiled arg values
type CustomAggregateFactory = (
  valueExtractor: (...args: Array<unknown>) => unknown,
  additionalArgs: Array<unknown>,
) => CustomAggregateImpl

Implementation notes

• The valueExtractor passed to the factory is the raw value extractor (preserves the original value without numeric coercion). Internally this is ([, namespacedRow]) => compiledExpr(namespacedRow) — but for users it's opaque.
• Additional arguments: Any args beyond the first in the Aggregate IR node (e.g. new Aggregate("group_concat", [toExpression(field), toExpression(separator)])) are compiled and evaluated against an empty context, then passed as additionalArgs[0], additionalArgs[1], etc. This allows parameterized aggregates (custom separators, precision values, flags, etc.).
• Values in the reduce callback's array are [V, number] where the number is multiplicity (how many rows mapped to that value). Users must iterate multiplicity times or account for it in their reduction.
• Type assertions (e.g. value as string) are needed in reduce since the generic types flow as unknown at the boundary.
• Registration is case-insensitive (names are lowercased internally).
• unregisterAggregate(name) returns boolean indicating whether the aggregate existed.
• getRegisteredAggregates() returns a ReadonlySet<string> of registered names.

Implementation sketch

Minimal change to getAggregateFunction in group-by.ts:

// Module-level registry
const customAggregates = new Map<string, CustomAggregateFactory>();

export function registerAggregate(name: string, factory: CustomAggregateFactory): void {
  customAggregates.set(name.toLowerCase(), factory);
}

export function unregisterAggregate(name: string): boolean {
  return customAggregates.delete(name.toLowerCase());
}

export function getRegisteredAggregates(): ReadonlySet<string> {
  return new Set(customAggregates.keys());
}

// In getAggregateFunction, replace the default case:
default: {
  const custom = customAggregates.get(aggExpr.name.toLowerCase());
  if (custom) {
    // Compile additional args (beyond the first) into static values
    const additionalArgs = aggExpr.args.slice(1).map((arg) => {
      const compiled = compileExpression(arg);
      return compiled({});
    });
    return custom(
      rawValueExtractor as (...args: Array<unknown>) => unknown,
      additionalArgs,
    );
  }
  throw new UnsupportedAggregateFunctionError(aggExpr.name);
}