wehub-resource-sync
177 lines
6.0 KiB
Python
177 lines
6.0 KiB
Python
import os
|
|
from typing import Any, Dict
|
|
|
|
import pyarrow
|
|
import pytest
|
|
from pyarrow.fs import LocalFileSystem
|
|
|
|
import ray
|
|
from ray.data.block import BlockAccessor
|
|
from ray.data.datasource import BlockBasedFileDatasink, RowBasedFileDatasink
|
|
|
|
|
|
class FlakyOutputStream:
|
|
def __init__(self, stream: pyarrow.NativeFile, num_attempts: int):
|
|
self._stream = stream
|
|
self._num_attempts = num_attempts
|
|
|
|
def __enter__(self):
|
|
return self._stream.__enter__()
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback):
|
|
if self._num_attempts < 2:
|
|
raise RuntimeError("AWS Error NETWORK_CONNECTION")
|
|
|
|
self._stream.__exit__(exc_type, exc_value, traceback)
|
|
|
|
|
|
def test_flaky_block_based_open_output_stream(ray_start_regular_shared, tmp_path):
|
|
class FlakyCSVDatasink(BlockBasedFileDatasink):
|
|
def __init__(self, path: str):
|
|
super().__init__(path)
|
|
self._num_attempts = 0
|
|
self._filesystem = LocalFileSystem()
|
|
|
|
def open_output_stream(self, path: str) -> "pyarrow.NativeFile":
|
|
stream = self._filesystem.open_output_stream(path)
|
|
flaky_stream = FlakyOutputStream(stream, self._num_attempts)
|
|
self._num_attempts += 1
|
|
return flaky_stream
|
|
|
|
def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
|
|
block.to_pandas().to_csv(file)
|
|
|
|
ds = ray.data.range(100)
|
|
|
|
ds.write_datasink(FlakyCSVDatasink(tmp_path))
|
|
|
|
expected_values = list(range(100))
|
|
written_values = [row["id"] for row in ray.data.read_csv(tmp_path).take_all()]
|
|
assert sorted(written_values) == sorted(expected_values)
|
|
|
|
|
|
def test_flaky_row_based_open_output_stream(ray_start_regular_shared, tmp_path):
|
|
class FlakyTextDatasink(RowBasedFileDatasink):
|
|
def __init__(self, path: str):
|
|
super().__init__(path)
|
|
self._num_attempts = 0
|
|
self._filesystem = LocalFileSystem()
|
|
|
|
def open_output_stream(self, path: str) -> "pyarrow.NativeFile":
|
|
stream = self._filesystem.open_output_stream(path)
|
|
flaky_stream = FlakyOutputStream(stream, self._num_attempts)
|
|
self._num_attempts += 1
|
|
return flaky_stream
|
|
|
|
def write_row_to_file(self, row: Dict[str, Any], file: "pyarrow.NativeFile"):
|
|
file.write(f"{row['id']}".encode())
|
|
|
|
ds = ray.data.range(100)
|
|
|
|
ds.write_datasink(FlakyTextDatasink(tmp_path))
|
|
|
|
expected_values = [str(i) for i in range(100)]
|
|
written_values = [row["text"] for row in ray.data.read_text(tmp_path).take_all()]
|
|
assert sorted(written_values) == sorted(expected_values)
|
|
|
|
|
|
def test_flaky_write_block_to_file(ray_start_regular_shared, tmp_path):
|
|
class FlakyCSVDatasink(BlockBasedFileDatasink):
|
|
def __init__(self, path: str):
|
|
super().__init__(path)
|
|
self._num_attempts = 0
|
|
|
|
def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
|
|
if self._num_attempts < 2:
|
|
self._num_attempts += 1
|
|
raise RuntimeError("AWS Error INTERNAL_FAILURE")
|
|
|
|
block.to_pandas().to_csv(file)
|
|
|
|
ds = ray.data.range(100)
|
|
|
|
ds.write_datasink(FlakyCSVDatasink(tmp_path))
|
|
|
|
expected_values = list(range(100))
|
|
written_values = [row["id"] for row in ray.data.read_csv(tmp_path).take_all()]
|
|
assert sorted(written_values) == sorted(expected_values)
|
|
|
|
|
|
def test_flaky_write_row_to_file(ray_start_regular_shared, tmp_path):
|
|
class FlakyTextDatasink(RowBasedFileDatasink):
|
|
def __init__(self, path: str):
|
|
super().__init__(path)
|
|
self._num_attempts = 0
|
|
|
|
def write_row_to_file(self, row: Dict[str, Any], file: "pyarrow.NativeFile"):
|
|
if self._num_attempts < 2:
|
|
self._num_attempts += 1
|
|
raise RuntimeError("AWS Error INTERNAL_FAILURE")
|
|
|
|
file.write(f"{row['id']}".encode())
|
|
|
|
ds = ray.data.range(100)
|
|
|
|
ds.write_datasink(FlakyTextDatasink(tmp_path))
|
|
|
|
expected_values = [str(i) for i in range(100)]
|
|
written_values = [row["text"] for row in ray.data.read_text(tmp_path).take_all()]
|
|
assert sorted(written_values) == sorted(expected_values)
|
|
|
|
|
|
@pytest.mark.parametrize("num_rows", [0, 1])
|
|
def test_write_preserves_user_directory(num_rows, tmp_path, ray_start_regular_shared):
|
|
class MockFileDatasink(BlockBasedFileDatasink):
|
|
def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
|
|
file.write(b"")
|
|
|
|
ds = ray.data.range(num_rows)
|
|
path = os.path.join(tmp_path, "test")
|
|
os.mkdir(path) # User-created directory
|
|
|
|
ds.write_datasink(MockFileDatasink(path=path))
|
|
|
|
assert os.path.isdir(path)
|
|
|
|
|
|
def test_write_creates_dir(tmp_path, ray_start_regular_shared):
|
|
class MockFileDatasink(BlockBasedFileDatasink):
|
|
def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
|
|
file.write(b"")
|
|
|
|
ds = ray.data.range(1)
|
|
path = os.path.join(tmp_path, "test")
|
|
|
|
ds.write_datasink(MockFileDatasink(path=path, try_create_dir=True))
|
|
|
|
assert os.path.isdir(path)
|
|
|
|
|
|
@pytest.mark.parametrize("min_rows_per_file", [5, 10, 50])
|
|
def test_write_min_rows_per_file(tmp_path, ray_start_regular_shared, min_rows_per_file):
|
|
class MockFileDatasink(BlockBasedFileDatasink):
|
|
def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
|
|
for _ in range(block.num_rows()):
|
|
file.write(b"row\n")
|
|
|
|
ds = ray.data.range(100, override_num_blocks=20)
|
|
|
|
ds.write_datasink(
|
|
MockFileDatasink(path=tmp_path, min_rows_per_file=min_rows_per_file)
|
|
)
|
|
|
|
num_rows_written_total = 0
|
|
for filename in os.listdir(tmp_path):
|
|
with open(os.path.join(tmp_path, filename), "r") as file:
|
|
num_rows_written = len(file.read().splitlines())
|
|
assert num_rows_written == min_rows_per_file
|
|
num_rows_written_total += num_rows_written
|
|
|
|
assert num_rows_written_total == 100
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|