squirreling 0.15.2 → 0.16.0

package/src/execute/streamingAggregate.js

@@ -1,4 +1,6 @@
+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 { collectColumnsFromExpr } from '../plan/columns.js'
 import { isAggregateFunc } from '../validation/functions.js'
@@ -8,7 +10,8 @@ import { keyify } from './utils.js'
 import { yieldToEventLoop } from './yield.js'
 
 /**
- * @import { AsyncCells, AsyncRow, ExecuteContext, ExprNode, FunctionNode, IdentifierNode, QueryResults, SelectColumn, SqlPrimitive } from '../types.js'
+ * @import { BatchAggregateInputs, CompiledBatchExpression } from '../internalTypes.js'
+ * @import { AsyncBatch, AsyncCells, AsyncRow, ColumnVector, ExecuteContext, ExprNode, FunctionNode, IdentifierNode, QueryResults, SelectColumn, SqlPrimitive } from '../types.js'
  * @import { HashAggregateNode, ScalarAggregateNode } from '../plan/types.js'
  * @import { Accumulator } from './accumulator.js'
  */
@@ -419,6 +422,141 @@ async function accumulateChunk({ chunk, groupBy, specs, groups, needsRow, contex
   }
 }
 
+/**
+ * Compiles the expressions an aggregate reads from each input batch. A
+ * filtered non-star aggregate retains the row path because its argument must
+ * only be evaluated for passing rows.
+ *
+ * @param {ExprNode[]} groupBy
+ * @param {StreamingAggSpec[]} specs
+ * @param {readonly string[]} columns
+ * @param {ExecuteContext} context
+ * @returns {BatchAggregateInputs | undefined}
+ */
+function compileBatchAggregateInputs(groupBy, specs, columns, context) {
+  /** @type {CompiledBatchExpression[]} */
+  const keys = []
+  for (const expression of groupBy) {
+    if (referencesRowScope(expression, context)) return undefined
+    const key = compileBatchExpression(expression, columns)
+    if (!key) return undefined
+    keys.push(key)
+  }
+
+  /** @type {(CompiledBatchExpression | undefined)[]} */
+  const filters = []
+  /** @type {(CompiledBatchExpression | undefined)[]} */
+  const args = []
+  for (const spec of specs) {
+    if (spec.node.filter && !spec.star) return undefined
+    if (spec.node.filter && referencesRowScope(spec.node.filter, context)) return undefined
+    if (!spec.star && referencesRowScope(spec.node.args[0], context)) return undefined
+    const filter = spec.node.filter
+      ? compileBatchExpression(spec.node.filter, columns)
+      : undefined
+    const argument = spec.star
+      ? undefined
+      : compileBatchExpression(spec.node.args[0], columns)
+    if (spec.node.filter && !filter || !spec.star && !argument) return undefined
+    filters.push(filter)
+    args.push(argument)
+  }
+  return {
+    keys,
+    filters,
+    args,
+  }
+}
+
+/**
+ * Returns whether an expression reads a qualified identifier whose table
+ * scope the batch compiler cannot distinguish from struct-field access.
+ *
+ * @param {ExprNode} expression
+ * @param {ExecuteContext} context
+ * @returns {boolean}
+ */
+function referencesRowScope(expression, context) {
+  /** @type {IdentifierNode[]} */
+  const identifiers = []
+  collectColumnsFromExpr(expression, identifiers)
+  return identifiers.some(function isScopedReference(identifier) {
+    return Boolean(identifier.prefix && (
+      context.scope?.includes(identifier.prefix) || context.outerAliases?.has(identifier.prefix)
+    ))
+  })
+}
+
+/**
+ * Resolves one set of compiled expressions against a batch.
+ *
+ * @param {(CompiledBatchExpression | undefined)[]} expressions
+ * @param {AsyncBatch} batch
+ * @param {ExecuteContext} context
+ * @param {number} rowOffset
+ * @returns {Promise<(ColumnVector | undefined)[]>}
+ */
+function evaluateBatchInputs(expressions, batch, context, rowOffset) {
+  return Promise.all(expressions.map(function evaluateInput(expression) {
+    return expression?.evaluate({ batch, selection: batch.selection, signal: context.signal, rowOffset })
+  }))
+}
+
+/**
+ * Folds a native batch directly into aggregate state without constructing
+ * rows, cells, or per-value promises.
+ *
+ * @param {object} options
+ * @param {AsyncBatch} options.batch
+ * @param {BatchAggregateInputs} options.inputs
+ * @param {StreamingAggSpec[]} options.specs
+ * @param {Map<unknown, StreamingGroup>} options.groups
+ * @param {ExecuteContext} options.context
+ * @param {number} options.rowOffset
+ * @returns {Promise<void>}
+ */
+async function accumulateBatch({ batch, inputs, specs, groups, context, rowOffset }) {
+  const [keys, filters, args] = await Promise.all([
+    evaluateBatchInputs(inputs.keys, batch, context, rowOffset),
+    evaluateBatchInputs(inputs.filters, batch, context, rowOffset),
+    evaluateBatchInputs(inputs.args, batch, context, rowOffset),
+  ])
+  const rowCount = selectedRowCount(batch.selection)
+  for (let rowIndex = 0; rowIndex < rowCount; rowIndex++) {
+    if (rowIndex > 0 && rowIndex % CHUNK_SIZE === 0) {
+      await yieldToEventLoop()
+      context.signal?.throwIfAborted()
+    }
+    const keyValues = keys.map(function keyValue(vector) {
+      if (!vector) throw new Error('Missing compiled group key')
+      return valueAt(vector, rowIndex)
+    })
+    const key = keyValues.length === 0
+      ? true
+      : keyValues.length === 1 ? keyify(keyValues[0]) : keyify(...keyValues)
+    let group = groups.get(key)
+    if (!group) {
+      group = {
+        firstRow: undefined,
+        keyValues,
+        accumulators: specs.map(spec => newAccumulator(spec.funcName, spec.node.distinct)),
+      }
+      groups.set(key, group)
+    }
+    for (let specIndex = 0; specIndex < specs.length; specIndex++) {
+      const filter = filters[specIndex]
+      if (filter && !valueAt(filter, rowIndex)) continue
+      const spec = specs[specIndex]
+      if (spec.star && spec.funcName === 'COUNT') {
+        group.accumulators[specIndex].count++
+      } else {
+        const argument = args[specIndex]
+        updateAccumulator(spec.funcName, group.accumulators[specIndex], argument ? valueAt(argument, rowIndex) : null)
+      }
+    }
+  }
+}
+
 /**
  * Consumes the child rows into per-group accumulators, holding at most one
  * chunk of rows at a time. Throws when aborted so partial accumulators are
@@ -435,6 +573,18 @@ async function accumulateChunk({ chunk, groupBy, specs, groups, needsRow, contex
 async function accumulateGroups({ child, groupBy, specs, needsRow, context }) {
   /** @type {Map<unknown, StreamingGroup>} */
   const groups = new Map()
