Use sets to speed up queryRenderedFeatures (#4777)

* Use sets to speed up queryRenderedFeatures * Add Set to documentation description * Add types to test instances * Add type to import * Improve upstream types * Fix more upstream types * Improve expectations of null/undefined values * Update src/source/source_cache.ts Co-authored-by: neodescis <[email protected]> * Add benchmark * Move type definition outside method * Move type declaration * Add changelog --------- Co-authored-by: Harel M <[email protected]> Co-authored-by: neodescis <[email protected]> Co-authored-by: Isaac Besora <[email protected]>
maplibre · Oct 8, 2024 · c4b083f · c4b083f
1 parent 5222719
commit c4b083f
Show file tree

Hide file tree

Showing 9 changed files with 113 additions and 50 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,7 +1,7 @@
 ## main
 
 ### ✨ Features and improvements
-- _...Add new stuff here..._
+- Improve performance of `queryRenderedFeatures` by using JavaScript `Set`s to assess layer membership internally.
 
 ### 🐞 Bug fixes
 - _...Add new stuff here..._

diff --git a/src/data/feature_index.ts b/src/data/feature_index.ts
@@ -9,7 +9,7 @@ import vt from '@mapbox/vector-tile';
 import Protobuf from 'pbf';
 import {GeoJSONFeature} from '../util/vectortile_to_geojson';
 import type {MapGeoJSONFeature} from '../util/vectortile_to_geojson';
-import {arraysIntersect, mapObject, extend} from '../util/util';
+import {mapObject, extend} from '../util/util';
 import {OverscaledTileID} from '../source/tile_id';
 import {register} from '../util/web_worker_transfer';
 import {EvaluationParameters} from '../style/evaluation_parameters';
@@ -33,9 +33,9 @@ type QueryParameters = {
     cameraQueryGeometry: Array<Point>;
     queryPadding: number;
     params: {
-        filter: FilterSpecification;
-        layers: Array<string>;
-        availableImages: Array<string>;
+        filter?: FilterSpecification;
+        layers?: Set<string> | null;
+        availableImages?: Array<string>;
     };
 };
 
@@ -113,9 +113,9 @@ export class FeatureIndex {
     ): {[_: string]: Array<{featureIndex: number; feature: GeoJSONFeature}>} {
         this.loadVTLayers();
 
-        const params = args.params || {} as { filter: any; layers: string[]; availableImages: string[] },
-            pixelsToTileUnits = EXTENT / args.tileSize / args.scale,
-            filter = featureFilter(params.filter);
+        const params = args.params;
+        const pixelsToTileUnits = EXTENT / args.tileSize / args.scale;
+        const filter = featureFilter(params.filter);
 
         const queryGeometry = args.queryGeometry;
         const queryPadding = args.queryPadding * pixelsToTileUnits;
@@ -183,7 +183,7 @@ export class FeatureIndex {
         sourceLayerIndex: number,
         featureIndex: number,
         filter: FeatureFilter,
-        filterLayerIDs: Array<string>,
+        filterLayerIDs: Set<string> | undefined,
         availableImages: Array<string>,
         styleLayers: {[_: string]: StyleLayer},
         serializedLayers: {[_: string]: any},
@@ -196,7 +196,7 @@ export class FeatureIndex {
         ) => boolean | number) {
 
         const layerIDs = this.bucketLayerIDs[bucketIndex];
-        if (filterLayerIDs && !arraysIntersect(filterLayerIDs, layerIDs))
+        if (filterLayerIDs && !layerIDs.some(id => filterLayerIDs.has(id)))
             return;
 
         const sourceLayerName = this.sourceLayerCoder.decode(sourceLayerIndex);
@@ -217,7 +217,7 @@ export class FeatureIndex {
         for (let l = 0; l < layerIDs.length; l++) {
             const layerID = layerIDs[l];
 
-            if (filterLayerIDs && filterLayerIDs.indexOf(layerID) < 0) {
+            if (filterLayerIDs && !filterLayerIDs.has(layerID)) {
                 continue;
             }
 
@@ -259,7 +259,7 @@ export class FeatureIndex {
         bucketIndex: number,
         sourceLayerIndex: number,
         filterSpec: FilterSpecification,
-        filterLayerIDs: Array<string>,
+        filterLayerIDs: Set<string> | null,
         availableImages: Array<string>,
         styleLayers: {[_: string]: StyleLayer}) {
         const result = {};

diff --git a/src/source/query_features.ts b/src/source/query_features.ts
@@ -13,10 +13,10 @@ import {mat4} from 'gl-matrix';
  */
 export type QueryRenderedFeaturesOptions = {
     /**
-     * An array of [style layer IDs](https://maplibre.org/maplibre-style-spec/#layer-id) for the query to inspect.
+     * An array or set of [style layer IDs](https://maplibre.org/maplibre-style-spec/#layer-id) for the query to inspect.
      * Only features within these layers will be returned. If this parameter is undefined, all layers will be checked.
      */
-    layers?: Array<string>;
+    layers?: Array<string> | Set<string>;
     /**
      * A [filter](https://maplibre.org/maplibre-style-spec/layers/#filter) to limit query results.
      */
@@ -31,6 +31,10 @@ export type QueryRenderedFeaturesOptions = {
     validate?: boolean;
 };
 
+export type QueryRenderedFeaturesOptionsStrict = Omit<QueryRenderedFeaturesOptions, 'layers'> & {
+    layers: Set<string> | null;
+}
+
 /**
  * The options object related to the {@link Map#querySourceFeatures} method
  */
@@ -65,7 +69,7 @@ function getPixelPosMatrix(transform, tileID) {
     }
 }
 
-function queryIncludes3DLayer(layers: Array<string>, styleLayers: {[_: string]: StyleLayer}, sourceID: string) {
+function queryIncludes3DLayer(layers: Set<string> | undefined, styleLayers: {[_: string]: StyleLayer}, sourceID: string) {
     if (layers) {
         for (const layerID of layers) {
             const layer = styleLayers[layerID];
@@ -89,11 +93,11 @@ export function queryRenderedFeatures(
     styleLayers: {[_: string]: StyleLayer},
     serializedLayers: {[_: string]: any},
     queryGeometry: Array<Point>,
-    params: QueryRenderedFeaturesOptions,
+    params: QueryRenderedFeaturesOptionsStrict | undefined,
     transform: IReadonlyTransform
 ): { [key: string]: Array<{featureIndex: number; feature: MapGeoJSONFeature}> } {
 
-    const has3DLayer = queryIncludes3DLayer(params && params.layers, styleLayers, sourceCache.id);
+    const has3DLayer = queryIncludes3DLayer(params?.layers ?? null, styleLayers, sourceCache.id);
     const maxPitchScaleFactor = transform.maxPitchScaleFactor();
     const tilesIn = sourceCache.tilesIn(queryGeometry, maxPitchScaleFactor, has3DLayer);
 
@@ -137,14 +141,14 @@ export function queryRenderedSymbols(styleLayers: {[_: string]: StyleLayer},
     serializedLayers: {[_: string]: StyleLayer},
     sourceCaches: {[_: string]: SourceCache},
     queryGeometry: Array<Point>,
-    params: QueryRenderedFeaturesOptions,
+    params: QueryRenderedFeaturesOptionsStrict,
     collisionIndex: CollisionIndex,
     retainedQueryData: {
         [_: number]: RetainedQueryData;
     }) {
     const result = {};
     const renderedSymbols = collisionIndex.queryRenderedSymbols(queryGeometry);
-    const bucketQueryData = [];
+    const bucketQueryData: RetainedQueryData[] = [];
     for (const bucketInstanceId of Object.keys(renderedSymbols).map(Number)) {
         bucketQueryData.push(retainedQueryData[bucketInstanceId]);
     }
@@ -205,7 +209,7 @@ export function queryRenderedSymbols(styleLayers: {[_: string]: StyleLayer},
     return result;
 }
 
-export function querySourceFeatures(sourceCache: SourceCache, params: QuerySourceFeatureOptions) {
+export function querySourceFeatures(sourceCache: SourceCache, params: QuerySourceFeatureOptions | undefined) {
     const tiles = sourceCache.getRenderableIds().map((id) => {
         return sourceCache.getTileByID(id);
     });

diff --git a/src/source/source_cache.ts b/src/source/source_cache.ts
@@ -24,6 +24,14 @@ import type {MapSourceDataEvent} from '../ui/events';
 import type {Terrain} from '../render/terrain';
 import type {CanvasSourceSpecification} from './canvas_source';
 
+type TileResult = {
+    tile: Tile;
+    tileID: OverscaledTileID;
+    queryGeometry: Array<Point>;
+    cameraQueryGeometry: Array<Point>;
+    scale: number;
+}
+
 /**
  * @internal
  * `SourceCache` is responsible for
@@ -948,9 +956,8 @@ export class SourceCache extends Evented {
      * @param pointQueryGeometry - coordinates of the corners of bounding rectangle
      * @returns result items have `{tile, minX, maxX, minY, maxY}`, where min/max bounding values are the given bounds transformed in into the coordinate space of this tile.
      */
-    tilesIn(pointQueryGeometry: Array<Point>, maxPitchScaleFactor: number, has3DLayer: boolean): any[] {
-
-        const tileResults = [];
+    tilesIn(pointQueryGeometry: Array<Point>, maxPitchScaleFactor: number, has3DLayer: boolean) {
+        const tileResults: TileResult[] = [];
 
         const transform = this.transform;
         if (!transform) return tileResults;

diff --git a/src/source/tile.ts b/src/source/tile.ts
@@ -33,6 +33,7 @@ import type Point from '@mapbox/point-geometry';
 import {mat4} from 'gl-matrix';
 import type {VectorTileLayer} from '@mapbox/vector-tile';
 import {ExpiryData} from '../util/ajax';
+import {QueryRenderedFeaturesOptionsStrict} from './query_features';
 
 /**
  * The tile's state, can be:
@@ -285,11 +286,7 @@ export class Tile {
         queryGeometry: Array<Point>,
         cameraQueryGeometry: Array<Point>,
         scale: number,
-        params: {
-            filter: FilterSpecification;
-            layers: Array<string>;
-            availableImages: Array<string>;
-        },
+        params: Pick<QueryRenderedFeaturesOptionsStrict, 'filter' | 'layers' | 'availableImages'> | undefined,
         transform: IReadonlyTransform,
         maxPitchScaleFactor: number,
         pixelPosMatrix: mat4

diff --git a/src/style/style.test.ts b/src/style/style.test.ts
@@ -16,6 +16,7 @@ import {StubMap, sleep} from '../util/test/util';
 import {RTLPluginLoadedEventName} from '../source/rtl_text_plugin_status';
 import {MessageType} from '../util/actor_messages';
 import {MercatorTransform} from '../geo/projection/mercator_transform';
+import {Tile} from '../source/tile';
 
 function createStyleJSON(properties?): StyleSpecification {
     return extend({
@@ -35,7 +36,7 @@ function createSource() {
     } as any as SourceSpecification;
 }
 
-function createGeoJSONSource() {
+function createGeoJSONSource(): GeoJSONSourceSpecification {
     return {
         'type': 'geojson',
         'data': {
@@ -1962,7 +1963,7 @@ describe('Style#setFilter', () => {
         style.loadJSON({
             version: 8,
             sources: {
-                geojson: createGeoJSONSource() as GeoJSONSourceSpecification
+                geojson: createGeoJSONSource()
             },
             layers: [
                 {id: 'symbol', type: 'symbol', source: 'geojson', filter: ['==', 'id', 0]}
@@ -2089,7 +2090,7 @@ describe('Style#setLayerZoomRange', () => {
         style.loadJSON({
             'version': 8,
             'sources': {
-                'geojson': createGeoJSONSource() as GeoJSONSourceSpecification
+                'geojson': createGeoJSONSource()
             },
             'layers': [{
                 'id': 'symbol',
@@ -2182,8 +2183,8 @@ describe('Style#getLayersOrder', () => {
 
 describe('Style#queryRenderedFeatures', () => {
 
-    let style;
-    let transform;
+    let style: Style;
+    let transform: MercatorTransform;
 
     beforeEach(() => new Promise<void>(callback => {
         style = new Style(getStubMap());
@@ -2221,7 +2222,7 @@ describe('Style#queryRenderedFeatures', () => {
 
             if (params.layers) {
                 for (const l in features) {
-                    if (params.layers.indexOf(l) < 0) {
+                    if (!params.layers.has(l)) {
                         delete features[l];
                     }
                 }
@@ -2282,10 +2283,11 @@ describe('Style#queryRenderedFeatures', () => {
         style.on('style.load', () => {
             style.sourceCaches.mapLibre.tilesIn = () => {
                 return [{
-                    tile: {queryRenderedFeatures: queryMapLibreFeatures},
+                    tile: {queryRenderedFeatures: queryMapLibreFeatures} as unknown as Tile,
                     tileID: new OverscaledTileID(0, 0, 0, 0, 0),
                     queryGeometry: [],
-                    scale: 1
+                    scale: 1,
+                    cameraQueryGeometry: []
                 }];
             };
             style.sourceCaches.other.tilesIn = () => {
@@ -2316,14 +2318,20 @@ describe('Style#queryRenderedFeatures', () => {
         expect(results).toHaveLength(2);
     });
 
+    test('filters by `layers` option as a Set', () => {
+        const results = style.queryRenderedFeatures([{x: 0, y: 0}], {layers: new Set(['land'])}, transform);
+        expect(results).toHaveLength(2);
+    });
+
     test('checks type of `layers` option', () => {
         let errors = 0;
         jest.spyOn(style, 'fire').mockImplementation((event) => {
-            if (event['error'] && event['error'].message.includes('parameters.layers must be an Array.')) {
+            if (event['error'] && event['error'].message.includes('parameters.layers must be an Array')) {
                 errors++;
             }
+            return style;
         });
-        style.queryRenderedFeatures([{x: 0, y: 0}], {layers: 'string'}, transform);
+        style.queryRenderedFeatures([{x: 0, y: 0}], {layers: 'string' as any}, transform);
         expect(errors).toBe(1);
     });
 
@@ -2347,19 +2355,21 @@ describe('Style#queryRenderedFeatures', () => {
     });
 
     test('include multiple layers', () => {
-        const results = style.queryRenderedFeatures([{x: 0, y: 0}], {layers: ['land', 'landref']}, transform);
+        const results = style.queryRenderedFeatures([{x: 0, y: 0}], {layers: new Set(['land', 'landref'])}, transform);
         expect(results).toHaveLength(3);
     });
 
     test('does not query sources not implicated by `layers` parameter', () => {
-        style.sourceCaches.mapLibre.queryRenderedFeatures = () => { expect(true).toBe(false); };
+        style.sourceCaches.mapLibre.map.queryRenderedFeatures = jest.fn();
         style.queryRenderedFeatures([{x: 0, y: 0}], {layers: ['land--other']}, transform);
+        expect(style.sourceCaches.mapLibre.map.queryRenderedFeatures).not.toHaveBeenCalled();
     });
 
     test('fires an error if layer included in params does not exist on the style', () => {
         let errors = 0;
         jest.spyOn(style, 'fire').mockImplementation((event) => {
             if (event['error'] && event['error'].message.includes('does not exist in the map\'s style and cannot be queried for features.')) errors++;
+            return style;
         });
         const results = style.queryRenderedFeatures([{x: 0, y: 0}], {layers: ['merp']}, transform);
         expect(errors).toBe(1);
@@ -2415,7 +2425,7 @@ describe('Style#query*Features', () => {
     // These tests only cover filter validation. Most tests for these methods
     // live in the integration tests.
 
-    let style;
+    let style: Style;
     let onError;
     let transform;
 
@@ -2444,12 +2454,12 @@ describe('Style#query*Features', () => {
     }));
 
     test('querySourceFeatures emits an error on incorrect filter', () => {
-        expect(style.querySourceFeatures([10, 100], {filter: 7}, transform)).toEqual([]);
+        expect(style.querySourceFeatures([10, 100] as any, {filter: 7} as any)).toEqual([]);
         expect(onError.mock.calls[0][0].error.message).toMatch(/querySourceFeatures\.filter/);
     });
 
     test('queryRenderedFeatures emits an error on incorrect filter', () => {
-        expect(style.queryRenderedFeatures([{x: 0, y: 0}], {filter: 7}, transform)).toEqual([]);
+        expect(style.queryRenderedFeatures([{x: 0, y: 0}], {filter: 7 as any}, transform)).toEqual([]);
         expect(onError.mock.calls[0][0].error.message).toMatch(/queryRenderedFeatures\.filter/);
     });
 
@@ -2459,17 +2469,19 @@ describe('Style#query*Features', () => {
             if (event['error']) {
                 errors++;
             }
+            return style;
         });
-        style.queryRenderedFeatures([{x: 0, y: 0}], {filter: 'invalidFilter', validate: false}, transform);
+        style.queryRenderedFeatures([{x: 0, y: 0}], {filter: 'invalidFilter' as any, validate: false}, transform);
         expect(errors).toBe(0);
     });
 
     test('querySourceFeatures not raise validation errors if validation was disabled', () => {
         let errors = 0;
         jest.spyOn(style, 'fire').mockImplementation((event) => {
             if (event['error']) errors++;
+            return style;
         });
-        style.querySourceFeatures([{x: 0, y: 0}], {filter: 'invalidFilter', validate: false}, transform);
+        style.querySourceFeatures([{x: 0, y: 0}] as any, {filter: 'invalidFilter' as any, validate: false});
         expect(errors).toBe(0);
     });