squirreling 0.5.0 → 0.6.1

package/README.md

@@ -40,7 +40,10 @@ const users = [
   // ...more rows
 ]
 
-type AsyncRow = Record<string, AsyncCell>
+interface AsyncRow {
+  columns: string[]
+  cells: Record<string, AsyncCell>
+}
 type AsyncCell = () => Promise<SqlPrimitive>
 
 // Returns an async iterable of rows with async cells
@@ -76,6 +79,7 @@ console.log(allUsers)
 - `GROUP BY` and `HAVING` clauses
 - Aggregate functions: `COUNT`, `SUM`, `AVG`, `MIN`, `MAX`, `JSON_ARRAYAGG`
 - String functions: `CONCAT`, `SUBSTRING`, `LENGTH`, `UPPER`, `LOWER`
+- Math functions: `ABS`, `CEIL`, `FLOOR`, `ROUND`, `MOD`, `RAND`, `RANDOM`, `LN`, `LOG10`, `EXP`, `POWER`, `SQRT`, `SIN`, `COS`, `TAN`, `COT`, `ASIN`, `ACOS`, `ATAN`, `ATAN2`, `DEGREES`, `RADIANS`, `PI`
 - Date functions: `CURRENT_DATE`, `CURRENT_TIME`, `CURRENT_TIMESTAMP`, `INTERVAL`
 - Json functions: `JSON_VALUE`, `JSON_QUERY`, `JSON_OBJECT`
 - Basic expressions and arithmetic operations

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "squirreling",
-  "version": "0.5.0",
+  "version": "0.6.1",
   "description": "Squirreling SQL Engine",
   "author": "Hyperparam",
   "homepage": "https://hyperparam.app",
@@ -37,11 +37,11 @@
     "test": "vitest run"
   },
   "devDependencies": {
-    "@types/node": "24.10.2",
-    "@vitest/coverage-v8": "4.0.15",
-    "eslint": "9.39.1",
+    "@types/node": "24.10.4",
+    "@vitest/coverage-v8": "4.0.16",
+    "eslint": "9.39.2",
     "eslint-plugin-jsdoc": "61.5.0",
     "typescript": "5.9.3",
-    "vitest": "4.0.15"
+    "vitest": "4.0.16"
   }
 }

package/src/backend/dataSource.js

@@ -1,8 +1,7 @@
 /**
- * @import { AsyncDataSource, AsyncRow, SqlPrimitive } from '../types.js'
+ * @import { AsyncCell, AsyncCells, AsyncDataSource, AsyncRow, ScanOptions, SqlPrimitive } from '../types.js'
  */
 
