squirreling 0.15.2 → 0.16.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "squirreling",
-  "version": "0.15.2",
+  "version": "0.16.0",
   "description": "Squirreling Async SQL Engine",
   "author": "Hyperparam",
   "homepage": "https://hyperparam.app",
@@ -42,7 +42,7 @@
     "@types/node": "26.2.0",
     "@vitest/coverage-v8": "4.1.10",
     "eslint": "9.39.4",
-    "eslint-plugin-jsdoc": "64.1.0",
+    "eslint-plugin-jsdoc": "64.2.1",
     "typescript": "7.0.2",
     "vitest": "4.1.10"
   }

package/src/backend/batch.js

@@ -0,0 +1,236 @@
+/**
+ * @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>> }>>>} */
+const batchCache = new WeakMap()
+
+/**
+ * Returns the number of selected rows.
+ *
+ * @param {RowSelection} selection
+ * @returns {number}
+ */
+export function selectedRowCount(selection) {
+  if (selection.type === 'all') return selection.length
+  if (selection.type === 'range') return selection.end - selection.start
+  return selection.indices.length
+}
+
+/**
+ * Creates base-aligned ordinal values for the rows in a selection.
+ *
+ * @param {RowSelection} selection
+ * @returns {ColumnVector}
+ */
+export function selectionOrdinals(selection) {
+  const values = new Uint32Array(selection.length)
+  const length = selectedRowCount(selection)
+  for (let ordinal = 0; ordinal < length; ordinal++) {
+    values[selectionIndexAt(selection, ordinal)] = ordinal
+  }
+  return { type: 'typed', values, length: selection.length }
+}
+
+/**
+ * Composes a selection over an already-selected domain with the selection
+ * that produced that domain.
+ *
+ * @param {RowSelection} outer - selection over the original base domain
+ * @param {RowSelection} inner - selection over the rows selected by `outer`
+ * @returns {RowSelection}
+ */
+export function composeSelections(outer, inner) {
+  const outerCount = selectedRowCount(outer)
+  if (inner.length !== outerCount) {
+    throw new Error(`Cannot compose selection of length ${inner.length} over ${outerCount} rows`)
+  }
+  if (inner.type === 'all') return outer
+  if (outer.type === 'all') return { ...inner, length: outer.length }
+
+  if (outer.type === 'range' && inner.type === 'range') {
+    return {
+      type: 'range',
+      start: outer.start + inner.start,
+      end: outer.start + inner.end,
+      length: outer.length,
+    }
+  }
+
+  const indices = new Uint32Array(selectedRowCount(inner))
+  for (let i = 0; i < indices.length; i++) {
+    indices[i] = selectionIndexAt(outer, selectionIndexAt(inner, i))
+  }
+  return { type: 'indices', indices, length: outer.length }
+}
+
+/**
+ * Reads one logical value from a vector.
+ *
+ * @param {ColumnVector} vector
+ * @param {number} index
+ * @returns {SqlPrimitive}
+ */
+export function valueAt(vector, index) {
+  if (!Number.isInteger(index) || index < 0 || index >= vector.length) {
+    throw new RangeError(`Column index ${index} is outside vector length ${vector.length}`)
+  }
+  if (vector.type === 'values') return vector.values[index]
+  if (vector.type === 'typed') {
+    if (vector.validity && vector.validity[index] === 0) return null
+    return vector.values[index]
+  }
+  if (vector.type === 'constant') return vector.value
+  return valueAt(vector.source, selectionIndexAt(vector.selection, index))
+}
+
+/**
+ * Creates a zero-copy view of a vector through a row selection.
+ *
+ * @param {ColumnVector} vector
+ * @param {RowSelection} selection
+ * @returns {ColumnVector}
+ */
+export function selectVector(vector, selection) {
+  if (selection.length !== vector.length) {
+    throw new Error(`Cannot select ${selection.length} rows from vector length ${vector.length}`)
+  }
+  if (selection.type === 'all') return vector
+  if (vector.type === 'selected') {
+    const composed = composeSelections(vector.selection, selection)
+    return {
+      type: 'selected',
+      source: vector.source,
+      selection: composed,
+      length: selectedRowCount(composed),
+    }
+  }
+  return {
+    type: 'selected',
+    source: vector,
+    selection,
+    length: selectedRowCount(selection),
+  }
+}
+
+/**
+ * Resolves one batch column for an effective selection.
+ *
+ * @param {ReadBatchColumnOptions} options
+ * @returns {ColumnResult}
+ */
+export function readBatchColumn({ batch, columnIndex, selection = batch.selection, signal }) {
+  const column = batch.columns[columnIndex]
+  if (!column) throw new RangeError(`Column index ${columnIndex} is outside batch columns`)
+  if (selection.length !== batch.selection.length) {
+    throw new Error(`Selection length ${selection.length} does not match batch length ${batch.selection.length}`)
+  }
+  if (!('read' in column)) {
+    return selectVector(column, selection)
+  }
+
+  let batchResults = batchCache.get(batch)
+  if (!batchResults) {
+    batchResults = new Map()
+    batchCache.set(batch, batchResults)
+  }
+  let columnResults = batchResults.get(columnIndex)
+  if (!columnResults) {
+    columnResults = new Map()
+    batchResults.set(columnIndex, columnResults)
+  }
+  let cache = columnResults.get(selection)
+  if (!cache) {
+    cache = { pending: new Map() }
+    columnResults.set(selection, cache)
+  }
+  const pending = cache.pending.get(signal)
+  if (pending) return pending
+  if (cache.resolved) return cache.resolved
+
+  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) {
+      cache.pending.delete(signal)
+      cache.resolved ??= vector
+      return vector
+    }, function evictRejected(error) {
+      cache.pending.delete(signal)
+      throw error
+    })
+    cache.pending.set(signal, settled)
+    return settled
+  }
+  cache.resolved = validated
+  return validated
+}
+
+/**
+ * Selects rows from a batch without reading or copying its columns.
+ *
+ * @param {AsyncBatch} batch
+ * @param {RowSelection} selection - selection over the batch's current rows
+ * @returns {AsyncBatch}
+ */
+export function selectBatch(batch, selection) {
+  const composed = composeSelections(batch.selection, selection)
+  return {
+    selection: composed,
+    columns: batch.columns,
+  }
+}
+
+/**
+ * Returns a base-domain row index for a logical selected-row index.
+ *
+ * @param {RowSelection} selection
+ * @param {number} index
+ * @returns {number}
+ */
+function selectionIndexAt(selection, index) {
+  const count = selectedRowCount(selection)
+  if (!Number.isInteger(index) || index < 0 || index >= count) {
+    throw new RangeError(`Selection index ${index} is outside selected length ${count}`)
+  }
+  if (selection.type === 'all') return index
+  if (selection.type === 'range') return selection.start + index
+  return selection.indices[index]
+}
+
+/**
+ * Validates the resolved vector length without forcing synchronous column
+ * implementations through a promise.
+ *
+ * @param {ColumnResult} result
+ * @param {number} expectedLength
+ * @returns {ColumnResult}
+ */
+function validateColumnResult(result, expectedLength) {
+  if (result instanceof Promise) {
+    return result.then(function validateResolvedVector(vector) {
+      return validateVectorLength(vector, expectedLength)
+    })
+  }
+  return validateVectorLength(result, expectedLength)
+}
+
+/**
+ * @param {ColumnVector} vector
+ * @param {number} expectedLength
+ * @returns {ColumnVector}
+ */
+function validateVectorLength(vector, expectedLength) {
+  if (vector.length !== expectedLength) {
+    throw new Error(`Column returned ${vector.length} rows, expected ${expectedLength}`)
+  }
+  return vector
+}

