Apply test review findings

Critical:
- test_concurrent_replace_table now uses identical new schemas across
  the two transactions and matches on 'last assigned field id' — this
  is the only path that actually validates AssertLastAssignedFieldId.

Coverage:
- Add test_replace_table_with_sort_order_changes: unsorted → sorted →
  unsorted with order_id reuse from history.
- Extend the v2 partition-spec helper test with a reuse-by-bucket case
  (matching field_id under a renamed partition field).

Brittleness fixes:
- Drop dead HEAD mock in test_replace_table_transaction_404_raises;
  no view-exists pre-flight to mock anymore.
- Pin wire-payload assertions to fixture metadata fields (no more
  magic 999 / spec_id / sort_order_id constants).
- Tighten set-default-spec / set-default-sort-order to check the
  emitted id, not just existence.
- Split test_replace_table_location_resolution into two non-stringly-
  discriminated tests.

Removed:
- test_replace_table_accepts_pyarrow_schema: redundant with the RTAS
  test which already exercises pa.Schema via df.schema.

Author: smaheshwar-pltr

tests/catalog/test_catalog_behaviors.py

@@ -474,14 +474,6 @@ def test_replace_table_schema_id_reuse(
     assert replaced.metadata.last_column_id == expected_last_col_id
 
 
-def test_replace_table_accepts_pyarrow_schema(catalog: Catalog, test_table_identifier: Identifier) -> None:
-    """The schema argument may be a PyArrow schema; it must be converted and have IDs assigned."""
-    _create_simple_table(catalog, test_table_identifier)
-    pa_schema = pa.schema([pa.field("id", pa.int64()), pa.field("data", pa.large_string()), pa.field("extra", pa.bool_())])
-    replaced = catalog.replace_table(test_table_identifier, schema=pa_schema)
-    assert {f.name for f in replaced.schema().fields} == {"id", "data", "extra"}
-
-
 def test_replace_table_preserves_identifier_field_ids(catalog: Catalog, test_table_identifier: Identifier) -> None:
     """Identifier-field IDs on the new schema are honored after reuse-by-name."""
     schema = Schema(
@@ -509,37 +501,42 @@ def test_replace_table_reuses_partition_spec_id(catalog: Catalog, test_table_ide
     assert replaced.metadata.default_spec_id == 0
 
 
-@pytest.mark.parametrize(
-    "location_kwarg, expected",
-    [
-        ("explicit", "explicit"),  # explicit location is used verbatim (minus any trailing slash)
-        (None, "existing"),  # no location → existing location is preserved
-        ("explicit-with-trailing-slash", "explicit-with-trailing-slash-stripped"),
-    ],
-    ids=["explicit", "inherited", "trailing-slash-stripped"],
-)
-def test_replace_table_location_resolution(
-    catalog: Catalog,
-    test_table_identifier: Identifier,
-    tmp_path: Path,
-    location_kwarg: str | None,
-    expected: str,
-) -> None:
-    """`location=None` keeps the existing location; an explicit location is used (sans trailing slash)."""
+def test_replace_table_with_sort_order_changes(catalog: Catalog, test_table_identifier: Identifier) -> None:
+    """Replace can change the sort order. The new sort order is appended to the history and
+    becomes the default; a follow-up replace back to unsorted reuses the unsorted order_id
+    rather than appending a duplicate."""
+    _, schema = _create_simple_table(catalog, test_table_identifier)
+    sort = SortOrder(SortField(source_id=1, transform=IdentityTransform(), direction=SortDirection.ASC))
+
+    # unsorted → sorted: a new order is added and becomes the default.
+    sorted_table = catalog.replace_table(test_table_identifier, schema=schema, sort_order=sort)
+    assert sorted_table.sort_order().fields == sort.fields
+    assert sorted_table.metadata.default_sort_order_id != 0
+
+    # sorted → unsorted: reuses the unsorted order_id 0 from history.
+    unsorted_table = catalog.replace_table(test_table_identifier, schema=schema)
+    assert unsorted_table.sort_order().is_unsorted
+    assert unsorted_table.metadata.default_sort_order_id == 0
+
+
+def test_replace_table_inherits_existing_location(catalog: Catalog, test_table_identifier: Identifier) -> None:
+    """`location=None` keeps the existing table's location."""
     _, schema = _create_simple_table(catalog, test_table_identifier)
     existing = catalog.load_table(test_table_identifier).metadata.location
+    replaced = catalog.replace_table(test_table_identifier, schema=schema)
+    assert replaced.metadata.location == existing
 
-    if location_kwarg is None:
-        replaced = catalog.replace_table(test_table_identifier, schema=schema)
-        assert replaced.metadata.location == existing
-    elif location_kwarg == "explicit":
-        new_location = f"file://{tmp_path}/relocated"
-        replaced = catalog.replace_table(test_table_identifier, schema=schema, location=new_location)
-        assert replaced.metadata.location == new_location
-    else:
-        new_location = f"file://{tmp_path}/with-slash"
-        replaced = catalog.replace_table(test_table_identifier, schema=schema, location=new_location + "/")
-        assert replaced.metadata.location == new_location
+
+@pytest.mark.parametrize("trailing_slash", [False, True], ids=["no-slash", "trailing-slash"])
+def test_replace_table_uses_explicit_location(
+    catalog: Catalog, test_table_identifier: Identifier, tmp_path: Path, trailing_slash: bool
+) -> None:
+    """An explicit `location` is used verbatim; trailing slash is stripped."""
+    _, schema = _create_simple_table(catalog, test_table_identifier)
+    bare = f"file://{tmp_path}/relocated"
+    arg = bare + "/" if trailing_slash else bare
+    replaced = catalog.replace_table(test_table_identifier, schema=schema, location=arg)
+    assert replaced.metadata.location == bare
 
 
 def test_replace_table_merges_properties_with_overrides_and_additions(
@@ -657,23 +654,22 @@ def run_failing_replace() -> None:
 
 
 def test_concurrent_replace_table(catalog: Catalog, test_table_identifier: Identifier) -> None:
-    """Two replace_table calls staged from the same base metadata cannot both commit.
-
-    The first replace updates the catalog pointer to a new metadata location; the second
-    replace was built against the now-stale base, so its commit must be rejected with
-    `CommitFailedException`."""
-    _, schema = _create_simple_table(catalog, test_table_identifier)
+    """Two concurrent replace_table calls staged from the same base both adding a column
+    must fail on the second commit with an `assert-last-assigned-field-id` violation —
+    proving the `AssertLastAssignedFieldId` requirement actually guards against duplicate
+    field-id assignment under concurrent writers."""
+    _create_simple_table(catalog, test_table_identifier)
     new_schema = Schema(
         NestedField(field_id=1, name="id", field_type=LongType(), required=False),
         NestedField(field_id=2, name="data", field_type=StringType(), required=False),
         NestedField(field_id=3, name="extra", field_type=BooleanType(), required=False),
     )
-    # Both transactions are built from the same base metadata.
+    # Both transactions build from the same base metadata with the same new schema.
     txn_a = catalog.replace_table_transaction(test_table_identifier, schema=new_schema)
-    txn_b = catalog.replace_table_transaction(test_table_identifier, schema=schema)
+    txn_b = catalog.replace_table_transaction(test_table_identifier, schema=new_schema)
 
     txn_a.commit_transaction()
-    with pytest.raises(CommitFailedException):
+    with pytest.raises(CommitFailedException, match="last assigned field id"):
         txn_b.commit_transaction()
 
 

tests/catalog/test_rest.py

@@ -2948,10 +2948,13 @@ def test_replace_table_transaction_wire_payload(
     catalog.replace_table_transaction(identifier=("fokko", "fokko2"), schema=new_schema).commit_transaction()
     request = rest_mock.last_request.json()
 
+    # Pin the requirement values to the fixture's last-column-id / last-partition-id so the
+    # assertion fails loudly if the fixture changes underneath us.
+    fixture_metadata = example_table_metadata_with_snapshot_v1_rest_json["metadata"]
     assert request["requirements"] == [
         {"type": "assert-table-uuid", "uuid": table_uuid},
-        {"type": "assert-last-assigned-field-id", "last-assigned-field-id": 2},
-        {"type": "assert-last-assigned-partition-id", "last-assigned-partition-id": 999},
+        {"type": "assert-last-assigned-field-id", "last-assigned-field-id": fixture_metadata["last-column-id"]},
+        {"type": "assert-last-assigned-partition-id", "last-assigned-partition-id": fixture_metadata["last-partition-id"]},
     ]
 
     # No duplicate actions in the payload (the dict-by-action collapses them — guard against it).
@@ -2969,18 +2972,14 @@ def test_replace_table_transaction_wire_payload(
     # Set* updates are emitted unconditionally even when their resulting id is reused.
     # schema-id=-1 is the wire sentinel meaning "the schema we just added in this commit".
     assert updates_by_action["set-current-schema"]["schema-id"] == -1
-    assert "set-default-spec" in updates_by_action
-    assert "set-default-sort-order" in updates_by_action
+    # Spec/sort-order are reused from the fixture (default ids) since the replace doesn't change them.
+    assert updates_by_action["set-default-spec"]["spec-id"] == fixture_metadata["default-spec-id"]
+    assert updates_by_action["set-default-sort-order"]["sort-order-id"] == fixture_metadata["default-sort-order-id"]
 
 
 def test_replace_table_transaction_404_raises(
     rest_mock: Mocker,
 ) -> None:
-    rest_mock.head(
-        f"{TEST_URI}v1/namespaces/fokko/views/missing",
-        status_code=404,
-        request_headers=TEST_HEADERS,
-    )
     rest_mock.get(
         f"{TEST_URI}v1/namespaces/fokko/tables/missing",
         json={"error": {"message": "Table not found", "type": "NoSuchTableException", "code": 404}},

tests/table/test_partitioning.py

@@ -313,29 +313,50 @@ def test_incompatible_transform_source_type() -> None:
 
 
 @pytest.mark.parametrize(
-    "existing_specs, last_partition_id, expected_field_id, expected_last_partition_id",
+    "new_spec, existing_specs, last_partition_id, expected_field_id, expected_last_partition_id",
     [
-        # Reuse: same (source_id, transform) already in an existing spec → reuse its field_id.
-        (
+        # Reuse-by-identity: same (source_id, IdentityTransform) already in an existing spec.
+        pytest.param(
+            PartitionSpec(PartitionField(source_id=1, field_id=999, transform=IdentityTransform(), name="id")),
             [PartitionSpec(PartitionField(source_id=1, field_id=1000, transform=IdentityTransform(), name="id"), spec_id=0)],
             1000,
             1000,
             1000,
+            id="reuse-identity",
+        ),
+        # Reuse-by-(source,bucket): same source_id + same BucketTransform, even under a renamed field.
+        pytest.param(
+            PartitionSpec(PartitionField(source_id=1, field_id=999, transform=BucketTransform(8), name="id_bucket_renamed")),
+            [
+                PartitionSpec(
+                    PartitionField(source_id=1, field_id=1042, transform=BucketTransform(8), name="id_bucket"), spec_id=0
+                )
+            ],
+            1042,
+            1042,
+            1042,
+            id="reuse-bucket-under-rename",
+        ),
+        # No match: fresh id above last_partition_id.
+        pytest.param(
+            PartitionSpec(PartitionField(source_id=1, field_id=999, transform=IdentityTransform(), name="id")),
+            [PartitionSpec(spec_id=0)],
+            999,
+            1000,
+            1000,
+            id="new-field-above-last-partition-id",
         ),
-        # New: no matching (source_id, transform) → fresh id above last_partition_id.
-        ([PartitionSpec(spec_id=0)], 999, 1000, 1000),
     ],
-    ids=["reuse-from-existing-spec", "new-field-above-last-partition-id"],
 )
 def test_assign_fresh_partition_spec_ids_for_replace_v2(
+    new_spec: PartitionSpec,
     existing_specs: list[PartitionSpec],
     last_partition_id: int,
     expected_field_id: int,
     expected_last_partition_id: int,
 ) -> None:
-    spec = PartitionSpec(PartitionField(source_id=1, field_id=999, transform=IdentityTransform(), name="id"))
     fresh_spec, new_last_pid = assign_fresh_partition_spec_ids_for_replace(
-        spec, _REPLACE_SCHEMA_FOR_PARTITION, _REPLACE_SCHEMA_FOR_PARTITION, existing_specs, last_partition_id
+        new_spec, _REPLACE_SCHEMA_FOR_PARTITION, _REPLACE_SCHEMA_FOR_PARTITION, existing_specs, last_partition_id
     )
     assert fresh_spec.fields[0].field_id == expected_field_id
     assert new_last_pid == expected_last_partition_id