@query-farm/vgi-rpc 0.1.0

package/src/types.ts

@@ -0,0 +1,162 @@
+// © Copyright 2025-2026, Query.Farm LLC - https://query.farm
+// SPDX-License-Identifier: Apache-2.0
+
+import { type Schema, RecordBatch, recordBatchFromArrays } from "apache-arrow";
+import { buildLogBatch, coerceInt64 } from "./wire/response.js";
+
+export enum MethodType {
+  UNARY = "unary",
+  STREAM = "stream",
+}
+
+/** Logging interface available to handlers. */
+export interface LogContext {
+  clientLog(level: string, message: string, extra?: Record<string, string>): void;
+}
+
+/** Handler for unary (request-response) RPC methods. */
+export type UnaryHandler = (
+  params: Record<string, any>,
+  ctx: LogContext,
+) => Promise<Record<string, any>> | Record<string, any>;
+
+/** Initialization function for producer streams. Returns the initial state object. */
+export type ProducerInit<S = any> = (
+  params: Record<string, any>,
+) => Promise<S> | S;
+/** Called repeatedly to produce output batches. Call `out.finish()` to end the stream. */
+export type ProducerFn<S = any> = (
+  state: S,
+  out: OutputCollector,
+) => Promise<void> | void;
+
+/** Initialization function for exchange streams. Returns the initial state object. */
+export type ExchangeInit<S = any> = (
+  params: Record<string, any>,
+) => Promise<S> | S;
+/** Called once per input batch. Must emit exactly one output batch per call. */
+export type ExchangeFn<S = any> = (
+  state: S,
+  input: RecordBatch,
+  out: OutputCollector,
+) => Promise<void> | void;
+
+/** Produces a header batch sent before the first output batch in a stream. */
+export type HeaderInit = (
+  params: Record<string, any>,
+  state: any,
+  ctx: LogContext,
+) => Record<string, any>;
+
+export interface MethodDefinition {
+  name: string;
+  type: MethodType;
+  paramsSchema: Schema;
+  resultSchema: Schema;
+  outputSchema?: Schema;
+  inputSchema?: Schema;
+  handler?: UnaryHandler;
+  producerInit?: ProducerInit;
+  producerFn?: ProducerFn;
+  exchangeInit?: ExchangeInit;
+  exchangeFn?: ExchangeFn;
+  headerSchema?: Schema;
+  headerInit?: HeaderInit;
+  doc?: string;
+  defaults?: Record<string, any>;
+  paramTypes?: Record<string, string>;
+}
+
+export interface EmittedBatch {
+  batch: RecordBatch;
+  metadata?: Map<string, string>;
+}
+
+/**
+ * Accumulates output batches during a produce/exchange call.
+ * Enforces that exactly one data batch is emitted per call (plus any number of log batches).
+ */
+export class OutputCollector implements LogContext {
+  private _batches: EmittedBatch[] = [];
+  private _dataBatchIdx: number | null = null;
+  private _finished = false;
+  private _producerMode: boolean;
+  private _outputSchema: Schema;
+  private _serverId: string;
+  private _requestId: string | null;
+
+  constructor(outputSchema: Schema, producerMode = true, serverId = "", requestId: string | null = null) {
+    this._outputSchema = outputSchema;
+    this._producerMode = producerMode;
+    this._serverId = serverId;
+    this._requestId = requestId;
+  }
+
+  get outputSchema(): Schema {
+    return this._outputSchema;
+  }
+
+  get finished(): boolean {
+    return this._finished;
+  }
+
+  get batches(): EmittedBatch[] {
+    return this._batches;
+  }
+
+  /** Emit a pre-built RecordBatch as the data batch for this call. */
+  emit(batch: RecordBatch, metadata?: Map<string, string>): void;
+  /** Emit a data batch from column arrays keyed by field name. Int64 Number values are coerced to BigInt. */
+  emit(columns: Record<string, any[]>): void;
+  emit(
+    batchOrColumns: RecordBatch | Record<string, any[]>,
+    metadata?: Map<string, string>,
+  ): void {
+    let batch: RecordBatch;
+    if (batchOrColumns instanceof RecordBatch) {
+      batch = batchOrColumns;
+    } else {
+      const coerced = coerceInt64(this._outputSchema, batchOrColumns);
+      batch = recordBatchFromArrays(coerced, this._outputSchema);
+    }
+    if (this._dataBatchIdx !== null) {
+      throw new Error("Only one data batch may be emitted per call");
+    }
+    this._dataBatchIdx = this._batches.length;
+    this._batches.push({ batch, metadata });
+  }
+
+  /** Single-row convenience. Wraps each value in `[value]` then calls `emit()`. */
+  emitRow(values: Record<string, any>): void {
+    const columns: Record<string, any[]> = {};
+    for (const [key, value] of Object.entries(values)) {
+      columns[key] = [value];
+    }
+    this.emit(columns);
+  }
+
+  /** Signal stream completion for producer streams. Throws if called on exchange streams. */
+  finish(): void {
+    if (!this._producerMode) {
+      throw new Error(
+        "finish() is not allowed on exchange streams; " +
+          "exchange streams must emit exactly one data batch per call",
+      );
+    }
+    this._finished = true;
+  }
+
+  /** Emit a zero-row client-directed log batch. */
+  clientLog(level: string, message: string, extra?: Record<string, string>): void {
+    const batch = buildLogBatch(
+      this._outputSchema,
+      level,
+      message,
+      extra,
+      this._serverId,
+      this._requestId,
+    );
+    this._batches.push({ batch });
+  }
+
+}

