wehub-resource-sync
284 lines
7.8 KiB
Python
284 lines
7.8 KiB
Python
import pyarrow as pa
|
|
import pyarrow.compute as pc
|
|
|
|
import ray
|
|
from ray.data._internal.compute import ActorPoolStrategy
|
|
from ray.data._internal.logical.interfaces import LogicalPlan
|
|
from ray.data._internal.logical.operators import (
|
|
CSE_TEMP_COLUMN_PREFIX,
|
|
Project,
|
|
)
|
|
from ray.data._internal.logical.operators.input_data_operator import InputData
|
|
from ray.data._internal.logical.optimizers import LogicalOptimizer
|
|
from ray.data._internal.logical.rules import CommonSubExprElimination
|
|
from ray.data.datatype import DataType
|
|
from ray.data.expressions import (
|
|
AliasExpr,
|
|
BinaryExpr,
|
|
PyArrowComputeUDFExpr,
|
|
UDFExpr,
|
|
col,
|
|
udf,
|
|
)
|
|
|
|
|
|
def _input_op():
|
|
return InputData(input_data=[])
|
|
|
|
|
|
def _apply_cse(project: Project) -> Project:
|
|
plan = LogicalPlan(project, ray.data.DataContext.get_current())
|
|
optimized = CommonSubExprElimination().apply(plan)
|
|
assert isinstance(optimized.dag, Project)
|
|
return optimized.dag
|
|
|
|
|
|
def _temp_names(project: Project) -> list[str]:
|
|
return [expr.name for expr in project.get_common_sub_exprs()]
|
|
|
|
|
|
@udf(return_dtype=DataType.int64())
|
|
def add_one(x: pa.Array) -> pa.Array:
|
|
return pc.add(x, 1)
|
|
|
|
|
|
def _unwrap_alias(expr):
|
|
return expr.expr if isinstance(expr, AliasExpr) else expr
|
|
|
|
|
|
def test_repeated_udf_in_one_project():
|
|
subexpr = add_one(col("a"))
|
|
project = Project(
|
|
exprs=[(subexpr + subexpr + subexpr).alias("y")],
|
|
input_dependencies=[_input_op()],
|
|
)
|
|
|
|
optimized = _apply_cse(project)
|
|
|
|
assert len(optimized.get_common_sub_exprs()) == 1
|
|
temp_name = _temp_names(optimized)[0]
|
|
assert temp_name.startswith(CSE_TEMP_COLUMN_PREFIX)
|
|
assert isinstance(_unwrap_alias(optimized.get_common_sub_exprs()[0]), UDFExpr)
|
|
assert isinstance(optimized.exprs[0], AliasExpr)
|
|
assert optimized.exprs[0].name == "y"
|
|
assert temp_name in repr(optimized.exprs[0])
|
|
|
|
|
|
def test_structurally_equal_separately_constructed_udf_calls():
|
|
project = Project(
|
|
exprs=[(add_one(col("a")) + add_one(col("a"))).alias("y")],
|
|
input_dependencies=[_input_op()],
|
|
)
|
|
|
|
optimized = _apply_cse(project)
|
|
|
|
assert len(optimized.get_common_sub_exprs()) == 1
|
|
|
|
|
|
def test_nested_common_expressions_materialize_bottom_up():
|
|
leaf_1 = add_one(col("a"))
|
|
leaf_2 = add_one(col("a"))
|
|
parent_1 = leaf_1 + leaf_2
|
|
|
|
leaf_3 = add_one(col("a"))
|
|
leaf_4 = add_one(col("a"))
|
|
parent_2 = leaf_3 + leaf_4
|
|
|
|
project = Project(
|
|
exprs=[(parent_1 + parent_2).alias("y")],
|
|
input_dependencies=[_input_op()],
|
|
)
|
|
|
|
optimized = _apply_cse(project)
|
|
|
|
common_exprs = optimized.get_common_sub_exprs()
|
|
assert len(common_exprs) == 2
|
|
first_temp, second_temp = _temp_names(optimized)
|
|
assert isinstance(_unwrap_alias(common_exprs[0]), UDFExpr)
|
|
assert first_temp in repr(common_exprs[1])
|
|
assert second_temp in repr(optimized.exprs[0])
|
|
|
|
|
|
def test_alias_root_is_ignored_but_child_is_extracted():
|
|
project = Project(
|
|
exprs=[
|
|
add_one(col("a")).alias("x"),
|
|
add_one(col("a")).alias("y"),
|
|
],
|
|
input_dependencies=[_input_op()],
|
|
)
|
|
|
|
optimized = _apply_cse(project)
|
|
|
|
assert len(optimized.get_common_sub_exprs()) == 1
|
|
assert isinstance(_unwrap_alias(optimized.get_common_sub_exprs()[0]), UDFExpr)
|
|
assert [expr.name for expr in optimized.exprs] == ["x", "y"]
|
|
|
|
|
|
def test_columns_and_literals_alone_are_not_materialized():
|
|
project = Project(
|
|
exprs=[
|
|
col("a").alias("a1"),
|
|
col("a").alias("a2"),
|
|
(col("b") + 1).alias("b1"),
|
|
(col("b") + 1).alias("b2"),
|
|
],
|
|
input_dependencies=[_input_op()],
|
|
)
|
|
|
|
optimized = _apply_cse(project)
|
|
|
|
assert len(optimized.get_common_sub_exprs()) == 1
|
|
common_inner = _unwrap_alias(optimized.get_common_sub_exprs()[0])
|
|
assert isinstance(common_inner, BinaryExpr)
|
|
|
|
|
|
def test_pyarrow_compute_udf_reuse():
|
|
project = Project(
|
|
exprs=[(col("a").abs() + col("a").abs()).alias("y")],
|
|
input_dependencies=[_input_op()],
|
|
)
|
|
|
|
optimized = _apply_cse(project)
|
|
|
|
assert len(optimized.get_common_sub_exprs()) == 1
|
|
common_inner = _unwrap_alias(optimized.get_common_sub_exprs()[0])
|
|
assert isinstance(common_inner, PyArrowComputeUDFExpr)
|
|
|
|
|
|
def test_output_schema_uses_visible_expressions_only(ray_start_regular_shared_2_cpus):
|
|
ds = ray.data.from_arrow(pa.table({"a": [1, 2]}))
|
|
ds = ds.with_column("x", add_one(col("a")))
|
|
ds = ds.with_column("y", col("x") + col("x"))
|
|
ds = ds.select_columns(["y"])
|
|
|
|
optimized = LogicalOptimizer().optimize(ds._logical_plan)
|
|
project = optimized.dag
|
|
assert isinstance(project, Project)
|
|
assert project.get_common_sub_exprs()
|
|
assert project.infer_schema().names == ["y"]
|
|
assert all(
|
|
not name.startswith(CSE_TEMP_COLUMN_PREFIX)
|
|
for name in project.infer_schema().names
|
|
)
|
|
|
|
|
|
def test_compute_strategy_is_preserved():
|
|
@udf(return_dtype=DataType.int64())
|
|
class AddOffset:
|
|
def __init__(self, offset):
|
|
self.offset = offset
|
|
|
|
def __call__(self, x: pa.Array) -> pa.Array:
|
|
return pc.add(x, self.offset)
|
|
|
|
add_ten = AddOffset(10)
|
|
subexpr = add_ten(col("a"))
|
|
project = Project(
|
|
exprs=[(subexpr + subexpr).alias("y")],
|
|
input_dependencies=[_input_op()],
|
|
)
|
|
|
|
optimized = _apply_cse(project)
|
|
|
|
assert isinstance(optimized.compute, ActorPoolStrategy)
|
|
|
|
|
|
def test_cse_rule_is_idempotent():
|
|
project = Project(
|
|
exprs=[(add_one(col("a")) + add_one(col("a"))).alias("y")],
|
|
input_dependencies=[_input_op()],
|
|
)
|
|
|
|
once = _apply_cse(project)
|
|
twice = CommonSubExprElimination._try_optimize_project(once)
|
|
|
|
assert twice is once
|
|
|
|
|
|
def test_logical_post_optimize_cse_executes_without_pushing_temps_into_read(
|
|
ray_start_regular_shared_2_cpus,
|
|
):
|
|
ds = ray.data.from_arrow(pa.table({"a": [1, 2]}))
|
|
ds = ds.with_column("x", add_one(col("a")))
|
|
ds = ds.with_column("y", col("x") + col("x"))
|
|
ds = ds.select_columns(["y"])
|
|
|
|
optimized = LogicalOptimizer().optimize(ds._logical_plan)
|
|
project = optimized.dag
|
|
assert isinstance(project, Project)
|
|
assert project.get_common_sub_exprs()
|
|
assert ds.take_all() == [{"y": 4}, {"y": 6}]
|
|
|
|
|
|
def test_udf_is_called_once_per_block(ray_start_regular_shared_2_cpus):
|
|
@ray.remote
|
|
class Counter:
|
|
def __init__(self):
|
|
self.count = 0
|
|
|
|
def inc(self, n):
|
|
self.count += n
|
|
|
|
def get(self):
|
|
return self.count
|
|
|
|
counter = Counter.remote()
|
|
|
|
@udf(return_dtype=DataType.int64())
|
|
def counted_add_one(x: pa.Array) -> pa.Array:
|
|
ray.get(counter.inc.remote(1))
|
|
return pc.add(x, 1)
|
|
|
|
expr = counted_add_one(col("id"))
|
|
ds = ray.data.range(6, override_num_blocks=3)
|
|
ds = ds.with_column("x", expr)
|
|
ds = ds.with_column("y", col("x") + col("x") + col("x"))
|
|
ds = ds.select_columns(["y"])
|
|
|
|
assert ds.take_all() == [
|
|
{"y": 3},
|
|
{"y": 6},
|
|
{"y": 9},
|
|
{"y": 12},
|
|
{"y": 15},
|
|
{"y": 18},
|
|
]
|
|
assert ray.get(counter.get.remote()) == 3
|
|
|
|
|
|
def test_callable_class_udf_still_initializes(ray_start_regular_shared_2_cpus):
|
|
@udf(return_dtype=DataType.int64())
|
|
class AddOffset:
|
|
def __init__(self, offset):
|
|
self.offset = offset
|
|
|
|
def __call__(self, x: pa.Array) -> pa.Array:
|
|
return pc.add(x, self.offset)
|
|
|
|
add_ten = AddOffset(10)
|
|
expr = add_ten(col("id"))
|
|
ds = ray.data.range(4, override_num_blocks=2)
|
|
ds = ds.with_column("x", expr)
|
|
ds = ds.with_column("y", col("x") + col("x"))
|
|
ds = ds.select_columns(["y"])
|
|
|
|
rows = ds.take_all()
|
|
assert rows == [
|
|
{"y": 20},
|
|
{"y": 22},
|
|
{"y": 24},
|
|
{"y": 26},
|
|
]
|
|
assert all(
|
|
not any(key.startswith(CSE_TEMP_COLUMN_PREFIX) for key in row) for row in rows
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|