@convex-dev/geospatial/dist/commonjs/component/lib/pointQuery.js

A geospatial index for Convex

import { Heap } from "heap-js";
import * as approximateCounter from "./approximateCounter.js";
import { cellCounterKey } from "../streams/cellRange.js";
import { decodeTupleKey } from "./tupleKey.js";
export class ClosestPointQuery {
    s2;
    logger;
    point;
    maxDistance;
    maxResults;
    minLevel;
    maxLevel;
    levelMod;
    // Min-heap of cells to process.
    toProcess;
    // Max-heap of results.
    results;
    maxDistanceChordAngle;
    constructor(s2, logger, point, maxDistance, maxResults, minLevel, maxLevel, levelMod) {
        this.s2 = s2;
        this.logger = logger;
        this.point = point;
        this.maxDistance = maxDistance;
        this.maxResults = maxResults;
        this.minLevel = minLevel;
        this.maxLevel = maxLevel;
        this.levelMod = levelMod;
        this.toProcess = new Heap((a, b) => a.distance - b.distance);
        this.results = new Heap((a, b) => b.distance - a.distance);
        this.maxDistanceChordAngle =
            this.maxDistance && this.s2.metersToChordAngle(this.maxDistance);
        for (const cellID of this.s2.initialCells(this.minLevel)) {
            const distance = this.s2.minDistanceToCell(this.point, cellID);
            const level = this.s2.cellIDLevel(cellID);
            this.addCandidate(cellID, level, distance);
        }
    }
    async execute(ctx) {
        while (true) {
            const candidate = this.popCandidate();
            this.logger.debug(`Processing candidate: ${candidate?.cellID}`);
            if (candidate === null) {
                break;
            }
            const canSubdivide = candidate.level < this.maxLevel;
            const cellIDToken = this.s2.cellIDToken(candidate.cellID);
            const sizeEstimate = await approximateCounter.estimateCount(ctx, cellCounterKey(cellIDToken));
            this.logger.debug(`Size estimate for ${cellIDToken}: ${sizeEstimate}`);
            if (canSubdivide && sizeEstimate >= approximateCounter.SAMPLING_RATE) {
                this.logger.debug(`Subdividing cell ${candidate.cellID}`);
                const nextLevel = Math.min(candidate.level + this.levelMod, this.maxLevel);
                for (const cellID of this.s2.cellIDChildren(candidate.cellID, nextLevel)) {
                    const distance = this.s2.minDistanceToCell(this.point, cellID);
                    this.addCandidate(cellID, nextLevel, distance);
                }
            }
            else {
                // Query the current cell and add its results in.
                const pointEntries = await ctx.db
                    .query("pointsByCell")
                    .withIndex("cell", (q) => q.eq("cell", cellIDToken))
                    .collect();
                this.logger.debug(`Found ${pointEntries.length} points in cell ${cellIDToken}`);
                const pointIds = pointEntries.map((entry) => decodeTupleKey(entry.tupleKey).pointId);
                const points = await Promise.all(pointIds.map((id) => ctx.db.get(id)));
                for (const point of points) {
                    if (!point) {
                        throw new Error("Point not found");
                    }
                    this.addResult(point._id, point.coordinates);
                }
            }
        }
        const entries = this.results
            .toArray()
            .sort((a, b) => a.distance - b.distance);
        const points = await Promise.all(entries.map((r) => ctx.db.get(r.pointID)));
        const results = [];
        for (let i = 0; i < entries.length; i++) {
            const point = points[i];
            if (!point) {
                throw new Error("Point not found");
            }
            results.push({
                key: point.key,
                coordinates: point.coordinates,
                distance: this.s2.chordAngleToMeters(entries[i].distance),
            });
        }
        return results;
    }
    addCandidate(cellID, level, distance) {
        if (this.maxDistanceChordAngle && distance > this.maxDistanceChordAngle) {
            return;
        }
        const threshold = this.distanceThreshold();
        if (threshold !== undefined && distance >= threshold) {
            return;
        }
        this.toProcess.push({ cellID, level, distance });
    }
    popCandidate() {
        const threshold = this.distanceThreshold();
        while (true) {
            const candidate = this.toProcess.pop();
            if (candidate === undefined) {
                break;
            }
            if (threshold === undefined || candidate.distance <= threshold) {
                return candidate;
            }
        }
        return null;
    }
    addResult(pointID, point) {
        const distance = this.s2.pointDistance(this.point, point);
        const threshold = this.distanceThreshold();
        if (this.maxDistanceChordAngle && distance > this.maxDistanceChordAngle) {
            return;
        }
        if (threshold !== undefined && distance >= threshold) {
            return;
        }
        while (this.results.size() >= this.maxResults) {
            this.results.pop();
        }
        this.results.push({ pointID, distance });
    }
    distanceThreshold() {
        const worstEntry = this.results.peek();
        if (worstEntry && this.results.size() >= this.maxResults) {
            if (this.maxDistanceChordAngle &&
                worstEntry.distance > this.maxDistanceChordAngle) {
                throw new Error("Max distance exceeded by entry in heap?");
            }
            return worstEntry.distance;
        }
        return this.maxDistanceChordAngle;
    }
}
//# sourceMappingURL=pointQuery.js.map