package/src/util/schema.ts

@@ -0,0 +1,31 @@
+// © Copyright 2025-2026, Query.Farm LLC - https://query.farm
+// SPDX-License-Identifier: Apache-2.0
+
+import { RecordBatchStreamWriter, type Schema } from "apache-arrow";
+
+/**
+ * Serialize a Schema to the Arrow IPC Schema message format.
+ * This produces bytes compatible with Python's `pa.ipc.read_schema()`.
+ *
+ * We serialize by writing an empty-batch IPC stream and extracting
+ * the bytes, which includes the schema message. Python's read_schema()
+ * uses `pa.ipc.read_schema(pa.py_buffer(bytes))` which expects
+ * the schema flatbuffer message bytes directly — but the Python side
+ * actually uses `schema.serialize()` which produces Schema message bytes.
+ *
+ * In arrow-js, we can get the equivalent by using Message.from(schema)
+ * and encoding it, or by serializing a zero-batch stream.
+ *
+ * The Python `schema.serialize()` produces the Schema flatbuffer message bytes,
+ * and `pa.ipc.read_schema()` expects an IPC stream containing a schema message.
+ * The actual format is: continuation marker (0xFFFFFFFF) + length + flatbuffer bytes.
+ */
+export function serializeSchema(schema: Schema): Uint8Array {
+  // Write a complete IPC stream with no batches.
+  // This writes: Schema message + EOS marker.
+  // Python's pa.ipc.read_schema() can read this format.
+  const writer = new RecordBatchStreamWriter();
+  writer.reset(undefined, schema);
+  writer.close();
+  return writer.toUint8Array(true);
+}

package/src/util/zstd.ts

@@ -0,0 +1,49 @@
+// © Copyright 2025-2026, Query.Farm LLC - https://query.farm
+// SPDX-License-Identifier: Apache-2.0
+
+/**
+ * Cross-runtime zstd compression/decompression.
+ *
+ * Uses Bun.zstd* when running on Bun, otherwise falls back to node:zlib
+ * (available on Node.js 22.15+ and Deno 2.6.9+).
+ */
+
+import * as zlib from "node:zlib";
+
+const isBun = typeof globalThis.Bun !== "undefined";
+
+/** Compress data with zstd at the given level (1-22). */
+export function zstdCompress(data: Uint8Array, level: number): Uint8Array<ArrayBuffer> {
+  if (isBun) {
+    return new Uint8Array(Bun.zstdCompressSync(data, { level }));
+  }
+  const fn = (zlib as any).zstdCompressSync;
+  if (typeof fn !== "function") {
+    throw new Error(
+      "zstd is not available in this runtime. " +
+        "Requires Bun, Node.js >= 22.15, or Deno >= 2.6.9.",
+    );
+  }
+  return new Uint8Array(
+    fn(data, {
+      params: {
+        [(zlib.constants as any).ZSTD_c_compressionLevel]: level,
+      },
+    }),
+  );
+}
+
+/** Decompress zstd-compressed data. */
+export function zstdDecompress(data: Uint8Array): Uint8Array<ArrayBuffer> {
+  if (isBun) {
+    return new Uint8Array(Bun.zstdDecompressSync(data));
+  }
+  const fn = (zlib as any).zstdDecompressSync;
+  if (typeof fn !== "function") {
+    throw new Error(
+      "zstd is not available in this runtime. " +
+        "Requires Bun, Node.js >= 22.15, or Deno >= 2.6.9.",
+    );
+  }
+  return new Uint8Array(fn(data));
+}

