firestore: fix slow queries when a collection has many NO_DOCUMENT tombstones

firebase-firestore/CHANGELOG.md

@@ -1,4 +1,6 @@
 # Unreleased
+* [changed] Improve the performance of queries in collections that contain many deleted documents.
+  [#7295](//github.com/firebase/firebase-android-sdk/issues/7295)
 
 
 # 26.0.0

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

@@ -65,6 +65,32 @@ public void setIndexManager(IndexManager indexManager) {
     this.indexManager = indexManager;
   }
 
+  private enum DocumentType {
+    NO_DOCUMENT(1),
+    FOUND_DOCUMENT(2),
+    UNKNOWN_DOCUMENT(3),
+    INVALID_DOCUMENT(4);
+
+    final int dbValue;
+
+    DocumentType(int dbValue) {
+      this.dbValue = dbValue;
+    }
+
+    static DocumentType forMutableDocument(MutableDocument document) {
+      if (document.isNoDocument()) {
+        return NO_DOCUMENT;
+      } else if (document.isFoundDocument()) {
+        return FOUND_DOCUMENT;
+      } else if (document.isUnknownDocument()) {
+        return UNKNOWN_DOCUMENT;
+      } else {
+        hardAssert(!document.isValidDocument(), "MutableDocument has an unknown type");
+        return INVALID_DOCUMENT;
+      }
+    }
+  }
+
   @Override
   public void add(MutableDocument document, SnapshotVersion readTime) {
     hardAssert(
@@ -77,12 +103,13 @@ public void add(MutableDocument document, SnapshotVersion readTime) {
 
     db.execute(
         "INSERT OR REPLACE INTO remote_documents "
-            + "(path, path_length, read_time_seconds, read_time_nanos, contents) "
-            + "VALUES (?, ?, ?, ?, ?)",
+            + "(path, path_length, read_time_seconds, read_time_nanos, document_type, contents) "
+            + "VALUES (?, ?, ?, ?, ?, ?)",
         EncodedPath.encode(documentKey.getPath()),
         documentKey.getPath().length(),
         timestamp.getSeconds(),
         timestamp.getNanoseconds(),
+        DocumentType.forMutableDocument(document).dbValue,
         message.toByteArray());
 
     indexManager.addToCollectionParentIndex(document.getKey().getCollectionPath());
@@ -182,6 +209,7 @@ private Map<DocumentKey, MutableDocument> getAll(
       List<ResourcePath> collections,
       IndexOffset offset,
       int count,
+      @Nullable DocumentType tryFilterDocumentType,
       @Nullable Function<MutableDocument, Boolean> filter,
       @Nullable QueryContext context) {
     Timestamp readTime = offset.getReadTime().getTimestamp();
@@ -192,20 +220,29 @@ private Map<DocumentKey, MutableDocument> getAll(
             "SELECT contents, read_time_seconds, read_time_nanos, path "
                 + "FROM remote_documents "
                 + "WHERE path >= ? AND path < ? AND path_length = ? "
+                + (tryFilterDocumentType == null
+                    ? ""
+                    : " AND (document_type IS NULL OR document_type = ?) ")
                 + "AND (read_time_seconds > ? OR ( "
                 + "read_time_seconds = ? AND read_time_nanos > ?) OR ( "
                 + "read_time_seconds = ? AND read_time_nanos = ? and path > ?)) ",
             collections.size(),
             " UNION ");
     sql.append("ORDER BY read_time_seconds, read_time_nanos, path LIMIT ?");
 
-    Object[] bindVars = new Object[BINDS_PER_STATEMENT * collections.size() + 1];
+    Object[] bindVars =
+        new Object
+            [(BINDS_PER_STATEMENT + (tryFilterDocumentType != null ? 1 : 0)) * collections.size()
+                + 1];
     int i = 0;
     for (ResourcePath collection : collections) {
       String prefixPath = EncodedPath.encode(collection);
       bindVars[i++] = prefixPath;
       bindVars[i++] = EncodedPath.prefixSuccessor(prefixPath);
       bindVars[i++] = collection.length() + 1;
+      if (tryFilterDocumentType != null) {
+        bindVars[i++] = tryFilterDocumentType.dbValue;
+      }
       bindVars[i++] = readTime.getSeconds();
       bindVars[i++] = readTime.getSeconds();
       bindVars[i++] = readTime.getNanoseconds();
@@ -235,7 +272,8 @@ private Map<DocumentKey, MutableDocument> getAll(
       IndexOffset offset,
       int count,
       @Nullable Function<MutableDocument, Boolean> filter) {
-    return getAll(collections, offset, count, filter, /*context*/ null);
+    return getAll(
+        collections, offset, count, /*tryFilterDocumentType*/ null, filter, /*context*/ null);
   }
 
   private void processRowInBackground(
@@ -278,6 +316,10 @@ public Map<DocumentKey, MutableDocument> getDocumentsMatchingQuery(
         Collections.singletonList(query.getPath()),
         offset,
         Integer.MAX_VALUE,
+        // Specify tryFilterDocumentType=FOUND_DOCUMENT to getAll() as an optimization, because
+        // query.matches(doc) will return false for all non-"found" document types anyways.
+        // See https://github.com/firebase/firebase-android-sdk/issues/7295
+        DocumentType.FOUND_DOCUMENT,
         (MutableDocument doc) -> query.matches(doc) || mutatedKeys.contains(doc.getKey()),
         context);
   }

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

@@ -48,7 +48,7 @@ class SQLiteSchema {
    * The version of the schema. Increase this by one for each migration added to runMigrations
    * below.
    */
-  static final int VERSION = 17;
+  static final int VERSION = 18;
 
   /**
    * The batch size for data migrations.
@@ -183,6 +183,10 @@ void runSchemaUpgrades(int fromVersion, int toVersion) {
       createGlobalsTable();
     }
 
+    if (fromVersion < 18 && toVersion >= 18) {
+      addDocumentType();
+    }
+
     /*
      * Adding a new schema upgrade? READ THIS FIRST!
      *
@@ -448,6 +452,18 @@ private void addPathLength() {
     }
   }
 
+  private void addDocumentType() {
+    // The new "document_type" column is a copy of the document type encoded in the "contents" blob.
+    // Its range of values are defined in the `SQLiteRemoteDocumentCache.DocumentType` enum.
+    // The "document_type" value for a given row must be equal to the document type encoded in the
+    // "contents" column for that row. The purpose of the "document_type" column is to enable
+    // efficient filtering on document type. But when using it as a filter, a null value must also
+    // be considered as their document type must be determined by parsing the the "contents" column.
+    if (!tableContainsColumn("remote_documents", "document_type")) {
+      db.execSQL("ALTER TABLE remote_documents ADD COLUMN document_type INTEGER");
+    }
+  }
+
   private boolean hasReadTime() {
     boolean hasReadTimeSeconds = tableContainsColumn("remote_documents", "read_time_seconds");
     boolean hasReadTimeNanos = tableContainsColumn("remote_documents", "read_time_nanos");

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

@@ -702,6 +702,27 @@ public void createsOverlaysAndMigrationTable() {
     assertEquals(Persistence.DATA_MIGRATION_BUILD_OVERLAYS, migrationName);
   }
 
+  @Test
+  public void existingDocumentsMatchAfterRemoteDocumentsDocumentTypeColumnAdded() {
+    schema.runSchemaUpgrades(0, 1);
+    for (int i = 0; i < 3; i++) {
+      db.execSQL(
+          "INSERT INTO remote_documents (path, contents) VALUES (?, ?)",
+          new Object[] {encode(path("coll/doc" + i)), createDummyDocument("coll/doc" + i)});
+    }
+
+    // The migration of interest is 18, but go to the latest migration to ensure compatibility with
+    // the SQLiteRemoteDocumentCache implementation.
+    schema.runSchemaUpgrades(2, VERSION);
+
+    SQLiteRemoteDocumentCache remoteDocumentCache = createRemoteDocumentCache();
+
+    Map<DocumentKey, MutableDocument> results =
+        remoteDocumentCache.getDocumentsMatchingQuery(
+            query("coll"), IndexOffset.NONE, new HashSet<DocumentKey>());
+    assertResultsContain(results, "coll/doc0", "coll/doc1", "coll/doc2");
+  }
+
   private SQLiteRemoteDocumentCache createRemoteDocumentCache() {
     SQLitePersistence persistence =
         new SQLitePersistence(serializer, LruGarbageCollector.Params.Default(), opener);