Index-Free: Only exclude limit queries if last document changed (#737)

Author: schmidt-sebastian

firebase-database-collection/src/main/java/com/google/firebase/database/collection/ImmutableSortedMap.java

@@ -56,14 +56,6 @@ public abstract class ImmutableSortedMap<K, V> implements Iterable<Map.Entry<K,
 
   public abstract Comparator<K> getComparator();
 
-  public ImmutableSortedMap<K, V> insertAll(ImmutableSortedMap<K, V> source) {
-    ImmutableSortedMap<K, V> result = this;
-    for (Map.Entry<K, V> entry : source) {
-      result = result.insert(entry.getKey(), entry.getValue());
-    }
-    return result;
-  }
-
   @Override
   @SuppressWarnings("unchecked")
   public boolean equals(Object o) {

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

@@ -21,10 +21,10 @@
 import com.google.firebase.database.collection.ImmutableSortedSet;
 import com.google.firebase.firestore.core.Query;
 import com.google.firebase.firestore.model.Document;
-import com.google.firebase.firestore.model.DocumentCollections;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.MaybeDocument;
 import com.google.firebase.firestore.model.SnapshotVersion;
+import java.util.Collections;
 import java.util.Map;
 
 /**
@@ -76,74 +76,90 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
       return executeFullCollectionScan(query);
     }
 
-    ImmutableSortedMap<DocumentKey, Document> result =
-        executeIndexFreeQuery(query, queryData, remoteKeys);
+    ImmutableSortedSet<Document> previousResults = getSortedPreviousResults(query, remoteKeys);
 
-    return result != null ? result : executeFullCollectionScan(query);
+    if (query.hasLimit()
+        && needsRefill(previousResults, remoteKeys, queryData.getLastLimboFreeSnapshotVersion())) {
+      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());
+    for (Document result : previousResults) {
+      updatedResults = updatedResults.insert(result.getKey(), result);
+    }
+
+    return 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, sorted by
+   * the query's comparator.
    */
-  private @Nullable ImmutableSortedMap<DocumentKey, Document> executeIndexFreeQuery(
-      Query query, QueryData queryData, ImmutableSortedSet<DocumentKey> remoteKeys) {
+  private ImmutableSortedSet<Document> getSortedPreviousResults(
+      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.
+    // Sort the documents and re-apply the query filter since previously matching documents do not
+    // necessarily still match the query.
+    ImmutableSortedSet<Document> results =
+        new ImmutableSortedSet<>(Collections.emptyList(), query.comparator());
     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;
+        results = results.insert(doc);
       }
     }
+    return results;
+  }
 
-    // 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());
+  /**
+   * Determines if a limit query needs to be refilled from cache, making it ineligible for
+   * index-free execution.
+   *
+   * @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 snapshot.
+   * @param limboFreeSnapshotVersion The version of the snapshot when the query was last
+   *     synchronized.
+   */
+  private boolean needsRefill(
+      ImmutableSortedSet<Document> sortedPreviousResults,
+      ImmutableSortedSet<DocumentKey> remoteKeys,
+      SnapshotVersion limboFreeSnapshotVersion) {
+    // The query needs to be refilled if a previously matching document no longer matches.
+    if (remoteKeys.size() != sortedPreviousResults.size()) {
+      return true;
+    }
 
-    results = results.insertAll(updatedResults);
+    // We don't need to find a better match from cache if no documents matched the query.
+    if (sortedPreviousResults.isEmpty()) {
+      return false;
+    }
 
-    return results;
+    // Limit queries are not eligible for index-free query execution if there is a potential that an
+    // older document from cache now sorts before a document that was previously part of the limit.
+    // This, however, can only happen if the last document of the limit sorts lower than it did when
+    // the query was last synchronized. If a document that is not the limit boundary sorts
+    // differently, the boundary of the limit itself did not change and documents from cache will
+    // continue to be "rejected" by this boundary. Therefore, we can ignore any modifications that
+    // don't affect the last document.
+    Document lastDocumentInLimit = sortedPreviousResults.getMaxEntry();
+    return lastDocumentInLimit.hasPendingWrites()
+        || lastDocumentInLimit.getVersion().compareTo(limboFreeSnapshotVersion) > 0;
   }
 
   @Override
   public void handleDocumentChange(MaybeDocument oldDocument, MaybeDocument newDocument) {
     // No indexes to update.
   }
 
-  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;
-      }
-    }
-
-    return false;
-  }
-
   private ImmutableSortedMap<DocumentKey, Document> executeFullCollectionScan(Query query) {
     return localDocumentsView.getDocumentsMatchingQuery(query, SnapshotVersion.NONE);
   }

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,