perf: avoid flattening and shifting result batches

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Reduce result chunking and `fetchAll()` overhead for large query results.
 - Thrift CloudFetch downloads are now bounded end to end by `retriesTimeout` (default 15 minutes), including every attempt and retry backoff. Previously a stalled download, where the response headers or body never arrived, blocked the fetch forever. A download still running when the deadline expires is aborted and fails with `RetryError` (`TIMEOUT_EXCEEDED`, message `CloudFetch result download failed: Retry timeout exceeded`). (#520)
 - Thrift connections made through `proxy` now apply `checkServerCertificate`, `customCaCert` and `clientCert`/`clientKey`. Previously the proxied path ignored them and verified against Node's process-wide trust store, so a certificate outside the configured trust set could be accepted. **Compatibility:** proxied connections now verify unless `checkServerCertificate: false` is set; as on direct connections, `NODE_TLS_REJECT_UNAUTHORIZED=0` no longer turns verification off. For an `https` proxy these options also apply to the proxy's own certificate. (#485)
 - Honor the client `useCloudFetch: false` default on the kernel backend so complete results can be fetched without cloud-storage downloads. (#574)

lib/DBSQLOperation.ts

@@ -14,6 +14,7 @@ import Status from './dto/Status';
 import { LogLevel } from './contracts/IDBSQLLogger';
 import OperationStateError, { OperationStateErrorCode } from './errors/OperationStateError';
 import { OperationChunksIterator, OperationRowsIterator } from './utils/OperationIterator';
+import { concatChunks } from './result/ResultSlicer';
 import IClientContext from './contracts/IClientContext';
 import IOperationBackend, { IOperationBackendWaitOptions } from './contracts/IOperationBackend';
 import { ResultFormat, ResultMetadata } from './contracts/ResultMetadata';
@@ -131,7 +132,7 @@ export default class DBSQLOperation implements IOperation {
     } while (await this.hasMoreRows()); // eslint-disable-line no-await-in-loop
     this.context.getLogger().log(LogLevel.debug, `Fetched all data from operation with id: ${this.id}`);
 
-    return data.flat();
+    return concatChunks(data);
   }
 
   /**

lib/result/ArrowResultConverter.ts

@@ -279,8 +279,7 @@ export default class ArrowResultConverter implements IResultsProvider<Array<any>
       const batch = this.prefetchedRecordBatch;
       this.prefetchedRecordBatch = undefined;
       // Get table rows, but not more than remaining count. When materialization
-      // is disabled, skip `toArray()`/`getRows()` and return `null` placeholders
-      // (filled, not sparse, so `fetchAll`'s `.flat()` keeps the count).
+      // is disabled, skip `toArray()`/`getRows()` and return dense `null` placeholders.
       const batchRows = Math.min(batch.numRows, this.remainingRows);
       let result: Array<any>;
       if (this.disableRowMaterialization) {

lib/result/ResultSlicer.ts

@@ -7,20 +7,71 @@ export interface ResultSlicerFetchNextOptions extends ResultsProviderFetchNextOp
   disableBuffering?: boolean;
 }
 
+export function concatChunks<T>(chunks: Array<Array<T>>): Array<T> {
+  if (chunks.length === 1) {
+    return chunks[0];
+  }
+
+  // Avoid concat: apache-arrow's Symbol.isConcatSpreadable forces V8's slow path.
+  let total = 0;
+  for (const chunk of chunks) {
+    total += chunk.length;
+  }
+
+  // Larger initial allocations use V8's slow dictionary elements; grow sequentially instead.
+  const result = new Array<T>(Math.min(total, 2 ** 25));
+  let offset = 0;
+  for (const chunk of chunks) {
+    for (let i = 0; i < chunk.length; i += 1) {
+      result[offset + i] = chunk[i];
+    }
+    offset += chunk.length;
+  }
+  return result;
+}
+
 export default class ResultSlicer<T> implements IResultsProvider<Array<T>> {
   private readonly context: IClientContext;
 
   private readonly source: IResultsProvider<Array<T>>;
 
-  private remainingResults: Array<T> = [];
+  private remainingResults: Array<T | undefined> = [];
+
+  private remainingOffset: number = 0;
 
   constructor(context: IClientContext, source: IResultsProvider<Array<T>>) {
     this.context = context;
     this.source = source;
   }
 
+  private get remainingCount(): number {
+    return this.remainingResults.length - this.remainingOffset;
+  }
+
+  private compactRemaining(): void {
+    if (this.remainingCount === 0) {
+      this.remainingResults = [];
+      this.remainingOffset = 0;
+    } else if (this.remainingOffset > this.remainingResults.length / 2) {
+      // Shrink the backing array without shifting the tail on every fetch.
+      this.remainingResults = this.remainingResults.slice(this.remainingOffset);
+      this.remainingOffset = 0;
+    }
+  }
+
+  private takeRemaining(maxCount: number): Array<T> {
+    const count = Number.isNaN(maxCount) ? this.remainingCount : Math.max(0, Math.trunc(maxCount));
+    const end = Math.min(this.remainingResults.length, this.remainingOffset + count);
+    // Unread slots contain T; only consumed slots are cleared.
+    const result = this.remainingResults.slice(this.remainingOffset, end) as Array<T>;
+    this.remainingResults.fill(undefined, this.remainingOffset, end);
+    this.remainingOffset = end;
+    this.compactRemaining();
+    return result;
+  }
+
   public async hasMore(): Promise<boolean> {
-    if (this.remainingResults.length > 0) {
+    if (this.remainingCount > 0) {
       return true;
     }
     return this.source.hasMore();
@@ -30,47 +81,52 @@ export default class ResultSlicer<T> implements IResultsProvider<Array<T>> {
     // If we're asked to not use buffer - first try to return whatever we have in buffer.
     // If buffer is empty - just proxy the call to underlying results provider
     if (options.disableBuffering) {
-      if (this.remainingResults.length > 0) {
-        const result = this.remainingResults;
-        this.remainingResults = [];
-        return result;
+      if (this.remainingCount > 0) {
+        return this.takeRemaining(this.remainingCount);
       }
 
       return this.source.fetchNext(options);
     }
 
+    const limit = Number(options.limit);
+    const sourceOptions = { ...options, limit };
     const result: Array<Array<T>> = [];
     let resultsCount = 0;
 
     // First, use remaining items from the previous fetch
-    if (this.remainingResults.length > 0) {
-      result.push(this.remainingResults);
-      resultsCount += this.remainingResults.length;
-      this.remainingResults = [];
+    if (this.remainingCount > 0) {
+      const remaining = this.takeRemaining(limit);
+      if (this.remainingCount > 0) {
+        return remaining;
+      }
+      result.push(remaining);
+      resultsCount += remaining.length;
     }
 
     // Fetch items from source results provider until we reach a requested count
-    while (resultsCount < options.limit) {
+    while (resultsCount < limit) {
       // eslint-disable-next-line no-await-in-loop
       const hasMore = await this.source.hasMore();
       if (!hasMore) {
         break;
       }
 
       // eslint-disable-next-line no-await-in-loop
-      const chunk = await this.source.fetchNext(options);
-      result.push(chunk);
-      resultsCount += chunk.length;
-    }
-
-    // If we collected more results than requested, slice the excess items and store them for the next time
-    if (resultsCount > options.limit) {
-      const lastChunk = result.pop() ?? [];
-      const neededCount = options.limit - (resultsCount - lastChunk.length);
-      result.push(lastChunk.splice(0, neededCount));
-      this.remainingResults = lastChunk;
+      const chunk = await this.source.fetchNext(sourceOptions);
+      const neededCount = limit - resultsCount;
+      if (chunk.length > neededCount) {
+        const taken = chunk.slice(0, neededCount);
+        result.push(taken);
+        // Own the tail so consumed references can be cleared without mutating the source.
+        this.remainingResults = chunk.slice(taken.length);
+        this.remainingOffset = 0;
+        resultsCount += neededCount;
+      } else {
+        result.push(chunk);
+        resultsCount += chunk.length;
+      }
     }
 
-    return result.flat();
+    return concatChunks(result);
   }
 }