Address feedback

Author: cherylEnkidu

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

@@ -284,7 +284,7 @@ ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
   /** Performs a simple document lookup for the given path. */
   private ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingDocumentQuery(
       ResourcePath path, QueryContext counter) {
-    counter.fullScanCount++;
+    counter.increaseDocumentCount();
     ImmutableSortedMap<DocumentKey, Document> result = emptyDocumentMap();
     // Just do a simple document lookup.
     Document doc = getDocument(DocumentKey.fromPath(path));

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

@@ -15,7 +15,15 @@
 package com.google.firebase.firestore.local;
 
 public class QueryContext {
-  QueryContext() {}
+  public QueryContext() {}
 
-  int fullScanCount = 0;
+  private int documentCount = 0;
+
+  public int getDocumentCount() {
+    return documentCount;
+  }
+
+  public void increaseDocumentCount() {
+    documentCount++;
+  }
 }

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

@@ -86,16 +86,12 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
       return result;
     }
 
-    QueryContext counter = new QueryContext();
-    result = performQueryUsingRemoteKeys(query, remoteKeys, lastLimboFreeSnapshotVersion, counter);
+    result = performQueryUsingRemoteKeys(query, remoteKeys, lastLimboFreeSnapshotVersion);
     if (result != null) {
-      if (autoIndexEnabled) {
-        CreateCacheIndices(query, counter, result.size());
-      }
       return result;
     }
 
-    counter = new QueryContext();
+    QueryContext counter = new QueryContext();
     result = executeFullCollectionScan(query, counter);
     if (result != null && autoIndexEnabled) {
       CreateCacheIndices(query, counter, result.size());
@@ -104,7 +100,7 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
   }
 
   private void CreateCacheIndices(Query query, QueryContext counter, int resultSize) {
-    if (counter.fullScanCount > 2 * resultSize) {
+    if (counter.getDocumentCount() > 2 * resultSize) {
       indexManager.createTargetIndices(query.toTarget());
     }
   }
@@ -163,8 +159,7 @@ private void CreateCacheIndices(Query query, QueryContext counter, int resultSiz
   private @Nullable ImmutableSortedMap<DocumentKey, Document> performQueryUsingRemoteKeys(
       Query query,
       ImmutableSortedSet<DocumentKey> remoteKeys,
-      SnapshotVersion lastLimboFreeSnapshotVersion,
-      QueryContext counter) {
+      SnapshotVersion lastLimboFreeSnapshotVersion) {
     if (query.matchesAllDocuments()) {
       // Don't use indexes for queries that can be executed by scanning the collection.
       return null;
@@ -177,7 +172,7 @@ private void CreateCacheIndices(Query query, QueryContext counter, int resultSiz
     }
 
     ImmutableSortedMap<DocumentKey, Document> documents =
-        localDocumentsView.getDocuments(remoteKeys, counter);
+        localDocumentsView.getDocuments(remoteKeys);
     ImmutableSortedSet<Document> previousResults = applyQuery(query, documents);
 
     if (needsRefill(query, remoteKeys.size(), previousResults, lastLimboFreeSnapshotVersion)) {
@@ -196,8 +191,7 @@ private void CreateCacheIndices(Query query, QueryContext counter, int resultSiz
         previousResults,
         query,
         IndexOffset.createSuccessor(
-            lastLimboFreeSnapshotVersion, FieldIndex.INITIAL_LARGEST_BATCH_ID),
-        counter);
+            lastLimboFreeSnapshotVersion, FieldIndex.INITIAL_LARGEST_BATCH_ID));
   }
 
   /** Applies the query filter and sorting to the provided documents. */
@@ -278,7 +272,7 @@ private ImmutableSortedMap<DocumentKey, Document> appendRemainingResults(
       Iterable<Document> indexedResults, Query query, IndexOffset offset, QueryContext counter) {
     // Retrieve all results for documents that were updated since the offset.
     ImmutableSortedMap<DocumentKey, Document> remainingResults =
-        localDocumentsView.getDocumentsMatchingQuery(query, offset, counter);
+        localDocumentsView.getDocumentsMatchingQuery(query, offset);
     for (Document entry : indexedResults) {
       remainingResults = remainingResults.insert(entry.getKey(), entry);
     }

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

@@ -239,7 +239,7 @@ public void createTargetIndices(Target target) {
 
     for (Target subTarget : getSubTargets(target)) {
       IndexType type = getIndexType(subTarget);
-      if (type == IndexType.NONE) {
+      if (type == IndexType.NONE || type == IndexType.PARTIAL) {
         TargetIndexMatcher targetIndexMatcher = new TargetIndexMatcher(subTarget);
         addFieldIndex(targetIndexMatcher.BuildTargetIndex());
       }

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

@@ -144,7 +144,7 @@ public Map<DocumentKey, MutableDocument> getAll(
           .forEach(
               row -> {
                 processRowInBackground(backgroundQueue, results, row, /*filter*/ null);
-                counter.fullScanCount++;
+                counter.increaseDocumentCount();
               });
     }
     backgroundQueue.drain();
@@ -233,7 +233,7 @@ private Map<DocumentKey, MutableDocument> getAll(
             row -> {
               QueryContext singleCounter = new QueryContext();
               processRowInBackground(backgroundQueue, results, row, filter);
-              counter.fullScanCount++;
+              counter.increaseDocumentCount();
             });
     backgroundQueue.drain();
     return results;

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

@@ -213,6 +213,14 @@ public FieldIndex BuildTargetIndex() {
     }
 
     for (OrderBy orderBy : orderBys) {
+      // Stop adding more segments if we see a order-by on key. Typically this is the default
+      // implicit order-by which is covered in the index_entry table as a separate column.
+      // If it is not the default order-by, the generated index will be missing some segments
+      // optimized for order-bys, which is probably fine.
+      if (orderBy.getField().isKeyField()) {
+        continue;
+      }
+
       if (uniqueFields.contains(orderBy.getField())) {
         continue;
       }