chore: import upstream snapshot with attribution

python/ray/data/tests/test_common_subexpr_elimination.py

@@ -0,0 +1,283 @@
+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__]))