wip

wip

Author: Yingjian Wu

pyiceberg/table/__init__.py

@@ -397,7 +397,9 @@ def _build_partition_predicate(self, partition_records: Set[Record]) -> BooleanE
             expr = Or(expr, match_partition_expression)
         return expr
 
-    def _append_snapshot_producer(self, snapshot_properties: Dict[str, str], branch: Optional[str]) -> _FastAppendFiles:
+    def _append_snapshot_producer(
+        self, snapshot_properties: Dict[str, str], branch: Optional[str] = MAIN_BRANCH
+    ) -> _FastAppendFiles:
         """Determine the append type based on table properties.
 
         Args:
@@ -431,19 +433,14 @@ def update_schema(self, allow_incompatible_changes: bool = False, case_sensitive
         )
 
     def update_snapshot(
-        self, snapshot_properties: Dict[str, str] = EMPTY_DICT, branch: Optional[str] = None, stage_only: bool = False
+        self, snapshot_properties: Dict[str, str] = EMPTY_DICT, branch: Optional[str] = MAIN_BRANCH
     ) -> UpdateSnapshot:
         """Create a new UpdateSnapshot to produce a new snapshot for the table.
 
         Returns:
             A new UpdateSnapshot
         """
-        if branch is None:
-            branch = MAIN_BRANCH
-
-        return UpdateSnapshot(
-            self, io=self._table.io, branch=branch, snapshot_properties=snapshot_properties, stage_only=stage_only
-        )
+        return UpdateSnapshot(self, io=self._table.io, branch=branch, snapshot_properties=snapshot_properties)
 
     def update_statistics(self) -> UpdateStatistics:
         """
@@ -454,7 +451,7 @@ def update_statistics(self) -> UpdateStatistics:
         """
         return UpdateStatistics(transaction=self)
 
-    def append(self, df: pa.Table, snapshot_properties: Dict[str, str] = EMPTY_DICT, branch: Optional[str] = None) -> None:
+    def append(self, df: pa.Table, snapshot_properties: Dict[str, str] = EMPTY_DICT, branch: Optional[str] = MAIN_BRANCH) -> None:
         """
         Shorthand API for appending a PyArrow table to a table transaction.
 
