diff --git a/data/src/main/java/org/apache/iceberg/data/PartitionStatsHandler.java b/data/src/main/java/org/apache/iceberg/data/PartitionStatsHandler.java
index a194a0825557..7e8eff294fa8 100644
--- a/data/src/main/java/org/apache/iceberg/data/PartitionStatsHandler.java
+++ b/data/src/main/java/org/apache/iceberg/data/PartitionStatsHandler.java
@@ -36,7 +36,6 @@
 import org.apache.iceberg.Partitioning;
 import org.apache.iceberg.Schema;
 import org.apache.iceberg.Snapshot;
-import org.apache.iceberg.SnapshotRef;
 import org.apache.iceberg.StructLike;
 import org.apache.iceberg.Table;
 import org.apache.iceberg.avro.Avro;
@@ -53,7 +52,6 @@
 import org.apache.iceberg.types.Types.LongType;
 import org.apache.iceberg.types.Types.NestedField;
 import org.apache.iceberg.types.Types.StructType;
-import org.apache.iceberg.util.SnapshotUtil;
 
 /**
  * Computes, writes and reads the {@link PartitionStatisticsFile}. Uses generic readers and writers
@@ -121,29 +119,27 @@ public static Schema schema(StructType unifiedPartitionType) {
    *     present.
    */
   public static PartitionStatisticsFile computeAndWriteStatsFile(Table table) throws IOException {
-    return computeAndWriteStatsFile(table, SnapshotRef.MAIN_BRANCH);
+    if (table.currentSnapshot() == null) {
+      return null;
+    }
+
+    return computeAndWriteStatsFile(table, table.currentSnapshot().snapshotId());
   }
 
   /**
-   * Computes and writes the {@link PartitionStatisticsFile} for a given table and branch.
+   * Computes and writes the {@link PartitionStatisticsFile} for a given table and snapshot.
    *
    * @param table The {@link Table} for which the partition statistics is computed.
-   * @param branch A branch information to select the required snapshot.
-   * @return {@link PartitionStatisticsFile} for the given branch, or null if no statistics are
+   * @param snapshotId snapshot for which partition statistics are computed.
+   * @return {@link PartitionStatisticsFile} for the given snapshot, or null if no statistics are
    *     present.
    */
-  public static PartitionStatisticsFile computeAndWriteStatsFile(Table table, String branch)
+  public static PartitionStatisticsFile computeAndWriteStatsFile(Table table, long snapshotId)
       throws IOException {
-    Snapshot currentSnapshot = SnapshotUtil.latestSnapshot(table, branch);
-    if (currentSnapshot == null) {
-      Preconditions.checkArgument(
-          branch == null || branch.equals(SnapshotRef.MAIN_BRANCH),
-          "Couldn't find the snapshot for the branch %s",
-          branch);
-      return null;
-    }
+    Snapshot snapshot = table.snapshot(snapshotId);
+    Preconditions.checkArgument(snapshot != null, "Snapshot not found: %s", snapshotId);
 
-    Collection<PartitionStats> stats = PartitionStatsUtil.computeStats(table, currentSnapshot);
+    Collection<PartitionStats> stats = PartitionStatsUtil.computeStats(table, snapshot);
     if (stats.isEmpty()) {
       return null;
     }
@@ -151,7 +147,7 @@ public static PartitionStatisticsFile computeAndWriteStatsFile(Table table, Stri
     StructType partitionType = Partitioning.partitionType(table);
     List<PartitionStats> sortedStats = PartitionStatsUtil.sortStats(stats, partitionType);
     return writePartitionStatsFile(
-        table, currentSnapshot.snapshotId(), schema(partitionType), sortedStats);
+        table, snapshot.snapshotId(), schema(partitionType), sortedStats);
   }
 
   @VisibleForTesting
diff --git a/data/src/test/java/org/apache/iceberg/data/TestPartitionStatsHandler.java b/data/src/test/java/org/apache/iceberg/data/TestPartitionStatsHandler.java
index 8256f80d9be9..1d84b8e229eb 100644
--- a/data/src/test/java/org/apache/iceberg/data/TestPartitionStatsHandler.java
+++ b/data/src/test/java/org/apache/iceberg/data/TestPartitionStatsHandler.java
@@ -104,17 +104,17 @@ public void testPartitionStatsOnEmptyTable() throws Exception {
   public void testPartitionStatsOnEmptyBranch() throws Exception {
     Table testTable = TestTables.create(tempDir("empty_branch"), "empty_branch", SCHEMA, SPEC, 2);
     testTable.manageSnapshots().createBranch("b1").commit();
-    assertThat(PartitionStatsHandler.computeAndWriteStatsFile(testTable, "b1")).isNull();
+    long branchSnapshot = testTable.refs().get("b1").snapshotId();
+    assertThat(PartitionStatsHandler.computeAndWriteStatsFile(testTable, branchSnapshot)).isNull();
   }
 
   @Test
   public void testPartitionStatsOnInvalidSnapshot() throws Exception {
     Table testTable =
         TestTables.create(tempDir("invalid_snapshot"), "invalid_snapshot", SCHEMA, SPEC, 2);
-    assertThatThrownBy(
-            () -> PartitionStatsHandler.computeAndWriteStatsFile(testTable, "INVALID_BRANCH"))
+    assertThatThrownBy(() -> PartitionStatsHandler.computeAndWriteStatsFile(testTable, 42L))
         .isInstanceOf(IllegalArgumentException.class)
-        .hasMessage("Couldn't find the snapshot for the branch INVALID_BRANCH");
+        .hasMessage("Snapshot not found: 42");
   }
 
   @Test