diff --git a/hypaware-core/plugins-workspace/ai-gateway/src/dataset.js b/hypaware-core/plugins-workspace/ai-gateway/src/dataset.js index d0100ba..471cc85 100644 --- a/hypaware-core/plugins-workspace/ai-gateway/src/dataset.js +++ b/hypaware-core/plugins-workspace/ai-gateway/src/dataset.js @@ -168,13 +168,35 @@ const SCHEMA_COLUMN_NAMES = AI_GATEWAY_SCHEMA_COLUMNS.map((c) => c.name) */ function withSchemaColumns(source) { const columns = Array.from(new Set([...source.columns, ...SCHEMA_COLUMN_NAMES])) - return { + /** @type {AsyncDataSource} */ + const wrapped = { columns, numRows: source.numRows, scan(options) { return source.scan(options) }, } + // Forward the column-stream hook so single-column aggregates stay on the + // engine's streaming fast path. A partition that physically lacks the + // requested column (the additive schema-drift case this wrapper exists + // for) surfaces its values as `undefined` holes in the chunk; normalize + // them to null, the same "this partition predates the column" value the + // row path reads, so accumulators see one representation either way. + // @ref LLP 0055 [implements]: withSchemaColumns forwards scanColumn; a partition lacking the column yields nulls, never throws + if (typeof source.scanColumn === 'function') { + const scanColumn = /** @type {NonNullable} */ (source.scanColumn) + wrapped.scanColumn = (options) => ({ + async *[Symbol.asyncIterator]() { + for await (const chunk of scanColumn(options)) { + for (let i = 0; i < chunk.length; i++) { + if (chunk[i] === undefined) /** @type {unknown[]} */ (chunk)[i] = null + } + yield chunk + } + }, + }) + } + return wrapped } /** diff --git a/llp/0054-bounded-query-execution.spec.md b/llp/0054-bounded-query-execution.spec.md index bb94220..d3dab91 100644 --- a/llp/0054-bounded-query-execution.spec.md +++ b/llp/0054-bounded-query-execution.spec.md @@ -161,6 +161,10 @@ adds options to `ExecuteSqlOptions`, not a new import surface or a shim. `ai_gateway_messages` volume on a representative box — deferred to [LLP 0057](./0057-bounded-query-execution.plan.md). Until measured, ship a safe-low default and let operators raise it. + *Measured 2026-07-10; resolved by [LLP 0097](./0097-heap-growth-query-budget.decision.md), + which enforces the budget as sampled heap growth (default 1GiB) from the + kernel while per-operator buffered-row/byte accounting remains the engine + follow-up.* - **Spill-to-disk.** A future external-merge / spilling path would let large `ORDER BY` / `GROUP BY` **complete** instead of refusing; named as a deferred follow-up in [LLP 0056](./0056-refuse-over-spill-or-truncate.decision.md), not diff --git a/llp/0057-bounded-query-execution.plan.md b/llp/0057-bounded-query-execution.plan.md index b765899..76d8f68 100644 --- a/llp/0057-bounded-query-execution.plan.md +++ b/llp/0057-bounded-query-execution.plan.md @@ -6,6 +6,7 @@ **Author:** Phil / Claude **Date:** 2026-06-30 **Related:** LLP 0054, LLP 0055, LLP 0056 +**Extended-by:** LLP 0097 (Phase 0 measurements; Phases 1-3 realized kernel-side against squirreling 0.14 / icebird 0.8.13, where the engine already streams aggregates and implements `scanColumn` at the leaf) > Turns the bounded-execution spec ([LLP 0054](./0054-bounded-query-execution.spec.md)) > and its two decisions ([LLP 0055](./0055-stream-aggregates-via-scancolumn.decision.md), diff --git a/llp/0097-heap-growth-query-budget.decision.md b/llp/0097-heap-growth-query-budget.decision.md new file mode 100644 index 0000000..eaef4f0 --- /dev/null +++ b/llp/0097-heap-growth-query-budget.decision.md @@ -0,0 +1,124 @@ +# LLP 0097: Heap-Growth Guard Enforces the Execution Budget from the Kernel + +**Type:** Decision +**Status:** Active +**Systems:** Query, Cache +**Author:** Phil / Claude +**Date:** 2026-07-10 +**Related:** LLP 0054, LLP 0055, LLP 0056, LLP 0057 + +> How the kernel bounds query execution memory TODAY, with the pinned engine: +> a sampled process-heap-growth guard checked inline on the scan path, refusing +> per [LLP 0056](./0056-refuse-over-spill-or-truncate.decision.md). Realizes +> [LLP 0054](./0054-bounded-query-execution.spec.md) `#memory-invariant` and +> sizes the default ceiling `#execution-budget` deferred to measurement. + +## Context + +[LLP 0057](./0057-bounded-query-execution.plan.md) Phase 0 (measure) ran on +2026-07-10 against the production cache (202k-row / 931MB-on-disk +`ai_gateway_messages`, 5.4k-node / 12k-edge context graph), one fresh CLI +process per query, median of 3, peak RSS via `/usr/bin/time -l`: + +| query class | wall | peak RSS | +|---|---|---| +| process floor (`LIMIT 1`) | 173ms | 172MB | +| `COUNT(*)` | 171ms | 179MB | +| `COUNT(DISTINCT session_id)` | 759ms | 510MB | +| `GROUP BY provider` / high-card `GROUP BY session_id` | ~780ms | 470-534MB | +| top-K `ORDER BY ... LIMIT 20` | 888ms | 444MB | +| full sort, narrow projection, no LIMIT | 987ms | 743MB | +| `COUNT(DISTINCT content_text)` | 1085ms | 685MB | +| `SELECT * ORDER BY` (no LIMIT, issue #9 class) | 18.3s | **6.1GB, died** | +| `hyp graph neighbors` depth 1-3 | ~170ms | 130-150MB | + +Two facts changed the implementation picture since the 0054/0055 docs were +authored against squirreling 0.12.24: + +- **The engine already streams.** squirreling 0.14.0 (pinned) streams scalar + and `GROUP BY` aggregates through accumulators, sorts top-K when a `LIMIT` + reaches the sort, threads an abort `signal`, and has the `scanColumn` + column-stream fast path; icebird 0.8.13 (pinned) implements `scanColumn` at + the leaf. The dormant pieces were all kernel wrappers, now lit + ([LLP 0055](./0055-stream-aggregates-via-scancolumn.decision.md) `@ref`s in + `src/core/cache/storage.js`, `src/core/query/union-source.js`, ai-gateway + `dataset.js`). +- **The remaining crasher is retained buffering** in blocking operators with + no bound: the engine-side buffered-row/byte accounting that + [LLP 0054](./0054-bounded-query-execution.spec.md) `#execution-budget` + specifies is an upstream squirreling change that has not landed. + +Waiting for the engine accounting would leave the daemon OOM-killable in the +meantime. The kernel needed an enforcement mechanism that works with the +pinned engine, entirely from the `hypaware/core/query` surface. + +## Options considered + +1. **Sampled process-heap-growth guard in the kernel.** (Chosen.) Sample + `process.memoryUsage().heapUsed` growth since query start; refuse when it + exceeds the budget. +2. **Wait for engine-side buffered-row/byte accounting** (LLP 0054 + `#execution-budget` as specified). Rejected as the only line of defense: + upstream latency leaves the crasher class live; the engine accounting + remains the intended refinement and composes with this guard when it lands. +3. **Timer-only watchdog** (setInterval + abort signal). Rejected as + insufficient alone, from evidence: a query whose reads resolve without real + I/O holds the event loop for its entire run, so timer callbacks never fire. + The measured issue-#9 crasher ran 8+ seconds to 4.8GB with a 50MB budget + and **zero** watchdog samples. + +## Decision + +`executeQuerySql` enforces a **per-query heap-growth budget** with two +coordinated layers: + +- **Inline guard (primary):** every table source is decorated so its row scans + check sampled heap growth every 4096 rows, and its column streams check per + chunk, from *inside* the loop that a blocking operator drives. Starvation- + proof by construction. +- **Interval watchdog (secondary):** a 100ms `setInterval` covers execution + phases that pull no further source rows (join amplification, output + finalization) but do yield to the event loop. + +Either tripping aborts the run through the threaded signal +([LLP 0054](./0054-bounded-query-execution.spec.md) `#signal-threading`) and +surfaces a typed `QueryExecutionBudgetError` (exported from +`hypaware/core/query`) carrying the limit and observed growth: a refusal, not +a truncation ([LLP 0056](./0056-refuse-over-spill-or-truncate.decision.md)). + +**Growth, not absolute:** the budget bounds heap growth attributable to the +query (sampled minus at-start baseline), so a long-lived daemon's resident +baseline neither eats the budget nor causes blanket refusals. + +**Default ceiling: 1GiB growth**, from the Phase 0 measurements: every +well-formed query in the measured set stays under ~500MB of growth (2x +headroom), while the crasher class blows past 4GB. Operators override with +`HYP_QUERY_MAX_HEAP_MB` (or the `maxHeapBytes` option on +`ExecuteSqlOptions`; `0` disables). With the default in place the measured +crasher refuses in 0.7-1.2s at ~1.4GB peak RSS instead of dying at 6GB. + +## Consequences + +- Every caller of `hypaware/core/query` (CLI, MCP `query_sql`, HypAware + Server `POST /v1/query`) inherits the bound with no per-surface work + ([LLP 0054](./0054-bounded-query-execution.spec.md) `#uniform-surface`). + The server can pass its own `maxHeapBytes` and map the typed error to a + 4xx (HypAware Server LLP 0020 owns that wiring). +- Heap growth is process-global. Concurrent queries in one process share the + observable, so a query can be refused partly on a neighbor's allocations; + conservative and safe in the direction we care about (protect the process). + Per-operator buffered-byte accounting (the LLP 0054 `#execution-budget` + letter, upstream in squirreling) remains the precise refinement; when it + lands, this guard stays as defense-in-depth. +- `heapUsed` includes not-yet-collected garbage, so a pathologically + garbage-heavy but well-bounded query could trip early; measured headroom + (2x over the worst legitimate query) and the scavenge-on-allocation + behavior of young-generation garbage make this unlikely, and the refusal + message names the override. +- Post-change measurements (same harness, same cache): every measured query + got faster (up to -26% wall) and none regressed; the speed budget for this + work ("within 10%") was met with margin. + +The code site (`src/core/query/sql.js`) carries `@ref`s to this decision and +to [LLP 0054](./0054-bounded-query-execution.spec.md) `#signal-threading` / +[LLP 0056](./0056-refuse-over-spill-or-truncate.decision.md). diff --git a/src/core/cache/storage.js b/src/core/cache/storage.js index cfd27c8..7601d46 100644 --- a/src/core/cache/storage.js +++ b/src/core/cache/storage.js @@ -29,7 +29,7 @@ import path from 'node:path' * @import { ColumnSpec, QueryScope, QueryStorageService, SinkContinuation } from '../../../hypaware-plugin-kernel-types.js' * @import { CachePartitioningDeclaration, ExtendedQueryStorageService } from '../../../src/core/cache/types.js' * @import { UsagePolicyResolver } from '../../../src/core/usage-policy/types.js' - * @import { AsyncCells } from 'squirreling' + * @import { AsyncDataSource } from 'squirreling' */ /** @@ -303,34 +303,40 @@ export function createQueryStorageService({ cacheRoot, getDeclaration, getSettle async dataSourceForTable(tablePath) { const source = await dataSourceForTable(resolveIcebergDir(tablePath)) if (!source) return null - return { + const publicColumns = source.columns.filter((c) => !INTERNAL_FIELDS.includes(c)) + /** @type {AsyncDataSource} */ + const wrapped = { numRows: source.numRows, - columns: source.columns.filter((c) => !INTERNAL_FIELDS.includes(c)), + columns: publicColumns, scan(options) { - const inner = source.scan({ - ...options, - columns: options.columns?.filter((c) => !INTERNAL_FIELDS.includes(c)), - }) + // Internal fields are hidden by PROJECTION, not by rebuilding every + // row: the inner scan is always given an explicit column list with + // the internal fields already stripped (the advertised public set + // when the caller asked for everything), so the rows it yields never + // carry an internal column and can be passed through untouched. The + // previous per-row rebuild (filter columns, re-key cells, re-spread + // resolved) allocated ~5 objects per row on every query, which + // dominated scan-side garbage on large datasets. + const requested = options?.columns + const columns = requested + ? requested.filter((c) => !INTERNAL_FIELDS.includes(c)) + : publicColumns + const inner = source.scan({ ...options, columns }) return { appliedWhere: inner.appliedWhere, appliedLimitOffset: inner.appliedLimitOffset, - async *rows() { - for await (const row of inner.rows()) { - const filteredColumns = row.columns.filter((c) => !INTERNAL_FIELDS.includes(c)) - const filteredResolved = row.resolved - ? Object.fromEntries(Object.entries(row.resolved).filter(([k]) => !INTERNAL_FIELDS.includes(k))) - : undefined - /** @type {AsyncCells} */ - const filteredCells = {} - for (const col of filteredColumns) { - if (row.cells && col in row.cells) filteredCells[col] = row.cells[col] - } - yield { ...row, columns: filteredColumns, cells: filteredCells, resolved: filteredResolved } - } - }, + rows: () => inner.rows(), } }, } + // @ref LLP 0055 [implements]: forward the column-stream hook so the + // engine's streaming-aggregate fast path stays lit through the storage + // wrapper; internal fields are not advertised, so the engine can never + // request one here. + if (typeof source.scanColumn === 'function') { + wrapped.scanColumn = (options) => /** @type {NonNullable} */ (source.scanColumn)(options) + } + return wrapped }, async flushTable(tablePath, opts = {}) { diff --git a/src/core/query/index.js b/src/core/query/index.js index 09eb098..731d164 100644 --- a/src/core/query/index.js +++ b/src/core/query/index.js @@ -4,7 +4,7 @@ // Reading parquet/Iceberg back from a BlobStore-backed query source is // built on top of these helpers. -export { executeQuerySql } from './sql.js' +export { executeQuerySql, QueryExecutionBudgetError } from './sql.js' export { parquetDataSource } from './parquet-source.js' export { whereToParquetFilter } from './parquet-pushdown.js' export { unionSources, emptySource } from './union-source.js' diff --git a/src/core/query/sql.js b/src/core/query/sql.js index bc68a6c..db8a76c 100644 --- a/src/core/query/sql.js +++ b/src/core/query/sql.js @@ -12,6 +12,124 @@ import { QUERY_FLUSH_DEBOUNCE_MS } from '../cache/spool.js' * @import { AsyncDataSource } from 'squirreling' */ +/** + * Default per-query heap-growth budget. Sized from the LLP 0057 Phase 0 + * measurement pass (2026-07-10, ~202k-row / 931MB ai_gateway_messages): + * every well-formed query in the measured set peaks under ~500MB of + * process heap growth, while the issue-#9 crasher class (unbounded wide + * ORDER BY) grows past 4GB before dying. 1GiB refuses the crasher class + * with roomy headroom for legitimate queries. + */ +const DEFAULT_MAX_HEAP_GROWTH_BYTES = 1024 * 1024 * 1024 + +/** How often the watchdog samples heap growth while a query runs. */ +const HEAP_WATCH_INTERVAL_MS = 100 + +/** + * Typed refusal for a query whose execution outgrew its heap budget. + * Refusal, not truncation: a partial sort or partial aggregate would be a + * silently wrong answer. Callers render it as actionable guidance and map + * it to a 4xx (server) or non-zero exit (CLI). + * + * @ref LLP 0056 [implements]: over-budget queries refuse with a distinct typed error carrying the limit that was hit + */ +export class QueryExecutionBudgetError extends Error { + /** + * @param {number} limitBytes + * @param {number} observedBytes + */ + constructor(limitBytes, observedBytes) { + const limitMb = Math.round(limitBytes / 1048576) + const observedMb = Math.round(observedBytes / 1048576) + super( + `query exceeded its execution memory budget (${observedMb}MB used of ${limitMb}MB) - ` + + 'add a WHERE/date filter, a LIMIT, or aggregate instead of selecting raw rows ' + + '(raise the budget with HYP_QUERY_MAX_HEAP_MB if this query truly needs more)' + ) + this.name = 'QueryExecutionBudgetError' + this.code = 'query_budget_exceeded' + this.limitBytes = limitBytes + this.observedBytes = observedBytes + } +} + +/** + * Resolve the effective heap-growth budget: explicit option, then the + * HYP_QUERY_MAX_HEAP_MB operator override, then the measured default. + * 0 (or a non-positive override) disables the watchdog entirely. + * + * @param {number | undefined} optionBytes + * @returns {number} + */ +export function resolveHeapBudgetBytes(optionBytes) { + if (optionBytes !== undefined) return optionBytes + // A set-but-blank var (`export HYP_QUERY_MAX_HEAP_MB=`, how many config + // systems render an unset optional) must NOT silently disable the guard: + // Number('') and Number(' ') are 0 (finite), which reads as "disabled". + // Only a non-empty value counts as an override; anything else (unset or + // blank) falls through to the measured default. + const raw = process.env.HYP_QUERY_MAX_HEAP_MB?.trim() + if (raw) { + const env = Number(raw) + if (Number.isFinite(env)) return env * 1024 * 1024 + } + return DEFAULT_MAX_HEAP_GROWTH_BYTES +} + +/** Rows between inline heap checks on a row scan. */ +const BUDGET_CHECK_ROW_STRIDE = 4096 + +/** + * Decorate a data source so its scans enforce the query's heap budget + * INLINE, from within the row loop itself. A timer-based watchdog alone is + * not enough: a query whose reads resolve without real I/O (warm cache, + * synchronous resolvers) can hold the event loop for its entire run, so a + * setInterval callback never fires while a blocking operator's buffer + * grows. The stride keeps the memoryUsage() sample cost far below one + * sample per row. All sources of one query share `guard`, so growth is + * judged per query, not per table. + * + * @param {AsyncDataSource} source + * @param {{ check: () => void }} guard + * @returns {AsyncDataSource} + */ +function withHeapBudget(source, guard) { + /** @type {AsyncDataSource} */ + const bounded = { + numRows: source.numRows, + columns: source.columns, + scan(options) { + const inner = source.scan(options) + return { + appliedWhere: inner.appliedWhere, + appliedLimitOffset: inner.appliedLimitOffset, + async *rows() { + let sinceCheck = 0 + for await (const row of inner.rows()) { + if (++sinceCheck >= BUDGET_CHECK_ROW_STRIDE) { + sinceCheck = 0 + guard.check() + } + yield row + } + }, + } + }, + } + if (typeof source.scanColumn === 'function') { + const scanColumn = /** @type {NonNullable} */ (source.scanColumn) + bounded.scanColumn = (options) => ({ + async *[Symbol.asyncIterator]() { + for await (const chunk of scanColumn(options)) { + guard.check() + yield chunk + } + }, + }) + } + return bounded +} + /** * Run a read-only SELECT against the kernel's dataset registry. The * caller (the `hyp query sql` command or a future server endpoint) @@ -112,18 +230,97 @@ export async function executeQuerySql(args) { tables[name] = source } - const results = squirrelExecuteSql({ tables, query: trimmed }) - const rows = await collect(results) - const columns = results.columns ?? [] - span.setAttribute('row_count', rows.length) + // Execution is bounded by a heap-growth watchdog: the engine and + // every data source already honor an abort signal on their hot + // loops, so tripping the budget aborts the run mid-stream instead + // of letting a blocking operator (unbounded ORDER BY / GROUP BY / + // DISTINCT buffering) grow until the process is OOM-killed. The + // sampled process-heap growth is a stand-in for the per-operator + // buffered-byte accounting that belongs upstream in the engine. + // @ref LLP 0054#signal-threading [implements]: the kernel constructs the signal and forwards it into squirrelExecuteSql, activating the operators' abort checks + // @ref LLP 0097 [implements]: heap-growth watchdog enforces the execution budget from the kernel while buffered-byte accounting stays an engine follow-up + const budgetBytes = resolveHeapBudgetBytes(args.maxHeapBytes) + const controller = new AbortController() + // Detach the linked-signal listener when the query settles: a + // long-lived upstream signal (one shared across many queries) would + // otherwise retain this controller closure per call. + /** @type {(() => void) | undefined} */ + let removeUpstreamAbort + if (args.signal) { + const upstream = args.signal + if (upstream.aborted) controller.abort(upstream.reason) + else { + const onUpstreamAbort = () => controller.abort(upstream.reason) + upstream.addEventListener('abort', onUpstreamAbort, { once: true }) + removeUpstreamAbort = () => upstream.removeEventListener('abort', onUpstreamAbort) + } + } + const baselineHeap = process.memoryUsage().heapUsed + /** @type {QueryExecutionBudgetError | undefined} */ + let budgetError + /** @type {NodeJS.Timeout | undefined} */ + let watchdog + const trip = (/** @type {number} */ growth) => { + if (!budgetError) { + budgetError = new QueryExecutionBudgetError(budgetBytes, growth) + controller.abort(budgetError) + } + if (watchdog) clearInterval(watchdog) + return budgetError + } + const guard = { + check() { + if (budgetBytes <= 0) return + if (budgetError) throw budgetError + const growth = process.memoryUsage().heapUsed - baselineHeap + if (growth > budgetBytes) throw trip(growth) + }, + } + if (budgetBytes > 0) { + // Second enforcement layer for execution phases that pull no + // further source rows (join amplification, output finalization) + // but do yield to the event loop. + watchdog = setInterval(() => { + const growth = process.memoryUsage().heapUsed - baselineHeap + if (growth > budgetBytes) trip(growth) + }, HEAP_WATCH_INTERVAL_MS) + watchdog.unref() + for (const name of Object.keys(tables)) { + tables[name] = withHeapBudget(tables[name], guard) + } + } - instruments.queryRunsTotal.add(1, { status: 'ok' }) - instruments.queryDurationMs.record(Date.now() - start, { status: 'ok' }) + try { + const results = squirrelExecuteSql({ tables, query: trimmed, signal: controller.signal }) + const rows = await collect(results) + // Terminal budget check: the inline guard only samples every + // BUDGET_CHECK_ROW_STRIDE rows (and per column chunk), and the + // interval watchdog cannot fire during a fully synchronous run, so + // growth concentrated in a sub-stride tail or in finalization would + // otherwise return a wrongly-successful result. One check after + // materialization closes that window before we record success. + guard.check() + const columns = results.columns ?? [] + span.setAttribute('row_count', rows.length) - return { columns, rows, datasets: datasetsUsed, freshnessMessages } + instruments.queryRunsTotal.add(1, { status: 'ok' }) + instruments.queryDurationMs.record(Date.now() - start, { status: 'ok' }) + + return { columns, rows, datasets: datasetsUsed, freshnessMessages } + } catch (err) { + // Any abort surfaced while the budget stands tripped maps to the + // typed refusal, whichever layer's abort check fired first. + if (budgetError) throw budgetError + throw err + } finally { + if (watchdog) clearInterval(watchdog) + if (removeUpstreamAbort) removeUpstreamAbort() + } } catch (err) { + const budgeted = err instanceof QueryExecutionBudgetError span.setAttribute('status', 'failed') - instruments.queryRunsTotal.add(1, { status: 'failed' }) + if (budgeted) span.setAttribute('error_kind', 'budget_exceeded') + instruments.queryRunsTotal.add(1, { status: 'failed', ...(budgeted ? { error_kind: 'budget_exceeded' } : {}) }) instruments.queryDurationMs.record(Date.now() - start, { status: 'failed' }) throw err } diff --git a/src/core/query/types.d.ts b/src/core/query/types.d.ts index 5a38bfe..08bbc4a 100644 --- a/src/core/query/types.d.ts +++ b/src/core/query/types.d.ts @@ -36,6 +36,17 @@ export interface ExecuteSqlOptions { scope?: QueryScope refresh?: RefreshMode log?: PluginLogger + /** Caller-supplied cancellation; linked into the signal the engine and data sources observe. */ + signal?: AbortSignal + /** + * Execution memory budget: the query is refused (typed + * QueryExecutionBudgetError) once its sampled process-heap growth exceeds + * this many bytes. 0 disables the bound. Distinct from ContextControls, + * which cap display/output bytes after materialization; this bounds the + * execution itself. Defaults to the kernel ceiling (overridable with + * HYP_QUERY_MAX_HEAP_MB). + */ + maxHeapBytes?: number } export interface ExecuteSqlResult { diff --git a/src/core/query/union-source.js b/src/core/query/union-source.js index ce80f20..bc4a9e1 100644 --- a/src/core/query/union-source.js +++ b/src/core/query/union-source.js @@ -38,7 +38,8 @@ export function unionSources(sources) { for (const col of s.columns) allColumns.add(col) totalRows += s.numRows ?? 0 } - return { + /** @type {AsyncDataSource} */ + const union = { columns: Array.from(allColumns), numRows: totalRows, scan(options) { @@ -67,6 +68,68 @@ export function unionSources(sources) { } }, } + // The column-stream hook is offered only when EVERY partition can stream + // the column; a mixed union stays row-based so the engine's fallback owns + // correctness. Unlike scan(), the scanColumn contract has no + // appliedLimitOffset escape hatch: the source must fully honor + // limit/offset itself, so the union applies them over the CONCATENATED + // stream (they are not distributive across partitions, the same + // discipline scan() applies). Only the remaining-limit bound is pushed + // per partition, as an upper-bound optimization that can never change + // the result. + // @ref LLP 0055 [implements]: union forwards scanColumn by concatenating per-partition column streams; limit/offset apply to the merged stream + if (sources.every((s) => typeof s.scanColumn === 'function')) { + union.scanColumn = ({ column, limit, offset, signal }) => ({ + async *[Symbol.asyncIterator]() { + let remainingSkip = offset ?? 0 + let remaining = limit ?? Infinity + for (const source of sources) { + if (remaining <= 0) return + signal?.throwIfAborted() + // A known-empty or fully-skippable partition needs no stream. + const numRows = source.numRows + if (numRows !== undefined && numRows <= remainingSkip) { + remainingSkip -= numRows + continue + } + const scanColumn = /** @type {NonNullable} */ (source.scanColumn) + const chunks = scanColumn({ + column, + // Per-partition upper bound: this partition can contribute at + // most the values still owed, including any skip not yet spent. + limit: remaining === Infinity ? undefined : remainingSkip + remaining, + signal, + }) + for await (const chunk of chunks) { + signal?.throwIfAborted() + let start = 0 + if (remainingSkip > 0) { + if (remainingSkip >= chunk.length) { + remainingSkip -= chunk.length + continue + } + start = remainingSkip + remainingSkip = 0 + } + const end = remaining === Infinity + ? chunk.length + : Math.min(chunk.length, start + remaining) + if (start === 0 && end === chunk.length) { + yield chunk + remaining -= chunk.length + } else if (end > start) { + const slice = [] + for (let i = start; i < end; i++) slice.push(chunk[i]) + yield slice + remaining -= slice.length + } + if (remaining <= 0) break + } + } + }, + }) + } + return union } /** diff --git a/test/core/ai-gateway-dataset.test.js b/test/core/ai-gateway-dataset.test.js index 6c0af5c..f12b26a 100644 --- a/test/core/ai-gateway-dataset.test.js +++ b/test/core/ai-gateway-dataset.test.js @@ -256,3 +256,44 @@ test('ai-gateway createDataSource pads declared schema columns absent from an ol await fs.rm(cacheRoot, { recursive: true, force: true }) } }) + +test('ai-gateway createDataSource streams scanColumn with nulls for a physically absent column', async () => { + // The column-stream analog of the schema-padding row test above: the + // engine's streaming-aggregate fast path consumes scanColumn, and a + // partition that predates a declared column must contribute nulls (never + // undefined, never a throw) so accumulators see the same value the row + // path reads. @ref LLP 0055 + const cacheRoot = await makeTmpDir('scan-column') + try { + await appendRowsToSourceTable( + cacheRoot, DATASET_NAME, ['source=claude'], + TEST_COLUMNS, [{ id: 1, date: '2026-05-26' }, { id: 2, date: '2026-05-27' }] + ) + + const storage = createQueryStorageService({ cacheRoot }) + /** @type {QueryScope} */ + const scope = { limit: 1000 } + const partitions = await discoverParts({ cacheDir: cacheRoot, scope, config: { version: 2 } }) + const source = await createDataSource(partitions, { scope, storage }) + + assert.equal(typeof source.scanColumn, 'function', 'the storage-backed source streams columns') + const scanColumn = /** @type {NonNullable} */ (source.scanColumn) + + /** @type {unknown[]} */ + const ids = [] + for await (const chunk of scanColumn({ column: 'id' })) { + for (let i = 0; i < chunk.length; i++) ids.push(chunk[i]) + } + assert.deepEqual([...ids].sort(), [1, 2], 'a physical column streams its values') + + /** @type {unknown[]} */ + const absent = [] + for await (const chunk of scanColumn({ column: 'git_remote' })) { + for (let i = 0; i < chunk.length; i++) absent.push(chunk[i]) + } + assert.equal(absent.length, 2) + for (const v of absent) assert.strictEqual(v, null, 'absent column streams null, not undefined') + } finally { + await fs.rm(cacheRoot, { recursive: true, force: true }) + } +}) diff --git a/test/core/query-sql-budget.test.js b/test/core/query-sql-budget.test.js new file mode 100644 index 0000000..110d76d --- /dev/null +++ b/test/core/query-sql-budget.test.js @@ -0,0 +1,157 @@ +// @ts-check + +import test from 'node:test' +import assert from 'node:assert/strict' + +import { asyncRow } from 'squirreling' +import { executeQuerySql, QueryExecutionBudgetError, resolveHeapBudgetBytes } from '../../src/core/query/sql.js' + +/** + * @import { AsyncDataSource, SqlPrimitive } from 'squirreling/src/types.js' + */ + +/** + * A memory-backed AsyncDataSource over `rows`, with an optional recording + * scanColumn hook so tests can prove the streaming-aggregate fast path + * stays reachable through the kernel's budget decoration. + * + * @param {Record[]} rows + * @param {{ scanColumnCalls?: string[] }} [opts] + * @returns {AsyncDataSource} + */ +function memorySource(rows, opts = {}) { + const columns = Object.keys(rows[0] ?? {}) + /** @type {AsyncDataSource} */ + const source = { + columns, + numRows: rows.length, + scan(options) { + const rowColumns = options?.columns ?? columns + return { + appliedWhere: false, + appliedLimitOffset: false, + async *rows() { + for (const row of rows) yield asyncRow(row, rowColumns) + }, + } + }, + } + if (opts.scanColumnCalls) { + const calls = opts.scanColumnCalls + source.scanColumn = ({ column, limit, offset }) => ({ + async *[Symbol.asyncIterator]() { + calls.push(column) + const start = offset ?? 0 + const end = limit === undefined ? rows.length : Math.min(rows.length, start + limit) + if (end > start) yield rows.slice(start, end).map((r) => r[column] ?? null) + }, + }) + } + return source +} + +/** @param {AsyncDataSource} source */ +function registryFor(source) { + const dataset = { + discoverPartitions: async () => [], + createDataSource: async () => source, + } + return /** @type {any} */ ({ getDataset: () => dataset, listDatasets: () => [] }) +} + +const storage = /** @type {any} */ ({ + cacheRoot: '/tmp/hypaware-test', + pendingInfo: async () => ({ pending: false }), +}) + +test('a query whose heap growth exceeds the execution budget refuses with the typed error', async () => { + // More rows than the inline check stride, so the guard samples mid-scan. + const rows = Array.from({ length: 6000 }, (_, i) => ({ a: `value-${i}` })) + await assert.rejects( + executeQuerySql({ + query: 'SELECT a FROM t ORDER BY a', + registry: registryFor(memorySource(rows)), + storage, + maxHeapBytes: 1, + }), + (err) => { + assert.ok(err instanceof QueryExecutionBudgetError, 'typed refusal, not a generic error') + assert.equal(err.code, 'query_budget_exceeded') + assert.equal(err.limitBytes, 1) + assert.ok(err.observedBytes > 1) + assert.match(err.message, /execution memory budget/) + assert.match(err.message, /WHERE|LIMIT|aggregate/, 'message carries actionable guidance') + return true + } + ) +}) + +test('resolveHeapBudgetBytes resolves the effective ceiling and never disables on a blank env', () => { + const DEFAULT = 1024 * 1024 * 1024 + const prev = process.env.HYP_QUERY_MAX_HEAP_MB + const withEnv = (/** @type {string | undefined} */ v, /** @type {() => void} */ body) => { + if (v === undefined) delete process.env.HYP_QUERY_MAX_HEAP_MB + else process.env.HYP_QUERY_MAX_HEAP_MB = v + body() + } + try { + // Explicit option always wins, including 0 (disable). + assert.equal(resolveHeapBudgetBytes(0), 0) + assert.equal(resolveHeapBudgetBytes(5 * 1024 * 1024), 5 * 1024 * 1024) + // Blank / whitespace-only must NOT resolve to 0 and disable the guard. + withEnv('', () => assert.equal(resolveHeapBudgetBytes(undefined), DEFAULT)) + withEnv(' ', () => assert.equal(resolveHeapBudgetBytes(undefined), DEFAULT)) + // Unset falls to the default. + withEnv(undefined, () => assert.equal(resolveHeapBudgetBytes(undefined), DEFAULT)) + // Garbage falls to the default (NaN is not finite). + withEnv('512mb', () => assert.equal(resolveHeapBudgetBytes(undefined), DEFAULT)) + // A real numeric override is honored, and an explicit 0 still disables. + withEnv('256', () => assert.equal(resolveHeapBudgetBytes(undefined), 256 * 1024 * 1024)) + withEnv('0', () => assert.equal(resolveHeapBudgetBytes(undefined), 0)) + } finally { + if (prev === undefined) delete process.env.HYP_QUERY_MAX_HEAP_MB + else process.env.HYP_QUERY_MAX_HEAP_MB = prev + } +}) + +test('maxHeapBytes 0 disables the budget entirely', async () => { + const rows = Array.from({ length: 6000 }, (_, i) => ({ a: i })) + const result = await executeQuerySql({ + query: 'SELECT COUNT(*) AS n FROM t', + registry: registryFor(memorySource(rows)), + storage, + maxHeapBytes: 0, + }) + assert.equal(result.rows[0].n, 6000) +}) + +test('a pre-aborted caller signal aborts execution before rows flow', async () => { + const rows = [{ a: 1 }] + const controller = new AbortController() + controller.abort() + await assert.rejects( + executeQuerySql({ + query: 'SELECT a FROM t', + registry: registryFor(memorySource(rows)), + storage, + signal: controller.signal, + }), + (err) => { + assert.equal(/** @type {Error} */ (err).name, 'AbortError') + return true + } + ) +}) + +test('the streaming-aggregate scanColumn fast path stays lit through the budget decoration', async () => { + const rows = Array.from({ length: 100 }, (_, i) => ({ a: `s${i % 7}` })) + /** @type {string[]} */ + const scanColumnCalls = [] + const result = await executeQuerySql({ + query: 'SELECT COUNT(DISTINCT a) AS n FROM t', + registry: registryFor(memorySource(rows, { scanColumnCalls })), + storage, + }) + assert.equal(result.rows[0].n, 7) + assert.deepEqual(scanColumnCalls, ['a'], 'the engine consumed the column stream, not buffered rows') +}) diff --git a/test/core/union-source.test.js b/test/core/union-source.test.js index ffed0aa..f463613 100644 --- a/test/core/union-source.test.js +++ b/test/core/union-source.test.js @@ -179,6 +179,84 @@ test('unionSources tolerates a scan with no options', async () => { assert.deepEqual(out, [{ id: 'a1' }]) }) +/** + * Add a recording `scanColumn` to a fake source, honoring its own + * limit/offset like the icebird-backed sources do (the scanColumn contract + * has no appliedLimitOffset escape hatch). + * + * @param {AsyncDataSource} source + * @param {Record[]} rows + * @param {{ column: string, limit?: number, offset?: number }[]} seenColumnScans + * @returns {AsyncDataSource} + */ +function withFakeScanColumn(source, rows, seenColumnScans) { + source.scanColumn = ({ column, limit, offset }) => ({ + async *[Symbol.asyncIterator]() { + seenColumnScans.push({ column, limit, offset }) + const start = offset ?? 0 + const end = limit === undefined ? rows.length : Math.min(rows.length, start + limit) + if (end > start) yield rows.slice(start, end).map((r) => r[column] ?? null) + }, + }) + return source +} + +test('unionSources omits scanColumn unless every partition can stream the column', () => { + const rows = [{ id: 'a1' }] + const withHook = withFakeScanColumn(fakeSource(rows, []), rows, []) + const withoutHook = fakeSource([{ id: 'b1' }], []) + assert.equal(typeof unionSources([withHook, withoutHook]).scanColumn, 'undefined', 'mixed union stays row-based') + assert.equal(typeof unionSources([withHook]).scanColumn, 'function') +}) + +test('unionSources scanColumn concatenates partitions and owns limit/offset over the merged stream', async () => { + /** @type {{ column: string, limit?: number, offset?: number }[]} */ + const seen = [] + const aRows = [{ v: 1 }, { v: 2 }, { v: 3 }] + const bRows = [{ v: 4 }, { v: 5 }, { v: 6 }] + const union = unionSources([ + withFakeScanColumn(fakeSource(aRows, []), aRows, seen), + withFakeScanColumn(fakeSource(bRows, []), bRows, seen), + ]) + + const scanColumn = /** @type {NonNullable} */ (union.scanColumn) + /** @type {SqlPrimitive[]} */ + const values = [] + for await (const chunk of scanColumn({ column: 'v', offset: 2, limit: 3 })) { + for (let i = 0; i < chunk.length; i++) values.push(chunk[i]) + } + + // Offset/limit apply to the CONCATENATED stream: skip 1,2 then take 3. + assert.deepEqual(values, [3, 4, 5]) + // Offset is never pushed per partition (not distributive); only the + // remaining-need upper bound is, so a partition never over-reads. + assert.deepEqual(seen, [ + { column: 'v', limit: 5, offset: undefined }, + { column: 'v', limit: 2, offset: undefined }, + ]) +}) + +test('unionSources scanColumn skips a whole partition its numRows proves is inside the offset', async () => { + /** @type {{ column: string, limit?: number, offset?: number }[]} */ + const seen = [] + const aRows = [{ v: 1 }, { v: 2 }] + const bRows = [{ v: 3 }, { v: 4 }] + const union = unionSources([ + withFakeScanColumn(fakeSource(aRows, []), aRows, seen), + withFakeScanColumn(fakeSource(bRows, []), bRows, seen), + ]) + + const scanColumn = /** @type {NonNullable} */ (union.scanColumn) + /** @type {SqlPrimitive[]} */ + const values = [] + for await (const chunk of scanColumn({ column: 'v', offset: 3 })) { + for (let i = 0; i < chunk.length; i++) values.push(chunk[i]) + } + + assert.deepEqual(values, [4]) + assert.deepEqual(seen, [{ column: 'v', limit: undefined, offset: undefined }], 'first partition never opened') +}) + test('emptySource advertises the given columns and yields no rows', async () => { const source = emptySource(['x', 'y']) assert.deepEqual(source.columns, ['x', 'y'])