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

Code before:
    setup_data_path = _unwrap_protocol(data_path)

    read_paths = []
    file_rows = {}
    for i in range(num_files):
        start = i * rows_per_file
        end = start + rows_per_file
        chunk = input_list[start:end]
        table = pa.Table.from_pandas(pd.DataFrame({"one": chunk}))
        filename = f"test{i}.parquet"
        write_path = os.path.join(setup_data_path, filename)
        pq.write_table(table, write_path, filesystem=fs)
        read_path = os.path.join(data_path, filename)
        read_paths.append(read_path)
        file_rows[read_path] = chunk

    # Pick a seed that produces a non-identity permutation for this file list.
    shuffle_seed = None
    shuffled_paths = None
    for seed in range(100):
        candidate_paths = list(read_paths)
        np.random.default_rng(seed).shuffle(candidate_paths)
        if candidate_paths != read_paths:
            shuffle_seed = seed
            shuffled_paths = candidate_paths
            break

    assert shuffle_seed is not None, "Failed to find a non-identity shuffle seed."
    assert shuffled_paths is not None

    # Use a fixed seed so the shuffled order is deterministic and testable.
    shuffle = FileShuffleConfig(seed=shuffle_seed, reseed_after_execution=False)
    ds = ray.data.read_parquet(read_paths, filesystem=fs, shuffle=shuffle)
    result = [row["one"] for row in ds.take_all()]

    # Expected row order is the concatenation of rows per shuffled file order.
    expected = [value for path in shuffled_paths for value in file_rows[path]]
    assert expected != input_list
    assert result == expected, (
        "Shuffle did not produce expected order.\n"
        f"Expected: {expected}\n"
        f"Got: {result}\n"
        f"Shuffled paths: {shuffled_paths}"
    )


@pytest.mark.parametrize(
    "fs,data_path",
    [
        (None, lazy_fixture("local_path")),
        (lazy_fixture("local_fs"), lazy_fixture("local_path")),

Review comment (bug):
This makes strong assumptions about the implementation. For example, if we don't use numpy for shuffling, I think this could be incorrect.

I think it also makes the test harder to read. This logic is non-trivial

Code after:
        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"),
        ),
    ],
)
def test_parquet_read_partitioned(
    ray_start_regular_shared, fs, data_path, target_max_block_size_infinite_or_default
):
    df = pd.DataFrame(
        {"one": [1, 1, 1, 3, 3, 3], "two": ["a", "b", "c", "e", "f", "g"]}
    )
    table = pa.Table.from_pandas(df)
    pq.write_to_dataset(
        table,
        root_path=_unwrap_protocol(data_path),
        partition_cols=["one"],
        filesystem=fs,
    )

    ds = ray.data.read_parquet(data_path, filesystem=fs)

    # Test metadata-only parquet ops.
    assert ds.count() == 6
    assert ds.size_bytes() > 0
    # Schema information and input files are available from Parquet metadata,
    # so we do not need to compute the first block.
    assert ds.schema() == Schema(pa.schema({"two": pa.string(), "one": pa.string()}))
    input_files = ds.input_files()
    assert len(input_files) == 2, input_files