@milaboratories/pf-driver 1.4.11 → 1.4.13

package/src/driver_impl.ts

@@ -71,7 +71,6 @@ import {
   PTableCachePlainOpsDefaults,
   type PTableCachePlainOps,
 } from "./ptable_cache_plain";
-import { createPFrame as createSpecFrame } from "@milaboratories/pframes-rs-wasm";
 
 export interface LocalBlobProvider<TreeEntry extends JsonSerializable>
   extends PoolLocalBlobProvider<TreeEntry>, AsyncDisposable {}
@@ -194,26 +193,31 @@ export class AbstractPFrameDriver<
   }
 
   public createPTable(rawDef: PTableDef<PColumn<PColumnData>>): PoolEntry<PTableHandle> {
-    using pFrameGuard = new PoolEntryGuard(this.createPFrame(extractAllColumns(rawDef.src)));
+    const columns = uniqueBy(extractAllColumns(rawDef.src), (c) => c.id);
+    using pFrameGuard = new PoolEntryGuard(this.createPFrame(columns));
     const sortedDef = sortPTableDef(
       migrateTableFilter(
         mapPTableDef(rawDef, (c) => c.id),
         this.logger,
       ),
     );
-    const pTableEntry = this.pTableDefs.acquire({
-      type: "v1",
-      def: sortedDef,
-      pFrameHandle: pFrameGuard.key,
-    });
+    const pFrameSpec = this.pFrames.getByKey(pFrameGuard.key).pFrameSpec;
+    using pTableGuard = new PoolEntryGuard(
+      this.pTableDefs.acquireFromLegacy({
+        pFrameHandle: pFrameGuard.key,
+        def: sortedDef,
+        pFrameSpec,
+      }).entry,
+    );
     if (logPFrames()) {
       this.logger(
         "info",
-        `Create PTable call (pFrameHandle = ${pFrameGuard.key}; pTableHandle = ${pTableEntry.key})`,
+        `Create PTable call (pFrameHandle = ${pFrameGuard.key}; pTableHandle = ${pTableGuard.key})`,
       );
     }
 
     const pFrameEntry = pFrameGuard.keep();
+    const pTableEntry = pTableGuard.keep();
     const unref = () => {
       pTableEntry.unref();
       pFrameEntry.unref();
@@ -228,38 +232,27 @@ export class AbstractPFrameDriver<
 
   public createPTableV2(def: PTableDefV2<PColumn<PColumnData>>): PoolEntry<PTableHandle> {
     const columns = uniqueBy(collectSpecQueryColumns(def.query), (c) => c.id);
-    const columnsMap = columns.reduce(
-      (acc, col) => ((acc[col.id] = col.spec), acc),
-      {} as Record<string, PColumnSpec>,
-    );
-
     using pFrameGuard = new PoolEntryGuard(this.createPFrame(columns));
-    const ValueTypes = new Set(Object.values(ValueType));
-    const specColumnsMap = Object.fromEntries(
-      Object.entries(columnsMap)
-        .filter(([, spec]) => ValueTypes.has(spec.valueType))
-        .map(([id, spec]) => [id, resolveAnnotationParents(spec)]),
-    );
-    const specFrame = createSpecFrame(specColumnsMap);
+    const pFrameSpec = this.pFrames.getByKey(pFrameGuard.key).pFrameSpec;
     const sortedQuery = sortSpecQuery(mapSpecQueryColumns(def.query, (c) => c.id));
-    const { tableSpec, dataQuery } = specFrame.evaluateQuery(sortedQuery);
+    const { tableSpec, dataQuery } = pFrameSpec.evaluateQuery(sortedQuery);
 
-    const pTableEntry = this.pTableDefs.acquire({
-      type: "v2",
-      pFrameHandle: pFrameGuard.key,
-      def: {
+    using pTableGuard = new PoolEntryGuard(
+      this.pTableDefs.acquire({
+        pFrameHandle: pFrameGuard.key,
         tableSpec,
         dataQuery,
-      },
-    });
+      }),
+    );
     if (logPFrames()) {
       this.logger(
         "info",
-        `Create PTable call (pFrameHandle = ${pFrameGuard.key}; pTableHandle = ${pTableEntry.key})`,
+        `Create PTable call (pFrameHandle = ${pFrameGuard.key}; pTableHandle = ${pTableGuard.key})`,
       );
     }
 
     const pFrameEntry = pFrameGuard.keep();
+    const pTableEntry = pTableGuard.keep();
     const unref = () => {
       pTableEntry.unref();
       pFrameEntry.unref();
@@ -293,7 +286,7 @@ export class AbstractPFrameDriver<
     return await this.tableConcurrencyLimiter.run(async () => {
       const shape = await pTable.getShape({ signal: combinedSignal });
       const clippedRange = clipRange(options.range, shape);
-      const specs = pTable.getSpec();
+      const specs = def.tableSpec;
       const separator = options.format === "tsv" ? "\t" : ",";
 
       const iterable = streamPTableRows({
@@ -371,10 +364,8 @@ export class AbstractPFrameDriver<
           : [],
     };
 
-    const { pFramePromise } = this.pFrames.getByKey(handle);
-    const pFrame = await pFramePromise;
-
-    const response = await pFrame.findColumns(iRequest);
+    const { pFrameSpec } = this.pFrames.getByKey(handle);
+    const response = pFrameSpec.findColumns(iRequest);
     return {
       hits: response.hits
         .filter(
@@ -395,15 +386,14 @@ export class AbstractPFrameDriver<
     handle: PFrameHandle,
     columnId: PObjectId,
   ): Promise<PColumnSpec | null> {
-    const { pFramePromise } = this.pFrames.getByKey(handle);
-    const pFrame = await pFramePromise;
-    return await pFrame.getColumnSpec(columnId);
+    const { pFrameSpec } = this.pFrames.getByKey(handle);
+    const info = pFrameSpec.listColumns().find((c) => c.columnId === columnId);
+    return info?.spec ?? null;
   }
 
   public async listColumns(handle: PFrameHandle): Promise<PColumnIdAndSpec[]> {
-    const { pFramePromise } = this.pFrames.getByKey(handle);
-    const pFrame = await pFramePromise;
-    return await pFrame.listColumns();
+    const { pFrameSpec } = this.pFrames.getByKey(handle);
+    return pFrameSpec.listColumns().map(({ columnId, spec }) => ({ columnId, spec }));
   }
 
   public async calculateTableData(
@@ -419,13 +409,15 @@ export class AbstractPFrameDriver<
       );
     }
 
-    using tableGuard = new PoolEntryGuard(
-      this.pTables.acquire({
-        type: "v1",
-        pFrameHandle: handle,
-        def: sortPTableDef(migrateTableFilter(request, this.logger)),
-      }),
-    );
+    const { pFrameSpec } = this.pFrames.getByKey(handle);
+    const sortedDef = sortPTableDef(migrateTableFilter(request, this.logger));
+    const { def, entry } = this.pTables.acquireFromLegacy({
+      pFrameHandle: handle,
+      def: sortedDef,
+      pFrameSpec,
+    });
+    using tableGuard = new PoolEntryGuard(entry);
+    const tableSpec = def.tableSpec;
     const { pTablePromise, disposeSignal } = tableGuard.resource;
     const pTable = await pTablePromise;
 
@@ -434,8 +426,7 @@ export class AbstractPFrameDriver<
       // TODO: throw error when more then 150k rows is requested
       // after pf-plots migration to stream API
 
-      const spec = pTable.getSpec();
-      const data = await pTable.getData([...spec.keys()], {
+      const data = await pTable.getData([...tableSpec.keys()], {
         range,
         signal: combinedSignal,
       });
@@ -445,7 +436,7 @@ export class AbstractPFrameDriver<
       });
       this.pTableCachePerFrame.cache(tableGuard.keep(), overallSize);
 
-      return spec.map((spec, i) => ({
+      return tableSpec.map((spec, i) => ({
         spec: spec,
         data: data[i],
       }));
@@ -464,12 +455,12 @@ export class AbstractPFrameDriver<
       );
     }
 
-    const { pFramePromise, disposeSignal } = this.pFrames.getByKey(handle);
-    const pFrame = await pFramePromise;
+    const { pFrameDataPromise, disposeSignal } = this.pFrames.getByKey(handle);
+    const pFrameData = await pFrameDataPromise;
 
     const combinedSignal = AbortSignal.any([signal, disposeSignal].filter((s) => !!s));
     return await this.frameConcurrencyLimiter.run(async () => {
-      return await pFrame.getUniqueValues(
+      return await pFrameData.getUniqueValues(
         {
           ...request,
           filters: migrateFilters(request.filters, this.logger),
@@ -487,12 +478,7 @@ export class AbstractPFrameDriver<
 
   public async getSpec(handle: PTableHandle): Promise<PTableColumnSpec[]> {
     const { def } = this.pTableDefs.getByKey(handle);
-    using table = this.pTables.acquire(def);
-
-    const { pTablePromise } = table.resource;
-    const pTable = await pTablePromise;
-
-    return pTable.getSpec();
+    return def.tableSpec;
   }
 
   public async getShape(handle: PTableHandle, signal?: AbortSignal): Promise<PTableShape> {

package/src/pframe_pool.ts

@@ -5,13 +5,17 @@ import {
   ensureError,
   mapPObjectData,
   PFrameDriverError,
+  ValueType,
+  resolveAnnotationParents,
   type JsonSerializable,
   type PColumn,
+  type PColumnSpec,
   type PFrameHandle,
 } from "@milaboratories/pl-model-common";
 import { hashJson, PFrameInternal } from "@milaboratories/pl-model-middle-layer";
 import { RefCountPoolBase, type PoolEntry } from "@milaboratories/helpers";
 import { PFrameFactory } from "@milaboratories/pframes-rs-node";
+import { createPFrame as createPFrameSpec } from "@milaboratories/pframes-rs-wasm";
 import { mapValues } from "es-toolkit";
 import { logPFrames } from "./logging";
 
@@ -26,7 +30,13 @@ export interface RemoteBlobProvider<TreeEntry extends JsonSerializable> {
 }
 
 export class PFrameHolder<TreeEntry extends JsonSerializable> implements Disposable {
-  public readonly pFramePromise: Promise<PFrameInternal.PFrameV13>;
+  public readonly pFrameDataPromise: Promise<PFrameInternal.PFrameV13>;
+  /**
+   * WASM-spec frame built from this PFrame's columns. Source of truth
+   * for spec-side operations: column discovery, selector resolution,
+   * legacy-query lowering.
+   */
+  public readonly pFrameSpec: PFrameInternal.PFrameWasmV3;
   private readonly abortController = new AbortController();
 
   private readonly localBlobs: PoolEntry<PFrameInternal.PFrameBlobId>[] = [];
@@ -40,86 +50,114 @@ export class PFrameHolder<TreeEntry extends JsonSerializable> implements Disposa
     private readonly spillPath: string,
     columns: PColumn<PFrameInternal.DataInfo<TreeEntry>>[],
   ) {
-    const makeLocalBlobId = (blob: TreeEntry): PFrameInternal.PFrameBlobId => {
-      const localBlob = this.localBlobProvider.acquire(blob);
-      this.localBlobs.push(localBlob);
-      return localBlob.key;
-    };
+    const ValueTypes = new Set(Object.values(ValueType));
+    const specColumnsMap: Record<string, PColumnSpec> = {};
+    for (const c of columns) {
+      if (ValueTypes.has(c.spec.valueType)) {
+        specColumnsMap[c.id] = resolveAnnotationParents(c.spec);
+      }
+    }
+    this.pFrameSpec = createPFrameSpec(specColumnsMap);
 
-    const makeRemoteBlobId = (blob: TreeEntry): PFrameInternal.PFrameBlobId => {
-      const remoteBlob = this.remoteBlobProvider.acquire(blob);
-      this.remoteBlobs.push(remoteBlob);
-      return `${remoteBlob.key}${PFrameInternal.ParquetExtension}` as PFrameInternal.PFrameBlobId;
-    };
+    try {
+      const makeLocalBlobId = (blob: TreeEntry): PFrameInternal.PFrameBlobId => {
+        const localBlob = this.localBlobProvider.acquire(blob);
+        this.localBlobs.push(localBlob);
+        return localBlob.key;
+      };
 
-    const mapColumnData = (
-      data: PFrameInternal.DataInfo<TreeEntry>,
-    ): PFrameInternal.DataInfo<PFrameInternal.PFrameBlobId> => {
-      switch (data.type) {
-        case "Json":
-          return { ...data };
-        case "JsonPartitioned":
-          return {
-            ...data,
-            parts: mapValues(data.parts, makeLocalBlobId),
-          };
-        case "BinaryPartitioned":
-          return {
-            ...data,
-            parts: mapValues(data.parts, (v) => ({
-              index: makeLocalBlobId(v.index),
-              values: makeLocalBlobId(v.values),
-            })),
-          };
-        case "ParquetPartitioned":
-          return {
-            ...data,
-            parts: mapValues(data.parts, (v) => ({
-              ...v,
-              data: makeRemoteBlobId(v.data),
-            })),
-          };
-        default:
-          assertNever(data);
-      }
-    };
+      const makeRemoteBlobId = (blob: TreeEntry): PFrameInternal.PFrameBlobId => {
+        const remoteBlob = this.remoteBlobProvider.acquire(blob);
+        this.remoteBlobs.push(remoteBlob);
+        return `${remoteBlob.key}${PFrameInternal.ParquetExtension}` as PFrameInternal.PFrameBlobId;
+      };
 
-    const jsonifiedColumns = columns.map((column) => ({
-      ...column,
-      data: mapColumnData(column.data),
-    }));
+      const mapColumnData = (
+        data: PFrameInternal.DataInfo<TreeEntry>,
+      ): PFrameInternal.DataInfo<PFrameInternal.PFrameBlobId> => {
+        switch (data.type) {
+          case "Json":
+            return { ...data };
+          case "JsonPartitioned":
+            return {
+              ...data,
+              parts: mapValues(data.parts, makeLocalBlobId),
+            };
+          case "BinaryPartitioned":
+            return {
+              ...data,
+              parts: mapValues(data.parts, (v) => ({
+                index: makeLocalBlobId(v.index),
+                values: makeLocalBlobId(v.values),
+              })),
+            };
+          case "ParquetPartitioned":
+            return {
+              ...data,
+              parts: mapValues(data.parts, (v) => ({
+                ...v,
+                data: makeRemoteBlobId(v.data),
+              })),
+            };
+          default:
+            assertNever(data);
+        }
+      };
 
-    try {
-      const pFrame = PFrameFactory.createPFrame({ frameId, spillPath: this.spillPath, logger });
-      pFrame.setDataSource({
-        ...this.localBlobProvider.makeDataSource(this.disposeSignal),
-        parquetServer: this.remoteBlobProvider.httpServerInfo(),
-      });
+      const jsonifiedColumns = columns.map((column) => ({
+        ...column,
+        data: mapColumnData(column.data),
+      }));
 
+      const pFrameData = PFrameFactory.createPFrame({ frameId, spillPath: this.spillPath, logger });
       const promises: Promise<void>[] = [];
-      for (const column of jsonifiedColumns) {
-        pFrame.addColumnSpec(column.id, column.spec);
-        promises.push(pFrame.setColumnData(column.id, column.data, { signal: this.disposeSignal }));
-      }
+      try {
+        pFrameData.setDataSource({
+          ...this.localBlobProvider.makeDataSource(this.disposeSignal),
+          parquetServer: this.remoteBlobProvider.httpServerInfo(),
+        });
 
-      this.pFramePromise = Promise.all(promises)
-        .then(() => pFrame)
-        .catch((err) => {
-          this.dispose();
-          pFrame.dispose();
-          const error = new PFrameDriverError("PFrame creation failed asynchronously");
-          error.cause = new Error(
-            `PFrame cannot be created from columns: ${JSON.stringify(jsonifiedColumns)}`,
-            { cause: ensureError(err) },
+        for (const column of jsonifiedColumns) {
+          pFrameData.addColumnSpec(column.id, column.spec);
+          promises.push(
+            pFrameData.setColumnData(column.id, column.data, { signal: this.disposeSignal }),
           );
-          throw error;
-        });
+        }
+
+        this.pFrameDataPromise = Promise.all(promises)
+          .then(() => pFrameData)
+          .catch((err) => {
+            this.dispose();
+            pFrameData.dispose();
+            const error = new PFrameDriverError("PFrame creation failed asynchronously");
+            error.cause = new Error(
+              `PFrame cannot be created from columns: ${JSON.stringify(jsonifiedColumns)}`,
+              { cause: ensureError(err) },
+            );
+            throw error;
+          });
+      } catch (err: unknown) {
+        // setDataSource / addColumnSpec / setColumnData threw synchronously
+        // before pFrameDataPromise was assigned — dispose the addon's
+        // pFrameData explicitly (the dispose() path can't reach it yet).
+        // `allSettled` attaches handlers to any in-flight setColumnData
+        // promises so they don't trigger unhandled-rejection warnings
+        // when the dispose below rejects them.
+        void Promise.allSettled(promises);
+        pFrameData.dispose();
+        throw err;
+      }
     } catch (err: unknown) {
+      // Release everything allocated so far in this constructor:
+      // acquired blobs, the WASM spec frame, and the abort signal.
+      this.abortController.abort();
+      this.localBlobs.forEach((entry) => entry.unref());
+      this.remoteBlobs.forEach((entry) => entry.unref());
+      this.pFrameSpec[Symbol.dispose]();
       const error = new PFrameDriverError("PFrame creation failed synchronously");
-      error.cause = new Error(
-        `PFrame cannot be created from columns: ${JSON.stringify(jsonifiedColumns)}`,
-        { cause: ensureError(err) },
-      );
+      error.cause = new Error(`PFrame cannot be created from columns: ${JSON.stringify(columns)}`, {
+        cause: ensureError(err),
+      });
       throw error;
     }
   }
@@ -132,15 +170,16 @@ export class PFrameHolder<TreeEntry extends JsonSerializable> implements Disposa
     this.abortController.abort();
     this.localBlobs.forEach((entry) => entry.unref());
     this.remoteBlobs.forEach((entry) => entry.unref());
+    this.pFrameSpec[Symbol.dispose]();
+    void this.pFrameDataPromise
+      .then((pFrameData) => pFrameData.dispose())
+      .catch(() => {
+        /* mute error */
+      });
   }
 
   [Symbol.dispose](): void {
     this.dispose();
-    void this.pFramePromise
-      .then((pFrame) => pFrame.dispose())
-      .catch(() => {
-        /* mute error */
-      });
   }
 }
 

package/src/ptable_def_pool.ts

@@ -1,9 +1,15 @@
-import { PFrameDriverError, type PTableHandle } from "@milaboratories/pl-model-common";
+import {
+  PFrameDriverError,
+  type PFrameHandle,
+  type PObjectId,
+  type PTableDef,
+  type PTableHandle,
+} from "@milaboratories/pl-model-common";
 import type { PFrameInternal } from "@milaboratories/pl-model-middle-layer";
-import { RefCountPoolBase } from "@milaboratories/helpers";
+import { RefCountPoolBase, type PoolEntry } from "@milaboratories/helpers";
 import { logPFrames } from "./logging";
 import type { FullPTableDef } from "./ptable_shared";
-import { stableKeyFromFullPTableDef } from "./ptable_shared";
+import { buildFullPTableDefFromLegacy, stableKeyFromFullPTableDef } from "./ptable_shared";
 
 export class PTableDefHolder implements Disposable {
   private readonly abortController = new AbortController();
@@ -52,4 +58,21 @@ export class PTableDefPool extends RefCountPoolBase<FullPTableDef, PTableHandle,
     }
     return resource;
   }
+
+  /**
+   * Acquire a def from a legacy `PTableDef` + column specs map.
+   * Lowers the input via WASM-spec and stores the resulting
+   * `{ tableSpec, dataQuery }` shape. Returns the lowered def
+   * alongside the pool entry — callers that need `tableSpec`
+   * (e.g. `calculateTableData`) read it from `def.tableSpec`
+   * without re-deriving.
+   */
+  public acquireFromLegacy(opts: {
+    pFrameHandle: PFrameHandle;
+    def: PTableDef<PObjectId>;
+    pFrameSpec: PFrameInternal.PFrameWasmV3;
+  }): { def: FullPTableDef; entry: PoolEntry<PTableHandle> } {
+    const def = buildFullPTableDefFromLegacy(opts);
+    return { def, entry: this.acquire(def) };
+  }
 }