squirreling 0.15.2 → 0.16.0

package/src/execute/execute.js

@@ -1,11 +1,16 @@
-import { memorySource } from '../backend/dataSource.js'
+import { selectedRowCount } from '../backend/batch.js'
+import { batchesToRows } from '../backend/batchAdapters.js'
+import { dataSourceColumns, memorySource } from '../backend/dataSource.js'
 import { derivedAlias } from '../expression/alias.js'
+import { compileBatchExpression } from '../expression/batch.js'
 import { evaluateExpr } from '../expression/evaluate.js'
 import { parseSql } from '../parse/parse.js'
 import { planSql, planStatement } from '../plan/plan.js'
-import { statementScope } from '../plan/columns.js'
+import { collectColumnsFromExpr, statementScope } from '../plan/columns.js'
 import { validateScan, validateTable } from '../validation/tables.js'
 import { executeHashAggregate, executeScalarAggregate } from './aggregates.js'
+import { batchResult } from './batchResults.js'
+import { distinctBatches, filterBatches, limitBatches, projectExpressionBatches } from './batches.js'
 import { executeHashJoin, executeNestedLoopJoin, executePositionalJoin } from './join.js'
 import { normalizeScanColumnResult } from './scanColumn.js'
 import { executeSort } from './sort.js'
@@ -14,7 +19,8 @@ import { executeWindow } from './window.js'
 import { yieldToEventLoop } from './yield.js'
 
 /**
- * @import { AsyncCells, AsyncDataSource, AsyncRow, DerivedColumn, ExecuteContext, ExecuteSqlOptions, ExprNode, IdentifierNode, QueryResults, SelectColumn, SqlPrimitive, Statement } from '../types.js'
+ * @import { BatchProjection, CompiledBatchExpression } from '../internalTypes.js'
+ * @import { AsyncBatch, AsyncCells, AsyncDataSource, AsyncRow, ColumnDemand, ColumnVector, DerivedColumn, ExecuteContext, ExecuteSqlOptions, ExprNode, IdentifierNode, PreparedScan, QueryResults, RelationSchema, ScanRequest, SelectColumn, SqlPrimitive, Statement } from '../types.js'
  * @import { CountNode, DistinctNode, FilterNode, LimitNode, ProjectNode, QueryPlan, ScanNode, SetOperationNode, TableFunctionNode } from '../plan/types.js'
  */
 
