feat: add loop operators

amber/src/main/python/core/architecture/packaging/output_manager.py

@@ -92,6 +92,11 @@ def __init__(self, worker_id: str):
             PortIdentity, typing.Tuple[Queue, PortStorageWriter, Thread]
         ] = dict()
 
+        # Loop-end operators have a single output port; remember its base
+        # URI so `reset_loopend_storage` can re-provision the iceberg
+        # tables on each loop iteration.
+        self._storage_uri_base: typing.Optional[str] = None
+
     def is_missing_output_ports(self):
         """
         This method is only used for ensuring correct region execution.
@@ -133,6 +138,9 @@ def set_up_port_storage_writer(self, port_id: PortIdentity, storage_uri_base: st
         state materialization on the same port. `storage_uri_base` is the
         port's base URI; the result and state URIs are derived from it.
         """
+        # Remember the base URI so `reset_loopend_storage` can re-provision
+        # the iceberg tables on subsequent loop iterations.
+        self._storage_uri_base = storage_uri_base
         document, _ = DocumentFactory.open_document(
             VFSURIFactory.result_uri(storage_uri_base)
         )
@@ -217,6 +225,19 @@ def save_state_to_storage_if_needed(self, state: State, port_id=None) -> None:
         elif port_id in self._port_state_writers:
             self._port_state_writers[port_id][0].put(element)
 
+    def reset_storage(self) -> None:
+        port_id = self.get_port_ids()[0]
+        storage_uri_base = self._storage_uri_base
+        self.close_port_storage_writers()
+        DocumentFactory.create_document(
+            VFSURIFactory.result_uri(storage_uri_base),
+            self._ports[port_id].get_schema(),
+        )
+        DocumentFactory.create_document(
+            VFSURIFactory.state_uri(storage_uri_base), State.SCHEMA
+        )
+        self.set_up_port_storage_writer(port_id, storage_uri_base)
+
     def close_port_storage_writers(self) -> None:
         """
         Flush the buffers of port storage writers and wait for all the

amber/src/main/python/core/models/operator.py

@@ -291,3 +291,30 @@ def process_table(self, table: Table, port: int) -> Iterator[Optional[TableLike]
             time, or None.
         """
         yield
+
+
+class LoopStartOperator(TableOperator):
+    @overrides.final
+    def process_state(self, state: State, port: int) -> Optional[State]:
+        if "LoopStartStateURI" in state:
+            state["loop_counter"] += 1
+            return state
+        self.state.update(state)
+        return None
+
+    @overrides.final
+    def produce_state_on_finish(self, port: int) -> State:
+        from pickle import dumps
+
+        self.state["table"] = dumps(Table(self._TableOperator__table_data[port]))
+        return dict(self.state)
+
+
+class LoopEndOperator(TableOperator):
+    @overrides.final
+    def process_table(self, table: Table, port: int) -> Iterator[Optional[TableLike]]:
+        yield table
+
+    @abstractmethod
+    def condition(self) -> bool:
+        pass

amber/src/main/python/core/runnables/main_loop.py

@@ -38,8 +38,11 @@
     ECMElement,
     InternalQueueElement,
 )
+from core.models.operator import LoopEndOperator, LoopStartOperator
 from core.models.state import State
 from core.runnables.data_processor import DataProcessor
+from core.storage.document_factory import DocumentFactory
+from core.storage.vfs_uri_factory import VFSURIFactory
 from core.util import StoppableQueueBlockingRunnable, get_one_of
 from core.util.console_message.timestamp import current_time_in_local_timezone
 from core.util.customized_queue.queue_base import QueueElement
@@ -48,6 +51,7 @@
     PortIdentity,
     ChannelIdentity,
     EmbeddedControlMessageIdentity,
