diff --git a/src/Document.php b/src/Document.php
index dc7098b..95cc609 100644
--- a/src/Document.php
+++ b/src/Document.php
@@ -313,14 +313,24 @@ public function getValue()
      *
      * @return Reference
      */
-    public function set($value): self
+    public function set($value, $merge = false): self
     {
-        $value = [
+        $payload = [
             'name' => basename($this->uri->getPath()),
-            'fields' => $this->valueMapper->encodeValues($value)
+            'fields' => $this->valueMapper->encodeValues($value),
         ];
 
-        $this->apiClient->set($this->uri, $value);
+        if ($merge) {
+            $paths = $this->valueMapper->encodeFieldPaths($value);
+            $prefix = '&updateMask.fieldPaths=';
+            $query = $prefix . implode($prefix, $paths);
+            $uri = $this->uri->withQuery("updateMask.fieldPaths=message$query");
+        }
+        else {
+            $uri = $this->uri;
+        }
+
+        $this->apiClient->set($uri, $payload);
 
         return $this;
     }
diff --git a/tests/DocumentTest.php b/tests/DocumentTest.php
index 3088bac..44c4b0d 100644
--- a/tests/DocumentTest.php
+++ b/tests/DocumentTest.php
@@ -33,20 +33,18 @@ public function testSetAndGet()
         }
     }
 
-    public function testUpdate()
+    public function testSetMergeAndGet()
     {
-        $this->markTestIncomplete('This test has not been implemented yet.');
-
-        $ref = $this->ref->getChild(__FUNCTION__);
-        $ref->set([
+        $doc = $this->collection->getDocument(__FUNCTION__);
+        $doc->set([
             'first' => 'value',
             'second' => 'value',
         ]);
 
-        $ref->update([
+        $doc->set([
             'first' => 'updated',
             'third' => 'new',
-        ]);
+        ], true);
 
         $expected = [
             'first' => 'updated',
@@ -54,7 +52,7 @@ public function testUpdate()
             'third' => 'new',
         ];
 
-        $this->assertEquals($expected, $ref->getValue());
+        $this->assertEquals($expected, $doc->getSnapshot()->data());
     }
 
     public function testPush()