Index-Free: Only exclude limit queries if last document changed

schmidt-sebastian · schmidt-sebastian · commit 6cb8e855b4bb · 2019-08-26T10:18:58.000-07:00
diff --git a/firebase-firestore/src/main/java/com/google/firebase/firestore/local/IndexFreeQueryEngine.java b/firebase-firestore/src/main/java/com/google/firebase/firestore/local/IndexFreeQueryEngine.java
@@ -25,6 +25,8 @@
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.MaybeDocument;
 import com.google.firebase.firestore.model.SnapshotVersion;
+import java.util.Comparator;
+import java.util.Iterator;
 import java.util.Map;
 
 /**
@@ -76,72 +78,92 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
       return executeFullCollectionScan(query);
     }
 
-    ImmutableSortedMap<DocumentKey, Document> result =
-        executeIndexFreeQuery(query, queryData, remoteKeys);
+    ImmutableSortedMap<DocumentKey, Document> previousResults =
+        getPreviousResults(query, remoteKeys);
 
-    return result != null ? result : executeFullCollectionScan(query);
+    if (query.hasLimit()
+        && needsRefill(
+            query, queryData.getLastLimboFreeSnapshotVersion(), previousResults, remoteKeys)) {
+      return executeFullCollectionScan(query);
+    }
+
+    // Retrieve all results for documents that were updated since the last limbo-document free
+    // remote snapshot.
+    ImmutableSortedMap<DocumentKey, Document> updatedResults =
+        localDocumentsView.getDocumentsMatchingQuery(
+            query, queryData.getLastLimboFreeSnapshotVersion());
+    return previousResults.insertAll(updatedResults);
   }
 
   /**
-   * Attempts index-free query execution. Returns the set of query results on success, otherwise
-   * returns null.
+   * Returns the documents for the specified remote keys if they still match the query.
    */
