pass write result to on_write_complete

python/ray/data/_internal/datasource/bigquery_datasink.py

@@ -3,7 +3,7 @@
 import tempfile
 import time
 import uuid
-from typing import Iterable, Optional
+from typing import Any, Iterable, Optional
 
 import pyarrow.parquet as pq
 
@@ -70,7 +70,7 @@ def write(
         self,
         blocks: Iterable[Block],
         ctx: TaskContext,
-    ) -> None:
+    ) -> Any:
         def _write_single_block(block: Block, project_id: str, dataset: str) -> None:
             from google.api_core import exceptions
             from google.cloud import bigquery

python/ray/data/_internal/datasource/mongo_datasink.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Iterable
+from typing import Any, Iterable
 
 from ray.data._internal.datasource.mongo_datasource import (
     _validate_database_collection_exist,
@@ -26,7 +26,7 @@ def write(
         self,
         blocks: Iterable[Block],
         ctx: TaskContext,
-    ) -> None:
+    ) -> Any:
         import pymongo
 
         _validate_database_collection_exist(

python/ray/data/_internal/datasource/parquet_datasink.py

@@ -57,7 +57,7 @@ def write(
         self,
         blocks: Iterable[Block],
         ctx: TaskContext,
-    ) -> None:
+    ) -> Any:
         import pyarrow as pa
         import pyarrow.parquet as pq
 

python/ray/data/_internal/datasource/sql_datasink.py

@@ -1,4 +1,4 @@
-from typing import Callable, Iterable
+from typing import Any, Callable, Iterable
 
 from ray.data._internal.datasource.sql_datasource import Connection, _connect
 from ray.data._internal.execution.interfaces import TaskContext
@@ -18,7 +18,7 @@ def write(
         self,
         blocks: Iterable[Block],
         ctx: TaskContext,
-    ) -> None:
+    ) -> Any:
         with _connect(self.connection_factory) as cursor:
             for block in blocks:
                 block_accessor = BlockAccessor.for_block(block)

python/ray/data/_internal/planner/plan_write_op.py

@@ -1,5 +1,6 @@
 import itertools
-from typing import Callable, Iterator, List, Union
+import pickle
+from typing import Callable, Iterable, List, Union
 
 from ray.data._internal.compute import TaskPoolStrategy
 from ray.data._internal.execution.interfaces import PhysicalOperator
@@ -8,6 +9,7 @@
 from ray.data._internal.execution.operators.map_transformer import (
     BlockMapTransformFn,
     MapTransformer,
+    MapTransformFn,
 )
 from ray.data._internal.logical.operators.write_operator import Write
 from ray.data.block import Block, BlockAccessor
@@ -18,30 +20,40 @@
 
 def generate_write_fn(
     datasink_or_legacy_datasource: Union[Datasink, Datasource], **write_args
-) -> Callable[[Iterator[Block], TaskContext], Iterator[Block]]:
-    def fn(blocks: Iterator[Block], ctx) -> Iterator[Block]:
+) -> Callable[[Iterable[Block], TaskContext], Iterable[Block]]:
+    stats_fn = generate_collect_write_stats_fn()
+
+    def fn(blocks: Iterable[Block], ctx) -> Iterable[Block]:
         """Writes the blocks to the given datasink or legacy datasource.
 
         Outputs the original blocks to be written."""
         # Create a copy of the iterator, so we can return the original blocks.
         it1, it2 = itertools.tee(blocks, 2)
         if isinstance(datasink_or_legacy_datasource, Datasink):
-            datasink_or_legacy_datasource.write(it1, ctx)
+            write_result = datasink_or_legacy_datasource.write(it1, ctx)
         else:
-            datasink_or_legacy_datasource.write(it1, ctx, **write_args)
-        return it2
+            write_result = datasink_or_legacy_datasource.write(it1, ctx, **write_args)
+
+        import pandas as pd
+
+        payload = pd.DataFrame({"payload": [pickle.dumps(write_result)]})
+
+        stats = list(stats_fn(it2, ctx))
+        assert len(stats) == 1
+        block = pd.concat([stats[0], payload], axis=1)
+        return iter([block])
 
     return fn
 
 
-def generate_collect_write_stats_fn() -> (
-    Callable[[Iterator[Block], TaskContext], Iterator[Block]]
-):
+def generate_collect_write_stats_fn() -> Callable[
+    [Iterable[Block], TaskContext], Iterable[Block]
+]:
     # If the write op succeeds, the resulting Dataset is a list of
     # one Block which contain stats/metrics about the write.
     # Otherwise, an error will be raised. The Datasource can handle
     # execution outcomes with `on_write_complete()`` and `on_write_failed()``.
-    def fn(blocks: Iterator[Block], ctx) -> Iterator[Block]:
+    def fn(blocks: Iterable[Block], ctx) -> Iterable[Block]:
         """Handles stats collection for block writes."""
         block_accessors = [BlockAccessor.for_block(block) for block in blocks]
         total_num_rows = sum(ba.num_rows() for ba in block_accessors)
