diff --git a/__tests__/types/DefaultMigrator.ts b/__tests__/types/DefaultMigrator.test-d.ts
similarity index 95%
rename from __tests__/types/DefaultMigrator.ts
rename to __tests__/types/DefaultMigrator.test-d.ts
index 97c1029..daab63a 100644
--- a/__tests__/types/DefaultMigrator.ts
+++ b/__tests__/types/DefaultMigrator.test-d.ts
@@ -60,8 +60,7 @@ defaultMigrator.setWithDerivedData((doc) => {
 });
 defaultMigrator.setWithDerivedData(
   (doc) => {
-    // TODO: We probably want firestore.QueryDocumentSnapshot<D> here
-    expectType<firestore.QueryDocumentSnapshot<Partial<D>>>(doc);
+    expectType<firestore.QueryDocumentSnapshot<D>>(doc);
     return { num: 0 };
   },
   { merge: true }
diff --git a/src/api/createBatchMigrator.ts b/src/api/createBatchMigrator.ts
index d0aea6a..185a92c 100644
--- a/src/api/createBatchMigrator.ts
+++ b/src/api/createBatchMigrator.ts
@@ -48,7 +48,7 @@ export function createBatchMigrator<D = firestore.DocumentData>(
   traversalConfig?: Partial<TraversalConfig>
 ): BatchMigrator<D> {
   const traverser = isTraverser(traversableOrTraverser)
-    ? traversableOrTraverser
-    : createTraverser(traversableOrTraverser, traversalConfig);
+    ? (traversableOrTraverser as Traverser<D>)
+    : createTraverser(traversableOrTraverser as Traversable<D>, traversalConfig);
   return new BasicBatchMigratorImpl(traverser);
 }
diff --git a/src/api/createMigrator.ts b/src/api/createMigrator.ts
index 3ce1150..442e257 100644
--- a/src/api/createMigrator.ts
+++ b/src/api/createMigrator.ts
@@ -37,7 +37,7 @@ export function createMigrator<D = firestore.DocumentData>(
   traversalConfig?: Partial<TraversalConfig>
 ): DefaultMigrator<D> {
   const traverser = isTraverser(traversableOrTraverser)
-    ? traversableOrTraverser
-    : createTraverser(traversableOrTraverser, traversalConfig);
+    ? (traversableOrTraverser as Traverser<D>)
+    : createTraverser(traversableOrTraverser as Traversable<D>, traversalConfig);
   return new BasicDefaultMigratorImpl(traverser);
 }
diff --git a/src/api/interfaces/Migrator.ts b/src/api/interfaces/Migrator.ts
index c06dc93..de6f245 100644
--- a/src/api/interfaces/Migrator.ts
+++ b/src/api/interfaces/Migrator.ts
@@ -5,6 +5,7 @@ import type {
   MigrationResult,
   SetDataGetter,
   SetOptions,
+  SetPartialDataGetter,
   Traverser,
   UpdateDataGetter,
   UpdateFieldValueGetter,
@@ -208,9 +209,10 @@ export interface Migrator<D = firestore.DocumentData> {
    * - _SC_(`traverser`): space complexity of the underlying traverser
    *
    * @param data - A data object with which to set each document.
+   * @param options - An object to configure the set behavior.
    * @returns A Promise resolving to an object representing the details of the migration.
    */
-  set(data: D): Promise<MigrationResult>;
+  set(data: firestore.PartialWithFieldValue<D>, options: SetOptions): Promise<MigrationResult>;
 
   /**
    * Sets all documents in this collection with the provided data.
@@ -232,10 +234,9 @@ export interface Migrator<D = firestore.DocumentData> {
    * - _SC_(`traverser`): space complexity of the underlying traverser
    *
    * @param data - A data object with which to set each document.
-   * @param options - An object to configure the set behavior.
    * @returns A Promise resolving to an object representing the details of the migration.
    */
-  set(data: Partial<D>, options: SetOptions): Promise<MigrationResult>;
+  set(data: firestore.WithFieldValue<D>): Promise<MigrationResult>;
 
   /**
    * Sets all documents in this collection with the provided data.
@@ -258,9 +259,13 @@ export interface Migrator<D = firestore.DocumentData> {
    *
    * @param getData - A function that takes a document snapshot and returns a data object with
    * which to set each document.
+   * @param options - An object to configure the set behavior.
    * @returns A Promise resolving to an object representing the details of the migration.
    */
-  setWithDerivedData(getData: SetDataGetter<D>): Promise<MigrationResult>;
+  setWithDerivedData(
+    getData: SetPartialDataGetter<D>,
+    options: SetOptions
+  ): Promise<MigrationResult>;
 
   /**
    * Sets all documents in this collection with the provided data.
@@ -283,13 +288,9 @@ export interface Migrator<D = firestore.DocumentData> {
    *
    * @param getData - A function that takes a document snapshot and returns a data object with
    * which to set each document.
-   * @param options - An object to configure the set behavior.
    * @returns A Promise resolving to an object representing the details of the migration.
    */
-  setWithDerivedData(
-    getData: SetDataGetter<Partial<D>>,
-    options: SetOptions
-  ): Promise<MigrationResult>;
+  setWithDerivedData(getData: SetDataGetter<D>): Promise<MigrationResult>;
 
   /**
    * Updates all documents in this collection with the provided data.
diff --git a/src/api/interfaces/SetDataGetter.ts b/src/api/interfaces/SetDataGetter.ts
index 2cb6fb5..249ff1e 100644
--- a/src/api/interfaces/SetDataGetter.ts
+++ b/src/api/interfaces/SetDataGetter.ts
@@ -5,4 +5,4 @@ import type { firestore } from 'firebase-admin';
  */
 export type SetDataGetter<D = firestore.DocumentData> = (
   doc: firestore.QueryDocumentSnapshot<D>
-) => D;
+) => firestore.WithFieldValue<D>;
diff --git a/src/api/interfaces/SetPartialDataGetter.ts b/src/api/interfaces/SetPartialDataGetter.ts
new file mode 100644
index 0000000..8f76a2d
--- /dev/null
+++ b/src/api/interfaces/SetPartialDataGetter.ts
@@ -0,0 +1,8 @@
+import type { firestore } from 'firebase-admin';
+
+/**
+ * A function that takes a document snapshot and derives the data with which to partially set that document.
+ */
+export type SetPartialDataGetter<D = firestore.DocumentData> = (
+  doc: firestore.QueryDocumentSnapshot<D>
+) => firestore.PartialWithFieldValue<D>;
diff --git a/src/api/interfaces/index.ts b/src/api/interfaces/index.ts
index 18a56b9..83f5ee7 100644
--- a/src/api/interfaces/index.ts
+++ b/src/api/interfaces/index.ts
@@ -7,6 +7,7 @@ export type { MigrationResult } from './MigrationResult';
 export type { Migrator } from './Migrator';
 export type { SetDataGetter } from './SetDataGetter';
 export type { SetOptions } from './SetOptions';
+export type { SetPartialDataGetter } from './SetPartialDataGetter';
 export type { Traversable } from './Traversable';
 export type { TraversalConfig } from './TraversalConfig';
 export type { TraversalResult } from './TraversalResult';
diff --git a/src/internal/implementations/BasicBatchMigratorImpl.ts b/src/internal/implementations/BasicBatchMigratorImpl.ts
index 564cbd7..5ac3fee 100644
--- a/src/internal/implementations/BasicBatchMigratorImpl.ts
+++ b/src/internal/implementations/BasicBatchMigratorImpl.ts
@@ -5,6 +5,7 @@ import type {
   MigrationResult,
   SetDataGetter,
   SetOptions,
+  SetPartialDataGetter,
   Traverser,
   UpdateDataGetter,
   UpdateFieldValueGetter,
@@ -53,41 +54,47 @@ export class BasicBatchMigratorImpl<D> extends AbstractMigrator<D> implements Ba
     );
   }
 
-  public set(data: D): Promise<MigrationResult>;
+  public set(
+    data: firestore.PartialWithFieldValue<D>,
+    options: SetOptions
+  ): Promise<MigrationResult>;
 
-  public set(data: Partial<D>, options: SetOptions): Promise<MigrationResult>;
+  public set(data: firestore.WithFieldValue<D>): Promise<MigrationResult>;
 
-  public async set(data: D | Partial<D>, options?: SetOptions): Promise<MigrationResult> {
+  public async set(
+    data: firestore.PartialWithFieldValue<D> | firestore.WithFieldValue<D>,
+    options?: SetOptions
+  ): Promise<MigrationResult> {
     return this.#migrate((writeBatch, doc) => {
       if (options === undefined) {
-        // Signature 1
-        writeBatch.set(doc.ref, data as D);
-      } else {
         // Signature 2
-        writeBatch.set(doc.ref, data as Partial<D>, options);
+        writeBatch.set(doc.ref, data as firestore.WithFieldValue<D>);
+      } else {
+        // Signature 1
+        writeBatch.set(doc.ref, data as firestore.PartialWithFieldValue<D>, options);
       }
     });
   }
 
