perf: prefetch kernel batches for async statement handles

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- Prefetch one kernel result batch for queries submitted asynchronously (`runAsync: true`, or default execution that remains pending after the initial response). Queries that finish in the initial response fetch on demand. Successful close and cancellation discard the prefetch without waiting for it; failed close requests preserve unread batches. An already-started native fetch may finish in the background. (#578)
 - CloudFetch retry-exhaustion errors now identify the download and the last failure: the `RetryError` message reads `CloudFetch result download failed: <original message> (last error: <cause>)`. `<cause>` is `HTTP <status>: <statusText>` (or `HTTP <status>` when the status text is empty) for HTTP failures, and `network error <code>`, `network error (<type>)` (e.g. `aborted`, when there is no `code`) or a bare `network error` for network failures (the raw network-error message, which can contain the presigned URL, is not copied). If no cause can be determined, the ` (last error: …)` suffix is omitted. The `RetryError` instance, `errorCode` and `payload` (including the node-fetch `Request` / `Response`) are unchanged. (#471)
 - Report untyped NULL columns as STRING in kernel result metadata, matching Thrift. (#470)
 - Reject non-scalar query parameter values at bind time with a `ParameterError` on both the Thrift and kernel backends, instead of stringifying them (which could leave a Thrift operation stuck before FINISHED). Newly rejected on the Thrift backend: arrays, plain objects, boxed primitives (e.g. `new String('x')`, `new Number(1)`), functions, symbols, and a `DBSQLParameter` that wraps another `DBSQLParameter`. Primitives, `Date`, `Int64`, and `DBSQLParameter` are still accepted.

lib/kernel/KernelOperationBackend.ts

@@ -206,6 +206,8 @@ export default class KernelOperationBackend implements IOperationBackend {
 
   private resultsProvider?: KernelResultsProvider;
 
+  private closePromise?: Promise<Status>;
+
   private metadata?: ResultMetadata;
 
   private metadataPromise?: Promise<ResultMetadata>;
@@ -309,7 +311,7 @@ export default class KernelOperationBackend implements IOperationBackend {
       // wedged, so nothing downstream forces another close). We still don't
       // mask the original fetch error, but log the close failure at warn so
       // the leak is diagnosable rather than completely invisible.
-      await kernelClose(this.lifecycle, this.lifecycleHandle, this.context, this._id).catch((closeErr) => {
+      await this.close().catch((closeErr) => {
         const cause = closeErr instanceof Error ? closeErr.message : String(closeErr);
         this.context
           .getLogger()
@@ -443,11 +445,26 @@ export default class KernelOperationBackend implements IOperationBackend {
   }
 
   public async cancel(): Promise<Status> {
-    return kernelCancel(this.lifecycle, this.lifecycleHandle, this.context, this._id);
+    const result = kernelCancel(this.lifecycle, this.lifecycleHandle, this.context, this._id);
+    // A pending close makes cancel a no-op without setting the cancelled flag.
+    if (this.lifecycle.isCancelled) {
+      this.resultsProvider?.dispose();
+    }
+    return result;
   }
 
-  public async close(): Promise<Status> {
-    return kernelClose(this.lifecycle, this.lifecycleHandle, this.context, this._id);
+  public close(): Promise<Status> {
+    if (!this.closePromise) {
+      this.closePromise = kernelClose(this.lifecycle, this.lifecycleHandle, this.context, this._id)
+        .then((status) => {
+          this.resultsProvider?.dispose();
+          return status;
+        })
+        .finally(() => {
+          this.closePromise = undefined;
+        });
+    }
+    return this.closePromise;
   }
 
   // ---------------------------------------------------------------------------
@@ -706,7 +723,7 @@ export default class KernelOperationBackend implements IOperationBackend {
    * original error; warn-logs a close failure so the leak is diagnosable.
    */
   private async bestEffortClose(): Promise<void> {
-    await kernelClose(this.lifecycle, this.lifecycleHandle, this.context, this._id).catch((closeErr) => {
+    await this.close().catch((closeErr) => {
       const cause = closeErr instanceof Error ? closeErr.message : String(closeErr);
       this.context
         .getLogger()
@@ -779,9 +796,16 @@ export default class KernelOperationBackend implements IOperationBackend {
     }
     const metadata = await this.getResultMetadata();
     const handle = await this.getFetchHandle();
-    // KernelResultsProvider consumes only `fetchNextBatch`; both the async result
-    // handle and the blocking statement satisfy that surface.
-    this.resultsProvider = new KernelResultsProvider(handle as unknown as KernelStatement);
+    failIfNotActive(this.lifecycle);
+    if (this.resultSlicer) {
+      return this.resultSlicer;
+    }
+    // Only AsyncResultHandle has a fetch lock separate from close/cancel.
+    this.resultsProvider = new KernelResultsProvider(
+      handle,
+      this.asyncStatement !== undefined,
+      this.context.getLogger(),
+    );
     // DECIMAL/BIGINT precision preservation is opt-in via the
     // `preserveBigNumericPrecision` connection option (default off). The kernel
     // always delivers native Arrow Decimal128 / Int64, so when enabled the

lib/kernel/KernelResultsProvider.ts

@@ -14,6 +14,7 @@
 
 import IResultsProvider, { ResultsProviderFetchNextOptions } from '../result/IResultsProvider';
 import { ArrowBatch } from '../result/utils';
+import IDBSQLLogger, { LogLevel } from '../contracts/IDBSQLLogger';
 import { countRowsInIpc, patchIpcBytes } from './KernelArrowIpc';
 
 /**
@@ -26,6 +27,8 @@ export interface KernelFetchHandle {
   fetchNextBatch(): Promise<{ ipcBytes: Buffer } | null>;
 }
 
+type Settled = { ok: true; value: { ipcBytes: Buffer } | null } | { ok: false; error: unknown };
+
 /**
  * `IResultsProvider<ArrowBatch>` that pulls Arrow IPC batches from the
  * kernel via the napi `Statement` handle and adapts them onto the
@@ -50,15 +53,44 @@ export interface KernelFetchHandle {
 export default class KernelResultsProvider implements IResultsProvider<ArrowBatch> {
   private readonly statement: KernelFetchHandle;
 
+  private readonly prefetchEnabled: boolean;
+
   // Prefetched next batch so `hasMore()` can be answered without an
   // extra round-trip. Set by `prime()` (lazy) and by `fetchNext`.
   private prefetched?: ArrowBatch;
 
   // Set once the kernel returns `null` from `fetchNextBatch()`.
   private exhausted = false;
 
-  constructor(statement: KernelFetchHandle) {
+  private inflight?: Promise<Settled>;
+
+  private primePromise?: Promise<void>;
+
+  private disposed = false;
+
+  private readonly logger?: IDBSQLLogger;
+
+  constructor(statement: KernelFetchHandle, prefetchEnabled = false, logger?: IDBSQLLogger) {
     this.statement = statement;
+    this.prefetchEnabled = prefetchEnabled;
+    this.logger = logger;
+  }
+
+  public dispose(): void {
+    if (this.disposed) {
+      return;
+    }
+    this.disposed = true;
+    this.exhausted = true;
+    this.prefetched = undefined;
+    // The binding cannot abort an AsyncResultHandle fetch already in progress.
+    const abandoned = this.inflight;
+    this.inflight = undefined;
+    abandoned?.then((settled) => {
+      if (!settled.ok) {
+        this.logger?.log(LogLevel.debug, 'Discarded a failed kernel result fetch after disposal');
+      }
+    });
   }
 
   public async hasMore(): Promise<boolean> {
@@ -84,21 +116,56 @@ export default class KernelResultsProvider implements IResultsProvider<ArrowBatc
     return out;
   }
 
+  private startFetch(): void {
+    if (this.inflight === undefined && !this.exhausted) {
+      // Reflect errors so an abandoned prefetch cannot reject unhandled.
+      try {
+        this.inflight = this.statement.fetchNextBatch().then(
+          (value): Settled => ({ ok: true, value }),
+          (error): Settled => ({ ok: false, error }),
+        );
+      } catch (error) {
+        this.inflight = Promise.resolve({ ok: false, error });
+      }
+    }
+  }
+
   // Pull the next batch from the kernel and stash it in `prefetched`,
   // or mark the stream exhausted. Used by both `hasMore` and `fetchNext`
   // to keep one batch buffered ahead so `hasMore` is accurate without
   // re-asking the kernel.
-  private async prime(): Promise<void> {
+  private prime(): Promise<void> {
+    if (!this.primePromise) {
+      this.primePromise = this.primeNext().finally(() => {
+        this.primePromise = undefined;
+      });
+    }
+    return this.primePromise;
+  }
+
+  private async primeNext(): Promise<void> {
     // Loop rather than self-recurse: a long run of empty batches drains
     // iteratively toward the null sentinel without building a deep promise
     // chain.
     while (!this.exhausted && this.prefetched === undefined) {
+      this.startFetch();
       // eslint-disable-next-line no-await-in-loop
-      const next = await this.statement.fetchNextBatch();
+      const settled = await this.inflight!;
+      this.inflight = undefined;
+      if (this.disposed) {
+        return;
+      }
+      if (!settled.ok) {
+        throw settled.error;
+      }
+      const next = settled.value;
       if (next === null) {
         this.exhausted = true;
         return;
       }
+      if (this.prefetchEnabled) {
+        this.startFetch();
+      }
       // Patch the raw bytes once: rewrite any Arrow `Duration` field to
       // `Int64` with a `databricks.arrow.duration_unit` marker, so that
       // apache-arrow@13 (which predates Duration support) can decode the