[Core] Extend Ray Data with IcebergDataSink for Distributed Writes to Iceberg Tables

[zszheng]

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

[richardliaw]

This proposal makes sense; would you be interested in contributing?

[jimmyxie-figma]

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.

[zszheng]

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

[zszheng]

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

[jimmyxie-figma]

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.

[jimmyxie-figma]

After some investigation, it seems that an implementation like this 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 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 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.

[jimmyxie-figma]

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)

[jimmyxie-figma]

@kevinjqliu when you get a chance can you take a quick look and let us know what you think about the proposed implementation?

[raulchen]

@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.