-  public setWithDerivedData(getData: SetDataGetter<D>): Promise<MigrationResult>;
-
   public setWithDerivedData(
-    getData: SetDataGetter<Partial<D>>,
+    getData: SetPartialDataGetter<D>,
     options: SetOptions
   ): Promise<MigrationResult>;
 
-  public setWithDerivedData(
-    getData: SetDataGetter<D> | SetDataGetter<Partial<D>>,
+  public setWithDerivedData(getData: SetDataGetter<D>): Promise<MigrationResult>;
+
+  public async setWithDerivedData(
+    getData: SetPartialDataGetter<D> | SetDataGetter<D>,
     options?: SetOptions
   ): Promise<MigrationResult> {
     return this.#migrate((writeBatch, doc) => {
       if (options === undefined) {
-        // Signature 1
+        // Signature 2
         const data = (getData as SetDataGetter<D>)(doc);
         writeBatch.set(doc.ref, data);
       } else {
-        // Signature 2
-        const data = (getData as SetDataGetter<D | Partial<D>>)(doc);
+        // Signature 1
+        const data = (getData as SetPartialDataGetter<D>)(doc);
         writeBatch.set(doc.ref, data, options);
       }
     });
diff --git a/src/internal/implementations/BasicDefaultMigratorImpl.ts b/src/internal/implementations/BasicDefaultMigratorImpl.ts
index e60f9e5..cb77a31 100644
--- a/src/internal/implementations/BasicDefaultMigratorImpl.ts
+++ b/src/internal/implementations/BasicDefaultMigratorImpl.ts
@@ -5,6 +5,7 @@ import type {
   MigrationResult,
   SetDataGetter,
   SetOptions,
+  SetPartialDataGetter,
   Traverser,
   UpdateDataGetter,
   UpdateFieldValueGetter,
@@ -35,41 +36,47 @@ export class BasicDefaultMigratorImpl<D> extends AbstractMigrator<D> implements
     );
   }
 
