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

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