squirreling 0.14.0 → 0.15.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "squirreling",
-  "version": "0.14.0",
+  "version": "0.15.0",
   "description": "Squirreling Async SQL Engine",
   "author": "Hyperparam",
   "homepage": "https://hyperparam.app",
@@ -39,11 +39,11 @@
     "test": "vitest run"
   },
   "devDependencies": {
-    "@types/node": "26.1.0",
+    "@types/node": "26.1.1",
     "@vitest/coverage-v8": "4.1.10",
     "eslint": "9.39.4",
-    "eslint-plugin-jsdoc": "63.0.12",
-    "typescript": "6.0.3",
+    "eslint-plugin-jsdoc": "63.0.13",
+    "typescript": "7.0.2",
     "vitest": "4.1.10"
   }
 }

package/src/execute/aggregates.js

@@ -1,7 +1,8 @@
 import { derivedAlias } from '../expression/alias.js'
 import { evaluateExpr } from '../expression/evaluate.js'
 import { finalizeAccumulator, newAccumulator, updateAccumulator } from './accumulator.js'
-import { executePlan, selectColumnNames } from './execute.js'
+import { executePlan, executeScan, selectColumnNames } from './execute.js'
+import { normalizeScanColumnResult } from './scanColumn.js'
 import { sortEntriesByTerms } from './sort.js'
 import { planStreamingAggregates, streamingHashAggregateRows, streamingScalarAggregateRows } from './streamingAggregate.js'
 import { keyify } from './utils.js'
