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

Code before:
        (
            lazy_fixture("s3_fs_with_space"),
            lazy_fixture("s3_path_with_space"),
        ),  # Path contains space.
        (
            lazy_fixture("s3_fs_with_anonymous_crendential"),
            lazy_fixture("s3_path_with_anonymous_crendential"),
        ),
    ],
)
def test_parquet_read_random_shuffle(
    ray_start_regular_shared,
    restore_data_context,
    fs,
    data_path,
    target_max_block_size_infinite_or_default,
):
    # NOTE: set preserve_order to True to allow consistent output behavior.
    context = ray.data.DataContext.get_current()
    context.execution_options.preserve_order = True

    num_files = 20
    input_list = list(range(num_files))
    setup_data_path = _unwrap_protocol(data_path)
    for i in range(num_files):
        table = pa.Table.from_pandas(pd.DataFrame({"one": [i]}))
        path = os.path.join(setup_data_path, f"test_{i}.parquet")
        pq.write_table(table, path, filesystem=fs)

    shuffle = FileShuffleConfig(seed=0)
    ds = ray.data.read_parquet(data_path, filesystem=fs, shuffle=shuffle)

    first = [row["one"] for row in ds.take_all()]
    second = [row["one"] for row in ds.take_all()]

    assert sorted(first) == input_list
    assert sorted(second) == input_list
    assert first != second


@pytest.mark.parametrize(
    "fs,data_path",
    [
        (None, lazy_fixture("local_path")),
        (lazy_fixture("local_fs"), lazy_fixture("local_path")),
        (lazy_fixture("s3_fs"), lazy_fixture("s3_path")),
        (
            lazy_fixture("s3_fs_with_anonymous_crendential"),
            lazy_fixture("s3_path_with_anonymous_crendential"),
        ),
    ],

Review comment (suggestion):
Might be simpler to construct the Table directly:
```suggestion
        table = pa.Table.from_pydict({"id": [i]})
```

Code after:
        (
            lazy_fixture("s3_fs_with_space"),
            lazy_fixture("s3_path_with_space"),
        ),  # Path contains space.
        (
            lazy_fixture("s3_fs_with_anonymous_crendential"),
            lazy_fixture("s3_path_with_anonymous_crendential"),
        ),
    ],
)
def test_parquet_read_random_shuffle(
    ray_start_regular_shared,
    restore_data_context,
    fs,
    data_path,
    target_max_block_size_infinite_or_default,
):
    # NOTE: set preserve_order to True to allow consistent output behavior.
    context = ray.data.DataContext.get_current()
    context.execution_options.preserve_order = True

    num_files = 10
    input_list = list(range(num_files))
    setup_data_path = _unwrap_protocol(data_path)
    for i in range(num_files):
        table = pa.Table.from_pydict({"id": [i]})
        path = os.path.join(setup_data_path, f"test_{i}.parquet")
        pq.write_table(table, path, filesystem=fs)

    shuffle = FileShuffleConfig(seed=0)
    ds = ray.data.read_parquet(data_path, filesystem=fs, shuffle=shuffle)

    first = [row["id"] for row in ds.take_all()]
    second = [row["id"] for row in ds.take_all()]

    assert sorted(first) == input_list
    assert sorted(second) == input_list
    assert first != second


@pytest.mark.parametrize(
    "fs,data_path",
    [
        (None, lazy_fixture("local_path")),
        (lazy_fixture("local_fs"), lazy_fixture("local_path")),
        (lazy_fixture("s3_fs"), lazy_fixture("s3_path")),
        (
            lazy_fixture("s3_fs_with_anonymous_crendential"),
            lazy_fixture("s3_path_with_anonymous_crendential"),
        ),
    ],