Table: assert every written ref in a multi-branch transaction

pyiceberg/table/__init__.py

@@ -305,9 +305,14 @@ def _stage(
 
         # For the requirements, it does not make sense to add a requirement more than once
         # For example, you cannot assert that the current schema has two different IDs
-        existing_requirements = {type(requirement) for requirement in self._requirements}
+        # Ref requirements are the exception: each ref needs its own assertion
+        def key(requirement: TableRequirement) -> tuple[type, str | None]:
+            return type(requirement), requirement.ref if isinstance(requirement, AssertRefSnapshotId) else None
+
+        existing_requirements = {key(requirement) for requirement in self._requirements}
         for new_requirement in requirements:
-            if type(new_requirement) not in existing_requirements:
+            if key(new_requirement) not in existing_requirements:
+                existing_requirements.add(key(new_requirement))
                 self._requirements = self._requirements + (new_requirement,)
 
         return self

tests/table/test_init.py

@@ -1437,6 +1437,16 @@ def test_assert_ref_snapshot_id(table_v2: Table) -> None:
     assert AssertRefSnapshotId(ref="main").model_dump_json() == expected_json
 
 
+def test_stage_keeps_ref_requirement_per_ref(table_v2: Table) -> None:
+    tx = table_v2.transaction()
+    main = AssertRefSnapshotId(ref=MAIN_BRANCH, snapshot_id=3055729675574597004)
+    test = AssertRefSnapshotId(ref="test", snapshot_id=3051729675574597004)
+
+    tx._stage((), (main, test, main, test))
+
+    assert tx._requirements == (main, test)
+
+
 def test_assert_last_assigned_field_id(table_v2: Table) -> None:
     base_metadata = table_v2.metadata
     AssertLastAssignedFieldId(last_assigned_field_id=base_metadata.last_column_id).validate(base_metadata)