This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/texera.git
The following commit(s) were added to refs/heads/main by this push:
new 878eb8a38b feat(amber): add loop-bookkeeping columns to materialized
State (dormant) (#5900)
878eb8a38b is described below
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()
+ }
+ }
}
