import sys from typing import TYPE_CHECKING, List, Optional import numpy as np import pandas as pd import pytest import ray if TYPE_CHECKING: from ray.data.context import DataContext from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer from ray.data._internal.execution.operators.map_operator import MapOperator from ray.data._internal.execution.operators.task_pool_map_operator import ( TaskPoolMapOperator, ) from ray.data._internal.logical.interfaces import LogicalPlan from ray.data._internal.logical.operators import ( Filter, FlatMap, FromArrow, FromItems, FromNumpy, FromPandas, MapBatches, MapRows, Project, ) from ray.data._internal.logical.optimizers import PhysicalOptimizer from ray.data._internal.planner import create_planner from ray.data.block import BlockMetadata from ray.data.context import DataContext from ray.data.datasource import Datasource from ray.data.datasource.datasource import ReadTask from ray.data.expressions import col from ray.data.tests.conftest import * # noqa from ray.data.tests.test_util import _check_usage_record, get_parquet_read_logical_op from ray.data.tests.util import column_udf, extract_values, named_values from ray.tests.conftest import * # noqa def test_read_operator(ray_start_regular_shared_2_cpus): ctx = DataContext.get_current() planner = create_planner() op = get_parquet_read_logical_op() plan = LogicalPlan(op, ctx) physical_plan, _ = planner.plan(plan) physical_op = physical_plan.dag assert op.name == "ReadParquet" assert isinstance(physical_op, MapOperator) assert len(physical_op.input_dependencies) == 1 assert isinstance(physical_op.input_dependencies[0], InputDataBuffer) # Check that the linked logical operator is the same the input op. assert physical_op._logical_operators == [op] assert physical_op.input_dependencies[0]._logical_operators == [op] def test_read_operator_emits_warning_for_large_read_tasks(): class StubDatasource(Datasource): def estimate_inmemory_data_size(self) -> Optional[int]: return None def get_read_tasks( self, parallelism: int, per_task_row_limit: Optional[int] = None, data_context: Optional["DataContext"] = None, ) -> List[ReadTask]: large_object = np.zeros((128, 1024, 1024), dtype=np.uint8) # 128 MiB def read_fn(): _ = large_object yield pd.DataFrame({"column": [0]}) return [ ReadTask( read_fn, BlockMetadata(1, None, None, None), per_task_row_limit=per_task_row_limit, ) ] with pytest.warns(UserWarning): ray.data.read_datasource(StubDatasource()).materialize() def test_split_blocks_operator(ray_start_regular_shared_2_cpus): ctx = DataContext.get_current() planner = create_planner() op = get_parquet_read_logical_op(parallelism=10) logical_plan = LogicalPlan(op, ctx) physical_plan, _ = planner.plan(logical_plan) physical_plan = PhysicalOptimizer().optimize(physical_plan) physical_op = physical_plan.dag assert physical_op.name == "ReadParquet->SplitBlocks(10)" assert isinstance(physical_op, MapOperator) assert len(physical_op.input_dependencies) == 1 assert isinstance(physical_op.input_dependencies[0], InputDataBuffer) assert physical_op._additional_split_factor == 10 # Test that split blocks prevents fusion. op = MapBatches( lambda x: x, input_dependencies=[op], ) logical_plan = LogicalPlan(op, ctx) physical_plan, _ = planner.plan(logical_plan) physical_plan = PhysicalOptimizer().optimize(physical_plan) physical_op = physical_plan.dag assert physical_op.name == "MapBatches()" assert len(physical_op.input_dependencies) == 1 up_physical_op = physical_op.input_dependencies[0] assert isinstance(up_physical_op, MapOperator) assert up_physical_op.name == "ReadParquet->SplitBlocks(10)" def test_from_operators(ray_start_regular_shared_2_cpus): ctx = DataContext.get_current() op_classes = [ FromArrow, FromItems, FromNumpy, FromPandas, ] for op_cls in op_classes: planner = create_planner() op = op_cls([], []) plan = LogicalPlan(op, ctx) physical_plan, _ = planner.plan(plan) physical_op = physical_plan.dag assert op.name == op_cls.__name__ assert isinstance(physical_op, InputDataBuffer) assert len(physical_op.input_dependencies) == 0 # Check that the linked logical operator is the same the input op. assert physical_op._logical_operators == [op] def test_from_items_e2e(ray_start_regular_shared_2_cpus): data = ["Hello", "World"] ds = ray.data.from_items(data) assert ds.take_all() == named_values("item", data), ds # Check that metadata fetch is included in stats. assert "FromItems" in ds.stats() assert ds._logical_plan.dag.name == "FromItems" _check_usage_record(["FromItems"]) def test_map_operator_udf_name(ray_start_regular_shared_2_cpus): # Test the name of the Map operator with different types of UDF. def normal_function(x): return x lambda_function = lambda x: x # noqa: E731 class CallableClass: def __call__(self, x): return x class NormalClass: def method(self, x): return x udf_list = [ # A nomral function. normal_function, # A lambda function lambda_function, # A callable class. CallableClass, # An instance of a callable class. CallableClass(), # A normal class method. NormalClass().method, ] expected_names = [ "normal_function", "", "CallableClass", "CallableClass", "NormalClass.method", ] for udf, expected_name in zip(udf_list, expected_names): op = MapRows( udf, input_dependencies=[get_parquet_read_logical_op()], ) assert op.name == f"Map({expected_name})" def test_map_batches_operator(ray_start_regular_shared_2_cpus): ctx = DataContext.get_current() planner = create_planner() read_op = get_parquet_read_logical_op() op = MapBatches( lambda x: x, input_dependencies=[read_op], ) plan = LogicalPlan(op, ctx) physical_plan, _ = planner.plan(plan) physical_op = physical_plan.dag assert op.name == "MapBatches()" assert isinstance(physical_op, MapOperator) assert len(physical_op.input_dependencies) == 1 assert isinstance(physical_op.input_dependencies[0], MapOperator) # Check that the linked logical operator is the same the input op. assert physical_op._logical_operators == [op] def test_map_batches_e2e(ray_start_regular_shared_2_cpus): ds = ray.data.range(5) ds = ds.map_batches(column_udf("id", lambda x: x)) assert sorted(extract_values("id", ds.take_all())) == list(range(5)), ds _check_usage_record(["ReadRange", "MapBatches"]) def test_map_rows_operator(ray_start_regular_shared_2_cpus): ctx = DataContext.get_current() planner = create_planner() read_op = get_parquet_read_logical_op() op = MapRows( lambda x: x, input_dependencies=[read_op], ) plan = LogicalPlan(op, ctx) physical_plan, _ = planner.plan(plan) physical_op = physical_plan.dag assert op.name == "Map()" assert isinstance(physical_op, MapOperator) assert len(physical_op.input_dependencies) == 1 assert isinstance(physical_op.input_dependencies[0], MapOperator) def test_map_rows_e2e(ray_start_regular_shared_2_cpus): ds = ray.data.range(5) ds = ds.map(column_udf("id", lambda x: x + 1)) expected = [1, 2, 3, 4, 5] actual = sorted(extract_values("id", ds.take_all())) assert actual == expected, f"Expected {expected}, but got {actual}" _check_usage_record(["ReadRange", "MapRows"]) def test_filter_operator(ray_start_regular_shared_2_cpus): ctx = DataContext.get_current() planner = create_planner() read_op = get_parquet_read_logical_op() op = Filter( fn=lambda x: x, input_dependencies=[read_op], ) plan = LogicalPlan(op, ctx) physical_plan, _ = planner.plan(plan) physical_op = physical_plan.dag assert op.name == "Filter()" assert isinstance(physical_op, MapOperator) assert len(physical_op.input_dependencies) == 1 assert isinstance(physical_op.input_dependencies[0], MapOperator) def test_filter_e2e(ray_start_regular_shared_2_cpus): ds = ray.data.range(5) ds = ds.filter(fn=lambda x: x["id"] % 2 == 0) assert sorted(extract_values("id", ds.take_all())) == [0, 2, 4], ds _check_usage_record(["ReadRange", "Filter"]) def test_project_operator_select(ray_start_regular_shared_2_cpus): """ Checks that the physical plan is properly generated for the Project operator from select columns. """ path = "example://iris.parquet" ds = ray.data.read_parquet(path) ds = ds.map_batches(lambda d: d) cols = ["sepal.length", "petal.width"] ds = ds.select_columns(cols) logical_plan = ds._logical_plan op = logical_plan.dag assert isinstance(op, Project), op.name assert op.exprs == [col("sepal.length"), col("petal.width")] physical_plan, _ = create_planner().plan(logical_plan) physical_plan = PhysicalOptimizer().optimize(physical_plan) physical_op = physical_plan.dag assert isinstance(physical_op, TaskPoolMapOperator) assert isinstance(physical_op.input_dependency, TaskPoolMapOperator) def test_project_operator_rename(ray_start_regular_shared_2_cpus): """ Checks that the physical plan is properly generated for the Project operator from rename columns. """ from ray.data.expressions import star path = "example://iris.parquet" ds = ray.data.read_parquet(path) ds = ds.map_batches(lambda d: d) cols_rename = {"sepal.length": "sepal_length", "petal.width": "pedal_width"} ds = ds.rename_columns(cols_rename) logical_plan = ds._logical_plan op = logical_plan.dag assert isinstance(op, Project), op.name assert op.exprs == [ star(), col("sepal.length").alias("sepal_length"), col("petal.width").alias("pedal_width"), ] physical_plan, _ = create_planner().plan(logical_plan) physical_plan = PhysicalOptimizer().optimize(physical_plan) physical_op = physical_plan.dag assert isinstance(physical_op, TaskPoolMapOperator) assert isinstance(physical_op.input_dependency, TaskPoolMapOperator) def test_flat_map(ray_start_regular_shared_2_cpus): ctx = DataContext.get_current() planner = create_planner() read_op = get_parquet_read_logical_op() op = FlatMap( lambda x: x, input_dependencies=[read_op], ) plan = LogicalPlan(op, ctx) physical_plan, _ = planner.plan(plan) physical_op = physical_plan.dag assert op.name == "FlatMap()" assert isinstance(physical_op, MapOperator) assert len(physical_op.input_dependencies) == 1 assert isinstance(physical_op.input_dependencies[0], MapOperator) def test_flat_map_e2e(ray_start_regular_shared_2_cpus): ds = ray.data.range(2) ds = ds.flat_map(fn=lambda x: [{"id": x["id"]}, {"id": x["id"]}]) assert extract_values("id", ds.take_all()) == [0, 0, 1, 1], ds _check_usage_record(["ReadRange", "FlatMap"]) def test_column_ops_e2e(ray_start_regular_shared_2_cpus): ds = ray.data.range(2) ds = ds.add_column(fn=lambda df: df.iloc[:, 0], col="new_col") assert ds.take_all() == [{"id": 0, "new_col": 0}, {"id": 1, "new_col": 1}], ds _check_usage_record(["ReadRange", "MapBatches"]) select_ds = ds.select_columns(cols=["new_col"]) assert select_ds.take_all() == [{"new_col": 0}, {"new_col": 1}] _check_usage_record(["ReadRange", "MapBatches"]) ds = ds.drop_columns(cols=["new_col"]) assert ds.take_all() == [{"id": 0}, {"id": 1}], ds _check_usage_record(["ReadRange", "MapBatches"]) def test_random_sample_e2e(ray_start_regular_shared_2_cpus): import math def ensure_sample_size_close(dataset, sample_percent=0.5): r1 = ds.random_sample(sample_percent) assert math.isclose( r1.count(), int(ds.count() * sample_percent), rel_tol=2, abs_tol=2 ) ds = ray.data.range(10, override_num_blocks=2) ensure_sample_size_close(ds) ds = ray.data.range(10, override_num_blocks=2) ensure_sample_size_close(ds) ds = ray.data.range_tensor(5, override_num_blocks=2, shape=(2, 2)) ensure_sample_size_close(ds) _check_usage_record(["ReadRange", "MapBatches"]) if __name__ == "__main__": sys.exit(pytest.main(["-v", __file__]))