diff --git a/packages/shared/src/btree-set.ts b/packages/shared/src/btree-set.ts index 968b9aaed6..a03537e92d 100644 --- a/packages/shared/src/btree-set.ts +++ b/packages/shared/src/btree-set.ts @@ -96,7 +96,7 @@ export class BTreeSet { return valuesFrom(this.#root, this.comparator, undefined, true); } - valuesFrom(lowestKey?: K, inclusive: boolean = true): IterableIterator { + valuesFrom(lowestKey?: K, inclusive: boolean = true): ValueIterator { return valuesFrom(this.#root, this.comparator, lowestKey, inclusive); } @@ -113,7 +113,7 @@ export class BTreeSet { valuesFromReversed( highestKey?: K, inclusive: boolean = true, - ): IterableIterator { + ): ValueIterator { return valuesFromReversed( this.#maxKey(), this.#root, @@ -187,7 +187,16 @@ export class BTreeSet { } } -class BTreeForwardIterator implements IterableIterator { +/** + * An iterator that can also hand back values without the iterator protocol's + * per-value result object. `nextValue()` returns `undefined` once exhausted, + * so it is only meaningful for sets whose keys are never `undefined`. + */ +export interface ValueIterator extends IterableIterator { + nextValue(): K | undefined; +} + +class BTreeForwardIterator implements ValueIterator { readonly #nodeQueue: BNode[][]; readonly #nodeIndex: number[]; #leaf: BNode; @@ -205,16 +214,17 @@ class BTreeForwardIterator implements IterableIterator { this.#i = startI; } - next(): IteratorResult { + /** Moves to the next key; false once exhausted. */ + #advance(): boolean { for (;;) { if (++this.#i < this.#leaf.keys.length) { - return {done: false, value: this.#leaf.keys[this.#i]}; + return true; } let level = -1; for (;;) { if (++level >= this.#nodeQueue.length) { - return {done: true, value: undefined as unknown as K}; + return false; } if (++this.#nodeIndex[level] < this.#nodeQueue[level].length) { break; @@ -231,12 +241,22 @@ class BTreeForwardIterator implements IterableIterator { } } + next(): IteratorResult { + return this.#advance() + ? {done: false, value: this.#leaf.keys[this.#i]} + : {done: true, value: undefined as unknown as K}; + } + + nextValue(): K | undefined { + return this.#advance() ? this.#leaf.keys[this.#i] : undefined; + } + [Symbol.iterator]() { return this; } } -class BTreeReverseIterator implements IterableIterator { +class BTreeReverseIterator implements ValueIterator { readonly #nodeQueue: BNode[][]; readonly #nodeIndex: number[]; #leaf: BNode; @@ -254,17 +274,18 @@ class BTreeReverseIterator implements IterableIterator { this.#i = startI; } - next(): IteratorResult { + /** Moves to the previous key; false once exhausted. */ + #advance(): boolean { for (;;) { if (--this.#i >= 0) { - return {done: false, value: this.#leaf.keys[this.#i]}; + return true; } let level; // Advance to the next leaf node for (level = -1; ;) { if (++level >= this.#nodeQueue.length) { - return {done: true, value: undefined as unknown as K}; + return false; } if (--this.#nodeIndex[level] >= 0) { break; @@ -281,6 +302,16 @@ class BTreeReverseIterator implements IterableIterator { } } + next(): IteratorResult { + return this.#advance() + ? {done: false, value: this.#leaf.keys[this.#i]} + : {done: true, value: undefined as unknown as K}; + } + + nextValue(): K | undefined { + return this.#advance() ? this.#leaf.keys[this.#i] : undefined; + } + [Symbol.iterator]() { return this; } @@ -291,10 +322,10 @@ function valuesFrom( comparator: Comparator, lowestKey: K | undefined, inclusive: boolean, -): IterableIterator { +): ValueIterator { const info = findPath(lowestKey, root, comparator); if (info === undefined) { - return iterator(() => ({done: true, value: undefined})); + return emptyValueIterator(); } let [nodeQueue, nodeIndex, leaf] = info; @@ -322,11 +353,11 @@ function valuesFromReversed( comparator: Comparator, highestKey: K | undefined, inclusive: boolean, -): IterableIterator { +): ValueIterator { if (highestKey === undefined) { highestKey = maxKey; if (highestKey === undefined) { - return iterator(() => ({done: true, value: undefined})); + return emptyValueIterator(); } // collection is empty } let [nodeQueue, nodeIndex, leaf] = @@ -371,9 +402,10 @@ function findPath( return [nodeQueue, nodeIndex, nextNode]; } -function iterator(next: () => IteratorResult): IterableIterator { +function emptyValueIterator(): ValueIterator { return { - next, + next: () => ({done: true, value: undefined as unknown as K}), + nextValue: () => undefined, [Symbol.iterator]() { return this; }, diff --git a/packages/zero-cache/src/auth/write-authorizer.ts b/packages/zero-cache/src/auth/write-authorizer.ts index 4858e97f57..87225daaf4 100644 --- a/packages/zero-cache/src/auth/write-authorizer.ts +++ b/packages/zero-cache/src/auth/write-authorizer.ts @@ -569,10 +569,13 @@ export class WriteAuthorizerImpl implements WriteAuthorizer { const input = buildPipeline(rowQueryAst, this.#builderDelegate, 'query-id'); try { const res = input.fetch({}); - for (const _ of res) { - // if any row is returned at all, the - // rule passes. - return true; + try { + // if any row is returned at all, the rule passes. + if (res.next() !== undefined) { + return true; + } + } finally { + res.close(); } } finally { input.destroy(); diff --git a/packages/zero-cache/src/services/run-ast.ts b/packages/zero-cache/src/services/run-ast.ts index 1c75f0e0ff..c008fe4ac5 100644 --- a/packages/zero-cache/src/services/run-ast.ts +++ b/packages/zero-cache/src/services/run-ast.ts @@ -102,8 +102,13 @@ export async function runAst( // triggering early return on Take's #initialFetch assertion. // The subquery AST already has limit: 1, so at most one row is produced. let node: Node | undefined; - for (const n of skipYields(input.fetch({}))) { - node ??= n; + const rows = skipYields(input.fetch({})); + try { + for (let n = rows.next(); n !== undefined; n = rows.next()) { + node ??= n; + } + } finally { + rows.close(); } input.destroy(); return node ? ((node.row[childField] as LiteralValue) ?? null) : undefined; diff --git a/packages/zero-cache/src/services/view-syncer/flipped-exists-fetch-filter.bench.ts b/packages/zero-cache/src/services/view-syncer/flipped-exists-fetch-filter.bench.ts index 3e901dfcb1..02aacaf9af 100644 --- a/packages/zero-cache/src/services/view-syncer/flipped-exists-fetch-filter.bench.ts +++ b/packages/zero-cache/src/services/view-syncer/flipped-exists-fetch-filter.bench.ts @@ -21,7 +21,8 @@ type Fetch = { function fetchRowCount(input: Input, filter: NoSubqueryCondition | undefined) { let count = 0; - for (const node of input.fetch({filter})) { + const stream = input.fetch({filter}); + for (let node = stream.next(); node !== undefined; node = stream.next()) { if (node !== 'yield') { count++; } diff --git a/packages/zero-cache/src/services/view-syncer/pipeline-driver.ts b/packages/zero-cache/src/services/view-syncer/pipeline-driver.ts index 3ef8023316..3cfc9f6307 100644 --- a/packages/zero-cache/src/services/view-syncer/pipeline-driver.ts +++ b/packages/zero-cache/src/services/view-syncer/pipeline-driver.ts @@ -15,7 +15,7 @@ import { import {ChangeIndex} from '../../../../zql/src/ivm/change-index.ts'; import {ChangeType} from '../../../../zql/src/ivm/change-type.ts'; import type {Change} from '../../../../zql/src/ivm/change.ts'; -import type {Node} from '../../../../zql/src/ivm/data.ts'; +import type {Node, RelationshipStream} from '../../../../zql/src/ivm/data.ts'; import { skipYields, throwOutput, @@ -33,6 +33,7 @@ import { makeSourceChangeEdit, makeSourceChangeRemove, } from '../../../../zql/src/ivm/source.ts'; +import {pullOf, type PullStream} from '../../../../zql/src/ivm/stream.ts'; import type {ConnectionCostModel} from '../../../../zql/src/planner/planner-connection.ts'; import {MeasurePushOperator} from '../../../../zql/src/query/measure-push-operator.ts'; import type {ClientGroupStorage} from '../../../../zqlite/src/database-storage.ts'; @@ -560,8 +561,13 @@ export class PipelineDriver { // triggering early return on Take's #initialFetch assertion. // The subquery AST already has limit: 1, so at most one row is produced. let node: Node | undefined; - for (const n of skipYields(input.fetch({}))) { - node ??= n; + const rows = skipYields(input.fetch({})); + try { + for (let n = rows.next(); n !== undefined; n = rows.next()) { + node ??= n; + } + } finally { + rows.close(); } if (!node) { return undefined; @@ -1369,9 +1375,9 @@ class Streamer { switch (type) { case ChangeType.REMOVE: case ChangeType.ADD: { - yield* this.#streamNodes(queryID, schema, type, () => [ - change[ChangeIndex.NODE], - ]); + yield* this.#streamNodes(queryID, schema, type, () => + pullOf([change[ChangeIndex.NODE]]), + ); break; } @@ -1385,9 +1391,9 @@ class Streamer { break; } case ChangeType.EDIT: - yield* this.#streamNodes(queryID, schema, type, () => [ - {row: change[ChangeIndex.NODE].row, relationships: {}}, - ]); + yield* this.#streamNodes(queryID, schema, type, () => + pullOf([{row: change[ChangeIndex.NODE].row, relationships: {}}]), + ); break; default: unreachable(change[ChangeIndex.TYPE]); @@ -1399,7 +1405,7 @@ class Streamer { queryID: string, schema: SourceSchema, op: ChangeType.ADD | ChangeType.REMOVE | ChangeType.EDIT, - nodes: () => Iterable, + nodes: () => RelationshipStream, ): Iterable { const {tableName: table, system} = schema; @@ -1412,36 +1418,46 @@ class Streamer { return; } - for (const node of nodes()) { - if (node === 'yield') { - yield node; - continue; - } - const {relationships} = node; - let {row} = node; - const rowKey = getRowKey(primaryKey, row); - if (op !== ChangeType.REMOVE) { - const rowVersion = row[ZERO_VERSION_COLUMN_NAME]; - if ( - typeof rowVersion === 'string' && - rowVersion < (spec.minRowVersion ?? '00') - ) { - row = {...row, [ZERO_VERSION_COLUMN_NAME]: spec.minRowVersion}; + const stream = nodes(); + try { + for (let node = stream.next(); node !== undefined; node = stream.next()) { + if (node === 'yield') { + yield node; + continue; + } + const {relationships} = node; + let {row} = node; + const rowKey = getRowKey(primaryKey, row); + if (op !== ChangeType.REMOVE) { + const rowVersion = row[ZERO_VERSION_COLUMN_NAME]; + if ( + typeof rowVersion === 'string' && + rowVersion < (spec.minRowVersion ?? '00') + ) { + row = {...row, [ZERO_VERSION_COLUMN_NAME]: spec.minRowVersion}; + } } - } - yield { - type: op, - queryID, - table, - rowKey, - row: op === ChangeType.REMOVE ? undefined : row, - } as RowChange; - - for (const [relationship, children] of Object.entries(relationships)) { - const childSchema = must(schema.relationships[relationship]); - yield* this.#streamNodes(queryID, childSchema, op, children); + yield { + type: op, + queryID, + table, + rowKey, + row: op === ChangeType.REMOVE ? undefined : row, + } as RowChange; + + for (const [relationship, children] of Object.entries(relationships)) { + const childSchema = must(schema.relationships[relationship]); + yield* this.#streamNodes( + queryID, + childSchema, + op, + children as () => RelationshipStream, + ); + } } + } finally { + stream.close(); } } } @@ -1481,7 +1497,7 @@ class QueryFailureLoggingOperator implements Input, Output { this.#input.destroy(); } - fetch(req: FetchRequest): Iterable { + fetch(req: FetchRequest): PullStream { return this.#input.fetch(req); } @@ -1525,13 +1541,24 @@ function logQueryFailure( queryLC.error?.(message, error); } -function* toAdds(nodes: Iterable): Iterable { - for (const node of nodes) { - if (node === 'yield') { - yield node; - continue; +function* toAdds( + nodes: PullStream, +): Iterable { + // `finally` is load-bearing: abandoning this generator -- an aborted or + // evicted hydration -- must still close the stream. `for...of` did that + // implicitly via `.return()`; the pull protocol makes it explicit, and a + // leaked SQLite cursor leaves later writes failing with "database + // connection is busy executing a query". + try { + for (let node = nodes.next(); node !== undefined; node = nodes.next()) { + if (node === 'yield') { + yield node; + continue; + } + yield [ChangeType.ADD, node, null]; } - yield [ChangeType.ADD, node, null]; + } finally { + nodes.close(); } } diff --git a/packages/zero-client/src/client/custom.test.ts b/packages/zero-client/src/client/custom.test.ts index 4e71a0ecae..2c7c0a8fac 100644 --- a/packages/zero-client/src/client/custom.test.ts +++ b/packages/zero-client/src/client/custom.test.ts @@ -12,6 +12,7 @@ import {zeroData} from '../../../replicache/src/transactions.ts'; import {createSilentLogContext} from '../../../shared/src/logging-test-utils.ts'; import {must} from '../../../shared/src/must.ts'; import {promiseUndefined} from '../../../shared/src/resolved-promises.ts'; +import {drainPull} from '../../../zql/src/ivm/stream.ts'; import {refCountSymbol} from '../../../zql/src/ivm/view-apply-change.ts'; import type {InsertValue} from '../../../zql/src/mutate/crud.ts'; import type {Transaction} from '../../../zql/src/mutate/custom.ts'; @@ -485,11 +486,13 @@ describe('rebasing custom mutators', () => { createdAt: 1743018138477, }); - expect([ - ...must(branch.getSource('issue')) - .connect([['id', 'asc']]) - .fetch({}), - ]).toMatchInlineSnapshot(` + expect( + drainPull( + must(branch.getSource('issue')) + .connect([['id', 'asc']]) + .fetch({}), + ), + ).toMatchInlineSnapshot(` [ { "relationships": {}, diff --git a/packages/zero-client/src/client/ivm-branch.test.ts b/packages/zero-client/src/client/ivm-branch.test.ts index ace2ecbb7b..76f613b831 100644 --- a/packages/zero-client/src/client/ivm-branch.test.ts +++ b/packages/zero-client/src/client/ivm-branch.test.ts @@ -17,7 +17,7 @@ import type {Hash} from '../../../replicache/src/hash.ts'; import type {Diff} from '../../../replicache/src/sync/patch.ts'; import {createSilentLogContext} from '../../../shared/src/logging-test-utils.ts'; import type {Node} from '../../../zql/src/ivm/data.ts'; -import {consume} from '../../../zql/src/ivm/stream.ts'; +import {consume, drainPull} from '../../../zql/src/ivm/stream.ts'; import {ENTITIES_KEY_PREFIX} from './keys.ts'; import {createDb} from './test/create-db.ts'; @@ -42,7 +42,7 @@ test('fork', () => { // Fork should have same initial data const fork = main.fork(); const forkConnection = fork.getSource('users')!.connect([['id', 'asc']]); - expect([...forkConnection.fetch({})]).toMatchInlineSnapshot(` + expect(drainPull(forkConnection.fetch({}))).toMatchInlineSnapshot(` [ { "relationships": {}, @@ -65,7 +65,7 @@ test('fork', () => { ); // Verify main and fork evolved independently - expect([...mainConnection.fetch({})]).toMatchInlineSnapshot(` + expect(drainPull(mainConnection.fetch({}))).toMatchInlineSnapshot(` [ { "relationships": {}, @@ -84,7 +84,7 @@ test('fork', () => { ] `); - expect([...forkConnection.fetch({})]).toMatchInlineSnapshot(` + expect(drainPull(forkConnection.fetch({}))).toMatchInlineSnapshot(` [ { "relationships": {}, @@ -127,11 +127,13 @@ describe('advance', () => { ); await initFromStore(branch, syncHash, dagStore); - expect([ - ...must(branch.getSource('issue')) - .connect([['id', 'asc']]) - .fetch({}), - ]).toMatchInlineSnapshot(` + expect( + drainPull( + must(branch.getSource('issue')) + .connect([['id', 'asc']]) + .fetch({}), + ), + ).toMatchInlineSnapshot(` [ { "relationships": {}, @@ -385,11 +387,13 @@ describe('advance', () => { const head = await w.commit(SYNC_HEAD_NAME); await branch.advance(syncHash, head, diffs); - expect([ - ...must(branch.getSource('issue')) - .connect([['id', 'asc']]) - .fetch({}), - ]).toEqual(expected); + expect( + drainPull( + must(branch.getSource('issue')) + .connect([['id', 'asc']]) + .fetch({}), + ), + ).toEqual(expected); }); }); @@ -475,11 +479,13 @@ describe('forkToHead', () => { ); await initFromStore(branch, syncHash, dagStore); await branch.forkToHead(dagStore, syncHash); - expect([ - ...must(branch.getSource('issue')) - .connect([['id', 'asc']]) - .fetch({}), - ]).toMatchInlineSnapshot(` + expect( + drainPull( + must(branch.getSource('issue')) + .connect([['id', 'asc']]) + .fetch({}), + ), + ).toMatchInlineSnapshot(` [ { "relationships": {}, @@ -521,11 +527,13 @@ describe('forkToHead', () => { const head = await w.commit(SYNC_HEAD_NAME); const fork = await branch.forkToHead(dagStore, head); - expect([ - ...must(fork.getSource('issue')) - .connect([['id', 'asc']]) - .fetch({}), - ]).toMatchInlineSnapshot(` + expect( + drainPull( + must(fork.getSource('issue')) + .connect([['id', 'asc']]) + .fetch({}), + ), + ).toMatchInlineSnapshot(` [ { "relationships": {}, @@ -542,11 +550,13 @@ describe('forkToHead', () => { // can also re-wind the fork to the original head const fork2 = await fork.forkToHead(dagStore, syncHash); - expect([ - ...must(fork2.getSource('issue')) - .connect([['id', 'asc']]) - .fetch({}), - ]).toMatchInlineSnapshot(` + expect( + drainPull( + must(fork2.getSource('issue')) + .connect([['id', 'asc']]) + .fetch({}), + ), + ).toMatchInlineSnapshot(` [] `); }); diff --git a/packages/zero-solid/src/solid-view.test.ts b/packages/zero-solid/src/solid-view.test.ts index fa19ba0945..0a94f09147 100644 --- a/packages/zero-solid/src/solid-view.test.ts +++ b/packages/zero-solid/src/solid-view.test.ts @@ -11,6 +11,7 @@ import { makeEditChange, makeRemoveChange, } from '../../zql/src/ivm/change.ts'; +import type {Node} from '../../zql/src/ivm/data.ts'; import {Join} from '../../zql/src/ivm/join.ts'; import {MemorySource} from '../../zql/src/ivm/memory-source.ts'; import {MemoryStorage} from '../../zql/src/ivm/memory-storage.ts'; @@ -21,7 +22,7 @@ import { makeSourceChangeEdit, makeSourceChangeRemove, } from '../../zql/src/ivm/source.ts'; -import {consume} from '../../zql/src/ivm/stream.ts'; +import {consume, emptyPullStream, pullOf} from '../../zql/src/ivm/stream.ts'; import {Take} from '../../zql/src/ivm/take.ts'; import {createSource} from '../../zql/src/ivm/test/source-factory.ts'; import {idSymbol, refCountSymbol, unreachable} from './bindings.ts'; @@ -706,7 +707,7 @@ test('collapse', () => { const input: Input = { fetch() { - return []; + return emptyPullStream(); }, destroy() {}, getSchema() { @@ -752,27 +753,29 @@ test('collapse', () => { name: 'issue', }, relationships: { - labels: () => [ - { - row: { - id: 1, - issueId: 1, - labelId: 1, - extra: 'a', - }, - relationships: { - labels: () => [ - { - row: { - id: 1, - name: 'label', - }, - relationships: {}, - }, - ], + labels: () => + pullOf([ + { + row: { + id: 1, + issueId: 1, + labelId: 1, + extra: 'a', + }, + relationships: { + labels: () => + pullOf([ + { + row: { + id: 1, + name: 'label', + }, + relationships: {}, + }, + ]), + }, }, - }, - ], + ]), }, }, } as const; @@ -820,46 +823,49 @@ test('collapse', () => { name: 'issue', }, relationships: { - labels: () => [ - { - row: { - id: 1, - issueId: 1, - labelId: 1, - extra: 'a', - }, - relationships: { - labels: () => [ - { - row: { - id: 1, - name: 'label', - }, - relationships: {}, - }, - ], - }, - }, - { - row: { - id: 2, - issueId: 1, - labelId: 2, - extra: 'b', + labels: () => + pullOf([ + { + row: { + id: 1, + issueId: 1, + labelId: 1, + extra: 'a', + }, + relationships: { + labels: () => + pullOf([ + { + row: { + id: 1, + name: 'label', + }, + relationships: {}, + }, + ]), + }, }, - relationships: { - labels: () => [ - { - row: { - id: 2, - name: 'label2', - }, - relationships: {}, - }, - ], + { + row: { + id: 2, + issueId: 1, + labelId: 2, + extra: 'b', + }, + relationships: { + labels: () => + pullOf([ + { + row: { + id: 2, + name: 'label2', + }, + relationships: {}, + }, + ]), + }, }, - }, - ], + ]), }, }, { @@ -872,15 +878,16 @@ test('collapse', () => { extra: 'b', }, relationships: { - labels: () => [ - { - row: { - id: 2, - name: 'label2', + labels: () => + pullOf([ + { + row: { + id: 2, + name: 'label2', + }, + relationships: {}, }, - relationships: {}, - }, - ], + ]), }, }), }, @@ -925,46 +932,49 @@ test('collapse', () => { name: 'issue', }, relationships: { - labels: () => [ - { - row: { - id: 1, - issueId: 1, - labelId: 1, - extra: 'a', - }, - relationships: { - labels: () => [ - { - row: { - id: 1, - name: 'label', - }, - relationships: {}, - }, - ], - }, - }, - { - row: { - id: 2, - issueId: 1, - labelId: 2, - extra: 'b2', + labels: () => + pullOf([ + { + row: { + id: 1, + issueId: 1, + labelId: 1, + extra: 'a', + }, + relationships: { + labels: () => + pullOf([ + { + row: { + id: 1, + name: 'label', + }, + relationships: {}, + }, + ]), + }, }, - relationships: { - labels: () => [ - { - row: { - id: 2, - name: 'label2', - }, - relationships: {}, - }, - ], + { + row: { + id: 2, + issueId: 1, + labelId: 2, + extra: 'b2', + }, + relationships: { + labels: () => + pullOf([ + { + row: { + id: 2, + name: 'label2', + }, + relationships: {}, + }, + ]), + }, }, - }, - ], + ]), }, }, { @@ -978,15 +988,16 @@ test('collapse', () => { extra: 'b2', }, relationships: { - labels: () => [ - { - row: { - id: 2, - name: 'label2', + labels: () => + pullOf([ + { + row: { + id: 2, + name: 'label2', + }, + relationships: {}, }, - relationships: {}, - }, - ], + ]), }, }, { @@ -997,15 +1008,16 @@ test('collapse', () => { extra: 'b', }, relationships: { - labels: () => [ - { - row: { - id: 2, - name: 'label2', + labels: () => + pullOf([ + { + row: { + id: 2, + name: 'label2', + }, + relationships: {}, }, - relationships: {}, - }, - ], + ]), }, }, ), @@ -1050,46 +1062,49 @@ test('collapse', () => { name: 'issue', }, relationships: { - labels: () => [ - { - row: { - id: 1, - issueId: 1, - labelId: 1, - extra: 'a', - }, - relationships: { - labels: () => [ - { - row: { - id: 1, - name: 'label', - }, - relationships: {}, - }, - ], - }, - }, - { - row: { - id: 2, - issueId: 1, - labelId: 2, - extra: 'b2', + labels: () => + pullOf([ + { + row: { + id: 1, + issueId: 1, + labelId: 1, + extra: 'a', + }, + relationships: { + labels: () => + pullOf([ + { + row: { + id: 1, + name: 'label', + }, + relationships: {}, + }, + ]), + }, }, - relationships: { - labels: () => [ - { - row: { - id: 2, - name: 'label2x', - }, - relationships: {}, - }, - ], + { + row: { + id: 2, + issueId: 1, + labelId: 2, + extra: 'b2', + }, + relationships: { + labels: () => + pullOf([ + { + row: { + id: 2, + name: 'label2x', + }, + relationships: {}, + }, + ]), + }, }, - }, - ], + ]), }, }, { @@ -1103,15 +1118,16 @@ test('collapse', () => { extra: 'b2', }, relationships: { - labels: () => [ - { - row: { - id: 2, - name: 'label2x', + labels: () => + pullOf([ + { + row: { + id: 2, + name: 'label2x', + }, + relationships: {}, }, - relationships: {}, - }, - ], + ]), }, }, { @@ -1211,10 +1227,10 @@ test('collapse-single', () => { const input = { cleanup() { - return []; + return emptyPullStream(); }, fetch() { - return []; + return emptyPullStream(); }, destroy() {}, getSchema() { @@ -1263,26 +1279,28 @@ test('collapse-single', () => { name: 'issue', }, relationships: { - labels: () => [ - { - row: { - id: 1, - issueId: 1, - labelId: 1, - }, - relationships: { - labels: () => [ - { - row: { - id: 1, - name: 'label', - }, - relationships: {}, - }, - ], + labels: () => + pullOf([ + { + row: { + id: 1, + issueId: 1, + labelId: 1, + }, + relationships: { + labels: () => + pullOf([ + { + row: { + id: 1, + name: 'label', + }, + relationships: {}, + }, + ]), + }, }, - }, - ], + ]), }, }, } as const; @@ -1910,7 +1928,7 @@ test('edit to preserve relationships', () => { return schema; }, fetch() { - return []; + return emptyPullStream(); }, setOutput() {}, destroy() { @@ -1952,12 +1970,13 @@ test('edit to preserve relationships', () => { makeAddChange({ row: {id: 1, title: 'issue1'}, relationships: { - labels: () => [ - { - row: {id: 1, name: 'label1'}, - relationships: {}, - }, - ], + labels: () => + pullOf([ + { + row: {id: 1, name: 'label1'}, + relationships: {}, + }, + ]), }, }), ); @@ -1968,12 +1987,13 @@ test('edit to preserve relationships', () => { makeAddChange({ row: {id: 2, title: 'issue2'}, relationships: { - labels: () => [ - { - row: {id: 2, name: 'label2'}, - relationships: {}, - }, - ], + labels: () => + pullOf([ + { + row: {id: 2, name: 'label2'}, + relationships: {}, + }, + ]), }, }), ); @@ -2130,7 +2150,7 @@ test('edit leaf', () => { const input: Input = { fetch() { - return []; + return emptyPullStream(); }, destroy() {}, getSchema() { @@ -2176,17 +2196,18 @@ test('edit leaf', () => { name: 'issue', }, relationships: { - labels: () => [ - { - row: { - id: 1, - issueId: 1, - labelId: 1, - extra: 'a', + labels: () => + pullOf([ + { + row: { + id: 1, + issueId: 1, + labelId: 1, + extra: 'a', + }, + relationships: {}, }, - relationships: {}, - }, - ], + ]), }, }, } as const; @@ -2235,26 +2256,27 @@ test('edit leaf', () => { name: 'issue', }, relationships: { - labels: () => [ - { - row: { - id: 1, - issueId: 1, - labelId: 1, - extra: 'a', + labels: () => + pullOf([ + { + row: { + id: 1, + issueId: 1, + labelId: 1, + extra: 'a', + }, + relationships: {}, }, - relationships: {}, - }, - { - row: { - id: 2, - issueId: 1, - labelId: 2, - extra: 'b', + { + row: { + id: 2, + issueId: 1, + labelId: 2, + extra: 'b', + }, + relationships: {}, }, - relationships: {}, - }, - ], + ]), }, }, { @@ -2312,26 +2334,27 @@ test('edit leaf', () => { name: 'issue', }, relationships: { - labels: () => [ - { - row: { - id: 1, - issueId: 1, - labelId: 1, - extra: 'a', + labels: () => + pullOf([ + { + row: { + id: 1, + issueId: 1, + labelId: 1, + extra: 'a', + }, + relationships: {}, }, - relationships: {}, - }, - { - row: { - id: 2, - issueId: 1, - labelId: 2, - extra: 'b2', + { + row: { + id: 2, + issueId: 1, + labelId: 2, + extra: 'b2', + }, + relationships: {}, }, - relationships: {}, - }, - ], + ]), }, }, { diff --git a/packages/zero-solid/src/solid-view.ts b/packages/zero-solid/src/solid-view.ts index 50ef412d75..2970eb13a4 100644 --- a/packages/zero-solid/src/solid-view.ts +++ b/packages/zero-solid/src/solid-view.ts @@ -2,6 +2,8 @@ import {produce, reconcile, type SetStoreFunction} from 'solid-js/store'; import {emptyArray} from '../../shared/src/sentinels.ts'; import {ChangeIndex} from '../../zql/src/ivm/change-index.ts'; import {ChangeType} from '../../zql/src/ivm/change-type.ts'; +import type {RelationshipStream} from '../../zql/src/ivm/data.ts'; +import {drainPull, pullOf} from '../../zql/src/ivm/stream.ts'; import { applyChange, idSymbol, @@ -21,7 +23,6 @@ import { type QueryErrorDetails, type QueryResultDetails, type Schema, - type Stream, type TTL, } from './zero.ts'; @@ -113,7 +114,7 @@ export class SolidView implements Output { const initialRoot = this.#createEmptyRoot(); this.#applyChangesToRoot( - skipYields(input.fetch({})), + drainPull(skipYields(input.fetch({}))), node => ({type: 'add', node}), initialRoot, ); @@ -286,13 +287,14 @@ function materializeRelationships(change: Change): ViewChange { } function materializeNodeRelationships(node: Node): Node { - const relationships: Record Stream> = {}; + const relationships: Record RelationshipStream> = {}; for (const relationship in node.relationships) { const materialized: Node[] = []; - for (const n of skipYields(node.relationships[relationship]())) { + const children = skipYields(node.relationships[relationship]()); + for (let n = children.next(); n !== undefined; n = children.next()) { materialized.push(materializeNodeRelationships(n)); } - relationships[relationship] = () => materialized; + relationships[relationship] = () => pullOf(materialized); } return { row: node.row, diff --git a/packages/zql-integration-tests/src/helpers/runner.ts b/packages/zql-integration-tests/src/helpers/runner.ts index da2863b8b5..65dfe04728 100644 --- a/packages/zql-integration-tests/src/helpers/runner.ts +++ b/packages/zql-integration-tests/src/helpers/runner.ts @@ -688,8 +688,17 @@ function gatherRows( _queryComplete, ) => { const schema = input.getSchema(); - for (const node of skipYields(input.fetch({}))) { - processNode(schema, node); + const stream = skipYields(input.fetch({})); + try { + for ( + let node = stream.next(); + node !== undefined; + node = stream.next() + ) { + processNode(schema, node); + } + } finally { + stream.close(); } return { @@ -715,8 +724,17 @@ function gatherRows( node.relationships, )) { const childSchema = must(schema.relationships[relationship]); - for (const child of skipYields(getChildren())) { - processNode(childSchema, child); + const children = skipYields(getChildren()); + try { + for ( + let child = children.next(); + child !== undefined; + child = children.next() + ) { + processNode(childSchema, child); + } + } finally { + children.close(); } } } diff --git a/packages/zql/src/ivm/array-view.test.ts b/packages/zql/src/ivm/array-view.test.ts index 37b7f24d84..eddf5ec2cb 100644 --- a/packages/zql/src/ivm/array-view.test.ts +++ b/packages/zql/src/ivm/array-view.test.ts @@ -17,16 +17,17 @@ import { makeEditChange, makeRemoveChange, } from './change.ts'; +import type {Node} from './data.ts'; import {Join} from './join.ts'; import {MemoryStorage} from './memory-storage.ts'; -import type {Input} from './operator.ts'; +import type {FetchRequest, Input} from './operator.ts'; import type {SourceSchema} from './schema.ts'; import { makeSourceChangeAdd, makeSourceChangeEdit, makeSourceChangeRemove, } from './source.ts'; -import {consume} from './stream.ts'; +import {consume, emptyPullStream, pullOf} from './stream.ts'; import {Take} from './take.ts'; import {createSource} from './test/source-factory.ts'; import {refCountSymbol} from './view-apply-change.ts'; @@ -653,8 +654,8 @@ test('collapse', () => { }; const input: Input = { - fetch() { - return []; + fetch(_req: FetchRequest) { + return emptyPullStream(); }, destroy() {}, getSchema() { @@ -685,27 +686,29 @@ test('collapse', () => { name: 'issue', }, relationships: { - labels: () => [ - { - row: { - id: 1, - issueId: 1, - labelId: 1, - extra: 'a', - }, - relationships: { - labels: () => [ - { - row: { - id: 1, - name: 'label', - }, - relationships: {}, - }, - ], + labels: () => + pullOf([ + { + row: { + id: 1, + issueId: 1, + labelId: 1, + extra: 'a', + }, + relationships: { + labels: () => + pullOf([ + { + row: { + id: 1, + name: 'label', + }, + relationships: {}, + }, + ]), + }, }, - }, - ], + ]), }, }, } as const; @@ -747,46 +750,49 @@ test('collapse', () => { name: 'issue', }, relationships: { - labels: () => [ - { - row: { - id: 1, - issueId: 1, - labelId: 1, - extra: 'a', - }, - relationships: { - labels: () => [ - { - row: { - id: 1, - name: 'label', - }, - relationships: {}, - }, - ], - }, - }, - { - row: { - id: 2, - issueId: 1, - labelId: 2, - extra: 'b', + labels: () => + pullOf([ + { + row: { + id: 1, + issueId: 1, + labelId: 1, + extra: 'a', + }, + relationships: { + labels: () => + pullOf([ + { + row: { + id: 1, + name: 'label', + }, + relationships: {}, + }, + ]), + }, }, - relationships: { - labels: () => [ - { - row: { - id: 2, - name: 'label2', - }, - relationships: {}, - }, - ], + { + row: { + id: 2, + issueId: 1, + labelId: 2, + extra: 'b', + }, + relationships: { + labels: () => + pullOf([ + { + row: { + id: 2, + name: 'label2', + }, + relationships: {}, + }, + ]), + }, }, - }, - ], + ]), }, }, { @@ -799,15 +805,16 @@ test('collapse', () => { extra: 'b', }, relationships: { - labels: () => [ - { - row: { - id: 2, - name: 'label2', + labels: () => + pullOf([ + { + row: { + id: 2, + name: 'label2', + }, + relationships: {}, }, - relationships: {}, - }, - ], + ]), }, }), }, @@ -848,46 +855,49 @@ test('collapse', () => { name: 'issue', }, relationships: { - labels: () => [ - { - row: { - id: 1, - issueId: 1, - labelId: 1, - extra: 'a', - }, - relationships: { - labels: () => [ - { - row: { - id: 1, - name: 'label', - }, - relationships: {}, - }, - ], - }, - }, - { - row: { - id: 2, - issueId: 1, - labelId: 2, - extra: 'b2', + labels: () => + pullOf([ + { + row: { + id: 1, + issueId: 1, + labelId: 1, + extra: 'a', + }, + relationships: { + labels: () => + pullOf([ + { + row: { + id: 1, + name: 'label', + }, + relationships: {}, + }, + ]), + }, }, - relationships: { - labels: () => [ - { - row: { - id: 2, - name: 'label2', - }, - relationships: {}, - }, - ], + { + row: { + id: 2, + issueId: 1, + labelId: 2, + extra: 'b2', + }, + relationships: { + labels: () => + pullOf([ + { + row: { + id: 2, + name: 'label2', + }, + relationships: {}, + }, + ]), + }, }, - }, - ], + ]), }, }, { @@ -901,15 +911,16 @@ test('collapse', () => { extra: 'b2', }, relationships: { - labels: () => [ - { - row: { - id: 2, - name: 'label2', + labels: () => + pullOf([ + { + row: { + id: 2, + name: 'label2', + }, + relationships: {}, }, - relationships: {}, - }, - ], + ]), }, }, { @@ -920,15 +931,16 @@ test('collapse', () => { extra: 'b', }, relationships: { - labels: () => [ - { - row: { - id: 2, - name: 'label2', + labels: () => + pullOf([ + { + row: { + id: 2, + name: 'label2', + }, + relationships: {}, }, - relationships: {}, - }, - ], + ]), }, }, ), @@ -970,46 +982,49 @@ test('collapse', () => { name: 'issue', }, relationships: { - labels: () => [ - { - row: { - id: 1, - issueId: 1, - labelId: 1, - extra: 'a', - }, - relationships: { - labels: () => [ - { - row: { - id: 1, - name: 'label', - }, - relationships: {}, - }, - ], - }, - }, - { - row: { - id: 2, - issueId: 1, - labelId: 2, - extra: 'b2', + labels: () => + pullOf([ + { + row: { + id: 1, + issueId: 1, + labelId: 1, + extra: 'a', + }, + relationships: { + labels: () => + pullOf([ + { + row: { + id: 1, + name: 'label', + }, + relationships: {}, + }, + ]), + }, }, - relationships: { - labels: () => [ - { - row: { - id: 2, - name: 'label2x', - }, - relationships: {}, - }, - ], + { + row: { + id: 2, + issueId: 1, + labelId: 2, + extra: 'b2', + }, + relationships: { + labels: () => + pullOf([ + { + row: { + id: 2, + name: 'label2x', + }, + relationships: {}, + }, + ]), + }, }, - }, - ], + ]), }, }, { @@ -1023,15 +1038,16 @@ test('collapse', () => { extra: 'b2', }, relationships: { - labels: () => [ - { - row: { - id: 2, - name: 'label2x', + labels: () => + pullOf([ + { + row: { + id: 2, + name: 'label2x', + }, + relationships: {}, }, - relationships: {}, - }, - ], + ]), }, }, { @@ -1128,8 +1144,8 @@ test('collapse-single', () => { }; const input = { - fetch() { - return []; + fetch(_req: FetchRequest) { + return emptyPullStream(); }, destroy() {}, getSchema() { @@ -1162,26 +1178,28 @@ test('collapse-single', () => { name: 'issue', }, relationships: { - labels: () => [ - { - row: { - id: 1, - issueId: 1, - labelId: 1, - }, - relationships: { - labels: () => [ - { - row: { - id: 1, - name: 'label', - }, - relationships: {}, - }, - ], + labels: () => + pullOf([ + { + row: { + id: 1, + issueId: 1, + labelId: 1, + }, + relationships: { + labels: () => + pullOf([ + { + row: { + id: 1, + name: 'label', + }, + relationships: {}, + }, + ]), + }, }, - }, - ], + ]), }, }, } as const; @@ -1601,8 +1619,8 @@ test('edit to preserve relationships', () => { getSchema() { return schema; }, - fetch() { - return []; + fetch(_req: FetchRequest) { + return emptyPullStream(); }, setOutput() {}, destroy() { @@ -1624,12 +1642,13 @@ test('edit to preserve relationships', () => { makeAddChange({ row: {id: 1, title: 'issue1'}, relationships: { - labels: () => [ - { - row: {id: 1, name: 'label1'}, - relationships: {}, - }, - ], + labels: () => + pullOf([ + { + row: {id: 1, name: 'label1'}, + relationships: {}, + }, + ]), }, }), ), @@ -1639,12 +1658,13 @@ test('edit to preserve relationships', () => { makeAddChange({ row: {id: 2, title: 'issue2'}, relationships: { - labels: () => [ - { - row: {id: 2, name: 'label2'}, - relationships: {}, - }, - ], + labels: () => + pullOf([ + { + row: {id: 2, name: 'label2'}, + relationships: {}, + }, + ]), }, }), ), diff --git a/packages/zql/src/ivm/array-view.ts b/packages/zql/src/ivm/array-view.ts index 9b8d0fa10c..9f9029d9f0 100644 --- a/packages/zql/src/ivm/array-view.ts +++ b/packages/zql/src/ivm/array-view.ts @@ -7,7 +7,7 @@ import type {Listener, ResultType, TypedView} from '../query/typed-view.ts'; import {ChangeIndex} from './change-index.ts'; import {ChangeType} from './change-type.ts'; import type {Change} from './change.ts'; -import {skipYields, type Input, type Output} from './operator.ts'; +import type {Input, Output} from './operator.ts'; import type {SourceSchema} from './schema.ts'; import {applyChange, type ViewChange} from './view-apply-change.ts'; import type {Entry, Format, View} from './view.ts'; @@ -144,19 +144,33 @@ export class ArrayView implements Output, TypedView { #hydrate() { this.#dirty = true; - for (const node of skipYields(this.#input.fetch({}))) { - this.#root = applyChange( - this.#root, - {type: 'add', node}, - this.#getSchema(), - '', - this.#format, - false /* withIDs */, - true /* mutate: #root is freshly created and not yet observed by any - consumer, so build it in place to avoid O(N^2) array copies. - Every later push() is immutable, preserving reference - stability for unchanged subtrees. */, - ); + // Pull protocol: no result object per node. 'yield' is skipped inline + // rather than through skipYields, which would put an iterator back. + const stream = this.#input.fetch({}); + try { + for (;;) { + const node = stream.next(); + if (node === undefined) { + break; + } + if (node === 'yield') { + continue; + } + this.#root = applyChange( + this.#root, + {type: 'add', node}, + this.#getSchema(), + '', + this.#format, + false /* withIDs */, + true /* mutate: #root is freshly created and not yet observed by any + consumer, so build it in place to avoid O(N^2) array copies. + Every later push() is immutable, preserving reference + stability for unchanged subtrees. */, + ); + } + } finally { + stream.close(); } this.flush(); } diff --git a/packages/zql/src/ivm/cap.ts b/packages/zql/src/ivm/cap.ts index 2590b59326..6e723d81bc 100644 --- a/packages/zql/src/ivm/cap.ts +++ b/packages/zql/src/ivm/cap.ts @@ -15,7 +15,12 @@ import { type Storage, } from './operator.ts'; import type {SourceSchema} from './schema.ts'; -import {type Stream} from './stream.ts'; +import { + type Stream, + emptyPullStream, + type PullStream, + limitedScan, +} from './stream.ts'; import { constraintMatchesPartitionKey, makePartitionKeyComparator, @@ -84,7 +89,7 @@ export class Cap implements Operator { return this.#input.getSchema(); } - *fetch(req: FetchRequest): Stream { + fetch(req: FetchRequest): PullStream { assert(!req.start, 'Cap does not support start'); assert(!req.reverse, 'Cap does not support reverse'); @@ -102,78 +107,45 @@ export class Cap implements Operator { const capStateKey = getCapStateKey(this.#partitionKey, req.constraint); const capState = this.#storage.get(capStateKey); if (!capState) { - yield* this.#initialFetch(req); - return; + return this.#initialFetch(req, capStateKey); } if (capState.size === 0) { - return; + return emptyPullStream(); } // PK-based point lookups: fetch each tracked row by its PK directly, // rather than scanning the partition and filtering. - for (const pk of capState.pks) { - const constraint = deserializePKToConstraint(pk, this.#primaryKey); - for (const inputNode of this.#input.fetch({constraint})) { - if (inputNode === 'yield') { - yield inputNode; - continue; - } - yield inputNode; - } - } + return new CapPointLookups(capState.pks, pk => + this.#input.fetch({ + constraint: deserializePKToConstraint(pk, this.#primaryKey), + }), + ); } - *#initialFetch(req: FetchRequest): Stream { + #initialFetch( + req: FetchRequest, + capStateKey: string, + ): PullStream { if (this.#limit === 0) { - return; + return emptyPullStream(); } - assert( constraintMatchesPartitionKey(req.constraint, this.#partitionKey), 'Constraint should match partition key', ); - - const capStateKey = getCapStateKey(this.#partitionKey, req.constraint); assert( this.#storage.get(capStateKey) === undefined, 'Cap state should be undefined', ); - - let size = 0; const pks: string[] = []; - let downstreamEarlyReturn = true; - let exceptionThrown = false; - try { - for (const inputNode of this.#input.fetch(req)) { - if (inputNode === 'yield') { - yield 'yield'; - continue; - } - yield inputNode; - pks.push(serializePK(inputNode.row, this.#primaryKey)); - size++; - if (size === this.#limit) { - break; - } - } - downstreamEarlyReturn = false; - } catch (e) { - exceptionThrown = true; - throw e; - } finally { - if (!exceptionThrown) { - this.#storage.set(capStateKey, {size, pks}); - // If it becomes necessary to support downstream early return, this - // assert should be removed, and replaced with code that consumes - // the input stream until limit is reached or the input stream is - // exhausted so that capState is properly hydrated. - assert( - !downstreamEarlyReturn, - 'Unexpected early return prevented full hydration', - ); - } - } + return limitedScan( + this.#input.fetch(req), + this.#limit, + node => node !== 'yield', + node => pks.push(serializePK((node as Node).row, this.#primaryKey)), + () => this.#storage.set(capStateKey, {size: pks.length, pks}), + () => assert(false, 'Unexpected early return prevented full hydration'), + ); } - *push(change: Change): Stream<'yield'> { if (change[ChangeIndex.TYPE] === ChangeType.EDIT) { yield* this.#pushEditChange(change); @@ -223,16 +195,26 @@ export class Cap implements Operator { : undefined; let replacement: Node | undefined; - for (const node of this.#input.fetch({constraint})) { - if (node === 'yield') { - yield node; - continue; - } - const nodePK = serializePK(node.row, this.#primaryKey); - if (!pkSet.has(nodePK)) { - replacement = node; - break; + + const candidates = this.#input.fetch({constraint}); + try { + for ( + let node = candidates.next(); + node !== undefined; + node = candidates.next() + ) { + if (node === 'yield') { + yield node; + continue; + } + const nodePK = serializePK(node.row, this.#primaryKey); + if (!pkSet.has(nodePK)) { + replacement = node; + break; + } } + } finally { + candidates.close(); } if (replacement) { @@ -327,3 +309,41 @@ function deserializePKToConstraint( } return constraint; } + +/** Flattens per-PK point lookups into one stream. */ +class CapPointLookups implements PullStream { + readonly #pks: readonly string[]; + readonly #fetch: (pk: string) => PullStream; + #i = 0; + #cur: PullStream | undefined; + + constructor( + pks: readonly string[], + fetch: (pk: string) => PullStream, + ) { + this.#pks = pks; + this.#fetch = fetch; + } + + next(): Node | 'yield' | undefined { + for (;;) { + if (this.#cur !== undefined) { + const v = this.#cur.next(); + if (v !== undefined) { + return v; + } + this.#cur = undefined; + } + if (this.#i >= this.#pks.length) { + return undefined; + } + this.#cur = this.#fetch(this.#pks[this.#i++]); + } + } + + close(): void { + this.#i = this.#pks.length; + this.#cur?.close(); + this.#cur = undefined; + } +} diff --git a/packages/zql/src/ivm/catch.ts b/packages/zql/src/ivm/catch.ts index 99c4a4d05c..5ddc2c1367 100644 --- a/packages/zql/src/ivm/catch.ts +++ b/packages/zql/src/ivm/catch.ts @@ -7,6 +7,7 @@ import {ChangeType} from './change-type.ts'; import type {Change} from './change.ts'; import type {Node} from './data.ts'; import {type FetchRequest, type Input, type Output} from './operator.ts'; +import {drainPullMap} from './stream.ts'; export type CaughtNode = | { @@ -63,12 +64,12 @@ export class Catch implements Output { } fetch(req: FetchRequest = {}) { - return Array.from(this.#input.fetch(req), expandNode); + return drainPullMap(this.#input.fetch(req), expandNode); } push(change: Change) { const fetch = this.#fetchOnPush - ? Array.from(this.#input.fetch({}), expandNode) + ? drainPullMap(this.#input.fetch({}), expandNode) : []; const expandedChange = expandChange(change); if (this.#fetchOnPush) { @@ -129,8 +130,18 @@ export function expandNode(node: Node | 'yield'): CaughtNode { row: node.row, relationships: mapValues(node.relationships, getChildren => { const children: CaughtNode[] = []; - for (const child of getChildren()) { - children.push(expandNode(child)); + + const childStream = getChildren(); + try { + for ( + let child = childStream.next(); + child !== undefined; + child = childStream.next() + ) { + children.push(expandNode(child)); + } + } finally { + childStream.close(); } return children; }), diff --git a/packages/zql/src/ivm/data.ts b/packages/zql/src/ivm/data.ts index 1c0c2892aa..e0465010e9 100644 --- a/packages/zql/src/ivm/data.ts +++ b/packages/zql/src/ivm/data.ts @@ -1,7 +1,14 @@ import {compareUTF8} from 'compare-utf8'; import type {Ordering} from '../../../zero-protocol/src/ast.ts'; import type {Row, Value} from '../../../zero-protocol/src/data.ts'; -import type {Stream} from './stream.ts'; +import type {PullStream} from './stream.ts'; + +/** + * What a relationship closure returns. The pull protocol only -- a relationship + * read is the hottest path in hydration, and accepting an iterable here would + * let a producer put the per-row result object back. + */ +export type RelationshipStream = PullStream; /** * A row flowing through the pipeline, plus its relationships. @@ -14,7 +21,7 @@ export type Node = { * The stream may contain 'yield' to indicate the operator has yielded control. * See {@linkcode Operator.fetch} for more details about yields. */ - relationships: Record Stream>; + relationships: Record RelationshipStream>; }; /** @@ -135,8 +142,17 @@ export function drainStreams(node: Node | 'yield') { return; } for (const stream of Object.values(node.relationships)) { - for (const node of stream()) { - drainStreams(node); + const children = stream(); + try { + for ( + let node = children.next(); + node !== undefined; + node = children.next() + ) { + drainStreams(node); + } + } finally { + children.close(); } } } diff --git a/packages/zql/src/ivm/deferred-input.ts b/packages/zql/src/ivm/deferred-input.ts index b9ad418f12..334b1870b1 100644 --- a/packages/zql/src/ivm/deferred-input.ts +++ b/packages/zql/src/ivm/deferred-input.ts @@ -3,7 +3,7 @@ import {makeAddChange} from './change.ts'; import type {Node} from './data.ts'; import type {FetchRequest, Input, Output} from './operator.ts'; import type {SourceSchema} from './schema.ts'; -import {consume, type Stream} from './stream.ts'; +import {consume, emptyPullStream, type PullStream} from './stream.ts'; /** * A placeholder `Input` for a view whose pipeline has not been built yet. @@ -53,10 +53,8 @@ export class DeferredInput implements Input { return this.#schema; } - *fetch(req: FetchRequest): Stream { - if (this.#input) { - yield* this.#input.fetch(req); - } + fetch(req: FetchRequest): PullStream { + return this.#input ? this.#input.fetch(req) : emptyPullStream(); } destroy(): void { @@ -85,7 +83,8 @@ export class DeferredInput implements Input { } try { input.setOutput(output); - for (const node of input.fetch({})) { + const stream = input.fetch({}); + for (let node = stream.next(); node !== undefined; node = stream.next()) { if (node === 'yield') { continue; } diff --git a/packages/zql/src/ivm/exists.ts b/packages/zql/src/ivm/exists.ts index 77eaebd49e..fa6d116879 100644 --- a/packages/zql/src/ivm/exists.ts +++ b/packages/zql/src/ivm/exists.ts @@ -12,7 +12,12 @@ import { type FilterOutput, } from './filter-operators.ts'; import type {SourceSchema} from './schema.ts'; -import {type Stream} from './stream.ts'; +import { + type Stream, + emptyPullStream, + pullOf, + type PullStream, +} from './stream.ts'; /** * The Exists operator filters data based on whether or not a relationship is @@ -78,25 +83,103 @@ export class Exists implements FilterOperator { this.#output.endFilter(); } - *filter(node: Node): Generator<'yield', boolean> { - let exists: boolean | undefined; - if (!this.#noSizeReuse && !this.#inPush) { - const key = this.#getCacheKey(node, this.#parentJoinKey); - exists = this.#cache.get(key); - if (exists === undefined) { - exists = yield* this.#fetchExists(node); - this.#cache.set(key, exists); - } else if (this.#cacheHitCountsForTesting) { - this.#cacheHitCountsForTesting.set( - key, - (this.#cacheHitCountsForTesting.get(key) ?? 0) + 1, - ); + /** + * The node `filterPull` is part-way through, and where it got to. + * + * Exists is the one filter that genuinely suspends: counting a relationship + * pulls a child stream that can emit 'yield'. The generator this replaces + * held that position implicitly; here it is explicit, so no generator is + * created per node. + */ + #pending: + | { + node: Node; + key: string | undefined; + count: {stream: PullStream; size: number} | undefined; + exists: boolean | undefined; + } + | undefined; + + filter(node: Node): boolean | 'yield' { + let p = this.#pending; + if (p === undefined || p.node !== node) { + p = {node, key: undefined, count: undefined, exists: undefined}; + this.#pending = p; + if (!this.#noSizeReuse && !this.#inPush) { + const key = this.#getCacheKey(node, this.#parentJoinKey); + p.key = key; + const cached = this.#cache.get(key); + if (cached !== undefined) { + p.exists = cached; + if (this.#cacheHitCountsForTesting) { + this.#cacheHitCountsForTesting.set( + key, + (this.#cacheHitCountsForTesting.get(key) ?? 0) + 1, + ); + } + } + } + } + + if (p.exists === undefined) { + p.count ??= this.#startCount(node); + const r = this.#countStep(p.count); + if (r === 'yield') { + return 'yield'; + } + p.count.stream.close(); + p.count = undefined; + p.exists = r > 0; + if (p.key !== undefined) { + this.#cache.set(p.key, p.exists); } } - const result = - (yield* this.#filter(node, exists)) && (yield* this.#output.filter(node)); - return result; + if (!(this.#not ? !p.exists : p.exists)) { + this.#pending = undefined; + return false; + } + const out = this.#output.filter(node); + if (out === 'yield') { + return 'yield'; + } + this.#pending = undefined; + return out; + } + + /** Opens the relationship stream whose rows are being counted. */ + #startCount(node: Node): { + stream: PullStream; + size: number; + } { + const relationship = node.relationships[this.#relationshipName]; + assert( + relationship, + () => + `Exists: relationship "${this.#relationshipName}" not found on node`, + ); + return {stream: relationship(), size: 0}; + } + + /** + * Counts until the stream yields or ends: 'yield' to forward, otherwise the + * final size. Shared by `filterPull` above and the push path's `#fetchSize`, + * so the counting rule exists once. + */ + #countStep(state: { + stream: PullStream; + size: number; + }): 'yield' | number { + for (;;) { + const n = state.stream.next(); + if (n === undefined) { + return state.size; + } + if (n === 'yield') { + return 'yield'; + } + state.size++; + } } destroy(): void { @@ -150,7 +233,8 @@ export class Exists implements FilterOperator { row: change[ChangeIndex.NODE].row, relationships: { ...change[ChangeIndex.NODE].relationships, - [this.#relationshipName]: () => [], + [this.#relationshipName]: () => + emptyPullStream(), }, }), this, @@ -183,11 +267,12 @@ export class Exists implements FilterOperator { row: change[ChangeIndex.NODE].row, relationships: { ...change[ChangeIndex.NODE].relationships, - [this.#relationshipName]: () => [ - change[ChangeIndex.CHILD_DATA].change[ - ChangeIndex.NODE - ], - ], + [this.#relationshipName]: () => + pullOf([ + change[ChangeIndex.CHILD_DATA].change[ + ChangeIndex.NODE + ], + ]), }, }), this, @@ -247,20 +332,18 @@ export class Exists implements FilterOperator { } *#fetchSize(node: Node): Generator<'yield', number> { - const relationship = node.relationships[this.#relationshipName]; - assert( - relationship, - () => - `Exists: relationship "${this.#relationshipName}" not found on node`, - ); - let size = 0; - for (const n of relationship()) { - if (n === 'yield') { - yield 'yield'; - } else { - size++; + const state = this.#startCount(node); + try { + for (;;) { + const r = this.#countStep(state); + if (r === 'yield') { + yield 'yield'; + continue; + } + return r; } + } finally { + state.stream.close(); } - return size; } } diff --git a/packages/zql/src/ivm/fan-in.ts b/packages/zql/src/ivm/fan-in.ts index f0c52fa59b..89ce4c3599 100644 --- a/packages/zql/src/ivm/fan-in.ts +++ b/packages/zql/src/ivm/fan-in.ts @@ -64,8 +64,8 @@ export class FanIn implements FilterOperator { this.#output.endFilter(); } - *filter(node: Node): Generator<'yield', boolean> { - return yield* this.#output.filter(node); + filter(node: Node): boolean | 'yield' { + return this.#output.filter(node); } push(change: Change) { diff --git a/packages/zql/src/ivm/fan-out.ts b/packages/zql/src/ivm/fan-out.ts index 6da271ff39..86fe60efa0 100644 --- a/packages/zql/src/ivm/fan-out.ts +++ b/packages/zql/src/ivm/fan-out.ts @@ -60,15 +60,24 @@ export class FanOut implements FilterOperator { } } - *filter(node: Node): Generator<'yield', boolean> { - let result = false; - for (const output of this.#outputs) { - result = (yield* output.filter(node)) || result; - if (result) { + /** Which output suspended on 'yield', so re-entry resumes there. */ + #filterIndex = 0; + + filter(node: Node): boolean | 'yield' { + const outputs = this.#outputs; + for (let i = this.#filterIndex; i < outputs.length; i++) { + const r = outputs[i].filter(node); + if (r === 'yield') { + this.#filterIndex = i; + return 'yield'; + } + if (r) { + this.#filterIndex = 0; return true; } } - return result; + this.#filterIndex = 0; + return false; } *push(change: Change) { diff --git a/packages/zql/src/ivm/filter-operators.test.ts b/packages/zql/src/ivm/filter-operators.test.ts index ac8a05a7e1..2e2ed3b29c 100644 --- a/packages/zql/src/ivm/filter-operators.test.ts +++ b/packages/zql/src/ivm/filter-operators.test.ts @@ -2,16 +2,18 @@ import {describe, expect, test, vi} from 'vitest'; import {FilterStart, type FilterOutput} from './filter-operators.ts'; import type {FetchRequest, Input} from './operator.ts'; import type {SourceSchema} from './schema.ts'; +import {drainGenerator, pullOf} from './stream.ts'; describe('FilterStart', () => { test('fetch calls endFilter even if stream is not fully consumed', () => { const mockInput: Input = { setOutput: vi.fn(), - fetch: function* (_req: FetchRequest) { - yield {row: {id: 1}, relationships: {}}; - yield {row: {id: 2}, relationships: {}}; - yield {row: {id: 3}, relationships: {}}; - }, + fetch: (_req: FetchRequest) => + pullOf([ + {row: {id: 1}, relationships: {}}, + {row: {id: 2}, relationships: {}}, + {row: {id: 3}, relationships: {}}, + ]), destroy: vi.fn(), getSchema: vi.fn(() => ({}) as SourceSchema), }; @@ -19,17 +21,22 @@ describe('FilterStart', () => { const mockFilterOutput: FilterOutput = { push: vi.fn(), beginFilter: vi.fn(), - filter: filterGenerator, + filter: () => drainGenerator(filterGenerator()), endFilter: vi.fn(), }; const filterStart = new FilterStart(mockInput); filterStart.setFilterOutput(mockFilterOutput); - for (const n of filterStart.fetch({} as FetchRequest)) { - expect(n).toEqual({row: {id: 1}, relationships: {}}); - // break after consuming 1 of the 3 nodes. - break; + const stream = filterStart.fetch({} as FetchRequest); + try { + for (let n = stream.next(); n !== undefined; n = stream.next()) { + expect(n).toEqual({row: {id: 1}, relationships: {}}); + // break after consuming 1 of the 3 nodes. + break; + } + } finally { + stream.close(); } expect(mockFilterOutput.beginFilter).toHaveBeenCalledTimes(1); diff --git a/packages/zql/src/ivm/filter-operators.ts b/packages/zql/src/ivm/filter-operators.ts index de0a14edb7..254b974f8d 100644 --- a/packages/zql/src/ivm/filter-operators.ts +++ b/packages/zql/src/ivm/filter-operators.ts @@ -2,9 +2,14 @@ import type {BuilderDelegate} from '../builder/builder.ts'; import type {NoSubqueryCondition} from '../builder/filter.ts'; import type {Change} from './change.ts'; import {type Node} from './data.ts'; -import type {FetchRequest, Input, InputBase, Output} from './operator.ts'; +import { + type FetchRequest, + type Input, + type InputBase, + type Output, +} from './operator.ts'; import type {SourceSchema} from './schema.ts'; -import {type Stream} from './stream.ts'; +import {LazyPullStream, type PullStream, type Stream} from './stream.ts'; /** * The `where` clause of a ZQL query is implemented using a sub-graph of @@ -35,7 +40,14 @@ export interface FilterOutput extends Output { // nodes. E.g., so the operator can cache results for the // duration of the loop. beginFilter(): void; - filter(node: Node): Generator<'yield', boolean>; + /** + * The verdict, or 'yield' to hand control back -- the caller must then call + * again with the same node until it gets a boolean. Delegates that never + * suspend return the boolean directly, so a chain of them costs no + * allocation per node. `Exists` is the one that does suspend, and holds its + * position in explicit state rather than in a generator. + */ + filter(node: Node): boolean | 'yield'; endFilter(): void; } @@ -51,7 +63,7 @@ export const throwFilterOutput: FilterOutput = { throw new Error('Output not set'); }, - *filter(_node: Node): Generator<'yield', boolean> { + filter(): boolean | 'yield' { throw new Error('Output not set'); }, @@ -86,26 +98,15 @@ export class FilterStart implements FilterInput, Output { yield* this.#output.push(change, this); } - *fetch(req: FetchRequest): Stream { + fetch(req: FetchRequest): PullStream { const mergedFilter = mergeFilters(req.filter, this.#condition); const childReq = mergedFilter === req.filter ? req : {...req, filter: mergedFilter}; - this.#output.beginFilter(); - try { - for (const node of this.#input.fetch(childReq)) { - if (node === 'yield') { - yield node; - continue; - } - if (yield* this.#output.filter(node)) { - yield node; - } - } - } finally { - // finally is important if an exception is thrown or - // if the stream is not fully consumed. - this.#output.endFilter(); - } + // Lazy so beginFilter() runs when iteration starts, as the generator did. + return new LazyPullStream(() => { + this.#output.beginFilter(); + return new FilterStartPull(this.#input.fetch(childReq), this.#output); + }); } } @@ -134,16 +135,14 @@ export class FilterEnd implements Input, FilterOutput { input.setFilterOutput(this); } - *fetch(req: FetchRequest): Stream { - for (const node of this.#start.fetch(req)) { - yield node; - } + fetch(req: FetchRequest): PullStream { + return this.#start.fetch(req); } beginFilter() {} endFilter() {} - *filter(_node: Node) { + filter(_node: Node): boolean { return true; } @@ -178,3 +177,71 @@ export function buildFilterPipeline( delegate.addEdge(middle, filterEnd); return filterEnd; } + +/** + * FilterStart's fetch in the pull protocol. Holds the node a delegate has + * suspended on so the same node is offered again after a 'yield'; calls + * endFilter() exactly once, on exhaustion, close, or throw. + */ +class FilterStartPull implements PullStream { + readonly #input: PullStream; + readonly #output: FilterOutput; + #pending: Node | undefined; + #ended = false; + + constructor(input: PullStream, output: FilterOutput) { + this.#input = input; + this.#output = output; + } + + next(): Node | 'yield' | undefined { + if (this.#ended) { + return undefined; + } + try { + for (;;) { + let node: Node; + const pending = this.#pending; + if (pending !== undefined) { + node = pending; + } else { + const v = this.#input.next(); + if (v === undefined) { + this.#end(); + return undefined; + } + if (v === 'yield') { + return v; + } + node = v; + } + const verdict = this.#output.filter(node); + if (verdict === 'yield') { + this.#pending = node; + return 'yield'; + } + this.#pending = undefined; + if (verdict) { + return node; + } + } + } catch (e) { + this.#end(); + throw e; + } + } + + #end(): void { + if (!this.#ended) { + this.#ended = true; + this.#output.endFilter(); + } + } + + close(): void { + if (!this.#ended) { + this.#input.close(); + this.#end(); + } + } +} diff --git a/packages/zql/src/ivm/filter.ts b/packages/zql/src/ivm/filter.ts index 9273018b49..56c37909cb 100644 --- a/packages/zql/src/ivm/filter.ts +++ b/packages/zql/src/ivm/filter.ts @@ -35,8 +35,8 @@ export class Filter implements FilterOperator { this.#output.endFilter(); } - *filter(node: Node): Generator<'yield', boolean> { - return this.#predicate(node.row) && (yield* this.#output.filter(node)); + filter(node: Node): boolean | 'yield' { + return this.#predicate(node.row) && this.#output.filter(node); } setFilterOutput(output: FilterOutput) { diff --git a/packages/zql/src/ivm/flipped-join.chunked.test.ts b/packages/zql/src/ivm/flipped-join.chunked.test.ts index da6789a74e..3fd82b2f7c 100644 --- a/packages/zql/src/ivm/flipped-join.chunked.test.ts +++ b/packages/zql/src/ivm/flipped-join.chunked.test.ts @@ -15,8 +15,7 @@ import { import type {FetchRequest, Input, Output} from './operator.ts'; import {Snitch, type FetchMessage, type SnitchMessage} from './snitch.ts'; import {makeSourceChangeAdd, makeSourceChangeRemove} from './source.ts'; -import type {Stream} from './stream.ts'; -import {consume} from './stream.ts'; +import {consume, type PullStream} from './stream.ts'; import {createSource} from './test/source-factory.ts'; type CaughtRow = Exclude; @@ -193,23 +192,14 @@ test('chunked fetch propagates .return() to sub-streams on early termination', ( getSchema: () => parentInput.getSchema(), setOutput: (o: Output) => parentInput.setOutput(o), destroy: () => parentInput.destroy(), - fetch: (req: FetchRequest): Stream => { + fetch: (req: FetchRequest): PullStream => { const idx = nextStreamIdx++; const inner = parentInput.fetch(req); return { - [Symbol.iterator]() { - const it = inner[Symbol.iterator](); - const wrapped: IterableIterator = { - next: () => it.next(), - return(value?: unknown): IteratorResult { - returnCalls.push(idx); - return it.return?.(value) ?? {done: true, value: undefined}; - }, - [Symbol.iterator]() { - return wrapped; - }, - }; - return wrapped; + next: () => inner.next(), + close: () => { + returnCalls.push(idx); + inner.close(); }, }; }, @@ -219,12 +209,9 @@ test('chunked fetch propagates .return() to sub-streams on early termination', ( // Manually pull from the generator and break early, so .return() is // invoked (the for-of doesn't optimize this away). const stream = fj.fetch({}); - const it = stream[Symbol.iterator](); - const first = it.next(); - expect(first.done).toBe(false); - // Early termination — JS calls it.return() under the hood for break in - // a for-of, but here we invoke it manually. - it.return?.(); + expect(stream.next()).toBeDefined(); + // Early termination is explicit under the pull protocol. + stream.close(); // 3 chunks → 3 sub-streams. The first one we partially consumed; the // remaining 2 were primed but not advanced past their first row. All @@ -244,11 +231,27 @@ test('chunked fetch forwards yields from parent and child sub-streams', () => { getSchema: () => inner.getSchema(), setOutput: (o: Output) => inner.setOutput(o), destroy: () => inner.destroy(), - *fetch(req: FetchRequest): Stream { - for (const node of inner.fetch(req)) { - yield 'yield'; - yield node; - } + fetch(req: FetchRequest): PullStream { + const src = inner.fetch(req); + // Emit 'yield' before each row: hold the row until the marker has + // been handed back. + let pending: Node | 'yield' | undefined; + return { + next() { + if (pending !== undefined) { + const p = pending; + pending = undefined; + return p; + } + const node = src.next(); + if (node === undefined) { + return undefined; + } + pending = node; + return 'yield'; + }, + close: () => src.close(), + }; }, }; } @@ -262,7 +265,8 @@ test('chunked fetch forwards yields from parent and child sub-streams', () => { // Collect both yields and rows so we can prove yields are forwarded. const yieldsAndRows: ('yield' | string)[] = []; - for (const node of fj.fetch({})) { + const fjStream = fj.fetch({}); + for (let node = fjStream.next(); node !== undefined; node = fjStream.next()) { yieldsAndRows.push(node === 'yield' ? 'yield' : String(node.row.id)); } diff --git a/packages/zql/src/ivm/flipped-join.ts b/packages/zql/src/ivm/flipped-join.ts index efa0904453..f8c01727c7 100644 --- a/packages/zql/src/ivm/flipped-join.ts +++ b/packages/zql/src/ivm/flipped-join.ts @@ -28,7 +28,14 @@ import { type Output, } from './operator.ts'; import type {SourceSchema} from './schema.ts'; -import {type Stream} from './stream.ts'; +import { + type Stream, + emptyPullStream, + type PullStream, + LazyPullStream, + pullOf, + drainPull, +} from './stream.ts'; /** * Maximum number of entries sent in a single batched `parent.fetch` @@ -158,7 +165,17 @@ export class FlippedJoin implements Input { return this.#schema; } - *fetch(req: FetchRequest): Stream { + fetch(req: FetchRequest): PullStream { + return new LazyPullStream(() => this.#startFetch(req)); + } + + /** + * Two phases: drain the child stream (forwarding its 'yield's) to collect + * the child nodes, then hand off to the batched parent fetch built from + * them. The generator this replaces expressed the handoff as a `yield*` + * after its collection loop. + */ + #startFetch(req: FetchRequest): PullStream { // Translate constraints for the parent on parts of the join key to // constraints for the child. const childConstraint: Record = {}; @@ -173,64 +190,64 @@ export class FlippedJoin implements Input { } } - const childNodes: Node[] = []; - for (const node of this.#child.fetch( + const childStream = this.#child.fetch( hasChildConstraint ? {constraint: childConstraint} : {}, - )) { - if (node === 'yield') { - yield node; - continue; + ); + const childNodes: Node[] = []; + let collecting = true; + let inner: PullStream | undefined; + + const startBatched = () => { + childStream.close(); + // FlippedJoin's split-push change overlay logic is largely + // the same as Join's with the exception of remove. For remove, + // the change is undone here, and then re-applied to parents with order + // less than or equal to change.position below. This is necessary + // because if the removed node was the last related child, the + // related parents with position greater than change.position + // (which should not yet have the node removed), would not even + // be fetched here, and would be absent from the output all together. + if ( + this.#inprogressChildChange?.[ChangeIndex.TYPE] === ChangeType.REMOVE + ) { + const removedNode = this.#inprogressChildChange[ChangeIndex.NODE]; + const compare = this.#child.getSchema().compareRows; + const insertPos = binarySearch(childNodes.length, i => + compare(removedNode.row, childNodes[i].row), + ); + childNodes.splice(insertPos, 0, removedNode); } - childNodes.push(node); - } - - // FlippedJoin's split-push change overlay logic is largely - // the same as Join's with the exception of remove. For remove, - // the change is undone here, and then re-applied to parents with order - // less than or equal to change.position below. This is necessary - // because if the removed node was the last related child, the - // related parents with position greater than change.position - // (which should not yet have the node removed), would not even - // be fetched here, and would be absent from the output all together. - if (this.#inprogressChildChange?.[ChangeIndex.TYPE] === ChangeType.REMOVE) { - const removedNode = this.#inprogressChildChange[ChangeIndex.NODE]; - const compare = this.#child.getSchema().compareRows; - const insertPos = binarySearch(childNodes.length, i => - compare(removedNode.row, childNodes[i].row), - ); - childNodes.splice(insertPos, 0, removedNode); - } + inner = this.#fetchBatched(req, childNodes); + }; - yield* this.#fetchBatched(req, childNodes); + return { + next: (): Node | 'yield' | undefined => { + while (collecting) { + const node = childStream.next(); + if (node === undefined) { + collecting = false; + startBatched(); + break; + } + if (node === 'yield') { + return node; + } + childNodes.push(node); + } + return inner?.next(); + }, + close: () => { + collecting = false; + childStream.close(); + inner?.close(); + }, + }; } - /** - * Fetches parents for `childNodes` in batched calls, using - * `multiConstraint` so the source can issue one query per chunk (e.g. - * SQL `IN` with index-aware seek) instead of N per-child cursors. - * - * Multi-constraint values are split into chunks of `CHUNK_SIZE`, so - * SQL `IN` lists stay bounded — predictable plans, statement-cache - * hits across calls of the same chunk size, well below SQLite's - * parameter limit. - * - * Within each chunk, the source returns parents in `compareRows` order. - * Across chunks, we merge with `mergeSortedStreams` so the overall - * stream is also in order. Note: the merge primes one row from every - * chunk before yielding the first output, so all chunks open their - * cursors up front. Early termination downstream then prevents any - * further work on un-advanced chunks (cursors get `.return()`'d via - * `mergeSortedStreams`'s finally block). - * - * Replaces the previous split between `#fetchMergeSort` and - * `#fetchQuicksort`. The unique-vs-not distinction is no longer needed: - * the source handles cardinality (single index seek for each value) and - * ordering (SQL `ORDER BY` / index walk). - */ - *#fetchBatched( + #fetchBatched( req: FetchRequest, childNodes: Node[], - ): Stream { + ): PullStream { const parentReqConstraint = req.constraint; const parentKey = this.#parentKey; const childKey = this.#childKey; @@ -263,7 +280,7 @@ export class FlippedJoin implements Input { } if (computedMulti.length === 0) { - return; + return emptyPullStream(); } // Source returns parents in compareRows order within each chunk. @@ -285,36 +302,47 @@ export class FlippedJoin implements Input { }) : this.#fetchChunked(req, incoming, computedMulti, compare); - for (const node of parentStream) { - if (node === 'yield') { - yield 'yield'; - continue; - } - const key = canonicalKey(node.row, parentKey); - const idxs = childIndexesByKey.get(key); - if (idxs === undefined) { - // This row's parent-key doesn't match any of our computed - // multi-constraint entries. Happens when our parent is an - // intermediate operator (e.g. a chained FlippedJoin) that passes - // multiConstraints through unchanged instead of filtering — see - // FetchRequest.multiConstraints contract. The lookup miss here - // performs the required filter, so just skip the row. - continue; - } - // Children retain their original input order within the group - // because we appended to `idxs` in iteration order. - const relatedChildNodes: Node[] = idxs.map(i => childNodes[i]); - yield* this.#yieldParentWithOverlay(node, relatedChildNodes); - } + return { + next: (): Node | 'yield' | undefined => { + for (;;) { + const node = parentStream.next(); + if (node === undefined) { + return undefined; + } + if (node === 'yield') { + return 'yield'; + } + const key = canonicalKey(node.row, parentKey); + const idxs = childIndexesByKey.get(key); + if (idxs === undefined) { + // This row's parent-key doesn't match any of our computed + // multi-constraint entries. Happens when our parent is an + // intermediate operator (e.g. a chained FlippedJoin) that passes + // multiConstraints through unchanged instead of filtering — see + // FetchRequest.multiConstraints contract. The lookup miss here + // performs the required filter, so just skip the row. + continue; + } + // Children retain their original input order within the group + // because we appended to `idxs` in iteration order. + const relatedChildNodes: Node[] = idxs.map(i => childNodes[i]); + const parent = this.#parentWithOverlay(node, relatedChildNodes); + if (parent !== undefined) { + return parent; + } + } + }, + close: () => parentStream.close(), + }; } - *#fetchChunked( + #fetchChunked( req: FetchRequest, incomingMultis: readonly MultiConstraint[], computedMulti: MultiConstraint, compare: (a: Node, b: Node) => number, - ): Stream { - const chunkStreams: Stream[] = []; + ): PullStream { + const chunkStreams: PullStream[] = []; for (let i = 0; i < computedMulti.length; i += multiConstraintChunkSize) { chunkStreams.push( this.#parent.fetch({ @@ -326,13 +354,14 @@ export class FlippedJoin implements Input { }), ); } - yield* mergeSortedStreams(chunkStreams, compare); + return mergeSortedStreams(chunkStreams, compare); } - *#yieldParentWithOverlay( + /** The parent with its overlaid children, or undefined if none remain. */ + #parentWithOverlay( minParentNode: Node, relatedChildNodes: Node[], - ): Stream { + ): Node | undefined { let overlaidRelatedChildNodes = relatedChildNodes; if ( this.#inprogressChildChange && @@ -360,26 +389,27 @@ export class FlippedJoin implements Input { ); } } else if (!hasInprogressChildChangeBeenPushedForMinParentNode) { - overlaidRelatedChildNodes = [ - ...generateWithOverlayNoYield( - relatedChildNodes, + overlaidRelatedChildNodes = drainPull( + generateWithOverlayNoYield( + pullOf(relatedChildNodes), this.#inprogressChildChange, this.#child.getSchema(), ), - ]; + ); } } - // yield node if after the overlay it still has relationship nodes - if (overlaidRelatedChildNodes.length > 0) { - yield { - ...minParentNode, - relationships: { - ...minParentNode.relationships, - [this.#relationshipName]: () => overlaidRelatedChildNodes, - }, - }; + // emit the node if after the overlay it still has relationship nodes + if (overlaidRelatedChildNodes.length === 0) { + return undefined; } + return { + ...minParentNode, + relationships: { + ...minParentNode.relationships, + [this.#relationshipName]: () => pullOf(overlaidRelatedChildNodes), + }, + }; } *#pushChild(change: Change): Stream<'yield'> { @@ -417,70 +447,95 @@ export class FlippedJoin implements Input { ); const parentNodeStream = constraint ? this.#parent.fetch({constraint}) - : []; - for (const parentNode of parentNodeStream) { - if (parentNode === 'yield') { - yield 'yield'; - continue; - } - this.#inprogressChildChange = change; - this.#inprogressChildChangePosition = parentNode.row; - const childNodeStream = () => { - const constraint = buildJoinConstraint( - parentNode.row, - this.#parentKey, - this.#childKey, - ); - return constraint ? this.#child.fetch({constraint}) : []; - }; - if (!exists) { - for (const childNode of childNodeStream()) { - if (childNode === 'yield') { - yield 'yield'; - continue; - } - if ( - this.#child - .getSchema() - .compareRows(childNode.row, change[ChangeIndex.NODE].row) !== 0 - ) { - exists = true; - break; + : emptyPullStream(); + + const parents = parentNodeStream; + try { + for ( + let parentNode = parents.next(); + parentNode !== undefined; + parentNode = parents.next() + ) { + if (parentNode === 'yield') { + yield 'yield'; + continue; + } + this.#inprogressChildChange = change; + this.#inprogressChildChangePosition = parentNode.row; + const childNodeStream = () => { + const constraint = buildJoinConstraint( + parentNode.row, + this.#parentKey, + this.#childKey, + ); + return constraint + ? this.#child.fetch({constraint}) + : emptyPullStream(); + }; + if (!exists) { + const children = childNodeStream(); + try { + for ( + let childNode = children.next(); + childNode !== undefined; + childNode = children.next() + ) { + if (childNode === 'yield') { + yield 'yield'; + continue; + } + if ( + this.#child + .getSchema() + .compareRows( + childNode.row, + change[ChangeIndex.NODE].row, + ) !== 0 + ) { + exists = true; + break; + } + } + } finally { + children.close(); } } - } - if (exists) { - yield* this.#output.push( - makeChildChange( - { - ...parentNode, - relationships: { - ...parentNode.relationships, - [this.#relationshipName]: childNodeStream, + if (exists) { + yield* this.#output.push( + makeChildChange( + { + ...parentNode, + relationships: { + ...parentNode.relationships, + [this.#relationshipName]: childNodeStream, + }, }, + { + relationshipName: this.#relationshipName, + change, + }, + ), + this, + ); + } else { + const newNode = { + ...parentNode, + relationships: { + ...parentNode.relationships, + [this.#relationshipName]: () => + pullOf([change[ChangeIndex.NODE]]), }, - { - relationshipName: this.#relationshipName, - change, - }, - ), - this, - ); - } else { - const newNode = { - ...parentNode, - relationships: { - ...parentNode.relationships, - [this.#relationshipName]: () => [change[ChangeIndex.NODE]], - }, - }; - yield* this.#output.push( - change[ChangeIndex.TYPE] === ChangeType.ADD - ? makeAddChange(newNode) - : makeRemoveChange(newNode), - this, - ); + }; + yield* this.#output.push( + change[ChangeIndex.TYPE] === ChangeType.ADD + ? makeAddChange(newNode) + : makeRemoveChange(newNode), + this, + ); + } } + } finally { + parents.close(); } } finally { this.#inprogressChildChange = undefined; @@ -494,7 +549,9 @@ export class FlippedJoin implements Input { this.#parentKey, this.#childKey, ); - return constraint ? this.#child.fetch({constraint}) : []; + return constraint + ? this.#child.fetch({constraint}) + : emptyPullStream(); }; const flip = (node: Node) => ({ @@ -507,14 +564,24 @@ export class FlippedJoin implements Input { // If no related child don't push as this is an inner join. let hasRelatedChild = false; - for (const node of childNodeStream(change[ChangeIndex.NODE])()) { - if (node === 'yield') { - yield 'yield'; - continue; - } else { - hasRelatedChild = true; - break; + + const children = childNodeStream(change[ChangeIndex.NODE])(); + try { + for ( + let node = children.next(); + node !== undefined; + node = children.next() + ) { + if (node === 'yield') { + yield 'yield'; + continue; + } else { + hasRelatedChild = true; + break; + } } + } finally { + children.close(); } if (!hasRelatedChild) { return; diff --git a/packages/zql/src/ivm/join-utils.test.ts b/packages/zql/src/ivm/join-utils.test.ts index 8d39d1b7d8..3491839bef 100644 --- a/packages/zql/src/ivm/join-utils.test.ts +++ b/packages/zql/src/ivm/join-utils.test.ts @@ -16,7 +16,7 @@ import { rowEqualsForCompoundKey, } from './join-utils.ts'; import type {SourceSchema} from './schema.ts'; -import type {Stream} from './stream.ts'; +import {drainPull, pullOf, type PullStream} from './stream.ts'; function makeNode(row: Row): Node { return {row, relationships: {}}; @@ -34,14 +34,16 @@ function makeSchema(primaryKey: readonly [string, ...string[]]): SourceSchema { }; } -function collectNodes(stream: Stream): (Node | 'yield')[] { - return [...stream]; +function collectNodes(stream: PullStream): (Node | 'yield')[] { + return drainPull(stream); } function collectRows( - stream: Stream, + stream: PullStream, ): Record[] { - return [...stream].filter((n): n is Node => n !== 'yield').map(n => n.row); + return drainPull(stream) + .filter((n): n is Node => n !== 'yield') + .map(n => n.row); } describe('generateWithOverlayUnordered', () => { @@ -49,24 +51,21 @@ describe('generateWithOverlayUnordered', () => { describe('remove', () => { test('yields overlay node first then all stream nodes', () => { - const stream: Stream = [ - makeNode({id: 1}), - makeNode({id: 2}), - ]; + const stream: (Node | 'yield')[] = [makeNode({id: 1}), makeNode({id: 2})]; const overlay: Change = makeRemoveChange(makeNode({id: 3})); const result = collectRows( - generateWithOverlayUnordered(stream, overlay, schema), + generateWithOverlayUnordered(pullOf(stream), overlay, schema), ); expect(result).toEqual([{id: 3}, {id: 1}, {id: 2}]); }); test('does not assert when overlay node is not in stream', () => { - const stream: Stream = []; + const stream: (Node | 'yield')[] = []; const overlay: Change = makeRemoveChange(makeNode({id: 99})); const result = collectRows( - generateWithOverlayUnordered(stream, overlay, schema), + generateWithOverlayUnordered(pullOf(stream), overlay, schema), ); expect(result).toEqual([{id: 99}]); }); @@ -74,7 +73,7 @@ describe('generateWithOverlayUnordered', () => { describe('add', () => { test('suppresses matching node from stream', () => { - const stream: Stream = [ + const stream: (Node | 'yield')[] = [ makeNode({id: 1}), makeNode({id: 2}), makeNode({id: 3}), @@ -82,17 +81,19 @@ describe('generateWithOverlayUnordered', () => { const overlay: Change = makeAddChange(makeNode({id: 2})); const result = collectRows( - generateWithOverlayUnordered(stream, overlay, schema), + generateWithOverlayUnordered(pullOf(stream), overlay, schema), ); expect(result).toEqual([{id: 1}, {id: 3}]); }); test('asserts if no matching node found in stream', () => { - const stream: Stream = [makeNode({id: 1})]; + const stream: (Node | 'yield')[] = [makeNode({id: 1})]; const overlay: Change = makeAddChange(makeNode({id: 99})); expect(() => - collectNodes(generateWithOverlayUnordered(stream, overlay, schema)), + collectNodes( + generateWithOverlayUnordered(pullOf(stream), overlay, schema), + ), ).toThrow( 'overlayGenerator: overlay was never applied to any fetched node', ); @@ -101,7 +102,7 @@ describe('generateWithOverlayUnordered', () => { describe('edit', () => { test('yields old node first and suppresses matching node from stream', () => { - const stream: Stream = [ + const stream: (Node | 'yield')[] = [ makeNode({id: 1}), makeNode({id: 2, val: 'new'}), ]; @@ -111,20 +112,22 @@ describe('generateWithOverlayUnordered', () => { ); const result = collectRows( - generateWithOverlayUnordered(stream, overlay, schema), + generateWithOverlayUnordered(pullOf(stream), overlay, schema), ); expect(result).toEqual([{id: 2, val: 'old'}, {id: 1}]); }); test('asserts if no matching node found in stream', () => { - const stream: Stream = [makeNode({id: 1})]; + const stream: (Node | 'yield')[] = [makeNode({id: 1})]; const overlay: Change = makeEditChange( makeNode({id: 99}), makeNode({id: 99}), ); expect(() => - collectNodes(generateWithOverlayUnordered(stream, overlay, schema)), + collectNodes( + generateWithOverlayUnordered(pullOf(stream), overlay, schema), + ), ).toThrow( 'overlayGenerator: overlay was never applied to any fetched node', ); @@ -139,15 +142,12 @@ describe('generateWithOverlayUnordered', () => { relationships: {items: childSchema}, }; - const stream: Stream = [ + const stream: (Node | 'yield')[] = [ makeNode({id: 1}), { row: {id: 2}, relationships: { - items: function* () { - yield makeNode({cid: 'a'}); - yield makeNode({cid: 'b'}); - }, + items: () => pullOf([makeNode({cid: 'a'}), makeNode({cid: 'b'})]), }, }, ]; @@ -159,7 +159,7 @@ describe('generateWithOverlayUnordered', () => { }); const result = collectNodes( - generateWithOverlayUnordered(stream, overlay, schemaWithRel), + generateWithOverlayUnordered(pullOf(stream), overlay, schemaWithRel), ); expect(result).toHaveLength(2); // First node passes through unchanged @@ -177,7 +177,7 @@ describe('generateWithOverlayUnordered', () => { relationships: {items: makeSchema(['cid'])}, }; - const stream: Stream = [makeNode({id: 1})]; + const stream: (Node | 'yield')[] = [makeNode({id: 1})]; const overlay: Change = makeChildChange(makeNode({id: 99}), { relationshipName: 'items', change: makeAddChange(makeNode({cid: 'c'})), @@ -185,7 +185,7 @@ describe('generateWithOverlayUnordered', () => { expect(() => collectNodes( - generateWithOverlayUnordered(stream, overlay, schemaWithRel), + generateWithOverlayUnordered(pullOf(stream), overlay, schemaWithRel), ), ).toThrow( 'overlayGenerator: overlay was never applied to any fetched node', @@ -197,7 +197,7 @@ describe('generateWithOverlayUnordered', () => { const compoundSchema = makeSchema(['a', 'b']); test('matches on all PK columns', () => { - const stream: Stream = [ + const stream: (Node | 'yield')[] = [ makeNode({a: 1, b: 1, val: 'x'}), makeNode({a: 1, b: 2, val: 'y'}), makeNode({a: 2, b: 1, val: 'z'}), @@ -205,7 +205,7 @@ describe('generateWithOverlayUnordered', () => { const overlay: Change = makeAddChange(makeNode({a: 1, b: 2})); const result = collectRows( - generateWithOverlayUnordered(stream, overlay, compoundSchema), + generateWithOverlayUnordered(pullOf(stream), overlay, compoundSchema), ); expect(result).toEqual([ {a: 1, b: 1, val: 'x'}, @@ -214,7 +214,7 @@ describe('generateWithOverlayUnordered', () => { }); test('does not match on partial PK', () => { - const stream: Stream = [ + const stream: (Node | 'yield')[] = [ makeNode({a: 1, b: 1}), makeNode({a: 1, b: 2}), ]; @@ -223,7 +223,7 @@ describe('generateWithOverlayUnordered', () => { expect(() => collectNodes( - generateWithOverlayUnordered(stream, overlay, compoundSchema), + generateWithOverlayUnordered(pullOf(stream), overlay, compoundSchema), ), ).toThrow( 'overlayGenerator: overlay was never applied to any fetched node', @@ -233,7 +233,7 @@ describe('generateWithOverlayUnordered', () => { describe('yield markers', () => { test('passes yield markers through unchanged', () => { - const stream: Stream = [ + const stream: (Node | 'yield')[] = [ makeNode({id: 1}), 'yield' as const, makeNode({id: 2}), @@ -243,7 +243,7 @@ describe('generateWithOverlayUnordered', () => { const overlay: Change = makeAddChange(makeNode({id: 2})); const result = collectNodes( - generateWithOverlayUnordered(stream, overlay, schema), + generateWithOverlayUnordered(pullOf(stream), overlay, schema), ); expect(result).toEqual([ expect.objectContaining({row: {id: 1}}), @@ -259,16 +259,13 @@ describe('generateWithOverlayNoYieldUnordered', () => { const schema = makeSchema(['id']); test('strips yield markers from output', () => { - function* stream(): Stream { - yield makeNode({id: 1}); - yield makeNode({id: 2}); - yield makeNode({id: 3}); - } + const stream = (): PullStream => + pullOf([makeNode({id: 1}), makeNode({id: 2}), makeNode({id: 3})]); const overlay: Change = makeAddChange(makeNode({id: 2})); - const result = [ - ...generateWithOverlayNoYieldUnordered(stream(), overlay, schema), - ]; + const result = drainPull( + generateWithOverlayNoYieldUnordered(stream(), overlay, schema), + ); expect(result).toHaveLength(2); expect(result.map(n => n.row)).toEqual([{id: 1}, {id: 3}]); }); diff --git a/packages/zql/src/ivm/join-utils.ts b/packages/zql/src/ivm/join-utils.ts index 78b638fef8..e571063d66 100644 --- a/packages/zql/src/ivm/join-utils.ts +++ b/packages/zql/src/ivm/join-utils.ts @@ -6,66 +6,89 @@ import {ChangeType} from './change-type.ts'; import type {Change} from './change.ts'; import {compareValues, valuesEqual, type Node} from './data.ts'; import type {SourceSchema} from './schema.ts'; -import type {Stream} from './stream.ts'; +import {type PullStream} from './stream.ts'; export function generateWithOverlayNoYield( - stream: Stream, + stream: PullStream, overlay: Change, schema: SourceSchema, -): Stream { - return generateWithOverlay(stream, overlay, schema) as Stream; +): PullStream { + return generateWithOverlay( + stream as PullStream, + overlay, + schema, + ) as PullStream; } -export function* generateWithOverlay( - stream: Stream, - overlay: Change, - schema: SourceSchema, -): Stream { - let applied = false; - let editOldApplied = false; - let editNewApplied = false; - for (const node of stream) { - if (node === 'yield') { - yield node; - continue; - } +/** + * Splices a pending change into a node stream, as a pull stream. + * + * One input node can produce two outputs -- the overlay and the node itself -- + * which the generator expressed as two `yield`s in one loop iteration. `#q` + * holds those so `next()` can hand them back one at a time. + */ +class JoinOverlay implements PullStream { + readonly #stream: PullStream; + readonly #overlay: Change; + readonly #schema: SourceSchema; + readonly #q: (Node | 'yield')[] = []; + #applied = false; + #editOldApplied = false; + #editNewApplied = false; + #exhausted = false; + #tailDone = false; + + constructor( + stream: PullStream, + overlay: Change, + schema: SourceSchema, + ) { + this.#stream = stream; + this.#overlay = overlay; + this.#schema = schema; + } + + #step(node: Node): void { + const overlay = this.#overlay; + const schema = this.#schema; + const q = this.#q; let yieldNode = true; - if (!applied) { + if (!this.#applied) { switch (overlay[ChangeIndex.TYPE]) { case ChangeType.ADD: { if ( schema.compareRows(overlay[ChangeIndex.NODE].row, node.row) === 0 ) { - applied = true; + this.#applied = true; yieldNode = false; } break; } case ChangeType.REMOVE: { if (schema.compareRows(overlay[ChangeIndex.NODE].row, node.row) < 0) { - applied = true; - yield overlay[ChangeIndex.NODE]; + this.#applied = true; + q.push(overlay[ChangeIndex.NODE]); } break; } case ChangeType.EDIT: { if ( - !editOldApplied && + !this.#editOldApplied && schema.compareRows(overlay[ChangeIndex.OLD_NODE].row, node.row) < 0 ) { - editOldApplied = true; - if (editNewApplied) { - applied = true; + this.#editOldApplied = true; + if (this.#editNewApplied) { + this.#applied = true; } - yield overlay[ChangeIndex.OLD_NODE]; + q.push(overlay[ChangeIndex.OLD_NODE]); } if ( - !editNewApplied && + !this.#editNewApplied && schema.compareRows(overlay[ChangeIndex.NODE].row, node.row) === 0 ) { - editNewApplied = true; - if (editOldApplied) { - applied = true; + this.#editNewApplied = true; + if (this.#editOldApplied) { + this.#applied = true; } yieldNode = false; } @@ -75,8 +98,8 @@ export function* generateWithOverlay( if ( schema.compareRows(overlay[ChangeIndex.NODE].row, node.row) === 0 ) { - applied = true; - yield { + this.#applied = true; + q.push({ row: node.row, relationships: { ...node.relationships, @@ -91,7 +114,7 @@ export function* generateWithOverlay( ], ), }, - }; + }); yieldNode = false; } break; @@ -99,58 +122,109 @@ export function* generateWithOverlay( } } if (yieldNode) { - yield node; + q.push(node); + } + } + + #tail(): void { + const overlay = this.#overlay; + if (!this.#applied) { + if (overlay[ChangeIndex.TYPE] === ChangeType.REMOVE) { + this.#applied = true; + this.#q.push(overlay[ChangeIndex.NODE]); + } else if (overlay[ChangeIndex.TYPE] === ChangeType.EDIT) { + assert( + this.#editNewApplied, + 'edit overlay: new node must be applied before old node', + ); + this.#editOldApplied = true; + this.#applied = true; + this.#q.push(overlay[ChangeIndex.OLD_NODE]); + } } + assert( + this.#applied, + 'overlayGenerator: overlay was never applied to any fetched node', + ); } - if (!applied) { - if (overlay[ChangeIndex.TYPE] === ChangeType.REMOVE) { - applied = true; - yield overlay[ChangeIndex.NODE]; - } else if (overlay[ChangeIndex.TYPE] === ChangeType.EDIT) { - assert( - editNewApplied, - 'edit overlay: new node must be applied before old node', - ); - editOldApplied = true; - applied = true; - yield overlay[ChangeIndex.OLD_NODE]; + + next(): Node | 'yield' | undefined { + for (;;) { + if (this.#q.length > 0) { + return this.#q.shift(); + } + if (this.#exhausted) { + if (!this.#tailDone) { + this.#tailDone = true; + this.#tail(); + continue; + } + return undefined; + } + const node = this.#stream.next(); + if (node === undefined) { + this.#exhausted = true; + continue; + } + if (node === 'yield') { + return node; + } + this.#step(node); } } - assert( - applied, - 'overlayGenerator: overlay was never applied to any fetched node', - ); + close(): void { + this.#exhausted = true; + this.#tailDone = true; + this.#q.length = 0; + this.#stream.close(); + } } -export function generateWithOverlayNoYieldUnordered( - stream: Stream, +export function generateWithOverlay( + stream: PullStream, overlay: Change, schema: SourceSchema, -): Stream { - return generateWithOverlayUnordered(stream, overlay, schema) as Stream; +): PullStream { + return new JoinOverlay(stream, overlay, schema); } -export function* generateWithOverlayUnordered( - stream: Stream, +export function generateWithOverlayNoYieldUnordered( + stream: PullStream, overlay: Change, schema: SourceSchema, -): Stream { - // Eager inject - if (overlay[ChangeIndex.TYPE] === ChangeType.REMOVE) { - yield overlay[ChangeIndex.NODE]; - } else if (overlay[ChangeIndex.TYPE] === ChangeType.EDIT) { - yield overlay[ChangeIndex.OLD_NODE]; +): PullStream { + return generateWithOverlayUnordered( + stream as PullStream, + overlay, + schema, + ) as PullStream; +} + +/** {@link JoinOverlay} for unordered streams: eager inject, inline suppress. */ +class JoinOverlayUnordered implements PullStream { + readonly #stream: PullStream; + readonly #overlay: Change; + readonly #schema: SourceSchema; + readonly #q: (Node | 'yield')[] = []; + #injected = false; + #suppressed = false; + #done = false; + + constructor( + stream: PullStream, + overlay: Change, + schema: SourceSchema, + ) { + this.#stream = stream; + this.#overlay = overlay; + this.#schema = schema; } - // Stream with inline suppress - let suppressed = false; - for (const node of stream) { - if (node === 'yield') { - yield node; - continue; - } - if (!suppressed) { + #step(node: Node): void { + const overlay = this.#overlay; + const schema = this.#schema; + if (!this.#suppressed) { if ( overlay[ChangeIndex.TYPE] === ChangeType.ADD || overlay[ChangeIndex.TYPE] === ChangeType.EDIT @@ -162,8 +236,8 @@ export function* generateWithOverlayUnordered( schema.primaryKey, ) ) { - suppressed = true; - continue; + this.#suppressed = true; + return; } } if (overlay[ChangeIndex.TYPE] === ChangeType.CHILD) { @@ -174,8 +248,8 @@ export function* generateWithOverlayUnordered( schema.primaryKey, ) ) { - suppressed = true; - yield { + this.#suppressed = true; + this.#q.push({ row: node.row, relationships: { ...node.relationships, @@ -190,17 +264,61 @@ export function* generateWithOverlayUnordered( ], ), }, - }; - continue; + }); + return; } } } - yield node; + this.#q.push(node); } - assert( - suppressed || overlay[ChangeIndex.TYPE] === ChangeType.REMOVE, - 'overlayGenerator: overlay was never applied to any fetched node', - ); + + next(): Node | 'yield' | undefined { + if (!this.#injected) { + this.#injected = true; + const overlay = this.#overlay; + if (overlay[ChangeIndex.TYPE] === ChangeType.REMOVE) { + this.#q.push(overlay[ChangeIndex.NODE]); + } else if (overlay[ChangeIndex.TYPE] === ChangeType.EDIT) { + this.#q.push(overlay[ChangeIndex.OLD_NODE]); + } + } + for (;;) { + if (this.#q.length > 0) { + return this.#q.shift(); + } + if (this.#done) { + return undefined; + } + const node = this.#stream.next(); + if (node === undefined) { + this.#done = true; + assert( + this.#suppressed || + this.#overlay[ChangeIndex.TYPE] === ChangeType.REMOVE, + 'overlayGenerator: overlay was never applied to any fetched node', + ); + return undefined; + } + if (node === 'yield') { + return node; + } + this.#step(node); + } + } + + close(): void { + this.#done = true; + this.#q.length = 0; + this.#stream.close(); + } +} + +export function generateWithOverlayUnordered( + stream: PullStream, + overlay: Change, + schema: SourceSchema, +): PullStream { + return new JoinOverlayUnordered(stream, overlay, schema); } export function rowEqualsForCompoundKey( diff --git a/packages/zql/src/ivm/join.ts b/packages/zql/src/ivm/join.ts index df1f88361d..4078f8b437 100644 --- a/packages/zql/src/ivm/join.ts +++ b/packages/zql/src/ivm/join.ts @@ -10,7 +10,7 @@ import { makeRemoveChange, type Change, } from './change.ts'; -import type {Node} from './data.ts'; +import type {Node, RelationshipStream} from './data.ts'; import { buildJoinConstraint, generateWithOverlay, @@ -25,7 +25,12 @@ import { type Output, } from './operator.ts'; import type {SourceSchema} from './schema.ts'; -import {type Stream} from './stream.ts'; +import { + emptyPullStream, + type PullStream, + type Stream, + mapPull, +} from './stream.ts'; type Args = { parent: Input; @@ -116,14 +121,15 @@ export class Join implements Input { return this.#schema; } - *fetch(req: FetchRequest): Stream { - for (const parentNode of this.#parent.fetch(req)) { - if (parentNode === 'yield') { - yield parentNode; - continue; - } - yield this.#processParentNode(parentNode.row, parentNode.relationships); - } + fetch(req: FetchRequest): PullStream { + // The parent spine in the pull protocol. Child lookups behind each node's + // relationships still go through `fetch`; converting those means + // converting what consumes relationships. + return mapPull(this.#parent.fetch(req), parentNode => + parentNode === 'yield' + ? parentNode + : this.#processParentNode(parentNode.row, parentNode.relationships), + ); } *#pushParent(change: Change): Stream<'yield'> { @@ -228,20 +234,29 @@ export class Join implements Input { this.#parentKey, ); if (constraint) { - for (const parentNode of this.#parent.fetch({constraint})) { - if (parentNode === 'yield') { - yield parentNode; - continue; + const parents = this.#parent.fetch({constraint}); + try { + for ( + let parentNode = parents.next(); + parentNode !== undefined; + parentNode = parents.next() + ) { + if (parentNode === 'yield') { + yield parentNode; + continue; + } + this.#inprogressChildChangePosition = parentNode.row; + const childChange = makeChildChange( + this.#processParentNode(parentNode.row, parentNode.relationships), + { + relationshipName: this.#relationshipName, + change, + }, + ); + yield* this.#output.push(childChange, this); } - this.#inprogressChildChangePosition = parentNode.row; - const childChange = makeChildChange( - this.#processParentNode(parentNode.row, parentNode.relationships), - { - relationshipName: this.#relationshipName, - change, - }, - ); - yield* this.#output.push(childChange, this); + } finally { + parents.close(); } } } finally { @@ -251,7 +266,7 @@ export class Join implements Input { #processParentNode( parentNodeRow: Row, - parentNodeRelations: Record Stream>, + parentNodeRelations: Record RelationshipStream>, ): Node { const childStream = () => { const constraint = buildJoinConstraint( @@ -259,7 +274,9 @@ export class Join implements Input { this.#parentKey, this.#childKey, ); - const stream = constraint ? this.#child.fetch({constraint}) : []; + const stream = constraint + ? this.#child.fetch({constraint}) + : emptyPullStream(); if ( this.#inprogressChildChange && @@ -276,18 +293,17 @@ export class Join implements Input { ) > 0 ) { const childSchema = this.#child.getSchema(); - if (childSchema.sort === undefined) { - return generateWithOverlayUnordered( - stream, - this.#inprogressChildChange, - childSchema, - ); - } - return generateWithOverlay( - stream, - this.#inprogressChildChange, - childSchema, - ); + return childSchema.sort === undefined + ? generateWithOverlayUnordered( + stream, + this.#inprogressChildChange, + childSchema, + ) + : generateWithOverlay( + stream, + this.#inprogressChildChange, + childSchema, + ); } return stream; }; diff --git a/packages/zql/src/ivm/memory-source.test.ts b/packages/zql/src/ivm/memory-source.test.ts index 22a7998bac..d6f05dec73 100644 --- a/packages/zql/src/ivm/memory-source.test.ts +++ b/packages/zql/src/ivm/memory-source.test.ts @@ -22,8 +22,8 @@ import { type Overlay, } from './memory-source.ts'; import type {MultiConstraint} from './operator.ts'; -import type {Stream} from './stream.ts'; -import {consume} from './stream.ts'; +import type {PullStream} from './stream.ts'; +import {consume, drainPull, pullOf} from './stream.ts'; import {compareRowsTest} from './test/compare-rows-test.ts'; import {createSource} from './test/source-factory.ts'; @@ -217,7 +217,7 @@ test('fetch during push edit change', () => { row: {a: 'a', b: 'b', c: 'c'}, relationships: {}, }); - fetchDuringPush = [...conn.fetch({})]; + fetchDuringPush = drainPull(conn.fetch({})); return emptyArray; }, }); @@ -258,8 +258,8 @@ describe('fetch with req.filter', () => { consume(ms.push(makeSourceChangeAdd({a: 'a3', b: 'x'}))); const conn = ms.connect([['a', 'asc']]); - const rows = [ - ...conn.fetch({ + const rows = drainPull( + conn.fetch({ filter: { type: 'simple', op: '=', @@ -267,7 +267,7 @@ describe('fetch with req.filter', () => { right: {type: 'literal', value: 'x'}, }, }), - ].filter(n => n !== 'yield'); + ).filter(n => n !== 'yield'); expect(rows.map(n => n.row)).toEqual([ {a: 'a1', b: 'x'}, @@ -289,8 +289,8 @@ describe('fetch with req.filter', () => { } const conn = ms.connect([['a', 'asc']]); - const rows = [ - ...conn.fetch({ + const rows = drainPull( + conn.fetch({ filter: { type: 'simple', op: '=', @@ -298,7 +298,7 @@ describe('fetch with req.filter', () => { right: {type: 'literal', value: 'id-42'}, }, }), - ].filter(n => n !== 'yield'); + ).filter(n => n !== 'yield'); expect(rows.map(n => n.row)).toEqual([{a: 'id-42', b: 'val-42'}]); conn.destroy(); @@ -327,8 +327,8 @@ describe('fetch with req.filter', () => { right: {type: 'literal', value: 'x'}, }); - const rows = [ - ...conn.fetch({ + const rows = drainPull( + conn.fetch({ filter: { type: 'simple', op: '=', @@ -336,7 +336,7 @@ describe('fetch with req.filter', () => { right: {type: 'literal', value: 'p'}, }, }), - ].filter(n => n !== 'yield'); + ).filter(n => n !== 'yield'); expect(rows.map(n => n.row)).toEqual([{a: '1', b: 'x', c: 'p'}]); conn.destroy(); @@ -357,8 +357,8 @@ describe('fetch with req.filter', () => { consume(ms.push(makeSourceChangeAdd({a: 'a5', b: 'x'}))); const conn = ms.connect([['a', 'asc']]); - const rows = [ - ...conn.fetch({ + const rows = drainPull( + conn.fetch({ reverse: true, filter: { type: 'simple', @@ -367,7 +367,7 @@ describe('fetch with req.filter', () => { right: {type: 'literal', value: 'x'}, }, }), - ].filter(n => n !== 'yield'); + ).filter(n => n !== 'yield'); expect(rows.map(n => n.row)).toEqual([ {a: 'a5', b: 'x'}, @@ -392,8 +392,8 @@ describe('fetch with req.filter', () => { consume(ms.push(makeSourceChangeAdd({a: 'a5', b: 'x'}))); const conn = ms.connect([['a', 'asc']]); - const rows = [ - ...conn.fetch({ + const rows = drainPull( + conn.fetch({ start: {row: {a: 'a2', b: 'y'}, basis: 'after'}, filter: { type: 'simple', @@ -402,7 +402,7 @@ describe('fetch with req.filter', () => { right: {type: 'literal', value: 'x'}, }, }), - ].filter(n => n !== 'yield'); + ).filter(n => n !== 'yield'); expect(rows.map(n => n.row)).toEqual([ {a: 'a3', b: 'x'}, @@ -430,8 +430,8 @@ describe('fetch with req.filter', () => { consume(ms.push(makeSourceChangeAdd({a: 'a4', b: 'x'}))); const conn = ms.connect([['a', 'asc']]); - const rows = [ - ...conn.fetch({ + const rows = drainPull( + conn.fetch({ multiConstraints: [[{a: 'a1'}, {a: 'a2'}, {a: 'a3'}]], filter: { type: 'simple', @@ -440,7 +440,7 @@ describe('fetch with req.filter', () => { right: {type: 'literal', value: 'x'}, }, }), - ].filter(n => n !== 'yield'); + ).filter(n => n !== 'yield'); // a1 (IN-list ✓, b=x ✓), a2 (IN-list ✓, b=y ✗), a3 (IN-list ✓, b=x ✓), // a4 (IN-list ✗) → expect [a1, a3]. @@ -472,7 +472,7 @@ describe('fetch with req.filter during push (overlay)', () => { let captured: (Node | 'yield')[] = []; conn.setOutput({ push(_change: Change) { - captured = [...conn.fetch({filter: bEqX})]; + captured = drainPull(conn.fetch({filter: bEqX})); return emptyArray; }, }); @@ -869,8 +869,8 @@ describe('generateWithOverlayInner', () => { expected: [rows[0], rows[1], {id: 2.5, s: 'c', n: 33}], }, ] as const)('$name', ({overlays, expected}) => { - const actual = generateWithOverlayInner(rows, overlays, compare); - expect(Array.from(actual, ({row}) => row)).toEqual(expected); + const actual = generateWithOverlayInner(pullOf(rows), overlays, compare); + expect(drainPull(actual).map(({row}) => row)).toEqual(expected); }); }); @@ -1115,10 +1115,9 @@ describe('generateWithOverlayInnerUnordered', () => { }, ] as const)('$name', c => { const input = 'rows' in c ? c.rows : rows; - const actual = Array.from( - generateWithOverlayInnerUnordered(input, c.overlays, pk), - ({row}) => row, - ); + const actual = drainPull( + generateWithOverlayInnerUnordered(pullOf(input), c.overlays, pk), + ).map(({row}) => row); expect(actual).toEqual(c.expected); }); @@ -1129,14 +1128,13 @@ describe('generateWithOverlayInnerUnordered', () => { {a: 2, b: 'x', v: 30}, ]; const compoundPK = ['a', 'b'] as const; - const actual = Array.from( + const actual = drainPull( generateWithOverlayInnerUnordered( - compoundRows, + pullOf(compoundRows), {add: undefined, remove: {a: 1, b: 'y', v: 20}}, compoundPK, ), - ({row}) => row, - ); + ).map(({row}) => row); expect(actual).toEqual([compoundRows[0], compoundRows[2]]); }); @@ -1147,14 +1145,13 @@ describe('generateWithOverlayInnerUnordered', () => { {a: 2, b: 'x', v: 30}, ]; const compoundPK = ['a', 'b'] as const; - const actual = Array.from( + const actual = drainPull( generateWithOverlayInnerUnordered( - compoundRows, + pullOf(compoundRows), {add: undefined, remove: {a: 1, b: 'z', v: 0}}, compoundPK, ), - ({row}) => row, - ); + ).map(({row}) => row); expect(actual).toEqual(compoundRows); }); }); @@ -1173,10 +1170,9 @@ describe('generateWithOverlayUnordered', () => { epoch: 5, change: makeSourceChangeAdd({id: 4, s: 'd', n: 44}), }; - const actual = Array.from( - generateWithOverlayUnordered(rows, undefined, overlay, 4, pk), - ({row}) => row, - ); + const actual = drainPull( + generateWithOverlayUnordered(pullOf(rows), undefined, overlay, 4, pk), + ).map(({row}) => row); expect(actual).toEqual(rows); }); @@ -1185,10 +1181,9 @@ describe('generateWithOverlayUnordered', () => { epoch: 5, change: makeSourceChangeAdd({id: 4, s: 'd', n: 44}), }; - const actual = Array.from( - generateWithOverlayUnordered(rows, undefined, overlay, 5, pk), - ({row}) => row, - ); + const actual = drainPull( + generateWithOverlayUnordered(pullOf(rows), undefined, overlay, 5, pk), + ).map(({row}) => row); expect(actual).toEqual([{id: 4, s: 'd', n: 44}, ...rows]); }); @@ -1197,10 +1192,9 @@ describe('generateWithOverlayUnordered', () => { epoch: 1, change: makeSourceChangeAdd({id: 4, s: 'd', n: 44}), }; - const actual = Array.from( - generateWithOverlayUnordered(rows, {s: 'a'}, overlay, 1, pk), - ({row}) => row, - ); + const actual = drainPull( + generateWithOverlayUnordered(pullOf(rows), {s: 'a'}, overlay, 1, pk), + ).map(({row}) => row); expect(actual).toEqual(rows); }); @@ -1209,17 +1203,16 @@ describe('generateWithOverlayUnordered', () => { epoch: 1, change: makeSourceChangeAdd({id: 4, s: 'd', n: 44}), }; - const actual = Array.from( + const actual = drainPull( generateWithOverlayUnordered( - rows, + pullOf(rows), undefined, overlay, 1, pk, (row: Row) => (row.n as number) < 40, ), - ({row}) => row, - ); + ).map(({row}) => row); expect(actual).toEqual(rows); }); @@ -1228,10 +1221,9 @@ describe('generateWithOverlayUnordered', () => { epoch: 1, change: makeSourceChangeAdd({id: 4, s: 'd', n: 44}), }; - const actual = Array.from( - generateWithOverlayUnordered(rows, undefined, overlay, 1, pk), - ({row}) => row, - ); + const actual = drainPull( + generateWithOverlayUnordered(pullOf(rows), undefined, overlay, 1, pk), + ).map(({row}) => row); expect(actual).toEqual([{id: 4, s: 'd', n: 44}, ...rows]); }); @@ -1240,10 +1232,9 @@ describe('generateWithOverlayUnordered', () => { epoch: 1, change: makeSourceChangeRemove({id: 2, s: 'b', n: 22}), }; - const actual = Array.from( - generateWithOverlayUnordered(rows, undefined, overlay, 1, pk), - ({row}) => row, - ); + const actual = drainPull( + generateWithOverlayUnordered(pullOf(rows), undefined, overlay, 1, pk), + ).map(({row}) => row); expect(actual).toEqual([rows[0], rows[2]]); }); @@ -1255,10 +1246,9 @@ describe('generateWithOverlayUnordered', () => { {id: 2, s: 'b', n: 22}, ), }; - const actual = Array.from( - generateWithOverlayUnordered(rows, undefined, overlay, 1, pk), - ({row}) => row, - ); + const actual = drainPull( + generateWithOverlayUnordered(pullOf(rows), undefined, overlay, 1, pk), + ).map(({row}) => row); expect(actual).toEqual([{id: 2, s: 'b2', n: 225}, rows[0], rows[2]]); }); }); @@ -1380,9 +1370,9 @@ describe('multiConstraints overlay handling — both helpers', () => { const input = iteratorRows ?? rows; test('unordered', () => { - const actual = Array.from( + const actual = drainPull( generateWithOverlayUnordered( - input, + pullOf(input), undefined, overlay, 1, @@ -1390,16 +1380,15 @@ describe('multiConstraints overlay handling — both helpers', () => { undefined, multiConstraints, ), - ({row}) => row, - ); + ).map(({row}) => row); expect(actual).toEqual(expectedUnordered); }); test('ordered', () => { - const actual = Array.from( + const actual = drainPull( generateWithOverlay( undefined, - input, + pullOf(input), undefined, overlay, 1, @@ -1408,8 +1397,7 @@ describe('multiConstraints overlay handling — both helpers', () => { undefined, multiConstraints, ), - ({row}) => row, - ); + ).map(({row}) => row); expect(actual).toEqual(expectedOrdered); }); }, @@ -1418,15 +1406,15 @@ describe('multiConstraints overlay handling — both helpers', () => { describe('mergeSortedStreams', () => { const node = (id: number): Node => ({row: {id}, relationships: {}}); - const ids = (xs: Iterable): (number | 'yield')[] => - Array.from(xs, x => (x === 'yield' ? 'yield' : (x.row.id as number))); + const ids = (xs: PullStream): (number | 'yield')[] => + drainPull(xs).map(x => (x === 'yield' ? 'yield' : (x.row.id as number))); const byId = (a: Node, b: Node) => (a.row.id as number) - (b.row.id as number); - function* gen(values: readonly (Node | 'yield')[]): Stream { - for (const v of values) { - yield v; - } + function gen( + values: readonly (Node | 'yield')[], + ): PullStream { + return pullOf(values); } test('no streams yields nothing', () => { @@ -1548,84 +1536,61 @@ describe('mergeSortedStreams', () => { test('.return() propagates to un-exhausted sub-iterators on early termination', () => { // Build streams that record whether .return() was invoked. - const returned: boolean[] = [false, false, false]; - const trackable = (i: number, values: readonly Node[]): Stream => ({ - [Symbol.iterator]() { - let idx = 0; - return { - next() { - if (idx < values.length) { - return {value: values[idx++], done: false}; - } - return {value: undefined, done: true}; - }, - return(v?: unknown) { - returned[i] = true; - return {value: v, done: true}; - }, - [Symbol.iterator]() { - return this; - }, - }; - }, - }); + const closed: boolean[] = [false, false, false]; + const trackable = ( + i: number, + values: readonly Node[], + ): PullStream => { + let idx = 0; + return { + next: () => (idx < values.length ? values[idx++] : undefined), + close: () => { + closed[i] = true; + }, + }; + }; const a = trackable(0, [node(1), node(4), node(7)]); const b = trackable(1, [node(2), node(5), node(8)]); const c = trackable(2, [node(3), node(6), node(9)]); const merged = mergeSortedStreams([a, b, c], byId); - const it = merged[Symbol.iterator](); - // Pull a couple values then break — JS would call .return() under - // for-of `break`; we invoke explicitly. - expect(it.next().value).toEqual(node(1)); - expect(it.next().value).toEqual(node(2)); - it.return?.(); - - // All three sub-iterators still had un-yielded rows, so all should - // have been .return()-d via mergeSortedStreams's finally block. - expect(returned).toEqual([true, true, true]); + // Take a couple of values, then stop early. Under the iterator protocol a + // `break` triggered `.return()`; the pull protocol makes that explicit. + expect(merged.next()).toEqual(node(1)); + expect(merged.next()).toEqual(node(2)); + merged.close(); + + // All three sub-streams still had unread rows, so all must be closed. + expect(closed).toEqual([true, true, true]); }); - test('.return() not called on already-exhausted sub-iterators', () => { - let aReturned = false; - let bReturned = false; + test('close() not called on already-exhausted sub-streams', () => { + let aClosed = false; + let bClosed = false; const trackable = ( values: readonly Node[], - onReturn: () => void, - ): Stream => ({ - [Symbol.iterator]() { - let idx = 0; - return { - next() { - if (idx < values.length) { - return {value: values[idx++], done: false}; - } - return {value: undefined, done: true}; - }, - return(v?: unknown) { - onReturn(); - return {value: v, done: true}; - }, - [Symbol.iterator]() { - return this; - }, - }; - }, - }); + onClose: () => void, + ): PullStream => { + let idx = 0; + return { + next: () => (idx < values.length ? values[idx++] : undefined), + close: onClose, + }; + }; - // `a` will be drained; `b` will still have rows when we early-exit. - const a = trackable([node(1)], () => (aReturned = true)); - const b = trackable([node(2), node(3), node(4)], () => (bReturned = true)); + // `a` will be drained; `b` will still have rows when we stop early. + const a = trackable([node(1)], () => (aClosed = true)); + const b = trackable([node(2), node(3), node(4)], () => (bClosed = true)); - const it = mergeSortedStreams([a, b], byId)[Symbol.iterator](); - expect(it.next().value).toEqual(node(1)); // from a — exhausts a - expect(it.next().value).toEqual(node(2)); // from b - it.return?.(); + const merged = mergeSortedStreams([a, b], byId); + expect(merged.next()).toEqual(node(1)); // from a -- exhausts a + expect(merged.next()).toEqual(node(2)); // from b + merged.close(); - // a was exhausted naturally; the merge marks heads[0]=null and skips - // .return() on it. b still had rows, so it must be .return()'d. - expect(aReturned).toBe(false); - expect(bReturned).toBe(true); + // `a` ran out on its own, so the merge marks it inactive and skips it; + // `b` still had rows, so it must be closed. + expect(aClosed).toBe(false); + expect(bClosed).toBe(true); }); }); diff --git a/packages/zql/src/ivm/memory-source.ts b/packages/zql/src/ivm/memory-source.ts index 50b00f5ded..be82b3a548 100644 --- a/packages/zql/src/ivm/memory-source.ts +++ b/packages/zql/src/ivm/memory-source.ts @@ -1,7 +1,7 @@ import {assert, unreachable} from '../../../shared/src/asserts.ts'; -import {BTreeSet} from '../../../shared/src/btree-set.ts'; +import {BTreeSet, type ValueIterator} from '../../../shared/src/btree-set.ts'; import {hasOwn} from '../../../shared/src/has-own.ts'; -import {once, toSorted} from '../../../shared/src/iterables.ts'; +import {toSorted} from '../../../shared/src/iterables.ts'; import {must} from '../../../shared/src/must.ts'; import type { Condition, @@ -54,7 +54,13 @@ import type { SourceInput, } from './source.ts'; import {makeSourceChangeAdd, makeSourceChangeRemove} from './source.ts'; -import type {Stream} from './stream.ts'; +import { + LazyPullStream, + type PullStream, + type Stream, + filterPull, + takeWhilePull, +} from './stream.ts'; export type Overlay = { epoch: number; @@ -94,6 +100,20 @@ export type Connection = { * This data is kept in sorted order as downstream pipelines will always expect * the data they receive from `pull` to be in sorted order. */ + +/** + * Rows from an index scan, wrapped as Nodes, stopping at the first row that + * fails `constraint`. + * + * Rows are sorted by the constraint key first, so matches are contiguous and + * the first miss ends the scan. This is `#fetch`'s hot path -- no overlay, no + * `start`, no filters -- which is what a plain scan and every join + * child-lookup take. + */ +type FetchPlan = + | {rows: ValueIterator; constraint: Constraint | undefined} + | {stream: PullStream}; + export class MemorySource implements Source { readonly #tableName: string; readonly #columns: Record; @@ -258,7 +278,24 @@ export class MemorySource implements Source { return [...this.#indexes.keys()]; } - *#fetch(req: FetchRequest, conn: Connection): Stream { + #fetch(req: FetchRequest, conn: Connection): PullStream { + // Lazy, as a generator body was: this reads `#overlay` and + // `conn.lastPushedEpoch`, and a caller may fetch then iterate after a push. + return new LazyPullStream(() => { + const plan = this.#prepareFetch(req, conn); + return 'rows' in plan + ? new ConstrainedRowPull(plan.rows, plan.constraint) + : plan.stream; + }); + } + + /** + * Everything `#fetch` decides before it produces a row, shared by both + * protocols so the fast-path condition and the overlay setup exist once. + * Returns the raw index scan for the hot path -- no overlay, no `start`, no + * filters -- and an assembled iterator for everything else. + */ + #prepareFetch(req: FetchRequest, conn: Connection): FetchPlan { // multiConstraints is handled by driving sub-fetches off the first // entry's values and post-filtering matches against any remaining // entries. TableSource implements multi-IN natively via SQL `AND` of @@ -269,8 +306,7 @@ export class MemorySource implements Source { req.multiConstraints && req.multiConstraints.some(mc => mc.length > 0) ) { - yield* this.#fetchMulti(req, conn); - return; + return {stream: this.#fetchMulti(req, conn)}; } const requestedSort = must(conn.sort); const {compareRows} = conn; @@ -369,19 +405,12 @@ export class MemorySource implements Source { const overlayActive = this.#overlay && conn.lastPushedEpoch >= this.#overlay.epoch; if (!overlayActive && !req.start && !conn.filters && !req.filter) { - const {constraint} = req; - for (const row of rowsIterable) { - if (constraint && !constraintMatchesRow(constraint, row)) { - break; - } - yield {row, relationships: {}}; - } - return; + return {rows: rowsIterable, constraint: req.constraint}; } const withOverlay = generateWithOverlay( startAt, - pkConstraint ? once(rowsIterable) : rowsIterable, + new RowScan(rowsIterable), // use `req.constraint` here and not `fetchOrPkConstraint` since `fetchOrPkConstraint` could be the // primary key constraint. The primary key constraint comes from filters and is acting as a filter // rather than as the fetch constraint. @@ -410,21 +439,27 @@ export class MemorySource implements Source { mergedFilterPredicate, ); - const withConstraint = generateWithConstraint( + // we use `req.constraint` and not `fetchOrPkConstraint` here because we + // need to AND the constraint with what could have been the primary key + // constraint. Rows are sorted by the constraint key first, so matches are + // contiguous and the first miss ends the scan. + const {constraint} = req; + const withConstraint = takeWhilePull( skipYields( generateWithStart(withOverlay, req.start, connectionComparator), ), - // we use `req.constraint` and not `fetchOrPkConstraint` here because we need to - // AND the constraint with what could have been the primary key constraint - req.constraint, + node => + constraint === undefined || constraintMatchesRow(constraint, node.row), ); - yield* mergedFilterPredicate - ? generateWithFilter(withConstraint, mergedFilterPredicate) - : withConstraint; + return { + stream: mergedFilterPredicate + ? filterPull(withConstraint, node => mergedFilterPredicate(node.row)) + : withConstraint, + }; } - *#fetchMulti(req: FetchRequest, conn: Connection): Stream { + #fetchMulti(req: FetchRequest, conn: Connection): PullStream { // Caller (`#fetch`) guards entry on `req.multiConstraints.some(mc => // mc.length > 0)`, so `multis` is guaranteed non-empty after the // empty-entry filter. Per the MultiConstraint contract (operator.ts), @@ -441,7 +476,7 @@ export class MemorySource implements Source { // MultiConstraint contract (see operator.ts), entries are unique and // key-compatible with `baseConstraint`, so we don't dedupe or check // compatibility here. - const subStreams: Stream[] = primary.map(c => { + const subStreams: PullStream[] = primary.map(c => { const merged: Constraint = baseConstraint ? {...baseConstraint, ...c} : c; return this.#fetch( {...req, constraint: merged, multiConstraints: undefined}, @@ -455,32 +490,9 @@ export class MemorySource implements Source { : (a, b) => conn.compareRows(a.row, b.row), ); - if (rest.length === 0) { - yield* merged; - return; - } - - for (const node of merged) { - if (node === 'yield') { - yield 'yield'; - continue; - } - let matchesAll = true; - for (const mc of rest) { - let any = false; - for (const c of mc) { - if (constraintMatchesRow(c, node.row)) { - any = true; - break; - } - } - if (!any) { - matchesAll = false; - break; - } - } - if (matchesAll) yield node; - } + return rest.length === 0 + ? merged + : filterPull(merged, node => node === 'yield' || matchesAll(node, rest)); } *push(change: SourceChange): Stream<'yield'> { @@ -563,22 +575,69 @@ function mergePredicates( return row => connPredicate(row) && reqPredicate(row); } -function* generateWithConstraint( - it: Stream, - constraint: Constraint | undefined, -) { - for (const node of it) { - if (constraint && !constraintMatchesRow(constraint, node.row)) { - break; +/** Stops at the first row that fails `constraint`; matches are contiguous. */ +/** + * Rows from an index scan wrapped as Nodes, stopping at the first row that + * fails `constraint`. Rows are sorted by the constraint key first, so matches + * are contiguous. This is `#fetch`'s hot path -- no overlay, no `start`, no + * filters -- which a plain scan and every join child-lookup take. + */ +class ConstrainedRowPull implements PullStream { + readonly #rows: ValueIterator; + readonly #constraint: Constraint | undefined; + #done = false; + + constructor(rows: ValueIterator, constraint: Constraint | undefined) { + this.#rows = rows; + this.#constraint = constraint; + } + + next(): Node | undefined { + if (this.#done) { + return undefined; + } + const row = this.#rows.nextValue(); + if (row === undefined) { + this.#done = true; + return undefined; } - yield node; + const constraint = this.#constraint; + if (constraint !== undefined && !constraintMatchesRow(constraint, row)) { + this.#done = true; + return undefined; + } + return {row, relationships: {}}; + } + + close(): void { + this.#done = true; } } -function* generateWithFilter(it: Stream, filter: (row: Row) => boolean) { - for (const node of it) { - if (filter(node.row)) { - yield node; +/** The index scan as a pull stream; `nextValue()` avoids a result object. */ +class RowScan implements PullStream { + readonly #rows: ValueIterator; + #done = false; + + constructor(rows: ValueIterator) { + this.#rows = rows; + } + + next(): Row | undefined { + if (this.#done) { + return undefined; + } + const row = this.#rows.nextValue(); + if (row === undefined) { + this.#done = true; + } + return row; + } + + close(): void { + if (!this.#done) { + this.#done = true; + this.#rows.return?.(); } } } @@ -711,36 +770,53 @@ function* genPush( setOverlay(undefined); } -export function* generateWithStart( - nodes: Iterable, - start: Start | undefined, - compare: (r1: Row, r2: Row) => number, -): Stream { - if (!start) { - yield* nodes; - return; +export class WithStart implements PullStream { + readonly #nodes: PullStream; + readonly #start: Start | undefined; + readonly #compare: (r1: Row, r2: Row) => number; + #started: boolean; + + constructor( + nodes: PullStream, + start: Start | undefined, + compare: (r1: Row, r2: Row) => number, + ) { + this.#nodes = nodes; + this.#start = start; + this.#compare = compare; + this.#started = start === undefined; } - let started = false; - for (const node of nodes) { - if (node === 'yield') { - yield node; - continue; - } - if (!started) { - if (start.basis === 'at') { - if (compare(node.row, start.row) >= 0) { - started = true; - } - } else if (start.basis === 'after') { - if (compare(node.row, start.row) > 0) { - started = true; - } + + next(): Node | 'yield' | undefined { + for (;;) { + const node = this.#nodes.next(); + if (node === undefined || node === 'yield') { + return node; + } + if (this.#started) { + return node; + } + // `start` is non-undefined here: #started begins true when it is not. + const start = this.#start as Start; + const c = this.#compare(node.row, start.row); + if (start.basis === 'at' ? c >= 0 : c > 0) { + this.#started = true; + return node; } } - if (started) { - yield node; - } } + + close(): void { + this.#nodes.close(); + } +} + +export function generateWithStart( + nodes: PullStream, + start: Start | undefined, + compare: (r1: Row, r2: Row) => number, +): PullStream { + return new WithStart(nodes, start, compare); } /** @@ -763,9 +839,9 @@ export function* generateWithStart( * is what #4926 fixed for `generateWithStart`; this parameter is the same * distinction for the overlay's own `startAt` pruning. */ -export function* generateWithOverlay( +export function generateWithOverlay( startAt: Row | undefined, - rows: Iterable, + rows: PullStream, constraint: Constraint | undefined, overlay: Overlay | undefined, lastPushedEpoch: number, @@ -773,7 +849,7 @@ export function* generateWithOverlay( startAtCompare: Comparator, filterPredicate?: (row: Row) => boolean | undefined, multiConstraints?: readonly MultiConstraint[] | undefined, -) { +): PullStream { let overlayToApply: Overlay | undefined = undefined; if (overlay && lastPushedEpoch >= overlay.epoch) { overlayToApply = overlay; @@ -786,7 +862,7 @@ export function* generateWithOverlay( filterPredicate, multiConstraints, ); - yield* generateWithOverlayInner(rows, overlays, compare); + return new OverlayInner(rows, overlays, compare); } function computeOverlays( @@ -918,51 +994,96 @@ function overlaysForFilterPredicate( }; } -export function* generateWithOverlayInner( - rowIterator: Iterable, - overlays: Overlays, - compare: (r1: Row, r2: Row) => number, -) { - let addOverlayYielded = false; - let removeOverlaySkipped = false; - for (const row of rowIterator) { - if (!addOverlayYielded && overlays.add) { - const cmp = compare(overlays.add, row); - if (cmp < 0) { - addOverlayYielded = true; - yield {row: overlays.add, relationships: {}}; - } - } +/** + * Splices the overlay rows into an ordered row stream. + * + * `#pending` holds a row that has been read but not yet emitted: when the add + * overlay sorts before it, that overlay is returned first and the row is kept + * for the following call. The generator this replaces expressed the same thing + * with two `yield`s in one loop iteration. + */ +export class OverlayInner implements PullStream { + readonly #rows: PullStream; + readonly #overlays: Overlays; + readonly #compare: (r1: Row, r2: Row) => number; + #pending: Row | undefined; + #addYielded = false; + #removeSkipped = false; + #done = false; - if (!removeOverlaySkipped && overlays.remove) { - const cmp = compare(overlays.remove, row); - if (cmp === 0) { - removeOverlaySkipped = true; + constructor( + rows: PullStream, + overlays: Overlays, + compare: (r1: Row, r2: Row) => number, + ) { + this.#rows = rows; + this.#overlays = overlays; + this.#compare = compare; + } + + next(): Node | undefined { + if (this.#done) { + return undefined; + } + const {add, remove} = this.#overlays; + for (;;) { + let row = this.#pending; + if (row === undefined) { + row = this.#rows.next(); + if (row === undefined) { + this.#done = true; + if (!this.#addYielded && add) { + this.#addYielded = true; + return {row: add, relationships: {}}; + } + return undefined; + } + this.#pending = row; + } + if (!this.#addYielded && add && this.#compare(add, row) < 0) { + this.#addYielded = true; + return {row: add, relationships: {}}; + } + if (!this.#removeSkipped && remove && this.#compare(remove, row) === 0) { + this.#removeSkipped = true; + this.#pending = undefined; continue; } + this.#pending = undefined; + return {row, relationships: {}}; } - yield {row, relationships: {}}; } - if (!addOverlayYielded && overlays.add) { - yield {row: overlays.add, relationships: {}}; + close(): void { + if (!this.#done) { + this.#done = true; + this.#rows.close(); + } } } +export function generateWithOverlayInner( + rows: PullStream, + overlays: Overlays, + compare: (r1: Row, r2: Row) => number, +): PullStream { + return new OverlayInner(rows, overlays, compare); +} + /** * Like {@link generateWithOverlay} but for unordered streams. * No `startAt` or comparator needed. Injects remove/old-edit rows eagerly * at the start, and suppresses add/new-edit rows inline by PK match. */ -export function* generateWithOverlayUnordered( - rows: Iterable, +export function generateWithOverlayUnordered( + rows: PullStream, constraint: Constraint | undefined, overlay: Overlay | undefined, lastPushedEpoch: number, primaryKey: PrimaryKey, filterPredicate?: (row: Row) => boolean, multiConstraints?: readonly MultiConstraint[] | undefined, -) { +): PullStream { let overlayToApply: Overlay | undefined = undefined; if (overlay && lastPushedEpoch >= overlay.epoch) { overlayToApply = overlay; @@ -995,33 +1116,73 @@ export function* generateWithOverlayUnordered( if (filterPredicate) { overlays = overlaysForFilterPredicate(overlays, filterPredicate); } - yield* generateWithOverlayInnerUnordered(rows, overlays, primaryKey); + return new OverlayInnerUnordered(rows, overlays, primaryKey); } -export function* generateWithOverlayInnerUnordered( - rowIterator: Iterable, - overlays: Overlays, - primaryKey: PrimaryKey, -) { - // Eager inject: yield the add overlay at the start (row not yet in storage) - if (overlays.add) { - yield {row: overlays.add, relationships: {}}; +/** {@link OverlayInner} for unordered streams: eager add, inline PK suppress. */ +export class OverlayInnerUnordered implements PullStream { + readonly #rows: PullStream; + readonly #overlays: Overlays; + readonly #primaryKey: PrimaryKey; + #addEmitted = false; + #removeSkipped = false; + #done = false; + + constructor( + rows: PullStream, + overlays: Overlays, + primaryKey: PrimaryKey, + ) { + this.#rows = rows; + this.#overlays = overlays; + this.#primaryKey = primaryKey; } - // Stream with inline suppress: skip the remove overlay (row still in storage) - let removeSkipped = false; - for (const row of rowIterator) { - if ( - !removeSkipped && - overlays.remove && - rowMatchesPK(overlays.remove, row, primaryKey) - ) { - removeSkipped = true; - continue; + + next(): Node | undefined { + if (this.#done) { + return undefined; + } + const {add, remove} = this.#overlays; + if (!this.#addEmitted) { + this.#addEmitted = true; + if (add) { + return {row: add, relationships: {}}; + } + } + for (;;) { + const row = this.#rows.next(); + if (row === undefined) { + this.#done = true; + return undefined; + } + if ( + !this.#removeSkipped && + remove && + rowMatchesPK(remove, row, this.#primaryKey) + ) { + this.#removeSkipped = true; + continue; + } + return {row, relationships: {}}; + } + } + + close(): void { + if (!this.#done) { + this.#done = true; + this.#rows.close(); } - yield {row, relationships: {}}; } } +export function generateWithOverlayInnerUnordered( + rows: PullStream, + overlays: Overlays, + primaryKey: PrimaryKey, +): PullStream { + return new OverlayInnerUnordered(rows, overlays, primaryKey); +} + function rowMatchesPK(a: Row, b: Row, primaryKey: PrimaryKey): boolean { for (const key of primaryKey) { if (!valuesEqual(a[key], b[key])) { @@ -1086,14 +1247,23 @@ function compareBounds(a: Bound, b: Bound): number { return compareValues(a, b); } -function* generateRows( +/** + * Rows from `scanStart` onwards. + * + * Returns the BTree's own iterator rather than delegating to it from a + * generator. `yield*` over an iterable that is already an `IterableIterator` + * buys nothing and costs a generator resume per row, which on Hermes is the + * most expensive way to move a value. Single-use semantics are unchanged: an + * `IterableIterator` returns itself from `[Symbol.iterator]()`, exactly as a + * generator does. + */ +function generateRows( data: BTreeSet, scanStart: RowBound | undefined, reverse: boolean | undefined, -) { - yield* data[reverse ? 'valuesFromReversed' : 'valuesFrom']( - scanStart as Row | undefined, - ); +): ValueIterator { + const from = scanStart as Row | undefined; + return reverse ? data.valuesFromReversed(from) : data.valuesFrom(from); } export function stringify(change: SourceChange) { @@ -1120,109 +1290,171 @@ export function stringify(change: SourceChange) { * leaves cursors open, causing later writes on the same connection to * fail with "database connection is busy executing a query". */ -export function* mergeSortedStreams( - streams: readonly Stream[], +export function mergeSortedStreams( + streams: readonly PullStream[], compare: (a: Node, b: Node) => number, -): Stream { - const iterators: Iterator[] = streams.map(s => - s[Symbol.iterator](), - ); - // True while iterators[i] hasn't yet returned `done`. The finally - // block uses this to skip already-exhausted streams when propagating - // `.return()`. - const active: boolean[] = new Array(iterators.length).fill(true); +): PullStream { + return new MergeSortedStreams(streams, compare); +} - // Min-heap of entries; `idx` tells us which stream to refill from - // after the entry's row is emitted. - type Entry = {row: Node; idx: number}; - const heap: Entry[] = []; +/** True when `node` satisfies every remaining `MultiConstraint` entry. */ +function matchesAll(node: Node, rest: readonly MultiConstraint[]): boolean { + for (const mc of rest) { + let any = false; + for (const c of mc) { + if (constraintMatchesRow(c, node.row)) { + any = true; + break; + } + } + if (!any) { + return false; + } + } + return true; +} - const siftUp = (start: number) => { +type MergeEntry = {row: Node; idx: number}; + +/** + * N-way merge of pre-sorted Node streams, as a pull stream. + * + * The generator this replaces suspended mid-prime and mid-refill to forward a + * 'yield'; the same points are now explicit state -- `#primeIdx` for priming, + * `#refill` for the stream owing a replacement for the root it just emitted -- + * so a 'yield' can be returned and the merge resumed exactly where it paused. + * + * Streams that are not exhausted are closed on completion or `close()`, so the + * underlying cursors are released; leaking one leaves later writes on the same + * connection failing with "database connection is busy executing a query". + */ +class MergeSortedStreams implements PullStream { + readonly #streams: readonly PullStream[]; + readonly #compare: (a: Node, b: Node) => number; + readonly #active: boolean[]; + readonly #heap: MergeEntry[] = []; + #priming = true; + #primeIdx = 0; + #refill: number | undefined; + #done = false; + + constructor( + streams: readonly PullStream[], + compare: (a: Node, b: Node) => number, + ) { + this.#streams = streams; + this.#compare = compare; + this.#active = new Array(streams.length).fill(true); + } + + #siftUp(start: number): void { + const heap = this.#heap; let i = start; while (i > 0) { const p = (i - 1) >> 1; - if (compare(heap[i].row, heap[p].row) >= 0) return; + if (this.#compare(heap[i].row, heap[p].row) >= 0) { + return; + } const t = heap[i]; heap[i] = heap[p]; heap[p] = t; i = p; } - }; + } - const siftDown = (start: number) => { + #siftDown(start: number): void { + const heap = this.#heap; let i = start; const n = heap.length; - while (true) { + for (;;) { const l = (i << 1) + 1; const r = l + 1; let smallest = i; - if (l < n && compare(heap[l].row, heap[smallest].row) < 0) smallest = l; - if (r < n && compare(heap[r].row, heap[smallest].row) < 0) smallest = r; - if (smallest === i) return; + if (l < n && this.#compare(heap[l].row, heap[smallest].row) < 0) { + smallest = l; + } + if (r < n && this.#compare(heap[r].row, heap[smallest].row) < 0) { + smallest = r; + } + if (smallest === i) { + return; + } const t = heap[i]; heap[i] = heap[smallest]; heap[smallest] = t; i = smallest; } - }; + } - // Pull the next Node from iterator `idx`, forwarding any 'yield's. - // Returns the Node, or `undefined` once the stream is exhausted. - const pullNext = function* ( - idx: number, - ): Generator<'yield', Node | undefined, undefined> { - while (true) { - const r = iterators[idx].next(); - if (r.done) { - active[idx] = false; - return undefined; + /** One value from stream `idx`: a Node, 'yield' to forward, or undefined. */ + #pullOne(idx: number): Node | 'yield' | undefined { + const v = this.#streams[idx].next(); + if (v === undefined) { + this.#active[idx] = false; + } + return v; + } + + next(): Node | 'yield' | undefined { + if (this.#done) { + return undefined; + } + for (;;) { + if (this.#priming) { + while (this.#primeIdx < this.#streams.length) { + const v = this.#pullOne(this.#primeIdx); + if (v === 'yield') { + return v; + } + if (v !== undefined) { + this.#heap.push({row: v, idx: this.#primeIdx}); + this.#siftUp(this.#heap.length - 1); + } + this.#primeIdx++; + } + this.#priming = false; } - if (r.value === 'yield') { - yield 'yield'; + + if (this.#refill !== undefined) { + const v = this.#pullOne(this.#refill); + if (v === 'yield') { + return v; + } + this.#refill = undefined; + if (v !== undefined) { + // The emitted row was captured by the caller, so replacing the + // root's row in place is safe. + this.#heap[0].row = v; + this.#siftDown(0); + } else { + const last = must(this.#heap.pop()); + if (this.#heap.length > 0) { + this.#heap[0] = last; + this.#siftDown(0); + } + } continue; } - return r.value; - } - }; - try { - // Prime: push the first row of each non-empty stream onto the heap. - for (let i = 0; i < iterators.length; i++) { - const row = yield* pullNext(i); - if (row !== undefined) { - heap.push({row, idx: i}); - siftUp(heap.length - 1); + if (this.#heap.length === 0) { + this.close(); + return undefined; } + const top = this.#heap[0]; + this.#refill = top.idx; + return top.row; } + } - while (heap.length > 0) { - // Root is the global min across all active streams. - const top = heap[0]; - yield top.row; - const next = yield* pullNext(top.idx); - if (next !== undefined) { - // Refill root in place (top === heap[0]) and sift down. The - // already-yielded `top.row` value is captured by the yield, so - // mutating it here doesn't affect what was emitted. - top.row = next; - siftDown(0); - } else { - // Stream exhausted. Move tail to root and shrink. Pop returns - // the last entry; if the heap had only one entry it was the - // root we just yielded, so we just leave the heap empty. - const last = must(heap.pop()); - if (heap.length > 0) { - heap[0] = last; - siftDown(0); - } - } + close(): void { + if (this.#done) { + return; } - } finally { - // Close any iterators that aren't already exhausted so their - // `finally` blocks (which release cursors / cached statements) run. - for (let i = 0; i < iterators.length; i++) { - if (active[i]) { - iterators[i].return?.(); + this.#done = true; + for (let i = 0; i < this.#streams.length; i++) { + if (this.#active[i]) { + this.#active[i] = false; + this.#streams[i].close(); } } } diff --git a/packages/zql/src/ivm/operator.ts b/packages/zql/src/ivm/operator.ts index 8d19b6ae6d..abb1e2facb 100644 --- a/packages/zql/src/ivm/operator.ts +++ b/packages/zql/src/ivm/operator.ts @@ -5,7 +5,7 @@ import type {Change} from './change.ts'; import type {Constraint} from './constraint.ts'; import type {Node} from './data.ts'; import type {SourceSchema} from './schema.ts'; -import type {Stream} from './stream.ts'; +import type {PullStream, Stream} from './stream.ts'; export {skipYields} from './skip-yields.ts'; @@ -41,7 +41,7 @@ export interface Input extends InputBase { * - During push: If a fetch to an input consumed by the push logic yields * 'yield', it must be yielded to the caller of push immediately. */ - fetch(req: FetchRequest): Stream; + fetch(req: FetchRequest): PullStream; } /** diff --git a/packages/zql/src/ivm/predicate-pushdown.test.ts b/packages/zql/src/ivm/predicate-pushdown.test.ts index c59a9bb810..dc64632c21 100644 --- a/packages/zql/src/ivm/predicate-pushdown.test.ts +++ b/packages/zql/src/ivm/predicate-pushdown.test.ts @@ -4,6 +4,7 @@ import {createSilentLogContext} from '../../../shared/src/logging-test-utils.ts' import type {SimpleCondition} from '../../../zero-protocol/src/ast.ts'; import type {NoSubqueryCondition} from '../builder/filter.ts'; import {Catch} from './catch.ts'; +import type {Node} from './data.ts'; import {FilterEnd, FilterStart} from './filter-operators.ts'; import {FlippedJoin} from './flipped-join.ts'; import {Join} from './join.ts'; @@ -12,8 +13,7 @@ import {type FetchRequest, type Input, type Output} from './operator.ts'; import type {SourceSchema} from './schema.ts'; import {Skip} from './skip.ts'; import {makeSourceChangeAdd} from './source.ts'; -import {consume} from './stream.ts'; -import type {Stream} from './stream.ts'; +import {consume, drainPull, type PullStream} from './stream.ts'; import {Take} from './take.ts'; import {createSource} from './test/source-factory.ts'; import {UnionFanIn} from './union-fan-in.ts'; @@ -40,13 +40,9 @@ class RecordingInput implements Input { this.#wrapped = wrapped; } - fetch( - req: FetchRequest, - ): Stream extends Stream ? T : never> { + fetch(req: FetchRequest): PullStream { this.received.push(req); - return this.#wrapped.fetch(req) as Stream< - ReturnType extends Stream ? T : never - >; + return this.#wrapped.fetch(req); } setOutput(output: Output): void { @@ -79,14 +75,14 @@ describe('req.filter contract (pass-through operators preserve it)', () => { const filterEnd = new FilterEnd(filterStart, filterStart); // Bare consumer fetch: req.filter is undefined. - [...filterEnd.fetch({})]; + drainPull(filterEnd.fetch({})); expect(recorder.received).toHaveLength(1); expect(recorder.received[0].filter).toEqual(filterCondition); // Consumer-provided req.filter should be AND-merged with FilterStart's // own condition. const incomingFilter: NoSubqueryCondition = cmpEq('a', 'y'); - [...filterEnd.fetch({filter: incomingFilter})]; + drainPull(filterEnd.fetch({filter: incomingFilter})); expect(recorder.received).toHaveLength(2); expect(recorder.received[1].filter).toEqual({ type: 'and', @@ -108,7 +104,7 @@ describe('req.filter contract (pass-through operators preserve it)', () => { const filterEnd = new FilterEnd(filterStart, filterStart); const incomingFilter: NoSubqueryCondition = cmpEq('a', 'y'); - [...filterEnd.fetch({filter: incomingFilter})]; + drainPull(filterEnd.fetch({filter: incomingFilter})); expect(recorder.received).toHaveLength(1); // The FetchRequest object itself should be passed through unchanged // (no clone/spread needed) when no merge happens. @@ -130,7 +126,7 @@ describe('req.filter contract (pass-through operators preserve it)', () => { const skip = new Skip(recorder, {row: {a: 'a0'}, exclusive: false}); const incomingFilter: NoSubqueryCondition = cmpEq('b', 'x'); - [...skip.fetch({filter: incomingFilter})]; + drainPull(skip.fetch({filter: incomingFilter})); expect(recorder.received).toHaveLength(1); expect(recorder.received[0].filter).toBe(incomingFilter); @@ -155,7 +151,7 @@ describe('req.filter contract (pass-through operators preserve it)', () => { const ufi = new UnionFanIn(ufo, [branch1, branch2]); const incomingFilter: NoSubqueryCondition = cmpEq('b', 'x'); - [...ufi.fetch({filter: incomingFilter})]; + drainPull(ufi.fetch({filter: incomingFilter})); // Each branch should have triggered a fetch with the same req.filter. expect(recorder.received.length).toBeGreaterThanOrEqual(2); @@ -179,7 +175,7 @@ describe('req.filter contract (pass-through operators preserve it)', () => { const take = new Take(recorder, new MemoryStorage(), 10); const incomingFilter: NoSubqueryCondition = cmpEq('b', 'x'); - [...take.fetch({filter: incomingFilter})]; + drainPull(take.fetch({filter: incomingFilter})); expect(recorder.received).toHaveLength(1); expect(recorder.received[0].filter).toBe(incomingFilter); @@ -224,7 +220,7 @@ describe('req.filter contract (pass-through operators preserve it)', () => { const incomingFilter: NoSubqueryCondition = cmpEq('status', 'open'); const sink = new Catch(join); - [...sink.fetch({filter: incomingFilter})]; + sink.fetch({filter: incomingFilter}); // Exactly one parent fetch, and it must carry the original filter. expect(recorder.received).toHaveLength(1); @@ -271,7 +267,7 @@ describe('req.filter contract (pass-through operators preserve it)', () => { const incomingFilter: NoSubqueryCondition = cmpEq('status', 'open'); const sink = new Catch(flippedJoin); - [...sink.fetch({filter: incomingFilter})]; + sink.fetch({filter: incomingFilter}); // FlippedJoin should have invoked at least one parent fetch with // child-derived multiConstraints AND the original req.filter intact. diff --git a/packages/zql/src/ivm/push-accumulated.test.ts b/packages/zql/src/ivm/push-accumulated.test.ts index c83035ce3f..f6532117a7 100644 --- a/packages/zql/src/ivm/push-accumulated.test.ts +++ b/packages/zql/src/ivm/push-accumulated.test.ts @@ -13,6 +13,7 @@ import { type EditChange, type RemoveChange, } from './change.js'; +import type {RelationshipStream} from './data.ts'; import type {InputBase, Output} from './operator.js'; import { pushAccumulatedChanges as genPushAccumulatedChanges, @@ -21,6 +22,7 @@ import { mergeRelationships, } from './push-accumulated.js'; import type {SourceSchema} from './schema.js'; +import {drainPull, emptyPullStream, pullOf} from './stream.ts'; const mockPusher: InputBase = { getSchema: () => mockSchema as any, @@ -104,8 +106,14 @@ describe('pushAccumulatedChanges', () => { test('multiple add changes collapse to single add', () => { const accumulatedPushes: Change[] = [ - makeAddChange({row: {id: 1}, relationships: {rel1: () => []}}), - makeAddChange({row: {id: 1}, relationships: {rel2: () => []}}), + makeAddChange({ + row: {id: 1}, + relationships: {rel1: () => emptyPullStream()}, + }), + makeAddChange({ + row: {id: 1}, + relationships: {rel2: () => emptyPullStream()}, + }), ]; pushAccumulatedChanges( @@ -161,8 +169,14 @@ describe('pushAccumulatedChanges', () => { test('multiple remove changes collapse to single remove', () => { const accumulatedPushes: Change[] = [ - makeRemoveChange({row: {id: 1}, relationships: {rel1: () => []}}), - makeRemoveChange({row: {id: 1}, relationships: {rel2: () => []}}), + makeRemoveChange({ + row: {id: 1}, + relationships: {rel1: () => emptyPullStream()}, + }), + makeRemoveChange({ + row: {id: 1}, + relationships: {rel2: () => emptyPullStream()}, + }), ]; pushAccumulatedChanges( @@ -268,16 +282,19 @@ describe('pushAccumulatedChanges', () => { test('edit supersedes add and remove when all three present', () => { const accumulatedPushes: Change[] = [ makeEditChange( - {row: {id: 1, value: 3}, relationships: {editRel: () => []}}, + { + row: {id: 1, value: 3}, + relationships: {editRel: () => emptyPullStream()}, + }, {row: {id: 1, value: 0}, relationships: {}}, ), makeAddChange({ row: {id: 1, value: 2}, - relationships: {addRel: () => []}, + relationships: {addRel: () => emptyPullStream()}, }), makeRemoveChange({ row: {id: 1, value: 1}, - relationships: {removeRel: () => []}, + relationships: {removeRel: () => emptyPullStream()}, }), ]; @@ -399,11 +416,11 @@ describe('mergeRelationships', () => { test('merges relationships from add changes', () => { const left: Change = makeAddChange({ row: {id: 1}, - relationships: {rel1: () => []}, + relationships: {rel1: () => emptyPullStream()}, }); const right: Change = makeAddChange({ row: {id: 1}, - relationships: {rel2: () => []}, + relationships: {rel2: () => emptyPullStream()}, }); const result = mergeRelationships(left, right); @@ -417,11 +434,11 @@ describe('mergeRelationships', () => { test('merges relationships from remove changes', () => { const left: Change = makeRemoveChange({ row: {id: 1}, - relationships: {rel1: () => []}, + relationships: {rel1: () => emptyPullStream()}, }); const right: Change = makeRemoveChange({ row: {id: 1}, - relationships: {rel2: () => []}, + relationships: {rel2: () => emptyPullStream()}, }); const result = mergeRelationships(left, right); @@ -434,12 +451,12 @@ describe('mergeRelationships', () => { test('merges relationships from edit changes', () => { const left: Change = makeEditChange( - {row: {id: 1}, relationships: {rel1: () => []}}, - {row: {id: 1}, relationships: {oldRel1: () => []}}, + {row: {id: 1}, relationships: {rel1: () => emptyPullStream()}}, + {row: {id: 1}, relationships: {oldRel1: () => emptyPullStream()}}, ); const right: Change = makeEditChange( - {row: {id: 1}, relationships: {rel2: () => []}}, - {row: {id: 1}, relationships: {oldRel2: () => []}}, + {row: {id: 1}, relationships: {rel2: () => emptyPullStream()}}, + {row: {id: 1}, relationships: {oldRel2: () => emptyPullStream()}}, ); const result = mergeRelationships(left, right) as EditChange; @@ -454,8 +471,8 @@ describe('mergeRelationships', () => { }); test('left takes precedence when same relationship exists', () => { - const rel1Left = () => []; - const rel1Right = () => []; + const rel1Left = () => pullOf([]); + const rel1Right = () => pullOf([]); const left: Change = makeAddChange({ row: {id: 1}, @@ -473,12 +490,12 @@ describe('mergeRelationships', () => { test('merges edit with add', () => { const left: Change = makeEditChange( - {row: {id: 1}, relationships: {editRel: () => []}}, + {row: {id: 1}, relationships: {editRel: () => emptyPullStream()}}, {row: {id: 1}, relationships: {}}, ); const right: Change = makeAddChange({ row: {id: 1}, - relationships: {addRel: () => []}, + relationships: {addRel: () => emptyPullStream()}, }); const result = mergeRelationships(left, right) as EditChange; @@ -492,11 +509,11 @@ describe('mergeRelationships', () => { test('merges edit with remove', () => { const left: Change = makeEditChange( {row: {id: 1}, relationships: {}}, - {row: {id: 1}, relationships: {editOldRel: () => []}}, + {row: {id: 1}, relationships: {editOldRel: () => emptyPullStream()}}, ); const right: Change = makeRemoveChange({ row: {id: 1}, - relationships: {removeRel: () => []}, + relationships: {removeRel: () => emptyPullStream()}, }); const result = mergeRelationships(left, right) as EditChange; @@ -513,11 +530,11 @@ describe('mergeRelationships', () => { relationshipName: 'childRel', }; const left: Change = makeChildChange( - {row: {id: 1}, relationships: {rel1: () => []}}, + {row: {id: 1}, relationships: {rel1: () => emptyPullStream()}}, childInfo, ); const right: Change = makeChildChange( - {row: {id: 1}, relationships: {rel2: () => []}}, + {row: {id: 1}, relationships: {rel2: () => emptyPullStream()}}, childInfo, ); @@ -543,8 +560,12 @@ describe('makeAddEmptyRelationships', () => { expect(Object.keys(result[ChangeIndex.NODE].relationships)).toEqual( expect.arrayContaining(['rel1', 'rel2']), ); - expect(result[ChangeIndex.NODE].relationships.rel1?.()).toEqual([]); - expect(result[ChangeIndex.NODE].relationships.rel2?.()).toEqual([]); + expect(drainPull(result[ChangeIndex.NODE].relationships.rel1?.()!)).toEqual( + [], + ); + expect(drainPull(result[ChangeIndex.NODE].relationships.rel2?.()!)).toEqual( + [], + ); }); test('adds empty relationships for remove change', () => { @@ -586,7 +607,7 @@ describe('makeAddEmptyRelationships', () => { const addEmptyRelationships = makeAddEmptyRelationships(schema); - const existingRel = () => [{row: {id: 2}, relationships: {}}]; + const existingRel = () => pullOf([{row: {id: 2}, relationships: {}}]); const change: Change = makeAddChange({ row: {id: 1}, relationships: {rel1: existingRel}, @@ -595,7 +616,9 @@ describe('makeAddEmptyRelationships', () => { const result = addEmptyRelationships(change) as AddChange; expect(result[ChangeIndex.NODE].relationships.rel1).toBe(existingRel); - expect(result[ChangeIndex.NODE].relationships.rel2?.()).toEqual([]); + expect(drainPull(result[ChangeIndex.NODE].relationships.rel2?.()!)).toEqual( + [], + ); }); test('does not modify child changes', () => { @@ -634,8 +657,8 @@ describe('makeAddEmptyRelationships', () => { describe('mergeEmpty', () => { test('adds empty streams for missing relationships', () => { - const relationships: Record any[]> = { - existing: () => [{id: 1}], + const relationships: Record RelationshipStream> = { + existing: () => pullOf([{row: {id: 1}, relationships: {}}]), }; mergeEmpty(relationships, ['existing', 'new1', 'new2']); @@ -643,20 +666,22 @@ describe('mergeEmpty', () => { expect(Object.keys(relationships)).toEqual( expect.arrayContaining(['existing', 'new1', 'new2']), ); - expect(relationships.existing()).toEqual([{id: 1}]); - expect(relationships.new1()).toEqual([]); - expect(relationships.new2()).toEqual([]); + expect(drainPull(relationships.existing())).toEqual([ + {row: {id: 1}, relationships: {}}, + ]); + expect(drainPull(relationships.new1())).toEqual([]); + expect(drainPull(relationships.new2())).toEqual([]); }); test('does not overwrite existing relationships', () => { - const existingFn = () => [{id: 1}]; - const relationships: Record any[]> = { + const existingFn = () => pullOf([{row: {id: 1}, relationships: {}}]); + const relationships: Record RelationshipStream> = { rel1: existingFn, }; mergeEmpty(relationships, ['rel1', 'rel2']); expect(relationships.rel1).toBe(existingFn); - expect(relationships.rel2()).toEqual([]); + expect(drainPull(relationships.rel2())).toEqual([]); }); }); diff --git a/packages/zql/src/ivm/push-accumulated.ts b/packages/zql/src/ivm/push-accumulated.ts index c7ee8af501..bcda121879 100644 --- a/packages/zql/src/ivm/push-accumulated.ts +++ b/packages/zql/src/ivm/push-accumulated.ts @@ -1,6 +1,5 @@ import {assert, unreachable} from '../../../shared/src/asserts.ts'; import {must} from '../../../shared/src/must.ts'; -import {emptyArray} from '../../../shared/src/sentinels.ts'; import {ChangeIndex} from './change-index.ts'; import {ChangeType} from './change-type.ts'; import { @@ -10,10 +9,10 @@ import { makeRemoveChange, type Change, } from './change.ts'; -import type {Node} from './data.ts'; +import type {RelationshipStream} from './data.ts'; import type {InputBase, Output} from './operator.ts'; import type {SourceSchema} from './schema.ts'; -import type {Stream} from './stream.ts'; +import {type Stream, emptyPullStream} from './stream.ts'; /** * # pushAccumulatedChanges @@ -419,12 +418,12 @@ export function makeAddEmptyRelationships( * This modifies the `relationships` object in place. */ export function mergeEmpty( - relationships: Record Stream>, + relationships: Record RelationshipStream>, relationshipNames: string[], ) { for (const relName of relationshipNames) { if (relationships[relName] === undefined) { - relationships[relName] = () => emptyArray; + relationships[relName] = () => emptyPullStream(); } } } diff --git a/packages/zql/src/ivm/skip-yields.ts b/packages/zql/src/ivm/skip-yields.ts index 4e802cb1d1..4aaf0d38e8 100644 --- a/packages/zql/src/ivm/skip-yields.ts +++ b/packages/zql/src/ivm/skip-yields.ts @@ -1,46 +1,14 @@ import type {Node} from './data.ts'; -import type {Stream} from './stream.ts'; - -// Implemented as a custom IterableIterator rather than a generator function to -// reduce allocations. A generator creates a new state-machine object each time -// it is called. By returning `this` from [Symbol.iterator](), the same object -// acts as both the Iterable and the Iterator, so iterating the stream incurs -// only one allocation (the SkipYieldsStream instance itself) instead of two. -// Additionally, the IteratorResult objects ({value, done}) from the inner -// iterator are returned directly rather than being recreated, avoiding further -// per-item allocations. -class SkipYieldsStream implements IterableIterator { - readonly #stream: Stream; - #it: Iterator | undefined = undefined; - - constructor(stream: Stream) { - this.#stream = stream; - } - - [Symbol.iterator](): IterableIterator { - this.#it = this.#stream[Symbol.iterator](); - return this; - } - - next(): IteratorResult { - // #it is always set before next() is called, as [Symbol.iterator]() must - // be called first (e.g. by a for-of loop). - // eslint-disable-next-line @typescript-eslint/no-non-null-assertion - const it = this.#it!; - for (;;) { - const r = it.next(); - if (r.done || r.value !== 'yield') { - return r as IteratorResult; - } - } - } - - return(value?: undefined): IteratorResult { - this.#it?.return?.(value); - return {done: true, value: undefined}; - } -} - -export function skipYields(stream: Stream): Stream { - return new SkipYieldsStream(stream); +import {filterPull, type PullStream} from './stream.ts'; + +/** + * Drops the 'yield' markers from a stream. + * + * The cast is the one place that knows dropping every 'yield' leaves only + * Nodes; `filterPull` cannot narrow its own element type. + */ +export function skipYields( + stream: PullStream, +): PullStream { + return filterPull(stream, v => v !== 'yield') as PullStream; } diff --git a/packages/zql/src/ivm/skip.ts b/packages/zql/src/ivm/skip.ts index 69160c1176..8db007ae7c 100644 --- a/packages/zql/src/ivm/skip.ts +++ b/packages/zql/src/ivm/skip.ts @@ -19,7 +19,12 @@ import { type Start, } from './operator.ts'; import type {SourceSchema} from './schema.ts'; -import type {Stream} from './stream.ts'; +import { + type Stream, + emptyPullStream, + type PullStream, + takeWhilePull, +} from './stream.ts'; export type Bound = { row: Row; @@ -50,28 +55,22 @@ export class Skip implements Operator { return this.#input.getSchema(); } - *fetch(req: FetchRequest): Stream { + fetch(req: FetchRequest): PullStream { const start = this.#getStart(req); if (start === 'empty') { - return; + return emptyPullStream(); } const nodes = this.#input.fetch({...req, start}); if (!req.reverse) { - yield* nodes; - return; - } - for (const node of nodes) { - if (node === 'yield') { - yield node; - continue; - } - if (!this.#shouldBePresent(node.row)) { - return; - } - yield node; + return nodes; } + // Reverse: rows arrive descending, so the first row that should not be + // present ends the stream. + return takeWhilePull( + nodes, + node => node === 'yield' || this.#shouldBePresent(node.row), + ); } - setOutput(output: Output): void { this.#output = output; } diff --git a/packages/zql/src/ivm/snitch.ts b/packages/zql/src/ivm/snitch.ts index f1ac480aa0..a333f4bf03 100644 --- a/packages/zql/src/ivm/snitch.ts +++ b/packages/zql/src/ivm/snitch.ts @@ -4,10 +4,10 @@ import {ChangeIndex} from './change-index.ts'; import {ChangeType} from './change-type.ts'; import type {Change} from './change.ts'; import type {Node} from './data.ts'; -import type { - FilterInput, - FilterOperator, - FilterOutput, +import { + type FilterInput, + type FilterOperator, + type FilterOutput, } from './filter-operators.ts'; import { type FetchRequest, @@ -16,7 +16,7 @@ import { type Output, } from './operator.ts'; import type {SourceSchema} from './schema.ts'; -import type {Stream} from './stream.ts'; +import type {PullStream, Stream} from './stream.ts'; /** * Snitch is an Operator that records all messages it receives. Useful for @@ -62,25 +62,43 @@ export class Snitch implements Operator { this.log.push(message); } - fetch(req: FetchRequest): Stream { + fetch(req: FetchRequest): PullStream { this.#log([this.#name, 'fetch', req]); - return this.fetchGenerator(req); - } - - *fetchGenerator(req: FetchRequest): Stream { + const input = this.#input.fetch(req); let count = 0; - try { - for (const node of this.#input.fetch(req)) { - if (node === 'yield') { - yield node; - continue; - } - count++; - yield node; + let done = false; + const finish = () => { + if (!done) { + done = true; + input.close(); + this.#log([this.#name, 'fetchCount', req, count]); } - } finally { - this.#log([this.#name, 'fetchCount', req, count]); - } + }; + return { + next: (): Node | 'yield' | undefined => { + if (done) { + return undefined; + } + let node: Node | 'yield' | undefined; + try { + node = input.next(); + } catch (e) { + // The generator logged the count from a `finally`, so a throw + // mid-scan still recorded what had been read. + finish(); + throw e; + } + if (node === undefined) { + finish(); + return undefined; + } + if (node !== 'yield') { + count++; + } + return node; + }, + close: finish, + }; } *push(change: Change): Stream<'yield'> { @@ -151,10 +169,20 @@ export class FilterSnitch implements FilterOperator { this.#output?.endFilter(); } - *filter(node: Node): Generator<'yield', boolean> { - this.#log([this.#name, 'filter', node.row]); + /** The node whose 'filter' has been logged but not yet resolved. */ + #logged: Node | undefined; + + filter(node: Node): boolean | 'yield' { + if (this.#logged !== node) { + this.#log([this.#name, 'filter', node.row]); + this.#logged = node; + } assert(this.#output, 'Snitch: output must be set before filter is called'); - return yield* this.#output.filter(node); + const r = this.#output.filter(node); + if (r !== 'yield') { + this.#logged = undefined; + } + return r; } destroy(): void { diff --git a/packages/zql/src/ivm/source.test.ts b/packages/zql/src/ivm/source.test.ts index a40b26bcee..5b3235123d 100644 --- a/packages/zql/src/ivm/source.test.ts +++ b/packages/zql/src/ivm/source.test.ts @@ -24,7 +24,7 @@ import { makeSourceChangeRemove, type SourceChange, } from './source.ts'; -import {consume} from './stream.ts'; +import {consume, drainPull} from './stream.ts'; import {createSource} from './test/source-factory.ts'; const lc = createSilentLogContext(); @@ -63,7 +63,7 @@ class OverlaySpy implements Output { } fetch(req: FetchRequest) { - this.fetches.push(Array.from(this.#input.fetch(req), expandNode)); + this.fetches.push(drainPull(this.#input.fetch(req)).map(expandNode)); } push() { @@ -3707,25 +3707,14 @@ test('streams-are-one-time-only', () => { const conn = source.connect([['a', 'asc']]); const stream = conn.fetch({}); - const it1 = stream[Symbol.iterator](); - const it2 = stream[Symbol.iterator](); - expect(it1.next()).toEqual({ - done: false, - value: {row: {a: 1}, relationships: {}}, - }); - expect(it2.next()).toEqual({ - done: false, - value: {row: {a: 2}, relationships: {}}, - }); - expect(it1.next()).toEqual({ - done: false, - value: {row: {a: 3}, relationships: {}}, - }); - expect(it2.next()).toEqual({done: true, value: undefined}); - expect(it1.next()).toEqual({done: true, value: undefined}); - - const it3 = stream[Symbol.iterator](); - expect(it3.next()).toEqual({done: true, value: undefined}); + // A pull stream *is* the cursor: there is no way to obtain a second, + // independent reader, so interleaved reads share one position and the + // stream stays exhausted once it ends. + expect(stream.next()).toEqual({row: {a: 1}, relationships: {}}); + expect(stream.next()).toEqual({row: {a: 2}, relationships: {}}); + expect(stream.next()).toEqual({row: {a: 3}, relationships: {}}); + expect(stream.next()).toBeUndefined(); + expect(stream.next()).toBeUndefined(); }); test('json is a valid type to read and write to/from a source', () => { diff --git a/packages/zql/src/ivm/stream.ts b/packages/zql/src/ivm/stream.ts index 3a8c929138..85f13003d4 100644 --- a/packages/zql/src/ivm/stream.ts +++ b/packages/zql/src/ivm/stream.ts @@ -34,3 +34,253 @@ export function drainGenerator( } return result.value; } + +/** + * The pull-function protocol. + * + * `next()` returns the next value directly, or `undefined` once the stream is + * exhausted; `close()` releases resources if the consumer stops early. This is + * the same pull model as {@link Stream} -- the consumer drives, so 'yield' and + * stream merging work unchanged -- minus the iterator protocol's costs: no + * `{done, value}` result object per value, and one object per stream rather + * than per stage. A four-deep pipeline measured ~26% faster than hand-written + * iterators on both Hermes and V8, within a few percent of push callbacks. + * + * `undefined` is the end marker, so a PullStream cannot carry `undefined` as a + * value. Nodes are objects and 'yield' is a string; neither can be. + */ +/** + * A stream read by calling `next()` until it returns `undefined`; `close()` + * releases resources if the consumer stops early. + * + * Deliberately NOT `Iterable`. If a pull stream could be `for...of`'d, every + * unconverted consumer would keep silently paying the iterator protocol -- a + * `{done, value}` object per row -- which is the cost this protocol exists to + * remove. There is no adapter back to an iterable: a consumer that wants + * values calls `next()`. + */ +export interface PullStream { + next(): T | undefined; + close(): void; +} + +const EMPTY: PullStream = { + next: () => undefined, + close: () => {}, +}; +/** + * Keeps the values `keep` accepts. + * + * These three cover most of what operators do to a stream, replacing a + * per-operator class each -- all the same pull/check/return shape. A node + * stream carrying 'yield' markers passes them through by accepting them in + * the predicate. + */ +export function filterPull( + stream: PullStream, + keep: (value: T) => boolean, +): PullStream { + return { + next() { + for (;;) { + const v = stream.next(); + if (v === undefined || keep(v)) { + return v; + } + } + }, + close: () => stream.close(), + }; +} + +/** Ends the stream at the first value `keep` rejects, closing the source. */ +export function takeWhilePull( + stream: PullStream, + keep: (value: T) => boolean, +): PullStream { + let done = false; + return { + next() { + if (done) { + return undefined; + } + const v = stream.next(); + if (v === undefined) { + done = true; + return undefined; + } + if (!keep(v)) { + done = true; + stream.close(); + return undefined; + } + return v; + }, + close() { + if (!done) { + done = true; + stream.close(); + } + }, + }; +} + +/** Applies `map` to each value. */ +export function mapPull( + stream: PullStream, + map: (value: T) => U, +): PullStream { + return { + next() { + const v = stream.next(); + return v === undefined ? undefined : map(v); + }, + close: () => stream.close(), + }; +} + +/** + * Emits at most `limit` values, calling `onValue` for each and `onComplete` + * once the scan finishes. + * + * `Take` and `Cap` both hydrate this way: read up to a limit, record what was + * seen, and treat a consumer that stops early as a bug -- their initial fetch + * must run to completion or the state they persist is wrong. Closing early + * still records, then raises, exactly as their generators' `finally` did. + */ +export function limitedScan( + stream: PullStream, + limit: number, + isValue: (v: T) => boolean, + onValue: (v: T) => void, + onComplete: () => void, + onEarlyClose: () => void, +): PullStream { + let seen = 0; + let done = false; + const complete = () => { + done = true; + stream.close(); + onComplete(); + }; + return { + next() { + if (done) { + return undefined; + } + if (seen === limit) { + complete(); + return undefined; + } + let v: T | undefined; + try { + v = stream.next(); + } catch (e) { + // As the generators did: an exception records no state. + done = true; + throw e; + } + if (v === undefined) { + complete(); + return undefined; + } + if (isValue(v)) { + onValue(v); + seen++; + } + return v; + }, + close() { + if (!done) { + complete(); + onEarlyClose(); + } + }, + }; +} + +/** A pull stream over a fixed list; for producers that already have an array. */ +class ArrayPull implements PullStream { + readonly #items: readonly T[]; + #i = 0; + constructor(items: readonly T[]) { + this.#items = items; + } + next(): T | undefined { + return this.#i < this.#items.length ? this.#items[this.#i++] : undefined; + } + close(): void { + this.#i = this.#items.length; + } +} + +/** + * Reads a pull stream to completion, applying `map` to each value as it + * arrives. + * + * The interleaving matters: `Array.from(iter, fn)` called `fn` between pulls, + * and callers such as `Catch` rely on that -- expanding a node's + * relationships triggers child fetches, so draining first and mapping second + * reorders those fetches relative to the parent scan. + */ +export function drainPullMap( + stream: PullStream, + map: (value: T) => U, +): U[] { + const out: U[] = []; + for (let v = stream.next(); v !== undefined; v = stream.next()) { + out.push(map(v)); + } + return out; +} + +/** Reads a pull stream to completion. For tests and for `Catch`. */ +export function drainPull(stream: PullStream): T[] { + const out: T[] = []; + for (let v = stream.next(); v !== undefined; v = stream.next()) { + out.push(v); + } + return out; +} + +export function pullOf(items: readonly T[]): PullStream { + return items.length === 0 ? emptyPullStream() : new ArrayPull(items); +} + +export function emptyPullStream(): PullStream { + return EMPTY; +} + +/** + * A pull stream whose work starts on the first `next()`, as a generator body + * does. Lets a source defer reading mutable state until iteration actually + * begins. + */ +export class LazyPullStream implements PullStream { + #start: (() => PullStream) | undefined; + #inner: PullStream | undefined; + + constructor(start: () => PullStream) { + this.#start = start; + } + + next(): T | undefined { + let inner = this.#inner; + if (inner === undefined) { + const start = this.#start; + if (start === undefined) { + return undefined; + } + this.#start = undefined; + inner = this.#inner = start(); + } + return inner.next(); + } + + close(): void { + this.#start = undefined; + const inner = this.#inner; + this.#inner = undefined; + inner?.close(); + } +} diff --git a/packages/zql/src/ivm/take.fetch.test.ts b/packages/zql/src/ivm/take.fetch.test.ts index bb25492279..45c6ab9d0a 100644 --- a/packages/zql/src/ivm/take.fetch.test.ts +++ b/packages/zql/src/ivm/take.fetch.test.ts @@ -12,8 +12,7 @@ import type {Node} from './data.ts'; import {MemoryStorage} from './memory-storage.ts'; import type {FetchRequest} from './operator.ts'; import {Snitch, type SnitchMessage} from './snitch.ts'; -import type {Stream} from './stream.ts'; -import {consume} from './stream.ts'; +import {consume, drainPull, type PullStream} from './stream.ts'; import {Take, type PartitionKey} from './take.ts'; import {createSource} from './test/source-factory.ts'; @@ -382,7 +381,7 @@ suite('take with no partition', () => { }); class ThrowingSnitch extends Snitch { - fetch(_: FetchRequest): Stream { + fetch(_: FetchRequest): PullStream { throw new Error('ThrowingSnitch error'); } } @@ -401,7 +400,7 @@ test('exception during hydrate', () => { const limit = 10; const take = new Take(snitch, storage, limit); - expect(() => [...take.fetch({})]).toThrow('ThrowingSnitch error'); + expect(() => drainPull(take.fetch({}))).toThrow('ThrowingSnitch error'); }); test('early return during hydrate', () => { @@ -424,11 +423,17 @@ test('early return during hydrate', () => { const take = new Take(snitch, storage, limit); expect(() => { let count = 0; - for (const _ of take.fetch({})) { - count++; - if (count > 1) { - break; + + const stream = take.fetch({}); + try { + for (let _ = stream.next(); _ !== undefined; _ = stream.next()) { + count++; + if (count > 1) { + break; + } } + } finally { + stream.close(); } }).toThrow('Unexpected early return prevented full hydration'); }); diff --git a/packages/zql/src/ivm/take.ts b/packages/zql/src/ivm/take.ts index 0595b3a12e..8428ad2939 100644 --- a/packages/zql/src/ivm/take.ts +++ b/packages/zql/src/ivm/take.ts @@ -22,7 +22,13 @@ import { type Storage, } from './operator.ts'; import type {SourceSchema} from './schema.ts'; -import {type Stream} from './stream.ts'; +import { + emptyPullStream, + LazyPullStream, + type PullStream, + type Stream, + limitedScan, +} from './stream.ts'; const MAX_BOUND_KEY = 'maxBound'; @@ -90,7 +96,14 @@ export class Take implements Operator { return this.#input.getSchema(); } - *fetch(req: FetchRequest): Stream { + fetch(req: FetchRequest): PullStream { + // Lazy: the generator read take state on first next(), and a push between + // fetch() and iteration must still be visible. + return new LazyPullStream(() => this.#startFetch(req)); + } + + #startFetch(req: FetchRequest): PullStream { + const compareRows = this.getSchema().compareRows; if ( !this.#partitionKey || (req.constraint && @@ -99,120 +112,75 @@ export class Take implements Operator { const takeStateKey = getTakeStateKey(this.#partitionKey, req.constraint); const takeState = this.#storage.get(takeStateKey); if (!takeState) { - yield* this.#initialFetch(req); - return; + return this.#initialFetch(req, takeStateKey); } if (takeState.bound === undefined) { - return; + return emptyPullStream(); } - for (const inputNode of this.#input.fetch(req)) { - if (inputNode === 'yield') { - yield inputNode; - continue; - } - if (this.getSchema().compareRows(takeState.bound, inputNode.row) < 0) { - return; - } - if ( - this.#rowHiddenFromFetch && - this.getSchema().compareRows( - this.#rowHiddenFromFetch, - inputNode.row, - ) === 0 - ) { - continue; - } - yield inputNode; - } - return; + const bound = takeState.bound; + const hidden = this.#rowHiddenFromFetch; + return new TakeScanPull(this.#input.fetch(req), node => + compareRows(bound, node.row) < 0 + ? 'stop' + : hidden && compareRows(hidden, node.row) === 0 + ? 'skip' + : 'emit', + ); } - // There is a partition key, but the fetch is not constrained or constrained - // on a different key. Thus we don't have a single take state to bound by. - // This currently only happens with nested sub-queries - // e.g. issues include issuelabels include label. We could remove this - // case if we added a translation layer (powered by some state) in join. - // Specifically we need joinKeyValue => parent constraint key const maxBound = this.#storage.get(MAX_BOUND_KEY); if (maxBound === undefined) { - return; + return emptyPullStream(); } - for (const inputNode of this.#input.fetch(req)) { - if (inputNode === 'yield') { - yield inputNode; - continue; - } - if (this.getSchema().compareRows(inputNode.row, maxBound) > 0) { - return; + return new TakeScanPull(this.#input.fetch(req), node => { + if (compareRows(node.row, maxBound) > 0) { + return 'stop'; } - const takeStateKey = getTakeStateKey(this.#partitionKey, inputNode.row); - const takeState = this.#storage.get(takeStateKey); - if ( - takeState?.bound !== undefined && - this.getSchema().compareRows(takeState.bound, inputNode.row) >= 0 - ) { - yield inputNode; - } - } + const takeState = this.#storage.get( + getTakeStateKey(this.#partitionKey, node.row), + ); + return takeState?.bound !== undefined && + compareRows(takeState.bound, node.row) >= 0 + ? 'emit' + : 'skip'; + }); } - *#initialFetch(req: FetchRequest): Stream { + #initialFetch( + req: FetchRequest, + takeStateKey: string, + ): PullStream { assert(req.start === undefined, 'Start should be undefined'); assert(!req.reverse, 'Reverse should be false'); - if (this.#limit === 0) { - return; + return emptyPullStream(); } - assert( constraintMatchesPartitionKey(req.constraint, this.#partitionKey), 'Constraint should match partition key', ); - - const takeStateKey = getTakeStateKey(this.#partitionKey, req.constraint); assert( this.#storage.get(takeStateKey) === undefined, 'Take state should be undefined', ); - let size = 0; let bound: Row | undefined; - let downstreamEarlyReturn = true; - let exceptionThrown = false; - try { - for (const inputNode of this.#input.fetch(req)) { - if (inputNode === 'yield') { - yield 'yield'; - continue; - } - yield inputNode; - bound = inputNode.row; + return limitedScan( + this.#input.fetch(req), + this.#limit, + node => node !== 'yield', + node => { + bound = (node as Node).row; size++; - if (size === this.#limit) { - break; - } - } - downstreamEarlyReturn = false; - } catch (e) { - exceptionThrown = true; - throw e; - } finally { - if (!exceptionThrown) { + }, + () => this.#setTakeState( takeStateKey, size, bound, this.#storage.get(MAX_BOUND_KEY), - ); - // If it becomes necessary to support downstream early return, this - // assert should be removed, and replaced with code that consumes - // the input stream until limit is reached or the input stream is - // exhausted so that takeState is properly hydrated. - assert( - !downstreamEarlyReturn, - 'Unexpected early return prevented full hydration', - ); - } - } + ), + () => assert(false, 'Unexpected early return prevented full hydration'), + ); } #getStateAndConstraint(row: Row) { @@ -283,38 +251,48 @@ export class Take implements Operator { let beforeBoundNode: Node | undefined; let boundNode: Node | undefined; if (this.#limit === 1) { - for (const node of this.#input.fetch({ + const rows = this.#input.fetch({ start: { row: takeState.bound, basis: 'at', }, constraint, - })) { - if (node === 'yield') { - yield node; - continue; + }); + try { + for (let node = rows.next(); node !== undefined; node = rows.next()) { + if (node === 'yield') { + yield node; + continue; + } + boundNode = node; + break; } - boundNode = node; - break; + } finally { + rows.close(); } } else { - for (const node of this.#input.fetch({ + const rows = this.#input.fetch({ start: { row: takeState.bound, basis: 'at', }, constraint, reverse: true, - })) { - if (node === 'yield') { - yield node; - continue; - } else if (boundNode === undefined) { - boundNode = node; - } else { - beforeBoundNode = node; - break; + }); + try { + for (let node = rows.next(); node !== undefined; node = rows.next()) { + if (node === 'yield') { + yield node; + continue; + } else if (boundNode === undefined) { + boundNode = node; + } else { + beforeBoundNode = node; + break; + } } + } finally { + rows.close(); } } assert( @@ -352,20 +330,26 @@ export class Take implements Operator { return; } let beforeBoundNode: Node | undefined; - for (const node of this.#input.fetch({ + + const rows = this.#input.fetch({ start: { row: takeState.bound, basis: 'after', }, constraint, reverse: true, - })) { - if (node === 'yield') { - yield node; - continue; + }); + try { + for (let node = rows.next(); node !== undefined; node = rows.next()) { + if (node === 'yield') { + yield node; + continue; + } + beforeBoundNode = node; + break; } - beforeBoundNode = node; - break; + } finally { + rows.close(); } let newBound: {node: Node; push: boolean} | undefined; @@ -377,25 +361,30 @@ export class Take implements Operator { }; } if (!newBound?.push) { - for (const node of this.#input.fetch({ + const rows = this.#input.fetch({ start: { row: takeState.bound, basis: 'at', }, constraint, - })) { - if (node === 'yield') { - yield node; - continue; - } - const push = compareRows(node.row, takeState.bound) > 0; - newBound = { - node, - push, - }; - if (push) { - break; + }); + try { + for (let node = rows.next(); node !== undefined; node = rows.next()) { + if (node === 'yield') { + yield node; + continue; + } + const push = compareRows(node.row, takeState.bound) > 0; + newBound = { + node, + push, + }; + if (push) { + break; + } } + } finally { + rows.close(); } } @@ -484,20 +473,26 @@ export class Take implements Operator { // bounds. let beforeBoundNode: Node | undefined; - for (const node of this.#input.fetch({ + + const rows = this.#input.fetch({ start: { row: takeState.bound, basis: 'after', }, constraint, reverse: true, - })) { - if (node === 'yield') { - yield node; - continue; + }); + try { + for (let node = rows.next(); node !== undefined; node = rows.next()) { + if (node === 'yield') { + yield node; + continue; + } + beforeBoundNode = node; + break; } - beforeBoundNode = node; - break; + } finally { + rows.close(); } assert( beforeBoundNode !== undefined, @@ -517,19 +512,25 @@ export class Take implements Operator { assert(newCmp > 0, 'New comparison must be greater than 0'); // Find the first item at the old bounds. This will be the new bounds. let newBoundNode: Node | undefined; - for (const node of this.#input.fetch({ + + const rows = this.#input.fetch({ start: { row: takeState.bound, basis: 'at', }, constraint, - })) { - if (node === 'yield') { - yield node; - continue; + }); + try { + for (let node = rows.next(); node !== undefined; node = rows.next()) { + if (node === 'yield') { + yield node; + continue; + } + newBoundNode = node; + break; } - newBoundNode = node; - break; + } finally { + rows.close(); } assert( newBoundNode !== undefined, @@ -572,23 +573,29 @@ export class Take implements Operator { let oldBoundNode: Node | undefined; let newBoundNode: Node | undefined; - for (const node of this.#input.fetch({ + + const rows = this.#input.fetch({ start: { row: takeState.bound, basis: 'at', }, constraint, reverse: true, - })) { - if (node === 'yield') { - yield node; - continue; - } else if (oldBoundNode === undefined) { - oldBoundNode = node; - } else { - newBoundNode = node; - break; + }); + try { + for (let node = rows.next(); node !== undefined; node = rows.next()) { + if (node === 'yield') { + yield node; + continue; + } else if (oldBoundNode === undefined) { + oldBoundNode = node; + } else { + newBoundNode = node; + break; + } } + } finally { + rows.close(); } assert( oldBoundNode !== undefined, @@ -632,19 +639,25 @@ export class Take implements Operator { // at this point we need to find the row after the bound and use that or // the newRow as the new bound. let afterBoundNode: Node | undefined; - for (const node of this.#input.fetch({ + + const rows = this.#input.fetch({ start: { row: takeState.bound, basis: 'after', }, constraint, - })) { - if (node === 'yield') { - yield node; - continue; + }); + try { + for (let node = rows.next(); node !== undefined; node = rows.next()) { + if (node === 'yield') { + yield node; + continue; + } + afterBoundNode = node; + break; } - afterBoundNode = node; - break; + } finally { + rows.close(); } assert( afterBoundNode !== undefined, @@ -755,3 +768,52 @@ export function makePartitionKeyComparator( return 0; }; } + +/** + * A Take scan in the pull protocol: forwards 'yield', and asks `decide` per + * node whether to emit it, skip it, or stop (closing the input). + */ +class TakeScanPull implements PullStream { + readonly #input: PullStream; + readonly #decide: (node: Node) => 'emit' | 'skip' | 'stop'; + #done = false; + + constructor( + input: PullStream, + decide: (node: Node) => 'emit' | 'skip' | 'stop', + ) { + this.#input = input; + this.#decide = decide; + } + + next(): Node | 'yield' | undefined { + if (this.#done) { + return undefined; + } + for (;;) { + const v = this.#input.next(); + if (v === undefined) { + this.#done = true; + return undefined; + } + if (v === 'yield') { + return v; + } + const d = this.#decide(v); + if (d === 'emit') { + return v; + } + if (d === 'stop') { + this.close(); + return undefined; + } + } + } + + close(): void { + if (!this.#done) { + this.#done = true; + this.#input.close(); + } + } +} diff --git a/packages/zql/src/ivm/test/mode-yield-source.ts b/packages/zql/src/ivm/test/mode-yield-source.ts index 262623c820..c002042df5 100644 --- a/packages/zql/src/ivm/test/mode-yield-source.ts +++ b/packages/zql/src/ivm/test/mode-yield-source.ts @@ -13,7 +13,7 @@ import type {DebugDelegate} from '../../builder/debug-delegate.ts'; import type {Node} from '../data.ts'; import type {FetchRequest} from '../operator.ts'; import type {Source, SourceChange, SourceInput} from '../source.ts'; -import type {Stream} from '../stream.ts'; +import type {PullStream, Stream} from '../stream.ts'; export type YieldMode = 'fetch' | 'push' | 'both'; @@ -67,16 +67,50 @@ export class ModeYieldSource implements Source { const originalFetch = sourceInput.fetch.bind(sourceInput); return { ...sourceInput, - *fetch(req: FetchRequest): Stream { - for (const item of originalFetch(req)) { - if (rng() < p) { - yield 'yield'; + fetch: (req: FetchRequest): PullStream => { + const src = originalFetch(req); + // A 'yield' is emitted before an item, so the item is held until the + // marker has been handed back; one more may follow the last item. + let pending: Node | 'yield' | undefined; + let exhausted = false; + let tailDone = false; + const tail = (): 'yield' | undefined => { + if (!tailDone) { + tailDone = true; + if (rng() < p) { + return 'yield'; + } } - yield item; - } - if (rng() < p) { - yield 'yield'; - } + return undefined; + }; + return { + next(): Node | 'yield' | undefined { + if (pending !== undefined) { + const held = pending; + pending = undefined; + return held; + } + if (exhausted) { + return tail(); + } + const item = src.next(); + if (item === undefined) { + exhausted = true; + return tail(); + } + if (rng() < p) { + pending = item; + return 'yield'; + } + return item; + }, + close() { + exhausted = true; + tailDone = true; + pending = undefined; + src.close(); + }, + }; }, }; } diff --git a/packages/zql/src/ivm/test/random-yield-source.ts b/packages/zql/src/ivm/test/random-yield-source.ts index 157ad7ffbe..b4087d17d3 100644 --- a/packages/zql/src/ivm/test/random-yield-source.ts +++ b/packages/zql/src/ivm/test/random-yield-source.ts @@ -4,7 +4,7 @@ import type {DebugDelegate} from '../../builder/debug-delegate.ts'; import type {Node} from '../data.ts'; import type {FetchRequest} from '../operator.ts'; import type {Source, SourceChange, SourceInput} from '../source.ts'; -import type {Stream} from '../stream.ts'; +import type {PullStream, Stream} from '../stream.ts'; /** * A source wrapper that randomly injects 'yield' values into fetch and push @@ -59,22 +59,53 @@ export class RandomYieldSource implements Source { const wrappedInput: SourceInput = { ...sourceInput, - *fetch(req: FetchRequest): Stream { - for (const item of originalFetch(req)) { - // Check for abort (can throw) - checkAbort?.(); - // Randomly yield before each item - if (rng() < yieldProbability) { - yield 'yield'; + fetch: (req: FetchRequest): PullStream => { + const src = originalFetch(req); + let pending: Node | 'yield' | undefined; + let exhausted = false; + let tailDone = false; + const tail = (): 'yield' | undefined => { + if (!tailDone) { + tailDone = true; + // Check for abort at the end (can throw) + checkAbort?.(); + if (rng() < yieldProbability) { + return 'yield'; + } } - yield item; - } - // Check for abort at the end - checkAbort?.(); - // Randomly yield at the end - if (rng() < yieldProbability) { - yield 'yield'; - } + return undefined; + }; + return { + next(): Node | 'yield' | undefined { + if (pending !== undefined) { + const held = pending; + pending = undefined; + return held; + } + if (exhausted) { + return tail(); + } + const item = src.next(); + if (item === undefined) { + exhausted = true; + return tail(); + } + // Check for abort (can throw) + checkAbort?.(); + // Randomly yield before each item + if (rng() < yieldProbability) { + pending = item; + return 'yield'; + } + return item; + }, + close() { + exhausted = true; + tailDone = true; + pending = undefined; + src.close(); + }, + }; }, }; diff --git a/packages/zql/src/ivm/union-fan-in.test.ts b/packages/zql/src/ivm/union-fan-in.test.ts index 8e8710e602..2a1fb7cd3b 100644 --- a/packages/zql/src/ivm/union-fan-in.test.ts +++ b/packages/zql/src/ivm/union-fan-in.test.ts @@ -3,6 +3,7 @@ import {describe, expect, test, vi} from 'vitest'; import type {Node} from './data.js'; import {skipYields, type FetchRequest, type Operator} from './operator.js'; import type {SourceSchema} from './schema.js'; +import {drainPull, pullOf} from './stream.ts'; import {UnionFanIn} from './union-fan-in.js'; import type {UnionFanOut} from './union-fan-out.js'; @@ -19,7 +20,7 @@ const mockSchema: SourceSchema = { const mockOperator = (schema: SourceSchema, data: Node[] = []): Operator => ({ getSchema: () => schema, - fetch: (_req: FetchRequest) => data, + fetch: (_req: FetchRequest) => pullOf(data), push: vi.fn(), setOutput: vi.fn(), destroy: vi.fn(), @@ -209,7 +210,7 @@ describe('UnionFanIn', () => { const input2 = mockOperator(mockSchema, data2); const fanIn = new UnionFanIn(fanOut, [input1, input2]); - const result = [...skipYields(fanIn.fetch({} as FetchRequest))]; + const result = drainPull(skipYields(fanIn.fetch({} as FetchRequest))); expect(result).toHaveLength(4); expect(result.map(n => n.row.id)).toEqual([1, 2, 3, 4]); @@ -219,7 +220,7 @@ describe('UnionFanIn', () => { const fanOut = mockUnionFanOut(mockSchema); const fanIn = new UnionFanIn(fanOut, []); - const result = [...fanIn.fetch({} as FetchRequest)]; + const result = drainPull(fanIn.fetch({} as FetchRequest)); expect(result).toHaveLength(0); }); @@ -238,7 +239,7 @@ describe('UnionFanIn', () => { const input2 = mockOperator(mockSchema, data2); const fanIn = new UnionFanIn(fanOut, [input1, input2]); - const result = [...skipYields(fanIn.fetch({} as FetchRequest))]; + const result = drainPull(skipYields(fanIn.fetch({} as FetchRequest))); expect(result).toHaveLength(3); expect(result.map(n => n.row.id)).toEqual([1, 2, 3]); @@ -252,7 +253,7 @@ describe('UnionFanIn', () => { ): Operator => ({ getSchema: () => schema, fetch: (req: FetchRequest) => - req.reverse ? data.toReversed() : [...data], + pullOf(req.reverse ? data.toReversed() : [...data]), push: vi.fn(), setOutput: vi.fn(), destroy: vi.fn(), @@ -277,9 +278,9 @@ describe('UnionFanIn', () => { const inputB = mockOrderedOperator(mockSchema, dataB); const fanIn = new UnionFanIn(fanOut, [inputA, inputB]); - const result = [ - ...skipYields(fanIn.fetch({reverse: true} as FetchRequest)), - ]; + const result = drainPull( + skipYields(fanIn.fetch({reverse: true} as FetchRequest)), + ); // BUG: UnionFanIn.fetch hands mergeFetches an ascending comparator // regardless of req.reverse, so the merge picks min across descending @@ -308,9 +309,9 @@ describe('UnionFanIn', () => { const inputB = mockOrderedOperator(mockSchema, dataB); const fanIn = new UnionFanIn(fanOut, [inputA, inputB]); - const result = [ - ...skipYields(fanIn.fetch({reverse: true} as FetchRequest)), - ]; + const result = drainPull( + skipYields(fanIn.fetch({reverse: true} as FetchRequest)), + ); expect(result.map(n => n.row.id)).toEqual([4, 3, 2, 1]); }); diff --git a/packages/zql/src/ivm/union-fan-in.ts b/packages/zql/src/ivm/union-fan-in.ts index 57bc1e533f..602e7f84d7 100644 --- a/packages/zql/src/ivm/union-fan-in.ts +++ b/packages/zql/src/ivm/union-fan-in.ts @@ -19,7 +19,7 @@ import { pushAccumulatedChanges, } from './push-accumulated.ts'; import type {SourceSchema} from './schema.ts'; -import type {Stream} from './stream.ts'; +import {type Stream, type PullStream} from './stream.ts'; import type {UnionFanOut} from './union-fan-out.ts'; export class UnionFanIn implements Operator { @@ -100,7 +100,7 @@ export class UnionFanIn implements Operator { } } - fetch(req: FetchRequest): Stream { + fetch(req: FetchRequest): PullStream { const iterables = this.#inputs.map(input => input.fetch(req)); const compareRows = this.#schema.compareRows; const compare = req.reverse @@ -178,13 +178,23 @@ export class UnionFanIn implements Operator { // looked like a branch holding the row, silently dropping the // add/remove and desyncing a downstream `Take`'s push and fetch paths. let otherBranchHasRow = false; - for (const node of fetchResult) { - if (node === 'yield') { - yield node; - continue; + + const branchRows = fetchResult; + try { + for ( + let node = branchRows.next(); + node !== undefined; + node = branchRows.next() + ) { + if (node === 'yield') { + yield node; + continue; + } + otherBranchHasRow = true; + break; } - otherBranchHasRow = true; - break; + } finally { + branchRows.close(); } if (otherBranchHasRow) { @@ -238,78 +248,123 @@ export class UnionFanIn implements Operator { } } -export function* mergeFetches( - fetches: Iterable[], +export function mergeFetches( + fetches: PullStream[], comparator: (l: Node, r: Node) => number, -): IterableIterator { - const iterators = fetches.map(i => i[Symbol.iterator]()); - let threw = false; - try { - const current: (Node | null)[] = []; - let lastNodeYielded: Node | undefined; - for (let i = 0; i < iterators.length; i++) { - const iter = iterators[i]; - let result = iter.next(); - // yield yields when initializing - while (!result.done && result.value === 'yield') { - yield result.value; - result = iter.next(); - } - current[i] = result.done ? null : (result.value as Node); +): PullStream { + return new MergeFetches(fetches, comparator); +} + +/** + * Linear-scan merge of pre-sorted branches, dropping duplicates that compare + * equal to the last emitted node. + * + * The generator this replaces suspended inside its "advance this branch" loop + * to forward a 'yield'. That point is now `#refill`: the branch owing a + * replacement for the node just emitted, so a 'yield' can be returned and the + * merge resumed at the same place. + */ +class MergeFetches implements PullStream { + readonly #streams: readonly PullStream[]; + readonly #comparator: (l: Node, r: Node) => number; + readonly #current: (Node | null)[]; + #lastEmitted: Node | undefined; + /** Node selected but not yet emitted: its branch is refilled first. */ + #held: Node | undefined; + #primeIdx = 0; + #priming = true; + #refill: number | undefined; + #done = false; + + constructor( + streams: readonly PullStream[], + comparator: (l: Node, r: Node) => number, + ) { + this.#streams = streams; + this.#comparator = comparator; + this.#current = new Array(streams.length).fill(null); + } + + /** A Node, 'yield' to forward, or undefined when the branch is spent. */ + #pullOne(idx: number): Node | 'yield' | undefined { + return this.#streams[idx].next(); + } + + next(): Node | 'yield' | undefined { + if (this.#done) { + return undefined; } - while (current.some(c => c !== null)) { - const min = current.reduce( - (acc: [Node, number] | undefined, c, i): [Node, number] | undefined => { - if (c === null) { - return acc; + for (;;) { + if (this.#priming) { + while (this.#primeIdx < this.#streams.length) { + const v = this.#pullOne(this.#primeIdx); + if (v === 'yield') { + return v; } - if (acc === undefined || comparator(c, acc[0]) < 0) { - return [c, i]; + this.#current[this.#primeIdx] = v === undefined ? null : v; + this.#primeIdx++; + } + this.#priming = false; + } + + if (this.#refill !== undefined) { + const idx = this.#refill; + const v = this.#pullOne(idx); + if (v === 'yield') { + return v; + } + this.#refill = undefined; + this.#current[idx] = v === undefined ? null : v; + // The branch has been advanced; now the held node may be emitted. + // Order matters: the generator forwarded a branch's 'yield's before + // emitting the node it had selected from that branch. + const held = this.#held; + this.#held = undefined; + if (held !== undefined) { + if ( + this.#lastEmitted !== undefined && + this.#comparator(this.#lastEmitted, held) === 0 + ) { + continue; } - return acc; - }, - undefined, - ); + this.#lastEmitted = held; + return held; + } + continue; + } - assert(min !== undefined, 'min is undefined'); - const [minNode, minIndex] = min; - const iter = iterators[minIndex]; - let result = iter.next(); - while (!result.done && result.value === 'yield') { - yield result.value; - result = iter.next(); + let minNode: Node | undefined; + let minIndex = -1; + for (let i = 0; i < this.#current.length; i++) { + const c = this.#current[i]; + if (c === null) { + continue; + } + if (minNode === undefined || this.#comparator(c, minNode) < 0) { + minNode = c; + minIndex = i; + } } - current[minIndex] = result.done ? null : (result.value as Node); - if ( - lastNodeYielded !== undefined && - comparator(lastNodeYielded, minNode) === 0 - ) { - continue; + if (minNode === undefined) { + this.close(); + return undefined; } - lastNodeYielded = minNode; - yield minNode; + + // Hold the node and advance its branch first; the duplicate check and + // the emit both happen once the refill completes. + this.#held = minNode; + this.#refill = minIndex; } - } catch (e) { - threw = true; - for (const iter of iterators) { - try { - iter.throw?.(e); - } catch (_cleanupError) { - // error in the iter.throw cleanup, - // catch so other iterators are cleaned up - } + } + + close(): void { + if (this.#done) { + return; } - throw e; - } finally { - if (!threw) { - for (const iter of iterators) { - try { - iter.return?.(); - } catch (_cleanupError) { - // error in the iter.return cleanup, - // catch so other iterators are cleaned up - } - } + this.#done = true; + this.#held = undefined; + for (const s of this.#streams) { + s.close(); } } } diff --git a/packages/zql/src/ivm/union-fan-out.test.ts b/packages/zql/src/ivm/union-fan-out.test.ts index 1cfb051873..c1f39fc5ef 100644 --- a/packages/zql/src/ivm/union-fan-out.test.ts +++ b/packages/zql/src/ivm/union-fan-out.test.ts @@ -2,7 +2,7 @@ import {expect, test, vi} from 'vitest'; import {testLogConfig} from '../../../otel/src/test-log-config.ts'; import {createSilentLogContext} from '../../../shared/src/logging-test-utils.ts'; import {Catch} from './catch.ts'; -import {consume} from './stream.ts'; +import {consume, drainPull} from './stream.ts'; import {createSource} from './test/source-factory.ts'; import {UnionFanOut} from './union-fan-out.ts'; @@ -146,7 +146,7 @@ test('fetch delegates to input', () => { const connector = s.connect([['a', 'asc']]); const fanOut = new UnionFanOut(connector); - const result = [...fanOut.fetch({})]; + const result = drainPull(fanOut.fetch({})); expect(result).toMatchInlineSnapshot(` [ { diff --git a/packages/zql/src/ivm/union-fan-out.ts b/packages/zql/src/ivm/union-fan-out.ts index ea8a7d14ff..4afdf6ed84 100644 --- a/packages/zql/src/ivm/union-fan-out.ts +++ b/packages/zql/src/ivm/union-fan-out.ts @@ -5,7 +5,7 @@ import type {Change} from './change.ts'; import type {Node} from './data.ts'; import type {FetchRequest, Input, Operator, Output} from './operator.ts'; import type {SourceSchema} from './schema.ts'; -import type {Stream} from './stream.ts'; +import type {PullStream, Stream} from './stream.ts'; import type {UnionFanIn} from './union-fan-in.ts'; export class UnionFanOut implements Operator { @@ -40,7 +40,7 @@ export class UnionFanOut implements Operator { return this.#input.getSchema(); } - fetch(req: FetchRequest): Stream { + fetch(req: FetchRequest): PullStream { return this.#input.fetch(req); } diff --git a/packages/zql/src/ivm/view-apply-change.test.ts b/packages/zql/src/ivm/view-apply-change.test.ts index 2a74f19997..7f7c549d98 100644 --- a/packages/zql/src/ivm/view-apply-change.test.ts +++ b/packages/zql/src/ivm/view-apply-change.test.ts @@ -1,6 +1,7 @@ import {describe, expect, test} from 'vitest'; import {makeComparator} from './data.ts'; import type {SourceSchema} from './schema.ts'; +import {emptyPullStream, pullOf} from './stream.ts'; import { applyChange, idSymbol, @@ -109,7 +110,7 @@ describe('applyChange', () => { name: 'Buffalo Big Board Classic', }, relationships: { - athletes: () => [], + athletes: () => emptyPullStream(), }, }, }, @@ -133,16 +134,17 @@ describe('applyChange', () => { disciplineID: 'd1', }, relationships: { - athletes: () => [ - { - row: { - country: 'USA', - id: 'a1', - name: 'Mason Ho', + athletes: () => + pullOf([ + { + row: { + country: 'USA', + id: 'a1', + name: 'Mason Ho', + }, + relationships: {}, }, - relationships: {}, - }, - ], + ]), }, }, }, @@ -168,16 +170,17 @@ describe('applyChange', () => { disciplineID: 'd2', }, relationships: { - athletes: () => [ - { - row: { - country: 'USA', - id: 'a1', - name: 'Mason Ho', + athletes: () => + pullOf([ + { + row: { + country: 'USA', + id: 'a1', + name: 'Mason Ho', + }, + relationships: {}, }, - relationships: {}, - }, - ], + ]), }, }, }, @@ -232,16 +235,17 @@ describe('applyChange', () => { disciplineID: 'd1', }, relationships: { - athletes: () => [ - { - row: { - country: 'USA', - id: 'a1', - name: 'Mason Ho', + athletes: () => + pullOf([ + { + row: { + country: 'USA', + id: 'a1', + name: 'Mason Ho', + }, + relationships: {}, }, - relationships: {}, - }, - ], + ]), }, }, }, @@ -296,16 +300,17 @@ describe('applyChange', () => { disciplineID: 'd2', }, relationships: { - athletes: () => [ - { - row: { - country: 'USA', - id: 'a1', - name: 'Mason Ho', + athletes: () => + pullOf([ + { + row: { + country: 'USA', + id: 'a1', + name: 'Mason Ho', + }, + relationships: {}, }, - relationships: {}, - }, - ], + ]), }, }, }, @@ -366,7 +371,7 @@ describe('applyChange', () => { name: 'Buffalo Big Board Classic', }, relationships: { - athletes: () => [], + athletes: () => emptyPullStream(), }, }, }, @@ -390,16 +395,17 @@ describe('applyChange', () => { disciplineID: 'd1', }, relationships: { - athletes: () => [ - { - row: { - country: 'USA', - id: 'a1', - name: 'Mason Ho', + athletes: () => + pullOf([ + { + row: { + country: 'USA', + id: 'a1', + name: 'Mason Ho', + }, + relationships: {}, }, - relationships: {}, - }, - ], + ]), }, }, }, @@ -425,16 +431,17 @@ describe('applyChange', () => { disciplineID: 'd2', }, relationships: { - athletes: () => [ - { - row: { - country: 'USA', - id: 'a1', - name: 'Mason Ho', + athletes: () => + pullOf([ + { + row: { + country: 'USA', + id: 'a1', + name: 'Mason Ho', + }, + relationships: {}, }, - relationships: {}, - }, - ], + ]), }, }, }, @@ -487,16 +494,17 @@ describe('applyChange', () => { disciplineID: 'd1', }, relationships: { - athletes: () => [ - { - row: { - country: 'USA', - id: 'a1', - name: 'Mason Ho', + athletes: () => + pullOf([ + { + row: { + country: 'USA', + id: 'a1', + name: 'Mason Ho', + }, + relationships: {}, }, - relationships: {}, - }, - ], + ]), }, }, }, @@ -549,16 +557,17 @@ describe('applyChange', () => { disciplineID: 'd2', }, relationships: { - athletes: () => [ - { - row: { - country: 'USA', - id: 'a1', - name: 'Mason Ho', + athletes: () => + pullOf([ + { + row: { + country: 'USA', + id: 'a1', + name: 'Mason Ho', + }, + relationships: {}, }, - relationships: {}, - }, - ], + ]), }, }, }, @@ -870,7 +879,7 @@ describe('applyChange', () => { type: 'add', node: { row: {id: '1', name: 'Aaron'}, - relationships: makeProtoRelationships(() => []), + relationships: makeProtoRelationships(() => pullOf([])), }, }, schema, @@ -1716,9 +1725,8 @@ describe('applyChange', () => { node: { row: {id: 'b', name: 'Bob'}, relationships: { - children: () => [ - {row: {id: 'c1', parentId: 'b'}, relationships: {}}, - ], + children: () => + pullOf([{row: {id: 'c1', parentId: 'b'}, relationships: {}}]), }, }, }); @@ -1728,9 +1736,8 @@ describe('applyChange', () => { node: { row: {id: 'd', name: 'Dave'}, relationships: { - children: () => [ - {row: {id: 'c2', parentId: 'd'}, relationships: {}}, - ], + children: () => + pullOf([{row: {id: 'c2', parentId: 'd'}, relationships: {}}]), }, }, }); @@ -1742,9 +1749,8 @@ describe('applyChange', () => { node: { row: {id: 'a', name: 'Alice'}, relationships: { - children: () => [ - {row: {id: 'c3', parentId: 'a'}, relationships: {}}, - ], + children: () => + pullOf([{row: {id: 'c3', parentId: 'a'}, relationships: {}}]), }, }, }); @@ -1808,7 +1814,7 @@ describe('applyChange', () => { type: 'add', node: { row: {id: 'a', name: 'Alice'}, - relationships: {children: () => []}, + relationships: {children: () => emptyPullStream()}, }, }); @@ -1816,7 +1822,7 @@ describe('applyChange', () => { type: 'add', node: { row: {id: 'c', name: 'Charlie'}, - relationships: {children: () => []}, + relationships: {children: () => emptyPullStream()}, }, }); @@ -1826,10 +1832,11 @@ describe('applyChange', () => { node: { row: {id: 'b', name: 'Bob'}, relationships: { - children: () => [ - {row: {id: 'child1', parentId: 'b'}, relationships: {}}, - {row: {id: 'child2', parentId: 'b'}, relationships: {}}, - ], + children: () => + pullOf([ + {row: {id: 'child1', parentId: 'b'}, relationships: {}}, + {row: {id: 'child2', parentId: 'b'}, relationships: {}}, + ]), }, }, }); @@ -2067,10 +2074,11 @@ describe('applyChange', () => { node: { row: {id: 'p1', name: 'Parent1'}, relationships: { - children: () => [ - {row: {id: 'c1', parentId: 'p1'}, relationships: {}}, - {row: {id: 'c2', parentId: 'p1'}, relationships: {}}, - ], + children: () => + pullOf([ + {row: {id: 'c1', parentId: 'p1'}, relationships: {}}, + {row: {id: 'c2', parentId: 'p1'}, relationships: {}}, + ]), }, }, }); @@ -2193,9 +2201,8 @@ describe('applyChange', () => { node: { row: {id: 'p1', name: 'Parent1'}, relationships: { - children: () => [ - {row: {id: 'c1', parentId: 'p1'}, relationships: {}}, - ], + children: () => + pullOf([{row: {id: 'c1', parentId: 'p1'}, relationships: {}}]), }, }, }); @@ -2205,9 +2212,8 @@ describe('applyChange', () => { node: { row: {id: 'p2', name: 'Parent2'}, relationships: { - children: () => [ - {row: {id: 'c2', parentId: 'p2'}, relationships: {}}, - ], + children: () => + pullOf([{row: {id: 'c2', parentId: 'p2'}, relationships: {}}]), }, }, }); @@ -2536,10 +2542,11 @@ describe('applyChange', () => { node: { row: {id: 'p1', name: 'Parent1'}, relationships: { - children: () => [ - {row: {id: 'c1', parentId: 'p1'}, relationships: {}}, - {row: {id: 'c2', parentId: 'p1'}, relationships: {}}, - ], + children: () => + pullOf([ + {row: {id: 'c1', parentId: 'p1'}, relationships: {}}, + {row: {id: 'c2', parentId: 'p1'}, relationships: {}}, + ]), }, }, }); @@ -2674,9 +2681,8 @@ describe('applyChange', () => { node: { row: {id: 'p1', name: 'Parent1'}, relationships: { - children: () => [ - {row: {id: 'c1', parentId: 'p1'}, relationships: {}}, - ], + children: () => + pullOf([{row: {id: 'c1', parentId: 'p1'}, relationships: {}}]), }, }, }); @@ -2686,9 +2692,8 @@ describe('applyChange', () => { node: { row: {id: 'p2', name: 'Parent2'}, relationships: { - children: () => [ - {row: {id: 'c2', parentId: 'p2'}, relationships: {}}, - ], + children: () => + pullOf([{row: {id: 'c2', parentId: 'p2'}, relationships: {}}]), }, }, }); diff --git a/packages/zql/src/ivm/view-apply-change.ts b/packages/zql/src/ivm/view-apply-change.ts index 62bb716cbf..29d57f2db4 100644 --- a/packages/zql/src/ivm/view-apply-change.ts +++ b/packages/zql/src/ivm/view-apply-change.ts @@ -9,6 +9,8 @@ import type {Writable} from '../../../shared/src/writable.ts'; import type {Row} from '../../../zero-protocol/src/data.ts'; import {type Comparator, type Node} from './data.ts'; import {skipYields} from './operator.ts'; +import {pullOf, type PullStream} from './stream.ts'; + import type {SourceSchema} from './schema.ts'; import type {Entry, Format} from './view.ts'; @@ -104,18 +106,21 @@ export interface RefCountMap { /** * Get child nodes from a relationship, handling both lazy (Node) and expanded (ExpandedNode). */ -function* getChildNodes( +/** + * Reads a relationship in the pull protocol, skipping 'yield'. Returns the + * children as a PullStream for ExpandedNode arrays and adapted streams too, so + * the three call sites have one loop shape and the hot case -- a native + * PullStream from a Join -- pays no iterator, no generator and no result + * object per child. + */ +function childNodes( node: ViewNode, relationship: string, -): Generator { +): PullStream { const children = node.relationships[relationship]; - if (Array.isArray(children)) { - // ExpandedNode: already an array - yield* children; - } else { - // Node: lazy generator function - yield* skipYields(children()); - } + return Array.isArray(children) + ? pullOf(children) + : (skipYields(children()) as PullStream); } type Mutate = boolean; @@ -225,7 +230,12 @@ export function applyChangeInternal( let currentParent = parentEntry; for (const relationship of Object.keys(change.node.relationships)) { const childSchema = must(schema.relationships[relationship]); - for (const node of getChildNodes(change.node, relationship)) { + const children = childNodes(change.node, relationship); + for ( + let node = children.next(); + node !== undefined; + node = children.next() + ) { currentParent = applyChangeInternal( currentParent, {type: change.type, node}, @@ -643,7 +653,12 @@ function initializeRelationshipsForNewEntryIfAny( : track([] as MutableMetaEntryList); result[relationship] = newView; - for (const childNode of getChildNodes(node, relationship)) { + const children = childNodes(node, relationship); + for ( + let childNode = children.next(); + childNode !== undefined; + childNode = children.next() + ) { applyChangeInternal( result, {type: 'add', node: childNode}, @@ -658,7 +673,12 @@ function initializeRelationshipsForNewEntryIfAny( // Plural non-hidden: build array in-place for efficiency const childArray: MutableMetaEntryList = track([]); - for (const childNode of getChildNodes(node, relationship)) { + const children = childNodes(node, relationship); + for ( + let childNode = children.next(); + childNode !== undefined; + childNode = children.next() + ) { const newEntry = makeNewMetaEntry( childNode.row, childSchema, diff --git a/packages/zql/src/ivm/yield.fetch.test.ts b/packages/zql/src/ivm/yield.fetch.test.ts index 76c24cf2b8..bc29e1f107 100644 --- a/packages/zql/src/ivm/yield.fetch.test.ts +++ b/packages/zql/src/ivm/yield.fetch.test.ts @@ -10,7 +10,7 @@ import type {FetchRequest, Input, Output} from './operator.ts'; import type {SourceSchema} from './schema.ts'; import {Skip} from './skip.ts'; import {Snitch} from './snitch.ts'; -import type {Stream} from './stream.ts'; +import {pullOf, type PullStream} from './stream.ts'; import {Take} from './take.ts'; import {UnionFanIn} from './union-fan-in.ts'; import {UnionFanOut} from './union-fan-out.ts'; @@ -41,11 +41,13 @@ class YieldSource implements Input { return this.#schema; } - *fetch(_req: FetchRequest): Stream { - yield 'yield'; - yield {row: {id: '1'}, relationships: {}}; - yield 'yield'; - yield {row: {id: '2'}, relationships: {}}; + fetch(_req: FetchRequest): PullStream { + return pullOf([ + 'yield', + {row: {id: '1'}, relationships: {}}, + 'yield', + {row: {id: '2'}, relationships: {}}, + ]); } destroy(): void {} @@ -376,9 +378,18 @@ describe('Yield Propagation', () => { test('Error propagation during fetch', () => { const source = new YieldSource(); const error = new Error('Fetch failed'); - source.fetch = function* (_req: FetchRequest) { - yield 'yield'; - throw error; + source.fetch = (_req: FetchRequest) => { + let first = true; + return { + next: () => { + if (first) { + first = false; + return 'yield' as const; + } + throw error; + }, + close: () => {}, + }; }; const catchOp = new Catch(source); diff --git a/packages/zql/src/ivm/yield.push.test.ts b/packages/zql/src/ivm/yield.push.test.ts index 4599eb8baf..3d7c514f01 100644 --- a/packages/zql/src/ivm/yield.push.test.ts +++ b/packages/zql/src/ivm/yield.push.test.ts @@ -23,7 +23,7 @@ import { type SourceInput, } from './source.ts'; import type {Stream} from './stream.ts'; -import {consume} from './stream.ts'; +import {consume, drainPull, type PullStream} from './stream.ts'; import {Take} from './take.ts'; import {UnionFanIn} from './union-fan-in.ts'; import {UnionFanOut} from './union-fan-out.ts'; @@ -31,6 +31,18 @@ import {UnionFanOut} from './union-fan-out.ts'; class YieldOutput implements FilterOutput { yields: boolean = false; + /** The node whose 'yield' has already been emitted. */ + #yielded: Node | undefined; + + filter(node: Node): boolean | 'yield' { + if (this.yields && this.#yielded !== node) { + this.#yielded = node; + return 'yield'; + } + this.#yielded = undefined; + return true; + } + *push(_change: Change | SourceChange, _pusher: InputBase): Stream<'yield'> { if (this.yields) yield 'yield'; // Consume change @@ -39,10 +51,6 @@ class YieldOutput implements FilterOutput { beginFilter() {} endFilter() {} - *filter(_node: Node): Generator<'yield', boolean> { - if (this.yields) yield 'yield'; - return true; - } } class YieldMemorySource extends MemorySource { @@ -71,14 +79,49 @@ class YieldMemorySource extends MemorySource { const originalFetch = input.fetch.bind(input); const source = this; - input.fetch = function* (req: FetchRequest): Stream { - for (const n of originalFetch(req)) { - if (source.yieldOnFetch) { - yield 'yield'; + input.fetch = (req: FetchRequest): PullStream => { + const inner = originalFetch(req); + let pending: Node | 'yield' | undefined; + let exhausted = false; + let tailDone = false; + const tail = (): 'yield' | undefined => { + if (!tailDone) { + tailDone = true; + if (source.yieldOnFetch) { + return 'yield'; + } } - yield n; - } - if (source.yieldOnFetch) yield 'yield'; + return undefined; + }; + return { + next(): Node | 'yield' | undefined { + if (pending !== undefined) { + const held = pending; + pending = undefined; + return held; + } + if (exhausted) { + return tail(); + } + const n = inner.next(); + if (n === undefined) { + exhausted = true; + inner.close(); + return tail(); + } + if (source.yieldOnFetch) { + pending = n; + return 'yield'; + } + return n; + }, + close() { + exhausted = true; + tailDone = true; + pending = undefined; + inner.close(); + }, + }; }; return input; } @@ -165,7 +208,7 @@ describe('Yield Propagation (Push)', () => { take.setOutput(output); // Initialize Take - consume(take.fetch({})); + drainPull(take.fetch({})); expect(collectPush(source, makeAdd('0'))).toEqual(['yield', 'yield']); }); @@ -183,7 +226,7 @@ describe('Yield Propagation (Push)', () => { take.setOutput(output); // Initialize Take - consume(take.fetch({})); + drainPull(take.fetch({})); expect(collectPush(source, makeRemove('0'))).toEqual(['yield', 'yield']); }); @@ -203,7 +246,7 @@ describe('Yield Propagation (Push)', () => { take.setOutput(output); // Initialize Take - consume(take.fetch({})); + drainPull(take.fetch({})); // Push add '0'. This should displace '1'. const result = collectPush(source, makeAdd('0')); @@ -228,7 +271,7 @@ describe('Yield Propagation (Push)', () => { take.setOutput(output); // Initialize Take - consume(take.fetch({})); + drainPull(take.fetch({})); // Push remove '0'. const result = collectPush(source, makeRemove('0')); @@ -251,7 +294,7 @@ describe('Yield Propagation (Push)', () => { take.setOutput(output); // Initialize Take - consume(take.fetch({})); + drainPull(take.fetch({})); // Edit '0' to '2' (move out of bounds). const result = collectPush(source, makeEdit('2', '0')); @@ -275,7 +318,7 @@ describe('Yield Propagation (Push)', () => { take.setOutput(output); // Initialize Take - consume(take.fetch({})); + drainPull(take.fetch({})); // Edit '2' to '0' (move into bounds). const result = collectPush(source, makeEdit('0', '2')); @@ -645,7 +688,7 @@ describe('Yield Propagation (Push)', () => { output.yields = true; take.setOutput(output); - consume(take.fetch({})); + drainPull(take.fetch({})); // Output yields 2. expect(collectPush(source, makeAdd('1'))).toEqual(['yield', 'yield']); diff --git a/packages/zql/src/query/measure-push-operator.test.ts b/packages/zql/src/query/measure-push-operator.test.ts index dbfe02c542..22d030c5ae 100644 --- a/packages/zql/src/query/measure-push-operator.test.ts +++ b/packages/zql/src/query/measure-push-operator.test.ts @@ -5,6 +5,7 @@ import {makeAddChange} from '../ivm/change.ts'; import type {Node} from '../ivm/data.ts'; import type {FetchRequest, Input, Output} from '../ivm/operator.ts'; import type {SourceSchema} from '../ivm/schema.ts'; +import {emptyPullStream} from '../ivm/stream.ts'; import {MeasurePushOperator} from './measure-push-operator.ts'; import type {MetricsDelegate} from './metrics-delegate.ts'; @@ -12,7 +13,7 @@ describe('MeasurePushOperator', () => { test('should pass through fetch calls', () => { const mockInput: Input = { setOutput: vi.fn(), - fetch: vi.fn(() => []), + fetch: vi.fn(() => emptyPullStream()), getSchema: vi.fn(() => ({}) as SourceSchema), destroy: vi.fn(), }; @@ -38,7 +39,7 @@ describe('MeasurePushOperator', () => { const schema = {} as SourceSchema; const mockInput: Input = { setOutput: vi.fn(), - fetch: vi.fn(() => []), + fetch: vi.fn(() => emptyPullStream()), getSchema: vi.fn(() => schema), destroy: vi.fn(), }; @@ -63,7 +64,7 @@ describe('MeasurePushOperator', () => { test('should pass through destroy calls', () => { const mockInput: Input = { setOutput: vi.fn(), - fetch: vi.fn(() => []), + fetch: vi.fn(() => emptyPullStream()), getSchema: vi.fn(() => ({}) as SourceSchema), destroy: vi.fn(), }; @@ -87,7 +88,7 @@ describe('MeasurePushOperator', () => { test('should measure push timing and record metric', () => { const mockInput: Input = { setOutput: vi.fn(), - fetch: vi.fn(() => []), + fetch: vi.fn(() => emptyPullStream()), getSchema: vi.fn(() => ({}) as SourceSchema), destroy: vi.fn(), }; @@ -123,7 +124,7 @@ describe('MeasurePushOperator', () => { test('should not record metric when output.push throws', () => { const mockInput: Input = { setOutput: vi.fn(), - fetch: vi.fn(() => []), + fetch: vi.fn(() => emptyPullStream()), getSchema: vi.fn(() => ({}) as SourceSchema), destroy: vi.fn(), }; diff --git a/packages/zql/src/query/measure-push-operator.ts b/packages/zql/src/query/measure-push-operator.ts index c52018a6c4..52592c90b5 100644 --- a/packages/zql/src/query/measure-push-operator.ts +++ b/packages/zql/src/query/measure-push-operator.ts @@ -8,7 +8,7 @@ import { type Output, } from '../ivm/operator.ts'; import type {SourceSchema} from '../ivm/schema.ts'; -import type {Stream} from '../ivm/stream.ts'; +import {type Stream, type PullStream} from '../ivm/stream.ts'; import type {MetricsDelegate} from './metrics-delegate.ts'; type MetricName = 'query-update-client' | 'query-update-server'; @@ -38,7 +38,7 @@ export class MeasurePushOperator implements Operator { this.#output = output; } - fetch(req: FetchRequest): Stream { + fetch(req: FetchRequest): PullStream { return this.#input.fetch(req); } diff --git a/packages/zqlite/src/table-source.test.ts b/packages/zqlite/src/table-source.test.ts index 19702b37fc..83a1d48845 100644 --- a/packages/zqlite/src/table-source.test.ts +++ b/packages/zqlite/src/table-source.test.ts @@ -22,7 +22,7 @@ import { makeSourceChangeEdit, makeSourceChangeRemove, } from '../../zql/src/ivm/source.ts'; -import {consume} from '../../zql/src/ivm/stream.ts'; +import {consume, drainPull, drainPullMap} from '../../zql/src/ivm/stream.ts'; import {Database, Statement} from './db.ts'; import {explainQueries} from './explain-queries.ts'; import {format} from './internal/sql.ts'; @@ -494,12 +494,12 @@ describe('fetched value types', () => { if (c.output) { expect( - Array.from(input.fetch({}), node => + drainPullMap(input.fetch({}), node => node === 'yield' ? node : node.row, ), ).toEqual([c.output]); } else { - expect(() => [...input.fetch({})]).toThrow(UnsupportedValueError); + expect(() => drainPull(input.fetch({}))).toThrow(UnsupportedValueError); } }); } @@ -1087,7 +1087,7 @@ describe('fromSQLiteTypes error messages', () => { ); const input = source.connect([['id', 'asc']]); - expect(() => [...input.fetch({})]).toThrow( + expect(() => drainPull(input.fetch({}))).toThrow( /value .* \(in test_table\.big_value\) is outside of supported bounds/, ); }); @@ -1114,7 +1114,7 @@ describe('fromSQLiteTypes error messages', () => { ); const input = source.connect([['id', 'asc']]); - expect(() => [...input.fetch({})]).toThrow( + expect(() => drainPull(input.fetch({}))).toThrow( /Failed to parse JSON for test_table\.json_data/, ); }); @@ -1143,7 +1143,8 @@ describe('fromSQLiteTypes error messages', () => { let caughtError: unknown; try { - for (const _ of input.fetch({})) { + const stream = input.fetch({}); + for (let _ = stream.next(); _ !== undefined; _ = stream.next()) { // Consume the iterator to trigger the error } } catch (error) { @@ -1182,7 +1183,7 @@ test('debug.recordExplain captures the plan SQLite picked for the real bindings' const input = source.connect([['id', 'asc']], undefined, undefined, debug); // Drain the iterator with a constraint that uses the email index. - [...input.fetch({constraint: {email: 'a@b'}})]; + drainPull(input.fetch({constraint: {email: 'a@b'}})); const plans = debug.getSQLitePlans(); const entries = Object.entries(plans); @@ -1237,7 +1238,7 @@ test('captured plan diverges from substituted-literal plan when bindings affect const debug = new Debug(); const input = source.connect([['id', 'asc']], likeFilter, undefined, debug); - [...input.fetch({})]; + drainPull(input.fetch({})); const plans = debug.getSQLitePlans(); const entries = Object.entries(plans); @@ -1319,7 +1320,7 @@ test('SQLite iterator is closed when an error occurs before #mapFromSQLiteTypes throwingDebug, ); - expect(() => [...input.fetch({})]).toThrow('initQuery error'); + expect(() => drainPull(input.fetch({}))).toThrow('initQuery error'); expect(iteratorReturnCalled).toBe(true); } finally { Statement.prototype.iterate = origIterate; diff --git a/packages/zqlite/src/table-source.ts b/packages/zqlite/src/table-source.ts index 96d328fd04..22ecdbbc91 100644 --- a/packages/zqlite/src/table-source.ts +++ b/packages/zqlite/src/table-source.ts @@ -37,6 +37,7 @@ import { type SourceChange, type SourceInput, } from '../../zql/src/ivm/source.ts'; +import {LazyPullStream, type PullStream} from '../../zql/src/ivm/stream.ts'; import type {Stream} from '../../zql/src/ivm/stream.ts'; import {assertOrderingIncludesPK} from '../../zql/src/query/complete-ordering.ts'; import type {Database, Statement} from './db.ts'; @@ -281,30 +282,108 @@ export class TableSource implements Source { ) as Row; } - *#fetch(req: FetchRequest, connection: Connection): Stream { - const {sort, debug} = connection; + #fetch( + req: FetchRequest, + connection: Connection, + ): PullStream { + // Lazy, as the generator body was: nothing runs until the first `next()`. + return new LazyPullStream(() => { + const {sort, debug} = connection; + + const query = this.#requestToSQL( + req, + connection.filters?.condition, + sort, + ); + const sqlAndBindings = format(query); - const query = this.#requestToSQL(req, connection.filters?.condition, sort); - const sqlAndBindings = format(query); + const cachedStatement = this.#stmts.cache.get(sqlAndBindings.text); + cachedStatement.statement.safeIntegers(true); + const rowIterator = cachedStatement.statement.iterate( + ...sqlAndBindings.values, + ); + const overlayPredicate = mergeOverlayPredicate( + connection.filters?.predicate, + req.filter, + ); - const cachedStatement = this.#stmts.cache.get(sqlAndBindings.text); - cachedStatement.statement.safeIntegers(true); - const rowIterator = cachedStatement.statement.iterate( - ...sqlAndBindings.values, - ); - const overlayPredicate = mergeOverlayPredicate( - connection.filters?.predicate, - req.filter, - ); - try { - debug?.initQuery(this.#table, sqlAndBindings.text); + // The generator's `finally` ran on exhaustion, early return, or throw. + // `onDone` is hoisted above the try so a throw from `initQuery` -- or + // from building the chain -- still closes the SQLite cursor. Leaking one + // leaves later writes failing with "database connection is busy". + const onDone = () => { + // Ensure the SQLite iterate() is closed. + rowIterator.return?.(); + if (debug) { + let totalNvisit = 0; + const planLines: string[] = []; + for (let i = 0; ; i++) { + const nvisit = cachedStatement.statement.scanStatus( + i, + SQLite3Database.SQLITE_SCANSTAT_NVISIT, + 1, + ); + if (nvisit === undefined) { + break; + } + totalNvisit += Number(nvisit); + const explain = cachedStatement.statement.scanStatus( + i, + SQLite3Database.SQLITE_SCANSTAT_EXPLAIN, + 1, + ); + if (typeof explain === 'string' && explain.length > 0) { + planLines.push(explain); + } + } + if (totalNvisit !== 0) { + debug.recordNVisit(this.#table, sqlAndBindings.text, totalNvisit); + } + if (planLines.length > 0) { + debug.recordExplain(this.#table, sqlAndBindings.text, planLines); + } + cachedStatement.statement.scanStatusReset(); + } + this.#stmts.cache.return(cachedStatement); + }; - if (sort) { - const comparator = makeComparator(sort, req.reverse); - yield* generateWithStart( + try { + debug?.initQuery(this.#table, sqlAndBindings.text); + if (sort) { + const comparator = makeComparator(sort, req.reverse); + return new FinallyPull( + generateWithStart( + generateWithYields( + generateWithOverlay( + req.start?.row, + this.#mapFromSQLiteTypes( + this.#columns, + rowIterator, + sqlAndBindings.text, + debug, + ), + req.constraint, + this.#overlay, + connection.lastPushedEpoch, + comparator, + // SQL does the ordering and constraining, so the row stream + // is already in the connection's sort order: the splice + // comparator and the `startAt` comparator coincide here. + comparator, + overlayPredicate, + req.multiConstraints, + ), + this.#shouldYield, + ), + req.start, + comparator, + ), + onDone, + ); + } + return new FinallyPull( generateWithYields( - generateWithOverlay( - req.start?.row, + generateWithOverlayUnordered( this.#mapFromSQLiteTypes( this.#columns, rowIterator, @@ -314,99 +393,60 @@ export class TableSource implements Source { req.constraint, this.#overlay, connection.lastPushedEpoch, - comparator, - // SQL does the ordering and constraining, so the row stream is - // already in the connection's sort order: the splice comparator - // and the `startAt` comparator coincide here. - comparator, + this.#primaryKey, overlayPredicate, req.multiConstraints, ), this.#shouldYield, ), - req.start, - comparator, + onDone, ); - } else { - yield* generateWithYields( - generateWithOverlayUnordered( - this.#mapFromSQLiteTypes( - this.#columns, - rowIterator, - sqlAndBindings.text, - debug, - ), - req.constraint, - this.#overlay, - connection.lastPushedEpoch, - this.#primaryKey, - overlayPredicate, - req.multiConstraints, - ), - this.#shouldYield, - ); - } - } finally { - // Ensure the SQLite iterate() is closed. - rowIterator.return?.(); - if (debug) { - let totalNvisit = 0; - const planLines: string[] = []; - for (let i = 0; ; i++) { - const nvisit = cachedStatement.statement.scanStatus( - i, - SQLite3Database.SQLITE_SCANSTAT_NVISIT, - 1, - ); - if (nvisit === undefined) { - break; - } - totalNvisit += Number(nvisit); - const explain = cachedStatement.statement.scanStatus( - i, - SQLite3Database.SQLITE_SCANSTAT_EXPLAIN, - 1, - ); - if (typeof explain === 'string' && explain.length > 0) { - planLines.push(explain); - } - } - if (totalNvisit !== 0) { - debug.recordNVisit(this.#table, sqlAndBindings.text, totalNvisit); - } - if (planLines.length > 0) { - debug.recordExplain(this.#table, sqlAndBindings.text, planLines); - } - cachedStatement.statement.scanStatusReset(); + } catch (e) { + onDone(); + throw e; } - this.#stmts.cache.return(cachedStatement); - } + }); } - *#mapFromSQLiteTypes( + #mapFromSQLiteTypes( valueTypes: Record, rowIterator: IterableIterator, query: string, debug: DebugDelegate | undefined, - ): IterableIterator { - let result; - do { - result = timeSampled( - this.#lc, - ++eventCount, - this.#logConfig.ivmSampling, - () => rowIterator.next(), - this.#logConfig.slowRowThreshold, - () => - `table-source.next took too long for ${query}. Are you missing an index?`, - ); - if (result.done) { - break; - } - const row = fromSQLiteTypes(valueTypes, result.value, this.#table); - debug?.rowVended(this.#table, query, row); - yield row; - } while (!result.done); + ): PullStream { + const lc = this.#lc; + const logConfig = this.#logConfig; + const table = this.#table; + let done = false; + return { + next(): Row | undefined { + if (done) { + return undefined; + } + const result = timeSampled( + lc, + ++eventCount, + logConfig.ivmSampling, + () => rowIterator.next(), + logConfig.slowRowThreshold, + () => + `table-source.next took too long for ${query}. Are you missing an index?`, + ); + if (result.done) { + done = true; + return undefined; + } + const row = fromSQLiteTypes(valueTypes, result.value, table); + debug?.rowVended(table, query, row); + return row; + }, + close(): void { + if (!done) { + done = true; + rowIterator.return?.(); + } + }, + }; } *push(change: SourceChange): Stream<'yield'> { @@ -714,11 +754,74 @@ function nonPrimaryKeys( return Object.keys(columns).filter(c => !primaryKey.includes(c)); } -function* generateWithYields(stream: Stream, shouldYield: () => boolean) { - for (const n of stream) { - if (shouldYield()) { - yield 'yield'; +/** Runs `onDone` once, when the stream ends, is closed, or throws. */ +class FinallyPull implements PullStream { + readonly #inner: PullStream; + readonly #onDone: () => void; + #done = false; + + constructor(inner: PullStream, onDone: () => void) { + this.#inner = inner; + this.#onDone = onDone; + } + + next(): Node | 'yield' | undefined { + if (this.#done) { + return undefined; + } + let v: Node | 'yield' | undefined; + try { + v = this.#inner.next(); + } catch (e) { + this.close(); + throw e; + } + if (v === undefined) { + this.close(); } - yield n; + return v; } + + close(): void { + if (!this.#done) { + this.#done = true; + this.#inner.close(); + this.#onDone(); + } + } +} + +/** + * Injects a 'yield' marker before a row whenever `shouldYield()` says so. + * + * The row is held until the marker has been handed back, which is the two + * `yield`s the generator emitted in one loop iteration. + */ +function generateWithYields( + stream: PullStream, + shouldYield: () => boolean, +): PullStream { + let pending: Node | undefined; + return { + next(): Node | 'yield' | undefined { + if (pending !== undefined) { + const held = pending; + pending = undefined; + return held; + } + const n = stream.next(); + if (n === undefined) { + return undefined; + } + if (shouldYield()) { + pending = n; + return 'yield'; + } + return n; + }, + close() { + pending = undefined; + stream.close(); + }, + }; }