squirreling 0.16.2 → 0.16.4

package/README.md

@@ -95,10 +95,10 @@ Because Squirreling uses lazy cell evaluation, the `AI_SCORE` function only exec
 
 ### Custom Data Sources
 
-Squirreling can work with any data source that implements the `AsyncDataSource` interface.
+Squirreling can work with any data source that implements the `AsyncDataSource` interface. The simplest form is a `ScannableDataSource`: column names plus a row scan.
 
 ```typescript
-interface AsyncDataSource {
+interface ScannableDataSource {
   numRows?: number
   columns: string[]
   scan(options: ScanOptions): ScanResults
@@ -144,7 +144,7 @@ Squirreling mostly follows the SQL standard. The following features are supporte
 - `WITH` clause for Common Table Expressions (CTEs)
 - Subqueries in `SELECT`, `FROM`, `WHERE`, and correlated subqueries
 - `JOIN` operations: `INNER JOIN`, `LEFT JOIN`, `RIGHT JOIN`, `FULL JOIN`, `CROSS JOIN`, `POSITIONAL JOIN`, `LATERAL VIEW [OUTER] EXPLODE(...)`, with `ON` or `USING (col, ...)` conditions
-- `GROUP BY` and `HAVING` clauses
+- `GROUP BY` and `HAVING` clauses, including `GROUP BY ALL`
 - Set operations: `UNION`, `UNION ALL`, `INTERSECT`, `INTERSECT ALL`, `EXCEPT`, `EXCEPT ALL`
 - Expressions: `CASE`, `CAST`, `TRY_CAST`, `BETWEEN`, `IN`, `LIKE`, `IS NULL`, `IS NOT NULL`, string concatenation `||`
 - Subscript access: zero-based array indexing `col[0]`, struct field access `col['field']`, and chains like `col[0].field`

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "squirreling",
-  "version": "0.16.2",
+  "version": "0.16.4",
   "description": "Squirreling Async SQL Engine",
   "author": "Hyperparam",
   "homepage": "https://hyperparam.app",
@@ -39,11 +39,13 @@
     "test": "vitest run"
   },
   "devDependencies": {
-    "@types/node": "26.2.0",
-    "@vitest/coverage-v8": "4.1.11",
-    "eslint": "9.39.4",
-    "eslint-plugin-jsdoc": "64.2.1",
+    "@eslint/js": "10.0.1",
+    "@types/node": "26.4.1",
+    "@vitest/coverage-v8": "5.0.0",
+    "eslint": "10.10.0",
+    "eslint-plugin-jsdoc": "64.3.6",
+    "globals": "17.12.0",
     "typescript": "7.0.2",
-    "vitest": "4.1.11"
+    "vitest": "5.0.0"
   }
 }

package/src/execute/aggregates.js

@@ -3,6 +3,7 @@ import { derivedAlias } from '../expression/alias.js'
 import { evaluateExpr } from '../expression/evaluate.js'
 import { finalizeAccumulator, newAccumulator, updateAccumulator } from './accumulator.js'
 import { executePlan, executeScan, selectColumnNames } from './execute.js'
+import { foldEvaluatedRows } from './fold.js'
 import { normalizeScanColumnResult } from './scanColumn.js'
 import { sortEntriesByTerms } from './sort.js'
 import { planStreamingAggregates, streamingHashAggregateRows, streamingScalarAggregateRows } from './streamingAggregate.js'
@@ -109,47 +110,42 @@ export function executeHashAggregate(plan, context) {
       }
       context.signal?.throwIfAborted()
 
-      // Group rows by GROUP BY keys.
-      // Each chunk dispatches all per-row key evaluations in parallel so
-      // async cells (e.g. lazy parquet decode) overlap; the await is at the
-      // chunk boundary. Synchronous cells stay cheap because we skip the
-      // inner Promise.all wrapper when there's a single GROUP BY expression.
+      // Group rows by GROUP BY keys. Keys are evaluated in adaptive chunks
+      // so async cells (e.g. lazy parquet decode) overlap while evaluated
+      // key values stay byte-bounded. The single-key branch skips the inner
+      // Promise.all wrapper so synchronous cells stay cheap.
       /** @type {Map<any, AsyncRow[]>} */
       const groups = new Map()
       const { groupBy } = plan
-      const singleKey = groupBy.length === 1
-      const singleExpr = singleKey ? groupBy[0] : null
+      const singleExpr = groupBy.length === 1 ? groupBy[0] : null
 
-      for (let chunkStart = 0; chunkStart < allRows.length; chunkStart += YIELD_INTERVAL) {
-        if (chunkStart > 0) {
-          await yieldToEventLoop()
-          context.signal?.throwIfAborted()
-        }
-        const chunkEnd = Math.min(chunkStart + YIELD_INTERVAL, allRows.length)
-        const chunkLen = chunkEnd - chunkStart
-        /** @type {Promise<any>[]} */
-        const pending = new Array(chunkLen)
-        if (singleKey) {
-          for (let j = 0; j < chunkLen; j++) {
-            pending[j] = evaluateExpr({ node: singleExpr, row: allRows[chunkStart + j], context })
-          }
-        } else {
-          for (let j = 0; j < chunkLen; j++) {
-            const row = allRows[chunkStart + j]
-            pending[j] = Promise.all(groupBy.map(expr => evaluateExpr({ node: expr, row, context })))
-          }
-        }
-        const chunkKeys = await Promise.all(pending)
-        for (let j = 0; j < chunkLen; j++) {
-          const key = singleKey ? keyify(chunkKeys[j]) : keyify(...chunkKeys[j])
-          const row = allRows[chunkStart + j]
-          let group = groups.get(key)
-          if (!group) {
-            group = []
-            groups.set(key, group)
-          }
-          group.push(row)
+      /**
+       * @param {any} key
+       * @param {number} index
+       */
+      function addToGroup(key, index) {
+        let group = groups.get(key)
+        if (!group) {
+          group = []
+          groups.set(key, group)
         }
+        group.push(allRows[index])
+      }
+
+      if (singleExpr) {
+        await foldEvaluatedRows({
+          rows: allRows,
+          signal: context.signal,
+          evaluate: row => evaluateExpr({ node: singleExpr, row, context }),
+          fold: (value, index) => addToGroup(keyify(value), index),
+        })
+      } else {
+        await foldEvaluatedRows({
+          rows: allRows,
+          signal: context.signal,
+          evaluate: row => Promise.all(groupBy.map(expr => evaluateExpr({ node: expr, row, context }))),
+          fold: (values, index) => addToGroup(keyify(...values), index),
+        })
       }
 
       /** @type {{ row: AsyncRow, rows: AsyncRow[], outputRow: AsyncRow }[]} */

package/src/execute/execute.js

@@ -588,7 +588,9 @@ function executeCount(plan, context) {
           throw new Error(`Data source "${plan.table}" does not implement scan()`)
         }
         let count = 0
-        const { rows } = table.scan({ signal })
+        // Counting needs row presence, not payload columns. Sources may still
+        // read columns required for their own visibility or delete filters.
+        const { rows } = table.scan({ columns: [], signal })
         // eslint-disable-next-line no-unused-vars
         for await (const _ of rows()) {
           signal?.throwIfAborted()

package/src/execute/fold.js

@@ -0,0 +1,74 @@
+import { yieldToEventLoop } from './yield.js'
+
+/**
+ * @import { AsyncRow } from '../types.js'
+ */
+
+// Chunk sizing for foldEvaluatedRows. Dispatch starts small so one chunk of
+// unexpectedly fat values cannot overshoot far, then adapts to the observed
+// result sizes: small values grow the chunk toward MAX_CHUNK_ROWS so async
+// cells still overlap, while large values (e.g. long string group keys)
+// shrink it so in-flight results stay near CHUNK_BYTE_BUDGET instead of
+// scaling with row count.
+const INITIAL_CHUNK_ROWS = 64
+const MAX_CHUNK_ROWS = 4000
+const CHUNK_BYTE_BUDGET = 16 * 1024 * 1024
+
+/**
+ * Approximate retained bytes of an evaluated value. Strings dominate the
+ * workloads where size matters; everything else counts as a small constant.
+ *
+ * @param {unknown} value
+ * @returns {number}
+ */
+function valueBytes(value) {
+  if (typeof value === 'string') return value.length * 2
+  if (Array.isArray(value)) {
+    let bytes = 0
+    for (const item of value) bytes += valueBytes(item)
+    return bytes
+  }
+  return 16
+}
+
+/**
+ * Evaluates a value for every row and folds each result in row order, holding
+ * at most one adaptively sized chunk of evaluated values. Rows in a chunk are
+ * dispatched together so async cells overlap, and the chunk boundary yields
+ * to the event loop so aborts can fire. Chunks are bounded by bytes, not row
+ * count: a fixed 4000-row chunk of ~90KB string keys would hold hundreds of
+ * megabytes of results at once.
+ *
+ * @template T
+ * @param {Object} options
+ * @param {AsyncRow[]} options.rows
+ * @param {(row: AsyncRow, index: number) => Promise<T>} options.evaluate
+ * @param {(value: T, index: number) => void} options.fold
+ * @param {AbortSignal} [options.signal]
+ * @returns {Promise<void>}
+ */
+export async function foldEvaluatedRows({ rows, evaluate, fold, signal }) {
+  let chunkSize = INITIAL_CHUNK_ROWS
+  /** @type {Promise<T>[]} */
+  const pending = []
+  for (let start = 0; start < rows.length;) {
+    if (start > 0) {
+      await yieldToEventLoop()
+      signal?.throwIfAborted()
+    }
+    const end = Math.min(start + chunkSize, rows.length)
+    pending.length = end - start
+    for (let i = start; i < end; i++) {
+      pending[i - start] = evaluate(rows[i], i)
+    }
+    const values = await Promise.all(pending)
+    let bytes = 0
+    for (let j = 0; j < values.length; j++) {
+      bytes += valueBytes(values[j])
+      fold(values[j], start + j)
+    }
+    start = end
+    const bytesPerRow = Math.max(1, bytes / values.length)
+    chunkSize = Math.min(MAX_CHUNK_ROWS, Math.max(1, Math.floor(CHUNK_BYTE_BUDGET / bytesPerRow)))
+  }
+}

package/src/execute/streamingAggregate.js

@@ -1,10 +1,11 @@
 import { selectedRowCount, valueAt } from '../backend/batch.js'
 import { derivedAlias } from '../expression/alias.js'
 import { compileBatchExpression } from '../expression/batch.js'
-import { evaluateAll, evaluateExpr } from '../expression/evaluate.js'
+import { evaluateExpr } from '../expression/evaluate.js'
 import { collectColumnsFromExpr } from '../plan/columns.js'
 import { isAggregateFunc } from '../validation/functions.js'
 import { finalizeAccumulator, newAccumulator, updateAccumulator } from './accumulator.js'
+import { foldEvaluatedRows } from './fold.js'
 import { referencesRowScope } from './rowScope.js'
 import { sortEntriesByTerms } from './sort.js'
 import { keyify } from './utils.js'
@@ -342,9 +343,11 @@ function substituteValues(node, values) {
 }
 
 /**
- * Folds one chunk of rows into the group accumulators. Group keys, FILTER
- * conditions, and aggregate arguments are each evaluated across the whole
- * chunk so async cells overlap; the chunk is released afterwards.
+ * Folds one chunk of rows into the group accumulators. Each row's group keys,
+ * FILTER conditions, and aggregate arguments are dispatched together so async
+ * cells overlap, and folded in row order as they resolve, so evaluated values
+ * (which can be large strings) are bounded by bytes instead of being held for
+ * the whole chunk.
  *
  * @param {object} options
  * @param {AsyncRow[]} options.chunk
@@ -355,72 +358,88 @@ function substituteValues(node, values) {
  * @param {ExecuteContext} options.context
  * @returns {Promise<void>}
  */
-async function accumulateChunk({ chunk, groupBy, specs, groups, needsRow, context }) {
-  /** @type {SqlPrimitive[][] | undefined} */
-  let keyColumns
-  if (groupBy.length) {
-    keyColumns = await Promise.all(groupBy.map(expr => evaluateAll(expr, chunk, context)))
+function accumulateChunk({ chunk, groupBy, specs, groups, needsRow, context }) {
+  /**
+   * @param {SqlPrimitive[]} keyValues
+   * @param {number} index
+   * @returns {StreamingGroup}
+   */
+  function newGroup(keyValues, index) {
+    return {
+      firstRow: needsRow ? chunk[index] : undefined,
+      keyValues,
+      accumulators: specs.map(spec => newAccumulator(spec.funcName, spec.node.distinct)),
+    }
   }
 
-  /** @type {(SqlPrimitive[] | undefined)[]} */
-  const filters = new Array(specs.length)
-  /** @type {(SqlPrimitive[] | undefined)[]} */
-  const args = new Array(specs.length)
-  for (let s = 0; s < specs.length; s++) {
-    const { node, star } = specs[s]
-    if (node.filter) {
-      const passes = await evaluateAll(node.filter, chunk, context)
-      filters[s] = passes
-      if (!star) {
-        // The buffered path filters the group before evaluating arguments,
-        // so only evaluate the argument for rows that pass the FILTER
-        /** @type {AsyncRow[]} */
-        const passingRows = []
-        /** @type {number[]} */
-        const passingIndices = []
-        for (let j = 0; j < chunk.length; j++) {
-          if (passes[j]) {
-            passingRows.push(chunk[j])
-            passingIndices.push(j)
-          }
+  // Fast path: a single group key and only bare star aggregates (the common
+  // COUNT(*) GROUP BY x shape) skip the per-row tuple wrapper so synchronous
+  // cells stay cheap.
+  const singleExpr = groupBy.length === 1 && specs.every(spec => spec.star && !spec.node.filter)
+    ? groupBy[0] : null
+  if (singleExpr) {
+    return foldEvaluatedRows({
+      rows: chunk,
+      signal: context.signal,
+      evaluate: row => evaluateExpr({ node: singleExpr, row, context }),
+      fold(value, index) {
+        const key = keyify(value)
+        let group = groups.get(key)
+        if (!group) {
+          group = newGroup([value], index)
+          groups.set(key, group)
         }
-        const values = await evaluateAll(node.args[0], passingRows, context)
-        const spread = new Array(chunk.length).fill(null)
-        for (let k = 0; k < passingIndices.length; k++) {
-          spread[passingIndices[k]] = values[k]
+        for (let s = 0; s < specs.length; s++) {
+          if (specs[s].funcName === 'COUNT') group.accumulators[s].count++
+          else updateAccumulator(specs[s].funcName, group.accumulators[s], null)
         }
-        args[s] = spread
-      }
-    } else {
-      args[s] = star ? undefined : await evaluateAll(node.args[0], chunk, context)
-    }
+      },
+    })
   }
 
-  for (let j = 0; j < chunk.length; j++) {
-    const key = keyColumns
-      ? keyColumns.length === 1 ? keyify(keyColumns[0][j]) : keyify(...keyColumns.map(c => c[j]))
-      : true
-    let group = groups.get(key)
-    if (!group) {
-      group = {
-        firstRow: needsRow ? chunk[j] : undefined,
-        keyValues: keyColumns ? keyColumns.map(c => c[j]) : [],
-        accumulators: specs.map(spec => newAccumulator(spec.funcName, spec.node.distinct)),
+  return foldEvaluatedRows({
+    rows: chunk,
+    signal: context.signal,
+    // One flat tuple per row: group key values, then per spec its FILTER
+    // result and its argument. The argument is only evaluated for rows that
+    // pass the FILTER, matching the buffered path.
+    evaluate(row) {
+      const pending = groupBy.map(expr => evaluateExpr({ node: expr, row, context }))
+      for (const { node, star } of specs) {
+        if (node.filter) {
+          const passes = evaluateExpr({ node: node.filter, row, context })
+          pending.push(passes)
+          if (!star) {
+            pending.push(passes.then(pass => pass ? evaluateExpr({ node: node.args[0], row, context }) : null))
+          }
+        } else if (!star) {
+          pending.push(evaluateExpr({ node: node.args[0], row, context }))
+        }
       }
-      groups.set(key, group)
-    }
-    for (let s = 0; s < specs.length; s++) {
-      const filter = filters[s]
-      if (filter && !filter[j]) continue
-      const spec = specs[s]
-      if (spec.star && spec.funcName === 'COUNT') {
-        group.accumulators[s].count++
-      } else {
-        const arg = args[s]
-        updateAccumulator(spec.funcName, group.accumulators[s], arg ? arg[j] : null)
+      return Promise.all(pending)
+    },
+    fold(values, index) {
+      const key = groupBy.length === 0 ? true
+        : groupBy.length === 1 ? keyify(values[0]) : keyify(...values.slice(0, groupBy.length))
+      let group = groups.get(key)
+      if (!group) {
+        group = newGroup(values.slice(0, groupBy.length), index)
+        groups.set(key, group)
       }
-    }
-  }
+      let slot = groupBy.length
+      for (let s = 0; s < specs.length; s++) {
+        const spec = specs[s]
+        const passes = spec.node.filter ? values[slot++] : true
+        const arg = spec.star ? null : values[slot++]
+        if (!passes) continue
+        if (spec.star && spec.funcName === 'COUNT') {
+          group.accumulators[s].count++
+        } else {
+          updateAccumulator(spec.funcName, group.accumulators[s], arg)
+        }
+      }
+    },
+  })
 }
 
 /**

package/src/execute/window.js

@@ -1,5 +1,6 @@
 import { evaluateExpr } from '../expression/evaluate.js'
 import { executePlan } from './execute.js'
+import { foldEvaluatedRows } from './fold.js'
 import { compareForTerm, keyify } from './utils.js'
 import { yieldToEventLoop } from './yield.js'
 
@@ -116,30 +117,24 @@ export function executeWindow(plan, context) {
  * @param {ExecuteContext} context
  */
 async function computeWindow(spec, rows, output, context) {
-  // Bucket row indices by partition key.
+  // Bucket row indices by partition key. Keys are evaluated in adaptive
+  // chunks so async cells overlap while evaluated key values stay byte-bounded.
   /** @type {Map<string | number | bigint | boolean, number[]>} */
   const partitions = new Map()
-  for (let chunkStart = 0; chunkStart < rows.length; chunkStart += YIELD_INTERVAL) {
-    if (chunkStart > 0) {
-      await yieldToEventLoop()
-      context.signal?.throwIfAborted()
-    }
-    const chunkEnd = Math.min(chunkStart + YIELD_INTERVAL, rows.length)
-    const chunkKeys = await Promise.all(
-      rows.slice(chunkStart, chunkEnd).map(row =>
-        Promise.all(spec.partitionBy.map(expr => evaluateExpr({ node: expr, row, context })))
-      )
-    )
-    for (let j = 0; j < chunkKeys.length; j++) {
-      const key = keyify(...chunkKeys[j])
+  await foldEvaluatedRows({
+    rows,
+    signal: context.signal,
+    evaluate: row => Promise.all(spec.partitionBy.map(expr => evaluateExpr({ node: expr, row, context }))),
+    fold(keyValues, index) {
+      const key = keyify(...keyValues)
       let bucket = partitions.get(key)
       if (!bucket) {
         bucket = []
         partitions.set(key, bucket)
       }
-      bucket.push(chunkStart + j)
-    }
-  }
+      bucket.push(index)
+    },
+  })
 
   for (const bucket of partitions.values()) {
     context.signal?.throwIfAborted()

package/src/expression/regexp.js

@@ -1,3 +1,4 @@
+import { stringify } from '../execute/utils.js'
 import { ArgValueError } from '../validation/executionErrors.js'
 
 /**
@@ -187,5 +188,9 @@ export function evaluateRegexpLike({ node, args, rowIndex, cache }) {
       cache.regex = regex
     }
   }
-  return regex.test(String(string))
+  // Objects, arrays, and Dates test against their JSON text, the same
+  // coercion LIKE and CAST(x AS VARCHAR) use, so the three agree on
+  // JSON-typed columns. String() would collapse every object to
+  // '[object Object]', making REGEXP_LIKE a silent never-match on them.
+  return regex.test(typeof string === 'object' ? stringify(string) : String(string))
 }

package/src/index.d.ts

@@ -1,8 +1,10 @@
-import type { AsyncBatch, AsyncDataSource, AsyncRow, ColumnResult, ColumnVector, ExecuteContext, ExecuteSqlOptions, ExprNode, ParseSqlOptions, PlanSqlOptions, QueryPlan, QueryResults, ReadBatchColumnOptions, RowsToBatchesOptions, RowSelection, SqlPrimitive, Statement, Token } from './types.js'
+import type { AsyncBatch, AsyncDataSource, AsyncRow, PreparableDataSource, ScannableDataSource, ColumnResult, ColumnVector, ExecuteContext, ExecuteSqlOptions, ExprNode, ParseSqlOptions, PlanSqlOptions, QueryPlan, QueryResults, ReadBatchColumnOptions, RowsToBatchesOptions, RowSelection, SqlPrimitive, Statement, Token } from './types.js'
 export type {
   AsyncBatch,
   AsyncCells,
   AsyncDataSource,
+  PreparableDataSource,
+  ScannableDataSource,
   AsyncRow,
   BatchColumn,
   ColumnDemand,

package/src/parse/parse.js

@@ -274,11 +274,16 @@ function parseSelect(state) {
 
   if (match(state, 'keyword', 'GROUP')) {
     expect(state, 'keyword', 'BY')
-    while (true) {
-      const expr = resolvePositionalRef(parseExpression(state), columns, 'GROUP BY')
-      expectNoAggregate(expr, 'GROUP BY')
-      groupBy.push(expr)
-      if (!match(state, 'comma')) break
+    if (current(state).type === 'keyword' && current(state).value === 'ALL') {
+      const tok = consume(state)
+      groupBy.push(...expandGroupByAll(columns, tok.positionStart, tok.positionEnd))
+    } else {
+      while (true) {
+        const expr = resolvePositionalRef(parseExpression(state), columns, 'GROUP BY')
+        expectNoAggregate(expr, 'GROUP BY')
+        groupBy.push(expr)
+        if (!match(state, 'comma')) break
+      }
     }
   }
 
@@ -371,6 +376,30 @@ function parseSelect(state) {
   }
 }
 
+/**
+ * Expand GROUP BY ALL into every SELECT expression that contains no aggregate.
+ *
+ * @param {SelectColumn[]} columns
+ * @param {number} positionStart
+ * @param {number} positionEnd
+ * @returns {ExprNode[]}
+ */
+function expandGroupByAll(columns, positionStart, positionEnd) {
+  /** @type {ExprNode[]} */
+  const groupBy = []
+  for (const col of columns) {
+    if (col.type === 'star') {
+      throw new ParseError({
+        message: `GROUP BY ALL does not support * at position ${positionStart}`,
+        positionStart,
+        positionEnd,
+      })
+    }
+    if (!findAggregate(col.expr)) groupBy.push(col.expr)
+  }
+  return groupBy
+}
+
 /**
  * Resolves a positive integer literal in GROUP BY or ORDER BY as a positional
  * reference to the Nth SELECT column. Non-integer and non-literal expressions

package/src/types.d.ts

@@ -209,30 +209,33 @@ export type AsyncCell = () => Promise<SqlPrimitive>
 
 export type Row = Record<string, SqlPrimitive>[]
 
-/**
- * Async data source for streaming SQL execution. A source must implement
- * either scan() or prepareScan().
- */
 interface AsyncDataSourceBase {
   numRows?: number
   // Optional method for fast column scans
   scanColumn?(options: ScanColumnOptions): AsyncIterable<ArrayLike<SqlPrimitive>> | ScanColumnResults
 }
 
-export type AsyncDataSource = AsyncDataSourceBase & (
-  | {
-      columns: string[]
-      scan(options: ScanOptions): ScanResults
-      schema?: RelationSchema
-      prepareScan?: PrepareScan
-    }
-  | {
-      columns?: string[]
-      scan?(options: ScanOptions): ScanResults
-      schema: RelationSchema
-      prepareScan: PrepareScan
-    }
-)
+/** A data source that always provides column names and a row scan. */
+export interface ScannableDataSource extends AsyncDataSourceBase {
+  columns: string[]
+  scan(options: ScanOptions): ScanResults
+  schema?: RelationSchema
+  prepareScan?: PrepareScan
+}
+
+/** A data source that always provides a schema and a prepared batch scan. */
+export interface PreparableDataSource extends AsyncDataSourceBase {
+  columns?: string[]
+  scan?(options: ScanOptions): ScanResults
+  schema: RelationSchema
+  prepareScan: PrepareScan
+}
+
+/**
+ * Async data source for streaming SQL execution. A source must implement
+ * either scan() or prepareScan().
+ */
+export type AsyncDataSource = ScannableDataSource | PreparableDataSource
 
 /**
  * Result of a scan: streaming rows and flags indicating which hints were