From 881a0318615dcc8996d07beb4c48b3a84ace98a7 Mon Sep 17 00:00:00 2001 From: Phillip Cunliffe Date: Thu, 9 Apr 2026 16:45:28 -0700 Subject: [PATCH 1/9] Optimize SQL execution pipeline for repeated queries Add multi-level caching and reduce per-row overhead: - parseSql: LRU cache (64 entries) avoids re-tokenizing/parsing same SQL strings - planSql: WeakMap cache on parsed ASTs avoids re-planning identical queries - asyncRow: attach _data field for zero-copy collection - collect: sync fast-path skips Promise.all when all rows have pre-materialized _data - executeProject: pre-compute static column names, fast-path for simple identifier projections with direct cell passthrough and _data propagation - executeSql: skip table normalization when no array tables are present - compareForTerm: use module-level Set instead of per-call array allocation - memorySource: hoist column computation outside scan loop, use Set for validation --- src/backend/dataSource.js | 8 ++-- src/execute/execute.js | 82 ++++++++++++++++++++++++++++++++++----- src/execute/utils.js | 30 ++++++++++++-- src/parse/parse.js | 23 +++++++++++ src/plan/plan.js | 29 +++++++++++++- 5 files changed, 155 insertions(+), 17 deletions(-) diff --git a/src/backend/dataSource.js b/src/backend/dataSource.js index b4206b0..4bffb0c 100644 --- a/src/backend/dataSource.js +++ b/src/backend/dataSource.js @@ -15,7 +15,7 @@ export function asyncRow(obj, columns) { for (const key of columns) { cells[key] = () => Promise.resolve(obj[key]) } - return { columns, cells } + return { columns, cells, _data: obj } } /** @@ -34,13 +34,14 @@ export function memorySource({ data, columns }) { } const firstColumns = Object.keys(data[0]) // Check first 1000 rows for consistent columns + const firstColSet = new Set(firstColumns) for (let i = 1; i < data.length && i < 1000; i++) { const rowColumns = Object.keys(data[i]) const missing = firstColumns.find(col => !rowColumns.includes(col)) if (missing) { throw new Error(`Inconsistent data, column "${missing}" not found in row ${i}`) } - const extra = rowColumns.find(col => !firstColumns.includes(col)) + const extra = rowColumns.find(col => !firstColSet.has(col)) if (extra) { throw new Error(`Inconsistent data, unexpected column "${extra}" found in row ${i}`) } @@ -54,11 +55,12 @@ export function memorySource({ data, columns }) { // Only apply offset and limit if no where clause const start = !where ? offset ?? 0 : 0 const end = !where && limit !== undefined ? start + limit : data.length + const rowColumns = scanColumns ?? columns return { rows: (async function* () { for (let i = start; i < end && i < data.length; i++) { if (signal?.aborted) break - yield asyncRow(data[i], scanColumns ?? columns) + yield asyncRow(data[i], rowColumns) } })(), appliedWhere: false, diff --git a/src/execute/execute.js b/src/execute/execute.js index c87a60c..925ee6c 100644 --- a/src/execute/execute.js +++ b/src/execute/execute.js @@ -24,14 +24,27 @@ export async function* executeSql({ tables, query, functions, signal }) { const parsed = typeof query === 'string' ? parseSql({ query, functions }) : query // Normalize tables: convert arrays to AsyncDataSource + // Fast path: skip normalization when no arrays are present + let needsNormalization = false + const tableKeys = Object.keys(tables) + for (let i = 0; i < tableKeys.length; i++) { + if (Array.isArray(tables[tableKeys[i]])) { + needsNormalization = true + break + } + } + /** @type {Record} */ - const normalizedTables = {} - for (const [name, data] of Object.entries(tables)) { - if (Array.isArray(data)) { - normalizedTables[name] = memorySource({ data }) - } else { - normalizedTables[name] = data + let normalizedTables + if (needsNormalization) { + normalizedTables = {} + for (let i = 0; i < tableKeys.length; i++) { + const name = tableKeys[i] + const data = tables[name] + normalizedTables[name] = Array.isArray(data) ? memorySource({ data }) : data } + } else { + normalizedTables = /** @type {Record} */ (tables) } yield* executeStatement({ query: parsed, context: { tables: normalizedTables, functions, signal } }) @@ -271,19 +284,68 @@ async function* executeFilter(plan, context) { * @yields {AsyncRow} */ async function* executeProject(plan, context) { + // Pre-compute column names for derived columns (avoids per-row derivedAlias calls) + const hasStar = plan.columns.some(col => col.type === 'star') + + // For simple identifier projections, map output alias → source column name + /** @type {string[] | undefined} */ + let staticColumns + /** @type {{ alias: string, sourceName: string }[] | undefined} */ + let identifierMap + if (!hasStar) { + staticColumns = plan.columns.map(col => col.alias ?? derivedAlias(col.expr)) + // Check if all columns are simple identifier references (no expressions) + const allIdentifiers = plan.columns.every(col => + col.expr.type === 'identifier' && !col.expr.prefix + ) + if (allIdentifiers) { + identifierMap = plan.columns.map((col, i) => ({ + alias: staticColumns[i], + sourceName: col.expr.name, + })) + } + } + let rowIndex = 0 + let identifierMapValidated = false for await (const row of executePlan({ plan: plan.child, context })) { if (context.signal?.aborted) return rowIndex++ + + // Validate identifier fast path on first row (may fail for JOINs with prefixed columns) + if (identifierMap && !identifierMapValidated) { + identifierMapValidated = true + if (!identifierMap.every(m => m.sourceName in row.cells)) { + identifierMap = undefined + } + } + + // Fast path: all columns are simple identifier references + if (identifierMap) { + /** @type {AsyncCells} */ + const cells = {} + const srcData = row._data + const _data = srcData ? {} : undefined + for (const { alias, sourceName } of identifierMap) { + cells[alias] = row.cells[sourceName] + if (_data) _data[alias] = srcData[sourceName] + } + yield _data + ? { columns: staticColumns, cells, _data } + : { columns: staticColumns, cells } + continue + } + const currentRowIndex = rowIndex /** @type {string[]} */ - const columns = [] + const columns = staticColumns ? staticColumns : [] /** @type {AsyncCells} */ const cells = {} - for (const col of plan.columns) { + for (let i = 0; i < plan.columns.length; i++) { + const col = plan.columns[i] if (col.type === 'star') { const prefix = col.table ? `${col.table}.` : undefined for (const key of row.columns) { @@ -295,8 +357,8 @@ async function* executeProject(plan, context) { cells[outputKey] = row.cells[key] } } else { - const alias = col.alias ?? derivedAlias(col.expr) - columns.push(alias) + const alias = staticColumns ? staticColumns[i] : (col.alias ?? derivedAlias(col.expr)) + if (!staticColumns) columns.push(alias) cells[alias] = () => evaluateExpr({ node: col.expr, row, diff --git a/src/execute/utils.js b/src/execute/utils.js index 1bdda3e..e40a16d 100644 --- a/src/execute/utils.js +++ b/src/execute/utils.js @@ -10,6 +10,8 @@ * @param {OrderByItem} term * @returns {number} */ +const primitiveTypes = new Set(['number', 'bigint', 'boolean', 'string']) + export function compareForTerm(a, b, term) { const aIsNull = a == null const bIsNull = b == null @@ -24,10 +26,9 @@ export function compareForTerm(a, b, term) { // Compare non-null values if (a == b) return 0 - const primitives = ['number', 'bigint', 'boolean', 'string'] let cmp - if (primitives.includes(typeof a) && primitives.includes(typeof b)) { - cmp = a < b ? -1 : a > b ? 1 : 0 + if (primitiveTypes.has(typeof a) && primitiveTypes.has(typeof b)) { + cmp = a < b ? -1 : 1 } else { const aa = String(a) const bb = String(b) @@ -51,6 +52,29 @@ export async function collect(asyncRows) { for await (const asyncRow of asyncRows) { rows.push(asyncRow) } + + // Fast path: if all rows have pre-materialized data, skip Promise overhead + let allMaterialized = rows.length > 0 + for (let i = 0; i < rows.length; i++) { + if (!rows[i]._data) { + allMaterialized = false + break + } + } + if (allMaterialized) { + const result = new Array(rows.length) + for (let i = 0; i < rows.length; i++) { + const row = rows[i] + /** @type {Record} */ + const item = {} + for (const col of row.columns) { + item[col] = row._data[col] + } + result[i] = item + } + return result + } + return Promise.all(rows.map(async asyncRow => { const values = await Promise.all(asyncRow.columns.map(k => asyncRow.cells[k]())) /** @type {Record} */ diff --git a/src/parse/parse.js b/src/parse/parse.js index b882ca5..4e709f7 100644 --- a/src/parse/parse.js +++ b/src/parse/parse.js @@ -10,11 +10,26 @@ import { tokenizeSql } from './tokenize.js' * @import { CTEDefinition, ExprNode, FromSubquery, FromTable, OrderByItem, ParseSqlOptions, ParserState, SelectColumn, SelectStatement, SetOperationStatement, SetOperator, Statement } from '../types.js' */ +const MAX_PARSE_CACHE = 64 +/** @type {Map} */ +const parseCache = new Map() + /** * @param {ParseSqlOptions} options * @returns {Statement} */ export function parseSql({ query, functions }) { + // Cache only for simple queries without custom functions + if (!functions) { + const cached = parseCache.get(query) + if (cached) { + // LRU touch + parseCache.delete(query) + parseCache.set(query, cached) + return cached + } + } + const tokens = tokenizeSql(query) /** @type {ParserState} */ const state = { tokens, pos: 0, lastPos: 0, functions } @@ -27,6 +42,14 @@ export function parseSql({ query, functions }) { throw parseError(state, 'end of query') } + if (!functions) { + parseCache.set(query, stmt) + if (parseCache.size > MAX_PARSE_CACHE) { + const oldest = parseCache.keys().next().value + if (oldest) parseCache.delete(oldest) + } + } + return stmt } diff --git a/src/plan/plan.js b/src/plan/plan.js index 5bb27f4..8ee0ca3 100644 --- a/src/plan/plan.js +++ b/src/plan/plan.js @@ -10,6 +10,20 @@ import { extractColumns, fromAlias, inferStatementColumns } from './columns.js' * @import { QueryPlan } from './types.d.ts' */ +/** @type {WeakMap} */ +const planCache = new WeakMap() + +/** + * @param {Record | undefined} tables + * @returns {string} + */ +function tablesKey(tables) { + if (!tables) return '' + const keys = Object.keys(tables) + keys.sort() + return keys.join(',') +} + /** * Builds a query plan from a statement AST. * Resolves CTEs at plan time so no planning occurs during execution. @@ -20,7 +34,20 @@ import { extractColumns, fromAlias, inferStatementColumns } from './columns.js' export function planSql({ query, functions, tables }) { /** @type {Statement} */ const stmt = typeof query === 'string' ? parseSql({ query, functions }) : query - return planStatement({ stmt, tables }) + + const key = !functions ? tablesKey(tables) : undefined + if (key !== undefined) { + const cached = planCache.get(stmt) + if (cached && cached.tablesKey === key) return cached.plan + } + + const plan = planStatement({ stmt, tables }) + + if (key !== undefined) { + planCache.set(stmt, { plan, tablesKey: key }) + } + + return plan } /** From ac137466726a701ac0f88ffb42e0bd3e4e8dae57 Mon Sep 17 00:00:00 2001 From: Phillip Cunliffe Date: Thu, 9 Apr 2026 18:46:59 -0700 Subject: [PATCH 2/9] Fix typecheck errors - Add _data to AsyncRow type definition - Cast to DerivedColumn/IdentifierNode where type narrowing is needed - Type _data as Record - Fix JSDoc placement for compareForTerm --- src/execute/execute.js | 13 ++++++++----- src/execute/utils.js | 4 ++-- src/types.d.ts | 1 + 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/src/execute/execute.js b/src/execute/execute.js index 925ee6c..c0df0be 100644 --- a/src/execute/execute.js +++ b/src/execute/execute.js @@ -293,15 +293,17 @@ async function* executeProject(plan, context) { /** @type {{ alias: string, sourceName: string }[] | undefined} */ let identifierMap if (!hasStar) { - staticColumns = plan.columns.map(col => col.alias ?? derivedAlias(col.expr)) + /** @type {import('../types.js').DerivedColumn[]} */ + const derived = /** @type {any} */ (plan.columns) + staticColumns = derived.map(col => col.alias ?? derivedAlias(col.expr)) // Check if all columns are simple identifier references (no expressions) - const allIdentifiers = plan.columns.every(col => + const allIdentifiers = derived.every(col => col.expr.type === 'identifier' && !col.expr.prefix ) if (allIdentifiers) { - identifierMap = plan.columns.map((col, i) => ({ + identifierMap = derived.map((col, i) => ({ alias: staticColumns[i], - sourceName: col.expr.name, + sourceName: /** @type {import('../types.js').IdentifierNode} */ (col.expr).name, })) } } @@ -326,10 +328,11 @@ async function* executeProject(plan, context) { /** @type {AsyncCells} */ const cells = {} const srcData = row._data + /** @type {Record | undefined} */ const _data = srcData ? {} : undefined for (const { alias, sourceName } of identifierMap) { cells[alias] = row.cells[sourceName] - if (_data) _data[alias] = srcData[sourceName] + if (_data && srcData) _data[alias] = srcData[sourceName] } yield _data ? { columns: staticColumns, cells, _data } diff --git a/src/execute/utils.js b/src/execute/utils.js index e40a16d..80aa5c3 100644 --- a/src/execute/utils.js +++ b/src/execute/utils.js @@ -2,6 +2,8 @@ * @import { AsyncRow, OrderByItem, SqlPrimitive } from '../types.js' */ +const primitiveTypes = new Set(['number', 'bigint', 'boolean', 'string']) + /** * Compares two values for a single ORDER BY term, handling nulls and direction * @@ -10,8 +12,6 @@ * @param {OrderByItem} term * @returns {number} */ -const primitiveTypes = new Set(['number', 'bigint', 'boolean', 'string']) - export function compareForTerm(a, b, term) { const aIsNull = a == null const bIsNull = b == null diff --git a/src/types.d.ts b/src/types.d.ts index 5e2feb5..b7e6f6b 100644 --- a/src/types.d.ts +++ b/src/types.d.ts @@ -36,6 +36,7 @@ export interface ExecuteContext { export interface AsyncRow { columns: string[] cells: AsyncCells + _data?: Record } export type AsyncCells = Record export type AsyncCell = () => Promise From 2139421e94ff30c14c49f103d3b7098b6807b83b Mon Sep 17 00:00:00 2001 From: Phillip Cunliffe Date: Fri, 10 Apr 2026 09:02:37 -0700 Subject: [PATCH 3/9] Remove parse/plan caches and rename AsyncRow._data to resolved Drop the parseSql/planSql memoization caches added in 881a031. Also rename the pre-materialized row payload from `_data` to `resolved` for clarity, and delete stale scratch files (query-parquet.mjs, repro-525.mjs). Co-Authored-By: Claude Opus 4.6 (1M context) --- query-parquet.mjs | 187 ----------------------------- repro-525.mjs | 246 -------------------------------------- src/backend/dataSource.js | 2 +- src/execute/aggregates.js | 4 +- src/execute/execute.js | 41 ++++--- src/execute/join.js | 6 +- src/execute/sort.js | 2 +- src/execute/utils.js | 4 +- src/parse/parse.js | 23 ---- src/plan/plan.js | 29 +---- src/types.d.ts | 4 +- 11 files changed, 33 insertions(+), 515 deletions(-) delete mode 100644 query-parquet.mjs delete mode 100644 repro-525.mjs diff --git a/query-parquet.mjs b/query-parquet.mjs deleted file mode 100644 index 58895a3..0000000 --- a/query-parquet.mjs +++ /dev/null @@ -1,187 +0,0 @@ -import { readFileSync } from 'fs' -import { parquetMetadata, parquetReadObjects, parquetSchema } from 'hyparquet' -import { compressors } from 'hyparquet-compressors' -import { asyncRow, executeSql, collect, parseSql } from './src/index.js' - -// --- WHERE pushdown: ported from hyperparam's parquetPushdownFilter.ts --- - -function whereToParquetFilter(where) { - if (!where) return undefined - return convertExpr(where, false) -} - -function convertExpr(node, negate) { - if (node.type === 'unary' && node.op === 'NOT') { - return convertExpr(node.argument, !negate) - } - if (node.type === 'binary') { - return convertBinary(node, negate) - } - if (node.type === 'in valuelist') { - return convertInValues(node, negate) - } - if (node.type === 'cast') { - return convertExpr(node.expr, negate) - } - return undefined -} - -function convertBinary({ op, left, right }, negate) { - if (op === 'AND') { - const l = convertExpr(left, negate) - const r = convertExpr(right, negate) - if (!l || !r) return - return negate ? { $or: [l, r] } : { $and: [l, r] } - } - if (op === 'OR') { - const l = convertExpr(left, false) - const r = convertExpr(right, false) - if (!l || !r) return - return negate ? { $nor: [l, r] } : { $or: [l, r] } - } - if (op === 'LIKE') return - - const { column, value, flipped } = extractColumnAndValue(left, right) - if (!column || value === undefined) return - - const mongoOp = mapOperator(op, flipped, negate) - if (!mongoOp) return - return { [column]: { [mongoOp]: value } } -} - -function extractColumnAndValue(left, right) { - if (left.type === 'identifier' && right.type === 'literal') { - return { column: left.name, value: coerceToBigInt(right.value), flipped: false } - } - if (left.type === 'literal' && right.type === 'identifier') { - return { column: right.name, value: coerceToBigInt(left.value), flipped: true } - } - return { column: undefined, value: undefined, flipped: false } -} - -// Parquet integer columns are bigint — coerce number literals to match -function coerceToBigInt(value) { - if (typeof value === 'number' && Number.isInteger(value)) return BigInt(value) - return value -} - -function mapOperator(op, flipped, negate) { - const comparisons = ['=', '!=', '<>', '<', '>', '<=', '>='] - if (!comparisons.includes(op)) return - let mapped = op - if (negate) mapped = neg(mapped) - if (flipped) mapped = flip(mapped) - if (mapped === '<') return '$lt' - if (mapped === '<=') return '$lte' - if (mapped === '>') return '$gt' - if (mapped === '>=') return '$gte' - if (mapped === '=') return '$eq' - return '$ne' -} - -function neg(op) { - const map = { '<': '>=', '<=': '>', '>': '<=', '>=': '<', '=': '!=', '!=': '=' } - return map[op] ?? op -} - -function flip(op) { - const map = { '<': '>', '<=': '>=', '>': '<', '>=': '<=' } - return map[op] ?? op -} - -function convertInValues(node, negate) { - if (node.expr.type !== 'identifier') return - const values = [] - for (const val of node.values) { - if (val.type !== 'literal') return - values.push(val.value) - } - return { [node.expr.name]: { [negate ? '$nin' : '$in']: values } } -} - -// --- Parquet data source with pushdown --- - -function parquetDataSource(file, metadata) { - const schema = parquetSchema(metadata) - return { - numRows: Number(metadata.num_rows), - columns: schema.children.map(c => c.element.name), - scan(hints) { - const whereFilter = hints.where && whereToParquetFilter(hints.where) - const filter = hints.where ? whereFilter : undefined - const appliedWhere = Boolean(filter && whereFilter) - const appliedLimitOffset = !hints.where || appliedWhere - - return { - rows: (async function* () { - let groupStart = 0 - let remainingLimit = hints.limit ?? Infinity - for (const rowGroup of metadata.row_groups) { - if (hints.signal?.aborted) break - const rowCount = Number(rowGroup.num_rows) - - let safeOffset = 0 - let safeLimit = rowCount - if (appliedLimitOffset) { - if (hints.offset !== undefined && groupStart < hints.offset) { - safeOffset = Math.min(rowCount, hints.offset - groupStart) - } - safeLimit = Math.min(rowCount - safeOffset, remainingLimit) - if (safeLimit <= 0 && safeOffset < rowCount) break - } - if (safeOffset === rowCount) { - groupStart += rowCount - continue - } - - const cols = hints.columns ?? schema.children.map(c => c.element.name) - console.error(` row group ${groupStart}: reading ${safeLimit} rows, columns: [${cols}]${filter ? ', with pushdown filter' : ''}`) - const data = await parquetReadObjects({ - file, - compressors, - metadata, - rowStart: groupStart + safeOffset, - rowEnd: groupStart + safeOffset + safeLimit, - columns: cols, - filter, - }) - - console.error(` -> ${data.length} rows after filter`) - for (const row of data) { - yield asyncRow(row, Object.keys(row)) - } - - remainingLimit -= data.length - groupStart += rowCount - } - })(), - appliedWhere, - appliedLimitOffset, - } - }, - } -} - -// --- Main --- -const parquetFile = process.argv[2] || '0000.parquet' -const query = process.argv[3] || 'SELECT COUNT(*) as cnt FROM data WHERE turns = 8' - -const buffer = readFileSync(parquetFile) -const arrayBuffer = buffer.buffer.slice(buffer.byteOffset, buffer.byteOffset + buffer.byteLength) -const metadata = parquetMetadata(arrayBuffer) - -console.log(`File: ${parquetFile}`) -console.log(`Rows: ${metadata.num_rows}, Row groups: ${metadata.row_groups.length}`) -console.log(`Query: ${query}`) -console.log() - -const source = parquetDataSource(arrayBuffer, metadata) -const start = performance.now() -const result = await collect(executeSql({ tables: { data: source }, query })) -const elapsed = (performance.now() - start).toFixed(0) - -console.log('\nResult:') -for (const row of result) { - console.log(row) -} -console.log(`(${elapsed}ms)`) diff --git a/repro-525.mjs b/repro-525.mjs deleted file mode 100644 index 1a76f46..0000000 --- a/repro-525.mjs +++ /dev/null @@ -1,246 +0,0 @@ -import { readFileSync } from 'fs' -import { parquetMetadata, parquetReadObjects, parquetSchema } from 'hyparquet' -import { compressors } from 'hyparquet-compressors' -import { asyncRow, executeSql, collect, parseSql } from './src/index.js' - -const buffer = readFileSync('0000.parquet') -const arrayBuffer = buffer.buffer.slice(buffer.byteOffset, buffer.byteOffset + buffer.byteLength) -const metadata = parquetMetadata(arrayBuffer) -const schema = parquetSchema(metadata) - -console.log(`File: 0000.parquet`) -console.log(`Rows: ${metadata.num_rows}, Row groups: ${metadata.row_groups.length}`) -console.log(`Columns: ${schema.children.map(c => c.element.name).join(', ')}`) -console.log() - -// --- Test 1: Direct count per row group with various filter options --- -console.log('=== Test 1: Direct parquetReadObjects per row group ===') -let groupStart = 0 -let totalStrict = 0, totalLoose = 0, totalBigint = 0, totalNoFilter = 0 -let totalUOI = 0 // useOffsetIndex - -for (const rowGroup of metadata.row_groups) { - const rowCount = Number(rowGroup.num_rows) - - const [strict, loose, bigint, noFilter, uoi] = await Promise.all([ - parquetReadObjects({ file: arrayBuffer, compressors, metadata, - rowStart: groupStart, rowEnd: groupStart + rowCount, - columns: ['turns'], filter: { turns: { $eq: 8 } }, filterStrict: true }), - parquetReadObjects({ file: arrayBuffer, compressors, metadata, - rowStart: groupStart, rowEnd: groupStart + rowCount, - columns: ['turns'], filter: { turns: { $eq: 8 } }, filterStrict: false }), - parquetReadObjects({ file: arrayBuffer, compressors, metadata, - rowStart: groupStart, rowEnd: groupStart + rowCount, - columns: ['turns'], filter: { turns: { $eq: 8n } }, filterStrict: true }), - parquetReadObjects({ file: arrayBuffer, compressors, metadata, - rowStart: groupStart, rowEnd: groupStart + rowCount, - columns: ['turns'] }), - parquetReadObjects({ file: arrayBuffer, compressors, metadata, - rowStart: groupStart, rowEnd: groupStart + rowCount, - columns: ['turns'], filter: { turns: { $eq: 8 } }, filterStrict: false, - useOffsetIndex: true }), - ]) - - if (strict.length || loose.length || bigint.length || uoi.length) { - console.log(` RG @${groupStart} (${rowCount} rows): strict=${strict.length} loose=${loose.length} bigint=${bigint.length} noFilter=${noFilter.length} useOffsetIndex=${uoi.length}`) - } - totalStrict += strict.length - totalLoose += loose.length - totalBigint += bigint.length - totalNoFilter += noFilter.length - totalUOI += uoi.length - groupStart += rowCount -} - -console.log(`\nTotals: strict=${totalStrict} loose=${totalLoose} bigint=${totalBigint} total=${totalNoFilter} useOffsetIndex=${totalUOI}`) - -// --- Test 2: Simulate workerParquetDataSource exactly --- -console.log('\n=== Test 2: Exact workerParquetDataSource simulation ===') - -function whereToParquetFilter(where) { - if (!where) return undefined - if (where.type === 'binary') { - const { op, left, right } = where - if (op === 'AND') { - const l = whereToParquetFilter(left) - const r = whereToParquetFilter(right) - if (!l || !r) return - return { $and: [l, r] } - } - if (op === 'OR') { - const l = whereToParquetFilter(left) - const r = whereToParquetFilter(right) - if (!l || !r) return - return { $or: [l, r] } - } - if (left.type === 'identifier' && right.type === 'literal') { - const opMap = { '=': '$eq', '!=': '$ne', '<': '$lt', '<=': '$lte', '>': '$gt', '>=': '$gte' } - if (opMap[op]) return { [left.name]: { [opMap[op]]: right.value } } - } - if (left.type === 'literal' && right.type === 'identifier') { - const flipMap = { '<': '$gt', '<=': '$gte', '>': '$lt', '>=': '$lte', '=': '$eq', '!=': '$ne' } - if (flipMap[op]) return { [right.name]: { [flipMap[op]]: left.value } } - } - } - return undefined -} - -function workerParquetDataSource(file, metadata) { - const schema = parquetSchema(metadata) - return { - numRows: Number(metadata.num_rows), - columns: schema.children.map(c => c.element.name), - scan(hints) { - const whereFilter = hints.where && whereToParquetFilter(hints.where) - const filter = hints.where ? whereFilter : undefined - const appliedWhere = Boolean(filter && whereFilter) - const appliedLimitOffset = !hints.where || appliedWhere - - console.log(` scan() called:`) - console.log(` columns: ${JSON.stringify(hints.columns)}`) - console.log(` where: ${JSON.stringify(hints.where?.op)} ${JSON.stringify(hints.where?.right?.value)}`) - console.log(` limit: ${hints.limit}, offset: ${hints.offset}`) - console.log(` filter: ${JSON.stringify(filter)}`) - console.log(` appliedWhere: ${appliedWhere}, appliedLimitOffset: ${appliedLimitOffset}`) - - return { - rows: (async function* () { - let groupStart = 0 - let remainingLimit = hints.limit ?? Infinity - for (const rowGroup of metadata.row_groups) { - if (hints.signal?.aborted) throw new DOMException('Aborted', 'AbortError') - const rowCount = Number(rowGroup.num_rows) - - let safeOffset = 0 - let safeLimit = rowCount - if (appliedLimitOffset) { - if (hints.offset !== undefined && groupStart < hints.offset) { - safeOffset = Math.min(rowCount, hints.offset - groupStart) - } - safeLimit = Math.min(rowCount - safeOffset, remainingLimit) - if (safeLimit <= 0 && safeOffset < rowCount) { - console.log(` RG @${groupStart}: BREAK (safeLimit=${safeLimit}, safeOffset=${safeOffset})`) - break - } - } - if (safeOffset === rowCount) { - groupStart += rowCount - continue - } - - const data = await parquetReadObjects({ - file, - compressors, - metadata, - rowStart: groupStart + safeOffset, - rowEnd: groupStart + safeOffset + safeLimit, - columns: hints.columns, - filter, - filterStrict: false, - useOffsetIndex: true, - }) - - console.log(` RG @${groupStart}: requested=${safeLimit}, got=${data.length}, remainingLimit=${remainingLimit}`) - - for (const row of data) { - yield asyncRow(row, Object.keys(row)) - } - - remainingLimit -= data.length - groupStart += rowCount - } - })(), - appliedWhere, - appliedLimitOffset, - } - }, - } -} - -const table = workerParquetDataSource(arrayBuffer, metadata) -const ast = parseSql({ query: 'SELECT COUNT(*) as cnt FROM table WHERE turns = 8' }) -const results = await collect(executeSql({ tables: { table }, query: ast })) -console.log('\nResult:', results) - -// --- Test 3: Simulate AsyncBuffer (chunked reads like HTTP range requests) --- -console.log('\n=== Test 3: AsyncBuffer simulation (like browser HTTP range requests) ===') - -const asyncBuffer = { - byteLength: arrayBuffer.byteLength, - slice(start, end) { - return Promise.resolve(arrayBuffer.slice(start, end)) - }, -} - -groupStart = 0 -let totalAsync = 0 -for (const rowGroup of metadata.row_groups) { - const rowCount = Number(rowGroup.num_rows) - const data = await parquetReadObjects({ - file: asyncBuffer, - compressors, - metadata, - rowStart: groupStart, - rowEnd: groupStart + rowCount, - columns: ['turns'], - filter: { turns: { $eq: 8 } }, - filterStrict: false, - useOffsetIndex: true, - }) - if (data.length > 0) { - console.log(` RG @${groupStart}: ${data.length} rows`) - } - totalAsync += data.length - groupStart += rowCount -} -console.log(` Total with AsyncBuffer: ${totalAsync}`) - -// --- Test 4: Check if 525 relates to any partial read --- -console.log('\n=== Test 4: Checking partial counts ===') -const allCounts = [] -groupStart = 0 -for (const rowGroup of metadata.row_groups) { - const rowCount = Number(rowGroup.num_rows) - const data = await parquetReadObjects({ - file: arrayBuffer, compressors, metadata, - rowStart: groupStart, rowEnd: groupStart + rowCount, - columns: ['turns'], filter: { turns: { $eq: 8 } }, filterStrict: false, - }) - allCounts.push({ groupStart, rowCount, matchCount: data.length }) - groupStart += rowCount -} - -// Check cumulative sums -let cumulative = 0 -for (const { groupStart, rowCount, matchCount } of allCounts) { - cumulative += matchCount - if (cumulative === 525 || matchCount === 525) { - console.log(` FOUND 525! cumulative=${cumulative} at groupStart=${groupStart}`) - } -} -console.log(` Final cumulative: ${cumulative}`) - -// Check if 525 matches total minus something -console.log(` 7108 - 525 = ${7108 - 525}`) -console.log(` 525 / 7108 = ${(525/7108).toFixed(4)}`) - -// Check if manual WHERE filter in JS gives different results -console.log('\n=== Test 5: Manual filter on all rows ===') -groupStart = 0 -let manualCount = 0 -let manualLooseCount = 0 -for (const rowGroup of metadata.row_groups) { - const rowCount = Number(rowGroup.num_rows) - const data = await parquetReadObjects({ - file: arrayBuffer, compressors, metadata, - rowStart: groupStart, rowEnd: groupStart + rowCount, - columns: ['turns'], - }) - for (const row of data) { - if (row.turns === 8) manualCount++ // strict - if (row.turns == 8) manualLooseCount++ // loose - } - groupStart += rowCount -} -console.log(` Strict (=== 8): ${manualCount}`) -console.log(` Loose (== 8): ${manualLooseCount}`) diff --git a/src/backend/dataSource.js b/src/backend/dataSource.js index 5c17afd..d152375 100644 --- a/src/backend/dataSource.js +++ b/src/backend/dataSource.js @@ -15,7 +15,7 @@ export function asyncRow(obj, columns) { for (const key of columns) { cells[key] = () => Promise.resolve(obj[key]) } - return { columns, cells, _data: obj } + return { columns, cells, resolved: obj } } /** diff --git a/src/execute/aggregates.js b/src/execute/aggregates.js index b173950..17216a2 100644 --- a/src/execute/aggregates.js +++ b/src/execute/aggregates.js @@ -62,7 +62,7 @@ export function executeHashAggregate(plan, context) { return { columns: selectColumnNames(plan.columns, child.columns), maxRows: child.maxRows, - async *rows () { + async *rows() { // Collect all rows /** @type {AsyncRow[]} */ const allRows = [] @@ -136,7 +136,7 @@ export function executeScalarAggregate(plan, context) { columns: selectColumnNames(plan.columns, child.columns), numRows: plan.having ? undefined : 1, maxRows: 1, - async *rows () { + async *rows() { // Collect all rows into single group /** @type {AsyncRow[]} */ const group = [] diff --git a/src/execute/execute.js b/src/execute/execute.js index 722853c..b73a3d5 100644 --- a/src/execute/execute.js +++ b/src/execute/execute.js @@ -10,7 +10,7 @@ import { executeSort } from './sort.js' import { addBounds, minBounds, stableRowKey } from './utils.js' /** - * @import { AsyncCells, AsyncDataSource, AsyncRow, ExecuteContext, ExecuteSqlOptions, ExprNode, QueryResults, SelectColumn, Statement } from '../types.js' + * @import { AsyncCells, AsyncDataSource, AsyncRow, DerivedColumn, ExecuteContext, ExecuteSqlOptions, ExprNode, IdentifierNode, QueryResults, SelectColumn, SqlPrimitive, Statement } from '../types.js' * @import { CountNode, DistinctNode, FilterNode, LimitNode, ProjectNode, QueryPlan, ScanNode, SetOperationNode } from '../plan/types.js' */ @@ -101,7 +101,7 @@ export function executePlan({ plan, context }) { } else if (plan.type === 'SetOperation') { return executeSetOperation(plan, context) } - return { columns: [], async *rows () {} } + return { columns: [], async *rows() {} } } /** @@ -155,7 +155,7 @@ function executeScan(plan, context) { columns: [column], numRows: scanRows, maxRows: scanRows, - async *rows () { + async *rows() { const columns = [column] for await (const chunk of chunks) { if (signal?.aborted) return @@ -185,7 +185,7 @@ function executeScan(plan, context) { columns: plan.hints.columns ?? table.columns, numRows: !plan.hints.where ? scanRows : undefined, maxRows: scanRows, - async *rows () { + async *rows() { let result = scanResult.rows() // Apply WHERE if data source did not @@ -218,7 +218,7 @@ function executeCount(plan, context) { columns: plan.columns.map(col => col.alias ?? derivedAlias(col.expr)), numRows: 1, maxRows: 1, - async *rows () { + async *rows() { // Use source numRows if available let count = table.numRows if (count === undefined) { @@ -366,8 +366,7 @@ function executeProject(plan, context) { /** @type {{ alias: string, sourceName: string }[] | undefined} */ let identifierMap if (!hasStar) { - /** @type {import('../types.js').DerivedColumn[]} */ - const derived = /** @type {any} */ (plan.columns) + const derived = /** @type {DerivedColumn[]} */ (plan.columns) staticColumns = derived.map(col => col.alias ?? derivedAlias(col.expr)) const allIdentifiers = derived.every(col => col.expr.type === 'identifier' && !col.expr.prefix @@ -375,7 +374,7 @@ function executeProject(plan, context) { if (allIdentifiers) { identifierMap = derived.map((col, i) => ({ alias: staticColumns[i], - sourceName: /** @type {import('../types.js').IdentifierNode} */ (col.expr).name, + sourceName: /** @type {IdentifierNode} */ (col.expr).name, })) } } @@ -384,7 +383,7 @@ function executeProject(plan, context) { columns: selectColumnNames(plan.columns, child.columns), numRows: child.numRows, maxRows: child.maxRows, - async *rows () { + async *rows() { let rowIndex = 0 let identifierMapValidated = false @@ -404,15 +403,15 @@ function executeProject(plan, context) { if (identifierMap) { /** @type {AsyncCells} */ const cells = {} - const srcData = row._data - /** @type {Record | undefined} */ - const _data = srcData ? {} : undefined + const source = row.resolved + /** @type {Record | undefined} */ + const resolved = source ? {} : undefined for (const { alias, sourceName } of identifierMap) { cells[alias] = row.cells[sourceName] - if (_data && srcData) _data[alias] = srcData[sourceName] + if (resolved && source) resolved[alias] = source[sourceName] } - yield _data - ? { columns: staticColumns, cells, _data } + yield resolved + ? { columns: staticColumns, cells, resolved } : { columns: staticColumns, cells } continue } @@ -420,7 +419,7 @@ function executeProject(plan, context) { const currentRowIndex = rowIndex /** @type {string[]} */ - const columns = staticColumns ? staticColumns : [] + const columns = staticColumns ?? [] /** @type {AsyncCells} */ const cells = {} @@ -465,7 +464,7 @@ function executeDistinct(plan, context) { return { columns: child.columns, maxRows: child.maxRows, - async *rows () { + async *rows() { const { signal } = context const MAX_CHUNK = 256 @@ -541,7 +540,7 @@ function executeSetOperation(plan, context) { columns: left.columns, numRows: addBounds(left.numRows, right.numRows), maxRows: addBounds(left.maxRows, right.maxRows), - async *rows () { + async *rows() { // UNION ALL: yield all rows from both sides yield* left.rows() yield* right.rows() @@ -553,7 +552,7 @@ function executeSetOperation(plan, context) { return { columns: left.columns, maxRows: addBounds(left.maxRows, right.maxRows), - async *rows () { + async *rows() { // UNION: yield deduplicated rows from both sides const seen = new Set() for await (const row of left.rows()) { @@ -581,7 +580,7 @@ function executeSetOperation(plan, context) { return { columns: left.columns, maxRows: minBounds(left.maxRows, right.maxRows), - async *rows () { + async *rows() { // Materialize right side keys /** @type {Map} */ const rightKeys = new Map() @@ -623,7 +622,7 @@ function executeSetOperation(plan, context) { return { columns: left.columns, maxRows: left.maxRows, - async *rows () { + async *rows() { // Materialize right side keys /** @type {Map} */ const rightKeys = new Map() diff --git a/src/execute/join.js b/src/execute/join.js index 4226f25..82e7b87 100644 --- a/src/execute/join.js +++ b/src/execute/join.js @@ -19,7 +19,7 @@ export function executeNestedLoopJoin(plan, context) { const right = executePlan({ plan: plan.right, context }) return { columns: mergeColumnNames(left.columns, right.columns, plan.leftAlias, plan.rightAlias), - async *rows () { + async *rows() { const leftTable = plan.leftAlias const rightTable = plan.rightAlias @@ -97,7 +97,7 @@ export function executePositionalJoin(plan, context) { columns: mergeColumnNames(left.columns, right.columns, plan.leftAlias, plan.rightAlias), numRows, maxRows: maxBounds(left.maxRows, right.maxRows), - async *rows () { + async *rows() { const { signal } = context const leftTable = plan.leftAlias const rightTable = plan.rightAlias @@ -143,7 +143,7 @@ export function executeHashJoin(plan, context) { const right = executePlan({ plan: plan.right, context }) return { columns: mergeColumnNames(left.columns, right.columns, plan.leftAlias, plan.rightAlias), - async *rows () { + async *rows() { const leftTable = plan.leftAlias const rightTable = plan.rightAlias diff --git a/src/execute/sort.js b/src/execute/sort.js index 1c193a0..eafe5a8 100644 --- a/src/execute/sort.js +++ b/src/execute/sort.js @@ -20,7 +20,7 @@ export function executeSort(plan, context) { columns: child.columns, numRows: child.numRows, maxRows: child.maxRows, - async *rows () { + async *rows() { // Buffer all rows /** @type {AsyncRow[]} */ const rows = [] diff --git a/src/execute/utils.js b/src/execute/utils.js index f69acd6..1302bdf 100644 --- a/src/execute/utils.js +++ b/src/execute/utils.js @@ -56,7 +56,7 @@ export async function collect(results) { // Fast path: if all rows have pre-materialized data, skip Promise overhead let allMaterialized = rows.length > 0 for (let i = 0; i < rows.length; i++) { - if (!rows[i]._data) { + if (!rows[i].resolved) { allMaterialized = false break } @@ -68,7 +68,7 @@ export async function collect(results) { /** @type {Record} */ const item = {} for (const col of row.columns) { - item[col] = row._data[col] + item[col] = row.resolved[col] } result[i] = item } diff --git a/src/parse/parse.js b/src/parse/parse.js index 4e709f7..b882ca5 100644 --- a/src/parse/parse.js +++ b/src/parse/parse.js @@ -10,26 +10,11 @@ import { tokenizeSql } from './tokenize.js' * @import { CTEDefinition, ExprNode, FromSubquery, FromTable, OrderByItem, ParseSqlOptions, ParserState, SelectColumn, SelectStatement, SetOperationStatement, SetOperator, Statement } from '../types.js' */ -const MAX_PARSE_CACHE = 64 -/** @type {Map} */ -const parseCache = new Map() - /** * @param {ParseSqlOptions} options * @returns {Statement} */ export function parseSql({ query, functions }) { - // Cache only for simple queries without custom functions - if (!functions) { - const cached = parseCache.get(query) - if (cached) { - // LRU touch - parseCache.delete(query) - parseCache.set(query, cached) - return cached - } - } - const tokens = tokenizeSql(query) /** @type {ParserState} */ const state = { tokens, pos: 0, lastPos: 0, functions } @@ -42,14 +27,6 @@ export function parseSql({ query, functions }) { throw parseError(state, 'end of query') } - if (!functions) { - parseCache.set(query, stmt) - if (parseCache.size > MAX_PARSE_CACHE) { - const oldest = parseCache.keys().next().value - if (oldest) parseCache.delete(oldest) - } - } - return stmt } diff --git a/src/plan/plan.js b/src/plan/plan.js index 8ee0ca3..5bb27f4 100644 --- a/src/plan/plan.js +++ b/src/plan/plan.js @@ -10,20 +10,6 @@ import { extractColumns, fromAlias, inferStatementColumns } from './columns.js' * @import { QueryPlan } from './types.d.ts' */ -/** @type {WeakMap} */ -const planCache = new WeakMap() - -/** - * @param {Record | undefined} tables - * @returns {string} - */ -function tablesKey(tables) { - if (!tables) return '' - const keys = Object.keys(tables) - keys.sort() - return keys.join(',') -} - /** * Builds a query plan from a statement AST. * Resolves CTEs at plan time so no planning occurs during execution. @@ -34,20 +20,7 @@ function tablesKey(tables) { export function planSql({ query, functions, tables }) { /** @type {Statement} */ const stmt = typeof query === 'string' ? parseSql({ query, functions }) : query - - const key = !functions ? tablesKey(tables) : undefined - if (key !== undefined) { - const cached = planCache.get(stmt) - if (cached && cached.tablesKey === key) return cached.plan - } - - const plan = planStatement({ stmt, tables }) - - if (key !== undefined) { - planCache.set(stmt, { plan, tablesKey: key }) - } - - return plan + return planStatement({ stmt, tables }) } /** diff --git a/src/types.d.ts b/src/types.d.ts index 4d39110..e2ede29 100644 --- a/src/types.d.ts +++ b/src/types.d.ts @@ -46,7 +46,9 @@ export interface ExecuteContext { export interface AsyncRow { columns: string[] cells: AsyncCells - _data?: Record + // Optional pre-materialized row values keyed by output column name. + // When present, consumers can skip the AsyncCell Promise roundtrip. + resolved?: Record } export type AsyncCells = Record export type AsyncCell = () => Promise From ad9d17aaefe6f0c5885194b95eba7004814a761b Mon Sep 17 00:00:00 2001 From: Phillip Cunliffe Date: Sun, 12 Apr 2026 13:47:53 -0700 Subject: [PATCH 4/9] Eagerly materialize row cells during sort buffering Resolves all cell values when rows are buffered for ORDER BY, replacing AsyncRow closures (which capture decompressed parquet row group data) with plain value-returning functions. The original closures become GC-eligible immediately. For tables with large text columns (~10KB/row), this reduces per-row buffer cost from ~10KB (closure over parquet data) to ~100B (plain value). --- src/execute/sort.js | 29 +++++++++++++++++++++++++++-- test/execute/expensive.test.js | 17 +++++++---------- 2 files changed, 34 insertions(+), 12 deletions(-) diff --git a/src/execute/sort.js b/src/execute/sort.js index eafe5a8..80af55c 100644 --- a/src/execute/sort.js +++ b/src/execute/sort.js @@ -7,6 +7,31 @@ import { compareForTerm } from './utils.js' * @import { SortNode } from '../plan/types.js' */ +/** + * Eagerly resolves all cell values in an AsyncRow, replacing closures with + * plain value-returning functions. This allows the original closures (which + * may capture large decompressed parquet data) to be garbage collected. + * + * @param {AsyncRow} row + * @returns {Promise} + */ +async function materializeRow(row) { + if (row.resolved) return row + const { columns } = row + /** @type {Record} */ + const resolved = {} + await Promise.all(columns.map(async col => { + resolved[col] = await row.cells[col]() + })) + /** @type {import('../types.js').AsyncCells} */ + const cells = {} + for (const col of columns) { + const val = resolved[col] + cells[col] = () => Promise.resolve(val) + } + return { columns, cells, resolved } +} + /** * Executes a sort operation (ORDER BY) * @@ -21,12 +46,12 @@ export function executeSort(plan, context) { numRows: child.numRows, maxRows: child.maxRows, async *rows() { - // Buffer all rows + // Buffer all rows, materializing cells to release closures over parquet data /** @type {AsyncRow[]} */ const rows = [] for await (const row of child.rows()) { if (context.signal?.aborted) return - rows.push(row) + rows.push(await materializeRow(row)) } if (rows.length === 0) return diff --git a/test/execute/expensive.test.js b/test/execute/expensive.test.js index 0d140c0..4d90b00 100644 --- a/test/execute/expensive.test.js +++ b/test/execute/expensive.test.js @@ -80,19 +80,17 @@ describe('expensive cell access', () => { }) it('should minimize expensive calls when limit + order by', async () => { + // Sort materializes all rows eagerly (releases closures over parquet data) await expect(countExpensiveCalls('SELECT * FROM data ORDER BY name DESC LIMIT 1')) - .resolves.toBe(1) + .resolves.toBe(5) }) it('should minimize expensive calls when sorting by multiple columns', async () => { - // ORDER BY cheap column, then expensive column - // Should only evaluate expensive column for rows that tie on cheap column - // With 5 unique names, no ties occur, so llm only evaluated for LIMIT rows + // Sort materializes all rows eagerly, so expensive column accessed once per row await expect(countExpensiveCalls('SELECT * FROM data ORDER BY name, llm LIMIT 1')) - .resolves.toBe(1) + .resolves.toBe(5) await expect(countExpensiveCalls('SELECT * FROM data ORDER BY name, llm LIMIT 2')) - .resolves.toBe(2) - // Without LIMIT, all rows need llm for final materialization + .resolves.toBe(5) await expect(countExpensiveCalls('SELECT * FROM data ORDER BY name, llm')) .resolves.toBe(5) }) @@ -140,9 +138,8 @@ describe('expensive cell access', () => { query: 'SELECT * FROM data ORDER BY name', })) - // With double-sorting bug: 15 accesses (2 sorts, 1 materialization) - // Without bug: 10 accesses (1 sort, 1 materialization) - expect(countingSource.getExpensiveCallCount()).toBe(10) + // Eager materialization during sort resolves each cell once: 5 rows × 1 expensive col = 5 + expect(countingSource.getExpensiveCallCount()).toBe(5) }) }) From ed8c9e1c1e2ccc7ee6f8115087676f4b3d908ec7 Mon Sep 17 00:00:00 2001 From: Phillip Cunliffe Date: Sun, 12 Apr 2026 13:51:26 -0700 Subject: [PATCH 5/9] Add TopN plan node for O(limit) ORDER BY + LIMIT MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fuses Sort + Limit into a TopN node that uses a bounded binary max-heap. ORDER BY x LIMIT N now buffers only N rows instead of the entire dataset. The planner detects two patterns: - Limit(Sort(child)) → TopN(child) - Limit(Project(Sort(child))) → Project(TopN(child)) --- src/execute/execute.js | 4 +- src/execute/sort.js | 95 ++++++++++++++++++- src/plan/plan.js | 14 +++ src/plan/types.d.ts | 8 ++ test/plan/plan.test.js | 209 +++++++++++++++++------------------------ 5 files changed, 203 insertions(+), 127 deletions(-) diff --git a/src/execute/execute.js b/src/execute/execute.js index b73a3d5..2243385 100644 --- a/src/execute/execute.js +++ b/src/execute/execute.js @@ -6,7 +6,7 @@ import { planSql } from '../plan/plan.js' import { validateScan, validateTable } from '../validation/tables.js' import { executeHashAggregate, executeScalarAggregate } from './aggregates.js' import { executeHashJoin, executeNestedLoopJoin, executePositionalJoin } from './join.js' -import { executeSort } from './sort.js' +import { executeSort, executeTopN } from './sort.js' import { addBounds, minBounds, stableRowKey } from './utils.js' /** @@ -94,6 +94,8 @@ export function executePlan({ plan, context }) { return executeScalarAggregate(plan, context) } else if (plan.type === 'Sort') { return executeSort(plan, context) + } else if (plan.type === 'TopN') { + return executeTopN(plan, context) } else if (plan.type === 'Distinct') { return executeDistinct(plan, context) } else if (plan.type === 'Limit') { diff --git a/src/execute/sort.js b/src/execute/sort.js index 80af55c..7aac189 100644 --- a/src/execute/sort.js +++ b/src/execute/sort.js @@ -4,7 +4,7 @@ import { compareForTerm } from './utils.js' /** * @import { AsyncRow, ExecuteContext, QueryResults, SqlPrimitive } from '../types.js' - * @import { SortNode } from '../plan/types.js' + * @import { SortNode, TopNNode } from '../plan/types.js' */ /** @@ -125,3 +125,96 @@ export function executeSort(plan, context) { }, } } + +/** + * Compares two entries by their sort keys across all ORDER BY terms. + * + * @param {SqlPrimitive[]} aKeys + * @param {SqlPrimitive[]} bKeys + * @param {import('../types.js').OrderByItem[]} orderBy + * @returns {number} + */ +function compareKeys(aKeys, bKeys, orderBy) { + for (let i = 0; i < orderBy.length; i++) { + const cmp = compareForTerm(aKeys[i], bKeys[i], orderBy[i]) + if (cmp !== 0) return cmp + } + return 0 +} + +/** + * Executes a TopN operation (ORDER BY + LIMIT fused) using a bounded heap. + * Memory usage is O(limit) instead of O(total rows). + * + * @param {TopNNode} plan + * @param {ExecuteContext} context + * @returns {QueryResults} + */ +export function executeTopN(plan, context) { + const child = executePlan({ plan: plan.child, context }) + return { + columns: child.columns, + numRows: Math.min(plan.limit, child.numRows ?? plan.limit), + maxRows: Math.min(plan.limit, child.maxRows ?? plan.limit), + async *rows() { + if (plan.limit <= 0) return + + // Bounded max-heap: heap[0] is the worst (largest for ASC) entry. + // When a new row is better than the worst, replace it. + /** @type {{ row: AsyncRow, keys: SqlPrimitive[] }[]} */ + const heap = [] + const limit = plan.limit + + function siftDown(i) { + const n = heap.length + while (true) { + let worst = i + const left = 2 * i + 1 + const right = 2 * i + 2 + if (left < n && compareKeys(heap[left].keys, heap[worst].keys, plan.orderBy) > 0) worst = left + if (right < n && compareKeys(heap[right].keys, heap[worst].keys, plan.orderBy) > 0) worst = right + if (worst === i) break + const tmp = heap[i] + heap[i] = heap[worst] + heap[worst] = tmp + i = worst + } + } + + function siftUp(i) { + while (i > 0) { + const parent = (i - 1) >> 1 + if (compareKeys(heap[i].keys, heap[parent].keys, plan.orderBy) <= 0) break + const tmp = heap[i] + heap[i] = heap[parent] + heap[parent] = tmp + i = parent + } + } + + for await (const row of child.rows()) { + if (context.signal?.aborted) return + + const keys = await Promise.all(plan.orderBy.map(term => + evaluateExpr({ node: term.expr, row, context }) + )) + + if (heap.length < limit) { + heap.push({ row: await materializeRow(row), keys }) + siftUp(heap.length - 1) + } else if (compareKeys(keys, heap[0].keys, plan.orderBy) < 0) { + // New row sorts before the worst in heap — replace it + heap[0] = { row: await materializeRow(row), keys } + siftDown(0) + } + // Otherwise discard — worse than everything in the heap + } + + // Extract in sorted order + const sorted = heap.sort((a, b) => compareKeys(a.keys, b.keys, plan.orderBy)) + for (const entry of sorted) { + yield entry.row + } + }, + } +} diff --git a/src/plan/plan.js b/src/plan/plan.js index 5bb27f4..9d795fc 100644 --- a/src/plan/plan.js +++ b/src/plan/plan.js @@ -245,6 +245,20 @@ function planSelect({ select, ctePlans, cteColumns, tables, parentColumns }) { } } + // Fuse Sort+Limit into TopN for O(limit) memory instead of O(n) + if (plan.type === 'Limit' && plan.limit !== undefined && !plan.offset) { + if (plan.child.type === 'Sort') { + plan = { type: 'TopN', limit: plan.limit, orderBy: plan.child.orderBy, child: plan.child.child } + } else if (plan.child.type === 'Project' && plan.child.child.type === 'Sort') { + const sort = plan.child.child + plan = { + type: 'Project', + columns: plan.child.columns, + child: { type: 'TopN', limit: plan.limit, orderBy: sort.orderBy, child: sort.child }, + } + } + } + return plan } diff --git a/src/plan/types.d.ts b/src/plan/types.d.ts index 9b63784..6db0e3f 100644 --- a/src/plan/types.d.ts +++ b/src/plan/types.d.ts @@ -6,6 +6,7 @@ export type QueryPlan = | FilterNode | ProjectNode | SortNode + | TopNNode | DistinctNode | LimitNode | HashAggregateNode @@ -48,6 +49,13 @@ export interface SortNode { child: QueryPlan } +export interface TopNNode { + type: 'TopN' + limit: number + orderBy: OrderByItem[] + child: QueryPlan +} + export interface DistinctNode { type: 'Distinct' child: QueryPlan diff --git a/test/plan/plan.test.js b/test/plan/plan.test.js index 3d4dede..df8be24 100644 --- a/test/plan/plan.test.js +++ b/test/plan/plan.test.js @@ -618,62 +618,60 @@ describe('planSql', () => { describe('complex queries', () => { it('plan for query with WHERE, ORDER BY, LIMIT', () => { const plan = planSql({ query: 'SELECT name FROM users WHERE age > 21 ORDER BY name LIMIT 10' }) + // TopN fuses Sort+Limit: Limit(Project(Sort(Scan))) → Project(TopN(Scan)) expect(plan).toEqual({ - type: 'Limit', - limit: 10, + type: 'Project', + columns: [ + { + type: 'derived', + expr: { + type: 'identifier', + name: 'name', + positionStart: 7, + positionEnd: 11, + }, + positionStart: 7, + positionEnd: 11, + }, + ], child: { - type: 'Project', - columns: [ + type: 'TopN', + limit: 10, + orderBy: [ { - type: 'derived', expr: { type: 'identifier', name: 'name', - positionStart: 7, - positionEnd: 11, + positionStart: 47, + positionEnd: 51, }, - positionStart: 7, - positionEnd: 11, + direction: 'ASC', + positionStart: 0, + positionEnd: 51, }, ], child: { - type: 'Sort', - orderBy: [ - { - expr: { + type: 'Scan', + table: 'users', + hints: { + columns: ['name', 'age'], + where: { + type: 'binary', + op: '>', + left: { type: 'identifier', - name: 'name', - positionStart: 47, - positionEnd: 51, - }, - direction: 'ASC', - positionStart: 0, - positionEnd: 51, - }, - ], - child: { - type: 'Scan', - table: 'users', - hints: { - columns: ['name', 'age'], - where: { - type: 'binary', - op: '>', - left: { - type: 'identifier', - name: 'age', - positionStart: 29, - positionEnd: 32, - }, - right: { - type: 'literal', - value: 21, - positionStart: 35, - positionEnd: 37, - }, + name: 'age', positionStart: 29, + positionEnd: 32, + }, + right: { + type: 'literal', + value: 21, + positionStart: 35, positionEnd: 37, }, + positionStart: 29, + positionEnd: 37, }, }, }, @@ -683,102 +681,63 @@ describe('planSql', () => { it('plan for grouped query with HAVING and ORDER BY', () => { const plan = planSql({ query: 'SELECT department, COUNT(*) as cnt FROM users GROUP BY department HAVING COUNT(*) > 5 ORDER BY cnt LIMIT 10' }) + // TopN fuses Sort+Limit: Limit(Sort(HashAggregate)) → TopN(HashAggregate) expect(plan).toEqual({ - type: 'Limit', + type: 'TopN', limit: 10, + orderBy: [{ + expr: { type: 'identifier', name: 'cnt', positionStart: 95, positionEnd: 98 }, + direction: 'ASC', + positionStart: 0, + positionEnd: 98, + }], child: { - type: 'Sort', - orderBy: [ + type: 'HashAggregate', + groupBy: [ + { type: 'identifier', name: 'department', positionStart: 55, positionEnd: 65 }, + ], + columns: [ { - expr: { - type: 'identifier', - name: 'cnt', - positionStart: 95, - positionEnd: 98, - }, - direction: 'ASC', - positionStart: 0, - positionEnd: 98, + type: 'derived', + expr: { type: 'identifier', name: 'department', positionStart: 7, positionEnd: 17 }, + positionStart: 7, + positionEnd: 17, }, - ], - child: { - type: 'HashAggregate', - groupBy: [ - { - type: 'identifier', - name: 'department', - positionStart: 55, - positionEnd: 65, - }, - ], - columns: [ - { - type: 'derived', - expr: { - type: 'identifier', - name: 'department', - positionStart: 7, - positionEnd: 17, - }, - positionStart: 7, - positionEnd: 17, - }, - { - type: 'derived', - expr: { - type: 'function', - funcName: 'COUNT', - args: [ - { - type: 'star', - positionStart: 25, - positionEnd: 26, - }, - ], - positionStart: 19, - positionEnd: 27, - }, - alias: 'cnt', - positionStart: 19, - positionEnd: 34, - }, - ], - having: { - type: 'binary', - op: '>', - left: { + { + type: 'derived', + expr: { type: 'function', funcName: 'COUNT', - args: [ - { - type: 'star', - positionStart: 79, - positionEnd: 80, - }, - ], - positionStart: 73, - positionEnd: 81, - }, - right: { - type: 'literal', - value: 5, - positionStart: 84, - positionEnd: 85, + args: [{ type: 'star', positionStart: 25, positionEnd: 26 }], + positionStart: 19, + positionEnd: 27, }, - positionStart: 73, - positionEnd: 85, + alias: 'cnt', + positionStart: 19, + positionEnd: 34, }, - child: { - type: 'Scan', - table: 'users', - hints: { - columns: ['department'], - }, + ], + having: { + type: 'binary', + op: '>', + left: { + type: 'function', + funcName: 'COUNT', + args: [{ type: 'star', positionStart: 79, positionEnd: 80 }], + positionStart: 73, + positionEnd: 81, }, + right: { type: 'literal', value: 5, positionStart: 84, positionEnd: 85 }, + positionStart: 73, + positionEnd: 85, + }, + child: { + type: 'Scan', + table: 'users', + hints: { columns: ['department'] }, }, }, }) }) }) - }) From de91774ea17393b9c0281526b313cf8a8e3b3b2e Mon Sep 17 00:00:00 2001 From: Phillip Cunliffe Date: Mon, 13 Apr 2026 12:17:09 -0700 Subject: [PATCH 6/9] Add missing JSDoc @param types for siftDown/siftUp Co-Authored-By: Claude Opus 4.6 (1M context) --- src/execute/sort.js | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/execute/sort.js b/src/execute/sort.js index 7aac189..59fc4e7 100644 --- a/src/execute/sort.js +++ b/src/execute/sort.js @@ -165,6 +165,7 @@ export function executeTopN(plan, context) { const heap = [] const limit = plan.limit + /** @param {number} i */ function siftDown(i) { const n = heap.length while (true) { @@ -181,6 +182,7 @@ export function executeTopN(plan, context) { } } + /** @param {number} i */ function siftUp(i) { while (i > 0) { const parent = (i - 1) >> 1 From 2660f383ad3a776cf47814b3e0354bdbfd36421c Mon Sep 17 00:00:00 2001 From: Kenny Daniel Date: Mon, 13 Apr 2026 13:37:08 -0700 Subject: [PATCH 7/9] Revert unnecessary table normalization optimization --- src/execute/execute.js | 25 ++++++------------------- 1 file changed, 6 insertions(+), 19 deletions(-) diff --git a/src/execute/execute.js b/src/execute/execute.js index 2065d33..68c4e15 100644 --- a/src/execute/execute.js +++ b/src/execute/execute.js @@ -25,27 +25,14 @@ export function executeSql({ tables, query, functions, signal }) { const parsed = typeof query === 'string' ? parseSql({ query, functions }) : query // Normalize tables: convert arrays to AsyncDataSource - // Fast path: skip normalization when no arrays are present - let needsNormalization = false - const tableKeys = Object.keys(tables) - for (let i = 0; i < tableKeys.length; i++) { - if (Array.isArray(tables[tableKeys[i]])) { - needsNormalization = true - break - } - } - /** @type {Record} */ - let normalizedTables - if (needsNormalization) { - normalizedTables = {} - for (let i = 0; i < tableKeys.length; i++) { - const name = tableKeys[i] - const data = tables[name] - normalizedTables[name] = Array.isArray(data) ? memorySource({ data }) : data + const normalizedTables = {} + for (const [name, data] of Object.entries(tables)) { + if (Array.isArray(data)) { + normalizedTables[name] = memorySource({ data }) + } else { + normalizedTables[name] = data } - } else { - normalizedTables = /** @type {Record} */ (tables) } const scope = statementScope(parsed) From d9712e8a429bc3904fa3d273b543497ea622e65b Mon Sep 17 00:00:00 2001 From: Kenny Daniel Date: Mon, 13 Apr 2026 17:10:41 -0700 Subject: [PATCH 8/9] Fix incorrect numRows on LIMIT when source numRows is unknown --- src/execute/sort.js | 32 ++++---- test/execute/execute.topn.test.js | 118 ++++++++++++++++++++++++++++++ test/execute/numRows.test.js | 38 ++++++++++ 3 files changed, 173 insertions(+), 15 deletions(-) create mode 100644 test/execute/execute.topn.test.js diff --git a/src/execute/sort.js b/src/execute/sort.js index 59fc4e7..8af4d83 100644 --- a/src/execute/sort.js +++ b/src/execute/sort.js @@ -3,7 +3,7 @@ import { executePlan } from './execute.js' import { compareForTerm } from './utils.js' /** - * @import { AsyncRow, ExecuteContext, QueryResults, SqlPrimitive } from '../types.js' + * @import { AsyncCells, AsyncRow, ExecuteContext, OrderByItem, QueryResults, SqlPrimitive } from '../types.js' * @import { SortNode, TopNNode } from '../plan/types.js' */ @@ -18,12 +18,12 @@ import { compareForTerm } from './utils.js' async function materializeRow(row) { if (row.resolved) return row const { columns } = row - /** @type {Record} */ + /** @type {Record} */ const resolved = {} await Promise.all(columns.map(async col => { resolved[col] = await row.cells[col]() })) - /** @type {import('../types.js').AsyncCells} */ + /** @type {AsyncCells} */ const cells = {} for (const col of columns) { const val = resolved[col] @@ -131,7 +131,7 @@ export function executeSort(plan, context) { * * @param {SqlPrimitive[]} aKeys * @param {SqlPrimitive[]} bKeys - * @param {import('../types.js').OrderByItem[]} orderBy + * @param {OrderByItem[]} orderBy * @returns {number} */ function compareKeys(aKeys, bKeys, orderBy) { @@ -152,18 +152,20 @@ function compareKeys(aKeys, bKeys, orderBy) { */ export function executeTopN(plan, context) { const child = executePlan({ plan: plan.child, context }) + const { limit, orderBy } = plan + const numRows = child.numRows !== undefined ? Math.min(limit, child.numRows) : undefined + const maxRows = Math.min(limit, child.maxRows ?? limit) return { columns: child.columns, - numRows: Math.min(plan.limit, child.numRows ?? plan.limit), - maxRows: Math.min(plan.limit, child.maxRows ?? plan.limit), + numRows, + maxRows, async *rows() { - if (plan.limit <= 0) return + if (limit <= 0) return // Bounded max-heap: heap[0] is the worst (largest for ASC) entry. // When a new row is better than the worst, replace it. /** @type {{ row: AsyncRow, keys: SqlPrimitive[] }[]} */ const heap = [] - const limit = plan.limit /** @param {number} i */ function siftDown(i) { @@ -172,8 +174,8 @@ export function executeTopN(plan, context) { let worst = i const left = 2 * i + 1 const right = 2 * i + 2 - if (left < n && compareKeys(heap[left].keys, heap[worst].keys, plan.orderBy) > 0) worst = left - if (right < n && compareKeys(heap[right].keys, heap[worst].keys, plan.orderBy) > 0) worst = right + if (left < n && compareKeys(heap[left].keys, heap[worst].keys, orderBy) > 0) worst = left + if (right < n && compareKeys(heap[right].keys, heap[worst].keys, orderBy) > 0) worst = right if (worst === i) break const tmp = heap[i] heap[i] = heap[worst] @@ -185,8 +187,8 @@ export function executeTopN(plan, context) { /** @param {number} i */ function siftUp(i) { while (i > 0) { - const parent = (i - 1) >> 1 - if (compareKeys(heap[i].keys, heap[parent].keys, plan.orderBy) <= 0) break + const parent = i - 1 >> 1 + if (compareKeys(heap[i].keys, heap[parent].keys, orderBy) <= 0) break const tmp = heap[i] heap[i] = heap[parent] heap[parent] = tmp @@ -197,14 +199,14 @@ export function executeTopN(plan, context) { for await (const row of child.rows()) { if (context.signal?.aborted) return - const keys = await Promise.all(plan.orderBy.map(term => + const keys = await Promise.all(orderBy.map(term => evaluateExpr({ node: term.expr, row, context }) )) if (heap.length < limit) { heap.push({ row: await materializeRow(row), keys }) siftUp(heap.length - 1) - } else if (compareKeys(keys, heap[0].keys, plan.orderBy) < 0) { + } else if (compareKeys(keys, heap[0].keys, orderBy) < 0) { // New row sorts before the worst in heap — replace it heap[0] = { row: await materializeRow(row), keys } siftDown(0) @@ -213,7 +215,7 @@ export function executeTopN(plan, context) { } // Extract in sorted order - const sorted = heap.sort((a, b) => compareKeys(a.keys, b.keys, plan.orderBy)) + const sorted = heap.sort((a, b) => compareKeys(a.keys, b.keys, orderBy)) for (const entry of sorted) { yield entry.row } diff --git a/test/execute/execute.topn.test.js b/test/execute/execute.topn.test.js new file mode 100644 index 0000000..8063fcc --- /dev/null +++ b/test/execute/execute.topn.test.js @@ -0,0 +1,118 @@ +import { describe, expect, it } from 'vitest' +import { memorySource } from '../../src/backend/dataSource.js' +import { collect, executeSql } from '../../src/index.js' +import { planSql } from '../../src/plan/plan.js' + +describe('TopN (ORDER BY + LIMIT fusion)', () => { + const users = memorySource({ data: [ + { id: 1, name: 'Alice', age: 30 }, + { id: 2, name: 'Bob', age: 25 }, + { id: 3, name: 'Charlie', age: 35 }, + { id: 4, name: 'Diana', age: 28 }, + { id: 5, name: 'Eve', age: 30 }, + ] }) + + it('should plan Sort+Limit as TopN', () => { + const plan = planSql({ query: 'SELECT * FROM users ORDER BY age LIMIT 3', tables: { users } }) + // SELECT * has no Project wrapper, so TopN is the top of the plan + expect(plan.type).toBe('TopN') + if (plan.type !== 'TopN') return + expect(plan.limit).toBe(3) + expect(plan.child.type).toBe('Scan') + }) + + it('should not fuse when OFFSET is present', () => { + const plan = planSql({ query: 'SELECT * FROM users ORDER BY age LIMIT 3 OFFSET 1', tables: { users } }) + expect(plan.type).toBe('Limit') + if (plan.type !== 'Limit') return + expect(plan.child.type).toBe('Sort') + }) + + it('should return fewer rows when limit exceeds input', async () => { + const result = await collect(executeSql({ tables: { users }, query: 'SELECT * FROM users ORDER BY age LIMIT 100' })) + expect(result).toHaveLength(5) + expect(result.map(r => r.age)).toEqual([25, 28, 30, 30, 35]) + }) + + it('should return empty for LIMIT 0', async () => { + const result = await collect(executeSql({ tables: { users }, query: 'SELECT * FROM users ORDER BY age LIMIT 0' })) + expect(result).toEqual([]) + }) + + it('should return top N ascending', async () => { + const result = await collect(executeSql({ tables: { users }, query: 'SELECT name, age FROM users ORDER BY age LIMIT 2' })) + expect(result.map(r => r.age)).toEqual([25, 28]) + expect(result[0].name).toBe('Bob') + expect(result[1].name).toBe('Diana') + }) + + it('should return top N descending', async () => { + const result = await collect(executeSql({ tables: { users }, query: 'SELECT name, age FROM users ORDER BY age DESC LIMIT 2' })) + expect(result.map(r => r.age)).toEqual([35, 30]) + expect(result[0].name).toBe('Charlie') + }) + + it('should match Sort+Limit exactly on boundary ties (LIMIT hits middle of tie group)', async () => { + // ages: 25, 28, 30, 30, 35 — LIMIT 3 forces a tie-break at age=30 + const tables = { users } + const topn = await collect(executeSql({ tables, query: 'SELECT id, age FROM users ORDER BY age LIMIT 3' })) + expect(topn).toHaveLength(3) + expect(topn.map(r => r.age)).toEqual([25, 28, 30]) + // The exact id for the third row may be 1 or 5 (both age 30); assert it's one of them. + expect([1, 5]).toContain(topn[2].id) + }) + + it('should sort by multiple columns with mixed directions', async () => { + const data = [ + { a: 1, b: 'y' }, + { a: 2, b: 'x' }, + { a: 1, b: 'x' }, + { a: 2, b: 'y' }, + { a: 1, b: 'z' }, + ] + const result = await collect(executeSql({ tables: { data }, query: 'SELECT * FROM data ORDER BY a ASC, b DESC LIMIT 3' })) + expect(result).toEqual([ + { a: 1, b: 'z' }, + { a: 1, b: 'y' }, + { a: 1, b: 'x' }, + ]) + }) + + it('should handle NULLs in sort keys', async () => { + const data = [ + { id: 1, v: 10 }, + { id: 2, v: null }, + { id: 3, v: 5 }, + { id: 4, v: null }, + { id: 5, v: 20 }, + ] + // NULLs sort first (NULLS FIRST is the default, regardless of direction) + const asc = await collect(executeSql({ tables: { data }, query: 'SELECT * FROM data ORDER BY v LIMIT 3' })) + expect(asc.map(r => r.v)).toEqual([null, null, 5]) + + const desc = await collect(executeSql({ tables: { data }, query: 'SELECT * FROM data ORDER BY v DESC LIMIT 3' })) + expect(desc.map(r => r.v)).toEqual([null, null, 20]) + }) + + it('should produce identical results to Sort+Limit for a larger random input', async () => { + // Generate deterministic pseudo-random data + const data = [] + let seed = 42 + for (let i = 0; i < 200; i++) { + seed = seed * 1103515245 + 12345 & 0x7fffffff + data.push({ id: i, k: seed % 50 }) + } + const topn = await collect(executeSql({ tables: { data }, query: 'SELECT k FROM data ORDER BY k LIMIT 10' })) + // Reference: full JS sort + const reference = [...data].sort((a, b) => a.k - b.k).slice(0, 10).map(r => ({ k: r.k })) + expect(topn).toEqual(reference) + }) + + it('should project after TopN when Limit(Project(Sort)) pattern applies', async () => { + const result = await collect(executeSql({ tables: { users }, query: 'SELECT name FROM users ORDER BY age DESC LIMIT 2' })) + expect(result).toHaveLength(2) + expect(result[0]).toEqual({ name: 'Charlie' }) + // Second is age-30 tie between Alice and Eve; TopN is not stable. + expect(['Alice', 'Eve']).toContain(result[1].name) + }) +}) diff --git a/test/execute/numRows.test.js b/test/execute/numRows.test.js index 6f1b107..d99278b 100644 --- a/test/execute/numRows.test.js +++ b/test/execute/numRows.test.js @@ -201,6 +201,44 @@ describe('numRows and maxRows', () => { }) }) + describe('topn', () => { + it('should compute numRows and maxRows for ORDER BY + LIMIT', () => { + const result = executeSql({ + tables: { users: memorySource({ data: users }) }, + query: 'SELECT * FROM users ORDER BY age LIMIT 2', + }) + expect(result.numRows).toBe(2) + expect(result.maxRows).toBe(2) + }) + + it('should cap numRows at child count when LIMIT exceeds input', () => { + const result = executeSql({ + tables: { users: memorySource({ data: users }) }, + query: 'SELECT * FROM users ORDER BY age LIMIT 100', + }) + expect(result.numRows).toBe(3) + expect(result.maxRows).toBe(3) + }) + + it('should leave numRows undefined when child numRows is unknown (WHERE + ORDER BY + LIMIT)', () => { + const result = executeSql({ + tables: { users: memorySource({ data: users }) }, + query: 'SELECT * FROM users WHERE age > 25 ORDER BY age LIMIT 2', + }) + expect(result.numRows).toBeUndefined() + expect(result.maxRows).toBe(2) + }) + + it('should leave numRows undefined when data source lacks numRows', () => { + const result = executeSql({ + tables: { data: noNumRowsSource }, + query: 'SELECT * FROM data ORDER BY x LIMIT 5', + }) + expect(result.numRows).toBeUndefined() + expect(result.maxRows).toBe(5) + }) + }) + describe('union', () => { it('should return numRows and maxRows for UNION ALL', () => { const result = executeSql({ From 8fcefd63d364ab05b8ea56bdca2ae9f134ca1298 Mon Sep 17 00:00:00 2001 From: Kenny Daniel Date: Mon, 13 Apr 2026 17:37:19 -0700 Subject: [PATCH 9/9] Restore late materialization on sorting and topN --- src/execute/sort.js | 151 +++++++++++++++++++++------------ test/execute/expensive.test.js | 17 ++-- 2 files changed, 107 insertions(+), 61 deletions(-) diff --git a/src/execute/sort.js b/src/execute/sort.js index 8af4d83..6c8e164 100644 --- a/src/execute/sort.js +++ b/src/execute/sort.js @@ -7,31 +7,6 @@ import { compareForTerm } from './utils.js' * @import { SortNode, TopNNode } from '../plan/types.js' */ -/** - * Eagerly resolves all cell values in an AsyncRow, replacing closures with - * plain value-returning functions. This allows the original closures (which - * may capture large decompressed parquet data) to be garbage collected. - * - * @param {AsyncRow} row - * @returns {Promise} - */ -async function materializeRow(row) { - if (row.resolved) return row - const { columns } = row - /** @type {Record} */ - const resolved = {} - await Promise.all(columns.map(async col => { - resolved[col] = await row.cells[col]() - })) - /** @type {AsyncCells} */ - const cells = {} - for (const col of columns) { - const val = resolved[col] - cells[col] = () => Promise.resolve(val) - } - return { columns, cells, resolved } -} - /** * Executes a sort operation (ORDER BY) * @@ -46,12 +21,12 @@ export function executeSort(plan, context) { numRows: child.numRows, maxRows: child.maxRows, async *rows() { - // Buffer all rows, materializing cells to release closures over parquet data + // Buffer all rows (cells stay lazy — see multi-pass below) /** @type {AsyncRow[]} */ const rows = [] for await (const row of child.rows()) { if (context.signal?.aborted) return - rows.push(await materializeRow(row)) + rows.push(row) } if (rows.length === 0) return @@ -127,24 +102,80 @@ export function executeSort(plan, context) { } /** - * Compares two entries by their sort keys across all ORDER BY terms. + * @typedef {{ row: AsyncRow, keys: SqlPrimitive[] }} HeapEntry + * `keys` grows lazily: keys[i] is populated only when the i-th ORDER BY term + * is actually needed for a comparison involving this entry. + */ + +/** + * Resolves the i-th sort key for a heap entry, memoizing it on the entry. + * Fills any earlier unresolved positions to keep keys.length === resolved count. * - * @param {SqlPrimitive[]} aKeys - * @param {SqlPrimitive[]} bKeys + * @param {HeapEntry} entry + * @param {number} i * @param {OrderByItem[]} orderBy - * @returns {number} + * @param {ExecuteContext} context + * @returns {Promise} */ -function compareKeys(aKeys, bKeys, orderBy) { +async function resolveKey(entry, i, orderBy, context) { + while (entry.keys.length <= i) { + const idx = entry.keys.length + entry.keys.push(await evaluateExpr({ node: orderBy[idx].expr, row: entry.row, context })) + } + return entry.keys[i] +} + +/** + * Compares two heap entries lazily across ORDER BY terms: resolves the i-th + * key for each entry only when earlier terms have tied. Already-resolved keys + * are reused via the entry's `keys` cache. + * + * @param {HeapEntry} a + * @param {HeapEntry} b + * @param {OrderByItem[]} orderBy + * @param {ExecuteContext} context + * @returns {Promise} + */ +async function compareLazy(a, b, orderBy, context) { for (let i = 0; i < orderBy.length; i++) { - const cmp = compareForTerm(aKeys[i], bKeys[i], orderBy[i]) + const av = await resolveKey(a, i, orderBy, context) + const bv = await resolveKey(b, i, orderBy, context) + const cmp = compareForTerm(av, bv, orderBy[i]) if (cmp !== 0) return cmp } return 0 } +/** + * Splices already-resolved sort keys back into the row's cells so downstream + * consumers reading the sort-key columns don't re-evaluate them. Only safe + * for identifier terms whose name is an output column of the row. + * + * @param {AsyncRow} row + * @param {OrderByItem[]} orderBy + * @param {SqlPrimitive[]} keys + * @returns {AsyncRow} + */ +function withResolvedKeys(row, orderBy, keys) { + /** @type {AsyncCells | undefined} */ + let cells + for (let i = 0; i < orderBy.length && i < keys.length; i++) { + const { expr } = orderBy[i] + if (expr.type === 'identifier' && row.columns.includes(expr.name)) { + if (!cells) cells = { ...row.cells } + const val = keys[i] + cells[expr.name] = () => Promise.resolve(val) + } + } + return cells ? { columns: row.columns, cells } : row +} + /** * Executes a TopN operation (ORDER BY + LIMIT fused) using a bounded heap. - * Memory usage is O(limit) instead of O(total rows). + * Memory usage is O(limit) instead of O(total rows). Sort keys are evaluated + * lazily per-entry, so multi-column ORDER BY only pays for later terms when + * earlier terms tie. Non-sort cells are never materialized by TopN — they + * stay lazy for the downstream consumer. * * @param {TopNNode} plan * @param {ExecuteContext} context @@ -162,20 +193,20 @@ export function executeTopN(plan, context) { async *rows() { if (limit <= 0) return - // Bounded max-heap: heap[0] is the worst (largest for ASC) entry. - // When a new row is better than the worst, replace it. - /** @type {{ row: AsyncRow, keys: SqlPrimitive[] }[]} */ + // Bounded max-heap: heap[0] is the worst entry (largest for ASC). + // When a new row beats the worst, replace it. + /** @type {HeapEntry[]} */ const heap = [] /** @param {number} i */ - function siftDown(i) { + async function siftDown(i) { const n = heap.length while (true) { let worst = i const left = 2 * i + 1 const right = 2 * i + 2 - if (left < n && compareKeys(heap[left].keys, heap[worst].keys, orderBy) > 0) worst = left - if (right < n && compareKeys(heap[right].keys, heap[worst].keys, orderBy) > 0) worst = right + if (left < n && await compareLazy(heap[left], heap[worst], orderBy, context) > 0) worst = left + if (right < n && await compareLazy(heap[right], heap[worst], orderBy, context) > 0) worst = right if (worst === i) break const tmp = heap[i] heap[i] = heap[worst] @@ -185,10 +216,10 @@ export function executeTopN(plan, context) { } /** @param {number} i */ - function siftUp(i) { + async function siftUp(i) { while (i > 0) { const parent = i - 1 >> 1 - if (compareKeys(heap[i].keys, heap[parent].keys, orderBy) <= 0) break + if (await compareLazy(heap[i], heap[parent], orderBy, context) <= 0) break const tmp = heap[i] heap[i] = heap[parent] heap[parent] = tmp @@ -199,25 +230,37 @@ export function executeTopN(plan, context) { for await (const row of child.rows()) { if (context.signal?.aborted) return - const keys = await Promise.all(orderBy.map(term => - evaluateExpr({ node: term.expr, row, context }) - )) + /** @type {HeapEntry} */ + const entry = { row, keys: [] } if (heap.length < limit) { - heap.push({ row: await materializeRow(row), keys }) - siftUp(heap.length - 1) - } else if (compareKeys(keys, heap[0].keys, orderBy) < 0) { + heap.push(entry) + await siftUp(heap.length - 1) + } else if (await compareLazy(entry, heap[0], orderBy, context) < 0) { // New row sorts before the worst in heap — replace it - heap[0] = { row: await materializeRow(row), keys } - siftDown(0) + heap[0] = entry + await siftDown(0) } // Otherwise discard — worse than everything in the heap } - // Extract in sorted order - const sorted = heap.sort((a, b) => compareKeys(a.keys, b.keys, orderBy)) - for (const entry of sorted) { - yield entry.row + // Final sort of survivors. Resolve any keys still missing so we can + // use a synchronous comparator. + for (const entry of heap) { + for (let i = 0; i < orderBy.length; i++) { + await resolveKey(entry, i, orderBy, context) + } + } + heap.sort((a, b) => { + for (let i = 0; i < orderBy.length; i++) { + const cmp = compareForTerm(a.keys[i], b.keys[i], orderBy[i]) + if (cmp !== 0) return cmp + } + return 0 + }) + + for (const entry of heap) { + yield withResolvedKeys(entry.row, orderBy, entry.keys) } }, } diff --git a/test/execute/expensive.test.js b/test/execute/expensive.test.js index 4d90b00..0d140c0 100644 --- a/test/execute/expensive.test.js +++ b/test/execute/expensive.test.js @@ -80,17 +80,19 @@ describe('expensive cell access', () => { }) it('should minimize expensive calls when limit + order by', async () => { - // Sort materializes all rows eagerly (releases closures over parquet data) await expect(countExpensiveCalls('SELECT * FROM data ORDER BY name DESC LIMIT 1')) - .resolves.toBe(5) + .resolves.toBe(1) }) it('should minimize expensive calls when sorting by multiple columns', async () => { - // Sort materializes all rows eagerly, so expensive column accessed once per row + // ORDER BY cheap column, then expensive column + // Should only evaluate expensive column for rows that tie on cheap column + // With 5 unique names, no ties occur, so llm only evaluated for LIMIT rows await expect(countExpensiveCalls('SELECT * FROM data ORDER BY name, llm LIMIT 1')) - .resolves.toBe(5) + .resolves.toBe(1) await expect(countExpensiveCalls('SELECT * FROM data ORDER BY name, llm LIMIT 2')) - .resolves.toBe(5) + .resolves.toBe(2) + // Without LIMIT, all rows need llm for final materialization await expect(countExpensiveCalls('SELECT * FROM data ORDER BY name, llm')) .resolves.toBe(5) }) @@ -138,8 +140,9 @@ describe('expensive cell access', () => { query: 'SELECT * FROM data ORDER BY name', })) - // Eager materialization during sort resolves each cell once: 5 rows × 1 expensive col = 5 - expect(countingSource.getExpensiveCallCount()).toBe(5) + // With double-sorting bug: 15 accesses (2 sorts, 1 materialization) + // Without bug: 10 accesses (1 sort, 1 materialization) + expect(countingSource.getExpensiveCallCount()).toBe(10) }) })