diff --git a/amber/src/main/python/core/models/operator.py b/amber/src/main/python/core/models/operator.py index 4e65fbb2ab2..acdc75040a8 100644 --- a/amber/src/main/python/core/models/operator.py +++ b/amber/src/main/python/core/models/operator.py @@ -366,11 +366,19 @@ class LoopStartOperator(TableOperator): @overrides.final def process_state(self, state: State, port: int) -> Optional[State]: - # First-entry only: merge upstream state into self.state. The nested - # pass-through (a frame already stamped with a LoopStartId) and all - # loop_counter bookkeeping are owned by the worker runtime - # (main_loop._process_state_frame), so this operator never sees the - # counter and never mutates the State it is handed. + # First-entry only: merge non-conflicting upstream state into self.state. + # The nested pass-through (a frame already stamped with a LoopStartId) and + # all loop_counter bookkeeping are owned by the worker runtime + # (main_loop._process_state_frame), so this operator never sees the counter + # and never mutates the State it is handed. + collisions = self.state.keys() & state.keys() + + if collisions: + raise ValueError( + f"Loop state variable(s) cannot be overwritten: " + f"{', '.join(sorted(collisions))}" + ) + self.state.update(state) return None diff --git a/amber/src/test/python/core/models/test_loop_operators.py b/amber/src/test/python/core/models/test_loop_operators.py index 6b496a5f17b..f9d8aea234e 100644 --- a/amber/src/test/python/core/models/test_loop_operators.py +++ b/amber/src/test/python/core/models/test_loop_operators.py @@ -134,12 +134,24 @@ def test_first_time_state_is_merged_into_self_state_and_none_is_returned(self): # nothing flows downstream of LoopStart until the table is in. op = _StubLoopStart() op.open() - op.state["i"] = 0 # simulate the user's initialization + op.state["i"] = 0 result = op.process_state(State({"upstream_key": "v"}), port=0) - assert result is None, "first-time state must not be forwarded" - assert op.state["upstream_key"] == "v", "state was not merged into self.state" + assert result is None + assert op.state["upstream_key"] == "v" + + def test_state_rejects_overwriting_loop_variable(self): + op = _StubLoopStart() + op.open() + + with pytest.raises( + ValueError, + match=r"Loop state variable\(s\) cannot be overwritten: i", + ): + op.process_state(State({"i": 999}), port=0) + + assert op.state["i"] == 0 # NOTE: LoopStart re-entry (+1) is owned by the worker runtime now, not the # operator (which only does the first-entry merge above). It and the nested @@ -313,8 +325,8 @@ def test_accumulator_persists_and_reserved_names_never_leak(self): # # Each pass of the while loop mimics one engine iteration: the # LoopStart region is re-executed (a fresh operator whose open() seeds - # the loop variables, the back-edge state overriding them, and the - # upstream table re-read), the produced state crosses the materialized + # the loop variables, the worker runtime restores the back-edge state, + # and the upstream table re-read), the produced state crosses the materialized # channel (a State to_tuple/from_tuple round-trip), the LoopEnd runs # the user update and evaluates the condition, and on continuation # only the user loop variables cross the back-edge. @@ -332,7 +344,7 @@ def test_accumulator_persists_and_reserved_names_never_leak(self): ) start.open() if back_edge is not None: - start.process_state(back_edge, port=0) + start.state.update(back_edge) for row in rows: list(start.process_tuple(row, port=0)) emitted.extend(o for o in start.on_finish(port=0) if o is not None)