Optimize query execution by deferring query evaluation (#3201)

Author: schmidt-sebastian

firebase-firestore/src/main/java/com/google/firebase/firestore/local/LocalDocumentsView.java

@@ -254,26 +254,21 @@ private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingCollection
   /** Queries the remote documents and overlays by doing a full collection scan. */
   private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingCollectionQuery(
       Query query, IndexOffset offset) {
-    ImmutableSortedMap<DocumentKey, MutableDocument> remoteDocuments =
-        remoteDocumentCache.getAllDocumentsMatchingQuery(query, offset);
+    Map<DocumentKey, MutableDocument> remoteDocuments =
+        remoteDocumentCache.getAll(query.getPath(), offset);
     Map<DocumentKey, Mutation> overlays = documentOverlayCache.getOverlays(query.getPath(), -1);
 
-    // As documents might match the query because of their overlay we need to include all documents
-    // in the result.
-    Set<DocumentKey> missingDocuments = new HashSet<>();
+    // As documents might match the query because of their overlay we need to include documents
+    // for all overlays in the initial document set.
     for (Map.Entry<DocumentKey, Mutation> entry : overlays.entrySet()) {
       if (!remoteDocuments.containsKey(entry.getKey())) {
-        missingDocuments.add(entry.getKey());
+        remoteDocuments.put(entry.getKey(), MutableDocument.newInvalidDocument(entry.getKey()));
       }
     }
-    for (Map.Entry<DocumentKey, MutableDocument> entry :
-        remoteDocumentCache.getAll(missingDocuments).entrySet()) {
-      remoteDocuments = remoteDocuments.insert(entry.getKey(), entry.getValue());
-    }
 
     // Apply the overlays and match against the query.
     ImmutableSortedMap<DocumentKey, Document> results = emptyDocumentMap();
-    for (Map.Entry<DocumentKey, MutableDocument> docEntry : remoteDocuments) {
+    for (Map.Entry<DocumentKey, MutableDocument> docEntry : remoteDocuments.entrySet()) {
       Mutation overlay = overlays.get(docEntry.getKey());
       if (overlay != null) {
         overlay.applyToLocalView(docEntry.getValue(), null, Timestamp.now());

firebase-firestore/src/main/java/com/google/firebase/firestore/local/MemoryRemoteDocumentCache.java

@@ -14,12 +14,10 @@
 
 package com.google.firebase.firestore.local;
 
-import static com.google.firebase.firestore.model.DocumentCollections.emptyMutableDocumentMap;
 import static com.google.firebase.firestore.util.Assert.hardAssert;
 
 import androidx.annotation.NonNull;
 import com.google.firebase.database.collection.ImmutableSortedMap;
-import com.google.firebase.firestore.core.Query;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.FieldIndex.IndexOffset;
 import com.google.firebase.firestore.model.MutableDocument;
@@ -89,41 +87,35 @@ public Map<DocumentKey, MutableDocument> getAll(
   }
 
   @Override
-  public ImmutableSortedMap<DocumentKey, MutableDocument> getAllDocumentsMatchingQuery(
-      Query query, IndexOffset offset) {
-    hardAssert(
-        !query.isCollectionGroupQuery(),
-        "CollectionGroup queries should be handled in LocalDocumentsView");
-    ImmutableSortedMap<DocumentKey, MutableDocument> result = emptyMutableDocumentMap();
+  public Map<DocumentKey, MutableDocument> getAll(ResourcePath collection, IndexOffset offset) {
+    Map<DocumentKey, MutableDocument> result = new HashMap<>();
 
     // Documents are ordered by key, so we can use a prefix scan to narrow down the documents
     // we need to match the query against.
-    ResourcePath queryPath = query.getPath();
-    DocumentKey prefix = DocumentKey.fromPath(queryPath.append(""));
+    DocumentKey prefix = DocumentKey.fromPath(collection.append(""));
     Iterator<Map.Entry<DocumentKey, MutableDocument>> iterator = docs.iteratorFrom(prefix);
 
     while (iterator.hasNext()) {
       Map.Entry<DocumentKey, MutableDocument> entry = iterator.next();
+      MutableDocument doc = entry.getValue();
 
       DocumentKey key = entry.getKey();
-      if (!queryPath.isPrefixOf(key.getPath())) {
+      if (!collection.isPrefixOf(key.getPath())) {
+        // We are now scanning the next collection. Abort.
         break;
       }
 
-      MutableDocument doc = entry.getValue();
-      if (!doc.isFoundDocument()) {
+      if (key.getPath().length() > collection.length() + 1) {
+        // Exclude entries from subcollections.
         continue;
       }
 
       if (IndexOffset.fromDocument(doc).compareTo(offset) <= 0) {
+        // The document sorts before the offset.
         continue;
       }
 
-      if (!query.matches(doc)) {
-        continue;
-      }
-
-      result = result.insert(doc.getKey(), doc.clone());
+      result.put(doc.getKey(), doc.clone());
     }
 
     return result;

firebase-firestore/src/main/java/com/google/firebase/firestore/local/RemoteDocumentCache.java

@@ -14,11 +14,10 @@
 
 package com.google.firebase.firestore.local;
 
-import com.google.firebase.database.collection.ImmutableSortedMap;
-import com.google.firebase.firestore.core.Query;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.FieldIndex.IndexOffset;
 import com.google.firebase.firestore.model.MutableDocument;
+import com.google.firebase.firestore.model.ResourcePath;
 import com.google.firebase.firestore.model.SnapshotVersion;
 import java.util.Map;
 
@@ -72,24 +71,18 @@ interface RemoteDocumentCache {
    * @param collectionGroup The collection group to scan.
    * @param offset The offset to start the scan at.
    * @param limit The maximum number of results to return.
-   * @return A map with next set of documents.
+   * @return A newly created map with next set of documents.
    */
   Map<DocumentKey, MutableDocument> getAll(String collectionGroup, IndexOffset offset, int limit);
 
   /**
-   * Executes a query against the cached Document entries
+   * Returns the documents from the provided collection.
    *
-   * <p>Implementations may return extra documents if convenient. The results should be re-filtered
-   * by the consumer before presenting them to the user.
-   *
-   * <p>Cached entries for non-existing documents have no bearing on query results.
-   *
-   * @param query The query to match documents against.
+   * @param collection The collection to read.
    * @param offset The read time and document key to start scanning at (exclusive).
-   * @return The set of matching documents.
+   * @return A newly created map with the set of documents in the collection.
    */
-  ImmutableSortedMap<DocumentKey, MutableDocument> getAllDocumentsMatchingQuery(
-      Query query, IndexOffset offset);
+  Map<DocumentKey, MutableDocument> getAll(ResourcePath collection, IndexOffset offset);
 
   /** Returns the latest read time of any document in the cache. */
   SnapshotVersion getLatestReadTime();

firebase-firestore/src/main/java/com/google/firebase/firestore/local/SQLiteRemoteDocumentCache.java

@@ -21,9 +21,6 @@
 
 import androidx.annotation.VisibleForTesting;
 import com.google.firebase.Timestamp;
-import com.google.firebase.database.collection.ImmutableSortedMap;
-import com.google.firebase.firestore.core.Query;
-import com.google.firebase.firestore.model.DocumentCollections;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.FieldIndex.IndexOffset;
 import com.google.firebase.firestore.model.MutableDocument;
@@ -234,23 +231,8 @@ private Map<DocumentKey, MutableDocument> getAll(
   }
 
   @Override
-  public ImmutableSortedMap<DocumentKey, MutableDocument> getAllDocumentsMatchingQuery(
-      final Query query, IndexOffset offset) {
-    hardAssert(
-        !query.isCollectionGroupQuery(),
-        "CollectionGroup queries should be handled in LocalDocumentsView");
-
-    Map<DocumentKey, MutableDocument> allDocuments =
-        getAll(Collections.singletonList(query.getPath()), offset, Integer.MAX_VALUE);
-
-    ImmutableSortedMap<DocumentKey, MutableDocument> matchingDocuments =
-        DocumentCollections.emptyMutableDocumentMap();
-    for (MutableDocument document : allDocuments.values()) {
-      if (document.isFoundDocument() && query.matches(document)) {
-        matchingDocuments = matchingDocuments.insert(document.getKey(), document);
-      }
-    }
-    return matchingDocuments;
+  public Map<DocumentKey, MutableDocument> getAll(ResourcePath collection, IndexOffset offset) {
+    return getAll(Collections.singletonList(collection), offset, Integer.MAX_VALUE);
   }
 
   @Override

firebase-firestore/src/main/java/com/google/firebase/firestore/model/MutableDocument.java

@@ -62,22 +62,25 @@ private enum DocumentState {
   private final DocumentKey key;
   private DocumentType documentType;
   private SnapshotVersion version;
-  private SnapshotVersion readTime = SnapshotVersion.NONE;
+  private SnapshotVersion readTime;
   private ObjectValue value;
   private DocumentState documentState;
 
   private MutableDocument(DocumentKey key) {
     this.key = key;
+    this.readTime = SnapshotVersion.NONE;
   }
 
   private MutableDocument(
       DocumentKey key,
       DocumentType documentType,
       SnapshotVersion version,
+      SnapshotVersion readTime,
       ObjectValue value,
       DocumentState documentState) {
     this.key = key;
     this.version = version;
+    this.readTime = readTime;
     this.documentType = documentType;
     this.documentState = documentState;
     this.value = value;
@@ -92,6 +95,7 @@ public static MutableDocument newInvalidDocument(DocumentKey documentKey) {
         documentKey,
         DocumentType.INVALID,
         SnapshotVersion.NONE,
+        SnapshotVersion.NONE,
         new ObjectValue(),
         DocumentState.SYNCED);
   }
@@ -226,7 +230,7 @@ public boolean isUnknownDocument() {
   @Override
   @NonNull
   public MutableDocument clone() {
-    return new MutableDocument(key, documentType, version, value.clone(), documentState);
+    return new MutableDocument(key, documentType, version, readTime, value.clone(), documentState);
   }
 
   @Override
@@ -238,6 +242,8 @@ public boolean equals(Object o) {
 
     if (!key.equals(document.key)) return false;
     if (!version.equals(document.version)) return false;
+    // TODO(mrschmidt): Include readTime (requires a lot of test updates)
+    // if (!readTime.equals(document.readTime)) return false;
     if (!documentType.equals(document.documentType)) return false;
     if (!documentState.equals(document.documentState)) return false;
     return value.equals(document.value);

firebase-firestore/src/test/java/com/google/firebase/firestore/local/CountingQueryEngine.java

@@ -23,6 +23,7 @@
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.FieldIndex.IndexOffset;
 import com.google.firebase.firestore.model.MutableDocument;
+import com.google.firebase.firestore.model.ResourcePath;
 import com.google.firebase.firestore.model.SnapshotVersion;
 import com.google.firebase.firestore.model.mutation.Mutation;
 import com.google.firebase.firestore.model.mutation.MutationBatch;
@@ -37,19 +38,19 @@
 class CountingQueryEngine implements QueryEngine {
   private final QueryEngine queryEngine;
 
-  private final int[] mutationsReadByQuery = new int[] {0};
+  private final int[] mutationsReadByCollection = new int[] {0};
   private final int[] mutationsReadByKey = new int[] {0};
-  private final int[] documentsReadByQuery = new int[] {0};
+  private final int[] documentsReadByCollection = new int[] {0};
   private final int[] documentsReadByKey = new int[] {0};
 
   CountingQueryEngine(QueryEngine queryEngine) {
     this.queryEngine = queryEngine;
   }
 
   void resetCounts() {
-    mutationsReadByQuery[0] = 0;
+    mutationsReadByCollection[0] = 0;
     mutationsReadByKey[0] = 0;
-    documentsReadByQuery[0] = 0;
+    documentsReadByCollection[0] = 0;
     documentsReadByKey[0] = 0;
   }
 
@@ -83,11 +84,11 @@ QueryEngine getSubject() {
   }
 
   /**
-   * Returns the number of documents returned by the RemoteDocumentCache's
-   * `getDocumentsMatchingQuery()` API (since the last call to `resetCounts()`)
+   * Returns the number of documents returned by the RemoteDocumentCache's `getAll()` API (since the
+   * last call to `resetCounts()`)
    */
-  int getDocumentsReadByQuery() {
-    return documentsReadByQuery[0];
+  int getDocumentsReadByCollection() {
+    return documentsReadByCollection[0];
   }
 
   /**
@@ -102,8 +103,8 @@ int getDocumentsReadByKey() {
    * Returns the number of mutations returned by the MutationQueue's
    * `getAllMutationBatchesAffectingQuery()` API (since the last call to `resetCounts()`)
    */
-  int getMutationsReadByQuery() {
-    return mutationsReadByQuery[0];
+  int getMutationsReadByCollection() {
+    return mutationsReadByCollection[0];
   }
 
   /**
@@ -153,16 +154,15 @@ public Map<DocumentKey, MutableDocument> getAll(Iterable<DocumentKey> documentKe
       public Map<DocumentKey, MutableDocument> getAll(
           String collectionGroup, IndexOffset offset, int limit) {
         Map<DocumentKey, MutableDocument> result = subject.getAll(collectionGroup, offset, limit);
-        documentsReadByQuery[0] += result.size();
+        documentsReadByCollection[0] += result.size();
+
         return result;
       }
 
       @Override
-      public ImmutableSortedMap<DocumentKey, MutableDocument> getAllDocumentsMatchingQuery(
-          Query query, IndexOffset offset) {
-        ImmutableSortedMap<DocumentKey, MutableDocument> result =
-            subject.getAllDocumentsMatchingQuery(query, offset);
-        documentsReadByQuery[0] += result.size();
+      public Map<DocumentKey, MutableDocument> getAll(ResourcePath collection, IndexOffset offset) {
+        Map<DocumentKey, MutableDocument> result = subject.getAll(collection, offset);
+        documentsReadByCollection[0] += result.size();
         return result;
       }
 
@@ -250,7 +250,7 @@ public List<MutationBatch> getAllMutationBatchesAffectingDocumentKeys(
       @Override
       public List<MutationBatch> getAllMutationBatchesAffectingQuery(Query query) {
         List<MutationBatch> result = subject.getAllMutationBatchesAffectingQuery(query);
-        mutationsReadByQuery[0] += result.size();
+        mutationsReadByCollection[0] += result.size();
         return result;
       }
 

firebase-firestore/src/test/java/com/google/firebase/firestore/local/LocalStoreTestCase.java

@@ -321,18 +321,21 @@ private void assertHasNamedQuery(NamedQuery expectedNamedQuery) {
    * Asserts the expected numbers of mutations read by the MutationQueue since the last call to
    * `resetPersistenceStats()`.
    */
-  private void assertMutationsRead(int byKey, int byQuery) {
-    assertEquals("Mutations read (by query)", byQuery, queryEngine.getMutationsReadByQuery());
+  private void assertMutationsRead(int byKey, int byCollection) {
+    assertEquals(
+        "Mutations read (by collection)", byCollection, queryEngine.getMutationsReadByCollection());
     assertEquals("Mutations read (by key)", byKey, queryEngine.getMutationsReadByKey());
   }
 
   /**
    * Asserts the expected numbers of documents read by the RemoteDocumentCache since the last call
    * to `resetPersistenceStats()`.
    */
-  private void assertRemoteDocumentsRead(int byKey, int byQuery) {
+  private void assertRemoteDocumentsRead(int byKey, int byCollection) {
     assertEquals(
-        "Remote documents read (by query)", byQuery, queryEngine.getDocumentsReadByQuery());
+        "Remote documents read (by collection)",
+        byCollection,
+        queryEngine.getDocumentsReadByCollection());
     assertEquals("Remote documents read (by key)", byKey, queryEngine.getDocumentsReadByKey());
   }
 
@@ -976,10 +979,9 @@ public void testReadsAllDocumentsForInitialCollectionQueries() {
     resetPersistenceStats();
 
     localStore.executeQuery(query, /* usePreviousResults= */ true);
-
-    assertRemoteDocumentsRead(/* byKey= */ 0, /* byQuery= */ 2);
+    assertRemoteDocumentsRead(/* byKey= */ 0, /* byCollection= */ 2);
     // No mutations are read because only overlay is needed.
-    assertMutationsRead(/* byKey= */ 0, /* byQuery= */ 0);
+    assertMutationsRead(/* byKey= */ 0, /* byCollection= */ 0);
   }
 
   @Test
@@ -1104,7 +1106,7 @@ public void testUsesTargetMappingToExecuteQueries() {
     // Execute the query, but note that we read all existing documents from the RemoteDocumentCache
     // since we do not yet have target mapping.
     executeQuery(query);
-    assertRemoteDocumentsRead(/* byKey= */ 0, /* byQuery= */ 2);
+    assertRemoteDocumentsRead(/* byKey= */ 0, /* byCollection= */ 3);
 
     // Issue a RemoteEvent to persist the target mapping.
     applyRemoteEvent(
@@ -1118,7 +1120,7 @@ public void testUsesTargetMappingToExecuteQueries() {
     // Execute the query again, this time verifying that we only read the two documents that match
     // the query.
     executeQuery(query);
-    assertRemoteDocumentsRead(/* byKey= */ 2, /* byQuery= */ 0);
+    assertRemoteDocumentsRead(/* byKey= */ 2, /* byCollection= */ 0);
     assertQueryReturned("foo/a", "foo/b");
   }