diff --git a/src/backend/batch.js b/src/backend/batch.js index 8acf5d5..a5796f6 100644 --- a/src/backend/batch.js +++ b/src/backend/batch.js @@ -1,6 +1,5 @@ /** - * @import { AsyncBatch, ColumnResult, ColumnVector, ReadBatchColumnOptions, RowSelection } from '../internalTypes.js' - * @import { SqlPrimitive } from '../types.js' + * @import { AsyncBatch, ColumnResult, ColumnVector, ReadBatchColumnOptions, RowSelection, SqlPrimitive } from '../types.js' */ /** @type {WeakMap> }>>>} */ @@ -126,7 +125,7 @@ export function readBatchColumn({ batch, columnIndex, selection = batch.selectio if (selection.length !== batch.selection.length) { throw new Error(`Selection length ${selection.length} does not match batch length ${batch.selection.length}`) } - if (column.type !== 'source' && column.type !== 'computed') { + if (!('read' in column)) { return selectVector(column, selection) } @@ -149,15 +148,15 @@ export function readBatchColumn({ batch, columnIndex, selection = batch.selectio if (pending) return pending if (cache.resolved) return cache.resolved - const result = column.type === 'source' - ? column.read({ selection, signal }) - : column.expression.evaluate({ - batch: column.input, - selection, - signal, - rowOffset: column.rowOffset, - rowOrdinals: selectVector(column.rowOrdinals, selection), - }) + 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) { @@ -185,7 +184,6 @@ export function readBatchColumn({ batch, columnIndex, selection = batch.selectio export function selectBatch(batch, selection) { const composed = composeSelections(batch.selection, selection) return { - columnNames: batch.columnNames, selection: composed, columns: batch.columns, } diff --git a/src/backend/batchAdapters.js b/src/backend/batchAdapters.js index 988802f..aee8154 100644 --- a/src/backend/batchAdapters.js +++ b/src/backend/batchAdapters.js @@ -2,8 +2,7 @@ import { readBatchColumn, selectVector, selectedRowCount, valueAt } from './batc import { yieldToEventLoop } from '../execute/yield.js' /** - * @import { AsyncBatch, ColumnResult, ColumnVector, RowsToBatchesOptions } from '../internalTypes.js' - * @import { AsyncCells, AsyncRow, SqlPrimitive } from '../types.js' + * @import { AsyncBatch, AsyncCells, AsyncRow, ColumnResult, ColumnVector, RowsToBatchesOptions, SqlPrimitive } from '../types.js' */ const DEFAULT_BATCH_ROWS = 1024 @@ -33,13 +32,13 @@ export async function* rowsToBatches(rows, columnNames, options) { } rowCount++ if (rowCount === batchRows) { - yield loadedBatch(columnNames, values, rowCount) + yield loadedBatch(values, rowCount) values = makeValueBuffers(columnNames.length) rowCount = 0 } } - if (rowCount > 0) yield loadedBatch(columnNames, values, rowCount) + if (rowCount > 0) yield loadedBatch(values, rowCount) options?.signal?.throwIfAborted() } @@ -48,14 +47,14 @@ export async function* rowsToBatches(rows, columnNames, options) { * remain lazy and are memoized once per batch/selection by `readBatchColumn`. * * @param {AsyncIterable} batches + * @param {string[]} columns * @param {AbortSignal} [signal] * @yields {AsyncRow} */ -export async function* batchesToRows(batches, signal) { +export async function* batchesToRows(batches, columns, signal) { for await (const batch of batches) { - const columns = batch.columnNames const loadedVectors = batch.columns.map(function loadedVector(column) { - if (column.type === 'source' || column.type === 'computed') return undefined + if ('read' in column) return undefined return selectVector(column, batch.selection) }) const rowCount = selectedRowCount(batch.selection) @@ -90,14 +89,14 @@ export async function* batchesToRows(batches, signal) { * constructing compatibility `AsyncRow` values. * * @param {AsyncIterable} batches + * @param {string[]} names * @param {AbortSignal} [signal] * @returns {Promise[]>} */ -export async function collectBatches(batches, signal) { +export async function collectBatches(batches, names, signal) { /** @type {Record[]} */ const rows = [] for await (const batch of batches) { - const names = batch.columnNames const results = batch.columns.map(function readColumn(_column, columnIndex) { return readBatchColumn({ batch, columnIndex, signal }) }) @@ -146,14 +145,12 @@ function makeValueBuffers(count) { } /** - * @param {string[]} columnNames * @param {SqlPrimitive[][]} values * @param {number} rowCount * @returns {AsyncBatch} */ -function loadedBatch(columnNames, values, rowCount) { +function loadedBatch(values, rowCount) { return { - columnNames, selection: { type: 'all', length: rowCount }, columns: values.map(function loadedColumn(columnValues) { return { type: 'values', values: columnValues, length: rowCount } diff --git a/src/backend/dataSource.js b/src/backend/dataSource.js index 9abac03..176b425 100644 --- a/src/backend/dataSource.js +++ b/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>>} */ 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) { diff --git a/src/execute/aggregates.js b/src/execute/aggregates.js index 36d35b5..996a45f 100644 --- a/src/execute/aggregates.js +++ b/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 diff --git a/src/execute/batchResults.js b/src/execute/batchResults.js index 4bd52c7..849a3c4 100644 --- a/src/execute/batchResults.js +++ b/src/execute/batchResults.js @@ -1,15 +1,12 @@ import { batchesToRows } from '../backend/batchAdapters.js' +import { bindQuerySignal } from './utils.js' /** - * @import { AsyncBatch, InternalBatchResults } from '../internalTypes.js' - * @import { QueryResults } from '../types.js' + * @import { AsyncBatch, QueryResults } from '../types.js' */ -/** @type {WeakMap} */ -const internalBatches = new WeakMap() - /** - * Creates a public row result backed by private batches. + * Creates query results backed by batches. * * @param {object} options * @param {string[]} options.columns @@ -19,23 +16,13 @@ const internalBatches = new WeakMap() * @param {AbortSignal} [options.signal] * @returns {QueryResults} */ -export function batchResult({ batches, signal, ...metadata }) { +export function batchResult({ batches: readBatches, signal, ...metadata }) { const results = { ...metadata, + batches: readBatches, rows() { - return batchesToRows(batches(), signal) + return batchesToRows(readBatches(), metadata.columns, signal) }, } - internalBatches.set(results, { columns: metadata.columns, batches, signal }) - return results -} - -/** - * Returns the private batch execution path for a result, when available. - * - * @param {QueryResults} results - * @returns {InternalBatchResults | undefined} - */ -export function batchResultsFor(results) { - return internalBatches.get(results) + return bindQuerySignal(results, signal) } diff --git a/src/execute/batches.js b/src/execute/batches.js index 720734b..1f1ac1a 100644 --- a/src/execute/batches.js +++ b/src/execute/batches.js @@ -3,8 +3,8 @@ import { keyify } from './utils.js' import { yieldToEventLoop } from './yield.js' /** - * @import { AsyncBatch, BatchColumn, BatchProjection, ColumnResult, ColumnVector, CompiledBatchExpression, RowSelection } from '../internalTypes.js' - * @import { SqlPrimitive } from '../types.js' + * @import { BatchProjection, CompiledBatchExpression } from '../internalTypes.js' + * @import { AsyncBatch, BatchColumn, ColumnResult, ColumnVector, ReadColumn, RowSelection, SqlPrimitive } from '../types.js' */ const INITIAL_FILTER_WINDOW_ROWS = 256 @@ -180,11 +180,10 @@ export async function* distinctBatches(batches, signal) { * so output columns remain aligned without hidden dependency columns. * * @param {AsyncIterable} batches - * @param {string[]} columnNames * @param {readonly BatchProjection[]} projections * @yields {AsyncBatch} */ -export async function* projectExpressionBatches(batches, columnNames, projections) { +export async function* projectExpressionBatches(batches, projections) { let rowOffset = 0 for await (const batch of batches) { const currentRowOffset = rowOffset @@ -192,13 +191,12 @@ export async function* projectExpressionBatches(batches, columnNames, projection /** @type {ColumnVector | undefined} */ let rowOrdinals yield { - columnNames, selection: batch.selection, columns: projections.map(function projectColumn(projection) { /** @type {BatchColumn} */ let column if (projection.type === 'column') { - column = batch.columns[projection.columnIndex] + column = projectedBatchColumn(batch, projection.columnIndex) } else if (projection.type === 'constant') { column = { type: 'constant', @@ -208,9 +206,8 @@ export async function* projectExpressionBatches(batches, columnNames, projection } else { rowOrdinals ??= selectionOrdinals(batch.selection) column = { - type: 'computed', + read: projection.expression.evaluate, input: batch, - expression: projection.expression, rowOffset: currentRowOffset, rowOrdinals, } @@ -221,6 +218,23 @@ export async function* projectExpressionBatches(batches, columnNames, projection } } +/** + * Preserves the source batch's deferred-read cache through projection. + * + * @param {AsyncBatch} batch + * @param {number} columnIndex + * @returns {BatchColumn} + */ +function projectedBatchColumn(batch, columnIndex) { + const column = batch.columns[columnIndex] + if (!('read' in column)) return column + /** @type {ReadColumn} */ + function readProjectedColumn({ selection, signal }) { + return readBatchColumn({ batch, columnIndex, selection, signal }) + } + return { read: readProjectedColumn } +} + /** * @param {ColumnVector} predicate * @param {number} rowCount diff --git a/src/execute/execute.js b/src/execute/execute.js index a497e41..76ef5e9 100644 --- a/src/execute/execute.js +++ b/src/execute/execute.js @@ -1,4 +1,6 @@ -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' @@ -7,7 +9,7 @@ import { planSql, planStatement } from '../plan/plan.js' import { collectColumnsFromExpr, statementScope } from '../plan/columns.js' import { validateScan, validateTable } from '../validation/tables.js' import { executeHashAggregate, executeScalarAggregate } from './aggregates.js' -import { batchResult, batchResultsFor } from './batchResults.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' @@ -17,8 +19,8 @@ import { executeWindow } from './window.js' import { yieldToEventLoop } from './yield.js' /** - * @import { AsyncBatch, BatchProjection, ColumnVector, CompiledBatchExpression } from '../internalTypes.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' */ @@ -284,6 +286,11 @@ export function executeScan(plan, context, existingColumnResult) { 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. const scanColumnOptions = plan.hints.columns?.length === 1 @@ -310,7 +317,7 @@ export function executeScan(plan, context, existingColumnResult) { /** @returns {AsyncIterable} */ function makeBatches() { /** @type {AsyncIterable} */ - let batches = columnBatches(columnResult.chunks(), columns, signal) + let batches = columnBatches(columnResult.chunks(), signal) if (residualFilter) { const targetRows = plan.hints.limit === undefined ? undefined @@ -368,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 @@ -379,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() { @@ -404,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} */ + function makeBatches() { + /** @type {AsyncIterable} */ + 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 * @@ -426,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 @@ -590,12 +749,11 @@ function referencesRowScope(expression, columns, context) { */ function executeFilter(plan, context) { const child = executePlan({ plan: plan.child, context }) - const childBatches = batchResultsFor(child) - const expression = childBatches - ? compileUnscopedBatchExpression(plan.condition, childBatches.columns, context) + const expression = child.batches + ? compileUnscopedBatchExpression(plan.condition, child.columns, context) : undefined - if (expression && childBatches) { - const readChildBatches = childBatches.batches + if (expression && child.batches) { + const readChildBatches = child.batches /** @returns {AsyncIterable} */ function makeBatches() { return filterBatches(readChildBatches(), expression, context.signal) @@ -634,15 +792,14 @@ function executeProject(plan, context) { return child.columns.includes(sourceName) }) - const childBatches = batchResultsFor(child) - const projection = childBatches + const projection = child.batches ? batchProjection(plan.columns, columns, child.columns, context) : undefined - if (projection && childBatches) { - const readChildBatches = childBatches.batches + if (projection && child.batches) { + const readChildBatches = child.batches /** @returns {AsyncIterable} */ function makeBatches() { - return projectExpressionBatches(readChildBatches(), columns, projection) + return projectExpressionBatches(readChildBatches(), projection) } return batchResult({ columns, @@ -741,9 +898,8 @@ function executeProject(plan, context) { */ function executeDistinct(plan, context) { const child = executePlan({ plan: plan.child, context }) - const childBatches = batchResultsFor(child) - if (childBatches) { - const readChildBatches = childBatches.batches + if (child.batches) { + const readChildBatches = child.batches /** @returns {AsyncIterable} */ function makeBatches() { return distinctBatches(readChildBatches(), context.signal) @@ -813,9 +969,8 @@ function executeDistinct(plan, context) { */ function executeLimit(plan, context) { const child = executePlan({ plan: plan.child, context }) - const childBatches = batchResultsFor(child) - if (childBatches) { - const readChildBatches = childBatches.batches + if (child.batches) { + const readChildBatches = child.batches /** @returns {AsyncIterable} */ function makeBatches() { return limitBatches(readChildBatches(), plan.limit, plan.offset, context.signal) @@ -841,17 +996,15 @@ function executeLimit(plan, context) { * arrays. Each source chunk remains the async scheduling unit. * * @param {AsyncIterable>} chunks - * @param {string[]} columnNames * @param {AbortSignal} [signal] * @yields {AsyncBatch} */ -async function* columnBatches(chunks, columnNames, signal) { +async function* columnBatches(chunks, signal) { for await (const chunk of chunks) { signal?.throwIfAborted() const vector = vectorFromChunk(chunk) /** @type {AsyncBatch} */ const batch = { - columnNames, selection: { type: 'all', length: vector.length }, columns: [vector], } @@ -880,7 +1033,7 @@ function vectorFromChunk(chunk) { /** * @param {ArrayLike} values - * @returns {values is import('../internalTypes.js').NumericArray} + * @returns {values is import('../types.js').NumericArray} */ function isNumericArray(values) { return values instanceof Int8Array @@ -941,9 +1094,6 @@ function batchProjection(planColumns, outputColumns, childColumns, context) { projections.push({ type: 'expression', expression }) } if (projections.length !== outputColumns.length) return undefined - for (let index = 0; index < projections.length; index++) { - projections[index] = projections[outputColumns.lastIndexOf(outputColumns[index])] - } return projections } @@ -957,7 +1107,7 @@ function identifierColumnIndex(identifier, childColumns) { const sourceName = identifier.prefix ? `${identifier.prefix}.${identifier.name}` : identifier.name - const index = childColumns.indexOf(sourceName) + const index = childColumns.lastIndexOf(sourceName) if (index >= 0) return index const suffix = `.${identifier.name}` diff --git a/src/execute/streamingAggregate.js b/src/execute/streamingAggregate.js index 1ac069a..23347e8 100644 --- a/src/execute/streamingAggregate.js +++ b/src/execute/streamingAggregate.js @@ -5,14 +5,13 @@ import { evaluateAll, evaluateExpr } from '../expression/evaluate.js' import { collectColumnsFromExpr } from '../plan/columns.js' import { isAggregateFunc } from '../validation/functions.js' import { finalizeAccumulator, newAccumulator, updateAccumulator } from './accumulator.js' -import { batchResultsFor } from './batchResults.js' import { sortEntriesByTerms } from './sort.js' import { keyify } from './utils.js' import { yieldToEventLoop } from './yield.js' /** - * @import { AsyncBatch, BatchAggregateInputs, ColumnVector, CompiledBatchExpression } from '../internalTypes.js' - * @import { AsyncCells, AsyncRow, ExecuteContext, ExprNode, FunctionNode, IdentifierNode, QueryResults, SelectColumn, SqlPrimitive } from '../types.js' + * @import { BatchAggregateInputs, CompiledBatchExpression } from '../internalTypes.js' + * @import { AsyncBatch, AsyncCells, AsyncRow, ColumnVector, ExecuteContext, ExprNode, FunctionNode, IdentifierNode, QueryResults, SelectColumn, SqlPrimitive } from '../types.js' * @import { HashAggregateNode, ScalarAggregateNode } from '../plan/types.js' * @import { Accumulator } from './accumulator.js' */ @@ -438,7 +437,7 @@ function compileBatchAggregateInputs(groupBy, specs, columns, context) { /** @type {CompiledBatchExpression[]} */ const keys = [] for (const expression of groupBy) { - if (referencesOuterScope(expression, context)) return undefined + if (referencesRowScope(expression, context)) return undefined const key = compileBatchExpression(expression, columns) if (!key) return undefined keys.push(key) @@ -450,8 +449,8 @@ function compileBatchAggregateInputs(groupBy, specs, columns, context) { const args = [] for (const spec of specs) { if (spec.node.filter && !spec.star) return undefined - if (spec.node.filter && referencesOuterScope(spec.node.filter, context)) return undefined - if (!spec.star && referencesOuterScope(spec.node.args[0], context)) return undefined + if (spec.node.filter && referencesRowScope(spec.node.filter, context)) return undefined + if (!spec.star && referencesRowScope(spec.node.args[0], context)) return undefined const filter = spec.node.filter ? compileBatchExpression(spec.node.filter, columns) : undefined @@ -470,21 +469,21 @@ function compileBatchAggregateInputs(groupBy, specs, columns, context) { } /** - * Returns whether an expression reads a qualified identifier from the - * enclosing query rather than the current aggregate input. + * Returns whether an expression reads a qualified identifier whose table + * scope the batch compiler cannot distinguish from struct-field access. * * @param {ExprNode} expression * @param {ExecuteContext} context * @returns {boolean} */ -function referencesOuterScope(expression, context) { +function referencesRowScope(expression, context) { /** @type {IdentifierNode[]} */ const identifiers = [] collectColumnsFromExpr(expression, identifiers) - return identifiers.some(function isOuterReference(identifier) { - return Boolean(identifier.prefix && - context.outerAliases?.has(identifier.prefix) && - !context.scope?.includes(identifier.prefix)) + return identifiers.some(function isScopedReference(identifier) { + return Boolean(identifier.prefix && ( + context.scope?.includes(identifier.prefix) || context.outerAliases?.has(identifier.prefix) + )) }) } @@ -574,13 +573,12 @@ async function accumulateBatch({ batch, inputs, specs, groups, context, rowOffse async function accumulateGroups({ child, groupBy, specs, needsRow, context }) { /** @type {Map} */ const groups = new Map() - const batchResults = batchResultsFor(child) - const batchInputs = batchResults && !needsRow - ? compileBatchAggregateInputs(groupBy, specs, batchResults.columns, context) + const batchInputs = child.batches && !needsRow + ? compileBatchAggregateInputs(groupBy, specs, child.columns, context) : undefined - if (batchInputs && batchResults) { + if (batchInputs && child.batches) { let rowOffset = 0 - for await (const batch of batchResults.batches()) { + for await (const batch of child.batches()) { await accumulateBatch({ batch, inputs: batchInputs, specs, groups, context, rowOffset }) rowOffset += selectedRowCount(batch.selection) context.signal?.throwIfAborted() diff --git a/src/execute/utils.js b/src/execute/utils.js index f281f09..53794db 100644 --- a/src/execute/utils.js +++ b/src/execute/utils.js @@ -1,5 +1,4 @@ import { collectBatches } from '../backend/batchAdapters.js' -import { batchResultsFor } from './batchResults.js' /** * @import { AsyncRow, OrderByItem, QueryResults, SqlPrimitive } from '../types.js' @@ -7,6 +6,37 @@ import { batchResultsFor } from './batchResults.js' const primitiveTypes = new Set(['number', 'bigint', 'boolean', 'string']) +/** @type {WeakMap} */ +const querySignals = new WeakMap() + +/** + * Associates an execution signal with results without expanding the public + * result shape solely for the collection adapter. + * + * @param {QueryResults} results + * @param {AbortSignal} [signal] + * @returns {QueryResults} + */ +export function bindQuerySignal(results, signal) { + if (!signal || querySignals.get(results) === signal) return results + + const { batches, rows } = results + results.rows = async function* boundRows() { + signal.throwIfAborted() + yield* rows.call(results) + signal.throwIfAborted() + } + if (batches) { + results.batches = async function* boundBatches() { + signal.throwIfAborted() + yield* batches.call(results) + signal.throwIfAborted() + } + } + querySignals.set(results, signal) + return results +} + /** * Compares two values for a single ORDER BY term, handling nulls and direction * @@ -52,8 +82,9 @@ export function compareForTerm(a, b, term) { * @returns {Promise[]>} array of all yielded values */ export async function collect(results) { - const batchResults = batchResultsFor(results) - if (batchResults) return await collectBatches(batchResults.batches(), batchResults.signal) + if (results.batches) { + return await collectBatches(results.batches(), results.columns, querySignals.get(results)) + } // Collect all rows first, then materialize cells concurrently // This enables dataloader-style batching of cell accessors diff --git a/src/expression/batch.js b/src/expression/batch.js index ebcc176..2106d7b 100644 --- a/src/expression/batch.js +++ b/src/expression/batch.js @@ -8,8 +8,8 @@ import { applyCast, evaluateJsonExtract } from './scalar.js' import { evaluateStringFunc } from './strings.js' /** - * @import { AsyncBatch, ColumnResult, ColumnVector, CompiledBatchExpression, CompileState, EvaluationContext, RowSelection, ValueKernel } from '../internalTypes.js' - * @import { ExprNode, FunctionNode, SqlPrimitive } from '../types.js' + * @import { CompiledBatchExpression, CompileState, ValueKernel } from '../internalTypes.js' + * @import { AsyncBatch, ColumnReadRequest, ColumnResult, ColumnVector, ExprNode, FunctionNode, RowSelection, SqlPrimitive } from '../types.js' */ const YIELD_INTERVAL = 4000 @@ -343,7 +343,7 @@ function compileCaseEvaluator(node, columns) { * @param {'AND' | 'OR'} operator * @param {CompiledBatchExpression} left * @param {CompiledBatchExpression} right - * @param {EvaluationContext} context + * @param {ColumnReadRequest} context * @returns {Promise} */ async function evaluateLogical(operator, left, right, context) { @@ -383,7 +383,7 @@ async function evaluateLogical(operator, left, right, context) { * @param {CompiledBatchExpression | undefined} caseExpression * @param {{ condition: CompiledBatchExpression, result: CompiledBatchExpression }[]} clauses * @param {CompiledBatchExpression | undefined} elseResult - * @param {EvaluationContext} context + * @param {ColumnReadRequest} context * @returns {Promise} */ async function evaluateCase(caseExpression, clauses, elseResult, context) { @@ -437,7 +437,7 @@ async function evaluateCase(caseExpression, clauses, elseResult, context) { * argument, matching COALESCE's lazy row semantics. * * @param {CompiledBatchExpression[]} arguments_ - * @param {EvaluationContext} context + * @param {ColumnReadRequest} context * @returns {Promise} */ async function evaluateCoalesce(arguments_, context) { @@ -463,9 +463,9 @@ async function evaluateCoalesce(arguments_, context) { } /** - * @param {EvaluationContext} context + * @param {ColumnReadRequest} context * @param {Uint32Array} indices - positions in the context's selected rows - * @returns {EvaluationContext} + * @returns {ColumnReadRequest} */ function subsetContext(context, indices) { const length = selectedRowCount(context.selection) @@ -492,7 +492,7 @@ function allIndices(length) { } /** - * @param {EvaluationContext} context + * @param {ColumnReadRequest} context * @param {(rowIndex: number, streamRowIndex: number) => SqlPrimitive} evaluate * @returns {Promise} */ @@ -522,7 +522,7 @@ async function visitRows(length, signal, visit) { } /** - * @param {EvaluationContext} context + * @param {ColumnReadRequest} context * @param {number} rowIndex * @returns {number} */ @@ -540,7 +540,7 @@ function resolveIdentifier(identifier, columns) { const sourceName = identifier.prefix ? `${identifier.prefix}.${identifier.name}` : identifier.name - const exact = columns.indexOf(sourceName) + const exact = columns.lastIndexOf(sourceName) if (exact >= 0) return [{ columnIndex: exact }] if (identifier.prefix) { @@ -562,7 +562,7 @@ function resolveIdentifier(identifier, columns) { if (baseMatches.length === 1) { accesses.push({ columnIndex: baseMatches[0], field: identifier.name }) } - const bare = columns.indexOf(identifier.name) + const bare = columns.lastIndexOf(identifier.name) if (bare >= 0) accesses.push({ columnIndex: bare }) return accesses.length > 0 ? accesses : undefined } diff --git a/src/index.d.ts b/src/index.d.ts index a10447d..d1463f4 100644 --- a/src/index.d.ts +++ b/src/index.d.ts @@ -1,19 +1,38 @@ -import type { AsyncDataSource, AsyncRow, ExecuteContext, ExecuteSqlOptions, ExprNode, ParseSqlOptions, PlanSqlOptions, QueryPlan, QueryResults, SqlPrimitive, Statement, Token } from './types.js' +import type { AsyncBatch, AsyncDataSource, AsyncRow, ColumnResult, ColumnVector, ExecuteContext, ExecuteSqlOptions, ExprNode, ParseSqlOptions, PlanSqlOptions, QueryPlan, QueryResults, ReadBatchColumnOptions, RowsToBatchesOptions, RowSelection, SqlPrimitive, Statement, Token } from './types.js' export type { + AsyncBatch, AsyncCells, AsyncDataSource, AsyncRow, + BatchColumn, + ColumnDemand, + ColumnReadRequest, + ColumnResult, + ColumnVector, ExecuteContext, ExecuteSqlOptions, ExprNode, + Field, + NumericArray, ParseSqlOptions, PlanSqlOptions, + PreparedScan, + PrepareScan, QueryPlan, QueryResults, + ReadBatchColumnOptions, + ReadColumn, + RelationSchema, + RowsToBatchesOptions, + RowSelection, ScanOptions, + ScanProperties, + ScanRequest, + ScanResidual, ScanResults, SelectStatement, SetOperationStatement, + SqlType, SqlPrimitive, Statement, Token, @@ -96,6 +115,30 @@ export function asyncRow(row: Record, columns: string[]): export function cachedDataSource(source: AsyncDataSource): AsyncDataSource +export function selectedRowCount(selection: RowSelection): number + +export function composeSelections(outer: RowSelection, inner: RowSelection): RowSelection + +export function valueAt(vector: ColumnVector, index: number): SqlPrimitive + +export function selectVector(vector: ColumnVector, selection: RowSelection): ColumnVector + +export function readBatchColumn(options: ReadBatchColumnOptions): ColumnResult + +export function selectBatch(batch: AsyncBatch, selection: RowSelection): AsyncBatch + +export function rowsToBatches( + rows: AsyncIterable, + columns: string[], + options?: RowsToBatchesOptions, +): AsyncIterable + +export function batchesToRows( + batches: AsyncIterable, + columns: string[], + signal?: AbortSignal, +): AsyncIterable + /** * Generates a default alias for a derived column expression. * Useful for generating column names pre-execution. diff --git a/src/index.js b/src/index.js index 504ff0d..890c9f1 100644 --- a/src/index.js +++ b/src/index.js @@ -5,4 +5,13 @@ export { planSql } from './plan/plan.js' export { tokenizeSql } from './parse/tokenize.js' export { collect } from './execute/utils.js' export { asyncRow, cachedDataSource } from './backend/dataSource.js' +export { + composeSelections, + readBatchColumn, + selectBatch, + selectedRowCount, + selectVector, + valueAt, +} from './backend/batch.js' +export { batchesToRows, rowsToBatches } from './backend/batchAdapters.js' export { derivedAlias } from './expression/alias.js' diff --git a/src/internalTypes.d.ts b/src/internalTypes.d.ts index 4368878..567fed0 100644 --- a/src/internalTypes.d.ts +++ b/src/internalTypes.d.ts @@ -1,47 +1,7 @@ -import type { SqlPrimitive } from './types.js' - -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 { - selection: RowSelection - signal?: AbortSignal -} - -export type ColumnResult = ColumnVector | Promise -export type ReadColumn = (request: ColumnReadRequest) => ColumnResult - -export interface EvaluationContext { - batch: AsyncBatch - selection: RowSelection - signal?: AbortSignal - rowOffset?: number - rowOrdinals?: ColumnVector -} +import type { ColumnReadRequest, ColumnResult, ColumnVector, SqlPrimitive } from './types.js' export interface CompiledBatchExpression { - evaluate(context: EvaluationContext): ColumnResult + evaluate(context: ColumnReadRequest): ColumnResult } export type ValueKernel = ( @@ -66,38 +26,3 @@ export type BatchProjection = | { type: 'column', columnIndex: number } | { type: 'constant', value: SqlPrimitive } | { type: 'expression', expression: CompiledBatchExpression } - -export type BatchColumn = - | ColumnVector - | { type: 'source', read: ReadColumn } - | { - type: 'computed' - input: AsyncBatch - expression: CompiledBatchExpression - rowOffset: number - rowOrdinals: ColumnVector - } - -export interface AsyncBatch { - columnNames: string[] - 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 InternalBatchResults { - columns: string[] - batches(): AsyncIterable - signal?: AbortSignal -} diff --git a/src/plan/columns.js b/src/plan/columns.js index 03e5960..ad29965 100644 --- a/src/plan/columns.js +++ b/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, + * tables?: Record, + * outerAliases?: Set, + * }} ColumnCollectionContext + */ + /** * @param {FromTable | FromSubquery | FromFunction | undefined} from * @returns {string | undefined} @@ -61,9 +70,12 @@ export function tableFunctionColumnNames(from) { * @param {SelectStatement} options.select * @param {IdentifierNode[]} [options.parentColumns] - columns needed by the parent query * @param {Set} [options.scopeColumns] - bare column names available in the current scope + * @param {Map} [options.cteColumns] + * @param {Record} [options.tables] + * @param {string[]} [options.outerScope] * @returns {Map} */ -export function extractColumns({ select, parentColumns, scopeColumns }) { +export function extractColumns({ select, parentColumns, scopeColumns, cteColumns, tables, outerScope }) { /** @type {Map} */ const result = new Map() @@ -75,6 +87,11 @@ export function extractColumns({ select, parentColumns, scopeColumns }) { 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)) { @@ -134,26 +151,26 @@ export function extractColumns({ select, parentColumns, scopeColumns }) { 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) { @@ -166,7 +183,7 @@ export function extractColumns({ select, parentColumns, scopeColumns }) { /** @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] }) } @@ -235,49 +252,50 @@ export function extractColumns({ select, parentColumns, scopeColumns }) { * @param {ExprNode} expr * @param {IdentifierNode[]} columns * @param {Set} [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. @@ -286,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) + } + } } /** @@ -349,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) } } @@ -483,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) : []) } /** diff --git a/src/plan/plan.js b/src/plan/plan.js index 73acee9..a5d4665 100644 --- a/src/plan/plan.js +++ b/src/plan/plan.js @@ -208,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, scopeColumns }) + 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 diff --git a/src/types.d.ts b/src/types.d.ts index ce45644..6455bbf 100644 --- a/src/types.d.ts +++ b/src/types.d.ts @@ -11,10 +11,148 @@ export { QueryPlan } from './plan/types.js' export interface QueryResults { columns: string[] rows(): AsyncGenerator + batches?(): AsyncIterable 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 +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 +} + +export type PrepareScan = (request: ScanRequest) => PreparedScan + // parseSql(options) export interface ParseSqlOptions { query: string @@ -72,16 +210,30 @@ export type AsyncCell = () => Promise export type Row = Record[] /** - * 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> | 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. diff --git a/src/validation/tables.js b/src/validation/tables.js index acb0113..ddd4bec 100644 --- a/src/validation/tables.js +++ b/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, }) diff --git a/test/asyncDataSource.test-d.ts b/test/asyncDataSource.test-d.ts new file mode 100644 index 0000000..9724fa8 --- /dev/null +++ b/test/asyncDataSource.test-d.ts @@ -0,0 +1,19 @@ +import type { AsyncDataSource, PrepareScan, ScanResults } from '../src/types.js' +import type { ReadColumn } from '../src/index.js' + +declare const prepareScan: PrepareScan +declare const readColumn: ReadColumn +declare const scan: () => ScanResults + +const legacySource: AsyncDataSource = { columns: ['id'], scan } +const preparedSource: AsyncDataSource = { + schema: { fields: [] }, + prepareScan, +} + +// @ts-expect-error A data source must implement a usable scan contract. +const emptySource: AsyncDataSource = {} +// @ts-expect-error prepareScan requires a schema. +const prepareOnlySource: AsyncDataSource = { prepareScan } + +void [legacySource, preparedSource, emptySource, prepareOnlySource, readColumn] diff --git a/test/backend/batch.test.js b/test/backend/batch.test.js index 319c1fd..a122c1a 100644 --- a/test/backend/batch.test.js +++ b/test/backend/batch.test.js @@ -9,11 +9,9 @@ import { } from '../../src/backend/batch.js' /** - * @import { AsyncBatch, ColumnVector, ReadColumn } from '../../src/internalTypes.js' + * @import { AsyncBatch, ColumnVector, ReadColumn } from '../../src/types.js' */ -const schema = ['value'] - describe('row selections', () => { it('counts every selection representation', () => { expect(selectedRowCount({ type: 'all', length: 5 })).toBe(5) @@ -105,9 +103,8 @@ describe('async batches', () => { const read = vi.fn(readColumn) /** @type {AsyncBatch} */ const batch = { - columnNames: schema, selection: { type: 'all', length: 5 }, - columns: [{ type: 'source', read }], + columns: [{ read }], } expect(read).not.toHaveBeenCalled() @@ -124,8 +121,11 @@ describe('async batches', () => { expect(await first).toEqual({ type: 'constant', value: 7, length: 2 }) expect(read).toHaveBeenCalledTimes(1) expect(read).toHaveBeenCalledWith({ + batch: selected, selection: { type: 'range', start: 1, end: 3, length: 5 }, signal: undefined, + rowOffset: undefined, + rowOrdinals: undefined, }) }) @@ -134,7 +134,6 @@ describe('async batches', () => { const vector = { type: 'values', values: [1, 2, 3], length: 3 } /** @type {AsyncBatch} */ const batch = { - columnNames: schema, selection: { type: 'all', length: 3 }, columns: [vector], } @@ -145,10 +144,8 @@ describe('async batches', () => { it('validates deferred vector alignment', async () => { /** @type {AsyncBatch} */ const batch = { - columnNames: schema, selection: { type: 'all', length: 3 }, columns: [{ - type: 'source', read() { return Promise.resolve({ type: 'values', values: [1], length: 1 }) }, @@ -179,9 +176,8 @@ describe('async batches', () => { const read = vi.fn(readColumn) /** @type {AsyncBatch} */ const batch = { - columnNames: schema, selection: { type: 'all', length: 1 }, - columns: [{ type: 'source', read }], + columns: [{ read }], } const rejected = readBatchColumn({ batch, columnIndex: 0, signal: first.signal }) diff --git a/test/backend/batchAdapters.test.js b/test/backend/batchAdapters.test.js index c2b862d..da40812 100644 --- a/test/backend/batchAdapters.test.js +++ b/test/backend/batchAdapters.test.js @@ -5,8 +5,7 @@ import { batchesToRows, rowsToBatches } from '../../src/backend/batchAdapters.js import { batchResult } from '../../src/execute/batchResults.js' /** - * @import { AsyncBatch, ColumnVector, ReadColumn } from '../../src/internalTypes.js' - * @import { AsyncRow } from '../../src/types.js' + * @import { AsyncBatch, AsyncRow, ColumnVector, ReadColumn } from '../../src/types.js' */ const schema = ['id', 'name'] @@ -28,7 +27,7 @@ describe('batch adapters', () => { expect(batches.map(function summarize(batch) { return { rowCount: selectedRowCount(batch.selection), - ids: batch.columns[0].type === 'values' + ids: 'type' in batch.columns[0] && batch.columns[0].type === 'values' ? batch.columns[0].values : [], } @@ -48,12 +47,11 @@ describe('batch adapters', () => { const read = vi.fn(readColumn) /** @type {AsyncBatch} */ const batch = { - columnNames: [schema[0]], selection: { type: 'all', length: 2 }, - columns: [{ type: 'source', read }], + columns: [{ read }], } - const iterator = batchesToRows(asyncValues([batch]))[Symbol.asyncIterator]() + const iterator = batchesToRows(asyncValues([batch]), [schema[0]])[Symbol.asyncIterator]() const first = await iterator.next() expect(read).not.toHaveBeenCalled() if (first.done || !first.value) throw new Error('expected a row') @@ -69,7 +67,6 @@ describe('batch adapters', () => { it('adapts loaded columns through the batch selection', async () => { /** @type {AsyncBatch} */ const batch = { - columnNames: schema, selection: { type: 'indices', indices: new Uint32Array([2, 0]), @@ -90,7 +87,7 @@ describe('batch adapters', () => { } const rows = [] - for await (const row of batchesToRows(asyncValues([batch]))) { + for await (const row of batchesToRows(asyncValues([batch]), schema)) { rows.push({ id: await row.cells.id(), name: await row.cells.name(), @@ -106,7 +103,6 @@ describe('batch adapters', () => { it('collects native batches without consuming the row adapter', async () => { /** @type {AsyncBatch} */ const batch = { - columnNames: schema, selection: { type: 'indices', indices: new Uint32Array([2, 0]), @@ -145,7 +141,6 @@ describe('batch adapters', () => { const reason = new Error('batch collection aborted') /** @type {AsyncBatch} */ const batch = { - columnNames: [schema[0]], selection: { type: 'all', length: 1 }, columns: [{ type: 'constant', diff --git a/test/execute/batchAggregate.test.js b/test/execute/batchAggregate.test.js index bdedf58..391c615 100644 --- a/test/execute/batchAggregate.test.js +++ b/test/execute/batchAggregate.test.js @@ -2,40 +2,86 @@ import { describe, expect, it, vi } from 'vitest' import { collect, executeSql } from '../../src/index.js' /** - * @import { AsyncDataSource, ScanColumnResults } from '../../src/types.js' + * @import { AsyncBatch, AsyncDataSource, PrepareScan, ReadColumn, RelationSchema, ScanColumnResults } from '../../src/types.js' */ -describe('private batch aggregate execution', () => { - it('groups computed values from the existing column scan API', async () => { - const source = columnSource([1, 2, 3, 4]) +/** @type {RelationSchema} */ +const schema = { + fields: [ + { id: 1, name: 'provider', dataType: { type: 'string' }, nullable: false }, + { id: 2, name: 'session_id', dataType: { type: 'string' }, nullable: false }, + { id: 3, name: 'attributes', dataType: { type: 'unknown' }, nullable: true }, + ], +} + +/** @type {RelationSchema} */ +const aliasSchema = { + fields: [ + { id: 1, name: 'd', dataType: { type: 'unknown' }, nullable: false }, + { id: 2, name: 'keep', dataType: { type: 'boolean' }, nullable: false }, + { id: 3, name: 'e', dataType: { type: 'unknown' }, nullable: false }, + ], +} + +const aliasValues = { + d: [{ keep: false }, { keep: false }], + keep: [true, true], + e: [{ keep: false }, { keep: false }], +} + +describe('batch aggregate execution', () => { + it('groups production token expressions without using the row adapter', async () => { + /** @type {ReadColumn} */ + function readAttributes({ selection }) { + expect(selection).toEqual({ type: 'all', length: 4 }) + return { + type: 'values', + values: [ + { usage: { input_tokens: 10 } }, + { usage: { input_tokens: 7 } }, + { usage: { input_tokens: 20 } }, + { usage: {} }, + ], + length: 4, + } + } + const attributes = vi.fn(readAttributes) + const batch = loadedBatch(attributes) + /** @type {PrepareScan} */ + function prepareScan() { + return { + schema, + residual: {}, + properties: { exactRows: 4, maxRows: 4 }, + async *batches() { yield batch }, + } + } + /** @type {AsyncDataSource} */ + const source = { + columns: schema.fields.map(function fieldName(field) { return field.name }), + schema, + prepareScan, + scan: vi.fn(function scan() { throw new Error('legacy row scan should not be called') }), + } const results = executeSql({ - tables: { data: source }, - query: `SELECT id % 2 AS bucket, + tables: { messages: source }, + query: `SELECT provider, COUNT(*) AS parts, - COUNT(DISTINCT id % 3) AS distinct_values, - SUM(id * 10) AS total, - COUNT(*) FILTER (WHERE id > 2) AS filtered - FROM data GROUP BY id % 2 ORDER BY bucket`, + COUNT(DISTINCT session_id) AS sessions, + COALESCE(SUM(CAST(JSON_EXTRACT(attributes, '$.usage.input_tokens') AS BIGINT)), 0) AS tokens + FROM messages GROUP BY provider ORDER BY provider`, }) expect(await collect(results)).toEqual([ - { bucket: 0, parts: 2, distinct_values: 2, total: 60, filtered: 1 }, - { bucket: 1, parts: 2, distinct_values: 2, total: 40, filtered: 1 }, + { provider: 'claude', parts: 2, sessions: 2, tokens: 30 }, + { provider: 'codex', parts: 2, sessions: 1, tokens: 7 }, ]) expect(source.scan).not.toHaveBeenCalled() + expect(attributes).toHaveBeenCalledTimes(1) }) - it('retains row semantics for filtered non-star aggregates', async () => { - const source = columnSource([1, 2, 3, 4]) - - await expect(collect(executeSql({ - tables: { data: source }, - query: 'SELECT SUM(id * 10) FILTER (WHERE id > 2) AS total FROM data', - }))).resolves.toEqual([{ total: 70 }]) - }) - - it('aggregates CASE, compound predicates, and NULLIF from private batches', async () => { + it('aggregates CASE, compound predicates, and NULLIF from native batches', async () => { const source = columnSource([1, 2, 3, 4]) await expect(collect(executeSql({ @@ -57,6 +103,36 @@ describe('private batch aggregate execution', () => { }))).resolves.toEqual([{ total: 5 }]) }) + it('preserves table-alias precedence in batch aggregate inputs', async () => { + const source = preparedSource(aliasSchema, aliasValues) + + await expect(collect(executeSql({ + tables: { data: source }, + query: `SELECT d.keep AS grouped, MAX(d) AS max_d, COUNTIF(d.keep) AS matching + FROM data d GROUP BY d.keep`, + }))).resolves.toEqual([{ + grouped: true, + max_d: { keep: false }, + matching: 2, + }]) + }) + + it('preserves table-alias precedence in compound batch aggregate inputs', async () => { + const source = preparedSource(aliasSchema, aliasValues) + + await expect(collect(executeSql({ + tables: { data: source }, + query: `SELECT d.keep AS grouped, MAX(d) AS max_d, COUNTIF(d.keep) AS matching + FROM data d GROUP BY d.keep + UNION ALL + SELECT e.keep AS grouped, MAX(e) AS max_d, COUNTIF(e.keep) AS matching + FROM data e GROUP BY e.keep`, + }))).resolves.toEqual([ + { grouped: true, max_d: { keep: false }, matching: 2 }, + { grouped: true, max_d: { keep: false }, matching: 2 }, + ]) + }) + it('preserves outer references in aggregate inputs', async () => { await expect(collect(executeSql({ tables: { @@ -73,6 +149,21 @@ describe('private batch aggregate execution', () => { }) }) +/** + * @param {ReadColumn} readAttributes + * @returns {AsyncBatch} + */ +function loadedBatch(readAttributes) { + return { + selection: { type: 'all', length: 4 }, + columns: [ + { type: 'values', values: ['claude', 'codex', 'claude', 'codex'], length: 4 }, + { type: 'values', values: ['a', 'b', 'c', 'b'], length: 4 }, + { read: readAttributes }, + ], + } +} + /** * @param {import('../../src/types.js').SqlPrimitive[]} values * @returns {AsyncDataSource} @@ -94,3 +185,37 @@ function columnSource(values) { }, } } + +/** + * @param {RelationSchema} sourceSchema + * @param {Record} values + * @returns {AsyncDataSource} + */ +function preparedSource(sourceSchema, values) { + const length = Object.values(values)[0]?.length ?? 0 + return { + schema: sourceSchema, + prepareScan(request) { + const fields = request.columns.map(function requestedField(demand) { + const field = sourceSchema.fields.find(function fieldById(candidate) { + return candidate.id === demand.field + }) + if (!field) throw new Error(`unknown field: ${demand.field}`) + return field + }) + return { + schema: { fields }, + residual: {}, + properties: { exactRows: length }, + async *batches() { + yield { + selection: { type: 'all', length }, + columns: fields.map(function loadedField(field) { + return { type: 'values', values: values[field.name], length } + }), + } + }, + } + }, + } +} diff --git a/test/execute/batchExecution.test.js b/test/execute/batchExecution.test.js index bfa3d2c..4e9f3e6 100644 --- a/test/execute/batchExecution.test.js +++ b/test/execute/batchExecution.test.js @@ -1,6 +1,5 @@ import { describe, expect, it, vi } from 'vitest' -import { collect, executeSql } from '../../src/index.js' -import { batchResultsFor } from '../../src/execute/batchResults.js' +import { collect, executeSql, readBatchColumn, valueAt } from '../../src/index.js' /** * @import { AsyncDataSource, ScanColumnResults } from '../../src/types.js' @@ -15,17 +14,16 @@ describe('native batch execution', () => { query: 'SELECT data.id AS value FROM data', }) - expect(Object.keys(results)).toEqual(['columns', 'numRows', 'maxRows', 'rows']) - const batchResults = batchResultsFor(results) - if (!batchResults) throw new Error('expected internal batches') - const iterator = batchResults.batches()[Symbol.asyncIterator]() + expect(results.batches).toBeTypeOf('function') + if (!results.batches) throw new Error('expected native batches') + const iterator = results.batches()[Symbol.asyncIterator]() const first = await iterator.next() if (first.done) throw new Error('expected a batch') const column = first.value.columns[0] - expect(column.type).toBe('typed') - if (column.type !== 'typed') throw new Error('expected a typed vector') + expect('type' in column && column.type).toBe('typed') + if (!('type' in column) || column.type !== 'typed') throw new Error('expected a typed vector') expect(column.values).toBe(values) - expect(first.value.columnNames[0]).toBe('value') + expect(results.columns[0]).toBe('value') }) it('applies limit and offset as a zero-copy selection across chunks', async () => { @@ -50,23 +48,20 @@ describe('native batch execution', () => { expect(rows).toEqual([4, 5]) }) - it('keeps computed projection private and lazy at column granularity', async () => { + it('keeps computed projection lazy at column granularity', async () => { const source = columnSource(function chunks() { return [['a', 'abcd', null]] }) const results = executeSql({ tables: { data: source }, query: 'SELECT LENGTH(id) + 1 AS size FROM data', }) - expect(Object.keys(results)).toEqual(['columns', 'numRows', 'maxRows', 'rows']) - const batchResults = batchResultsFor(results) - if (!batchResults) throw new Error('expected internal batches') - const iterator = batchResults.batches()[Symbol.asyncIterator]() + expect(results.batches).toBeTypeOf('function') + if (!results.batches) throw new Error('expected native batches') + const iterator = results.batches()[Symbol.asyncIterator]() const first = await iterator.next() if (first.done) throw new Error('expected a batch') const [column] = first.value.columns - expect(column.type).toBe('computed') - if (column.type !== 'computed') throw new Error('expected a computed column') - expect(column.expression).toBeDefined() + expect('read' in column).toBe(true) expect(await collect(results)).toEqual([{ size: 2 }, { size: 5 }, { size: null }]) }) @@ -92,23 +87,20 @@ describe('native batch execution', () => { query: 'SELECT CASE WHEN id IS NULL THEN 0 ELSE id END AS value FROM data', }) - expect(batchResultsFor(results)).toBeDefined() + expect(results.batches).toBeTypeOf('function') expect(await collect(results)).toEqual([{ value: 0 }, { value: 2 }]) }) - it('turns a residual predicate into a private batch selection', async () => { + it('turns a residual predicate into a batch selection', async () => { const source = columnSource(function chunks() { return [[null, 2, 3, 4]] }, false) const results = executeSql({ tables: { data: source }, query: 'SELECT id FROM data WHERE id > 2', }) - expect(Object.keys(results)).toEqual(['columns', 'numRows', 'maxRows', 'rows']) - expect(Object.hasOwn(results, 'batches')).toBe(false) - expect(Object.hasOwn(results, 'schema')).toBe(false) - const batchResults = batchResultsFor(results) - if (!batchResults) throw new Error('expected internal batches') - const iterator = batchResults.batches()[Symbol.asyncIterator]() + expect(results.batches).toBeTypeOf('function') + if (!results.batches) throw new Error('expected native batches') + const iterator = results.batches()[Symbol.asyncIterator]() const first = await iterator.next() if (first.done) throw new Error('expected a batch') expect(first.value.selection).toEqual({ @@ -118,7 +110,6 @@ describe('native batch execution', () => { }) expect(await collect(results)).toEqual([{ id: 3 }, { id: 4 }]) }) - it('applies a residual predicate before limit and offset', async () => { const source = columnSource(function chunks() { return [[1, 2, 3, 4], [5, 6, 7, 8]] }, false) const results = executeSql({ @@ -126,7 +117,7 @@ describe('native batch execution', () => { query: 'SELECT id FROM data WHERE id % 2 = 0 LIMIT 2 OFFSET 1', }) - expect(batchResultsFor(results)).toBeDefined() + expect(results.batches).toBeTypeOf('function') expect(await collect(results)).toEqual([{ id: 4 }, { id: 6 }]) }) @@ -137,7 +128,7 @@ describe('native batch execution', () => { query: 'SELECT id FROM data WHERE (id > 1 AND id < 4) OR id = 5', }) - expect(batchResultsFor(results)).toBeDefined() + expect(results.batches).toBeTypeOf('function') expect(await collect(results)).toEqual([{ id: 2 }, { id: 3 }, { id: 5 }]) }) @@ -148,20 +139,18 @@ describe('native batch execution', () => { query: 'SELECT id FROM data WHERE CASE WHEN id IS NULL THEN 0 ELSE id END > 2', }) - expect(batchResultsFor(results)).toBeDefined() + expect(results.batches).toBeTypeOf('function') expect(await collect(results)).toEqual([{ id: 3 }]) }) - it('filters a computed subquery through private batches', async () => { + it('filters a computed subquery through batches', async () => { const source = columnSource(function chunks() { return [[1, 2, 3, 4]] }) const results = executeSql({ tables: { data: source }, query: 'SELECT value FROM (SELECT id * 3 AS value FROM data) WHERE value > 6', }) - expect(Object.hasOwn(results, 'batches')).toBe(false) - expect(Object.hasOwn(results, 'schema')).toBe(false) - expect(batchResultsFor(results)).toBeDefined() + expect(results.batches).toBeTypeOf('function') expect(await collect(results)).toEqual([{ value: 9 }, { value: 12 }]) }) @@ -172,7 +161,7 @@ describe('native batch execution', () => { query: 'SELECT value FROM (SELECT id AS value FROM data) WHERE CASE WHEN value IS NULL THEN 0 ELSE value END > 2', }) - expect(batchResultsFor(results)).toBeDefined() + expect(results.batches).toBeTypeOf('function') expect(await collect(results)).toEqual([{ value: 3 }]) }) @@ -204,30 +193,58 @@ describe('native batch execution', () => { ]) }) - it('executes distinct over private computed batches', async () => { + it('executes distinct over computed batches', async () => { const source = columnSource(function chunks() { return [[1, 2, 3], [4, 5]] }) const results = executeSql({ tables: { data: source }, query: 'SELECT DISTINCT id % 2 AS value FROM data', }) - expect(Object.hasOwn(results, 'batches')).toBe(false) - expect(Object.hasOwn(results, 'schema')).toBe(false) - expect(batchResultsFor(results)).toBeDefined() + expect(results.batches).toBeTypeOf('function') expect(await collect(results)).toEqual([{ value: 1 }, { value: 0 }]) }) - it('falls back to row projection for duplicate aliases', async () => { + it('preserves duplicate alias projections by batch position', async () => { + const source = columnSource(function chunks() { return [[1, 2]] }) + const results = executeSql({ + tables: { data: source }, + query: 'SELECT id AS x, id + 1 AS x FROM data', + }) + + expect(results.batches).toBeTypeOf('function') + if (!results.batches) throw new Error('expected native batches') + const iterator = results.batches()[Symbol.asyncIterator]() + const first = await iterator.next() + if (first.done) throw new Error('expected a batch') + const left = await readBatchColumn({ batch: first.value, columnIndex: 0 }) + const right = await readBatchColumn({ batch: first.value, columnIndex: 1 }) + + expect([valueAt(left, 0), valueAt(left, 1)]).toEqual([1, 2]) + expect([valueAt(right, 0), valueAt(right, 1)]).toEqual([2, 3]) + }) + + it('uses the last duplicate alias in downstream batch expressions', async () => { const source = columnSource(function chunks() { return [[1, 2]] }) const results = executeSql({ tables: { data: source }, query: 'SELECT DISTINCT x FROM (SELECT id AS x, 1 AS x FROM data)', }) - expect(batchResultsFor(results)).toBeDefined() + expect(results.batches).toBeTypeOf('function') expect(await collect(results)).toEqual([{ x: 1 }]) }) + it('uses the last duplicate alias for qualified downstream batch expressions', async () => { + const source = columnSource(function chunks() { return [[1, 2, 3]] }) + const results = executeSql({ + tables: { data: source }, + query: 'SELECT id FROM (SELECT id, id AS x, 1 AS x FROM data) q WHERE q.x = 1', + }) + + expect(results.batches).toBeTypeOf('function') + expect(await collect(results)).toEqual([{ id: 1 }, { id: 2 }, { id: 3 }]) + }) + it('preserves projected row positions through a later offset', async () => { const source = columnSource(function chunks() { return [[1, 2, 3, { value: 4 }]] }) const results = executeSql({ diff --git a/test/execute/batches.test.js b/test/execute/batches.test.js index 5232cb2..75a82f2 100644 --- a/test/execute/batches.test.js +++ b/test/execute/batches.test.js @@ -5,8 +5,7 @@ import { compileBatchExpression } from '../../src/expression/batch.js' import { parseSql } from '../../src/parse/parse.js' /** - * @import { AsyncBatch, ReadColumn } from '../../src/internalTypes.js' - * @import { ExprNode } from '../../src/types.js' + * @import { AsyncBatch, ExprNode, ReadColumn } from '../../src/types.js' */ const schema = ['keep', 'payload'] @@ -20,11 +19,10 @@ describe('batch operators', () => { const read = vi.fn(readPayload) /** @type {AsyncBatch} */ const batch = { - columnNames: schema, selection: { type: 'all', length: 4 }, columns: [ { type: 'values', values: [true, false, true, false], length: 4 }, - { type: 'source', read }, + { read }, ], } const expression = compileBatchExpression(parseExpression('keep'), schema) @@ -129,6 +127,28 @@ describe('batch operators', () => { length: 2, }) }) + + it('yields while deduplicating a large loaded batch so timer aborts can fire', async () => { + const controller = new AbortController() + const batch = valueBatch(Array.from({ length: 20_000 }, function value(_item, index) { + return index + })) + const timer = setTimeout(function abortDistinct() { + controller.abort(new Error('distinct timed out')) + }, 0) + + async function consumeDistinct() { + for await (const result of distinctBatches(asyncValues([batch]), controller.signal)) { + selectedRowCount(result.selection) + } + } + + try { + await expect(consumeDistinct()).rejects.toThrow('distinct timed out') + } finally { + clearTimeout(timer) + } + }) }) /** @@ -158,7 +178,6 @@ async function* asyncValues(values) { */ function valueBatch(values) { return { - columnNames: ['value'], selection: { type: 'all', length: values.length }, columns: [{ type: 'values', diff --git a/test/execute/execute.aggregate.test.js b/test/execute/execute.aggregate.test.js index df633c1..6cb7fd8 100644 --- a/test/execute/execute.aggregate.test.js +++ b/test/execute/execute.aggregate.test.js @@ -1,4 +1,4 @@ -import { describe, expect, it } from 'vitest' +import { describe, expect, it, vi } from 'vitest' import { collect, executeSql } from '../../src/index.js' import { memorySource } from '../../src/backend/dataSource.js' @@ -1100,6 +1100,22 @@ describe('executeSql', () => { expect(result).toEqual([{ count_all: 5 }]) }) + it('should ignore prepareScan without a schema', async () => { + const legacy = memorySource({ data: [{ id: 1 }, { id: 2 }] }) + if (!legacy.columns || !legacy.scan) throw new Error('expected legacy source') + const scan = vi.fn(legacy.scan) + const prepareScan = vi.fn(function prepareScan() { + throw new Error('prepareScan should not be called without a schema') + }) + + await expect(collect(executeSql({ + tables: { data: { columns: legacy.columns, scan, prepareScan } }, + query: 'SELECT COUNT(*) FROM data', + }))).resolves.toEqual([{ count_all: 2 }]) + expect(scan).toHaveBeenCalledTimes(1) + expect(prepareScan).not.toHaveBeenCalled() + }) + it('should not optimize COUNT(column) with numRows', async () => { const data = [ { id: 1, name: 'Alice' }, diff --git a/test/execute/expensive.test.js b/test/execute/expensive.test.js index 0d140c0..bbd7fdf 100644 --- a/test/execute/expensive.test.js +++ b/test/execute/expensive.test.js @@ -21,6 +21,27 @@ const other = [ ] describe('expensive cell access', () => { + it('preserves the source receiver through the cache wrapper', async () => { + const base = memorySource({ data: [{ id: 1 }] }) + const source = { + ...base, + marker: 'configured', + /** + * @param {ScanOptions} options + * @returns {ScanResults} + */ + scan(options) { + expect(this.marker).toBe('configured') + return base.scan(options) + }, + } + + await expect(collect(executeSql({ + tables: { data: cachedDataSource(source) }, + query: 'SELECT id FROM data', + }))).resolves.toEqual([{ id: 1 }]) + }) + it('should make no expensive calls when not accessing expensive columns', async () => { await expect(countExpensiveCalls('SELECT id, name FROM data')).resolves.toBe(0) }) diff --git a/test/execute/preparedScan.test.js b/test/execute/preparedScan.test.js new file mode 100644 index 0000000..56ab819 --- /dev/null +++ b/test/execute/preparedScan.test.js @@ -0,0 +1,995 @@ +import { describe, expect, it, vi } from 'vitest' +import { selectedRowCount } from '../../src/backend/batch.js' +import { cachedDataSource, collect, executeSql } from '../../src/index.js' + +/** + * @import { AsyncBatch, AsyncDataSource, PreparedScan, PrepareScan, ReadColumn, RelationSchema, RowSelection, ScanRequest, SqlPrimitive } from '../../src/types.js' + */ + +/** @type {RelationSchema} */ +const schema = { + fields: [ + { id: 10, name: 'keep', dataType: { type: 'boolean' }, nullable: false }, + { id: 20, name: 'payload', dataType: { type: 'string' }, nullable: false }, + ], +} + +describe('prepared scans', () => { + it('preserves alias scope in correlated residual filters', async () => { + /** @type {RelationSchema} */ + const outerSchema = { + fields: [ + { id: 1, name: 'id', dataType: { type: 'number' }, nullable: false }, + { id: 2, name: 'payload', dataType: { type: 'string' }, nullable: false }, + ], + } + /** @type {RelationSchema} */ + const innerSchema = { + fields: [ + { id: 3, name: 'id', dataType: { type: 'number' }, nullable: false }, + ], + } + + await expect(collect(executeSql({ + tables: { + a: preparedValuesSource(outerSchema, { id: [1, 2], payload: ['yes', 'no'] }), + b: preparedValuesSource(innerSchema, { id: [1] }), + }, + query: `SELECT aa.payload + FROM a aa + WHERE EXISTS (SELECT 1 FROM b bb WHERE bb.id = aa.id)`, + }))).resolves.toEqual([{ payload: 'yes' }]) + }) + + it('preserves outer scope in filters above prepared derived tables', async () => { + /** @type {RelationSchema} */ + const ordersSchema = { + fields: [ + { id: 1, name: 'user_id', dataType: { type: 'number' }, nullable: false }, + ], + } + + await expect(collect(executeSql({ + tables: { + users: [{ id: 1 }, { id: 2 }], + orders: preparedValuesSource(ordersSchema, { user_id: [1] }), + }, + query: `SELECT u.id FROM users u + WHERE EXISTS ( + SELECT 1 FROM (SELECT user_id AS id FROM orders) q + WHERE q.id = u.id + ) + ORDER BY u.id`, + }))).resolves.toEqual([{ id: 1 }]) + }) + + it('preserves outer scope in prepared projections', async () => { + /** @type {RelationSchema} */ + const innerSchema = { + fields: [ + { id: 1, name: 'u', dataType: { type: 'unknown' }, nullable: false }, + ], + } + + await expect(collect(executeSql({ + tables: { + users: [{ id: 1 }, { id: 2 }], + items: preparedValuesSource(innerSchema, { u: [{ id: 99 }] }), + }, + query: `SELECT u.id, (SELECT u.id FROM items) AS projected + FROM users u ORDER BY u.id`, + }))).resolves.toEqual([ + { id: 1, projected: 1 }, + { id: 2, projected: 2 }, + ]) + }) + + it('preserves table alias precedence in prepared projections', async () => { + /** @type {RelationSchema} */ + const aliasSchema = { + fields: [ + { id: 1, name: 'd', dataType: { type: 'unknown' }, nullable: false }, + { id: 2, name: 'keep', dataType: { type: 'boolean' }, nullable: false }, + ], + } + + await expect(collect(executeSql({ + tables: { + data: preparedValuesSource(aliasSchema, { + d: [{ keep: false }], + keep: [true], + }), + }, + query: 'SELECT d AS value, d.keep AS projected FROM data d', + }))).resolves.toEqual([{ value: { keep: false }, projected: true }]) + }) + + it('resolves a table alias before an object field with the same name', async () => { + /** @type {RelationSchema} */ + const aliasSchema = { + fields: [ + { id: 1, name: 'd', dataType: { type: 'unknown' }, nullable: false }, + { id: 2, name: 'keep', dataType: { type: 'boolean' }, nullable: false }, + ], + } + const source = preparedValuesSource(aliasSchema, { + d: [{ keep: false }], + keep: [true], + }) + + await expect(collect(executeSql({ + tables: { data: source }, + query: 'SELECT d FROM data d WHERE d.keep', + }))).resolves.toEqual([{ d: { keep: false } }]) + }) + + it('preserves table alias precedence inside a CTE', async () => { + /** @type {RelationSchema} */ + const aliasSchema = { + fields: [ + { + id: 1, + name: 't', + dataType: { + type: 'struct', + fields: [{ id: 2, name: 'keep', dataType: { type: 'boolean' }, nullable: false }], + }, + nullable: false, + }, + { id: 3, name: 'keep', dataType: { type: 'boolean' }, nullable: false }, + { id: 4, name: 'payload', dataType: { type: 'string' }, nullable: false }, + ], + } + + await expect(collect(executeSql({ + tables: { + data: preparedValuesSource(aliasSchema, { + t: [{ keep: false }], + keep: [true], + payload: ['yes'], + }), + }, + query: `WITH q AS (SELECT payload FROM data t WHERE t.keep) + SELECT * FROM q`, + }))).resolves.toEqual([{ payload: 'yes' }]) + + await expect(collect(executeSql({ + tables: { + data: preparedValuesSource(aliasSchema, { + t: [{ keep: false }], + keep: [true], + payload: ['yes'], + }), + }, + query: `WITH q AS ( + SELECT payload FROM data t WHERE t.keep + UNION ALL + SELECT payload FROM data t WHERE t.keep + ) SELECT * FROM q`, + }))).resolves.toEqual([{ payload: 'yes' }, { payload: 'yes' }]) + }) + + it('shares a deferred read between direct and computed projections', async () => { + /** @type {RelationSchema} */ + const payloadSchema = { + fields: [ + { id: 1, name: 'payload', dataType: { type: 'string' }, nullable: false }, + ], + } + /** @type {ReadColumn} */ + function readPayload() { + return { type: 'values', values: ['value'], length: 1 } + } + const payload = vi.fn(readPayload) + /** @type {AsyncDataSource} */ + const source = { + schema: payloadSchema, + prepareScan() { + return { + schema: payloadSchema, + residual: {}, + properties: { exactRows: 1 }, + async *batches() { + yield { + selection: { type: 'all', length: 1 }, + columns: [{ read: payload }], + } + }, + } + }, + } + + await expect(collect(executeSql({ + tables: { data: source }, + query: 'SELECT payload, LENGTH(payload) AS length FROM data', + }))).resolves.toEqual([{ payload: 'value', length: 5 }]) + expect(payload).toHaveBeenCalledTimes(1) + + payload.mockClear() + await expect(collect(executeSql({ + tables: { data: source }, + query: 'SELECT payload AS first, payload AS second FROM data', + }))).resolves.toEqual([{ first: 'value', second: 'value' }]) + expect(payload).toHaveBeenCalledTimes(1) + }) + + it('passes the final filtered range to a deferred output column', async () => { + const controller = new AbortController() + /** @type {ReadColumn} */ + function readPayload({ selection, signal }) { + expect(signal).toBe(controller.signal) + expect(selection).toEqual({ + type: 'indices', + indices: new Uint32Array([1]), + length: 3, + }) + return { type: 'values', values: ['second'], length: 1 } + } + const payload = vi.fn(readPayload) + /** @type {AsyncBatch} */ + const batch = { + selection: { type: 'all', length: 3 }, + columns: [ + { type: 'values', values: [false, true, true], length: 3 }, + { read: payload }, + ], + } + /** @type {PrepareScan} */ + function prepareScan(request) { + return { + schema, + residual: { + filter: request.filter, + limit: request.limit, + offset: request.offset, + }, + properties: { + exactRows: 3, + maxRows: 3, + }, + async *batches() { + yield batch + }, + } + } + const prepare = vi.fn(prepareScan) + const source = preparedSource(prepare) + + const results = executeSql({ + tables: { data: source }, + query: 'SELECT payload FROM data WHERE keep LIMIT 1', + signal: controller.signal, + }) + + expect(prepare).toHaveBeenCalledTimes(1) + const [request] = prepare.mock.calls[0] + expect(request.columns).toEqual([ + { field: 20, phase: 1, purpose: 'output', mode: 'deferred' }, + { field: 10, phase: 0, purpose: 'filter', mode: 'required' }, + ]) + expect(payload).not.toHaveBeenCalled() + expect(await collect(results)).toEqual([{ payload: 'second' }]) + expect(payload).toHaveBeenCalledTimes(1) + }) + + it('bounds residual predicate reads to the limited range', async () => { + const length = 1_000 + /** @type {ReadColumn} */ + function readKeep({ selection }) { + return { type: 'constant', value: true, length: selectedRowCount(selection) } + } + const keep = vi.fn(readKeep) + /** @type {AsyncBatch} */ + const batch = { + selection: { type: 'all', length }, + columns: [ + { read: keep }, + { type: 'constant', value: 'match', length }, + ], + } + const source = preparedSource(function prepareScan(request) { + return { + schema, + residual: { + filter: request.filter, + limit: request.limit, + offset: request.offset, + }, + properties: { exactRows: length, maxRows: length }, + async *batches() { yield batch }, + } + }) + + await expect(collect(executeSql({ + tables: { data: source }, + query: 'SELECT payload FROM data WHERE keep LIMIT 1 OFFSET 300', + }))).resolves.toEqual([{ payload: 'match' }]) + expect(keep).toHaveBeenCalledTimes(1) + expect(keep.mock.calls[0][0].selection).toEqual({ + type: 'range', + start: 0, + end: 301, + length, + }) + }) + + it('keeps CASE residuals in native batches', async () => { + /** @type {AsyncBatch} */ + const batch = { + selection: { type: 'all', length: 2 }, + columns: [ + { type: 'values', values: [false, true], length: 2 }, + { type: 'values', values: ['first', 'second'], length: 2 }, + ], + } + const source = preparedSource(function prepareScan(request) { + return { + schema, + residual: { filter: request.filter }, + properties: { maxRows: 2 }, + async *batches() { yield batch }, + } + }) + const results = executeSql({ + tables: { data: source }, + query: 'SELECT payload FROM data WHERE CASE WHEN keep THEN TRUE ELSE FALSE END', + }) + + expect(results.batches).toBeTypeOf('function') + expect(await collect(results)).toEqual([{ payload: 'second' }]) + }) + + it('uses the row compatibility boundary for an unsupported residual', async () => { + /** @type {AsyncBatch} */ + const batch = { + selection: { type: 'all', length: 2 }, + columns: [ + { type: 'values', values: [false, true], length: 2 }, + { type: 'values', values: ['first', 'second'], length: 2 }, + ], + } + const source = preparedSource(function prepareScan(request) { + return { + schema, + residual: { filter: request.filter }, + properties: { maxRows: 2 }, + async *batches() { yield batch }, + } + }) + const results = executeSql({ + tables: { data: source }, + query: 'SELECT payload FROM data WHERE REGEXP_LIKE(payload, \'^second$\')', + }) + + expect(results.batches).toBeUndefined() + expect(await collect(results)).toEqual([{ payload: 'second' }]) + }) + + it('preserves short-circuiting before deferred column reads', async () => { + const payload = vi.fn(function readPayload() { + throw new Error('payload should not be read') + }) + /** @type {AsyncBatch} */ + const batch = { + selection: { type: 'all', length: 2 }, + columns: [ + { type: 'values', values: [true, false], length: 2 }, + { read: payload }, + ], + } + const source = preparedSource(function prepareScan(request) { + return { + schema, + residual: { filter: request.filter }, + properties: { maxRows: 2 }, + async *batches() { yield batch }, + } + }) + + await expect(collect(executeSql({ + tables: { data: source }, + query: 'SELECT keep FROM data WHERE FALSE AND LENGTH(payload) > 0', + }))).resolves.toEqual([]) + expect(payload).not.toHaveBeenCalled() + }) + + it('counts prepared zero-column batches without a legacy scan', async () => { + /** @type {RelationSchema} */ + const emptySchema = { fields: [] } + /** @type {PrepareScan} */ + function prepareCount(request) { + expect(request).toEqual({ columns: [] }) + return { + schema: emptySchema, + residual: {}, + properties: {}, + async *batches() { + yield { selection: { type: 'all', length: 2 }, columns: [] } + yield { selection: { type: 'all', length: 1 }, columns: [] } + }, + } + } + const prepareScan = vi.fn(prepareCount) + const source = { + columns: ['keep'], + schema: { fields: [schema.fields[0]] }, + prepareScan, + } + + await expect(collect(executeSql({ + tables: { data: source }, + query: 'SELECT COUNT(*) AS count FROM data', + }))).resolves.toEqual([{ count: 3 }]) + expect(prepareScan).toHaveBeenCalledTimes(1) + }) + + it('reports batch expression errors with stream-global row numbers', async () => { + /** @type {AsyncBatch[]} */ + const batches = [ + { + selection: { type: 'all', length: 2 }, + columns: [ + { type: 'values', values: [true, true], length: 2 }, + { type: 'values', values: ['a', 'bb'], length: 2 }, + ], + }, + { + selection: { type: 'all', length: 1 }, + columns: [ + { type: 'values', values: [true], length: 1 }, + { type: 'values', values: [3], length: 1 }, + ], + }, + ] + const source = preparedSource(function prepareScan() { + return { + schema, + residual: {}, + properties: { exactRows: 3 }, + async *batches() { yield* batches }, + } + }) + + await expect(collect(executeSql({ + tables: { data: source }, + query: 'SELECT LENGTH(payload) AS length FROM data', + }))).rejects.toThrow( + 'LENGTH(string): expected string or array, got number. Use CAST to convert to a string first. (row 3)' + ) + }) + + it('rejects an applied range while a residual filter remains', () => { + const source = preparedSource(function prepareScan(request) { + return { + schema, + residual: { filter: request.filter }, + properties: { maxRows: 1 }, + async *batches() {}, + } + }) + + expect(() => executeSql({ + tables: { data: source }, + query: 'SELECT payload FROM data WHERE keep LIMIT 1', + })).toThrow('Data source "data" applied limit/offset without applying where') + }) + + it('applies a source-applied range to fallback row metadata', async () => { + /** @type {RelationSchema} */ + const rangeSchema = { + fields: [ + { id: 1, name: 'payload', dataType: { type: 'string' }, nullable: false }, + ], + } + const source = { + numRows: 100, + schema: rangeSchema, + /** @type {PrepareScan} */ + prepareScan(request) { + expect(request.limit).toBe(1) + expect(request.offset).toBe(5) + return { + schema: rangeSchema, + residual: {}, + properties: {}, + async *batches() { + yield { + selection: { type: 'all', length: 1 }, + columns: [{ type: 'constant', value: 'sixth', length: 1 }], + } + }, + } + }, + } + + const results = executeSql({ + tables: { data: source }, + query: 'SELECT payload FROM data LIMIT 1 OFFSET 5', + }) + + expect(results.numRows).toBe(1) + expect(results.maxRows).toBe(1) + await expect(collect(results)).resolves.toEqual([{ payload: 'sixth' }]) + }) + + it('preserves logical schema order while assigning filter phases', async () => { + /** @type {RelationSchema} */ + const orderedSchema = { + fields: [ + { id: 1, name: 'first', dataType: { type: 'string' }, nullable: false }, + { id: 2, name: 'keep', dataType: { type: 'boolean' }, nullable: false }, + { id: 3, name: 'last', dataType: { type: 'string' }, nullable: false }, + ], + } + /** @type {PrepareScan} */ + function prepareScan(request) { + const fields = request.columns.map(function requestedField(demand) { + const field = orderedSchema.fields.find(function fieldById(candidate) { + return candidate.id === demand.field + }) + if (!field) throw new Error(`unknown field: ${demand.field}`) + return field + }) + /** @type {Record} */ + const values = { + first: ['a'], + keep: [true], + last: ['z'], + } + return { + schema: { fields }, + residual: { filter: request.filter }, + properties: { exactRows: 1 }, + async *batches() { + yield { + selection: { type: 'all', length: 1 }, + columns: fields.map(function loadedField(field) { + return { type: 'values', values: values[field.name], length: 1 } + }), + } + }, + } + } + const prepare = vi.fn(prepareScan) + const source = { + schema: orderedSchema, + prepareScan: prepare, + } + + const results = executeSql({ + tables: { data: source }, + query: 'SELECT * FROM data WHERE keep', + }) + + expect(results.columns).toEqual(['first', 'keep', 'last']) + expect(prepare.mock.calls[0][0].columns).toEqual([ + { field: 1, phase: 1, purpose: 'output', mode: 'deferred' }, + { field: 2, phase: 0, purpose: 'filter', mode: 'required' }, + { field: 3, phase: 1, purpose: 'output', mode: 'deferred' }, + ]) + await expect(collect(results)).resolves.toEqual([{ first: 'a', keep: true, last: 'z' }]) + }) + + it('requests a struct base field used by a residual filter', async () => { + /** @type {RelationSchema} */ + const structSchema = { + fields: [ + { + id: 1, + name: 'obj', + dataType: { + type: 'struct', + fields: [{ id: 2, name: 'x', dataType: { type: 'number' }, nullable: false }], + }, + nullable: false, + }, + { id: 3, name: 'payload', dataType: { type: 'string' }, nullable: false }, + ], + } + /** @type {PrepareScan} */ + function prepareScan(request) { + const fields = request.columns.map(function requestedField(demand) { + const field = structSchema.fields.find(function fieldById(candidate) { + return candidate.id === demand.field + }) + if (!field) throw new Error(`unknown field: ${demand.field}`) + return field + }) + return { + schema: { fields }, + residual: { filter: request.filter }, + properties: { exactRows: 2 }, + async *batches() { + yield { + selection: { type: 'all', length: 2 }, + columns: fields.map(function loadedField(field) { + const values = field.name === 'obj' ? [{ x: 1 }, { x: 2 }] : ['yes', 'no'] + return { type: 'values', values, length: 2 } + }), + } + }, + } + } + const prepare = vi.fn(prepareScan) + const source = { + columns: ['obj', 'payload'], + schema: structSchema, + prepareScan: prepare, + } + + await expect(collect(executeSql({ + tables: { data: source }, + query: 'SELECT payload FROM data WHERE obj.x = 1', + }))).resolves.toEqual([{ payload: 'yes' }]) + expect(prepare.mock.calls[0][0].columns).toEqual([ + { field: 3, phase: 1, purpose: 'output', mode: 'deferred' }, + { field: 1, phase: 0, purpose: 'filter', mode: 'required' }, + ]) + }) + + it('requests a struct base field used by a projection', async () => { + /** @type {RelationSchema} */ + const structSchema = { + fields: [ + { + id: 1, + name: 'obj', + dataType: { + type: 'struct', + fields: [{ id: 2, name: 'x', dataType: { type: 'number' }, nullable: false }], + }, + nullable: false, + }, + ], + } + const prepareScan = vi.fn(preparedValuesSource(structSchema, { + obj: [{ x: 1 }, { x: 2 }], + }).prepareScan) + const source = { schema: structSchema, prepareScan } + + await expect(collect(executeSql({ + tables: { data: source }, + query: 'SELECT obj.x FROM data', + }))).resolves.toEqual([{ x: 1 }, { x: 2 }]) + expect(prepareScan.mock.calls[0][0].columns).toEqual([ + { field: 1, phase: 1, purpose: 'output', mode: 'deferred' }, + ]) + }) + + it('requests an object base field used only by ordering', async () => { + /** @type {RelationSchema} */ + const objectSchema = { + fields: [ + { id: 1, name: 'obj', dataType: { type: 'unknown' }, nullable: false }, + { id: 2, name: 'payload', dataType: { type: 'string' }, nullable: false }, + ], + } + const prepareScan = vi.fn(preparedValuesSource(objectSchema, { + obj: [{ x: 2 }, { x: 1 }], + payload: ['second', 'first'], + }).prepareScan) + const source = { schema: objectSchema, prepareScan } + + await expect(collect(executeSql({ + tables: { data: source }, + query: 'SELECT payload FROM data ORDER BY obj.x', + }))).resolves.toEqual([{ payload: 'first' }, { payload: 'second' }]) + expect(prepareScan.mock.calls[0][0].columns).toEqual([ + { field: 2, phase: 1, purpose: 'output', mode: 'deferred' }, + { field: 1, phase: 1, purpose: 'output', mode: 'deferred' }, + ]) + }) + + it('requests an unknown object base field used by a residual filter', async () => { + /** @type {RelationSchema} */ + const objectSchema = { + fields: [ + { id: 1, name: 'obj', dataType: { type: 'unknown' }, nullable: false }, + { id: 2, name: 'payload', dataType: { type: 'string' }, nullable: false }, + ], + } + const prepareScan = vi.fn(preparedValuesSource(objectSchema, { + obj: [{ x: 1 }, { x: 2 }], + payload: ['yes', 'no'], + }).prepareScan) + const source = { schema: objectSchema, prepareScan } + + await expect(collect(executeSql({ + tables: { data: source }, + query: 'SELECT payload FROM data WHERE obj.x = 1', + }))).resolves.toEqual([{ payload: 'yes' }]) + expect(prepareScan.mock.calls[0][0].columns).toEqual([ + { field: 2, phase: 1, purpose: 'output', mode: 'deferred' }, + { field: 1, phase: 0, purpose: 'filter', mode: 'required' }, + ]) + }) + + it('does not request qualified columns owned by a predicate subquery', async () => { + /** @type {RelationSchema} */ + const outerSchema = { + fields: [ + { id: 1, name: 'id', dataType: { type: 'number' }, nullable: false }, + { id: 2, name: 'payload', dataType: { type: 'string' }, nullable: false }, + ], + } + /** @type {PrepareScan} */ + function prepareScan(request) { + const fields = request.columns.map(function requestedField(demand) { + const field = outerSchema.fields.find(function fieldById(candidate) { + return candidate.id === demand.field + }) + if (!field) throw new Error(`unknown field: ${demand.field}`) + return field + }) + return { + schema: { fields }, + residual: { filter: request.filter }, + properties: { exactRows: 2 }, + async *batches() { + yield { + selection: { type: 'all', length: 2 }, + columns: fields.map(function loadedField(field) { + const values = field.name === 'id' ? [1, 2] : ['yes', 'no'] + return { type: 'values', values, length: 2 } + }), + } + }, + } + } + const prepare = vi.fn(prepareScan) + const source = { + columns: ['id', 'payload'], + schema: outerSchema, + prepareScan: prepare, + } + + await expect(collect(executeSql({ + tables: { data: source, other: [{ foreign_id: 1 }] }, + query: 'SELECT payload FROM data WHERE id IN (SELECT o.foreign_id FROM other o)', + }))).resolves.toEqual([{ payload: 'yes' }]) + expect(prepare.mock.calls[0][0].columns).toEqual([ + { field: 2, phase: 1, purpose: 'output', mode: 'deferred' }, + { field: 1, phase: 0, purpose: 'filter', mode: 'required' }, + ]) + }) + + it('does not request a struct column owned by a predicate subquery', async () => { + /** @type {RelationSchema} */ + const parentSchema = { + fields: [ + { id: 1, name: 'obj', dataType: { type: 'unknown' }, nullable: false }, + { id: 2, name: 'payload', dataType: { type: 'string' }, nullable: false }, + ], + } + /** @type {RelationSchema} */ + const childSchema = { + fields: [ + { id: 3, name: 'obj', dataType: { type: 'unknown' }, nullable: false }, + ], + } + const parent = preparedValuesSource(parentSchema, { + obj: [{ x: 0 }], + payload: ['yes'], + }) + const prepareScan = vi.fn(parent.prepareScan) + + await expect(collect(executeSql({ + tables: { + parents: { schema: parentSchema, prepareScan }, + children: preparedValuesSource(childSchema, { obj: [{ x: 1 }] }), + }, + query: `SELECT payload FROM parents + WHERE EXISTS (SELECT * FROM children c WHERE obj.x = 1)`, + }))).resolves.toEqual([{ payload: 'yes' }]) + expect(prepareScan.mock.calls[0][0].columns).toEqual([ + { field: 2, phase: 1, purpose: 'output', mode: 'deferred' }, + ]) + }) + + it('resolves a filter prefix as a table alias before a struct column', () => { + /** @type {RelationSchema} */ + const aliasSchema = { + fields: [ + { + id: 1, + name: 'd', + dataType: { + type: 'struct', + fields: [{ id: 2, name: 'x', dataType: { type: 'number' }, nullable: false }], + }, + nullable: false, + }, + { id: 3, name: 'keep', dataType: { type: 'boolean' }, nullable: false }, + { id: 4, name: 'payload', dataType: { type: 'string' }, nullable: false }, + ], + } + /** @type {PrepareScan} */ + function prepareScan(request) { + const fields = request.columns.map(function requestedField(demand) { + const field = aliasSchema.fields.find(function fieldById(candidate) { + return candidate.id === demand.field + }) + if (!field) throw new Error(`unknown field: ${demand.field}`) + return field + }) + return { + schema: { fields }, + residual: {}, + properties: { exactRows: 0 }, + async *batches() {}, + } + } + const prepare = vi.fn(prepareScan) + const source = { + columns: ['d', 'keep', 'payload'], + schema: aliasSchema, + prepareScan: prepare, + } + + executeSql({ + tables: { data: source }, + query: 'SELECT payload FROM data d WHERE d.keep', + }) + + expect(prepare.mock.calls[0][0].columns).toEqual([ + { field: 4, phase: 1, purpose: 'output', mode: 'deferred' }, + { field: 3, phase: 0, purpose: 'filter', mode: 'required' }, + ]) + }) + + it('preserves prototype prepareScan methods through the cache wrapper', async () => { + class PreparedSource { + constructor() { + /** @type {RelationSchema} */ + this.schema = { + fields: [ + { id: 1, name: 'id', dataType: { type: 'number' }, nullable: false }, + ], + } + } + + /** + * @param {ScanRequest} request + * @returns {PreparedScan} + */ + prepareScan(request) { + expect(this).toBe(source) + expect(request.columns).toEqual([ + { field: 1, phase: 1, purpose: 'output', mode: 'deferred' }, + ]) + return { + schema: this.schema, + residual: {}, + properties: { exactRows: 1 }, + async *batches() { + yield { + selection: { type: 'all', length: 1 }, + columns: [{ type: 'constant', value: 1, length: 1 }], + } + }, + } + } + + /** @returns {never} */ + scan() { + throw new Error('legacy row scan should not be called') + } + } + const source = new PreparedSource() + + await expect(collect(executeSql({ + tables: { data: cachedDataSource(source) }, + query: 'SELECT id FROM data', + }))).resolves.toEqual([{ id: 1 }]) + }) + + it('surfaces cooperative aborts from direct batch iteration', async () => { + const { results } = abortingPreparedResults() + if (!results.batches) throw new Error('expected native batches') + async function consumeBatches() { + const batches = [] + for await (const batch of results.batches()) batches.push(batch) + return batches + } + + await expect(consumeBatches()).rejects.toThrow('prepared scan aborted') + }) + + it('surfaces cooperative aborts from direct row iteration', async () => { + const { results } = abortingPreparedResults() + async function consumeRows() { + const rows = [] + for await (const row of results.rows()) rows.push(row) + return rows + } + + await expect(consumeRows()).rejects.toThrow('prepared scan aborted') + }) + + it('surfaces cooperative aborts while collecting batches', async () => { + const { results } = abortingPreparedResults() + await expect(collect(results)).rejects.toThrow('prepared scan aborted') + }) +}) + +/** + * @returns {{ results: import('../../src/types.js').QueryResults }} + */ +function abortingPreparedResults() { + const controller = new AbortController() + const source = preparedSource(function prepareScan() { + return { + schema, + residual: {}, + properties: { exactRows: 1 }, + async *batches({ signal } = {}) { + yield { + selection: { type: 'all', length: 1 }, + columns: [ + { type: 'constant', value: true, length: 1 }, + { type: 'constant', value: 'value', length: 1 }, + ], + } + controller.abort(new Error('prepared scan aborted')) + if (!signal?.aborted) throw new Error('expected prepared signal to be aborted') + }, + } + }) + return { + results: executeSql({ + tables: { data: source }, + query: 'SELECT payload FROM data', + signal: controller.signal, + }), + } +} + +/** + * @param {PrepareScan} prepareScan + * @returns {AsyncDataSource} + */ +function preparedSource(prepareScan) { + return { + columns: ['keep', 'payload'], + schema, + prepareScan, + scan() { + throw new Error('legacy row scan should not be called') + }, + } +} + +/** + * @param {RelationSchema} sourceSchema + * @param {Record} values + * @returns {AsyncDataSource} + */ +function preparedValuesSource(sourceSchema, values) { + const length = Object.values(values)[0]?.length ?? 0 + return { + numRows: length, + schema: sourceSchema, + prepareScan(request) { + const fields = request.columns.map(function requestedField(demand) { + const field = sourceSchema.fields.find(function fieldById(candidate) { + return candidate.id === demand.field + }) + if (!field) throw new Error(`unknown field: ${demand.field}`) + return field + }) + return { + schema: { fields }, + residual: { filter: request.filter }, + properties: { exactRows: length }, + async *batches() { + yield { + selection: { type: 'all', length }, + columns: fields.map(function loadedField(field) { + return { type: 'values', values: values[field.name], length } + }), + } + }, + } + }, + } +} diff --git a/test/expression/batch.test.js b/test/expression/batch.test.js index 3963f02..996d620 100644 --- a/test/expression/batch.test.js +++ b/test/expression/batch.test.js @@ -3,8 +3,7 @@ import { compileBatchExpression } from '../../src/expression/batch.js' import { parseSql } from '../../src/parse/parse.js' /** - * @import { AsyncBatch, ReadColumn, RowSelection } from '../../src/internalTypes.js' - * @import { ExprNode } from '../../src/types.js' + * @import { AsyncBatch, ExprNode, ReadColumn, RowSelection } from '../../src/types.js' */ const schema = ['n', 'text'] @@ -53,11 +52,10 @@ describe('batch expressions', () => { if (!compiled) throw new Error('expected expression to compile') /** @type {AsyncBatch} */ const batch = { - columnNames: schema, selection: { type: 'all', length: 2 }, columns: [ { type: 'values', values: [1, 2], length: 2 }, - { type: 'source', read }, + { read }, ], } @@ -84,11 +82,10 @@ describe('batch expressions', () => { if (!compiled) throw new Error('expected expression to compile') /** @type {AsyncBatch} */ const batch = { - columnNames: schema, selection: { type: 'all', length: 3 }, columns: [ { type: 'values', values: [-1, 1, -2], length: 3 }, - { type: 'source', read }, + { read }, ], } @@ -115,11 +112,10 @@ describe('batch expressions', () => { if (!compiled) throw new Error('expected expression to compile') /** @type {AsyncBatch} */ const batch = { - columnNames: schema, selection: { type: 'all', length: 3 }, columns: [ { type: 'values', values: [1, null, null], length: 3 }, - { type: 'source', read }, + { read }, ], } @@ -150,11 +146,10 @@ describe('batch expressions', () => { if (!compiled) throw new Error('expected expression to compile') /** @type {AsyncBatch} */ const batch = { - columnNames: schema, selection: { type: 'all', length: 3 }, columns: [ { type: 'values', values: [1, 0, -1], length: 3 }, - { type: 'source', read }, + { read }, ], } @@ -181,11 +176,10 @@ describe('batch expressions', () => { if (!compiled) throw new Error('expected expression to compile') /** @type {AsyncBatch} */ const batch = { - columnNames: schema, selection: { type: 'indices', indices: new Uint32Array([1, 3]), length: 4 }, columns: [ { type: 'values', values: [0, -1, 0, 1], length: 4 }, - { type: 'source', read }, + { read }, ], } @@ -249,7 +243,6 @@ describe('batch expressions', () => { if (!compiled) throw new Error('expected expression to compile') /** @type {AsyncBatch} */ const batch = { - columnNames: structSchema, selection: { type: 'all', length: 2 }, columns: [ { type: 'values', values: [{ x: 2 }, 7], length: 2 }, @@ -329,7 +322,6 @@ function expression(sql) { */ function loadedBatch(numbers, texts) { return { - columnNames: schema, selection: { type: 'all', length: numbers.length }, columns: [ { type: 'values', values: numbers, length: numbers.length }, diff --git a/test/internalBatchTypes.d.ts b/test/internalBatchTypes.d.ts deleted file mode 100644 index 02217d1..0000000 --- a/test/internalBatchTypes.d.ts +++ /dev/null @@ -1,15 +0,0 @@ -import type { AsyncDataSource, QueryResults, ScanOptions, ScanResults } from '../src/index.js' - -type Assert = T -type Not = T extends true ? false : true -type HasKey = K extends keyof T ? true : false - -type LegacySource = { - columns: string[] - scan(options: ScanOptions): ScanResults -} - -type QueryResultsHideBatches = Assert>> -type QueryResultsHideSchema = Assert>> -type DataSourcesKeepLegacyContract = Assert -type DataSourcesHidePreparedScans = Assert>>