Performance: Don't deserialize full document Proto for Query execution

firebase-firestore/CHANGELOG.md

@@ -1,4 +1,6 @@
 # Unreleased
+- [changed] Improved performance for queries with filters that only return a
+  small subset of the documents in a collection.
 - [changed] Instead of failing silently, Firestore now crashes the client app
   if it fails to load SSL Ciphers. To avoid these crashes, you must bundle 
   Conscrypt to support non-GMSCore devices on Android KitKat or JellyBean (see

firebase-firestore/ktx/src/test/java/com/google/firebase/firestore/testutil/TestUtil.java

@@ -175,21 +175,21 @@ public static SnapshotVersion version(long versionMicros) {
 
   public static Document doc(String key, long version, Map<String, Object> data) {
     return new Document(
-        key(key), version(version), wrapObject(data), Document.DocumentState.SYNCED);
+        key(key), version(version), Document.DocumentState.SYNCED, wrapObject(data));
   }
 
   public static Document doc(DocumentKey key, long version, Map<String, Object> data) {
-    return new Document(key, version(version), wrapObject(data), Document.DocumentState.SYNCED);
+    return new Document(key, version(version), Document.DocumentState.SYNCED, wrapObject(data));
   }
 
   public static Document doc(
       String key, long version, ObjectValue data, Document.DocumentState documentState) {
-    return new Document(key(key), version(version), data, documentState);
+    return new Document(key(key), version(version), documentState, data);
   }
 
   public static Document doc(
       String key, long version, Map<String, Object> data, Document.DocumentState documentState) {
-    return new Document(key(key), version(version), wrapObject(data), documentState);
+    return new Document(key(key), version(version), documentState, wrapObject(data));
   }
 
   public static NoDocument deletedDoc(String key, long version) {

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

@@ -112,15 +112,15 @@ private com.google.firestore.v1.Document encodeDocument(Document document) {
   private Document decodeDocument(
       com.google.firestore.v1.Document document, boolean hasCommittedMutations) {
     DocumentKey key = rpcSerializer.decodeKey(document.getName());
-    ObjectValue value = rpcSerializer.decodeFields(document.getFieldsMap());
     SnapshotVersion version = rpcSerializer.decodeVersion(document.getUpdateTime());
     return new Document(
         key,
         version,
-        value,
         hasCommittedMutations
             ? Document.DocumentState.COMMITTED_MUTATIONS
-            : Document.DocumentState.SYNCED);
+            : Document.DocumentState.SYNCED,
+        document,
+        rpcSerializer::decodeValue);
   }
 
   /** Encodes a NoDocument value to the equivalent proto. */

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

@@ -14,9 +14,16 @@
 
 package com.google.firebase.firestore.model;
 
+import static com.google.firebase.firestore.util.Assert.hardAssert;
+
+import com.google.common.base.Function;
 import com.google.firebase.firestore.model.value.FieldValue;
 import com.google.firebase.firestore.model.value.ObjectValue;
+import com.google.firestore.v1.Value;
 import java.util.Comparator;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
 /**
@@ -36,58 +43,107 @@ public enum DocumentState {
   }
 
   private static final Comparator<Document> KEY_COMPARATOR =
-      new Comparator<Document>() {
-        @Override
-        public int compare(Document left, Document right) {
-          return left.getKey().compareTo(right.getKey());
-        }
-      };
+      (left, right) -> left.getKey().compareTo(right.getKey());
 
   /** A document comparator that returns document by key and key only. */
   public static Comparator<Document> keyComparator() {
     return KEY_COMPARATOR;
   }
 
-  private final ObjectValue data;
-
   private final DocumentState documentState;
+  private @Nullable final com.google.firestore.v1.Document proto;
+  private @Nullable final Function<Value, FieldValue> converter;
+  private @Nullable ObjectValue objectValue;
 
-  /**
-   * Memoized serialized form of the document for optimization purposes (avoids repeated
-   * serialization). Might be null.
-   */
-  private final com.google.firestore.v1.Document proto;
-
-  public @Nullable com.google.firestore.v1.Document getProto() {
-    return proto;
-  }
+  /** A cache for FieldValues that have already been deserialized in `getField()`. */
+  private @Nullable Map<FieldPath, FieldValue> fieldValueCache;
 
   public Document(
-      DocumentKey key, SnapshotVersion version, ObjectValue data, DocumentState documentState) {
+      DocumentKey key,
+      SnapshotVersion version,
+      DocumentState documentState,
+      ObjectValue objectValue) {
     super(key, version);
-    this.data = data;
     this.documentState = documentState;
+    this.objectValue = objectValue;
     this.proto = null;
+    this.converter = null;
   }
 
   public Document(
       DocumentKey key,
       SnapshotVersion version,
-      ObjectValue data,
       DocumentState documentState,
-      com.google.firestore.v1.Document proto) {
+      com.google.firestore.v1.Document proto,
+      Function<com.google.firestore.v1.Value, FieldValue> converter) {
     super(key, version);
-    this.data = data;
     this.documentState = documentState;
     this.proto = proto;
+    this.converter = converter;
   }
 
+  /**
+   * Memoized serialized form of the document for optimization purposes (avoids repeated
+   * serialization). Might be null.
+   */
+  public @Nullable com.google.firestore.v1.Document getProto() {
+    return proto;
+  }
+
+  @Nonnull
   public ObjectValue getData() {
-    return data;
+    if (objectValue == null) {
+      hardAssert(proto != null && converter != null, "Expected proto and converter to be non-null");
+
+      ObjectValue result = ObjectValue.emptyObject();
+      for (Map.Entry<String, com.google.firestore.v1.Value> entry :
+          proto.getFieldsMap().entrySet()) {
+        FieldPath path = FieldPath.fromSingleSegment(entry.getKey());
+        FieldValue value = converter.apply(entry.getValue());
+        result = result.set(path, value);
+      }
+      objectValue = result;
+
+      // Once objectValue is computed, values inside the fieldValueCache are no longer accessed.
+      fieldValueCache = null;
+    }
+
+    return objectValue;
   }
 
   public @Nullable FieldValue getField(FieldPath path) {
-    return data.get(path);
+    if (objectValue != null) {
+      return objectValue.get(path);
+    } else {
+      hardAssert(proto != null && converter != null, "Expected proto and converter to be non-null");
+
+      if (fieldValueCache == null) {
+        // TODO(b/136090445): Remove the cache when `getField` is no longer called during Query
+        // ordering.
+        fieldValueCache = new ConcurrentHashMap<>();
+      }
+
+      FieldValue fieldValue = fieldValueCache.get(path);
+      if (fieldValue == null) {
+        // Instead of deserializing the full Document proto, we only deserialize the value at
+        // the requested field path. This speeds up Query execution as query filters can discard
+        // documents based on a single field.
+        Value protoValue = proto.getFieldsMap().get(path.getFirstSegment());
+        for (int i = 1; protoValue != null && i < path.length(); ++i) {
+          if (protoValue.getValueTypeCase() != Value.ValueTypeCase.MAP_VALUE) {
+            return null;
+          }
+          protoValue = protoValue.getMapValue().getFieldsMap().get(path.getSegment(i));
+        }
+
+        if (protoValue != null) {
+          fieldValue = converter.apply(protoValue);
+          fieldValueCache.put(path, fieldValue);
+        }
+      }
+
+      return fieldValue;
+    }
   }
 
   public @Nullable Object getFieldValue(FieldPath path) {
@@ -113,7 +169,7 @@ public boolean equals(Object o) {
     if (this == o) {
       return true;
     }
-    if (o == null || getClass() != o.getClass()) {
+    if (!(o instanceof Document)) {
       return false;
     }
 
@@ -122,13 +178,13 @@ public boolean equals(Object o) {
     return getVersion().equals(document.getVersion())
         && getKey().equals(document.getKey())
         && documentState.equals(document.documentState)
-        && data.equals(document.data);
+        && getData().equals(document.getData());
   }
 
   @Override
   public int hashCode() {
+    // Note: We deliberately decided to omit `getData()` since its computation is expensive.
     int result = getKey().hashCode();
-    result = 31 * result + data.hashCode();
     result = 31 * result + getVersion().hashCode();
     result = 31 * result + documentState.hashCode();
     return result;
@@ -140,7 +196,7 @@ public String toString() {
         + "key="
         + getKey()
         + ", data="
-        + data
+        + getData()
         + ", version="
         + getVersion()
         + ", documentState="

firebase-firestore/src/main/java/com/google/firebase/firestore/model/mutation/PatchMutation.java

@@ -112,7 +112,7 @@ public MaybeDocument applyToRemoteDocument(
 
     SnapshotVersion version = mutationResult.getVersion();
     ObjectValue newData = patchDocument(maybeDoc);
-    return new Document(getKey(), version, newData, Document.DocumentState.COMMITTED_MUTATIONS);
+    return new Document(getKey(), version, Document.DocumentState.COMMITTED_MUTATIONS, newData);
   }
 
   @Nullable
@@ -127,7 +127,7 @@ public MaybeDocument applyToLocalView(
 
     SnapshotVersion version = getPostMutationVersion(maybeDoc);
     ObjectValue newData = patchDocument(maybeDoc);
-    return new Document(getKey(), version, newData, Document.DocumentState.LOCAL_MUTATIONS);
+    return new Document(getKey(), version, Document.DocumentState.LOCAL_MUTATIONS, newData);
   }
 
   @Nullable

firebase-firestore/src/main/java/com/google/firebase/firestore/model/mutation/SetMutation.java

@@ -72,7 +72,7 @@ public MaybeDocument applyToRemoteDocument(
     // accepted the mutation so the precondition must have held.
 
     SnapshotVersion version = mutationResult.getVersion();
-    return new Document(getKey(), version, value, Document.DocumentState.COMMITTED_MUTATIONS);
+    return new Document(getKey(), version, Document.DocumentState.COMMITTED_MUTATIONS, value);
   }
 
   @Nullable
@@ -86,7 +86,7 @@ public MaybeDocument applyToLocalView(
     }
 
     SnapshotVersion version = getPostMutationVersion(maybeDoc);
-    return new Document(getKey(), version, value, Document.DocumentState.LOCAL_MUTATIONS);
+    return new Document(getKey(), version, Document.DocumentState.LOCAL_MUTATIONS, value);
   }
 
   @Nullable

firebase-firestore/src/main/java/com/google/firebase/firestore/model/mutation/TransformMutation.java

@@ -103,7 +103,7 @@ public MaybeDocument applyToRemoteDocument(
         serverTransformResults(doc, mutationResult.getTransformResults());
     ObjectValue newData = transformObject(doc.getData(), transformResults);
     return new Document(
-        getKey(), mutationResult.getVersion(), newData, Document.DocumentState.COMMITTED_MUTATIONS);
+        getKey(), mutationResult.getVersion(), Document.DocumentState.COMMITTED_MUTATIONS, newData);
   }
 
   @Nullable
@@ -120,7 +120,7 @@ public MaybeDocument applyToLocalView(
     List<FieldValue> transformResults = localTransformResults(localWriteTime, baseDoc);
     ObjectValue newData = transformObject(doc.getData(), transformResults);
     return new Document(
-        getKey(), doc.getVersion(), newData, Document.DocumentState.LOCAL_MUTATIONS);
+        getKey(), doc.getVersion(), Document.DocumentState.LOCAL_MUTATIONS, newData);
   }
 
   @Override

firebase-firestore/src/main/java/com/google/firebase/firestore/remote/RemoteSerializer.java

@@ -407,11 +407,11 @@ private Document decodeFoundDocument(BatchGetDocumentsResponse response) {
         response.getResultCase().equals(ResultCase.FOUND),
         "Tried to deserialize a found document from a missing document.");
     DocumentKey key = decodeKey(response.getFound().getName());
-    ObjectValue value = decodeFields(response.getFound().getFieldsMap());
     SnapshotVersion version = decodeVersion(response.getFound().getUpdateTime());
     hardAssert(
         !version.equals(SnapshotVersion.NONE), "Got a document response with no snapshot version");
-    return new Document(key, version, value, Document.DocumentState.SYNCED, response.getFound());
+    return new Document(
+        key, version, Document.DocumentState.SYNCED, response.getFound(), this::decodeValue);
   }
 
   private NoDocument decodeMissingDocument(BatchGetDocumentsResponse response) {
@@ -1055,12 +1055,13 @@ public WatchChange decodeWatchChange(ListenResponse protoChange) {
         SnapshotVersion version = decodeVersion(docChange.getDocument().getUpdateTime());
         hardAssert(
             !version.equals(SnapshotVersion.NONE), "Got a document change without an update time");
-        ObjectValue data = decodeFields(docChange.getDocument().getFieldsMap());
-        // The document may soon be re-serialized back to protos in order to store it in local
-        // persistence. Memoize the encoded form to avoid encoding it again.
         Document document =
             new Document(
-                key, version, data, Document.DocumentState.SYNCED, docChange.getDocument());
+                key,
+                version,
+                Document.DocumentState.SYNCED,
+                docChange.getDocument(),
+                this::decodeValue);
         watchChange = new WatchChange.DocumentChange(added, removed, document.getKey(), document);
         break;
       case DOCUMENT_DELETE:

firebase-firestore/src/test/java/com/google/firebase/firestore/DocumentSnapshotTest.java

@@ -45,12 +45,14 @@ public void testEquals() {
     assertNotEquals(base, differentData);
     assertNotEquals(base, fromCache);
 
+    // The assertions below that hash codes of different values are not equal is not something that
+    // we guarantee. In particular `base` and `differentData` have a hash collision because we
+    // don't use data in the hashCode.
     assertEquals(base.hashCode(), baseDup.hashCode());
     assertEquals(noData.hashCode(), noDataDup.hashCode());
     assertNotEquals(base.hashCode(), noData.hashCode());
     assertNotEquals(noData.hashCode(), base.hashCode());
     assertNotEquals(base.hashCode(), differentPath.hashCode());
-    assertNotEquals(base.hashCode(), differentData.hashCode());
     assertNotEquals(base.hashCode(), fromCache.hashCode());
   }
 }

firebase-firestore/src/test/java/com/google/firebase/firestore/QuerySnapshotTest.java

@@ -71,9 +71,10 @@ public void testEquals() {
     assertNotEquals(foo, noPendingWrites);
     assertNotEquals(foo, fromCache);
 
+    // Note: `foo` and `differentDoc` have the same hash code since we no longer take document
+    // contents into account.
     assertEquals(foo.hashCode(), fooDup.hashCode());
     assertNotEquals(foo.hashCode(), differentPath.hashCode());
-    assertNotEquals(foo.hashCode(), differentDoc.hashCode());
     assertNotEquals(foo.hashCode(), noPendingWrites.hashCode());
     assertNotEquals(foo.hashCode(), fromCache.hashCode());
   }

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

@@ -229,15 +229,15 @@ public void testHandlesSetMutation() {
     assertChanged(doc("foo/bar", 0, map("foo", "bar"), Document.DocumentState.LOCAL_MUTATIONS));
     assertContains(doc("foo/bar", 0, map("foo", "bar"), Document.DocumentState.LOCAL_MUTATIONS));
 
-    acknowledgeMutation(0);
-    assertChanged(doc("foo/bar", 0, map("foo", "bar"), Document.DocumentState.COMMITTED_MUTATIONS));
+    acknowledgeMutation(1);
+    assertChanged(doc("foo/bar", 1, map("foo", "bar"), Document.DocumentState.COMMITTED_MUTATIONS));
     if (garbageCollectorIsEager()) {
       // Nothing is pinning this anymore, as it has been acknowledged and there are no targets
       // active.
       assertNotContains("foo/bar");
     } else {
       assertContains(
-          doc("foo/bar", 0, map("foo", "bar"), Document.DocumentState.COMMITTED_MUTATIONS));
+          doc("foo/bar", 1, map("foo", "bar"), Document.DocumentState.COMMITTED_MUTATIONS));
     }
   }
 

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

@@ -553,7 +553,7 @@ public void testRemoveTargetsThenGC() {
         () -> {
           SnapshotVersion newVersion = version(3);
           Document doc =
-              new Document(middleDocToUpdate, newVersion, testValue, Document.DocumentState.SYNCED);
+              new Document(middleDocToUpdate, newVersion, Document.DocumentState.SYNCED, testValue);
           documentCache.add(doc);
           updateTargetInTransaction(middleTarget);
         });

firebase-firestore/src/test/java/com/google/firebase/firestore/model/DocumentTest.java

@@ -36,10 +36,10 @@
 public class DocumentTest {
 
   @Test
-  public void testConstructor() {
+  public void testInstantiation() {
     Document document =
         new Document(
-            key("messages/first"), version(1), wrapObject("a", 1), Document.DocumentState.SYNCED);
+            key("messages/first"), version(1), Document.DocumentState.SYNCED, wrapObject("a", 1));
 
     assertEquals(key("messages/first"), document.getKey());
     assertEquals(version(1), document.getVersion());
@@ -56,7 +56,7 @@ public void testExtractFields() {
             "owner",
             map("name", "Jonny", "title", "scallywag"));
     Document document =
-        new Document(key("rooms/eros"), version(1), data, Document.DocumentState.SYNCED);
+        new Document(key("rooms/eros"), version(1), Document.DocumentState.SYNCED, data);
 
     assertEquals("Discuss all the project related stuff", document.getFieldValue(field("desc")));
     assertEquals("scallywag", document.getFieldValue(field("owner.title")));