feat(python): add a Ray Data VortexDatasink

docs/user-guide/ray.md

@@ -35,3 +35,11 @@ and then run `make -C docs doctest` -->
 <BLANKLINE>
 [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"))
+```

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)

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}]