From f7717d20f802cf04b752ac457a2cd6556dc7c554 Mon Sep 17 00:00:00 2001 From: Kenny Daniel Date: Thu, 2 Jul 2026 16:51:31 -0700 Subject: [PATCH 1/7] Stream aggregates with bounded memory instead of buffering all rows Aggregate executors buffered every input row before evaluating, pinning entire column batches (and whole files) in memory. Queries whose aggregates are all COUNT, COUNTIF, SUM, AVG, MIN, or MAX now accumulate incrementally in 4000-row chunks, retaining one detached representative row per group. Finalized values are substituted into select, HAVING, and ORDER BY expressions as literals, keyed by aggregate node identity so same-named aggregates with different FILTERs stay separate. Queries with other aggregates (MEDIAN, ARRAY_AGG, ...) or subqueries fall back to the buffered path. SUM(LENGTH(text)) over a 414MB parquet file drops from 1.3GB to 155MB. --- .claude/CLAUDE.md | 1 + src/execute/aggregates.js | 18 + src/execute/streamingAggregate.js | 509 ++++++++++++++++++++++++ src/expression/evaluate.js | 2 +- test/execute/streamingAggregate.test.js | 92 +++++ 5 files changed, 621 insertions(+), 1 deletion(-) create mode 100644 src/execute/streamingAggregate.js create mode 100644 test/execute/streamingAggregate.test.js diff --git a/.claude/CLAUDE.md b/.claude/CLAUDE.md index c796222..2db768a 100644 --- a/.claude/CLAUDE.md +++ b/.claude/CLAUDE.md @@ -56,6 +56,7 @@ The main flow inside `executeSql({ tables, query, functions, signal })`: **Executor** (`src/execute/`): - `execute.js` - `executeSql()` and `executePlan()`, dispatches to node-specific executors - `join.js`, `aggregates.js`, `sort.js` - executors for join, aggregate, and sort nodes +- `streamingAggregate.js` - bounded-memory aggregation for streamable aggregates (COUNT, COUNTIF, SUM, AVG, MIN, MAX); other aggregates buffer rows in `aggregates.js` **Expression Evaluator** (`src/expression/`): - `evaluate.js` - evaluates expression AST nodes against rows diff --git a/src/execute/aggregates.js b/src/execute/aggregates.js index 5a92793..41c1093 100644 --- a/src/execute/aggregates.js +++ b/src/execute/aggregates.js @@ -2,6 +2,7 @@ import { derivedAlias } from '../expression/alias.js' import { evaluateExpr } from '../expression/evaluate.js' import { executePlan, selectColumnNames } from './execute.js' import { sortEntriesByTerms } from './sort.js' +import { planStreamingAggregates, streamingHashAggregateRows, streamingScalarAggregateRows } from './streamingAggregate.js' import { keyify } from './utils.js' import { yieldToEventLoop } from './yield.js' @@ -80,6 +81,14 @@ function aggregateContextRow(group, aggregateRow) { */ export function executeHashAggregate(plan, context) { const child = executePlan({ plan: plan.child, context }) + const specs = planStreamingAggregates(plan) + if (specs) { + return { + columns: selectColumnNames(plan.columns, child.columns), + maxRows: child.maxRows, + rows: streamingHashAggregateRows({ plan, specs, child, context }), + } + } return { columns: selectColumnNames(plan.columns, child.columns), maxRows: child.maxRows, @@ -196,6 +205,15 @@ export function executeScalarAggregate(plan, context) { } const child = executePlan({ plan: plan.child, context }) + const specs = planStreamingAggregates(plan) + if (specs) { + return { + columns: selectColumnNames(plan.columns, child.columns), + numRows: plan.having ? undefined : 1, + maxRows: 1, + rows: streamingScalarAggregateRows({ plan, specs, child, context }), + } + } return { columns: selectColumnNames(plan.columns, child.columns), numRows: plan.having ? undefined : 1, diff --git a/src/execute/streamingAggregate.js b/src/execute/streamingAggregate.js new file mode 100644 index 0000000..6088c3e --- /dev/null +++ b/src/execute/streamingAggregate.js @@ -0,0 +1,509 @@ +import { derivedAlias } from '../expression/alias.js' +import { evaluateAll, evaluateExpr } from '../expression/evaluate.js' +import { isAggregateFunc } from '../validation/functions.js' +import { compareForTerm, keyify } from './utils.js' +import { yieldToEventLoop } from './yield.js' + +/** + * @import { AsyncCells, AsyncRow, ExecuteContext, ExprNode, FunctionNode, QueryResults, SelectColumn, SqlPrimitive } from '../types.js' + * @import { HashAggregateNode, ScalarAggregateNode } from '../plan/types.js' + */ + +// Accumulate rows in chunks of this size so aborts can fire and async cells overlap +const CHUNK_SIZE = 4000 + +// Aggregate functions whose state can be accumulated one row at a time with +// bounded memory. Aggregates outside this set (MEDIAN, ARRAY_AGG, STDDEV, ...) +// need the full value set, so their queries buffer rows instead. +const STREAMABLE_FUNCS = new Set(['COUNT', 'COUNTIF', 'SUM', 'AVG', 'MIN', 'MAX']) + +/** + * Specs are keyed by aggregate node identity, so two aggregates that differ + * only in FILTER accumulate separately even though their derived aliases match. + * + * @typedef {{ + * node: FunctionNode, + * funcName: string, + * star: boolean, + * }} StreamingAggSpec + */ + +/** + * @typedef {{ + * count: number, + * sum: number, + * min: SqlPrimitive, + * max: SqlPrimitive, + * seen: Set | null, + * }} StreamingAccumulator + */ + +/** + * @typedef {{ + * firstRow: AsyncRow | undefined, + * accumulators: StreamingAccumulator[], + * }} StreamingGroup + */ + +/** + * Extracts the aggregate calls an aggregate node needs so they can be + * computed incrementally, without buffering the group's rows. Returns + * undefined when any expression needs a buffered group: an aggregate outside + * STREAMABLE_FUNCS, an aggregate over a non-scalar argument, or a subquery. + * + * @param {Pick & Partial>} plan + * @returns {StreamingAggSpec[] | undefined} + */ +export function planStreamingAggregates({ columns, having, orderBy }) { + /** @type {StreamingAggSpec[]} */ + const specs = [] + for (const col of columns) { + if (col.type === 'star') continue + if (!collectAggregates(col.expr, specs)) return + } + if (having && !collectAggregates(having, specs)) return + for (const term of orderBy ?? []) { + if (!collectAggregates(term.expr, specs)) return + } + return specs +} + +/** + * Walks an expression collecting streamable aggregate calls into specs. + * Returns false if the expression cannot be evaluated from precomputed + * aggregate values plus a single representative row. + * + * @param {ExprNode} node + * @param {StreamingAggSpec[]} specs + * @returns {boolean} + */ +function collectAggregates(node, specs) { + switch (node.type) { + case 'literal': + case 'identifier': + case 'star': + case 'interval': + return true + case 'unary': + return collectAggregates(node.argument, specs) + case 'binary': + return collectAggregates(node.left, specs) && collectAggregates(node.right, specs) + case 'cast': + return collectAggregates(node.expr, specs) + case 'case': + return (!node.caseExpr || collectAggregates(node.caseExpr, specs)) && + node.whenClauses.every(w => collectAggregates(w.condition, specs) && collectAggregates(w.result, specs)) && + (!node.elseResult || collectAggregates(node.elseResult, specs)) + case 'in valuelist': + return collectAggregates(node.expr, specs) && node.values.every(v => collectAggregates(v, specs)) + case 'function': { + const funcName = node.funcName.toUpperCase() + if (!isAggregateFunc(funcName)) { + return node.args.every(arg => collectAggregates(arg, specs)) + } + if (!STREAMABLE_FUNCS.has(funcName)) return false + const star = node.args[0]?.type === 'star' + if (!star && !node.args.every(arg => isScalarExpr(arg))) return false + if (node.filter && !isScalarExpr(node.filter)) return false + if (!specs.some(spec => spec.node === node)) { + specs.push({ node, funcName, star }) + } + return true + } + default: + // subqueries, EXISTS, IN (subquery), window functions + return false + } +} + +/** + * Reports whether an expression is a plain scalar over row values: no + * aggregates and no subqueries, so it can be evaluated per input row. + * + * @param {ExprNode} node + * @returns {boolean} + */ +function isScalarExpr(node) { + switch (node.type) { + case 'literal': + case 'identifier': + case 'star': + case 'interval': + return true + case 'unary': + return isScalarExpr(node.argument) + case 'binary': + return isScalarExpr(node.left) && isScalarExpr(node.right) + case 'cast': + return isScalarExpr(node.expr) + case 'case': + return (!node.caseExpr || isScalarExpr(node.caseExpr)) && + node.whenClauses.every(w => isScalarExpr(w.condition) && isScalarExpr(w.result)) && + (!node.elseResult || isScalarExpr(node.elseResult)) + case 'in valuelist': + return isScalarExpr(node.expr) && node.values.every(v => isScalarExpr(v)) + case 'function': + return !isAggregateFunc(node.funcName.toUpperCase()) && node.args.every(arg => isScalarExpr(arg)) + default: + return false + } +} + +/** + * Replaces each aggregate call in an expression with its computed value as a + * literal, so the rest of the expression can be evaluated against a single + * representative row. Nodes without aggregates are returned unchanged. + * + * @param {ExprNode} node + * @param {Map} values - computed value per aggregate node + * @returns {ExprNode} + */ +function substituteAggregates(node, values) { + switch (node.type) { + case 'unary': { + const argument = substituteAggregates(node.argument, values) + return argument === node.argument ? node : { ...node, argument } + } + case 'binary': { + const left = substituteAggregates(node.left, values) + const right = substituteAggregates(node.right, values) + return left === node.left && right === node.right ? node : { ...node, left, right } + } + case 'cast': { + const expr = substituteAggregates(node.expr, values) + return expr === node.expr ? node : { ...node, expr } + } + case 'case': { + const caseExpr = node.caseExpr && substituteAggregates(node.caseExpr, values) + const whenClauses = node.whenClauses.map(w => { + const condition = substituteAggregates(w.condition, values) + const result = substituteAggregates(w.result, values) + return condition === w.condition && result === w.result ? w : { ...w, condition, result } + }) + const elseResult = node.elseResult && substituteAggregates(node.elseResult, values) + return { ...node, caseExpr, whenClauses, elseResult } + } + case 'in valuelist': { + const expr = substituteAggregates(node.expr, values) + const valueNodes = node.values.map(v => substituteAggregates(v, values)) + return { ...node, expr, values: valueNodes } + } + case 'function': { + if (values.has(node)) { + return { + type: 'literal', + value: values.get(node) ?? null, + positionStart: node.positionStart, + positionEnd: node.positionEnd, + } + } + const args = node.args.map(arg => substituteAggregates(arg, values)) + return args.every((arg, i) => arg === node.args[i]) ? node : { ...node, args } + } + default: + return node + } +} + +/** + * @param {StreamingAggSpec} spec + * @returns {StreamingAccumulator} + */ +function newAccumulator(spec) { + return { + count: 0, + sum: 0, + min: null, + max: null, + seen: spec.funcName === 'COUNT' && spec.node.distinct ? new Set() : null, + } +} + +/** + * Folds one value into an accumulator, matching the buffered semantics in + * evaluate.js: COUNT counts non-null, COUNTIF counts truthy, MIN/MAX compare + * raw values, SUM/AVG only accumulate finite numbers. + * + * @param {StreamingAggSpec} spec + * @param {StreamingAccumulator} acc + * @param {SqlPrimitive} value + */ +function updateAccumulator(spec, acc, value) { + switch (spec.funcName) { + case 'COUNT': + if (spec.star) acc.count++ + else if (value != null) { + if (acc.seen) acc.seen.add(keyify(value)) + else acc.count++ + } + break + case 'COUNTIF': + if (value) acc.count++ + break + default: { // SUM, AVG, MIN, MAX + if (value == null) break + if (acc.min === null || value < acc.min) acc.min = value + if (acc.max === null || value > acc.max) acc.max = value + const num = Number(value) + if (Number.isFinite(num)) { + acc.sum += num + acc.count++ + } + } + } +} + +/** + * @param {StreamingAggSpec} spec + * @param {StreamingAccumulator} acc + * @returns {SqlPrimitive} + */ +function finalizeAccumulator(spec, acc) { + switch (spec.funcName) { + case 'COUNT': return acc.seen ? acc.seen.size : acc.count + case 'COUNTIF': return acc.count + case 'SUM': return acc.count === 0 ? null : acc.sum + case 'AVG': return acc.count === 0 ? null : acc.sum / acc.count + case 'MIN': return acc.min + case 'MAX': return acc.max + default: return null + } +} + +/** + * Folds one chunk of rows into the group accumulators. Group keys, FILTER + * conditions, and aggregate arguments are each evaluated across the whole + * chunk so async cells overlap; the chunk is released afterwards. + * + * @param {object} options + * @param {AsyncRow[]} options.chunk + * @param {ExprNode[]} options.groupBy + * @param {StreamingAggSpec[]} options.specs + * @param {Map} options.groups + * @param {ExecuteContext} options.context + * @returns {Promise} + */ +async function accumulateChunk({ chunk, groupBy, specs, groups, context }) { + /** @type {unknown[] | undefined} */ + let keys + if (groupBy.length === 1) { + const values = await evaluateAll(groupBy[0], chunk, context) + keys = values.map(v => keyify(v)) + } else if (groupBy.length > 1) { + const columns = await Promise.all(groupBy.map(expr => evaluateAll(expr, chunk, context))) + keys = chunk.map((_, j) => keyify(...columns.map(c => c[j]))) + } + + /** @type {(SqlPrimitive[] | undefined)[]} */ + const filters = new Array(specs.length) + /** @type {(SqlPrimitive[] | undefined)[]} */ + const args = new Array(specs.length) + for (let s = 0; s < specs.length; s++) { + const { node, star } = specs[s] + filters[s] = node.filter ? await evaluateAll(node.filter, chunk, context) : undefined + args[s] = star ? undefined : await evaluateAll(node.args[0], chunk, context) + } + + for (let j = 0; j < chunk.length; j++) { + const key = keys ? keys[j] : true + let group = groups.get(key) + if (!group) { + group = { firstRow: chunk[j], accumulators: specs.map(spec => newAccumulator(spec)) } + groups.set(key, group) + } + for (let s = 0; s < specs.length; s++) { + const filter = filters[s] + if (filter && !filter[j]) continue + const arg = args[s] + updateAccumulator(specs[s], group.accumulators[s], arg ? arg[j] : null) + } + } +} + +/** + * Consumes the child rows into per-group accumulators, holding at most one + * chunk of rows at a time. + * + * @param {object} options + * @param {QueryResults} options.child + * @param {ExprNode[]} options.groupBy + * @param {StreamingAggSpec[]} options.specs + * @param {ExecuteContext} options.context + * @returns {Promise>} + */ +async function accumulateGroups({ child, groupBy, specs, context }) { + /** @type {Map} */ + const groups = new Map() + /** @type {AsyncRow[]} */ + let chunk = [] + for await (const row of child.rows()) { + chunk.push(row) + if (chunk.length >= CHUNK_SIZE) { + await accumulateChunk({ chunk, groupBy, specs, groups, context }) + chunk = [] + await yieldToEventLoop() + context.signal?.throwIfAborted() + } + } + if (chunk.length) { + await accumulateChunk({ chunk, groupBy, specs, groups, context }) + } + context.signal?.throwIfAborted() + return groups +} + +/** + * Builds a group's output row and the context row visible to HAVING and + * grouped ORDER BY, by substituting the group's finalized aggregate values + * into the select expressions and evaluating them against the group's + * representative row. + * + * @param {object} options + * @param {SelectColumn[]} options.selectColumns + * @param {StreamingAggSpec[]} options.specs + * @param {StreamingGroup} options.group + * @param {ExecuteContext} options.context + * @returns {{ contextRow: AsyncRow, outputRow: AsyncRow, values: Map }} + */ +function finalizeGroup({ selectColumns, specs, group, context }) { + const firstRow = group.firstRow ?? { columns: [], cells: {} } + + /** @type {Map} */ + const values = new Map() + for (let s = 0; s < specs.length; s++) { + values.set(specs[s].node, finalizeAccumulator(specs[s], group.accumulators[s])) + } + + /** @type {string[]} */ + const columns = [] + /** @type {AsyncCells} */ + const cells = {} + for (const col of selectColumns) { + if (col.type === 'star') { + if (group.firstRow) { + const prefix = col.table ? `${col.table}.` : undefined + for (const key of firstRow.columns) { + if (prefix && !key.startsWith(prefix)) continue + const dotIndex = key.indexOf('.') + const outputKey = prefix ? key.substring(prefix.length) : dotIndex >= 0 ? key.substring(dotIndex + 1) : key + columns.push(outputKey) + cells[outputKey] = firstRow.cells[key] + } + } + } else { + const alias = col.alias ?? derivedAlias(col.expr) + const expr = substituteAggregates(col.expr, values) + columns.push(alias) + cells[alias] = () => evaluateExpr({ node: expr, row: firstRow, context }) + } + } + /** @type {AsyncRow} */ + const outputRow = { columns, cells } + + // Row visible to HAVING and grouped ORDER BY: the group's columns plus the + // select output aliases, mirroring the buffered aggregate context row. + /** @type {AsyncRow} */ + const contextRow = { + columns: [...firstRow.columns, ...columns], + cells: { ...firstRow.cells, ...cells }, + } + return { contextRow, outputRow, values } +} + +/** + * Streaming GROUP BY execution: accumulates aggregates incrementally instead + * of buffering every input row, then applies HAVING and grouped ORDER BY + * against the finalized aggregate values. + * + * @param {object} options + * @param {HashAggregateNode} options.plan + * @param {StreamingAggSpec[]} options.specs + * @param {QueryResults} options.child + * @param {ExecuteContext} options.context + * @returns {() => AsyncGenerator} + */ +export function streamingHashAggregateRows({ plan, specs, child, context }) { + return async function* () { + const groups = await accumulateGroups({ child, groupBy: plan.groupBy, specs, context }) + const { orderBy } = plan + + /** @type {{ outputRow: AsyncRow, contextRow: AsyncRow, values: Map, orderValues: SqlPrimitive[] }[]} */ + const outputRows = [] + for (const group of groups.values()) { + const { contextRow, outputRow, values } = finalizeGroup({ selectColumns: plan.columns, specs, group, context }) + if (plan.having) { + const passes = await evaluateExpr({ + node: substituteAggregates(plan.having, values), + row: contextRow, + context, + }) + if (!passes) continue + } + outputRows.push({ outputRow, contextRow, values, orderValues: [] }) + } + + if (orderBy?.length) { + // Evaluate each sort key across all groups in concurrent chunks so + // async cells and UDFs overlap + for (let t = 0; t < orderBy.length; t++) { + const term = orderBy[t] + for (let start = 0; start < outputRows.length; start += CHUNK_SIZE) { + if (start > 0) { + await yieldToEventLoop() + context.signal?.throwIfAborted() + } + const chunk = outputRows.slice(start, start + CHUNK_SIZE) + const termValues = await Promise.all(chunk.map(entry => evaluateExpr({ + node: substituteAggregates(term.expr, entry.values), + row: entry.contextRow, + context, + }))) + for (let j = 0; j < chunk.length; j++) { + chunk[j].orderValues[t] = termValues[j] + } + } + } + outputRows.sort((a, b) => { + for (let i = 0; i < orderBy.length; i++) { + const cmp = compareForTerm(a.orderValues[i], b.orderValues[i], orderBy[i]) + if (cmp) return cmp + } + return 0 + }) + } + + for (const { outputRow } of outputRows) { + yield outputRow + } + } +} + +/** + * Streaming scalar aggregate execution: the whole input is one group, + * accumulated incrementally with bounded memory. + * + * @param {object} options + * @param {ScalarAggregateNode} options.plan + * @param {StreamingAggSpec[]} options.specs + * @param {QueryResults} options.child + * @param {ExecuteContext} options.context + * @returns {() => AsyncGenerator} + */ +export function streamingScalarAggregateRows({ plan, specs, child, context }) { + return async function* () { + const groups = await accumulateGroups({ child, groupBy: [], specs, context }) + /** @type {StreamingGroup} */ + const group = groups.get(true) ?? { firstRow: undefined, accumulators: specs.map(spec => newAccumulator(spec)) } + + const { contextRow, outputRow, values } = finalizeGroup({ selectColumns: plan.columns, specs, group, context }) + if (plan.having) { + const passes = await evaluateExpr({ + node: substituteAggregates(plan.having, values), + row: contextRow, + context, + }) + if (!passes) return + } + yield outputRow + } +} diff --git a/src/expression/evaluate.js b/src/expression/evaluate.js index 1d43d36..7e90727 100644 --- a/src/expression/evaluate.js +++ b/src/expression/evaluate.js @@ -29,7 +29,7 @@ const YIELD_INTERVAL = 4000 * @param {ExecuteContext} context * @returns {Promise} */ -async function evaluateAll(node, rows, context) { +export async function evaluateAll(node, rows, context) { /** @type {SqlPrimitive[]} */ const results = new Array(rows.length) /** @type {Promise[]} */ diff --git a/test/execute/streamingAggregate.test.js b/test/execute/streamingAggregate.test.js new file mode 100644 index 0000000..07af335 --- /dev/null +++ b/test/execute/streamingAggregate.test.js @@ -0,0 +1,92 @@ +import { describe, expect, it } from 'vitest' +import { memorySource } from '../../src/backend/dataSource.js' +import { executeSql } from '../../src/execute/execute.js' +import { collect } from '../../src/index.js' + +// 10000 rows spans multiple 4000-row accumulation chunks, so these tests pin +// that group state carries across chunk boundaries in the streaming path. +const N = 10000 +const data = new Array(N) +for (let i = 0; i < N; i++) { + data[i] = { g: 'g' + i % 3, v: i % 100 } +} +const big = memorySource({ data }) + +describe('streaming aggregates', () => { + it('accumulates grouped aggregates across chunk boundaries', async () => { + const result = await collect(executeSql({ + tables: { big }, + query: 'SELECT g, COUNT(*) AS n, SUM(v) AS s, MIN(v) AS lo, MAX(v) AS hi FROM big GROUP BY g ORDER BY g', + })) + expect(result).toEqual([ + { g: 'g0', n: 3334, s: 165033, lo: 0, hi: 99 }, + { g: 'g1', n: 3333, s: 164967, lo: 0, hi: 99 }, + { g: 'g2', n: 3333, s: 165000, lo: 0, hi: 99 }, + ]) + }) + + it('accumulates scalar aggregates across chunk boundaries', async () => { + const result = await collect(executeSql({ + tables: { big }, + query: 'SELECT COUNT(*) AS n, AVG(v) AS a, COUNT(DISTINCT v) AS d FROM big', + })) + expect(result).toEqual([{ n: 10000, a: 49.5, d: 100 }]) + }) + + it('applies FILTER clauses across chunk boundaries', async () => { + const result = await collect(executeSql({ + tables: { big }, + query: 'SELECT COUNT(*) FILTER (WHERE v < 10) AS low, COUNT(*) FILTER (WHERE v >= 90) AS high FROM big', + })) + expect(result).toEqual([{ low: 1000, high: 1000 }]) + }) + + it('evaluates HAVING and ORDER BY on aggregates not in the select list', async () => { + const result = await collect(executeSql({ + tables: { big }, + query: 'SELECT g FROM big GROUP BY g HAVING COUNT(*) > 3333 ORDER BY SUM(v) DESC', + })) + expect(result).toEqual([{ g: 'g0' }]) + }) + + it('evaluates expressions over aggregates from substituted values', async () => { + const result = await collect(executeSql({ + tables: { big }, + query: 'SELECT g, SUM(v) / COUNT(*) AS mean, COUNT(*) + 1 AS n1 FROM big GROUP BY g ORDER BY g LIMIT 1', + })) + expect(result).toEqual([{ g: 'g0', mean: 49.5, n1: 3335 }]) + }) + + it('streams GROUP BY with no aggregates in the select list', async () => { + const result = await collect(executeSql({ + tables: { big }, + query: 'SELECT g FROM big GROUP BY g ORDER BY g', + })) + expect(result).toEqual([{ g: 'g0' }, { g: 'g1' }, { g: 'g2' }]) + }) + + it('returns aggregate defaults for empty input', async () => { + const empty = memorySource({ data: [], columns: ['v'] }) + const result = await collect(executeSql({ + tables: { empty }, + query: 'SELECT COUNT(*) AS n, SUM(v) AS s, MIN(v) AS lo FROM empty', + })) + expect(result).toEqual([{ n: 0, s: null, lo: null }]) + }) + + it('falls back to buffered aggregation for non-streamable aggregates', async () => { + const result = await collect(executeSql({ + tables: { big }, + query: 'SELECT g, MEDIAN(v) AS m, COUNT(*) AS n FROM big GROUP BY g ORDER BY g LIMIT 1', + })) + expect(result).toEqual([{ g: 'g0', m: 49.5, n: 3334 }]) + }) + + it('falls back to buffered aggregation when a subquery appears in the select list', async () => { + const result = await collect(executeSql({ + tables: { big, other: [{ x: 7 }] }, + query: 'SELECT (SELECT MAX(x) FROM other) AS mx, COUNT(*) AS n FROM big', + })) + expect(result).toEqual([{ mx: 7, n: 10000 }]) + }) +}) From f1c78b2e8f54cfaae7d248caea36d859e60fbfa1 Mon Sep 17 00:00:00 2001 From: Kenny Daniel Date: Fri, 3 Jul 2026 15:20:37 -0700 Subject: [PATCH 2/7] Drop per-group row retention when group keys cover all references The streaming aggregate path retained each group's first row as a representative, so a high-cardinality GROUP BY still held every input row. Track group key references during planning and substitute their values as literals alongside the aggregates; a representative row is only retained when an expression references columns outside the group keys (or SELECT *). Groups also finalize and yield one at a time when there is no ORDER BY, and FILTERed aggregate arguments are only evaluated for rows that pass the filter, matching buffered semantics. GROUP BY over 495k distinct keys now completes in a 256MB heap. --- src/execute/aggregates.js | 12 +- src/execute/streamingAggregate.js | 382 +++++++++++++++--------- test/execute/streamingAggregate.test.js | 164 +++++++++- 3 files changed, 418 insertions(+), 140 deletions(-) diff --git a/src/execute/aggregates.js b/src/execute/aggregates.js index 41c1093..41a2a7e 100644 --- a/src/execute/aggregates.js +++ b/src/execute/aggregates.js @@ -81,12 +81,12 @@ function aggregateContextRow(group, aggregateRow) { */ export function executeHashAggregate(plan, context) { const child = executePlan({ plan: plan.child, context }) - const specs = planStreamingAggregates(plan) - if (specs) { + const streaming = planStreamingAggregates(plan) + if (streaming) { return { columns: selectColumnNames(plan.columns, child.columns), maxRows: child.maxRows, - rows: streamingHashAggregateRows({ plan, specs, child, context }), + rows: streamingHashAggregateRows({ plan, streaming, child, context }), } } return { @@ -205,13 +205,13 @@ export function executeScalarAggregate(plan, context) { } const child = executePlan({ plan: plan.child, context }) - const specs = planStreamingAggregates(plan) - if (specs) { + const streaming = planStreamingAggregates(plan) + if (streaming) { return { columns: selectColumnNames(plan.columns, child.columns), numRows: plan.having ? undefined : 1, maxRows: 1, - rows: streamingScalarAggregateRows({ plan, specs, child, context }), + rows: streamingScalarAggregateRows({ plan, streaming, child, context }), } } return { diff --git a/src/execute/streamingAggregate.js b/src/execute/streamingAggregate.js index 6088c3e..7a0aeda 100644 --- a/src/execute/streamingAggregate.js +++ b/src/execute/streamingAggregate.js @@ -41,79 +41,146 @@ const STREAMABLE_FUNCS = new Set(['COUNT', 'COUNTIF', 'SUM', 'AVG', 'MIN', 'MAX' /** * @typedef {{ * firstRow: AsyncRow | undefined, + * keyValues: SqlPrimitive[], * accumulators: StreamingAccumulator[], * }} StreamingGroup */ +/** + * The streaming plan for an aggregate node: which aggregate calls to + * accumulate, which expression nodes are group key references (substituted + * from the group's key values), and whether any expression still needs a + * representative row from the group. When needsRow is false, no input rows + * are retained at all, so memory is bounded by the number of groups even for + * high-cardinality GROUP BY. + * + * @typedef {{ + * specs: StreamingAggSpec[], + * keyRefs: Map, + * needsRow: boolean, + * }} StreamingAggPlan + */ + +/** + * Structural signature of an expression node, ignoring source positions, so + * an expression repeated in SELECT and GROUP BY compares equal. + * + * @param {ExprNode} node + * @returns {string} + */ +function exprSig(node) { + return JSON.stringify(node, (key, value) => + key === 'positionStart' || key === 'positionEnd' ? undefined : value + ) +} + /** * Extracts the aggregate calls an aggregate node needs so they can be * computed incrementally, without buffering the group's rows. Returns * undefined when any expression needs a buffered group: an aggregate outside * STREAMABLE_FUNCS, an aggregate over a non-scalar argument, or a subquery. * - * @param {Pick & Partial>} plan - * @returns {StreamingAggSpec[] | undefined} + * @param {Pick & Partial>} plan + * @returns {StreamingAggPlan | undefined} */ -export function planStreamingAggregates({ columns, having, orderBy }) { +export function planStreamingAggregates({ columns, having, orderBy, groupBy }) { + const groupExprs = groupBy ?? [] + const groupSigs = groupExprs.map(exprSig) /** @type {StreamingAggSpec[]} */ const specs = [] - for (const col of columns) { - if (col.type === 'star') continue - if (!collectAggregates(col.expr, specs)) return - } - if (having && !collectAggregates(having, specs)) return - for (const term of orderBy ?? []) { - if (!collectAggregates(term.expr, specs)) return + /** @type {Map} */ + const keyRefs = new Map() + let needsRow = false + + /** + * Index of the group key the expression refers to, or -1. An unqualified + * column reference matches a qualified group key and vice versa. + * + * @param {ExprNode} node + * @returns {number} + */ + function matchGroupKey(node) { + const signature = exprSig(node) + for (let i = 0; i < groupExprs.length; i++) { + if (groupSigs[i] === signature) return i + } + if (node.type === 'identifier') { + for (let i = 0; i < groupExprs.length; i++) { + const expr = groupExprs[i] + if (expr.type === 'identifier' && expr.name === node.name && + (!expr.prefix || !node.prefix || expr.prefix === node.prefix)) return i + } + } + return -1 } - return specs -} -/** - * Walks an expression collecting streamable aggregate calls into specs. - * Returns false if the expression cannot be evaluated from precomputed - * aggregate values plus a single representative row. - * - * @param {ExprNode} node - * @param {StreamingAggSpec[]} specs - * @returns {boolean} - */ -function collectAggregates(node, specs) { - switch (node.type) { - case 'literal': - case 'identifier': - case 'star': - case 'interval': - return true - case 'unary': - return collectAggregates(node.argument, specs) - case 'binary': - return collectAggregates(node.left, specs) && collectAggregates(node.right, specs) - case 'cast': - return collectAggregates(node.expr, specs) - case 'case': - return (!node.caseExpr || collectAggregates(node.caseExpr, specs)) && - node.whenClauses.every(w => collectAggregates(w.condition, specs) && collectAggregates(w.result, specs)) && - (!node.elseResult || collectAggregates(node.elseResult, specs)) - case 'in valuelist': - return collectAggregates(node.expr, specs) && node.values.every(v => collectAggregates(v, specs)) - case 'function': { - const funcName = node.funcName.toUpperCase() - if (!isAggregateFunc(funcName)) { - return node.args.every(arg => collectAggregates(arg, specs)) + /** + * Walks an expression collecting streamable aggregate calls and group key + * references. Returns false if the expression cannot be evaluated from + * precomputed values plus a representative row. + * + * @param {ExprNode} node + * @returns {boolean} + */ + function walk(node) { + const keyIndex = matchGroupKey(node) + if (keyIndex >= 0) { + keyRefs.set(node, keyIndex) + return true } - if (!STREAMABLE_FUNCS.has(funcName)) return false - const star = node.args[0]?.type === 'star' - if (!star && !node.args.every(arg => isScalarExpr(arg))) return false - if (node.filter && !isScalarExpr(node.filter)) return false - if (!specs.some(spec => spec.node === node)) { - specs.push({ node, funcName, star }) + switch (node.type) { + case 'literal': + case 'interval': + return true + case 'identifier': + case 'star': + // resolves against the group's representative row + needsRow = true + return true + case 'unary': + return walk(node.argument) + case 'binary': + return walk(node.left) && walk(node.right) + case 'cast': + return walk(node.expr) + case 'case': + return (!node.caseExpr || walk(node.caseExpr)) && + node.whenClauses.every(w => walk(w.condition) && walk(w.result)) && + (!node.elseResult || walk(node.elseResult)) + case 'in valuelist': + return walk(node.expr) && node.values.every(walk) + case 'function': { + const funcName = node.funcName.toUpperCase() + if (!isAggregateFunc(funcName)) { + return node.args.every(walk) + } + if (!STREAMABLE_FUNCS.has(funcName)) return false + const star = node.args[0]?.type === 'star' + if (!star && !node.args.every(arg => isScalarExpr(arg))) return false + if (node.filter && !isScalarExpr(node.filter)) return false + if (!specs.some(spec => spec.node === node)) { + specs.push({ node, funcName, star }) + } + return true + } + default: + // subqueries, EXISTS, IN (subquery), window functions + return false } - return true } - default: - // subqueries, EXISTS, IN (subquery), window functions - return false + + for (const col of columns) { + if (col.type === 'star') { + needsRow = true + continue + } + if (!walk(col.expr)) return } + if (having && !walk(having)) return + for (const term of orderBy ?? []) { + if (!walk(term.expr)) return + } + return { specs, keyRefs, needsRow } } /** @@ -150,54 +217,55 @@ function isScalarExpr(node) { } /** - * Replaces each aggregate call in an expression with its computed value as a - * literal, so the rest of the expression can be evaluated against a single - * representative row. Nodes without aggregates are returned unchanged. + * Replaces each precomputed node in an expression with its value as a + * literal: aggregate calls (keyed by node identity) and group key references, + * so the rest of the expression can be evaluated against a representative + * row. Nodes without precomputed values are returned unchanged. * * @param {ExprNode} node - * @param {Map} values - computed value per aggregate node + * @param {Map} values - computed value per substituted node * @returns {ExprNode} */ -function substituteAggregates(node, values) { +function substituteValues(node, values) { + if (values.has(node)) { + return { + type: 'literal', + value: values.get(node) ?? null, + positionStart: node.positionStart, + positionEnd: node.positionEnd, + } + } switch (node.type) { case 'unary': { - const argument = substituteAggregates(node.argument, values) + const argument = substituteValues(node.argument, values) return argument === node.argument ? node : { ...node, argument } } case 'binary': { - const left = substituteAggregates(node.left, values) - const right = substituteAggregates(node.right, values) + const left = substituteValues(node.left, values) + const right = substituteValues(node.right, values) return left === node.left && right === node.right ? node : { ...node, left, right } } case 'cast': { - const expr = substituteAggregates(node.expr, values) + const expr = substituteValues(node.expr, values) return expr === node.expr ? node : { ...node, expr } } case 'case': { - const caseExpr = node.caseExpr && substituteAggregates(node.caseExpr, values) + const caseExpr = node.caseExpr && substituteValues(node.caseExpr, values) const whenClauses = node.whenClauses.map(w => { - const condition = substituteAggregates(w.condition, values) - const result = substituteAggregates(w.result, values) + const condition = substituteValues(w.condition, values) + const result = substituteValues(w.result, values) return condition === w.condition && result === w.result ? w : { ...w, condition, result } }) - const elseResult = node.elseResult && substituteAggregates(node.elseResult, values) + const elseResult = node.elseResult && substituteValues(node.elseResult, values) return { ...node, caseExpr, whenClauses, elseResult } } case 'in valuelist': { - const expr = substituteAggregates(node.expr, values) - const valueNodes = node.values.map(v => substituteAggregates(v, values)) + const expr = substituteValues(node.expr, values) + const valueNodes = node.values.map(v => substituteValues(v, values)) return { ...node, expr, values: valueNodes } } case 'function': { - if (values.has(node)) { - return { - type: 'literal', - value: values.get(node) ?? null, - positionStart: node.positionStart, - positionEnd: node.positionEnd, - } - } - const args = node.args.map(arg => substituteAggregates(arg, values)) + const args = node.args.map(arg => substituteValues(arg, values)) return args.every((arg, i) => arg === node.args[i]) ? node : { ...node, args } } default: @@ -280,18 +348,15 @@ function finalizeAccumulator(spec, acc) { * @param {ExprNode[]} options.groupBy * @param {StreamingAggSpec[]} options.specs * @param {Map} options.groups + * @param {boolean} options.needsRow - retain each group's first row? * @param {ExecuteContext} options.context * @returns {Promise} */ -async function accumulateChunk({ chunk, groupBy, specs, groups, context }) { - /** @type {unknown[] | undefined} */ - let keys - if (groupBy.length === 1) { - const values = await evaluateAll(groupBy[0], chunk, context) - keys = values.map(v => keyify(v)) - } else if (groupBy.length > 1) { - const columns = await Promise.all(groupBy.map(expr => evaluateAll(expr, chunk, context))) - keys = chunk.map((_, j) => keyify(...columns.map(c => c[j]))) +async function accumulateChunk({ chunk, groupBy, specs, groups, needsRow, context }) { + /** @type {SqlPrimitive[][] | undefined} */ + let keyColumns + if (groupBy.length) { + keyColumns = await Promise.all(groupBy.map(expr => evaluateAll(expr, chunk, context))) } /** @type {(SqlPrimitive[] | undefined)[]} */ @@ -300,15 +365,45 @@ async function accumulateChunk({ chunk, groupBy, specs, groups, context }) { const args = new Array(specs.length) for (let s = 0; s < specs.length; s++) { const { node, star } = specs[s] - filters[s] = node.filter ? await evaluateAll(node.filter, chunk, context) : undefined - args[s] = star ? undefined : await evaluateAll(node.args[0], chunk, context) + if (node.filter) { + const passes = await evaluateAll(node.filter, chunk, context) + filters[s] = passes + if (!star) { + // The buffered path filters the group before evaluating arguments, + // so only evaluate the argument for rows that pass the FILTER + /** @type {AsyncRow[]} */ + const passingRows = [] + /** @type {number[]} */ + const passingIndices = [] + for (let j = 0; j < chunk.length; j++) { + if (passes[j]) { + passingRows.push(chunk[j]) + passingIndices.push(j) + } + } + const values = await evaluateAll(node.args[0], passingRows, context) + const spread = new Array(chunk.length).fill(null) + for (let k = 0; k < passingIndices.length; k++) { + spread[passingIndices[k]] = values[k] + } + args[s] = spread + } + } else { + args[s] = star ? undefined : await evaluateAll(node.args[0], chunk, context) + } } for (let j = 0; j < chunk.length; j++) { - const key = keys ? keys[j] : true + const key = keyColumns + ? keyColumns.length === 1 ? keyify(keyColumns[0][j]) : keyify(...keyColumns.map(c => c[j])) + : true let group = groups.get(key) if (!group) { - group = { firstRow: chunk[j], accumulators: specs.map(spec => newAccumulator(spec)) } + group = { + firstRow: needsRow ? chunk[j] : undefined, + keyValues: keyColumns ? keyColumns.map(c => c[j]) : [], + accumulators: specs.map(spec => newAccumulator(spec)), + } groups.set(key, group) } for (let s = 0; s < specs.length; s++) { @@ -328,10 +423,11 @@ async function accumulateChunk({ chunk, groupBy, specs, groups, context }) { * @param {QueryResults} options.child * @param {ExprNode[]} options.groupBy * @param {StreamingAggSpec[]} options.specs + * @param {boolean} options.needsRow * @param {ExecuteContext} options.context * @returns {Promise>} */ -async function accumulateGroups({ child, groupBy, specs, context }) { +async function accumulateGroups({ child, groupBy, specs, needsRow, context }) { /** @type {Map} */ const groups = new Map() /** @type {AsyncRow[]} */ @@ -339,40 +435,44 @@ async function accumulateGroups({ child, groupBy, specs, context }) { for await (const row of child.rows()) { chunk.push(row) if (chunk.length >= CHUNK_SIZE) { - await accumulateChunk({ chunk, groupBy, specs, groups, context }) + await accumulateChunk({ chunk, groupBy, specs, groups, needsRow, context }) chunk = [] await yieldToEventLoop() context.signal?.throwIfAborted() } } if (chunk.length) { - await accumulateChunk({ chunk, groupBy, specs, groups, context }) + await accumulateChunk({ chunk, groupBy, specs, groups, needsRow, context }) } context.signal?.throwIfAborted() return groups } /** - * Builds a group's output row and the context row visible to HAVING and - * grouped ORDER BY, by substituting the group's finalized aggregate values - * into the select expressions and evaluating them against the group's - * representative row. + * Builds a group's output row by substituting the group's finalized + * aggregate and group key values into the select expressions and evaluating + * them against the group's representative row (an empty row when no + * expression needs one). * * @param {object} options * @param {SelectColumn[]} options.selectColumns * @param {StreamingAggSpec[]} options.specs + * @param {Map} options.keyRefs * @param {StreamingGroup} options.group * @param {ExecuteContext} options.context - * @returns {{ contextRow: AsyncRow, outputRow: AsyncRow, values: Map }} + * @returns {{ outputRow: AsyncRow, values: Map }} */ -function finalizeGroup({ selectColumns, specs, group, context }) { +function finalizeGroup({ selectColumns, specs, keyRefs, group, context }) { const firstRow = group.firstRow ?? { columns: [], cells: {} } - /** @type {Map} */ + /** @type {Map} */ const values = new Map() for (let s = 0; s < specs.length; s++) { values.set(specs[s].node, finalizeAccumulator(specs[s], group.accumulators[s])) } + for (const [node, keyIndex] of keyRefs) { + values.set(node, group.keyValues[keyIndex]) + } /** @type {string[]} */ const columns = [] @@ -392,22 +492,31 @@ function finalizeGroup({ selectColumns, specs, group, context }) { } } else { const alias = col.alias ?? derivedAlias(col.expr) - const expr = substituteAggregates(col.expr, values) + const expr = substituteValues(col.expr, values) columns.push(alias) cells[alias] = () => evaluateExpr({ node: expr, row: firstRow, context }) } } /** @type {AsyncRow} */ const outputRow = { columns, cells } + return { outputRow, values } +} - // Row visible to HAVING and grouped ORDER BY: the group's columns plus the - // select output aliases, mirroring the buffered aggregate context row. - /** @type {AsyncRow} */ - const contextRow = { - columns: [...firstRow.columns, ...columns], - cells: { ...firstRow.cells, ...cells }, +/** + * Builds the row visible to HAVING and grouped ORDER BY: the group's + * representative columns plus the select output aliases, mirroring the + * buffered aggregate context row. + * + * @param {StreamingGroup} group + * @param {AsyncRow} outputRow + * @returns {AsyncRow} + */ +function groupContextRow(group, outputRow) { + const firstRow = group.firstRow ?? { columns: [], cells: {} } + return { + columns: [...firstRow.columns, ...outputRow.columns], + cells: { ...firstRow.cells, ...outputRow.cells }, } - return { contextRow, outputRow, values } } /** @@ -417,32 +526,39 @@ function finalizeGroup({ selectColumns, specs, group, context }) { * * @param {object} options * @param {HashAggregateNode} options.plan - * @param {StreamingAggSpec[]} options.specs + * @param {StreamingAggPlan} options.streaming * @param {QueryResults} options.child * @param {ExecuteContext} options.context * @returns {() => AsyncGenerator} */ -export function streamingHashAggregateRows({ plan, specs, child, context }) { +export function streamingHashAggregateRows({ plan, streaming, child, context }) { + const { specs, keyRefs, needsRow } = streaming return async function* () { - const groups = await accumulateGroups({ child, groupBy: plan.groupBy, specs, context }) - const { orderBy } = plan + const groups = await accumulateGroups({ child, groupBy: plan.groupBy, specs, needsRow, context }) + const { orderBy, having } = plan - /** @type {{ outputRow: AsyncRow, contextRow: AsyncRow, values: Map, orderValues: SqlPrimitive[] }[]} */ - const outputRows = [] + // Without ORDER BY, groups finalize and yield one at a time so output + // rows are never all held at once; sorting needs the full set below. + /** @type {{ outputRow: AsyncRow, contextRow: AsyncRow, values: Map, orderValues: SqlPrimitive[] }[] | undefined} */ + const outputRows = orderBy?.length ? [] : undefined for (const group of groups.values()) { - const { contextRow, outputRow, values } = finalizeGroup({ selectColumns: plan.columns, specs, group, context }) - if (plan.having) { + const { outputRow, values } = finalizeGroup({ selectColumns: plan.columns, specs, keyRefs, group, context }) + if (having) { const passes = await evaluateExpr({ - node: substituteAggregates(plan.having, values), - row: contextRow, + node: substituteValues(having, values), + row: groupContextRow(group, outputRow), context, }) if (!passes) continue } - outputRows.push({ outputRow, contextRow, values, orderValues: [] }) + if (outputRows) { + outputRows.push({ outputRow, contextRow: groupContextRow(group, outputRow), values, orderValues: [] }) + } else { + yield outputRow + } } - if (orderBy?.length) { + if (outputRows && orderBy) { // Evaluate each sort key across all groups in concurrent chunks so // async cells and UDFs overlap for (let t = 0; t < orderBy.length; t++) { @@ -454,7 +570,7 @@ export function streamingHashAggregateRows({ plan, specs, child, context }) { } const chunk = outputRows.slice(start, start + CHUNK_SIZE) const termValues = await Promise.all(chunk.map(entry => evaluateExpr({ - node: substituteAggregates(term.expr, entry.values), + node: substituteValues(term.expr, entry.values), row: entry.contextRow, context, }))) @@ -470,10 +586,9 @@ export function streamingHashAggregateRows({ plan, specs, child, context }) { } return 0 }) - } - - for (const { outputRow } of outputRows) { - yield outputRow + for (const { outputRow } of outputRows) { + yield outputRow + } } } } @@ -484,22 +599,23 @@ export function streamingHashAggregateRows({ plan, specs, child, context }) { * * @param {object} options * @param {ScalarAggregateNode} options.plan - * @param {StreamingAggSpec[]} options.specs + * @param {StreamingAggPlan} options.streaming * @param {QueryResults} options.child * @param {ExecuteContext} options.context * @returns {() => AsyncGenerator} */ -export function streamingScalarAggregateRows({ plan, specs, child, context }) { +export function streamingScalarAggregateRows({ plan, streaming, child, context }) { + const { specs, keyRefs, needsRow } = streaming return async function* () { - const groups = await accumulateGroups({ child, groupBy: [], specs, context }) + const groups = await accumulateGroups({ child, groupBy: [], specs, needsRow, context }) /** @type {StreamingGroup} */ - const group = groups.get(true) ?? { firstRow: undefined, accumulators: specs.map(spec => newAccumulator(spec)) } + const group = groups.get(true) ?? { firstRow: undefined, keyValues: [], accumulators: specs.map(spec => newAccumulator(spec)) } - const { contextRow, outputRow, values } = finalizeGroup({ selectColumns: plan.columns, specs, group, context }) + const { outputRow, values } = finalizeGroup({ selectColumns: plan.columns, specs, keyRefs, group, context }) if (plan.having) { const passes = await evaluateExpr({ - node: substituteAggregates(plan.having, values), - row: contextRow, + node: substituteValues(plan.having, values), + row: groupContextRow(group, outputRow), context, }) if (!passes) return diff --git a/test/execute/streamingAggregate.test.js b/test/execute/streamingAggregate.test.js index 07af335..73107b4 100644 --- a/test/execute/streamingAggregate.test.js +++ b/test/execute/streamingAggregate.test.js @@ -1,7 +1,8 @@ import { describe, expect, it } from 'vitest' import { memorySource } from '../../src/backend/dataSource.js' import { executeSql } from '../../src/execute/execute.js' -import { collect } from '../../src/index.js' +import { planStreamingAggregates } from '../../src/execute/streamingAggregate.js' +import { collect, planSql } from '../../src/index.js' // 10000 rows spans multiple 4000-row accumulation chunks, so these tests pin // that group state carries across chunk boundaries in the streaming path. @@ -90,3 +91,164 @@ describe('streaming aggregates', () => { expect(result).toEqual([{ mx: 7, n: 10000 }]) }) }) + +const sales = [ + { region: 'east', product: 'apple', amount: 100, qty: 1 }, + { region: 'east', product: 'banana', amount: 200, qty: 2 }, + { region: 'west', product: 'apple', amount: 150, qty: null }, + { region: 'west', product: 'apple', amount: null, qty: 4 }, + { region: 'south', product: 'cherry', amount: 50, qty: 5 }, +] +const tables = { sales: memorySource({ data: sales }) } + +describe('streaming aggregate row retention', () => { + /** + * @param {string} query + * @returns {ReturnType} + */ + function streamingPlan(query) { + const plan = planSql({ query, tables }) + if (plan.type !== 'HashAggregate' && plan.type !== 'ScalarAggregate') { + throw new Error(`expected aggregate plan, got ${plan.type}`) + } + return planStreamingAggregates(plan) + } + + it('retains no rows when expressions only use group keys and aggregates', () => { + const streaming = streamingPlan('SELECT region, count(*) FROM sales GROUP BY region HAVING count(*) > 1 ORDER BY sum(amount)') + expect(streaming?.needsRow).toBe(false) + expect(streaming?.keyRefs.size).toBe(1) + }) + + it('retains no rows for function group keys', () => { + const streaming = streamingPlan('SELECT upper(region) AS r, count(*) FROM sales GROUP BY upper(region)') + expect(streaming?.needsRow).toBe(false) + expect(streaming?.keyRefs.size).toBe(1) + }) + + it('retains a representative row for columns outside the group keys', () => { + const streaming = streamingPlan('SELECT product, count(*) FROM sales GROUP BY region') + expect(streaming?.needsRow).toBe(true) + }) + + it('does not stream buffering aggregates like median', () => { + expect(streamingPlan('SELECT region, median(amount) FROM sales GROUP BY region')).toBeUndefined() + }) + + it('does not stream array_agg', () => { + expect(streamingPlan('SELECT array_agg(product) FROM sales')).toBeUndefined() + }) +}) + +describe('streaming aggregate results', () => { + it('computes grouped aggregates with nulls', async () => { + const result = await collect(executeSql({ + tables, + query: `SELECT region, count(*) AS c, count(amount) AS ca, sum(amount) AS s, min(amount) AS mn, max(amount) AS mx + FROM sales GROUP BY region ORDER BY region`, + })) + expect(result).toEqual([ + { region: 'east', c: 2, ca: 2, s: 300, mn: 100, mx: 200 }, + { region: 'south', c: 1, ca: 1, s: 50, mn: 50, mx: 50 }, + { region: 'west', c: 2, ca: 1, s: 150, mn: 150, mx: 150 }, + ]) + }) + + it('computes grouped count distinct', async () => { + const result = await collect(executeSql({ + tables, + query: 'SELECT region, count(DISTINCT product) AS p FROM sales GROUP BY region ORDER BY region', + })) + expect(result).toEqual([ + { region: 'east', p: 2 }, + { region: 'south', p: 1 }, + { region: 'west', p: 1 }, + ]) + }) + + it('computes multiple scalar count distinct', async () => { + const result = await collect(executeSql({ + tables, + query: 'SELECT count(DISTINCT region) AS r, count(DISTINCT product) AS p, count(DISTINCT qty) AS q FROM sales', + })) + expect(result).toEqual([{ r: 3, p: 3, q: 4 }]) + }) + + it('keeps same-column COUNT and COUNT DISTINCT separate', async () => { + const dup = memorySource({ data: [{ v: 1 }, { v: 1 }, { v: 2 }] }) + const result = await collect(executeSql({ + tables: { dup }, + query: 'SELECT count(v) AS a, count(DISTINCT v) AS b FROM dup', + })) + expect(result).toEqual([{ a: 3, b: 2 }]) + }) + + it('evaluates expressions over finalized aggregates', async () => { + const result = await collect(executeSql({ + tables, + query: 'SELECT region, count(*) * 2 AS c2, round(avg(amount), 1) AS a FROM sales GROUP BY region ORDER BY region', + })) + expect(result).toEqual([ + { region: 'east', c2: 4, a: 150 }, + { region: 'south', c2: 2, a: 50 }, + { region: 'west', c2: 4, a: 150 }, + ]) + }) + + it('groups by function expressions', async () => { + const result = await collect(executeSql({ + tables, + query: 'SELECT upper(region) AS r, count(*) AS c FROM sales GROUP BY upper(region) ORDER BY r', + })) + expect(result).toEqual([ + { r: 'EAST', c: 2 }, + { r: 'SOUTH', c: 1 }, + { r: 'WEST', c: 2 }, + ]) + }) + + it('computes grouped countif', async () => { + const result = await collect(executeSql({ + tables, + query: 'SELECT region, countif(amount > 100) AS c FROM sales GROUP BY region ORDER BY region', + })) + expect(result).toEqual([ + { region: 'east', c: 1 }, + { region: 'south', c: 0 }, + { region: 'west', c: 1 }, + ]) + }) + + it('orders groups by aggregate alias', async () => { + const result = await collect(executeSql({ + tables, + query: 'SELECT region, count(*) AS c FROM sales GROUP BY region ORDER BY c DESC, region', + })) + expect(result).toEqual([ + { region: 'east', c: 2 }, + { region: 'west', c: 2 }, + { region: 'south', c: 1 }, + ]) + }) + + it('projects columns outside the group keys from a representative row', async () => { + const result = await collect(executeSql({ + tables, + query: 'SELECT region, product, count(*) AS c FROM sales GROUP BY region, product ORDER BY region, product', + })) + expect(result).toEqual([ + { region: 'east', product: 'apple', c: 1 }, + { region: 'east', product: 'banana', c: 1 }, + { region: 'south', product: 'cherry', c: 1 }, + { region: 'west', product: 'apple', c: 2 }, + ]) + }) + + it('aggregates an empty filter result', async () => { + const result = await collect(executeSql({ + tables, + query: 'SELECT count(*) AS c, count(amount) AS ca, sum(amount) AS s, min(amount) AS m FROM sales WHERE amount > 1000', + })) + expect(result).toEqual([{ c: 0, ca: 0, s: null, m: null }]) + }) +}) From 6961daeb9a49117e5ea3aa20d8bf06fdee3888d2 Mon Sep 17 00:00:00 2001 From: Kenny Daniel Date: Fri, 3 Jul 2026 15:41:32 -0700 Subject: [PATCH 3/7] Restrict group key substitution to exact matches and handle BigInt literals The structural signature used JSON.stringify, which throws on BigInt literal values, so any aggregate query containing a bigint literal failed during streaming planning. The signature replacer now encodes bigints safely. Group key matching also treated bare and qualified identifiers with the same name as interchangeable. In a join, a bare identifier can resolve to a different column than a qualified group key, so substituting the key value changed results. Matching now requires exact structural equality; mixed qualification retains the representative row and evaluates the identifier against it, preserving buffered semantics. --- src/execute/streamingAggregate.js | 22 ++++++------- test/execute/streamingAggregate.test.js | 44 +++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 12 deletions(-) diff --git a/src/execute/streamingAggregate.js b/src/execute/streamingAggregate.js index 7a0aeda..b518d8d 100644 --- a/src/execute/streamingAggregate.js +++ b/src/execute/streamingAggregate.js @@ -69,9 +69,12 @@ const STREAMABLE_FUNCS = new Set(['COUNT', 'COUNTIF', 'SUM', 'AVG', 'MIN', 'MAX' * @returns {string} */ function exprSig(node) { - return JSON.stringify(node, (key, value) => - key === 'positionStart' || key === 'positionEnd' ? undefined : value - ) + return JSON.stringify(node, (key, value) => { + if (key === 'positionStart' || key === 'positionEnd') return undefined + // JSON.stringify throws on BigInt literal values; wrap so 1n !== '1n' + if (typeof value === 'bigint') return { bigint: value.toString() } + return value + }) } /** @@ -93,8 +96,10 @@ export function planStreamingAggregates({ columns, having, orderBy, groupBy }) { let needsRow = false /** - * Index of the group key the expression refers to, or -1. An unqualified - * column reference matches a qualified group key and vice versa. + * Index of the group key the expression structurally matches, or -1. + * Only exact matches qualify: a bare identifier and a qualified group key + * (or vice versa) can resolve to different columns in a join, so mixed + * qualification falls back to evaluating against the representative row. * * @param {ExprNode} node * @returns {number} @@ -104,13 +109,6 @@ export function planStreamingAggregates({ columns, having, orderBy, groupBy }) { for (let i = 0; i < groupExprs.length; i++) { if (groupSigs[i] === signature) return i } - if (node.type === 'identifier') { - for (let i = 0; i < groupExprs.length; i++) { - const expr = groupExprs[i] - if (expr.type === 'identifier' && expr.name === node.name && - (!expr.prefix || !node.prefix || expr.prefix === node.prefix)) return i - } - } return -1 } diff --git a/test/execute/streamingAggregate.test.js b/test/execute/streamingAggregate.test.js index 73107b4..74c30ec 100644 --- a/test/execute/streamingAggregate.test.js +++ b/test/execute/streamingAggregate.test.js @@ -131,6 +131,20 @@ describe('streaming aggregate row retention', () => { expect(streaming?.needsRow).toBe(true) }) + it('retains a representative row when qualification differs from the group key', () => { + // A bare identifier can resolve to a different column than a qualified + // group key in a join, so it is not treated as a group key reference + const streaming = streamingPlan('SELECT region, count(*) FROM sales GROUP BY sales.region') + expect(streaming?.needsRow).toBe(true) + expect(streaming?.keyRefs.size).toBe(0) + }) + + it('plans streaming aggregates for queries with bigint literals', () => { + const streaming = streamingPlan('SELECT count(*) + 0n AS c, 2n AS two FROM sales') + expect(streaming?.needsRow).toBe(false) + expect(streaming?.specs.length).toBe(1) + }) + it('does not stream buffering aggregates like median', () => { expect(streamingPlan('SELECT region, median(amount) FROM sales GROUP BY region')).toBeUndefined() }) @@ -244,6 +258,36 @@ describe('streaming aggregate results', () => { ]) }) + it('resolves a bare identifier against the joined row, not a qualified group key', async () => { + const result = await collect(executeSql({ + tables: { a: [{ id: 1 }, { id: 2 }], b: [{ id: 10 }] }, + query: 'SELECT id, count(*) AS c FROM a JOIN b ON TRUE GROUP BY b.id, a.id ORDER BY id', + })) + expect(result).toEqual([ + { id: 1, c: 1 }, + { id: 2, c: 1 }, + ]) + }) + + it('evaluates bigint literals alongside streamed aggregates', async () => { + const result = await collect(executeSql({ + tables, + query: 'SELECT count(*) AS c, 2n AS two FROM sales', + })) + expect(result).toEqual([{ c: 5, two: 2n }]) + }) + + it('compares bigint literals in HAVING', async () => { + const result = await collect(executeSql({ + tables, + query: 'SELECT region, count(*) AS c FROM sales GROUP BY region HAVING count(*) > 1n ORDER BY region', + })) + expect(result).toEqual([ + { region: 'east', c: 2 }, + { region: 'west', c: 2 }, + ]) + }) + it('aggregates an empty filter result', async () => { const result = await collect(executeSql({ tables, From 3da7797e9122b0ee567261ba09a22006271437f1 Mon Sep 17 00:00:00 2001 From: Kenny Daniel Date: Fri, 3 Jul 2026 16:06:51 -0700 Subject: [PATCH 4/7] Preserve short-circuit and tie-aware evaluation in streaming aggregates Aggregates under short-circuited positions (AND/OR right sides, CASE branches past the first WHEN condition) may never be evaluated by the buffered path, so accumulating them eagerly could evaluate expressions the query never asks for. The streaming planner now falls back to buffered aggregation when an aggregate appears in such a position. Grouped ORDER BY previously evaluated every sort term for every group before sorting. Streaming results now go through sortEntriesByTerms, which evaluates later terms only within ties on earlier terms; sort entries carry pre-substituted per-group expressions so finalized aggregate values still replace aggregate calls. --- src/execute/sort.js | 9 ++- src/execute/streamingAggregate.js | 87 ++++++++++++------------- test/execute/streamingAggregate.test.js | 52 +++++++++++++++ 3 files changed, 99 insertions(+), 49 deletions(-) diff --git a/src/execute/sort.js b/src/execute/sort.js index d4c4ff3..14b6d30 100644 --- a/src/execute/sort.js +++ b/src/execute/sort.js @@ -4,16 +4,21 @@ import { executePlan } from './execute.js' import { compareForTerm } from './utils.js' /** - * @import { AsyncRow, ExecuteContext, OrderByItem, QueryResults, SqlPrimitive } from '../types.js' + * @import { AsyncRow, ExecuteContext, ExprNode, OrderByItem, QueryResults, SqlPrimitive } from '../types.js' * @import { SortNode } from '../plan/types.js' */ const MAX_CHUNK = 256 /** + * When exprs is set, exprs[i] is evaluated for ORDER BY term i instead of the + * term's own expression, so callers can pre-substitute per-entry values (for + * example, finalized aggregates) while keeping tie-aware term evaluation. + * * @typedef {{ * row: AsyncRow, * rows?: AsyncRow[], + * exprs?: ExprNode[], * }} SortEntry */ @@ -63,7 +68,7 @@ export async function sortEntriesByTerms({ entries, orderBy, context, cacheValue const chunk = missing.slice(start, start + chunkSize) const values = await Promise.all(chunk.map(idx => evaluateExpr({ - node: term.expr, + node: entries[idx].exprs?.[orderByIdx] ?? term.expr, row: entries[idx].row, rows: entries[idx].rows, context, diff --git a/src/execute/streamingAggregate.js b/src/execute/streamingAggregate.js index b518d8d..b14231a 100644 --- a/src/execute/streamingAggregate.js +++ b/src/execute/streamingAggregate.js @@ -1,7 +1,8 @@ import { derivedAlias } from '../expression/alias.js' import { evaluateAll, evaluateExpr } from '../expression/evaluate.js' import { isAggregateFunc } from '../validation/functions.js' -import { compareForTerm, keyify } from './utils.js' +import { sortEntriesByTerms } from './sort.js' +import { keyify } from './utils.js' import { yieldToEventLoop } from './yield.js' /** @@ -115,12 +116,17 @@ export function planStreamingAggregates({ columns, having, orderBy, groupBy }) { /** * Walks an expression collecting streamable aggregate calls and group key * references. Returns false if the expression cannot be evaluated from - * precomputed values plus a representative row. + * precomputed values plus a representative row. `lazy` marks positions the + * evaluator can skip (short-circuited AND/OR right sides and CASE + * branches): an aggregate there may never be evaluated by the buffered + * path, so accumulating it eagerly could evaluate expressions the query + * never asks for, and the query falls back to buffered aggregation. * * @param {ExprNode} node + * @param {boolean} lazy * @returns {boolean} */ - function walk(node) { + function walk(node, lazy) { const keyIndex = matchGroupKey(node) if (keyIndex >= 0) { keyRefs.set(node, keyIndex) @@ -136,22 +142,29 @@ export function planStreamingAggregates({ columns, having, orderBy, groupBy }) { needsRow = true return true case 'unary': - return walk(node.argument) + return walk(node.argument, lazy) case 'binary': - return walk(node.left) && walk(node.right) + if (node.op === 'AND' || node.op === 'OR') { + // the right side is skipped when the left side short-circuits + return walk(node.left, lazy) && walk(node.right, true) + } + return walk(node.left, lazy) && walk(node.right, lazy) case 'cast': - return walk(node.expr) + return walk(node.expr, lazy) case 'case': - return (!node.caseExpr || walk(node.caseExpr)) && - node.whenClauses.every(w => walk(w.condition) && walk(w.result)) && - (!node.elseResult || walk(node.elseResult)) + // only the first WHEN condition is always evaluated; later + // conditions, results, and ELSE run only when reached + return (!node.caseExpr || walk(node.caseExpr, lazy)) && + node.whenClauses.every((w, i) => walk(w.condition, lazy || i > 0) && walk(w.result, true)) && + (!node.elseResult || walk(node.elseResult, true)) case 'in valuelist': - return walk(node.expr) && node.values.every(walk) + return walk(node.expr, lazy) && node.values.every(v => walk(v, lazy)) case 'function': { const funcName = node.funcName.toUpperCase() if (!isAggregateFunc(funcName)) { - return node.args.every(walk) + return node.args.every(arg => walk(arg, lazy)) } + if (lazy) return false if (!STREAMABLE_FUNCS.has(funcName)) return false const star = node.args[0]?.type === 'star' if (!star && !node.args.every(arg => isScalarExpr(arg))) return false @@ -172,11 +185,11 @@ export function planStreamingAggregates({ columns, having, orderBy, groupBy }) { needsRow = true continue } - if (!walk(col.expr)) return + if (!walk(col.expr, false)) return } - if (having && !walk(having)) return + if (having && !walk(having, false)) return for (const term of orderBy ?? []) { - if (!walk(term.expr)) return + if (!walk(term.expr, false)) return } return { specs, keyRefs, needsRow } } @@ -537,8 +550,8 @@ export function streamingHashAggregateRows({ plan, streaming, child, context }) // Without ORDER BY, groups finalize and yield one at a time so output // rows are never all held at once; sorting needs the full set below. - /** @type {{ outputRow: AsyncRow, contextRow: AsyncRow, values: Map, orderValues: SqlPrimitive[] }[] | undefined} */ - const outputRows = orderBy?.length ? [] : undefined + /** @type {{ row: AsyncRow, exprs: ExprNode[], outputRow: AsyncRow }[] | undefined} */ + const entries = orderBy?.length ? [] : undefined for (const group of groups.values()) { const { outputRow, values } = finalizeGroup({ selectColumns: plan.columns, specs, keyRefs, group, context }) if (having) { @@ -549,42 +562,22 @@ export function streamingHashAggregateRows({ plan, streaming, child, context }) }) if (!passes) continue } - if (outputRows) { - outputRows.push({ outputRow, contextRow: groupContextRow(group, outputRow), values, orderValues: [] }) + if (entries && orderBy) { + entries.push({ + row: groupContextRow(group, outputRow), + exprs: orderBy.map(term => substituteValues(term.expr, values)), + outputRow, + }) } else { yield outputRow } } - if (outputRows && orderBy) { - // Evaluate each sort key across all groups in concurrent chunks so - // async cells and UDFs overlap - for (let t = 0; t < orderBy.length; t++) { - const term = orderBy[t] - for (let start = 0; start < outputRows.length; start += CHUNK_SIZE) { - if (start > 0) { - await yieldToEventLoop() - context.signal?.throwIfAborted() - } - const chunk = outputRows.slice(start, start + CHUNK_SIZE) - const termValues = await Promise.all(chunk.map(entry => evaluateExpr({ - node: substituteValues(term.expr, entry.values), - row: entry.contextRow, - context, - }))) - for (let j = 0; j < chunk.length; j++) { - chunk[j].orderValues[t] = termValues[j] - } - } - } - outputRows.sort((a, b) => { - for (let i = 0; i < orderBy.length; i++) { - const cmp = compareForTerm(a.orderValues[i], b.orderValues[i], orderBy[i]) - if (cmp) return cmp - } - return 0 - }) - for (const { outputRow } of outputRows) { + if (entries && orderBy) { + // The shared sorter evaluates later ORDER BY terms only within ties + // on earlier terms, so expensive sort keys are skipped when possible + const sorted = await sortEntriesByTerms({ entries, orderBy, context }) + for (const { outputRow } of sorted) { yield outputRow } } diff --git a/test/execute/streamingAggregate.test.js b/test/execute/streamingAggregate.test.js index 74c30ec..99b0abd 100644 --- a/test/execute/streamingAggregate.test.js +++ b/test/execute/streamingAggregate.test.js @@ -4,6 +4,20 @@ import { executeSql } from '../../src/execute/execute.js' import { planStreamingAggregates } from '../../src/execute/streamingAggregate.js' import { collect, planSql } from '../../src/index.js' +/** + * @import { UserDefinedFunction } from '../../src/index.js' + */ + +/** @type {Record} */ +const boomFunctions = { + BOOM: { + apply() { + throw new Error('should not be called') + }, + arguments: { min: 1, max: 1 }, + }, +} + // 10000 rows spans multiple 4000-row accumulation chunks, so these tests pin // that group state carries across chunk boundaries in the streaming path. const N = 10000 @@ -149,6 +163,21 @@ describe('streaming aggregate row retention', () => { expect(streamingPlan('SELECT region, median(amount) FROM sales GROUP BY region')).toBeUndefined() }) + it('does not stream aggregates in short-circuited CASE branches', () => { + // The buffered evaluator never evaluates the ELSE branch when the first + // WHEN condition matches, so its aggregate must not accumulate eagerly + expect(streamingPlan('SELECT CASE WHEN count(*) > 0 THEN 1 ELSE sum(amount) END AS c FROM sales')).toBeUndefined() + }) + + it('does not stream aggregates on the short-circuited side of AND', () => { + expect(streamingPlan('SELECT region FROM sales GROUP BY region HAVING count(*) > 1 AND sum(amount) > 100')).toBeUndefined() + }) + + it('streams aggregates in the first WHEN condition', () => { + const streaming = streamingPlan('SELECT CASE WHEN count(*) > 0 THEN 1 ELSE 2 END AS c FROM sales') + expect(streaming?.specs.length).toBe(1) + }) + it('does not stream array_agg', () => { expect(streamingPlan('SELECT array_agg(product) FROM sales')).toBeUndefined() }) @@ -258,6 +287,29 @@ describe('streaming aggregate results', () => { ]) }) + it('never evaluates aggregates in unreachable CASE branches', async () => { + const result = await collect(executeSql({ + tables, + functions: boomFunctions, + query: 'SELECT CASE WHEN count(*) > 0 THEN 1 ELSE min(BOOM(amount)) END AS r FROM sales', + })) + expect(result).toEqual([{ r: 1 }]) + }) + + it('evaluates later ORDER BY terms only for groups that tie on earlier terms', async () => { + const t = memorySource({ data: [{ g: 'a' }, { g: 'b' }, { g: 'b' }, { g: 'c' }, { g: 'c' }, { g: 'c' }] }) + const result = await collect(executeSql({ + tables: { t }, + functions: boomFunctions, + query: 'SELECT g, count(*) AS c FROM t GROUP BY g ORDER BY count(*) DESC, BOOM(g)', + })) + expect(result).toEqual([ + { g: 'c', c: 3 }, + { g: 'b', c: 2 }, + { g: 'a', c: 1 }, + ]) + }) + it('resolves a bare identifier against the joined row, not a qualified group key', async () => { const result = await collect(executeSql({ tables: { a: [{ id: 1 }, { id: 2 }], b: [{ id: 10 }] }, From 4ad94153e8e1f3f5adda10d38e985602669c5ef4 Mon Sep 17 00:00:00 2001 From: Kenny Daniel Date: Fri, 3 Jul 2026 16:48:44 -0700 Subject: [PATCH 5/7] Preserve buffered semantics for pruned columns, lazy positions, and aborts Streaming aggregation evaluates aggregate arguments eagerly, but projection pushdown prunes columns whose output cells are never read, so a subquery whose aggregate output goes unread crashed with ColumnNotFoundError. planStreamingAggregates now receives the child's columns and falls back to buffered aggregation when an aggregate references a column the child does not produce. IN value lists short-circuit once an earlier value matches, and the sorter evaluates later ORDER BY terms only within ties, so aggregates in those positions are now treated as lazy and fall back to buffered aggregation instead of accumulating eagerly. Aborting mid-accumulation now ends the row stream silently, matching the buffered aggregate path and the other executors; group keys over missing columns stay undefined instead of being coerced to null. COUNT/COUNTIF/SUM/AVG/MIN/MAX fold semantics move to a shared accumulator module used by both the streaming path and the scanColumn fast path, so the copies cannot drift. The streaming plan docs no longer overstate the memory bound for COUNT(DISTINCT ...). --- src/execute/accumulator.js | 89 +++++++++++++ src/execute/aggregates.js | 69 ++-------- src/execute/streamingAggregate.js | 159 ++++++++++-------------- src/plan/columns.js | 2 +- test/execute/streamingAggregate.test.js | 80 ++++++++++++ 5 files changed, 246 insertions(+), 153 deletions(-) create mode 100644 src/execute/accumulator.js diff --git a/src/execute/accumulator.js b/src/execute/accumulator.js new file mode 100644 index 0000000..d013a3d --- /dev/null +++ b/src/execute/accumulator.js @@ -0,0 +1,89 @@ +import { keyify } from './utils.js' + +/** + * @import { SqlPrimitive } from '../types.js' + */ + +/** + * Incremental state for one streamable aggregate. Shared by the streaming + * aggregate executor and the scanColumn fast path so COUNT/COUNTIF/SUM/AVG/ + * MIN/MAX fold semantics live in one place. + * + * @typedef {{ + * count: number, + * sum: number, + * min: SqlPrimitive, + * max: SqlPrimitive, + * seen: Set | null, + * }} Accumulator + */ + +/** + * @param {string} funcName + * @param {boolean} [distinct] + * @returns {Accumulator} + */ +export function newAccumulator(funcName, distinct) { + return { + count: 0, + sum: 0, + min: null, + max: null, + seen: funcName === 'COUNT' && distinct ? new Set() : null, + } +} + +/** + * Folds one value into an accumulator, matching the buffered semantics in + * evaluate.js: COUNT counts non-null, COUNTIF counts truthy, MIN/MAX compare + * raw values, SUM/AVG only accumulate finite numbers. + * + * @param {string} funcName + * @param {Accumulator} acc + * @param {SqlPrimitive} value + */ +export function updateAccumulator(funcName, acc, value) { + switch (funcName) { + case 'COUNT': + if (value == null) break + if (acc.seen) acc.seen.add(keyify(value)) + else acc.count++ + break + case 'COUNTIF': + if (value) acc.count++ + break + case 'MIN': + if (value != null && (acc.min === null || value < acc.min)) acc.min = value + break + case 'MAX': + if (value != null && (acc.max === null || value > acc.max)) acc.max = value + break + default: { // SUM, AVG + if (value == null) break + const num = Number(value) + if (Number.isFinite(num)) { + acc.sum += num + acc.count++ + } + } + } +} + +/** + * Reduces an accumulator to its final aggregate value. + * + * @param {string} funcName + * @param {Accumulator} acc + * @returns {SqlPrimitive} + */ +export function finalizeAccumulator(funcName, acc) { + switch (funcName) { + case 'COUNT': return acc.seen ? acc.seen.size : acc.count + case 'COUNTIF': return acc.count + case 'SUM': return acc.count === 0 ? null : acc.sum + case 'AVG': return acc.count === 0 ? null : acc.sum / acc.count + case 'MIN': return acc.min + case 'MAX': return acc.max + default: return null + } +} diff --git a/src/execute/aggregates.js b/src/execute/aggregates.js index 41a2a7e..82ae970 100644 --- a/src/execute/aggregates.js +++ b/src/execute/aggregates.js @@ -1,5 +1,6 @@ import { derivedAlias } from '../expression/alias.js' import { evaluateExpr } from '../expression/evaluate.js' +import { finalizeAccumulator, newAccumulator, updateAccumulator } from './accumulator.js' import { executePlan, selectColumnNames } from './execute.js' import { sortEntriesByTerms } from './sort.js' import { planStreamingAggregates, streamingHashAggregateRows, streamingScalarAggregateRows } from './streamingAggregate.js' @@ -81,7 +82,7 @@ function aggregateContextRow(group, aggregateRow) { */ export function executeHashAggregate(plan, context) { const child = executePlan({ plan: plan.child, context }) - const streaming = planStreamingAggregates(plan) + const streaming = planStreamingAggregates(plan, child.columns) if (streaming) { return { columns: selectColumnNames(plan.columns, child.columns), @@ -205,7 +206,7 @@ export function executeScalarAggregate(plan, context) { } const child = executePlan({ plan: plan.child, context }) - const streaming = planStreamingAggregates(plan) + const streaming = planStreamingAggregates(plan, child.columns) if (streaming) { return { columns: selectColumnNames(plan.columns, child.columns), @@ -362,17 +363,6 @@ function extractColumnAggSpec({ expr, alias }) { } } -/** - * @typedef {{ - * spec: ColumnAggSpec, - * count: number, - * sum: number, - * min: SqlPrimitive, - * max: SqlPrimitive, - * seen: Set | null, - * }} AggAccumulator - */ - /** * Scans a column once and computes every aggregate over it in a single pass. * All specs share the one scanColumn walk, so MIN(x)/MAX(x)/AVG(x) decode x once. @@ -389,41 +379,15 @@ async function scanColumnGroup({ table, specs, limit, offset, signal }) { const { column } = specs[0] const values = table.scanColumn({ column, limit, offset, signal }) - /** @type {AggAccumulator[]} */ - const accs = specs.map(spec => /** @type {AggAccumulator} */ ({ - spec, - count: 0, - sum: 0, - min: null, - max: null, - seen: spec.funcName === 'COUNT' && spec.distinct ? new Set() : null, - })) + const accs = specs.map(spec => ({ spec, acc: newAccumulator(spec.funcName, spec.distinct) })) for await (const chunk of values) { signal?.throwIfAborted() for (let i = 0; i < chunk.length; i++) { const v = chunk[i] if (v == null) continue - for (const acc of accs) { - switch (acc.spec.funcName) { - case 'COUNT': - if (acc.seen) acc.seen.add(keyify(v)) - else acc.count++ - break - case 'MIN': - if (acc.min === null || v < acc.min) acc.min = v - break - case 'MAX': - if (acc.max === null || v > acc.max) acc.max = v - break - default: { // SUM, AVG - const num = Number(v) - if (Number.isFinite(num)) { - acc.sum += num - acc.count++ - } - } - } + for (const { spec, acc } of accs) { + updateAccumulator(spec.funcName, acc, v) } } } @@ -431,25 +395,8 @@ async function scanColumnGroup({ table, specs, limit, offset, signal }) { /** @type {Map} */ const result = new Map() - for (const acc of accs) { - result.set(acc.spec.alias, finalizeAgg(acc)) + for (const { spec, acc } of accs) { + result.set(spec.alias, finalizeAccumulator(spec.funcName, acc)) } return result } - -/** - * Reduces one accumulator to its final aggregate value. - * - * @param {AggAccumulator} acc - * @returns {SqlPrimitive} - */ -function finalizeAgg(acc) { - switch (acc.spec.funcName) { - case 'COUNT': return acc.seen ? acc.seen.size : acc.count - case 'SUM': return acc.count === 0 ? null : acc.sum - case 'AVG': return acc.count === 0 ? null : acc.sum / acc.count - case 'MIN': return acc.min - case 'MAX': return acc.max - default: return null - } -} diff --git a/src/execute/streamingAggregate.js b/src/execute/streamingAggregate.js index b14231a..33ee38d 100644 --- a/src/execute/streamingAggregate.js +++ b/src/execute/streamingAggregate.js @@ -1,13 +1,16 @@ import { derivedAlias } from '../expression/alias.js' import { evaluateAll, evaluateExpr } from '../expression/evaluate.js' +import { collectColumnsFromExpr } from '../plan/columns.js' import { isAggregateFunc } from '../validation/functions.js' +import { finalizeAccumulator, newAccumulator, updateAccumulator } from './accumulator.js' import { sortEntriesByTerms } from './sort.js' import { keyify } from './utils.js' import { yieldToEventLoop } from './yield.js' /** - * @import { AsyncCells, AsyncRow, ExecuteContext, ExprNode, FunctionNode, QueryResults, SelectColumn, SqlPrimitive } from '../types.js' + * @import { AsyncCells, AsyncRow, ExecuteContext, ExprNode, FunctionNode, IdentifierNode, QueryResults, SelectColumn, SqlPrimitive } from '../types.js' * @import { HashAggregateNode, ScalarAggregateNode } from '../plan/types.js' + * @import { Accumulator } from './accumulator.js' */ // Accumulate rows in chunks of this size so aborts can fire and async cells overlap @@ -29,21 +32,11 @@ const STREAMABLE_FUNCS = new Set(['COUNT', 'COUNTIF', 'SUM', 'AVG', 'MIN', 'MAX' * }} StreamingAggSpec */ -/** - * @typedef {{ - * count: number, - * sum: number, - * min: SqlPrimitive, - * max: SqlPrimitive, - * seen: Set | null, - * }} StreamingAccumulator - */ - /** * @typedef {{ * firstRow: AsyncRow | undefined, * keyValues: SqlPrimitive[], - * accumulators: StreamingAccumulator[], + * accumulators: Accumulator[], * }} StreamingGroup */ @@ -52,7 +45,8 @@ const STREAMABLE_FUNCS = new Set(['COUNT', 'COUNTIF', 'SUM', 'AVG', 'MIN', 'MAX' * accumulate, which expression nodes are group key references (substituted * from the group's key values), and whether any expression still needs a * representative row from the group. When needsRow is false, no input rows - * are retained at all, so memory is bounded by the number of groups even for + * are retained, so memory is bounded by the number of groups — plus, for + * COUNT(DISTINCT ...), each group's set of distinct values — even for * high-cardinality GROUP BY. * * @typedef {{ @@ -83,11 +77,15 @@ function exprSig(node) { * computed incrementally, without buffering the group's rows. Returns * undefined when any expression needs a buffered group: an aggregate outside * STREAMABLE_FUNCS, an aggregate over a non-scalar argument, or a subquery. + * Also returns undefined when an aggregate references a column the child + * does not produce: projection pushdown prunes columns whose output cells + * are never read, and only the buffered path defers evaluation of those cells. * * @param {Pick & Partial>} plan + * @param {string[]} [childColumns] - columns produced by the child plan * @returns {StreamingAggPlan | undefined} */ -export function planStreamingAggregates({ columns, having, orderBy, groupBy }) { +export function planStreamingAggregates({ columns, having, orderBy, groupBy }, childColumns) { const groupExprs = groupBy ?? [] const groupSigs = groupExprs.map(exprSig) /** @type {StreamingAggSpec[]} */ @@ -158,7 +156,8 @@ export function planStreamingAggregates({ columns, having, orderBy, groupBy }) { node.whenClauses.every((w, i) => walk(w.condition, lazy || i > 0) && walk(w.result, true)) && (!node.elseResult || walk(node.elseResult, true)) case 'in valuelist': - return walk(node.expr, lazy) && node.values.every(v => walk(v, lazy)) + // values after the first are skipped once an earlier value matches + return walk(node.expr, lazy) && node.values.every((v, i) => walk(v, lazy || i > 0)) case 'function': { const funcName = node.funcName.toUpperCase() if (!isAggregateFunc(funcName)) { @@ -188,12 +187,43 @@ export function planStreamingAggregates({ columns, having, orderBy, groupBy }) { if (!walk(col.expr, false)) return } if (having && !walk(having, false)) return - for (const term of orderBy ?? []) { - if (!walk(term.expr, false)) return + const orderTerms = orderBy ?? [] + for (let i = 0; i < orderTerms.length; i++) { + // the sorter evaluates later terms only to break ties on earlier terms + if (!walk(orderTerms[i].expr, i > 0)) return } + if (childColumns && !specsResolvable(specs, childColumns)) return return { specs, keyRefs, needsRow } } +/** + * Reports whether every identifier the streaming path evaluates eagerly + * (aggregate arguments and FILTER conditions) can resolve against the + * child's output columns, using the same exact-then-suffix matching as + * identifier evaluation. Anything unresolvable means projection pushdown + * pruned the column, so the query must buffer instead. + * + * @param {StreamingAggSpec[]} specs + * @param {string[]} childColumns + * @returns {boolean} + */ +function specsResolvable(specs, childColumns) { + /** @type {IdentifierNode[]} */ + const identifiers = [] + for (const spec of specs) { + collectColumnsFromExpr(spec.node, identifiers) + } + return identifiers.every(({ prefix, name }) => { + if (childColumns.includes(prefix ? `${prefix}.${name}` : name)) return true + if (!prefix) return childColumns.some(col => col.endsWith('.' + name)) + // a qualified name may also resolve as struct access on a base column, + // or fall back to the bare column part + return childColumns.includes(prefix) || + childColumns.some(col => col.endsWith('.' + prefix)) || + childColumns.includes(name) + }) +} + /** * Reports whether an expression is a plain scalar over row values: no * aggregates and no subqueries, so it can be evaluated per input row. @@ -239,9 +269,13 @@ function isScalarExpr(node) { */ function substituteValues(node, values) { if (values.has(node)) { + // group keys over missing columns evaluate to undefined; preserve it + // so streaming output matches the buffered path's evaluation result + // eslint-disable-next-line no-extra-parens + const value = /** @type {SqlPrimitive} */ (values.get(node)) return { type: 'literal', - value: values.get(node) ?? null, + value, positionStart: node.positionStart, positionEnd: node.positionEnd, } @@ -284,71 +318,6 @@ function substituteValues(node, values) { } } -/** - * @param {StreamingAggSpec} spec - * @returns {StreamingAccumulator} - */ -function newAccumulator(spec) { - return { - count: 0, - sum: 0, - min: null, - max: null, - seen: spec.funcName === 'COUNT' && spec.node.distinct ? new Set() : null, - } -} - -/** - * Folds one value into an accumulator, matching the buffered semantics in - * evaluate.js: COUNT counts non-null, COUNTIF counts truthy, MIN/MAX compare - * raw values, SUM/AVG only accumulate finite numbers. - * - * @param {StreamingAggSpec} spec - * @param {StreamingAccumulator} acc - * @param {SqlPrimitive} value - */ -function updateAccumulator(spec, acc, value) { - switch (spec.funcName) { - case 'COUNT': - if (spec.star) acc.count++ - else if (value != null) { - if (acc.seen) acc.seen.add(keyify(value)) - else acc.count++ - } - break - case 'COUNTIF': - if (value) acc.count++ - break - default: { // SUM, AVG, MIN, MAX - if (value == null) break - if (acc.min === null || value < acc.min) acc.min = value - if (acc.max === null || value > acc.max) acc.max = value - const num = Number(value) - if (Number.isFinite(num)) { - acc.sum += num - acc.count++ - } - } - } -} - -/** - * @param {StreamingAggSpec} spec - * @param {StreamingAccumulator} acc - * @returns {SqlPrimitive} - */ -function finalizeAccumulator(spec, acc) { - switch (spec.funcName) { - case 'COUNT': return acc.seen ? acc.seen.size : acc.count - case 'COUNTIF': return acc.count - case 'SUM': return acc.count === 0 ? null : acc.sum - case 'AVG': return acc.count === 0 ? null : acc.sum / acc.count - case 'MIN': return acc.min - case 'MAX': return acc.max - default: return null - } -} - /** * Folds one chunk of rows into the group accumulators. Group keys, FILTER * conditions, and aggregate arguments are each evaluated across the whole @@ -413,22 +382,28 @@ async function accumulateChunk({ chunk, groupBy, specs, groups, needsRow, contex group = { firstRow: needsRow ? chunk[j] : undefined, keyValues: keyColumns ? keyColumns.map(c => c[j]) : [], - accumulators: specs.map(spec => newAccumulator(spec)), + accumulators: specs.map(spec => newAccumulator(spec.funcName, spec.node.distinct)), } groups.set(key, group) } for (let s = 0; s < specs.length; s++) { const filter = filters[s] if (filter && !filter[j]) continue - const arg = args[s] - updateAccumulator(specs[s], group.accumulators[s], arg ? arg[j] : null) + const spec = specs[s] + if (spec.star && spec.funcName === 'COUNT') { + group.accumulators[s].count++ + } else { + const arg = args[s] + updateAccumulator(spec.funcName, group.accumulators[s], arg ? arg[j] : null) + } } } } /** * Consumes the child rows into per-group accumulators, holding at most one - * chunk of rows at a time. + * chunk of rows at a time. Returns undefined when aborted mid-collection so + * callers end the row stream silently, matching the buffered path. * * @param {object} options * @param {QueryResults} options.child @@ -436,7 +411,7 @@ async function accumulateChunk({ chunk, groupBy, specs, groups, needsRow, contex * @param {StreamingAggSpec[]} options.specs * @param {boolean} options.needsRow * @param {ExecuteContext} options.context - * @returns {Promise>} + * @returns {Promise | undefined>} */ async function accumulateGroups({ child, groupBy, specs, needsRow, context }) { /** @type {Map} */ @@ -449,7 +424,7 @@ async function accumulateGroups({ child, groupBy, specs, needsRow, context }) { await accumulateChunk({ chunk, groupBy, specs, groups, needsRow, context }) chunk = [] await yieldToEventLoop() - context.signal?.throwIfAborted() + if (context.signal?.aborted) return } } if (chunk.length) { @@ -479,7 +454,7 @@ function finalizeGroup({ selectColumns, specs, keyRefs, group, context }) { /** @type {Map} */ const values = new Map() for (let s = 0; s < specs.length; s++) { - values.set(specs[s].node, finalizeAccumulator(specs[s], group.accumulators[s])) + values.set(specs[s].node, finalizeAccumulator(specs[s].funcName, group.accumulators[s])) } for (const [node, keyIndex] of keyRefs) { values.set(node, group.keyValues[keyIndex]) @@ -546,6 +521,7 @@ export function streamingHashAggregateRows({ plan, streaming, child, context }) const { specs, keyRefs, needsRow } = streaming return async function* () { const groups = await accumulateGroups({ child, groupBy: plan.groupBy, specs, needsRow, context }) + if (!groups) return const { orderBy, having } = plan // Without ORDER BY, groups finalize and yield one at a time so output @@ -599,8 +575,9 @@ export function streamingScalarAggregateRows({ plan, streaming, child, context } const { specs, keyRefs, needsRow } = streaming return async function* () { const groups = await accumulateGroups({ child, groupBy: [], specs, needsRow, context }) + if (!groups) return /** @type {StreamingGroup} */ - const group = groups.get(true) ?? { firstRow: undefined, keyValues: [], accumulators: specs.map(spec => newAccumulator(spec)) } + const group = groups.get(true) ?? { firstRow: undefined, keyValues: [], accumulators: specs.map(spec => newAccumulator(spec.funcName, spec.node.distinct)) } const { outputRow, values } = finalizeGroup({ selectColumns: plan.columns, specs, keyRefs, group, context }) if (plan.having) { diff --git a/src/plan/columns.js b/src/plan/columns.js index 32e18cc..cd4411d 100644 --- a/src/plan/columns.js +++ b/src/plan/columns.js @@ -226,7 +226,7 @@ export function extractColumns({ select, parentColumns }) { * @param {IdentifierNode[]} columns * @param {Set} [aliases] - aliases to exclude from columns */ -function collectColumnsFromExpr(expr, columns, aliases) { +export function collectColumnsFromExpr(expr, columns, aliases) { if (!expr) return if (expr.type === 'identifier') { if (expr.prefix || !aliases?.has(expr.name)) { diff --git a/test/execute/streamingAggregate.test.js b/test/execute/streamingAggregate.test.js index 99b0abd..930d114 100644 --- a/test/execute/streamingAggregate.test.js +++ b/test/execute/streamingAggregate.test.js @@ -104,6 +104,37 @@ describe('streaming aggregates', () => { })) expect(result).toEqual([{ mx: 7, n: 10000 }]) }) + + it('falls back to buffered aggregation when pushdown prunes an aggregate argument', async () => { + // The outer query never reads s, so projection pushdown prunes v from the + // scan; only the buffered path defers sum(v) to the never-read output cell + const result = await collect(executeSql({ + tables: { t: [{ g: 'a', v: 1 }, { g: 'b', v: 2 }] }, + query: 'SELECT g2 FROM (SELECT g AS g2, sum(v) AS s FROM t GROUP BY g) q ORDER BY g2', + })) + expect(result).toEqual([{ g2: 'a' }, { g2: 'b' }]) + }) + + it('ends the stream silently when aborted during accumulation', async () => { + const controller = new AbortController() + /** @type {Record} */ + const functions = { + ABORT_NOW: { + apply(v) { + controller.abort() + return v + }, + arguments: { min: 1, max: 1 }, + }, + } + const result = await collect(executeSql({ + tables: { big }, + functions, + query: 'SELECT g, count(ABORT_NOW(v)) AS c FROM big GROUP BY g', + signal: controller.signal, + })) + expect(result).toEqual([]) + }) }) const sales = [ @@ -173,6 +204,23 @@ describe('streaming aggregate row retention', () => { expect(streamingPlan('SELECT region FROM sales GROUP BY region HAVING count(*) > 1 AND sum(amount) > 100')).toBeUndefined() }) + it('does not stream aggregates in later IN list values', () => { + // IN short-circuits once an earlier value matches, so a later aggregate + // may never be evaluated by the buffered path + expect(streamingPlan('SELECT 1 IN (1, sum(amount)) AS x FROM sales')).toBeUndefined() + }) + + it('streams aggregates in the first IN list value', () => { + const streaming = streamingPlan('SELECT 1 IN (sum(qty), 2) AS x FROM sales') + expect(streaming?.specs.length).toBe(1) + }) + + it('does not stream aggregates in ORDER BY tie-breaker terms', () => { + // The sorter evaluates later ORDER BY terms only within ties on earlier + // terms, so a tie-breaker aggregate may never be evaluated + expect(streamingPlan('SELECT region, count(*) AS c FROM sales GROUP BY region ORDER BY c, sum(amount)')).toBeUndefined() + }) + it('streams aggregates in the first WHEN condition', () => { const streaming = streamingPlan('SELECT CASE WHEN count(*) > 0 THEN 1 ELSE 2 END AS c FROM sales') expect(streaming?.specs.length).toBe(1) @@ -296,6 +344,38 @@ describe('streaming aggregate results', () => { expect(result).toEqual([{ r: 1 }]) }) + it('never evaluates aggregates in later IN list values', async () => { + const result = await collect(executeSql({ + tables, + functions: boomFunctions, + query: 'SELECT 100 IN (100, min(BOOM(amount))) AS x FROM sales', + })) + expect(result).toEqual([{ x: true }]) + }) + + it('never evaluates aggregate tie-breakers when earlier sort keys are unique', async () => { + const t = memorySource({ data: [{ g: 'a', v: 1 }, { g: 'b', v: 2 }, { g: 'b', v: 3 }] }) + const result = await collect(executeSql({ + tables: { t }, + functions: boomFunctions, + query: 'SELECT g, count(*) AS c FROM t GROUP BY g ORDER BY c DESC, min(BOOM(v))', + })) + expect(result).toEqual([ + { g: 'b', c: 2 }, + { g: 'a', c: 1 }, + ]) + }) + + it('leaves group keys over missing columns undefined like the buffered path', async () => { + // toEqual ignores undefined properties: this asserts g is undefined, not null + const t = memorySource({ data: [{ v: 1 }, { v: 2 }], columns: ['g', 'v'] }) + const result = await collect(executeSql({ + tables: { t }, + query: 'SELECT g, count(*) AS c FROM t GROUP BY g', + })) + expect(result).toEqual([{ c: 2 }]) + }) + it('evaluates later ORDER BY terms only for groups that tie on earlier terms', async () => { const t = memorySource({ data: [{ g: 'a' }, { g: 'b' }, { g: 'b' }, { g: 'c' }, { g: 'c' }, { g: 'c' }] }) const result = await collect(executeSql({ From 68ca15ed4c99e3189f5e16a45417bb1e1ea4db09 Mon Sep 17 00:00:00 2001 From: Kenny Daniel Date: Fri, 3 Jul 2026 17:41:01 -0700 Subject: [PATCH 6/7] Prune unread aggregate outputs and end aborted final chunks silently --- src/execute/streamingAggregate.js | 3 ++ src/plan/plan.js | 40 +++++++++++++++- test/execute/streamingAggregate.test.js | 63 +++++++++++++++++++++++-- 3 files changed, 102 insertions(+), 4 deletions(-) diff --git a/src/execute/streamingAggregate.js b/src/execute/streamingAggregate.js index 33ee38d..6182d0b 100644 --- a/src/execute/streamingAggregate.js +++ b/src/execute/streamingAggregate.js @@ -429,6 +429,9 @@ async function accumulateGroups({ child, groupBy, specs, needsRow, context }) { } if (chunk.length) { await accumulateChunk({ chunk, groupBy, specs, groups, needsRow, context }) + // An abort during the final partial chunk ends the stream silently, + // consistent with the full-chunk check above + if (context.signal?.aborted) return } context.signal?.throwIfAborted() return groups diff --git a/src/plan/plan.js b/src/plan/plan.js index 9c91ec5..fa53209 100644 --- a/src/plan/plan.js +++ b/src/plan/plan.js @@ -4,7 +4,7 @@ import { findAggregate } from '../validation/aggregates.js' import { ParseError } from '../validation/parseErrors.js' import { ColumnNotFoundError, TableNotFoundError } from '../validation/tables.js' import { validateNoIdentifiers, validateScan, validateTableRefs } from '../validation/tables.js' -import { collectScopeColumns, extractColumns, fromAlias, inferSelectSourceColumns, inferStatementColumns, statementScope, tableFunctionColumnNames } from './columns.js' +import { collectColumnsFromExpr, collectScopeColumns, extractColumns, fromAlias, inferSelectSourceColumns, inferStatementColumns, statementScope, tableFunctionColumnNames } from './columns.js' /** * @import { AsyncDataSource, DerivedColumn, ExprNode, FromFunction, IdentifierNode, JoinClause, OrderByItem, PlanSqlOptions, ScanOptions, SelectColumn, SelectStatement, SetOperationStatement, Statement, WindowFunctionNode } from '../types.js' @@ -208,6 +208,10 @@ function planSelect({ select, ctePlans, cteColumns, tables, parentColumns, outer const hints = {} const perTableColumns = extractColumns({ select: originalSelect, parentColumns }) if (sourceAlias !== undefined) hints.columns = perTableColumns.get(sourceAlias) + // Capture what the parent reads from a FROM subquery before the reset + // below, so aggregate outputs it never reads can still be pruned when the + // parent reads nothing at all (an empty array here means exactly that). + const subqueryNeeds = select.from?.type === 'subquery' ? hints.columns : undefined // Empty columns array means no columns were referenced, but a FROM subquery // still needs its own columns (e.g. for DISTINCT). Treat empty as unrestricted. if (hints.columns?.length === 0 && select.from?.type === 'subquery') { @@ -224,6 +228,7 @@ function planSelect({ select, ctePlans, cteColumns, tables, parentColumns, outer // Start with the data source (FROM clause) /** @type {QueryPlan} */ let plan = planFrom({ select, ctePlans, cteColumns, hints, tables, outerScope }) + pruneAggregateColumns(plan, subqueryNeeds) // Add JOINs if (select.joins.length) { @@ -349,6 +354,38 @@ function pushLimitIntoSort(plan, limit, offset) { } } +/** + * Drops derived columns from a subquery's aggregate node when the parent + * query never reads them. The buffered path defers aggregate cells so unread + * outputs are never evaluated, but the streaming path accumulates every + * planned aggregate eagerly, so unread aggregates must be removed at plan + * time to preserve that lazy behavior. Columns referenced by HAVING or the + * aggregate's ORDER BY are kept: both evaluate against the group context + * row, which exposes output aliases. Descends only through Subquery and + * Limit wrappers; Sort and Distinct depend on the full column set, so + * anything below them is left untouched. + * + * @param {QueryPlan} plan + * @param {string[] | undefined} needed - output column names the parent reads + */ +function pruneAggregateColumns(plan, needed) { + if (!needed) return + let node = plan + while (node.type === 'Subquery' || node.type === 'Limit') node = node.child + if (node.type !== 'HashAggregate' && node.type !== 'ScalarAggregate') return + /** @type {IdentifierNode[]} */ + const identifiers = [] + collectColumnsFromExpr(node.having, identifiers) + if (node.type === 'HashAggregate' && node.orderBy) { + for (const term of node.orderBy) collectColumnsFromExpr(term.expr, identifiers) + } + const keep = new Set(needed) + for (const { name } of identifiers) keep.add(name) + node.columns = node.columns.filter(col => + col.type === 'star' || keep.has(col.alias ?? derivedAlias(col.expr)) + ) +} + /** * @param {object} options * @param {SelectStatement} options.select @@ -487,6 +524,7 @@ function planJoin({ left, joins, leftTable, ctePlans, cteColumns, perTableColumn outerScope, parentColumns: subColumns?.map(name => ({ type: 'identifier', name, positionStart: 0, positionEnd: 0 })), }) + pruneAggregateColumns(subPlan, perTableColumns.get(rightTable)) const availableColumns = inferStatementColumns({ stmt: join.subquery.query, cteColumns, tables }) if (subColumns && availableColumns.length) { const missingColumn = subColumns.find(col => !availableColumns.includes(col)) diff --git a/test/execute/streamingAggregate.test.js b/test/execute/streamingAggregate.test.js index 930d114..1fc2f1e 100644 --- a/test/execute/streamingAggregate.test.js +++ b/test/execute/streamingAggregate.test.js @@ -105,9 +105,9 @@ describe('streaming aggregates', () => { expect(result).toEqual([{ mx: 7, n: 10000 }]) }) - it('falls back to buffered aggregation when pushdown prunes an aggregate argument', async () => { - // The outer query never reads s, so projection pushdown prunes v from the - // scan; only the buffered path defers sum(v) to the never-read output cell + it('prunes aggregate outputs the parent query never reads', async () => { + // The outer query never reads s, so the planner drops it from the + // aggregate node and sum(v) is never accumulated const result = await collect(executeSql({ tables: { t: [{ g: 'a', v: 1 }, { g: 'b', v: 2 }] }, query: 'SELECT g2 FROM (SELECT g AS g2, sum(v) AS s FROM t GROUP BY g) q ORDER BY g2', @@ -115,6 +115,42 @@ describe('streaming aggregates', () => { expect(result).toEqual([{ g2: 'a' }, { g2: 'b' }]) }) + it('does not evaluate unread scalar aggregate arguments in a subquery', async () => { + const result = await collect(executeSql({ + tables: { t: [{ v: 1 }, { v: 2 }] }, + functions: boomFunctions, + query: 'SELECT 1 AS one FROM (SELECT COUNT(BOOM(v)) AS c FROM t) q', + })) + expect(result).toEqual([{ one: 1 }]) + }) + + it('does not evaluate unread grouped aggregate arguments in a subquery', async () => { + const result = await collect(executeSql({ + tables: { t: [{ g: 'a', v: 1 }, { g: 'b', v: 2 }] }, + functions: boomFunctions, + query: 'SELECT g2 FROM (SELECT g AS g2, COUNT(BOOM(v)) AS c FROM t GROUP BY g) q ORDER BY g2', + })) + expect(result).toEqual([{ g2: 'a' }, { g2: 'b' }]) + }) + + it('does not evaluate unread aggregate arguments in a joined subquery', async () => { + const result = await collect(executeSql({ + tables: { t: [{ id: 1 }], u: [{ id: 1, v: 5 }] }, + functions: boomFunctions, + query: 'SELECT t.id FROM t JOIN (SELECT id, COUNT(BOOM(v)) AS c FROM u GROUP BY id) q ON t.id = q.id', + })) + expect(result).toEqual([{ id: 1 }]) + }) + + it('applies HAVING while pruning unread aggregate outputs', async () => { + const result = await collect(executeSql({ + tables: { t: [{ g: 'a', v: 1 }, { g: 'a', v: 2 }, { g: 'b', v: 3 }] }, + functions: boomFunctions, + query: 'SELECT g2 FROM (SELECT g AS g2, COUNT(BOOM(v)) AS c FROM t GROUP BY g HAVING COUNT(*) > 1) q', + })) + expect(result).toEqual([{ g2: 'a' }]) + }) + it('ends the stream silently when aborted during accumulation', async () => { const controller = new AbortController() /** @type {Record} */ @@ -135,6 +171,27 @@ describe('streaming aggregates', () => { })) expect(result).toEqual([]) }) + + it('ends the stream silently when aborted during the final partial chunk', async () => { + const controller = new AbortController() + /** @type {Record} */ + const functions = { + ABORT_NOW: { + apply(v) { + controller.abort() + return v + }, + arguments: { min: 1, max: 1 }, + }, + } + const result = await collect(executeSql({ + tables: { t: [{ g: 'a', v: 1 }, { g: 'b', v: 2 }] }, + functions, + query: 'SELECT g, count(ABORT_NOW(v)) AS c FROM t GROUP BY g', + signal: controller.signal, + })) + expect(result).toEqual([]) + }) }) const sales = [ From d6fbd41f3d003eb637234bb22b6111c7a4462037 Mon Sep 17 00:00:00 2001 From: Kenny Daniel Date: Fri, 3 Jul 2026 18:04:26 -0700 Subject: [PATCH 7/7] Keep SELECT and ORDER BY aggregates lazy when HAVING can reject groups --- src/execute/streamingAggregate.js | 21 +++++++---- test/execute/streamingAggregate.test.js | 46 ++++++++++++++++++++++++- 2 files changed, 59 insertions(+), 8 deletions(-) diff --git a/src/execute/streamingAggregate.js b/src/execute/streamingAggregate.js index 6182d0b..30726dd 100644 --- a/src/execute/streamingAggregate.js +++ b/src/execute/streamingAggregate.js @@ -115,10 +115,13 @@ export function planStreamingAggregates({ columns, having, orderBy, groupBy }, c * Walks an expression collecting streamable aggregate calls and group key * references. Returns false if the expression cannot be evaluated from * precomputed values plus a representative row. `lazy` marks positions the - * evaluator can skip (short-circuited AND/OR right sides and CASE - * branches): an aggregate there may never be evaluated by the buffered - * path, so accumulating it eagerly could evaluate expressions the query - * never asks for, and the query falls back to buffered aggregation. + * evaluator can skip (short-circuited AND/OR right sides, CASE branches, + * and select or sort expressions of a group HAVING may reject): an + * aggregate there may never be evaluated by the buffered path, so + * accumulating it eagerly could evaluate expressions the query never asks + * for, and the query falls back to buffered aggregation. A FILTER-less + * star aggregate is exempt: accumulating it evaluates nothing against + * input rows, so eager accumulation is unobservable. * * @param {ExprNode} node * @param {boolean} lazy @@ -163,9 +166,9 @@ export function planStreamingAggregates({ columns, having, orderBy, groupBy }, c if (!isAggregateFunc(funcName)) { return node.args.every(arg => walk(arg, lazy)) } - if (lazy) return false if (!STREAMABLE_FUNCS.has(funcName)) return false const star = node.args[0]?.type === 'star' + if (lazy && !(star && !node.filter)) return false if (!star && !node.args.every(arg => isScalarExpr(arg))) return false if (node.filter && !isScalarExpr(node.filter)) return false if (!specs.some(spec => spec.node === node)) { @@ -179,18 +182,22 @@ export function planStreamingAggregates({ columns, having, orderBy, groupBy }, c } } + // HAVING can reject a group before its output cells are ever read, so + // with HAVING present the buffered path may never evaluate SELECT or + // ORDER BY aggregates; treat those positions as lazy + const rejectable = Boolean(having) for (const col of columns) { if (col.type === 'star') { needsRow = true continue } - if (!walk(col.expr, false)) return + if (!walk(col.expr, rejectable)) return } if (having && !walk(having, false)) return const orderTerms = orderBy ?? [] for (let i = 0; i < orderTerms.length; i++) { // the sorter evaluates later terms only to break ties on earlier terms - if (!walk(orderTerms[i].expr, i > 0)) return + if (!walk(orderTerms[i].expr, rejectable || i > 0)) return } if (childColumns && !specsResolvable(specs, childColumns)) return return { specs, keyRefs, needsRow } diff --git a/test/execute/streamingAggregate.test.js b/test/execute/streamingAggregate.test.js index 1fc2f1e..1c644ec 100644 --- a/test/execute/streamingAggregate.test.js +++ b/test/execute/streamingAggregate.test.js @@ -151,6 +151,33 @@ describe('streaming aggregates', () => { expect(result).toEqual([{ g2: 'a' }]) }) + it('does not evaluate SELECT aggregate arguments for groups HAVING rejects', async () => { + const result = await collect(executeSql({ + tables: { t: [{ g: 'a', v: 1 }] }, + functions: boomFunctions, + query: 'SELECT g, count(BOOM(v)) AS c FROM t GROUP BY g HAVING count(*) > 1', + })) + expect(result).toEqual([]) + }) + + it('does not evaluate scalar SELECT aggregate arguments when HAVING rejects the row', async () => { + const result = await collect(executeSql({ + tables: { t: [{ v: 1 }] }, + functions: boomFunctions, + query: 'SELECT count(BOOM(v)) AS c FROM t HAVING FALSE', + })) + expect(result).toEqual([]) + }) + + it('does not evaluate ORDER BY aggregate arguments for groups HAVING rejects', async () => { + const result = await collect(executeSql({ + tables: { t: [{ g: 'a', v: 1 }] }, + functions: boomFunctions, + query: 'SELECT g FROM t GROUP BY g HAVING count(*) > 1 ORDER BY min(BOOM(v))', + })) + expect(result).toEqual([]) + }) + it('ends the stream silently when aborted during accumulation', async () => { const controller = new AbortController() /** @type {Record} */ @@ -217,11 +244,28 @@ describe('streaming aggregate row retention', () => { } it('retains no rows when expressions only use group keys and aggregates', () => { - const streaming = streamingPlan('SELECT region, count(*) FROM sales GROUP BY region HAVING count(*) > 1 ORDER BY sum(amount)') + const streaming = streamingPlan('SELECT region, count(*) FROM sales GROUP BY region HAVING count(*) > 1 ORDER BY count(*) DESC') expect(streaming?.needsRow).toBe(false) expect(streaming?.keyRefs.size).toBe(1) }) + it('does not stream non-star SELECT aggregates when HAVING can reject the group', () => { + // HAVING can reject a group before its output cells are read, so the + // buffered path may never evaluate sum(amount); star aggregates like the + // count(*) in HAVING stay streamable since accumulating them evaluates nothing + expect(streamingPlan('SELECT region, sum(amount) AS s FROM sales GROUP BY region HAVING count(*) > 1')).toBeUndefined() + }) + + it('does not stream non-star ORDER BY aggregates when HAVING can reject the group', () => { + expect(streamingPlan('SELECT region, count(*) AS c FROM sales GROUP BY region HAVING count(*) > 1 ORDER BY sum(amount)')).toBeUndefined() + }) + + it('does not stream star aggregates with FILTER when HAVING can reject the group', () => { + // the FILTER condition is evaluated per input row, which the buffered + // path skips entirely for groups HAVING rejects + expect(streamingPlan('SELECT region, count(*) FILTER (WHERE amount > 100) AS c FROM sales GROUP BY region HAVING count(*) > 1')).toBeUndefined() + }) + it('retains no rows for function group keys', () => { const streaming = streamingPlan('SELECT upper(region) AS r, count(*) FROM sales GROUP BY upper(region)') expect(streaming?.needsRow).toBe(false)