Repository: ray-project/ray
Pull request: [data] add iceberg write support through pyiceberg
<!-- Thank you for your contribution! Please review https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. -->

## Why are these changes needed?

Add iceberg write support through pyiceberg. Ray will distributedly write the dataset using the Datasink interface through blocks to an iceberg table leveraging pyarrow as parquet files, and commit append snapshot to the iceberg table. There are couple of limitations that's worth pointing out
- Supports only appending to an existing Iceberg table.
- Only supports the Parquet data file format.
- Ray will validate schema compatibility between its blocks and the Iceberg table and will evolve the Iceberg table schema if needed.
- A single block is bin-packed into the target output size; however, the sink does not bin-pack across different blocks.
  - This is a limitation of [the PyIceberg API](https://github.com/apache/iceberg-python/blob/main/pyiceberg/io/pyarrow.py#L2535C5-L2535C29), and cross-block bin packing can be addressed in the future by extending the API on the Pyiceberg side

<!-- Please give a short summary of the change and the problem this solves. -->

## Related issue number

<!-- For example: "Closes #1234" -->
Closes #49032

## Checks

- [x] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR.
- [x] I've run `scripts/format.sh` to lint the changes in this PR.
- [x] I've included any doc changes needed for https://docs.ray.io/en/master/.
    - [x] I've added any new APIs to the API Reference. For example, if I added a
           method in Tune, I've added it in `doc/source/data/api/` under the
           corresponding `.rst` file.
- [x] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [x] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(


Problem statement:
[data] Extend Ray Data with IcebergDataSink for Distributed Writes to Iceberg Tables
### Description

Currently, Ray supports distributed reads from Iceberg tables but lacks support for distributed writes. To enhance the functionality of Ray in handling Iceberg tables, it is proposed to extend the datasink module within Ray Data by introducing an IcebergDataSink. This would allow users to perform distributed writes on Iceberg tables, significantly improving the utility and flexibility of Ray in big data processing workflows.

The current limitation of Ray's Iceberg table integration only supporting read operations restricts its applicability in scenarios where real-time or batch updates to Iceberg tables are required. Additionally, PyIceberg's existing implementation tightly couples file generation and transaction commitment, which is not suitable for a distributed write operation managed by Ray.

Proposed Solution
To achieve distributed writes to Iceberg tables, we propose the following steps:

1. Modify PyIceberg Transaction Handling:
- Decouple the transaction management from the append operation.
- Introduce a method generate_data_files that prepares Parquet files for writing without committing them immediately.
- Provide a mechanism to commit transactions manually after all nodes have completed their writes.

2. Develop IcebergDataSink in Ray:
- Implement the IcebergDataSink class within Ray Data to handle distributed writes.
- Utilize Ray's distributed computing capabilities to write Parquet files across multiple nodes.
- Coordinate the commit process to maintain data consistency and integrity.

3. Handle Failed Writes:
- Implement a cleanup mechanism to remove any partially written or failed data files (data_file) to prevent zombie data from cluttering the storage system.
- Ensure this cleanup process is idempotent and can be safely retried if necessary.

### Use case

_No response_

Hints from issue discussion:
This proposal makes sense; would you be interested in contributing?
A couple of questions about the proposal:
- The current proposal suggests that we will only support append operations, not delete operations (whether quality deletes or positional deletes). Or does pyiceberg support merge-on-read, and we can take care of compaction asynchronously 
- We probably shouldn’t clean up failed writes. Generally, data from partial writes to Iceberg (no sucessful committed) is treated as orphan files (files with reference count of zero). These should be handled by Iceberg table maintenance rather than through explicitly orchestrated cleanup.

> This proposal makes sense; would you be interested in contributing?

Sorry for the late reply. I can contribute to the implementation of this part, but Pyiceberg needs to make some extensions. This has already been implemented internally, and I can submit the relevant code next week
> A couple of questions about the proposal:
> 
> * The current proposal suggests that we will only support append operations, not delete operations (whether quality deletes or positional deletes). Or does pyiceberg support merge-on-read, and we can take care of compaction asynchronously
> * We probably shouldn’t clean up failed writes. Generally, data from partial writes to Iceberg (no sucessful committed) is treated as orphan files (files with reference count of zero). These should be handled by Iceberg table maintenance rather than through explicitly orchestrated cleanup.

Our internal scenario is mainly focused on adding, and the Ray layer mainly implements distributed writing and commit metadata submission. Cleaning is indeed more suitable for Iceberg to maintain
I had an conversation with the PyIceberg maintainer. Solution one is a bit more complex and involved to implement in practice and they mentioned it goes against the current philosophy of what PyIceberg is which is a set of API to manage iceberg table instead of a processing engine.

They recommend using Ray’s `write_parquet` API and then calling `add_files` to commit the files. Under the hood, this will perform checks for duplicates, missing files, and schema compatibility.

Additionally, there is no need for a clean-up process on failure because only in-use data files are referenced in the manifest files. Table maintenance will clean up any orphaned files from expired snapshots and commit failures. This is also how failures are handled in the JVM Iceberg path.
After some investigation, it seems that an implementation like [this](https://github.com/ray-project/ray/pull/49961) would not be a clean solution, as it bypasses significant portions of Iceberg’s internal logic, such as handling partitions, data file bin-packing, sorting, and other critical functions.

Since both PyIceberg and Ray Data rely heavily on PyArrow under the hood, we should aim to converge around that. PyIceberg provides [an append API](https://github.com/apache/iceberg-python/blob/f0346472e4301f2ea3679e0793bb8623f2bb80f1/pyiceberg/table/__init__.py#L424) that allows appending a single `pa.Table` to an Iceberg table within a transaction. This approach is advantageous because it manages many of Iceberg’s internal checks and handles the write operation for us. 

I would like to propose the following implementation for the `class IcebergDatasink(Datasink[DataFile])`:
- In `on_write_start`, generate a `UUID.uuid4()` within the Datasink instance for the data write.
- Extend the `_dataframe_to_data_files` [function](https://github.com/apache/iceberg-python/blob/f0346472e4301f2ea3679e0793bb8623f2bb80f1/pyiceberg/io/pyarrow.py#L2482) on the pyiceberg side to accept an iterator of `pa.Table`.
  - It might be better to create a new function based on the existing one.
- In the `write` method, call the updated `_dataframes_to_data_files` function with the pre-generated UUID and return a list of Iceberg `DataFile` paths after written.
  - Extend the Datasink with `DataFile` as its generic type to enable the write function to return additional write information.
- In `on_write_complete`, open a transaction with a new FastAppend snapshot producer providing it with the `UUID`, and commit the transaction using the returned DataFile paths, similar to [this implementation](https://github.com/apache/iceberg-python/blob/main/pyiceberg/table/__init__.py#L453-L460).
Here is the pseudocode implementation:
```
class IcebergDatasink(Generic[List[DataFile]]):

    def __init__(
        self,
        table_identifier: str,
        catalog_kwargs: Optional[Dict[str, Any]] = None,
        snapshot_properties: Optional[Dict[str, str]] = None,
    ):
        self._uuid: str|None = None
        self._snapshot_properties = snapshot_properties
        self.table_identifier = table_identifier
        self._table = None
        self._datasink = None
  
    def on_write_start(self) -> None:
        """Prepare for the transaction"""
        self._uuid = uuid.uuid4()

    def _get_catalog(self) -> "Catalog":
        from pyiceberg import catalog
        return catalog.load_catalog(self._catalog_name, **self._catalog_kwargs)

    @property
    def table(self) -> "Table":
        if self._table is None:
            catalog = self._get_catalog()
            self._table = catalog.load_table(self.table_identifier)
        return self._table

    def write(self, blocks: Iterable[Block], ctx: TaskContext) -> List[DataFile]:
        ...
        tables = [BlockAccessor.for_block(block).to_arrow() for block in blocks]
        data_files = _dataframes_to_data_files(tables, ctx)
        return DataFiles

    def on_write_complete(self, write_result: WriteResult[List[DataFile]]):
         ...
         fast_append = _FastAppendFiles(
            operation=Operation.APPEND, transaction=self._table._transaction, io=self._table.io, snapshot_properties=self._snapshot_properties
        )
      
         with fast_append(uuid=self._uuid) as append_files:
            for write_return in write_result.write_returns:
                for data_file in write_return:
                    append_files.append_data_file(data_file)
```
@kevinjqliu when you get a chance can you take a quick look and let us know what you think about the proposed implementation?
@jimmyxie-figma I'm not familiar with PyIceberg. Just to confirm, the actual write is performed inside `_dataframes_to_data_files`, and DataFiles only includes the metadata, right? 
If so, the approach looks good to me. 
Hey folks! I'm one of the contributors to pyiceberg. Great to see interests in adding iceberg write capabilities to Ray! I'm not (too) familiar with Ray and its internals but I can comment a bit on pyiceberg and what I would recommend as an integration point. 

Generally speaking, you can think of Iceberg tables as a collection of parquet files with some metadata files and a pointer in a catalog. In a distributed environment, such as Spark, writes are handled by the different workers. These workers track statistics about the files and are table/partition aware. When the files are all written to object store, the Spark driver then creates the metadata portion and commits to the catalog. (This is my general understanding of the process in Spark.) The key part is the distributed writes, "collecting/merging" the metadata into the driver, and the single driver committing to the catalog. 
I would recommend taking the similar approach with distributed writing on worker nodes and then finally committing the table through a single "driver" node. 

In terms of pyiceberg's API, all the components are available to commit the data files to the table using a catalog. 
I would recommend looking at the `add_files` API, where data files can be appended to a table. The [implementation](https://github.com/apache/iceberg-python/blob/2cd4e789d0fd7d6aab261037e8229ca94c218923/pyiceberg/table/__init__.py#L697-L702) creates a new snapshot, adds the data files to it, and commits to the iceberg table. 
There's also an `overwrite` API to write pyarrow dataframes to iceberg. The dataframe is partitioned and written to object storage using `_dataframe_to_data_files`. And then similarly, a new snapshot is created. 
For [deletes, pyiceberg currently uses "copy-on-write" ](https://github.com/apache/iceberg-python/blob/2cd4e789d0fd7d6aab261037e8229ca94c218923/pyiceberg/table/__init__.py#L598-L665)where data files are completely rewritten on commit. There's also "merge-on-read" delete, this requires generating delete files and isn't implemented yet

@jimmyxie-figma the proposed implementation looks good to me. I assume `on_write_complete` is called once in the "driver" node. `_dataframes_to_data_files` should take care of partitioning. I'm not sure if `uuid` is necessary since one will be generated if not provided. I'd also like to mention that "FastAppend" is one of the ways to generate metadata, there's also [`MergeAppend`](https://github.com/apache/iceberg-python/blob/2cd4e789d0fd7d6aab261037e8229ca94c218923/pyiceberg/table/update/snapshot.py#L618-L626). You can read more about FastAppend from the [iceberg spec](https://iceberg.apache.org/spec/#snapshots)


Happy to provide more guidance on the pyiceberg side! Feel free to ping me or raise an issue at https://github.com/apache/iceberg-python

Commit message: fix write iceberg table

Signed-off-by: Jimmy Xie <rxie@figma.com>

Patch to review:
diff --git a/doc/source/data/api/input_output.rst b/doc/source/data/api/input_output.rst
index 6487f06a968b..781d4e038471 100644
--- a/doc/source/data/api/input_output.rst
+++ b/doc/source/data/api/input_output.rst
@@ -212,6 +212,7 @@ Iceberg
    :toctree: doc/
 
    read_iceberg
+   Dataset.write_iceberg
 
 Lance
 -----
diff --git a/python/ray/data/_internal/datasource/iceberg_datasink.py b/python/ray/data/_internal/datasource/iceberg_datasink.py
new file mode 100644
index 000000000000..502291ef5ae3
--- /dev/null
+++ b/python/ray/data/_internal/datasource/iceberg_datasink.py
@@ -0,0 +1,153 @@
+"""
+Module to write a Ray Dataset into an iceberg table, by using the Ray Datasink API.
+"""
+import logging
+
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional
+
+from ray.data.datasource.datasink import Datasink
+from ray.util.annotations import DeveloperAPI
+from ray.data.block import BlockAccessor, Block
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data.datasource.datasink import WriteResult
+import uuid
+
+if TYPE_CHECKING:
+    from pyiceberg.catalog import Catalog
+    from pyiceberg.manifest import DataFile
+
+
+logger = logging.getLogger(__name__)
+
+
+@DeveloperAPI
+class IcebergDatasink(Datasink[List["DataFile"]]):
+    """
+    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
+    `ray.data.Datasink`.
+
+    """
+
+    def __init__(
+        self,
+        table_identifier: str,
+        catalog_kwargs: Optional[Dict[str, Any]] = None,
+        snapshot_properties: Optional[Dict[str, str]] = None,
+    ):
+        """
+        Initialize the IcebergDatasink
+
+        Args:
+            table_identifier: The identifier of the table to read
+            catalog_kwargs: Optional arguments to use when setting up the Iceberg
+                catalog
+            snapshot_properties: Optional snapshot properties to
+                use when reading the table
+        """
+
+        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 {}
+        )
+
+        if "name" in self._catalog_kwargs:
+            self._catalog_name = self._catalog_kwargs.pop("name")
+        else:
+            self._catalog_name = "default"
+
+        self._uuid = None
+        self._io = None
+        self._txn = None
+        self._table_metadata = None
+
+    def __getstate__(self) -> dict:
+        """Exclude `_txn` during pickling."""
+        state = self.__dict__.copy()
+        del state["_txn"]
+        return state
+
+    def __setstate__(self, state: dict) -> None:
+        self.__dict__.update(state)
+        self._txn = 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"""
+        from pyiceberg.table import PropertyUtil, TableProperties
+
+        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._manifest_merge_enabled = PropertyUtil.property_as_bool(
+            self._table_metadata.properties,
+            TableProperties.MANIFEST_MERGE_ENABLED,
+            TableProperties.MANIFEST_MERGE_ENABLED_DEFAULT,
+        )
+
+    def write(self, blocks: Iterable[Block], ctx: TaskContext):
+        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
+
+        data_files_list = []
+        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.shape[0] <= 0:
+                continue
+
+            data_files = _dataframe_to_data_files(
+                self._table_metadata, pa_table, self._io, self._uuid
+            )
+            data_files_list.extend(data_files)
+
+        return data_files_list
+
+    def on_write_complete(self, write_result: WriteResult[List["DataFile"]]):
+        update_snapshot = self._txn.update_snapshot(
+            snapshot_properties=self._snapshot_properties
+        )
+        append_method = (
+            update_snapshot.merge_append
+            if self._manifest_merge_enabled
+            else update_snapshot.fast_append
+        )
+
+        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)
+
+        self._txn.commit_transaction()
diff --git a/python/ray/data/dataset.py b/python/ray/data/dataset.py
index 78cee5179da7..6b3de24b5f41 100644
--- a/python/ray/data/dataset.py
+++ b/python/ray/data/dataset.py
@@ -36,6 +36,7 @@
 from ray.data._internal.compute import ComputeStrategy
 from ray.data._internal.datasource.bigquery_datasink import BigQueryDatasink
 from ray.data._internal.datasource.csv_datasink import CSVDatasink
+from ray.data._internal.datasource.iceberg_datasink import IcebergDatasink
 from ray.data._internal.datasource.image_datasink import ImageDatasink
 from ray.data._internal.datasource.json_datasink import JSONDatasink
 from ray.data._internal.datasource.lance_datasink import LanceDatasink
@@ -3300,6 +3301,61 @@ def write_json(
             concurrency=concurrency,
         )
 
+    @ConsumptionAPI
+    @PublicAPI(stability="alpha", api_group=IOC_API_GROUP)
+    def write_iceberg(
+        self,
+        table_identifier: str,
+        catalog_kwargs: Optional[Dict[str, Any]] = None,
+        snapshot_properties: Optional[Dict[str, str]] = None,
+        ray_remote_args: Dict[str, Any] = None,
+        concurrency: Optional[int] = None,
+    ) -> None:
+        """Writes the :class:`~ray.data.Dataset` to an Iceberg table.
+
+        .. tip::
+            For more details on PyIceberg, see
+            - URI: https://py.iceberg.apache.org/
+
+        Examples:
+             .. testcode::
+                :skipif: True
+
+                import ray
+                import pandas as pd
+                docs = [{"title": "Iceberg data sink test"} for key in range(4)]
+                ds = ray.data.from_pandas(pd.DataFrame(docs))
+                ds.write_iceberg(
+                    table_identifier="db_name.table_name",
+                    catalog_kwargs={"name": "default", "type": "sql"}
+                )
+
+        Args:
+            table_identifier: Fully qualified table identifier (``db_name.table_name``)
+            catalog_kwargs: Optional arguments to pass to PyIceberg's catalog.load_catalog()
+                function (e.g., name, type, etc.). For the function definition, see
+                `pyiceberg catalog
+                <https://py.iceberg.apache.org/reference/pyiceberg/catalog/\
+                #pyiceberg.catalog.load_catalog>`_.
+            snapshot_properties: custom properties write to snapshot when committing
+            to an iceberg table.
+            ray_remote_args: kwargs passed to :func:`ray.remote` in the write tasks.
+            concurrency: The maximum number of Ray tasks to run concurrently. Set this
+                to control number of tasks to run concurrently. This doesn't change the
+                total number of tasks run. By default, concurrency is dynamically
+                decided based on the available resources.
+        """
+
+        datasink = IcebergDatasink(
+            table_identifier, catalog_kwargs, snapshot_properties
+        )
+
+        self.write_datasink(
+            datasink,
+            ray_remote_args=ray_remote_args,
+            concurrency=concurrency,
+        )
+
     @PublicAPI(stability="alpha", api_group=IOC_API_GROUP)
     @ConsumptionAPI
     def write_images(
diff --git a/python/ray/data/tests/test_iceberg.py b/python/ray/data/tests/test_iceberg.py
index c4988cfad14b..3cf43aad129d 100644
--- a/python/ray/data/tests/test_iceberg.py
+++ b/python/ray/data/tests/test_iceberg.py
@@ -28,8 +28,27 @@
     "warehouse": f"file://{_WAREHOUSE_PATH}",
 }
 
+_SCHEMA = pa.schema(
+    [
+        pa.field("col_a", pa.int32()),
+        pa.field("col_b", pa.string()),
+        pa.field("col_c", pa.int16()),
+    ]
+)
+
 
-@pytest.fixture(autouse=True, scope="session")
+def create_pa_table():
+    return pa.Table.from_pydict(
+        mapping={
+            "col_a": list(range(120)),
+            "col_b": random.choices(["a", "b", "c", "d"], k=120),
+            "col_c": random.choices(list(range(10)), k=120),
+        },
+        schema=_SCHEMA,
+    )
+
+
+@pytest.fixture(autouse=True, scope="function")
 def pyiceberg_table():
     from pyiceberg.catalog.sql import SqlCatalog
 
@@ -43,21 +62,7 @@ def pyiceberg_table():
         },
     )
 
-    schema = pa.schema(
-        [
-            pa.field("col_a", pa.int32()),
-            pa.field("col_b", pa.string()),
-            pa.field("col_c", pa.int16()),
-        ]
-    )
-    pya_table = pa.Table.from_pydict(
-        mapping={
-            "col_a": list(range(120)),
-            "col_b": random.choices(["a", "b", "c", "d"], k=120),
-            "col_c": random.choices(list(range(10)), k=120),
-        },
-        schema=schema,
-    )
+    pya_table = create_pa_table()
 
     if (_DB_NAME,) not in dummy_catalog.list_namespaces():
         dummy_catalog.create_namespace(_DB_NAME)
@@ -256,6 +261,43 @@ def test_read_basic():
     assert orig_table_p.equals(table_p)
 
 
+@pytest.mark.skipif(
+    parse_version(_get_pyarrow_version()) < parse_version("14.0.0"),
+    reason="PyIceberg 0.7.0 fails on pyarrow <= 14.0.0",
+)
+def test_write_basic():
+    # NOTE: Iceberg only works with PyArrow 9 or above.
+    pyarrow_version = _get_pyarrow_version()
+    if pyarrow_version is not None:
+        pyarrow_version = parse_version(pyarrow_version)
+    if pyarrow_version is not None and pyarrow_version < parse_version("9.0.0"):
+        return
+
+    sql_catalog = pyi_catalog.load_catalog(**_CATALOG_KWARGS)
+    table = sql_catalog.load_table(f"{_DB_NAME}.{_TABLE_NAME}")
+    table.delete()
+
+    ds = ray.data.from_arrow(create_pa_table())
+    ds.write_iceberg(
+        table_identifier=f"{_DB_NAME}.{_TABLE_NAME}",
+        catalog_kwargs=_CATALOG_KWARGS.copy(),
+    )
+
+    # Read the raw table from PyIceberg after writing
+    table = sql_catalog.load_table(f"{_DB_NAME}.{_TABLE_NAME}")
+    orig_table_p = (
+        table.scan()
+        .to_pandas()
+        .sort_values(["col_a", "col_b", "col_c"])
+        .reset_index(drop=True)
+    )
+
+    table_p = (
+        ds.to_pandas().sort_values(["col_a", "col_b", "col_c"]).reset_index(drop=True)
+    )
+    assert orig_table_p.equals(table_p)
+
+
 if __name__ == "__main__":
     import sys
 


Code review comments:
diff_hunk: @@ -0,0 +1,153 @@
+"""
+Module to write a Ray Dataset into an iceberg table, by using the Ray Datasink API.
+"""
+import logging
+
+from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional
+
+from ray.data.datasource.datasink import Datasink
+from ray.util.annotations import DeveloperAPI
+from ray.data.block import BlockAccessor, Block
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data.datasource.datasink import WriteResult
+import uuid
+
+if TYPE_CHECKING:
+    from pyiceberg.catalog import Catalog
+    from pyiceberg.manifest import DataFile
+
+
+logger = logging.getLogger(__name__)
+
+
+@DeveloperAPI
+class IcebergDatasink(Datasink[List["DataFile"]]):
+    """
+    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
+    `ray.data.Datasink`.
+
+    """
+
+    def __init__(
+        self,
+        table_identifier: str,
+        catalog_kwargs: Optional[Dict[str, Any]] = None,
+        snapshot_properties: Optional[Dict[str, str]] = None,
+    ):
+        """
+        Initialize the IcebergDatasink
+
+        Args:
+            table_identifier: The identifier of the table to read
+            catalog_kwargs: Optional arguments to use when setting up the Iceberg
+                catalog
+            snapshot_properties: Optional snapshot properties to
+                use when reading the table
+        """
+
+        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 {}
+        )
+
+        if "name" in self._catalog_kwargs:
+            self._catalog_name = self._catalog_kwargs.pop("name")
+        else:
+            self._catalog_name = "default"
+
+        self._uuid = None
+        self._io = None
+        self._txn = None
+        self._table_metadata = None
+
+    def __getstate__(self) -> dict:
+        """Exclude `_txn` during pickling."""
+        state = self.__dict__.copy()
+        del state["_txn"]
+        return state
+
+    def __setstate__(self, state: dict) -> None:
+        self.__dict__.update(state)
+        self._txn = 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"""
+        from pyiceberg.table import PropertyUtil, TableProperties
+
+        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._manifest_merge_enabled = PropertyUtil.property_as_bool(
+            self._table_metadata.properties,
+            TableProperties.MANIFEST_MERGE_ENABLED,
+            TableProperties.MANIFEST_MERGE_ENABLED_DEFAULT,
+        )
+
+    def write(self, blocks: Iterable[Block], ctx: TaskContext):
+        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
+
+        data_files_list = []
path: python/ray/data/_internal/datasource/iceberg_datasink.py
text: @user1:
Please add type annotations

diff_hunk: @@ -256,6 +261,43 @@ def test_read_basic():
     assert orig_table_p.equals(table_p)
 
 
+@pytest.mark.skipif(
+    parse_version(_get_pyarrow_version()) < parse_version("14.0.0"),
+    reason="PyIceberg 0.7.0 fails on pyarrow <= 14.0.0",
+)
+def test_write_basic():
+    # NOTE: Iceberg only works with PyArrow 9 or above.
+    pyarrow_version = _get_pyarrow_version()
+    if pyarrow_version is not None:
+        pyarrow_version = parse_version(pyarrow_version)
+    if pyarrow_version is not None and pyarrow_version < parse_version("9.0.0"):
+        return
path: python/ray/data/tests/test_iceberg.py
text: @user1:
This is not necessary anymore, right?