diff --git a/README.md b/README.md index 3572906..02afb15 100644 --- a/README.md +++ b/README.md @@ -269,15 +269,24 @@ const example = query({ ctx, { latitude: 40.7813, longitude: -73.9737 }, maxResults, - maxDistance, + { + maxDistance, + filter: (q) => q.eq("category", "coffee"), + }, ); return result; }, }); ``` -The `maxDistance` parameter is optional, but providing it can greatly speed up -searching the index. +The fourth argument can either be a numeric `maxDistance` (for backwards +compatibility) or an options object. When you pass an options object you can +combine `maxDistance` with the same filter builder used by `query`, including +`eq`, `in`, `gte`, and `lt` conditions. These filters are enforced through the +indexed `pointsByFilterKey` range before documents are loaded, so the database +does the heavy lifting and the query avoids reading unrelated points. Pairing +that with a sensible `maxDistance` further constrains the search space and can +greatly speed up searching the index. ## Example diff --git a/package-lock.json b/package-lock.json index b32863b..59debc3 100644 --- a/package-lock.json +++ b/package-lock.json @@ -184,7 +184,6 @@ "integrity": "sha512-e7jT4DxYvIDLk1ZHmU/m/mB19rex9sv0c2ftBtjSBv+kVM/902eh0fINUzD7UwLLNR+jU585GxUJ8/EBfAM5fw==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "@babel/code-frame": "^7.27.1", "@babel/generator": "^7.28.5", @@ -482,7 +481,6 @@ "integrity": "sha512-LqPw+yaSPpCNnVZl5XoHQAySEzlnZiC9gReUuQHMh9GI03KKqwpVqWkIK1UfK116Yww7f2WZuAgnY/nhHwTsJA==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "@edge-runtime/primitives": "5.1.1" }, @@ -1802,7 +1800,6 @@ "integrity": "sha512-FE5u0ezmi6y9OZEzlJfg37mqqf6ZDSF2V/NLjUyGrR9uTZ7Sb9F7bLNZ03S4XVUNRWGA7Ck4c1kK+YnuWjl+DA==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "undici-types": "~6.21.0" } @@ -1813,7 +1810,6 @@ "integrity": "sha512-6mDvHUFSjyT2B2yeNx2nUgMxh9LtOWvkhIU3uePn2I2oyNymUAX1NIsdgviM4CH+JSrp2D2hsMvJOkxY+0wNRA==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "csstype": "^3.0.2" } @@ -1864,7 +1860,6 @@ "integrity": "sha512-BnOroVl1SgrPLywqxyqdJ4l3S2MsKVLDVxZvjI1Eoe8ev2r3kGDo+PcMihNmDE+6/KjkTubSJnmqGZZjQSBq/g==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "@typescript-eslint/scope-manager": "8.46.2", "@typescript-eslint/types": "8.46.2", @@ -2213,7 +2208,6 @@ "integrity": "sha512-NZyJarBfL7nWwIq+FDL6Zp/yHEhePMNnnJ0y3qfieCrmNvYct8uvtiV41UvlSe6apAfk0fY1FbWx+NwfmpvtTg==", "dev": true, "license": "MIT", - "peer": true, "bin": { "acorn": "bin/acorn" }, @@ -2450,7 +2444,6 @@ } ], "license": "MIT", - "peer": true, "dependencies": { "baseline-browser-mapping": "^2.8.19", "caniuse-lite": "^1.0.30001751", @@ -2706,7 +2699,6 @@ "integrity": "sha512-uoIPXRKIp2eLCkkR9WJ2vc9NtgQtx8Pml59WPUahwbrd5EuW2WLI/cf2E7XrUzOSifdQC3kJZepisk4wJNTJaA==", "dev": true, "license": "Apache-2.0", - "peer": true, "dependencies": { "esbuild": "0.25.4", "prettier": "^3.0.0" @@ -2795,8 +2787,7 @@ "resolved": "https://registry.npmjs.org/dayjs/-/dayjs-1.11.18.tgz", "integrity": "sha512-zFBQ7WFRvVRhKcWoUh+ZA1g2HVgUbsZm9sbddh8EC5iv93sui8DVVz1Npvz+r6meo9VKfa8NyLWBsQK1VvIKPA==", "dev": true, - "license": "MIT", - "peer": true + "license": "MIT" }, "node_modules/debug": { "version": "4.4.3", @@ -2948,7 +2939,6 @@ "deprecated": "This version is no longer supported. Please see https://eslint.org/version-support for other options.", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "@eslint-community/eslint-utils": "^4.2.0", "@eslint-community/regexpp": "^4.6.1", @@ -3762,8 +3752,7 @@ "resolved": "https://registry.npmjs.org/leaflet/-/leaflet-1.9.4.tgz", "integrity": "sha512-nxS1ynzJOmOlHp+iL3FyWqK89GtNL8U8rvlMOsQdTTssxZwCXh8N2NB3GDQOL+YR3XnWyZAxwQixURb+FA74PA==", "dev": true, - "license": "BSD-2-Clause", - "peer": true + "license": "BSD-2-Clause" }, "node_modules/levn": { "version": "0.4.1", @@ -4956,7 +4945,6 @@ "integrity": "sha512-wS+hAgJShR0KhEvPJArfuPVN1+Hz1t0Y6n5jLrGQbkb4urgPE/0Rve+1kMB1v/oWgHgm4WIcV+i7F2pTVj+2iQ==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "loose-envify": "^1.1.0" }, @@ -4970,7 +4958,6 @@ "integrity": "sha512-5m4nQKp+rZRb09LNH59GM4BxTh9251/ylbKIbpe7TpGxfJ+9kv6BLkLBXIjjspbgbnIBNqlI23tRnTWT0snUIw==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "loose-envify": "^1.1.0", "scheduler": "^0.23.2" @@ -5456,7 +5443,6 @@ "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", "dev": true, "license": "MIT", - "peer": true, "engines": { "node": ">=12" }, @@ -5559,7 +5545,6 @@ "integrity": "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==", "dev": true, "license": "Apache-2.0", - "peer": true, "bin": { "tsc": "bin/tsc", "tsserver": "bin/tsserver" @@ -5646,7 +5631,6 @@ "integrity": "sha512-+Oxm7q9hDoLMyJOYfUYBuHQo+dkAloi33apOPP56pzj+vsdJDzr+j1NISE5pyaAuKL4A3UD34qd0lx5+kfKp2g==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "esbuild": "^0.25.0", "fdir": "^6.4.4", @@ -5763,7 +5747,6 @@ "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", "dev": true, "license": "MIT", - "peer": true, "engines": { "node": ">=12" }, @@ -5777,7 +5760,6 @@ "integrity": "sha512-LUCP5ev3GURDysTWiP47wRRUpLKMOfPh+yKTx3kVIEiu5KOMeqzpnYNsKyOoVrULivR8tLcks4+lga33Whn90A==", "dev": true, "license": "MIT", - "peer": true, "dependencies": { "@types/chai": "^5.2.2", "@vitest/expect": "3.2.4", @@ -6112,7 +6094,6 @@ "integrity": "sha512-JInaHOamG8pt5+Ey8kGmdcAcg3OL9reK8ltczgHTAwNhMys/6ThXHityHxVV2p3fkw/c+MAvBHFVYHFZDmjMCQ==", "dev": true, "license": "MIT", - "peer": true, "funding": { "url": "https://github.com/sponsors/colinhacks" } diff --git a/package.json b/package.json index f6e6218..179bf2d 100644 --- a/package.json +++ b/package.json @@ -106,4 +106,4 @@ }, "types": "./dist/client/index.d.ts", "module": "./dist/client/index.js" -} +} \ No newline at end of file diff --git a/src/client/index.ts b/src/client/index.ts index 995a5b0..31d6256 100644 --- a/src/client/index.ts +++ b/src/client/index.ts @@ -37,6 +37,13 @@ export type GeospatialDocument< sortKey: number; }; +export type QueryNearestOptions< + Doc extends GeospatialDocument = GeospatialDocument, +> = { + maxDistance?: number; + filter?: NonNullable["filter"]>; +}; + export interface GeospatialIndexOptions { /** * The minimum S2 cell level to use when querying. Defaults to 4. @@ -209,23 +216,50 @@ export class GeospatialIndex< * @param ctx - The Convex query context. * @param point - The point to query for. * @param maxResults - The maximum number of results to return. - * @param maxDistance - The maximum distance to return results within in meters. + * @param maxDistanceOrOptions - Either the maximum distance in meters or an options object containing filtering logic. + * @param maybeOptions - Additional options when the maximum distance is provided separately. * @returns - An array of objects with the key-coordinate pairs and their distance from the query point in meters. */ async queryNearest( ctx: QueryCtx, point: Point, maxResults: number, - maxDistance?: number, + maxDistanceOrOptions?: + | number + | QueryNearestOptions>, + maybeOptions?: QueryNearestOptions>, ) { + let options: + | QueryNearestOptions> + | undefined; + let maxDistance: number | undefined; + if ( + typeof maxDistanceOrOptions === "object" && + maxDistanceOrOptions !== null + ) { + options = maxDistanceOrOptions; + } else { + maxDistance = maxDistanceOrOptions; + options = maybeOptions; + } + + const filterBuilder = new FilterBuilderImpl< + GeospatialDocument + >(); + if (options?.filter) { + options.filter(filterBuilder); + } + const resp = await ctx.runQuery(this.component.query.nearestPoints, { point, - maxDistance, + maxDistance: options?.maxDistance ?? maxDistance, maxResults, minLevel: this.minLevel, maxLevel: this.maxLevel, levelMod: this.levelMod, logLevel: this.logLevel, + filtering: filterBuilder.filterConditions, + sorting: { interval: filterBuilder.interval ?? {} }, }); return resp as { key: Key; coordinates: Point; distance: number }[]; } diff --git a/src/component/_generated/api.d.ts b/src/component/_generated/api.d.ts new file mode 100644 index 0000000..36a6bf5 --- /dev/null +++ b/src/component/_generated/api.d.ts @@ -0,0 +1,225 @@ +/* prettier-ignore-start */ + +/* eslint-disable */ +/** + * Generated `api` utility. + * + * THIS CODE IS AUTOMATICALLY GENERATED. + * + * To regenerate, run `npx convex dev`. + * @module + */ + +import type * as document from "../document.js"; +import type * as lib_approximateCounter from "../lib/approximateCounter.js"; +import type * as lib_d64 from "../lib/d64.js"; +import type * as lib_goRuntime from "../lib/goRuntime.js"; +import type * as lib_interval from "../lib/interval.js"; +import type * as lib_logging from "../lib/logging.js"; +import type * as lib_pointQuery from "../lib/pointQuery.js"; +import type * as lib_primitive from "../lib/primitive.js"; +import type * as lib_s2Bindings from "../lib/s2Bindings.js"; +import type * as lib_s2wasm from "../lib/s2wasm.js"; +import type * as lib_tupleKey from "../lib/tupleKey.js"; +import type * as lib_xxhash from "../lib/xxhash.js"; +import type * as query from "../query.js"; +import type * as streams_cellRange from "../streams/cellRange.js"; +import type * as streams_databaseRange from "../streams/databaseRange.js"; +import type * as streams_filterKeyRange from "../streams/filterKeyRange.js"; +import type * as streams_intersection from "../streams/intersection.js"; +import type * as streams_union from "../streams/union.js"; +import type * as streams_zigzag from "../streams/zigzag.js"; +import type * as types from "../types.js"; + +import type { + ApiFromModules, + FilterApi, + FunctionReference, +} from "convex/server"; +/** + * A utility for referencing Convex functions in your app's API. + * + * Usage: + * ```js + * const myFunctionReference = api.myModule.myFunction; + * ``` + */ +declare const fullApi: ApiFromModules<{ + document: typeof document; + "lib/approximateCounter": typeof lib_approximateCounter; + "lib/d64": typeof lib_d64; + "lib/goRuntime": typeof lib_goRuntime; + "lib/interval": typeof lib_interval; + "lib/logging": typeof lib_logging; + "lib/pointQuery": typeof lib_pointQuery; + "lib/primitive": typeof lib_primitive; + "lib/s2Bindings": typeof lib_s2Bindings; + "lib/s2wasm": typeof lib_s2wasm; + "lib/tupleKey": typeof lib_tupleKey; + "lib/xxhash": typeof lib_xxhash; + query: typeof query; + "streams/cellRange": typeof streams_cellRange; + "streams/databaseRange": typeof streams_databaseRange; + "streams/filterKeyRange": typeof streams_filterKeyRange; + "streams/intersection": typeof streams_intersection; + "streams/union": typeof streams_union; + "streams/zigzag": typeof streams_zigzag; + types: typeof types; +}>; +export type Mounts = { + document: { + get: FunctionReference< + "query", + "public", + { key: string }, + { + coordinates: { latitude: number; longitude: number }; + filterKeys: Record< + string, + | string + | number + | boolean + | null + | bigint + | Array + >; + key: string; + sortKey: number; + } | null + >; + insert: FunctionReference< + "mutation", + "public", + { + document: { + coordinates: { latitude: number; longitude: number }; + filterKeys: Record< + string, + | string + | number + | boolean + | null + | bigint + | Array + >; + key: string; + sortKey: number; + }; + levelMod: number; + maxCells: number; + maxLevel: number; + minLevel: number; + }, + null + >; + remove: FunctionReference< + "mutation", + "public", + { + key: string; + levelMod: number; + maxCells: number; + maxLevel: number; + minLevel: number; + }, + boolean + >; + }; + query: { + debugCells: FunctionReference< + "query", + "public", + { + levelMod: number; + maxCells: number; + maxLevel: number; + minLevel: number; + rectangle: { east: number; north: number; south: number; west: number }; + }, + Array<{ + token: string; + vertices: Array<{ latitude: number; longitude: number }>; + }> + >; + execute: FunctionReference< + "query", + "public", + { + cursor?: string; + levelMod: number; + logLevel: "DEBUG" | "INFO" | "WARN" | "ERROR"; + maxCells: number; + maxLevel: number; + minLevel: number; + query: { + filtering: Array<{ + filterKey: string; + filterValue: string | number | boolean | null | bigint; + occur: "should" | "must"; + }>; + maxResults: number; + rectangle: { + east: number; + north: number; + south: number; + west: number; + }; + sorting: { + interval: { endExclusive?: number; startInclusive?: number }; + }; + }; + }, + { + nextCursor?: string; + results: Array<{ + coordinates: { latitude: number; longitude: number }; + key: string; + }>; + } + >; + nearestPoints: FunctionReference< + "query", + "public", + { + levelMod: number; + logLevel: "DEBUG" | "INFO" | "WARN" | "ERROR"; + maxDistance?: number; + maxLevel: number; + maxResults: number; + minLevel: number; + nextCursor?: string; + point: { latitude: number; longitude: number }; + filtering: Array<{ + filterKey: string; + filterValue: string | number | boolean | null | bigint; + occur: "should" | "must"; + }>; + sorting: { + interval: { endExclusive?: number; startInclusive?: number }; + }; + }, + Array<{ + coordinates: { latitude: number; longitude: number }; + distance: number; + key: string; + }> + >; + }; +}; +// For now fullApiWithMounts is only fullApi which provides +// jump-to-definition in component client code. +// Use Mounts for the same type without the inference. +declare const fullApiWithMounts: typeof fullApi; + +export declare const api: FilterApi< + typeof fullApiWithMounts, + FunctionReference +>; +export declare const internal: FilterApi< + typeof fullApiWithMounts, + FunctionReference +>; + +export declare const components: {}; + +/* prettier-ignore-end */ diff --git a/src/component/_generated/component.ts b/src/component/_generated/component.ts index 8dd0668..8370781 100644 --- a/src/component/_generated/component.ts +++ b/src/component/_generated/component.ts @@ -155,6 +155,14 @@ export type ComponentApi = minLevel: number; nextCursor?: string; point: { latitude: number; longitude: number }; + filtering: Array<{ + filterKey: string; + filterValue: string | number | boolean | null | bigint; + occur: "should" | "must"; + }>; + sorting: { + interval: { endExclusive?: number; startInclusive?: number }; + }; }, Array<{ coordinates: { latitude: number; longitude: number }; diff --git a/src/component/lib/pointQuery.ts b/src/component/lib/pointQuery.ts index d988fed..944ad58 100644 --- a/src/component/lib/pointQuery.ts +++ b/src/component/lib/pointQuery.ts @@ -1,12 +1,24 @@ import { Heap } from "heap-js"; -import type { ChordAngle, Meters, Point } from "../types.js"; -import type { Id } from "../_generated/dataModel.js"; +import type { ChordAngle, Meters, Point, Primitive } from "../types.js"; +import type { Doc, Id } from "../_generated/dataModel.js"; import { S2Bindings } from "./s2Bindings.js"; import type { QueryCtx } from "../_generated/server.js"; import * as approximateCounter from "./approximateCounter.js"; -import { cellCounterKey } from "../streams/cellRange.js"; +import { cellCounterKey, CellRange } from "../streams/cellRange.js"; +import { FilterKeyRange } from "../streams/filterKeyRange.js"; +import { Union } from "../streams/union.js"; +import { Intersection } from "../streams/intersection.js"; +import type { PointSet, Stats } from "../streams/zigzag.js"; +import { PREFETCH_SIZE } from "../streams/constants.js"; import { decodeTupleKey } from "./tupleKey.js"; import type { Logger } from "./logging.js"; +import type { Interval } from "./interval.js"; + +type FilterCondition = { + filterKey: string; + filterValue: Primitive; + occur: "must" | "should"; +}; export class ClosestPointQuery { // Min-heap of cells to process. @@ -16,6 +28,12 @@ export class ClosestPointQuery { results: Heap; maxDistanceChordAngle?: ChordAngle; + private mustFilters: FilterCondition[]; + private shouldFilters: FilterCondition[]; + private sortInterval: Interval; + private readonly checkFilters: boolean; + private static readonly FILTER_SUBDIVIDE_THRESHOLD = 8; + private cellStreams = new Map(); constructor( private s2: S2Bindings, @@ -26,11 +44,20 @@ export class ClosestPointQuery { private minLevel: number, private maxLevel: number, private levelMod: number, + filtering: FilterCondition[] = [], + interval: Interval = {}, ) { 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); + this.mustFilters = filtering.filter((filter) => filter.occur === "must"); + this.shouldFilters = filtering.filter( + (filter) => filter.occur === "should", + ); + this.sortInterval = interval; + this.checkFilters = + this.mustFilters.length > 0 || this.shouldFilters.length > 0; for (const cellID of this.s2.initialCells(this.minLevel)) { const distance = this.s2.minDistanceToCell(this.point, cellID); @@ -55,7 +82,16 @@ export class ClosestPointQuery { ); this.logger.debug(`Size estimate for ${cellIDToken}: ${sizeEstimate}`); - if (canSubdivide && sizeEstimate >= approximateCounter.SAMPLING_RATE) { + const approxRows = Math.floor( + sizeEstimate / approximateCounter.SAMPLING_RATE, + ); + const shouldSubdivide = + canSubdivide && + (approxRows >= 1 || + (this.checkFilters && + approxRows >= ClosestPointQuery.FILTER_SUBDIVIDE_THRESHOLD)); + + if (shouldSubdivide) { this.logger.debug(`Subdividing cell ${candidate.cellID}`); const nextLevel = Math.min( candidate.level + this.levelMod, @@ -69,23 +105,35 @@ export class ClosestPointQuery { 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) { + const streamState = this.getOrCreateStreamForCell(ctx, cellIDToken); + while (!streamState.done) { + if (this.shouldStopProcessingCell(candidate.distance)) { + break; + } + const tupleKey = await streamState.stream.current(); + if (tupleKey === null) { + streamState.done = true; + break; + } + const { pointId, sortKey } = decodeTupleKey(tupleKey); + if (!this.withinSortInterval(sortKey)) { + const next = await streamState.stream.advance(); + if (next === null) { + streamState.done = true; + } + continue; + } + const point = await ctx.db.get(pointId); if (!point) { throw new Error("Point not found"); } - this.addResult(point._id, point.coordinates); + if (this.matchesFilters(point)) { + this.addResult(point._id, point.coordinates); + } + const nextTuple = await streamState.stream.advance(); + if (nextTuple === null) { + streamState.done = true; + } } } } @@ -99,15 +147,156 @@ export class ClosestPointQuery { if (!point) { throw new Error("Point not found"); } + if (!this.matchesFilters(point)) { + continue; + } results.push({ key: point.key, coordinates: point.coordinates, distance: this.s2.chordAngleToMeters(entries[i].distance), }); } + this.cellStreams.clear(); return results; } + private shouldStopProcessingCell(candidateDistance: ChordAngle): boolean { + if (this.results.size() < this.maxResults) { + return false; + } + const threshold = this.distanceThreshold(); + if (threshold === undefined) { + return false; + } + return threshold <= candidateDistance; + } + + private withinSortInterval(sortKey: number): boolean { + if ( + this.sortInterval.startInclusive !== undefined && + sortKey < this.sortInterval.startInclusive + ) { + return false; + } + if ( + this.sortInterval.endExclusive !== undefined && + sortKey >= this.sortInterval.endExclusive + ) { + return false; + } + return true; + } + + private getOrCreateStreamForCell( + ctx: QueryCtx, + cellIDToken: string, + ): CellStreamState { + const existing = this.cellStreams.get(cellIDToken); + if (existing) { + return existing; + } + const stats: Stats = { + cells: 1, + queriesIssued: 0, + rowsRead: 0, + rowsPostFiltered: 0, + }; + const ranges: PointSet[] = [ + new CellRange( + ctx, + this.logger, + cellIDToken, + undefined, + this.sortInterval, + PREFETCH_SIZE, + stats, + ), + ]; + for (const filter of this.mustFilters) { + ranges.push( + new FilterKeyRange( + ctx, + this.logger, + filter.filterKey, + filter.filterValue, + undefined, + this.sortInterval, + PREFETCH_SIZE, + stats, + ), + ); + } + if (this.shouldFilters.length > 0) { + const shouldRanges = this.shouldFilters.map( + (filter) => + new FilterKeyRange( + ctx, + this.logger, + filter.filterKey, + filter.filterValue, + undefined, + this.sortInterval, + PREFETCH_SIZE, + stats, + ), + ); + ranges.push( + shouldRanges.length === 1 ? shouldRanges[0] : new Union(shouldRanges), + ); + } + const stream = ranges.length === 1 ? ranges[0] : new Intersection(ranges); + const state: CellStreamState = { stream, done: false }; + this.cellStreams.set(cellIDToken, state); + return state; + } + + private matchesFilters(point: Doc<"points">): boolean { + if ( + this.sortInterval.startInclusive !== undefined && + point.sortKey < this.sortInterval.startInclusive + ) { + return false; + } + if ( + this.sortInterval.endExclusive !== undefined && + point.sortKey >= this.sortInterval.endExclusive + ) { + return false; + } + + for (const filter of this.mustFilters) { + if (!this.pointMatchesCondition(point, filter)) { + return false; + } + } + + if (this.shouldFilters.length > 0) { + let anyMatch = false; + for (const filter of this.shouldFilters) { + if (this.pointMatchesCondition(point, filter)) { + anyMatch = true; + break; + } + } + if (!anyMatch) { + return false; + } + } + + return true; + } + + private pointMatchesCondition(point: Doc<"points">, filter: FilterCondition) { + const value = point.filterKeys[filter.filterKey]; + if (value === undefined) { + return false; + } + if (Array.isArray(value)) { + return value.some((candidate) => candidate === filter.filterValue); + } + return value === filter.filterValue; + } + addCandidate(cellID: bigint, level: number, distance: ChordAngle) { if (this.maxDistanceChordAngle && distance > this.maxDistanceChordAngle) { return; @@ -173,3 +362,8 @@ type Result = { pointID: Id<"points">; distance: ChordAngle; }; + +type CellStreamState = { + stream: PointSet; + done: boolean; +}; diff --git a/src/component/query.ts b/src/component/query.ts index 33e47ee..8e2314c 100644 --- a/src/component/query.ts +++ b/src/component/query.ts @@ -13,8 +13,9 @@ import type { Doc } from "./_generated/dataModel.js"; import { createLogger, logLevel } from "./lib/logging.js"; import { S2Bindings } from "./lib/s2Bindings.js"; import { ClosestPointQuery } from "./lib/pointQuery.js"; +import { PREFETCH_SIZE } from "./streams/constants.js"; -export const PREFETCH_SIZE = 16; +export { PREFETCH_SIZE } from "./streams/constants.js"; const equalityCondition = v.object({ occur: v.union(v.literal("should"), v.literal("must")), @@ -274,6 +275,10 @@ export const nearestPoints = query({ maxLevel: v.number(), levelMod: v.number(), nextCursor: v.optional(v.string()), + filtering: v.array(equalityCondition), + sorting: v.object({ + interval, + }), logLevel, }, returns: v.array(queryResultWithDistance), @@ -292,6 +297,8 @@ export const nearestPoints = query({ args.minLevel, args.maxLevel, args.levelMod, + args.filtering, + args.sorting.interval, ); const results = await query.execute(ctx); return results; diff --git a/src/component/streams/constants.ts b/src/component/streams/constants.ts new file mode 100644 index 0000000..97ae041 --- /dev/null +++ b/src/component/streams/constants.ts @@ -0,0 +1,2 @@ +export const PREFETCH_SIZE = 16; + diff --git a/src/component/streams/intersection.ts b/src/component/streams/intersection.ts index 4634ece..0f75c6e 100644 --- a/src/component/streams/intersection.ts +++ b/src/component/streams/intersection.ts @@ -1,5 +1,5 @@ import type { TupleKey } from "../lib/tupleKey.js"; -import { PREFETCH_SIZE } from "../query.js"; +import { PREFETCH_SIZE } from "./constants.js"; import type { PointSet } from "./zigzag.js"; export class Intersection implements PointSet { diff --git a/src/component/tests/pointQuery.test.ts b/src/component/tests/pointQuery.test.ts index fb41b7a..0f584db 100644 --- a/src/component/tests/pointQuery.test.ts +++ b/src/component/tests/pointQuery.test.ts @@ -27,19 +27,19 @@ test("closest point query - basic functionality", async () => { key: "point1", coordinates: { latitude: 0, longitude: 0 }, sortKey: 1, - filterKeys: {}, + filterKeys: { category: "coffee" }, }, { key: "point2", coordinates: { latitude: 1, longitude: 1 }, sortKey: 2, - filterKeys: {}, + filterKeys: { category: "tea" }, }, { key: "point3", coordinates: { latitude: -1, longitude: -1 }, sortKey: 3, - filterKeys: {}, + filterKeys: { category: "coffee" }, }, ]; @@ -99,6 +99,96 @@ test("closest point query - basic functionality", async () => { const result3 = await query3.execute(ctx); expect(result3.length).toBe(1); expect(result3[0].key).toBe("point1"); + + // Test must filter + const query4 = new ClosestPointQuery( + s2, + logger, + { latitude: 0, longitude: 0 }, + 10000000, + 3, + opts.minLevel, + opts.maxLevel, + opts.levelMod, + [ + { + occur: "must", + filterKey: "category", + filterValue: "coffee", + }, + ], + ); + const result4 = await query4.execute(ctx); + expect(result4.length).toBe(2); + expect(result4.map((r) => r.key).sort()).toEqual(["point1", "point3"]); + + // Test should filter (must match at least one) + const query5 = new ClosestPointQuery( + s2, + logger, + { latitude: 0, longitude: 0 }, + 10000000, + 3, + opts.minLevel, + opts.maxLevel, + opts.levelMod, + [ + { + occur: "should", + filterKey: "category", + filterValue: "tea", + }, + ], + ); + const result5 = await query5.execute(ctx); + expect(result5.length).toBe(1); + expect(result5[0].key).toBe("point2"); + + // Test sort key interval + const query6 = new ClosestPointQuery( + s2, + logger, + { latitude: 0, longitude: 0 }, + 10000000, + 3, + opts.minLevel, + opts.maxLevel, + opts.levelMod, + [], + { startInclusive: 3 }, + ); + const result6 = await query6.execute(ctx); + expect(result6.length).toBe(1); + expect(result6[0].key).toBe("point3"); + + // Test multiple should filters + const query7 = new ClosestPointQuery( + s2, + logger, + { latitude: 0, longitude: 0 }, + 10000000, + 3, + opts.minLevel, + opts.maxLevel, + opts.levelMod, + [ + { + occur: "should", + filterKey: "category", + filterValue: "tea", + }, + { + occur: "should", + filterKey: "category", + filterValue: "coffee", + }, + ], + ); + const result7 = await query7.execute(ctx); + expect(result7.length).toBe(3); + expect(new Set(result7.map((r) => r.key))).toEqual( + new Set(["point1", "point2", "point3"]), + ); }); });