Line data    Source code

       1              : import * as arrow from 'apache-arrow';
       2              : 
       3              : import type { Logger } from '../../platform/logger/logger.js';
       4              : import { createQueryResponseStreamMetrics, QueryExecutionMetrics, QueryExecutionProgress, QueryExecutionResponseStream, QueryExecutionStatus } from '../../query/query_execution_state.js';
       5              : import type { AsyncConsumer } from '../../utils/async_consumer.js';
       6              : import { translateAnyRowsToArrowBatch } from '../trino/arrow_conversion.js';
       7              : import type { ClientFunctionResult } from './clfn_runtime.js';
       8              : 
       9              : /// The result of a client function call as a query result. A call has finished once its table exists, so this holds one batch.
      10            1 : export class ClientFunctionResultStream implements QueryExecutionResponseStream {
      11            1 :     private readonly metrics: QueryExecutionMetrics = createQueryResponseStreamMetrics();
      12            1 :     private status: QueryExecutionStatus = QueryExecutionStatus.RECEIVED_FIRST_BATCH;
      13              : 
      14            1 :     constructor(private readonly logger: Logger, private readonly result: ClientFunctionResult) { }
      15              : 
      16            1 :     getMetadata(): Map<string, string> {
      17            1 :         return new Map();
      18            1 :     }
      19            1 :     getMetrics(): QueryExecutionMetrics {
      20            1 :         return this.metrics;
      21            1 :     }
      22            0 :     getStatus(): QueryExecutionStatus {
      23            0 :         return this.status;
      24            0 :     }
      25            1 :     async getSchema(): Promise<arrow.Schema | null> {
      26            1 :         return this.result.schema;
      27            1 :     }
      28            1 :     async produce(batches: AsyncConsumer<QueryExecutionResponseStream, arrow.RecordBatch>, _progress: AsyncConsumer<QueryExecutionResponseStream, QueryExecutionProgress>, abort?: AbortSignal): Promise<void> {
      29            1 :         abort?.throwIfAborted();
      30            1 :         if (this.result.rows.length > 0) {
      31            1 :             const rows = this.result.rows as any[][];
      32            1 :             batches.resolve(this, translateAnyRowsToArrowBatch(this.result.schema, rows, this.logger));
      33            1 :             this.metrics.totalBatchesReceived += 1;
      34            1 :             this.metrics.totalRowsReceived += rows.length;
      35            1 :         }
      36            1 :         this.status = QueryExecutionStatus.RECEIVED_ALL_BATCHES;
      37            1 :     }
      38            1 : }