@milaboratories/pf-driver 1.3.10 → 1.4.0

package/src/csv_writer.ts

@@ -0,0 +1,154 @@
+import {
+  Annotation,
+  isValueNA,
+  readAnnotation,
+  ValueType,
+  type PTableColumnSpec,
+  type PTableVector,
+  type TableRange,
+} from "@milaboratories/pl-model-common";
+import { isNil } from "@milaboratories/helpers";
+
+/** Minimal subset of PTableV8 required by streamPTableRows. */
+export interface PTableDataSource {
+  getData(
+    columnIndices: number[],
+    options?: { range?: TableRange; signal?: AbortSignal },
+  ): Promise<PTableVector[]>;
+}
+
+// ── Public API (high-level → low-level) ──────────────────────────────
+
+/** Format a CSV/TSV header row from column specs. Line ending is CRLF. */
+export function formatHeader(specs: PTableColumnSpec[], separator: string): string {
+  return specs.map((spec) => escapeField(columnLabel(spec), separator)).join(separator) + "\r\n";
+}
+
+/** Format a single data row from parallel vectors. Line ending is CRLF. */
+export function formatRow(vectors: PTableVector[], rowIndex: number, separator: string): string {
+  return (
+    vectors
+      .map((vector) => escapeField(serializeValue(vector, rowIndex), separator))
+      .join(separator) + "\r\n"
+  );
+}
+
+/**
+ * Async generator that streams CSV/TSV content chunk by chunk.
+ *
+ * The caller is responsible for providing a concrete `range` (already clipped
+ * to the table shape). When `range` is undefined the generator does nothing
+ * beyond emitting an optional BOM and header.
+ */
+export interface StreamPTableRowsOptions {
+  pTable: PTableDataSource;
+  specs: PTableColumnSpec[];
+  columnIndices: number[];
+  range?: TableRange;
+  chunkSize: number;
+  separator: string;
+  includeHeader: boolean;
+  bom: boolean;
+  signal?: AbortSignal;
+}
+export async function* streamPTableRows(options: StreamPTableRowsOptions): AsyncIterable<string> {
+  const { pTable, columnIndices, range, chunkSize, separator, signal, specs, includeHeader, bom } =
+    options;
+
+  if (bom) {
+    yield "\uFEFF";
+  }
+
+  if (includeHeader) {
+    const selectedSpecs = columnIndices.map((index) => specs[index]);
+    yield formatHeader(selectedSpecs, separator);
+  }
+
+  if (isNil(range)) {
+    return;
+  }
+
+  const end = range.offset + range.length;
+
+  for (let from = range.offset; from < end; from += chunkSize) {
+    signal?.throwIfAborted();
+
+    const length = Math.min(chunkSize, end - from);
+    const subRange: TableRange = { offset: from, length };
+
+    const vectors = await pTable.getData(columnIndices, { range: subRange, signal });
+
+    const rows: string[] = [];
+    for (let rowIndex = 0; rowIndex < length; rowIndex++) {
+      rows.push(formatRow(vectors, rowIndex, separator));
+    }
+    yield rows.join("");
+  }
+}
+
+// ── Helpers (low-level) ──────────────────────────────────────────────
+
+/** Extract a human-readable label from a PTableColumnSpec. */
+function columnLabel(spec: PTableColumnSpec): string {
+  const annotation = readAnnotation(spec.spec, Annotation.Label);
+  return isNil(annotation) ? spec.spec.name : annotation.trim();
+}
+
+/**
+ * RFC 4180 field escaping.
+ * Quote if the field contains the separator, a double-quote, CR, or LF.
+ * Embedded `"` are doubled.
+ */
+function escapeField(value: string, separator: string): string {
+  return needsQuoting(value, separator) ? '"' + value.replace(/"/g, '""') + '"' : value;
+}
+
+/** Returns true when the value must be wrapped in double-quotes. */
+function needsQuoting(value: string, separator: string): boolean {
+  return (
+    value.includes(separator) || value.includes('"') || value.includes("\r") || value.includes("\n")
+  );
+}
+
+/**
+ * Serialize one cell value from a typed vector to its string representation.
+ *
+ * - `null` / `undefined`  -> `""`
+ * - `bigint`              -> `String(x)`
+ * - `NaN` / `+Inf` / `-Inf` -> `""`
+ */
+function serializeValue(vector: PTableVector, rowIndex: number): string {
+  const rawValue = vector.data[rowIndex];
+
+  if (isNil(rawValue)) {
+    return "";
+  }
+
+  if (isValueNA(vector, rowIndex)) {
+    return "";
+  }
+
+  switch (vector.type) {
+    case ValueType.Long: {
+      // BigInt64Array element — may be stored as bigint
+      return String(rawValue);
+    }
+    case ValueType.Float:
+    case ValueType.Double: {
+      const numeric = rawValue as number;
+      return Number.isNaN(numeric) || !Number.isFinite(numeric) ? "" : String(numeric);
+    }
+    case ValueType.Int: {
+      return String(rawValue);
+    }
+    case ValueType.String: {
+      return rawValue as string;
+    }
+    case ValueType.Bytes: {
+      return "";
+    }
+    default: {
+      return String(rawValue);
+    }
+  }
+}

package/src/driver_decl.ts

@@ -14,9 +14,17 @@ import type {
   TableRange,
   UniqueValuesRequest,
   UniqueValuesResponse,
+  WritePTableToFsOptions,
+  WritePTableToFsResult,
 } from "@milaboratories/pl-model-common";
 import type { PoolEntry } from "@milaboratories/helpers";
 
+export type {
+  WritePTableToFsOptions,
+  WritePTableToFsResult,
+  PTableDownloadFormat,
+} from "@milaboratories/pl-model-common";
+
 /**
  * Extends public and safe SDK's driver API with methods used internally in the middle
  * layer and in tests.
@@ -75,4 +83,10 @@ export interface AbstractInternalPFrameDriver<PColumnData> extends PFrameDriver,
     range: TableRange | undefined,
     signal?: AbortSignal,
   ): Promise<PTableVector[]>;
+
+  /** Download PTable data to a file in CSV or TSV format. */
+  writePTableToFs(
+    handle: PTableHandle,
+    options: WritePTableToFsOptions,
+  ): Promise<WritePTableToFsResult>;
 }