-
 /**
  * Wraps an async generator of plain objects into an AsyncDataSource
  *
@@ -11,8 +10,11 @@
  */
 export function generatorSource(gen) {
   return {
-    async *getRows() {
-      yield* gen
+    async *scan({ signal }) {
+      for await (const row of gen) {
+        if (signal?.aborted) break
+        yield row
+      }
     },
   }
 }
@@ -24,12 +26,12 @@ export function generatorSource(gen) {
  * @returns {AsyncRow} a row accessor interface
  */
 function asyncRow(obj) {
-  /** @type {AsyncRow} */
-  const row = {}
+  /** @type {AsyncCells} */
+  const cells = {}
   for (const [key, value] of Object.entries(obj)) {
-    row[key] = () => Promise.resolve(value)
+    cells[key] = () => Promise.resolve(value)
   }
-  return row
+  return { columns: Object.keys(obj), cells }
 }
 
 /**
@@ -40,8 +42,9 @@ function asyncRow(obj) {
  */
 export function memorySource(data) {
   return {
-    async *getRows() {
+    async *scan({ signal }) {
       for (const item of data) {
+        if (signal?.aborted) break
         yield asyncRow(item)
       }
     },
@@ -58,17 +61,21 @@ export function cachedDataSource(source) {
   const cache = new Map()
   return {
     /**
+     * @param {ScanOptions} options
      * @yields {AsyncRow}
      */
-    async *getRows() {
+    async *scan(options) {
+      const { signal } = options
       let index = 0
-      for await (const row of source.getRows()) {
+      for await (const row of source.scan(options)) {
+        if (signal?.aborted) break
         const rowIndex = index
-        /** @type {AsyncRow} */
-        const out = {}
-        for (const [key, cell] of Object.entries(row)) {
+        /** @type {AsyncCells} */
+        const cells = {}
+        for (const key of row.columns) {
+          const cell = row.cells[key]
           // Wrap the cell to cache accesses
-          out[key] = () => {
+          cells[key] = () => {
             const cacheKey = `${rowIndex}:${key}`
             let value = cache.get(cacheKey)
             if (!value) {
@@ -78,7 +85,7 @@ export function cachedDataSource(source) {
             return value
           }
         }
-        yield out
+        yield { columns: row.columns, cells }
         index++
       }
     },

package/src/execute/columns.js

@@ -1,7 +1,46 @@
+import { isAggregateFunc } from '../validation.js'
+
 /**
  * @import { ExprNode, SelectStatement, SelectColumn } from '../types.js'
  */
 
+/**
+ * Checks if an expression contains any aggregate function calls
+ *
+ * @param {ExprNode | undefined} expr
+ * @returns {boolean}
+ */
+export function containsAggregate(expr) {
+  if (!expr) return false
+  if (expr.type === 'function' && isAggregateFunc(expr.name.toUpperCase())) {
+    return true
+  }
+  if (expr.type === 'binary') {
+    return containsAggregate(expr.left) || containsAggregate(expr.right)
+  }
+  if (expr.type === 'unary') {
+    return containsAggregate(expr.argument)
+  }
+  if (expr.type === 'cast') {
+    return containsAggregate(expr.expr)
+  }
+  if (expr.type === 'case') {
+    if (expr.caseExpr && containsAggregate(expr.caseExpr)) return true
+    for (const when of expr.whenClauses) {
+      if (containsAggregate(when.condition) || containsAggregate(when.result)) return true
+    }
+    if (containsAggregate(expr.elseResult)) return true
+  }
+  if (expr.type === 'in valuelist') {
+    if (containsAggregate(expr.expr)) return true
+    for (const val of expr.values) {
+      if (containsAggregate(val)) return true
+    }
+  }
+  // Note: Don't recurse into subqueries - they have their own aggregate scope
+  return false
+}
+
 /**
  * Extracts column names needed from a SELECT statement.
  *
@@ -50,11 +89,6 @@ export function extractColumns(select) {
 function collectColumnsFromSelectColumn(col, columns) {
   if (col.kind === 'derived') {
     collectColumnsFromExpr(col.expr, columns)
-  } else if (col.kind === 'aggregate') {
-    if (col.arg.kind === 'expression') {
-      collectColumnsFromExpr(col.arg.expr, columns)
-    }
-    // 'star' aggregate (COUNT(*)) doesn't reference specific columns
   }
   // 'star' columns handled separately (returns undefined for all columns)
 }

package/src/execute/execute.js

@@ -2,15 +2,14 @@ import { missingClauseError } from '../parseErrors.js'
 import { tableNotFoundError, unsupportedOperationError } from '../executionErrors.js'
 import { generatorSource, memorySource } from '../backend/dataSource.js'
 import { parseSql } from '../parse/parse.js'
-import { defaultAggregateAlias, evaluateAggregate } from './aggregates.js'
-import { extractColumns } from './columns.js'
+import { containsAggregate, extractColumns } from './columns.js'
 import { evaluateExpr } from './expression.js'
 import { evaluateHavingExpr } from './having.js'
 import { executeJoins } from './join.js'
 import { compareForTerm, defaultDerivedAlias, stringify } from './utils.js'
 
 /**
- * @import { AsyncDataSource, AsyncRow, ExecuteSqlOptions, OrderByItem, QueryHints, SelectStatement, SqlPrimitive } from '../types.js'
+ * @import { AsyncCells, AsyncDataSource, AsyncRow, ExecuteSqlOptions, OrderByItem, QueryHints, SelectStatement, SqlPrimitive } from '../types.js'
  */
 
 /**
@@ -19,8 +18,8 @@ import { compareForTerm, defaultDerivedAlias, stringify } from './utils.js'
  * @param {ExecuteSqlOptions} options - the execution options
  * @yields {AsyncRow} async generator yielding result rows
  */
-export async function* executeSql({ tables, query }) {
-  const select = parseSql(query)
+export async function* executeSql({ tables, query, signal }) {
+  const select = typeof query === 'string' ? parseSql(query) : query
 
   // Check for unsupported operations
   if (!select.from) {
@@ -41,17 +40,23 @@ export async function* executeSql({ tables, query }) {
     }
   }
 
-  yield* executeSelect(select, normalizedTables)
+  yield* executeSelect({ select, tables: normalizedTables, signal })
 }
 
+/**
+ * @typedef {Object} ExecuteSelectOptions
+ * @property {SelectStatement} select
+ * @property {Record<string, AsyncDataSource>} tables
+ * @property {AbortSignal} [signal]
+ */
+
 /**
  * Executes a SELECT query against the provided tables
  *
- * @param {SelectStatement} select
- * @param {Record<string, AsyncDataSource>} tables
+ * @param {ExecuteSelectOptions} options
  * @yields {AsyncRow}
  */
-export async function* executeSelect(select, tables) {
+export async function* executeSelect({ select, tables, signal }) {
   /** @type {AsyncDataSource} */
   let dataSource
   /** @type {string} */
@@ -67,7 +72,7 @@ export async function* executeSelect(select, tables) {
   } else {
     // Nested subquery - recursively resolve
     fromTableName = select.from.alias
-    dataSource = generatorSource(executeSelect(select.from.query, tables))
+    dataSource = generatorSource(executeSelect({ select: select.from.query, tables, signal }))
   }
 
   // Execute JOINs if present
@@ -75,21 +80,21 @@ export async function* executeSelect(select, tables) {
     dataSource = await executeJoins(dataSource, select.joins, fromTableName, tables)
   }
 
-  yield* evaluateSelectAst(select, dataSource, tables)
+  yield* evaluateSelectAst({ select, dataSource, tables, signal })
 }
 
 /**
  * Creates a stable string key for a row to enable deduplication
  *
- * @param {AsyncRow} row
+ * @param {AsyncCells} cells
  * @returns {Promise<string>} a stable string representation of the row
  */
-async function stableRowKey(row) {
-  const keys = Object.keys(row).sort()
+async function stableRowKey(cells) {
+  const keys = Object.keys(cells).sort()
   /** @type {string[]} */
   const parts = []
   for (const k of keys) {
-    const v = await row[k]()
+    const v = await cells[k]()
     parts.push(k + ':' + stringify(v))
   }
   return parts.join('|')
@@ -109,7 +114,7 @@ async function applyDistinct(rows, distinct) {
   /** @type {AsyncRow[]} */
   const result = []
   for (const row of rows) {
-    const key = await stableRowKey(row)
+    const key = await stableRowKey(row.cells)
     if (seen.has(key)) continue
     seen.add(key)
     result.push(row)
@@ -201,40 +206,52 @@ async function sortRows(rows, orderBy, tables) {
   return groups.flat().map(i => rows[i])
 }
 
+/**
+ * @typedef {Object} EvaluateSelectAstOptions
+ * @property {SelectStatement} select
+ * @property {AsyncDataSource} dataSource
+ * @property {Record<string, AsyncDataSource>} tables
+ * @property {AbortSignal} [signal]
+ */
+
 /**
  * Evaluates a select with a resolved FROM data source
  *
- * @param {SelectStatement} select
- * @param {AsyncDataSource} dataSource
- * @param {Record<string, AsyncDataSource>} tables
+ * @param {EvaluateSelectAstOptions} options
  * @yields {AsyncRow}
  */
-async function* evaluateSelectAst(select, dataSource, tables) {
+async function* evaluateSelectAst({ select, dataSource, tables, signal }) {
   // SQL priority: from, where, group by, having, select, order by, offset, limit
 
-  const hasAggregate = select.columns.some(col => col.kind === 'aggregate')
+  const hasAggregate = select.columns.some(col => col.kind === 'derived' && containsAggregate(col.expr))
   const useGrouping = hasAggregate || select.groupBy.length > 0
   const needsBuffering = useGrouping || select.orderBy.length > 0
 
   if (needsBuffering) {
     // BUFFERING PATH: Collect all rows, process, then yield
-    yield* evaluateBuffered(select, dataSource, tables, hasAggregate, useGrouping)
+    yield* evaluateBuffered({ select, dataSource, tables, hasAggregate, useGrouping, signal })
   } else {
     // STREAMING PATH: Yield rows one by one
-    yield* evaluateStreaming(select, dataSource, tables)
+    yield* evaluateStreaming({ select, dataSource, tables, signal })
   }
 }
 
+/**
+ * @typedef {Object} EvaluateStreamingOptions
+ * @property {SelectStatement} select
+ * @property {AsyncDataSource} dataSource
+ * @property {Record<string, AsyncDataSource>} tables
+ * @property {AbortSignal} [signal]
+ */
+
 /**
  * Streaming evaluation for simple queries (no ORDER BY or GROUP BY)
  * Supports DISTINCT by tracking seen row keys without buffering full rows
  *
- * @param {SelectStatement} select
- * @param {AsyncDataSource} dataSource
- * @param {Record<string, AsyncDataSource>} tables
+ * @param {EvaluateStreamingOptions} options
  * @yields {AsyncRow}
  */
-async function* evaluateStreaming(select, dataSource, tables) {
+async function* evaluateStreaming({ select, dataSource, tables, signal }) {
   let rowsYielded = 0
   let rowsSkipped = 0
   let rowIndex = 0
@@ -255,7 +272,7 @@ async function* evaluateStreaming(select, dataSource, tables) {
     offset: select.offset,
   }
 
-  for await (const row of dataSource.getRows(hints)) {
+  for await (const row of dataSource.scan({ hints, signal })) {
     rowIndex++
     // WHERE filter
     if (select.where) {
@@ -270,27 +287,27 @@ async function* evaluateStreaming(select, dataSource, tables) {
     }
 
     // SELECT projection
-    /** @type {AsyncRow} */
-    const outRow = {}
+    /** @type {string[]} */
+    const columns = []
+    /** @type {AsyncCells} */
+    const cells = {}
     const currentRowIndex = rowIndex
     for (const col of select.columns) {
       if (col.kind === 'star') {
-        for (const [key, cell] of Object.entries(row)) {
-          outRow[key] = cell
+        for (const key of row.columns) {
+          columns.push(key)
+          cells[key] = row.cells[key]
         }
       } else if (col.kind === 'derived') {
         const alias = col.alias ?? defaultDerivedAlias(col.expr)
-        outRow[alias] = () => evaluateExpr({ node: col.expr, row, tables, rowIndex: currentRowIndex })
-      } else if (col.kind === 'aggregate') {
-        throw new Error(
-          'Aggregate functions require GROUP BY or will act on the whole dataset; add GROUP BY or remove aggregates'
-        )
+        columns.push(alias)
+        cells[alias] = () => evaluateExpr({ node: col.expr, row, tables, rowIndex: currentRowIndex })
       }
     }
 
     // DISTINCT: skip duplicate rows
     if (seen) {
-      const key = await stableRowKey(outRow)
+      const key = await stableRowKey(cells)
       if (seen.has(key)) continue
       seen.add(key)
       // OFFSET applies to distinct rows
@@ -300,7 +317,7 @@ async function* evaluateStreaming(select, dataSource, tables) {
       }
     }
 
-    yield outRow
+    yield { columns, cells }
     rowsYielded++
     if (rowsYielded >= limit) {
       break
@@ -308,17 +325,23 @@ async function* evaluateStreaming(select, dataSource, tables) {
   }
 }
 
+/**
+ * @typedef {Object} EvaluateBufferedOptions
+ * @property {SelectStatement} select
+ * @property {AsyncDataSource} dataSource
+ * @property {Record<string, AsyncDataSource>} tables
+ * @property {boolean} hasAggregate
+ * @property {boolean} useGrouping
+ * @property {AbortSignal} [signal]
+ */
+
 /**
  * Buffered evaluation for complex queries (with ORDER BY or GROUP BY)
  *
- * @param {SelectStatement} select
- * @param {AsyncDataSource} dataSource
- * @param {Record<string, AsyncDataSource>} tables
- * @param {boolean} hasAggregate
- * @param {boolean} useGrouping
+ * @param {EvaluateBufferedOptions} options
  * @yields {AsyncRow}
  */
-async function* evaluateBuffered(select, dataSource, tables, hasAggregate, useGrouping) {
+async function* evaluateBuffered({ select, dataSource, tables, hasAggregate, useGrouping, signal }) {
   // Build hints for data source optimization
   // Note: limit/offset not passed here since buffering needs all rows for sorting/grouping
   /** @type {QueryHints} */
@@ -330,7 +353,7 @@ async function* evaluateBuffered(select, dataSource, tables, hasAggregate, useGr
   // Step 1: Collect all rows from data source
   /** @type {AsyncRow[]} */
   const working = []
-  for await (const row of dataSource.getRows(hints)) {
+  for await (const row of dataSource.scan({ hints, signal })) {
     working.push(row)
   }
 
@@ -392,14 +415,16 @@ async function* evaluateBuffered(select, dataSource, tables, hasAggregate, useGr
     }
 
     for (const group of groups) {
-      /** @type {AsyncRow} */
-      const resultRow = {}
+      const columns = []
+      /** @type {AsyncCells} */
+      const cells = {}
       for (const col of select.columns) {
         if (col.kind === 'star') {
           const firstRow = group[0]
           if (firstRow) {
-            for (const [key, cell] of Object.entries(firstRow)) {
-              resultRow[key] = cell
+            for (const key of firstRow.columns) {
+              columns.push(key)
+              cells[key] = firstRow.cells[key]
             }
           }
           continue
@@ -407,29 +432,23 @@ async function* evaluateBuffered(select, dataSource, tables, hasAggregate, useGr
 
         if (col.kind === 'derived') {
           const alias = col.alias ?? defaultDerivedAlias(col.expr)
-          if (group.length > 0) {
-            resultRow[alias] = () => evaluateExpr({ node: col.expr, row: group[0], tables })
-          } else {
-            delete resultRow[alias]
-          }
-          continue
-        }
-
-        if (col.kind === 'aggregate') {
-          const alias = col.alias ?? defaultAggregateAlias(col)
-          resultRow[alias] = () => evaluateAggregate({ col, rows: group, tables })
+          columns.push(alias)
+          // Pass group to evaluateExpr so it can handle aggregate functions within expressions
+          // For empty groups, still provide an empty row context for aggregates to return appropriate values
+          cells[alias] = () => evaluateExpr({ node: col.expr, row: group[0] ?? { columns: [], cells: {} }, tables, rows: group })
           continue
         }
       }
+      const asyncRow = { columns, cells }
 
       // Apply HAVING filter before adding to projected results
       if (select.having) {
-        if (!await evaluateHavingExpr(select.having, resultRow, group, tables)) {
+        if (!await evaluateHavingExpr(select.having, asyncRow, group, tables)) {
           continue
         }
       }
 
-      projected.push(resultRow)
+      projected.push(asyncRow)
     }
   } else {
     // No grouping, simple projection
@@ -446,19 +465,22 @@ async function* evaluateBuffered(select, dataSource, tables, hasAggregate, useGr
     }
 
     for (const row of rowsToProject) {
-      /** @type {AsyncRow} */
-      const outRow = {}
+      const columns = []
+      /** @type {AsyncCells} */
+      const cells = {}
       for (const col of select.columns) {
         if (col.kind === 'star') {
-          for (const [key, cell] of Object.entries(row)) {
-            outRow[key] = cell
+          for (const key of row.columns) {
+            columns.push(key)
+            cells[key] = row.cells[key]
           }
         } else if (col.kind === 'derived') {
           const alias = col.alias ?? defaultDerivedAlias(col.expr)
-          outRow[alias] = () => evaluateExpr({ node: col.expr, row, tables })
+          columns.push(alias)
+          cells[alias] = () => evaluateExpr({ node: col.expr, row, tables })
         }
       }
-      projected.push(outRow)
+      projected.push({ columns, cells })
     }
   }