[Data] - Iceberg support upsert tables + schema update + overwrite tables (#58270)

## Description
- Support upserting iceberg tables for IcebergDatasink
- Update schema on APPEND and UPSERT
- Enable overwriting the entire table

Upgrades to pyicberg 0.10.0 because it now supports upsert and overwrite
functionality. Also for append, the library now handles the transaction
logic implicitly so that burden can be lifted from Ray Data.

## Related issues
> Link related issues: "Fixes #1234", "Closes #1234", or "Related to
#1234".

## Additional information
> Optional: Add implementation details, API changes, usage examples,
screenshots, etc.

---------

Signed-off-by: Goutam <[email protected]>

Author: goutamvenkat-anyscale

ci/lint/pydoclint-baseline.txt

@@ -1018,10 +1018,6 @@ python/ray/data/_internal/block_batching/util.py
     DOC402: Function `finalize_batches` has "yield" statements, but the docstring does not have a "Yields" section
     DOC404: Function `finalize_batches` yield type(s) in docstring not consistent with the return annotation. Return annotation exists, but docstring "yields" section does not exist or has 0 type(s).
 --------------------
-python/ray/data/_internal/datasource/iceberg_datasink.py
-    DOC102: Method `IcebergDatasink.__init__`: Docstring contains more arguments than in function signature.
-    DOC103: Method `IcebergDatasink.__init__`: Docstring arguments are different from function arguments. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Arguments in the docstring but not in the function signature: [to an iceberg table, e.g. {"commit_time": ].
---------------------
 python/ray/data/_internal/datasource/lance_datasink.py
     DOC101: Method `LanceDatasink.__init__`: Docstring contains fewer arguments than in function signature.
     DOC103: Method `LanceDatasink.__init__`: Docstring arguments are different from function arguments. (Or could be other formatting issues: https://jsh9.github.io/pydoclint/violation_codes.html#notes-on-doc103 ). Arguments in the function signature but not in the docstring: [**kwargs: , *args: , max_rows_per_file: int, min_rows_per_file: int, mode: Literal['create', 'append', 'overwrite'], schema: Optional[pa.Schema], storage_options: Optional[Dict[str, Any]], uri: str]. Arguments in the docstring but not in the function signature: [max_rows_per_file : , min_rows_per_file : , mode : , schema : , storage_options : , uri : ].

python/ray/data/__init__.py

@@ -22,6 +22,7 @@
     FileShuffleConfig,
     ReadTask,
     RowBasedFileDatasink,
+    SaveMode,
 )
 from ray.data.iterator import DataIterator, DatasetIterator
 from ray.data.preprocessor import Preprocessor
@@ -131,6 +132,7 @@
     "NodeIdStr",
     "ReadTask",
     "RowBasedFileDatasink",
+    "SaveMode",
     "Schema",
     "SinkMode",
     "TaskPoolStrategy",

python/ray/data/_internal/datasource/iceberg_datasink.py

@@ -2,26 +2,27 @@
 Module to write a Ray Dataset into an iceberg table, by using the Ray Datasink API.
 """
 import logging
-import uuid
 from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional
 
-from packaging import version
-
 from ray.data._internal.execution.interfaces import TaskContext
+from ray.data._internal.savemode import SaveMode
 from ray.data.block import Block, BlockAccessor
 from ray.data.datasource.datasink import Datasink, WriteResult
 from ray.util.annotations import DeveloperAPI
 
 if TYPE_CHECKING:
+    import pyarrow as pa
     from pyiceberg.catalog import Catalog
-    from pyiceberg.manifest import DataFile
+    from pyiceberg.table import Table
+
+    from ray.data.expressions import Expr
 
 
 logger = logging.getLogger(__name__)
 
 
 @DeveloperAPI
-class IcebergDatasink(Datasink[List["DataFile"]]):
+class IcebergDatasink(Datasink[List["pa.Table"]]):
     """
     Iceberg datasink to write a Ray Dataset into an existing Iceberg table. This module
     heavily uses PyIceberg to write to iceberg table. All the routines in this class override
@@ -34,137 +35,223 @@ def __init__(
         table_identifier: str,
         catalog_kwargs: Optional[Dict[str, Any]] = None,
         snapshot_properties: Optional[Dict[str, str]] = None,
+        mode: SaveMode = SaveMode.APPEND,
+        overwrite_filter: Optional["Expr"] = None,
+        upsert_kwargs: Optional[Dict[str, Any]] = None,
+        overwrite_kwargs: Optional[Dict[str, Any]] = None,
     ):
         """
         Initialize the IcebergDatasink
 
         Args:
-            table_identifier: The identifier of the table to read e.g. `default.taxi_dataset`
+            table_identifier: The identifier of the table to read such as `default.taxi_dataset`
             catalog_kwargs: Optional arguments to use when setting up the Iceberg
                 catalog
-            snapshot_properties: custom properties write to snapshot when committing
-            to an iceberg table, e.g. {"commit_time": "2021-01-01T00:00:00Z"}
+            snapshot_properties: Custom properties to write to snapshot summary, such as commit metadata
+            mode: Write mode - APPEND, UPSERT, or OVERWRITE. Defaults to APPEND.
+                - APPEND: Add new data without checking for duplicates
+                - UPSERT: Update existing rows or insert new ones based on a join condition
+                - OVERWRITE: Replace table data (all data or filtered subset)
+            overwrite_filter: Optional filter for OVERWRITE mode to perform partial overwrites.
+                Must be a Ray Data expression from `ray.data.expressions`. Only rows matching
+                this filter are replaced. If None with OVERWRITE mode, replaces all table data.
+            upsert_kwargs: Optional arguments to pass through to PyIceberg's table.upsert()
+                method. Supported parameters include join_cols (List[str]),
+                when_matched_update_all (bool), when_not_matched_insert_all (bool),
+                case_sensitive (bool), branch (str). See PyIceberg documentation for details.
+            overwrite_kwargs: Optional arguments to pass through to PyIceberg's table.overwrite()
+                method. Supported parameters include case_sensitive (bool) and branch (str).
+                See PyIceberg documentation for details.
+
+        Note:
+            Schema evolution is automatically enabled. New columns in the incoming data
+            are automatically added to the table schema.
         """
 
-        from pyiceberg.io import FileIO
-        from pyiceberg.table import Transaction
-        from pyiceberg.table.metadata import TableMetadata
-
         self.table_identifier = table_identifier
-        self._catalog_kwargs = catalog_kwargs if catalog_kwargs is not None else {}
-        self._snapshot_properties = (
-            snapshot_properties if snapshot_properties is not None else {}
-        )
+        self._catalog_kwargs = (catalog_kwargs or {}).copy()
+        self._snapshot_properties = snapshot_properties or {}
+        self._mode = mode
+        self._overwrite_filter = overwrite_filter
+        self._upsert_kwargs = (upsert_kwargs or {}).copy()
+        self._overwrite_kwargs = (overwrite_kwargs or {}).copy()
+
+        # Validate kwargs are only set for relevant modes
+        if self._upsert_kwargs and self._mode != SaveMode.UPSERT:
+            raise ValueError(
+                f"upsert_kwargs can only be specified when mode is SaveMode.UPSERT, "
+                f"but mode is {self._mode}"
+            )
+        if self._overwrite_kwargs and self._mode != SaveMode.OVERWRITE:
+            raise ValueError(
+                f"overwrite_kwargs can only be specified when mode is SaveMode.OVERWRITE, "
+                f"but mode is {self._mode}"
+            )
 
         if "name" in self._catalog_kwargs:
             self._catalog_name = self._catalog_kwargs.pop("name")
         else:
             self._catalog_name = "default"
 
-        self._uuid: str = None
-        self._io: FileIO = None
-        self._txn: Transaction = None
-        self._table_metadata: TableMetadata = None
+        self._table: "Table" = None
 
-    # Since iceberg transaction is not pickle-able, because of the table and catalog properties
-    # we need to exclude the transaction object during serialization and deserialization during pickle
+    # Since iceberg table is not pickle-able, we need to exclude it during serialization
     def __getstate__(self) -> dict:
-        """Exclude `_txn` during pickling."""
+        """Exclude `_table` during pickling."""
         state = self.__dict__.copy()
-        del state["_txn"]
+        state.pop("_table", None)
         return state
 
     def __setstate__(self, state: dict) -> None:
         self.__dict__.update(state)
-        self._txn = None
+        self._table = None
 
     def _get_catalog(self) -> "Catalog":
         from pyiceberg import catalog
 
         return catalog.load_catalog(self._catalog_name, **self._catalog_kwargs)
 
-    def on_write_start(self) -> None:
-        """Prepare for the transaction"""
-        import pyiceberg
-        from pyiceberg.table import TableProperties
+    def _update_schema(self, incoming_schema: "pa.Schema") -> None:
+        """
+        Update the table schema to accommodate incoming data using union-by-name semantics.
 
-        if version.parse(pyiceberg.__version__) >= version.parse("0.9.0"):
-            from pyiceberg.utils.properties import property_as_bool
-        else:
-            from pyiceberg.table import PropertyUtil
+        This automatically handles:
+        - Adding new columns from the incoming schema
+        - Type promotion (e.g., int32 -> int64) where compatible
+        - Preserving existing columns not in the incoming schema
 
-            property_as_bool = PropertyUtil.property_as_bool
+        Args:
+            incoming_schema: The PyArrow schema from the incoming data
+        """
+        # Use PyIceberg's update_schema API
+        with self._table.update_schema() as update:
+            update.union_by_name(incoming_schema)
 
+        # Reload table completely after schema evolution
         catalog = self._get_catalog()
-        table = catalog.load_table(self.table_identifier)
-        self._txn = table.transaction()
-        self._io = self._txn._table.io
-        self._table_metadata = self._txn.table_metadata
-        self._uuid = uuid.uuid4()
-
-        if unsupported_partitions := [
-            field
-            for field in self._table_metadata.spec().fields
-            if not field.transform.supports_pyarrow_transform
-        ]:
-            raise ValueError(
-                f"Not all partition types are supported for writes. Following partitions cannot be written using pyarrow: {unsupported_partitions}."
-            )
+        self._table = catalog.load_table(self.table_identifier)
 
-        self._manifest_merge_enabled = property_as_bool(
-            self._table_metadata.properties,
-            TableProperties.MANIFEST_MERGE_ENABLED,
-            TableProperties.MANIFEST_MERGE_ENABLED_DEFAULT,
-        )
+    def on_write_start(self) -> None:
+        """Initialize table for writing."""
+        catalog = self._get_catalog()
+        self._table = catalog.load_table(self.table_identifier)
 
-    def write(
-        self, blocks: Iterable[Block], ctx: TaskContext
-    ) -> WriteResult[List["DataFile"]]:
-        from pyiceberg.io.pyarrow import (
-            _check_pyarrow_schema_compatible,
-            _dataframe_to_data_files,
-        )
-        from pyiceberg.table import DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE
-        from pyiceberg.utils.config import Config
+    def _collect_tables_from_blocks(self, blocks: Iterable[Block]) -> List["pa.Table"]:
+        """Collect PyArrow tables from blocks."""
+        collected_tables = []
 
-        data_files_list: WriteResult[List["DataFile"]] = []
         for block in blocks:
             pa_table = BlockAccessor.for_block(block).to_arrow()
 
-            downcast_ns_timestamp_to_us = (
-                Config().get_bool(DOWNCAST_NS_TIMESTAMP_TO_US_ON_WRITE) or False
-            )
-            _check_pyarrow_schema_compatible(
-                self._table_metadata.schema(),
-                provided_schema=pa_table.schema,
-                downcast_ns_timestamp_to_us=downcast_ns_timestamp_to_us,
-            )
+            if pa_table.num_rows > 0:
+                collected_tables.append(pa_table)
 
-            if pa_table.shape[0] <= 0:
-                continue
+        return collected_tables
 
-            task_uuid = uuid.uuid4()
-            data_files = _dataframe_to_data_files(
-                self._table_metadata, pa_table, self._io, task_uuid
-            )
-            data_files_list.extend(data_files)
+    def write(self, blocks: Iterable[Block], ctx: TaskContext) -> List["pa.Table"]:
+        """Collect blocks as PyArrow tables for all write modes."""
+        return self._collect_tables_from_blocks(blocks)
+
+    def _collect_and_concat_tables(
+        self, write_result: WriteResult[List["pa.Table"]]
+    ) -> Optional["pa.Table"]:
+        """Collect and concatenate all PyArrow tables from write results."""
+        import pyarrow as pa
+
+        all_tables = []
+        for tables_batch in write_result.write_returns:
+            all_tables.extend(tables_batch)
 
-        return data_files_list
+        if not all_tables:
+            logger.warning("No data to write")
+            return None
 
-    def on_write_complete(self, write_result: WriteResult[List["DataFile"]]):
-        update_snapshot = self._txn.update_snapshot(
-            snapshot_properties=self._snapshot_properties
+        return pa.concat_tables(all_tables)
+
+    def _complete_append(self, combined_table: "pa.Table") -> None:
+        """Complete APPEND mode write using PyIceberg's append API."""
+        self._table.append(
+            df=combined_table,
+            snapshot_properties=self._snapshot_properties,
         )
-        append_method = (
-            update_snapshot.merge_append
-            if self._manifest_merge_enabled
-            else update_snapshot.fast_append
+        logger.info(
+            f"Appended {combined_table.num_rows} rows to {self.table_identifier}"
         )
 
-        with append_method() as append_files:
-            append_files.commit_uuid = self._uuid
-            for data_files in write_result.write_returns:
-                for data_file in data_files:
-                    append_files.append_data_file(data_file)
+    def _complete_upsert(self, combined_table: "pa.Table") -> None:
+        """Complete UPSERT mode write using PyIceberg's upsert API."""
+        self._table.upsert(df=combined_table, **self._upsert_kwargs)
 
-        self._txn.commit_transaction()
+        join_cols = self._upsert_kwargs.get("join_cols")
+        if join_cols:
+            logger.info(
+                f"Upserted {combined_table.num_rows} rows to {self.table_identifier} "
+                f"using join columns: {join_cols}"
+            )
+        else:
+            logger.info(
+                f"Upserted {combined_table.num_rows} rows to {self.table_identifier} "
+                f"using table-defined identifier-field-ids"
+            )
+
+    def _complete_overwrite(self, combined_table: "pa.Table") -> None:
+        """Complete OVERWRITE mode write using PyIceberg's overwrite API."""
+        # Warn if user passed overwrite_filter via overwrite_kwargs
+        if "overwrite_filter" in self._overwrite_kwargs:
+            self._overwrite_kwargs.pop("overwrite_filter")
+            logger.warning(
+                "Use Ray Data's Expressions for overwrite filter instead of passing "
+                "it via PyIceberg's overwrite_filter parameter"
+            )
+
+        if self._overwrite_filter:
+            # Partial overwrite with filter
+            from ray.data._internal.datasource.iceberg_datasource import (
+                _IcebergExpressionVisitor,
+            )
+
+            iceberg_filter = _IcebergExpressionVisitor().visit(self._overwrite_filter)
+            self._table.overwrite(
+                df=combined_table,
+                overwrite_filter=iceberg_filter,
+                snapshot_properties=self._snapshot_properties,
+                **self._overwrite_kwargs,
+            )
+            logger.info(
+                f"Overwrote {combined_table.num_rows} rows in {self.table_identifier} "
+                f"matching filter: {self._overwrite_filter}"
+            )
+        else:
+            # Full table overwrite
+            self._table.overwrite(
+                df=combined_table,
+                snapshot_properties=self._snapshot_properties,
+                **self._overwrite_kwargs,
+            )
+            logger.info(
+                f"Overwrote entire table {self.table_identifier} "
+                f"with {combined_table.num_rows} rows"
+            )
+
+    def on_write_complete(self, write_result: WriteResult[List["pa.Table"]]) -> None:
+        """Complete the write operation based on the configured mode."""
+        # Collect and concatenate all PyArrow tables
+        combined_table = self._collect_and_concat_tables(write_result)
+        if combined_table is None:
+            return
+
+        # Apply schema evolution for all modes (PyIceberg doesn't handle this automatically)
+        self._update_schema(combined_table.schema)
+
+        # Execute the appropriate write operation
+        if self._mode == SaveMode.APPEND:
+            self._complete_append(combined_table)
+        elif self._mode == SaveMode.UPSERT:
+            self._complete_upsert(combined_table)
+        elif self._mode == SaveMode.OVERWRITE:
+            self._complete_overwrite(combined_table)
+        else:
+            raise ValueError(
+                f"Unsupported write mode: {self._mode}. "
+                f"Supported modes are: APPEND, UPSERT, OVERWRITE"
+            )

python/ray/data/_internal/savemode.py

@@ -5,7 +5,21 @@
 
 @PublicAPI(stability="alpha")
 class SaveMode(str, Enum):
+    """Enum of possible modes for saving/writing data."""
+
     APPEND = "append"
+    """Add new data without modifying existing data."""
+
     OVERWRITE = "overwrite"
+    """Replace all existing data with new data."""
+
     IGNORE = "ignore"
+    """Don't write if data already exists."""
+
     ERROR = "error"
+    """Raise an error if data already exists."""
+
+    UPSERT = "upsert"
+    """Update existing rows that match on key fields, or insert new rows.
+    Requires identifier/key fields to be specified.
+    """