package/src/driver_impl.ts

@@ -35,11 +35,24 @@ import {
   resolveAnnotationParents,
 } from "@milaboratories/pl-model-common";
 import type { PFrameInternal } from "@milaboratories/pl-model-middle-layer";
-import { ConcurrencyLimitingExecutor } from "@milaboratories/ts-helpers";
-import { PoolEntryGuard, type PoolEntry } from "@milaboratories/helpers";
+import {
+  ConcurrencyLimitingExecutor,
+  createPathAtomically,
+  type MiLogger,
+} from "@milaboratories/ts-helpers";
+import { isNil, PoolEntryGuard, type PoolEntry } from "@milaboratories/helpers";
 import { PFrameFactory } from "@milaboratories/pframes-rs-node";
 import { tmpdir } from "node:os";
-import type { AbstractInternalPFrameDriver } from "./driver_decl";
+import * as fs from "node:fs";
+import { Readable } from "node:stream";
+import { pipeline } from "node:stream/promises";
+import * as zlib from "node:zlib";
+import { streamPTableRows } from "./csv_writer";
+import type {
+  AbstractInternalPFrameDriver,
+  WritePTableToFsOptions,
+  WritePTableToFsResult,
+} from "./driver_decl";
 import { logPFrames } from "./logging";
 import {
   PFramePool,
@@ -259,6 +272,80 @@ export class AbstractPFrameDriver<
     };
   }
 
