squirreling 0.6.0 → 0.7.0

package/src/execute/join.js

@@ -4,20 +4,22 @@ import { evaluateExpr } from './expression.js'
 import { stringify } from './utils.js'
 
 /**
- * @import { AsyncRow, AsyncDataSource, JoinClause, ExprNode, AsyncCells } from '../types.js'
+ * @import { AsyncRow, AsyncDataSource, JoinClause, ExprNode, AsyncCells, UserDefinedFunction } from '../types.js'
  */
 
 /**
  * Executes JOIN operations against a base data source
  *
- * @param {AsyncDataSource} leftSource - the left side of the join (FROM table)
- * @param {JoinClause[]} joins - array of join clauses to execute
- * @param {string} leftTableName - name of the left table (for column prefixing)
- * @param {Record<string, AsyncDataSource>} tables - all available tables
+ * @param {Object} options
+ * @param {AsyncDataSource} options.leftSource - the left side of the join (FROM table)
+ * @param {JoinClause[]} options.joins - array of join clauses to execute
+ * @param {string} options.leftTable - name of the left table (for column prefixing)
+ * @param {Record<string, AsyncDataSource>} options.tables - all available tables
+ * @param {Record<string, UserDefinedFunction>} [options.functions]
  * @returns {Promise<AsyncDataSource>} data source yielding joined rows
  */