@@ -278,6 +284,12 @@ export function executeScan(plan, context, existingColumnResult) {
   const table = validateTable({ ...plan, tables })
   validateScan({ ...plan, tables })
   const hasLimitOffset = plan.hints.limit !== undefined || plan.hints.offset // 0 offset is noop
+  const scanContext = { ...context, scope: [plan.alias ?? plan.table] }
+
+  if (!existingColumnResult && table.prepareScan && table.schema) {
+    const prepared = table.prepareScan(scanRequest(plan, table.schema, scanContext))
+    return executePreparedScan({ plan, prepared, context: scanContext, table })
+  }
 
   // Fast path: single column scan. As with scan(), hints the source did not
   // apply are handled by the engine over the returned column values.
@@ -294,9 +306,43 @@ export function executeScan(plan, context, existingColumnResult) {
   if (columnResult && scanColumnOptions) {
     const column = plan.hints.columns[0]
     const scanRows = computeScanRows(table.numRows, plan.hints.limit, plan.hints.offset)
+    const columns = [column]
+    const appliedLimitOffset = plan.hints.where
+      ? false
+      : columnResult.appliedLimitOffset
+    const residualFilter = plan.hints.where && !columnResult.appliedWhere
+      ? compileBatchExpression(plan.hints.where, columns)
+      : undefined
+
+    /** @returns {AsyncIterable<AsyncBatch>} */
+    function makeBatches() {
+      /** @type {AsyncIterable<AsyncBatch>} */
+      let batches = columnBatches(columnResult.chunks(), signal)
+      if (residualFilter) {
+        const targetRows = plan.hints.limit === undefined
+          ? undefined
+          : plan.hints.limit + (plan.hints.offset ?? 0)
+        batches = filterBatches(batches, residualFilter, signal, targetRows)
+      }
+      if (!appliedLimitOffset && hasLimitOffset) {
+        batches = limitBatches(batches, plan.hints.limit, plan.hints.offset, signal)
+      }
+      return batches
+    }
+
+    const canUseBatches = !plan.hints.where || columnResult.appliedWhere || residualFilter !== undefined
+    if (canUseBatches) {
+      return batchResult({
+        columns,
+        numRows: plan.hints.where ? undefined : scanRows,
+        maxRows: scanRows,
+        batches: makeBatches,
+        signal,
+      })
+    }
     return {
-      columns: [column],
-      numRows: plan.hints.where ? undefined : scanRows,
+      columns,
+      numRows: undefined,
       maxRows: scanRows,
       async *rows() {
         const columns = [column]
@@ -315,13 +361,8 @@ export function executeScan(plan, context, existingColumnResult) {
           }
         })()
 
-        if (!columnResult.appliedWhere && plan.hints.where) {
-          result = filterRows(result, plan.hints.where, context, plan.hints.limit)
-        }
+        result = filterRows(result, plan.hints.where, scanContext, 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)
         }
@@ -334,6 +375,10 @@ export function executeScan(plan, context, existingColumnResult) {
     }
   }
 
+  if (!table.scan) {
+    throw new Error(`Data source "${plan.table}" does not implement scan()`)
+  }
+
   // do the scan
   const scanResult = table.scan({ ...plan.hints, signal })
   const { appliedWhere, appliedLimitOffset } = scanResult
@@ -345,7 +390,7 @@ export function executeScan(plan, context, existingColumnResult) {
 
   const scanRows = computeScanRows(table.numRows, plan.hints.limit, plan.hints.offset)
   return {
-    columns: plan.hints.columns ?? table.columns,
+    columns: plan.hints.columns ?? dataSourceColumns(table),
     numRows: !plan.hints.where ? scanRows : undefined,
     maxRows: scanRows,
     async *rows() {
@@ -353,7 +398,7 @@ export function executeScan(plan, context, existingColumnResult) {
 
       // Apply WHERE if data source did not
       if (!appliedWhere && plan.hints.where) {
-        result = filterRows(result, plan.hints.where, context, plan.hints.limit)
+        result = filterRows(result, plan.hints.where, scanContext, plan.hints.limit)
       }
 
       // Apply LIMIT/OFFSET if data source did not
@@ -370,6 +415,137 @@ export function executeScan(plan, context, existingColumnResult) {
   }
 }
 
+/**
+ * Executes a prepared native-batch scan and applies only the residual work
+ * reported by the source.
+ *
+ * @param {Object} options
+ * @param {ScanNode} options.plan
+ * @param {PreparedScan} options.prepared
+ * @param {ExecuteContext} options.context
+ * @param {AsyncDataSource} options.table
+ * @returns {QueryResults}
+ */
+function executePreparedScan({ plan, prepared, context, table }) {
+  const { signal } = context
+  const { residual, properties, schema } = prepared
+  const hasRequestedRange = plan.hints.limit !== undefined || Boolean(plan.hints.offset)
+  if (residual.filter && hasRequestedRange && (
+    residual.limit !== plan.hints.limit || (residual.offset ?? 0) !== (plan.hints.offset ?? 0)
+  )) {
+    throw new Error(`Data source "${plan.table}" applied limit/offset without applying where`)
+  }
+  const columns = schema.fields.map(function fieldName(field) { return field.name })
+  const residualFilter = residual.filter
+    ? compileUnscopedBatchExpression(residual.filter, columns, context)
+    : undefined
+  const canUseBatches = !residual.filter || residualFilter !== undefined
+
+  /** @returns {AsyncIterable<AsyncBatch>} */
+  function makeBatches() {
+    /** @type {AsyncIterable<AsyncBatch>} */
+    let batches = prepared.batches({ signal })
+    if (residualFilter) {
+      const targetRows = residual.limit === undefined
+        ? undefined
+        : residual.limit + (residual.offset ?? 0)
+      batches = filterBatches(batches, residualFilter, signal, targetRows)
+    }
+    if (residual.limit !== undefined || residual.offset) {
+      batches = limitBatches(batches, residual.limit, residual.offset, signal)
+    }
+    return batches
+  }
+
+  const exactRows = properties.exactRows === undefined
+    ? computeScanRows(plan.hints.where ? undefined : table.numRows, plan.hints.limit, plan.hints.offset)
+    : computeScanRows(properties.exactRows, residual.limit, residual.offset)
+  const preparedMaxRows = properties.maxRows ?? properties.exactRows
+  const maxRows = preparedMaxRows === undefined
+    ? computeScanRows(table.numRows, plan.hints.limit, plan.hints.offset)
+    : computeScanRows(preparedMaxRows, residual.limit, residual.offset)
+  const metadata = {
+    columns,
+    numRows: residual.filter ? undefined : exactRows,
+    maxRows,
+  }
+  if (canUseBatches) {
+    return batchResult({ ...metadata, batches: makeBatches, signal })
+  }
+  return {
+    ...metadata,
+    async *rows() {
+      let result = batchesToRows(prepared.batches({ signal }), columns, signal)
+      if (residual.filter) result = filterRows(result, residual.filter, context, residual.limit)
+      if (residual.limit !== undefined || residual.offset) {
+        result = limitRows(result, residual.limit, residual.offset, signal)
+      }
+      yield* result
+      signal?.throwIfAborted()
+    },
+  }
+}
+
+/**
+ * Builds a generic demand schedule from the scan's logical columns. Predicate
+ * fields are required in phase zero; remaining output fields stay deferred.
+ *
+ * @param {ScanNode} plan
+ * @param {RelationSchema} schema
+ * @param {ExecuteContext} context
+ * @returns {ScanRequest}
+ */
+function scanRequest(plan, schema, context) {
+  const currentScope = context.scope ?? [plan.table]
+  /** @type {IdentifierNode[]} */
+  const predicateIdentifiers = []
+  collectColumnsFromExpr(plan.hints.where, predicateIdentifiers, undefined, {
+    cteColumns: context.cteColumns,
+    tables: context.tables,
+    outerAliases: new Set([...currentScope, ...context.outerAliases ?? []]),
+  })
+  const predicateNames = new Set()
+  for (const identifier of predicateIdentifiers) {
+    if (!identifier.prefix) {
+      predicateNames.add(identifier.name)
+      continue
+    }
+    if (currentScope.includes(identifier.prefix)) {
+      predicateNames.add(identifier.name)
+      continue
+    }
+    if (context.outerAliases?.has(identifier.prefix)) continue
+    const baseField = schema.fields.some(function fieldOwnsPrefix(field) {
+      return field.name === identifier.prefix
+    })
+    if (baseField) predicateNames.add(identifier.prefix)
+  }
+  const requestedNames = plan.hints.columns
+    ? [...plan.hints.columns]
+    : schema.fields.map(function fieldName(field) { return field.name })
+  for (const name of predicateNames) {
+    if (!requestedNames.includes(name)) requestedNames.push(name)
+  }
+  /** @type {ColumnDemand[]} */
+  const columns = requestedNames.map(function columnDemand(name) {
+    const field = schema.fields.find(function fieldName(candidate) { return candidate.name === name })
+    if (!field) throw new Error(`Prepared source schema does not contain column "${name}"`)
+    const predicate = predicateNames.has(name)
+    return {
+      field: field.id,
+      phase: predicate ? 0 : 1,
+      purpose: predicate ? 'filter' : 'output',
+      mode: predicate ? 'required' : 'deferred',
+    }
+  })
+  return {
+    columns,
+    filter: plan.hints.where,
+    limit: plan.hints.limit,
+    offset: plan.hints.offset,
+  }
+}
+
 /**
  * Executes a Count node using numRows when available, falling back to scan
  *
@@ -392,7 +568,24 @@ function executeCount(plan, context) {
         // Use source numRows if available
         if (table.numRows !== undefined) return table.numRows
 
+        if (table.prepareScan && table.schema) {
+          const prepared = table.prepareScan({ columns: [] })
+          if (prepared.properties.exactRows !== undefined) {
+            return prepared.properties.exactRows
+          }
+          let count = 0
+          for await (const batch of prepared.batches({ signal })) {
+            signal?.throwIfAborted()
+            count += selectedRowCount(batch.selection)
+          }
+          signal?.throwIfAborted()
+          return count
+        }
+
         // Fall back to counting rows via scan
+        if (!table.scan) {
+          throw new Error(`Data source "${plan.table}" does not implement scan()`)
+        }
         let count = 0
         const { rows } = table.scan({ signal })
         // eslint-disable-next-line no-unused-vars
@@ -512,6 +705,41 @@ async function* limitRows(rows, limit = Infinity, offset = 0, signal) {
   }
 }
 
+/**
+ * Compiles a batch expression only when qualified identifiers do not depend
+ * on current or outer row scope.
+ *
+ * @param {ExprNode} expression
+ * @param {string[]} columns
+ * @param {ExecuteContext} context
+ * @returns {CompiledBatchExpression | undefined}
+ */
+function compileUnscopedBatchExpression(expression, columns, context) {
+  return referencesRowScope(expression, columns, context)
+    ? undefined
+    : compileBatchExpression(expression, columns)
+}
+
+/**
+ * Returns whether an expression reads a qualified identifier from row scope.
+ *
+ * @param {ExprNode} expression
+ * @param {string[]} columns
+ * @param {ExecuteContext} context
+ * @returns {boolean}
+ */
+function referencesRowScope(expression, columns, context) {
+  /** @type {IdentifierNode[]} */
+  const identifiers = []
+  collectColumnsFromExpr(expression, identifiers)
+  return identifiers.some(function scopedIdentifier(identifier) {
+    return Boolean(identifier.prefix && (
+      context.outerAliases?.has(identifier.prefix) ||
+      context.scope?.includes(identifier.prefix) && columns.includes(identifier.prefix)
+    ))
+  })
+}
+
 /**
  * Executes a filter operation (WHERE clause)
  *
@@ -521,6 +749,22 @@ async function* limitRows(rows, limit = Infinity, offset = 0, signal) {
  */
 function executeFilter(plan, context) {
   const child = executePlan({ plan: plan.child, context })
+  const expression = child.batches
+    ? compileUnscopedBatchExpression(plan.condition, child.columns, context)
+    : undefined
+  if (expression && child.batches) {
+    const readChildBatches = child.batches
+    /** @returns {AsyncIterable<AsyncBatch>} */
+    function makeBatches() {
+      return filterBatches(readChildBatches(), expression, context.signal)
+    }
+    return batchResult({
+      columns: child.columns,
+      maxRows: child.maxRows,
+      batches: makeBatches,
+      signal: context.signal,
+    })
+  }
   return {
     columns: child.columns,
     maxRows: child.maxRows,
@@ -539,9 +783,32 @@ function executeProject(plan, context) {
   const child = executePlan({ plan: plan.child, context })
   const columns = selectColumnNames(plan.columns, child.columns)
 
-  const resolveable = plan.columns.every(col =>
-    col.type === 'star' || col.type === 'derived' && col.expr.type === 'identifier'
-  )
+  const resolveable = plan.columns.every(function resolvesDirectly(column) {
+    if (column.type === 'star') return true
+    if (column.expr.type !== 'identifier') return false
+    const sourceName = column.expr.prefix
+      ? `${column.expr.prefix}.${column.expr.name}`
+      : column.expr.name
+    return child.columns.includes(sourceName)
+  })
+
+  const projection = child.batches
+    ? batchProjection(plan.columns, columns, child.columns, context)
+    : undefined
+  if (projection && child.batches) {
+    const readChildBatches = child.batches
+    /** @returns {AsyncIterable<AsyncBatch>} */
+    function makeBatches() {
+      return projectExpressionBatches(readChildBatches(), projection)
+    }
+    return batchResult({
+      columns,
+      numRows: child.numRows,
+      maxRows: child.maxRows,
+      batches: makeBatches,
+      signal: context.signal,
+    })
+  }
 
   return {
     columns,
@@ -631,6 +898,19 @@ function executeProject(plan, context) {
  */
 function executeDistinct(plan, context) {
   const child = executePlan({ plan: plan.child, context })
+  if (child.batches) {
+    const readChildBatches = child.batches
+    /** @returns {AsyncIterable<AsyncBatch>} */
+    function makeBatches() {
+      return distinctBatches(readChildBatches(), context.signal)
+    }
+    return batchResult({
+      columns: child.columns,
+      maxRows: child.maxRows,
+      batches: makeBatches,
+      signal: context.signal,
+    })
+  }
   return {
     columns: child.columns,
     maxRows: child.maxRows,
@@ -689,6 +969,20 @@ function executeDistinct(plan, context) {
  */
 function executeLimit(plan, context) {
   const child = executePlan({ plan: plan.child, context })
+  if (child.batches) {
+    const readChildBatches = child.batches
+    /** @returns {AsyncIterable<AsyncBatch>} */
+    function makeBatches() {
+      return limitBatches(readChildBatches(), plan.limit, plan.offset, context.signal)
+    }
+    return batchResult({
+      columns: child.columns,
+      numRows: computeScanRows(child.numRows, plan.limit, plan.offset),
+      maxRows: computeScanRows(child.maxRows, plan.limit, plan.offset),
+      batches: makeBatches,
+      signal: context.signal,
+    })
+  }
   return {
     columns: child.columns,
     numRows: computeScanRows(child.numRows, plan.limit, plan.offset),
@@ -697,6 +991,132 @@ function executeLimit(plan, context) {
   }
 }
 
+/**
+ * Wraps one-column source chunks as loaded batches without copying their
+ * arrays. Each source chunk remains the async scheduling unit.
+ *
+ * @param {AsyncIterable<ArrayLike<SqlPrimitive>>} chunks
+ * @param {AbortSignal} [signal]
+ * @yields {AsyncBatch}
+ */
+async function* columnBatches(chunks, signal) {
+  for await (const chunk of chunks) {
+    signal?.throwIfAborted()
+    const vector = vectorFromChunk(chunk)
+    /** @type {AsyncBatch} */
+    const batch = {
+      selection: { type: 'all', length: vector.length },
+      columns: [vector],
+    }
+    yield batch
+  }
+  signal?.throwIfAborted()
+}
+
+/**
+ * @param {ArrayLike<SqlPrimitive>} chunk
+ * @returns {ColumnVector}
+ */
+function vectorFromChunk(chunk) {
+  if (Array.isArray(chunk)) {
+    return { type: 'values', values: chunk, length: chunk.length }
+  }
+  if (isNumericArray(chunk)) {
+    return {
+      type: 'typed',
+      values: chunk,
+      length: chunk.length,
+    }
+  }
+  return { type: 'values', values: Array.from(chunk), length: chunk.length }
+}
+
+/**
+ * @param {ArrayLike<SqlPrimitive>} values
+ * @returns {values is import('../types.js').NumericArray}
+ */
+function isNumericArray(values) {
+  return values instanceof Int8Array
+    || values instanceof Uint8Array
+    || values instanceof Uint8ClampedArray
+    || values instanceof Int16Array
+    || values instanceof Uint16Array
+    || values instanceof Int32Array
+    || values instanceof Uint32Array
+    || values instanceof Float32Array
+    || values instanceof Float64Array
+    || values instanceof BigInt64Array
+    || values instanceof BigUint64Array
+}
+
+/**
+ * Compiles a projection to direct, constant, or computed batch columns.
+ *
+ * @param {SelectColumn[]} planColumns
+ * @param {string[]} outputColumns
+ * @param {string[]} childColumns
+ * @param {ExecuteContext} context
+ * @returns {BatchProjection[] | undefined}
+ */
+function batchProjection(planColumns, outputColumns, childColumns, context) {
+  /** @type {BatchProjection[]} */
+  const projections = []
+  for (const column of planColumns) {
+    if (column.type === 'star') {
+      const prefix = column.table ? `${column.table}.` : undefined
+      for (let i = 0; i < childColumns.length; i++) {
+        if (!prefix || childColumns[i].startsWith(prefix)) {
+          projections.push({ type: 'column', columnIndex: i })
+        }
+      }
+      continue
+    }
+
+    if (referencesRowScope(column.expr, childColumns, context)) return undefined
+
+    if (column.expr.type === 'literal') {
+      projections.push({ type: 'constant', value: column.expr.value })
+      continue
+    }
+    if (column.expr.type === 'identifier') {
+      const index = identifierColumnIndex(column.expr, childColumns)
+      if (index !== undefined) {
+        projections.push({ type: 'column', columnIndex: index })
+        continue
+      }
+      const expression = compileBatchExpression(column.expr, childColumns)
+      if (!expression) return undefined
+      projections.push({ type: 'expression', expression })
+      continue
+    }
+    const expression = compileBatchExpression(column.expr, childColumns)
+    if (!expression) return undefined
+    projections.push({ type: 'expression', expression })
+  }
+  if (projections.length !== outputColumns.length) return undefined
+
+  return projections
+}
+
+/**
+ * @param {IdentifierNode} identifier
+ * @param {string[]} childColumns
+ * @returns {number | undefined}
+ */
+function identifierColumnIndex(identifier, childColumns) {
+  const sourceName = identifier.prefix
+    ? `${identifier.prefix}.${identifier.name}`
+    : identifier.name
+  const index = childColumns.lastIndexOf(sourceName)
+  if (index >= 0) return index
+
+  const suffix = `.${identifier.name}`
+  const matches = childColumns
+    .map(function childColumn(name, childIndex) { return { name, childIndex } })
+    .filter(function suffixMatch(candidate) { return candidate.name.endsWith(suffix) })
+  return matches.length === 1 ? matches[0].childIndex : undefined
+}
+
 /**
  * Executes a set operation (UNION, INTERSECT, EXCEPT)
  *
@@ -706,11 +1126,13 @@ function executeLimit(plan, context) {
  */
 function executeSetOperation(plan, context) {
   const { signal } = context
+  const leftContext = plan.leftScope === undefined ? context : { ...context, scope: plan.leftScope }
+  const rightContext = plan.rightScope === undefined ? context : { ...context, scope: plan.rightScope }
+  const left = executePlan({ plan: plan.left, context: leftContext })
+  const right = executePlan({ plan: plan.right, context: rightContext })
 
   if (plan.operator === 'UNION') {
     if (plan.all) {
-      const left = executePlan({ plan: plan.left, context })
-      const right = executePlan({ plan: plan.right, context })
       return {
         columns: left.columns,
         numRows: addBounds(left.numRows, right.numRows),
@@ -722,8 +1144,6 @@ function executeSetOperation(plan, context) {
         },
       }
     } else {
-      const left = executePlan({ plan: plan.left, context })
-      const right = executePlan({ plan: plan.right, context })
       return {
         columns: left.columns,
         maxRows: addBounds(left.maxRows, right.maxRows),
@@ -759,8 +1179,6 @@ function executeSetOperation(plan, context) {
       }
     }
   } else if (plan.operator === 'INTERSECT') {
-    const left = executePlan({ plan: plan.left, context })
-    const right = executePlan({ plan: plan.right, context })
     return {
       columns: left.columns,
       maxRows: minBounds(left.maxRows, right.maxRows),
@@ -814,8 +1232,6 @@ function executeSetOperation(plan, context) {
     }
   } else {
     // EXCEPT
-    const left = executePlan({ plan: plan.left, context })
-    const right = executePlan({ plan: plan.right, context })
     return {
       columns: left.columns,
       maxRows: left.maxRows,