squirreling 0.15.2 → 0.16.0

package/src/plan/columns.js

@@ -1,10 +1,19 @@
-import { tableFunctionDefaultColumns } from '../parse/parse.js'
+import { dataSourceColumns } from '../backend/dataSource.js'
 import { derivedAlias } from '../expression/alias.js'
+import { tableFunctionDefaultColumns } from '../parse/parse.js'
 
 /**
  * @import { AsyncDataSource, ExprNode, FromFunction, FromSubquery, FromTable, IdentifierNode, SelectStatement, Statement } from '../types.js'
  */
 
+/**
+ * @typedef {{
+ *   cteColumns?: Map<string, string[]>,
+ *   tables?: Record<string, AsyncDataSource>,
+ *   outerAliases?: Set<string>,
+ * }} ColumnCollectionContext
+ */
+
 /**
  * @param {FromTable | FromSubquery | FromFunction | undefined} from
  * @returns {string | undefined}
@@ -60,9 +69,13 @@ export function tableFunctionColumnNames(from) {
  * @param {object} options
  * @param {SelectStatement} options.select
  * @param {IdentifierNode[]} [options.parentColumns] - columns needed by the parent query
+ * @param {Set<string>} [options.scopeColumns] - bare column names available in the current scope
+ * @param {Map<string, string[]>} [options.cteColumns]
+ * @param {Record<string, AsyncDataSource>} [options.tables]
+ * @param {string[]} [options.outerScope]
  * @returns {Map<string, string[] | undefined>}
  */
