squirreling 0.11.4 → 0.12.0

package/src/execute/execute.js

@@ -7,10 +7,10 @@ import { validateScan, validateTable } from '../validation/tables.js'
 import { executeHashAggregate, executeScalarAggregate } from './aggregates.js'
 import { executeHashJoin, executeNestedLoopJoin, executePositionalJoin } from './join.js'
 import { executeSort } from './sort.js'
-import { stableRowKey } from './utils.js'
+import { addBounds, minBounds, stableRowKey } from './utils.js'
 
 /**
- * @import { AsyncCells, AsyncDataSource, AsyncRow, ExecuteContext, ExecuteSqlOptions, ExprNode, Statement } from '../types.js'
+ * @import { AsyncCells, AsyncDataSource, AsyncRow, ExecuteContext, ExecuteSqlOptions, ExprNode, QueryResults, Statement } from '../types.js'
  * @import { CountNode, DistinctNode, FilterNode, LimitNode, ProjectNode, QueryPlan, ScanNode, SetOperationNode } from '../plan/types.js'
  */
 
@@ -18,9 +18,9 @@ import { stableRowKey } from './utils.js'
  * Executes a SQL SELECT query against tables
  *
  * @param {ExecuteSqlOptions} options
- * @yields {AsyncRow}
+ * @returns {QueryResults}
  */
-export async function* executeSql({ tables, query, functions, signal }) {
+export function executeSql({ tables, query, functions, signal }) {
   const parsed = typeof query === 'string' ? parseSql({ query, functions }) : query
 
   // Normalize tables: convert arrays to AsyncDataSource
@@ -34,7 +34,9 @@ export async function* executeSql({ tables, query, functions, signal }) {
     }
   }
 
-  yield* executeStatement({ query: parsed, context: { tables: normalizedTables, functions, signal } })
+  const context = { tables: normalizedTables, functions, signal }
+  const plan = planSql({ query: parsed, functions, tables: normalizedTables })
+  return executePlan({ plan, context })
 }
 
 /**
@@ -43,60 +45,62 @@ export async function* executeSql({ tables, query, functions, signal }) {
  * @param {Object} options
  * @param {Statement} options.query
  * @param {ExecuteContext} options.context
- * @yields {AsyncRow}
+ * @returns {QueryResults}
  */
