Add support for filtering to queryNearest

README.md

@@ -251,14 +251,17 @@ 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. Filtering helps constrain the search space and can speed up lookups.
 
 ## Example
 

example/convex/_generated/api.d.ts

@@ -172,6 +172,14 @@ export declare const components: {
           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 };

package.json

@@ -23,8 +23,8 @@
   "scripts": {
     "build": "npm run build:esm && npm run build:cjs",
     "build:go": "cd src/s2-bindings && python build.py",
-    "build:esm": "tsc --project ./esm.json && echo '{\\n  \"type\": \"module\"\\n}' > dist/esm/package.json",
-    "build:cjs": "tsc --project ./commonjs.json && echo '{\\n  \"type\": \"commonjs\"\\n}' > dist/commonjs/package.json",
+    "build:esm": "tsc --project ./esm.json && node -e \"require('fs').writeFileSync('dist/esm/package.json', JSON.stringify({type:'module'}, null, 2) + '\\n')\"",
+    "build:cjs": "tsc --project ./commonjs.json && node -e \"require('fs').writeFileSync('dist/commonjs/package.json', JSON.stringify({type:'commonjs'}, null, 2) + '\\n')\"",
     "dev": "cd example; npm run dev",
     "typecheck": "tsc --noEmit",
     "prepare": "npm run build",
@@ -94,4 +94,4 @@
   "main": "./dist/commonjs/client/index.js",
   "types": "./dist/commonjs/client/index.d.ts",
   "module": "./dist/esm/client/index.js"
-}
+}

src/client/index.ts

@@ -39,6 +39,13 @@ export type GeospatialDocument<
   sortKey: number;
 };
 
+export type QueryNearestOptions<
+  Doc extends GeospatialDocument = GeospatialDocument,
+> = {
+  maxDistance?: number;
+  filter?: NonNullable<GeospatialQuery<Doc>["filter"]>;
+};
+
 export interface GeospatialIndexOptions {
   /**
    * The minimum S2 cell level to use when querying. Defaults to 4.
@@ -211,23 +218,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<GeospatialDocument<Key, Filters>>,
+    maybeOptions?: QueryNearestOptions<GeospatialDocument<Key, Filters>>,
   ) {
+    let options:
+      | QueryNearestOptions<GeospatialDocument<Key, Filters>>
+      | undefined;
+    let maxDistance: number | undefined;
+    if (
+      typeof maxDistanceOrOptions === "object" &&
+      maxDistanceOrOptions !== null
+    ) {
+      options = maxDistanceOrOptions;
+    } else {
+      maxDistance = maxDistanceOrOptions;
+      options = maybeOptions;
+    }
+
+    const filterBuilder = new FilterBuilderImpl<
+      GeospatialDocument<Key, Filters>
+    >();
+    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 }[];
   }

src/component/_generated/api.d.ts

@@ -189,6 +189,14 @@ export type Mounts = {
         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 };

src/component/lib/pointQuery.ts

@@ -1,12 +1,20 @@
 import { Heap } from "heap-js";
+import type { Primitive } from "../types.js";
 import { ChordAngle, Meters, Point } from "../types.js";
-import { Id } from "../_generated/dataModel.js";
+import { Doc, Id } from "../_generated/dataModel.js";
 import { S2Bindings } from "./s2Bindings.js";
 import { QueryCtx } from "../_generated/server.js";
 import * as approximateCounter from "./approximateCounter.js";
 import { cellCounterKey } from "../streams/cellRange.js";
 import { decodeTupleKey } from "./tupleKey.js";
 import { 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 +24,9 @@ export class ClosestPointQuery {
   results: Heap<Result>;
 
   maxDistanceChordAngle?: ChordAngle;
+  private mustFilters: FilterCondition[];
+  private shouldFilters: FilterCondition[];
+  private sortInterval: Interval;
 
   constructor(
     private s2: S2Bindings,
@@ -26,11 +37,18 @@ export class ClosestPointQuery {
     private minLevel: number,
     private maxLevel: number,
     private levelMod: number,
+    filtering: FilterCondition[] = [],
+    interval: Interval = {},
   ) {
     this.toProcess = new Heap<CellCandidate>((a, b) => a.distance - b.distance);
     this.results = new Heap<Result>((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;
 
     for (const cellID of this.s2.initialCells(this.minLevel)) {
       const distance = this.s2.minDistanceToCell(this.point, cellID);
@@ -85,7 +103,9 @@ export class ClosestPointQuery {
           if (!point) {
             throw new Error("Point not found");
           }
-          this.addResult(point._id, point.coordinates);
+          if (this.matchesFilters(point)) {
+            this.addResult(point._id, point.coordinates);
+          }
         }
       }
     }
@@ -99,6 +119,9 @@ export class ClosestPointQuery {
       if (!point) {
         throw new Error("Point not found");
       }
+      if (!this.matchesFilters(point)) {
+        continue;
+      }
       results.push({
         key: point.key,
         coordinates: point.coordinates,
@@ -108,6 +131,53 @@ export class ClosestPointQuery {
     return results;
   }
 
+  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;

src/component/query.ts

@@ -274,6 +274,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 +296,8 @@ export const nearestPoints = query({
       args.minLevel,
       args.maxLevel,
       args.levelMod,
+      args.filtering,
+      args.sorting.interval,
     );
     const results = await query.execute(ctx);
     return results;

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,67 @@ 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");
   });
 });