Re-run limit queries without limit if they need to be re-filled (#3657)

* Re-run limit queries without limit if they need to be re-filled

This changes the query engine to also drop the limit on limit queries that need to be refilled. It also cleans up some of the Target.getLimit() logic that leads to repetitive code.

* Fix test

* Debugging tests via CI is not fun

* Add test

* Maybe test fix? Wish I could run it locally

* Next try

* Fix logic

* Format

Author: schmidt-sebastian

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

@@ -14,6 +14,7 @@
 
 package com.google.firebase.firestore;
 
+import static com.google.firebase.firestore.core.Query.LimitType.LIMIT_TO_LAST;
 import static com.google.firebase.firestore.util.Assert.fail;
 import static com.google.firebase.firestore.util.Assert.hardAssert;
 import static com.google.firebase.firestore.util.Preconditions.checkNotNull;
@@ -1216,7 +1217,7 @@ private ListenerRegistration addSnapshotListenerInternal(
   }
 
   private void validateHasExplicitOrderByForLimitToLast() {
-    if (query.hasLimitToLast() && query.getExplicitOrderBy().isEmpty()) {
+    if (query.getLimitType().equals(LIMIT_TO_LAST) && query.getExplicitOrderBy().isEmpty()) {
       throw new IllegalStateException(
           "limitToLast() queries require specifying at least one orderBy() clause");
     }

firebase-firestore/src/main/java/com/google/firebase/firestore/core/Query.java

@@ -147,34 +147,16 @@ public List<Filter> getFilters() {
     return filters;
   }
 
-  /**
-   * The maximum number of results to return. If there is no limit on the query, then this will
-   * cause an assertion failure.
-   */
-  public long getLimitToFirst() {
-    hardAssert(hasLimitToFirst(), "Called getLimitToFirst when no limit was set");
-    return limit;
-  }
-
-  public boolean hasLimitToFirst() {
-    return limitType == LimitType.LIMIT_TO_FIRST && limit != Target.NO_LIMIT;
-  }
-
-  /**
-   * The maximum number of last-matching results to return. If there is no limit on the query, then
-   * this will cause an assertion failure.
-   */
-  public long getLimitToLast() {
-    hardAssert(hasLimitToLast(), "Called getLimitToLast when no limit was set");
+  /** The maximum number of results to return or {@link Target#NO_LIMIT} if there is no limit. */
+  public long getLimit() {
     return limit;
   }
 
-  public boolean hasLimitToLast() {
-    return limitType == LimitType.LIMIT_TO_LAST && limit != Target.NO_LIMIT;
+  public boolean hasLimit() {
+    return limit != Target.NO_LIMIT;
   }
 
   public LimitType getLimitType() {
-    hardAssert(hasLimitToLast() || hasLimitToFirst(), "Called getLimitType when no limit was set");
     return limitType;
   }
 

firebase-firestore/src/main/java/com/google/firebase/firestore/core/View.java

@@ -14,6 +14,8 @@
 
 package com.google.firebase.firestore.core;
 
+import static com.google.firebase.firestore.core.Query.LimitType.LIMIT_TO_FIRST;
+import static com.google.firebase.firestore.core.Query.LimitType.LIMIT_TO_LAST;
 import static com.google.firebase.firestore.util.Assert.hardAssert;
 import static com.google.firebase.firestore.util.Util.compareIntegers;
 
@@ -147,11 +149,11 @@ public DocumentChanges computeDocChanges(
     // Note that this should never get used in a refill (when previousChanges is set), because there
     // will only be adds -- no deletes or updates.
     Document lastDocInLimit =
-        (query.hasLimitToFirst() && oldDocumentSet.size() == query.getLimitToFirst())
+        (query.getLimitType().equals(LIMIT_TO_FIRST) && oldDocumentSet.size() == query.getLimit())
             ? oldDocumentSet.getLastDocument()
             : null;
     Document firstDocInLimit =
-        (query.hasLimitToLast() && oldDocumentSet.size() == query.getLimitToLast())
+        (query.getLimitType().equals(LIMIT_TO_LAST) && oldDocumentSet.size() == query.getLimit())
             ? oldDocumentSet.getFirstDocument()
             : null;
 
@@ -222,11 +224,10 @@ public DocumentChanges computeDocChanges(
     }
 
     // Drop documents out to meet limitToFirst/limitToLast requirement.
-    if (query.hasLimitToFirst() || query.hasLimitToLast()) {
-      long limit = query.hasLimitToFirst() ? query.getLimitToFirst() : query.getLimitToLast();
-      for (long i = newDocumentSet.size() - limit; i > 0; --i) {
+    if (query.hasLimit()) {
+      for (long i = newDocumentSet.size() - query.getLimit(); i > 0; --i) {
         Document oldDoc =
-            query.hasLimitToFirst()
+            query.getLimitType().equals(LIMIT_TO_FIRST)
                 ? newDocumentSet.getLastDocument()
                 : newDocumentSet.getFirstDocument();
         newDocumentSet = newDocumentSet.remove(oldDoc.getKey());

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

@@ -15,7 +15,6 @@
 package com.google.firebase.firestore.local;
 
 import static com.google.firebase.firestore.util.Assert.hardAssert;
-import static com.google.firebase.firestore.util.Util.values;
 
 import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.database.collection.ImmutableSortedSet;
@@ -117,15 +116,14 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
       return null;
     }
 
-    if (indexType.equals(IndexType.PARTIAL)) {
+    if (query.hasLimit() && indexType.equals(IndexType.PARTIAL)) {
       // We cannot apply a limit for targets that are served using a partial index.
       // If a partial index will be used to serve the target, the query may return a superset of
       // documents that match the target (e.g. if the index doesn't include all the target's
       // filters), or may return the correct set of documents in the wrong order (e.g. if the index
       // doesn't include a segment for one of the orderBys). Therefore a limit should not be applied
       // in such cases.
-      query = query.limitToFirst(Target.NO_LIMIT);
-      target = query.toTarget();
+      return performQueryUsingIndex(query.limitToFirst(Target.NO_LIMIT));
     }
 
     List<DocumentKey> keys = indexManager.getDocumentsMatchingTarget(target);
@@ -136,12 +134,15 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
     IndexOffset offset = indexManager.getMinOffset(target);
 
     ImmutableSortedSet<Document> previousResults = applyQuery(query, indexedDocuments);
-    if ((query.hasLimitToFirst() || query.hasLimitToLast())
-        && needsRefill(query.getLimitType(), keys.size(), previousResults, offset.getReadTime())) {
-      return null;
+    if (needsRefill(query, keys.size(), previousResults, offset.getReadTime())) {
+      // A limit query whose boundaries change due to local edits can be re-run against the cache
+      // by excluding the limit. This ensures that all documents that match the query's filters are
+      // included in the result set. The SDK can then apply the limit once all local edits are
+      // incorporated.
+      return performQueryUsingIndex(query.limitToFirst(Target.NO_LIMIT));
     }
 
-    return appendRemainingResults(values(indexedDocuments), query, offset);
+    return appendRemainingResults(previousResults, query, offset);
   }
 
   /**
@@ -167,12 +168,7 @@ && needsRefill(query.getLimitType(), keys.size(), previousResults, offset.getRea
         localDocumentsView.getDocuments(remoteKeys);
     ImmutableSortedSet<Document> previousResults = applyQuery(query, documents);
 
-    if ((query.hasLimitToFirst() || query.hasLimitToLast())
-        && needsRefill(
-            query.getLimitType(),
-            remoteKeys.size(),
-            previousResults,
-            lastLimboFreeSnapshotVersion)) {
+    if (needsRefill(query, remoteKeys.size(), previousResults, lastLimboFreeSnapshotVersion)) {
       return null;
     }
 
@@ -211,7 +207,7 @@ private ImmutableSortedSet<Document> applyQuery(
    * Determines if a limit query needs to be refilled from cache, making it ineligible for
    * index-free execution.
    *
-   * @param limitType The type of limit query for refill calculation.
+   * @param query The query.
    * @param expectedDocumentCount The number of documents keys that matched the query at the last
    *     snapshot.
    * @param sortedPreviousResults The documents that match the query based on the previous result,
@@ -221,12 +217,17 @@ private ImmutableSortedSet<Document> applyQuery(
    *     synchronized.
    */
   private boolean needsRefill(
-      Query.LimitType limitType,
+      Query query,
       int expectedDocumentCount,
       ImmutableSortedSet<Document> sortedPreviousResults,
       SnapshotVersion limboFreeSnapshotVersion) {
-    // The query needs to be refilled if a previously matching document no longer matches.
+    if (!query.hasLimit()) {
+      // Queries without limits do not need to be refilled.
+      return false;
+    }
+
     if (expectedDocumentCount != sortedPreviousResults.size()) {
+      // The query needs to be refilled if a previously matching document no longer matches.
       return true;
     }
 
@@ -237,7 +238,7 @@ private boolean needsRefill(
     // 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 documentAtLimitEdge =
-        limitType == Query.LimitType.LIMIT_TO_FIRST
+        query.getLimitType() == Query.LimitType.LIMIT_TO_FIRST
             ? sortedPreviousResults.getMaxEntry()
             : sortedPreviousResults.getMinEntry();
     if (documentAtLimitEdge == null) {

firebase-firestore/src/test/java/com/google/firebase/firestore/bundle/BundleSerializerTest.java

@@ -712,30 +712,26 @@ private void assertDecodesNamedQuery(String json, Query query) throws JSONExcept
             + json
             + ",\n"
             + "    limitType: '"
-            + (query.hasLimitToLast() ? "LAST" : "FIRST")
+            + (query.getLimitType().equals(Query.LimitType.LIMIT_TO_LAST) ? "LAST" : "FIRST")
             + "'\n"
             + "   },\n"
             + " readTime: '2020-01-01T00:00:01.000000001Z'\n"
             + "}";
     NamedQuery actualNamedQuery = serializer.decodeNamedQuery(new JSONObject(queryJson));
 
-    long limit =
-        query.hasLimitToFirst()
-            ? query.getLimitToFirst()
-            : (query.hasLimitToLast() ? query.getLimitToLast() : Target.NO_LIMIT);
     Target target =
         new Target(
             query.getPath(),
             query.getCollectionGroup(),
             query.getFilters(),
             query.getExplicitOrderBy(),
-            limit,
+            query.getLimit(),
             query.getStartAt(),
             query.getEndAt());
     BundledQuery bundledQuery =
         new BundledQuery(
             target,
-            query.hasLimitToLast()
+            query.getLimitType().equals(Query.LimitType.LIMIT_TO_LAST)
                 ? Query.LimitType.LIMIT_TO_LAST
                 : Query.LimitType.LIMIT_TO_FIRST);
     NamedQuery expectedNamedQuery =

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

@@ -213,7 +213,7 @@ public void testUsesPartiallyIndexedOverlaysWhenAvailable() {
   }
 
   @Test
-  public void testDoesNotUseIndexForLimitQueryWhenIndexIsOutdated() {
+  public void testDoesNotUseLimitWhenIndexIsOutdated() {
     FieldIndex index =
         fieldIndex("coll", 0, FieldIndex.INITIAL_STATE, "count", FieldIndex.Segment.Kind.ASCENDING);
     configureFieldIndexes(singletonList(index));
@@ -234,10 +234,10 @@ public void testDoesNotUseIndexForLimitQueryWhenIndexIsOutdated() {
     writeMutation(deleteMutation("coll/b"));
 
     executeQuery(query);
-    // The query engine first reads the documents by key and then discards the results, which means
-    // that we read both by key and by collection.
-    assertRemoteDocumentsRead(/* byKey= */ 2, /* byCollection= */ 3);
-    assertOverlaysRead(/* byKey= */ 2, /* byCollection= */ 1);
+
+    // The query engine first reads the documents by key and then re-runs the query without limit.
+    assertRemoteDocumentsRead(/* byKey= */ 5, /* byCollection= */ 0);
+    assertOverlaysRead(/* byKey= */ 5, /* byCollection= */ 1);
     assertQueryReturned("coll/a", "coll/c");
   }
 

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

@@ -23,6 +23,8 @@
 import static com.google.firebase.firestore.testutil.TestUtil.fieldIndex;
 import static com.google.firebase.firestore.testutil.TestUtil.filter;
 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.patchMutation;
 import static com.google.firebase.firestore.testutil.TestUtil.query;
 import static com.google.firebase.firestore.testutil.TestUtil.setMutation;
 import static org.junit.Assert.assertEquals;
@@ -46,7 +48,7 @@ Persistence getPersistence() {
   }
 
   @Test
-  public void combinesIndexedWithNonIndexedResults() throws Exception {
+  public void testCombinesIndexedWithNonIndexedResults() throws Exception {
     MutableDocument doc1 = doc("coll/a", 1, map("foo", true));
     MutableDocument doc2 = doc("coll/b", 2, map("foo", true));
     MutableDocument doc3 = doc("coll/c", 3, map("foo", true));
@@ -70,7 +72,7 @@ public void combinesIndexedWithNonIndexedResults() throws Exception {
   }
 
   @Test
-  public void usesPartialIndexForLimitQueries() throws Exception {
+  public void testUsesPartialIndexForLimitQueries() throws Exception {
     MutableDocument doc1 = doc("coll/1", 1, map("a", 1, "b", 0));
     MutableDocument doc2 = doc("coll/2", 1, map("a", 1, "b", 1));
     MutableDocument doc3 = doc("coll/3", 1, map("a", 1, "b", 2));
@@ -87,4 +89,23 @@ public void usesPartialIndexForLimitQueries() throws Exception {
     DocumentSet result = expectOptimizedCollectionScan(() -> runQuery(query, SnapshotVersion.NONE));
     assertEquals(docSet(query.comparator(), doc2), result);
   }
+
+  @Test
+  public void testRefillsIndexedLimitQueries() throws Exception {
+    MutableDocument doc1 = doc("coll/1", 1, map("a", 1));
+    MutableDocument doc2 = doc("coll/2", 1, map("a", 2));
+    MutableDocument doc3 = doc("coll/3", 1, map("a", 3));
+    MutableDocument doc4 = doc("coll/4", 1, map("a", 4));
+    addDocument(doc1, doc2, doc3, doc4);
+
+    indexManager.addFieldIndex(fieldIndex("coll", "a", Kind.ASCENDING));
+    indexManager.updateIndexEntries(docMap(doc1, doc2, doc3, doc4));
+    indexManager.updateCollectionGroup("coll", IndexOffset.fromDocument(doc4));
+
+    addMutation(patchMutation("coll/3", map("a", 5)));
+
+    Query query = query("coll").orderBy(orderBy("a")).limitToFirst(3);
+    DocumentSet result = expectOptimizedCollectionScan(() -> runQuery(query, SnapshotVersion.NONE));
+    assertEquals(docSet(query.comparator(), doc1, doc2, doc4), result);
+  }
 }