Optimize query execution by deferring query evaluation

Author: schmidt-sebastian

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

@@ -322,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());
@@ -358,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.
@@ -378,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());
@@ -398,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/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;
@@ -90,26 +88,11 @@ public Map<DocumentKey, MutableDocument> getAll(
     List<ResourcePath> collectionParents = indexManager.getCollectionParents(collectionGroup);
     Set<MutableDocument> allDocuments =
         new TreeSet<>((l, r) -> IndexOffset.fromDocument(l).compareTo(IndexOffset.fromDocument(r)));
-    Map<DocumentKey, MutableDocument> matchingDocuments = new HashMap<>();
-
     for (ResourcePath collectionParent : collectionParents) {
-      ResourcePath documentParent = collectionParent.append(collectionGroup);
-      Iterator<Map.Entry<DocumentKey, MutableDocument>> iterator =
-          docs.iteratorFrom(DocumentKey.fromPath(documentParent.append("")));
-      while (iterator.hasNext()) {
-        Map.Entry<DocumentKey, MutableDocument> entry = iterator.next();
-        DocumentKey key = entry.getKey();
-
-        if (!documentParent.isPrefixOf(key.getPath())) {
-          break;
-        }
-
-        if (IndexOffset.fromDocument(entry.getValue()).compareTo(offset) > 0) {
-          allDocuments.add(entry.getValue());
-        }
-      }
+      allDocuments.addAll(getAll(collectionParent.append(collectionGroup), offset).values());
     }
 
+    Map<DocumentKey, MutableDocument> matchingDocuments = new HashMap<>();
     Iterator<MutableDocument> it = allDocuments.iterator();
     while (it.hasNext() && matchingDocuments.size() < count) {
       MutableDocument document = it.next();
@@ -120,41 +103,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;
 
@@ -80,19 +79,13 @@ interface RemoteDocumentCache {
   Map<DocumentKey, MutableDocument> getAll(String collectionGroup, IndexOffset offset, int count);
 
   /**
-   * 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.
    */
-  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/SQLiteIndexManager.java

@@ -24,7 +24,6 @@
 import androidx.annotation.Nullable;
 import com.google.firebase.Timestamp;
 import com.google.firebase.database.collection.ImmutableSortedMap;
-import com.google.firebase.database.collection.LLRBNode;
 import com.google.firebase.firestore.auth.User;
 import com.google.firebase.firestore.core.Bound;
 import com.google.firebase.firestore.core.FieldFilter;

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

@@ -20,9 +20,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;
@@ -229,23 +226,9 @@ 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(
+      final ResourcePath collection, IndexOffset offset) {
+    return getAll(Collections.singletonList(collection), offset, Integer.MAX_VALUE);
   }
 
   @Override

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

@@ -49,7 +49,7 @@ class SQLiteSchema {
    * The version of the schema. Increase this by one for each migration added to runMigrations
    * below.
    */
-  static final int VERSION = 13;
+  static final int VERSION = 14;
 
   static final int OVERLAY_SUPPORT_VERSION = VERSION + 1;
 
@@ -178,6 +178,10 @@ void runSchemaUpgrades(int fromVersion, int toVersion) {
       ensurePathLength();
     }
 
+    if (fromVersion < 14 && toVersion >= 14) {
+      ensureReadTime();
+    }
+
     /*
      * Adding a new migration? READ THIS FIRST!
      *
@@ -658,6 +662,12 @@ private void ensurePathLength() {
     } while (resultsRemaining[0]);
   }
 
+  /** Initialize the remote_document's read_time column with 0 values if they are not set. */
+  private void ensureReadTime() {
+    db.execSQL(
+        "UPDATE remote_documents SET read_time_seconds = 0, read_time_nanos = 0 WHERE read_time_seconds IS NULL");
+  }
+
   private void createBundleCache() {
     ifTablesDontExist(
         new String[] {"bundles", "named_queries"},

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);