Files
ray-project--ray/python/ray/data/tests/unit/test_filename_provider.py
T
2026-07-13 13:17:40 +08:00

68 lines
2.1 KiB
Python

import pytest
from ray.data.datasource.filename_provider import FilenameProvider
@pytest.fixture(params=["csv", None])
def filename_provider(request):
yield FilenameProvider(dataset_uuid="", file_format=request.param)
def test_default_filename_for_task_includes_task_index(filename_provider):
filename_0 = filename_provider.get_filename_for_task(
write_uuid="spam", task_index=0
)
filename_1 = filename_provider.get_filename_for_task(
write_uuid="spam", task_index=1
)
assert filename_0 != filename_1
assert "000000" in filename_0
assert "000001" in filename_1
def test_default_get_filename_for_task_is_deterministic(filename_provider):
"""Test the new get_filename_for_task() method is deterministic."""
first_filename = filename_provider.get_filename_for_task(
write_uuid="spam", task_index=0
)
second_filename = filename_provider.get_filename_for_task(
write_uuid="spam", task_index=0
)
assert first_filename == second_filename
def test_default_row_filenames_derived_from_task_are_unique(filename_provider):
"""Row filenames derived from task filename with block/row index are unique."""
task_filename = filename_provider.get_filename_for_task(
write_uuid="spam", task_index=0
)
filenames = []
for block_index in range(2):
for row_index in range(4):
if "." in task_filename:
base, ext = task_filename.rsplit(".", 1)
filenames.append(f"{base}_{block_index:06}_{row_index:06}.{ext}")
else:
filenames.append(f"{task_filename}_{block_index:06}_{row_index:06}")
assert len(set(filenames)) == len(filenames)
def test_default_get_filename_for_task_is_unique(filename_provider):
"""Test the new get_filename_for_task() method generates unique filenames."""
filenames = [
filename_provider.get_filename_for_task(
write_uuid="spam",
task_index=task_index,
)
for task_index in range(4)
]
assert len(set(filenames)) == len(filenames)
if __name__ == "__main__":
import sys
sys.exit(pytest.main(["-v", __file__]))