@convex-dev/geospatial/dist/component/streams/union.js

A geospatial index for Convex

import { Heap } from "heap-js";
export class Union {
    streams;
    heap;
    constructor(streams) {
        this.streams = streams;
    }
    async initializeHeap() {
        if (this.heap) {
            return this.heap;
        }
        const promises = this.streams.map((stream) => stream.current());
        const results = await Promise.all(promises);
        const entries = [];
        for (let i = 0; i < this.streams.length; i++) {
            const result = results[i];
            if (result !== null) {
                entries.push({ tuple: result, stream: this.streams[i] });
            }
        }
        const heap = new Heap((a, b) => a.tuple < b.tuple ? -1 : a.tuple > b.tuple ? 1 : 0);
        heap.init(entries);
        this.heap = heap;
        return heap;
    }
    async current() {
        const heap = await this.initializeHeap();
        const smallest = heap.peek();
        if (smallest === undefined) {
            return null;
        }
        return smallest.tuple;
    }
    async advance() {
        const heap = await this.initializeHeap();
        const smallest = heap.pop();
        if (smallest === undefined) {
            return null;
        }
        const toRefill = [smallest.stream];
        while (true) {
            const next = heap.peek();
            if (next === undefined) {
                break;
            }
            if (smallest.tuple === next.tuple) {
                heap.pop();
                toRefill.push(next.stream);
                continue;
            }
            break;
        }
        for (const stream of toRefill) {
            const result = await stream.advance();
            if (result !== null) {
                heap.push({ tuple: result, stream });
            }
        }
        return smallest.tuple;
    }
    async seek(tuple) {
        const heap = await this.initializeHeap();
        // Remove the heap entries that are less than `pos`.
        const toRefill = [];
        while (true) {
            const next = heap.peek();
            if (next === undefined) {
                break;
            }
            if (next.tuple < tuple) {
                heap.pop();
                toRefill.push(next.stream);
                continue;
            }
            break;
        }
        const seekPromises = toRefill.map((stream) => stream.seek(tuple));
        await Promise.all(seekPromises);
        const currentPromises = toRefill.map(async (stream) => {
            const result = await stream.current();
            return { result, stream };
        });
        const currentResults = await Promise.all(currentPromises);
        for (const { result, stream } of currentResults) {
            if (result !== null) {
                heap.push({ tuple: result, stream });
            }
        }
    }
    async sizeHint() {
        // Assume the underlying streams are disjoint (which is true for our covering).
        const promises = this.streams.map((stream) => stream.sizeHint());
        const results = await Promise.all(promises);
        return results.reduce((a, b) => a + b, 0);
    }
    setPrefetch(prefetch) {
        for (const stream of this.streams) {
            stream.setPrefetch(prefetch);
        }
    }
}
//# sourceMappingURL=union.js.map