package/src/wire/reader.ts

@@ -0,0 +1,113 @@
+// © Copyright 2025-2026, Query.Farm LLC - https://query.farm
+// SPDX-License-Identifier: Apache-2.0
+
+import { RecordBatchReader, type RecordBatch, type Schema } from "apache-arrow";
+
+export interface StreamMessage {
+  schema: Schema;
+  batches: RecordBatch[];
+}
+
+/**
+ * Reads sequential IPC streams from a byte source (e.g., process.stdin).
+ * Uses autoDestroy: false + reset/open pattern to read multiple streams
+ * from the same underlying byte source.
+ */
+export class IpcStreamReader {
+  private reader: RecordBatchReader;
+  private initialized = false;
+  /** True once readNextBatch() returns null (EOS reached for current stream). */
+  private streamEnded = false;
+
+  private constructor(reader: RecordBatchReader) {
+    this.reader = reader;
+  }
+
+  static async create(
+    input: ReadableStream<Uint8Array> | NodeJS.ReadableStream,
+  ): Promise<IpcStreamReader> {
+    const reader = await RecordBatchReader.from(input as any);
+    await reader.open({ autoDestroy: false });
+    if (reader.closed) {
+      throw new Error("Input stream closed before first IPC message");
+    }
+    return new IpcStreamReader(reader);
+  }
+
+  /**
+   * Read one complete IPC stream (schema + all batches).
+   * Returns null on EOF (no more streams).
+   */
+  async readStream(): Promise<StreamMessage | null> {
+    if (this.initialized) {
+      // Advance to next stream
+      await this.reader.reset().open();
+      if (this.reader.closed) {
+        return null;
+      }
+    }
+    this.initialized = true;
+
+    const schema = this.reader.schema;
+    if (!schema) {
+      return null;
+    }
+
+    const batches: RecordBatch[] = [];
+    while (true) {
+      const result = await this.reader.next();
+      if (result.done) break;
+      // Skip Arrow-JS synthetic placeholder for empty streams
+      if (result.value.constructor.name === "_InternalEmptyPlaceholderRecordBatch") break;
+      batches.push(result.value);
+    }
+
+    return { schema, batches };
+  }
+
+  /**
+   * Open the next IPC stream and return its schema.
+   * Use readNextBatch() to read batches one at a time.
+   * Returns null on EOF.
+   */
+  async openNextStream(): Promise<Schema | null> {
+    if (this.initialized) {
+      await this.reader.reset().open();
+      if (this.reader.closed) {
+        return null;
+      }
+    }
+    this.initialized = true;
+    this.streamEnded = false;
+    return this.reader.schema ?? null;
+  }
+
+  /**
+   * Read the next batch from the currently open IPC stream.
+   * Returns null when the stream ends (EOS).
+   *
+   * Once EOS is reached, subsequent calls return null immediately without
+   * reading from the underlying byte source. This prevents the Arrow-JS
+   * reader from consuming bytes that belong to the next IPC stream.
+   */
+  async readNextBatch(): Promise<RecordBatch | null> {
+    if (this.streamEnded) return null;
+    const result = await this.reader.next();
+    if (result.done) {
+      this.streamEnded = true;
+      return null;
+    }
+    // Arrow-JS synthesizes a placeholder batch for streams with a schema but
+    // zero real batches. Treat it as EOS so callers don't block trying to
+    // read more bytes from a stream that has already ended.
+    if (result.value.constructor.name === "_InternalEmptyPlaceholderRecordBatch") {
+      this.streamEnded = true;
+      return null;
+    }
+    return result.value;
+  }
+
+  async cancel(): Promise<void> {
+    await this.reader.cancel();
+  }
+}

package/src/wire/request.ts

