feat(db): support custom aggregate functions

.changeset/smart-pugs-listen.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/db': minor
+---
+
+Add support for custom aggregate functions. `createAggregate(name, factory)` registers an aggregate and returns a typed helper for use in `select()`, and the lower-level `registerAggregate` / `unregisterAggregate` / `getRegisteredAggregates` APIs are available for dynamic registration. Custom aggregates work anywhere built-ins do, including `having` and `orderBy` via `$selected`.

docs/guides/live-queries.md

@@ -1495,6 +1495,74 @@ const orderStats = createCollection(liveQueryCollectionOptions({
 
 See the [Aggregate Functions](#aggregate-functions) section for a complete list of available aggregate functions.
 
+### Custom Aggregate Functions
+
+If the built-in aggregates aren't enough, register your own with `createAggregate`. It registers the aggregate and returns a typed helper you can call in `select`:
+
+```ts
+import { createAggregate, createCollection, liveQueryCollectionOptions } from '@tanstack/db'
+
+// Concatenates the values of a group, ordered by row key
+const groupConcat = createAggregate<string, [separator?: string]>(
+  'group_concat',
+  (ctx, [separator = ',']) => ({
+    // Pair each value with its row key so rows stay distinct
+    preMap: (entry) => [ctx.key(entry), String(ctx.value(entry) ?? '')],
+    reduce: (values) => {
+      const rows: Array<[string, string]> = []
+      for (const [row, multiplicity] of values) {
+        for (let i = 0; i < multiplicity; i++) rows.push(row)
+      }
+      rows.sort(([a], [b]) => (a < b ? -1 : a > b ? 1 : 0))
+      return rows.map(([, text]) => text).join(separator)
+    },
+  })
+)
+
+const listSummaries = createCollection(liveQueryCollectionOptions({
+  query: (q) =>
+    q
+      .from({ todo: todosCollection })
+      .groupBy(({ todo }) => todo.listId)
+      .select(({ todo }) => ({
+        listId: todo.listId,
+        allNames: groupConcat(todo.text, ' | '), // typed as string
+      }))
+}))
+```
+
+An implementation has three parts:
+
+- `preMap(entry)` — maps one row to the value that gets aggregated. Use `ctx.value(entry)` for the raw value of the first argument (no numeric coercion) and `ctx.key(entry)` for the row's key.
+- `reduce(values)` — receives the **entire group** on every change as `[value, multiplicity]` pairs and returns the reduced value. It is a full recompute, not a delta, so no accumulator bookkeeping is needed.
+- `postMap(result)` — optional final transformation of the reduced value.
+
+Arguments after the first one become the `params` tuple passed to your factory. They are evaluated once at query-compile time and must be constants — referencing a column throws `NonConstantAggregateArgumentError`.
+
+> [!IMPORTANT]
+> Values returned by `preMap` are consolidated by value: two rows producing the same value become a single entry with a multiplicity of `2`, and iteration order is not row order. Ignoring `multiplicity` silently drops duplicates. When order or per-row identity matters, include `ctx.key(entry)` in the `preMap` output (as above) and sort in `reduce`.
+
+For dynamic scenarios there is also a lower-level API:
+
+```ts
+import { registerAggregate, unregisterAggregate, getRegisteredAggregates, IR, toExpression } from '@tanstack/db'
+
+registerAggregate('bit_or', (ctx) => ({
+  preMap: (entry) => Number(ctx.value(entry)) | 0,
+  reduce: (values) => values.reduce((acc, [value]) => acc | value, 0),
+}))
+
+// Build the IR node yourself
+const bitOr = (arg) => new IR.Aggregate('bit_or', [toExpression(arg)])
+
+getRegisteredAggregates() // ReadonlySet<string> of registered names
+unregisterAggregate('bit_or') // true if a registration existed
+```
+
+Registration is global and case-insensitive. Registering a name that already exists — including a built-in like `sum` — replaces it for queries compiled *afterwards* and logs a warning in development. Queries already compiled keep the implementation they were compiled with, so overriding built-ins can produce inconsistent results across your app; prefer a distinct name. Unregistering a name that shadowed a built-in restores the built-in.
+
+Custom aggregates work anywhere built-ins do, including `having` and ordering by `$selected.<alias>`. Because factories are plain functions, they cannot be serialized: with SSR, make sure the same registrations run on both the server and the client.
+
 ### Having Clauses
 
 Filter aggregated results using `having` - this is similar to the `where` clause, but is applied after the aggregation has been performed.

packages/db/src/errors.ts

@@ -631,8 +631,26 @@ export class NonAggregateExpressionNotInGroupByError extends GroupByError {
 }
 
 export class UnsupportedAggregateFunctionError extends GroupByError {
-  constructor(functionName: string) {
-    super(`Unsupported aggregate function: ${functionName}`)
+  constructor(functionName: string, registeredNames?: Iterable<string>) {
+    const registered = registeredNames ? [...registeredNames] : []
+    super(
+      `Unsupported aggregate function: ${functionName}` +
+        (registered.length > 0
+          ? `. Registered custom aggregates: ${registered.join(`, `)}`
+          : ``),
+    )
+  }
+}
+
+/**
+ * Error thrown when an argument after the first of an aggregate expression
+ * is not a constant (e.g. it references a column).
+ */
+export class NonConstantAggregateArgumentError extends GroupByError {
+  constructor(functionName: string, argIndex: number) {
+    super(
+      `Argument ${argIndex} of aggregate function '${functionName}' must be a constant expression, not a column reference`,
+    )
   }
 }
 

packages/db/src/query/aggregates.ts

@@ -0,0 +1,171 @@
+import { Aggregate } from './ir.js'
+import { toExpression } from './builder/ref-proxy.js'
+import type { ExpressionLike } from './builder/functions.js'
+import type { NamespacedRow } from '../types.js'
+
+/**
+ * A single row as seen by an aggregate: `[rowKey, namespacedRow]`.
+ */
+export type AggregateEntry = [string, NamespacedRow]
+
+/**
+ * Accessors handed to a custom aggregate factory.
+ */
+export type AggregateContext = {
+  /**
+   * Raw value of the aggregate's first argument for this row.
+   * No numeric coercion is applied.
+   */
+  value: (entry: AggregateEntry) => unknown
+  /**
+   * Stable per-row key. Use it to keep rows distinct (values emitted by
+   * `preMap` are consolidated by hash) or to order deterministically.
+   */
+  key: (entry: AggregateEntry) => string
+}
+
+/**
+ * Implementation of a custom aggregate, mirroring db-ivm's basic aggregate contract.
+ *
+ * `reduce` receives the complete consolidated multiset for the group on every
+ * change, as `[value, multiplicity]` pairs — it is a full recompute, not a delta.
+ * Ignoring `multiplicity` under-counts duplicate values.
+ */
+export type CustomAggregateImpl<TValue = unknown, TResult = unknown> = {
+  preMap: (entry: AggregateEntry) => TValue
+  reduce: (values: Array<[TValue, number]>) => TValue
+  postMap?: (result: TValue) => TResult
+}
+
+/**
+ * Factory that builds a custom aggregate implementation for one compiled query.
+ *
+ * `additionalArgs` holds the evaluated values of any arguments after the first
+ * one in the aggregate expression; they must be constant expressions.
+ */
+export type CustomAggregateFactory<TValue = any, TResult = unknown> = (
+  ctx: AggregateContext,
+  additionalArgs: Array<unknown>,
+) => CustomAggregateImpl<TValue, TResult>
+
+// `any` for the value type: it is existential from the registry's point of view,
+// and `unknown` would make user implementations non-assignable (contravariance).
+type AnyCustomAggregateFactory = CustomAggregateFactory<any, unknown>
+
+/** Aggregate names implemented natively by the group-by compiler. */
+export const BUILTIN_AGGREGATE_NAMES: ReadonlySet<string> = new Set([
+  `sum`,
+  `count`,
+  `avg`,
+  `min`,
+  `max`,
+])
+
+const customAggregates = new Map<string, AnyCustomAggregateFactory>()
+
+const DEV =
+  typeof process !== `undefined` && process.env.NODE_ENV !== `production`
+
+/**
+ * Registers a custom aggregate function under `name` (case-insensitive).
+ *
+ * Re-registering a name — including a built-in — replaces the previous
+ * implementation for queries compiled afterwards and warns in development.
+ * Already-compiled live queries keep the implementation they were compiled with.
+ */
+export function registerAggregate(
+  name: string,
+  factory: AnyCustomAggregateFactory,
+): void {
+  const normalized = name.toLowerCase()
+
+  if (DEV) {
+    if (BUILTIN_AGGREGATE_NAMES.has(normalized)) {
+      console.warn(
+        `[@tanstack/db] registerAggregate("${name}") overrides the built-in ` +
+          `aggregate "${normalized}". This affects every query compiled afterwards, ` +
+          `app-wide. Already-compiled queries keep the built-in behavior.`,
+      )
+    } else if (customAggregates.has(normalized)) {
+      console.warn(
+        `[@tanstack/db] registerAggregate("${name}") replaces an existing custom ` +
+          `aggregate registration. Queries compiled before this call keep the ` +
+          `previous implementation.`,
+      )
+    }
+  }
+
+  customAggregates.set(normalized, factory)
+}
+
+/**
+ * Removes a custom aggregate registration.
+ *
+ * If the name shadowed a built-in, the built-in becomes active again because
+ * the compiler falls back to it when no registration exists.
+ *
+ * @returns whether a registration existed for the name
+ */
+export function unregisterAggregate(name: string): boolean {
+  return customAggregates.delete(name.toLowerCase())
+}
+
+/** Names of all currently registered custom aggregates. */
+export function getRegisteredAggregates(): ReadonlySet<string> {
+  return new Set(customAggregates.keys())
+}
+
+/** Looks up a registered custom aggregate factory. Used by the compiler. */
+export function getCustomAggregate(
+  name: string,
+): AnyCustomAggregateFactory | undefined {
+  return customAggregates.get(name.toLowerCase())
+}
+
+/**
+ * Registers a custom aggregate and returns a typed builder function for use in
+ * `select()` callbacks.
+ *
+ * @param name - Aggregate name (case-insensitive)
+ * @param factory - Builds the aggregate implementation from the row accessors
+ *   and the evaluated extra parameters
+ * @returns a function taking the aggregated expression plus the extra parameters
+ *
+ * @example
+ * ```ts
+ * const groupConcat = createAggregate<string, [separator?: string]>(
+ *   `group_concat`,
+ *   (ctx, [separator = `,`]) => ({
+ *     preMap: (entry) => [ctx.key(entry), String(ctx.value(entry) ?? ``)],
+ *     reduce: (values) =>
+ *       values
+ *         .filter(([, multiplicity]) => multiplicity > 0)
+ *         .sort(([a], [b]) => (a[0] < b[0] ? -1 : a[0] > b[0] ? 1 : 0))
+ *         .map(([[, text]]) => text)
+ *         .join(separator),
+ *   }),
+ * )
+ *
+ * query.groupBy(({ todo }) => todo.listId).select(({ todo }) => ({
+ *   listId: todo.listId,
+ *   names: groupConcat(todo.text, ` | `),
+ * }))
+ * ```
+ */
+export function createAggregate<TResult, TParams extends Array<unknown> = []>(
+  name: string,
+  factory: (
+    ctx: AggregateContext,
+    params: TParams,
+  ) => CustomAggregateImpl<any, TResult>,
+): (arg: ExpressionLike, ...params: TParams) => Aggregate<TResult> {
+  registerAggregate(name, (ctx, additionalArgs) =>
+    factory(ctx, additionalArgs as TParams),
+  )
+
+  return (arg: ExpressionLike, ...params: TParams) =>
+    new Aggregate<TResult>(name, [
+      toExpression(arg),
+      ...params.map((param) => toExpression(param)),
+    ])
+}

packages/db/src/query/builder/functions.ts

@@ -45,7 +45,7 @@ type ComparisonOperandPrimitive<T extends string | number | boolean> =
   | null
 
 // Helper type for values that can be lowered to expressions.
-type ExpressionLike =
+export type ExpressionLike =
   | Aggregate
   | BasicExpression
   | RefProxy<any>

packages/db/src/query/compiler/group-by.ts

@@ -15,9 +15,11 @@ import {
 import {
   AggregateFunctionNotInSelectError,
   NonAggregateExpressionNotInGroupByError,
+  NonConstantAggregateArgumentError,
   UnknownHavingExpressionTypeError,
   UnsupportedAggregateFunctionError,
 } from '../../errors.js'
+import { getCustomAggregate, getRegisteredAggregates } from '../aggregates.js'
 import {
   compileExpression,
   isCaseWhenConditionTrue,
@@ -551,6 +553,15 @@ function getAggregateFunction(aggExpr: Aggregate) {
     return compiledExpr(namespacedRow)
   }
 
+  // Custom registrations take precedence so that built-ins can be overridden
+  const custom = getCustomAggregate(aggExpr.name)
+  if (custom) {
+    return custom(
+      { value: rawValueExtractor, key: ([key]) => key },
+      compileAdditionalAggregateArgs(aggExpr),
+    )
+  }
+
   // Return the appropriate aggregate function
   switch (aggExpr.name.toLowerCase()) {
     case `sum`:
@@ -564,8 +575,38 @@ function getAggregateFunction(aggExpr: Aggregate) {
     case `max`:
       return max(valueExtractorForMinMax)
     default:
-      throw new UnsupportedAggregateFunctionError(aggExpr.name)
+      throw new UnsupportedAggregateFunctionError(
+        aggExpr.name,
+        getRegisteredAggregates(),
+      )
+  }
+}
+
+/**
+ * Evaluates the arguments after the first one of an aggregate expression into
+ * static values passed to a custom aggregate factory. They must be constant,
+ * since they are evaluated once at compile time against an empty row.
+ */
+function compileAdditionalAggregateArgs(aggExpr: Aggregate): Array<unknown> {
+  return aggExpr.args.slice(1).map((arg, index) => {
+    if (containsRef(arg)) {
+      throw new NonConstantAggregateArgumentError(aggExpr.name, index + 1)
+    }
+    return compileExpression(arg)({})
+  })
+}
+
+/**
+ * Whether an expression references a column, making it non-constant.
+ */
+function containsRef(expr: BasicExpression): boolean {
+  if (expr.type === `ref`) {
+    return true
   }
+  if (expr.type === `func`) {
+    return expr.args.some((arg) => containsRef(arg))
+  }
+  return false
 }
 
 /**

packages/db/src/query/index.ts

@@ -75,6 +75,24 @@ export {
 
 // Ref proxy utilities
 export type { Ref } from './builder/types.js'
+export { toExpression } from './builder/ref-proxy.js'
+export type { ExpressionLike } from './builder/functions.js'
+
+// Custom aggregate functions
+export type { Aggregate } from './ir.js'
+export {
+  registerAggregate,
+  unregisterAggregate,
+  getRegisteredAggregates,
+  createAggregate,
+  BUILTIN_AGGREGATE_NAMES,
+} from './aggregates.js'
+export type {
+  AggregateContext,
+  AggregateEntry,
+  CustomAggregateImpl,
+  CustomAggregateFactory,
+} from './aggregates.js'
 
 // Compiler
 export { compileQuery } from './compiler/index.js'