This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-5900-322babf6205cf72277d13a4f2d16c5954a4dfbd9 in repository https://gitbox.apache.org/repos/asf/texera.git
commit 878eb8a38b720dfd1131fa6f04f863a813f2bb53 Author: Xinyuan Lin <[email protected]> AuthorDate: Wed Jul 1 16:06:22 2026 -0700 feat(amber): add loop-bookkeeping columns to materialized State (dormant) (#5900) ### What changes were proposed in this PR? Extends the cross-region **State materialization** format from a single `content` column to **3 columns** — `content`, `loop_counter`, `loop_start_id` — promoting loop bookkeeping to first-class columns (never inside the content JSON). The transport carries them end to end: the `OutputManager` state writer + `emit_state`, the Python network sender/receiver, the materialization reader, and the Scala `state.toTuple` call sites. In memory the two loop fields ride on the `StateFrame` envelope; they are materialized/serialized as their own columns (parallel to `content`), and `from_tuple` / `fromTuple` read only `content` back into the `State`. The loop-back write address (LoopStart's input-port URI) is **intentionally not** carried in State. It's constant per-execution config, not per-iteration data, so it will be delivered to Loop End at **setup** in the loop PR rather than round-tripping through every State row. (An earlier revision of this PR carried a `loop_start_state_uri` column; it was dropped after review — better than shipping a dormant column and removing it later.) On the Python side the column-name → value mapping is defined once in `State.to_columns` and reused by both `to_tuple` (iceberg) and the network sender's `StateFrame` branch, so adding a column later is a single-line change in one place rather than an edit in every serializer. **Dormant on `main`** — nothing observable changes without the loop operators: - `to_tuple()` / `toTuple()` and `OutputManager.save_state_to_storage_if_needed` / `emit_state` default the two loop columns to `0` / `""`, so every existing non-loop caller is unchanged. - `from_tuple` / `fromTuple` read only the `content` column, so round-trip identity is preserved and the extra columns are inert. No backward-compatible read of old 1-column State is needed: State materialization is **intra-execution only** — the iceberg state-document URI is execution-scoped (`…/eid/{executionId}/`) and recreated fresh each run, and State tuples are never replayed across executions or engine versions, so a 1-column tuple can never reach the 3-column reader. This is the state-format prerequisite the loop operators build on; the columns carry non-default values only once Loop Start/End set them (follow-up PR). ### Any related issues, documentation, discussions? Extracted from #5700 (loop operators) per @Xiao-zhen-Liu's split request; part of #4442 ("Introduce for loop"). ### How was this PR tested? - **Format / round-trip:** `test_state.py` (loop columns are their own columns, never in content JSON, default to `0` / `""`), Scala `StateSpec` (both loop columns round-trip through a tuple with non-default values, not just `content`), `ArrowUtilsSpec` (3-column Arrow vector round-trip), `IcebergDocumentSpec` (iceberg state-doc round-trip). - **Transport:** `test_network_receiver.py`, `test_input_port_materialization_reader_runnable.py`, and `test_state_materialization_e2e.py` — the e2e (hermetic sqlite catalog) writes non-default values for both loop columns end-to-end and asserts they replay both on the `StateFrame` and on the raw iceberg row, exercising the real Tuple/Schema/iceberg path. - **Dormancy:** `test_output_manager.py::test_defaults_loop_columns_when_omitted` pins that a no-loop caller (no `loop_counter`) still produces a valid 3-column tuple with the loop columns at `0` / `""`. - Local: `workflow-core` + `amber` compile; `StateSpec` + `ArrowUtilsSpec` pass; Python state + transport + e2e tests pass; scalafmt + scalafix + black clean. (`IcebergDocumentSpec` needs the iceberg catalog backend, so it runs in CI.) ### Was this PR authored or co-authored using generative AI tooling? Co-authored with Claude Opus 4.8 in compliance with ASF. --- .../core/architecture/packaging/output_manager.py | 107 +++--- amber/src/main/python/core/models/payload.py | 8 + amber/src/main/python/core/models/state.py | 43 ++- .../main/python/core/runnables/network_receiver.py | 6 +- .../main/python/core/runnables/network_sender.py | 15 +- .../input_port_materialization_reader_runnable.py | 17 +- .../messaginglayer/OutputManager.scala | 2 +- .../pythonworker/PythonProxyClient.scala | 2 +- .../result/iceberg/IcebergDocumentSpec.scala | 64 ++-- .../architecture/packaging/test_output_manager.py | 22 +- .../packaging/test_state_materialization_e2e.py | 400 +++++++++++---------- amber/src/test/python/core/models/test_state.py | 24 +- .../python/core/runnables/test_network_receiver.py | 13 +- ...t_input_port_materialization_reader_runnable.py | 23 +- .../org/apache/texera/amber/core/state/State.scala | 20 +- .../apache/texera/amber/core/state/StateSpec.scala | 16 +- .../apache/texera/amber/util/ArrowUtilsSpec.scala | 45 +++ 17 files changed, 510 insertions(+), 317 deletions(-) diff --git a/amber/src/main/python/core/architecture/packaging/output_manager.py b/amber/src/main/python/core/architecture/packaging/output_manager.py index b85e3e39bf..90ceda2eb9 100644 --- a/amber/src/main/python/core/architecture/packaging/output_manager.py +++ b/amber/src/main/python/core/architecture/packaging/output_manager.py @@ -133,47 +133,31 @@ class OutputManager: state materialization on the same port. `storage_uri_base` is the port's base URI; the result and state URIs are derived from it. """ - document, _ = DocumentFactory.open_document( - VFSURIFactory.result_uri(storage_uri_base) - ) - buffered_item_writer = document.writer(str(get_worker_index(self.worker_id))) - writer_queue = Queue() - port_storage_writer = PortStorageWriter( - buffered_item_writer=buffered_item_writer, queue=writer_queue - ) - writer_thread = threading.Thread( - target=port_storage_writer.run, - daemon=True, - name=f"port_storage_writer_thread_{port_id}", - ) - writer_thread.start() - self._port_storage_writers[port_id] = ( - writer_queue, - port_storage_writer, - writer_thread, - ) - state_document, _ = DocumentFactory.open_document( - VFSURIFactory.state_uri(storage_uri_base) - ) - state_buffered_item_writer = state_document.writer( - str(get_worker_index(self.worker_id)) - ) - state_writer_queue = Queue() - state_port_writer = PortStorageWriter( - buffered_item_writer=state_buffered_item_writer, - queue=state_writer_queue, - ) - state_writer_thread = threading.Thread( - target=state_port_writer.run, - daemon=True, - name=f"port_state_writer_thread_{port_id}", + def start_writer(uri: str, name_prefix: str, registry: dict) -> None: + document, _ = DocumentFactory.open_document(uri) + writer_queue = Queue() + writer = PortStorageWriter( + buffered_item_writer=document.writer( + str(get_worker_index(self.worker_id)) + ), + queue=writer_queue, + ) + thread = threading.Thread( + target=writer.run, daemon=True, name=f"{name_prefix}_{port_id}" + ) + thread.start() + registry[port_id] = (writer_queue, writer, thread) + + start_writer( + VFSURIFactory.result_uri(storage_uri_base), + "port_storage_writer_thread", + self._port_storage_writers, ) - state_writer_thread.start() - self._port_state_writers[port_id] = ( - state_writer_queue, - state_port_writer, - state_writer_thread, + start_writer( + VFSURIFactory.state_uri(storage_uri_base), + "port_state_writer_thread", + self._port_state_writers, ) def get_port(self, port_id=None) -> WorkerPort: @@ -203,14 +187,22 @@ class OutputManager: PortStorageWriterElement(data_tuple=tuple_) ) - def save_state_to_storage_if_needed(self, state: State, port_id=None) -> None: + def save_state_to_storage_if_needed( + self, + state: State, + loop_counter: int = 0, + loop_start_id: str = "", + port_id=None, + ) -> None: # When port_id is omitted the same state row is fanned out to # every output port's state table. This mirrors the # broadcast-to-all-workers behavior on the emit side: state is # shared context, not per-key data, so every downstream operator # (and every worker reading the materialization) needs the full # set. - element = PortStorageWriterElement(data_tuple=state.to_tuple()) + element = PortStorageWriterElement( + data_tuple=state.to_tuple(loop_counter, loop_start_id) + ) if port_id is None: for writer_queue, _, _ in self._port_state_writers.values(): writer_queue.put(element) @@ -223,18 +215,16 @@ class OutputManager: writer threads to finish, which indicates the port storage writing are finished. """ - for _, writer, _ in self._port_storage_writers.values(): - # This non-blocking stop call will let the storage writers - # flush the remaining buffer - writer.stop() - for _, _, writer_thread in self._port_storage_writers.values(): - # This blocking call will wait for all the writer to finish commit - writer_thread.join() - for _, state_writer, _ in self._port_state_writers.values(): - state_writer.stop() - for _, _, state_writer_thread in self._port_state_writers.values(): - state_writer_thread.join() - self._port_state_writers.clear() + for registry in (self._port_storage_writers, self._port_state_writers): + # Non-blocking stop lets each writer flush its remaining buffer; + # the join then waits for the commit to finish. + for _, writer, _ in registry.values(): + writer.stop() + for _, _, thread in registry.values(): + thread.join() + # Drop the stopped writers so a later close doesn't act on + # stale entries. + registry.clear() def add_partitioning(self, tag: PhysicalLink, partitioning: Partitioning) -> None: """ @@ -290,7 +280,10 @@ class OutputManager: ) def emit_state( - self, state: State + self, + state: State, + loop_counter: int = 0, + loop_start_id: str = "", ) -> Iterable[typing.Tuple[ActorVirtualIdentity, DataPayload]]: return chain( *( @@ -298,7 +291,11 @@ class OutputManager: ( receiver, ( - StateFrame(payload) + StateFrame( + payload, + loop_counter=loop_counter, + loop_start_id=loop_start_id, + ) if isinstance(payload, State) else self.tuple_to_frame(payload) ), diff --git a/amber/src/main/python/core/models/payload.py b/amber/src/main/python/core/models/payload.py index 61a3329488..6fc3580679 100644 --- a/amber/src/main/python/core/models/payload.py +++ b/amber/src/main/python/core/models/payload.py @@ -34,3 +34,11 @@ class DataFrame(DataPayload): @dataclass class StateFrame(DataPayload): frame: State + # Loop-control bookkeeping owned by the worker runtime, carried alongside + # the State payload (not inside it) so it never collides with user state. + # Defaults are the "no loop" values for all non-loop state. + loop_counter: int = 0 + # Which LoopStart to jump back to. Set by the runtime on a LoopStart's + # output, consumed by the matching LoopEnd. Empty for non-loop / + # not-yet-stamped state. + loop_start_id: str = "" diff --git a/amber/src/main/python/core/models/state.py b/amber/src/main/python/core/models/state.py index 003aaa212a..559d6eca45 100644 --- a/amber/src/main/python/core/models/state.py +++ b/amber/src/main/python/core/models/state.py @@ -25,13 +25,50 @@ from .tuple import Tuple class State(dict): CONTENT = "content" - SCHEMA = Schema(raw_schema={CONTENT: "STRING"}) + # Loop-control bookkeeping owned by the worker runtime, NOT user state -- it + # never appears in the content JSON. In memory it rides on the StateFrame + # envelope; it is materialized/serialized as its own column (parallel to + # content) by to_tuple(...). from_tuple() returns the bare State; callers + # that need these values read the corresponding columns off the tuple. + LOOP_COUNTER = "loop_counter" + LOOP_START_ID = "loop_start_id" + SCHEMA = Schema( + raw_schema={ + CONTENT: "STRING", + LOOP_COUNTER: "LONG", + LOOP_START_ID: "STRING", + } + ) def to_json(self) -> str: return json.dumps(_to_json_value(self), separators=(",", ":")) - def to_tuple(self) -> Tuple: - return Tuple({State.CONTENT: self.to_json()}, schema=State.SCHEMA) + @staticmethod + def to_columns( + content_json: str, + loop_counter: int = 0, + loop_start_id: str = "", + ) -> dict: + """The single column-name -> value mapping for the State wire/storage + format. Both ``to_tuple`` (iceberg materialization) and the network + sender build from this, so adding a column is a one-line change here + rather than in every serializer. + """ + return { + State.CONTENT: content_json, + State.LOOP_COUNTER: int(loop_counter), + State.LOOP_START_ID: loop_start_id, + } + + def to_tuple( + self, + loop_counter: int = 0, + loop_start_id: str = "", + ) -> Tuple: + return Tuple( + State.to_columns(self.to_json(), loop_counter, loop_start_id), + schema=State.SCHEMA, + ) @classmethod def from_json(cls, payload: str) -> "State": diff --git a/amber/src/main/python/core/runnables/network_receiver.py b/amber/src/main/python/core/runnables/network_receiver.py index 8ba4fbe147..c9f64c1fca 100644 --- a/amber/src/main/python/core/runnables/network_receiver.py +++ b/amber/src/main/python/core/runnables/network_receiver.py @@ -96,7 +96,11 @@ class NetworkReceiver(Runnable, Stoppable): "Data", lambda _: DataFrame(table), "State", - lambda _: StateFrame(State.from_json(table[State.CONTENT][0].as_py())), + lambda _: StateFrame( + State.from_json(table[State.CONTENT][0].as_py()), + loop_counter=int(table[State.LOOP_COUNTER][0].as_py()), + loop_start_id=table[State.LOOP_START_ID][0].as_py(), + ), "ECM", lambda _: EmbeddedControlMessage().parse(table["payload"][0].as_py()), ) diff --git a/amber/src/main/python/core/runnables/network_sender.py b/amber/src/main/python/core/runnables/network_sender.py index d8e3889ac1..3dcffb98b7 100644 --- a/amber/src/main/python/core/runnables/network_sender.py +++ b/amber/src/main/python/core/runnables/network_sender.py @@ -20,7 +20,13 @@ from loguru import logger from overrides import overrides from typing import Optional -from core.models import DataPayload, InternalQueue, DataFrame, State, StateFrame +from core.models import ( + DataPayload, + InternalQueue, + DataFrame, + State, + StateFrame, +) from core.models.internal_queue import ( InternalQueueElement, DataElement, @@ -99,8 +105,13 @@ class NetworkSender(StoppableQueueBlockingRunnable): self._proxy_client.send_data(bytes(data_header), data_payload.frame) elif isinstance(data_payload, StateFrame): data_header = PythonDataHeader(tag=to, payload_type="State") + columns = State.to_columns( + data_payload.frame.to_json(), + data_payload.loop_counter, + data_payload.loop_start_id, + ) table = pa.Table.from_pydict( - {State.CONTENT: [data_payload.frame.to_json()]}, + {name: [value] for name, value in columns.items()}, schema=State.SCHEMA.as_arrow_schema(), ) self._proxy_client.send_data(bytes(data_header), table) diff --git a/amber/src/main/python/core/storage/runnables/input_port_materialization_reader_runnable.py b/amber/src/main/python/core/storage/runnables/input_port_materialization_reader_runnable.py index 3e0e2d48ab..4bc4e68539 100644 --- a/amber/src/main/python/core/storage/runnables/input_port_materialization_reader_runnable.py +++ b/amber/src/main/python/core/storage/runnables/input_port_materialization_reader_runnable.py @@ -34,7 +34,14 @@ from core.architecture.sendsemantics.range_based_shuffle_partitioner import ( from core.architecture.sendsemantics.round_robin_partitioner import ( RoundRobinPartitioner, ) -from core.models import Tuple, InternalQueue, DataFrame, DataPayload, State, StateFrame +from core.models import ( + Tuple, + InternalQueue, + DataFrame, + DataPayload, + State, + StateFrame, +) from core.models.internal_queue import DataElement, ECMElement from core.storage.document_factory import DocumentFactory from core.storage.vfs_uri_factory import VFSURIFactory @@ -152,7 +159,13 @@ class InputPortMaterializationReaderRunnable(Runnable, Stoppable): VFSURIFactory.state_uri(self.uri) ) for state_row in state_document.get(): - self.emit_payload(StateFrame(State.from_tuple(state_row))) + self.emit_payload( + StateFrame( + State.from_tuple(state_row), + loop_counter=state_row[State.LOOP_COUNTER], + loop_start_id=state_row[State.LOOP_START_ID], + ) + ) storage_iterator = self.materialization.get() # Iterate and process tuples. diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/OutputManager.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/OutputManager.scala index 030fa3a3bb..80cc24780b 100644 --- a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/OutputManager.scala +++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/OutputManager.scala @@ -242,7 +242,7 @@ class OutputManager( // emit side: state is shared context, not per-key data, so every // downstream operator (and every worker reading the materialization) // needs the full set. - stateWriterThreads.values.foreach(_.queue.put(Left(state.toTuple))) + stateWriterThreads.values.foreach(_.queue.put(Left(state.toTuple()))) } /** diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyClient.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyClient.scala index 6618e857b1..144c3ac57a 100644 --- a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyClient.scala +++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyClient.scala @@ -125,7 +125,7 @@ class PythonProxyClient(portNumberPromise: Promise[Int], val actorId: ActorVirtu case DataFrame(frame) => writeArrowStream(mutable.Queue(ArraySeq.unsafeWrapArray(frame): _*), from, "Data") case StateFrame(state) => - writeArrowStream(mutable.Queue(state.toTuple), from, "State") + writeArrowStream(mutable.Queue(state.toTuple()), from, "State") } } diff --git a/amber/src/test/integration/org/apache/texera/amber/storage/result/iceberg/IcebergDocumentSpec.scala b/amber/src/test/integration/org/apache/texera/amber/storage/result/iceberg/IcebergDocumentSpec.scala index 78999c7ea3..aa8011e553 100644 --- a/amber/src/test/integration/org/apache/texera/amber/storage/result/iceberg/IcebergDocumentSpec.scala +++ b/amber/src/test/integration/org/apache/texera/amber/storage/result/iceberg/IcebergDocumentSpec.scala @@ -213,7 +213,7 @@ class IcebergDocumentSpec extends VirtualDocumentSpec[Tuple] with BeforeAndAfter DocumentFactory.openDocument(stateUri)._1.asInstanceOf[VirtualDocument[Tuple]] val state = State( Map( - "loop_counter" -> 3, + "i" -> 3, "name" -> "outer-loop", "payload" -> Array[Byte](0, 1, 2, 3), "nested" -> Map("enabled" -> true, "values" -> List(1, 2, 3)) @@ -222,13 +222,17 @@ class IcebergDocumentSpec extends VirtualDocumentSpec[Tuple] with BeforeAndAfter val writer = stateDocument.writer(UUID.randomUUID().toString) writer.open() - writer.putOne(state.toTuple) + writer.putOne(state.toTuple(loopCounter = 7L, loopStartId = "ls")) writer.close() val storedRows = stateDocument.get().toList assert(storedRows.length == 1) + // Loop bookkeeping is materialized as its own columns, not in the content JSON. + assert(storedRows.head.getField[java.lang.Long]("loop_counter").toLong == 7L) + assert(storedRows.head.getField[String]("loop_start_id") == "ls") + // User state round-trips through the content column (fromTuple reads only content). val deserialized = State.fromTuple(storedRows.head).values - assert(deserialized("loop_counter") == 3L) + assert(deserialized("i") == 3L) assert(deserialized("name") == "outer-loop") assert(deserialized("payload").asInstanceOf[Array[Byte]].sameElements(Array[Byte](0, 1, 2, 3))) assert(deserialized("nested").asInstanceOf[Map[String, Any]]("enabled") == true) @@ -240,45 +244,45 @@ class IcebergDocumentSpec extends VirtualDocumentSpec[Tuple] with BeforeAndAfter DocumentFactory.createDocument(stateUri, State.schema) val stateDocument = DocumentFactory.openDocument(stateUri)._1.asInstanceOf[VirtualDocument[Tuple]] - val states: List[State] = List( - State(Map("loop_counter" -> 0, "i" -> 1, "payload" -> Array[Byte](1, 2, 3))), - State( - Map( - "loop_counter" -> 1, - "i" -> 2, - "payload" -> Array[Byte](4, 5, 6), - "nested" -> Map("values" -> List(3, 4)) - ) + // (user state, loopCounter) -- the counter is written to its own column. + val states: List[(State, Long)] = List( + (State(Map("i" -> 1, "payload" -> Array[Byte](1, 2, 3))), 0L), + ( + State( + Map( + "i" -> 2, + "payload" -> Array[Byte](4, 5, 6), + "nested" -> Map("values" -> List(3, 4)) + ) + ), + 1L ) ) val writer = stateDocument.writer(UUID.randomUUID().toString) writer.open() - states.foreach(state => writer.putOne(state.toTuple)) + states.foreach { case (state, loopCounter) => writer.putOne(state.toTuple(loopCounter)) } writer.close() - val deserializedStates = - stateDocument - .get() - .toList - .map(State.fromTuple) - .sortBy(_.values("loop_counter").asInstanceOf[Long]) - assert(deserializedStates.length == states.length) - deserializedStates.zip(states).foreach { - case (actual, expected) => - assert( - actual.values("loop_counter") == expected.values("loop_counter").asInstanceOf[Int].toLong - ) - assert(actual.values("i") == expected.values("i").asInstanceOf[Int].toLong) + val storedRows = + stateDocument.get().toList.sortBy(_.getField[java.lang.Long]("loop_counter").toLong) + assert(storedRows.length == states.length) + storedRows.zip(states).foreach { + case (row, (expectedState, expectedLoopCounter)) => + // loop_counter is its own column... + assert(row.getField[java.lang.Long]("loop_counter").toLong == expectedLoopCounter) + // ...and the user state round-trips through the content column. + val actual = State.fromTuple(row).values + assert(actual("i") == expectedState.values("i").asInstanceOf[Int].toLong) assert( - actual - .values("payload") + actual("payload") .asInstanceOf[Array[Byte]] - .sameElements(expected.values("payload").asInstanceOf[Array[Byte]]) + .sameElements(expectedState.values("payload").asInstanceOf[Array[Byte]]) ) } assert( - deserializedStates(1) + State + .fromTuple(storedRows(1)) .values("nested") .asInstanceOf[Map[String, Any]]("values") == List(3L, 4L) ) diff --git a/amber/src/test/python/core/architecture/packaging/test_output_manager.py b/amber/src/test/python/core/architecture/packaging/test_output_manager.py index dcf7ccde67..3eae42e9c1 100644 --- a/amber/src/test/python/core/architecture/packaging/test_output_manager.py +++ b/amber/src/test/python/core/architecture/packaging/test_output_manager.py @@ -49,15 +49,15 @@ class TestSaveStateToStorageIfNeeded: @pytest.fixture def state(self): - return State({"loop_counter": 1, "i": 2}) + return State({"i": 2}) def test_no_state_writers_is_a_noop(self, output_manager, state): # With no port set up, save_state_to_storage_if_needed must not # touch any writer. - output_manager.save_state_to_storage_if_needed(state) # no-op + output_manager.save_state_to_storage_if_needed(state, 0) # no-op def test_unknown_port_id_is_a_noop(self, output_manager, state, port_a): - output_manager.save_state_to_storage_if_needed(state, port_id=port_a) + output_manager.save_state_to_storage_if_needed(state, 0, port_id=port_a) # No assertion needed -- the absence of any writer means nothing # was attempted. @@ -67,7 +67,7 @@ class TestSaveStateToStorageIfNeeded: queue_a, _, _ = _stub_state_writer(output_manager, port_a) queue_b, _, _ = _stub_state_writer(output_manager, port_b) - output_manager.save_state_to_storage_if_needed(state) + output_manager.save_state_to_storage_if_needed(state, 0) # Each port's writer queue receives one PortStorageWriterElement. # Critically, save is non-blocking -- the call must not invoke @@ -84,7 +84,7 @@ class TestSaveStateToStorageIfNeeded: queue_a, _, _ = _stub_state_writer(output_manager, port_a) queue_b, _, _ = _stub_state_writer(output_manager, port_b) - output_manager.save_state_to_storage_if_needed(state, port_id=port_a) + output_manager.save_state_to_storage_if_needed(state, 0, port_id=port_a) assert queue_a.put.call_count == 1 queue_b.put.assert_not_called() @@ -105,3 +105,15 @@ class TestSaveStateToStorageIfNeeded: thread_a.join.assert_called_once() thread_b.join.assert_called_once() assert output_manager._port_state_writers == {} + + def test_defaults_loop_columns_when_omitted(self, output_manager, state, port_a): + # Dormancy: callers that pass no loop bookkeeping (every non-loop + # caller, e.g. MainLoop.process_input_state) still produce a valid + # 3-column state tuple with the loop columns at their no-loop defaults. + queue_a, _, _ = _stub_state_writer(output_manager, port_a) + + output_manager.save_state_to_storage_if_needed(state) # no loop_counter + + data_tuple = queue_a.put.call_args.args[0].data_tuple + assert data_tuple[State.LOOP_COUNTER] == 0 + assert data_tuple[State.LOOP_START_ID] == "" diff --git a/amber/src/test/python/core/architecture/packaging/test_state_materialization_e2e.py b/amber/src/test/python/core/architecture/packaging/test_state_materialization_e2e.py index 8613be95b1..b77bdca3f5 100644 --- a/amber/src/test/python/core/architecture/packaging/test_state_materialization_e2e.py +++ b/amber/src/test/python/core/architecture/packaging/test_state_materialization_e2e.py @@ -25,7 +25,7 @@ This test wires: OutputManager.set_up_port_storage_writer(port, base_uri) → real PortStorageWriter thread - → real IcebergTableWriter (postgres-backed JdbcCatalog) + → real IcebergTableWriter (sqlite-backed SqlCatalog) → state document at VFSURIFactory.state_uri(base_uri) → InputPortMaterializationReaderRunnable.run() → DataElement(StateFrame) on the consumer's input queue @@ -33,20 +33,14 @@ This test wires: and asserts that a state put through `save_state_to_storage_if_needed` on the producer side actually arrives at the consumer's queue, with the same payload. - -Marked @integration so the CI runner that has postgres + iceberg -catalog DB provisioned (amber-integration) picks it up via -`pytest -m integration`. Earlier versions of this test substituted a -sqlite-backed SqlCatalog to dodge that infra dependency; that diverged -from the prod catalog code path, so we now exercise the real one. """ -import os import tempfile import threading import uuid import pytest +from pyiceberg.catalog.sql import SqlCatalog from core.architecture.packaging.output_manager import OutputManager from core.models import State, StateFrame @@ -74,203 +68,219 @@ from proto.org.apache.texera.amber.engine.architecture.sendsemantics import ( ) [email protected] -class TestStateMaterializationE2E: - @pytest.fixture(autouse=True, scope="class") - def _init_storage_config(self): - """Initialize StorageConfig + IcebergCatalogInstance for the real - postgres-backed catalog in the `amber-integration` CI job. - - Critical detail: the Scala integration tests that run earlier in - the same job connect to the iceberg catalog DB as user - `postgres/postgres` (the storage.conf default for - `STORAGE_ICEBERG_CATALOG_POSTGRES_USERNAME/PASSWORD`). pyiceberg - creates the catalog's `iceberg_tables` metadata table on first - use, owned by whoever wrote first — so it ends up owned by - `postgres`. We MUST connect as the same user, otherwise we hit - `permission denied for table iceberg_tables`. - - Why the reset: `test_iceberg_document.py` also calls - `StorageConfig.initialize` at module import time (with a - different `texera/password` user that works for it because no - Scala writes first in the `pyamber` job where it runs). pytest - imports every test module during collection, even ones whose - tests will be deselected by `-m integration`, so that - initialization happens here too. We force-reset the singletons - and re-init with the prod-correct credentials; safe because - test_iceberg_document's tests are deselected from this run. - - All catalog + S3 settings read the same `STORAGE_*` env vars - the production code consumes (via storage.conf), so the test - matches whichever identity the Scala side uses in the same job - and stays aligned with the bucket / endpoint the workflow - provisions. Defaults mirror storage.conf so a local sbt run - without those vars exported still works. - - Class-scoped so the reset + tempdir allocation happens once - per class; the two tests in this class share state through the - same StorageConfig singleton anyway. - """ - StorageConfig._initialized = False - IcebergCatalogInstance._instance = None - large_binaries_bucket = os.environ.get( - "STORAGE_S3_LARGE_BINARIES_BUCKET", "texera-large-binaries" - ) +# Module-level scratch dir for the sqlite catalog + iceberg warehouse. +# We don't initialize `StorageConfig` here: other test modules (e.g. +# test_iceberg_document.py) also call `StorageConfig.initialize` at +# import time, and the class rejects re-initialization with +# RuntimeError. Whichever module gets collected first wins; we adopt +# its namespaces below. +_WAREHOUSE_DIR = tempfile.mkdtemp(prefix="texera-state-e2e-warehouse-") + + [email protected](scope="module", autouse=True) +def sqlite_iceberg_catalog(): + """Inject a sqlite-backed SqlCatalog so the test runs without external + iceberg infra (postgres/minio). + + Note: the other iceberg-backed tests (e.g. test_iceberg_document.py) use a + postgres/REST catalog to mirror production. This e2e deliberately diverges + to a hermetic sqlite catalog so the writer→storage→reader join can run as a + fast, infra-free unit test -- the materialization logic it exercises is + catalog-agnostic, so the sqlite backend exercises the same code path. + + Module-scoped so all tests in this file share one warehouse, and so + namespace creation only happens once. We save/restore the original + `IcebergCatalogInstance` singleton so other test modules that expect + a real postgres-backed catalog (e.g. test_iceberg_document.py) are + not affected by our replacement. + """ + # Some other test module may have initialized StorageConfig already + # (it has a single-init lock). If nothing has initialized it yet, + # do it here with arbitrary values -- we replace the catalog + # instance below so the postgres/rest fields are never exercised. + if not StorageConfig._initialized: StorageConfig.initialize( catalog_type="postgres", - postgres_uri_without_scheme=os.environ.get( - "STORAGE_ICEBERG_CATALOG_POSTGRES_URI_WITHOUT_SCHEME", - "localhost:5432/texera_iceberg_catalog", - ), - postgres_username=os.environ.get( - "STORAGE_ICEBERG_CATALOG_POSTGRES_USERNAME", "postgres" - ), - postgres_password=os.environ.get( - "STORAGE_ICEBERG_CATALOG_POSTGRES_PASSWORD", "postgres" - ), - rest_catalog_uri="http://localhost:8181/catalog/";, - rest_catalog_warehouse_name="texera", + postgres_uri_without_scheme="unused", + postgres_username="unused", + postgres_password="unused", + rest_catalog_uri="unused", + rest_catalog_warehouse_name="unused", table_result_namespace="operator-port-result", table_state_namespace="operator-port-state", - directory_path=tempfile.mkdtemp(prefix="texera-state-e2e-warehouse-"), + directory_path=_WAREHOUSE_DIR, commit_batch_size=4096, - s3_endpoint=os.environ.get("STORAGE_S3_ENDPOINT", "http://localhost:9000";), - s3_region=os.environ.get("STORAGE_S3_REGION", "us-west-2"), - s3_auth_username=os.environ.get("STORAGE_S3_AUTH_USERNAME", "texera_minio"), - s3_auth_password=os.environ.get("STORAGE_S3_AUTH_PASSWORD", "password"), - s3_large_binaries_base_uri=f"s3://{large_binaries_bucket}/objects/0/", + s3_endpoint="unused", + s3_region="unused", + s3_auth_username="unused", + s3_auth_password="unused", + s3_large_binaries_base_uri="s3://texera-large-binaries/objects/0/", ) - @pytest.fixture - def base_uri(self) -> str: - """A unique port-base URI per test so tables don't collide.""" - return VFSURIFactory.create_port_base_uri( - WorkflowIdentity(id=0), - ExecutionIdentity(id=0), - GlobalPortIdentity( - op_id=PhysicalOpIdentity( - logical_op_id=OperatorIdentity(id=f"e2e-{uuid.uuid4().hex}"), - layer_name="main", - ), - port_id=PortIdentity(id=0, internal=False), - input=False, + original_instance = IcebergCatalogInstance._instance + db_path = f"{_WAREHOUSE_DIR}/catalog.sqlite" + catalog = SqlCatalog( + "texera_iceberg_e2e", + **{ + "uri": f"sqlite:///{db_path}", + "warehouse": f"file://{_WAREHOUSE_DIR}", + }, + ) + # Adopt whatever namespaces StorageConfig already has -- those are + # the ones DocumentFactory will route into. + catalog.create_namespace_if_not_exists(StorageConfig.ICEBERG_TABLE_RESULT_NAMESPACE) + catalog.create_namespace_if_not_exists(StorageConfig.ICEBERG_TABLE_STATE_NAMESPACE) + IcebergCatalogInstance.replace_instance(catalog) + try: + yield catalog + finally: + IcebergCatalogInstance.replace_instance(original_instance) + + +def _fresh_base_uri() -> str: + """A unique port-base URI per test so tables don't collide.""" + return VFSURIFactory.create_port_base_uri( + WorkflowIdentity(id=0), + ExecutionIdentity(id=0), + GlobalPortIdentity( + op_id=PhysicalOpIdentity( + logical_op_id=OperatorIdentity(id=f"e2e-{uuid.uuid4().hex}"), + layer_name="main", ), - ) + port_id=PortIdentity(id=0, internal=False), + input=False, + ), + ) - @pytest.fixture - def producer(self, base_uri): - """An OutputManager wired to the iceberg result + state documents - at `base_uri`. Closes its writer threads on teardown so cached - buffers are flushed even if a test errors out before - `close_port_storage_writers()`. - """ - # RegionExecutionCoordinator's responsibility in prod: provision - # result + state documents at the port base URI before any - # worker starts. We emulate that here. - DocumentFactory.create_document( - VFSURIFactory.result_uri(base_uri), State.SCHEMA - ) - DocumentFactory.create_document(VFSURIFactory.state_uri(base_uri), State.SCHEMA) - mgr = OutputManager(worker_id="Worker:WF0-test-producer-main-0") - mgr.add_output_port( - PortIdentity(id=0, internal=False), - schema=State.SCHEMA, - storage_uri_base=base_uri, - ) - try: - yield mgr - finally: - # close_port_storage_writers is idempotent — fine to call - # again here if the test already closed. - try: - mgr.close_port_storage_writers() - except Exception: - pass - - def test_state_written_by_output_manager_is_replayed_by_reader( - self, base_uri, producer - ): - """Producer side writes a state via OutputManager; consumer side - reads it via InputPortMaterializationReaderRunnable. The state - must arrive on the consumer's input queue intact. - """ - # Drive a state through the producer-side path. - state = State({"flag": True, "loop_counter": 7, "name": "outer"}) - producer.save_state_to_storage_if_needed(state) - - # Force the writer threads to flush + commit by closing them. - # Without this, the iceberg buffer holds the state in memory - # and nothing is durable yet. - producer.close_port_storage_writers() - - # Consumer side: spin up the materialization reader against the - # same base URI. Each reader needs a partitioning even when no - # real downstream worker exists — supply a OneToOnePartitioning - # whose only receiver is the consumer worker itself. - consumer_worker = ActorVirtualIdentity(name="consumer-worker-0") - consumer_queue = InternalQueue() - partitioning = Partitioning( - one_to_one_partitioning=OneToOnePartitioning( - batch_size=400, - channels=[ - ChannelIdentity( - from_worker_id=ActorVirtualIdentity(name="producer-worker-0"), - to_worker_id=consumer_worker, - is_control=False, - ) - ], - ) - ) - reader = InputPortMaterializationReaderRunnable( - uri=base_uri, - queue=consumer_queue, - worker_actor_id=consumer_worker, - partitioning=partitioning, - ) +def test_state_written_by_output_manager_is_replayed_by_reader(): + """Producer side writes a state via OutputManager; consumer side reads + it via InputPortMaterializationReaderRunnable. The state must arrive + on the consumer's input queue intact. + """ + base_uri = _fresh_base_uri() + port_id = PortIdentity(id=0, internal=False) + worker_schema_for_result = State.SCHEMA # producer-side: only state matters - # Run the reader on a worker thread so we can time out cleanly - # if something goes wrong. - reader_thread = threading.Thread(target=reader.run, daemon=True) - reader_thread.start() - reader_thread.join(timeout=30) - assert not reader_thread.is_alive(), "reader did not finish within timeout" - assert reader.finished(), "reader exited but did not mark itself finished" - - # Drain the consumer's queue and find the StateFrame(s). - state_frames: list[State] = [] - while not consumer_queue.is_empty(): - elem = consumer_queue.get() - if isinstance(elem, DataElement) and isinstance(elem.payload, StateFrame): - state_frames.append(elem.payload.frame) - - assert len(state_frames) == 1, ( - f"expected exactly one State to flow through writer→iceberg→reader; " - f"got {len(state_frames)}: {state_frames}" - ) - assert state_frames[0] == state, ( - f"replayed state did not match what was written; " - f"wrote={state}, read={state_frames[0]}" - ) + # 1. RegionExecutionCoordinator's responsibility: provision result + + # state documents at the port base URI before any worker starts. + # We emulate that here. + DocumentFactory.create_document( + VFSURIFactory.result_uri(base_uri), worker_schema_for_result + ) + DocumentFactory.create_document(VFSURIFactory.state_uri(base_uri), State.SCHEMA) - def test_state_table_persists_across_writer_close(self, base_uri, producer): - """Independently verify the iceberg state table contains the row. - If this passes but the reader test above fails, the bug is in - the reader / consumer wiring; if this fails, the bug is in the - writer / storage layer. - """ - state = State({"flag": False, "checkpoint": 42}) - producer.save_state_to_storage_if_needed(state) - producer.close_port_storage_writers() - - # Read directly from the iceberg state document, bypassing the - # reader. - state_document, _ = DocumentFactory.open_document( - VFSURIFactory.state_uri(base_uri) - ) - rows = list(state_document.get()) - assert len(rows) == 1, ( - f"expected exactly one row in the iceberg state table after " - f"the writer was closed; got {len(rows)} rows" + # 2. Producer side: spin up an OutputManager, set up real state + + # result writer threads against the iceberg storage. + producer = OutputManager(worker_id="Worker:WF0-test-producer-main-0") + producer.add_output_port( + port_id, schema=worker_schema_for_result, storage_uri_base=base_uri + ) + + # 3. Drive a state through the producer-side path. The loop bookkeeping + # rides alongside the State (not inside it) and is materialized as its own + # set of columns. Use non-default values for both so a regression in + # either column's plumbing is caught, not just loop_counter's. + state = State({"flag": True, "name": "outer"}) + producer.save_state_to_storage_if_needed( + state, + loop_counter=7, + loop_start_id="outer-loop", + ) + + # 4. Force the writer threads to flush + commit by closing them. + # Without this, the iceberg buffer holds the state in memory and + # nothing is durable yet. + producer.close_port_storage_writers() + + # 5. Consumer side: spin up the materialization reader against the + # same base URI. Each reader needs a partitioning even when no real + # downstream worker exists -- supply a OneToOnePartitioning whose + # only receiver is the consumer worker itself. + consumer_worker = ActorVirtualIdentity(name="consumer-worker-0") + consumer_queue = InternalQueue() + partitioning = Partitioning( + one_to_one_partitioning=OneToOnePartitioning( + batch_size=400, + channels=[ + ChannelIdentity( + from_worker_id=ActorVirtualIdentity(name="producer-worker-0"), + to_worker_id=consumer_worker, + is_control=False, + ) + ], ) - assert State.from_tuple(rows[0]) == state + ) + reader = InputPortMaterializationReaderRunnable( + uri=base_uri, + queue=consumer_queue, + worker_actor_id=consumer_worker, + partitioning=partitioning, + ) + + # Run the reader on a worker thread so we can time out cleanly if + # something goes wrong. + reader_thread = threading.Thread(target=reader.run, daemon=True) + reader_thread.start() + reader_thread.join(timeout=30) + assert not reader_thread.is_alive(), "reader did not finish within timeout" + assert reader.finished(), "reader exited but did not mark itself finished" + + # 6. Drain the consumer's queue and find the StateFrame(s). + state_frames: list[StateFrame] = [] + while not consumer_queue.is_empty(): + elem = consumer_queue.get() + if isinstance(elem, DataElement) and isinstance(elem.payload, StateFrame): + state_frames.append(elem.payload) + + assert len(state_frames) == 1, ( + f"expected exactly one State to flow through writer→iceberg→reader; " + f"got {len(state_frames)}: {state_frames}" + ) + assert state_frames[0].frame == state, ( + f"replayed state did not match what was written; " + f"wrote={state}, read={state_frames[0].frame}" + ) + assert state_frames[0].loop_counter == 7, ( + f"loop_counter must round-trip through its own column; " + f"got {state_frames[0].loop_counter}" + ) + assert state_frames[0].loop_start_id == "outer-loop", ( + f"loop_start_id must round-trip through its own column; " + f"got {state_frames[0].loop_start_id!r}" + ) + + +def test_state_table_persists_across_writer_close(): + """Independently verify the iceberg state table contains the row. + If this passes but the reader test above fails, the bug is in the + reader / consumer wiring; if this fails, the bug is in the writer / + storage layer. + """ + base_uri = _fresh_base_uri() + port_id = PortIdentity(id=0, internal=False) + + DocumentFactory.create_document(VFSURIFactory.result_uri(base_uri), State.SCHEMA) + DocumentFactory.create_document(VFSURIFactory.state_uri(base_uri), State.SCHEMA) + + producer = OutputManager(worker_id="Worker:WF0-test-producer2-main-0") + producer.add_output_port(port_id, schema=State.SCHEMA, storage_uri_base=base_uri) + + state = State({"flag": False, "checkpoint": 42}) + producer.save_state_to_storage_if_needed( + state, + loop_counter=3, + loop_start_id="inner-loop", + ) + producer.close_port_storage_writers() + + # Read directly from the iceberg state document, bypassing the reader. + state_document, _ = DocumentFactory.open_document(VFSURIFactory.state_uri(base_uri)) + rows = list(state_document.get()) + assert len(rows) == 1, ( + f"expected exactly one row in the iceberg state table after the " + f"writer was closed; got {len(rows)} rows" + ) + assert State.from_tuple(rows[0]) == state + assert rows[0][State.LOOP_COUNTER] == 3 + assert rows[0][State.LOOP_START_ID] == "inner-loop" diff --git a/amber/src/test/python/core/models/test_state.py b/amber/src/test/python/core/models/test_state.py index aef2297130..b1c9fad35a 100644 --- a/amber/src/test/python/core/models/test_state.py +++ b/amber/src/test/python/core/models/test_state.py @@ -29,7 +29,15 @@ class TestState: def test_class_attributes(self): assert State.CONTENT == "content" - assert State.SCHEMA.get_attr_names() == ["content"] + assert State.LOOP_COUNTER == "loop_counter" + assert State.LOOP_START_ID == "loop_start_id" + # The loop-control columns are runtime-owned bookkeeping, sibling to + # content, not part of the user state JSON. + assert State.SCHEMA.get_attr_names() == [ + "content", + "loop_counter", + "loop_start_id", + ] def test_json_round_trip_primitives(self): original = State( @@ -83,14 +91,20 @@ class TestState: State({"bad": Custom()}).to_json() def test_tuple_round_trip(self): - original = State({"loop_counter": 3, "label": "outer", "blob": b"\x01\x02"}) + original = State({"i": 3, "label": "outer", "blob": b"\x01\x02"}) decoded = State.from_tuple(original.to_tuple()) assert decoded == original - def test_to_tuple_uses_state_schema(self): - tuple_ = State({"x": 1}).to_tuple() - # Single STRING column whose value is the JSON serialization. + def test_to_tuple_writes_content_and_loop_counter_columns(self): + tuple_ = State({"x": 1}).to_tuple(7) + # content holds the JSON serialization; loop_counter is its own column. assert tuple_[State.CONTENT] == '{"x":1}' + assert tuple_[State.LOOP_COUNTER] == 7 + # loop_counter must not leak into the content JSON. + assert "loop_counter" not in tuple_[State.CONTENT] + + def test_to_tuple_defaults_loop_counter_to_zero(self): + assert State({"x": 1}).to_tuple()[State.LOOP_COUNTER] == 0 def test_nested_dict_decodes_to_plain_dict(self): # Top-level returns a State; nested dicts come back as plain dict. diff --git a/amber/src/test/python/core/runnables/test_network_receiver.py b/amber/src/test/python/core/runnables/test_network_receiver.py index bf890e4a2f..49ba2408ef 100644 --- a/amber/src/test/python/core/runnables/test_network_receiver.py +++ b/amber/src/test/python/core/runnables/test_network_receiver.py @@ -152,16 +152,19 @@ class TestNetworkReceiver: worker_id = ActorVirtualIdentity(name="test") channel_id = ChannelIdentity(worker_id, worker_id, False) + # loop_counter rides the StateFrame envelope (its own Arrow column), + # not the user content. Use a non-zero counter so the round-trip + # actually exercises the second column over the sender->receiver wire. input_queue.put( DataElement( tag=channel_id, - payload=StateFrame(State({"loop_counter": 0, "i": 1})), + payload=StateFrame(State({"i": 1}), loop_counter=0), ) ) input_queue.put( DataElement( tag=channel_id, - payload=StateFrame(State({"loop_counter": 1, "i": 2})), + payload=StateFrame(State({"i": 2}), loop_counter=5), ) ) @@ -169,11 +172,13 @@ class TestNetworkReceiver: second_element: DataElement = output_queue.get() assert isinstance(first_element.payload, StateFrame) - assert first_element.payload.frame == {"loop_counter": 0, "i": 1} + assert first_element.payload.frame == {"i": 1} + assert first_element.payload.loop_counter == 0 assert first_element.tag == channel_id assert isinstance(second_element.payload, StateFrame) - assert second_element.payload.frame == {"loop_counter": 1, "i": 2} + assert second_element.payload.frame == {"i": 2} + assert second_element.payload.loop_counter == 5 assert second_element.tag == channel_id @pytest.mark.timeout(10) diff --git a/amber/src/test/python/core/storage/runnables/test_input_port_materialization_reader_runnable.py b/amber/src/test/python/core/storage/runnables/test_input_port_materialization_reader_runnable.py index 5016c2df2f..288ee6374b 100644 --- a/amber/src/test/python/core/storage/runnables/test_input_port_materialization_reader_runnable.py +++ b/amber/src/test/python/core/storage/runnables/test_input_port_materialization_reader_runnable.py @@ -60,16 +60,25 @@ class TestRunStateReadingBlock: return instance def test_state_rows_are_emitted_as_state_frames(self, runnable): - state_a = State({"loop_counter": 0}) - state_b = State({"loop_counter": 1}) + state_a = State({"i": 0}) + state_b = State({"i": 1}) - # The state document yields opaque tuples; from_tuple deserializes - # them. Patch from_tuple so we don't have to wire a real - # serialization. + # The state document yields opaque multi-column tuples. State.from_tuple + # (patched) deserializes the content column; the reader reads the + # loop-control columns directly off the row and carries them onto the + # emitted StateFrame envelope. + row_a = { + State.LOOP_COUNTER: 0, + State.LOOP_START_ID: "loop-a", + } + row_b = { + State.LOOP_COUNTER: 1, + State.LOOP_START_ID: "loop-b", + } result_doc = MagicMock() result_doc.get.return_value = iter([]) # No materialized tuples. state_doc = MagicMock() - state_doc.get.return_value = iter(["row-a", "row-b"]) + state_doc.get.return_value = iter([row_a, row_b]) with ( patch( @@ -96,4 +105,6 @@ class TestRunStateReadingBlock: and isinstance(call.args[0].payload, StateFrame) ] assert [sf.payload.frame for sf in state_frames] == [state_a, state_b] + assert [sf.payload.loop_counter for sf in state_frames] == [0, 1] + assert [sf.payload.loop_start_id for sf in state_frames] == ["loop-a", "loop-b"] assert runnable._finished is True diff --git a/common/workflow-core/src/main/scala/org/apache/texera/amber/core/state/State.scala b/common/workflow-core/src/main/scala/org/apache/texera/amber/core/state/State.scala index ba146f1d57..f2896c8e19 100644 --- a/common/workflow-core/src/main/scala/org/apache/texera/amber/core/state/State.scala +++ b/common/workflow-core/src/main/scala/org/apache/texera/amber/core/state/State.scala @@ -31,18 +31,32 @@ final case class State(values: Map[String, Any]) { def toJson: String = objectMapper.writeValueAsString(State.toJsonValue(values)) - def toTuple: Tuple = - Tuple.builder(State.schema).addSequentially(Array(toJson)).build() + def toTuple( + loopCounter: Long = 0L, + loopStartId: String = "" + ): Tuple = + Tuple + .builder(State.schema) + .addSequentially(Array(toJson, Long.box(loopCounter), loopStartId)) + .build() } object State { private val Content = "content" + // loop-control bookkeeping owned by the (Python) worker runtime; not user + // state and never in the content JSON. Materialized as its own columns, + // parallel to content. Scala never originates loop state (loop operators are + // Python-only), so toTuple defaults these to the "no loop" values. + private val LoopCounter = "loop_counter" + private val LoopStartId = "loop_start_id" private val BytesTypeMarker = "__texera_type__" private val BytesValue = "bytes" private val PayloadMarker = "payload" val schema: Schema = new Schema( - new Attribute(Content, AttributeType.STRING) + new Attribute(Content, AttributeType.STRING), + new Attribute(LoopCounter, AttributeType.LONG), + new Attribute(LoopStartId, AttributeType.STRING) ) def fromJson(payload: String): State = diff --git a/common/workflow-core/src/test/scala/org/apache/texera/amber/core/state/StateSpec.scala b/common/workflow-core/src/test/scala/org/apache/texera/amber/core/state/StateSpec.scala index 976a585e31..f0af0b30c8 100644 --- a/common/workflow-core/src/test/scala/org/apache/texera/amber/core/state/StateSpec.scala +++ b/common/workflow-core/src/test/scala/org/apache/texera/amber/core/state/StateSpec.scala @@ -98,21 +98,29 @@ class StateSpec extends AnyFlatSpec { it should "tuple-round-trip" in { val original = State( Map( - "loop_counter" -> 3L, + "i" -> 3L, "label" -> "outer", "blob" -> Array[Byte](1, 2) ) ) - val decoded = State.fromTuple(original.toTuple) - assert(decoded.values("loop_counter") == 3L) + val tuple = original.toTuple(5L, "outer-loop") + + // Content round-trips through fromTuple, which reads only the content column. + val decoded = State.fromTuple(tuple) + assert(decoded.values("i") == 3L) assert(decoded.values("label") == "outer") assert( decoded.values("blob").asInstanceOf[Array[Byte]].sameElements(Array[Byte](1, 2)) ) + + // The loop bookkeeping is carried in its own columns (not the content + // JSON, and not surfaced by fromTuple), so assert it off the raw tuple. + assert(tuple.getField[java.lang.Long]("loop_counter").toLong == 5L) + assert(tuple.getField[String]("loop_start_id") == "outer-loop") } it should "produce a tuple whose payload is the JSON serialization" in { - val tuple = State(Map("x" -> 1L)).toTuple + val tuple = State(Map("x" -> 1L)).toTuple() assert(tuple.getSchema == State.schema) assert(tuple.getField[String]("content") == """{"x":1}""") } diff --git a/common/workflow-core/src/test/scala/org/apache/texera/amber/util/ArrowUtilsSpec.scala b/common/workflow-core/src/test/scala/org/apache/texera/amber/util/ArrowUtilsSpec.scala index 62b10a6686..4321008d46 100644 --- a/common/workflow-core/src/test/scala/org/apache/texera/amber/util/ArrowUtilsSpec.scala +++ b/common/workflow-core/src/test/scala/org/apache/texera/amber/util/ArrowUtilsSpec.scala @@ -19,8 +19,11 @@ package org.apache.texera.amber.util +import org.apache.arrow.memory.RootAllocator +import org.apache.arrow.vector.VectorSchemaRoot import org.apache.arrow.vector.types.{FloatingPointPrecision, TimeUnit} import org.apache.arrow.vector.types.pojo.{ArrowType, Field, FieldType} +import org.apache.texera.amber.core.state.State import org.apache.texera.amber.core.tuple.{Attribute, AttributeType, Schema} import org.apache.texera.amber.core.tuple.AttributeTypeUtils.AttributeTypeException import org.scalatest.flatspec.AnyFlatSpec @@ -282,4 +285,46 @@ class ArrowUtilsSpec extends AnyFlatSpec with Matchers { any.getMetadata.get("texera_type") shouldBe "ANY" Option(name.getMetadata).map(_.containsKey("texera_type")).getOrElse(false) shouldBe false } + + // ----- Tuple <-> Arrow data round-trip (the State wire-hop contract) ----- + + "tuple round-trip through Arrow vectors" should "preserve every column of a multi-column State tuple" in { + // The Python<->Scala state wire hop goes Tuple -> setTexeraTuple -> Arrow + // (PythonProxyClient.writeArrowStream) on one side and + // Arrow -> getTexeraTuple -> Tuple (PythonProxyServer) on the other. + // The schema-only round-trip tests above don't exercise the per-row data + // encode/decode, so a column dropped or mistyped there would slip through. + // Pin that the full multi-column State tuple (content STRING + the + // loop-control columns loop_counter LONG, loop_start_id STRING) survives a + // real setTexeraTuple -> Arrow vectors -> getTexeraTuple round-trip with + // every column intact -- the property the wire hop relies on. + val original = + State(Map("i" -> 5L, "label" -> "outer")).toTuple(3L, "outer-loop") + + val allocator = new RootAllocator() + val root = VectorSchemaRoot.create(ArrowUtils.fromTexeraSchema(original.getSchema), allocator) + try { + root.allocateNew() + ArrowUtils.setTexeraTuple(original, 0, root) + root.setRowCount(1) + + val recovered = ArrowUtils.getTexeraTuple(0, root) + + // Every column survives the encode/decode, with names and types intact. + recovered.getSchema.getAttributes.toList.map(a => (a.getName, a.getType)) shouldBe + List( + ("content", AttributeType.STRING), + ("loop_counter", AttributeType.LONG), + ("loop_start_id", AttributeType.STRING) + ) + // content (the user State JSON) round-trips... + State.fromTuple(recovered).values shouldBe Map("i" -> 5L, "label" -> "outer") + // ...and so do the loop-control columns. + recovered.getField[java.lang.Long]("loop_counter").toLong shouldBe 3L + recovered.getField[String]("loop_start_id") shouldBe "outer-loop" + } finally { + root.close() + allocator.close() + } + } }
