add mode

pola-rs · Jul 2, 2024 · c14870a · c14870a
1 parent 0801012
commit c14870a
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 4 deletions.
diff --git a/py-polars/polars/dataframe/frame.py b/py-polars/polars/dataframe/frame.py
@@ -3835,14 +3835,23 @@ def unpack_table_name(name: str) -> tuple[str | None, str | None, str]:
             msg = f"unrecognised connection type {connection!r}"
             raise TypeError(msg)
 
-    def write_iceberg(self, target: str | Path) -> pyiceberg.table.Table:
+    def write_iceberg(
+        self,
+        target: str | Path,
+        mode: Literal["append", "overwrite"],
+    ) -> pyiceberg.table.Table:
         """
         Write DataFrame to an Iceberg table.
 
         Parameters
         ----------
         target : str | Path
             The target path or identifier for the Iceberg table.
+        mode : {'append', 'overwrite'}
+            How to handle existing data.
+
+            - If 'append', will add new data.
+            - If 'overwrite', will replace table with new data.
 
         Returns
         -------
@@ -3861,8 +3870,10 @@ def write_iceberg(self, target: str | Path) -> pyiceberg.table.Table:
             "default.table",
             schema=schema,
         )
-
-        table.overwrite(data)
+        if mode == "append":
+            table.append(data)
+        else:
+            table.overwrite(data)
         return table
 
     @overload

diff --git a/py-polars/tests/unit/io/test_iceberg.py b/py-polars/tests/unit/io/test_iceberg.py
@@ -178,7 +178,7 @@ def test_write_iceberg(tmp_path: Path) -> None:
             "ham": ["a", "b", "c", "d", "e"],
         }
     )
-    iceberg_table = df.write_iceberg(tmp_path)
+    iceberg_table = df.write_iceberg(tmp_path, mode="overwrite")
     iceberg_path = iceberg_table.metadata_location
     new_df = pl.scan_iceberg(iceberg_path).collect()
     assert len(df) == len(new_df)