-export async function executeJoins(leftSource, joins, leftTableName, tables) {
-  let currentLeftTable = leftTableName
+export async function executeJoins({ leftSource, joins, leftTable, tables, functions }) {
+  let currentLeftTable = leftTable
 
   // Single join optimization: stream left rows without buffering
   if (joins.length === 1) {
@@ -30,23 +32,26 @@ 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.scan()) {
+    for await (const row of rightSource.scan({})) {
       rightRows.push(row)
     }
 
     // Use alias for column prefixing if present
-    const rightTableName = join.alias ?? join.table
+    const rightTable = join.alias ?? join.table
 
     // Return streaming data source - left rows stream through without buffering
     return {
-      async *scan() {
+      async *scan(options) {
+        const { signal } = options
         yield* hashJoin({
-          leftRows: leftSource.scan(), // Stream directly, not buffered
+          leftRows: leftSource.scan(options), // Stream directly, not buffered
           rightRows,
           join,
           leftTable: currentLeftTable,
-          rightTable: rightTableName,
+          rightTable,
           tables,
+          functions,
+          signal,
         })
       },
     }
@@ -55,7 +60,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.scan()) {
+  for await (const row of leftSource.scan({})) {
     leftRows.push(row)
   }
 
@@ -69,12 +74,12 @@ export async function executeJoins(leftSource, joins, leftTableName, tables) {
 
     /** @type {AsyncRow[]} */
     const rightRows = []
-    for await (const row of rightSource.scan()) {
+    for await (const row of rightSource.scan({})) {
       rightRows.push(row)
     }
 
     // Use alias for column prefixing if present
-    const rightTableName = join.alias ?? join.table
+    const rightTable = join.alias ?? join.table
 
     // Collect intermediate results into array for next join
     /** @type {AsyncRow[]} */
@@ -84,8 +89,9 @@ export async function executeJoins(leftSource, joins, leftTableName, tables) {
       rightRows,
       join,
       leftTable: currentLeftTable,
-      rightTable: rightTableName,
+      rightTable,
       tables,
+      functions,
     })
     for await (const row of joined) {
       newLeftRows.push(row)
@@ -93,34 +99,37 @@ export async function executeJoins(leftSource, joins, leftTableName, tables) {
     leftRows = newLeftRows
 
     // After join, the "left" table for the next join includes all joined tables
-    currentLeftTable = `${currentLeftTable}_${rightTableName}`
+    currentLeftTable = `${currentLeftTable}_${rightTable}`
   }
 
   // Final join: stream the results
-  const lastJoin = joins[joins.length - 1]
-  const rightSource = tables[lastJoin.table]
+  const join = joins[joins.length - 1]
+  const rightSource = tables[join.table]
   if (rightSource === undefined) {
-    throw tableNotFoundError({ tableName: lastJoin.table })
+    throw tableNotFoundError({ tableName: join.table })
   }
 
   /** @type {AsyncRow[]} */
   const rightRows = []
-  for await (const row of rightSource.scan()) {
+  for await (const row of rightSource.scan({})) {
     rightRows.push(row)
   }
 
   // Use alias for column prefixing if present
-  const lastRightTableName = lastJoin.alias ?? lastJoin.table
+  const rightTable = join.alias ?? join.table
 
   return {
-    async *scan() {
+    async *scan(options) {
+      const { signal } = options
       yield* hashJoin({
         leftRows,
         rightRows,
-        join: lastJoin,
+        join,
         leftTable: currentLeftTable,
-        rightTable: lastRightTableName,
+        rightTable,
         tables,
+        functions,
+        signal,
       })
     },
   }
@@ -232,9 +241,11 @@ 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 {Record<string, UserDefinedFunction>} [params.functions]
+ * @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, functions, signal }) {
   const { joinType, on: onCondition } = join
 
   if (!onCondition) {
@@ -264,7 +275,7 @@ async function* hashJoin({ leftRows, rightRows, join, leftTable, rightTable, tab
     // BUILD PHASE: Index right rows by join key
     // Skip null keys - SQL semantics: NULL != NULL
     for (const rightRow of rightRows) {
-      const keyValue = await evaluateExpr({ node: keys.rightKey, row: rightRow, tables })
+      const keyValue = await evaluateExpr({ node: keys.rightKey, row: rightRow, tables, functions })
       if (keyValue == null) continue // NULL keys never match
       const keyStr = stringify(keyValue)
       let bucket = hashMap.get(keyStr)
@@ -281,6 +292,7 @@ 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) {
         leftPrefixedCols = leftRow.columns.flatMap(col =>
@@ -288,7 +300,7 @@ async function* hashJoin({ leftRows, rightRows, join, leftTable, rightTable, tab
         )
       }
 
-      const keyValue = await evaluateExpr({ node: keys.leftKey, row: leftRow, tables })
+      const keyValue = await evaluateExpr({ node: keys.leftKey, row: leftRow, tables, functions })
       const keyStr = stringify(keyValue)
 
       const matchingRightRows = hashMap.get(keyStr)
@@ -322,6 +334,7 @@ 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) {
         leftPrefixedCols = leftRow.columns.flatMap(col =>
@@ -333,7 +346,7 @@ async function* hashJoin({ leftRows, rightRows, join, leftTable, rightTable, tab
 
       for (const rightRow of rightRows) {
         const tempMerged = mergeRows(leftRow, rightRow, leftTable, rightTable)
-        const matches = await evaluateExpr({ node: onCondition, row: tempMerged, tables })
+        const matches = await evaluateExpr({ node: onCondition, row: tempMerged, tables, functions })
 
         if (matches) {
           hasMatch = true

package/src/execute/math.js

@@ -1,7 +1,6 @@
 /**
  * @import { MathFunc, SqlPrimitive } from '../types.js'
  */
-import { argCountError } from '../validationErrors.js'
 
 /**
  * Evaluate a math function
@@ -9,71 +8,28 @@ import { argCountError } from '../validationErrors.js'
  * @param {Object} options
  * @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
- * @param {number} [options.rowNumber] - 1-based row number for error reporting
  * @returns {SqlPrimitive} Result
  */
-export function evaluateMathFunc({ funcName, args, positionStart, positionEnd, rowNumber }) {
+export function evaluateMathFunc({ funcName, args }) {
   if (funcName === 'FLOOR') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'FLOOR',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Math.floor(Number(val))
   }
 
   if (funcName === 'CEIL' || funcName === 'CEILING') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName,
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Math.ceil(Number(val))
   }
 
   if (funcName === 'ABS') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'ABS',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Math.abs(Number(val))
   }
 
   if (funcName === 'MOD') {
-    if (args.length !== 2) {
-      throw argCountError({
-        funcName: 'MOD',
-        expected: 2,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const dividend = args[0]
     const divisor = args[1]
     if (dividend == null || divisor == null) return null
@@ -81,64 +37,24 @@ export function evaluateMathFunc({ funcName, args, positionStart, positionEnd, r
   }
 
   if (funcName === 'EXP') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'EXP',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Math.exp(Number(val))
   }
 
   if (funcName === 'LN') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'LN',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Math.log(Number(val))
   }
 
   if (funcName === 'LOG10') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'LOG10',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Math.log10(Number(val))
   }
 
   if (funcName === 'POWER') {
-    if (args.length !== 2) {
-      throw argCountError({
-        funcName: 'POWER',
-        expected: 2,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const base = args[0]
     const exponent = args[1]
     if (base == null || exponent == null) return null
@@ -146,144 +62,61 @@ export function evaluateMathFunc({ funcName, args, positionStart, positionEnd, r
   }
 
   if (funcName === 'SQRT') {
-    if (args.length !== 1) {
-      throw argCountError({
-        funcName: 'SQRT',
-        expected: 1,
-        received: args.length,
-        positionStart,
-        positionEnd,
-        rowNumber,
-      })
-    }
     const val = args[0]
     if (val == null) return null
     return Math.sqrt(Number(val))
   }
 
   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,
-      })
+    if (args.length === 1) {
+      const val = args[0]
+      if (val == null) return null
+      return Math.atan(Number(val))
+    } else {
+      const y = args[0]
+      const x = args[1]
+      if (y == null || x == null) return null
+      return Math.atan2(Number(y), Number(x))
     }
-    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
@@ -291,50 +124,18 @@ export function evaluateMathFunc({ funcName, args, positionStart, positionEnd, r
   }
 
   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
 }

package/src/execute/utils.js

@@ -132,6 +132,10 @@ export function defaultDerivedAlias(expr) {
     return defaultDerivedAlias(expr.left) + '_' + expr.op + '_' + defaultDerivedAlias(expr.right)
   }
   if (expr.type === 'function') {
+    // Handle aggregate functions with star (COUNT(*) -> count_all)
+    if (expr.args.length === 1 && expr.args[0].type === 'identifier' && expr.args[0].name === '*') {
+      return expr.name.toLowerCase() + '_all'
+    }
     return expr.name.toLowerCase() + '_' + expr.args.map(defaultDerivedAlias).join('_')
   }
   if (expr.type === 'interval') {

package/src/index.d.ts

@@ -14,10 +14,11 @@ export function executeSql(options: ExecuteSqlOptions): AsyncGenerator<AsyncRow>
 /**
  * Parses a SQL query string into an abstract syntax tree
  *
- * @param query - SQL query string to parse
+ * @param options
+ * @param options.query - SQL query string to parse
  * @returns parsed SQL select statement
  */
-export function parseSql(query: string): SelectStatement
+export function parseSql(options: { query: string }): SelectStatement
 
 /**
  * Collects all results from an async generator into an array

package/src/parse/expression.js

@@ -1,10 +1,11 @@
 import {
+  argCountParseError,
   invalidLiteralError,
   missingClauseError,
   syntaxError,
   unknownFunctionError,
 } from '../parseErrors.js'
-import { isAggregateFunc, isIntervalUnit, isMathFunc, isStringFunc } from '../validation.js'
+import { isIntervalUnit, isKnownFunction, validateFunctionArgCount } from '../validation.js'
 import { parseComparison } from './comparison.js'
 import { parseSelectInternal } from './parse.js'
 import { consume, current, expect, expectIdentifier, lastPosition, match, peekToken } from './state.js'
@@ -122,10 +123,15 @@ export function parsePrimary(state) {
     // function call
     if (next.type === 'paren' && next.value === '(') {
       const funcName = tok.value
-
-      // validate function names
-      if (!isStringFunc(funcName) && !isAggregateFunc(funcName) && !isMathFunc(funcName)) {
-        throw unknownFunctionError({ funcName, positionStart: tok.positionStart, positionEnd: tok.positionEnd })
+      const funcNameUpper = funcName.toUpperCase()
+
+      // Validate function existence early for better error messages
+      if (!isKnownFunction(funcNameUpper, state.functions)) {
+        throw unknownFunctionError({
+          funcName,
+          positionStart,
+          positionEnd: tok.positionEnd,
+        })
       }
 
       consume(state) // function name
@@ -133,6 +139,15 @@ export function parsePrimary(state) {
 
       /** @type {ExprNode[]} */
       const args = []
+      let distinct = false
+
+      // Check for DISTINCT or ALL keyword (for aggregate functions like COUNT(DISTINCT x))
+      if (current(state).type === 'keyword' && current(state).value === 'DISTINCT') {
+        consume(state) // consume DISTINCT
+        distinct = true
+      } else if (current(state).type === 'keyword' && current(state).value === 'ALL') {
+        consume(state) // consume ALL (default behavior, just consume it)
+      }
 
       if (current(state).type !== 'paren' || current(state).value !== ')') {
         while (true) {
@@ -156,10 +171,23 @@ export function parsePrimary(state) {
 
       expect(state, 'paren', ')')
 
+      // Validate argument count at parse time
+      const validation = validateFunctionArgCount(funcNameUpper, args.length)
+      if (!validation.valid) {
+        throw argCountParseError({
+          funcName,
+          expected: validation.expected,
+          received: args.length,
+          positionStart,
+          positionEnd: lastPosition(state),
+        })
+      }
+
       return {
         type: 'function',
         name: funcName,
         args,
+        distinct: distinct || undefined,
         positionStart,
         positionEnd: lastPosition(state),
       }