@@ -499,7 +496,7 @@ def append(self, df: pa.Table, snapshot_properties: Dict[str, str] = EMPTY_DICT,
                     append_files.append_data_file(data_file)
 
     def dynamic_partition_overwrite(
-        self, df: pa.Table, snapshot_properties: Dict[str, str] = EMPTY_DICT, branch: Optional[str] = None
+        self, df: pa.Table, snapshot_properties: Dict[str, str] = EMPTY_DICT, branch: Optional[str] = MAIN_BRANCH
     ) -> None:
         """
         Shorthand for overwriting existing partitions with a PyArrow table.
@@ -566,7 +563,7 @@ def overwrite(
         overwrite_filter: Union[BooleanExpression, str] = ALWAYS_TRUE,
         snapshot_properties: Dict[str, str] = EMPTY_DICT,
         case_sensitive: bool = True,
-        branch: Optional[str] = None,
+        branch: Optional[str] = MAIN_BRANCH,
     ) -> None:
         """
         Shorthand for adding a table overwrite with a PyArrow table to the transaction.
@@ -632,7 +629,7 @@ def delete(
         delete_filter: Union[str, BooleanExpression],
         snapshot_properties: Dict[str, str] = EMPTY_DICT,
         case_sensitive: bool = True,
-        branch: Optional[str] = None,
+        branch: Optional[str] = MAIN_BRANCH,
     ) -> None:
         """
         Shorthand for deleting record from a table.
@@ -735,7 +732,7 @@ def upsert(
         when_matched_update_all: bool = True,
         when_not_matched_insert_all: bool = True,
         case_sensitive: bool = True,
-        branch: Optional[str] = None,
+        branch: Optional[str] = MAIN_BRANCH,
     ) -> UpsertResult:
         """Shorthand API for performing an upsert to an iceberg table.
 
@@ -820,7 +817,7 @@ def upsert(
             case_sensitive=case_sensitive,
         )
 
-        if branch is not None:
+        if branch in self.table_metadata.refs:
             matched_iceberg_record_batches_scan = matched_iceberg_record_batches_scan.use_ref(branch)
 
         matched_iceberg_record_batches = matched_iceberg_record_batches_scan.to_arrow_batch_reader()
@@ -1311,7 +1308,7 @@ def upsert(
         when_matched_update_all: bool = True,
         when_not_matched_insert_all: bool = True,
         case_sensitive: bool = True,
-        branch: Optional[str] = None,
+        branch: Optional[str] = MAIN_BRANCH,
     ) -> UpsertResult:
         """Shorthand API for performing an upsert to an iceberg table.
 
@@ -1358,7 +1355,7 @@ def upsert(
                 branch=branch,
             )
 
-    def append(self, df: pa.Table, snapshot_properties: Dict[str, str] = EMPTY_DICT, branch: Optional[str] = None) -> None:
+    def append(self, df: pa.Table, snapshot_properties: Dict[str, str] = EMPTY_DICT, branch: Optional[str] = MAIN_BRANCH) -> None:
         """
         Shorthand API for appending a PyArrow table to the table.
 
@@ -1371,7 +1368,7 @@ def append(self, df: pa.Table, snapshot_properties: Dict[str, str] = EMPTY_DICT,
             tx.append(df=df, snapshot_properties=snapshot_properties, branch=branch)
 
     def dynamic_partition_overwrite(
-        self, df: pa.Table, snapshot_properties: Dict[str, str] = EMPTY_DICT, branch: Optional[str] = None
+        self, df: pa.Table, snapshot_properties: Dict[str, str] = EMPTY_DICT, branch: Optional[str] = MAIN_BRANCH
     ) -> None:
         """Shorthand for dynamic overwriting the table with a PyArrow table.
 
@@ -1390,7 +1387,7 @@ def overwrite(
         overwrite_filter: Union[BooleanExpression, str] = ALWAYS_TRUE,
         snapshot_properties: Dict[str, str] = EMPTY_DICT,
         case_sensitive: bool = True,
-        branch: Optional[str] = None,
+        branch: Optional[str] = MAIN_BRANCH,
     ) -> None:
         """
         Shorthand for overwriting the table with a PyArrow table.
@@ -1423,7 +1420,7 @@ def delete(
         delete_filter: Union[BooleanExpression, str] = ALWAYS_TRUE,
         snapshot_properties: Dict[str, str] = EMPTY_DICT,
         case_sensitive: bool = True,
-        branch: Optional[str] = None,
+        branch: Optional[str] = MAIN_BRANCH,
     ) -> None:
         """
         Shorthand for deleting rows from the table.

pyiceberg/table/metadata.py

@@ -295,8 +295,10 @@ def new_snapshot_id(self) -> int:
 
         return snapshot_id
 
-    def snapshot_by_name(self, name: str) -> Optional[Snapshot]:
+    def snapshot_by_name(self, name: Optional[str]) -> Optional[Snapshot]:
         """Return the snapshot referenced by the given name or null if no such reference exists."""
+        if name is None:
+            name = MAIN_BRANCH
         if ref := self.refs.get(name):
             return self.snapshot_by_id(ref.snapshot_id)
         return None

pyiceberg/table/update/snapshot.py

@@ -108,8 +108,7 @@ class _SnapshotProducer(UpdateTableMetadata[U], Generic[U]):
     _manifest_num_counter: itertools.count[int]
     _deleted_data_files: Set[DataFile]
     _compression: AvroCompressionCodec
-    _target_branch = MAIN_BRANCH
-    _stage_only = False
+    _target_branch: Optional[str]
 
     def __init__(
         self,
@@ -118,8 +117,7 @@ def __init__(
         io: FileIO,
         commit_uuid: Optional[uuid.UUID] = None,
         snapshot_properties: Dict[str, str] = EMPTY_DICT,
-        branch: str = MAIN_BRANCH,
-        stage_only: bool = False,
+        branch: Optional[str] = MAIN_BRANCH,
     ) -> None:
         super().__init__(transaction)
         self.commit_uuid = commit_uuid or uuid.uuid4()
@@ -139,16 +137,14 @@ def __init__(
         self._parent_snapshot_id = (
             snapshot.snapshot_id if (snapshot := self._transaction.table_metadata.snapshot_by_name(self._target_branch)) else None
         )
-        self._stage_only = stage_only
 
-    def _validate_target_branch(self, branch: str) -> str:
+    def _validate_target_branch(self, branch: Optional[str]) -> Optional[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")
+        if branch is not None:
+            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]:
@@ -297,7 +293,7 @@ def _commit(self) -> UpdatesAndRequirements:
 
         add_snapshot_update = AddSnapshotUpdate(snapshot=snapshot)
 
-        if self._stage_only:
+        if self._target_branch is None:
             return (
                 (add_snapshot_update,),
                 (),
@@ -368,12 +364,11 @@ def __init__(
         operation: Operation,
         transaction: Transaction,
         io: FileIO,
-        branch: str,
+        branch: Optional[str] = MAIN_BRANCH,
         commit_uuid: Optional[uuid.UUID] = None,
         snapshot_properties: Dict[str, str] = EMPTY_DICT,
-        stage_only: bool = False,
     ):
-        super().__init__(operation, transaction, io, commit_uuid, snapshot_properties, branch, stage_only)
+        super().__init__(operation, transaction, io, commit_uuid, snapshot_properties, branch)
         self._predicate = AlwaysFalse()
         self._case_sensitive = True
 
@@ -539,14 +534,13 @@ def __init__(
         operation: Operation,
         transaction: Transaction,
         io: FileIO,
-        branch: str,
+        branch: Optional[str] = MAIN_BRANCH,
         commit_uuid: Optional[uuid.UUID] = None,
         snapshot_properties: Dict[str, str] = EMPTY_DICT,
-        stage_only: bool = False,
     ) -> None:
         from pyiceberg.table import TableProperties
 
-        super().__init__(operation, transaction, io, commit_uuid, snapshot_properties, branch, stage_only)
+        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,
@@ -661,23 +655,20 @@ def _get_entries(manifest: ManifestFile) -> List[ManifestEntry]:
 class UpdateSnapshot:
     _transaction: Transaction
     _io: FileIO
-    _branch: str
-    _stage_only: bool
+    _branch: Optional[str]
     _snapshot_properties: Dict[str, str]
 
     def __init__(
         self,
         transaction: Transaction,
         io: FileIO,
-        branch: str,
-        stage_only: bool = False,
+        branch: Optional[str] = MAIN_BRANCH,
         snapshot_properties: Dict[str, str] = EMPTY_DICT,
     ) -> None:
         self._transaction = transaction
         self._io = io
         self._snapshot_properties = snapshot_properties
         self._branch = branch
-        self._stage_only = stage_only
 
     def fast_append(self) -> _FastAppendFiles:
         return _FastAppendFiles(
@@ -686,7 +677,6 @@ def fast_append(self) -> _FastAppendFiles:
             io=self._io,
             branch=self._branch,
             snapshot_properties=self._snapshot_properties,
-            stage_only=self._stage_only,
         )
 
     def merge_append(self) -> _MergeAppendFiles:
@@ -696,7 +686,6 @@ def merge_append(self) -> _MergeAppendFiles:
             io=self._io,
             branch=self._branch,
             snapshot_properties=self._snapshot_properties,
-            stage_only=self._stage_only,
         )
 
     def overwrite(self, commit_uuid: Optional[uuid.UUID] = None) -> _OverwriteFiles:
@@ -709,7 +698,6 @@ def overwrite(self, commit_uuid: Optional[uuid.UUID] = None) -> _OverwriteFiles:
             io=self._io,
             branch=self._branch,
             snapshot_properties=self._snapshot_properties,
-            stage_only=self._stage_only,
         )
 
     def delete(self) -> _DeleteFiles:
@@ -719,7 +707,6 @@ def delete(self) -> _DeleteFiles:
             io=self._io,
             branch=self._branch,
             snapshot_properties=self._snapshot_properties,
-            stage_only=self._stage_only,
         )
 
 

tests/integration/test_writes/test_writes.py

@@ -2287,7 +2287,7 @@ def test_stage_only_delete(
     assert len(files_to_delete) > 0
 
     with tbl.transaction() as txn:
-        with txn.update_snapshot(stage_only=True).delete() as delete:
+        with txn.update_snapshot(branch=None).delete() as delete:
             delete.delete_by_predicate(EqualTo("int", 9))
 
     # a new delete snapshot is added
@@ -2324,7 +2324,7 @@ def test_stage_only_fast_append(
     assert original_count == 3
 
     with tbl.transaction() as txn:
-        with txn.update_snapshot(stage_only=True).fast_append() as fast_append:
+        with txn.update_snapshot(branch=None).fast_append() as fast_append:
             for data_file in _dataframe_to_data_files(
                 table_metadata=txn.table_metadata, df=arrow_table_with_null, io=txn._table.io
             ):
@@ -2364,7 +2364,7 @@ def test_stage_only_merge_append(
     assert original_count == 3
 
     with tbl.transaction() as txn:
-        with txn.update_snapshot(stage_only=True).merge_append() as merge_append:
+        with txn.update_snapshot(branch=None).merge_append() as merge_append:
             for data_file in _dataframe_to_data_files(
                 table_metadata=txn.table_metadata, df=arrow_table_with_null, io=txn._table.io
             ):
@@ -2409,7 +2409,7 @@ def test_stage_only_overwrite_files(
     assert len(files_to_delete) > 0
 
     with tbl.transaction() as txn:
-        with txn.update_snapshot(stage_only=True).overwrite() as overwrite:
+        with txn.update_snapshot(branch=None).overwrite() as overwrite:
             for data_file in _dataframe_to_data_files(
                 table_metadata=txn.table_metadata, df=arrow_table_with_null, io=txn._table.io
             ):