@@ -0,0 +1,98 @@
+// © Copyright 2025-2026, Query.Farm LLC - https://query.farm
+// SPDX-License-Identifier: Apache-2.0
+
+import { DataType, type Schema, type RecordBatch } from "apache-arrow";
+import {
+  RPC_METHOD_KEY,
+  REQUEST_VERSION_KEY,
+  REQUEST_VERSION,
+  REQUEST_ID_KEY,
+} from "../constants.js";
+import { RpcError, VersionError } from "../errors.js";
+
+export interface ParsedRequest {
+  methodName: string;
+  requestVersion: string;
+  requestId: string | null;
+  schema: Schema;
+  params: Record<string, any>;
+  rawMetadata: Map<string, string>;
+}
+
+/**
+ * Parse a request from a RecordBatch with metadata.
+ * Extracts method name, version, and params from the batch.
+ */
+export function parseRequest(
+  schema: Schema,
+  batch: RecordBatch,
+): ParsedRequest {
+  const metadata: Map<string, string> = batch.metadata ?? new Map();
+
+  const methodName = metadata.get(RPC_METHOD_KEY);
+  if (methodName === undefined) {
+    throw new RpcError(
+      "ProtocolError",
+      "Missing 'vgi_rpc.method' in request batch custom_metadata. " +
+        "Each request batch must carry a 'vgi_rpc.method' key in its Arrow IPC custom_metadata " +
+        "with the method name as a UTF-8 string.",
+      "",
+    );
+  }
+
+  const version = metadata.get(REQUEST_VERSION_KEY);
+  if (version === undefined) {
+    throw new VersionError(
+      "Missing 'vgi_rpc.request_version' in request batch custom_metadata. " +
+        `Set the 'vgi_rpc.request_version' custom_metadata value to '${REQUEST_VERSION}'.`,
+    );
+  }
+  if (version !== REQUEST_VERSION) {
+    throw new VersionError(
+      `Unsupported request version '${version}', expected '${REQUEST_VERSION}'. ` +
+        `Set the 'vgi_rpc.request_version' custom_metadata value to '${REQUEST_VERSION}'.`,
+    );
+  }
+
+  const requestId = metadata.get(REQUEST_ID_KEY) ?? null;
+
+  // Extract params from single-row batch
+  const params: Record<string, any> = {};
+  if (schema.fields.length > 0 && batch.numRows !== 1) {
+    throw new RpcError(
+      "ProtocolError",
+      `Expected 1 row in request batch, got ${batch.numRows}. ` +
+        "Each parameter is a column (not a row). The batch should have exactly 1 row.",
+      "",
+    );
+  }
+
+  for (let i = 0; i < schema.fields.length; i++) {
+    const field = schema.fields[i];
+    // Map_ columns have a broken .get() in arrow-js — pass through raw Data
+    if (DataType.isMap(field.type)) {
+      params[field.name] = batch.getChildAt(i)!.data[0];
+      continue;
+    }
+    let value = batch.getChildAt(i)?.get(0);
+    // Convert BigInt to Number when safe
+    if (typeof value === "bigint") {
+      if (
+        value >= BigInt(Number.MIN_SAFE_INTEGER) &&
+        value <= BigInt(Number.MAX_SAFE_INTEGER)
+      ) {
+        value = Number(value);
+      }
+    }
+    params[field.name] = value;
+  }
+
+  return {
+    methodName,
+    requestVersion: version,
+    requestId,
+    schema,
+    params,
+    rawMetadata: metadata,
+  };
+}

package/src/wire/response.ts

