Line data    Source code

       1              : import * as arrow from 'apache-arrow';
       2              : import {
       3              :     QueryExecutionProgress,
       4              :     QueryExecutionResponseStream, QueryExecutionMetrics,
       5              :     QueryExecutionStatus,
       6              :     createQueryResponseStreamMetrics,
       7              : } from '../query/query_execution_state.js';
       8              : import { sleep } from '../utils/sleep.js';
       9              : import { AsyncConsumer } from '../utils/async_consumer.js';
      10              : 
      11            1 : export class QueryExecutionResponseStreamMock implements QueryExecutionResponseStream {
      12            1 :     schema: arrow.Schema;
      13            1 :     batchCount: number;
      14            1 :     batchesWritten: number;
      15            1 :     metrics: QueryExecutionMetrics;
      16              : 
      17            1 :     constructor() {
      18            1 :         this.schema = new arrow.Schema([
      19            1 :             new arrow.Field('A', new arrow.Int32()),
      20            1 :             new arrow.Field('B', new arrow.Int32()),
      21            1 :             new arrow.Field('C', new arrow.Int32()),
      22            1 :         ]);
      23            1 :         this.batchCount = 1;
      24            1 :         this.batchesWritten = 0;
      25            1 :         this.metrics = createQueryResponseStreamMetrics();
      26            1 :     }
      27              : 
      28              :     /// Get the metrics
      29            1 :     getMetadata(): Map<string, string> {
      30            1 :         return new Map();
      31            1 :     }
      32              :     /// Get the metrics
      33            1 :     getMetrics(): QueryExecutionMetrics {
      34            1 :         return this.metrics;
      35            1 :     }
      36              :     /// Get the status
      37            1 :     getStatus() {
      38            1 :         return QueryExecutionStatus.RUNNING;
      39            1 :     }
      40              :     /// Get the arrow schema
      41            1 :     async getSchema(): Promise<arrow.Schema | null> {
      42            1 :         await sleep(200);
      43            1 :         return this.schema;
      44            1 :     }
      45              :     /// Await the next record batch
      46              :     /// Produce the result batches
      47            1 :     async produce(batches: AsyncConsumer<QueryExecutionResponseStream, arrow.RecordBatch>, _progress: AsyncConsumer<QueryExecutionResponseStream, QueryExecutionProgress>, abort?: AbortSignal): Promise<void> {
      48            1 :         while (this.batchesWritten < this.batchCount) {
      49            1 :             await sleep(400);
      50            1 :             abort?.throwIfAborted();
      51              : 
      52            1 :             const length = 1000;
      53            1 :             const columnA = Int32Array.from({ length }, () => Number((Math.random() * 1000).toFixed(0)));
      54            1 :             const columnB = Int32Array.from({ length }, () => Number((Math.random() * 1000).toFixed(0)));
      55            1 :             const columnC = Int32Array.from({ length }, () => Number((Math.random() * 1000).toFixed(0)));
      56              : 
      57            1 :             const vecA = arrow.makeData({ type: new arrow.Int32(), data: columnA });
      58            1 :             const vecB = arrow.makeData({ type: new arrow.Int32(), data: columnB });
      59            1 :             const vecC = arrow.makeData({ type: new arrow.Int32(), data: columnC });
      60              : 
      61            1 :             const data = arrow.makeData({
      62            1 :                 type: new arrow.Struct(this.schema.fields),
      63            1 :                 length,
      64            1 :                 children: [vecA, vecB, vecC],
      65            1 :                 nullCount: 0,
      66            1 :             });
      67            1 :             batches.resolve(this, new arrow.RecordBatch(this.schema, data));
      68            1 :             this.batchesWritten += 1;
      69            1 :         }
      70            1 :     }
      71            1 : }