Add option to allow SDK create cache indexes automatically to improve query execution locally

firebase-firestore/src/main/java/com/google/firebase/firestore/FirebaseFirestore.java

@@ -49,6 +49,7 @@
 import com.google.firebase.firestore.model.ResourcePath;
 import com.google.firebase.firestore.remote.FirestoreChannel;
 import com.google.firebase.firestore.remote.GrpcMetadataProvider;
+import com.google.firebase.firestore.util.Assert;
 import com.google.firebase.firestore.util.AsyncQueue;
 import com.google.firebase.firestore.util.ByteBufferInputStream;
 import com.google.firebase.firestore.util.Executors;
@@ -63,6 +64,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.Executor;
+import java.util.concurrent.atomic.AtomicReference;
 import org.json.JSONArray;
 import org.json.JSONException;
 import org.json.JSONObject;
@@ -105,7 +107,9 @@ public interface InstanceRegistry {
   private volatile FirestoreClient client;
   private final GrpcMetadataProvider metadataProvider;
 
-  @Nullable private PersistentCacheIndexManager persistentCacheIndexManager;
+  @Nullable
+  private AtomicReference<PersistentCacheIndexManager> persistentCacheIndexManager =
+      new AtomicReference<>();
 
   @NonNull
   private static FirebaseApp getDefaultFirebaseApp() {
@@ -418,14 +422,15 @@ public Task<Void> setIndexConfiguration(@NonNull String json) {
   @Nullable
   public PersistentCacheIndexManager getPersistentCacheIndexManager() {
     ensureClientConfigured();
-    if (persistentCacheIndexManager != null) {
-      return persistentCacheIndexManager;
-    }
+    Assert.hardAssertNonNull(
+        persistentCacheIndexManager, "persistentCacheIndexManager has not been initialized.");
+
     if (settings.isPersistenceEnabled()
         || settings.getCacheSettings() instanceof PersistentCacheSettings) {
-      persistentCacheIndexManager = new PersistentCacheIndexManager(client);
+      persistentCacheIndexManager.compareAndSet(null, new PersistentCacheIndexManager(client));
     }
-    return persistentCacheIndexManager;
+
+    return persistentCacheIndexManager.get();
   }
 
   /**

firebase-firestore/src/main/java/com/google/firebase/firestore/PersistentCacheIndexManager.java

@@ -22,16 +22,15 @@
  * A {@code PersistentCacheIndexManager} which you can config persistent cache indexes used for
  * local query execution.
  *
- * <p>To use, calling {@link FirebaseFirestore#getPersistentCacheIndexManager()} to get an instance.
+ * <p>To use, call {@link FirebaseFirestore#getPersistentCacheIndexManager()} to get an instance.
  */
 // TODO(csi): Remove the `hide` and scope annotations.
 /** @hide */
 @RestrictTo(RestrictTo.Scope.LIBRARY)
-public class PersistentCacheIndexManager {
+public final class PersistentCacheIndexManager {
   @NonNull private FirestoreClient client;
 
-  @RestrictTo(RestrictTo.Scope.LIBRARY)
-  public PersistentCacheIndexManager(FirestoreClient client) {
+  PersistentCacheIndexManager(FirestoreClient client) {
     this.client = client;
   }
 

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

@@ -281,7 +281,7 @@ ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
    */
   ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
       Query query, IndexOffset offset) {
-    return getDocumentsMatchingQuery(query, offset, null);
+    return getDocumentsMatchingQuery(query, offset, /*context*/ null);
   }
 
   /** Performs a simple document lookup for the given path. */

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

@@ -142,7 +142,7 @@ public Map<DocumentKey, MutableDocument> getDocumentsMatchingQuery(
   @Override
   public Map<DocumentKey, MutableDocument> getDocumentsMatchingQuery(
       Query query, IndexOffset offset, @Nonnull Set<DocumentKey> mutatedKeys) {
-    return getDocumentsMatchingQuery(query, offset, mutatedKeys, null);
+    return getDocumentsMatchingQuery(query, offset, mutatedKeys, /*context*/ null);
   }
 
   Iterable<Document> getDocuments() {

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

@@ -16,16 +16,14 @@
 
 /** A tracker to keep a record of important details during database local query execution. */
 public class QueryContext {
-  public QueryContext() {}
-
   /** Counts the number of documents passed through during local query execution. */
   private int documentReadCount = 0;
 
   public int getDocumentReadCount() {
     return documentReadCount;
   }
 
-  public void increaseDocumentCount() {
+  public void incrementDocumentReadCount() {
     documentReadCount++;
   }
 }

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

@@ -16,6 +16,7 @@
 
 import static com.google.firebase.firestore.util.Assert.hardAssert;
 
+import androidx.annotation.VisibleForTesting;
 import com.google.firebase.database.collection.ImmutableSortedMap;
 import com.google.firebase.database.collection.ImmutableSortedSet;
 import com.google.firebase.firestore.core.Query;
@@ -61,12 +62,27 @@
 public class QueryEngine {
   private static final String LOG_TAG = "QueryEngine";
 
+  private static final int MIN_COLLECTION_SIZE_TO_AUTO_CREATE_INDEX = 100;
+
+  // TODO(csi): update the temp cost.
+  /**
+   * This cost represents the evaluation result of (([index, docKey] + [docKey, docContent]) per
+   * document in the result set) / ([docKey, docContent] per documents in full collection scan)
+   * coming from experiment https://github.com/firebase/firebase-android-sdk/pull/5064.
+   */
+  private static final int RELATIVE_INDEX_READ_COST = 3;
+
   private LocalDocumentsView localDocumentsView;
   private IndexManager indexManager;
   private boolean initialized;
 
   private boolean automaticIndexingEnabled = false;
 
+  /** SDK only decides whether it should create index when collection size is larger than this. */
+  private int minCollectionSizeToAutoCreateIndex = MIN_COLLECTION_SIZE_TO_AUTO_CREATE_INDEX;
+
+  private int relativeIndexReadCost = RELATIVE_INDEX_READ_COST;
+
   public void initialize(LocalDocumentsView localDocumentsView, IndexManager indexManager) {
     this.localDocumentsView = localDocumentsView;
     this.indexManager = indexManager;
@@ -100,7 +116,7 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
     QueryContext context = new QueryContext();
     result = executeFullCollectionScan(query, context);
     if (result != null && automaticIndexingEnabled) {
-      CreateCacheIndexes(query, context, result.size());
+      createCacheIndexes(query, context, result.size());
     }
     return result;
   }
@@ -110,10 +126,26 @@ public ImmutableSortedMap<DocumentKey, Document> getDocumentsMatchingQuery(
    * context and query result size.
    */
   // TODO(csi): Auto experiment data.
-  private void CreateCacheIndexes(Query query, QueryContext context, int resultSize) {
+  private void createCacheIndexes(Query query, QueryContext context, int resultSize) {
+    if (context.getDocumentReadCount() < minCollectionSizeToAutoCreateIndex) {
+      Logger.debug(
+          LOG_TAG,
+          "SDK will only creates cache indexes for collection contains more than or equal to "
+              + "%s documents.",
+          minCollectionSizeToAutoCreateIndex);
+      return;
+    }
+
+    if (Logger.isDebugEnabled()) {
+      Logger.debug(
+          LOG_TAG,
+          "Query scans %s local documents and returns %s documents as results.",
+          context.getDocumentReadCount(),
+          resultSize);
+    }
+
     String decisionStr = "";
-    // If evaluation is updated, please update tests in SQLiteLocalStoreTest.java
-    if (context.getDocumentReadCount() > 2 * resultSize) {
+    if (context.getDocumentReadCount() > relativeIndexReadCost * resultSize) {
       indexManager.createTargetIndices(query.toTarget());
     } else {
       decisionStr = " not";
@@ -122,11 +154,8 @@ private void CreateCacheIndexes(Query query, QueryContext context, int resultSiz
     if (Logger.isDebugEnabled()) {
       Logger.debug(
           LOG_TAG,
-          "Query ran locally using a full collection scan, walking through %s documents in total "
-              + "and returning %s documents. The SDK has decided%s to create cache indexes "
-              + "for this query, as using cache indexes may%s help improve performance.",
-          context.getDocumentReadCount(),
-          resultSize,
+          "The SDK decides%s to create cache indexes for this query, as using cache indexes "
+              + "may%s help improve performance.",
           decisionStr,
           decisionStr);
     }
@@ -305,4 +334,14 @@ private ImmutableSortedMap<DocumentKey, Document> appendRemainingResults(
     }
     return remainingResults;
   }
+
+  @VisibleForTesting
+  void setMinCollectionSizeToAutoCreateIndex(int newMin) {
+    minCollectionSizeToAutoCreateIndex = newMin;
+  }
+
+  @VisibleForTesting
+  void setRelativeIndexReadCost(int newCost) {
+    relativeIndexReadCost = newCost;
+  }
 }

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

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

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

@@ -224,7 +224,7 @@ private Map<DocumentKey, MutableDocument> getAll(
               processRowInBackground(backgroundQueue, results, row, filter);
               if (context != null) {
                 // Increases the counter by 1 for every document processed.
-                context.increaseDocumentCount();
+                context.incrementDocumentReadCount();
               }
             });
     backgroundQueue.drain();
@@ -236,7 +236,7 @@ private Map<DocumentKey, MutableDocument> getAll(
       IndexOffset offset,
       int count,
       @Nullable Function<MutableDocument, Boolean> filter) {
-    return getAll(collections, offset, count, filter, null);
+    return getAll(collections, offset, count, filter, /*context*/ null);
   }
 
   private void processRowInBackground(
@@ -266,7 +266,7 @@ private void processRowInBackground(
   @Override
   public Map<DocumentKey, MutableDocument> getDocumentsMatchingQuery(
       Query query, IndexOffset offset, @Nonnull Set<DocumentKey> mutatedKeys) {
-    return getDocumentsMatchingQuery(query, offset, mutatedKeys, null);
+    return getDocumentsMatchingQuery(query, offset, mutatedKeys, /*context*/ null);
   }
 
   @Override

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

@@ -193,7 +193,7 @@ public boolean servedByIndex(FieldIndex index) {
   }
 
   /** Returns a full matched field index for this target. */
-  public FieldIndex BuildTargetIndex() {
+  public FieldIndex buildTargetIndex() {
     // We want to make sure only one segment created for one field. For example, in case like
     // a == 3 and a > 2, index, a ASCENDING, will only be created once.
     Set<FieldPath> uniqueFields = new HashSet<>();

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

@@ -97,6 +97,16 @@ public void disableIndexAutoCreation() {
     queryEngine.disableIndexAutoCreation();
   }
 
+  @Override
+  public void setMinCollectionSizeToAutoCreateIndex(int newMin) {
+    queryEngine.setMinCollectionSizeToAutoCreateIndex(newMin);
+  }
+
+  @Override
+  public void setRelativeIndexReadCost(int newCost) {
+    queryEngine.setRelativeIndexReadCost(newCost);
+  }
+
   /**
    * Returns the number of documents returned by the RemoteDocumentCache's `getAll()` API (since the
    * last call to `resetCounts()`)
@@ -181,7 +191,7 @@ public Map<DocumentKey, MutableDocument> getAll(
       @Override
       public Map<DocumentKey, MutableDocument> getDocumentsMatchingQuery(
           Query query, IndexOffset offset, @NonNull Set<DocumentKey> mutatedKeys) {
-        return getDocumentsMatchingQuery(query, offset, mutatedKeys, null);
+        return getDocumentsMatchingQuery(query, offset, mutatedKeys, /*context*/ null);
       }
 
       @Override

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

@@ -232,6 +232,14 @@ protected void disableIndexAutoCreation() {
     queryEngine.disableIndexAutoCreation();
   }
 
+  protected void setMinCollectionSizeToAutoCreateIndex(int newMin) {
+    queryEngine.setMinCollectionSizeToAutoCreateIndex(newMin);
+  }
+
+  protected void setRelativeIndexReadCost(int newCost) {
+    queryEngine.setRelativeIndexReadCost(newCost);
+  }
+
   private void releaseTarget(int targetId) {
     localStore.releaseTarget(targetId);
   }