chore: import upstream snapshot with attribution

python/ray/data/tests/unit/test_logical_plan.py

@@ -0,0 +1,74 @@
+from ray.data._internal.logical.interfaces import LogicalOperator, LogicalPlan
+from ray.data.context import DataContext
+
+
+class DummyLogicalOperator(LogicalOperator):
+    def __init__(self, input_dependencies, name=None):
+        object.__setattr__(self, "_input_dependencies", input_dependencies)
+        if name is not None:
+            object.__setattr__(self, "_name", name)
+
+
+def test_sources_singleton():
+    ctx = DataContext.get_current()
+
+    source = DummyLogicalOperator([], name="source")
+    assert LogicalPlan(source, ctx).sources() == [source]
+
+
+def test_sources_chain():
+    ctx = DataContext.get_current()
+
+    source = DummyLogicalOperator([], name="source")
+    sink = DummyLogicalOperator([source], name="sink")
+    assert LogicalPlan(sink, ctx).sources() == [source]
+
+
+def test_sources_multiple_sources():
+    ctx = DataContext.get_current()
+
+    source1 = DummyLogicalOperator([], name="source1")
+    source2 = DummyLogicalOperator([], name="source2")
+    sink = DummyLogicalOperator([source1, source2], name="sink")
+    assert LogicalPlan(sink, ctx).sources() == [source1, source2]
+
+
+def test_logical_operator_defaults_name_to_class_name():
+    op = DummyLogicalOperator([])
+    assert op.name == "DummyLogicalOperator"
+    assert op.dag_str == "DummyLogicalOperator[DummyLogicalOperator]"
+
+
+def test_logical_operator_does_not_track_output_dependencies():
+    source = DummyLogicalOperator([], name="source")
+    sink = DummyLogicalOperator([source], name="sink")
+
+    # Logical operators should not maintain reverse output-dependency state.
+    assert not hasattr(source, "_output_dependencies")
+    assert not hasattr(sink, "_output_dependencies")
+
+    transformed = sink._apply_transform(lambda op: op)
+    assert transformed is sink
+    assert not hasattr(source, "_output_dependencies")
+    assert not hasattr(sink, "_output_dependencies")
+
+
+def test_logical_operator_transform_supports_custom_subclasses():
+    source = DummyLogicalOperator([], name="source")
+    replacement = DummyLogicalOperator([], name="replacement")
+    sink = DummyLogicalOperator([source], name="sink")
+
+    transformed = sink._apply_transform(lambda op: replacement if op is source else op)
+
+    assert transformed is not sink
+    assert transformed.name == "sink"
+    assert transformed.input_dependencies == [replacement]
+    assert sink.input_dependencies == [source]
+
+
+if __name__ == "__main__":
+    import sys
+
+    import pytest
+
+    sys.exit(pytest.main(["-v", __file__]))