diff --git a/packages/cubejs-base-driver/src/queue-driver.interface.ts b/packages/cubejs-base-driver/src/queue-driver.interface.ts index bbded5bb199fb..5980031cd59b2 100644 --- a/packages/cubejs-base-driver/src/queue-driver.interface.ts +++ b/packages/cubejs-base-driver/src/queue-driver.interface.ts @@ -1,8 +1,6 @@ export type QueryDef = any; // Primary key of Queue item export type QueueId = string | number | bigint; -// The lock token of a retrieval, always the item's queueId. Only the memory driver compares it. -export type ProcessingId = string | number | bigint; export type QueryKey = (string | [string, any[]]) & { persistent?: true, }; @@ -13,21 +11,21 @@ export type GetActiveAndToProcessResponse = [active: QueryKeysTuple[], toProcess export type QueryStageStateResponse = [active: string[], toProcess: string[]] | [active: string[], toProcess: string[], defs: Record]; export type RetrieveForProcessingSuccess = [ added: unknown, - // QueueId is required for Cube Store, other providers don't support it + // Identifies the retrieved generation of the queue item. queueId: QueueId | null, active: QueryKeyHash[], pending: number, def: QueryDef, - lockAquired: true + retrieved: true ]; export type RetrieveForProcessingFail = [ added: unknown, - // QueueId is required for Cube Store, other providers don't support it + // Null when no queue item was retrieved. queueId: QueueId | null, active: QueryKeyHash[], pending: number, def: null, - lockAquired: false + retrieved: false ]; export type RetrieveForProcessingResponse = RetrieveForProcessingSuccess | RetrieveForProcessingFail | null; export type AddToQueueResponse = [ @@ -106,14 +104,13 @@ export interface QueueDriverConnectionInterface { getStalledQueries(): Promise; getQueryStageState(onlyKeys: boolean): Promise; updateHeartBeat(hash: QueryKeyHash, queueId: QueueId | null): Promise; - // Trying to acquire a lock for processing a queue item, this method can return null when - // multiple nodes tries to process the same query - retrieveForProcessing(hash: QueryKeyHash, processingId: ProcessingId): Promise; - freeProcessingLock(hash: QueryKeyHash, processingId: ProcessingId, activated: unknown): Promise; - optimisticQueryUpdate(hash: QueryKeyHash, toUpdate: unknown, processingId: ProcessingId, queueId: QueueId | null): Promise; + // Atomically moves a queue item to active. Returns null when another node is already + // processing the query or the concurrency budget is full. + retrieveForProcessing(hash: QueryKeyHash, queueId: QueueId): Promise; + optimisticQueryUpdate(hash: QueryKeyHash, toUpdate: unknown, queueId: QueueId): Promise; cancelQuery(queryKey: QueryKey, queueId: QueueId | null): Promise; getQueryAndRemove(hash: QueryKeyHash, queueId: QueueId | null): Promise<[QueryDef]>; - setResultAndRemoveQuery(hash: QueryKeyHash, executionResult: any, processingId: ProcessingId, queueId: QueueId | null): Promise; + setResultAndRemoveQuery(hash: QueryKeyHash, executionResult: any, queueId: QueueId): Promise; release(): void; // getQueriesToCancel(): Promise diff --git a/packages/cubejs-cubestore-driver/src/CubeStoreQueueDriver.ts b/packages/cubejs-cubestore-driver/src/CubeStoreQueueDriver.ts index 609fe1f3fef93..95563eb392041 100644 --- a/packages/cubejs-cubestore-driver/src/CubeStoreQueueDriver.ts +++ b/packages/cubejs-cubestore-driver/src/CubeStoreQueueDriver.ts @@ -12,7 +12,6 @@ import { AddToQueueResponse, QueryKey, QueryKeyHash, - ProcessingId, QueueId, GetActiveAndToProcessResponse, QueryKeysTuple, @@ -182,10 +181,6 @@ export class CubestoreQueueDriverConnection implements QueueDriverConnectionInte return null; } - public async freeProcessingLock(_hash: QueryKeyHash, _processingId: string, _activated: unknown): Promise { - // nothing to do - } - public async getActiveQueries(): Promise { const rows = await this.driver.query('QUEUE ACTIVE ?', [ this.options.redisQueuePrefix @@ -336,7 +331,7 @@ export class CubestoreQueueDriverConnection implements QueueDriverConnectionInte return null; } - public async optimisticQueryUpdate(hash: QueryKeyHash, toUpdate: unknown, _processingId: ProcessingId, queueId: QueueId): Promise { + public async optimisticQueryUpdate(hash: QueryKeyHash, toUpdate: unknown, queueId: QueueId): Promise { await this.driver.query('QUEUE MERGE_EXTRA ? ?', [ // queryKeyHash as compatibility fallback queueId || this.prefixKey(hash), @@ -372,7 +367,7 @@ export class CubestoreQueueDriverConnection implements QueueDriverConnectionInte ]; } - public async retrieveForProcessing(hash: QueryKeyHash, _processingId: string): Promise { + public async retrieveForProcessing(hash: QueryKeyHash, _queueId: QueueId): Promise { const rows = await this.driver.query('QUEUE RETRIEVE EXTENDED CONCURRENCY ? ?', [ this.options.concurrency, this.prefixKey(hash), @@ -404,7 +399,7 @@ export class CubestoreQueueDriverConnection implements QueueDriverConnectionInte return null; } - public async setResultAndRemoveQuery(hash: QueryKeyHash, executionResult: unknown, _processingId: ProcessingId, queueId: QueueId): Promise { + public async setResultAndRemoveQuery(hash: QueryKeyHash, executionResult: unknown, queueId: QueueId): Promise { const rows = await this.driver.query('QUEUE ACK ? ?', [ // queryKeyHash as compatibility fallback queueId || this.prefixKey(hash), diff --git a/packages/cubejs-query-orchestrator/DEVELOPMENT.md b/packages/cubejs-query-orchestrator/DEVELOPMENT.md index e8141fe59ebc8..f51d405843f47 100644 --- a/packages/cubejs-query-orchestrator/DEVELOPMENT.md +++ b/packages/cubejs-query-orchestrator/DEVELOPMENT.md @@ -182,12 +182,13 @@ sequenceDiagram ## Background execution: `processQuery` → `executeQuery` `processQuery` retrieves the item and nothing else. What it hands to `sendProcessMessageFn` is a -`RetrievedQuery` — `{ queryKeyHash, queueId, processingId, queueSize, query }`, plain data on +`RetrievedQuery` — `{ queryKeyHash, queueId, queueSize, query }`, plain data on purpose, so a custom implementation can serialize it and let another process run `executeQuery`. The default implementation calls `executeQuery` in-process. -`processingId` is the lock token of the retrieval and always carries the `queueId`. Only the -memory driver compares it, Cube Store ignores it and keys every command off the `queueId`. +`queueId` identifies the specific generation of a queue item. The memory driver compares it +with the active entry before updating or acknowledging a query, while Cube Store keys those +commands directly off the `queueId`. `sendProcessMessageFn` must resolve once the hand-off is done, **not** once the query is executed: reconcile awaits it, and `executeQuery` ends with `reconcileQueue`, which is @@ -200,8 +201,8 @@ Two consequences of retrieving before the hand-off: `@` suffix matches, so `sendProcessMessageFn` is always called on the owning process for them — it just must not route them elsewhere. - A retrieved item is already active. If a custom hand-off loses the message, the item is only - recovered by the stalled-heartbeat / `TO_CANCEL` path; `freeProcessingLock` is a no-op on - Cube Store, so the retrieval cannot be cheaply undone. + recovered by the stalled-heartbeat / `TO_CANCEL` path; a successful retrieval is not undone + when the hand-off fails. A stream query is dispatched while `executeInQueue` is still running, so `waitForQueryStream` subscribes to `streamStarted` *before* the dispatch — a handler which starts fast would @@ -222,10 +223,10 @@ sequenceDiagram QueryQueue->>QueueDriver: retrieveForProcessing QueueDriver->>CubeStore: QUEUE RETRIEVE EXTENDED CONCURRENCY ?n ?path CubeStore-->>QueueDriver: RetrieveResponse - QueueDriver-->>QueryQueue: [added, queueId, activeKeys, queueSize, def, lockAcquired] + QueueDriver-->>QueryQueue: [added, queueId, activeKeys, queueSize, def, retrieved] Note over QueueDriver,CubeStore: The retrieval is atomic in Cube Store:
only one node moves the item to active - alt def && added && activeKeys includes our key && lockAcquired + alt def && added && activeKeys includes our key && retrieved QueryQueue-)Background: sendProcessMessageFn(RetrievedQuery) Note over QueryQueue,Background: Detached from here on: the hand-off returns,
the execution keeps running @@ -254,8 +255,7 @@ sequenceDiagram Background->>Background: reconcileQueue Note over Background: The freed concurrency slot is
immediately given to the next query else the retrieval did not succeed - QueryQueue->>QueueDriver: freeProcessingLock - Note over QueryQueue,QueueDriver: Another node is running it, or the
concurrency budget is full. No-op for Cube Store + Note over QueryQueue,QueueDriver: Another node is running it, or the
concurrency budget is full. Queue state is unchanged end ``` diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/LocalQueueDriverConnection.ts b/packages/cubejs-query-orchestrator/src/orchestrator/LocalQueueDriverConnection.ts index f2c0a14a542c3..55d4211440364 100644 --- a/packages/cubejs-query-orchestrator/src/orchestrator/LocalQueueDriverConnection.ts +++ b/packages/cubejs-query-orchestrator/src/orchestrator/LocalQueueDriverConnection.ts @@ -4,7 +4,6 @@ import { QueryKey, QueryKeyHash, QueueId, - ProcessingId, QueryDef, AddToQueueQuery, AddToQueueOptions, @@ -54,8 +53,6 @@ export class LocalQueueDriverConnectionState { public active: Record = {}; public heartBeat: Record = {}; - - public processingLocks: Record = {}; } export class LocalQueueDriverConnection implements QueueDriverConnectionInterface { @@ -223,7 +220,6 @@ export class LocalQueueDriverConnection implements QueueDriverConnectionInterfac delete this.state.toProcess[queryKeyHash]; delete this.state.recent[queryKeyHash]; delete this.state.queryDef[queryKeyHash]; - delete this.state.processingLocks[queryKeyHash]; return [query]; } @@ -233,8 +229,8 @@ export class LocalQueueDriverConnection implements QueueDriverConnectionInterfac return query; } - public async setResultAndRemoveQuery(queryKeyHash: QueryKeyHash, executionResult: any, processingId: ProcessingId, _queueId?: QueueId | null): Promise { - if (this.state.processingLocks[queryKeyHash] !== processingId) { + public async setResultAndRemoveQuery(queryKeyHash: QueryKeyHash, executionResult: any, queueId: QueueId): Promise { + if (this.state.active[queryKeyHash]?.queueId !== queueId) { return false; } @@ -245,7 +241,6 @@ export class LocalQueueDriverConnection implements QueueDriverConnectionInterfac delete this.state.toProcess[queryKeyHash]; delete this.state.recent[queryKeyHash]; delete this.state.queryDef[queryKeyHash]; - delete this.state.processingLocks[queryKeyHash]; promise.resolved = true; if (promise.resolve) { @@ -277,48 +272,44 @@ export class LocalQueueDriverConnection implements QueueDriverConnectionInterfac } } - public async retrieveForProcessing(queryKeyHash: QueryKeyHash, processingId: ProcessingId): Promise { - let lockAcquired = false; - - if (!this.state.processingLocks[queryKeyHash]) { - this.state.processingLocks[queryKeyHash] = processingId; - lockAcquired = true; - } else { - return null; + public async retrieveForProcessing(queryKeyHash: QueryKeyHash, queueId: QueueId): Promise { + const query = this.state.queryDef[queryKeyHash]; + const activeKeys = this.queueArray(this.state.active) as QueryKeyHash[]; + + if ( + !query || + query.queueId !== queueId || + this.state.toProcess[queryKeyHash]?.queueId !== queueId || + this.state.active[queryKeyHash] || + activeKeys.length >= this.concurrency + ) { + return [ + 0, + null, + activeKeys, + Object.keys(this.state.toProcess).length, + null, + false + ]; } - let added = 0; - - if (Object.keys(this.state.active).length < this.concurrency && !this.state.active[queryKeyHash]) { - this.state.active[queryKeyHash] = { key: queryKeyHash, order: Number(processingId), queueId: processingId }; - delete this.state.toProcess[queryKeyHash]; - - added = 1; - } + this.state.active[queryKeyHash] = { key: queryKeyHash, order: Number(queueId), queueId }; + delete this.state.toProcess[queryKeyHash]; - this.state.heartBeat[queryKeyHash] = { key: queryKeyHash, order: new Date().getTime(), queueId: processingId }; + this.state.heartBeat[queryKeyHash] = { key: queryKeyHash, order: new Date().getTime(), queueId }; return [ - added, - this.state.queryDef[queryKeyHash]?.queueId ?? null, + 1, + query.queueId, this.queueArray(this.state.active) as QueryKeyHash[], Object.keys(this.state.toProcess).length, - this.state.queryDef[queryKeyHash], - lockAcquired + query, + true ]; } - public async freeProcessingLock(queryKeyHash: QueryKeyHash, processingId: ProcessingId, activated: any): Promise { - if (this.state.processingLocks[queryKeyHash] === processingId) { - delete this.state.processingLocks[queryKeyHash]; - if (activated) { - delete this.state.active[queryKeyHash]; - } - } - } - - public async optimisticQueryUpdate(queryKeyHash: QueryKeyHash, toUpdate: any, processingId: ProcessingId, _queueId?: QueueId | null): Promise { - if (this.state.processingLocks[queryKeyHash] !== processingId) { + public async optimisticQueryUpdate(queryKeyHash: QueryKeyHash, toUpdate: any, queueId: QueueId): Promise { + if (this.state.active[queryKeyHash]?.queueId !== queueId || !this.state.queryDef[queryKeyHash]) { return false; } diff --git a/packages/cubejs-query-orchestrator/src/orchestrator/QueryQueue.ts b/packages/cubejs-query-orchestrator/src/orchestrator/QueryQueue.ts index 13d39e518e0f3..d36558735f70c 100644 --- a/packages/cubejs-query-orchestrator/src/orchestrator/QueryQueue.ts +++ b/packages/cubejs-query-orchestrator/src/orchestrator/QueryQueue.ts @@ -9,7 +9,6 @@ import { QueryStageStateResponse, AddToQueueOptions, QueuePriority, - ProcessingId, RetrieveForProcessingSuccess } from '@cubejs-backend/base-driver'; import { CubeStoreQueueDriver } from '@cubejs-backend/cubestore-driver'; @@ -29,7 +28,6 @@ export type QueryHandlersMap = Record; export type RetrievedQuery = { queryKeyHash: QueryKeyHash; queueId: QueueId; - processingId: ProcessingId; queueSize: number; query: QueryDef; }; @@ -373,7 +371,6 @@ export class QueryQueue { return { queryKeyHash, queueId, - processingId: queueId, queueSize, query, }; @@ -817,9 +814,8 @@ export class QueryQueue { } /** - * Acquires the processing lock for the query specified by the `queryKeyHashed` and moves it to - * the active set. Returns `null` when the retrieval didn't succeed, which means another node is - * already running the query or the concurrency budget is full. + * Atomically moves the query specified by `queryKeyHashed` to the active set. Returns `null` + * when another node is already running the query or the concurrency budget is full. */ protected async retrieveQueryForProcessing(queryKeyHashed: QueryKeyHash, queueId: QueueId): Promise { const queueConnection = await this.queueDriver.createConnection(); @@ -828,16 +824,12 @@ export class QueryQueue { let activeKeys; let queueSize; let query; - let processingLockAcquired; + let retrievalSucceeded; try { - // The lock token is the queueId, every call which releases the lock has to be handed - // the same value retrieveForProcessing got - const processingId = queueId; - - const retrieveResult = await queueConnection.retrieveForProcessing(queryKeyHashed, processingId); + const retrieveResult = await queueConnection.retrieveForProcessing(queryKeyHashed, queueId); if (retrieveResult) { - [insertedCount, , activeKeys, queueSize, query, processingLockAcquired] = retrieveResult; + [insertedCount, , activeKeys, queueSize, query, retrievalSucceeded] = retrieveResult; } const activated = activeKeys && activeKeys.indexOf(queryKeyHashed) !== -1; @@ -845,7 +837,7 @@ export class QueryQueue { query = await queueConnection.getQueryDef(queryKeyHashed, null); } - if (!query || !insertedCount || !activated || !processingLockAcquired) { + if (!query || !insertedCount || !activated || !retrievalSucceeded) { // TODO Ideally streaming queries should reconcile queue here after waiting on open slot however in practice continue wait timeout reconciles faster CPU-wise // if (query?.queryHandler === 'stream') { // const [active] = await queueConnection.getQueryStageState(true); @@ -857,26 +849,22 @@ export class QueryQueue { this.logger('Skip processing', { queueId, - processingId, queryKey: query && query.queryKey || queryKeyHashed, requestId: query && query.requestId, queuePrefix: this.redisQueuePrefix, - processingLockAcquired, + retrievalSucceeded, query, insertedCount, activeKeys, activated, queryExists: !!query }); - await queueConnection.freeProcessingLock(queryKeyHashed, processingId, activated); - return null; } return { queryKeyHash: queryKeyHashed, queueId, - processingId, queueSize, query, }; @@ -901,7 +889,7 @@ export class QueryQueue { * implementation which hands the query over to another process. */ public async executeQuery(retrieved: RetrievedQuery): Promise { - const { queryKeyHash: queryKeyHashed, queueId, processingId, queueSize, query } = retrieved; + const { queryKeyHash: queryKeyHashed, queueId, queueSize, query } = retrieved; const queueConnection = await this.queueDriver.createConnection(); @@ -915,7 +903,6 @@ export class QueryQueue { const timeInQueue = (new Date()).getTime() - query.addedToQueueTime; this.logger('Performing query', { queueId, - processingId, queueSize, queryKey: query.queryKey, queuePrefix: this.redisQueuePrefix, @@ -927,7 +914,7 @@ export class QueryQueue { preAggregation: query.query?.preAggregation, addedToQueueTime: query.addedToQueueTime, }); - await queueConnection.optimisticQueryUpdate(queryKeyHashed, { startQueryTime }, processingId, queueId); + await queueConnection.optimisticQueryUpdate(queryKeyHashed, { startQueryTime }, queueId); let queryProcessHeartbeat = Date.now(); const heartBeatTimer = setInterval( @@ -1011,7 +998,7 @@ export class QueryQueue { async (cancelHandler) => { localCancelHandler = cancelHandler; try { - await queueConnection.optimisticQueryUpdate(queryKeyHashed, { cancelHandler }, processingId, queueId); + await queueConnection.optimisticQueryUpdate(queryKeyHashed, { cancelHandler }, queueId); } catch (e: any) { this.logger('Error while query update', { queueId, @@ -1035,7 +1022,6 @@ export class QueryQueue { this.logger('Performing query completed', { queueId, - processingId, queueSize, duration: ((new Date()).getTime() - startQueryTime), queryKey: query.queryKey, @@ -1054,7 +1040,6 @@ export class QueryQueue { }; this.logger('Error while querying', { queueId, - processingId, queueSize, duration: ((new Date()).getTime() - startQueryTime), queryKey: query.queryKey, @@ -1073,7 +1058,6 @@ export class QueryQueue { if (queryWithCancelHandle) { this.logger('Cancelling query due to timeout', { queueId, - processingId, queryKey: queryWithCancelHandle.queryKey, queuePrefix: this.redisQueuePrefix, requestId: queryWithCancelHandle.requestId, @@ -1092,11 +1076,10 @@ export class QueryQueue { clearInterval(heartBeatTimer); } - if (!(await queueConnection.setResultAndRemoveQuery(queryKeyHashed, executionResult, processingId, queueId))) { + if (!(await queueConnection.setResultAndRemoveQuery(queryKeyHashed, executionResult, queueId))) { this.logger('Orphaned execution result', { queueId, - processingId, - warn: 'Result for query was not set due to processing lock wasn\'t acquired', + warn: 'Result for query was not set because the queue item is no longer active', queryKey: query.queryKey, queuePrefix: this.redisQueuePrefix, requestId: query.requestId, diff --git a/packages/cubejs-query-orchestrator/test/benchmarks/QueueBench.abstract.ts b/packages/cubejs-query-orchestrator/test/benchmarks/QueueBench.abstract.ts index 2a3b0e63adecb..f558043c5e217 100644 --- a/packages/cubejs-query-orchestrator/test/benchmarks/QueueBench.abstract.ts +++ b/packages/cubejs-query-orchestrator/test/benchmarks/QueueBench.abstract.ts @@ -42,7 +42,6 @@ function patchQueueDriverConnectionForTrack(connection: QueueDriverConnectionInt setResultAndRemoveQuery: wrapAsyncMethod('setResultAndRemoveQuery'), getQueryStageState: wrapAsyncMethod('getQueryStageState'), getResultBlocking: wrapAsyncMethod('getResultBlocking'), - freeProcessingLock: wrapAsyncMethod('freeProcessingLock'), optimisticQueryUpdate: wrapAsyncMethod('optimisticQueryUpdate'), getQueryAndRemove: wrapAsyncMethod('getQueryAndRemove'), release: connection.release, diff --git a/packages/cubejs-query-orchestrator/test/unit/QueryQueue.abstract.ts b/packages/cubejs-query-orchestrator/test/unit/QueryQueue.abstract.ts index 92fdcee8f8bb5..2dc205f4c1be7 100644 --- a/packages/cubejs-query-orchestrator/test/unit/QueryQueue.abstract.ts +++ b/packages/cubejs-query-orchestrator/test/unit/QueryQueue.abstract.ts @@ -465,76 +465,111 @@ export const QueryQueueTest = (name: string, options: QueryQueueTestOptions) => } }); - onlyLocalTest('queue driver lock obtain race condition', async () => { - const connection: any = await queue.queueDriver.createConnection(); - const connection2: any = await queue.queueDriver.createConnection(); + onlyLocalTest('an active query cannot be retrieved twice', async () => { + const connection = await queue.queueDriver.createConnection(); + const connection2 = await queue.queueDriver.createConnection(); const priority = 10; + const key = 'active-retrieval' as any; - await queue.reconcileQueue(); - - const [, raceQueueId] = await connection.addToQueue( - 'race', 'handler', ['select'], priority, { queueId: queue.generateQueueId(), stageQueryKey: 'race' } - ); - - const [, race2QueueId] = await connection.addToQueue( - 'race2', 'handler2', ['select2'], priority, { queueId: queue.generateQueueId(), stageQueryKey: 'race2' } - ); - - // Neither is locked yet, so both releases are no-ops - await connection.freeProcessingLock('race', raceQueueId, true); - await connection.freeProcessingLock('race2', race2QueueId, true); - - await connection2.retrieveForProcessing('race2', race2QueueId); - - const retrieve6 = await connection.retrieveForProcessing('race', raceQueueId); - console.log(retrieve6); - expect(!!retrieve6[5]).toBe(true); + try { + const [, queueId] = await connection.addToQueue( + key, 'handler', ['select'], priority, { + queueId: queue.generateQueueId(), stageQueryKey: key, requestId: '1' + } + ); - console.log(await connection.getQueryAndRemove('race')); - console.log(await connection.getQueryAndRemove('race2')); + const firstRetrieval = await connection.retrieveForProcessing(key, queueId); + expect(firstRetrieval?.[5]).toBe(true); - await queue.queueDriver.release(connection); - await queue.queueDriver.release(connection2); + const secondRetrieval = await connection2.retrieveForProcessing(key, queueId); + expect(secondRetrieval).toStrictEqual([0, null, [key], 0, null, false]); + } finally { + await connection.getQueryAndRemove(key, null); + queue.queueDriver.release(connection); + queue.queueDriver.release(connection2); + } }); - onlyLocalTest('activated but lock is not acquired', async () => { + onlyLocalTest('a failed retrieval does not reserve a pending query', async () => { const connection = await queue.queueDriver.createConnection(); const connection2 = await queue.queueDriver.createConnection(); const priority = 10; + const firstKey = 'concurrency-first' as any; + const secondKey = 'concurrency-second' as any; - await queue.reconcileQueue(); + try { + const [, firstQueueId] = await connection.addToQueue( + firstKey, 'handler', ['select'], priority, { + queueId: queue.generateQueueId(), stageQueryKey: firstKey, requestId: '1' + } + ); + const [, secondQueueId] = await connection.addToQueue( + secondKey, 'handler2', ['select2'], priority, { + queueId: queue.generateQueueId(), stageQueryKey: secondKey, requestId: '1' + } + ); - const [, activated1QueueId] = await connection.addToQueue( - 'activated1', 'handler', ['select'], priority, { queueId: queue.generateQueueId(), stageQueryKey: 'race', requestId: '1' } - ); + expect((await connection.retrieveForProcessing(firstKey, firstQueueId))?.[5]).toBe(true); + expect(await connection2.retrieveForProcessing(secondKey, secondQueueId)).toStrictEqual([ + 0, null, [firstKey], 1, null, false + ]); + expect(await connection.getToProcessQueries()).toStrictEqual([[secondKey, secondQueueId]]); - const [, activated2QueueId] = await connection.addToQueue( - 'activated2', 'handler2', ['select2'], priority, { queueId: queue.generateQueueId(), stageQueryKey: 'race2', requestId: '1' } - ); + await connection.getQueryAndRemove(firstKey, firstQueueId); - const retrieve1 = await connection.retrieveForProcessing('activated1' as any, activated1QueueId); - console.log(retrieve1); - const retrieve2 = await connection2.retrieveForProcessing('activated2' as any, activated2QueueId); - console.log(retrieve2); - console.log(await connection.freeProcessingLock('activated1' as any, activated1QueueId, retrieve1 && retrieve1[2].indexOf('activated1' as any) !== -1)); + const secondRetrieval = await connection2.retrieveForProcessing(secondKey, secondQueueId); + expect(secondRetrieval?.[0]).toBe(1); + expect(secondRetrieval?.[5]).toBe(true); + } finally { + await connection.getQueryAndRemove(firstKey, null); + await connection.getQueryAndRemove(secondKey, null); + queue.queueDriver.release(connection); + queue.queueDriver.release(connection2); + } + }); - // Another node reaches the same item, so it comes with the same lock token and loses - const retrieve3 = await connection.retrieveForProcessing('activated2' as any, activated2QueueId); - expect(retrieve3).toBeNull(); + test('stale queueId cannot update or acknowledge a requeued query', async () => { + const connection = await queue.queueDriver.createConnection(); + const queryKey = 'requeued-query' as QueryKey; + const key = connection.redisHash(queryKey); + const priority = 10; - console.log(retrieve2[2].indexOf('activated2' as any) !== -1); - console.log(await connection2.freeProcessingLock('activated2' as any, activated2QueueId, retrieve2 && retrieve2[2].indexOf('activated2' as any) !== -1)); + try { + const [, staleQueueId] = await connection.addToQueue( + queryKey, 'handler', ['old'], priority, { + queueId: queue.generateQueueId(), stageQueryKey: key, requestId: '1' + } + ); + expect((await connection.retrieveForProcessing(key, staleQueueId))?.[5]).toBe(true); + await connection.getQueryAndRemove(key, staleQueueId); - const retrieve4 = await connection.retrieveForProcessing('activated2' as any, activated2QueueId); - console.log(retrieve4); - expect(retrieve4[0]).toBe(1); - expect(!!retrieve4[5]).toBe(true); + const [, currentQueueId] = await connection.addToQueue( + queryKey, 'handler', ['new'], priority, { + queueId: queue.generateQueueId(), stageQueryKey: key, requestId: '2' + } + ); - console.log(await connection.getQueryAndRemove('activated1' as any, null)); - console.log(await connection.getQueryAndRemove('activated2' as any, null)); + if (options.cacheAndQueueDriver !== 'cubestore') { + expect(await connection.retrieveForProcessing(key, staleQueueId)).toStrictEqual([ + 0, null, [], 1, null, false + ]); + expect(await connection.getToProcessQueries()).toStrictEqual([[key, currentQueueId]]); + } + expect((await connection.retrieveForProcessing(key, currentQueueId))?.[5]).toBe(true); - await queue.queueDriver.release(connection); - await queue.queueDriver.release(connection2); + const staleUpdateResult = await connection.optimisticQueryUpdate(key, { stale: true }, staleQueueId); + if (options.cacheAndQueueDriver !== 'cubestore') { + expect(staleUpdateResult).toBe(false); + } + expect(await connection.setResultAndRemoveQuery(key, { result: 'stale' }, staleQueueId)).toBe(false); + const currentQuery = await connection.getQueryDef(key, currentQueueId); + expect(currentQuery).toMatchObject({ query: ['new'] }); + expect(currentQuery).not.toHaveProperty('stale'); + expect(await connection.getActiveQueries()).toStrictEqual([[key, currentQueueId]]); + } finally { + await connection.getQueryAndRemove(key, null); + queue.queueDriver.release(connection); + } }); // eslint-disable-next-line no-unused-expressions diff --git a/rust/cubestore/cubestore/src/queryplanner/group_by_limit_aggregate/mod.rs b/rust/cubestore/cubestore/src/queryplanner/group_by_limit_aggregate/mod.rs index bd5a937239bc9..cc37a2b620ff2 100644 --- a/rust/cubestore/cubestore/src/queryplanner/group_by_limit_aggregate/mod.rs +++ b/rust/cubestore/cubestore/src/queryplanner/group_by_limit_aggregate/mod.rs @@ -4,7 +4,7 @@ use datafusion::arrow::compute::SortOptions; use datafusion::arrow::datatypes::SchemaRef; use datafusion::common::stats::Precision; use datafusion::common::Statistics; -use datafusion::error::Result as DFResult; +use datafusion::error::{DataFusionError, Result as DFResult}; use datafusion::execution::TaskContext; use datafusion::physical_expr::aggregate::AggregateFunctionExpr; use datafusion::physical_expr::{Distribution, LexRequirement}; @@ -16,7 +16,7 @@ use datafusion::physical_plan::{ PlanProperties, SendableRecordBatchStream, }; use std::any::Any; -use std::fmt::Debug; +use std::fmt::{Debug, Formatter}; use std::sync::Arc; /// Worker-side partial hash aggregate that trims its output to the top-k groups by a total order, @@ -32,7 +32,7 @@ use std::sync::Arc; /// columns), expressed as `(partial-output column index, sort options)`. A total order is required /// for correctness: the same group key can live on multiple workers, and a consistent cut across /// workers guarantees every partial state the router selects reaches it. -#[derive(Debug, Clone)] +#[derive(Clone)] pub struct GroupByLimitAggregateExec { group_by: PhysicalGroupBy, aggr_expr: Vec>, @@ -48,6 +48,28 @@ pub struct GroupByLimitAggregateExec { factor: usize, /// Total order over the partial output columns. order: Vec<(usize, SortOptions)>, + /// The aggregate this node replaced. Kept so that swapping the input can be delegated to it: + /// DataFusion recomputes the plan properties there and preserves the output schema. + source: Arc, +} + +/// Skips [GroupByLimitAggregateExec::source]: it carries a copy of the same input subtree the node +/// already prints, so deriving would double the output of every nested aggregate. +impl Debug for GroupByLimitAggregateExec { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GroupByLimitAggregateExec") + .field("group_by", &self.group_by) + .field("aggr_expr", &self.aggr_expr) + .field("filter_expr", &self.filter_expr) + .field("input", &self.input) + .field("schema", &self.schema) + .field("input_schema", &self.input_schema) + .field("cache", &self.cache) + .field("k", &self.k) + .field("factor", &self.factor) + .field("order", &self.order) + .finish_non_exhaustive() + } } impl GroupByLimitAggregateExec { @@ -77,10 +99,8 @@ impl GroupByLimitAggregateExec { return None; } let input = aggregate.input().clone(); - // A partial aggregate preserves its input's partitioning (it runs once per input partition). - // Derive the output partitioning from the input rather than copying the wrapped aggregate's - // cached value, which can be stale: a later pass may swap our input for one with a different - // partition count via `with_new_children` without the cache following, and a too-low count + // A partial aggregate runs once per input partition and so preserves its partitioning. + // `aggregate.cache()` already says that; spell it out anyway, since a wrong count here // makes the parent coalesce read only some partitions and silently drop the rest. let cache = aggregate .cache() @@ -97,6 +117,7 @@ impl GroupByLimitAggregateExec { k, factor, order, + source: Arc::new(aggregate.clone()), }) } @@ -177,24 +198,39 @@ impl ExecutionPlan for GroupByLimitAggregateExec { self: Arc, children: Vec>, ) -> DFResult> { - let input = children[0].clone(); - // Track the (possibly changed) input's partitioning; a partial aggregate preserves it. - let cache = self - .cache - .clone() - .with_partitioning(input.output_partitioning().clone()); - Ok(Arc::new(Self { - group_by: self.group_by.clone(), - aggr_expr: self.aggr_expr.clone(), - filter_expr: self.filter_expr.clone(), - input, - schema: self.schema.clone(), - input_schema: self.input_schema.clone(), - cache, - k: self.k, - factor: self.factor, - order: self.order.clone(), - })) + // Swapping the input is delegated to the aggregate this node was built from, so that + // DataFusion recomputes the whole of the plan properties -- output partitioning above all, + // a stale count there makes the parent read only some of the input's partitions and + // silently drop the rest -- and preserves the output schema. + let rebuilt = Arc::clone(&self.source).with_new_children(children)?; + let Some(aggregate) = rebuilt.as_any().downcast_ref::() else { + return Err(DataFusionError::Internal(format!( + "AggregateExec::with_new_children returned {}", + rebuilt.name() + ))); + }; + Ok( + match Self::try_new_from_partial(aggregate, self.k, self.factor, self.order.clone()) { + Some(trimmed) => Arc::new(trimmed), + None => { + // Correct but a perf cliff: the worker goes from `factor * k` rows per + // partition to the full group cardinality. Say it out loud so the loss of the + // trim is diagnosable instead of showing up only as a slow query. + log::warn!( + "Rebuilt aggregate no longer fits the group-by-limit trim (input order \ + mode {:?}), dropping it for group by [{}]", + aggregate.input_order_mode(), + self.group_by + .expr() + .iter() + .map(|(_, name)| name.as_str()) + .collect::>() + .join(", ") + ); + rebuilt + } + }, + ) } fn execute( @@ -453,6 +489,10 @@ mod tests { let three = MemorySourceConfig::try_new(&three_parts, input_schema(), None).unwrap(); let three_input: Arc = Arc::new(DataSourceExec::new(Arc::new(three))); let exec3 = exec.with_new_children(vec![three_input]).unwrap(); + assert!( + exec3.as_any().is::(), + "re-childing must keep the trimming exec, otherwise this test no longer pins it" + ); assert_eq!( exec3.output_partitioning().partition_count(), 3, diff --git a/rust/cubestore/cubestore/src/queryplanner/inline_aggregate/mod.rs b/rust/cubestore/cubestore/src/queryplanner/inline_aggregate/mod.rs index 8a58d1a8c0dba..e0f031c99435d 100644 --- a/rust/cubestore/cubestore/src/queryplanner/inline_aggregate/mod.rs +++ b/rust/cubestore/cubestore/src/queryplanner/inline_aggregate/mod.rs @@ -9,7 +9,7 @@ pub use sorted_group_values_rows::SortedGroupValuesRows; use datafusion::arrow::datatypes::{DataType, SchemaRef}; use datafusion::common::stats::Precision; use datafusion::common::Statistics; -use datafusion::error::Result as DFResult; +use datafusion::error::{DataFusionError, Result as DFResult}; use datafusion::execution::TaskContext; use datafusion::physical_expr::aggregate::AggregateFunctionExpr; use datafusion::physical_expr::{Distribution, LexRequirement}; @@ -22,7 +22,7 @@ use datafusion::physical_plan::{ SendableRecordBatchStream, }; use std::any::Any; -use std::fmt::Debug; +use std::fmt::{Debug, Formatter}; use std::sync::Arc; #[derive(Debug, Copy, Clone, PartialEq, Eq)] @@ -31,7 +31,7 @@ pub enum InlineAggregateMode { Final, } -#[derive(Debug, Clone)] +#[derive(Clone)] pub struct InlineAggregateExec { mode: InlineAggregateMode, /// Group by expressions @@ -55,6 +55,29 @@ pub struct InlineAggregateExec { pub input_schema: SchemaRef, cache: PlanProperties, required_input_ordering: Vec>, + /// The aggregate this node replaced. Kept so that swapping the input can be delegated to it: + /// DataFusion recomputes the plan properties there and preserves the output schema, which it + /// derives from aggregate expression names that its own rules are free to rewrite. + source: Arc, +} + +/// Skips [InlineAggregateExec::source]: it carries a copy of the same input subtree the node +/// already prints, so deriving would double the output of every nested aggregate. +impl Debug for InlineAggregateExec { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("InlineAggregateExec") + .field("mode", &self.mode) + .field("group_by", &self.group_by) + .field("aggr_expr", &self.aggr_expr) + .field("filter_expr", &self.filter_expr) + .field("limit", &self.limit) + .field("input", &self.input) + .field("schema", &self.schema) + .field("input_schema", &self.input_schema) + .field("cache", &self.cache) + .field("required_input_ordering", &self.required_input_ordering) + .finish_non_exhaustive() + } } impl InlineAggregateExec { @@ -101,6 +124,7 @@ impl InlineAggregateExec { input_schema, cache, required_input_ordering, + source: Arc::new(aggregate.clone()), }) } @@ -192,19 +216,46 @@ impl ExecutionPlan for InlineAggregateExec { self: Arc, children: Vec>, ) -> DFResult> { - let result = Self { - mode: self.mode, - group_by: self.group_by.clone(), - aggr_expr: self.aggr_expr.clone(), - filter_expr: self.filter_expr.clone(), - limit: self.limit.clone(), - input: children[0].clone(), - schema: self.schema.clone(), - input_schema: self.input_schema.clone(), - cache: self.cache.clone(), - required_input_ordering: self.required_input_ordering.clone(), + // Swapping the input is delegated to the aggregate this node was built from, so that + // DataFusion recomputes the plan properties -- output partitioning above all. Carrying + // them over leaves a stale partition count, and the parent then executes only that many + // of the input's partitions and silently drops the rows of the rest. + let rebuilt = Arc::clone(&self.source).with_new_children(children)?; + let Some(aggregate) = rebuilt.as_any().downcast_ref::() else { + return Err(DataFusionError::Internal(format!( + "AggregateExec::with_new_children returned {}", + rebuilt.name() + ))); }; - Ok(Arc::new(result)) + + // The streaming implementation is only valid while the new input stays sorted on the + // group keys; if it no longer is, keep the plain hash aggregate. The limit stays behind + // with it: here it counts complete groups off a sorted input, while a hash aggregate + // reads it as a cap on distinct groups and would stop mid-way through arbitrary ones. + // Dropping it only costs the early exit. + Ok(match Self::try_new_from_aggregate(aggregate) { + Some(inline) => Arc::new(inline.with_limit(self.limit)), + None => { + // No pass in the pipeline is known to de-sort an inline aggregate's input, and a + // parent merge built around the old ordering would not notice the switch, so say + // it out loud rather than degrade quietly. + log::warn!( + "Input of an inline {:?} aggregate is no longer sorted on the group keys, \ + falling back to a hash aggregate. Group by [{}], new input order mode {:?}, \ + new input ordering {:?}", + self.mode, + self.group_by + .expr() + .iter() + .map(|(_, name)| name.as_str()) + .collect::>() + .join(", "), + aggregate.input_order_mode(), + aggregate.input().properties().output_ordering() + ); + Arc::new(aggregate.clone().with_limit(None)) + } + }) } fn execute( @@ -408,6 +459,41 @@ mod tests { } /// A group continuing in the next input batch must not be emitted early with a partial sum. + /// A partial aggregate runs once per input partition, so its reported partitioning must + /// follow a re-childed input. A stale single-partition count makes the parent coalesce + /// execute only partition 0 and silently drop the rows of the rest -- the row loss this node + /// was fixed for. Built over a 1-partition input so its cache says 1, then re-childed onto 3. + #[test] + fn output_partitioning_follows_rechilded_input() { + let schema = test_schema(); + let one = sorted_source( + &schema, + vec![vec![make_batch(&schema, &[(1, 10), (2, 20)])]], + ); + let exec = partial_sum_inline_aggregate(one, None); + assert_eq!(exec.properties().output_partitioning().partition_count(), 1); + + let three: Vec> = (0..3) + .map(|i| vec![make_batch(&schema, &[(i, 10), (i + 10, 20)])]) + .collect(); + let rechilded = exec + .with_new_children(vec![sorted_source(&schema, three)]) + .unwrap(); + + assert!( + rechilded.as_any().is::(), + "re-childing a still-sorted input must keep the streaming exec" + ); + assert_eq!( + rechilded + .properties() + .output_partitioning() + .partition_count(), + 3, + "exec must report its re-childed input's partition count, not a stale 1" + ); + } + #[test] fn limit_emits_only_closed_groups() { let schema = test_schema(); diff --git a/rust/cubestore/cubestore/src/queryplanner/partition_filter.rs b/rust/cubestore/cubestore/src/queryplanner/partition_filter.rs index d59b2f396fcf8..a619a4fe2ad6d 100644 --- a/rust/cubestore/cubestore/src/queryplanner/partition_filter.rs +++ b/rust/cubestore/cubestore/src/queryplanner/partition_filter.rs @@ -1531,6 +1531,131 @@ mod tests { Vec::new() } } + + /// Exhaustive check that [MinMaxCondition] never prunes a range that holds a matching row. + /// + /// Enumerates every per-column bound combination over a small integer domain against every + /// ordered pair of min/max rows, and compares the verdict with a direct search for a tuple + /// that both satisfies the bounds and lies inside the range. A false negative here means a + /// query silently returning fewer rows than it should. + fn assert_no_false_negatives(n: usize, d: i64) { + // Rows and bounds use 0..d; candidate tuples reach one step further on both sides to + // stand in for the unbounded parts of the domain. + let rows = cross(n, &(0..d).collect::>()); + let candidates = cross(n, &(-1..=d).collect::>()); + + for cond in column_bounds(n, d) { + let min_max = MinMaxCondition { + min: cond.iter().map(|(mn, _)| mn.map(TableValue::Int)).collect(), + max: cond.iter().map(|(_, mx)| mx.map(TableValue::Int)).collect(), + }; + let matching = candidates + .iter() + .filter(|t| satisfies(&cond, t)) + .collect::>(); + + for mn in &rows { + for mx in &rows { + if mn > mx { + continue; + } + let min_row = mn.iter().cloned().map(TableValue::Int).collect::>(); + let max_row = mx.iter().cloned().map(TableValue::Int).collect::>(); + + if matching.iter().any(|t| *t >= mn && *t <= mx) { + assert!( + min_max.can_match(&min_row, &max_row), + "pruned a matching range: cond {:?}, min_row {:?}, max_row {:?}", + cond, + mn, + mx + ); + } + if matching.iter().any(|t| *t >= mn) { + assert!( + min_max.can_match_min(&min_row), + "pruned a matching range with unbounded max: cond {:?}, min_row {:?}", + cond, + mn + ); + } + if matching.iter().any(|t| *t <= mx) { + assert!( + min_max.can_match_max(&max_row), + "pruned a matching range with unbounded min: cond {:?}, max_row {:?}", + cond, + mx + ); + } + } + } + } + } + + /// All tuples of length `n` over `vals`. + fn cross(n: usize, vals: &[i64]) -> Vec> { + let mut r = vec![vec![]]; + for _ in 0..n { + r = r + .iter() + .flat_map(|t| { + vals.iter().map(move |v| { + let mut t = t.clone(); + t.push(*v); + t + }) + }) + .collect(); + } + r + } + + /// All per-column `(min, max)` bound combinations over `0..d`, `None` meaning unbounded. + fn column_bounds(n: usize, d: i64) -> Vec, Option)>> { + let mut per_col = Vec::new(); + for mn in 0..=d { + for mx in 0..=d { + let mn = (mn != d).then_some(mn); + let mx = (mx != d).then_some(mx); + if let (Some(mn), Some(mx)) = (mn, mx) { + if mn > mx { + continue; + } + } + per_col.push((mn, mx)); + } + } + let mut r = vec![vec![]]; + for _ in 0..n { + r = r + .iter() + .flat_map(|t| { + per_col.iter().map(move |c| { + let mut t = t.clone(); + t.push(*c); + t + }) + }) + .collect(); + } + r + } + + fn satisfies(cond: &[(Option, Option)], t: &[i64]) -> bool { + cond.iter() + .zip(t) + .all(|((mn, mx), v)| mn.map_or(true, |mn| mn <= *v) && mx.map_or(true, |mx| *v <= mx)) + } + + #[test] + fn no_false_negatives_two_columns() { + assert_no_false_negatives(2, 3); + } + + #[test] + fn no_false_negatives_three_columns() { + assert_no_false_negatives(3, 3); + } } struct ColumnStat { diff --git a/rust/cubestore/cubestore/src/queryplanner/rolling.rs b/rust/cubestore/cubestore/src/queryplanner/rolling.rs index 7eac8f08dc4aa..61c2223819bb9 100644 --- a/rust/cubestore/cubestore/src/queryplanner/rolling.rs +++ b/rust/cubestore/cubestore/src/queryplanner/rolling.rs @@ -709,6 +709,9 @@ impl ExecutionPlan for RollingWindowAggExec { ) -> Result, DataFusionError> { assert_eq!(children.len(), 1); Ok(Arc::new(RollingWindowAggExec { + // Safe to carry over: these are built from this node's own output schema plus + // constants, nothing in them is derived from the input. A node whose properties do + // follow the input -- output partitioning above all -- must recompute them here. properties: self.properties.clone(), sorted_input: children.remove(0), group_key: self.group_key.clone(), diff --git a/rust/cubestore/cubestore/src/queryplanner/topk/execute.rs b/rust/cubestore/cubestore/src/queryplanner/topk/execute.rs index eca963b705d41..67bfee7cfb1bc 100644 --- a/rust/cubestore/cubestore/src/queryplanner/topk/execute.rs +++ b/rust/cubestore/cubestore/src/queryplanner/topk/execute.rs @@ -190,6 +190,9 @@ impl ExecutionPlan for AggregateTopKExec { having: self.having.clone(), cluster, schema: self.schema.clone(), + // Safe to carry over: built from this node's own output schema plus constants, and + // the input's partition count is read live in `execute`. A node whose properties do + // follow the input -- output partitioning above all -- must recompute them here. cache: self.cache.clone(), sort_requirement: self.sort_requirement.clone(), merge_version: self.merge_version, diff --git a/rust/cubestore/cubestore/src/sql/mod.rs b/rust/cubestore/cubestore/src/sql/mod.rs index 0ce3abad3b156..3c777294e136c 100644 --- a/rust/cubestore/cubestore/src/sql/mod.rs +++ b/rust/cubestore/cubestore/src/sql/mod.rs @@ -6826,6 +6826,119 @@ mod tests { .await; Ok(()) } + /// A grouped aggregate over an index whose sort key does not start with the filtered + /// columns: every group column is pinned to a single value by the filter, which makes the + /// aggregate sorted and hands it to the streaming implementation. Its output partitioning + /// must follow the scan's, otherwise only the first partition is read and the rows living + /// in the others are silently dropped. + #[tokio::test] + async fn single_value_equals_scans_every_partition() -> Result<(), CubeError> { + Config::test("single_value_equals_scans_every_partition") + .update_config(|mut c| { + c.compaction_chunks_count_threshold = 0; + c.partition_split_threshold = 400; + c.max_partition_split_threshold = 400; + c + }) + .start_test(async move |services| { + let service = services.sql_service; + service.exec_query("CREATE SCHEMA pa").await?.collect().await?; + // The time dimension leads the sort key, so the equality-filtered columns are + // not an index prefix -- the shape of a partitioned rollup. + service + .exec_query( + "CREATE TABLE pa.rollup (date_code timestamp, equipment_type text, \ + month_code text, site_id text, hours int)", + ) + .await? + .collect() + .await?; + + let mut rows = Vec::new(); + for month in 1..=12 { + for day in 1..=8 { + for equipment_type in &["QC", "RTG", "STS"] { + for site_id in &["ABC", "PSE", "XYZ"] { + rows.push(format!( + "('2026-{:02}-{:02}T00:00:00.000', '{}', '2026{:02}', '{}', {})", + month, day, equipment_type, month, site_id, 10 + )); + } + } + } + } + for chunk in rows.chunks(300) { + service + .exec_query(&format!( + "INSERT INTO pa.rollup (date_code, equipment_type, month_code, \ + site_id, hours) VALUES {}", + chunk.join(", ") + )) + .await? + .collect() + .await?; + } + // Rows reachable only through the first of several scanned partitions is what + // the regression is about, so wait for the data to be split across partitions + // and fail loudly rather than silently testing a single-partition scan. + // Compaction settles in a few iterations; the bound is a ceiling, not a wait. + let mut scanned_partitions = TableValue::Int(0); + for _ in 0..1200 { + tokio::time::sleep(Duration::from_millis(50)).await; + let partitions = service + .exec_query( + "SELECT count(*) FROM system.partitions \ + WHERE active = true AND main_table_row_count > 0", + ) + .await? + .collect() + .await?; + scanned_partitions = partitions.get_rows()[0].values()[0].clone(); + if scanned_partitions > TableValue::Int(1) { + break; + } + } + assert!( + scanned_partitions > TableValue::Int(1), + "the scan must span several partitions, got {:?}", + scanned_partitions + ); + + // The last month of the range, so the matching rows sit in the last partition + // the scan visits and never in the first one -- reading only the first partition + // has to show up as a missing row rather than by chance returning the right one. + let query = "SELECT site_id, CAST(month_code AS INT) mc, sum(hours) v \ + FROM pa.rollup \ + WHERE (equipment_type = 'QC') \ + AND (CAST(month_code AS INT) = '202612') \ + AND (site_id = 'PSE') GROUP BY 1, 2"; + + // Pinning every group column to a single value is what hands the aggregate to + // the streaming implementation, and only that one carried a stale partition + // count. Without this the test would keep passing on a plan that never exercises + // the fix. + let worker_plan = pp_phys_plan(service.plan_query(query).await?.worker.as_ref()); + assert!( + worker_plan.contains("InlinePartialAggregate"), + "expected a streaming partial aggregate, got:\n{}", + worker_plan + ); + + let single_value = service.exec_query(query).await?.collect().await?; + assert_eq!( + single_value.get_rows(), + &vec![Row::new(vec![ + TableValue::String("PSE".to_string()), + TableValue::Int(202612), + TableValue::Int(80) + ])] + ); + + Ok::<(), CubeError>(()) + }) + .await; + Ok(()) + } } impl SqlServiceImpl {