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

firebase-firestore/CHANGELOG.md

@@ -1,4 +1,6 @@
 # Unreleased
+- [changed] Improved performance for queries that only return a subset of the
+  data 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

@@ -174,22 +174,23 @@ public static SnapshotVersion version(long versionMicros) {
   }
 
   public static Document doc(String key, long version, Map<String, Object> data) {
-    return new Document(
+    return Document.fromObjectValue(
         key(key), version(version), wrapObject(data), Document.DocumentState.SYNCED);
   }
 
   public static Document doc(DocumentKey key, long version, Map<String, Object> data) {
-    return new Document(key, version(version), wrapObject(data), Document.DocumentState.SYNCED);
+    return Document.fromObjectValue(
+        key, version(version), wrapObject(data), Document.DocumentState.SYNCED);
   }
 
   public static Document doc(
       String key, long version, ObjectValue data, Document.DocumentState documentState) {
-    return new Document(key(key), version(version), data, documentState);
+    return Document.fromObjectValue(key(key), version(version), data, documentState);
   }
 
   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 Document.fromObjectValue(key(key), version(version), wrapObject(data), documentState);
   }
 
   public static NoDocument deletedDoc(String key, long version) {

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

@@ -111,13 +111,9 @@ private com.google.firestore.v1.Document encodeDocument(Document document) {
   /** Decodes a Document proto to the equivalent model. */
   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,
+    return Document.fromProto(
+        this.rpcSerializer,
+        document,
         hasCommittedMutations
             ? Document.DocumentState.COMMITTED_MUTATIONS
             : Document.DocumentState.SYNCED);

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

@@ -14,16 +14,20 @@
 
 package com.google.firebase.firestore.model;
 
+import static com.google.firebase.firestore.util.Assert.hardAssert;
+
 import com.google.firebase.firestore.model.value.FieldValue;
 import com.google.firebase.firestore.model.value.ObjectValue;
+import com.google.firebase.firestore.remote.RemoteSerializer;
+import com.google.firestore.v1.Value;
 import java.util.Comparator;
 import javax.annotation.Nullable;
 
 /**
  * Represents a document in Firestore with a key, version, data and whether the data has local
  * mutations applied to it.
  */
-public final class Document extends MaybeDocument {
+public abstract class Document extends MaybeDocument {
 
   /** Describes the `hasPendingWrites` state of a document. */
   public enum DocumentState {
@@ -36,59 +40,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;
+  /** Creates a new Document from an existing ObjectValue. */
+  public static Document fromObjectValue(
+      DocumentKey key, SnapshotVersion version, ObjectValue data, DocumentState documentState) {
+    return new Document(key, version, documentState) {
+      @Nullable
+      @Override
+      public com.google.firestore.v1.Document getProto() {
+        return null;
+      }
+
+      @Override
+      public ObjectValue getData() {
+        return data;
+      }
+
+      @Nullable
+      @Override
+      public FieldValue getField(FieldPath path) {
+        return data.get(path);
+      }
+    };
+  }
 
   /**
-   * Memoized serialized form of the document for optimization purposes (avoids repeated
-   * serialization). Might be null.
+   * Creates a new Document from a Proto representation.
+   *
+   * <p>The Proto is only converted to an ObjectValue if the consumer calls `getData()`.
    */
-  private final com.google.firestore.v1.Document proto;
-
-  public @Nullable com.google.firestore.v1.Document getProto() {
-    return proto;
+  public static Document fromProto(
+      RemoteSerializer serializer,
+      com.google.firestore.v1.Document proto,
+      DocumentState documentState) {
+    DocumentKey key = serializer.decodeKey(proto.getName());
+    SnapshotVersion version = serializer.decodeVersion(proto.getUpdateTime());
+
+    hardAssert(
+        !version.equals(SnapshotVersion.NONE), "Found a document Proto with no snapshot version");
+
+    return new Document(key, version, documentState) {
+      private ObjectValue memoizedData = null;
+
+      @Override
+      public com.google.firestore.v1.Document getProto() {
+        return proto;
+      }
+
+      @Override
+      public ObjectValue getData() {
+        if (memoizedData != null) {
+          return memoizedData;
+        } else {
+          memoizedData = serializer.decodeFields(proto.getFieldsMap());
+          return memoizedData;
+        }
+      }
+
+      @Nullable
+      @Override
+      public FieldValue getField(FieldPath path) {
+        if (memoizedData != null) {
+          return memoizedData.get(path);
+        } else {
+          // 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 value = proto.getFieldsMap().get(path.getFirstSegment());
+          for (int i = 1; value != null && i < path.length(); ++i) {
+            if (value.getValueTypeCase() != Value.ValueTypeCase.MAP_VALUE) {
+              return null;
+            }
+            value = value.getMapValue().getFieldsMap().get(path.getSegment(i));
+          }
+          return value != null ? serializer.decodeValue(value) : null;
+        }
+      }
+    };
   }
 
-  public Document(
-      DocumentKey key, SnapshotVersion version, ObjectValue data, DocumentState documentState) {
-    super(key, version);
-    this.data = data;
-    this.documentState = documentState;
-    this.proto = null;
-  }
+  private final DocumentState documentState;
 
-  public Document(
-      DocumentKey key,
-      SnapshotVersion version,
-      ObjectValue data,
-      DocumentState documentState,
-      com.google.firestore.v1.Document proto) {
+  private Document(DocumentKey key, SnapshotVersion version, DocumentState documentState) {
     super(key, version);
-    this.data = data;
     this.documentState = documentState;
-    this.proto = proto;
   }
 
-  public ObjectValue getData() {
-    return data;
-  }
+  /**
+   * Memoized serialized form of the document for optimization purposes (avoids repeated
+   * serialization). Might be null.
+   */
+  public abstract @Nullable com.google.firestore.v1.Document getProto();
 
-  public @Nullable FieldValue getField(FieldPath path) {
-    return data.get(path);
-  }
+  public abstract ObjectValue getData();
+
+  public abstract @Nullable FieldValue getField(FieldPath path);
 
   public @Nullable Object getFieldValue(FieldPath path) {
     FieldValue value = getField(path);
@@ -113,7 +165,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 +174,12 @@ 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() {
     int result = getKey().hashCode();
-    result = 31 * result + data.hashCode();
     result = 31 * result + getVersion().hashCode();
     result = 31 * result + documentState.hashCode();
     return result;
@@ -140,7 +191,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,8 @@ public MaybeDocument applyToRemoteDocument(
 
     SnapshotVersion version = mutationResult.getVersion();
     ObjectValue newData = patchDocument(maybeDoc);
-    return new Document(getKey(), version, newData, Document.DocumentState.COMMITTED_MUTATIONS);
+    return Document.fromObjectValue(
+        getKey(), version, newData, Document.DocumentState.COMMITTED_MUTATIONS);
   }
 
   @Nullable
@@ -127,7 +128,8 @@ public MaybeDocument applyToLocalView(
 
     SnapshotVersion version = getPostMutationVersion(maybeDoc);
     ObjectValue newData = patchDocument(maybeDoc);
-    return new Document(getKey(), version, newData, Document.DocumentState.LOCAL_MUTATIONS);
+    return Document.fromObjectValue(
+        getKey(), version, newData, Document.DocumentState.LOCAL_MUTATIONS);
   }
 
   @Nullable

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

@@ -72,7 +72,8 @@ 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 Document.fromObjectValue(
+        getKey(), version, value, Document.DocumentState.COMMITTED_MUTATIONS);
   }
 
   @Nullable
@@ -86,7 +87,8 @@ public MaybeDocument applyToLocalView(
     }
 
     SnapshotVersion version = getPostMutationVersion(maybeDoc);
-    return new Document(getKey(), version, value, Document.DocumentState.LOCAL_MUTATIONS);
+    return Document.fromObjectValue(
+        getKey(), version, value, Document.DocumentState.LOCAL_MUTATIONS);
   }
 
   @Nullable

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

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

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

@@ -405,12 +405,7 @@ private Document decodeFoundDocument(BatchGetDocumentsResponse response) {
     Assert.hardAssert(
         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 Document.fromProto(this, response.getFound(), Document.DocumentState.SYNCED);
   }
 
   private NoDocument decodeMissingDocument(BatchGetDocumentsResponse response) {
@@ -1041,24 +1036,16 @@ public WatchChange decodeWatchChange(ListenResponse protoChange) {
         DocumentChange docChange = protoChange.getDocumentChange();
         List<Integer> added = docChange.getTargetIdsList();
         List<Integer> removed = docChange.getRemovedTargetIdsList();
-        DocumentKey key = decodeKey(docChange.getDocument().getName());
-        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());
+            Document.fromProto(this, docChange.getDocument(), Document.DocumentState.SYNCED);
         watchChange = new WatchChange.DocumentChange(added, removed, document.getKey(), document);
         break;
       case DOCUMENT_DELETE:
         DocumentDelete docDelete = protoChange.getDocumentDelete();
         removed = docDelete.getRemovedTargetIdsList();
-        key = decodeKey(docDelete.getDocument());
+        DocumentKey key = decodeKey(docDelete.getDocument());
         // Note that version might be unset in which case we use SnapshotVersion.NONE
-        version = decodeVersion(docDelete.getReadTime());
+        SnapshotVersion version = decodeVersion(docDelete.getReadTime());
         NoDocument doc = new NoDocument(key, version, /*hasCommittedMutations=*/ false);
         watchChange =
             new WatchChange.DocumentChange(Collections.emptyList(), removed, doc.getKey(), doc);

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

@@ -45,12 +45,13 @@ public void testEquals() {
     assertNotEquals(base, differentData);
     assertNotEquals(base, fromCache);
 
+    // Note: `base` and `differentData` have the same hash code since we no longer take document
+    // contents into account.
     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));
     }
   }