@query-farm/vgi-rpc 0.6.3 → 0.7.0

package/src/wire/writer.ts

@@ -1,18 +1,50 @@
 // © Copyright 2025-2026, Query.Farm LLC - https://query.farm
 // SPDX-License-Identifier: Apache-2.0
 
-import { writeSync } from "node:fs";
-import { type RecordBatch, RecordBatchStreamWriter, type Schema } from "@query-farm/apache-arrow";
+import type { Socket } from "node:net";
+import type { IncrementalEncoder, VgiBatch, VgiSchema } from "../arrow/index.js";
+import { createIncrementalEncoder, serializeBatches } from "../arrow/index.js";
 
 const STDOUT_FD = 1;
 
+// Resolve node:fs via indirect-string require so esbuild/wrangler don't
+// statically pull node:fs into the bundle. Workers (Cloudflare workerd) never
+// instantiate IpcStreamWriter (no stdio transport on workers), so the
+// runtime-time require("node:fs") is unreachable in those builds.
+//
+// `globalThis.require` is undefined in both Bun ESM and Node ESM, so we try
+// `import.meta.require` (Bun) first, then fall back to globalThis.require
+// (Node CJS). Node ESM consumers must polyfill require if they need the
+// subprocess transport.
+const _NODE_FS_MOD = "node:fs";
+let _writeSync: ((fd: number, data: Uint8Array, offset?: number, len?: number) => number) | null = null;
+function _loadWriteSync(): (fd: number, data: Uint8Array, offset?: number, len?: number) => number {
+  if (_writeSync) return _writeSync;
+  const req: any = (import.meta as any).require ?? (globalThis as any).require ?? null;
+  if (!req) {
+    throw new Error(
+      "IpcStreamWriter requires Bun or Node.js CJS for sync node:fs.writeSync. " +
+        "Subprocess transport is not available in this runtime.",
+    );
+  }
+  const fs = req(_NODE_FS_MOD);
+  _writeSync = fs.writeSync.bind(fs);
+  return _writeSync!;
+}
+
 /**
  * Write all bytes to a file descriptor, looping on partial writes.
- * Handles EAGAIN (pipe buffer full) by busy-waiting with Atomics.wait().
- * writeSync() can return fewer bytes than requested when the pipe buffer
- * is full (e.g., 64KB limit), and throws EAGAIN on non-blocking fds.
+ *
+ * Stdio path only. The fd here is a pipe (typically stdout) inherited from
+ * the launcher in subprocess transport. Pipes set up by fork/exec are
+ * blocking by default, so writeSync blocks the kernel (not the event loop in
+ * a bad way) until the parent reader drains — that is the *desired* lockstep
+ * with the client. The AF_UNIX path does not come through here; sockets are
+ * non-blocking and EAGAIN-spinning them would freeze the event loop and
+ * starve every other connection. See `socketWriteAll` below.
  */
 function writeAll(fd: number, data: Uint8Array): void {
+  const writeSync = _loadWriteSync();
   let offset = 0;
   while (offset < data.length) {
     try {
@@ -21,7 +53,9 @@ function writeAll(fd: number, data: Uint8Array): void {
       offset += written;
     } catch (e: any) {
       if (e.code === "EAGAIN") {
-        // Pipe buffer full — busy-wait briefly then retry
+        // A non-blocking pipe — unexpected for stdio, but handle defensively.
+        // Yielding is not possible from a synchronous function; busy-wait
+        // briefly. AF_UNIX sockets must NOT use this path.
         Atomics.wait(new Int32Array(new SharedArrayBuffer(4)), 0, 0, 1);
         continue;
       }
@@ -31,103 +65,165 @@ function writeAll(fd: number, data: Uint8Array): void {
 }
 
 /**
- * Writes sequential IPC streams to a file descriptor (e.g., stdout).
- * Each call to writeStream() writes a complete IPC stream: schema + batches + EOS.
+ * Asynchronously write all bytes to a Node net.Socket, honouring backpressure.
+ *
+ * `socket.write(buf)` queues `buf` into libuv's per-socket outbound buffer
+ * and returns false when that buffer has grown past `highWaterMark`. We then
+ * `await` the `'drain'` event so the JS thread yields — other connections
+ * can be dispatched while the kernel drains the send buffer in the
+ * background. Without this yield, a slow consumer on connection A would
+ * freeze every other connection's handler on the shared event loop.
+ *
+ * Resolves cleanly on socket close/error so the caller's try/finally can
+ * unwind without dangling listeners.
+ */
+async function socketWriteAll(socket: Socket, data: Uint8Array): Promise<void> {
+  if (socket.destroyed || socket.writableEnded) {
+    throw new Error("socketWriteAll: socket is already closed");
+  }
+  const ok = socket.write(data);
+  if (ok) return;
+  await new Promise<void>((resolve, reject) => {
+    const cleanup = () => {
+      socket.off("drain", onDrain);
+      socket.off("error", onError);
+      socket.off("close", onClose);
+    };
+    const onDrain = () => {
+      cleanup();
+      resolve();
+    };
+    const onError = (err: Error) => {
+      cleanup();
+      reject(err);
+    };
+    const onClose = () => {
+      // Peer hung up before drain. Resolve so the caller's try/finally
+      // runs; the next read on this connection will surface EOF.
+      cleanup();
+      resolve();
+    };
+    socket.once("drain", onDrain);
+    socket.once("error", onError);
+    socket.once("close", onClose);
+  });
+}
+
+type WriterTarget = { kind: "fd"; fd: number } | { kind: "socket"; socket: Socket };
+
+/**
+ * Writes sequential IPC streams to either an fd (stdio subprocess transport)
+ * or a Node Socket (AF_UNIX transport). Each call to writeStream() writes a
+ * complete IPC stream: schema + batches + EOS.
  *
- * All writes use synchronous I/O (writeSync) to avoid deadlocks when
- * interleaving stdout writes with blocking stdin reads.
+ * All public methods are async. The fd path resolves immediately after a
+ * synchronous writeSync; the socket path awaits real `'drain'` events on
+ * backpressure so the event loop stays responsive to other connections.
  */
 export class IpcStreamWriter {
-  private readonly fd: number;
+  private readonly target: WriterTarget;
 
-  constructor(fd: number = STDOUT_FD) {
-    this.fd = fd;
+  /**
+   * Construct from a file descriptor (stdio transport) or a Node net.Socket
+   * (AF_UNIX transport). The default targets stdout for legacy stdio servers
+   * that didn't pass an fd.
+   */
+  constructor(fdOrSocket: number | Socket = STDOUT_FD) {
+    if (typeof fdOrSocket === "number") {
+      this.target = { kind: "fd", fd: fdOrSocket };
+    } else {
+      this.target = { kind: "socket", socket: fdOrSocket };
+    }
   }
 
   /**
    * Write a complete IPC stream with the given schema and batches.
    * Creates schema message, writes all batches (with their metadata), writes EOS.
    */
-  writeStream(schema: Schema, batches: RecordBatch[]): void {
-    const writer = new RecordBatchStreamWriter();
-    writer.reset(undefined, schema);
-    for (const batch of batches) {
-      // Use _writeRecordBatch to bypass schema comparison (see IncrementalStream.write)
-      (writer as any)._writeRecordBatch(batch);
+  async writeStream(schema: VgiSchema, batches: VgiBatch[]): Promise<void> {
+    // Delegate to the Arrow facade so the bytes-on-the-wire match the
+    // active backend (arrow-js by default, flechette under the `workerd`/
+    // `worker`/`browser` package.json conditions). The incremental stream
+    // below still uses arrow-js directly — flechette has no equivalent
+    // streaming-writer surface and only the stdio server uses the
+    // incremental path.
+    const bytes = serializeBatches(schema, batches);
+    if (this.target.kind === "fd") {
+      writeAll(this.target.fd, bytes);
+    } else {
+      await socketWriteAll(this.target.socket, bytes);
     }
-    writer.close();
-    const bytes = writer.toUint8Array(true);
-    writeAll(this.fd, bytes);
   }
 
   /**
    * Open an incremental IPC stream for writing batches one at a time.
-   * Used for streaming methods where output batches are produced incrementally.
-   * Bytes are written synchronously after each batch.
    */
-  openStream(schema: Schema): IncrementalStream {
-    return new IncrementalStream(this.fd, schema);
+  openStream(schema: VgiSchema): IncrementalStream {
+    return new IncrementalStream(this.target, schema);
   }
 }
 
 /**
  * An open IPC stream that supports incremental batch writes.
  *
- * Uses RecordBatchStreamWriter with internal buffering (no pipe to stdout).
- * After each operation, drains the writer's internal AsyncByteQueue buffer
- * and writes bytes synchronously via writeAll(). This avoids deadlocks
- * caused by Node.js async stream piping when stdin reads block before
- * stdout writes flush through the event loop.
+ * Drives a backend {@link IncrementalEncoder} and flushes its bytes through
+ * the same target (fd or socket) as the parent IpcStreamWriter. The write()
+ * and close() methods are async so the socket path can yield on backpressure
+ * — critical under AF_UNIX where the kernel send buffer (~8 KB on macOS)
+ * fills quickly and any synchronous busy-wait would starve every other
+ * connection sharing this event loop.
+ *
+ * The encoder is obtained from the Arrow facade, so this file no longer
+ * imports arrow-js directly — keeping arrow-js out of the flechette
+ * (workerd/browser) bundle. The flechette encoder throws on construction;
+ * the stdio exchange protocol is lockstep (the client reads each response
+ * batch before sending the next input) which needs an incremental writer
+ * flechette doesn't provide. workerd/browser deployments use HTTP (no
+ * stdio), so the flechette path is never reached there; `flechette-pipe`
+ * conformance is xfailed for streams.
  */
 export class IncrementalStream {
-  private writer: RecordBatchStreamWriter;
-  private readonly fd: number;
+  private readonly encoder: IncrementalEncoder;
+  private readonly target: WriterTarget;
   private closed = false;
+  // Chains async writes so concurrent callers can't interleave bytes on the
+  // wire. Each operation appends to this promise; awaiting it preserves
+  // FIFO order even when the caller doesn't await between writes.
+  private writeChain: Promise<void> = Promise.resolve();
 
-  constructor(fd: number, schema: Schema) {
-    this.fd = fd;
-    this.writer = new RecordBatchStreamWriter();
-    // Buffer internally (no sink) — we drain manually via writeAll
-    this.writer.reset(undefined, schema);
-    this.drain();
+  constructor(target: WriterTarget, schema: VgiSchema) {
+    this.target = target;
+    this.encoder = createIncrementalEncoder(schema);
+    // Schema bytes — synchronous on the fd path, queued on the socket path.
+    // Callers don't await openStream(), so we serialize via writeChain.
+    this.enqueue(this.encoder.start());
   }
 
-  /**
-   * Write a single batch to the stream. Bytes are flushed synchronously.
-   *
-   * Uses _writeRecordBatch() directly to bypass the Arrow writer's schema
-   * comparison in write(). The public write() method calls compareSchemas()
-   * and auto-closes the writer if the batch's schema differs (e.g., in
-   * nullability), silently dropping the batch. Since our output schema is
-   * set at stream open time and all batches are structurally compatible,
-   * we skip the comparison.
-   */
-  write(batch: RecordBatch): void {
+  /** Write a single batch. Resolves once the bytes are queued/flushed. */
+  async write(batch: VgiBatch): Promise<void> {
     if (this.closed) throw new Error("Stream already closed");
-    (this.writer as any)._writeRecordBatch(batch);
-    this.drain();
+    return this.enqueue(this.encoder.writeBatch(batch));
   }
 
-  /**
-   * Close the stream (writes EOS marker synchronously).
-   */
-  close(): void {
+  /** Close the stream (writes EOS marker). */
+  async close(): Promise<void> {
     if (this.closed) return;
     this.closed = true;
-    // EOS marker: continuation (0xFFFFFFFF) + metadata length (0x00000000)
-    const eos = new Uint8Array(new Int32Array([-1, 0]).buffer);
-    writeAll(this.fd, eos);
+    return this.enqueue(this.encoder.finish());
   }
 
-  /**
-   * Drain buffered bytes from the Arrow writer's internal queue
-   * and write them synchronously to the output fd.
-   */
-  private drain(): void {
-    const values = (this.writer as any)._sink._values as Uint8Array[];
-    for (const chunk of values) {
-      writeAll(this.fd, chunk);
-    }
-    values.length = 0;
+  private enqueue(bytes: Uint8Array): Promise<void> {
+    const next = this.writeChain.then(() => {
+      if (this.target.kind === "fd") {
+        writeAll(this.target.fd, bytes);
+        return;
+      }
+      return socketWriteAll(this.target.socket, bytes);
+    });
+    // Swallow rejections on the chain so a single failure doesn't poison
+    // every subsequent enqueue with an unhandled rejection. The returned
+    // `next` still propagates the error to the caller that triggered it.
+    this.writeChain = next.catch(() => undefined);
+    return next;
   }
 }

package/dist/util/conform.d.ts

@@ -1,18 +0,0 @@
-import { RecordBatch, type Schema } from "@query-farm/apache-arrow";
-/**
- * Rebuild a batch's data to match the given schema's field types.
- *
- * Batches deserialized from IPC streams (e.g., from PyArrow) may use generic
- * types (Float) instead of specific ones (Float64).  Arrow-JS's
- * RecordBatchStreamWriter silently drops batches whose child Data types don't
- * match the writer's schema.  Cloning each child Data with the schema's field
- * type fixes the type metadata while preserving the underlying buffers.
- *
- * This is also used to cast compatible input types (e.g., int32→float64,
- * float32→float64) when the input batch schema doesn't exactly match the
- * method's declared input schema.  When the underlying buffer layout differs
- * (e.g., 4-byte int32 vs 8-byte float64), we read the values and build a
- * new vector with the target type.
- */
-export declare function conformBatchToSchema(batch: RecordBatch, schema: Schema): RecordBatch;
-//# sourceMappingURL=conform.d.ts.map

package/dist/util/conform.d.ts.map

@@ -1 +0,0 @@
-{"version":3,"file":"conform.d.ts","sourceRoot":"","sources":["../../src/util/conform.ts"],"names":[],"mappings":"AAGA,OAAO,EAGL,WAAW,EACX,KAAK,MAAM,EAIZ,MAAM,0BAA0B,CAAC;AAkBlC;;;;;;;;;;;;;;GAcG;AACH,wBAAgB,oBAAoB,CAAC,KAAK,EAAE,WAAW,EAAE,MAAM,EAAE,MAAM,GAAG,WAAW,CAiDpF"}

package/src/util/conform.ts

@@ -1,94 +0,0 @@
-// © Copyright 2025-2026, Query.Farm LLC - https://query.farm
-// SPDX-License-Identifier: Apache-2.0
-
-import {
-  type DataType,
-  makeData,
-  RecordBatch,
-  type Schema,
-  Struct,
-  Type,
-  vectorFromArray,
-} from "@query-farm/apache-arrow";
-
-/** Return true when the source type's values can be losslessly read and
- *  re-encoded into the target type (e.g., int32 → float64). */
-function needsValueCast(src: DataType, dst: DataType): boolean {
-  if (src.typeId === dst.typeId) return false;
-  // Same broad family (e.g. Float → Float64) — clone is sufficient.
-  if (src.constructor === dst.constructor) return false;
-  return true;
-}
-
-/** Check if a type is a numeric type we can cast between.
- *  Uses typeId instead of instanceof because IPC-deserialized types
- *  may be generic (e.g., Int_ instead of Int64). */
-function isNumeric(t: DataType): boolean {
-  return t.typeId === Type.Int || t.typeId === Type.Float;
-}
-
-/**
- * Rebuild a batch's data to match the given schema's field types.
- *
- * Batches deserialized from IPC streams (e.g., from PyArrow) may use generic
- * types (Float) instead of specific ones (Float64).  Arrow-JS's
- * RecordBatchStreamWriter silently drops batches whose child Data types don't
- * match the writer's schema.  Cloning each child Data with the schema's field
- * type fixes the type metadata while preserving the underlying buffers.
- *
- * This is also used to cast compatible input types (e.g., int32→float64,
- * float32→float64) when the input batch schema doesn't exactly match the
- * method's declared input schema.  When the underlying buffer layout differs
- * (e.g., 4-byte int32 vs 8-byte float64), we read the values and build a
- * new vector with the target type.
- */
-export function conformBatchToSchema(batch: RecordBatch, schema: Schema): RecordBatch {
-  if (batch.numRows === 0) return batch;
-
-  // Validate field count and names match before attempting any cast.
-  if (batch.schema.fields.length !== schema.fields.length) {
-    throw new TypeError(`Field count mismatch: expected ${schema.fields.length}, got ${batch.schema.fields.length}`);
-  }
-  for (let i = 0; i < schema.fields.length; i++) {
-    if (batch.schema.fields[i].name !== schema.fields[i].name) {
-      throw new TypeError(
-        `Field name mismatch at index ${i}: expected '${schema.fields[i].name}', got '${batch.schema.fields[i].name}'`,
-      );
-    }
-  }
-
-  const children = schema.fields.map((f, i) => {
-    const srcChild = batch.data.children[i];
-    const srcType = srcChild.type;
-    const dstType = f.type;
-
-    if (!needsValueCast(srcType, dstType)) {
-      return srcChild.clone(dstType);
-    }
-
-    // Numeric → numeric: read values and rebuild with target type.
-    if (isNumeric(srcType) && isNumeric(dstType)) {
-      // Read source values via the batch's column vector.
-      const col = batch.getChildAt(i)!;
-      const values: number[] = [];
-      for (let r = 0; r < batch.numRows; r++) {
-        const v = col.get(r);
-        values.push(typeof v === "bigint" ? Number(v) : (v as number));
-      }
-      return vectorFromArray(values, dstType).data[0];
-    }
-
-    // Fallback: clone type metadata (works for same-layout types).
-    return srcChild.clone(dstType);
-  });
-
-  const structType = new Struct(schema.fields);
-  const data = makeData({
-    type: structType,
-    length: batch.numRows,
-    children,
-    nullCount: batch.data.nullCount,
-    nullBitmap: batch.data.nullBitmap,
-  });
-  return new RecordBatch(schema, data, batch.metadata);
-}