Review

Author: schmidt-sebastian

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

@@ -16,6 +16,7 @@
 
 import static com.google.firebase.firestore.model.DocumentCollections.emptyMutableDocumentMap;
 import static com.google.firebase.firestore.util.Assert.hardAssert;
+import static com.google.firebase.firestore.util.Util.trimMap;
 
 import androidx.annotation.NonNull;
 import com.google.firebase.database.collection.ImmutableSortedMap;
@@ -86,10 +87,11 @@ public Map<DocumentKey, MutableDocument> getAll(Iterable<DocumentKey> keys) {
 
   @Override
   public Map<DocumentKey, MutableDocument> getAll(
-      String collectionGroup, IndexOffset offset, int count) {
+      String collectionGroup, IndexOffset offset, int limit) {
+    // Note: This method is pretty inefficient, but t is not called since the method is only used
+    // during index backfill, which is not supported by memory persistence.
+
     List<ResourcePath> collectionParents = indexManager.getCollectionParents(collectionGroup);
-    Set<MutableDocument> allDocuments =
-        new TreeSet<>((l, r) -> IndexOffset.fromDocument(l).compareTo(IndexOffset.fromDocument(r)));
     Map<DocumentKey, MutableDocument> matchingDocuments = new HashMap<>();
 
     for (ResourcePath collectionParent : collectionParents) {
@@ -99,24 +101,17 @@ public Map<DocumentKey, MutableDocument> getAll(
       while (iterator.hasNext()) {
         Map.Entry<DocumentKey, MutableDocument> entry = iterator.next();
         DocumentKey key = entry.getKey();
-
+        MutableDocument document = entry.getValue();
         if (!documentParent.isPrefixOf(key.getPath())) {
           break;
         }
-
-        if (IndexOffset.fromDocument(entry.getValue()).compareTo(offset) > 0) {
-          allDocuments.add(entry.getValue());
+        if (IndexOffset.fromDocument(document).compareTo(offset) > 0) {
+          matchingDocuments.put(key, document);
         }
       }
     }
 
-    Iterator<MutableDocument> it = allDocuments.iterator();
-    while (it.hasNext() && matchingDocuments.size() < count) {
-      MutableDocument document = it.next();
-      matchingDocuments.put(document.getKey(), document);
-    }
-
-    return matchingDocuments;
+    return trimMap(matchingDocuments, limit, IndexOffset.DOCUMENT_COMPARATOR);
   }
 
   @Override

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

@@ -66,18 +66,15 @@ interface RemoteDocumentCache {
   Map<DocumentKey, MutableDocument> getAll(Iterable<DocumentKey> documentKeys);
 
   /**
-   * Looks up the next {@code count} documents for a collection group based on the provided offset.
+   * Looks up the next {@code limit} documents for a collection group based on the provided offset.
    * The ordering is based on the document's read time and key.
    *
-   * <p>This method may return more results than requested if a collection group scan scans more
-   * than 100 collections.
-   *
    * @param collectionGroup The collection group to scan.
    * @param offset The offset to start the scan at.
-   * @param count The number of results to return.
+   * @param limit The maximum number of results to return.
    * @return A map with next set of documents.
    */
-  Map<DocumentKey, MutableDocument> getAll(String collectionGroup, IndexOffset offset, int count);
+  Map<DocumentKey, MutableDocument> getAll(String collectionGroup, IndexOffset offset, int limit);
 
   /**
    * Executes a query against the cached Document entries

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

@@ -17,6 +17,7 @@
 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.Util.repeatSequence;
+import static com.google.firebase.firestore.util.Util.trimMap;
 
 import androidx.annotation.VisibleForTesting;
 import com.google.firebase.Timestamp;
@@ -136,7 +137,7 @@ public Map<DocumentKey, MutableDocument> getAll(Iterable<DocumentKey> documentKe
 
   @Override
   public Map<DocumentKey, MutableDocument> getAll(
-      String collectionGroup, IndexOffset offset, int count) {
+      String collectionGroup, IndexOffset offset, int limit) {
     List<ResourcePath> collectionParents = indexManager.getCollectionParents(collectionGroup);
     List<ResourcePath> collections = new ArrayList<>(collectionParents.size());
     for (ResourcePath collectionParent : collectionParents) {
@@ -146,17 +147,17 @@ public Map<DocumentKey, MutableDocument> getAll(
     if (collections.isEmpty()) {
       return Collections.emptyMap();
     } else if (BINDS_PER_STATEMENT * collections.size() < SQLitePersistence.MAX_ARGS) {
-      return getAll(collections, offset, count);
+      return getAll(collections, offset, limit);
     } else {
       // We need to fan out our collection scan since SQLite only supports 999 binds per statement.
       Map<DocumentKey, MutableDocument> results = new HashMap<>();
       int pageSize = SQLitePersistence.MAX_ARGS / BINDS_PER_STATEMENT;
       for (int i = 0; i < collections.size(); i += pageSize) {
         results.putAll(
             getAll(
-                collections.subList(i, Math.min(collections.size(), i + pageSize)), offset, count));
+                collections.subList(i, Math.min(collections.size(), i + pageSize)), offset, limit));
       }
-      return results;
+      return trimMap(results, limit, IndexOffset.DOCUMENT_COMPARATOR);
     }
   }
 

firebase-firestore/src/main/java/com/google/firebase/firestore/model/FieldIndex.java

@@ -116,6 +116,8 @@ public static IndexState create(
   /** Stores the latest read time and document that were processed for an index. */
   @AutoValue
   public abstract static class IndexOffset implements Comparable<IndexOffset> {
+    public static final Comparator<MutableDocument> DOCUMENT_COMPARATOR = (l,r)->IndexOffset.fromDocument(l).compareTo(IndexOffset.fromDocument(r));
+
     public static final IndexOffset NONE = create(SnapshotVersion.NONE, DocumentKey.empty());
 
     /**

firebase-firestore/src/main/java/com/google/firebase/firestore/util/Util.java

@@ -31,8 +31,10 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.Random;
 import java.util.SortedSet;
 
@@ -366,4 +368,19 @@ private static <T> void diffCollections(
   private static <T> T advanceIterator(Iterator<T> it) {
     return it.hasNext() ? it.next() : null;
   }
+
+  /** Returns a map with the first {#code count} elements of {#code data} when sorted by comp. */
+  public static <K, V> Map<K, V> trimMap(Map<K, V> data, int count, Comparator<V> comp) {
+    if (data.size() <= count) {
+      return data;
+    } else {
+      List<Map.Entry<K, V>> sortedVlaues = new ArrayList<>(data.entrySet());
+      Collections.sort(sortedVlaues, (l, r) -> comp.compare(l.getValue(), r.getValue()));
+      Map<K, V> result = new HashMap<>();
+      for (int i = 0; i < count; ++i) {
+        result.put(sortedVlaues.get(i).getKey(), sortedVlaues.get(i).getValue());
+      }
+      return result;
+    }
+  }
 }

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

@@ -151,8 +151,8 @@ public Map<DocumentKey, MutableDocument> getAll(Iterable<DocumentKey> documentKe
 
       @Override
       public Map<DocumentKey, MutableDocument> getAll(
-          String collectionGroup, IndexOffset offset, int count) {
-        Map<DocumentKey, MutableDocument> result = subject.getAll(collectionGroup, offset, count);
+          String collectionGroup, IndexOffset offset, int limit) {
+        Map<DocumentKey, MutableDocument> result = subject.getAll(collectionGroup, offset, limit);
         documentsReadByQuery[0] += result.size();
         return result;
       }

firebase-firestore/src/test/java/com/google/firebase/firestore/util/UtilTest.java

@@ -15,14 +15,18 @@
 package com.google.firebase.firestore.util;
 
 import static com.google.common.truth.Truth.assertThat;
+import static com.google.firebase.firestore.util.Util.trimMap;
 import static org.junit.Assert.assertEquals;
 
 import com.google.firebase.firestore.testutil.TestUtil;
 import com.google.protobuf.ByteString;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
+
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
@@ -69,6 +73,17 @@ public void testDiffCollectionsWithEmptyLists() {
     validateDiffCollection(Collections.emptyList(), Collections.emptyList());
   }
 
+
+  @Test
+  public void testTrimMap() {
+    Map<Integer, Integer> data = new HashMap<>();
+    data.put(1, 1);
+    data.put(3, 3);
+    data.put(2, 2);
+    data = trimMap(data, 2, Integer::compare);
+    assertThat(data).containsExactly(1, 1, 2,2);
+  }
+
   private void validateDiffCollection(List<String> before, List<String> after) {
     List<String> result = new ArrayList<>(before);
     Util.diffCollections(before, after, String::compareTo, result::add, result::remove);