Don't use Indexes for outdated limit queries (#3549)

Author: schmidt-sebastian

firebase-firestore/CHANGELOG.md

@@ -2,6 +2,11 @@ Android changes are not released automatically. Ensure that changes are released
 by opting into a release at 
 [go/firebase-android-release](http:go/firebase-android-release) (Googlers only).
 
+# Unreleased
+- [changed] Fixed an issue in the experimental index engine that might have
+  caused Firestore to exclude document results for limit queries with local
+  modifications.
+
 # 24.1.0
 - [feature] Added experimental support for indexed query execution. Indexes can
   be enabled by invoking `FirebaseFirestore.setIndexConfiguration()` with the

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

@@ -117,8 +117,15 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
 
     ImmutableSortedMap<DocumentKey, Document> indexedDocuments =
         localDocumentsView.getDocuments(keys);
-    return appendRemainingResults(
-        values(indexedDocuments), query, indexManager.getMinOffset(target));
+    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;
+    }
+
+    return appendRemainingResults(values(indexedDocuments), query, offset);
   }
 
   /**
@@ -146,7 +153,10 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
 
     if ((query.hasLimitToFirst() || query.hasLimitToLast())
         && needsRefill(
-            query.getLimitType(), previousResults, remoteKeys, lastLimboFreeSnapshotVersion)) {
+            query.getLimitType(),
+            remoteKeys.size(),
+            previousResults,
+            lastLimboFreeSnapshotVersion)) {
       return null;
     }
 
@@ -186,19 +196,21 @@ private ImmutableSortedSet<Document> applyQuery(
    * index-free execution.
    *
    * @param limitType The type of limit query for refill calculation.
-   * @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 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,
+   *     sorted by the query's comparator. The size of the result set may be different from
+   *     `expectedDocumentCount` if documents cease to match the query.
    * @param limboFreeSnapshotVersion The version of the snapshot when the query was last
    *     synchronized.
    */
   private boolean needsRefill(
       Query.LimitType limitType,
+      int expectedDocumentCount,
       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()) {
+    if (expectedDocumentCount != sortedPreviousResults.size()) {
       return true;
     }
 

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

@@ -72,11 +72,6 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
     return queryEngine.getDocumentsMatchingQuery(query, lastLimboFreeSnapshotVersion, remoteKeys);
   }
 
-  /** Returns the query engine that is used as the backing implementation. */
-  QueryEngine getSubject() {
-    return queryEngine;
-  }
-
   /**
    * Returns the number of documents returned by the RemoteDocumentCache's `getAll()` API (since the
    * last call to `resetCounts()`)
@@ -147,7 +142,6 @@ public Map<DocumentKey, MutableDocument> getAll(
           String collectionGroup, IndexOffset offset, int limit) {
         Map<DocumentKey, MutableDocument> result = subject.getAll(collectionGroup, offset, limit);
         documentsReadByCollection[0] += result.size();
-
         return result;
       }
 

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

@@ -16,6 +16,7 @@
 
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.firebase.firestore.testutil.TestUtil.addedRemoteEvent;
+import static com.google.firebase.firestore.testutil.TestUtil.deleteMutation;
 import static com.google.firebase.firestore.testutil.TestUtil.deletedDoc;
 import static com.google.firebase.firestore.testutil.TestUtil.doc;
 import static com.google.firebase.firestore.testutil.TestUtil.fieldIndex;
@@ -35,6 +36,7 @@
 import com.google.firebase.firestore.model.FieldIndex;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.Collections;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.robolectric.RobolectricTestRunner;
@@ -141,15 +143,15 @@ public void testDeletedDocumentRemovesIndex() {
     backfillIndexes();
 
     executeQuery(query);
-    assertRemoteDocumentsRead(/* byKey= */ 1, /* byQuery= */ 0);
+    assertRemoteDocumentsRead(/* byKey= */ 1, /* byCollection= */ 0);
     assertQueryReturned("coll/a");
 
     applyRemoteEvent(
         updateRemoteEvent(deletedDoc("coll/a", 0), singletonList(targetId), emptyList()));
 
     // No backfill needed for deleted document.
     executeQuery(query);
-    assertRemoteDocumentsRead(/* byKey= */ 0, /* byQuery= */ 0);
+    assertRemoteDocumentsRead(/* byKey= */ 0, /* byCollection= */ 0);
     assertQueryReturned();
   }
 
@@ -168,7 +170,7 @@ public void testUsesIndexes() {
     backfillIndexes();
 
     executeQuery(query);
-    assertRemoteDocumentsRead(/* byKey= */ 1, /* byQuery= */ 0);
+    assertRemoteDocumentsRead(/* byKey= */ 1, /* byCollection= */ 0);
     assertQueryReturned("coll/a");
   }
 
@@ -188,7 +190,7 @@ public void testUsesPartiallyIndexedRemoteDocumentsWhenAvailable() {
     applyRemoteEvent(addedRemoteEvent(doc("coll/b", 20, map("matches", true)), targetId));
 
     executeQuery(query);
-    assertRemoteDocumentsRead(/* byKey= */ 1, /* byQuery= */ 1);
+    assertRemoteDocumentsRead(/* byKey= */ 1, /* byCollection= */ 1);
     assertQueryReturned("coll/a", "coll/b");
   }
 
@@ -210,6 +212,61 @@ public void testUsesPartiallyIndexedOverlaysWhenAvailable() {
     assertQueryReturned("coll/a", "coll/b");
   }
 
+  @Test
+  public void testDoesNotUseIndexForLimitQueryWhenIndexIsOutdated() {
+    FieldIndex index =
+        fieldIndex("coll", 0, FieldIndex.INITIAL_STATE, "count", FieldIndex.Segment.Kind.ASCENDING);
+    configureFieldIndexes(singletonList(index));
+
+    Query query = query("coll").orderBy(orderBy("count")).limitToFirst(2);
+    int targetId = allocateQuery(query);
+
+    applyRemoteEvent(
+        addedRemoteEvent(
+            Arrays.asList(
+                doc("coll/a", 10, map("count", 1)),
+                doc("coll/b", 10, map("count", 2)),
+                doc("coll/c", 10, map("count", 3))),
+            Collections.singletonList(targetId),
+            Collections.emptyList()));
+    backfillIndexes();
+
+    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);
+    assertQueryReturned("coll/a", "coll/c");
+  }
+
+  @Test
+  public void testUsesIndexForLimitQueryWhenIndexIsUpdated() {
+    FieldIndex index =
+        fieldIndex("coll", 0, FieldIndex.INITIAL_STATE, "count", FieldIndex.Segment.Kind.ASCENDING);
+    configureFieldIndexes(singletonList(index));
+
+    Query query = query("coll").orderBy(orderBy("count")).limitToFirst(2);
+    int targetId = allocateQuery(query);
+
+    applyRemoteEvent(
+        addedRemoteEvent(
+            Arrays.asList(
+                doc("coll/a", 10, map("count", 1)),
+                doc("coll/b", 10, map("count", 2)),
+                doc("coll/c", 10, map("count", 3))),
+            Collections.singletonList(targetId),
+            Collections.emptyList()));
+    writeMutation(deleteMutation("coll/b"));
+    backfillIndexes();
+
+    executeQuery(query);
+    assertRemoteDocumentsRead(/* byKey= */ 2, /* byCollection= */ 0);
+    assertOverlaysRead(/* byKey= */ 2, /* byCollection= */ 0);
+    assertQueryReturned("coll/a", "coll/c");
+  }
+
   @Test
   public void testIndexesServerTimestamps() {
     FieldIndex index =