[PECO-238] Initial CloudFetch implementation (#158)

* [PECO-238] Initial CloudFetch implementation

Signed-off-by: Levko Kravets <levko.ne@gmail.com>

* Fix tests

Signed-off-by: Levko Kravets <levko.ne@gmail.com>

* Handle expired links

Signed-off-by: Levko Kravets <levko.ne@gmail.com>

* Improve cloud fetch result handler: better performance and fixed memory consumption issues

Signed-off-by: Levko Kravets <levko.ne@gmail.com>

* Extract some cloudfetch settings to global config; add e2e tests

Signed-off-by: Levko Kravets <levko.ne@gmail.com>

* Add unit tests

Signed-off-by: Levko Kravets <levko.ne@gmail.com>

* Make e2e tests more stable

Signed-off-by: Levko Kravets <levko.ne@gmail.com>

* Increase Nodejs heap size for E2E tests

Signed-off-by: Levko Kravets <levko.ne@gmail.com>

---------

Signed-off-by: Levko Kravets <levko.ne@gmail.com>

Author: kravets-levko

.github/workflows/main.yml

@@ -59,7 +59,7 @@ jobs:
           E2E_TABLE_SUFFIX: ${{github.sha}}
         run: |
           npm ci
-          npm run e2e
+          NODE_OPTIONS="--max-old-space-size=4096" npm run e2e
       - name: Coverage
         uses: codecov/codecov-action@v3
         with:

lib/DBSQLOperation/SchemaHelper.ts

@@ -4,6 +4,7 @@ import Status from '../dto/Status';
 import IOperationResult from '../result/IOperationResult';
 import JsonResult from '../result/JsonResult';
 import ArrowResult from '../result/ArrowResult';
+import CloudFetchResult from '../result/CloudFetchResult';
 import HiveDriverError from '../errors/HiveDriverError';
 import { definedOrError } from '../utils';
 
@@ -14,6 +15,8 @@ export default class SchemaHelper {
 
   private metadata?: TGetResultSetMetadataResp;
 
+  private resultHandler?: IOperationResult;
+
   constructor(driver: HiveDriver, operationHandle: TOperationHandle, metadata?: TGetResultSetMetadataResp) {
     this.driver = driver;
     this.operationHandle = operationHandle;
@@ -41,13 +44,27 @@ export default class SchemaHelper {
     const metadata = await this.fetchMetadata();
     const resultFormat = definedOrError(metadata.resultFormat);
 
-    switch (resultFormat) {
-      case TSparkRowSetType.COLUMN_BASED_SET:
-        return new JsonResult(metadata.schema);
-      case TSparkRowSetType.ARROW_BASED_SET:
-        return new ArrowResult(metadata.schema, metadata.arrowSchema);
-      default:
-        throw new HiveDriverError(`Unsupported result format: ${TSparkRowSetType[resultFormat]}`);
+    if (!this.resultHandler) {
+      switch (resultFormat) {
+        case TSparkRowSetType.COLUMN_BASED_SET:
+          this.resultHandler = new JsonResult(metadata.schema);
+          break;
+        case TSparkRowSetType.ARROW_BASED_SET:
+          this.resultHandler = new ArrowResult(metadata.schema, metadata.arrowSchema);
+          break;
+        case TSparkRowSetType.URL_BASED_SET:
+          this.resultHandler = new CloudFetchResult(metadata.schema);
+          break;
+        default:
+          this.resultHandler = undefined;
+          break;
+      }
+    }
+
+    if (!this.resultHandler) {
+      throw new HiveDriverError(`Unsupported result format: ${TSparkRowSetType[resultFormat]}`);
     }
+
+    return this.resultHandler;
   }
 }

lib/DBSQLOperation/index.ts

@@ -106,7 +106,7 @@ export default class DBSQLOperation implements IOperation {
       this._data.fetch(options?.maxRows || defaultMaxRows),
     ]);
 
-    const result = resultHandler.getValue(data ? [data] : []);
+    const result = await resultHandler.getValue(data ? [data] : []);
     this.logger?.log(
       LogLevel.debug,
       `Fetched chunk of size: ${options?.maxRows || defaultMaxRows} from operation with id: ${this.getId()}`,
@@ -147,10 +147,19 @@ export default class DBSQLOperation implements IOperation {
   }
 
   public async hasMoreRows(): Promise<boolean> {
+    // If operation is closed or cancelled - we should not try to get data from it
     if (this._completeOperation.closed || this._completeOperation.cancelled) {
       return false;
     }
-    return this._data.hasMoreRows;
+
+    // Return early if there are still data available for fetching
+    if (this._data.hasMoreRows) {
+      return true;
+    }
+
+    // If we fetched all the data from server - check if there's anything buffered in result handler
+    const resultHandler = await this._schema.getResultHandler();
+    return resultHandler.hasPendingData();
   }
 
   public async getSchema(options?: GetSchemaOptions): Promise<TTableSchema | null> {

lib/DBSQLSession.ts

@@ -125,6 +125,7 @@ export default class DBSQLSession implements IDBSQLSession {
       runAsync: options.runAsync || false,
       ...getDirectResultsOptions(options.maxRows),
       ...getArrowOptions(),
+      canDownloadResult: options.useCloudFetch ?? globalConfig.useCloudFetch,
     });
 
     return this.createOperation(response);

lib/contracts/IDBSQLSession.ts

@@ -7,6 +7,7 @@ export type ExecuteStatementOptions = {
   queryTimeout?: Int64;
   runAsync?: boolean;
   maxRows?: number | null;
+  useCloudFetch?: boolean;
 };
 
 export type TypeInfoRequest = {

lib/globalConfig.ts

@@ -7,6 +7,9 @@ interface GlobalConfig {
   retriesTimeout: number; // in milliseconds
   retryDelayMin: number; // in milliseconds
   retryDelayMax: number; // in milliseconds
+
+  useCloudFetch: boolean;
+  cloudFetchConcurrentDownloads: number;
 }
 
 export default {
@@ -18,4 +21,7 @@ export default {
   retriesTimeout: 900 * 1000,
   retryDelayMin: 1 * 1000,
   retryDelayMax: 60 * 1000,
+
+  useCloudFetch: false,
+  cloudFetchConcurrentDownloads: 10,
 } satisfies GlobalConfig;

lib/result/ArrowResult.ts

@@ -30,12 +30,16 @@ export default class ArrowResult implements IOperationResult {
     this.arrowSchema = arrowSchema;
   }
 
-  getValue(data?: Array<TRowSet>) {
+  async hasPendingData() {
+    return false;
+  }
+
+  async getValue(data?: Array<TRowSet>) {
     if (this.schema.length === 0 || !this.arrowSchema || !data) {
       return [];
     }
 
-    const batches = this.getBatches(data);
+    const batches = await this.getBatches(data);
     if (batches.length === 0) {
       return [];
     }
@@ -44,7 +48,7 @@ export default class ArrowResult implements IOperationResult {
     return this.getRows(table.schema, table.toArray());
   }
 
-  private getBatches(data: Array<TRowSet>): Array<Buffer> {
+  protected async getBatches(data: Array<TRowSet>): Promise<Array<Buffer>> {
     const result: Array<Buffer> = [];
 
     data.forEach((rowSet) => {

lib/result/CloudFetchResult.ts

@@ -0,0 +1,56 @@
+import { Buffer } from 'buffer';
+import fetch, { RequestInfo, RequestInit } from 'node-fetch';
+import { TRowSet, TSparkArrowResultLink, TTableSchema } from '../../thrift/TCLIService_types';
+import ArrowResult from './ArrowResult';
+import globalConfig from '../globalConfig';
+
+export default class CloudFetchResult extends ArrowResult {
+  private pendingLinks: Array<TSparkArrowResultLink> = [];
+
+  private downloadedBatches: Array<Buffer> = [];
+
+  constructor(schema?: TTableSchema) {
+    // Arrow schema returned in metadata is not needed for CloudFetch results:
+    // each batch already contains schema and could be decoded as is
+    super(schema, Buffer.alloc(0));
+  }
+
+  async hasPendingData() {
+    return this.pendingLinks.length > 0 || this.downloadedBatches.length > 0;
+  }
+
+  protected async getBatches(data: Array<TRowSet>): Promise<Array<Buffer>> {
+    data.forEach((item) => {
+      item.resultLinks?.forEach((link) => {
+        this.pendingLinks.push(link);
+      });
+    });
+
+    if (this.downloadedBatches.length === 0) {
+      const links = this.pendingLinks.splice(0, globalConfig.cloudFetchConcurrentDownloads);
+      const tasks = links.map((link) => this.downloadLink(link));
+      const batches = await Promise.all(tasks);
+      this.downloadedBatches.push(...batches);
+    }
+
+    return this.downloadedBatches.splice(0, 1);
+  }
+
+  private async downloadLink(link: TSparkArrowResultLink): Promise<Buffer> {
+    if (Date.now() >= link.expiryTime.toNumber()) {
+      throw new Error('CloudFetch link has expired');
+    }
+
+    const response = await this.fetch(link.fileLink);
+    if (!response.ok) {
+      throw new Error(`CloudFetch HTTP error ${response.status} ${response.statusText}`);
+    }
+
+    const result = await response.arrayBuffer();
+    return Buffer.from(result);
+  }
+
+  private async fetch(url: RequestInfo, init?: RequestInit) {
+    return fetch(url, init);
+  }
+}

lib/result/IOperationResult.ts

@@ -1,5 +1,7 @@
 import { TRowSet } from '../../thrift/TCLIService_types';
 
 export default interface IOperationResult {
-  getValue(data?: Array<TRowSet>): any;
+  getValue(data?: Array<TRowSet>): Promise<any>;
+
+  hasPendingData(): Promise<boolean>;
 }

lib/result/JsonResult.ts

@@ -10,7 +10,11 @@ export default class JsonResult implements IOperationResult {
     this.schema = getSchemaColumns(schema);
   }
 
-  getValue(data?: Array<TRowSet>): Array<object> {
+  async hasPendingData() {
+    return false;
+  }
+
+  async getValue(data?: Array<TRowSet>): Promise<Array<object>> {
     if (this.schema.length === 0 || !data) {
       return [];
     }