chore: import upstream snapshot with attribution

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

@@ -0,0 +1,48 @@
+import logging
+from typing import Iterable
+
+from ray.data._internal.datasource.mongo_datasource import (
+    _validate_database_collection_exist,
+)
+from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+from ray.data._internal.execution.interfaces import TaskContext
+from ray.data._internal.util import _check_import
+from ray.data.block import Block, BlockAccessor
+from ray.data.datasource.datasink import Datasink
+
+logger = logging.getLogger(__name__)
+
+
+class MongoDatasink(Datasink[None]):
+    def __init__(self, uri: str, database: str, collection: str) -> None:
+        _check_import(self, module="pymongo", package="pymongo")
+        _check_import(self, module="pymongoarrow", package="pymongoarrow")
+
+        self.uri = uri
+        self.database = database
+        self.collection = collection
+
+    def write(
+        self,
+        blocks: Iterable[Block],
+        ctx: TaskContext,
+    ) -> None:
+        import pymongo
+
+        _validate_database_collection_exist(
+            pymongo.MongoClient(self.uri), self.database, self.collection
+        )
+
+        def write_block(uri: str, database: str, collection: str, block: Block):
+            from pymongoarrow.api import write
+
+            block = BlockAccessor.for_block(block).to_arrow()
+            client = pymongo.MongoClient(uri)
+            write(client[database][collection], block)
+
+        builder = DelegatingBlockBuilder()
+        for block in blocks:
+            builder.add_block(block)
+        block = builder.build()
+
+        write_block(self.uri, self.database, self.collection, block)