wehub-resource-sync
295 lines
10 KiB
Python
295 lines
10 KiB
Python
import json
|
|
import unittest
|
|
from typing import TYPE_CHECKING, Optional
|
|
from unittest import mock
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
from delta_sharing.protocol import Table
|
|
from delta_sharing.rest_client import DataSharingRestClient
|
|
|
|
from ray.data._internal.datasource.delta_sharing_datasource import (
|
|
DeltaSharingDatasource,
|
|
_parse_delta_sharing_url,
|
|
)
|
|
from ray.data.block import BlockMetadata
|
|
from ray.data.dataset import Dataset
|
|
from ray.data.datasource.datasource import ReadTask
|
|
from ray.data.read_api import read_delta_sharing_tables
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.data.context import DataContext
|
|
|
|
|
|
class TestDeltaSharingDatasource(unittest.TestCase):
|
|
def setUp(self):
|
|
self.url = "path/to/profile#share.schema.table"
|
|
self.limit = 1000
|
|
self.version = 1
|
|
self.json_predicate_hints = '{"column":"value"}'
|
|
self.table = Table(name="table", share="share", schema="schema")
|
|
|
|
self.mock_rest_client = mock.create_autospec(DataSharingRestClient)
|
|
self.mock_response = mock.Mock()
|
|
self.mock_rest_client.list_files_in_table.return_value = self.mock_response
|
|
|
|
self.mock_response.add_files = [
|
|
{"url": "file1", "id": "1"},
|
|
{"url": "file2", "id": "2"},
|
|
]
|
|
self.mock_response.metadata.schema_string = json.dumps(
|
|
{
|
|
"type": "struct",
|
|
"fields": [
|
|
{
|
|
"name": "column1",
|
|
"type": "string",
|
|
"nullable": True,
|
|
"metadata": {},
|
|
}
|
|
],
|
|
}
|
|
)
|
|
|
|
@patch(
|
|
"ray.data._internal.datasource.delta_sharing_datasource.DeltaSharingDatasource."
|
|
"setup_delta_sharing_connections"
|
|
)
|
|
def test_init(self, mock_setup_delta_sharing_connections):
|
|
mock_setup_delta_sharing_connections.return_value = (
|
|
self.table,
|
|
self.mock_rest_client,
|
|
)
|
|
datasource = DeltaSharingDatasource(
|
|
url=self.url,
|
|
json_predicate_hints=self.json_predicate_hints,
|
|
limit=self.limit,
|
|
version=self.version,
|
|
timestamp=None,
|
|
)
|
|
|
|
self.assertEqual(datasource._url, self.url)
|
|
self.assertEqual(datasource._json_predicate_hints, self.json_predicate_hints)
|
|
self.assertEqual(datasource._limit, self.limit)
|
|
self.assertEqual(datasource._version, self.version)
|
|
self.assertEqual(datasource._timestamp, None)
|
|
|
|
@patch(
|
|
"ray.data._internal.datasource.delta_sharing_datasource.DeltaSharingDatasource."
|
|
"setup_delta_sharing_connections"
|
|
)
|
|
def test_get_read_tasks(self, mock_setup_delta_sharing_connections):
|
|
mock_setup_delta_sharing_connections.return_value = (
|
|
self.table,
|
|
self.mock_rest_client,
|
|
)
|
|
datasource = DeltaSharingDatasource(
|
|
url=self.url,
|
|
json_predicate_hints=self.json_predicate_hints,
|
|
limit=self.limit,
|
|
version=self.version,
|
|
timestamp=None,
|
|
)
|
|
|
|
read_tasks = datasource.get_read_tasks(parallelism=2)
|
|
self.assertEqual(len(read_tasks), 2)
|
|
self.assertTrue(all(isinstance(task, ReadTask) for task in read_tasks))
|
|
|
|
for task in read_tasks:
|
|
metadata = task.metadata
|
|
self.assertIsInstance(metadata, BlockMetadata)
|
|
self.assertEqual(len(metadata.input_files), 1)
|
|
self.assertTrue(metadata.input_files[0]["url"] in ["file1", "file2"])
|
|
self.assertEqual(metadata.num_rows, None)
|
|
self.assertEqual(metadata.size_bytes, None)
|
|
self.assertEqual(task.schema, None)
|
|
self.assertEqual(metadata.exec_stats, None)
|
|
|
|
|
|
class TestParseDeltaSharingUrl(unittest.TestCase):
|
|
def test_valid_url(self):
|
|
url = "profile#share.schema.table"
|
|
expected_result = ("profile", "share", "schema", "table")
|
|
self.assertEqual(_parse_delta_sharing_url(url), expected_result)
|
|
|
|
def test_missing_hash(self):
|
|
url = "profile-share.schema.table"
|
|
with self.assertRaises(ValueError) as context:
|
|
_parse_delta_sharing_url(url)
|
|
self.assertEqual(str(context.exception), f"Invalid 'url': {url}")
|
|
|
|
def test_missing_fragments(self):
|
|
url = "profile#share.schema"
|
|
with self.assertRaises(ValueError) as context:
|
|
_parse_delta_sharing_url(url)
|
|
self.assertEqual(str(context.exception), f"Invalid 'url': {url}")
|
|
|
|
def test_empty_profile(self):
|
|
url = "#share.schema.table"
|
|
with self.assertRaises(ValueError) as context:
|
|
_parse_delta_sharing_url(url)
|
|
self.assertEqual(str(context.exception), f"Invalid 'url': {url}")
|
|
|
|
def test_empty_share(self):
|
|
url = "profile#.schema.table"
|
|
with self.assertRaises(ValueError) as context:
|
|
_parse_delta_sharing_url(url)
|
|
self.assertEqual(str(context.exception), f"Invalid 'url': {url}")
|
|
|
|
def test_empty_schema(self):
|
|
url = "profile#share..table"
|
|
with self.assertRaises(ValueError) as context:
|
|
_parse_delta_sharing_url(url)
|
|
self.assertEqual(str(context.exception), f"Invalid 'url': {url}")
|
|
|
|
def test_empty_table(self):
|
|
url = "profile#share.schema."
|
|
with self.assertRaises(ValueError) as context:
|
|
_parse_delta_sharing_url(url)
|
|
self.assertEqual(str(context.exception), f"Invalid 'url': {url}")
|
|
|
|
|
|
class MockDeltaSharingDatasource:
|
|
def __init__(
|
|
self, url, json_predicate_hints=None, limit=None, version=None, timestamp=None
|
|
):
|
|
self._url = url
|
|
self._json_predicate_hints = json_predicate_hints
|
|
self._limit = limit
|
|
self._version = version
|
|
self._timestamp = timestamp
|
|
|
|
def setup_delta_sharing_connections(self, url):
|
|
# Return mock objects for table and rest_client
|
|
table = MagicMock()
|
|
rest_client = MagicMock()
|
|
|
|
# Mock the rest_client's list_files_in_table method
|
|
rest_client.list_files_in_table.return_value = MagicMock(
|
|
add_files=[
|
|
{
|
|
"url": "https://s3-bucket-name.s3.us-west-2.amazonaws.com/delta-exchange-test/table2/date%3D2021-04-28/part-00000-591723a8-6a27-4240-a90e-57426f4736d2.c000.snappy.parquet", # noqa E501
|
|
"id": "591723a8-6a27-4240-a90e-57426f4736d2",
|
|
"size": 573,
|
|
"partitionValues": {"date": "2021-04-28"},
|
|
"stats": '{"numRecords":1,"minValues":{"eventTime":"2021-04-28T23:33:48.719Z"},"maxValues":{"eventTime":"2021-04-28T23:33:48.719Z"},"nullCount":{"eventTime":0}}', # noqa E501
|
|
"expirationTimestamp": 1652140800000,
|
|
}
|
|
],
|
|
metadata=MagicMock(
|
|
schema_string='{"type":"struct","fields":[{"name":"eventTime","type":"timestamp","nullable":true,"metadata":{}},{"name":"date","type":"date","nullable":true,"metadata":{}}]}' # noqa E501
|
|
),
|
|
)
|
|
return table, rest_client
|
|
|
|
def get_read_tasks(
|
|
self,
|
|
parallelism: int,
|
|
per_task_row_limit: Optional[int] = None,
|
|
data_context: Optional["DataContext"] = None,
|
|
):
|
|
self._table, self._rest_client = self.setup_delta_sharing_connections(self._url)
|
|
response = self._rest_client.list_files_in_table(
|
|
self._table,
|
|
jsonPredicateHints=self._json_predicate_hints,
|
|
limitHint=self._limit,
|
|
version=self._version,
|
|
timestamp=self._timestamp,
|
|
)
|
|
|
|
read_tasks = []
|
|
for _ in range(parallelism):
|
|
read_task = MagicMock()
|
|
read_task.metadata = MagicMock(
|
|
num_rows=1,
|
|
schema=None,
|
|
input_files=[file["url"] for file in response.add_files],
|
|
size_bytes=573,
|
|
exec_stats=None,
|
|
)
|
|
read_task.data = MagicMock(
|
|
return_value=[
|
|
{
|
|
"eventTime": "2021-04-28T23:33:48.719Z",
|
|
"date": "2021-04-28",
|
|
}
|
|
]
|
|
)
|
|
read_task.per_task_row_limit = per_task_row_limit
|
|
read_tasks.append(read_task)
|
|
return read_tasks
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_delta_sharing_datasource(mocker):
|
|
mock_datasource = mocker.patch(
|
|
"ray.data._internal.datasource.delta_sharing_datasource.DeltaSharingDatasource",
|
|
new=MockDeltaSharingDatasource,
|
|
)
|
|
return mock_datasource
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_ray_data_read_datasource(mocker):
|
|
mock_read_datasource = mocker.patch("ray.data.read_datasource")
|
|
mock_read_datasource.return_value = MagicMock(spec=Dataset)
|
|
return mock_read_datasource
|
|
|
|
|
|
@pytest.fixture
|
|
def setup_profile_file(tmpdir):
|
|
profile_content = {
|
|
"shareCredentialsVersion": 1,
|
|
"endpoint": "https://sharing.delta.io/delta-sharing/",
|
|
"bearerToken": "<token>",
|
|
"expirationTime": "2021-11-12T00:12:29.0Z",
|
|
}
|
|
profile_file = tmpdir.join("profile.json")
|
|
profile_file.write(json.dumps(profile_content))
|
|
return str(profile_file)
|
|
|
|
|
|
def test_read_delta_sharing_tables(
|
|
mock_delta_sharing_datasource, mock_ray_data_read_datasource, setup_profile_file
|
|
):
|
|
url = f"{setup_profile_file}#share.schema.table"
|
|
limit = 100
|
|
version = 1
|
|
timestamp = "2021-01-01T00:00:00Z"
|
|
json_predicate_hints = '{"eventTime": "2021-04-28T23:33:48.719Z"}'
|
|
ray_remote_args = {"num_cpus": 2}
|
|
concurrency = 4
|
|
override_num_blocks = 2
|
|
|
|
# Call the function under test
|
|
result = read_delta_sharing_tables(
|
|
url=url,
|
|
limit=limit,
|
|
version=version,
|
|
timestamp=timestamp,
|
|
json_predicate_hints=json_predicate_hints,
|
|
ray_remote_args=ray_remote_args,
|
|
concurrency=concurrency,
|
|
override_num_blocks=override_num_blocks,
|
|
)
|
|
|
|
# Assert the result and interactions
|
|
assert isinstance(result, Dataset)
|
|
mock_ray_data_read_datasource.assert_called_once()
|
|
args, kwargs = mock_ray_data_read_datasource.call_args
|
|
datasource = kwargs["datasource"]
|
|
assert datasource._url == url
|
|
assert datasource._json_predicate_hints == json_predicate_hints
|
|
assert datasource._limit == limit
|
|
assert datasource._version == version
|
|
assert datasource._timestamp == timestamp
|
|
assert kwargs["ray_remote_args"] == ray_remote_args
|
|
assert kwargs["concurrency"] == concurrency
|
|
assert kwargs["override_num_blocks"] == override_num_blocks
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|