wehub-resource-sync
38 lines
1.2 KiB
Python
38 lines
1.2 KiB
Python
import pytest
|
|
|
|
from ray.data._internal.execution.operators.base_physical_operator import (
|
|
AllToAllOperator,
|
|
)
|
|
from ray.data._internal.execution.operators.map_operator import MapOperator
|
|
from ray.data._internal.logical.interfaces import LogicalPlan
|
|
from ray.data._internal.logical.operators import RandomizeBlocks
|
|
from ray.data._internal.planner import create_planner
|
|
from ray.data._internal.random_config import RandomSeedConfig
|
|
from ray.data.context import DataContext
|
|
from ray.data.tests.test_util import get_parquet_read_logical_op
|
|
|
|
|
|
def test_randomize_blocks_operator(ray_start_regular_shared):
|
|
ctx = DataContext.get_current()
|
|
|
|
planner = create_planner()
|
|
read_op = get_parquet_read_logical_op()
|
|
op = RandomizeBlocks(
|
|
seed_config=RandomSeedConfig(seed=0),
|
|
input_dependencies=[read_op],
|
|
)
|
|
plan = LogicalPlan(op, ctx)
|
|
physical_plan, _ = planner.plan(plan)
|
|
physical_op = physical_plan.dag
|
|
|
|
assert op.name == "RandomizeBlockOrder"
|
|
assert isinstance(physical_op, AllToAllOperator)
|
|
assert len(physical_op.input_dependencies) == 1
|
|
assert isinstance(physical_op.input_dependencies[0], MapOperator)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|