Advanced: Read and write custom file types

This advanced guide shows you how to extend Ray Data to read and write file types that it doesn’t support natively. The guide relies on developer APIs, whose interfaces might change across minor releases of Ray.

Ray Data already supports images with the read_images() and write_images() APIs. This example shows you how to implement them anyway, as an illustration.

Read data from files

Tip

If you’re not contributing to Ray Data, you don’t need to create a Datasource. Instead, call read_binary_files() and decode files with map().

The core abstraction for reading files is FileBasedDatasource. It provides file-specific behavior on top of the Datasource interface.

To subclass FileBasedDatasource, implement the constructor and _read_stream.

Implement the datasource constructor

Call the superclass constructor and specify the files you want to read. Optionally, specify valid file extensions. Ray Data ignores files with other extensions.

from ray.data.datasource import FileBasedDatasource

class ImageDatasource(FileBasedDatasource):
    def __init__(self, paths: Union[str, List[str]], *, mode: str):
        super().__init__(
            paths,
            file_extensions=["png", "jpg", "jpeg", "bmp", "gif", "tiff"],
        )

        self.mode = mode  # Specify read options in the constructor

Implement _read_stream

_read_stream is a generator that yields one or more blocks of data from a file.

Blocks are Ray Data’s internal abstraction for a collection of rows. They can be PyArrow tables or pandas DataFrames.

Don’t create a block directly. Instead, add rows of data to a DelegatingBlockBuilder.

    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
        import io
        import numpy as np
        from PIL import Image
        from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder

        data = f.readall()
        image = Image.open(io.BytesIO(data))
        image = image.convert(self.mode)

        # Each block contains one row
        builder = DelegatingBlockBuilder()
        array = np.asarray(image)
        item = {"image": array}
        builder.add(item)
        yield builder.build()

Read your data

After you implement ImageDatasource, call read_datasource() to read images into a Dataset. Ray Data reads your files in parallel.

import ray

ds = ray.data.read_datasource(
    ImageDatasource("s3://anonymous@ray-example-data/batoidea", mode="RGB")
)

Write data to files

Note

The write interface is under active development and might change. To request a feature, open a GitHub issue.

The core abstractions for writing data to files are RowBasedFileDatasink and BlockBasedFileDatasink. They provide file-specific behavior on top of the Datasink interface.

To write one row per file, subclass RowBasedFileDatasink. Otherwise, subclass BlockBasedFileDatasink.

This example writes one image per file, so it subclasses RowBasedFileDatasink. To subclass RowBasedFileDatasink, implement the constructor and write_row_to_file().

Implement the datasink constructor

Call the superclass constructor and specify the folder to write to. Optionally, specify a string for the file format, such as "png". Ray Data uses the file format as the file extension.

from ray.data.datasource import RowBasedFileDatasink

class ImageDatasink(RowBasedFileDatasink):
    def __init__(self, path: str, column: str, file_format: str):
        super().__init__(path, file_format=file_format)

        self.column = column
        self.file_format = file_format  # Specify write options in the constructor

Implement write_row_to_file

write_row_to_file writes a row of data to a file. Each row is a dictionary that maps column names to values.

    def write_row_to_file(self, row: Dict[str, Any], file: pyarrow.NativeFile):
        import io
        from PIL import Image

        # PIL can't write to a NativeFile, so we have to write to a buffer first.
        image = Image.fromarray(row[self.column])
        buffer = io.BytesIO()
        image.save(buffer, format=self.file_format)
        file.write(buffer.getvalue())

Write your data

After you implement ImageDatasink, call write_datasink() to write images to files. Ray Data writes to multiple files in parallel.

ds.write_datasink(ImageDatasink("/tmp/results", column="image", file_format="png"))