@@ -195,17 +196,21 @@ export function executeHashAggregate(plan, context) {
  */
 export function executeScalarAggregate(plan, context) {
   // Fast path: use scanColumn when available
-  const fast = tryColumnScanAggregate(plan, context)
-  if (fast) {
+  const columnScan = tryColumnScanAggregate(plan, context)
+  if (columnScan?.rows) {
     return {
       columns: selectColumnNames(plan.columns, []),
       numRows: 1,
       maxRows: 1,
-      rows: fast,
+      rows: columnScan.rows,
     }
   }
 
-  const child = executePlan({ plan: plan.child, context })
+  // A declined pushdown still returned a usable one-column scan. Transfer it
+  // to the ordinary scan path instead of calling the source a second time.
+  const child = columnScan?.fallback
+    ? executeScan(columnScan.fallback.plan, context, columnScan.fallback.result)
+    : executePlan({ plan: plan.child, context })
   const streaming = planStreamingAggregates(plan, child.columns)
   if (streaming) {
     return {
@@ -263,6 +268,7 @@ export function executeScalarAggregate(plan, context) {
  *   column: string,
  *   alias: string,
  *   distinct?: boolean,
+ *   star?: boolean,
  * }} ColumnAggSpec
  */
 
@@ -272,7 +278,13 @@ export function executeScalarAggregate(plan, context) {
  *
  * @param {ScalarAggregateNode} plan
  * @param {ExecuteContext} context
- * @returns {(() => AsyncGenerator<AsyncRow>) | undefined}
+ * @returns {{
+ *   rows?: () => AsyncGenerator<AsyncRow>,
+ *   fallback?: {
+ *     plan: import('../plan/types.js').ScanNode,
+ *     result: import('../types.js').ScanColumnResults,
+ *   },
+ * } | undefined}
  */
 function tryColumnScanAggregate(plan, { tables, signal }) {
   // No HAVING support in fast path
@@ -281,61 +293,96 @@ function tryColumnScanAggregate(plan, { tables, signal }) {
   if (plan.child.type !== 'Scan') return
   const scanNode = plan.child
   const { limit, offset, where } = scanNode.hints
-  // scanColumn doesn't support filtering
-  if (where) return
 
   const table = tables[scanNode.table]
   if (!table?.scanColumn) return
 
+  // 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]
+  if (!starColumn) return
+
   // All columns must be simple aggregates on plain identifiers
   /** @type {ColumnAggSpec[]} */
   const specs = []
   for (const col of plan.columns) {
     if (col.type !== 'derived') return
-    const spec = extractColumnAggSpec(col)
+    const spec = extractColumnAggSpec(col, starColumn)
     if (!spec) return
     specs.push(spec)
   }
 
-  return async function* () {
-    /** @type {string[]} */
-    const columns = []
-    /** @type {AsyncCells} */
-    const cells = {}
-
-    // Group specs by column so each column is scanned at most once no matter how
-    // many aggregates read it (e.g. MIN(x), MAX(x), AVG(x) share one pass).
-    /** @type {Map<string, ColumnAggSpec[]>} */
-    const specsByColumn = new Map()
-    for (const spec of specs) {
-      const group = specsByColumn.get(spec.column)
-      if (group) group.push(spec)
-      else specsByColumn.set(spec.column, [spec])
+  const physicalColumns = new Set(specs.map(spec => spec.column))
+  const hasPushdown = where || limit !== undefined || offset !== undefined
+  if (hasPushdown) {
+    // If negotiation fails, the returned column stream can only replace the
+    // ordinary scan when that scan needs the same single column. Otherwise a
+    // probe could allocate work that no correct fallback is able to consume.
+    const scanColumns = scanNode.hints.columns
+    if (physicalColumns.size !== 1 || scanColumns?.length !== 1 ||
+      !physicalColumns.has(scanColumns[0])) return
+  }
+
+  // Ask once per physical column and retain the returned chunks. Sources
+  // report applied hints just like scan(); declined hints use the normal path.
+  /** @type {Map<string, import('../types.js').ScanColumnResults>} */
+  const columnScans = new Map()
+  for (const { column } of specs) {
+    if (columnScans.has(column)) continue
+    const options = { column, where, limit, offset, signal }
+    const result = normalizeScanColumnResult(table.scanColumn(options), options)
+    if (where && !result.appliedWhere) {
+      return { fallback: { plan: scanNode, result } }
+    }
+    if ((limit !== undefined || offset !== undefined) && !result.appliedLimitOffset) {
+      return { fallback: { plan: scanNode, result } }
     }
+    columnScans.set(column, result)
+  }
 
-    // Each column's single pass is computed once and shared by all its cells;
-    // a column is only scanned if one of its aggregates is actually read.
-    /** @type {Map<string, Promise<Map<string, SqlPrimitive>>>} */
-    const passes = new Map()
-    /**
-     * @param {string} column
-     * @returns {Promise<Map<string, SqlPrimitive>>}
-     */
-    function scanOnce(column) {
-      let pass = passes.get(column)
-      if (!pass) {
-        pass = scanColumnGroup({ table, specs: specsByColumn.get(column) ?? [], limit, offset, signal })
-        passes.set(column, pass)
+  return {
+    async *rows() {
+      /** @type {string[]} */
+      const columns = []
+      /** @type {AsyncCells} */
+      const cells = {}
+
+      // Group specs by column so each column is scanned at most once no matter how
+      // many aggregates read it (e.g. MIN(x), MAX(x), AVG(x) share one pass).
+      /** @type {Map<string, ColumnAggSpec[]>} */
+      const specsByColumn = new Map()
+      for (const spec of specs) {
+        const group = specsByColumn.get(spec.column)
+        if (group) group.push(spec)
+        else specsByColumn.set(spec.column, [spec])
       }
-      return pass
-    }
 
-    for (const spec of specs) {
-      columns.push(spec.alias)
-      cells[spec.alias] = async () => (await scanOnce(spec.column)).get(spec.alias) ?? null
-    }
+      // Each column's single pass is computed once and shared by all its cells;
+      // a column is only scanned if one of its aggregates is actually read.
+      /** @type {Map<string, Promise<Map<string, SqlPrimitive>>>} */
+      const passes = new Map()
+      /**
+       * @param {string} column
+       * @returns {Promise<Map<string, SqlPrimitive>>}
+       */
+      function scanOnce(column) {
+        let pass = passes.get(column)
+        if (!pass) {
+          const scan = columnScans.get(column)
+          if (!scan) throw new Error(`missing scanColumn result for ${column}`)
+          pass = scanColumnGroup({ values: scan.chunks(), specs: specsByColumn.get(column) ?? [], signal })
+          passes.set(column, pass)
+        }
+        return pass
+      }
 
-    yield { columns, cells }
+      for (const spec of specs) {
+        columns.push(spec.alias)
+        cells[spec.alias] = async () => (await scanOnce(spec.column)).get(spec.alias) ?? null
+      }
+
+      yield { columns, cells }
+    },
   }
 }
 
@@ -344,16 +391,27 @@ function tryColumnScanAggregate(plan, { tables, signal }) {
  * Returns undefined if the expression is not a supported simple aggregate.
  *
  * @param {DerivedColumn} col
+ * @param {string} starColumn
  * @returns {ColumnAggSpec | undefined}
  */
-function extractColumnAggSpec({ expr, alias }) {
+function extractColumnAggSpec({ expr, alias }, starColumn) {
   if (expr.type !== 'function') return
   if (expr.filter) return // FILTER not supported in fast path
   const funcName = expr.funcName.toUpperCase()
   if (!['COUNT', 'SUM', 'AVG', 'MIN', 'MAX'].includes(funcName)) return
 
-  // Argument must be a plain column identifier
+  // Argument must be a plain column identifier, except COUNT(*), which counts
+  // the lengths of filtered chunks from an arbitrary physical column.
   const arg = expr.args[0]
+  if (arg.type === 'star') {
+    if (funcName !== 'COUNT' || expr.distinct) return
+    return {
+      funcName,
+      column: starColumn,
+      alias: alias ?? derivedAlias(expr),
+      star: true,
+    }
+  }
   if (arg.type !== 'identifier') return
   return {
     funcName,
@@ -368,26 +426,23 @@ function extractColumnAggSpec({ expr, alias }) {
  * All specs share the one scanColumn walk, so MIN(x)/MAX(x)/AVG(x) decode x once.
  *
  * @param {Object} options
- * @param {AsyncDataSource} options.table
+ * @param {AsyncIterable<ArrayLike<SqlPrimitive>>} options.values
  * @param {ColumnAggSpec[]} options.specs - aggregates over the same column
- * @param {number} [options.limit]
- * @param {number} [options.offset]
  * @param {AbortSignal} [options.signal]
  * @returns {Promise<Map<string, SqlPrimitive>>} alias → aggregate value
  */
-async function scanColumnGroup({ table, specs, limit, offset, signal }) {
-  const { column } = specs[0]
-  const values = table.scanColumn({ column, limit, offset, signal })
-
+async function scanColumnGroup({ values, specs, signal }) {
   const accs = specs.map(spec => ({ spec, acc: newAccumulator(spec.funcName, spec.distinct) }))
 
   for await (const chunk of values) {
     signal?.throwIfAborted()
     for (let i = 0; i < chunk.length; i++) {
       const v = chunk[i]
-      if (v == null) continue
       for (const { spec, acc } of accs) {
-        updateAccumulator(spec.funcName, acc, v)
+        // COUNT(*) counts matching rows even when the arbitrary carrier column
+        // selected for scanColumn is null. Other aggregates ignore nulls.
+        if (spec.star) acc.count++
+        else if (v != null) updateAccumulator(spec.funcName, acc, v)
       }
     }
   }

package/src/execute/execute.js

@@ -7,6 +7,7 @@ import { statementScope } from '../plan/columns.js'
 import { validateScan, validateTable } from '../validation/tables.js'
 import { executeHashAggregate, executeScalarAggregate } from './aggregates.js'
 import { executeHashJoin, executeNestedLoopJoin, executePositionalJoin } from './join.js'
+import { normalizeScanColumnResult } from './scanColumn.js'
 import { executeSort } from './sort.js'
 import { addBounds, minBounds, stableRowKey } from './utils.js'
 import { executeWindow } from './window.js'
@@ -269,40 +270,63 @@ export function selectColumnNames(selectColumns, childColumns) {
 /**
  * @param {ScanNode} plan
  * @param {ExecuteContext} context
+ * @param {import('../types.js').ScanColumnResults} [existingColumnResult]
  * @returns {QueryResults}
  */
-function executeScan(plan, context) {
+export function executeScan(plan, context, existingColumnResult) {
   const { tables, signal } = context
   const table = validateTable({ ...plan, tables })
   validateScan({ ...plan, tables })
   const hasLimitOffset = plan.hints.limit !== undefined || plan.hints.offset // 0 offset is noop
 
-  // Fast path: single column scan without WHERE
-  if (table.scanColumn && plan.hints.columns?.length === 1 && !plan.hints.where) {
+  // Fast path: single column scan. As with scan(), hints the source did not
+  // apply are handled by the engine over the returned column values.
+  const scanColumnOptions = plan.hints.columns?.length === 1
+    ? plan.hints.where
+      // Do not push a filtered range until WHERE is known to be applied: an
+      // older source may ignore WHERE but eagerly apply LIMIT/OFFSET.
+      ? { column: plan.hints.columns[0], where: plan.hints.where, signal }
+      : { column: plan.hints.columns[0], ...plan.hints, signal }
+    : undefined
+  const columnResult = existingColumnResult ?? (table.scanColumn && scanColumnOptions
+    ? normalizeScanColumnResult(table.scanColumn(scanColumnOptions), scanColumnOptions)
+    : undefined)
+  if (columnResult && scanColumnOptions) {
     const column = plan.hints.columns[0]
-    const chunks = table.scanColumn({
-      column,
-      limit: plan.hints.limit,
-      offset: plan.hints.offset,
-      signal,
-    })
     const scanRows = computeScanRows(table.numRows, plan.hints.limit, plan.hints.offset)
     return {
       columns: [column],
-      numRows: scanRows,
+      numRows: plan.hints.where ? undefined : scanRows,
       maxRows: scanRows,
       async *rows() {
         const columns = [column]
-        for await (const chunk of chunks) {
-          signal?.throwIfAborted()
-          for (let i = 0; i < chunk.length; i++) {
-            const value = chunk[i]
-            yield {
-              columns,
-              cells: { [column]: () => Promise.resolve(value) },
+        let result = (async function* () {
+          // Creating the chunk stream may itself start I/O, so leave it until
+          // the returned rows are actually consumed.
+          for await (const chunk of columnResult.chunks()) {
+            signal?.throwIfAborted()
+            for (let i = 0; i < chunk.length; i++) {
+              const value = chunk[i]
+              yield {
+                columns,
+                cells: { [column]: () => Promise.resolve(value) },
+              }
             }
           }
+        })()
+
+        if (!columnResult.appliedWhere && plan.hints.where) {
+          result = filterRows(result, plan.hints.where, context, plan.hints.limit)
         }
+        // Filtered LIMIT/OFFSET was intentionally not passed to scanColumn.
+        const appliedLimitOffset = plan.hints.where
+          ? false
+          : columnResult.appliedLimitOffset
+        if (!appliedLimitOffset && hasLimitOffset) {
+          result = limitRows(result, plan.hints.limit, plan.hints.offset, signal)
+        }
+
+        yield* result
         // A data source may end its stream cooperatively on abort; surface
         // the abort so a truncated scan is not mistaken for a complete one
         signal?.throwIfAborted()

package/src/execute/scanColumn.js

@@ -0,0 +1,17 @@
+/**
+ * Normalizes scanColumn's legacy AsyncIterable return into ScanColumnResults.
+ * Legacy implementations are valid for unfiltered scans only: they cannot
+ * claim to have applied a WHERE predicate they predate.
+ *
+ * @param {AsyncIterable<ArrayLike<import('../types.js').SqlPrimitive>> | import('../types.js').ScanColumnResults} result
+ * @param {import('../types.js').ScanColumnOptions} options
+ * @returns {import('../types.js').ScanColumnResults}
+ */
+export function normalizeScanColumnResult(result, options) {
+  if ('chunks' in result) return result
+  return {
+    chunks: () => result,
+    appliedWhere: !options.where,
+    appliedLimitOffset: !options.where,
+  }
+}

package/src/types.d.ts

@@ -79,7 +79,7 @@ export interface AsyncDataSource {
   columns: string[]
   scan(options: ScanOptions): ScanResults
   // Optional method for fast column scans
-  scanColumn?(options: ScanColumnOptions): AsyncIterable<ArrayLike<SqlPrimitive>>
+  scanColumn?(options: ScanColumnOptions): AsyncIterable<ArrayLike<SqlPrimitive>> | ScanColumnResults
 }
 
 /**
@@ -113,11 +113,19 @@ export interface ScanOptions {
  */
 export interface ScanColumnOptions {
   column: string
+  where?: ExprNode
   limit?: number
   offset?: number
   signal?: AbortSignal
 }
 
+/** Result of a column scan, mirroring the ScanResults hint flags. */
+export interface ScanColumnResults {
+  chunks(): AsyncIterable<ArrayLike<SqlPrimitive>>
+  appliedWhere: boolean
+  appliedLimitOffset: boolean
+}
+
 export interface FunctionSignature {
   min: number
   max?: number