Add getNext()

Author: schmidt-sebastian

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

@@ -19,19 +19,15 @@
 import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 import com.google.firebase.database.collection.ImmutableSortedMap;
-import com.google.firebase.firestore.core.Query;
 import com.google.firebase.firestore.model.Document;
 import com.google.firebase.firestore.model.DocumentKey;
 import com.google.firebase.firestore.model.FieldIndex;
 import com.google.firebase.firestore.model.FieldIndex.IndexOffset;
-import com.google.firebase.firestore.model.ResourcePath;
 import com.google.firebase.firestore.util.AsyncQueue;
 import com.google.firebase.firestore.util.Logger;
-import java.util.ArrayList;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.HashSet;
-import java.util.List;
+import java.util.Iterator;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
@@ -138,22 +134,18 @@ private int writeIndexEntries(LocalDocumentsView localDocumentsView) {
   /** Writes entries for the fetched field indexes. */
   private int writeEntriesForCollectionGroup(
       LocalDocumentsView localDocumentsView, String collectionGroup, int entriesRemainingUnderCap) {
-    Query query = new Query(ResourcePath.EMPTY, collectionGroup);
+    // TODO(indexing): Support mutation batch Ids when sorting and writing indexes.
 
     // Use the earliest offset of all field indexes to query the local cache.
     IndexOffset existingOffset = getExistingOffset(indexManager.getFieldIndexes(collectionGroup));
-
-    // TODO(indexing): Use limit queries to only fetch the required number of entries.
-    // TODO(indexing): Support mutation batch Ids when sorting and writing indexes.
     ImmutableSortedMap<DocumentKey, Document> documents =
-        localDocumentsView.getDocumentsMatchingQuery(query, existingOffset);
-
-    List<Document> oldestDocuments = getOldestDocuments(documents, entriesRemainingUnderCap);
-    indexManager.updateIndexEntries(oldestDocuments);
+        localDocumentsView.getDocuments(collectionGroup, existingOffset, entriesRemainingUnderCap);
+    indexManager.updateIndexEntries(documents);
 
-    IndexOffset newOffset = getNewOffset(oldestDocuments, existingOffset);
+    IndexOffset newOffset = getNewOffset(documents, existingOffset);
     indexManager.updateCollectionGroup(collectionGroup, newOffset);
-    return oldestDocuments.size();
+
+    return documents.size();
   }
 
   /** Returns the lowest offset for the provided index group. */
@@ -168,37 +160,22 @@ private IndexOffset getExistingOffset(Collection<FieldIndex> fieldIndexes) {
     return lowestOffset == null ? IndexOffset.NONE : lowestOffset;
   }
 
-  /**
-   * Returns the offset for the index based on the newly indexed documents.
-   *
-   * @param documents a list of documents sorted by read time and key (ascending)
-   * @param currentOffset the current offset of the index group
-   */
-  private IndexOffset getNewOffset(List<Document> documents, IndexOffset currentOffset) {
-    IndexOffset latestOffset =
-        documents.isEmpty()
-            ? IndexOffset.create(remoteDocumentCache.getLatestReadTime())
-            : IndexOffset.create(
-                documents.get(documents.size() - 1).getReadTime(),
-                documents.get(documents.size() - 1).getKey());
-    // Make sure the index does not go back in time
-    latestOffset = latestOffset.compareTo(currentOffset) > 0 ? latestOffset : currentOffset;
-    return latestOffset;
-  }
-
-  /** Returns up to {@code count} documents sorted by read time and key. */
-  private List<Document> getOldestDocuments(
-      ImmutableSortedMap<DocumentKey, Document> documents, int count) {
-    List<Document> oldestDocuments = new ArrayList<>();
-    for (Map.Entry<DocumentKey, Document> entry : documents) {
-      oldestDocuments.add(entry.getValue());
+  /** Returns the offset for the index based on the newly indexed documents. */
+  private IndexOffset getNewOffset(
+      ImmutableSortedMap<DocumentKey, Document> documents, IndexOffset currentOffset) {
+    if (documents.isEmpty()) {
+      return IndexOffset.create(remoteDocumentCache.getLatestReadTime());
+    } else {
+      IndexOffset latestOffset = currentOffset;
+      Iterator<Map.Entry<DocumentKey, Document>> it = documents.iterator();
+      while (it.hasNext()) {
+        IndexOffset newOffset = IndexOffset.fromDocument(it.next().getValue());
+        if (newOffset.compareTo(latestOffset) > 0) {
+          latestOffset = newOffset;
+        }
+      }
+      return latestOffset;
     }
-    Collections.sort(
-        oldestDocuments,
-        (l, r) ->
-            IndexOffset.create(l.getReadTime(), l.getKey())
-                .compareTo(IndexOffset.create(r.getReadTime(), r.getKey())));
-    return oldestDocuments.subList(0, Math.min(count, oldestDocuments.size()));
   }
 
   @VisibleForTesting

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

@@ -15,6 +15,7 @@
 package com.google.firebase.firestore.local;
 
 import androidx.annotation.Nullable;
+import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.firestore.core.Target;
 import com.google.firebase.firestore.model.Document;
 import com.google.firebase.firestore.model.DocumentKey;
@@ -97,5 +98,5 @@ public interface IndexManager {
   void updateCollectionGroup(String collectionGroup, FieldIndex.IndexOffset offset);
 
   /** Updates the index entries for the provided documents. */
-  void updateIndexEntries(Collection<Document> documents);
+  void updateIndexEntries(ImmutableSortedMap<DocumentKey, Document> documents);
 }

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

@@ -237,10 +237,15 @@ void recalculateAndSaveOverlays(Set<DocumentKey> documentKeys) {
     recalculateAndSaveOverlays(docs);
   }
 
-  // TODO: The Querying implementation here should move 100% to the query engines.
-  // Instead, we should just provide a getCollectionDocuments() method here that return all the
-  // documents in a given collection so that query engine can do that and then filter in
-  // memory.
+  /**
+   * Gets the local view of the next {@code count} documents based on their read time.
+   */
+  ImmutableSortedMap<DocumentKey, Document> getDocuments(
+      String collectionGroup, IndexOffset offset, int count) {
+    Map<DocumentKey, MutableDocument> docs =
+        remoteDocumentCache.getAll(collectionGroup, offset, count);
+    return getLocalViewOfDocuments(docs, new HashSet<>());
+  }
 
   /**
    * Performs a query against the local view of all documents.

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

@@ -16,6 +16,7 @@
 import static com.google.firebase.firestore.util.Assert.hardAssert;
 
 import androidx.annotation.Nullable;
+import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.firestore.core.Target;
 import com.google.firebase.firestore.model.Document;
 import com.google.firebase.firestore.model.DocumentKey;
@@ -98,7 +99,7 @@ public Collection<FieldIndex> getFieldIndexes() {
   }
 
   @Override
-  public void updateIndexEntries(Collection<Document> documents) {
+  public void updateIndexEntries(ImmutableSortedMap<DocumentKey, Document> documents) {
     // Field indices are not supported with memory persistence.
   }
 

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

@@ -27,7 +27,10 @@
 import com.google.firebase.firestore.model.SnapshotVersion;
 import java.util.HashMap;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
+import java.util.Set;
+import java.util.TreeSet;
 
 /** In-memory cache of remote documents. */
 final class MemoryRemoteDocumentCache implements RemoteDocumentCache {
@@ -81,6 +84,38 @@ public Map<DocumentKey, MutableDocument> getAll(Iterable<DocumentKey> keys) {
     return result;
   }
 
+  @Override
+  public Map<DocumentKey, MutableDocument> getAll(
+      String collectionGroup, IndexOffset offset, int count) {
+    List<ResourcePath> collectionParents = indexManager.getCollectionParents(collectionGroup);
+
+    Set<MutableDocument> allDocuments =
+        new TreeSet<>((l, r) -> IndexOffset.fromDocument(l).compareTo(IndexOffset.fromDocument(r)));
+    for (ResourcePath collectionParent : collectionParents) {
+      ResourcePath documentParent = collectionParent.append(collectionGroup);
+      Iterator<Map.Entry<DocumentKey, MutableDocument>> iterator =
+          docs.iteratorFrom(DocumentKey.fromPath(documentParent.append("")));
+      while (iterator.hasNext()) {
+        Map.Entry<DocumentKey, MutableDocument> entry = iterator.next();
+        DocumentKey key = entry.getKey();
+        if (!documentParent.isPrefixOf(key.getPath())) {
+          break;
+        }
+        if (IndexOffset.fromDocument(entry.getValue()).compareTo(offset) > 0) {
+          allDocuments.add(entry.getValue());
+        }
+      }
+    }
+
+    Map<DocumentKey, MutableDocument> matchingDocuments = new HashMap<>();
+    Iterator<MutableDocument> it = allDocuments.iterator();
+    while (it.hasNext() && matchingDocuments.size() < count) {
+      MutableDocument document = it.next();
+      matchingDocuments.put(document.getKey(), document);
+    }
+    return matchingDocuments;
+  }
+
   @Override
   public ImmutableSortedMap<DocumentKey, MutableDocument> getAllDocumentsMatchingQuery(
       Query query, IndexOffset offset) {
@@ -108,7 +143,7 @@ public ImmutableSortedMap<DocumentKey, MutableDocument> getAllDocumentsMatchingQ
         continue;
       }
 
-      if (IndexOffset.create(doc.getReadTime(), doc.getKey()).compareTo(offset) <= 0) {
+      if (IndexOffset.fromDocument(doc).compareTo(offset) <= 0) {
         continue;
       }
 

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

@@ -17,7 +17,7 @@
 import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.firestore.core.Query;
 import com.google.firebase.firestore.model.DocumentKey;
-import com.google.firebase.firestore.model.FieldIndex;
+import com.google.firebase.firestore.model.FieldIndex.IndexOffset;
 import com.google.firebase.firestore.model.MutableDocument;
 import com.google.firebase.firestore.model.SnapshotVersion;
 import java.util.Map;
@@ -65,6 +65,20 @@ 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.
+   * 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 125 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.
+   * @return A map with next set of documents.
+   */
+  Map<DocumentKey, MutableDocument> getAll(String collectionGroup, IndexOffset offset, int count);
+
   /**
    * Executes a query against the cached Document entries
    *
@@ -78,7 +92,7 @@ interface RemoteDocumentCache {
    * @return The set of matching documents.
    */
   ImmutableSortedMap<DocumentKey, MutableDocument> getAllDocumentsMatchingQuery(
-      Query query, FieldIndex.IndexOffset offset);
+      Query query, IndexOffset offset);
 
   /** Returns the latest read time of any document in the cache. */
   SnapshotVersion getLatestReadTime();

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

@@ -23,6 +23,8 @@
 
 import androidx.annotation.Nullable;
 import com.google.firebase.Timestamp;
+import com.google.firebase.database.collection.ImmutableSortedMap;
+import com.google.firebase.database.collection.LLRBNode;
 import com.google.firebase.firestore.auth.User;
 import com.google.firebase.firestore.core.Bound;
 import com.google.firebase.firestore.core.FieldFilter;
@@ -222,19 +224,22 @@ public void deleteFieldIndex(FieldIndex index) {
   }
 
   @Override
-  public void updateIndexEntries(Collection<Document> documents) {
+  public void updateIndexEntries(ImmutableSortedMap<DocumentKey, Document> documents) {
     hardAssert(started, "IndexManager not started");
-    for (Document document : documents) {
-      Collection<FieldIndex> fieldIndexes = getFieldIndexes(document.getKey().getCollectionGroup());
-      for (FieldIndex fieldIndex : fieldIndexes) {
-        SortedSet<IndexEntry> existingEntries =
-            getExistingIndexEntries(document.getKey(), fieldIndex);
-        SortedSet<IndexEntry> newEntries = computeIndexEntries(document, fieldIndex);
-        if (!existingEntries.equals(newEntries)) {
-          updateEntries(document, existingEntries, newEntries);
-        }
-      }
-    }
+    documents.inOrderTraversal(
+        new LLRBNode.NodeVisitor<DocumentKey, Document>() {
+          @Override
+          public void visitEntry(DocumentKey key, Document document) {
+            Collection<FieldIndex> fieldIndexes = getFieldIndexes(key.getCollectionGroup());
+            for (FieldIndex fieldIndex : fieldIndexes) {
+              SortedSet<IndexEntry> existingEntries = getExistingIndexEntries(key, fieldIndex);
+              SortedSet<IndexEntry> newEntries = computeIndexEntries(document, fieldIndex);
+              if (!existingEntries.equals(newEntries)) {
+                updateEntries(document, existingEntries, newEntries);
+              }
+            }
+          }
+        });
   }
 
   /**

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

@@ -54,6 +54,10 @@
  * helper routines that make dealing with SQLite much more pleasant.
  */
 public final class SQLitePersistence extends Persistence {
+  /**
+   * The maximum number of bind args for a single statement. Set to 900 instead of 999 for safety.
+   */
+  public static final int MAX_ARGS = 900;
 
   /**
    * Creates the database name that is used to identify the database to be used with a Firestore