Feature: Write to branches

pyiceberg/cli/console.py

@@ -33,7 +33,7 @@
 from pyiceberg.cli.output import ConsoleOutput, JsonOutput, Output
 from pyiceberg.exceptions import NoSuchNamespaceError, NoSuchPropertyException, NoSuchTableError
 from pyiceberg.table import TableProperties
-from pyiceberg.table.refs import SnapshotRef
+from pyiceberg.table.refs import SnapshotRef, SnapshotRefType
 from pyiceberg.utils.properties import property_as_int
 
 
@@ -419,7 +419,7 @@ def list_refs(ctx: Context, identifier: str, type: str, verbose: bool) -> None:
     refs = table.refs()
     if type:
         type = type.lower()
-        if type not in {"branch", "tag"}:
+        if type not in {SnapshotRefType.BRANCH, SnapshotRefType.TAG}:
             raise ValueError(f"Type must be either branch or tag, got: {type}")
 
     relevant_refs = [
@@ -433,7 +433,7 @@ def list_refs(ctx: Context, identifier: str, type: str, verbose: bool) -> None:
 
 def _retention_properties(ref: SnapshotRef, table_properties: Dict[str, str]) -> Dict[str, str]:
     retention_properties = {}
-    if ref.snapshot_ref_type == "branch":
+    if ref.snapshot_ref_type == SnapshotRefType.BRANCH:
         default_min_snapshots_to_keep = property_as_int(
             table_properties,
             TableProperties.MIN_SNAPSHOTS_TO_KEEP,

pyiceberg/table/update/__init__.py

@@ -29,7 +29,7 @@
 from pyiceberg.partitioning import PARTITION_FIELD_ID_START, PartitionSpec
 from pyiceberg.schema import Schema
 from pyiceberg.table.metadata import SUPPORTED_TABLE_FORMAT_VERSION, TableMetadata, TableMetadataUtil
-from pyiceberg.table.refs import MAIN_BRANCH, SnapshotRef
+from pyiceberg.table.refs import MAIN_BRANCH, SnapshotRef, SnapshotRefType
 from pyiceberg.table.snapshots import (
     MetadataLogEntry,
     Snapshot,
@@ -139,7 +139,7 @@ class AddSnapshotUpdate(IcebergBaseModel):
 class SetSnapshotRefUpdate(IcebergBaseModel):
     action: Literal["set-snapshot-ref"] = Field(default="set-snapshot-ref")
     ref_name: str = Field(alias="ref-name")
-    type: Literal["tag", "branch"]
+    type: Literal[SnapshotRefType.TAG, SnapshotRefType.BRANCH]
     snapshot_id: int = Field(alias="snapshot-id")
     max_ref_age_ms: Annotated[Optional[int], Field(alias="max-ref-age-ms", default=None)]
     max_snapshot_age_ms: Annotated[Optional[int], Field(alias="max-snapshot-age-ms", default=None)]
@@ -702,6 +702,10 @@ class AssertRefSnapshotId(ValidatableTableRequirement):
     def validate(self, base_metadata: Optional[TableMetadata]) -> None:
         if base_metadata is None:
             raise CommitFailedException("Requirement failed: current table metadata is missing")
+        elif len(base_metadata.snapshots) == 0 and self.ref != MAIN_BRANCH:
+            raise CommitFailedException(
+                f"Requirement failed: Table has no snapshots and can only be written to the {MAIN_BRANCH} BRANCH."
+            )
         elif snapshot_ref := base_metadata.refs.get(self.ref):
             ref_type = snapshot_ref.snapshot_ref_type
             if self.snapshot_id is None:

pyiceberg/table/update/snapshot.py

@@ -55,6 +55,7 @@
 from pyiceberg.partitioning import (
     PartitionSpec,
 )
+from pyiceberg.table.refs import MAIN_BRANCH, SnapshotRefType
 from pyiceberg.table.snapshots import (
     Operation,
     Snapshot,
@@ -104,6 +105,7 @@ class _SnapshotProducer(UpdateTableMetadata[U], Generic[U]):
     _added_data_files: List[DataFile]
     _manifest_num_counter: itertools.count[int]
     _deleted_data_files: Set[DataFile]
+    _target_branch = MAIN_BRANCH
 
     def __init__(
         self,
@@ -112,20 +114,31 @@ def __init__(
         io: FileIO,
         commit_uuid: Optional[uuid.UUID] = None,
         snapshot_properties: Dict[str, str] = EMPTY_DICT,
+        branch: str = MAIN_BRANCH,
     ) -> None:
         super().__init__(transaction)
         self.commit_uuid = commit_uuid or uuid.uuid4()
         self._io = io
         self._operation = operation
         self._snapshot_id = self._transaction.table_metadata.new_snapshot_id()
-        # Since we only support the main branch for now
-        self._parent_snapshot_id = (
-            snapshot.snapshot_id if (snapshot := self._transaction.table_metadata.current_snapshot()) else None
-        )
         self._added_data_files = []
         self._deleted_data_files = set()
         self.snapshot_properties = snapshot_properties
         self._manifest_num_counter = itertools.count(0)
+        self._target_branch = self._validate_target_branch(branch=branch)
+        self._parent_snapshot_id = (
+            snapshot.snapshot_id if (snapshot := self._transaction.table_metadata.snapshot_by_name(self._target_branch)) else None
+        )
+
+    def _validate_target_branch(self, branch: str) -> str:
+        # Default is already set to MAIN_BRANCH. So branch name can't be None.
+        if branch is None:
+            raise ValueError("Invalid branch name: null")
+        if branch in self._transaction.table_metadata.refs:
+            ref = self._transaction.table_metadata.refs[branch]
+            if ref.snapshot_ref_type != SnapshotRefType.BRANCH:
+                raise ValueError(f"{branch} is a tag, not a branch. Tags cannot be targets for producing snapshots")
+        return branch
 
     def append_data_file(self, data_file: DataFile) -> _SnapshotProducer[U]:
         self._added_data_files.append(data_file)
@@ -271,10 +284,20 @@ def _commit(self) -> UpdatesAndRequirements:
             (
                 AddSnapshotUpdate(snapshot=snapshot),
                 SetSnapshotRefUpdate(
-                    snapshot_id=self._snapshot_id, parent_snapshot_id=self._parent_snapshot_id, ref_name="main", type="branch"
+                    snapshot_id=self._snapshot_id,
+                    parent_snapshot_id=self._parent_snapshot_id,
+                    ref_name=self._target_branch,
+                    type=SnapshotRefType.BRANCH,
+                ),
+            ),
+            (
+                AssertRefSnapshotId(
+                    snapshot_id=self._transaction.table_metadata.refs[self._target_branch].snapshot_id
+                    if self._target_branch in self._transaction.table_metadata.refs
+                    else None,
+                    ref=self._target_branch,
                 ),
             ),
-            (AssertRefSnapshotId(snapshot_id=self._transaction.table_metadata.current_snapshot_id, ref="main"),),
         )
 
     @property
@@ -321,10 +344,11 @@ def __init__(
         operation: Operation,
         transaction: Transaction,
         io: FileIO,
+        branch: str,
         commit_uuid: Optional[uuid.UUID] = None,
         snapshot_properties: Dict[str, str] = EMPTY_DICT,
     ):
-        super().__init__(operation, transaction, io, commit_uuid, snapshot_properties)
+        super().__init__(operation, transaction, io, commit_uuid, snapshot_properties, branch)
         self._predicate = AlwaysFalse()
         self._case_sensitive = True
 
@@ -384,46 +408,52 @@ def _copy_with_new_status(entry: ManifestEntry, status: ManifestEntryStatus) ->
         total_deleted_entries = []
         partial_rewrites_needed = False
         self._deleted_data_files = set()
-        if snapshot := self._transaction.table_metadata.current_snapshot():
-            for manifest_file in snapshot.manifests(io=self._io):
-                if manifest_file.content == ManifestContent.DATA:
-                    if not manifest_evaluators[manifest_file.partition_spec_id](manifest_file):
-                        # If the manifest isn't relevant, we can just keep it in the manifest-list
-                        existing_manifests.append(manifest_file)
-                    else:
-                        # It is relevant, let's check out the content
-                        deleted_entries = []
-                        existing_entries = []
-                        for entry in manifest_file.fetch_manifest_entry(io=self._io, discard_deleted=True):
-                            if strict_metrics_evaluator(entry.data_file) == ROWS_MUST_MATCH:
-                                # Based on the metadata, it can be dropped right away
-                                deleted_entries.append(_copy_with_new_status(entry, ManifestEntryStatus.DELETED))
-                                self._deleted_data_files.add(entry.data_file)
-                            else:
-                                # Based on the metadata, we cannot determine if it can be deleted
-                                existing_entries.append(_copy_with_new_status(entry, ManifestEntryStatus.EXISTING))
-                                if inclusive_metrics_evaluator(entry.data_file) != ROWS_MIGHT_NOT_MATCH:
-                                    partial_rewrites_needed = True
-
-                        if len(deleted_entries) > 0:
-                            total_deleted_entries += deleted_entries
-
-                            # Rewrite the manifest
-                            if len(existing_entries) > 0:
-                                with write_manifest(
-                                    format_version=self._transaction.table_metadata.format_version,
-                                    spec=self._transaction.table_metadata.specs()[manifest_file.partition_spec_id],
-                                    schema=self._transaction.table_metadata.schema(),
-                                    output_file=self.new_manifest_output(),
-                                    snapshot_id=self._snapshot_id,
-                                ) as writer:
-                                    for existing_entry in existing_entries:
-                                        writer.add_entry(existing_entry)
-                                existing_manifests.append(writer.to_manifest_file())
-                        else:
+
+        # Determine the snapshot to read manifests from for deletion
+        # Should be the current tip of the _target_branch
+        parent_snapshot_id_for_delete_source = self._parent_snapshot_id
+        if parent_snapshot_id_for_delete_source is not None:
+            snapshot = self._transaction.table_metadata.snapshot_by_id(parent_snapshot_id_for_delete_source)
+            if snapshot:  # Ensure snapshot is found
+                for manifest_file in snapshot.manifests(io=self._io):
+                    if manifest_file.content == ManifestContent.DATA:
+                        if not manifest_evaluators[manifest_file.partition_spec_id](manifest_file):
+                            # If the manifest isn't relevant, we can just keep it in the manifest-list
                             existing_manifests.append(manifest_file)
-                else:
-                    existing_manifests.append(manifest_file)
+                        else:
+                            # It is relevant, let's check out the content
+                            deleted_entries = []
+                            existing_entries = []
+                            for entry in manifest_file.fetch_manifest_entry(io=self._io, discard_deleted=True):
+                                if strict_metrics_evaluator(entry.data_file) == ROWS_MUST_MATCH:
+                                    # Based on the metadata, it can be dropped right away
+                                    deleted_entries.append(_copy_with_new_status(entry, ManifestEntryStatus.DELETED))
+                                    self._deleted_data_files.add(entry.data_file)
+                                else:
+                                    # Based on the metadata, we cannot determine if it can be deleted
+                                    existing_entries.append(_copy_with_new_status(entry, ManifestEntryStatus.EXISTING))
+                                    if inclusive_metrics_evaluator(entry.data_file) != ROWS_MIGHT_NOT_MATCH:
+                                        partial_rewrites_needed = True
+
+                            if len(deleted_entries) > 0:
+                                total_deleted_entries += deleted_entries
+
+                                # Rewrite the manifest
+                                if len(existing_entries) > 0:
+                                    with write_manifest(
+                                        format_version=self._transaction.table_metadata.format_version,
+                                        spec=self._transaction.table_metadata.specs()[manifest_file.partition_spec_id],
+                                        schema=self._transaction.table_metadata.schema(),
+                                        output_file=self.new_manifest_output(),
+                                        snapshot_id=self._snapshot_id,
+                                    ) as writer:
+                                        for existing_entry in existing_entries:
+                                            writer.add_entry(existing_entry)
+                                    existing_manifests.append(writer.to_manifest_file())
+                            else:
+                                existing_manifests.append(manifest_file)
+                    else:
+                        existing_manifests.append(manifest_file)
 
         return existing_manifests, total_deleted_entries, partial_rewrites_needed
 
@@ -483,12 +513,13 @@ def __init__(
         operation: Operation,
         transaction: Transaction,
         io: FileIO,
+        branch: str,
         commit_uuid: Optional[uuid.UUID] = None,
         snapshot_properties: Dict[str, str] = EMPTY_DICT,
     ) -> None:
         from pyiceberg.table import TableProperties
 
-        super().__init__(operation, transaction, io, commit_uuid, snapshot_properties)
+        super().__init__(operation, transaction, io, commit_uuid, snapshot_properties, branch)
         self._target_size_bytes = property_as_int(
             self._transaction.table_metadata.properties,
             TableProperties.MANIFEST_TARGET_SIZE_BYTES,
@@ -534,7 +565,7 @@ def _existing_manifests(self) -> List[ManifestFile]:
         """Determine if there are any existing manifest files."""
         existing_files = []
 
-        if snapshot := self._transaction.table_metadata.current_snapshot():
+        if snapshot := self._transaction.table_metadata.snapshot_by_name(name=self._target_branch):
             for manifest_file in snapshot.manifests(io=self._io):
                 entries = manifest_file.fetch_manifest_entry(io=self._io, discard_deleted=True)
                 found_deleted_data_files = [entry.data_file for entry in entries if entry.data_file in self._deleted_data_files]
@@ -551,19 +582,17 @@ def _existing_manifests(self) -> List[ManifestFile]:
                             output_file=self.new_manifest_output(),
                             snapshot_id=self._snapshot_id,
                         ) as writer:
-                            [
-                                writer.add_entry(
-                                    ManifestEntry.from_args(
-                                        status=ManifestEntryStatus.EXISTING,
-                                        snapshot_id=entry.snapshot_id,
-                                        sequence_number=entry.sequence_number,
-                                        file_sequence_number=entry.file_sequence_number,
-                                        data_file=entry.data_file,
+                            for entry in entries:
+                                if entry.data_file not in found_deleted_data_files:
+                                    writer.add_entry(
+                                        ManifestEntry.from_args(
+                                            status=ManifestEntryStatus.EXISTING,
+                                            snapshot_id=entry.snapshot_id,
+                                            sequence_number=entry.sequence_number,
+                                            file_sequence_number=entry.file_sequence_number,
+                                            data_file=entry.data_file,
+                                        )
                                     )
-                                )
-                                for entry in entries
-                                if entry.data_file not in found_deleted_data_files
-                            ]
                         existing_files.append(writer.to_manifest_file())
         return existing_files
 
@@ -604,21 +633,37 @@ def _get_entries(manifest: ManifestFile) -> List[ManifestEntry]:
 class UpdateSnapshot:
     _transaction: Transaction
     _io: FileIO
+    _branch: str
     _snapshot_properties: Dict[str, str]
 
-    def __init__(self, transaction: Transaction, io: FileIO, snapshot_properties: Dict[str, str] = EMPTY_DICT) -> None:
+    def __init__(
+        self,
+        transaction: Transaction,
+        io: FileIO,
+        branch: str,
+        snapshot_properties: Dict[str, str] = EMPTY_DICT,
+    ) -> None:
         self._transaction = transaction
         self._io = io
         self._snapshot_properties = snapshot_properties
+        self._branch = branch
 
     def fast_append(self) -> _FastAppendFiles:
         return _FastAppendFiles(
-            operation=Operation.APPEND, transaction=self._transaction, io=self._io, snapshot_properties=self._snapshot_properties
+            operation=Operation.APPEND,
+            transaction=self._transaction,
+            io=self._io,
+            branch=self._branch,
+            snapshot_properties=self._snapshot_properties,
         )
 
     def merge_append(self) -> _MergeAppendFiles:
         return _MergeAppendFiles(
-            operation=Operation.APPEND, transaction=self._transaction, io=self._io, snapshot_properties=self._snapshot_properties
+            operation=Operation.APPEND,
+            transaction=self._transaction,
+            io=self._io,
+            branch=self._branch,
+            snapshot_properties=self._snapshot_properties,
         )
 
     def overwrite(self, commit_uuid: Optional[uuid.UUID] = None) -> _OverwriteFiles:
@@ -629,6 +674,7 @@ def overwrite(self, commit_uuid: Optional[uuid.UUID] = None) -> _OverwriteFiles:
             else Operation.APPEND,
             transaction=self._transaction,
             io=self._io,
+            branch=self._branch,
             snapshot_properties=self._snapshot_properties,
         )
 
@@ -637,6 +683,7 @@ def delete(self) -> _DeleteFiles:
             operation=Operation.DELETE,
             transaction=self._transaction,
             io=self._io,
+            branch=self._branch,
             snapshot_properties=self._snapshot_properties,
         )
 

pyiceberg/utils/concurrent.py

@@ -25,6 +25,11 @@
 class ExecutorFactory:
     _instance: Optional[Executor] = None
 
+    @staticmethod
+    def max_workers() -> Optional[int]:
+        """Return the max number of workers configured."""
+        return Config().get_int("max-workers")
+
     @staticmethod
     def get_or_create() -> Executor:
         """Return the same executor in each call."""
@@ -33,8 +38,3 @@ def get_or_create() -> Executor:
             ExecutorFactory._instance = ThreadPoolExecutor(max_workers=max_workers)
 
         return ExecutorFactory._instance
-
-    @staticmethod
-    def max_workers() -> Optional[int]:
-        """Return the max number of workers configured."""
-        return Config().get_int("max-workers")