Improvements

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: smaheshwar-pltr

pyiceberg/catalog/rest/__init__.py

@@ -977,12 +977,12 @@ def replace_table_transaction(
 
         iceberg_schema = self._convert_schema_if_needed(
             schema,
-            int(properties.get(TableProperties.FORMAT_VERSION, TableProperties.DEFAULT_FORMAT_VERSION)),  # type: ignore
+            int(properties.get(TableProperties.FORMAT_VERSION, existing_metadata.format_version)),  # type: ignore
         )
 
-        # Assign fresh schema IDs, reusing IDs from the existing schema by field name
+        # Assign fresh schema IDs, reusing IDs from all existing schemas by field name
         fresh_schema, _ = assign_fresh_schema_ids_for_replace(
-            iceberg_schema, existing_metadata.schema(), existing_metadata.last_column_id
+            iceberg_schema, existing_metadata.schemas, existing_metadata.last_column_id
         )
 
         # Assign fresh partition spec IDs, reusing IDs from existing specs

pyiceberg/schema.py

@@ -1403,22 +1403,32 @@ def _get_and_increment(self, current_id: int) -> int:
         return new_id
 
 
-def assign_fresh_schema_ids_for_replace(schema: Schema, base_schema: Schema, last_column_id: int) -> tuple[Schema, int]:
-    """Assign fresh IDs to a schema for a replace operation, reusing IDs from the base schema.
+def assign_fresh_schema_ids_for_replace(schema: Schema, all_schemas: list[Schema], last_column_id: int) -> tuple[Schema, int]:
+    """Assign fresh IDs to a schema for a replace operation, reusing IDs from existing schemas.
 
     For each field in the new schema, if a field with the same full path name exists
-    in the base schema, its ID is reused. New fields get IDs starting from
+    in any of the existing schemas, its ID is reused. New fields get IDs starting from
     last_column_id + 1.
 
     Args:
         schema: The new schema to assign IDs to.
-        base_schema: The existing table's schema (IDs are reused from here by name).
+        all_schemas: All schemas from the existing table metadata (IDs are reused from here by name).
         last_column_id: The current table's last_column_id (new IDs start above this).
 
     Returns:
         A tuple of (fresh_schema, new_last_column_id).
     """
-    base_name_to_id = index_by_name(base_schema)
+    # N.B. We diverge from the Java implementation by using ALL historical schemas for
+    # field ID reuse, not just the current schema. Java's TypeUtil.assignFreshIds only
+    # uses the current schema as the base, so a replace A→B→A where A and B have
+    # disjoint fields would create a 3rd schema (field IDs from A are lost when B is
+    # current). By using all schemas, we can recover those IDs and correctly deduplicate
+    # the schema on the way back. This is safe because Iceberg guarantees field
+    # name-to-ID consistency across schema evolution.
+    base_name_to_id: dict[str, int] = {}
+    for existing_schema in all_schemas:
+        base_name_to_id.update(index_by_name(existing_schema))
+
     new_id_to_name = index_name_by_id(schema)
 
     old_id_to_base_id: dict[int, int] = {}

pyiceberg/table/__init__.py

@@ -1031,11 +1031,18 @@ def _initial_changes(
         # Remove the main branch ref to clear the current snapshot
         self._updates += (RemoveSnapshotRefUpdate(ref_name=MAIN_BRANCH),)
 
-        # Set the new schema (always added since field IDs are freshly assigned)
-        self._updates += (
-            AddSchemaUpdate(schema_=new_schema),
-            SetCurrentSchemaUpdate(schema_id=-1),
-        )
+        # Set the new schema, reusing an existing schema if structurally identical.
+        # Schema.__eq__ compares fields and identifier_field_ids (ignoring schema_id),
+        # matching Java's sameSchema() behavior.
+        existing_schema_id = self._find_matching_schema_id(table_metadata, new_schema)
+        if existing_schema_id is not None:
+            if existing_schema_id != table_metadata.current_schema_id:
+                self._updates += (SetCurrentSchemaUpdate(schema_id=existing_schema_id),)
+        else:
+            self._updates += (
+                AddSchemaUpdate(schema_=new_schema),
+                SetCurrentSchemaUpdate(schema_id=-1),
+            )
 
         # Set the new partition spec.
         # Only emit AddPartitionSpecUpdate + SetDefaultSpecUpdate(-1) when the spec
@@ -1072,6 +1079,14 @@ def _initial_changes(
         if new_properties:
             self._updates += (SetPropertiesUpdate(updates=new_properties),)
 
+    @staticmethod
+    def _find_matching_schema_id(table_metadata: TableMetadata, schema: Schema) -> int | None:
+        """Find an existing schema structurally equal to the given one, returning its schema_id or None."""
+        for existing in table_metadata.schemas:
+            if existing == schema:
+                return existing.schema_id
+        return None
+
     @staticmethod
     def _find_matching_spec_id(table_metadata: TableMetadata, spec: PartitionSpec) -> int | None:
         """Find an existing partition spec with the same fields, returning its spec_id or None."""

tests/catalog/test_rest.py

@@ -2891,6 +2891,12 @@ def test_replace_table_transaction_with_partition_spec(
     assert spec_fields[0]["transform"] == "truncate[3]"
     assert spec_fields[0]["name"] == "id_trunc"
 
+    # set-default-spec should also be present, pointing to the newly added spec
+    actions = [u["action"] for u in updates]
+    assert "set-default-spec" in actions
+    set_default_spec = next(u for u in updates if u["action"] == "set-default-spec")
+    assert set_default_spec["spec-id"] == -1
+
 
 def test_replace_table_404(
     rest_mock: Mocker,
@@ -2993,3 +2999,135 @@ def test_replace_table_transaction_same_location_no_set_location(
     actions = [u["action"] for u in updates]
     # set-location should NOT be present since location didn't change
     assert "set-location" not in actions
+
+
+def test_replace_table_transaction_same_schema_skips_add_schema(
+    rest_mock: Mocker,
+    example_table_metadata_with_snapshot_v1_rest_json: dict[str, Any],
+    example_table_metadata_no_snapshot_v1_rest_json: dict[str, Any],
+) -> None:
+    """Test that replacing with the same schema skips add-schema and set-current-schema."""
+    table_uuid = example_table_metadata_with_snapshot_v1_rest_json["metadata"]["table-uuid"]
+    example_table_metadata_no_snapshot_v1_rest_json["metadata"]["table-uuid"] = table_uuid
+
+    rest_mock.get(
+        f"{TEST_URI}v1/namespaces/fokko/tables/fokko2",
+        json=example_table_metadata_with_snapshot_v1_rest_json,
+        status_code=200,
+        request_headers=TEST_HEADERS,
+    )
+    rest_mock.post(
+        f"{TEST_URI}v1/namespaces/fokko/tables/fokko2",
+        json=example_table_metadata_no_snapshot_v1_rest_json,
+        status_code=200,
+        request_headers=TEST_HEADERS,
+    )
+    catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)
+
+    # Use the exact same schema as the existing table (id: int, data: string)
+    same_schema = Schema(
+        NestedField(field_id=1, name="id", field_type=IntegerType(), required=False),
+        NestedField(field_id=2, name="data", field_type=StringType(), required=False),
+    )
+    txn = catalog.replace_table_transaction(
+        identifier=("fokko", "fokko2"),
+        schema=same_schema,
+    )
+    txn.commit_transaction()
+
+    updates = rest_mock.last_request.json()["updates"]
+    actions = [u["action"] for u in updates]
+
+    # Since the schema is unchanged, add-schema and set-current-schema should be skipped
+    assert "add-schema" not in actions
+    assert "set-current-schema" not in actions
+
+    # The only update should be remove-snapshot-ref
+    assert actions == ["remove-snapshot-ref"]
+
+
+def test_replace_table_transaction_different_schema_adds_schema(
+    rest_mock: Mocker,
+    example_table_metadata_with_snapshot_v1_rest_json: dict[str, Any],
+    example_table_metadata_no_snapshot_v1_rest_json: dict[str, Any],
+) -> None:
+    """Test that replacing with a genuinely new schema includes add-schema and set-current-schema."""
+    table_uuid = example_table_metadata_with_snapshot_v1_rest_json["metadata"]["table-uuid"]
+    example_table_metadata_no_snapshot_v1_rest_json["metadata"]["table-uuid"] = table_uuid
+
+    rest_mock.get(
+        f"{TEST_URI}v1/namespaces/fokko/tables/fokko2",
+        json=example_table_metadata_with_snapshot_v1_rest_json,
+        status_code=200,
+        request_headers=TEST_HEADERS,
+    )
+    rest_mock.post(
+        f"{TEST_URI}v1/namespaces/fokko/tables/fokko2",
+        json=example_table_metadata_no_snapshot_v1_rest_json,
+        status_code=200,
+        request_headers=TEST_HEADERS,
+    )
+    catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)
+
+    # A new schema with a different field (new_col instead of data)
+    new_schema = Schema(
+        NestedField(field_id=1, name="id", field_type=IntegerType(), required=False),
+        NestedField(field_id=3, name="new_col", field_type=BooleanType(), required=False),
+    )
+    txn = catalog.replace_table_transaction(
+        identifier=("fokko", "fokko2"),
+        schema=new_schema,
+    )
+    txn.commit_transaction()
+
+    updates = rest_mock.last_request.json()["updates"]
+    actions = [u["action"] for u in updates]
+
+    # Since the schema is different, add-schema and set-current-schema must be present
+    assert "add-schema" in actions
+    assert "set-current-schema" in actions
+
+    # set-current-schema should reference -1 (the last added schema)
+    set_schema = next(u for u in updates if u["action"] == "set-current-schema")
+    assert set_schema["schema-id"] == -1
+
+
+def test_replace_table_transaction_with_sort_order(
+    rest_mock: Mocker,
+    example_table_metadata_with_snapshot_v1_rest_json: dict[str, Any],
+    example_table_metadata_no_snapshot_v1_rest_json: dict[str, Any],
+) -> None:
+    """Test that replacing with a custom sort order includes add-sort-order and set-default-sort-order."""
+    table_uuid = example_table_metadata_with_snapshot_v1_rest_json["metadata"]["table-uuid"]
+    example_table_metadata_no_snapshot_v1_rest_json["metadata"]["table-uuid"] = table_uuid
+
+    rest_mock.get(
+        f"{TEST_URI}v1/namespaces/fokko/tables/fokko2",
+        json=example_table_metadata_with_snapshot_v1_rest_json,
+        status_code=200,
+        request_headers=TEST_HEADERS,
+    )
+    rest_mock.post(
+        f"{TEST_URI}v1/namespaces/fokko/tables/fokko2",
+        json=example_table_metadata_no_snapshot_v1_rest_json,
+        status_code=200,
+        request_headers=TEST_HEADERS,
+    )
+    catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)
+
+    txn = catalog.replace_table_transaction(
+        identifier=("fokko", "fokko2"),
+        schema=Schema(
+            NestedField(field_id=1, name="id", field_type=IntegerType(), required=False),
+            NestedField(field_id=2, name="data", field_type=StringType(), required=False),
+        ),
+        sort_order=SortOrder(SortField(source_id=1, transform=IdentityTransform())),
+    )
+    txn.commit_transaction()
+
+    updates = rest_mock.last_request.json()["updates"]
+    actions = [u["action"] for u in updates]
+
+    # Should include add-sort-order and set-default-sort-order
+    assert "add-sort-order" in actions
+    assert "set-default-sort-order" in actions