partition field names validation against schema field conflicts

Author: rutb327

pyiceberg/table/update/spec.py

@@ -174,16 +174,21 @@ def _commit(self) -> UpdatesAndRequirements:
         return updates, requirements
 
     def _apply(self) -> PartitionSpec:
-        def _check_and_add_partition_name(schema: Schema, name: str, source_id: int, partition_names: Set[str]) -> None:
+        def _check_and_add_partition_name(
+            schema: Schema, name: str, source_id: int, transform: Transform[Any, Any], partition_names: Set[str]
+        ) -> None:
             try:
                 field = schema.find_field(name)
             except ValueError:
                 field = None
 
-            if source_id is not None and field is not None and field.field_id != source_id:
-                raise ValueError(f"Cannot create identity partition from a different field in the schema {name}")
-            elif field is not None and source_id != field.field_id:
-                raise ValueError(f"Cannot create partition from name that exists in schema {name}")
+            if field is not None:
+                if isinstance(transform, (IdentityTransform, VoidTransform)):
+                    # For identity transforms allow name conflict only if sourced from the same schema field
+                    if field.field_id != source_id:
+                        raise ValueError(f"Cannot create identity partition from a different field in the schema: {name}")
+                else:
+                    raise ValueError(f"Cannot create partition from name that exists in schema: {name}")
             if not name:
                 raise ValueError("Undefined name")
             if name in partition_names:
@@ -193,7 +198,7 @@ def _check_and_add_partition_name(schema: Schema, name: str, source_id: int, par
         def _add_new_field(
             schema: Schema, source_id: int, field_id: int, name: str, transform: Transform[Any, Any], partition_names: Set[str]
         ) -> PartitionField:
-            _check_and_add_partition_name(schema, name, source_id, partition_names)
+            _check_and_add_partition_name(schema, name, source_id, transform, partition_names)
             return PartitionField(source_id, field_id, transform, name)
 
         partition_fields = []
@@ -244,6 +249,13 @@ def _add_new_field(
                 partition_fields.append(new_field)
 
         for added_field in self._adds:
+            _check_and_add_partition_name(
+                self._transaction.table_metadata.schema(),
+                added_field.name,
+                added_field.source_id,
+                added_field.transform,
+                partition_names,
+            )
             new_field = PartitionField(
                 source_id=added_field.source_id,
                 field_id=added_field.field_id,

tests/integration/test_partition_evolution.py

@@ -564,3 +564,28 @@ def _validate_new_partition_fields(
     assert len(spec.fields) == len(expected_partition_fields)
     for i in range(len(spec.fields)):
         assert spec.fields[i] == expected_partition_fields[i]
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
+def test_partition_schema_field_name_conflict(catalog: Catalog) -> None:
+    schema = Schema(
+        NestedField(1, "id", LongType(), required=False),
+        NestedField(2, "event_ts", TimestampType(), required=False),
+        NestedField(3, "another_ts", TimestampType(), required=False),
+        NestedField(4, "str", StringType(), required=False),
+    )
+    table = _create_table_with_schema(catalog, schema, "2")
+
+    with pytest.raises(ValueError, match="Cannot create partition from name that exists in schema: another_ts"):
+        table.update_spec().add_field("event_ts", YearTransform(), "another_ts").commit()
+    with pytest.raises(ValueError, match="Cannot create partition from name that exists in schema: id"):
+        table.update_spec().add_field("event_ts", DayTransform(), "id").commit()
+
+    with pytest.raises(ValueError, match="Cannot create identity partition from a different field in the schema: another_ts"):
+        table.update_spec().add_field("event_ts", IdentityTransform(), "another_ts").commit()
+    with pytest.raises(ValueError, match="Cannot create identity partition from a different field in the schema: str"):
+        table.update_spec().add_field("id", IdentityTransform(), "str").commit()
+
+    table.update_spec().add_field("id", IdentityTransform(), "id").commit()
+    table.update_spec().add_field("event_ts", YearTransform(), "event_year").commit()