-  public set(data: D): Promise<MigrationResult>;
+  public set(
+    data: firestore.PartialWithFieldValue<D>,
+    options: SetOptions
+  ): Promise<MigrationResult>;
 
-  public set(data: Partial<D>, options: SetOptions): Promise<MigrationResult>;
+  public set(data: firestore.WithFieldValue<D>): Promise<MigrationResult>;
 
-  public async set(data: D | Partial<D>, options?: SetOptions): Promise<MigrationResult> {
+  public async set(
+    data: firestore.PartialWithFieldValue<D> | firestore.WithFieldValue<D>,
+    options?: SetOptions
+  ): Promise<MigrationResult> {
     return this.#migrate(async (doc) => {
       if (options === undefined) {
-        // Signature 1
-        await doc.ref.set(data as D);
-      } else {
         // Signature 2
-        await doc.ref.set(data as Partial<D>, options);
+        await doc.ref.set(data as firestore.WithFieldValue<D>);
+      } else {
+        // Signature 1
+        await doc.ref.set(data as firestore.PartialWithFieldValue<D>, options);
       }
     });
   }
 
-  public setWithDerivedData(getData: SetDataGetter<D>): Promise<MigrationResult>;
-
   public setWithDerivedData(
-    getData: SetDataGetter<Partial<D>>,
+    getData: SetPartialDataGetter<D>,
     options: SetOptions
   ): Promise<MigrationResult>;
 
+  public setWithDerivedData(getData: SetDataGetter<D>): Promise<MigrationResult>;
+
   public async setWithDerivedData(
-    getData: SetDataGetter<D> | SetDataGetter<Partial<D>>,
+    getData: SetPartialDataGetter<D> | SetDataGetter<D>,
     options?: SetOptions
   ): Promise<MigrationResult> {
     return this.#migrate(async (doc) => {
       if (options === undefined) {
-        // Signature 1
+        // Signature 2
         const data = (getData as SetDataGetter<D>)(doc);
         await doc.ref.set(data);
       } else {
-        // Signature 2
-        const data = (getData as SetDataGetter<D | Partial<D>>)(doc);
+        // Signature 1
+        const data = (getData as SetPartialDataGetter<D>)(doc);
         await doc.ref.set(data, options);
       }
     });
diff --git a/src/internal/implementations/abstract/AbstractMigrator.ts b/src/internal/implementations/abstract/AbstractMigrator.ts
index c5263fb..c46056d 100644
--- a/src/internal/implementations/abstract/AbstractMigrator.ts
+++ b/src/internal/implementations/abstract/AbstractMigrator.ts
@@ -6,6 +6,7 @@ import type {
   Migrator,
   SetDataGetter,
   SetOptions,
+  SetPartialDataGetter,
   Traverser,
   UpdateDataGetter,
   UpdateFieldValueGetter,
@@ -109,17 +110,20 @@ export abstract class AbstractMigrator<D> implements Migrator<D> {
 
   public abstract withTraverser(traverser: Traverser<D>): Migrator<D>;
 
-  public abstract set(data: D): Promise<MigrationResult>;
-
-  public abstract set(data: Partial<D>, options: SetOptions): Promise<MigrationResult>;
+  public abstract set(
+    data: firestore.PartialWithFieldValue<D>,
+    options: SetOptions
+  ): Promise<MigrationResult>;
 
-  public abstract setWithDerivedData(getData: SetDataGetter<D>): Promise<MigrationResult>;
+  public abstract set(data: firestore.WithFieldValue<D>): Promise<MigrationResult>;
 
   public abstract setWithDerivedData(
-    getData: SetDataGetter<Partial<D>>,
+    getData: SetPartialDataGetter<D>,
     options: SetOptions
   ): Promise<MigrationResult>;
 
+  public abstract setWithDerivedData(getData: SetDataGetter<D>): Promise<MigrationResult>;
+
   public abstract update<D>(
     data: firestore.UpdateData<D>,
     precondition?: firestore.Precondition