@@ -67,11 +79,9 @@ def plan_write_op(
     input_physical_dag = physical_children[0]
 
     write_fn = generate_write_fn(op._datasink_or_legacy_datasource, **op._write_args)
-    collect_stats_fn = generate_collect_write_stats_fn()
     # Create a MapTransformer for a write operator
-    transform_fns = [
+    transform_fns: List[MapTransformFn] = [
         BlockMapTransformFn(write_fn),
-        BlockMapTransformFn(collect_stats_fn),
     ]
     map_transformer = MapTransformer(transform_fns)
     return MapOperator.create(

python/ray/data/datasource/datasink.py

@@ -1,6 +1,6 @@
 import logging
 from dataclasses import dataclass, fields
-from typing import Iterable, List, Optional
+from typing import Any, Iterable, List, Optional
 
 import ray
 from ray.data._internal.execution.interfaces import TaskContext
@@ -67,7 +67,7 @@ def write(
         self,
         blocks: Iterable[Block],
         ctx: TaskContext,
-    ) -> None:
+    ) -> Any:
         """Write blocks. This is used by a single write task.
 
         Args:
@@ -82,6 +82,7 @@ def on_write_complete(self, write_result_blocks: List[Block]) -> WriteResult:
         This can be used to "commit" a write output. This method must
         succeed prior to ``write_datasink()`` returning to the user. If this
         method fails, then ``on_write_failed()`` is called.
+        Return value of write function is stored in payload column.
 
         Args:
             write_result_blocks: The blocks resulting from executing
@@ -165,7 +166,7 @@ def __init__(self):
                 self.rows_written = 0
                 self.enabled = True
 
-            def write(self, block: Block) -> None:
+            def write(self, block: Block) -> Any:
                 block = BlockAccessor.for_block(block)
                 self.rows_written += block.num_rows()
 
@@ -181,7 +182,7 @@ def write(
         self,
         blocks: Iterable[Block],
         ctx: TaskContext,
-    ) -> None:
+    ) -> Any:
         tasks = []
         if not self.enabled:
             raise ValueError("disabled")

python/ray/data/tests/test_datasink.py

@@ -1,11 +1,13 @@
-from typing import Iterable
+import pickle
+from typing import Any, Iterable, List
 
 import pytest
 
 import ray
 from ray.data._internal.execution.interfaces import TaskContext
 from ray.data.block import Block
 from ray.data.datasource import Datasink
+from ray.data.datasource.datasink import WriteResult
 
 
 @pytest.mark.parametrize("num_rows_per_write", [5, 10, 50])
@@ -14,7 +16,7 @@ class MockDatasink(Datasink):
         def __init__(self, num_rows_per_write):
             self._num_rows_per_write = num_rows_per_write
 
-        def write(self, blocks: Iterable[Block], ctx: TaskContext) -> None:
+        def write(self, blocks: Iterable[Block], ctx: TaskContext) -> Any:
             assert sum(len(block) for block in blocks) == self._num_rows_per_write
 
         @property
@@ -26,6 +28,34 @@ def num_rows_per_write(self):
     )
 
 
+@pytest.mark.parametrize("num_rows_per_write", [5, 10, 50])
+def test_on_write_complete(tmp_path, ray_start_regular_shared, num_rows_per_write):
+    class MockDatasink(Datasink):
+        def __init__(self, num_rows_per_write):
+            self._num_rows_per_write = num_rows_per_write
+            self.payloads = None
+
+        def write(self, blocks: Iterable[Block], ctx: TaskContext) -> Any:
+            assert sum(len(block) for block in blocks) == self._num_rows_per_write
+            return f"task-{ctx.task_idx}"
+
+        def on_write_complete(self, write_result_blocks: List[Block]) -> WriteResult:
+            self.payloads = [
+                pickle.loads(result["payload"].iloc[0])
+                for result in write_result_blocks
+            ]
+            return super().on_write_complete(write_result_blocks)
+
+        @property
+        def num_rows_per_write(self):
+            return self._num_rows_per_write
+
+    sink = MockDatasink(num_rows_per_write)
+    ray.data.range(100, override_num_blocks=20).write_datasink(sink)
+    expect = [f"task-{i}" for i in range(100 // num_rows_per_write)]
+    assert sink.payloads == expect
+
+
 if __name__ == "__main__":
     import sys
 

python/ray/data/tests/test_formats.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from typing import Iterable, List
+from typing import Any, Iterable, List
 
 import pandas as pd
 import pyarrow as pa
@@ -236,7 +236,7 @@ def __init__(self):
                 self.rows_written = 0
                 self.node_ids = set()
 
-            def write(self, node_id: str, block: Block) -> str:
+            def write(self, node_id: str, block: Block) -> Any:
                 block = BlockAccessor.for_block(block)
                 self.rows_written += block.num_rows()
                 self.node_ids.add(node_id)
@@ -255,7 +255,7 @@ def write(
         self,
         blocks: Iterable[Block],
         ctx: TaskContext,
-    ) -> None:
+    ) -> Any:
         data_sink = self.data_sink
 
         def write(b):