Add native async columnar batch execution

src/backend/batch.js

@@ -1,6 +1,5 @@
 /**
- * @import { AsyncBatch, ColumnResult, ColumnVector, ReadBatchColumnOptions, RowSelection } from '../internalTypes.js'
- * @import { SqlPrimitive } from '../types.js'
+ * @import { AsyncBatch, ColumnResult, ColumnVector, ReadBatchColumnOptions, RowSelection, SqlPrimitive } from '../types.js'
  */
 
 /** @type {WeakMap<AsyncBatch, Map<number, Map<RowSelection, { resolved?: ColumnVector, pending: Map<AbortSignal | undefined, Promise<ColumnVector>> }>>>} */
@@ -126,7 +125,7 @@ export function readBatchColumn({ batch, columnIndex, selection = batch.selectio
   if (selection.length !== batch.selection.length) {
     throw new Error(`Selection length ${selection.length} does not match batch length ${batch.selection.length}`)
   }
-  if (column.type !== 'source' && column.type !== 'computed') {
+  if (!('read' in column)) {
     return selectVector(column, selection)
   }
 
@@ -149,15 +148,15 @@ export function readBatchColumn({ batch, columnIndex, selection = batch.selectio
   if (pending) return pending
   if (cache.resolved) return cache.resolved
 
-  const result = column.type === 'source'
-    ? column.read({ selection, signal })
-    : column.expression.evaluate({
-      batch: column.input,
-      selection,
-      signal,
-      rowOffset: column.rowOffset,
-      rowOrdinals: selectVector(column.rowOrdinals, selection),
-    })
+  const result = column.read({
+    batch: column.input ?? batch,
+    selection,
+    signal,
+    rowOffset: column.rowOffset,
+    rowOrdinals: column.rowOrdinals
+      ? selectVector(column.rowOrdinals, selection)
+      : undefined,
+  })
   const validated = validateColumnResult(result, selectedRowCount(selection))
   if (validated instanceof Promise) {
     const settled = validated.then(function cacheResolved(vector) {
@@ -185,7 +184,6 @@ export function readBatchColumn({ batch, columnIndex, selection = batch.selectio
 export function selectBatch(batch, selection) {
   const composed = composeSelections(batch.selection, selection)
   return {
-    columnNames: batch.columnNames,
     selection: composed,
     columns: batch.columns,
   }

src/backend/batchAdapters.js

@@ -2,8 +2,7 @@ import { readBatchColumn, selectVector, selectedRowCount, valueAt } from './batc
 import { yieldToEventLoop } from '../execute/yield.js'
 
 /**
- * @import { AsyncBatch, ColumnResult, ColumnVector, RowsToBatchesOptions } from '../internalTypes.js'
- * @import { AsyncCells, AsyncRow, SqlPrimitive } from '../types.js'
+ * @import { AsyncBatch, AsyncCells, AsyncRow, ColumnResult, ColumnVector, RowsToBatchesOptions, SqlPrimitive } from '../types.js'
  */
 
 const DEFAULT_BATCH_ROWS = 1024
@@ -33,13 +32,13 @@ export async function* rowsToBatches(rows, columnNames, options) {
     }
     rowCount++
     if (rowCount === batchRows) {
-      yield loadedBatch(columnNames, values, rowCount)
+      yield loadedBatch(values, rowCount)
       values = makeValueBuffers(columnNames.length)
       rowCount = 0
     }
   }
 
-  if (rowCount > 0) yield loadedBatch(columnNames, values, rowCount)
+  if (rowCount > 0) yield loadedBatch(values, rowCount)
   options?.signal?.throwIfAborted()
 }
 
@@ -48,14 +47,14 @@ export async function* rowsToBatches(rows, columnNames, options) {
  * remain lazy and are memoized once per batch/selection by `readBatchColumn`.
  *
  * @param {AsyncIterable<AsyncBatch>} batches
+ * @param {string[]} columns
  * @param {AbortSignal} [signal]
  * @yields {AsyncRow}
  */
-export async function* batchesToRows(batches, signal) {
+export async function* batchesToRows(batches, columns, signal) {
   for await (const batch of batches) {
-    const columns = batch.columnNames
     const loadedVectors = batch.columns.map(function loadedVector(column) {
-      if (column.type === 'source' || column.type === 'computed') return undefined
+      if ('read' in column) return undefined
       return selectVector(column, batch.selection)
     })
     const rowCount = selectedRowCount(batch.selection)
@@ -90,14 +89,14 @@ export async function* batchesToRows(batches, signal) {
  * constructing compatibility `AsyncRow` values.
  *
  * @param {AsyncIterable<AsyncBatch>} batches
+ * @param {string[]} names
  * @param {AbortSignal} [signal]
  * @returns {Promise<Record<string, SqlPrimitive>[]>}
  */
-export async function collectBatches(batches, signal) {
+export async function collectBatches(batches, names, signal) {
   /** @type {Record<string, SqlPrimitive>[]} */
   const rows = []
   for await (const batch of batches) {
-    const names = batch.columnNames
     const results = batch.columns.map(function readColumn(_column, columnIndex) {
       return readBatchColumn({ batch, columnIndex, signal })
     })
@@ -146,14 +145,12 @@ function makeValueBuffers(count) {
 }
 
 /**
- * @param {string[]} columnNames
  * @param {SqlPrimitive[][]} values
  * @param {number} rowCount
  * @returns {AsyncBatch}
  */
-function loadedBatch(columnNames, values, rowCount) {
+function loadedBatch(values, rowCount) {
   return {
-    columnNames,
     selection: { type: 'all', length: rowCount },
     columns: values.map(function loadedColumn(columnValues) {
       return { type: 'values', values: columnValues, length: rowCount }

src/backend/dataSource.js

@@ -18,6 +18,19 @@ export function asyncRow(obj, columns) {
   return { columns, cells, resolved: obj }
 }
 
+/**
+ * Returns a source's authoritative logical column names.
+ *
+ * @param {AsyncDataSource} source
+ * @returns {string[]}
+ */
+export function dataSourceColumns(source) {
+  if (source.prepareScan && source.schema) {
+    return source.schema.fields.map(function fieldName(field) { return field.name })
+  }
+  return source.columns ?? []
+}
+
 /**
  * Creates an async memory-backed data source from an array of plain objects
  *
@@ -78,14 +91,17 @@ export function memorySource({ data, columns }) {
  * @returns {AsyncDataSource}
  */
 export function cachedDataSource(source) {
+  if (source.prepareScan && source.schema) return source
+  const { scan } = source
+  if (!scan) return source
   /** @type {WeakMap<object, Map<string, Promise<SqlPrimitive>>>} */
   const cache = new WeakMap()
   return {
     ...source,
     scan(options) {
       // Does re-run the scan, but cache avoids re-computing expensive async cells
       // TODO: check cache first to avoid re-scanning when possible
-      const { rows, appliedWhere, appliedLimitOffset } = source.scan(options)
+      const { rows, appliedWhere, appliedLimitOffset } = scan.call(source, options)
 
       // Applied where clause changes which rows are returned so can't be cached
       if (appliedWhere && options.where) {

src/execute/aggregates.js

@@ -1,3 +1,4 @@
+import { dataSourceColumns } from '../backend/dataSource.js'
 import { derivedAlias } from '../expression/alias.js'
 import { evaluateExpr } from '../expression/evaluate.js'
 import { finalizeAccumulator, newAccumulator, updateAccumulator } from './accumulator.js'
@@ -299,7 +300,7 @@ function tryColumnScanAggregate(plan, { tables, signal }) {
 
   // COUNT(*) needs a physical column whose filtered chunk lengths can be
   // counted. Prefer a predicate/projection column, then any table column.
-  const starColumn = scanNode.hints.columns?.[0] ?? table.columns[0]
+  const starColumn = scanNode.hints.columns?.[0] ?? dataSourceColumns(table)[0]
   if (!starColumn) return
 
   // All columns must be simple aggregates on plain identifiers

src/execute/batchResults.js

@@ -1,15 +1,12 @@
 import { batchesToRows } from '../backend/batchAdapters.js'
+import { bindQuerySignal } from './utils.js'
 
 /**
- * @import { AsyncBatch, InternalBatchResults } from '../internalTypes.js'
- * @import { QueryResults } from '../types.js'
+ * @import { AsyncBatch, QueryResults } from '../types.js'
  */
 
-/** @type {WeakMap<QueryResults, InternalBatchResults>} */
-const internalBatches = new WeakMap()
-
 /**
- * Creates a public row result backed by private batches.
+ * Creates query results backed by batches.
  *
  * @param {object} options
  * @param {string[]} options.columns
@@ -19,23 +16,13 @@ const internalBatches = new WeakMap()
  * @param {AbortSignal} [options.signal]
  * @returns {QueryResults}
  */
-export function batchResult({ batches, signal, ...metadata }) {
+export function batchResult({ batches: readBatches, signal, ...metadata }) {
   const results = {
     ...metadata,
+    batches: readBatches,
     rows() {
-      return batchesToRows(batches(), signal)
+      return batchesToRows(readBatches(), metadata.columns, signal)
     },
   }
-  internalBatches.set(results, { columns: metadata.columns, batches, signal })
-  return results
-}
-
-/**
- * Returns the private batch execution path for a result, when available.
- *
- * @param {QueryResults} results
- * @returns {InternalBatchResults | undefined}
- */
-export function batchResultsFor(results) {
-  return internalBatches.get(results)
+  return bindQuerySignal(results, signal)
 }

src/execute/batches.js

@@ -3,8 +3,8 @@ import { keyify } from './utils.js'
 import { yieldToEventLoop } from './yield.js'
 
 /**
- * @import { AsyncBatch, BatchColumn, BatchProjection, ColumnResult, ColumnVector, CompiledBatchExpression, RowSelection } from '../internalTypes.js'
- * @import { SqlPrimitive } from '../types.js'
+ * @import { BatchProjection, CompiledBatchExpression } from '../internalTypes.js'
+ * @import { AsyncBatch, BatchColumn, ColumnResult, ColumnVector, ReadColumn, RowSelection, SqlPrimitive } from '../types.js'
  */
 
 const INITIAL_FILTER_WINDOW_ROWS = 256
@@ -180,25 +180,23 @@ export async function* distinctBatches(batches, signal) {
  * so output columns remain aligned without hidden dependency columns.
  *
  * @param {AsyncIterable<AsyncBatch>} batches
- * @param {string[]} columnNames
  * @param {readonly BatchProjection[]} projections
  * @yields {AsyncBatch}
  */
-export async function* projectExpressionBatches(batches, columnNames, projections) {
+export async function* projectExpressionBatches(batches, projections) {
   let rowOffset = 0
   for await (const batch of batches) {
     const currentRowOffset = rowOffset
     rowOffset += selectedRowCount(batch.selection)
     /** @type {ColumnVector | undefined} */
     let rowOrdinals
     yield {
-      columnNames,
       selection: batch.selection,
       columns: projections.map(function projectColumn(projection) {
         /** @type {BatchColumn} */
         let column
         if (projection.type === 'column') {
-          column = batch.columns[projection.columnIndex]
+          column = projectedBatchColumn(batch, projection.columnIndex)
         } else if (projection.type === 'constant') {
           column = {
             type: 'constant',
@@ -208,9 +206,8 @@ export async function* projectExpressionBatches(batches, columnNames, projection
         } else {
           rowOrdinals ??= selectionOrdinals(batch.selection)
           column = {
-            type: 'computed',
+            read: projection.expression.evaluate,
             input: batch,
-            expression: projection.expression,
             rowOffset: currentRowOffset,
             rowOrdinals,
           }
@@ -221,6 +218,23 @@ export async function* projectExpressionBatches(batches, columnNames, projection
   }
 }
 
+/**
+ * Preserves the source batch's deferred-read cache through projection.
+ *
+ * @param {AsyncBatch} batch
+ * @param {number} columnIndex
+ * @returns {BatchColumn}
+ */
+function projectedBatchColumn(batch, columnIndex) {
+  const column = batch.columns[columnIndex]
+  if (!('read' in column)) return column
+  /** @type {ReadColumn} */
+  function readProjectedColumn({ selection, signal }) {
+    return readBatchColumn({ batch, columnIndex, selection, signal })
+  }
+  return { read: readProjectedColumn }
+}
+
 /**
  * @param {ColumnVector} predicate
  * @param {number} rowCount