From b76ee99d35c869f6513ae1f4b59d10417a3965bf Mon Sep 17 00:00:00 2001 From: Jake Boone Date: Tue, 28 Jul 2026 12:52:50 -0700 Subject: [PATCH 1/2] feat(db): support custom aggregate functions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a global, case-insensitive registry for user-defined aggregates. The group-by compiler now looks up registrations before the built-in switch, so custom names work anywhere built-ins do — select, having, and orderBy via $selected — and built-ins can be overridden (warned about in dev) and restored by unregistering. Public API: - createAggregate(name, factory): registers and returns a typed builder, so the aggregate name is declared once and its result type flows into select() - registerAggregate / unregisterAggregate / getRegisteredAggregates for dynamic registration - toExpression, ExpressionLike and the Aggregate type are now exported for the low-level path Factories receive the raw value extractor plus the row key, letting aggregates such as group_concat stay deterministic despite preMap outputs being consolidated by value hash. Arguments after the first are evaluated once at compile time and must be constant, otherwise NonConstantAggregateArgumentError is thrown; UnsupportedAggregateFunctionError now lists registered names. Closes #1558 --- .changeset/smart-pugs-listen.md | 5 + docs/guides/live-queries.md | 68 ++ packages/db/src/errors.ts | 22 +- packages/db/src/query/aggregates.ts | 171 +++++ packages/db/src/query/builder/functions.ts | 2 +- packages/db/src/query/compiler/group-by.ts | 43 +- packages/db/src/query/index.ts | 18 + .../tests/query/custom-aggregates.test-d.ts | 82 +++ .../db/tests/query/custom-aggregates.test.ts | 591 ++++++++++++++++++ 9 files changed, 998 insertions(+), 4 deletions(-) create mode 100644 .changeset/smart-pugs-listen.md create mode 100644 packages/db/src/query/aggregates.ts create mode 100644 packages/db/tests/query/custom-aggregates.test-d.ts create mode 100644 packages/db/tests/query/custom-aggregates.test.ts diff --git a/.changeset/smart-pugs-listen.md b/.changeset/smart-pugs-listen.md new file mode 100644 index 0000000000..31432c80da --- /dev/null +++ b/.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`. diff --git a/docs/guides/live-queries.md b/docs/guides/live-queries.md index 9b6ee18b06..9f6d74113a 100644 --- a/docs/guides/live-queries.md +++ b/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( + '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 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.`. 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. diff --git a/packages/db/src/errors.ts b/packages/db/src/errors.ts index 0bfd2f9969..0a65b04af9 100644 --- a/packages/db/src/errors.ts +++ b/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) { + 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`, + ) } } diff --git a/packages/db/src/query/aggregates.ts b/packages/db/src/query/aggregates.ts new file mode 100644 index 0000000000..22201c79e6 --- /dev/null +++ b/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 = { + 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 = ( + ctx: AggregateContext, + additionalArgs: Array, +) => CustomAggregateImpl + +// `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 + +/** Aggregate names implemented natively by the group-by compiler. */ +export const BUILTIN_AGGREGATE_NAMES: ReadonlySet = new Set([ + `sum`, + `count`, + `avg`, + `min`, + `max`, +]) + +const customAggregates = new Map() + +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 { + 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( + * `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 = []>( + name: string, + factory: ( + ctx: AggregateContext, + params: TParams, + ) => CustomAggregateImpl, +): (arg: ExpressionLike, ...params: TParams) => Aggregate { + registerAggregate(name, (ctx, additionalArgs) => + factory(ctx, additionalArgs as TParams), + ) + + return (arg: ExpressionLike, ...params: TParams) => + new Aggregate(name, [ + toExpression(arg), + ...params.map((param) => toExpression(param)), + ]) +} diff --git a/packages/db/src/query/builder/functions.ts b/packages/db/src/query/builder/functions.ts index 84c26e5ca4..c21a7e0579 100644 --- a/packages/db/src/query/builder/functions.ts +++ b/packages/db/src/query/builder/functions.ts @@ -45,7 +45,7 @@ type ComparisonOperandPrimitive = | null // Helper type for values that can be lowered to expressions. -type ExpressionLike = +export type ExpressionLike = | Aggregate | BasicExpression | RefProxy diff --git a/packages/db/src/query/compiler/group-by.ts b/packages/db/src/query/compiler/group-by.ts index c670de9649..b22dfeb35d 100644 --- a/packages/db/src/query/compiler/group-by.ts +++ b/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 { + 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 } /** diff --git a/packages/db/src/query/index.ts b/packages/db/src/query/index.ts index 8cda812b3f..019e69ae0e 100644 --- a/packages/db/src/query/index.ts +++ b/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' diff --git a/packages/db/tests/query/custom-aggregates.test-d.ts b/packages/db/tests/query/custom-aggregates.test-d.ts new file mode 100644 index 0000000000..85f4e55be1 --- /dev/null +++ b/packages/db/tests/query/custom-aggregates.test-d.ts @@ -0,0 +1,82 @@ +import { describe, expectTypeOf, test } from 'vitest' +import { createLiveQueryCollection } from '../../src/query/index.js' +import { createCollection } from '../../src/collection/index.js' +import { createAggregate } from '../../src/query/aggregates.js' +import { mockSyncCollectionOptions } from '../utils.js' +import type { Aggregate } from '../../src/query/ir.js' + +type Todo = { + id: number + listId: number + text: string + points: number +} + +const todosCollection = createCollection( + mockSyncCollectionOptions({ + id: `custom-agg-type-todos`, + getKey: (todo) => todo.id, + initialData: [], + autoIndex: `eager`, + }), +) + +const groupConcat = createAggregate( + `group_concat_types`, + (ctx, [separator = `,`]) => ({ + preMap: (entry) => String(ctx.value(entry) ?? ``), + reduce: (values) => values.map(([text]) => text).join(separator), + }), +) + +const pointsRange = createAggregate<[number, number]>( + `points_range_types`, + (ctx) => ({ + preMap: (entry) => Number(ctx.value(entry)), + reduce: (values) => values.reduce((acc, [value]) => acc + value, 0), + postMap: (result) => [result, result] as [number, number], + }), +) + +describe(`custom aggregate types`, () => { + test(`createAggregate returns a typed Aggregate builder`, () => { + expectTypeOf(groupConcat(`x`)).toEqualTypeOf>() + expectTypeOf(pointsRange(1)).toEqualTypeOf>() + }) + + test(`extra params are typed and optional params may be omitted`, () => { + groupConcat(`x`) + groupConcat(`x`, `|`) + // @ts-expect-error separator must be a string + groupConcat(`x`, 1) + // @ts-expect-error no third parameter is declared + groupConcat(`x`, `|`, `extra`) + // @ts-expect-error pointsRange declares no extra parameters + pointsRange(1, 2) + }) + + test(`result type flows into select() inference`, () => { + const summary = createLiveQueryCollection({ + startSync: false, + query: (q) => + q + .from({ todo: todosCollection }) + .groupBy(({ todo }) => todo.listId) + .select(({ todo }) => ({ + listId: todo.listId, + names: groupConcat(todo.text, ` | `), + range: pointsRange(todo.points), + })), + }) + + const row = summary.get(1) + expectTypeOf(row).toMatchTypeOf< + | { + listId: number + names: string + range: [number, number] + } + | undefined + >() + }) +}) diff --git a/packages/db/tests/query/custom-aggregates.test.ts b/packages/db/tests/query/custom-aggregates.test.ts new file mode 100644 index 0000000000..1db8b07c1b --- /dev/null +++ b/packages/db/tests/query/custom-aggregates.test.ts @@ -0,0 +1,591 @@ +import { afterEach, describe, expect, test, vi } from 'vitest' +import { createLiveQueryCollection } from '../../src/query/index.js' +import { createCollection } from '../../src/collection/index.js' +import { + createAggregate, + getRegisteredAggregates, + registerAggregate, + unregisterAggregate, +} from '../../src/query/aggregates.js' +import { Aggregate } from '../../src/query/ir.js' +import { toExpression } from '../../src/query/builder/ref-proxy.js' +import { count, gt, sum, upper } from '../../src/query/builder/functions.js' +import { + NonConstantAggregateArgumentError, + UnsupportedAggregateFunctionError, +} from '../../src/errors.js' +import { mockSyncCollectionOptions } from '../utils.js' +import type { + AggregateContext, + AggregateEntry, +} from '../../src/query/aggregates.js' + +type Todo = { + id: number + listId: number + text: string + points: number +} + +const sampleTodos: Array = [ + { id: 1, listId: 1, text: `alpha`, points: 1 }, + { id: 2, listId: 1, text: `beta`, points: 2 }, + { id: 3, listId: 1, text: `alpha`, points: 3 }, + { id: 4, listId: 2, text: `gamma`, points: 4 }, +] + +function createTodosCollection(id = `custom-agg-todos`) { + return createCollection( + mockSyncCollectionOptions({ + id, + getKey: (todo) => todo.id, + initialData: sampleTodos, + autoIndex: `eager`, + }), + ) +} + +/** + * Deterministic group_concat: pairs each value with its row key so rows are never + * consolidated by hash, then sorts by row key before joining. + */ +function groupConcatImpl(ctx: AggregateContext, separator: string) { + return { + preMap: (entry: AggregateEntry): [string, string] => [ + ctx.key(entry), + String(ctx.value(entry) ?? ``), + ], + reduce: (values: Array<[[string, string], number]>) => { + const parts: Array<[string, string]> = [] + for (const [[key, text], multiplicity] of values) { + for (let i = 0; i < multiplicity; i++) { + parts.push([key, text]) + } + } + parts.sort(([a], [b]) => (a < b ? -1 : a > b ? 1 : 0)) + return parts.map(([, text]) => text).join(separator) + }, + } +} + +const registeredInTest = new Set() + +function register(name: string, ...args: Array) { + registeredInTest.add(name.toLowerCase()) + return (registerAggregate as any)(name, ...args) +} + +afterEach(() => { + for (const name of registeredInTest) { + unregisterAggregate(name) + } + registeredInTest.clear() + vi.restoreAllMocks() +}) + +describe(`custom aggregate functions`, () => { + describe(`registry`, () => { + test(`register, list and unregister`, () => { + expect(getRegisteredAggregates().has(`my_agg`)).toBe(false) + + register(`my_agg`, (ctx: AggregateContext) => ({ + preMap: (entry: AggregateEntry) => ctx.value(entry), + reduce: (values: Array<[unknown, number]>) => values.length, + })) + + expect(getRegisteredAggregates().has(`my_agg`)).toBe(true) + expect(unregisterAggregate(`my_agg`)).toBe(true) + expect(unregisterAggregate(`my_agg`)).toBe(false) + expect(getRegisteredAggregates().has(`my_agg`)).toBe(false) + }) + + test(`names are case-insensitive`, () => { + register(`MixedCase`, (ctx: AggregateContext) => ({ + preMap: (entry: AggregateEntry) => ctx.value(entry), + reduce: () => 0, + })) + + expect(getRegisteredAggregates().has(`mixedcase`)).toBe(true) + expect(unregisterAggregate(`MIXEDCASE`)).toBe(true) + }) + + test(`getRegisteredAggregates returns a snapshot, not a live view`, () => { + const before = getRegisteredAggregates() + register(`snapshot_agg`, () => ({ + preMap: () => 0, + reduce: () => 0, + })) + + expect(before.has(`snapshot_agg`)).toBe(false) + expect(getRegisteredAggregates().has(`snapshot_agg`)).toBe(true) + }) + + test(`re-registering an existing custom aggregate warns`, () => { + const warn = vi.spyOn(console, `warn`).mockImplementation(() => {}) + const impl = () => ({ preMap: () => 0, reduce: () => 0 }) + + register(`dup_agg`, impl) + expect(warn).not.toHaveBeenCalled() + + register(`dup_agg`, impl) + expect(warn).toHaveBeenCalledOnce() + expect(warn.mock.calls[0]![0]).toContain(`replaces an existing custom`) + }) + + test(`overriding a built-in warns`, () => { + const warn = vi.spyOn(console, `warn`).mockImplementation(() => {}) + + register(`sum`, (ctx: AggregateContext) => ({ + preMap: (entry: AggregateEntry) => Number(ctx.value(entry)), + reduce: () => 0, + })) + + expect(warn).toHaveBeenCalledOnce() + expect(warn.mock.calls[0]![0]).toContain(`overrides the built-in`) + }) + }) + + describe(`queries`, () => { + test(`group_concat with a custom separator`, () => { + const groupConcat = createAggregate( + `group_concat`, + (ctx, [separator = `,`]) => groupConcatImpl(ctx, separator), + ) + registeredInTest.add(`group_concat`) + + const todos = createTodosCollection() + const result = createLiveQueryCollection({ + startSync: true, + query: (q) => + q + .from({ todo: todos }) + .groupBy(({ todo }) => todo.listId) + .select(({ todo }) => ({ + listId: todo.listId, + names: groupConcat(todo.text, ` | `), + })), + }) + + // Duplicate values ("alpha" twice) are preserved, ordered by row key + expect(result.get(1)?.names).toBe(`alpha | beta | alpha`) + expect(result.get(2)?.names).toBe(`gamma`) + }) + + test(`default parameter value is used when omitted`, () => { + const groupConcat = createAggregate( + `group_concat_default`, + (ctx, [separator = `,`]) => groupConcatImpl(ctx, separator), + ) + registeredInTest.add(`group_concat_default`) + + const todos = createTodosCollection() + const result = createLiveQueryCollection({ + startSync: true, + query: (q) => + q + .from({ todo: todos }) + .groupBy(({ todo }) => todo.listId) + .select(({ todo }) => ({ + listId: todo.listId, + names: groupConcat(todo.text), + })), + }) + + expect(result.get(1)?.names).toBe(`alpha,beta,alpha`) + }) + + test(`multiplicity is reported for consolidated duplicate values`, () => { + // preMap ignores the row key, so identical values consolidate + const distinctCount = createAggregate( + `distinct_count`, + (ctx) => ({ + preMap: (entry) => ctx.value(entry), + reduce: (values) => + values.filter(([, multiplicity]) => multiplicity > 0).length, + }), + ) + const totalCount = createAggregate(`total_count`, (ctx) => ({ + preMap: (entry) => ctx.value(entry), + reduce: (values) => + values.reduce((acc, [, multiplicity]) => acc + multiplicity, 0), + })) + registeredInTest.add(`distinct_count`) + registeredInTest.add(`total_count`) + + const todos = createTodosCollection() + const result = createLiveQueryCollection({ + startSync: true, + query: (q) => + q + .from({ todo: todos }) + .groupBy(({ todo }) => todo.listId) + .select(({ todo }) => ({ + listId: todo.listId, + distinct: distinctCount(todo.text), + total: totalCount(todo.text), + })), + }) + + // list 1 has alpha, beta, alpha + expect(result.get(1)?.distinct).toBe(2) + expect(result.get(1)?.total).toBe(3) + }) + + test(`postMap transforms the reduced value`, () => { + const range = createAggregate(`points_range`, (ctx) => ({ + preMap: (entry) => { + const value = Number(ctx.value(entry)) + return [value, value] as [number, number] + }, + reduce: (values) => { + let low = Infinity + let high = -Infinity + for (const [[min, max], multiplicity] of values) { + if (multiplicity <= 0) continue + low = Math.min(low, min) + high = Math.max(high, max) + } + return [low, high] as [number, number] + }, + postMap: ([low, high]) => `${low}-${high}`, + })) + registeredInTest.add(`points_range`) + + const todos = createTodosCollection() + const result = createLiveQueryCollection({ + startSync: true, + query: (q) => + q + .from({ todo: todos }) + .groupBy(({ todo }) => todo.listId) + .select(({ todo }) => ({ + listId: todo.listId, + range: range(todo.points), + })), + }) + + expect(result.get(1)?.range).toBe(`1-3`) + expect(result.get(2)?.range).toBe(`4-4`) + }) + + test(`aggregates the raw value without numeric coercion`, () => { + const firstValue = createAggregate(`first_value`, (ctx) => ({ + preMap: (entry) => + [ctx.key(entry), ctx.value(entry)] as [string, unknown], + reduce: (values) => { + const sorted = [...values] + .filter(([, multiplicity]) => multiplicity > 0) + .sort(([a], [b]) => (a[0] < b[0] ? -1 : a[0] > b[0] ? 1 : 0)) + return sorted[0]?.[0][1] + }, + })) + registeredInTest.add(`first_value`) + + const todos = createTodosCollection() + const result = createLiveQueryCollection({ + startSync: true, + query: (q) => + q + .from({ todo: todos }) + .groupBy(({ todo }) => todo.listId) + .select(({ todo }) => ({ + listId: todo.listId, + first: firstValue(todo.text), + })), + }) + + expect(result.get(1)?.first).toBe(`alpha`) + }) + + test(`the aggregated expression may be a computed function`, () => { + const groupConcat = createAggregate( + `group_concat_expr`, + (ctx, [separator = `,`]) => groupConcatImpl(ctx, separator), + ) + registeredInTest.add(`group_concat_expr`) + + const todos = createTodosCollection() + const result = createLiveQueryCollection({ + startSync: true, + query: (q) => + q + .from({ todo: todos }) + .groupBy(({ todo }) => todo.listId) + .select(({ todo }) => ({ + listId: todo.listId, + names: groupConcat(upper(todo.text)), + })), + }) + + expect(result.get(1)?.names).toBe(`ALPHA,BETA,ALPHA`) + }) + }) + + describe(`incremental updates`, () => { + test(`recomputes on insert, update and delete`, () => { + const groupConcat = createAggregate( + `group_concat_live`, + (ctx, [separator = `,`]) => groupConcatImpl(ctx, separator), + ) + registeredInTest.add(`group_concat_live`) + + const todos = createTodosCollection(`custom-agg-live`) + const result = createLiveQueryCollection({ + startSync: true, + query: (q) => + q + .from({ todo: todos }) + .groupBy(({ todo }) => todo.listId) + .select(({ todo }) => ({ + listId: todo.listId, + names: groupConcat(todo.text), + })), + }) + + expect(result.get(1)?.names).toBe(`alpha,beta,alpha`) + + todos.utils.begin() + todos.utils.write({ + type: `insert`, + value: { id: 5, listId: 1, text: `delta`, points: 5 }, + }) + todos.utils.commit() + expect(result.get(1)?.names).toBe(`alpha,beta,alpha,delta`) + + todos.utils.begin() + todos.utils.write({ + type: `update`, + value: { id: 2, listId: 1, text: `BETA`, points: 2 }, + }) + todos.utils.commit() + expect(result.get(1)?.names).toBe(`alpha,BETA,alpha,delta`) + + todos.utils.begin() + todos.utils.write({ + type: `delete`, + value: { id: 3, listId: 1, text: `alpha`, points: 3 }, + }) + todos.utils.commit() + expect(result.get(1)?.names).toBe(`alpha,BETA,delta`) + + // Removing the last row of a group removes the group + todos.utils.begin() + todos.utils.write({ + type: `delete`, + value: { id: 4, listId: 2, text: `gamma`, points: 4 }, + }) + todos.utils.commit() + expect(result.get(2)).toBeUndefined() + }) + }) + + describe(`integration with other clauses`, () => { + test(`works in a HAVING clause`, () => { + const totalPoints = createAggregate(`total_points`, (ctx) => ({ + preMap: (entry) => [ctx.key(entry), Number(ctx.value(entry))] as const, + reduce: (values) => + values.reduce( + (acc, [[, points], multiplicity]) => acc + points * multiplicity, + 0, + ), + })) + registeredInTest.add(`total_points`) + + const todos = createTodosCollection() + const result = createLiveQueryCollection({ + startSync: true, + query: (q) => + q + .from({ todo: todos }) + .groupBy(({ todo }) => todo.listId) + .select(({ todo }) => ({ + listId: todo.listId, + points: totalPoints(todo.points), + })) + .having(({ todo }) => gt(totalPoints(todo.points), 5)), + }) + + // list 1 => 6, list 2 => 4 + expect(result.size).toBe(1) + expect(result.get(1)?.points).toBe(6) + }) + + test(`works nested inside an expression`, () => { + const totalPoints = createAggregate(`nested_points`, (ctx) => ({ + preMap: (entry) => [ctx.key(entry), Number(ctx.value(entry))] as const, + reduce: (values) => + values.reduce( + (acc, [[, points], multiplicity]) => acc + points * multiplicity, + 0, + ), + })) + registeredInTest.add(`nested_points`) + + const todos = createTodosCollection() + const result = createLiveQueryCollection({ + startSync: true, + query: (q) => + q + .from({ todo: todos }) + .groupBy(({ todo }) => todo.listId) + .select(({ todo }) => ({ + listId: todo.listId, + doubled: gt(totalPoints(todo.points), 5), + withCount: count(todo.id), + })), + }) + + expect(result.get(1)?.doubled).toBe(true) + expect(result.get(2)?.doubled).toBe(false) + expect(result.get(1)?.withCount).toBe(3) + }) + + test(`can be ordered by via $selected`, () => { + const totalPoints = createAggregate(`ordered_points`, (ctx) => ({ + preMap: (entry) => [ctx.key(entry), Number(ctx.value(entry))] as const, + reduce: (values) => + values.reduce( + (acc, [[, points], multiplicity]) => acc + points * multiplicity, + 0, + ), + })) + registeredInTest.add(`ordered_points`) + + const todos = createTodosCollection() + const result = createLiveQueryCollection({ + startSync: true, + query: (q) => + q + .from({ todo: todos }) + .groupBy(({ todo }) => todo.listId) + .select(({ todo }) => ({ + listId: todo.listId, + points: totalPoints(todo.points), + })) + .orderBy(({ $selected }) => $selected.points, `desc`), + }) + + expect(result.toArray.map((row) => row.listId)).toEqual([1, 2]) + }) + }) + + describe(`built-in overrides`, () => { + test(`custom registration takes precedence and unregistering restores the built-in`, () => { + vi.spyOn(console, `warn`).mockImplementation(() => {}) + + const todos = createTodosCollection() + const builtInResult = createLiveQueryCollection({ + startSync: true, + query: (q) => + q + .from({ todo: todos }) + .groupBy(({ todo }) => todo.listId) + .select(({ todo }) => ({ + listId: todo.listId, + points: sum(todo.points), + })), + }) + expect(builtInResult.get(1)?.points).toBe(6) + + register(`sum`, (ctx: AggregateContext) => ({ + preMap: (entry: AggregateEntry) => Number(ctx.value(entry)), + reduce: () => -1, + })) + + const overriddenResult = createLiveQueryCollection({ + startSync: true, + query: (q) => + q + .from({ todo: todos }) + .groupBy(({ todo }) => todo.listId) + .select(({ todo }) => ({ + listId: todo.listId, + points: sum(todo.points), + })), + }) + expect(overriddenResult.get(1)?.points).toBe(-1) + + // Already-compiled queries keep the implementation they compiled with + expect(builtInResult.get(1)?.points).toBe(6) + + expect(unregisterAggregate(`sum`)).toBe(true) + const restoredResult = createLiveQueryCollection({ + startSync: true, + query: (q) => + q + .from({ todo: todos }) + .groupBy(({ todo }) => todo.listId) + .select(({ todo }) => ({ + listId: todo.listId, + points: sum(todo.points), + })), + }) + expect(restoredResult.get(1)?.points).toBe(6) + }) + }) + + describe(`errors`, () => { + test(`unknown aggregate throws and lists registered names`, () => { + register(`known_agg`, () => ({ preMap: () => 0, reduce: () => 0 })) + + const todos = createTodosCollection() + expect(() => + createLiveQueryCollection({ + startSync: true, + query: (q) => + q + .from({ todo: todos }) + .groupBy(({ todo }) => todo.listId) + .select(({ todo }) => ({ + listId: todo.listId, + value: new Aggregate(`nope`, [ + toExpression(todo.points), + ]) as any, + })), + }), + ).toThrow(UnsupportedAggregateFunctionError) + + try { + createLiveQueryCollection({ + startSync: true, + query: (q) => + q + .from({ todo: todos }) + .groupBy(({ todo }) => todo.listId) + .select(({ todo }) => ({ + listId: todo.listId, + value: new Aggregate(`nope`, [ + toExpression(todo.points), + ]) as any, + })), + }) + } catch (error) { + expect((error as Error).message).toContain(`known_agg`) + } + }) + + test(`non-constant additional argument throws`, () => { + const groupConcat = createAggregate( + `group_concat_bad_arg`, + (ctx, [separator = `,`]) => groupConcatImpl(ctx, separator), + ) + registeredInTest.add(`group_concat_bad_arg`) + + const todos = createTodosCollection() + expect(() => + createLiveQueryCollection({ + startSync: true, + query: (q) => + q + .from({ todo: todos }) + .groupBy(({ todo }) => todo.listId) + .select(({ todo }) => ({ + listId: todo.listId, + // separator references a column, which is not constant + names: groupConcat(todo.text, todo.text as unknown as string), + })), + }), + ).toThrow(NonConstantAggregateArgumentError) + }) + }) +}) From ff07e1234c852429d4143499ced32e191519c339 Mon Sep 17 00:00:00 2001 From: Jake Boone Date: Tue, 28 Jul 2026 13:18:46 -0700 Subject: [PATCH 2/2] test(db): address review nits in custom aggregate tests --- .../db/tests/query/custom-aggregates.test.ts | 29 +++++++------------ 1 file changed, 10 insertions(+), 19 deletions(-) diff --git a/packages/db/tests/query/custom-aggregates.test.ts b/packages/db/tests/query/custom-aggregates.test.ts index 1db8b07c1b..923ce5ae96 100644 --- a/packages/db/tests/query/custom-aggregates.test.ts +++ b/packages/db/tests/query/custom-aggregates.test.ts @@ -70,9 +70,12 @@ function groupConcatImpl(ctx: AggregateContext, separator: string) { const registeredInTest = new Set() -function register(name: string, ...args: Array) { +function register( + name: string, + factory: Parameters[1], +): void { registeredInTest.add(name.toLowerCase()) - return (registerAggregate as any)(name, ...args) + registerAggregate(name, factory) } afterEach(() => { @@ -529,22 +532,7 @@ describe(`custom aggregate functions`, () => { register(`known_agg`, () => ({ preMap: () => 0, reduce: () => 0 })) const todos = createTodosCollection() - expect(() => - createLiveQueryCollection({ - startSync: true, - query: (q) => - q - .from({ todo: todos }) - .groupBy(({ todo }) => todo.listId) - .select(({ todo }) => ({ - listId: todo.listId, - value: new Aggregate(`nope`, [ - toExpression(todo.points), - ]) as any, - })), - }), - ).toThrow(UnsupportedAggregateFunctionError) - + let caught: unknown try { createLiveQueryCollection({ startSync: true, @@ -560,8 +548,11 @@ describe(`custom aggregate functions`, () => { })), }) } catch (error) { - expect((error as Error).message).toContain(`known_agg`) + caught = error } + + expect(caught).toBeInstanceOf(UnsupportedAggregateFunctionError) + expect((caught as Error).message).toContain(`known_agg`) }) test(`non-constant additional argument throws`, () => {