-export function extractColumns({ select, parentColumns }) {
+export function extractColumns({ select, parentColumns, scopeColumns, cteColumns, tables, outerScope }) {
   /** @type {Map<string, string[] | undefined>} */
   const result = new Map()
 
@@ -74,6 +87,11 @@ export function extractColumns({ select, parentColumns }) {
   for (const join of select.joins) {
     aliases.push(join.alias ?? join.table)
   }
+  const collectionContext = {
+    cteColumns,
+    tables,
+    outerAliases: new Set([...aliases, ...outerScope ?? []]),
+  }
 
   // If any unqualified SELECT * exists, all tables need all columns
   if (select.columns.some(col => col.type === 'star' && !col.table)) {
@@ -133,26 +151,26 @@ export function extractColumns({ select, parentColumns }) {
         if (!parentColumns.some(id => id.name === outputName)) continue
       }
       // Exclude earlier SELECT aliases so they aren't treated as source columns
-      collectColumnsFromExpr(col.expr, identifiers, selectAliases)
+      collectColumnsFromExpr(col.expr, identifiers, selectAliases, collectionContext)
       if (col.alias) {
         selectAliases.add(col.alias)
       }
     }
   }
-  collectColumnsFromExpr(select.where, identifiers)
+  collectColumnsFromExpr(select.where, identifiers, undefined, collectionContext)
 
   for (const item of select.orderBy) {
-    collectColumnsFromExpr(item.expr, identifiers, selectAliases)
+    collectColumnsFromExpr(item.expr, identifiers, selectAliases, collectionContext)
   }
   for (const expr of select.groupBy) {
-    collectColumnsFromExpr(expr, identifiers, selectAliases)
+    collectColumnsFromExpr(expr, identifiers, selectAliases, collectionContext)
   }
-  collectColumnsFromExpr(select.having, identifiers, selectAliases)
+  collectColumnsFromExpr(select.having, identifiers, selectAliases, collectionContext)
   /** @type {string[]} */
   const visibleLateralAliases = []
   if (sourceAlias !== undefined) visibleLateralAliases.push(sourceAlias)
   for (const join of select.joins) {
-    collectColumnsFromExpr(join.on, identifiers)
+    collectColumnsFromExpr(join.on, identifiers, undefined, collectionContext)
     // USING columns are equi-join keys on both sides; keep them in every
     // table's needed set so projection pushdown can't prune the join key.
     if (join.using) {
@@ -165,7 +183,7 @@ export function extractColumns({ select, parentColumns }) {
       /** @type {IdentifierNode[]} */
       const lateralArgIdentifiers = []
       for (const arg of join.fromFunction.args) {
-        collectColumnsFromExpr(arg, lateralArgIdentifiers)
+        collectColumnsFromExpr(arg, lateralArgIdentifiers, undefined, collectionContext)
       }
       lateralArgGroups.push({ identifiers: lateralArgIdentifiers, visibleAliases: [...visibleLateralAliases] })
     }
@@ -175,9 +193,18 @@ export function extractColumns({ select, parentColumns }) {
   // Partition identifiers by table prefix
   for (const { prefix, name } of identifiers) {
     if (prefix) {
-      // Qualified: add to matching table only
-      const set = perTable.get(prefix)
-      if (set) set.add(name)
+      if (perTable.has(prefix)) {
+        // Table-qualified: add to the matching table only
+        perTable.get(prefix)?.add(name)
+      } else if (scopeColumns?.has(prefix)) {
+        // Struct-field access: the prefix is the source column being read.
+        if (aliases.length === 1) {
+          perTable.get(aliases[0])?.add(prefix)
+        } else {
+          // As with an unqualified column in a join, its owner is ambiguous.
+          for (const alias of aliases) perTable.set(alias, undefined)
+        }
+      }
     } else if (aliases.length > 1) {
       // Unqualified in a JOIN: can't disambiguate, request all columns from all tables
       for (const alias of aliases) {
@@ -225,49 +252,50 @@ export function extractColumns({ select, parentColumns }) {
  * @param {ExprNode} expr
  * @param {IdentifierNode[]} columns
  * @param {Set<string>} [aliases] - aliases to exclude from columns
+ * @param {ColumnCollectionContext} [context]
  */
-export function collectColumnsFromExpr(expr, columns, aliases) {
+export function collectColumnsFromExpr(expr, columns, aliases, context) {
   if (!expr) return
   if (expr.type === 'identifier') {
     if (expr.prefix || !aliases?.has(expr.name)) {
       columns.push(expr)
     }
   } else if (expr.type === 'binary') {
-    collectColumnsFromExpr(expr.left, columns, aliases)
-    collectColumnsFromExpr(expr.right, columns, aliases)
+    collectColumnsFromExpr(expr.left, columns, aliases, context)
+    collectColumnsFromExpr(expr.right, columns, aliases, context)
   } else if (expr.type === 'unary') {
-    collectColumnsFromExpr(expr.argument, columns, aliases)
+    collectColumnsFromExpr(expr.argument, columns, aliases, context)
   } else if (expr.type === 'function') {
     for (const arg of expr.args) {
-      collectColumnsFromExpr(arg, columns, aliases)
+      collectColumnsFromExpr(arg, columns, aliases, context)
     }
-    collectColumnsFromExpr(expr.filter, columns, aliases)
+    collectColumnsFromExpr(expr.filter, columns, aliases, context)
   } else if (expr.type === 'window') {
-    for (const arg of expr.args) collectColumnsFromExpr(arg, columns, aliases)
-    for (const p of expr.partitionBy) collectColumnsFromExpr(p, columns, aliases)
-    for (const o of expr.orderBy) collectColumnsFromExpr(o.expr, columns, aliases)
+    for (const arg of expr.args) collectColumnsFromExpr(arg, columns, aliases, context)
+    for (const p of expr.partitionBy) collectColumnsFromExpr(p, columns, aliases, context)
+    for (const o of expr.orderBy) collectColumnsFromExpr(o.expr, columns, aliases, context)
   } else if (expr.type === 'cast') {
-    collectColumnsFromExpr(expr.expr, columns, aliases)
+    collectColumnsFromExpr(expr.expr, columns, aliases, context)
   } else if (expr.type === 'in valuelist') {
-    collectColumnsFromExpr(expr.expr, columns, aliases)
+    collectColumnsFromExpr(expr.expr, columns, aliases, context)
     for (const val of expr.values) {
-      collectColumnsFromExpr(val, columns, aliases)
+      collectColumnsFromExpr(val, columns, aliases, context)
     }
   } else if (expr.type === 'in') {
-    collectColumnsFromExpr(expr.expr, columns, aliases)
+    collectColumnsFromExpr(expr.expr, columns, aliases, context)
   } else if (expr.type === 'subscript') {
-    collectColumnsFromExpr(expr.expr, columns, aliases)
-    collectColumnsFromExpr(expr.index, columns, aliases)
+    collectColumnsFromExpr(expr.expr, columns, aliases, context)
+    collectColumnsFromExpr(expr.index, columns, aliases, context)
   } else if (expr.type === 'case') {
     if (expr.caseExpr) {
-      collectColumnsFromExpr(expr.caseExpr, columns, aliases)
+      collectColumnsFromExpr(expr.caseExpr, columns, aliases, context)
     }
     for (const when of expr.whenClauses) {
-      collectColumnsFromExpr(when.condition, columns, aliases)
-      collectColumnsFromExpr(when.result, columns, aliases)
+      collectColumnsFromExpr(when.condition, columns, aliases, context)
+      collectColumnsFromExpr(when.result, columns, aliases, context)
     }
     if (expr.elseResult) {
-      collectColumnsFromExpr(expr.elseResult, columns, aliases)
+      collectColumnsFromExpr(expr.elseResult, columns, aliases, context)
     }
   }
   // Subqueries: collect prefixed identifiers for correlated column detection.
@@ -276,60 +304,76 @@ export function collectColumnsFromExpr(expr, columns, aliases) {
   // from the inner query would incorrectly be attributed to the outer table.
   if (expr.type === 'subquery' || expr.type === 'in' || expr.type === 'exists' || expr.type === 'not exists') {
     if (expr.type === 'in') {
-      collectColumnsFromExpr(expr.expr, columns, aliases)
+      collectColumnsFromExpr(expr.expr, columns, aliases, context)
     }
     const sub = expr.subquery
     if (sub) {
-      /** @type {IdentifierNode[]} */
-      const inner = []
-      collectColumnsFromStatement(sub, inner)
-      for (const id of inner) {
-        if (id.prefix) columns.push(id)
-      }
+      collectCorrelatedColumnsFromStatement(sub, columns, context)
       // FROM-function args (e.g. UNNEST in the subquery's FROM) are evaluated
       // against the outer scope — the table function is itself the FROM, so
       // any identifier inside its args must be correlated. Push them even if
       // unprefixed so the outer scan reads the columns they reference.
-      collectFromFunctionArgs(sub, columns)
+      collectFromFunctionArgs(sub, columns, context)
     }
   }
   // No columns: count(*), literal, interval
 }
 
 /**
- * Collects identifiers from a subquery statement for correlated column detection.
+ * Collects qualified identifiers that escape a subquery statement's scope.
+ * Each SELECT filters its own aliases before identifiers propagate outward,
+ * preserving ownership through nested and compound subqueries.
  *
  * @param {Statement} stmt
  * @param {IdentifierNode[]} columns
+ * @param {ColumnCollectionContext} [context]
  */
-function collectColumnsFromStatement(stmt, columns) {
+function collectCorrelatedColumnsFromStatement(stmt, columns, context) {
   if (stmt.type === 'compound') {
-    collectColumnsFromStatement(stmt.left, columns)
-    collectColumnsFromStatement(stmt.right, columns)
+    collectCorrelatedColumnsFromStatement(stmt.left, columns, context)
+    collectCorrelatedColumnsFromStatement(stmt.right, columns, context)
     return
   }
   if (stmt.type === 'with') {
-    collectColumnsFromStatement(stmt.query, columns)
+    collectCorrelatedColumnsFromStatement(stmt.query, columns, context)
     return
   }
+  const scope = statementScope(stmt) ?? []
+  const nestedContext = context && {
+    ...context,
+    outerAliases: new Set([...context.outerAliases ?? [], ...scope]),
+  }
+  /** @type {IdentifierNode[]} */
+  const identifiers = []
   for (const col of stmt.columns) {
-    if (col.type === 'derived') collectColumnsFromExpr(col.expr, columns)
+    if (col.type === 'derived') collectColumnsFromExpr(col.expr, identifiers, undefined, nestedContext)
   }
-  collectColumnsFromExpr(stmt.where, columns)
+  collectColumnsFromExpr(stmt.where, identifiers, undefined, nestedContext)
   if (stmt.from && stmt.from.type === 'subquery') {
-    collectColumnsFromStatement(stmt.from.query, columns)
+    collectCorrelatedColumnsFromStatement(stmt.from.query, identifiers, nestedContext)
   }
   for (const join of stmt.joins) {
-    collectColumnsFromExpr(join.on, columns)
+    collectColumnsFromExpr(join.on, identifiers, undefined, nestedContext)
     if (join.fromFunction) {
       for (const arg of join.fromFunction.args) {
-        collectColumnsFromExpr(arg, columns)
+        collectColumnsFromExpr(arg, identifiers, undefined, nestedContext)
       }
     }
   }
-  for (const expr of stmt.groupBy) collectColumnsFromExpr(expr, columns)
-  collectColumnsFromExpr(stmt.having, columns)
-  for (const item of stmt.orderBy) collectColumnsFromExpr(item.expr, columns)
+  for (const expr of stmt.groupBy) collectColumnsFromExpr(expr, identifiers, undefined, nestedContext)
+  collectColumnsFromExpr(stmt.having, identifiers, undefined, nestedContext)
+  for (const item of stmt.orderBy) collectColumnsFromExpr(item.expr, identifiers, undefined, nestedContext)
+
+  const localColumns = context
+    ? collectScopeColumns({ select: stmt, cteColumns: context.cteColumns, tables: context.tables })
+    : new Set()
+  for (const identifier of identifiers) {
+    if (identifier.prefix &&
+      !scope.includes(identifier.prefix) &&
+      (context?.outerAliases?.has(identifier.prefix) || !localColumns.has(identifier.prefix))) {
+      columns.push(identifier)
+    }
+  }
 }
 
 /**
@@ -339,23 +383,24 @@ function collectColumnsFromStatement(stmt, columns) {
  *
  * @param {Statement} stmt
  * @param {IdentifierNode[]} columns
+ * @param {ColumnCollectionContext} [context]
  */
-function collectFromFunctionArgs(stmt, columns) {
+function collectFromFunctionArgs(stmt, columns, context) {
   if (stmt.type === 'compound') {
-    collectFromFunctionArgs(stmt.left, columns)
-    collectFromFunctionArgs(stmt.right, columns)
+    collectFromFunctionArgs(stmt.left, columns, context)
+    collectFromFunctionArgs(stmt.right, columns, context)
     return
   }
   if (stmt.type === 'with') {
-    collectFromFunctionArgs(stmt.query, columns)
+    collectFromFunctionArgs(stmt.query, columns, context)
     return
   }
   if (stmt.from?.type === 'function') {
     for (const arg of stmt.from.args) {
-      collectColumnsFromExpr(arg, columns)
+      collectColumnsFromExpr(arg, columns, undefined, context)
     }
   } else if (stmt.from?.type === 'subquery') {
-    collectFromFunctionArgs(stmt.from.query, columns)
+    collectFromFunctionArgs(stmt.from.query, columns, context)
   }
 }
 
@@ -473,7 +518,8 @@ export function inferSelectSourceColumns({ select, cteColumns, tables }) {
  * @returns {string[]}
  */
 function lookupTableColumns(table, cteColumns, tables) {
-  return cteColumns?.get(table.toLowerCase()) ?? tables?.[table]?.columns ?? []
+  const source = tables?.[table]
+  return cteColumns?.get(table.toLowerCase()) ?? (source ? dataSourceColumns(source) : [])
 }
 
 /**

package/src/plan/plan.js

@@ -82,6 +82,8 @@ function planSetOperation({ compound, ctePlans, cteColumns, tables, parentColumn
     all: compound.all,
     left,
     right,
+    leftScope: statementScope(compound.left),
+    rightScope: statementScope(compound.right),
   }
 
   if (compound.orderBy.length) {
@@ -206,7 +208,14 @@ function planSelect({ select, ctePlans, cteColumns, tables, parentColumns, outer
   // included so they are only applied to fresh scans, not CTE/subquery plans)
   /** @type {ScanOptions} */
   const hints = {}
-  const perTableColumns = extractColumns({ select: originalSelect, parentColumns })
+  const perTableColumns = extractColumns({
+    select: originalSelect,
+    parentColumns,
+    scopeColumns,
+    cteColumns,
+    tables,
+    outerScope,
+  })
   if (sourceAlias !== undefined) hints.columns = perTableColumns.get(sourceAlias)
   // Capture what the parent reads from a FROM subquery before the reset
   // below, so aggregate outputs it never reads can still be pruned when the
@@ -406,7 +415,7 @@ function planFrom({ select, ctePlans, cteColumns, hints, tables, outerScope }) {
       return ctePlan
     }
     validateScan({ ...select.from, hints, tables })
-    return { type: 'Scan', table: select.from.table, hints }
+    return { type: 'Scan', table: select.from.table, alias: select.from.alias, hints }
   } else if (select.from.type === 'function') {
     for (const arg of select.from.args) {
       validateNoIdentifiers(arg, select.from.funcName, outerScope)
@@ -545,7 +554,7 @@ function planJoin({ left, joins, leftTable, ctePlans, cteColumns, perTableColumn
         // For CTE joins, use CTE column metadata for hints
         rightHints.columns = perTableColumns.get(rightTable) ?? cteColumns?.get(join.table.toLowerCase())
       }
-      rightScan = ctePlan ?? { type: 'Scan', table: join.table, hints: rightHints }
+      rightScan = ctePlan ?? { type: 'Scan', table: join.table, alias: join.alias, hints: rightHints }
     }
 
     if (join.joinType === 'POSITIONAL') {

package/src/plan/types.d.ts

@@ -23,6 +23,7 @@ export type QueryPlan =
 export interface ScanNode {
   type: 'Scan'
   table: string
+  alias?: string
   hints: ScanOptions
 }
 
@@ -130,6 +131,8 @@ export interface SetOperationNode {
   all: boolean
   left: QueryPlan
   right: QueryPlan
+  leftScope?: string[]
+  rightScope?: string[]
 }
 
 // Wraps a derived-table or CTE subplan with the lexical alias scope of its

package/src/types.d.ts

@@ -11,10 +11,148 @@ export { QueryPlan } from './plan/types.js'
 export interface QueryResults {
   columns: string[]
   rows(): AsyncGenerator<AsyncRow>
+  batches?(): AsyncIterable<AsyncBatch>
   numRows?: number
   maxRows?: number
 }
 
+export type SqlType =
+  | { type: 'unknown' }
+  | { type: 'string' }
+  | { type: 'number' }
+  | { type: 'bigint' }
+  | { type: 'boolean' }
+  | { type: 'date' }
+  | { type: 'array', items: SqlType }
+  | { type: 'struct', fields: readonly Field[] }
+
+export interface Field {
+  id: number
+  name: string
+  dataType: SqlType
+  nullable: boolean
+}
+
+export interface RelationSchema {
+  fields: readonly Field[]
+}
+
+/**
+ * A selection over a base domain of `length` rows.
+ */
+export type RowSelection =
+  | { type: 'all', length: number }
+  | { type: 'range', start: number, end: number, length: number }
+  | { type: 'indices', indices: Uint32Array, length: number }
+
+export type NumericArray =
+  | Int8Array
+  | Uint8Array
+  | Uint8ClampedArray
+  | Int16Array
+  | Uint16Array
+  | Int32Array
+  | Uint32Array
+  | Float32Array
+  | Float64Array
+  | BigInt64Array
+  | BigUint64Array
+
+export type ColumnVector =
+  | {
+      type: 'values'
+      values: readonly SqlPrimitive[]
+      length: number
+    }
+  | {
+      type: 'typed'
+      values: NumericArray
+      validity?: Uint8Array
+      length: number
+    }
+  | {
+      type: 'constant'
+      value: SqlPrimitive
+      length: number
+    }
+  | {
+      type: 'selected'
+      source: ColumnVector
+      selection: RowSelection
+      length: number
+    }
+
+export interface ColumnReadRequest {
+  batch: AsyncBatch
+  selection: RowSelection
+  signal?: AbortSignal
+  rowOffset?: number
+  rowOrdinals?: ColumnVector
+}
+
+export type ColumnResult = ColumnVector | Promise<ColumnVector>
+export type ReadColumn = (request: ColumnReadRequest) => ColumnResult
+
+export type BatchColumn =
+  | ColumnVector
+  | {
+      read: ReadColumn
+      input?: AsyncBatch
+      rowOffset?: number
+      rowOrdinals?: ColumnVector
+    }
+
+export interface AsyncBatch {
+  selection: RowSelection
+  columns: readonly BatchColumn[]
+}
+
+export interface ReadBatchColumnOptions {
+  batch: AsyncBatch
+  columnIndex: number
+  selection?: RowSelection
+  signal?: AbortSignal
+}
+
+export interface RowsToBatchesOptions {
+  batchRows?: number
+  signal?: AbortSignal
+}
+
+export interface ColumnDemand {
+  field: number
+  phase: number
+  purpose: 'filter' | 'output'
+  mode: 'required' | 'deferred'
+}
+
+export interface ScanRequest {
+  columns: readonly ColumnDemand[]
+  filter?: ExprNode
+  limit?: number
+  offset?: number
+}
+
+export interface ScanProperties {
+  exactRows?: number
+  maxRows?: number
+}
+
+export interface ScanResidual {
+  filter?: ExprNode
+  limit?: number
+  offset?: number
+}
+
+export interface PreparedScan {
+  schema: RelationSchema
+  residual: ScanResidual
+  properties: ScanProperties
+  batches(options?: { signal?: AbortSignal }): AsyncIterable<AsyncBatch>
+}
+
+export type PrepareScan = (request: ScanRequest) => PreparedScan
+
 // parseSql(options)
 export interface ParseSqlOptions {
   query: string
@@ -72,16 +210,30 @@ export type AsyncCell = () => Promise<SqlPrimitive>
 export type Row = Record<string, SqlPrimitive>[]
 
 /**
- * Async data source for streaming SQL execution.
+ * Async data source for streaming SQL execution. A source must implement
+ * either scan() or prepareScan().
  */
-export interface AsyncDataSource {
+interface AsyncDataSourceBase {
   numRows?: number
-  columns: string[]
-  scan(options: ScanOptions): ScanResults
   // Optional method for fast column scans
   scanColumn?(options: ScanColumnOptions): AsyncIterable<ArrayLike<SqlPrimitive>> | ScanColumnResults
 }
 
+export type AsyncDataSource = AsyncDataSourceBase & (
+  | {
+      columns: string[]
+      scan(options: ScanOptions): ScanResults
+      schema?: RelationSchema
+      prepareScan?: PrepareScan
+    }
+  | {
+      columns?: string[]
+      scan?(options: ScanOptions): ScanResults
+      schema: RelationSchema
+      prepareScan: PrepareScan
+    }
+)
+
 /**
  * Result of a scan: streaming rows and flags indicating which hints were
  * applied by the data source.

package/src/validation/tables.js

@@ -1,3 +1,4 @@
+import { dataSourceColumns } from '../backend/dataSource.js'
 import { ExecutionError } from './executionErrors.js'
 
 /**
@@ -34,11 +35,12 @@ export function validateTable({ table, qualified, tables, positionStart, positio
 export function validateScan({ table, hints, tables, positionStart, positionEnd }) {
   if (!tables) return
   const resolved = validateTable({ table, tables, positionStart, positionEnd })
-  const missingColumn = hints.columns?.find(col => !resolved.columns.includes(col))
+  const columns = dataSourceColumns(resolved)
+  const missingColumn = hints.columns?.find(col => !columns.includes(col))
   if (missingColumn) {
     throw new ColumnNotFoundError({
       missingColumn,
-      availableColumns: resolved.columns,
+      availableColumns: columns,
       positionStart,
       positionEnd,
     })