package/src/backend/batchAdapters.js

@@ -0,0 +1,178 @@
+import { readBatchColumn, selectVector, selectedRowCount, valueAt } from './batch.js'
+import { yieldToEventLoop } from '../execute/yield.js'
+
+/**
+ * @import { AsyncBatch, AsyncCells, AsyncRow, ColumnResult, ColumnVector, RowsToBatchesOptions, SqlPrimitive } from '../types.js'
+ */
+
+const DEFAULT_BATCH_ROWS = 1024
+
+/**
+ * Materializes a legacy row stream into aligned batches. This is a source or
+ * public compatibility boundary, not an operator implementation strategy.
+ *
+ * @param {AsyncIterable<AsyncRow>} rows
+ * @param {string[]} columnNames
+ * @param {RowsToBatchesOptions} [options]
+ * @yields {AsyncBatch}
+ */
+export async function* rowsToBatches(rows, columnNames, options) {
+  const batchRows = options?.batchRows ?? DEFAULT_BATCH_ROWS
+  if (!Number.isInteger(batchRows) || batchRows <= 0) {
+    throw new RangeError(`batchRows must be a positive integer, got ${batchRows}`)
+  }
+  let values = makeValueBuffers(columnNames.length)
+  let rowCount = 0
+
+  for await (const row of rows) {
+    options?.signal?.throwIfAborted()
+    for (let columnIndex = 0; columnIndex < columnNames.length; columnIndex++) {
+      const name = columnNames[columnIndex]
+      values[columnIndex].push(await readRowCell(row, name))
+    }
+    rowCount++
+    if (rowCount === batchRows) {
+      yield loadedBatch(values, rowCount)
+      values = makeValueBuffers(columnNames.length)
+      rowCount = 0
+    }
+  }
+
+  if (rowCount > 0) yield loadedBatch(values, rowCount)
+  options?.signal?.throwIfAborted()
+}
+
+/**
+ * Exposes batches through the legacy lazy-cell row interface. Column reads
+ * 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, columns, signal) {
+  for await (const batch of batches) {
+    const loadedVectors = batch.columns.map(function loadedVector(column) {
+      if ('read' in column) return undefined
+      return selectVector(column, batch.selection)
+    })
+    const rowCount = selectedRowCount(batch.selection)
+    for (let rowIndex = 0; rowIndex < rowCount; rowIndex++) {
+      signal?.throwIfAborted()
+      /** @type {AsyncCells} */
+      const cells = {}
+      for (let columnIndex = 0; columnIndex < columns.length; columnIndex++) {
+        const currentColumn = columnIndex
+        const currentRow = rowIndex
+        const loadedVector = loadedVectors[columnIndex]
+        if (loadedVector) {
+          const value = valueAt(loadedVector, currentRow)
+          cells[columns[columnIndex]] = function readLoadedCell() {
+            return Promise.resolve(value)
+          }
+        } else {
+          cells[columns[columnIndex]] = async function readCell() {
+            const vector = await readBatchColumn({ batch, columnIndex: currentColumn, signal })
+            return valueAt(vector, currentRow)
+          }
+        }
+      }
+      yield { columns, cells }
+    }
+  }
+  signal?.throwIfAborted()
+}
+
+/**
+ * Collects batches directly into the existing object-row result shape without
+ * constructing compatibility `AsyncRow` values.
+ *
+ * @param {AsyncIterable<AsyncBatch>} batches
+ * @param {string[]} names
+ * @param {AbortSignal} [signal]
+ * @returns {Promise<Record<string, SqlPrimitive>[]>}
+ */
+export async function collectBatches(batches, names, signal) {
+  /** @type {Record<string, SqlPrimitive>[]} */
+  const rows = []
+  for await (const batch of batches) {
+    const results = batch.columns.map(function readColumn(_column, columnIndex) {
+      return readBatchColumn({ batch, columnIndex, signal })
+    })
+    const vectors = await resolveColumns(results)
+    const rowCount = selectedRowCount(batch.selection)
+    for (let rowIndex = 0; rowIndex < rowCount; rowIndex++) {
+      if (signal && rowIndex % 4000 === 0) {
+        if (rowIndex > 0) await yieldToEventLoop()
+        signal.throwIfAborted()
+      }
+      /** @type {Record<string, SqlPrimitive>} */
+      const row = {}
+      for (let columnIndex = 0; columnIndex < names.length; columnIndex++) {
+        row[names[columnIndex]] = valueAt(vectors[columnIndex], rowIndex)
+      }
+      rows.push(row)
+    }
+  }
+  signal?.throwIfAborted()
+  return rows
+}
+
+/**
+ * @param {AsyncRow} row
+ * @param {string} name
+ * @returns {Promise<SqlPrimitive>}
+ */
+async function readRowCell(row, name) {
+  if (row.resolved && Object.prototype.hasOwnProperty.call(row.resolved, name)) {
+    return row.resolved[name]
+  }
+  const cell = row.cells[name]
+  if (!cell) throw new Error(`Row does not contain column "${name}"`)
+  return await cell()
+}
+
+/**
+ * @param {number} count
+ * @returns {SqlPrimitive[][]}
+ */
+function makeValueBuffers(count) {
+  /** @type {SqlPrimitive[][]} */
+  const values = []
+  for (let i = 0; i < count; i++) values.push([])
+  return values
+}
+
+/**
+ * @param {SqlPrimitive[][]} values
+ * @param {number} rowCount
+ * @returns {AsyncBatch}
+ */
+function loadedBatch(values, rowCount) {
+  return {
+    selection: { type: 'all', length: rowCount },
+    columns: values.map(function loadedColumn(columnValues) {
+      return { type: 'values', values: columnValues, length: rowCount }
+    }),
+  }
+}
+
+/**
+ * Avoids a promise boundary when every column is already loaded synchronously.
+ *
+ * @param {ColumnResult[]} results
+ * @returns {ColumnVector[] | Promise<ColumnVector[]>}
+ */
+function resolveColumns(results) {
+  if (results.some(function isPromise(result) { return result instanceof Promise })) {
+    return Promise.all(results)
+  }
+  /** @type {ColumnVector[]} */
+  const vectors = []
+  for (const result of results) {
+    if (result instanceof Promise) throw new Error('Unexpected asynchronous column result')
+    vectors.push(result)
+  }
+  return vectors
+}

package/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,6 +91,9 @@ 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 {
@@ -85,7 +101,7 @@ export function cachedDataSource(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) {

package/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

package/src/execute/batchResults.js

@@ -0,0 +1,28 @@
+import { batchesToRows } from '../backend/batchAdapters.js'
+import { bindQuerySignal } from './utils.js'
+
+/**
+ * @import { AsyncBatch, QueryResults } from '../types.js'
+ */
+
+/**
+ * Creates query results backed by batches.
+ *
+ * @param {object} options
+ * @param {string[]} options.columns
+ * @param {number} [options.numRows]
+ * @param {number} [options.maxRows]
+ * @param {() => AsyncIterable<AsyncBatch>} options.batches
+ * @param {AbortSignal} [options.signal]
+ * @returns {QueryResults}
+ */
+export function batchResult({ batches: readBatches, signal, ...metadata }) {
+  const results = {
+    ...metadata,
+    batches: readBatches,
+    rows() {
+      return batchesToRows(readBatches(), metadata.columns, signal)
+    },
+  }
+  return bindQuerySignal(results, signal)
+}