feat(python): give a port that carried no rows its declared columns

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

@@ -22,7 +22,7 @@
 from collections import defaultdict
 from typing import Iterator, List, Mapping, Optional, Union, MutableMapping, Protocol
 
-from . import Table, TableLike, Tuple, TupleLike, Batch, BatchLike
+from . import Table, TableLike, Tuple, TupleLike, Batch, BatchLike, Schema
 from .state import State
 from .table import all_output_to_tuple
 
@@ -92,6 +92,23 @@ def is_source(self) -> bool:
     def is_source(self, value: bool) -> None:
         self.__internal_is_source = value
 
+    __internal_input_schemas: Optional[MutableMapping[int, Schema]] = None
+
+    @property
+    @overrides.final
+    def input_schemas(self) -> MutableMapping[int, Schema]:
+        """
+        What each input port was declared to carry, keyed by port index and
+        written by the runtime before that port's data is handed over.
+
+        The tuples themselves say this too, so an operator only needs to ask
+        when there are none: a port can finish having carried no rows at all,
+        and its schema is then the only record of what its columns were.
+        """
+        if self.__internal_input_schemas is None:
+            self.__internal_input_schemas = {}
+        return self.__internal_input_schemas
+
     def open(self) -> None:
         """
         Open a context of the operator. Usually can be used for loading/initiating some
@@ -276,7 +293,11 @@ def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike
         yield
 
     def on_finish(self, port: int) -> Iterator[Optional[TableLike]]:
-        table = Table(self.__table_data[port])
+        rows = self.__table_data[port]
+        schema = self.input_schemas.get(port)
+        # A port that carried no rows has no tuples to read column names off,
+        # and a table of no columns fails every operator that names one.
+        table = Table(rows) if rows or schema is None else Table.empty_of(schema)
         yield from self.process_table(table, port)
 
     @abstractmethod

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

@@ -16,6 +16,7 @@
 # under the License.
 
 import pandas
+import pyarrow as pa
 from pampy import match
 from typing import Iterator, TypeVar, List
 
@@ -25,6 +26,22 @@
 
 
 class Table(pandas.DataFrame):
+    @staticmethod
+    def empty_of(schema) -> "Table":
+        """
+        The declared columns with no rows under them.
+
+        ``from_tuple_likes`` reads the column names off the tuples it is given,
+        so with none to read it produces a frame of no columns at all, and an
+        operator naming any of its own columns raises KeyError. A port that
+        carried no rows still has a schema, and this is what it looks like as a
+        table. Building it through Arrow gives each column the dtype it would
+        have had with rows in it.
+        """
+        return Table(
+            pa.Table.from_pylist([], schema=schema.as_arrow_schema()).to_pandas()
+        )
+
     @staticmethod
     def from_table(table):
         return table

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

@@ -74,8 +74,22 @@ def process_internal_marker(self, internal_marker: InternalMarker) -> None:
                 # Flush the state to MainLoop before producing tuples so the
                 # state and the tuple stream don't share a single switch.
                 self._switch_context()
+                self._declare_input_schema(executor, port_id)
                 self._set_output_tuple(executor.on_finish(port_id))
 
+    def _declare_input_schema(self, executor, port_id: int) -> None:
+        """
+        Tell the executor what the finishing port was declared to carry, so an
+        operator handed no rows can still say what its columns were. A source
+        has no input port to ask, so it is left alone.
+        """
+        port_identity = self._context.tuple_processing_manager.current_input_port_id
+        if port_identity is None:
+            return
+        executor.input_schemas[port_id] = self._context.input_manager.get_port(
+            port_identity
+        ).get_schema()
+
     def process_state(self, state: State) -> None:
         """
         Process an input marker by invoking appropriate state

amber/src/test/python/core/models/test_operator.py

@@ -22,6 +22,7 @@
 
 from core.models import (
     BatchOperator,
+    Schema,
     SourceOperator,
     State,
     Table,
@@ -386,6 +387,24 @@ def test_on_finish_with_no_buffered_tuples_yields_empty_table(self):
         assert len(op.received_tables) == 1
         assert list(op.received_tables[0].as_tuples()) == []
 
+    def test_on_finish_with_no_rows_keeps_the_declared_columns(self):
+        # The column names are read off the tuples, so a port that carried none
+        # left the operator a frame of no columns and every operator naming one
+        # of its own raised KeyError. The port's schema is the only record left.
+        op = _ConcreteTable()
+        op.input_schemas[0] = Schema(raw_schema={"x": "INTEGER", "y": "STRING"})
+
+        list(op.on_finish(port=0))
+
+        table = op.received_tables[0]
+        # Still a Table, so an operator reading it with as_tuples() keeps working.
+        assert isinstance(table, Table)
+        assert list(table.as_tuples()) == []
+        assert list(table.columns) == ["x", "y"]
+        assert table.empty
+        # The dtype each column would have had with rows under it.
+        assert table["x"].dtype == "int32"
+
     def test_buffers_are_keyed_by_port(self):
         # Each input port has its own tuple buffer; on_finish for one port
         # must not surface tuples written through a different port.

amber/src/test/python/core/runnables/test_main_loop.py

@@ -215,6 +215,10 @@ def state_processing_executor(self):
         # from `produce_state_on_finish` so EndChannel handling can be
         # observed.
         class StateProcessingExecutor:
+            # What Operator gives a real executor, for DataProcessor to record
+            # the finishing port's declared schema in.
+            input_schemas: dict = {}
+
             @staticmethod
             def process_tuple(tuple_, port):
                 yield tuple_