wehub-resource-sync
139 lines
4.7 KiB
Python
139 lines
4.7 KiB
Python
import logging
|
|
from json import loads
|
|
from typing import TYPE_CHECKING, List, Optional, Tuple
|
|
|
|
import numpy as np
|
|
|
|
from ray.data._internal.util import _check_import
|
|
from ray.data.block import BlockMetadata
|
|
from ray.data.datasource.datasource import Datasource, ReadTask
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.data.context import DataContext
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class DeltaSharingDatasource(Datasource):
|
|
def __init__(
|
|
self,
|
|
url: str,
|
|
json_predicate_hints: Optional[str] = None,
|
|
limit: Optional[int] = None,
|
|
version: Optional[int] = None,
|
|
timestamp: Optional[str] = None,
|
|
):
|
|
_check_import(self, module="delta_sharing", package="delta-sharing")
|
|
|
|
if limit is not None:
|
|
assert (
|
|
isinstance(limit, int) and limit >= 0
|
|
), "'limit' must be a non-negative int"
|
|
|
|
self._url = url
|
|
self._json_predicate_hints = json_predicate_hints
|
|
self._limit = limit
|
|
self._version = version
|
|
self._timestamp = timestamp
|
|
|
|
def estimate_inmemory_data_size(self) -> Optional[int]:
|
|
return None
|
|
|
|
def _read_files(self, files, converters):
|
|
"""Read files with Delta Sharing."""
|
|
from delta_sharing.reader import DeltaSharingReader
|
|
|
|
for file in files:
|
|
yield DeltaSharingReader._to_pandas(
|
|
action=file, converters=converters, for_cdf=False, limit=None
|
|
)
|
|
|
|
def setup_delta_sharing_connections(self, url: str):
|
|
"""
|
|
Set up delta sharing connections based on the url.
|
|
|
|
Args:
|
|
url: A URL under the format "<profile>#<share>.<schema>.<table>"
|
|
|
|
Returns:
|
|
A tuple of (table, rest_client) where table is a delta_sharing Table
|
|
object and rest_client is a DataSharingRestClient instance.
|
|
"""
|
|
from delta_sharing.protocol import DeltaSharingProfile, Table
|
|
from delta_sharing.rest_client import DataSharingRestClient
|
|
|
|
profile_str, share, schema, table_str = _parse_delta_sharing_url(url)
|
|
table = Table(name=table_str, share=share, schema=schema)
|
|
|
|
profile = DeltaSharingProfile.read_from_file(profile_str)
|
|
rest_client = DataSharingRestClient(profile)
|
|
return table, rest_client
|
|
|
|
def get_read_tasks(
|
|
self,
|
|
parallelism: int,
|
|
per_task_row_limit: Optional[int] = None,
|
|
data_context: Optional["DataContext"] = None,
|
|
) -> List[ReadTask]:
|
|
assert parallelism > 0, f"Invalid parallelism {parallelism}"
|
|
from delta_sharing.converter import to_converters
|
|
|
|
self._table, self._rest_client = self.setup_delta_sharing_connections(self._url)
|
|
self._response = self._rest_client.list_files_in_table(
|
|
self._table,
|
|
jsonPredicateHints=self._json_predicate_hints,
|
|
limitHint=self._limit,
|
|
version=self._version,
|
|
timestamp=self._timestamp,
|
|
)
|
|
|
|
if len(self._response.add_files) == 0 or self._limit == 0:
|
|
logger.warning("No files found from the delta sharing table or limit is 0")
|
|
|
|
schema_json = loads(self._response.metadata.schema_string)
|
|
self._converters = to_converters(schema_json)
|
|
|
|
read_tasks = []
|
|
# get file list to be read in this task and preserve original chunk order
|
|
for files in np.array_split(self._response.add_files, parallelism):
|
|
files = files.tolist()
|
|
metadata = BlockMetadata(
|
|
num_rows=None,
|
|
input_files=files,
|
|
size_bytes=None,
|
|
exec_stats=None,
|
|
)
|
|
converters = self._converters
|
|
read_task = ReadTask(
|
|
lambda f=files: self._read_files(f, converters),
|
|
metadata,
|
|
per_task_row_limit=per_task_row_limit,
|
|
)
|
|
read_tasks.append(read_task)
|
|
|
|
return read_tasks
|
|
|
|
|
|
def _parse_delta_sharing_url(url: str) -> Tuple[str, str, str, str]:
|
|
"""
|
|
Developed from delta_sharing's _parse_url function.
|
|
https://github.com/delta-io/delta-sharing/blob/main/python/delta_sharing/delta_sharing.py#L36
|
|
|
|
Args:
|
|
url: a url under the format "<profile>#<share>.<schema>.<table>"
|
|
|
|
Returns:
|
|
a tuple with parsed (profile, share, schema, table)
|
|
"""
|
|
shape_index = url.rfind("#")
|
|
if shape_index < 0:
|
|
raise ValueError(f"Invalid 'url': {url}")
|
|
profile = url[0:shape_index]
|
|
fragments = url[shape_index + 1 :].split(".")
|
|
if len(fragments) != 3:
|
|
raise ValueError(f"Invalid 'url': {url}")
|
|
share, schema, table = fragments
|
|
if len(profile) == 0 or len(share) == 0 or len(schema) == 0 or len(table) == 0:
|
|
raise ValueError(f"Invalid 'url': {url}")
|
|
return (profile, share, schema, table)
|