@@ -0,0 +1,181 @@
+// © Copyright 2025-2026, Query.Farm LLC - https://query.farm
+// SPDX-License-Identifier: Apache-2.0
+
+import {
+  RecordBatch,
+  Schema,
+  Field,
+  Data,
+  DataType,
+  makeData,
+  Struct,
+  vectorFromArray,
+} from "apache-arrow";
+import {
+  LOG_LEVEL_KEY,
+  LOG_MESSAGE_KEY,
+  LOG_EXTRA_KEY,
+  SERVER_ID_KEY,
+  REQUEST_ID_KEY,
+} from "../constants.js";
+
+/**
+ * Coerce values for Int64 schema fields from Number to BigInt.
+ * Handles both single values and arrays. Returns a new record with coerced values.
+ */
+export function coerceInt64(schema: Schema, values: Record<string, any>): Record<string, any> {
+  const result: Record<string, any> = { ...values };
+  for (const field of schema.fields) {
+    const val = result[field.name];
+    if (val === undefined) continue;
+    if (!DataType.isInt(field.type) || (field.type as any).bitWidth !== 64) continue;
+
+    if (Array.isArray(val)) {
+      result[field.name] = val.map((v: any) => (typeof v === "number" ? BigInt(v) : v));
+    } else if (typeof val === "number") {
+      result[field.name] = BigInt(val);
+    }
+  }
+  return result;
+}
+
+/**
+ * Build a 1-row result batch with optional metadata.
+ * For unary methods, `values` maps field names to single values.
+ */
+export function buildResultBatch(
+  schema: Schema,
+  values: Record<string, any>,
+  serverId: string,
+  requestId: string | null,
+): RecordBatch {
+  const metadata = new Map<string, string>();
+  metadata.set(SERVER_ID_KEY, serverId);
+  if (requestId !== null) {
+    metadata.set(REQUEST_ID_KEY, requestId);
+  }
+
+  if (schema.fields.length === 0) {
+    return buildEmptyBatch(schema, metadata);
+  }
+
+  // Validate required fields
+  for (const field of schema.fields) {
+    if (values[field.name] === undefined && !field.nullable) {
+      const got = Object.keys(values);
+      throw new TypeError(
+        `Handler result missing required field '${field.name}'. Got keys: [${got.join(", ")}]`,
+      );
+    }
+  }
+
+  const coerced = coerceInt64(schema, values);
+
+  const children = schema.fields.map((f: Field) => {
+    let val = coerced[f.name];
+    // Raw Data passthrough for Map_ types (whose .get() is broken in arrow-js)
+    if (val instanceof Data) {
+      return val;
+    }
+    const arr = vectorFromArray([val], f.type);
+    return arr.data[0];
+  });
+
+  const structType = new Struct(schema.fields);
+  const data = makeData({
+    type: structType,
+    length: 1,
+    children,
+    nullCount: 0,
+  });
+
+  return new RecordBatch(schema, data, metadata);
+}
+
+/**
+ * Build a 0-row error batch with EXCEPTION metadata matching Python's Message.from_exception().
+ */
+export function buildErrorBatch(
+  schema: Schema,
+  error: Error,
+  serverId: string,
+  requestId: string | null,
+): RecordBatch {
+  const metadata = new Map<string, string>();
+  metadata.set(LOG_LEVEL_KEY, "EXCEPTION");
+  metadata.set(LOG_MESSAGE_KEY, `${error.constructor.name}: ${error.message}`);
+
+  const extra: Record<string, any> = {
+    exception_type: error.constructor.name,
+    exception_message: error.message,
+    traceback: error.stack ?? "",
+  };
+  metadata.set(LOG_EXTRA_KEY, JSON.stringify(extra));
+  metadata.set(SERVER_ID_KEY, serverId);
+  if (requestId !== null) {
+    metadata.set(REQUEST_ID_KEY, requestId);
+  }
+
+  return buildEmptyBatch(schema, metadata);
+}
+
+/**
+ * Build a 0-row log batch.
+ */
+export function buildLogBatch(
+  schema: Schema,
+  level: string,
+  message: string,
+  extra?: Record<string, any>,
+  serverId?: string,
+  requestId?: string | null,
+): RecordBatch {
+  const metadata = new Map<string, string>();
+  metadata.set(LOG_LEVEL_KEY, level);
+  metadata.set(LOG_MESSAGE_KEY, message);
+  if (extra) {
+    metadata.set(LOG_EXTRA_KEY, JSON.stringify(extra));
+  }
+  if (serverId != null) {
+    metadata.set(SERVER_ID_KEY, serverId);
+  }
+  if (requestId != null) {
+    metadata.set(REQUEST_ID_KEY, requestId);
+  }
+
+  return buildEmptyBatch(schema, metadata);
+}
+
+/**
+ * Build a 0-row batch from a schema with metadata.
+ * Used for error/log batches.
+ */
+export function buildEmptyBatch(
+  schema: Schema,
+  metadata?: Map<string, string>,
+): RecordBatch {
+  const children = schema.fields.map((f: Field) => {
+    return makeData({ type: f.type, length: 0, nullCount: 0 });
+  });
+
+  if (schema.fields.length === 0) {
+    const structType = new Struct(schema.fields);
+    const data = makeData({
+      type: structType,
+      length: 0,
+      children: [],
+      nullCount: 0,
+    });
+    return new RecordBatch(schema, data, metadata);
+  }
+
+  const structType = new Struct(schema.fields);
+  const data = makeData({
+    type: structType,
+    length: 0,
+    children,
+    nullCount: 0,
+  });
+
+  return new RecordBatch(schema, data, metadata);
+}