chore: ray 2.45

awswrangler/distributed/ray/datasources/file_datasink.py

@@ -10,7 +10,7 @@
 from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
 from ray.data._internal.execution.interfaces import TaskContext
 from ray.data.block import Block, BlockAccessor
-from ray.data.datasource.datasink import Datasink
+from ray.data.datasource.datasink import Datasink, WriteResult
 from ray.data.datasource.filename_provider import FilenameProvider
 from ray.types import ObjectRef
 
@@ -21,7 +21,7 @@
 _logger: logging.Logger = logging.getLogger(__name__)
 
 
-class _BlockFileDatasink(Datasink):
+class _BlockFileDatasink(Datasink[str]):
     def __init__(
         self,
         path: str,
@@ -91,12 +91,12 @@ def write_block(self, file: Any, block: BlockAccessor) -> None:
     # and is meant to be used for singular actions like
     # [committing a transaction](https://docs.ray.io/en/latest/data/api/doc/ray.data.Datasource.html).
     # As deceptive as it may look, there is no race condition here.
-    def on_write_complete(self, write_results: list[Any], **_: Any) -> None:
+    def on_write_complete(self, write_results: WriteResult[str]) -> None:
         """Execute callback after all write tasks complete."""
         _logger.debug("Write complete %s.", write_results)
 
         # Collect and return all write task paths
-        self._write_paths.extend(write_results)
+        self._write_paths.extend(write_results.write_returns)
 
     def get_write_paths(self) -> list[str]:
         """Return S3 paths of where the results have been written."""