@convex-dev/geospatial
Version:
A geospatial index for Convex
102 lines • 3.27 kB
JavaScript
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