diff --git a/src/execute/execute.js b/src/execute/execute.js index eb40f86..68c4e15 100644 --- a/src/execute/execute.js +++ b/src/execute/execute.js @@ -7,7 +7,7 @@ import { fromAlias } from '../plan/columns.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' /** @@ -98,6 +98,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 eafe5a8..6c8e164 100644 --- a/src/execute/sort.js +++ b/src/execute/sort.js @@ -3,8 +3,8 @@ import { executePlan } from './execute.js' import { compareForTerm } from './utils.js' /** - * @import { AsyncRow, ExecuteContext, QueryResults, SqlPrimitive } from '../types.js' - * @import { SortNode } from '../plan/types.js' + * @import { AsyncCells, AsyncRow, ExecuteContext, OrderByItem, QueryResults, SqlPrimitive } from '../types.js' + * @import { SortNode, TopNNode } from '../plan/types.js' */ /** @@ -21,7 +21,7 @@ export function executeSort(plan, context) { numRows: child.numRows, maxRows: child.maxRows, async *rows() { - // Buffer all rows + // Buffer all rows (cells stay lazy — see multi-pass below) /** @type {AsyncRow[]} */ const rows = [] for await (const row of child.rows()) { @@ -100,3 +100,168 @@ export function executeSort(plan, context) { }, } } + +/** + * @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 {HeapEntry} entry + * @param {number} i + * @param {OrderByItem[]} orderBy + * @param {ExecuteContext} context + * @returns {Promise} + */ +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 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). 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 + * @returns {QueryResults} + */ +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, + maxRows, + async *rows() { + if (limit <= 0) return + + // 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 */ + 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 && 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] + heap[worst] = tmp + i = worst + } + } + + /** @param {number} i */ + async function siftUp(i) { + while (i > 0) { + const parent = i - 1 >> 1 + if (await compareLazy(heap[i], heap[parent], orderBy, context) <= 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 + + /** @type {HeapEntry} */ + const entry = { row, keys: [] } + + if (heap.length < limit) { + 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] = entry + await siftDown(0) + } + // Otherwise discard — worse than everything in the heap + } + + // 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/src/plan/plan.js b/src/plan/plan.js index f03c2fd..0b7550c 100644 --- a/src/plan/plan.js +++ b/src/plan/plan.js @@ -253,6 +253,20 @@ function planSelect({ select, ctePlans, cteColumns, tables, parentColumns, outer } } + // 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/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({ 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'] }, }, }, }) }) }) - })