wehub-resource-sync
49 lines
1.5 KiB
Python
49 lines
1.5 KiB
Python
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)
|