Use Output Ports of an Operator to Write Storage

core/amber/src/main/protobuf/edu/uci/ics/amber/engine/architecture/rpc/controlcommands.proto

@@ -227,6 +227,7 @@ message AssignPortRequest {
   core.PortIdentity portId = 1 [(scalapb.field).no_box = true];
   bool input = 2;
   map<string, string> schema = 3;
+  string storageUri = 4;
 }
 
 message FinalizeCheckpointRequest {

core/amber/src/main/python/core/architecture/handlers/control/assign_port_handler.py

@@ -14,7 +14,10 @@ async def assign_port(self, req: AssignPortRequest) -> EmptyReturn:
                 req.port_id, Schema(raw_schema=req.schema)
             )
         else:
+            storage_uri = None
+            if req.storage_uri != "":
+                storage_uri = req.storage_uri
             self.context.output_manager.add_output_port(
-                req.port_id, Schema(raw_schema=req.schema)
+                req.port_id, Schema(raw_schema=req.schema), storage_uri
             )
         return EmptyReturn()

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

@@ -1,9 +1,13 @@
+import threading
 import typing
 from collections import OrderedDict
 from itertools import chain
+from queue import Queue
+from threading import Thread
+from typing import Iterable, Iterator
+
 from loguru import logger
 from pyarrow import Table
-from typing import Iterable, Iterator
 
 from core.architecture.packaging.input_manager import WorkerPort, Channel
 from core.architecture.sendsemantics.broad_cast_partitioner import (
@@ -23,7 +27,19 @@
 from core.models import Tuple, Schema, MarkerFrame
 from core.models.marker import Marker
 from core.models.payload import DataPayload, DataFrame
+from core.storage.document_factory import DocumentFactory
+from core.storage.runnables.port_storage_writer import (
+    PortStorageWriter,
+    PortStorageWriterElement,
+)
 from core.util import get_one_of
+from core.util.virtual_identity import get_worker_index
+from proto.edu.uci.ics.amber.core import (
+    ActorVirtualIdentity,
+    PhysicalLink,
+    PortIdentity,
+    ChannelIdentity,
+)
 from proto.edu.uci.ics.amber.engine.architecture.rpc import ChannelMarkerPayload
 from proto.edu.uci.ics.amber.engine.architecture.sendsemantics import (
     HashBasedShufflePartitioning,
@@ -33,12 +49,6 @@
     RangeBasedShufflePartitioning,
     BroadcastPartitioning,
 )
-from proto.edu.uci.ics.amber.core import (
-    ActorVirtualIdentity,
-    PhysicalLink,
-    PortIdentity,
-    ChannelIdentity,
-)
 
 
 class OutputManager:
@@ -56,23 +66,89 @@ def __init__(self, worker_id: str):
         }
         self._ports: typing.Dict[PortIdentity, WorkerPort] = dict()
         self._channels: typing.Dict[ChannelIdentity, Channel] = dict()
+        self._port_storage_writers: typing.Dict[
+            PortIdentity, typing.Tuple[Queue, PortStorageWriter, Thread]
+        ] = dict()
 
-    def add_output_port(self, port_id: PortIdentity, schema: Schema) -> None:
+    def add_output_port(
+        self,
+        port_id: PortIdentity,
+        schema: Schema,
+        storage_uri: typing.Optional[str] = None,
+    ) -> None:
         if port_id.id is None:
             port_id.id = 0
         if port_id.internal is None:
             port_id.internal = False
 
+        if storage_uri is not None:
+            self.set_up_port_storage_writer(port_id, storage_uri)
+
         # each port can only be added and initialized once.
         if port_id not in self._ports:
             self._ports[port_id] = WorkerPort(schema)
 
+    def set_up_port_storage_writer(self, port_id: PortIdentity, storage_uri: str):
+        """
+        Create a separate thread for saving output tuples of a port
+        to storage in batch.
+        """
+        document, _ = DocumentFactory.open_document(storage_uri)
+        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,
+        )
+
     def get_port(self, port_id=None) -> WorkerPort:
         return list(self._ports.values())[0]
 
     def get_output_channel_ids(self):
         return self._channels.keys()
 
+    def save_tuple_to_storage_if_needed(self, tuple_: Tuple, port_id=None) -> None:
+        """
+        Optionally write the tuple to storage if the specified output port
+        is determined by the scheduler to need storage. This method is not blocking
+        because a separate thread is used to flush the tuple to storage in batch.
+        :param tuple_: A tuple produced by the data processor.
+        :param port_id: If not specified, the tuple will be written to all
+        output ports that need storage.
+        :return:
+        """
+        if port_id is None:
+            for writer_queue, _, _ in self._port_storage_writers.values():
+                writer_queue.put(PortStorageWriterElement(data_tuple=tuple_))
+        elif port_id in self._port_storage_writers.keys():
+            self._port_storage_writers[port_id][0].put(
+                PortStorageWriterElement(data_tuple=tuple_)
+            )
+
+    def close_port_storage_writers(self) -> None:
+        """
+        Flush the buffers of port storage writers and wait for all the
+        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()
+
     def add_partitioning(self, tag: PhysicalLink, partitioning: Partitioning) -> None:
         """
         Add down stream operator and its transfer policy

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

@@ -189,6 +189,9 @@ def process_input_tuple(self) -> None:
                             payload=batch,
                         )
                     )
+                self.context.output_manager.save_tuple_to_storage_if_needed(
+                    output_tuple
+                )
 
     def process_input_state(self) -> None:
         self._switch_context()
