Repository: ray-project/ray
File: python/ray/data/_internal/execution/interfaces/physical_operator.py (Python)

Code before:
            target = self._copy_for_transform()
            for input_op in self.input_dependencies:
                assert isinstance(input_op, PhysicalOperator), input_op
                input_op._output_dependencies = [
                    dep for dep in input_op._output_dependencies if dep is not self
                ]
            target._input_dependencies = transformed_input_ops
            target._wire_output_deps(transformed_input_ops)
        else:
            target = self

        # 3) Apply transform on the current node itself. If transform replaces
        # the current node, rewire reverse dependencies from old node inputs
        # to the returned replacement node inputs.
        transformed_target = transform(target)
        if transformed_target is not target:
            for input_op in target.input_dependencies:
                assert isinstance(input_op, PhysicalOperator), input_op
                input_op._output_dependencies = [
                    dep for dep in input_op._output_dependencies if dep is not target
                ]
            transformed_target._rewire_output_deps(
                target, transformed_target.input_dependencies
            )
        return transformed_target

    def _copy_for_transform(self) -> "PhysicalOperator":
        target = object.__new__(type(self))
        target.__dict__ = self.__dict__.copy()
        # The copied node belongs to a new transformed DAG. Reverse edges are
        # rewired by parents, so avoid carrying stale downstream references.
        target._output_dependencies = []
        return target

    def _rewire_output_deps(
        self,
        source_op: "PhysicalOperator",
        input_dependencies: List["PhysicalOperator"],
    ) -> None:
        for input_op in input_dependencies:
            assert isinstance(input_op, PhysicalOperator), input_op
            input_op._output_dependencies = [
                dep for dep in input_op._output_dependencies if dep is not source_op
            ]
            if self not in input_op._output_dependencies:
                input_op._output_dependencies.append(self)

    def _wire_output_deps(self, input_dependencies: List["PhysicalOperator"]) -> None:
        for input_op in input_dependencies:
            assert isinstance(input_op, PhysicalOperator), input_op
            input_op._output_dependencies.append(self)

Review comment (suggestion):
```suggestion
        if transformed_target is not target:
            for input_op in target.input_dependencies:
                assert isinstance(input_op, PhysicalOperator), input_op
                input_op._output_dependencies = [
                    dep for dep in input_op._output_dependencies if dep is not target
                ]
            transformed_target._rewire_output_deps(
                target, transformed_target.input_dependencies
            )
        else:
            assert (
                transformed_target.input_dependencies == original_inputs
            ), "In-place input mutation is not supported; return a new node instead."    
```

Code after:
            target = self._copy_for_transform()
            for input_op in self.input_dependencies:
                assert isinstance(input_op, PhysicalOperator), input_op
                input_op._output_dependencies = [
                    dep for dep in input_op._output_dependencies if dep is not self
                ]
            target._input_dependencies = transformed_input_ops
            target._wire_output_deps(transformed_input_ops)
        else:
            target = self

        # 3) Apply transform on the current node itself. If transform replaces
        # the current node, rewire reverse dependencies from old node inputs
        # to the returned replacement node inputs.
        # Returning the same node must not mutate inputs in-place.
        original_inputs = tuple(target.input_dependencies)
        transformed_target = transform(target)
        if transformed_target is not target:
            for input_op in original_inputs:
                assert isinstance(input_op, PhysicalOperator), input_op
                input_op._output_dependencies = [
                    dep for dep in input_op._output_dependencies if dep is not target
                ]
            transformed_target._rewire_output_deps(
                target, transformed_target.input_dependencies
            )
        else:
            assert (
                tuple(transformed_target.input_dependencies) == original_inputs
            ), "In-place input mutation is not supported; return a new node instead."
        return transformed_target

    def _copy_for_transform(self) -> "PhysicalOperator":
        # copy.copy() is not safe here because PhysicalOperator.__reduce__()
        # intentionally raises. Use a side-effect-free shallow copy to avoid
        # re-running __init__ wiring/ID/metrics initialization during transform.
        target = object.__new__(type(self))
        target.__dict__ = self.__dict__.copy()
        # The transformed node should have a distinct identity and metrics owner.
        target._id = str(uuid.uuid4())
        target._metrics = OpRuntimeMetrics(target)
        # The copied node belongs to a new transformed DAG. Reverse edges are
        # rewired by parents, so avoid carrying stale downstream references.
        target._output_dependencies = []
        return target

    def _rewire_output_deps(
        self,
        source_op: "PhysicalOperator",
        input_dependencies: List["PhysicalOperator"],
    ) -> None: