Add TopN plan node for O(limit) ORDER BY + LIMIT

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') {

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<SqlPrimitive>}
+ */
+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<number>}
+ */
+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)
+      }
+    },
+  }
+}

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
 }
 

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

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)
+  })
+})