diff --git a/mkdocs/docs/api.md b/mkdocs/docs/api.md index a732d2c0a3..b7afecc362 100644 --- a/mkdocs/docs/api.md +++ b/mkdocs/docs/api.md @@ -185,6 +185,26 @@ with catalog.create_table_transaction(identifier="docs_example.bids", schema=sch txn.set_properties(test_a="test_aa", test_b="test_b", test_c="test_c") ``` +## Replace a table + +Use `replace_table_transaction` to replace an existing table's definition and data in one atomic commit: + +```python +with catalog.replace_table_transaction(identifier="docs_example.bids", schema=df.schema) as txn: + txn.append(df) +``` + +The existing table remains visible until the transaction commits. The replacement can use an incompatible +schema, such as changing a column from an integer to a string. Field IDs are reused by name, while new +fields receive IDs above the table's last assigned field ID. + +The table UUID, historical schemas, snapshots, and other branches and tags are preserved. Replacement data +starts a new snapshot lineage. Committing without appending data leaves an empty table with no current snapshot. + +By default, the replacement is unpartitioned and unsorted. Supply `partition_spec` and `sort_order` to set +the new layout. Existing properties are retained and merged with `properties`, and the existing location is +retained unless `location` is supplied. + ## Register a table To register a table using existing metadata: diff --git a/pyiceberg/avro/resolver.py b/pyiceberg/avro/resolver.py index 3f19ae9878..430e3bb171 100644 --- a/pyiceberg/avro/resolver.py +++ b/pyiceberg/avro/resolver.py @@ -147,7 +147,7 @@ def struct(self, struct: StructType, field_results: list[Writer]) -> Writer: return StructWriter(tuple((pos, result) for pos, result in enumerate(field_results))) def field(self, field: NestedField, field_result: Writer) -> Writer: - return field_result if field.required else OptionWriter(field_result) + return field_result if field.required or isinstance(field.field_type, UnknownType) else OptionWriter(field_result) def list(self, list_type: ListType, element_result: Writer) -> Writer: return ListWriter(element_result) @@ -302,7 +302,9 @@ def struct(self, file_schema: StructType, record_struct: IcebergType | None, fil return StructWriter(field_writers=tuple(results)) def field(self, file_field: NestedField, record_type: IcebergType | None, field_writer: Writer) -> Writer: - return field_writer if file_field.required else OptionWriter(field_writer) + return ( + field_writer if file_field.required or isinstance(file_field.field_type, UnknownType) else OptionWriter(field_writer) + ) def list(self, file_list_type: ListType, file_list: IcebergType | None, element_writer: Writer) -> Writer: return ListWriter(element_writer if file_list_type.element_required else OptionWriter(element_writer)) diff --git a/pyiceberg/catalog/__init__.py b/pyiceberg/catalog/__init__.py index 022783b08a..855a959e38 100644 --- a/pyiceberg/catalog/__init__.py +++ b/pyiceberg/catalog/__init__.py @@ -43,12 +43,13 @@ from pyiceberg.io import FileIO, load_file_io from pyiceberg.manifest import ManifestFile from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionSpec -from pyiceberg.schema import Schema +from pyiceberg.schema import Schema, assign_fresh_schema_ids from pyiceberg.serializers import ToOutputFile from pyiceberg.table import ( DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE, CommitTableResponse, CreateTableTransaction, + ReplaceTableTransaction, StagedTable, Table, TableProperties, @@ -416,6 +417,31 @@ def create_table_transaction( CreateTableTransaction: createTableTransaction instance. """ + def replace_table_transaction( + self, + identifier: str | Identifier, + schema: Schema | pa.Schema, + location: str | None = None, + partition_spec: PartitionSpec = UNPARTITIONED_PARTITION_SPEC, + sort_order: SortOrder = UNSORTED_SORT_ORDER, + properties: Properties = EMPTY_DICT, + ) -> ReplaceTableTransaction: + """Replace an existing table atomically, preserving its UUID and history. + + The replacement starts empty, with the supplied schema, partition spec, and + sort order. Properties are merged with existing properties. The location + defaults to the existing table location. Append data to the transaction + before committing to replace both the table definition and its contents. + + Raises: + NoSuchTableError: If the table does not exist. + """ + table = self.load_table(identifier) + iceberg_schema = self._convert_schema_if_needed(schema, table.metadata.format_version) + if not isinstance(schema, Schema): + iceberg_schema = assign_fresh_schema_ids(iceberg_schema) + return ReplaceTableTransaction(table, iceberg_schema, location, partition_spec, sort_order, properties) + def create_table_if_not_exists( self, identifier: str | Identifier, diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py index 0cab26309b..9cb606ca9c 100644 --- a/pyiceberg/io/pyarrow.py +++ b/pyiceberg/io/pyarrow.py @@ -153,7 +153,7 @@ from pyiceberg.table.locations import load_location_provider from pyiceberg.table.metadata import TableMetadata from pyiceberg.table.name_mapping import NameMapping, apply_name_mapping -from pyiceberg.transforms import IdentityTransform, TruncateTransform +from pyiceberg.transforms import IdentityTransform, TruncateTransform, VoidTransform from pyiceberg.typedef import EMPTY_DICT, Properties, Record, TableVersion from pyiceberg.types import ( BinaryType, @@ -3169,6 +3169,9 @@ def _determine_partitions(spec: PartitionSpec, schema: Schema, arrow_table: pa.T partition_fields = [f"_partition_{field.name}" for field in spec.fields] for partition, name in zip(spec.fields, partition_fields, strict=True): + if isinstance(partition.transform, VoidTransform): + arrow_table = arrow_table.append_column(name, pa.nulls(arrow_table.num_rows, type=pa.int32())) + continue source_field = schema.find_field(partition.source_id) full_field_name = schema.find_column_name(partition.source_id) if full_field_name is None: diff --git a/pyiceberg/partitioning.py b/pyiceberg/partitioning.py index 3d06287b6c..a52bbb32fa 100644 --- a/pyiceberg/partitioning.py +++ b/pyiceberg/partitioning.py @@ -490,6 +490,8 @@ def partition_record_value(partition_field: PartitionField, value: Any, schema: Then the corresponding PartitionField's transform is applied to return the final partition record value. """ + if isinstance(partition_field.transform, VoidTransform): + return None iceberg_type = schema.find_field(name_or_id=partition_field.source_id).field_type return _to_partition_representation(iceberg_type, value) diff --git a/pyiceberg/schema.py b/pyiceberg/schema.py index 99f983074b..e9d92c3b9c 100644 --- a/pyiceberg/schema.py +++ b/pyiceberg/schema.py @@ -1310,9 +1310,13 @@ def build_position_accessors(schema_or_type: Schema | IcebergType) -> dict[int, return visit(schema_or_type, _BuildPositionAccessors()) -def assign_fresh_schema_ids(schema_or_type: Schema | IcebergType, next_id: Callable[[], int] | None = None) -> Schema: - """Traverses the schema, and sets new IDs.""" - return pre_order_visit(schema_or_type, _SetFreshIDs(next_id_func=next_id)) +def assign_fresh_schema_ids( + schema_or_type: Schema | IcebergType, + next_id: Callable[[], int] | None = None, + base_schema: Schema | None = None, +) -> Schema: + """Assign fresh IDs, reusing IDs by name from a base schema when supplied.""" + return pre_order_visit(schema_or_type, _SetFreshIDs(next_id_func=next_id, base_schema=base_schema)) class _SetFreshIDs(PreOrderSchemaVisitor[IcebergType]): @@ -1320,17 +1324,21 @@ class _SetFreshIDs(PreOrderSchemaVisitor[IcebergType]): old_id_to_new_id: dict[int, int] - def __init__(self, next_id_func: Callable[[], int] | None = None) -> None: + def __init__(self, next_id_func: Callable[[], int] | None = None, base_schema: Schema | None = None) -> None: self.old_id_to_new_id = {} - counter = itertools.count(1) + self._base_name_to_id = index_by_name(base_schema) if base_schema is not None else {} + self._id_to_name: dict[int, str] = {} + counter = itertools.count(base_schema.highest_field_id + 1 if base_schema is not None else 1) self.next_id_func = next_id_func if next_id_func is not None else lambda: next(counter) def _get_and_increment(self, current_id: int) -> int: - new_id = self.next_id_func() + name = self._id_to_name.get(current_id) + new_id = self._base_name_to_id[name] if name is not None and name in self._base_name_to_id else self.next_id_func() self.old_id_to_new_id[current_id] = new_id return new_id def schema(self, schema: Schema, struct_result: Callable[[], StructType]) -> Schema: + self._id_to_name = index_name_by_id(schema) if self._base_name_to_id else {} return Schema( *struct_result().fields, identifier_field_ids=[self.old_id_to_new_id[field_id] for field_id in schema.identifier_field_ids], @@ -1340,15 +1348,7 @@ def struct(self, struct: StructType, field_results: builtins.list[Callable[[], I new_ids = [self._get_and_increment(field.field_id) for field in struct.fields] new_fields = [] for field_id, field, field_type in zip(new_ids, struct.fields, field_results, strict=True): - new_fields.append( - NestedField( - field_id=field_id, - name=field.name, - field_type=field_type(), - required=field.required, - doc=field.doc, - ) - ) + new_fields.append(field.model_copy(update={"field_id": field_id, "field_type": field_type()})) return StructType(*new_fields) def field(self, field: NestedField, field_result: Callable[[], IcebergType]) -> IcebergType: diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index 7da0e59d26..8682fb9678 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -26,7 +26,7 @@ from collections.abc import Callable, Iterable, Iterator from dataclasses import dataclass from functools import cached_property -from itertools import chain +from itertools import chain, count from types import TracebackType from typing import TYPE_CHECKING, Any, TypeVar @@ -52,8 +52,15 @@ ManifestEntryStatus, ManifestFile, ) -from pyiceberg.partitioning import PARTITION_FIELD_ID_START, UNPARTITIONED_PARTITION_SPEC, PartitionKey, PartitionSpec -from pyiceberg.schema import Schema +from pyiceberg.partitioning import ( + PARTITION_FIELD_ID_START, + UNPARTITIONED_PARTITION_SPEC, + PartitionField, + PartitionKey, + PartitionSpec, + assign_fresh_partition_spec_ids, +) +from pyiceberg.schema import Schema, assign_fresh_schema_ids from pyiceberg.table.delete_file_index import DeleteFileIndex from pyiceberg.table.inspect import InspectTable from pyiceberg.table.locations import LocationProvider, load_location_provider @@ -69,16 +76,22 @@ is_ancestor_of, is_parent_ancestor_of, ) -from pyiceberg.table.sorting import UNSORTED_SORT_ORDER, SortOrder +from pyiceberg.table.sorting import UNSORTED_SORT_ORDER, SortOrder, assign_fresh_sort_order_ids from pyiceberg.table.update import ( AddPartitionSpecUpdate, AddSchemaUpdate, AddSortOrderUpdate, AssertCreate, + AssertCurrentSchemaId, + AssertDefaultSortOrderId, + AssertDefaultSpecId, + AssertLastAssignedFieldId, + AssertLastAssignedPartitionId, AssertRefSnapshotId, AssertTableUUID, AssignUUIDUpdate, RemovePropertiesUpdate, + RemoveSnapshotRefUpdate, SetCurrentSchemaUpdate, SetDefaultSortOrderUpdate, SetDefaultSpecUpdate, @@ -96,7 +109,7 @@ from pyiceberg.table.update.sorting import UpdateSortOrder from pyiceberg.table.update.spec import UpdateSpec from pyiceberg.table.update.statistics import UpdateStatistics -from pyiceberg.transforms import IdentityTransform +from pyiceberg.transforms import IdentityTransform, VoidTransform from pyiceberg.typedef import ( EMPTY_DICT, IcebergBaseModel, @@ -1233,6 +1246,118 @@ def _rebuild_snapshot_updates(self) -> None: self._stage(updates, requirements) +class ReplaceTableTransaction(Transaction): + """A transaction that atomically replaces a table while retaining its history.""" + + def __init__( + self, + table: Table, + schema: Schema, + location: str | None = None, + partition_spec: PartitionSpec = UNPARTITIONED_PARTITION_SPEC, + sort_order: SortOrder = UNSORTED_SORT_ORDER, + properties: Properties = EMPTY_DICT, + ): + super().__init__(table) + base = table.metadata + schema.check_format_version_compatibility(int(properties.get(TableProperties.FORMAT_VERSION, base.format_version))) + next_id = count(base.last_column_id + 1) + fresh_schema = assign_fresh_schema_ids(schema, lambda: next(next_id), base.schema()) + existing_schema = next((old for old in base.schemas if old == fresh_schema), None) + updates: tuple[TableUpdate, ...] = (RemoveSnapshotRefUpdate(ref_name=MAIN_BRANCH),) + if existing_schema is None: + fresh_schema = fresh_schema.model_copy(update={"schema_id": max(old.schema_id for old in base.schemas) + 1}) + updates += (AddSchemaUpdate(schema_=fresh_schema),) + else: + fresh_schema = existing_schema + updates += (SetCurrentSchemaUpdate(schema_id=fresh_schema.schema_id),) + + fresh_spec = assign_fresh_partition_spec_ids(partition_spec, schema, fresh_schema) + fresh_spec = self._replacement_spec(fresh_spec) + fresh_spec.check_compatible(fresh_schema) + if fresh_spec.spec_id not in base.specs(): + updates += (AddPartitionSpecUpdate(spec=fresh_spec),) + updates += (SetDefaultSpecUpdate(spec_id=fresh_spec.spec_id),) + + fresh_order = assign_fresh_sort_order_ids(sort_order, schema, fresh_schema) + fresh_order.check_compatible(fresh_schema) + existing_order = next((old for old in base.sort_orders if old.fields == fresh_order.fields), None) + if existing_order is None: + if not fresh_order.is_unsorted: + fresh_order = fresh_order.model_copy( + update={"order_id": max((old.order_id for old in base.sort_orders), default=0) + 1} + ) + updates += (AddSortOrderUpdate(sort_order=fresh_order),) + else: + fresh_order = existing_order + updates += (SetDefaultSortOrderUpdate(sort_order_id=fresh_order.order_id),) + + persisted_properties = dict(properties) + if format_version := persisted_properties.pop(TableProperties.FORMAT_VERSION, None): + updates += (UpgradeFormatVersionUpdate(format_version=int(format_version)),) + if location is not None: + updates += (SetLocationUpdate(location=location.rstrip("/")),) + updates += (SetPropertiesUpdate(updates=persisted_properties),) + self._stage( + updates, + ( + AssertLastAssignedFieldId(last_assigned_field_id=base.last_column_id), + AssertCurrentSchemaId(current_schema_id=base.current_schema_id), + AssertLastAssignedPartitionId(last_assigned_partition_id=base.last_partition_id), + AssertDefaultSpecId(default_spec_id=base.default_spec_id), + AssertDefaultSortOrderId(default_sort_order_id=base.default_sort_order_id), + AssertRefSnapshotId(ref=MAIN_BRANCH, snapshot_id=base.current_snapshot_id), + ), + ) + + def _replacement_spec(self, spec: PartitionSpec) -> PartitionSpec: + base = self._table.metadata + next_id = count((base.last_partition_id or PARTITION_FIELD_ID_START - 1) + 1) + fields: list[PartitionField] = [] + if base.format_version == 1: + remaining = {(field.source_id, str(field.transform)): field for field in spec.fields} + new_names = {field.name for field in spec.fields} + for old in base.spec().fields: + if replacement := remaining.pop((old.source_id, str(old.transform)), None): + fields.append(replacement.model_copy(update={"field_id": old.field_id})) + else: + fields.append( + old.model_copy( + update={ + "transform": VoidTransform(), + "name": f"{old.name}_{old.field_id}" if old.name in new_names else old.name, + } + ) + ) + fields.extend(field.model_copy(update={"field_id": next(next_id)}) for field in remaining.values()) + else: + historical_ids: dict[tuple[int, str], int] = {} + for old_spec in base.partition_specs: + for field in old_spec.fields: + key = (field.source_id, str(field.transform)) + historical_ids[key] = max(historical_ids.get(key, 0), field.field_id) + for field in spec.fields: + key = (field.source_id, str(field.transform)) + if key not in historical_ids: + historical_ids[key] = next(next_id) + fields.append(field.model_copy(update={"field_id": historical_ids[key]})) + existing_spec = next((old for old in base.partition_specs if old.fields == tuple(fields)), None) + return existing_spec or PartitionSpec(*fields, spec_id=max(old.spec_id for old in base.partition_specs) + 1) + + def _stage(self, updates: tuple[TableUpdate, ...], requirements: tuple[TableRequirement, ...] = ()) -> Transaction: + # Main is cleared in staged metadata, so guard the published head instead. + replaces_main = any(isinstance(req, AssertRefSnapshotId) and req.ref == MAIN_BRANCH for req in requirements) + super()._stage( + updates, + tuple(req for req in requirements if not (isinstance(req, AssertRefSnapshotId) and req.ref == MAIN_BRANCH)), + ) + if replaces_main and not any( + isinstance(req, AssertRefSnapshotId) and req.ref == MAIN_BRANCH for req in self._requirements + ): + self._requirements += (AssertRefSnapshotId(ref=MAIN_BRANCH, snapshot_id=self._table.metadata.current_snapshot_id),) + return self + + class CreateTableTransaction(Transaction): """A transaction that involves the creation of a new table.""" diff --git a/pyiceberg/utils/schema_conversion.py b/pyiceberg/utils/schema_conversion.py index 732439d34b..6a932b72e4 100644 --- a/pyiceberg/utils/schema_conversion.py +++ b/pyiceberg/utils/schema_conversion.py @@ -540,7 +540,7 @@ def field(self, field: NestedField, field_result: AvroType) -> AvroType: result = { "name": sanitized_name, FIELD_ID_PROP: field.field_id, - "type": field_result if field.required else ["null", field_result], + "type": field_result if field.required or field_result == "null" else ["null", field_result], } if original_name != sanitized_name: diff --git a/tests/avro/test_file.py b/tests/avro/test_file.py index 6cc190c33b..ab8bd4bf01 100644 --- a/tests/avro/test_file.py +++ b/tests/avro/test_file.py @@ -56,6 +56,7 @@ TimestamptzNanoType, TimestamptzType, TimeType, + UnknownType, UUIDType, ) from pyiceberg.utils.schema_conversion import AvroSchemaConversion @@ -571,3 +572,19 @@ def test_avro_output_file_tell_includes_buffered_records() -> None: out.write_block([manifest_entry(0)]) assert out.tell() > after_header + + +@pytest.mark.parametrize("with_record_schema", [False, True]) +def test_write_optional_unknown_field(with_record_schema: bool) -> None: + schema = Schema(NestedField(1, "removed_partition", UnknownType()), NestedField(2, "value", LongType())) + with TemporaryDirectory() as tmpdir: + path = tmpdir + "/null.avro" + with avro.AvroOutputFile[Record]( + output_file=PyArrowFileIO().new_output(path), + file_schema=schema, + record_schema=schema if with_record_schema else None, + schema_name="null_field", + ) as out: + out.write_block([Record(None, 42)]) + with open(path, "rb") as source: + assert list(reader(source)) == [{"removed_partition": None, "value": 42}] diff --git a/tests/catalog/test_catalog_behaviors.py b/tests/catalog/test_catalog_behaviors.py index 1148ca8527..fe9d22a039 100644 --- a/tests/catalog/test_catalog_behaviors.py +++ b/tests/catalog/test_catalog_behaviors.py @@ -1409,3 +1409,273 @@ def test_record_batch_reader_schema_mismatch_writes_no_files(catalog: Catalog) - # _append_snapshot_producer opens. assert tbl.metadata.current_snapshot() is None assert len(tbl.scan().to_arrow()) == 0 + + +@pytest.mark.parametrize("format_version", [1, 2]) +def test_replace_table_transaction(catalog: Catalog, format_version: int) -> None: + catalog.create_namespace("replace_test") + original = pa.table({"id": [1, 2], "old": ["a", "b"]}) + table = catalog.create_table("replace_test.tbl", original.schema, properties={"format-version": str(format_version)}) + table.append(original) + snapshot = table.current_snapshot() + assert snapshot is not None + table.manage_snapshots().create_tag(snapshot.snapshot_id, "before_replace").commit() + original_metadata = table.metadata + replacement = pa.table({"id": ["new"], "added": [True]}) + + with catalog.replace_table_transaction("replace_test.tbl", replacement.schema, properties={"owner": "new-owner"}) as tx: + tx.append(replacement) + published = catalog.load_table("replace_test.tbl") + assert published.metadata == original_metadata + assert published.scan().to_arrow().to_pylist() == original.to_pylist() + + replaced = catalog.load_table("replace_test.tbl") + assert replaced.scan().to_arrow().to_pylist() == replacement.to_pylist() + assert replaced.metadata.table_uuid == original_metadata.table_uuid + assert replaced.location() == table.location() + assert replaced.metadata.format_version == format_version + assert replaced.properties["owner"] == "new-owner" + assert replaced.schema().find_field("id").field_id == table.schema().find_field("id").field_id + assert replaced.schema().find_field("added").field_id > original_metadata.last_column_id + assert replaced.metadata.snapshot_by_id(snapshot.snapshot_id) == snapshot + assert replaced.scan(snapshot_id=snapshot.snapshot_id).to_arrow().to_pylist() == original.to_pylist() + assert replaced.metadata.refs["before_replace"].snapshot_id == snapshot.snapshot_id + current_snapshot = replaced.current_snapshot() + assert current_snapshot is not None + assert current_snapshot.parent_snapshot_id is None + assert len(replaced.snapshots()) == 2 + + +@pytest.mark.parametrize("format_version", [1, 2]) +def test_replace_table_transaction_empty_and_abort(catalog: Catalog, format_version: int) -> None: + catalog.create_namespace("replace_test") + data = pa.table({"id": [1]}) + table = catalog.create_table("replace_test.tbl", data.schema, properties={"format-version": str(format_version)}) + table.append(data) + original_metadata = table.metadata + with pytest.raises(ValueError, match="abort replacement"): + with catalog.replace_table_transaction("replace_test.tbl", data.schema) as tx: + tx.append(pa.table({"id": [2]})) + if tx.table_metadata.current_snapshot() is not None: + raise ValueError("abort replacement") + assert catalog.load_table("replace_test.tbl").metadata == original_metadata + + with catalog.replace_table_transaction("replace_test.tbl", Schema(NestedField(8, "empty", StringType()))): + pass + replaced = catalog.load_table("replace_test.tbl") + assert replaced.scan().to_arrow().num_rows == 0 + assert replaced.current_snapshot() is None + assert replaced.metadata.table_uuid == original_metadata.table_uuid + assert replaced.snapshots() == original_metadata.snapshots + assert replaced.schema().find_field("empty").field_type == StringType() + + +def test_replace_table_transaction_missing(catalog: Catalog) -> None: + with pytest.raises(NoSuchTableError): + catalog.replace_table_transaction("missing.tbl", Schema()) + + +@pytest.mark.parametrize("format_version", [1, 2]) +def test_replace_table_transaction_reuses_metadata(catalog: Catalog, format_version: int) -> None: + catalog.create_namespace("replace_test") + schema = Schema(NestedField(10, "id", LongType()), NestedField(20, "value", StringType())) + spec = PartitionSpec(PartitionField(10, 1000, IdentityTransform(), "id")) + order = SortOrder(SortField(source_id=20, transform=IdentityTransform(), direction=SortDirection.DESC)) + table = catalog.create_table( + "replace_test.tbl", + schema, + partition_spec=spec, + sort_order=order, + properties={"format-version": str(format_version), "owner": "old", "keep": "yes"}, + ) + metadata = table.metadata + for _ in range(2): + with catalog.replace_table_transaction( + "replace_test.tbl", schema, partition_spec=spec, sort_order=order, properties={"owner": "new"} + ): + pass + replaced = catalog.load_table("replace_test.tbl") + assert replaced.metadata.schemas == metadata.schemas + assert replaced.metadata.partition_specs == metadata.partition_specs + assert replaced.metadata.sort_orders == metadata.sort_orders + assert replaced.properties == {"owner": "new", "keep": "yes"} + + +@pytest.mark.parametrize("format_version", [1, 2]) +def test_replace_table_transaction_partition_evolution(catalog: Catalog, format_version: int) -> None: + from pyiceberg.transforms import BucketTransform, VoidTransform + + catalog.create_namespace("replace_test") + schema = Schema(NestedField(10, "id", LongType()), NestedField(20, "value", StringType())) + table = catalog.create_table( + "replace_test.tbl", + schema, + partition_spec=PartitionSpec(PartitionField(10, 1000, IdentityTransform(), "id")), + sort_order=SortOrder(SortField(source_id=20, transform=IdentityTransform())), + properties={"format-version": str(format_version)}, + ) + original_partition_id = table.spec().fields[0].field_id + spec = PartitionSpec(PartitionField(20, 1000, BucketTransform(8), "bucket")) + data = pa.table({"id": [1], "value": ["replacement"]}) + with catalog.replace_table_transaction("replace_test.tbl", schema, partition_spec=spec) as tx: + tx.append(data) + replaced = catalog.load_table("replace_test.tbl") + assert replaced.scan().to_arrow().to_pylist() == data.to_pylist() + assert replaced.sort_order().is_unsorted + assert replaced.sort_order().order_id == 0 + assert replaced.spec().fields[-1].source_id == replaced.schema().find_field("value").field_id + assert replaced.spec().fields[-1].field_id > original_partition_id + if format_version == 1: + assert isinstance(replaced.spec().fields[0].transform, VoidTransform) + else: + assert len(replaced.spec().fields) == 1 + + with catalog.replace_table_transaction("replace_test.tbl", schema): + pass + reset = catalog.load_table("replace_test.tbl") + assert reset.sort_order().is_unsorted + assert all(isinstance(field.transform, VoidTransform) for field in reset.spec().fields) + + +@pytest.mark.parametrize("format_version", [1, 2]) +def test_replace_table_transaction_partition_field_reused(catalog: Catalog, format_version: int) -> None: + catalog.create_namespace("replace_test") + schema = Schema(NestedField(10, "id", LongType())) + table = catalog.create_table( + "replace_test.tbl", + schema, + partition_spec=PartitionSpec(PartitionField(10, 1000, IdentityTransform(), "id")), + properties={"format-version": str(format_version)}, + ) + old = table.spec().fields[0] + with catalog.replace_table_transaction( + "replace_test.tbl", schema, partition_spec=PartitionSpec(PartitionField(10, 1007, IdentityTransform(), "renamed")) + ): + pass + field = catalog.load_table("replace_test.tbl").spec().fields[0] + assert field.field_id == old.field_id + assert field.source_id == old.source_id + assert field.name == "renamed" + + +def test_replace_table_transaction_recreated_table(catalog: Catalog) -> None: + catalog.create_namespace("replace_test") + data = pa.table({"id": [1]}) + catalog.create_table("replace_test.tbl", data.schema, properties={"commit.retry.num-retries": "0"}) + tx = catalog.replace_table_transaction("replace_test.tbl", data.schema) + tx.append(data) + catalog.drop_table("replace_test.tbl") + recreated = catalog.create_table("replace_test.tbl", data.schema) + with pytest.raises(CommitFailedException, match="UUID"): + tx.commit_transaction() + assert catalog.load_table("replace_test.tbl").metadata == recreated.metadata + + +def test_replace_table_transaction_concurrent_append(catalog: Catalog) -> None: + catalog.create_namespace("replace_test") + data = pa.table({"id": [1]}) + table = catalog.create_table("replace_test.tbl", data.schema, properties={"commit.retry.min-wait-ms": "0"}) + table.append(data) + tx = catalog.replace_table_transaction("replace_test.tbl", data.schema) + tx.append(pa.table({"id": [3]})) + table.append(pa.table({"id": [2]})) + tx.commit_transaction() + replaced = catalog.load_table("replace_test.tbl") + assert replaced.scan().to_arrow().to_pylist() == [{"id": 3}] + assert len(replaced.snapshots()) == 3 + assert replaced.metadata.last_sequence_number == 3 + current_snapshot = replaced.current_snapshot() + assert current_snapshot is not None + assert current_snapshot.parent_snapshot_id is None + + +def test_replace_table_transaction_concurrent_schema_change(catalog: Catalog) -> None: + catalog.create_namespace("replace_test") + table = catalog.create_table("replace_test.tbl", Schema(NestedField(1, "id", LongType()))) + tx = catalog.replace_table_transaction("replace_test.tbl", Schema(NestedField(1, "replacement", StringType()))) + table.update_schema().add_column("concurrent", LongType()).commit() + with pytest.raises(CommitFailedException): + tx.commit_transaction() + assert catalog.load_table("replace_test.tbl").schema() == table.schema() + + +def test_replace_table_transaction_new_location_and_format(catalog: Catalog, tmp_path: Path) -> None: + catalog.create_namespace("replace_test") + table = catalog.create_table("replace_test.tbl", Schema(), properties={"format-version": "1"}) + properties = {"format-version": "2", "owner": "replacement"} + location = f"file://{tmp_path}/replacement" + with catalog.replace_table_transaction("replace_test.tbl", Schema(), location=location + "/", properties=properties): + pass + replaced = catalog.load_table("replace_test.tbl") + assert replaced.location() == location + assert replaced.metadata.format_version == 2 + assert replaced.properties == {"owner": "replacement"} + assert properties == {"format-version": "2", "owner": "replacement"} + assert replaced.metadata.table_uuid == table.metadata.table_uuid + with pytest.raises(ValueError, match="Cannot downgrade"): + catalog.replace_table_transaction("replace_test.tbl", Schema(), properties={"format-version": "1"}).commit_transaction() + assert catalog.load_table("replace_test.tbl").metadata == replaced.metadata + + +@pytest.mark.parametrize("format_version", [1, 2]) +def test_replace_table_transaction_drops_partition_source(catalog: Catalog, format_version: int) -> None: + catalog.create_namespace("replace_test") + catalog.create_table( + "replace_test.tbl", + Schema(NestedField(1, "old", LongType())), + partition_spec=PartitionSpec(PartitionField(1, 1000, IdentityTransform(), "old")), + properties={"format-version": str(format_version)}, + ) + data = pa.table({"new": ["replacement"]}) + with catalog.replace_table_transaction("replace_test.tbl", data.schema) as tx: + tx.append(data) + assert catalog.load_table("replace_test.tbl").scan().to_arrow().to_pylist() == data.to_pylist() + + +def test_replace_table_transaction_reuses_historical_schema(catalog: Catalog) -> None: + catalog.create_namespace("replace_test") + original_schema = Schema(NestedField(1, "id", LongType())) + original = catalog.create_table("replace_test.tbl", original_schema) + with catalog.replace_table_transaction("replace_test.tbl", Schema(NestedField(10, "id", StringType()))): + pass + with catalog.replace_table_transaction("replace_test.tbl", original_schema): + pass + table = catalog.load_table("replace_test.tbl") + assert len(table.metadata.schemas) == 2 + assert table.metadata.current_schema_id == original.metadata.current_schema_id + + +def test_replace_table_transaction_preserves_field_id_high_water_mark(catalog: Catalog) -> None: + catalog.create_namespace("replace_test") + table = catalog.create_table( + "replace_test.tbl", Schema(NestedField(1, "id", LongType()), NestedField(2, "removed", LongType())) + ) + removed_id = table.schema().find_field("removed").field_id + table.update_schema().delete_column("removed").commit() + with catalog.replace_table_transaction("replace_test.tbl", Schema(NestedField(9, "new", StringType()))) as tx: + with tx.update_schema() as update: + update.add_column("extra", LongType()) + tx.append(pa.table({"new": ["replacement"], "extra": [42]})) + replaced = catalog.load_table("replace_test.tbl") + assert replaced.schema().find_field("new").field_id > removed_id + assert replaced.schema().find_field("extra").field_id > replaced.schema().find_field("new").field_id + assert replaced.scan().to_arrow().to_pylist() == [{"new": "replacement", "extra": 42}] + + +def test_replace_table_transaction_invalid_layout(catalog: Catalog) -> None: + from pyiceberg.exceptions import ValidationError + from pyiceberg.transforms import DayTransform + from pyiceberg.types import TimestampNanoType + + catalog.create_namespace("replace_test") + schema = Schema(NestedField(1, "id", LongType())) + table = catalog.create_table("replace_test.tbl", schema) + with pytest.raises(ValidationError, match="Invalid source field"): + catalog.replace_table_transaction( + "replace_test.tbl", schema, partition_spec=PartitionSpec(PartitionField(1, 1000, DayTransform(), "day")) + ) + with pytest.raises(ValidationError, match="Invalid source field"): + catalog.replace_table_transaction("replace_test.tbl", schema, sort_order=SortOrder(SortField(1, DayTransform()))) + with pytest.raises(ValueError, match="only supported in 3 or higher"): + catalog.replace_table_transaction("replace_test.tbl", Schema(NestedField(1, "ts", TimestampNanoType()))) + assert catalog.load_table("replace_test.tbl").metadata == table.metadata diff --git a/tests/catalog/test_rest.py b/tests/catalog/test_rest.py index ad1e3c3296..3667d9ccc0 100644 --- a/tests/catalog/test_rest.py +++ b/tests/catalog/test_rest.py @@ -70,12 +70,13 @@ from pyiceberg.io import ARROW_FILE_IO, FSSPEC_FILE_IO, PY_IO_IMPL, load_file_io from pyiceberg.partitioning import PartitionField, PartitionSpec from pyiceberg.schema import Schema -from pyiceberg.table import Table +from pyiceberg.table import CommitTableRequest, Table from pyiceberg.table.metadata import TableMetadataV1 from pyiceberg.table.sorting import SortField, SortOrder +from pyiceberg.table.update import update_table_metadata from pyiceberg.transforms import IdentityTransform, TruncateTransform from pyiceberg.typedef import RecursiveDict -from pyiceberg.types import StringType +from pyiceberg.types import NestedField, StringType from pyiceberg.utils.config import Config from pyiceberg.view import View from pyiceberg.view.metadata import ViewMetadata, ViewVersion @@ -3713,3 +3714,55 @@ def test_load_table_without_storage_credentials( ) assert actual.metadata.model_dump() == expected.metadata.model_dump() assert actual == expected + + +def test_replace_table_transaction_uses_commit_endpoint( + rest_mock: Mocker, example_table_metadata_with_snapshot_v1_rest_json: dict[str, Any] +) -> None: + response = example_table_metadata_with_snapshot_v1_rest_json + base = TableMetadataV1.model_validate(response["metadata"]) + url = f"{TEST_URI}v1/namespaces/fokko/tables/table" + rest_mock.get(url, json=response, request_headers=TEST_HEADERS) + catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN) + tx = catalog.replace_table_transaction("fokko.table", Schema(NestedField(88, "replacement", StringType()))) + updated_metadata = tx.table_metadata + rest_mock.post( + url, + json={ + "metadata": updated_metadata.model_dump(mode="json", by_alias=True), + "metadata-location": response["metadata-location"], + }, + request_headers=TEST_HEADERS, + ) + tx.commit_transaction() + + request = CommitTableRequest.model_validate(rest_mock.last_request.json()) + assert request.identifier.namespace.root == ["fokko"] + assert request.identifier.name == "table" + for requirement in request.requirements: + requirement.validate(base) + reconstructed = update_table_metadata(base, request.updates) + assert reconstructed.model_dump(exclude={"last_updated_ms"}) == updated_metadata.model_dump(exclude={"last_updated_ms"}) + assert reconstructed.table_uuid == base.table_uuid + assert reconstructed.current_snapshot_id is None + assert reconstructed.snapshots == base.snapshots + assert reconstructed.schema().find_field("replacement").field_id > base.last_column_id + assert [req.type for req in request.requirements].count("assert-table-uuid") == 1 + assert [entry.method for entry in rest_mock.request_history] == ["GET", "GET", "POST"] + + +def test_replace_table_transaction_without_sort_order_history( + rest_mock: Mocker, example_table_metadata_with_snapshot_v1_rest_json: dict[str, Any] +) -> None: + response = example_table_metadata_with_snapshot_v1_rest_json + metadata = TableMetadataV1.model_validate(response["metadata"]).to_v2() + metadata = metadata.model_copy(update={"sort_orders": [], "default_sort_order_id": 0}) + response["metadata"] = metadata.model_dump(mode="json", by_alias=True) + rest_mock.get(f"{TEST_URI}v1/namespaces/fokko/tables/table", json=response) + catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN) + tx = catalog.replace_table_transaction( + "fokko.table", + Schema(NestedField(1, "replacement", StringType())), + sort_order=SortOrder(SortField(1, IdentityTransform())), + ) + assert tx.table_metadata.sort_order().order_id == 1 diff --git a/tests/integration/test_catalog.py b/tests/integration/test_catalog.py index ab9589cb85..84e7943ed9 100644 --- a/tests/integration/test_catalog.py +++ b/tests/integration/test_catalog.py @@ -88,7 +88,15 @@ def sqlite_catalog_file(warehouse: Path) -> Generator[Catalog, None, None]: @pytest.fixture(scope="function") def rest_catalog() -> Generator[Catalog, None, None]: - test_catalog = RestCatalog("rest", uri="http://localhost:8181") + test_catalog = RestCatalog( + "rest", + **{ + "uri": "http://localhost:8181", + "s3.endpoint": "http://localhost:9000", + "s3.access-key-id": "admin", + "s3.secret-access-key": "password", + }, + ) yield test_catalog @@ -1058,3 +1066,48 @@ def test_load_missing_table(test_catalog: Catalog, database_name: str, table_nam with pytest.raises(NoSuchTableError): test_catalog.load_table(identifier) + + +@pytest.mark.integration +@pytest.mark.parametrize("test_catalog", CATALOGS) +def test_replace_table_transaction(test_catalog: Catalog, table_name: str, database_name: str) -> None: + import pyarrow as pa + + test_catalog.create_namespace(database_name) + original = pa.table({"id": [1], "value": ["original"]}) + for format_version in (1, 2): + identifier = (database_name, f"{table_name}_v{format_version}") + table = test_catalog.create_table( + identifier, + Schema(NestedField(1, "id", LongType()), NestedField(2, "value", StringType())), + partition_spec=PartitionSpec(PartitionField(1, 1000, IdentityTransform(), "id")), + sort_order=SortOrder(SortField(source_id=2, transform=IdentityTransform())), + properties={"format-version": str(format_version)}, + ) + table.append(original) + old_snapshot = table.current_snapshot() + assert old_snapshot is not None + table.manage_snapshots().create_tag(old_snapshot.snapshot_id, "original").commit() + original_metadata = table.metadata + + replacement = pa.table({"id": ["replacement"], "new": [True]}) + with test_catalog.replace_table_transaction(identifier, replacement.schema) as tx: + tx.append(replacement) + assert test_catalog.load_table(identifier).scan().to_arrow().to_pylist() == original.to_pylist() + + replaced = test_catalog.load_table(identifier) + assert replaced.scan().to_arrow().to_pylist() == replacement.to_pylist() + assert replaced.metadata.table_uuid == original_metadata.table_uuid + assert replaced.metadata.snapshot_by_id(old_snapshot.snapshot_id) == old_snapshot + assert replaced.metadata.refs["original"].snapshot_id == old_snapshot.snapshot_id + assert replaced.sort_order().order_id == 0 + new_snapshot = replaced.current_snapshot() + assert new_snapshot is not None + assert new_snapshot.parent_snapshot_id is None + + with test_catalog.replace_table_transaction(identifier, replacement.schema): + pass + empty = test_catalog.load_table(identifier) + assert empty.current_snapshot() is None + assert len(empty.snapshots()) == 2 + assert empty.scan().to_arrow().num_rows == 0 diff --git a/tests/test_schema.py b/tests/test_schema.py index 94931277b4..ba529321dd 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -26,6 +26,7 @@ Accessor, Schema, _check_schema_compatible, + assign_fresh_schema_ids, build_position_accessors, index_by_id, index_by_name, @@ -1844,3 +1845,52 @@ def test_check_schema_compatible_optional_map_field_present() -> None: ) # Should not raise - schemas match _check_schema_compatible(requested_schema, provided_schema) + + +def test_assign_fresh_schema_ids_reuses_nested_names() -> None: + from itertools import count + + base = Schema( + NestedField(1, "id", LongType(), required=True), + NestedField(2, "nested", StructType(NestedField(3, "kept", StringType()), NestedField(4, "removed", LongType()))), + NestedField(5, "items", ListType(6, StructType(NestedField(7, "value", LongType())))), + NestedField(8, "mapping", MapType(9, StringType(), 10, StructType(NestedField(11, "value", LongType())))), + identifier_field_ids=[1], + ) + replacement = Schema( + NestedField(21, "id", StringType(), required=True, doc="replacement", initial_default="initial", write_default="write"), + NestedField(22, "nested", StructType(NestedField(23, "kept", LongType()), NestedField(24, "added", StringType()))), + NestedField(25, "items", ListType(26, StructType(NestedField(27, "value", StringType())))), + NestedField(28, "mapping", MapType(29, StringType(), 30, StructType(NestedField(31, "value", StringType())))), + identifier_field_ids=[21], + ) + next_id = count(12) + actual = assign_fresh_schema_ids(replacement, lambda: next(next_id), base_schema=base) + assert actual.find_field("id").field_id == 1 + assert actual.find_field("id").field_type == StringType() + assert actual.find_field("id").doc == "replacement" + assert actual.find_field("id").initial_default == "initial" + assert actual.find_field("id").write_default == "write" + assert actual.identifier_field_ids == [1] + assert actual.find_field("nested.kept").field_id == 3 + assert actual.find_field("nested.added").field_id == 12 + assert actual.find_field("items.element").field_id == 6 + assert actual.find_field("items.element.value").field_id == 7 + assert actual.find_field("mapping.key").field_id == 9 + assert actual.find_field("mapping.value").field_id == 10 + assert actual.find_field("mapping.value.value").field_id == 11 + assert next(next_id) == 13 + assert assign_fresh_schema_ids(replacement, base_schema=base) == actual + + +def test_assign_fresh_schema_ids_without_base() -> None: + schema = Schema( + NestedField(9, "id", LongType(), required=True), + NestedField(12, "nested", StructType(NestedField(25, "value", StringType()))), + identifier_field_ids=[9], + ) + actual = assign_fresh_schema_ids(schema) + assert actual.find_field("id").field_id == 1 + assert actual.find_field("nested").field_id == 2 + assert actual.find_field("nested.value").field_id == 3 + assert actual.identifier_field_ids == [1] diff --git a/tests/utils/test_schema_conversion.py b/tests/utils/test_schema_conversion.py index a51b07075c..cb5e3a3ec4 100644 --- a/tests/utils/test_schema_conversion.py +++ b/tests/utils/test_schema_conversion.py @@ -414,3 +414,12 @@ def test_iceberg_to_avro_manifest(avro_schema_manifest_entry: dict[str, Any]) -> iceberg_schema = AvroSchemaConversion().avro_to_iceberg(avro_schema_manifest_entry) avro_result = AvroSchemaConversion().iceberg_to_avro(iceberg_schema, schema_name="manifest_entry") assert avro_schema_manifest_entry == avro_result + + +def test_optional_unknown_type_uses_avro_null() -> None: + converter = AvroSchemaConversion() + avro = converter.iceberg_to_avro(Schema(NestedField(1000, "removed_partition", UnknownType()))) + assert isinstance(avro, dict) + assert avro["fields"][0]["type"] == "null" + restored = converter.avro_to_iceberg(avro) + assert restored.find_field("removed_partition").field_type == UnknownType()