[Data] Fix DAG inconsistency in Operator._apply_transform

[veeceey]

Closes #57825

When _apply_transform is called and only some input dependencies are transformed, the resulting DAG becomes inconsistent. Unchanged operators still point to the original parent in their output_dependencies, while the new (shallow-copied) parent lists them in input_dependencies. This means you can't reliably traverse the DAG bidirectionally.

The root cause is on the old line 85: target._wire_output_deps(new_ops) only wires new operators to the copy, but unchanged operators still reference the old parent.

The fix does three things when a copy is needed:

1. Creates a fresh output_dependencies list on the copy (avoids sharing the list with the original)
2. Removes stale output_dependencies references from unchanged input ops
3. Wires all transformed input ops (not just new ones) to the new target

This is a fresh implementation based on the analysis in #57825 and the stalled PR #57887.

Test plan:

• Partial transform: only some inputs change, verify bidirectional consistency
• Identity transform: no changes, verify same object returned
• Diamond DAG: shared dependency, verify correct wiring
• Chain transform: cascading copies through a linear DAG

[veeceey]

Test results

Ran standalone tests (ray doesn't have Python 3.14 wheels yet, so tested by direct module loading):

Test 1: Partial transformation DAG consistency
  PASSED
Test 2: Identity transform
  PASSED
Test 3: Diamond DAG
  PASSED
Test 4: Chain transform
  PASSED

ALL TESTS PASSED

Test 1 is the exact reproduction from the issue:

a = Operator("A", [])
b = Operator("B", [])
c = Operator("C", [a, b])

def transform_b(op):
    if op.name == "B":
        return Operator("Transformed B", op.input_dependencies)
    return op

c_transformed = c._apply_transform(transform_b)

# Before fix: FAILS (A points to old C)
# After fix: PASSES
assert c_transformed is c_transformed.input_dependencies[0].output_dependencies[0]

Also added test_apply_transform_dag_consistency and test_apply_transform_no_change to the test suite.

[gemini-code-assist]

Code Review

The pull request effectively addresses the DAG inconsistency issue in Operator._apply_transform by ensuring bidirectional consistency when only some input dependencies are transformed. The changes correctly handle the shallow copy, fresh output dependencies, and removal of stale references, which are crucial for maintaining the integrity of the DAG. The added test cases provide good coverage for the identified scenarios, including partial transform, identity transform, and diamond DAGs, which helps validate the fix.

[gemini-code-assist]

The loop iterates over self.input_dependencies to remove stale references. However, self.input_dependencies is the original list of input operators. If transformed_input_op is a new operator, self might not be in its _output_dependencies list, leading to incorrect behavior or missed updates for the new operators. The intent is to update the output_dependencies of the original input operators that were not transformed.

Suggested change

	for input_op in self.input_dependencies:
	if self in input_op._output_dependencies:
	input_op._output_dependencies.remove(self)
	for i, input_op in enumerate(self.input_dependencies):
	if transformed_input_ops[i] is input_op and self in input_op._output_dependencies:
	input_op._output_dependencies.remove(self)

[gemini-code-assist]

The variable input_changed is initialized to False but is not used until line 77. It's good practice to initialize variables closer to their first use to improve readability and reduce potential for errors, especially in larger functions.

Suggested change

	input_changed = False
	transformed_input_ops = []
	for input_op in self.input_dependencies:
	transformed_input_op = input_op._apply_transform(transform)
	transformed_input_ops.append(transformed_input_op)
	# Keep track of whether any input ops changed
	if transformed_input_op is not input_op:
	input_changed = True
	input_changed = False

[peterxcli]

It appears that the only difference between this patch and #57887 is this specific line. Could you clarify the exact cases you are aiming to handle here? Providing an example test case would be really helpful for understanding the context better. Thanks!

[veeceey]

@peterxcli Good question — the two PRs tackle the same root issue (#57825) but the fix here handles one additional detail that #57887 misses.

Both PRs share the same core change: tracking input_changed instead of new_ops, calling _wire_output_deps(transformed_input_ops) with all transformed inputs (not just the new ones), and removing stale references from unchanged inputs' output_dependencies.

The difference is this line in my PR:

target._output_dependencies = []

Without it, copy.copy(self) produces a shallow copy where target._output_dependencies is the same list object as self._output_dependencies. That means any downstream operator that already references self through that list will see mutations on the copy's list and vice versa — they share identity.

Here's a concrete example where this matters:

a = Operator("A", [])
b = Operator("B", [a])
c = Operator("C", [b])   # c is in b's output_dependencies

def transform_a(op):
    if op.name == "A":
        return Operator("A2", [])
    return op

b2 = b._apply_transform(transform_a)
# b2 is a shallow copy of b
# Without `_output_dependencies = []`, b2._output_dependencies IS b._output_dependencies
# So b2 still appears to have c as a downstream, even though c was never rewired to b2.
# Worse, if anything later appends to b2._output_dependencies, it also modifies b's list.

In #57887, _purge_references cleans up the input side correctly, but the shared _output_dependencies list on the copy means the output side can still have stale or cross-contaminated references. Reinitializing it to an empty list on the copy ensures the new operator starts clean and only accumulates output deps that are explicitly wired to it.

I also added a test_apply_transform_no_change test case to cover the identity-transform path, which #57887 doesn't test.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

[veeceey]

gentle ping on this one!

[veeceey]

@peterxcli @cursor[bot] thanks for the feedback — just pushed a fix addressing both points.

Re: _output_dependencies not initialized on base Operator — good catch. I added _output_dependencies: List["Operator"] = [] to Operator.__init__, plus an output_dependencies property and a _wire_output_deps method on the base class. This way the stale-reference cleanup in _apply_transform works for any Operator subclass, not just PhysicalOperator (which has its own override and never hits this code path, but LogicalOperator delegates to it via super()).

Re: concrete test case showing the difference from #57887 — added test_apply_transform_copy_isolates_output_dependencies. It sets up a chain A -> B -> C with output deps wired, then transforms A. Without the target._output_dependencies = [] line, copy.copy(self) produces a shallow copy where b_transformed._output_dependencies is literally the same list object as b._output_dependencies — so the copy still appears to have C as a downstream, even though C was never rewired to it. The test asserts that the copy's output_dependencies don't leak stale references from the original DAG:

a = Operator("A", [])
b = Operator("B", [a])
c = Operator("C", [b])
a._output_dependencies = [b]
b._output_dependencies = [c]

def transform_a(op):
    if op.name == "A":
        return Operator("A2", [])
    return op

b_transformed = b._apply_transform(transform_a)
# Without target._output_dependencies = [], this fails:
# b_transformed._output_dependencies is b._output_dependencies, still containing [c]
assert c not in b_transformed._output_dependencies

PR #57887's _purge_references cleans the input side correctly, but doesn't address this shared-list issue on the output side.

[github-actions]

This pull request has been automatically marked as stale because it has not had
any activity for 14 days. It will be closed in another 14 days if no further activity occurs.
Thank you for your contributions.

You can always ask for help on our discussion forum or Ray's public slack channel.

If you'd like to keep this open, just leave any comment, and the stale label will be removed.

[veeceey]

hey, this is still relevant and ready for review -- removing the stale label. the fix handles an edge case that #57887 doesn't cover (stale _output_dependencies). happy to rebase if needed.