@milaboratories/pl-middle-layer 1.42.1 → 1.42.3

package/dist/pool/driver.js

@@ -0,0 +1,598 @@
+import { isPlTreeNodeAccessor } from '@milaboratories/pl-tree';
+import { PFrameDriverError, isPFrameDriverError, ensureError, mapPObjectData, mapDataInfo, isDataInfo, extractAllColumns, mapPTableDef, isAbortError } from '@platforma-sdk/model';
+import { LRUCache } from 'lru-cache';
+import { RefCountResourcePool } from './ref_count_pool.js';
+import { parseDataInfoResource, makeDataInfoFromPColumnValues, traverseParquetPartitionedResource, allBlobs, mapBlobs } from './data.js';
+import { createHash } from 'node:crypto';
+import { emptyDir, ConcurrencyLimitingExecutor, assertNever } from '@milaboratories/ts-helpers';
+import canonicalize from 'canonicalize';
+import { PFrame } from '@milaboratories/pframes-rs-node';
+import * as fs from 'node:fs/promises';
+import * as path from 'node:path';
+import { getDebugFlags } from '../debug/index.js';
+
+function blobKey(res) {
+    return String(res.id);
+}
+const valueTypes = ['Int', 'Long', 'Float', 'Double', 'String', 'Bytes'];
+function migrateFilters(filters) {
+    const filtersV1 = [];
+    const filtersV2 = [];
+    for (const filter of filters) {
+        if (filter.type === 'bySingleColumn') {
+            filtersV1.push(filter);
+            filtersV2.push({
+                ...filter,
+                type: 'bySingleColumnV2',
+            });
+        }
+        else {
+            filtersV2.push(filter);
+        }
+    }
+    if (filtersV1.length > 0) {
+        const filtersV1Json = JSON.stringify(filtersV1);
+        console.warn(`type overriten from 'bySingleColumn' to 'bySingleColumnV2' for filters: ${filtersV1Json}`);
+    }
+    return filtersV2;
+}
+function migratePTableFilters(def) {
+    if (!('partitionFilters' in def)) {
+        // For old blocks assume all axes filters to be partition filters
+        return {
+            ...def,
+            partitionFilters: migrateFilters(def.filters.filter((f) => f.column.type === 'axis')),
+            filters: migrateFilters(def.filters.filter((f) => f.column.type === 'column')),
+        };
+    }
+    return {
+        ...def,
+        partitionFilters: migrateFilters(def.partitionFilters),
+        filters: migrateFilters(def.filters),
+    };
+}
+const bigintReplacer = (_, v) => (typeof v === 'bigint' ? v.toString() : v);
+class PTableCache {
+    logger;
+    ops;
+    perFrame = new Map();
+    global;
+    disposeListeners = new Map();
+    constructor(logger, ops) {
+        this.logger = logger;
+        this.ops = ops;
+        this.global = new LRUCache({
+            maxSize: this.ops.pFramesCacheMaxSize,
+            dispose: (resource, key, reason) => {
+                if (reason === 'evict') {
+                    this.perFrame.get(resource.resource.pFrame)?.delete(key);
+                }
+                if (this.perFrame.get(resource.resource.pFrame)?.size === 0) {
+                    this.perFrame.delete(resource.resource.pFrame);
+                }
+                const disposeListener = this.disposeListeners.get(key);
+                this.disposeListeners.delete(key);
+                resource.resource.disposeSignal.removeEventListener('abort', disposeListener);
+                resource.unref();
+                if (getDebugFlags().logPFrameRequests) {
+                    this.logger.info(`calculateTableData cache - removed PTable ${key}`);
+                }
+            },
+        });
+    }
+    cache(resource, size) {
+        const key = resource.key;
+        if (getDebugFlags().logPFrameRequests) {
+            this.logger.info(`calculateTableData cache - added PTable ${key} with size ${size}`);
+        }
+        this.global.set(key, resource, { size });
+        let perFrame = this.perFrame.get(resource.resource.pFrame);
+        if (!perFrame) {
+            perFrame = new LRUCache({
+                max: this.ops.pFrameCacheMaxCount,
+                dispose: (_resource, key, reason) => {
+                    if (reason === 'evict') {
+                        this.global.delete(key);
+                    }
+                },
+            });
+            this.perFrame.set(resource.resource.pFrame, perFrame);
+        }
+        perFrame.set(key, resource);
+        const disposeListener = () => {
+            this.perFrame.get(resource.resource.pFrame)?.delete(key);
+            this.global.delete(key);
+        };
+        this.disposeListeners.set(key, disposeListener);
+        resource.resource.disposeSignal.addEventListener('abort', disposeListener);
+    }
+}
+class PFrameHolder {
+    blobDriver;
+    logger;
+    spillPath;
+    pFramePromise;
+    abortController = new AbortController();
+    blobIdToResource = new Map();
+    blobHandleComputables = new Map();
+    constructor(blobDriver, logger, spillPath, columns) {
+        this.blobDriver = blobDriver;
+        this.logger = logger;
+        this.spillPath = spillPath;
+        const logFunc = (level, message) => this.logger[level](message);
+        for (const column of columns) {
+            for (const blob of allBlobs(column.data)) {
+                this.blobIdToResource.set(blobKey(blob), blob);
+            }
+        }
+        const distinctСolumns = [
+            ...new Map(columns.map((column) => ({
+                ...column,
+                data: mapBlobs(column.data, blobKey),
+            })).map((item) => [canonicalize(item), item])).values(),
+        ];
+        try {
+            const pFrame = new PFrame(this.spillPath, logFunc);
+            pFrame.setDataSource(this);
+            const promises = [];
+            for (const column of distinctСolumns) {
+                pFrame.addColumnSpec(column.id, column.spec);
+                promises.push(pFrame.setColumnData(column.id, column.data, { signal: this.disposeSignal }));
+            }
+            this.pFramePromise = Promise.all(promises)
+                .then(() => pFrame)
+                .catch((err) => {
+                this.dispose();
+                pFrame.dispose();
+                throw new PFrameDriverError(`PFrame creation failed asynchronously, columns: ${JSON.stringify(distinctСolumns)}, error: ${ensureError(err)}`);
+            });
+        }
+        catch (err) {
+            throw new PFrameDriverError(`PFrame creation failed synchronously, columns: ${JSON.stringify(distinctСolumns)}, error: ${ensureError(err)}`);
+        }
+    }
+    getOrCreateComputableForBlob(blobId) {
+        let computable = this.blobHandleComputables.get(blobId);
+        if (computable !== undefined)
+            return computable;
+        const blobResource = this.blobIdToResource.get(blobId);
+        if (blobResource === undefined)
+            throw new PFrameDriverError(`Blob with id ${blobId} not found.`);
+        // precalculation of value tree will trigger the download proecess right away
+        computable = this.blobDriver.getDownloadedBlob(blobResource).withPreCalculatedValueTree();
+        this.blobHandleComputables.set(blobId, computable);
+        return computable;
+    }
+    preloadBlob = async (blobIds) => {
+        const computables = blobIds.map((blobId) => this.getOrCreateComputableForBlob(blobId));
+        for (const computable of computables) {
+            try {
+                await computable.awaitStableFullValue(this.disposeSignal);
+            }
+            catch (err) {
+                if (isAbortError(err)) {
+                    break; // silence abort errors
+                }
+                throw err;
+            }
+        }
+    };
+    resolveBlobContent = async (blobId) => {
+        const computable = this.getOrCreateComputableForBlob(blobId);
+        const path = this.blobDriver.getLocalPath((await computable.awaitStableValue(this.disposeSignal)).handle);
+        return await fs.readFile(path);
+    };
+    get disposeSignal() {
+        return this.abortController.signal;
+    }
+    dispose() {
+        this.abortController.abort();
+        for (const computable of this.blobHandleComputables.values())
+            computable.resetState();
+    }
+    async [Symbol.asyncDispose]() {
+        this.dispose();
+        await this.pFramePromise
+            .then((pFrame) => pFrame.dispose())
+            .catch(() => { });
+    }
+}
+class PTableHolder {
+    pFrame;
+    pTablePromise;
+    predecessor;
+    abortController = new AbortController();
+    combinedDisposeSignal;
+    constructor(pFrame, pFrameDisposeSignal, pTablePromise, predecessor) {
+        this.pFrame = pFrame;
+        this.pTablePromise = pTablePromise;
+        this.predecessor = predecessor;
+        this.combinedDisposeSignal = AbortSignal.any([pFrameDisposeSignal, this.abortController.signal]);
+    }
+    get disposeSignal() {
+        return this.combinedDisposeSignal;
+    }
+    async [Symbol.asyncDispose]() {
+        this.abortController.abort();
+        await this.pTablePromise
+            .then((pTable) => pTable.dispose())
+            .catch(() => { });
+        this.predecessor?.unref();
+    }
+}
+class PFrameDriver {
+    blobDriver;
+    logger;
+    spillPath;
+    pFrames;
+    pTables;
+    pTableCache;
+    frameConcurrencyLimiter;
+    tableConcurrencyLimiter;
+    async pprofDump() {
+        return await PFrame.pprofDump();
+    }
+    static async init(blobDriver, logger, spillPath, ops) {
+        const resolvedSpillPath = path.resolve(spillPath);
+        await emptyDir(resolvedSpillPath);
+        return new PFrameDriver(blobDriver, logger, resolvedSpillPath, ops);
+    }
+    constructor(blobDriver, logger, spillPath, ops) {
+        this.blobDriver = blobDriver;
+        this.logger = logger;
+        this.spillPath = spillPath;
+        const concurrencyLimiter = new ConcurrencyLimitingExecutor(ops.pFrameConcurrency);
+        this.frameConcurrencyLimiter = concurrencyLimiter;
+        this.tableConcurrencyLimiter = new ConcurrencyLimitingExecutor(ops.pTableConcurrency);
+        this.pTableCache = new PTableCache(this.logger, ops);
+        this.pFrames = new (class extends RefCountResourcePool {
+            blobDriver;
+            logger;
+            spillPath;
+            constructor(blobDriver, logger, spillPath) {
+                super();
+                this.blobDriver = blobDriver;
+                this.logger = logger;
+                this.spillPath = spillPath;
+            }
+            acquire(params) {
+                return super.acquire(params);
+            }
+            getByKey(key) {
+                const resource = super.tryGetByKey(key);
+                if (!resource)
+                    throw new PFrameDriverError(`PFrame not found, handle = ${key}`);
+                return resource;
+            }
+            createNewResource(params) {
+                if (getDebugFlags().logPFrameRequests)
+                    logger.info(`PFrame creation (pFrameHandle = ${this.calculateParamsKey(params)}): ${JSON.stringify(params, bigintReplacer)}`);
+                return new PFrameHolder(this.blobDriver, this.logger, this.spillPath, params);
+            }
+            calculateParamsKey(params) {
+                try {
+                    return stableKeyFromPFrameData(params);
+                }
+                catch (err) {
+                    if (isPFrameDriverError(err))
+                        throw err;
+                    throw new PFrameDriverError(`PFrame handle calculation failed, request: ${JSON.stringify(params, bigintReplacer)}, error: ${ensureError(err)}`);
+                }
+            }
+        })(this.blobDriver, this.logger, this.spillPath);
+        this.pTables = new (class extends RefCountResourcePool {
+            pFrames;
+            constructor(pFrames) {
+                super();
+                this.pFrames = pFrames;
+            }
+            getByKey(key) {
+                const resource = super.tryGetByKey(key);
+                if (!resource)
+                    throw new PFrameDriverError(`PTable not found, handle = ${key}`);
+                return resource;
+            }
+            createNewResource(params) {
+                if (getDebugFlags().logPFrameRequests) {
+                    logger.info(`PTable creation (pTableHandle = ${this.calculateParamsKey(params)}): ${JSON.stringify(params, bigintReplacer)}`);
+                }
+                const handle = params.pFrameHandle;
+                const { pFramePromise, disposeSignal } = this.pFrames.getByKey(handle);
+                // 3. Sort
+                if (params.def.sorting.length > 0) {
+                    const predecessor = this.acquire({
+                        ...params,
+                        def: {
+                            ...params.def,
+                            sorting: [],
+                        },
+                    });
+                    const { resource: { pTablePromise } } = predecessor;
+                    const sortedTable = pTablePromise.then((pTable) => pTable.sort(params.def.sorting));
+                    return new PTableHolder(handle, disposeSignal, sortedTable, predecessor);
+                }
+                // 2. Filter
+                if (params.def.filters.length > 0) {
+                    const predecessor = this.acquire({
+                        ...params,
+                        def: {
+                            ...params.def,
+                            filters: [],
+                        },
+                    });
+                    const { resource: { pTablePromise } } = predecessor;
+                    const filteredTable = pTablePromise.then((pTable) => pTable.filter(params.def.filters));
+                    return new PTableHolder(handle, disposeSignal, filteredTable, predecessor);
+                }
+                // 1. Join
+                const table = pFramePromise.then((pFrame) => pFrame.createTable({
+                    src: joinEntryToInternal(params.def.src),
+                    filters: params.def.partitionFilters,
+                }));
+                return new PTableHolder(handle, disposeSignal, table);
+            }
+            calculateParamsKey(params) {
+                try {
+                    return stableKeyFromFullPTableDef(params);
+                }
+                catch (err) {
+                    throw new PFrameDriverError(`PTable handle calculation failed, request: ${JSON.stringify(params)}, error: ${ensureError(err)}`);
+                }
+            }
+        })(this.pFrames);
+    }
+    //
+    // Internal / Config API Methods
+    //
+    createPFrame(def, ctx) {
+        const internalData = def
+            .filter((c) => valueTypes.find((t) => t === c.spec.valueType))
+            .map((c) => mapPObjectData(c, (d) => isPlTreeNodeAccessor(d)
+            ? parseDataInfoResource(d)
+            : isDataInfo(d)
+                ? d.type === 'ParquetPartitioned'
+                    ? mapDataInfo(d, (a) => traverseParquetPartitionedResource(a))
+                    : mapDataInfo(d, (a) => a.resourceInfo)
+                : makeDataInfoFromPColumnValues(c.spec, d)));
+        const res = this.pFrames.acquire(internalData);
+        ctx.addOnDestroy(res.unref);
+        return res.key;
+    }
+    createPTable(rawDef, ctx) {
+        const def = migratePTableFilters(rawDef);
+        const pFrameHandle = this.createPFrame(extractAllColumns(def.src), ctx);
+        const defIds = mapPTableDef(def, (c) => c.id);
+        const res = this.pTables.acquire({ def: defIds, pFrameHandle });
+        if (getDebugFlags().logPFrameRequests)
+            this.logger.info(`Create PTable call (pFrameHandle = ${pFrameHandle}; pTableHandle = ${JSON.stringify(res)}): ${JSON.stringify(mapPTableDef(def, (c) => c.spec), bigintReplacer)}`);
+        ctx.addOnDestroy(res.unref); // in addition to pframe unref added in createPFrame above
+        return res.key;
+    }
+    //
+    // PFrame istance methods
+    //
+    async findColumns(handle, request) {
+        const iRequest = {
+            ...request,
+            compatibleWith: request.compatibleWith.length !== 0
+                ? [{
+                        axesSpec: [
+                            ...new Map(request.compatibleWith.map((item) => [canonicalize(item), item])).values(),
+                        ],
+                        qualifications: [],
+                    }]
+                : [],
+        };
+        const { pFramePromise } = this.pFrames.getByKey(handle);
+        const pFrame = await pFramePromise;
+        const responce = await pFrame.findColumns(iRequest);
+        return {
+            hits: responce.hits
+                .filter((h) => // only exactly matching columns
+             h.mappingVariants.length === 0
+                || h.mappingVariants.some((v) => v.qualifications.forHit.length === 0
+                    && v.qualifications.forQueries.every((q) => q.length === 0)))
+                .map((h) => h.hit),
+        };
+    }
+    async getColumnSpec(handle, columnId) {
+        const { pFramePromise } = this.pFrames.getByKey(handle);
+        const pFrame = await pFramePromise;
+        return await pFrame.getColumnSpec(columnId);
+    }
+    async listColumns(handle) {
+        const { pFramePromise } = this.pFrames.getByKey(handle);
+        const pFrame = await pFramePromise;
+        return await pFrame.listColumns();
+    }
+    async calculateTableData(handle, request, range, signal) {
+        if (getDebugFlags().logPFrameRequests) {
+            this.logger.info(`Call calculateTableData, handle = ${handle}, request = ${JSON.stringify(request, bigintReplacer)}`);
+        }
+        const table = this.pTables.acquire({
+            pFrameHandle: handle,
+            def: migratePTableFilters(request),
+        });
+        const { resource: { pTablePromise, disposeSignal } } = table;
+        const pTable = await pTablePromise;
+        const combinedSignal = AbortSignal.any([signal, disposeSignal].filter((s) => !!s));
+        return await this.frameConcurrencyLimiter.run(async () => {
+            try {
+                const spec = pTable.getSpec();
+                const data = await pTable.getData([...spec.keys()], {
+                    range,
+                    signal: combinedSignal,
+                });
+                const size = await pTable.getFootprint({
+                    withPredecessors: true,
+                    signal: combinedSignal,
+                });
+                this.pTableCache.cache(table, size);
+                return spec.map((spec, i) => ({
+                    spec: spec,
+                    data: data[i],
+                }));
+            }
+            catch (err) {
+                table.unref();
+                throw err;
+            }
+        });
+    }
+    async getUniqueValues(handle, request, signal) {
+        if (getDebugFlags().logPFrameRequests) {
+            this.logger.info(`Call getUniqueValues, handle = ${handle}, request = ${JSON.stringify(request, bigintReplacer)}`);
+        }
+        const { pFramePromise, disposeSignal } = this.pFrames.getByKey(handle);
+        const pFrame = await pFramePromise;
+        const combinedSignal = AbortSignal.any([signal, disposeSignal].filter((s) => !!s));
+        return await this.frameConcurrencyLimiter.run(async () => {
+            return await pFrame.getUniqueValues({
+                ...request,
+                filters: migrateFilters(request.filters),
+            }, {
+                signal: combinedSignal,
+            });
+        });
+    }
+    //
+    // PTable istance methods
+    //
+    async getSpec(handle) {
+        const { pTablePromise } = this.pTables.getByKey(handle);
+        const pTable = await pTablePromise;
+        return pTable.getSpec();
+    }
+    async getShape(handle, signal) {
+        const { pTablePromise, disposeSignal } = this.pTables.getByKey(handle);
+        const pTable = await pTablePromise;
+        const combinedSignal = AbortSignal.any([signal, disposeSignal].filter((s) => !!s));
+        return await this.tableConcurrencyLimiter.run(async () => {
+            return await pTable.getShape({
+                signal: combinedSignal,
+            });
+        });
+    }
+    async getData(handle, columnIndices, range, signal) {
+        const { pTablePromise, disposeSignal } = this.pTables.getByKey(handle);
+        const pTable = await pTablePromise;
+        const combinedSignal = AbortSignal.any([signal, disposeSignal].filter((s) => !!s));
+        return await this.tableConcurrencyLimiter.run(async () => {
+            return await pTable.getData(columnIndices, {
+                range,
+                signal: combinedSignal,
+            });
+        });
+    }
+}
+function joinEntryToInternal(entry) {
+    switch (entry.type) {
+        case 'column':
+            return {
+                type: 'column',
+                columnId: entry.column,
+            };
+        case 'slicedColumn':
+            return {
+                type: 'slicedColumn',
+                columnId: entry.column,
+                newId: entry.newId,
+                axisFilters: entry.axisFilters,
+            };
+        case 'inlineColumn':
+            return {
+                type: 'inlineColumn',
+                newId: entry.column.id,
+                spec: entry.column.spec,
+                dataInfo: {
+                    type: 'Json',
+                    keyLength: entry.column.spec.axesSpec.length,
+                    data: entry.column.data.reduce((acc, row) => {
+                        acc[JSON.stringify(row.key)] = row.val;
+                        return acc;
+                    }, {}),
+                },
+            };
+        case 'inner':
+        case 'full':
+            return {
+                type: entry.type,
+                entries: entry.entries.map((col) => joinEntryToInternal(col)),
+            };
+        case 'outer':
+            return {
+                type: 'outer',
+                primary: joinEntryToInternal(entry.primary),
+                secondary: entry.secondary.map((col) => joinEntryToInternal(col)),
+            };
+        default:
+            assertNever(entry);
+    }
+}
+function stableKeyFromFullPTableDef(data) {
+    const hash = createHash('sha256');
+    hash.update(canonicalize(data));
+    return hash.digest().toString('hex');
+}
+function stableKeyFromPFrameData(data) {
+    const orderedData = [...data].map((column) => mapPObjectData(column, (r) => {
+        let result;
+        const type = r.type;
+        switch (type) {
+            case 'Json':
+                result = {
+                    type: r.type,
+                    keyLength: r.keyLength,
+                    payload: Object.entries(r.data).map(([part, value]) => ({
+                        key: part,
+                        value,
+                    })),
+                };
+                break;
+            case 'JsonPartitioned':
+                result = {
+                    type: r.type,
+                    keyLength: r.partitionKeyLength,
+                    payload: Object.entries(r.parts).map(([part, info]) => ({
+                        key: part,
+                        value: blobKey(info),
+                    })),
+                };
+                break;
+            case 'BinaryPartitioned':
+                result = {
+                    type: r.type,
+                    keyLength: r.partitionKeyLength,
+                    payload: Object.entries(r.parts).map(([part, info]) => ({
+                        key: part,
+                        value: [blobKey(info.index), blobKey(info.values)],
+                    })),
+                };
+                break;
+            case 'ParquetPartitioned':
+                result = {
+                    type: r.type,
+                    keyLength: r.partitionKeyLength,
+                    payload: Object.entries(r.parts).map(([part, info]) => ({
+                        key: part,
+                        value: info.dataDigest || [
+                            blobKey(info.data),
+                            JSON.stringify({ axes: info.axes, column: info.column }),
+                        ],
+                    })),
+                };
+                break;
+            default:
+                throw new PFrameDriverError(`unsupported resource type: ${JSON.stringify(type)}`);
+        }
+        result.payload.sort((lhs, rhs) => lhs.key.localeCompare(rhs.key));
+        return result;
+    }));
+    orderedData.sort((lhs, rhs) => lhs.id.localeCompare(rhs.id));
+    const hash = createHash('sha256');
+    hash.update(canonicalize(orderedData));
+    return hash.digest().toString('hex');
+}
+
+export { PFrameDriver };
+//# sourceMappingURL=driver.js.map