Repository: ray-project/ray
File: python/ray/data/tests/datasource/test_file_based_datasource.py (Python)

Code before:


def test_read_dir_and_file_mixed(ray_start_regular_shared, tmp_path):
    dir_path = os.path.join(tmp_path, "dir")
    os.mkdir(dir_path)

    p1 = os.path.join(dir_path, "a.txt")
    with open(p1, "wb") as f:
        f.write(b"a")

    p2 = os.path.join(tmp_path, "c.txt")
    with open(p2, "wb") as f:
        f.write(b"c")

    datasource = MockFileBasedDatasource([dir_path, p2])
    rows = execute_read_tasks(datasource.get_read_tasks(1))

    assert sorted(rows, key=lambda r: r["data"]) == [{"data": b"a"}, {"data": b"c"}]


def test_single_file_infinite_target_max_block_size(
    ray_start_regular_shared, target_max_block_size_infinite_or_default, tmp_path
):
    path = os.path.join(tmp_path, "file.txt")
    with open(path, "wb") as f:
        f.write(b"spam")

    datasource = MockFileBasedDatasource(path)
    rows = execute_read_tasks(datasource.get_read_tasks(1))

    assert rows == [{"data": b"spam"}]


def test_partitioning_hive(ray_start_regular_shared, tmp_path):
    path = os.path.join(tmp_path, "country=us")
    os.mkdir(path)
    with open(os.path.join(path, "file.txt"), "wb") as file:
        file.write(b"")

    datasource = MockFileBasedDatasource(tmp_path, partitioning=Partitioning("hive"))

    tasks = datasource.get_read_tasks(1)
    rows = execute_read_tasks(tasks)

    assert rows == [{"data": b"", "country": "us"}]


def test_partition_filter_hive(ray_start_regular_shared, tmp_path):
    for country in ["us", "jp"]:
        path = os.path.join(tmp_path, f"country={country}")
        os.mkdir(path)

Review comment (suggestion):
Nit: Here and elsewhere -- using the `Path` `/` operation and `write_bytes` might be a bit cleaner 
```suggestion
    path = tmp_path / "file.txt"
    path.write_bytes(b"spam")

    datasource = MockFileBasedDatasource(path)
    rows = execute_read_tasks(datasource.get_read_tasks(1))

    assert rows == [{"data": b"spam"}]
```

Code after:
def test_read_dir_and_file_mixed(ray_start_regular_shared, tmp_path):
    dir_path = tmp_path / "dir"
    dir_path.mkdir()

    p1 = dir_path / "a.txt"
    p1.write_bytes(b"a")

    p2 = tmp_path / "c.txt"
    p2.write_bytes(b"c")

    datasource = MockFileBasedDatasource([str(dir_path), str(p2)])
    rows = execute_read_tasks(datasource.get_read_tasks(1))

    assert sorted(rows, key=lambda r: r["data"]) == [{"data": b"a"}, {"data": b"c"}]


def test_single_file_infinite_target_max_block_size(
    ray_start_regular_shared, target_max_block_size_infinite_or_default, tmp_path
):
    path = tmp_path / "file.txt"
    path.write_bytes(b"spam")

    datasource = MockFileBasedDatasource(path)
    rows = execute_read_tasks(datasource.get_read_tasks(1))

    assert rows == [{"data": b"spam"}]


def test_partitioning_hive(ray_start_regular_shared, tmp_path):
    path = os.path.join(tmp_path, "country=us")
    os.mkdir(path)
    with open(os.path.join(path, "file.txt"), "wb") as file:
        file.write(b"")

    datasource = MockFileBasedDatasource(tmp_path, partitioning=Partitioning("hive"))

    tasks = datasource.get_read_tasks(1)
    rows = execute_read_tasks(tasks)

    assert rows == [{"data": b"", "country": "us"}]


def test_partition_filter_hive(ray_start_regular_shared, tmp_path):
    for country in ["us", "jp"]:
        path = os.path.join(tmp_path, f"country={country}")
        os.mkdir(path)
        with open(os.path.join(path, "file.txt"), "wb") as file:
            file.write(b"")

    filter = PathPartitionFilter.of(
        style=PartitionStyle.HIVE,