squirreling 0.5.0 → 0.6.1

package/src/execute/expression.js

@@ -1,11 +1,12 @@
 import { unknownFunctionError } from '../parseErrors.js'
 import { invalidContextError } from '../executionErrors.js'
 import {
+  aggregateError,
   argCountError,
   argValueError,
   castError,
 } from '../validationErrors.js'
-import { isMathFunc } from '../validation.js'
+import { isAggregateFunc, isMathFunc } from '../validation.js'
 import { applyIntervalToDate } from './date.js'
 import { executeSelect } from './execute.js'
 import { evaluateMathFunc } from './math.js'
@@ -23,23 +24,24 @@ import { applyBinaryOp, stringify } from './utils.js'
  * @param {AsyncRow} params.row - The data row to evaluate against
  * @param {Record<string, AsyncDataSource>} params.tables
  * @param {number} [params.rowIndex] - 1-based row index for error reporting
+ * @param {AsyncRow[]} [params.rows] - Group of rows for aggregate functions
  * @returns {Promise<SqlPrimitive>} The result of the evaluation
  */
-export async function evaluateExpr({ node, row, tables, rowIndex }) {
+export async function evaluateExpr({ node, row, tables, rowIndex, rows }) {
   if (node.type === 'literal') {
     return node.value
   }
 
   if (node.type === 'identifier') {
     // Try exact match first (handles both qualified and unqualified names)
-    if (row[node.name]) {
-      return row[node.name]()
+    if (row.cells[node.name]) {
+      return row.cells[node.name]()
     }
     // For qualified names like 'users.id', also try just the column part
     if (node.name.includes('.')) {
       const colName = node.name.split('.').pop()
-      if (colName && row[colName]) {
-        return row[colName]()
+      if (colName && row.cells[colName]) {
+        return row.cells[colName]()
       }
     }
     return null
@@ -47,29 +49,26 @@ export async function evaluateExpr({ node, row, tables, rowIndex }) {
 
   // Scalar subquery - returns a single value
   if (node.type === 'subquery') {
-    const gen = executeSelect(node.subquery, tables)
-    const first = await gen.next() // Start the generator
+    const gen = executeSelect({ select: node.subquery, tables })
+    const { value } = await gen.next() // Start the generator
     gen.return(undefined) // Stop further execution
-    if (!first.value) return null
-    /** @type {AsyncRow} */
-    const firstRow = first.value
-    const firstKey = Object.keys(firstRow)[0]
-    return firstRow[firstKey]()
+    if (!value) return null
+    return value.cells[value.columns[0]]()
   }
 
   // Unary operators
   if (node.type === 'unary') {
     if (node.op === 'NOT') {
-      return !await evaluateExpr({ node: node.argument, row, tables, rowIndex })
+      return !await evaluateExpr({ node: node.argument, row, tables, rowIndex, rows })
     }
     if (node.op === 'IS NULL') {
-      return await evaluateExpr({ node: node.argument, row, tables, rowIndex }) == null
+      return await evaluateExpr({ node: node.argument, row, tables, rowIndex, rows }) == null
     }
     if (node.op === 'IS NOT NULL') {
-      return await evaluateExpr({ node: node.argument, row, tables, rowIndex }) != null
+      return await evaluateExpr({ node: node.argument, row, tables, rowIndex, rows }) != null
     }
     if (node.op === '-') {
-      const val = await evaluateExpr({ node: node.argument, row, tables, rowIndex })
+      const val = await evaluateExpr({ node: node.argument, row, tables, rowIndex, rows })
       if (val == null) return null
       return -val
     }
@@ -79,15 +78,15 @@ export async function evaluateExpr({ node, row, tables, rowIndex }) {
   if (node.type === 'binary') {
     // Handle date +/- interval at AST level
     if ((node.op === '+' || node.op === '-') && node.right.type === 'interval') {
-      const dateVal = await evaluateExpr({ node: node.left, row, tables, rowIndex })
+      const dateVal = await evaluateExpr({ node: node.left, row, tables, rowIndex, rows })
       return applyIntervalToDate(dateVal, node.right.value, node.right.unit, node.op)
     }
     if (node.op === '+' && node.left.type === 'interval') {
-      const dateVal = await evaluateExpr({ node: node.right, row, tables, rowIndex })
+      const dateVal = await evaluateExpr({ node: node.right, row, tables, rowIndex, rows })
       return applyIntervalToDate(dateVal, node.left.value, node.left.unit, '+')
     }
 
-    const left = await evaluateExpr({ node: node.left, row, tables, rowIndex })
+    const left = await evaluateExpr({ node: node.left, row, tables, rowIndex, rows })
 
     // Short-circuit evaluation for AND and OR
     if (node.op === 'AND') {
@@ -97,15 +96,120 @@ export async function evaluateExpr({ node, row, tables, rowIndex }) {
       if (left) return true
     }
 
-    const right = await evaluateExpr({ node: node.right, row, tables, rowIndex })
+    const right = await evaluateExpr({ node: node.right, row, tables, rowIndex, rows })
     return applyBinaryOp(node.op, left, right)
   }
 
   // Function calls
   if (node.type === 'function') {
     const funcName = node.name.toUpperCase()
+
+    // Handle aggregate functions
+    if (isAggregateFunc(funcName)) {
+      if (!rows) {
+        throw aggregateError({
+          funcName,
+          issue: 'requires GROUP BY or will act on the whole dataset',
+        })
+      }
+
+      // Check for star argument (COUNT(*))
+      if (node.args.length === 1 && node.args[0].type === 'identifier' && node.args[0].name === '*') {
+        if (funcName === 'COUNT') {
+          return rows.length
+        }
+        throw aggregateError({
+          funcName,
+          issue: '(*) is not supported, use a column name',
+        })
+      }
+
+      if (node.args.length !== 1) {
+        throw argCountError({
+          funcName,
+          expected: 1,
+          received: node.args.length,
+          positionStart: node.positionStart,
+          positionEnd: node.positionEnd,
+          rowNumber: rowIndex,
+        })
+      }
+
+      const argNode = node.args[0]
+
+      if (funcName === 'COUNT') {
+        if (node.distinct) {
+          const seen = new Set()
+          for (const r of rows) {
+            const v = await evaluateExpr({ node: argNode, row: r, tables })
+            if (v != null) seen.add(v)
+          }
+          return seen.size
+        }
+        let count = 0
+        for (const r of rows) {
+          const v = await evaluateExpr({ node: argNode, row: r, tables })
+          if (v != null) count++
+        }
+        return count
+      }
+
+      if (funcName === 'SUM' || funcName === 'AVG' || funcName === 'MIN' || funcName === 'MAX') {
+        let sum = 0
+        let count = 0
+        /** @type {number | null} */
+        let min = null
+        /** @type {number | null} */
+        let max = null
+
+        for (const r of rows) {
+          const raw = await evaluateExpr({ node: argNode, row: r, tables })
+          if (raw == null) continue
+          const num = Number(raw)
+          if (!Number.isFinite(num)) continue
+
+          if (count === 0) {
+            min = num
+            max = num
+          } else {
+            if (min == null || num < min) min = num
+            if (max == null || num > max) max = num
+          }
+          sum += num
+          count++
+        }
+
+        if (funcName === 'SUM') return sum
+        if (funcName === 'AVG') return count === 0 ? null : sum / count
+        if (funcName === 'MIN') return min
+        if (funcName === 'MAX') return max
+      }
+
+      if (funcName === 'JSON_ARRAYAGG') {
+        /** @type {SqlPrimitive[]} */
+        const values = []
+        if (node.distinct) {
+          const seen = new Set()
+          for (const r of rows) {
+            const v = await evaluateExpr({ node: argNode, row: r, tables })
+            const key = stringify(v)
+            if (!seen.has(key)) {
+              seen.add(key)
+              values.push(v)
+            }
+          }
+        } else {
+          for (const r of rows) {
+            const v = await evaluateExpr({ node: argNode, row: r, tables })
+            values.push(v)
+          }
+        }
+        return values
+      }
+    }
+
     /** @type {SqlPrimitive[]} */
-    const args = await Promise.all(node.args.map(arg => evaluateExpr({ node: arg, row, tables, rowIndex })))
+    const args = await Promise.all(node.args.map(arg => evaluateExpr({ node: arg, row, tables, rowIndex, rows })))
 
     if (funcName === 'UPPER') {
       if (args.length !== 1) {
@@ -430,7 +534,7 @@ export async function evaluateExpr({ node, row, tables, rowIndex }) {
   }
 
   if (node.type === 'cast') {
-    const val = await evaluateExpr({ node: node.expr, row, tables, rowIndex })
+    const val = await evaluateExpr({ node: node.expr, row, tables, rowIndex, rows })
     if (val == null) return null
     const toType = node.toType.toUpperCase()
     if (toType === 'TEXT' || toType === 'STRING' || toType === 'VARCHAR') {
@@ -473,62 +577,59 @@ export async function evaluateExpr({ node, row, tables, rowIndex }) {
 
   // IN and NOT IN with value lists
   if (node.type === 'in valuelist') {
-    const exprVal = await evaluateExpr({ node: node.expr, row, tables, rowIndex })
+    const exprVal = await evaluateExpr({ node: node.expr, row, tables, rowIndex, rows })
     for (const valueNode of node.values) {
-      const val = await evaluateExpr({ node: valueNode, row, tables, rowIndex })
+      const val = await evaluateExpr({ node: valueNode, row, tables, rowIndex, rows })
       if (exprVal === val) return true
     }
     return false
   }
   // IN with subqueries
   if (node.type === 'in') {
-    const exprVal = await evaluateExpr({ node: node.expr, row, tables, rowIndex })
-    const results = executeSelect(node.subquery, tables)
-    /** @type {SqlPrimitive[]} */
-    const values = []
+    const exprVal = await evaluateExpr({ node: node.expr, row, tables, rowIndex, rows })
+    const results = executeSelect({ select: node.subquery, tables })
     for await (const resRow of results) {
-      const firstKey = Object.keys(resRow)[0]
-      const val = await resRow[firstKey]()
-      values.push(val)
+      const value = await resRow.cells[resRow.columns[0]]()
+      if (exprVal === value) return true
     }
-    return values.includes(exprVal)
+    return false
   }
 
   // EXISTS and NOT EXISTS with subqueries
   if (node.type === 'exists') {
-    const results = await executeSelect(node.subquery, tables).next()
+    const results = await executeSelect({ select: node.subquery, tables }).next()
     return results.done === false
   }
   if (node.type === 'not exists') {
-    const results = await executeSelect(node.subquery, tables).next()
+    const results = await executeSelect({ select: node.subquery, tables }).next()
     return results.done === true
   }
 
   // CASE expressions
   if (node.type === 'case') {
     // For simple CASE: evaluate the case expression once
-    const caseValue = node.caseExpr && await evaluateExpr({ node: node.caseExpr, row, tables, rowIndex })
+    const caseValue = node.caseExpr && await evaluateExpr({ node: node.caseExpr, row, tables, rowIndex, rows })
 
     // Iterate through WHEN clauses
     for (const whenClause of node.whenClauses) {
       let conditionResult
       if (caseValue !== undefined) {
         // Simple CASE: compare caseValue with condition
-        const whenValue = await evaluateExpr({ node: whenClause.condition, row, tables, rowIndex })
+        const whenValue = await evaluateExpr({ node: whenClause.condition, row, tables, rowIndex, rows })
         conditionResult = caseValue === whenValue
       } else {
         // Searched CASE: evaluate condition as boolean
-        conditionResult = await evaluateExpr({ node: whenClause.condition, row, tables, rowIndex })
+        conditionResult = await evaluateExpr({ node: whenClause.condition, row, tables, rowIndex, rows })
       }
 
       if (conditionResult) {
-        return evaluateExpr({ node: whenClause.result, row, tables, rowIndex })
+        return evaluateExpr({ node: whenClause.result, row, tables, rowIndex, rows })
       }
     }
 
     // No WHEN clause matched, return ELSE result or NULL
     if (node.elseResult) {
-      return evaluateExpr({ node: node.elseResult, row, tables, rowIndex })
+      return evaluateExpr({ node: node.elseResult, row, tables, rowIndex, rows })
     }
     return null
   }

package/src/execute/join.js

@@ -4,7 +4,7 @@ import { evaluateExpr } from './expression.js'
 import { stringify } from './utils.js'
 
 /**
- * @import { AsyncRow, AsyncDataSource, JoinClause, ExprNode } from '../types.js'
+ * @import { AsyncRow, AsyncDataSource, JoinClause, ExprNode, AsyncCells } from '../types.js'
  */
 
 /**
@@ -30,7 +30,7 @@ export async function executeJoins(leftSource, joins, leftTableName, tables) {
     // Buffer right rows for hash index (required for hash join)
     /** @type {AsyncRow[]} */
     const rightRows = []
-    for await (const row of rightSource.getRows()) {
+    for await (const row of rightSource.scan({})) {
       rightRows.push(row)
     }
 
@@ -39,14 +39,16 @@ export async function executeJoins(leftSource, joins, leftTableName, tables) {
 
     // Return streaming data source - left rows stream through without buffering
     return {
-      async *getRows() {
+      async *scan(options) {
+        const { signal } = options
         yield* hashJoin({
-          leftRows: leftSource.getRows(), // Stream directly, not buffered
+          leftRows: leftSource.scan(options), // Stream directly, not buffered
           rightRows,
           join,
           leftTable: currentLeftTable,
           rightTable: rightTableName,
           tables,
+          signal,
         })
       },
     }
@@ -55,7 +57,7 @@ export async function executeJoins(leftSource, joins, leftTableName, tables) {
   // Multiple joins: buffer intermediate results, stream final join
   /** @type {AsyncRow[]} */
   let leftRows = []
-  for await (const row of leftSource.getRows()) {
+  for await (const row of leftSource.scan({})) {
     leftRows.push(row)
   }
 
@@ -69,7 +71,7 @@ export async function executeJoins(leftSource, joins, leftTableName, tables) {
 
     /** @type {AsyncRow[]} */
     const rightRows = []
-    for await (const row of rightSource.getRows()) {
+    for await (const row of rightSource.scan({})) {
       rightRows.push(row)
     }
 
@@ -105,7 +107,7 @@ export async function executeJoins(leftSource, joins, leftTableName, tables) {
 
   /** @type {AsyncRow[]} */
   const rightRows = []
-  for await (const row of rightSource.getRows()) {
+  for await (const row of rightSource.scan({})) {
     rightRows.push(row)
   }
 
@@ -113,7 +115,8 @@ export async function executeJoins(leftSource, joins, leftTableName, tables) {
   const lastRightTableName = lastJoin.alias ?? lastJoin.table
 
   return {
-    async *getRows() {
+    async *scan(options) {
+      const { signal } = options
       yield* hashJoin({
         leftRows,
         rightRows,
@@ -121,6 +124,7 @@ export async function executeJoins(leftSource, joins, leftTableName, tables) {
         leftTable: currentLeftTable,
         rightTable: lastRightTableName,
         tables,
+        signal,
       })
     },
   }
@@ -172,12 +176,12 @@ function extractJoinKeys(onCondition, leftTable, rightTable) {
  * @returns {AsyncRow}
  */
 function createNullRow(columnNames) {
-  /** @type {AsyncRow} */
-  const row = {}
+  /** @type {AsyncCells} */
+  const cells = {}
   for (const col of columnNames) {
-    row[col] = () => Promise.resolve(null)
+    cells[col] = () => Promise.resolve(null)
   }
-  return row
+  return { columns: columnNames, cells }
 }
 
 /**
@@ -190,33 +194,35 @@ function createNullRow(columnNames) {
  * @returns {AsyncRow}
  */
 function mergeRows(leftRow, rightRow, leftTable, rightTable) {
-  /** @type {AsyncRow} */
-  const merged = {}
+  const columns = []
+  /** @type {AsyncCells} */
+  const cells = {}
 
   // Add left table columns with prefix
-  for (const [key, cell] of Object.entries(leftRow)) {
+  for (const [key, cell] of Object.entries(leftRow.cells)) {
     // Skip already-prefixed keys (from previous joins)
     if (!key.includes('.')) {
-      merged[`${leftTable}.${key}`] = cell
-    } else {
-      merged[key] = cell
+      const alias = `${leftTable}.${key}`
+      cells[alias] = cell
     }
-    // Also keep unqualified name for convenience (may be overwritten if ambiguous)
-    merged[key] = cell
+    // Also keep unqualified name for convenience
+    columns.push(key)
+    cells[key] = cell
   }
 
   // Add right table columns with prefix
-  for (const [key, cell] of Object.entries(rightRow)) {
+  for (const [key, cell] of Object.entries(rightRow.cells)) {
     if (!key.includes('.')) {
-      merged[`${rightTable}.${key}`] = cell
+      cells[`${rightTable}.${key}`] = cell
     } else {
-      merged[key] = cell
+      cells[key] = cell
     }
     // Unqualified name (overwrites if same name exists in left table)
-    merged[key] = cell
+    columns.push(key)
+    cells[key] = cell
   }
 
-  return merged
+  return { columns, cells }
 }
 
 /**
@@ -230,9 +236,10 @@ function mergeRows(leftRow, rightRow, leftTable, rightTable) {
  * @param {string} params.leftTable - name of left table (for column prefixing)
  * @param {string} params.rightTable - name of right table (for column prefixing, may be alias)
  * @param {Record<string, AsyncDataSource>} params.tables - all tables for expression evaluation
+ * @param {AbortSignal} [params.signal] - abort signal for cancellation
  * @yields {AsyncRow} joined rows
  */
-async function* hashJoin({ leftRows, rightRows, join, leftTable, rightTable, tables }) {
+async function* hashJoin({ leftRows, rightRows, join, leftTable, rightTable, tables, signal }) {
   const { joinType, on: onCondition } = join
 
   if (!onCondition) {
@@ -245,7 +252,7 @@ async function* hashJoin({ leftRows, rightRows, join, leftTable, rightTable, tab
   const keys = extractJoinKeys(onCondition, leftTable, rightTable)
 
   // Get column names for NULL row generation (right side is always buffered)
-  const rightCols = rightRows.length ? Object.keys(rightRows[0]) : []
+  const rightCols = rightRows.length ? rightRows[0].columns : []
   const rightPrefixedCols = rightCols.flatMap(col =>
     col.includes('.') ? [col] : [`${rightTable}.${col}`, col]
   )
@@ -279,10 +286,10 @@ async function* hashJoin({ leftRows, rightRows, join, leftTable, rightTable, tab
 
     // PROBE PHASE: Stream through left rows, yield matches immediately
     for await (const leftRow of leftRows) {
+      if (signal?.aborted) break
       // Capture left column info from first row (for NULL row generation)
       if (!leftPrefixedCols) {
-        const leftCols = Object.keys(leftRow)
-        leftPrefixedCols = leftCols.flatMap(col =>
+        leftPrefixedCols = leftRow.columns.flatMap(col =>
           col.includes('.') ? [col] : [`${leftTable}.${col}`, col]
         )
       }
@@ -321,10 +328,10 @@ async function* hashJoin({ leftRows, rightRows, join, leftTable, rightTable, tab
     const matchedRightRows = joinType === 'RIGHT' || joinType === 'FULL' ? new Set() : null
 
     for await (const leftRow of leftRows) {
+      if (signal?.aborted) break
       // Capture left column info from first row (for NULL row generation)
       if (!leftPrefixedCols) {
-        const leftCols = Object.keys(leftRow)
-        leftPrefixedCols = leftCols.flatMap(col =>
+        leftPrefixedCols = leftRow.columns.flatMap(col =>
           col.includes('.') ? [col] : [`${leftTable}.${col}`, col]
         )
       }

package/src/execute/math.js

@@ -1,5 +1,5 @@
 /**
- * @import { SqlPrimitive } from '../types.js'
+ * @import { MathFunc, SqlPrimitive } from '../types.js'
  */
 import { argCountError } from '../validationErrors.js'
 
@@ -7,7 +7,7 @@ import { argCountError } from '../validationErrors.js'
  * Evaluate a math function
  *
  * @param {Object} options
- * @param {string} options.funcName - Uppercase function name
+ * @param {MathFunc} options.funcName - Uppercase function name
  * @param {SqlPrimitive[]} options.args - Function arguments
  * @param {number} options.positionStart - Start position in query
  * @param {number} options.positionEnd - End position in query
@@ -161,5 +161,180 @@ export function evaluateMathFunc({ funcName, args, positionStart, positionEnd, r
     return Math.sqrt(Number(val))
   }
 
-  return undefined
+  if (funcName === 'SIN') {
+    if (args.length !== 1) {
+      throw argCountError({
+        funcName: 'SIN',
+        expected: 1,
+        received: args.length,
+        positionStart,
+        positionEnd,
+        rowNumber,
+      })
+    }
+    const val = args[0]
+    if (val == null) return null
+    return Math.sin(Number(val))
+  }
+
+  if (funcName === 'COS') {
+    if (args.length !== 1) {
+      throw argCountError({
+        funcName: 'COS',
+        expected: 1,
+        received: args.length,
+        positionStart,
+        positionEnd,
+        rowNumber,
+      })
+    }
+    const val = args[0]
+    if (val == null) return null
+    return Math.cos(Number(val))
+  }
+
+  if (funcName === 'TAN') {
+    if (args.length !== 1) {
+      throw argCountError({
+        funcName: 'TAN',
+        expected: 1,
+        received: args.length,
+        positionStart,
+        positionEnd,
+        rowNumber,
+      })
+    }
+    const val = args[0]
+    if (val == null) return null
+    return Math.tan(Number(val))
+  }
+
+  if (funcName === 'COT') {
+    if (args.length !== 1) {
+      throw argCountError({
+        funcName: 'COT',
+        expected: 1,
+        received: args.length,
+        positionStart,
+        positionEnd,
+        rowNumber,
+      })
+    }
+    const val = args[0]
+    if (val == null) return null
+    return 1 / Math.tan(Number(val))
+  }
+
+  if (funcName === 'ASIN') {
+    if (args.length !== 1) {
+      throw argCountError({
+        funcName: 'ASIN',
+        expected: 1,
+        received: args.length,
+        positionStart,
+        positionEnd,
+        rowNumber,
+      })
+    }
+    const val = args[0]
+    if (val == null) return null
+    return Math.asin(Number(val))
+  }
+
+  if (funcName === 'ACOS') {
+    if (args.length !== 1) {
+      throw argCountError({
+        funcName: 'ACOS',
+        expected: 1,
+        received: args.length,
+        positionStart,
+        positionEnd,
+        rowNumber,
+      })
+    }
+    const val = args[0]
+    if (val == null) return null
+    return Math.acos(Number(val))
+  }
+
+  if (funcName === 'ATAN') {
+    if (args.length !== 1) {
+      throw argCountError({
+        funcName: 'ATAN',
+        expected: 1,
+        received: args.length,
+        positionStart,
+        positionEnd,
+        rowNumber,
+      })
+    }
+    const val = args[0]
+    if (val == null) return null
+    return Math.atan(Number(val))
+  }
+
+  if (funcName === 'ATAN2') {
+    if (args.length !== 2) {
+      throw argCountError({
+        funcName: 'ATAN2',
+        expected: 2,
+        received: args.length,
+        positionStart,
+        positionEnd,
+        rowNumber,
+      })
+    }
+    const y = args[0]
+    const x = args[1]
+    if (y == null || x == null) return null
+    return Math.atan2(Number(y), Number(x))
+  }
+
+  if (funcName === 'DEGREES') {
+    if (args.length !== 1) {
+      throw argCountError({
+        funcName: 'DEGREES',
+        expected: 1,
+        received: args.length,
+        positionStart,
+        positionEnd,
+        rowNumber,
+      })
+    }
+    const val = args[0]
+    if (val == null) return null
+    return Number(val) * 180 / Math.PI
+  }
+
+  if (funcName === 'RADIANS') {
+    if (args.length !== 1) {
+      throw argCountError({
+        funcName: 'RADIANS',
+        expected: 1,
+        received: args.length,
+        positionStart,
+        positionEnd,
+        rowNumber,
+      })
+    }
+    const val = args[0]
+    if (val == null) return null
+    return Number(val) * Math.PI / 180
+  }
+
+  if (funcName === 'PI') {
+    if (args.length !== 0) {
+      throw argCountError({
+        funcName: 'PI',
+        expected: 0,
+        received: args.length,
+        positionStart,
+        positionEnd,
+        rowNumber,
+      })
+    }
+    return Math.PI
+  }
+
+  return null
 }