+  public async writePTableToFs(
+    handle: PTableHandle,
+    options: WritePTableToFsOptions,
+  ): Promise<WritePTableToFsResult> {
+    this.logger(
+      "info",
+      `[WritePTableToFs] ENTER (handle = ${handle}, path = ${options.path}, format = ${options.format}, compression = ${options.compression ?? "auto"}, columns = ${options.columnIndices.length})`,
+    );
+    const startTime = performance.now();
+    const { def, disposeSignal: defDisposeSignal } = this.pTableDefs.getByKey(handle);
+    using tableGuard = new PoolEntryGuard(this.pTables.acquire(def));
+    const { pTablePromise, disposeSignal } = tableGuard.resource;
+    const pTable = await pTablePromise;
+
+    const combinedSignal = AbortSignal.any(
+      [options.signal, disposeSignal].filter((s): s is AbortSignal => !isNil(s)),
+    );
+
+    return await this.tableConcurrencyLimiter.run(async () => {
+      const shape = await pTable.getShape({ signal: combinedSignal });
+      const clippedRange = clipRange(options.range, shape);
+      const specs = pTable.getSpec();
+      const separator = options.format === "tsv" ? "\t" : ",";
+
+      const iterable = streamPTableRows({
+        pTable,
+        specs,
+        columnIndices: options.columnIndices,
+        range: clippedRange,
+        chunkSize: options.chunkSize ?? 50_000,
+        separator,
+        includeHeader: options.includeHeader ?? true,
+        bom: options.bom ?? true,
+        signal: combinedSignal,
+      });
+
+      const miLogger: MiLogger = {
+        info: (msg) => this.logger("info", String(msg)),
+        warn: (msg) => this.logger("warn", String(msg)),
+        error: (msg) => this.logger("error", String(msg)),
+      };
+
+      let bytesWritten = 0;
+      await createPathAtomically(miLogger, options.path, async (tempPath) => {
+        const writeStream = fs.createWriteStream(tempPath, { flags: "wx" });
+        const source = Readable.from(iterable, { objectMode: false });
+        if (options.compression?.type === "gzip") {
+          const gzip = zlib.createGzip({ level: options.compression.level ?? 6 });
+          await pipeline(source, gzip, writeStream, { signal: combinedSignal });
+        } else {
+          await pipeline(source, writeStream, { signal: combinedSignal });
+        }
+        bytesWritten = writeStream.bytesWritten;
+      });
+
+      const overallSize = await pTable.getFootprint({ signal: combinedSignal });
+      this.pTableCachePlain.cache(tableGuard.keep(), overallSize, defDisposeSignal);
+
+      // rowsWritten equals the clipped range length — the generator streams the
+      // entire effective range without early termination, so this is accurate.
+      const rowsWritten = clippedRange.length;
+
+      if (logPFrames()) {
+        const durationMs = Math.round(performance.now() - startTime);
+        this.logger(
+          "info",
+          `[WritePTableToFs] complete (handle = ${handle}, columns = ${options.columnIndices.length}, rows = ${rowsWritten}, bytes = ${bytesWritten}, duration = ${durationMs}ms)`,
+        );
+      }
+
+      return { path: options.path, rowsWritten, bytesWritten };
+    });
+  }
+
   //
   // PFrame instance methods
   //
@@ -463,6 +550,16 @@ export class AbstractPFrameDriver<
   }
 }
 
+/** Clamp range to table shape. When range is undefined, returns full table range. */
+function clipRange(range: undefined | TableRange, shape: PTableShape): TableRange {
+  if (isNil(range)) {
+    return { offset: 0, length: shape.rows };
+  }
+  const clampedOffset = Math.min(range.offset, shape.rows);
+  const clampedLength = Math.min(range.length, shape.rows - clampedOffset);
+  return { offset: clampedOffset, length: clampedLength };
+}
+
 function migrateFilters(
   filters: PTableRecordFilter[],
   logger: PFrameInternal.Logger,