@@ -288,6 +291,8 @@ def _process_end_of_output_ports(self, _: EndOfOutputPorts) -> None:
 
         :param _: EndOfOutputPorts
         """
+        self.context.output_manager.close_port_storage_writers()
+
         for to, batch in self.context.output_manager.emit_marker(EndOfInputChannel()):
             self._output_queue.put(
                 DataElement(

core/amber/src/main/python/core/storage/document_factory.py

@@ -1,3 +1,4 @@
+import typing
 from urllib.parse import urlparse
 
 from typing import Optional
@@ -65,7 +66,7 @@ def create_document(uri: str, schema: Schema) -> VirtualDocument:
             )
 
     @staticmethod
-    def open_document(uri: str) -> (VirtualDocument, Optional[Schema]):
+    def open_document(uri: str) -> typing.Tuple[VirtualDocument, Optional[Schema]]:
         parsed_uri = urlparse(uri)
         if parsed_uri.scheme == "vfs":
             _, _, _, _, _, resource_type = VFSURIFactory.decode_uri(uri)

core/amber/src/main/python/core/storage/runnables/port_storage_writer.py

@@ -0,0 +1,34 @@
+from dataclasses import dataclass
+
+from overrides import overrides
+
+from core.models import Tuple
+from core.storage.model.buffered_item_writer import BufferedItemWriter
+from core.util import StoppableQueueBlockingRunnable, IQueue
+from core.util.customized_queue.queue_base import QueueElement
+
+
+@dataclass
+class PortStorageWriterElement(QueueElement):
+    data_tuple: Tuple
+
+
+class PortStorageWriter(StoppableQueueBlockingRunnable):
+    def __init__(self, buffered_item_writer: BufferedItemWriter, queue: IQueue):
+        super().__init__(name=self.__class__.__name__, queue=queue)
+        self.buffered_item_writer: BufferedItemWriter = buffered_item_writer
+
+    @overrides
+    def receive(self, next_entry: QueueElement) -> None:
+        if isinstance(next_entry, PortStorageWriterElement):
+            self.buffered_item_writer.put_one(next_entry.data_tuple)
+        else:
+            raise TypeError(f"Unexpected entry {next_entry}")
+
+    @overrides
+    def pre_start(self) -> None:
+        self.buffered_item_writer.open()
+
+    @overrides
+    def post_stop(self) -> None:
+        self.buffered_item_writer.close()

core/amber/src/main/python/core/storage/storage_config.py

@@ -34,7 +34,7 @@ def initialize(
         cls.ICEBERG_POSTGRES_CATALOG_PASSWORD = postgres_password
         cls.ICEBERG_TABLE_RESULT_NAMESPACE = table_result_namespace
         cls.ICEBERG_FILE_STORAGE_DIRECTORY_PATH = directory_path
-        cls.ICEBERG_TABLE_COMMIT_BATCH_SIZE = commit_batch_size
+        cls.ICEBERG_TABLE_COMMIT_BATCH_SIZE = int(commit_batch_size)
         cls._initialized = True
 
     def __new__(cls, *args, **kwargs):

core/amber/src/main/python/core/util/virtual_identity/__init__.py

@@ -0,0 +1,10 @@
+import re
+
+worker_name_pattern = re.compile(r"Worker:WF\d+-.+-(\w+)-(\d+)")
+
+
+def get_worker_index(worker_id: str) -> int:
+    match = worker_name_pattern.match(worker_id)
+    if match:
+        return int(match.group(2))
+    raise ValueError("Invalid worker ID format")

core/amber/src/main/python/texera_run_python_worker.py

@@ -3,6 +3,7 @@
 from loguru import logger
 
 from core.python_worker import PythonWorker
+from core.storage.storage_config import StorageConfig
 
 
 def init_loguru_logger(stream_log_level) -> None:
@@ -22,8 +23,28 @@ def init_loguru_logger(stream_log_level) -> None:
 
 
 if __name__ == "__main__":
-    _, worker_id, output_port, logger_level, r_path = sys.argv
+    (
+        _,
+        worker_id,
+        output_port,
+        logger_level,
+        r_path,
+        iceberg_postgres_catalog_uri_without_scheme,
+        iceberg_postgres_catalog_username,
+        iceberg_postgres_catalog_password,
+        iceberg_table_namespace,
+        iceberg_file_storage_directory_path,
+        iceberg_table_commit_batch_size,
+    ) = sys.argv
     init_loguru_logger(logger_level)
+    StorageConfig.initialize(
+        iceberg_postgres_catalog_uri_without_scheme,
+        iceberg_postgres_catalog_username,
+        iceberg_postgres_catalog_password,
+        iceberg_table_namespace,
+        iceberg_file_storage_directory_path,
+        iceberg_table_commit_batch_size,
+    )
 
     # Setting R_HOME environment variable for R-UDF usage
     if r_path:

core/amber/src/main/scala/edu/uci/ics/amber/engine/architecture/controller/promisehandlers/PortCompletedHandler.scala

@@ -2,6 +2,7 @@ package edu.uci.ics.amber.engine.architecture.controller.promisehandlers
 
 import com.twitter.util.Future
 import edu.uci.ics.amber.core.WorkflowRuntimeException
+import edu.uci.ics.amber.core.workflow.GlobalPortIdentity
 import edu.uci.ics.amber.engine.architecture.controller.{
   ControllerAsyncRPCHandlerInitializer,
   FatalError
@@ -12,7 +13,6 @@ import edu.uci.ics.amber.engine.architecture.rpc.controlcommands.{
   QueryStatisticsRequest
 }
 import edu.uci.ics.amber.engine.architecture.rpc.controlreturns.EmptyReturn
-import edu.uci.ics.amber.engine.architecture.scheduling.GlobalPortIdentity
 import edu.uci.ics.amber.engine.common.virtualidentity.util.CONTROLLER
 import edu.uci.ics.amber.util.VirtualIdentityUtils