Don't deserialize full Document proto for Query execution (#2115)

Author: schmidt-sebastian

packages/firestore/CHANGELOG.md

@@ -1,4 +1,3 @@
-
 # Unreleased (1.5.0)
 - [feature] Added a `Firestore.waitForPendingWrites()` method that
   allows users to wait until all pending writes are acknowledged by the
@@ -7,6 +6,8 @@
   the instance, releasing any held resources. Once it completes, you can
   optionally call `Firestore.clearPersistence()` to wipe persisted Firestore
   data from disk.
+- [changed] Improved performance for queries with filters that only return a
+  small subset of the documents in a collection.
 
 # 1.4.10
 - [changed] Transactions now perform exponential backoff before retrying.
@@ -25,6 +26,7 @@
   match the query (https://github.com/firebase/firebase-android-sdk/issues/155).
 
 # 1.4.4
+>>>>>>> master
 - [fixed] Fixed an internal assertion that was triggered when an update
    with a `FieldValue.serverTimestamp()` and an update with a
   `FieldValue.increment()` were pending for the same document.

packages/firestore/src/api/database.ts

@@ -1276,7 +1276,7 @@ export class DocumentSnapshot implements firestore.DocumentSnapshot {
     return !this._document
       ? undefined
       : this.convertObject(
-          this._document.data,
+          this._document.data(),
           FieldValueOptions.fromSnapshotOptions(
             options,
             this._firestore._areTimestampsInSnapshotsEnabled()
@@ -1291,9 +1291,9 @@ export class DocumentSnapshot implements firestore.DocumentSnapshot {
     validateBetweenNumberOfArgs('DocumentSnapshot.get', arguments, 1, 2);
     options = validateSnapshotOptions('DocumentSnapshot.get', options);
     if (this._document) {
-      const value = this._document.data.field(
-        fieldPathFromArgument('DocumentSnapshot.get', fieldPath)
-      );
+      const value = this._document
+        .data()
+        .field(fieldPathFromArgument('DocumentSnapshot.get', fieldPath));
       if (value !== null) {
         return this.convertValue(
           value,

packages/firestore/src/core/query.ts

@@ -759,7 +759,7 @@ export class Bound {
       } else {
         const docValue = doc.field(orderByComponent.field);
         assert(
-          docValue !== undefined,
+          docValue !== null,
           'Field should exist since document matched the orderBy already.'
         );
         comparison = component.compareTo(docValue!);

packages/firestore/src/core/view.ts

@@ -169,7 +169,7 @@ export class View {
 
         // Calculate change
         if (oldDoc && newDoc) {
-          const docsEqual = oldDoc.data.isEqual(newDoc.data);
+          const docsEqual = oldDoc.data().isEqual(newDoc.data());
           if (!docsEqual) {
             if (!this.shouldWaitForSyncedDocument(oldDoc, newDoc)) {
               changeSet.track({

packages/firestore/src/model/document.ts

@@ -16,13 +16,14 @@
  */
 
 import { SnapshotVersion } from '../core/snapshot_version';
-import { fail } from '../util/assert';
+import { assert, fail } from '../util/assert';
 
 import { DocumentKey } from './document_key';
 import { FieldValue, JsonObject, ObjectValue } from './field_value';
 import { FieldPath } from './path';
 
 import * as api from '../protos/firestore_proto_api';
+import * as obj from '../util/obj';
 
 export interface DocumentOptions {
   hasLocalMutations?: boolean;
@@ -59,43 +60,91 @@ export class Document extends MaybeDocument {
   readonly hasLocalMutations: boolean;
   readonly hasCommittedMutations: boolean;
 
+  /**
+   * A cache of canonicalized FieldPaths to FieldValues that have already been
+   * deserialized in `getField()`.
+   */
+  private fieldValueCache?: Map<string, FieldValue | null>;
+
   constructor(
     key: DocumentKey,
     version: SnapshotVersion,
-    readonly data: ObjectValue,
     options: DocumentOptions,
-    /**
-     * Memoized serialized form of the document for optimization purposes (avoids repeated
-     * serialization). Might be undefined.
-     */
-    readonly proto?: api.Document
+    private objectValue?: ObjectValue,
+    readonly proto?: api.Document,
+    private readonly converter?: (value: api.Value) => FieldValue
   ) {
     super(key, version);
+    assert(
+      this.objectValue !== undefined ||
+        (this.proto !== undefined && this.converter !== undefined),
+      'If objectValue is not defined, proto and converter need to be set.'
+    );
+
     this.hasLocalMutations = !!options.hasLocalMutations;
     this.hasCommittedMutations = !!options.hasCommittedMutations;
   }
 
   field(path: FieldPath): FieldValue | null {
-    return this.data.field(path);
+    if (this.objectValue) {
+      return this.objectValue.field(path);
+    } else {
+      if (!this.fieldValueCache) {
+        // TODO(b/136090445): Remove the cache when `getField` is no longer
+        // called during Query ordering.
+        this.fieldValueCache = new Map<string, FieldValue>();
+      }
+
+      const canonicalPath = path.canonicalString();
+
+      let fieldValue = this.fieldValueCache.get(canonicalPath);
+
+      if (fieldValue === undefined) {
+        // 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.
+        const protoValue = this.getProtoField(path);
+        if (protoValue === undefined) {
+          fieldValue = null;
+        } else {
+          fieldValue = this.converter!(protoValue);
+        }
+        this.fieldValueCache.set(canonicalPath, fieldValue);
+      }
+
+      return fieldValue!;
+    }
   }
 
-  fieldValue(path: FieldPath): unknown {
-    const field = this.field(path);
-    return field ? field.value() : undefined;
+  data(): ObjectValue {
+    if (!this.objectValue) {
+      let result = ObjectValue.EMPTY;
+      obj.forEach(this.proto!.fields, (key: string, value: api.Value) => {
+        result = result.set(new FieldPath([key]), this.converter!(value));
+      });
+      this.objectValue = result;
+
+      // Once objectValue is computed, values inside the fieldValueCache are no
+      // longer accessed.
+      this.fieldValueCache = undefined;
+    }
+
+    return this.objectValue;
   }
 
   value(): JsonObject<unknown> {
-    return this.data.value();
+    return this.data().value();
   }
 
   isEqual(other: MaybeDocument | null | undefined): boolean {
     return (
       other instanceof Document &&
       this.key.isEqual(other.key) &&
       this.version.isEqual(other.version) &&
-      this.data.isEqual(other.data) &&
       this.hasLocalMutations === other.hasLocalMutations &&
-      this.hasCommittedMutations === other.hasCommittedMutations
+      this.hasCommittedMutations === other.hasCommittedMutations &&
+      this.data().isEqual(other.data())
     );
   }
 
@@ -111,6 +160,29 @@ export class Document extends MaybeDocument {
     return this.hasLocalMutations || this.hasCommittedMutations;
   }
 
+  /**
+   * Returns the nested Protobuf value for 'path`. Can only be called if
+   * `proto` was provided at construction time.
+   */
+  private getProtoField(path: FieldPath): api.Value | undefined {
+    assert(
+      this.proto !== undefined,
+      'Can only call getProtoField() when proto is defined'
+    );
+
+    let protoValue: api.Value | undefined = this.proto!.fields[
+      path.firstSegment()
+    ];
+    for (let i = 1; i < path.length; ++i) {
+      if (!protoValue || !protoValue.mapValue) {
+        return undefined;
+      }
+      protoValue = protoValue.mapValue.fields[path.get(i)];
+    }
+
+    return protoValue;
+  }
+
   static compareByField(field: FieldPath, d1: Document, d2: Document): number {
     const v1 = d1.field(field);
     const v2 = d2.field(field);

packages/firestore/src/model/mutation.ts

@@ -353,9 +353,14 @@ export class SetMutation extends Mutation {
     // have held.
 
     const version = mutationResult.version;
-    return new Document(this.key, version, this.value, {
-      hasCommittedMutations: true
-    });
+    return new Document(
+      this.key,
+      version,
+      {
+        hasCommittedMutations: true
+      },
+      this.value
+    );
   }
 
   applyToLocalView(
@@ -370,9 +375,14 @@ export class SetMutation extends Mutation {
     }
 
     const version = Mutation.getPostMutationVersion(maybeDoc);
-    return new Document(this.key, version, this.value, {
-      hasLocalMutations: true
-    });
+    return new Document(
+      this.key,
+      version,
+      {
+        hasLocalMutations: true
+      },
+      this.value
+    );
   }
 
   extractBaseValue(maybeDoc: MaybeDocument | null): null {
@@ -434,9 +444,14 @@ export class PatchMutation extends Mutation {
     }
 
     const newData = this.patchDocument(maybeDoc);
-    return new Document(this.key, mutationResult.version, newData, {
-      hasCommittedMutations: true
-    });
+    return new Document(
+      this.key,
+      mutationResult.version,
+      {
+        hasCommittedMutations: true
+      },
+      newData
+    );
   }
 
   applyToLocalView(
@@ -452,9 +467,14 @@ export class PatchMutation extends Mutation {
 
     const version = Mutation.getPostMutationVersion(maybeDoc);
     const newData = this.patchDocument(maybeDoc);
-    return new Document(this.key, version, newData, {
-      hasLocalMutations: true
-    });
+    return new Document(
+      this.key,
+      version,
+      {
+        hasLocalMutations: true
+      },
+      newData
+    );
   }
 
   extractBaseValue(maybeDoc: MaybeDocument | null): null {
@@ -478,7 +498,7 @@ export class PatchMutation extends Mutation {
   private patchDocument(maybeDoc: MaybeDocument | null): ObjectValue {
     let data: ObjectValue;
     if (maybeDoc instanceof Document) {
-      data = maybeDoc.data;
+      data = maybeDoc.data();
     } else {
       data = ObjectValue.EMPTY;
     }
@@ -550,10 +570,15 @@ export class TransformMutation extends Mutation {
     );
 
     const version = mutationResult.version;
-    const newData = this.transformObject(doc.data, transformResults);
-    return new Document(this.key, version, newData, {
-      hasCommittedMutations: true
-    });
+    const newData = this.transformObject(doc.data(), transformResults);
+    return new Document(
+      this.key,
+      version,
+      {
+        hasCommittedMutations: true
+      },
+      newData
+    );
   }
 
   applyToLocalView(
@@ -573,10 +598,15 @@ export class TransformMutation extends Mutation {
       maybeDoc,
       baseDoc
     );
-    const newData = this.transformObject(doc.data, transformResults);
-    return new Document(this.key, doc.version, newData, {
-      hasLocalMutations: true
-    });
+    const newData = this.transformObject(doc.data(), transformResults);
+    return new Document(
+      this.key,
+      doc.version,
+      {
+        hasLocalMutations: true
+      },
+      newData
+    );
   }
 
   extractBaseValue(maybeDoc: MaybeDocument | null): ObjectValue | null {