diff --git a/docs/user-guide/ray.md b/docs/user-guide/ray.md index 0b831e23a1b..79dbaa628ba 100644 --- a/docs/user-guide/ray.md +++ b/docs/user-guide/ray.md @@ -35,3 +35,11 @@ and then run `make -C docs doctest` --> [3000 rows x 19 columns] ``` + +Use `VortexDatasink` to write a Dataset without collecting it on the driver. Each Ray write task streams its Arrow batches to one Vortex file. + +```python +from vortex.ray.datasink import VortexDatasink + +ds.write_datasink(VortexDatasink("ray_data_output")) +``` diff --git a/vortex-python/python/vortex/ray/datasink.py b/vortex-python/python/vortex/ray/datasink.py new file mode 100644 index 00000000000..cf35a1996b5 --- /dev/null +++ b/vortex-python/python/vortex/ray/datasink.py @@ -0,0 +1,61 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +from __future__ import annotations + +import itertools +import posixpath +from collections.abc import Iterable +from typing import final +from urllib.parse import quote + +import pyarrow as pa +from ray.data._internal.execution.interfaces import TaskContext +from ray.data._internal.planner.plan_write_op import WRITE_UUID_KWARG_NAME +from ray.data._internal.util import RetryingPyFileSystem +from ray.data.block import Block, BlockAccessor +from ray.data.datasource.file_datasink import BlockBasedFileDatasink +from typing_extensions import override + +from .. import io as vx_io +from .datasource import _url_scheme + + +@final +class VortexDatasink(BlockBasedFileDatasink): + """Write a Ray Dataset as one Vortex file per write task.""" + + def __init__(self, path: str) -> None: + super().__init__(path, file_format="vortex") + + @override + def write(self, blocks: Iterable[Block], ctx: TaskContext) -> None: + blocks = iter(blocks) + try: + first = next(blocks) + except StopIteration: + return + + first_table = BlockAccessor.for_block(first).to_arrow() + tables = itertools.chain( + (first_table,), + (BlockAccessor.for_block(block).to_arrow() for block in blocks), + ) + batches = (batch for table in tables for batch in table.to_batches()) + reader = pa.RecordBatchReader.from_batches(first_table.schema, batches) + + write_uuid = ctx.kwargs[WRITE_UUID_KWARG_NAME] + if hasattr(self.filename_provider, "get_filename_for_task"): + filename = self.filename_provider.get_filename_for_task(write_uuid, ctx.task_idx) + else: + # Ray 2.48 only provides the block-based filename hook. + filename = self.filename_provider.get_filename_for_block(first_table, write_uuid, ctx.task_idx, 0) + + write_path = posixpath.join(self.path, filename) + filesystem = self.filesystem + if isinstance(filesystem, RetryingPyFileSystem): + filesystem = filesystem.unwrap() + scheme = _url_scheme(filesystem) + if scheme is not None: + write_path = f"{scheme}://{quote(write_path)}" + vx_io.write(reader, write_path) diff --git a/vortex-python/test/test_datasource.py b/vortex-python/test/test_datasource.py index 5adca691f27..b99993477d1 100644 --- a/vortex-python/test/test_datasource.py +++ b/vortex-python/test/test_datasource.py @@ -14,10 +14,13 @@ import ray from pyarrow.fs import GcsFileSystem from ray.data import read_datasource +from ray.data._internal.execution.interfaces import TaskContext +from ray.data._internal.planner.plan_write_op import WRITE_UUID_KWARG_NAME from typing_extensions import override import vortex as vx import vortex.expr as ve +from vortex.ray.datasink import VortexDatasink from vortex.ray.datasource import VortexDatasource, partition from .test_file import record @@ -175,3 +178,66 @@ def test_url_scheme_of_gcs_is_gs() -> None: from vortex.ray.datasource import _url_scheme assert _url_scheme(GcsFileSystem(anonymous=True)) == "gs" + + +def test_vortex_datasink_writes_and_reads_dataset(ray_init, tmp_path: Path) -> None: + output = tmp_path / "output" + ds = ray.data.from_items([record(x) for x in range(20)], override_num_blocks=2) + + ds.write_datasink(VortexDatasink(str(output))) + + assert len(list(output.glob("*.vortex"))) == 2 + actual = read_datasource(VortexDatasource(url=str(output))).sort("index") + table = pa.concat_tables(pa.Table.from_pydict(batch) for batch in actual.iter_batches()) + expected = pa.Table.from_pylist([record(x) for x in range(20)], schema=table.schema) + assert table == expected + + +def test_vortex_datasink_writes_arrow_nulls(ray_init, tmp_path: Path) -> None: + output = tmp_path / "output" + rows = [{"x": 1, "label": "a"}, {"x": None, "label": None}, {"x": 3, "label": "c"}] + source = pa.Table.from_pylist(rows) + + ray.data.from_arrow(source).write_datasink(VortexDatasink(str(output))) + + files = sorted(output.glob("*.vortex")) + table = pa.concat_tables([vx.open(str(path)).to_arrow().read_all() for path in files]) + assert table.to_pylist() == rows + + +def test_vortex_datasink_appends_unique_files(ray_init, tmp_path: Path) -> None: + output = tmp_path / "output" + ds = ray.data.from_items([{"x": 1}]) + + ds.write_datasink(VortexDatasink(str(output))) + ds.write_datasink(VortexDatasink(str(output))) + + assert len(list(output.glob("*.vortex"))) == 2 + assert read_datasource(VortexDatasource(url=str(output))).count() == 2 + + +def test_vortex_datasink_empty_dataset(ray_init, tmp_path: Path) -> None: + output = tmp_path / "output" + source = pa.Table.from_batches([], schema=pa.schema([("x", pa.int64())])) + + ray.data.from_arrow(source).write_datasink(VortexDatasink(str(output))) + + assert list(output.glob("*.vortex")) == [] + + +def test_vortex_datasink_s3_url(s3_root: Path, monkeypatch: pytest.MonkeyPatch) -> None: + writes = {} + + def write(reader: pa.RecordBatchReader, path: str) -> None: + writes[path] = reader.read_all() + + monkeypatch.setattr("vortex.ray.datasink.vx_io.write", write) + sink = VortexDatasink("s3://ray-datasource-test/data") + ctx = TaskContext(task_idx=0, op_name="Write", kwargs={WRITE_UUID_KWARG_NAME: "test"}) + + sink.write([pa.Table.from_pylist([{"x": 1}])], ctx) + + path = next(iter(writes)) + assert path.startswith("s3://ray-datasource-test/data/test_") + assert path.endswith(".vortex") + assert writes[path].to_pylist() == [{"x": 1}]