[Data] Add Datasource.on_write_start

python/ray/data/BUILD

@@ -169,6 +169,14 @@ py_test(
     deps = ["//:ray_lib", ":conftest"],
 )
 
+py_test(
+    name = "test_file_based_datasource",
+    size = "small",
+    srcs = ["tests/test_file_based_datasource.py"],
+    tags = ["team:data", "exclusive"],
+    deps = ["//:ray_lib", ":conftest"],
+)
+
 py_test(
     name = "test_image",
     size = "small",

python/ray/data/dataset.py

@@ -3317,6 +3317,7 @@ def write_fn_wrapper(blocks: Iterator[Block], ctx, fn) -> Iterator[Block]:
             try:
                 import pandas as pd
 
+                datasource.on_write_start(**write_args)
                 self._write_ds = Dataset(
                     plan, self._epoch, self._lazy, logical_plan
                 ).materialize()
@@ -3326,7 +3327,7 @@ def write_fn_wrapper(blocks: Iterator[Block], ctx, fn) -> Iterator[Block]:
                     for block in blocks
                 )
                 write_results = [block["write_result"][0] for block in blocks]
-                datasource.on_write_complete(write_results)
+                datasource.on_write_complete(write_results, **write_args)
             except Exception as e:
                 datasource.on_write_failed([], e)
                 raise

python/ray/data/datasource/datasource.py

@@ -51,6 +51,17 @@ def prepare_read(self, parallelism: int, **read_args) -> List["ReadTask"]:
         """Deprecated: Please implement create_reader() instead."""
         raise NotImplementedError
 
+    def on_write_start(self, **write_args) -> None:
+        """Callback for when a write job starts.
+
+        Use this method to perform setup for write tasks. For example, creating a
+        staging bucket in S3.
+
+        Args:
+            write_args: Additional kwargs to pass to the datasource impl.
+        """
+        pass
+
     def write(
         self,
         blocks: Iterable[Block],

python/ray/data/datasource/file_based_datasource.py

@@ -261,6 +261,32 @@ def _read_file(self, f: "pyarrow.NativeFile", path: str, **reader_args) -> Block
             "Subclasses of FileBasedDatasource must implement _read_file()."
         )
 
+    def on_write_start(
+        self,
+        path: str,
+        try_create_dir: bool = True,
+        filesystem: Optional["pyarrow.fs.FileSystem"] = None,
+        **write_args,
+    ) -> None:
+        """Create a directory to write files to.
+
+        If ``try_create_dir`` is ``False``, this method is a no-op.
+        """
+        from pyarrow.fs import FileType
+
+        self.has_created_dir = False
+        if try_create_dir:
+            paths, filesystem = _resolve_paths_and_filesystem(path, filesystem)
+            assert len(paths) == 1, len(paths)
+            path = paths[0]
+
+            if filesystem.get_file_info(path).type is FileType.NotFound:
+                # Arrow's S3FileSystem doesn't allow creating buckets by default, so we
+                # add a query arg enabling bucket creation if an S3 URI is provided.
+                tmp = _add_creatable_buckets_param_if_s3_uri(path)
+                filesystem.create_dir(tmp, recursive=True)
+                self.has_created_dir = True
+
     def write(
         self,
         blocks: Iterable[Block],
@@ -306,15 +332,6 @@ def write(
             if block.num_rows() == 0:
                 continue
 
-            if block_idx == 0:
-                # On the first non-empty block, try to create the directory.
-                if try_create_dir:
-                    # Arrow's S3FileSystem doesn't allow creating buckets by
-                    # default, so we add a query arg enabling bucket creation
-                    # if an S3 URI is provided.
-                    tmp = _add_creatable_buckets_param_if_s3_uri(path)
-                    filesystem.create_dir(tmp, recursive=True)
-
             fs = _unwrap_s3_serialization_workaround(filesystem)
 
             if self._WRITE_FILE_PER_ROW:
@@ -367,6 +384,23 @@ def write(
         # succeeds.
         return "ok"
 
+    def on_write_complete(
+        self,
+        write_results: List[WriteResult],
+        path: Optional[str] = None,
+        filesystem: Optional["pyarrow.fs.FileSystem"] = None,
+        **kwargs,
+    ) -> None:
+        if not self.has_created_dir:
+            return
+
+        paths, filesystem = _resolve_paths_and_filesystem(path, filesystem)
+        assert len(paths) == 1, len(paths)
+        path = paths[0]
+
+        if all(write_results == "skip" for write_results in write_results):
+            filesystem.delete_dir(path)
+
     def _write_block(
         self,
         f: "pyarrow.NativeFile",

python/ray/data/tests/test_file_based_datasource.py

@@ -0,0 +1,43 @@
+import os
+
+import pyarrow
+import pytest
+
+import ray
+from ray.data.block import BlockAccessor
+from ray.data.datasource import FileBasedDatasource
+
+
+class MockFileBasedDatasource(FileBasedDatasource):
+    def _write_block(
+        self, f: "pyarrow.NativeFile", block: BlockAccessor, **writer_args
+    ):
+        f.write(b"")
+
+
+@pytest.mark.parametrize("num_rows", [0, 1])
+def test_write_preserves_user_directory(num_rows, tmp_path, ray_start_regular_shared):
+    ds = ray.data.range(num_rows)
+    path = os.path.join(tmp_path, "test")
+    os.mkdir(path)  # User-created directory
+
+    ds.write_datasource(MockFileBasedDatasource(), dataset_uuid=ds._uuid, path=path)
+
+    assert os.path.isdir(path)
+
+
+def test_write_creates_dir(tmp_path, ray_start_regular_shared):
+    ds = ray.data.range(1)
+    path = os.path.join(tmp_path, "test")
+
+    ds.write_datasource(
+        MockFileBasedDatasource(), dataset_uuid=ds._uuid, path=path, try_create_dir=True
+    )
+
+    assert os.path.isdir(path)
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main(["-v", __file__]))