+  const batchInputs = child.batches && !needsRow
+    ? compileBatchAggregateInputs(groupBy, specs, child.columns, context)
+    : undefined
+  if (batchInputs && child.batches) {
+    let rowOffset = 0
+    for await (const batch of child.batches()) {
+      await accumulateBatch({ batch, inputs: batchInputs, specs, groups, context, rowOffset })
+      rowOffset += selectedRowCount(batch.selection)
+      context.signal?.throwIfAborted()
+    }
+    return groups
+  }
   /** @type {AsyncRow[]} */
   let chunk = []
   for await (const row of child.rows()) {

package/src/execute/utils.js

@@ -1,9 +1,42 @@
+import { collectBatches } from '../backend/batchAdapters.js'
+
 /**
  * @import { AsyncRow, OrderByItem, QueryResults, SqlPrimitive } from '../types.js'
  */
 
 const primitiveTypes = new Set(['number', 'bigint', 'boolean', 'string'])
 
+/** @type {WeakMap<QueryResults, AbortSignal>} */
+const querySignals = new WeakMap()
+
+/**
+ * Associates an execution signal with results without expanding the public
+ * result shape solely for the collection adapter.
+ *
+ * @param {QueryResults} results
+ * @param {AbortSignal} [signal]
+ * @returns {QueryResults}
+ */
+export function bindQuerySignal(results, signal) {
+  if (!signal || querySignals.get(results) === signal) return results
+
+  const { batches, rows } = results
+  results.rows = async function* boundRows() {
+    signal.throwIfAborted()
+    yield* rows.call(results)
+    signal.throwIfAborted()
+  }
+  if (batches) {
+    results.batches = async function* boundBatches() {
+      signal.throwIfAborted()
+      yield* batches.call(results)
+      signal.throwIfAborted()
+    }
+  }
+  querySignals.set(results, signal)
+  return results
+}
+
 /**
  * Compares two values for a single ORDER BY term, handling nulls and direction
  *
@@ -49,6 +82,10 @@ export function compareForTerm(a, b, term) {
  * @returns {Promise<Record<string, SqlPrimitive>[]>} array of all yielded values
  */
 export async function collect(results) {
+  if (results.batches) {
+    return await collectBatches(results.batches(), results.columns, querySignals.get(results))
+  }
+
   // Collect all rows first, then materialize cells concurrently
   // This enables dataloader-style batching of cell accessors
   /** @type {AsyncRow[]} */