Port QueryEngine

packages/firestore/src/core/query.ts

@@ -155,7 +155,7 @@ export function asCollectionQueryAtPath(
  * Returns true if this query does not specify any query constraints that
  * could remove results.
  */
-export function matchesAllDocuments(query: Query): boolean {
+export function queryMatchesAllDocuments(query: Query): boolean {
   return (
     query.filters.length === 0 &&
     query.limit === null &&
@@ -393,7 +393,7 @@ export function queryWithAddedOrderBy(query: Query, orderBy: OrderBy): Query {
 
 export function queryWithLimit(
   query: Query,
-  limit: number,
+  limit: number | null,
   limitType: LimitType
 ): Query {
   return new QueryImpl(

packages/firestore/src/local/index_manager.ts

@@ -155,4 +155,13 @@ export interface IndexManager {
     transaction: PersistenceTransaction,
     documents: DocumentMap
   ): PersistencePromise<void>;
+
+  /**
+   * Iterates over all field indexes that are used to serve the given target,
+   * and returns the minimum offset of them all.
+   */
+  getMinOffset(
+    transaction: PersistenceTransaction,
+    target: Target
+  ): PersistencePromise<IndexOffset>;
 }

packages/firestore/src/local/indexeddb_index_manager.ts

@@ -972,6 +972,16 @@ export class IndexedDbIndexManager implements IndexManager {
     }
     return ranges;
   }
+
+  getMinOffset(
+    transaction: PersistenceTransaction,
+    target: Target
+  ): PersistencePromise<IndexOffset> {
+    // TODO(orqueries): Get the minimum offset for all subqueries
+    return this.getFieldIndex(transaction, target).next(index =>
+      index ? index.indexState.offset : IndexOffset.min()
+    );
+  }
 }
 
 /**

packages/firestore/src/local/indexeddb_schema.ts

@@ -29,7 +29,7 @@ import { EncodedResourcePath } from './encoded_resource_path';
 import { DbTimestampKey } from './indexeddb_sentinels';
 
 // TODO(indexing): Remove this constant
-const INDEXING_ENABLED = false;
+export const INDEXING_ENABLED = false;
 
 export const INDEXING_SCHEMA_VERSION = 14;
 

packages/firestore/src/local/local_documents_view.ts

@@ -52,7 +52,7 @@ export class LocalDocumentsView {
   constructor(
     readonly remoteDocumentCache: RemoteDocumentCache,
     readonly mutationQueue: MutationQueue,
-    readonly indexManager: IndexManager
+    private readonly indexManager: IndexManager
   ) {}
 
   /**

packages/firestore/src/local/local_store_impl.ts

@@ -203,7 +203,7 @@ class LocalStoreImpl implements LocalStore {
       this.indexManager
     );
     this.remoteDocuments.setIndexManager(this.indexManager);
-    this.queryEngine.setLocalDocumentsView(this.localDocuments);
+    this.queryEngine.initialize(this.localDocuments, this.indexManager);
   }
 
   collectGarbage(garbageCollector: LruGarbageCollector): Promise<LruResults> {

packages/firestore/src/local/memory_index_manager.ts

@@ -97,6 +97,13 @@ export class MemoryIndexManager implements IndexManager {
     return PersistencePromise.resolve<string | null>(null);
   }
 
+  getMinOffset(
+    transaction: PersistenceTransaction,
+    target: Target
+  ): PersistencePromise<IndexOffset> {
+    return PersistencePromise.resolve(IndexOffset.min());
+  }
+
   updateCollectionGroup(
     transaction: PersistenceTransaction,
     collectionGroup: string,

packages/firestore/src/local/query_engine.ts

@@ -19,33 +19,62 @@ import {
   hasLimitToFirst,
   hasLimitToLast,
   LimitType,
-  matchesAllDocuments,
   newQueryComparator,
   Query,
   queryMatches,
+  queryMatchesAllDocuments,
+  queryToTarget,
+  queryWithLimit,
   stringifyQuery
 } from '../core/query';
 import { SnapshotVersion } from '../core/snapshot_version';
-import { DocumentKeySet, DocumentMap } from '../model/collections';
+import {
+  documentKeySet,
+  DocumentKeySet,
+  DocumentMap
+} from '../model/collections';
 import { Document } from '../model/document';
 import {
   IndexOffset,
   INITIAL_LARGEST_BATCH_ID,
   newIndexOffsetSuccessorFromReadTime
 } from '../model/field_index';
 import { debugAssert } from '../util/assert';
-import { getLogLevel, LogLevel, logDebug } from '../util/log';
+import { getLogLevel, logDebug, LogLevel } from '../util/log';
+import { ValueIterable, values } from '../util/misc';
 import { SortedSet } from '../util/sorted_set';
 
+import { IndexManager, IndexType } from './index_manager';
+import { INDEXING_ENABLED } from './indexeddb_schema';
 import { LocalDocumentsView } from './local_documents_view';
 import { PersistencePromise } from './persistence_promise';
 import { PersistenceTransaction } from './persistence_transaction';
 
 /**
- * A query engine that takes advantage of the target document mapping in the
- * QueryCache. Query execution is optimized by only reading the documents that
- * previously matched a query plus any documents that were edited after the
- * query was last listened to.
+ * The Firestore query engine.
+ *
+ * Firestore queries can be executed in three modes. The Query Engine determines
+ * what mode to use based on what data is persisted. The mode only determines
+ * the runtime complexity of the query - the result set is equivalent across all
+ * implementations.
+ *
+ * The Query engine will use indexed-based execution if a user has configured
+ * any index that can be used to execute query (via `setIndexConfiguration()`).
+ * Otherwise, the engine will try to optimize the query by re-using a previously
+ * persisted query result. If that is not possible, the query will be executed
+ * via a full collection scan.
+ *
+ * Index-based execution is the default when available. The query engine
+ * supports partial indexed execution and merges the result from the index
+ * lookup with documents that have not yet been indexed. The index evaluation
+ * matches the backend's format and as such, the SDK can use indexing for all
+ * queries that the backend supports.
+ *
+ * If no index exists, the query engine tries to take advantage of the target
+ * document mapping in the TargetCache. These mappings exists for all queries
+ * that have been synced with the backend at least once and allow the query
+ * engine to only read documents that previously matched a query plus any
+ * documents that were edited after the query was last listened to.
  *
  * There are some cases when this optimization is not guaranteed to produce
  * the same results as full collection scans. In these cases, query
@@ -60,11 +89,18 @@ import { PersistenceTransaction } from './persistence_transaction';
  * - Queries that have never been CURRENT or free of limbo documents.
  */
 export class QueryEngine {
-  private localDocumentsView: LocalDocumentsView | undefined;
+  private localDocumentsView!: LocalDocumentsView;
+  private indexManager!: IndexManager;
+  private initialized = false;
 
   /** Sets the document view to query against. */
-  setLocalDocumentsView(localDocuments: LocalDocumentsView): void {
+  initialize(
+    localDocuments: LocalDocumentsView,
+    indexManager: IndexManager
+  ): void {
     this.localDocumentsView = localDocuments;
+    this.indexManager = indexManager;
+    this.initialized = true;
   }
 
   /** Returns all local documents matching the specified query. */
@@ -74,15 +110,122 @@ export class QueryEngine {
     lastLimboFreeSnapshotVersion: SnapshotVersion,
     remoteKeys: DocumentKeySet
   ): PersistencePromise<DocumentMap> {
-    debugAssert(
-      this.localDocumentsView !== undefined,
-      'setLocalDocumentsView() not called'
-    );
+    debugAssert(this.initialized, 'initialize() not called');
+
+    return this.performQueryUsingIndex(transaction, query)
+      .next(result =>
+        result
+          ? result
+          : this.performQueryUsingRemoteKeys(
+              transaction,
+              query,
+              remoteKeys,
+              lastLimboFreeSnapshotVersion
+            )
+      )
+      .next(result =>
+        result ? result : this.executeFullCollectionScan(transaction, query)
+      );
+  }
+
+  /**
+   * Performs an indexed query that evaluates the query based on a collection's
+   * persisted index values. Returns `null` if an index is not available.
+   */
+  private performQueryUsingIndex(
+    transaction: PersistenceTransaction,
+    query: Query
+  ): PersistencePromise<DocumentMap | null> {
+    if (!INDEXING_ENABLED) {
+      return PersistencePromise.resolve<DocumentMap | null>(null);
+    }
+
+    if (queryMatchesAllDocuments(query)) {
+      // Don't use indexes for queries that can be executed by scanning the
+      // collection.
+      return PersistencePromise.resolve<DocumentMap | null>(null);
+    }
+
+    let target = queryToTarget(query);
+    return this.indexManager
+      .getIndexType(transaction, target)
+      .next(indexType => {
+        if (indexType === IndexType.NONE) {
+          // The target cannot be served from any index.
+          return null;
+        }
+
+        if (indexType === IndexType.PARTIAL) {
+          // We cannot apply a limit for targets that are served using a partial
+          // index. If a partial index will be used to serve the target, the
+          // query may return a superset of documents that match the target
+          // (e.g. if the index doesn't include all the target's filters), or
+          // may return the correct set of documents in the wrong order (e.g. if
+          // the index doesn't include a segment for one of the orderBys).
+          // Therefore, a limit should not be applied in such cases.
+          query = queryWithLimit(query, null, LimitType.First);
+          target = queryToTarget(query);
+        }
 
+        return this.indexManager
+          .getDocumentsMatchingTarget(transaction, target)
+          .next(keys => {
+            debugAssert(
+              !!keys,
+              'Index manager must return results for partial and full indexes.'
+            );
+            const sortedKeys = documentKeySet(...keys);
+            return this.localDocumentsView
+              .getDocuments(transaction, sortedKeys)
+              .next(indexedDocuments => {
+                return this.indexManager
+                  .getMinOffset(transaction, target)
+                  .next(offset => {
+                    const previousResults = this.applyQuery(
+                      query,
+                      indexedDocuments
+                    );
+
+                    if (
+                      (hasLimitToFirst(query) || hasLimitToLast(query)) &&
+                      this.needsRefill(
+                        query.limitType,
+                        previousResults,
+                        sortedKeys,
+                        offset.readTime
+                      )
+                    ) {
+                      return PersistencePromise.resolve<DocumentMap | null>(
+                        null
+                      );
+                    }
+
+                    return this.appendRemainingResults(
+                      transaction,
+                      values(indexedDocuments),
+                      query,
+                      offset
+                    ) as PersistencePromise<DocumentMap | null>;
+                  });
+              });
+          });
+      });
+  }
+
+  /**
+   * Performs a query based on the target's persisted query mapping. Returns
+   * `null` if the mapping is not available or cannot be used.
+   */
+  private performQueryUsingRemoteKeys(
+    transaction: PersistenceTransaction,
+    query: Query,
+    remoteKeys: DocumentKeySet,
+    lastLimboFreeSnapshotVersion: SnapshotVersion
+  ): PersistencePromise<DocumentMap> {
     // Queries that match all documents don't benefit from using
     // key-based lookups. It is more efficient to scan all documents in a
     // collection, rather than to perform individual lookups.
-    if (matchesAllDocuments(query)) {
+    if (queryMatchesAllDocuments(query)) {
       return this.executeFullCollectionScan(transaction, query);
     }
 
@@ -119,22 +262,15 @@ export class QueryEngine {
 
         // Retrieve all results for documents that were updated since the last
         // limbo-document free remote snapshot.
-        return this.localDocumentsView!.getDocumentsMatchingQuery(
+        return this.appendRemainingResults(
           transaction,
+          previousResults,
           query,
           newIndexOffsetSuccessorFromReadTime(
             lastLimboFreeSnapshotVersion,
             INITIAL_LARGEST_BATCH_ID
           )
-        ).next(updatedResults => {
-          // We merge `previousResults` into `updateResults`, since
-          // `updateResults` is already a DocumentMap. If a document is
-          // contained in both lists, then its contents are the same.
-          previousResults.forEach(doc => {
-            updatedResults = updatedResults.insert(doc.key, doc);
-          });
-          return updatedResults;
-        });
+        );
       }
     );
   }
@@ -159,6 +295,7 @@ export class QueryEngine {
    * Determines if a limit query needs to be refilled from cache, making it
    * ineligible for index-free execution.
    *
+   * @param limitType The limit type used by the query.
    * @param sortedPreviousResults - The documents that matched the query when it
    * was last synchronized, sorted by the query's comparator.
    * @param remoteKeys - The document keys that matched the query at the last
@@ -218,4 +355,26 @@ export class QueryEngine {
       IndexOffset.min()
     );
   }
+
+  /**
+   * Combines the results from an indexed execution with the remaining documents
+   * that have not yet been indexed.
+   */
+  private appendRemainingResults(
+    transaction: PersistenceTransaction,
+    indexedResults: ValueIterable<Document>,
+    query: Query,
+    offset: IndexOffset
+  ): PersistencePromise<DocumentMap> {
+    // Retrieve all results for documents that were updated since the offset.
+    return this.localDocumentsView
+      .getDocumentsMatchingQuery(transaction, query, offset)
+      .next(remainingResults => {
+        // Merge with existing results
+        indexedResults.forEach(d => {
+          remainingResults = remainingResults.insert(d.key, d);
+        });
+        return remainingResults;
+      });
+  }
 }

packages/firestore/src/model/collections.ts

@@ -46,8 +46,12 @@ export type DocumentMap = SortedMap<DocumentKey, Document>;
 const EMPTY_DOCUMENT_MAP = new SortedMap<DocumentKey, Document>(
   DocumentKey.comparator
 );
-export function documentMap(): DocumentMap {
-  return EMPTY_DOCUMENT_MAP;
+export function documentMap(...docs: Document[]): DocumentMap {
+  let map = EMPTY_DOCUMENT_MAP;
+  for (const doc of docs) {
+    map = map.insert(doc.key, doc);
+  }
+  return map;
 }
 
 export type OverlayMap = ObjectMap<DocumentKey, Overlay>;