-export async function* executeStatement({ query, context }) {
+export function executeStatement({ query, context }) {
   const plan = planSql({ query, functions: context.functions, tables: context.tables })
-  yield* executePlan({ plan, context })
+  return executePlan({ plan, context })
 }
 
 /**
- * Executes a query plan and yields result rows
+ * Executes a query plan and returns query results with row count estimates
  *
  * @param {Object} options
  * @param {QueryPlan} options.plan - the query plan to execute
  * @param {ExecuteContext} options.context - execution context
- * @returns {AsyncGenerator<AsyncRow>}
+ * @returns {QueryResults}
  */
-export async function* executePlan({ plan, context }) {
+export function executePlan({ plan, context }) {
   if (plan.type === 'Scan') {
-    yield* executeScan(plan, context)
+    return executeScan(plan, context)
   } else if (plan.type === 'Count') {
-    yield* executeCount(plan, context)
+    return executeCount(plan, context)
   } else if (plan.type === 'Filter') {
-    yield* executeFilter(plan, context)
+    return executeFilter(plan, context)
   } else if (plan.type === 'Project') {
-    yield* executeProject(plan, context)
+    return executeProject(plan, context)
   } else if (plan.type === 'HashJoin') {
-    yield* executeHashJoin(plan, context)
+    return executeHashJoin(plan, context)
   } else if (plan.type === 'NestedLoopJoin') {
-    yield* executeNestedLoopJoin(plan, context)
+    return executeNestedLoopJoin(plan, context)
   } else if (plan.type === 'PositionalJoin') {
-    yield* executePositionalJoin(plan, context)
+    return executePositionalJoin(plan, context)
   } else if (plan.type === 'HashAggregate') {
-    yield* executeHashAggregate(plan, context)
+    return executeHashAggregate(plan, context)
   } else if (plan.type === 'ScalarAggregate') {
-    yield* executeScalarAggregate(plan, context)
+    return executeScalarAggregate(plan, context)
   } else if (plan.type === 'Sort') {
-    yield* executeSort(plan, context)
+    return executeSort(plan, context)
   } else if (plan.type === 'Distinct') {
-    yield* executeDistinct(plan, context)
+    return executeDistinct(plan, context)
   } else if (plan.type === 'Limit') {
-    yield* executeLimit(plan, context)
+    return executeLimit(plan, context)
   } else if (plan.type === 'SetOperation') {
-    yield* executeSetOperation(plan, context)
+    return executeSetOperation(plan, context)
   }
+  return { async *rows () {} }
 }
 
 /**
  * @param {ScanNode} plan
  * @param {ExecuteContext} context
- * @yields {AsyncRow}
+ * @returns {QueryResults}
  */
-async function* executeScan(plan, context) {
+function executeScan(plan, context) {
   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) {
@@ -107,42 +111,55 @@ async function* executeScan(plan, context) {
       offset: plan.hints.offset,
       signal,
     })
-    const columns = [column]
-    for await (const chunk of chunks) {
-      if (signal?.aborted) return
-      for (let i = 0; i < chunk.length; i++) {
-        const value = chunk[i]
-        yield {
-          columns,
-          cells: { [column]: () => Promise.resolve(value) },
+    const scanRows = computeScanRows(table.numRows, plan.hints.limit, plan.hints.offset)
+    return {
+      numRows: scanRows,
+      maxRows: scanRows,
+      async *rows () {
+        const columns = [column]
+        for await (const chunk of chunks) {
+          if (signal?.aborted) return
+          for (let i = 0; i < chunk.length; i++) {
+            const value = chunk[i]
+            yield {
+              columns,
+              cells: { [column]: () => Promise.resolve(value) },
+            }
+          }
         }
-      }
+      },
     }
-    return
   }
 
   // do the scan
-  const { rows, appliedWhere, appliedLimitOffset } = table.scan({ ...plan.hints, signal })
+  const scanResult = table.scan({ ...plan.hints, signal })
+  const { appliedWhere, appliedLimitOffset } = scanResult
 
   // Applied limit/offset without applied where is invalid
-  const hasLimitOffset = plan.hints.limit !== undefined || plan.hints.offset // 0 offset is noop
   if (!appliedWhere && appliedLimitOffset && plan.hints.where && hasLimitOffset) {
     throw new Error(`Data source "${plan.table}" applied limit/offset without applying where`)
   }
 
-  let result = rows
+  const scanRows = computeScanRows(table.numRows, plan.hints.limit, plan.hints.offset)
+  return {
+    numRows: !plan.hints.where ? scanRows : undefined,
+    maxRows: scanRows,
+    async *rows () {
+      let result = scanResult.rows()
 
-  // Apply WHERE if data source did not
-  if (!appliedWhere && plan.hints.where) {
-    result = filterRows(result, plan.hints.where, context, plan.hints.limit)
-  }
+      // Apply WHERE if data source did not
+      if (!appliedWhere && plan.hints.where) {
+        result = filterRows(result, plan.hints.where, context, plan.hints.limit)
+      }
 
-  // Apply LIMIT/OFFSET if data source did not
-  if (!appliedLimitOffset && hasLimitOffset) {
-    result = limitRows(result, plan.hints.limit, plan.hints.offset, signal)
-  }
+      // Apply LIMIT/OFFSET if data source did not
+      if (!appliedLimitOffset && hasLimitOffset) {
+        result = limitRows(result, plan.hints.limit, plan.hints.offset, signal)
+      }
 
-  yield* result
+      yield* result
+    },
+  }
 }
 
 /**
@@ -150,34 +167,55 @@ async function* executeScan(plan, context) {
  *
  * @param {CountNode} plan
  * @param {ExecuteContext} context
- * @yields {AsyncRow}
+ * @returns {QueryResults}
  */
-async function* executeCount(plan, { tables, signal }) {
+function executeCount(plan, context) {
+  const { tables, signal } = context
   const table = validateTable({ ...plan, tables })
 
-  // Use source numRows if available
-  let count = table.numRows
-  if (count === undefined) {
-    // Fall back to counting rows via scan
-    count = 0
-    const { rows } = table.scan({ signal })
-    // eslint-disable-next-line no-unused-vars
-    for await (const _ of rows) {
-      if (signal?.aborted) return
-      count++
-    }
-  }
+  return {
+    numRows: 1,
+    maxRows: 1,
+    async *rows () {
+      // Use source numRows if available
+      let count = table.numRows
+      if (count === undefined) {
+        // Fall back to counting rows via scan
+        count = 0
+        const { rows } = table.scan({ signal })
+        // eslint-disable-next-line no-unused-vars
+        for await (const _ of rows()) {
+          if (signal?.aborted) return
+          count++
+        }
+      }
 
-  /** @type {string[]} */
-  const columns = []
-  /** @type {AsyncCells} */
-  const cells = {}
-  for (const col of plan.columns) {
-    const alias = col.alias ?? derivedAlias(col.expr)
-    columns.push(alias)
-    cells[alias] = () => Promise.resolve(count)
+      /** @type {string[]} */
+      const columns = []
+      /** @type {AsyncCells} */
+      const cells = {}
+      for (const col of plan.columns) {
+        const alias = col.alias ?? derivedAlias(col.expr)
+        columns.push(alias)
+        cells[alias] = () => Promise.resolve(count)
+      }
+      yield { columns, cells }
+    },
   }
-  yield { columns, cells }
+}
+
+/**
+ * Computes numRows for a scan when the table provides numRows and there is no WHERE.
+ *
+ * @param {number | undefined} tableNumRows
+ * @param {number} [limit]
+ * @param {number} [offset]
+ * @returns {number | undefined}
+ */
+function computeScanRows(tableNumRows, limit, offset) {
+  if (tableNumRows === undefined) return undefined
+  const afterOffset = Math.max(0, tableNumRows - (offset ?? 0))
+  return limit !== undefined ? Math.min(limit, afterOffset) : afterOffset
 }
 
 /**
@@ -257,10 +295,14 @@ async function* limitRows(rows, limit, offset, signal) {
  *
  * @param {FilterNode} plan
  * @param {ExecuteContext} context
- * @yields {AsyncRow}
+ * @returns {QueryResults}
  */
-async function* executeFilter(plan, context) {
-  yield* filterRows(executePlan({ plan: plan.child, context }), plan.condition, context)
+function executeFilter(plan, context) {
+  const child = executePlan({ plan: plan.child, context })
+  return {
+    maxRows: child.maxRows,
+    rows: () => filterRows(child.rows(), plan.condition, context),
+  }
 }
 
 /**
@@ -268,45 +310,52 @@ async function* executeFilter(plan, context) {
  *
  * @param {ProjectNode} plan
  * @param {ExecuteContext} context
- * @yields {AsyncRow}
+ * @returns {QueryResults}
  */
-async function* executeProject(plan, context) {
-  let rowIndex = 0
-
-  for await (const row of executePlan({ plan: plan.child, context })) {
-    if (context.signal?.aborted) return
-    rowIndex++
-    const currentRowIndex = rowIndex
-
-    /** @type {string[]} */
-    const columns = []
-    /** @type {AsyncCells} */
-    const cells = {}
-
-    for (const col of plan.columns) {
-      if (col.type === 'star') {
-        const prefix = col.table ? `${col.table}.` : undefined
-        for (const key of row.columns) {
-          if (prefix && !key.startsWith(prefix)) continue
-          // Strip table prefix for output column names
-          const dotIndex = key.indexOf('.')
-          const outputKey = prefix ? key.substring(prefix.length) : dotIndex >= 0 ? key.substring(dotIndex + 1) : key
-          columns.push(outputKey)
-          cells[outputKey] = row.cells[key]
+function executeProject(plan, context) {
+  const child = executePlan({ plan: plan.child, context })
+  return {
+    numRows: child.numRows,
+    maxRows: child.maxRows,
+    async *rows () {
+      let rowIndex = 0
+
+      for await (const row of child.rows()) {
+        if (context.signal?.aborted) return
+        rowIndex++
+        const currentRowIndex = rowIndex
+
+        /** @type {string[]} */
+        const columns = []
+        /** @type {AsyncCells} */
+        const cells = {}
+
+        for (const col of plan.columns) {
+          if (col.type === 'star') {
+            const prefix = col.table ? `${col.table}.` : undefined
+            for (const key of row.columns) {
+              if (prefix && !key.startsWith(prefix)) continue
+              // Strip table prefix for output column names
+              const dotIndex = key.indexOf('.')
+              const outputKey = prefix ? key.substring(prefix.length) : dotIndex >= 0 ? key.substring(dotIndex + 1) : key
+              columns.push(outputKey)
+              cells[outputKey] = row.cells[key]
+            }
+          } else {
+            const alias = col.alias ?? derivedAlias(col.expr)
+            columns.push(alias)
+            cells[alias] = () => evaluateExpr({
+              node: col.expr,
+              row,
+              rowIndex: currentRowIndex,
+              context,
+            })
+          }
         }
-      } else {
-        const alias = col.alias ?? derivedAlias(col.expr)
-        columns.push(alias)
-        cells[alias] = () => evaluateExpr({
-          node: col.expr,
-          row,
-          rowIndex: currentRowIndex,
-          context,
-        })
-      }
-    }
 
-    yield { columns, cells }
+        yield { columns, cells }
+      }
+    },
   }
 }
 
@@ -315,44 +364,50 @@ async function* executeProject(plan, context) {
  *
  * @param {DistinctNode} plan
  * @param {ExecuteContext} context
- * @yields {AsyncRow}
+ * @returns {QueryResults}
  */
-async function* executeDistinct(plan, context) {
-  const { signal } = context
-  const MAX_CHUNK = 256
-
-  const seen = new Set()
+function executeDistinct(plan, context) {
+  const child = executePlan({ plan: plan.child, context })
+  return {
+    maxRows: child.maxRows,
+    async *rows () {
+      const { signal } = context
+      const MAX_CHUNK = 256
 
-  /** @type {AsyncRow[]} */
-  let buffer = []
+      const seen = new Set()
 
-  for await (const row of executePlan({ plan: plan.child, context })) {
-    if (signal?.aborted) return
-    buffer.push(row)
+      /** @type {AsyncRow[]} */
+      let buffer = []
 
-    if (buffer.length >= MAX_CHUNK) {
-      const keys = buffer.map(stableRowKey)
-      for (let i = 0; i < buffer.length; i++) {
-        const key = await keys[i]
-        if (!seen.has(key)) {
-          seen.add(key)
-          yield buffer[i]
+      for await (const row of child.rows()) {
+        if (signal?.aborted) return
+        buffer.push(row)
+
+        if (buffer.length >= MAX_CHUNK) {
+          const keys = buffer.map(stableRowKey)
+          for (let i = 0; i < buffer.length; i++) {
+            const key = await keys[i]
+            if (!seen.has(key)) {
+              seen.add(key)
+              yield buffer[i]
+            }
+          }
+          buffer = []
         }
       }
-      buffer = []
-    }
-  }
 
-  // Flush remaining
-  if (buffer.length > 0) {
-    const keys = buffer.map(stableRowKey)
-    for (let i = 0; i < buffer.length; i++) {
-      const key = await keys[i]
-      if (!seen.has(key)) {
-        seen.add(key)
-        yield buffer[i]
+      // Flush remaining
+      if (buffer.length > 0) {
+        const keys = buffer.map(stableRowKey)
+        for (let i = 0; i < buffer.length; i++) {
+          const key = await keys[i]
+          if (!seen.has(key)) {
+            seen.add(key)
+            yield buffer[i]
+          }
+        }
       }
-    }
+    },
   }
 }
 
@@ -361,10 +416,15 @@ async function* executeDistinct(plan, context) {
  *
  * @param {LimitNode} plan
  * @param {ExecuteContext} context
- * @yields {AsyncRow}
+ * @returns {QueryResults}
  */
-async function* executeLimit(plan, context) {
-  yield* limitRows(executePlan({ plan: plan.child, context }), plan.limit, plan.offset, context.signal)
+function executeLimit(plan, context) {
+  const child = executePlan({ plan: plan.child, context })
+  return {
+    numRows: computeScanRows(child.numRows, plan.limit, plan.offset),
+    maxRows: computeScanRows(child.maxRows, plan.limit, plan.offset),
+    rows: () => limitRows(child.rows(), plan.limit, plan.offset, context.signal),
+  }
 }
 
 /**
@@ -372,102 +432,132 @@ async function* executeLimit(plan, context) {
  *
  * @param {SetOperationNode} plan
  * @param {ExecuteContext} context
- * @yields {AsyncRow}
+ * @returns {QueryResults}
  */
-async function* executeSetOperation(plan, context) {
+function executeSetOperation(plan, context) {
   const { signal } = context
 
   if (plan.operator === 'UNION') {
     if (plan.all) {
-      // UNION ALL: yield all rows from both sides
-      yield* executePlan({ plan: plan.left, context })
-      yield* executePlan({ plan: plan.right, context })
-    } else {
-      // UNION: yield deduplicated rows from both sides
-      const seen = new Set()
-      for await (const row of executePlan({ plan: plan.left, context })) {
-        if (signal?.aborted) return
-        const key = await stableRowKey(row)
-        if (!seen.has(key)) {
-          seen.add(key)
-          yield row
-        }
+      const left = executePlan({ plan: plan.left, context })
+      const right = executePlan({ plan: plan.right, context })
+      return {
+        numRows: addBounds(left.numRows, right.numRows),
+        maxRows: addBounds(left.maxRows, right.maxRows),
+        async *rows () {
+          // UNION ALL: yield all rows from both sides
+          yield* left.rows()
+          yield* right.rows()
+        },
       }
-      for await (const row of executePlan({ plan: plan.right, context })) {
-        if (signal?.aborted) return
-        const key = await stableRowKey(row)
-        if (!seen.has(key)) {
-          seen.add(key)
-          yield row
-        }
+    } else {
+      const left = executePlan({ plan: plan.left, context })
+      const right = executePlan({ plan: plan.right, context })
+      return {
+        maxRows: addBounds(left.maxRows, right.maxRows),
+        async *rows () {
+          // UNION: yield deduplicated rows from both sides
+          const seen = new Set()
+          for await (const row of left.rows()) {
+            if (signal?.aborted) return
+            const key = await stableRowKey(row)
+            if (!seen.has(key)) {
+              seen.add(key)
+              yield row
+            }
+          }
+          for await (const row of right.rows()) {
+            if (signal?.aborted) return
+            const key = await stableRowKey(row)
+            if (!seen.has(key)) {
+              seen.add(key)
+              yield row
+            }
+          }
+        },
       }
     }
   } else if (plan.operator === 'INTERSECT') {
-    // Materialize right side keys
-    /** @type {Map<any, number>} */
-    const rightKeys = new Map()
-    for await (const row of executePlan({ plan: plan.right, context })) {
-      if (signal?.aborted) return
-      const key = await stableRowKey(row)
-      rightKeys.set(key, (rightKeys.get(key) ?? 0) + 1)
-    }
-
-    if (plan.all) {
-      // INTERSECT ALL: yield each left row that matches, consuming right counts
-      for await (const row of executePlan({ plan: plan.left, context })) {
-        if (signal?.aborted) return
-        const key = await stableRowKey(row)
-        const count = rightKeys.get(key)
-        if (count) {
-          rightKeys.set(key, count - 1)
-          yield row
+    const left = executePlan({ plan: plan.left, context })
+    const right = executePlan({ plan: plan.right, context })
+    return {
+      maxRows: minBounds(left.maxRows, right.maxRows),
+      async *rows () {
+        // Materialize right side keys
+        /** @type {Map<any, number>} */
+        const rightKeys = new Map()
+        for await (const row of right.rows()) {
+          if (signal?.aborted) return
+          const key = await stableRowKey(row)
+          rightKeys.set(key, (rightKeys.get(key) ?? 0) + 1)
         }
-      }
-    } else {
-      // INTERSECT: yield deduplicated rows present in both
-      const seen = new Set()
-      for await (const row of executePlan({ plan: plan.left, context })) {
-        if (signal?.aborted) return
-        const key = await stableRowKey(row)
-        if (rightKeys.has(key) && !seen.has(key)) {
-          seen.add(key)
-          yield row
+
+        if (plan.all) {
+          // INTERSECT ALL: yield each left row that matches, consuming right counts
+          for await (const row of left.rows()) {
+            if (signal?.aborted) return
+            const key = await stableRowKey(row)
+            const count = rightKeys.get(key)
+            if (count) {
+              rightKeys.set(key, count - 1)
+              yield row
+            }
+          }
+        } else {
+          // INTERSECT: yield deduplicated rows present in both
+          const seen = new Set()
+          for await (const row of left.rows()) {
+            if (signal?.aborted) return
+            const key = await stableRowKey(row)
+            if (rightKeys.has(key) && !seen.has(key)) {
+              seen.add(key)
+              yield row
+            }
+          }
         }
-      }
-    }
-  } else if (plan.operator === 'EXCEPT') {
-    // Materialize right side keys
-    /** @type {Map<any, number>} */
-    const rightKeys = new Map()
-    for await (const row of executePlan({ plan: plan.right, context })) {
-      if (signal?.aborted) return
-      const key = await stableRowKey(row)
-      rightKeys.set(key, (rightKeys.get(key) ?? 0) + 1)
+      },
     }
+  } else {
+    // EXCEPT
+    const left = executePlan({ plan: plan.left, context })
+    const right = executePlan({ plan: plan.right, context })
+    return {
+      maxRows: left.maxRows,
+      async *rows () {
+        // Materialize right side keys
+        /** @type {Map<any, number>} */
+        const rightKeys = new Map()
+        for await (const row of right.rows()) {
+          if (signal?.aborted) return
+          const key = await stableRowKey(row)
+          rightKeys.set(key, (rightKeys.get(key) ?? 0) + 1)
+        }
 
-    if (plan.all) {
-      // EXCEPT ALL: yield left rows, consuming right counts
-      for await (const row of executePlan({ plan: plan.left, context })) {
-        if (signal?.aborted) return
-        const key = await stableRowKey(row)
-        const count = rightKeys.get(key)
-        if (count) {
-          rightKeys.set(key, count - 1)
+        if (plan.all) {
+          // EXCEPT ALL: yield left rows, consuming right counts
+          for await (const row of left.rows()) {
+            if (signal?.aborted) return
+            const key = await stableRowKey(row)
+            const count = rightKeys.get(key)
+            if (count) {
+              rightKeys.set(key, count - 1)
+            } else {
+              yield row
+            }
+          }
         } else {
-          yield row
-        }
-      }
-    } else {
-      // EXCEPT: yield deduplicated left rows not in right
-      const seen = new Set()
-      for await (const row of executePlan({ plan: plan.left, context })) {
-        if (signal?.aborted) return
-        const key = await stableRowKey(row)
-        if (!rightKeys.has(key) && !seen.has(key)) {
-          seen.add(key)
-          yield row
+          // EXCEPT: yield deduplicated left rows not in right
+          const seen = new Set()
+          for await (const row of left.rows()) {
+            if (signal?.aborted) return
+            const key = await stableRowKey(row)
+            if (!rightKeys.has(key) && !seen.has(key)) {
+              seen.add(key)
+              yield row
+            }
+          }
         }
-      }
+      },
     }
   }
 }