-  private @Nullable ImmutableSortedMap<DocumentKey, Document> executeIndexFreeQuery(
-      Query query, QueryData queryData, ImmutableSortedSet<DocumentKey> remoteKeys) {
+  private ImmutableSortedMap<DocumentKey, Document> getPreviousResults(
+      Query query, ImmutableSortedSet<DocumentKey> remoteKeys) {
     // Fetch the documents that matched the query at the last snapshot.
     ImmutableSortedMap<DocumentKey, MaybeDocument> previousResults =
         localDocumentsView.getDocuments(remoteKeys);
 
-    // Limit queries are not eligible for index-free query execution if any part of the result was
-    // modified after we received the last query snapshot. This makes sure that we re-populate the
-    // view with older documents that may sort before the modified document.
-    if (query.hasLimit()
-        && containsUpdatesSinceSnapshotVersion(previousResults, queryData.getSnapshotVersion())) {
-      return null;
-    }
-
-    ImmutableSortedMap<DocumentKey, Document> results = DocumentCollections.emptyDocumentMap();
-
     // Re-apply the query filter since previously matching documents do not necessarily still
     // match the query.
+    ImmutableSortedMap<DocumentKey, Document> results = DocumentCollections.emptyDocumentMap();
     for (Map.Entry<DocumentKey, MaybeDocument> entry : previousResults) {
       MaybeDocument maybeDoc = entry.getValue();
       if (maybeDoc instanceof Document && query.matches((Document) maybeDoc)) {
         Document doc = (Document) maybeDoc;
         results = results.insert(entry.getKey(), doc);
-      } else if (query.hasLimit()) {
-        // Limit queries with documents that no longer match need to be re-filled from cache.
-        return null;
       }
     }
 
-    // Retrieve all results for documents that were updated since the last limbo-document free
-    // remote snapshot.
-    ImmutableSortedMap<DocumentKey, Document> updatedResults =
-        localDocumentsView.getDocumentsMatchingQuery(
-            query, queryData.getLastLimboFreeSnapshotVersion());
-
-    results = results.insertAll(updatedResults);
-
     return results;
   }
 
-  @Override
-  public void handleDocumentChange(MaybeDocument oldDocument, MaybeDocument newDocument) {
-    // No indexes to update.
-  }
+  /**
+   * Determines if a limit query needs to be refilled from cache, making it ineligible for
+   * index-free execution.
+   */
+  private boolean needsRefill(
+      Query query,
+      SnapshotVersion limboFreeSnapshotVersion,
+      ImmutableSortedMap<DocumentKey, Document> previousResults,
+      ImmutableSortedSet<DocumentKey> remoteKeys) {
+    hardAssert(query.hasLimit(), "Only limit queries should be refilled");
+
+    // The query doesn't need to be refilled if there were never any documents that matched the
+    // query constraint.
+    if (remoteKeys.isEmpty()) {
+      return false;
+    }
 
-  private boolean containsUpdatesSinceSnapshotVersion(
-      ImmutableSortedMap<DocumentKey, MaybeDocument> previousResults,
-      SnapshotVersion sinceSnapshotVersion) {
-    for (Map.Entry<DocumentKey, MaybeDocument> doc : previousResults) {
-      if (doc.getValue().hasPendingWrites()
-          || doc.getValue().getVersion().compareTo(sinceSnapshotVersion) > 0) {
-        return true;
+    // The query needs to be refilled if a previously matching document no longer matches.
+    if (remoteKeys.size() != previousResults.size()) {
+      return true;
+    }
+
+    // Limit queries are not eligible for index-free query execution if an older document from cache
+    // may sort before a document that was previously part of the limit. This can happen if the last
+    // document of the limit could sort lower than it did when the query was last synchronized. To
+    // verify this, we use the query's comparator to determine the last document that fit the
+    // limit. If the last document was edited after the query was last synchronized, there is a
+    // chance that another document from cache sorts higher, in which case we have to perform a full
+    // cache scan.
+    Comparator<Document> comparator = query.comparator();
+    Iterator<Map.Entry<DocumentKey, Document>> it = previousResults.iterator();
+    Document lastDocumentInLimit = it.next().getValue();
+    while (it.hasNext()) {
+      Document doc = it.next().getValue();
+      if (comparator.compare(lastDocumentInLimit, doc) < 0) {
+        lastDocumentInLimit = doc;
       }
     }
 
-    return false;
+    return lastDocumentInLimit.hasPendingWrites()
+        || lastDocumentInLimit.getVersion().compareTo(limboFreeSnapshotVersion) > 0;
+  }
+
+  @Override
+  public void handleDocumentChange(MaybeDocument oldDocument, MaybeDocument newDocument) {
+    // No indexes to update.
   }
 
   private ImmutableSortedMap<DocumentKey, Document> executeFullCollectionScan(Query query) {
diff --git a/firebase-firestore/src/test/java/com/google/firebase/firestore/local/IndexFreeQueryEngineTest.java b/firebase-firestore/src/test/java/com/google/firebase/firestore/local/IndexFreeQueryEngineTest.java
@@ -17,6 +17,7 @@
 import static com.google.firebase.firestore.testutil.TestUtil.doc;
 import static com.google.firebase.firestore.testutil.TestUtil.docSet;
 import static com.google.firebase.firestore.testutil.TestUtil.filter;
+import static com.google.firebase.firestore.testutil.TestUtil.key;
 import static com.google.firebase.firestore.testutil.TestUtil.map;
 import static com.google.firebase.firestore.testutil.TestUtil.orderBy;
 import static com.google.firebase.firestore.testutil.TestUtil.query;
@@ -61,8 +62,6 @@ public class IndexFreeQueryEngineTest {
       doc("coll/a", 11, map("matches", true, "order", 1), Document.DocumentState.SYNCED);
   private static final Document MATCHING_DOC_B =
       doc("coll/b", 1, map("matches", true, "order", 2), Document.DocumentState.SYNCED);
-  private static final Document NON_MATCHING_DOC_B =
-      doc("coll/b", 1, map("matches", false, "order", 2), Document.DocumentState.SYNCED);
   private static final Document UPDATED_MATCHING_DOC_B =
       doc("coll/b", 11, map("matches", true, "order", 2), Document.DocumentState.SYNCED);
 
@@ -102,13 +101,13 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
   }
 
   /** Adds the provided documents to the query target mapping. */
-  private void persistQueryMapping(Document... docs) {
+  private void persistQueryMapping(DocumentKey... documentKeys) {
     persistence.runTransaction(
         "persistQueryMapping",
         () -> {
           ImmutableSortedSet<DocumentKey> remoteKeys = DocumentKey.emptyKeySet();
-          for (Document doc : docs) {
-            remoteKeys = remoteKeys.insert(doc.getKey());
+          for (DocumentKey documentKey : documentKeys) {
+            remoteKeys = remoteKeys.insert(documentKey);
           }
           queryCache.addMatchingKeys(remoteKeys, TEST_TARGET_ID);
         });
@@ -162,7 +161,7 @@ public void usesTargetMappingForInitialView() throws Exception {
     QueryData queryData = queryData(query, /* hasLimboFreeSnapshot= */ true);
 
     addDocument(MATCHING_DOC_A, MATCHING_DOC_B);
-    persistQueryMapping(MATCHING_DOC_A, MATCHING_DOC_B);
+    persistQueryMapping(MATCHING_DOC_A.getKey(), MATCHING_DOC_B.getKey());
 
     DocumentSet docs = expectIndexFreeQuery(() -> runQuery(query, queryData));
     assertEquals(docSet(query.comparator(), MATCHING_DOC_A, MATCHING_DOC_B), docs);
@@ -174,7 +173,7 @@ public void filtersNonMatchingInitialResults() throws Exception {
     QueryData queryData = queryData(query, /* hasLimboFreeSnapshot= */ true);
 
     addDocument(MATCHING_DOC_A, MATCHING_DOC_B);
-    persistQueryMapping(MATCHING_DOC_A, MATCHING_DOC_B);
+    persistQueryMapping(MATCHING_DOC_A.getKey(), MATCHING_DOC_B.getKey());
 
     // Add a mutated document that is not yet part of query's set of remote keys.
     addDocument(PENDING_NON_MATCHING_DOC_A);
@@ -189,7 +188,7 @@ public void includesChangesSinceInitialResults() throws Exception {
     QueryData originalQueryData = queryData(query, /* hasLimboFreeSnapshot= */ true);
 
     addDocument(MATCHING_DOC_A, MATCHING_DOC_B);
-    persistQueryMapping(MATCHING_DOC_A, MATCHING_DOC_B);
+    persistQueryMapping(MATCHING_DOC_A.getKey(), MATCHING_DOC_B.getKey());
 
     DocumentSet docs = expectIndexFreeQuery(() -> runQuery(query, originalQueryData));
     assertEquals(docSet(query.comparator(), MATCHING_DOC_A, MATCHING_DOC_B), docs);
@@ -225,7 +224,7 @@ public void doesNotUseInitialResultsForLimitQueryWithDocumentRemoval() throws Ex
     // While the backend would never add DocA to the set of remote keys, this allows us to easily
     // simulate what would happen when a document no longer matches due to an out-of-band update.
     addDocument(NON_MATCHING_DOC_A);
-    persistQueryMapping(NON_MATCHING_DOC_A);
+    persistQueryMapping(NON_MATCHING_DOC_A.getKey());
 
     addDocument(MATCHING_DOC_B);
 
@@ -235,7 +234,8 @@ public void doesNotUseInitialResultsForLimitQueryWithDocumentRemoval() throws Ex
   }
 
   @Test
-  public void doesNotUseInitialResultsForLimitQueryWithPendingWrite() throws Exception {
+  public void doesNotUseInitialResultsForLimitQueryWhenLastDocumentHasPendingWrite()
+      throws Exception {
     Query query =
         query("coll")
             .filter(filter("matches", "==", true))
@@ -245,7 +245,7 @@ public void doesNotUseInitialResultsForLimitQueryWithPendingWrite() throws Excep
     // Add a query mapping for a document that matches, but that sorts below another document due to
     // a pending write.
     addDocument(PENDING_MATCHING_DOC_A);
-    persistQueryMapping(PENDING_MATCHING_DOC_A);
+    persistQueryMapping(PENDING_MATCHING_DOC_A.getKey());
 
     QueryData queryData = queryData(query, /* hasLimboFreeSnapshot= */ true);
 
@@ -256,7 +256,7 @@ public void doesNotUseInitialResultsForLimitQueryWithPendingWrite() throws Excep
   }
 
   @Test
-  public void doesNotUseInitialResultsForLimitQueryWithDocumentThatHasBeenUpdatedOutOfBand()
+  public void doesNotUseInitialResultsForLimitQueryWhenLastDocumentHasBeenUpdatedOutOfBand()
       throws Exception {
     Query query =
         query("coll")
@@ -267,7 +267,7 @@ public void doesNotUseInitialResultsForLimitQueryWithDocumentThatHasBeenUpdatedO
     // Add a query mapping for a document that matches, but that sorts below another document based
     // due to an update that the SDK received after the query's snapshot was persisted.
     addDocument(UDPATED_DOC_A);
-    persistQueryMapping(UDPATED_DOC_A);
+    persistQueryMapping(UDPATED_DOC_A.getKey());
 
     QueryData queryData = queryData(query, /* hasLimboFreeSnapshot= */ true);
 
@@ -277,6 +277,30 @@ public void doesNotUseInitialResultsForLimitQueryWithDocumentThatHasBeenUpdatedO
     assertEquals(docSet(query.comparator(), MATCHING_DOC_B), docs);
   }
 
+  @Test
+  public void limitQueriesUseInitialResultsIfLastDocumentInLimitIsUnchanged() throws Exception {
+    Query query = query("coll").orderBy(orderBy("order")).limit(2);
+
+    addDocument(doc("coll/a", 1, map("order", 1)));
+    addDocument(doc("coll/b", 1, map("order", 3)));
+    persistQueryMapping(key("coll/a"), key("coll/b"));
+    QueryData queryData = queryData(query, /* hasLimboFreeSnapshot= */ true);
+
+    // Update "coll/a" but make sure it still sorts before "coll/b"
+    addDocument(doc("coll/a", 1, map("order", 2), Document.DocumentState.LOCAL_MUTATIONS));
+
+    // Since the last document in the limit didn't change (and hence we know that all documents
+    // written prior to query execution still sort after "coll/b"), we should use an Index-Free
+    // query.
+    DocumentSet docs = expectIndexFreeQuery(() -> runQuery(query, queryData));
+    assertEquals(
+        docSet(
+            query.comparator(),
+            doc("coll/a", 1, map("order", 2), Document.DocumentState.LOCAL_MUTATIONS),
+            doc("coll/b", 1, map("order", 3))),
+        docs);
+  }
+
   private QueryData queryData(Query query, boolean hasLimboFreeSnapshot) {
     return new QueryData(
         query,