+    OperatorIdentity,
 )
 from proto.org.apache.texera.amber.engine.architecture.rpc import (
     ConsoleMessage,
@@ -61,6 +65,7 @@
     EmbeddedControlMessage,
     AsyncRpcContext,
     ControlRequest,
+    JumpToOperatorRegionRequest,
 )
 from proto.org.apache.texera.amber.engine.architecture.worker import (
     WorkerState,
@@ -87,19 +92,54 @@ def __init__(
             target=self.data_processor.run, daemon=True, name="data_processor_thread"
         ).start()
 
+    def _attach_loop_start_id(self, output_state: State) -> None:
+        if "LoopStartId" in output_state:
+            return
+        output_state["LoopStartId"] = self.context.worker_id.split("-", 1)[1].rsplit(
+            "-main-0", 1
+        )[0]
+        # The URI lives on the upstream operator's output port (which
+        # LoopStart's first materialization reader is reading from).
+        reader_runnables = (
+            self.context.input_manager.get_input_port_mat_reader_threads()
+        )
+        output_state["LoopStartStateURI"] = VFSURIFactory.state_uri(
+            next(iter(reader_runnables.values()))[0].uri
+        )
+
+    def _jump_to_loop_start(
+        self, executor: LoopEndOperator, controller_interface
+    ) -> None:
+        state = executor.state
+        controller_interface.jump_to_operator_region(
+            JumpToOperatorRegionRequest(OperatorIdentity(state["LoopStartId"]))
+        )
+        uri = state["LoopStartStateURI"]
+        # Strip the per-iteration scratch (`table`, `output`) and the
+        # loop metadata (`LoopStartId`, `LoopStartStateURI`) so only the
+        # user-visible loop state is written back to LoopStart's input.
+        for key in ("table", "output", "LoopStartId", "LoopStartStateURI"):
+            state.pop(key, None)
+        writer = DocumentFactory.create_document(uri, State.SCHEMA).writer("0")
+        writer.put_one(State(state).to_tuple())
+        writer.close()
+
     def complete(self) -> None:
         """
         Complete the DataProcessor, marking state to COMPLETED, and notify the
         controller.
         """
         # flush the buffered console prints
         self._check_and_report_console_messages(force_flush=True)
-        self.context.executor_manager.executor.close()
+        controller_interface = self._async_rpc_client.controller_stub()
+        executor = self.context.executor_manager.executor
+        if isinstance(executor, LoopEndOperator) and executor.condition():
+            self._jump_to_loop_start(executor, controller_interface)
+        executor.close()
         # stop the data processing thread
         self.data_processor.stop()
         self.context.state_manager.transit_to(WorkerState.COMPLETED)
         self.context.statistics_manager.update_total_execution_time(time.time_ns())
-        controller_interface = self._async_rpc_client.controller_stub()
         controller_interface.worker_execution_completed(EmptyRequest())
         self.context.close()
 
@@ -193,6 +233,11 @@ def process_input_state(self) -> None:
         self._switch_context()
         output_state = self.context.state_processing_manager.get_output_state()
         if output_state is not None:
+            executor = self.context.executor_manager.executor
+            if isinstance(executor, LoopEndOperator):
+                self.context.output_manager.reset_storage()
+            elif isinstance(executor, LoopStartOperator):
+                self._attach_loop_start_id(output_state)
             for to, batch in self.context.output_manager.emit_state(output_state):
                 self._output_queue.put(
                     DataElement(

amber/src/main/python/pytexera/__init__.py

@@ -19,6 +19,7 @@
 from overrides import overrides
 from typing import Iterator, Optional, Union
 
+from core.models.operator import LoopStartOperator, LoopEndOperator
 from pyamber import *
 from .storage.dataset_file_document import DatasetFileDocument
 from .storage.large_binary_input_stream import LargeBinaryInputStream
@@ -43,6 +44,8 @@
     "UDFTableOperator",
     "UDFBatchOperator",
     "UDFSourceOperator",
+    "LoopStartOperator",
+    "LoopEndOperator",
     "DatasetFileDocument",
     "largebinary",
     "LargeBinaryInputStream",

amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionCoordinator.scala

@@ -576,8 +576,17 @@ class RegionExecutionCoordinator(
           region.getOperator(outputPortId.opId).outputPorts(outputPortId.portId)._3
         val schema =
           schemaOptional.getOrElse(throw new IllegalStateException("Schema is missing"))
-        DocumentFactory.createDocument(resultURI, schema)
-        DocumentFactory.createDocument(stateURI, State.schema)
+        // LoopEnd operators may re-execute the region multiple times; on
+        // subsequent iterations the result/state documents already exist,
+        // and `createDocument` (overrideIfExists=true) would clobber them.
+        // Skip the create call when the document is already there.
+        val isLoopEndRegion = region.getOperators.exists(_.isLoopEnd)
+        if (!isLoopEndRegion || !DocumentFactory.documentExists(resultURI)) {
+          DocumentFactory.createDocument(resultURI, schema)
+        }
+        if (!isLoopEndRegion || !DocumentFactory.documentExists(stateURI)) {
+          DocumentFactory.createDocument(stateURI, State.schema)
+        }
         if (!isRestart) {
           val (_, eid, _, _) = decodeURI(resultURI)
           WorkflowExecutionsResource.insertOperatorPortResultUri(

amber/src/main/scala/org/apache/texera/web/service/WorkflowExecutionService.scala

@@ -21,13 +21,14 @@ package org.apache.texera.web.service
 
 import com.typesafe.scalalogging.LazyLogging
 import org.apache.texera.amber.core.virtualidentity.{ExecutionIdentity, WorkflowIdentity}
-import org.apache.texera.amber.core.workflow.WorkflowContext
+import org.apache.texera.amber.core.workflow.{ExecutionMode, WorkflowContext}
 import org.apache.texera.amber.engine.architecture.controller.{ControllerConfig, Workflow}
 import org.apache.texera.amber.engine.architecture.rpc.controlcommands.EmptyRequest
 import org.apache.texera.amber.engine.architecture.rpc.controlreturns.WorkflowAggregatedState._
 import org.apache.texera.amber.engine.common.Utils
 import org.apache.texera.amber.engine.common.client.AmberClient
 import org.apache.texera.amber.engine.common.executionruntimestate.ExecutionMetadataStore
+import org.apache.texera.amber.operator.loop.LoopStartOpDesc
 import org.apache.texera.web.model.websocket.event.{
   TexeraWebSocketEvent,
   WorkflowErrorEvent,
@@ -66,7 +67,12 @@ class WorkflowExecutionService(
 ) extends SubscriptionManager
     with LazyLogging {
 
-  workflowContext.workflowSettings = request.workflowSettings
+  workflowContext.workflowSettings =
+    if (request.logicalPlan.operators.exists(_.isInstanceOf[LoopStartOpDesc])) {
+      request.workflowSettings.copy(executionMode = ExecutionMode.MATERIALIZED)
+    } else {
+      request.workflowSettings
+    }
   val wsInput = new WebsocketInput(errorHandler)
 
   addSubscription(