wehub-resource-sync
74 lines
2.3 KiB
Python
74 lines
2.3 KiB
Python
import pytest
|
|
|
|
import ray
|
|
from ray.data.context import DataContext
|
|
from ray.tests.conftest import * # noqa
|
|
|
|
|
|
class TestPreserveHashShuffleBlocks:
|
|
"""Test that hash shuffle repartition preserves block structure."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def setup(self):
|
|
"""Setup test context with small target_max_block_size."""
|
|
ctx = DataContext.get_current()
|
|
# Very small to force splitting if enabled
|
|
ctx.target_max_block_size = 1
|
|
yield
|
|
|
|
def test_repartition_preserves_blocks(
|
|
self, ray_start_regular_shared_2_cpus, shutdown_only
|
|
):
|
|
"""Test that repartition with keys preserves block count."""
|
|
# Create a dataset with multiple blocks
|
|
ds = ray.data.range(10, override_num_blocks=10)
|
|
|
|
# Repartition using hash shuffle with keys
|
|
result = ds.repartition(2, keys=["id"]).materialize()
|
|
|
|
# Should have exactly 2 blocks (one per partition)
|
|
assert result.num_blocks() == 2
|
|
|
|
def test_map_groups_works(self, ray_start_regular_shared_2_cpus, shutdown_only):
|
|
"""Test that map_groups works correctly."""
|
|
ds = ray.data.from_items(
|
|
[
|
|
{"group": 1, "value": 1},
|
|
{"group": 1, "value": 2},
|
|
{"group": 2, "value": 3},
|
|
{"group": 2, "value": 4},
|
|
]
|
|
)
|
|
|
|
# map_groups should work correctly
|
|
result = (
|
|
ds.groupby("group")
|
|
.map_groups(lambda g: {"count": [len(g["value"])]})
|
|
.take_all()
|
|
)
|
|
|
|
assert len(result) == 2
|
|
counts = sorted([r["count"] for r in result])
|
|
assert counts == [2, 2]
|
|
|
|
def test_join_does_not_preserve_blocks(
|
|
self, ray_start_regular_shared_2_cpus, shutdown_only
|
|
):
|
|
"""Test that join does not preserve block structure (default behavior)."""
|
|
# Create a dataset with one large block
|
|
ds = ray.data.range(10, override_num_blocks=2)
|
|
|
|
# Join operation uses hash shuffle but doesn't set disallow_block_splitting
|
|
result = ds.join(
|
|
ds, on=("id",), join_type="inner", num_partitions=2
|
|
).materialize()
|
|
|
|
# Should have more than 2 blocks due to block splitting
|
|
assert result.num_blocks() >= 3
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|