Optimize query execution by deferring query evaluation

firebase-firestore/CHANGELOG.md

@@ -3,6 +3,10 @@ by opting into a release at
 [go/firebase-android-release](http:go/firebase-android-release) (Googlers only).
 
 # Unreleased
+- [changed] Improved performance for queries with collections that contain
+  subcollections.
+
+# 24.0.0
 - [feature] Added support for Firebase AppCheck.
 
 # 23.0.4

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/ConformanceTest.java

@@ -284,6 +284,6 @@ private com.google.firebase.firestore.FieldPath formatFieldPath(String serverFor
   }
 
   private String getId(Document document) {
-    return DocumentKey.fromName(document.getName()).getPath().getLastSegment();
+    return DocumentKey.fromName(document.getName()).getDocumentId();
   }
 }

firebase-firestore/src/androidTest/java/com/google/firebase/firestore/conformance/TestCaseConverter.java

@@ -48,7 +48,7 @@ private List<Collection> convertDatabaseContents(DatastoreTestTrace.FirestoreV1A
       if (response.hasDocument()) {
         Document document = response.getDocument();
         DocumentKey documentKey = DocumentKey.fromName(document.getName());
-        String collectionId = documentKey.getPath().popLast().getLastSegment();
+        String collectionId = documentKey.getCollectionGroup();
         List<Document> documents =
             collections.computeIfAbsent(collectionId, name -> new ArrayList<>());
         documents.add(document);

firebase-firestore/src/main/java/com/google/firebase/firestore/DocumentReference.java

@@ -96,7 +96,7 @@ public FirebaseFirestore getFirestore() {
 
   @NonNull
   public String getId() {
-    return key.getPath().getLastSegment();
+    return key.getDocumentId();
   }
 
   /**
@@ -106,7 +106,7 @@ public String getId() {
    */
   @NonNull
   public CollectionReference getParent() {
-    return new CollectionReference(key.getPath().popLast(), firestore);
+    return new CollectionReference(key.getCollectionPath(), firestore);
   }
 
   /**

firebase-firestore/src/main/java/com/google/firebase/firestore/DocumentSnapshot.java

@@ -102,7 +102,7 @@ static DocumentSnapshot fromNoDocument(
   /** @return The id of the document. */
   @NonNull
   public String getId() {
-    return key.getPath().getLastSegment();
+    return key.getDocumentId();
   }
 
   /** @return The metadata for this document snapshot. */

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

@@ -19,19 +19,15 @@
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 import com.google.firebase.database.collection.ImmutableSortedMap;
-import com.google.firebase.firestore.core.Query;
 import com.google.firebase.firestore.model.Document;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.FieldIndex;
 import com.google.firebase.firestore.model.FieldIndex.IndexOffset;
-import com.google.firebase.firestore.model.ResourcePath;
 import com.google.firebase.firestore.util.AsyncQueue;
 import com.google.firebase.firestore.util.Logger;
-import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.HashSet;
-import java.util.List;
+import java.util.Iterator;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
@@ -138,22 +134,18 @@ private int writeIndexEntries(LocalDocumentsView localDocumentsView) {
   /** Writes entries for the fetched field indexes. */
   private int writeEntriesForCollectionGroup(
       LocalDocumentsView localDocumentsView, String collectionGroup, int entriesRemainingUnderCap) {
-    Query query = new Query(ResourcePath.EMPTY, collectionGroup);
+    // TODO(indexing): Support mutation batch Ids when sorting and writing indexes.
 
     // Use the earliest offset of all field indexes to query the local cache.
     IndexOffset existingOffset = getExistingOffset(indexManager.getFieldIndexes(collectionGroup));
-
-    // TODO(indexing): Use limit queries to only fetch the required number of entries.
-    // TODO(indexing): Support mutation batch Ids when sorting and writing indexes.
     ImmutableSortedMap<DocumentKey, Document> documents =
-        localDocumentsView.getDocumentsMatchingQuery(query, existingOffset);
-
-    List<Document> oldestDocuments = getOldestDocuments(documents, entriesRemainingUnderCap);
-    indexManager.updateIndexEntries(oldestDocuments);
+        localDocumentsView.getDocuments(collectionGroup, existingOffset, entriesRemainingUnderCap);
+    indexManager.updateIndexEntries(documents);
 
-    IndexOffset newOffset = getNewOffset(oldestDocuments, existingOffset);
+    IndexOffset newOffset = getNewOffset(documents, existingOffset);
     indexManager.updateCollectionGroup(collectionGroup, newOffset);
-    return oldestDocuments.size();
+
+    return documents.size();
   }
 
   /** Returns the lowest offset for the provided index group. */
@@ -168,37 +160,22 @@ private IndexOffset getExistingOffset(Collection<FieldIndex> fieldIndexes) {
     return lowestOffset == null ? IndexOffset.NONE : lowestOffset;
   }
 
-  /**
-   * Returns the offset for the index based on the newly indexed documents.
-   *
-   * @param documents a list of documents sorted by read time and key (ascending)
-   * @param currentOffset the current offset of the index group
-   */
-  private IndexOffset getNewOffset(List<Document> documents, IndexOffset currentOffset) {
-    IndexOffset latestOffset =
-        documents.isEmpty()
-            ? IndexOffset.create(remoteDocumentCache.getLatestReadTime())
-            : IndexOffset.create(
-                documents.get(documents.size() - 1).getReadTime(),
-                documents.get(documents.size() - 1).getKey());
-    // Make sure the index does not go back in time
-    latestOffset = latestOffset.compareTo(currentOffset) > 0 ? latestOffset : currentOffset;
-    return latestOffset;
-  }
-
-  /** Returns up to {@code count} documents sorted by read time and key. */
-  private List<Document> getOldestDocuments(
-      ImmutableSortedMap<DocumentKey, Document> documents, int count) {
-    List<Document> oldestDocuments = new ArrayList<>();
-    for (Map.Entry<DocumentKey, Document> entry : documents) {
-      oldestDocuments.add(entry.getValue());
+  /** Returns the offset for the index based on the newly indexed documents. */
+  private IndexOffset getNewOffset(
+      ImmutableSortedMap<DocumentKey, Document> documents, IndexOffset currentOffset) {
+    if (documents.isEmpty()) {
+      return IndexOffset.create(remoteDocumentCache.getLatestReadTime());
+    } else {
+      IndexOffset latestOffset = currentOffset;
+      Iterator<Map.Entry<DocumentKey, Document>> it = documents.iterator();
+      while (it.hasNext()) {
+        IndexOffset newOffset = IndexOffset.fromDocument(it.next().getValue());
+        if (newOffset.compareTo(latestOffset) > 0) {
+          latestOffset = newOffset;
+        }
+      }
+      return latestOffset;
     }
-    Collections.sort(
-        oldestDocuments,
-        (l, r) ->
-            IndexOffset.create(l.getReadTime(), l.getKey())
-                .compareTo(IndexOffset.create(r.getReadTime(), r.getKey())));
-    return oldestDocuments.subList(0, Math.min(count, oldestDocuments.size()));
   }
 
   @VisibleForTesting

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

@@ -15,6 +15,7 @@
 package com.google.firebase.firestore.local;
 
 import androidx.annotation.Nullable;
+import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.firestore.core.Target;
 import com.google.firebase.firestore.model.Document;
 import com.google.firebase.firestore.model.DocumentKey;
@@ -97,5 +98,5 @@ public interface IndexManager {
   void updateCollectionGroup(String collectionGroup, FieldIndex.IndexOffset offset);
 
   /** Updates the index entries for the provided documents. */
-  void updateIndexEntries(Collection<Document> documents);
+  void updateIndexEntries(ImmutableSortedMap<DocumentKey, Document> documents);
 }

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

@@ -237,10 +237,13 @@ void recalculateAndSaveOverlays(Set<DocumentKey> documentKeys) {
     recalculateAndSaveOverlays(docs);
   }
 
-  // TODO: The Querying implementation here should move 100% to the query engines.
-  // Instead, we should just provide a getCollectionDocuments() method here that return all the
-  // documents in a given collection so that query engine can do that and then filter in
-  // memory.
+  /** Gets the local view of the next {@code count} documents based on their read time. */
+  ImmutableSortedMap<DocumentKey, Document> getDocuments(
+      String collectionGroup, IndexOffset offset, int count) {
+    Map<DocumentKey, MutableDocument> docs =
+        remoteDocumentCache.getAll(collectionGroup, offset, count);
+    return getLocalViewOfDocuments(docs, new HashSet<>());
+  }
 
   /**
    * Performs a query against the local view of all documents.
@@ -319,26 +322,21 @@ private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingCollection
   /** Queries the remote documents and overlays by doing a full collection scan. */
   private ImmutableSortedMap<DocumentKey, Document>
       getDocumentsMatchingCollectionQueryFromOverlayCache(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());
@@ -355,14 +353,12 @@ private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingCollection
   /** Queries the remote documents and mutation queue, by doing a full collection scan. */
   private ImmutableSortedMap<DocumentKey, Document>
       getDocumentsMatchingCollectionQueryFromMutationQueue(Query query, IndexOffset offset) {
-    ImmutableSortedMap<DocumentKey, MutableDocument> remoteDocuments =
-        remoteDocumentCache.getAllDocumentsMatchingQuery(query, offset);
+    Map<DocumentKey, MutableDocument> remoteDocuments =
+        remoteDocumentCache.getAll(query.getPath(), offset);
 
     // TODO(indexing): We should plumb sinceReadTime through to the mutation queue
     List<MutationBatch> matchingBatches = mutationQueue.getAllMutationBatchesAffectingQuery(query);
 
-    remoteDocuments = addMissingBaseDocuments(matchingBatches, remoteDocuments);
-
     for (MutationBatch batch : matchingBatches) {
       for (Mutation mutation : batch.getMutations()) {
         // Only process documents belonging to the collection.
@@ -375,18 +371,18 @@ private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingCollection
         if (document == null) {
           // Create invalid document to apply mutations on top of
           document = MutableDocument.newInvalidDocument(key);
-          remoteDocuments = remoteDocuments.insert(key, document);
+          remoteDocuments.put(key, document);
         }
         mutation.applyToLocalView(
             document, FieldMask.fromSet(new HashSet<>()), batch.getLocalWriteTime());
         if (!document.isFoundDocument()) {
-          remoteDocuments = remoteDocuments.remove(key);
+          remoteDocuments.remove(key);
         }
       }
     }
 
     ImmutableSortedMap<DocumentKey, Document> results = emptyDocumentMap();
-    for (Map.Entry<DocumentKey, MutableDocument> docEntry : remoteDocuments) {
+    for (Map.Entry<DocumentKey, MutableDocument> docEntry : remoteDocuments.entrySet()) {
       // Finally, insert the documents that still match the query
       if (query.matches(docEntry.getValue())) {
         results = results.insert(docEntry.getKey(), docEntry.getValue());
@@ -395,35 +391,4 @@ private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingCollection
 
     return results;
   }
-
-  /**
-   * It is possible that a {@code PatchMutation} can make a document match a query, even if the
-   * version in the {@code RemoteDocumentCache} is not a match yet (waiting for server to ack). To
-   * handle this, we find all document keys affected by the {@code PatchMutation}s that are not in
-   * {@code existingDocs} yet, and back fill them via {@code remoteDocumentCache.getAll}, otherwise
-   * those {@code PatchMutation}s will be ignored because no base document can be found, and lead to
-   * missing results for the query.
-   */
-  private ImmutableSortedMap<DocumentKey, MutableDocument> addMissingBaseDocuments(
-      List<MutationBatch> matchingBatches,
-      ImmutableSortedMap<DocumentKey, MutableDocument> existingDocs) {
-    HashSet<DocumentKey> missingDocKeys = new HashSet<>();
-    for (MutationBatch batch : matchingBatches) {
-      for (Mutation mutation : batch.getMutations()) {
-        if (mutation instanceof PatchMutation && !existingDocs.containsKey(mutation.getKey())) {
-          missingDocKeys.add(mutation.getKey());
-        }
-      }
-    }
-
-    ImmutableSortedMap<DocumentKey, MutableDocument> mergedDocs = existingDocs;
-    Map<DocumentKey, MutableDocument> missingDocs = remoteDocumentCache.getAll(missingDocKeys);
-    for (Map.Entry<DocumentKey, MutableDocument> entry : missingDocs.entrySet()) {
-      if (entry.getValue().isFoundDocument()) {
-        mergedDocs = mergedDocs.insert(entry.getKey(), entry.getValue());
-      }
-    }
-
-    return mergedDocs;
-  }
 }

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

@@ -16,6 +16,7 @@
 import static com.google.firebase.firestore.util.Assert.hardAssert;
 
 import androidx.annotation.Nullable;
+import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.firestore.core.Target;
 import com.google.firebase.firestore.model.Document;
 import com.google.firebase.firestore.model.DocumentKey;
@@ -98,7 +99,7 @@ public Collection<FieldIndex> getFieldIndexes() {
   }
 
   @Override
-  public void updateIndexEntries(Collection<Document> documents) {
+  public void updateIndexEntries(ImmutableSortedMap<DocumentKey, Document> documents) {
     // Field indices are not supported with memory persistence.
   }
 

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

@@ -152,7 +152,7 @@ public MutationBatch addMutationBatch(
       batchesByDocumentKey =
           batchesByDocumentKey.insert(new DocumentReference(mutation.getKey(), batchId));
 
-      indexManager.addToCollectionParentIndex(mutation.getKey().getPath().popLast());
+      indexManager.addToCollectionParentIndex(mutation.getKey().getCollectionPath());
     }
 
     return batch;